Sharing memory between java processes with java.lang.foreign (and jextract)
The java.lang.foreign contains powerful APIs that you can use to take advantage
of various OS-specific optimizations. Often though, due to high usage of the
C/C++ preprocessor and its powerful macros it can be cumbersome to write and
maintain the equivalent java.lang.foreign code.
Help is available in the form of jextract - a JDK tool that reads C/C++ header files (*.h)
and generates Java code that can be used exactly like the original C/C++ macros, structs,
typedefs, etc.
In this post we’ll use jextract to implement one such OS-specific optimization:
short-circuit reads and writes, a zero-copy IPC technique based on shared
memory. We’ll also show how jextract-generated code is much easier to read
and maintain.
Sharing memory between processes is perhaps the fastest inter-process communication (IPC)
mechanism available on modern operating systems. As the name suggests, it literally
allows one process to modify a segment of memory and the changes will be immediately
visible within another process. No networking is involved, no TCP/IP stack, it’s just
two (or more) processes reading from and writing to exactly the same memory.
Shared memory segments are often used by high-performance systems to achieve data locality -
to run algorithms on data that is practically in local memory, without the overhead of
transferring it. Data may still need to be loaded from disk (if not cached already),
but the extra syscalls and context switches that might be incurred when using sockets
are completely avoided. There’s no copying of data either, which saves a huge number
of CPU cycles.
Additional information:“Short-circuit” reads and writes
Within large-scale distributed systems the term “short-circuiting”
(like short-circuit read or short-circuit write) is often used to
describe the technique of using shared memory to achieve zero-copy
data locality.
The DataNode of Apache HDFS is a good example of a system that uses short-circuiting
to get zero-copy data locality for Apache Spark or MapReduce jobs. When a job is submitted,
the driver (application master, etc) program asks HDFS which DataNodes store
copies of the desired input. Based on that information it then tries to launch
tasks (workers, etc) on the same DataNodes. Once that happens - the worker
requests short-circuiting from the DataNode and starts working directly on the
memory-mapped blocks that the DataNode shares with it.
Another typical scenario is when multiple applications, written in different programming
languages, need to cooperate. For exmpale, a database engine implemented in some
high-performance language that’s being used interactively by researchers with Python
notebooks. When the researcher needs to analyze a large dataset, the notebook’s
DataFrame can ask the database engine to do most of the heavy lifting (like in
a predicte pushdown manner), the database prepares data in a shared memory segment,
the notebook can then complete the disired computation entirely in memory.
There are wonderful articles on the web that explain how memory sharing is implemented in
an operating system. In here, we’re not going to cover that, we’re just going
to include a little reminder of the steps a C programmer would follow in order
to implement short-circuiting on the Linux operating system. After all, we would like to
do the same in Java:
First, create a shared memory object and get its file descriptor:
int shared_memory_fd = shm_open("/my_shared_memory_object",
O_CREAT | O_RDWR,
S_IRUSR | S_IWUSR);
if (shared_memory_fd == -1) {
perror("failed to create a shared memory object");
exit(EXIT_FAILURE);
}
Then, allocate some memory (size bytes) in the shared object:
if (ftruncate(shared_memory_fd, size) == -1) {
perror("failed to allocate shared memory");
exit(EXIT_FAILURE);
}
Next, “map” the newly allocated shared memory within our own address space:
void *shared_memory_ptr = mmap(NULL, size,
PROT_READ | PROT_WRITE,
MAP_SHARED,
shared_memory_fd, 0);
if (shared_memory_ptr == MAP_FAILED) {
perror("failed to map shared memory");
exit(EXIT_FAILURE);
}
Now we can modify the memory pointed to by shared_memory_ptr to whatever we want to
share with other processes. Changes these processes make to the same memory will be
immediately visible to us too.
So far so good - we’ve allocated a shared memory segment and we can read and write to it. We
haven’t actually shared it though. To do that we simply pass the shared_memory_fd (the file
descriptor behind the shared memory object) to whichever process we want to.
On the Linux operating system this is done with the sendmsg syscall on a Unix domain socket:
// server side:
int socket_fd = socket(AF_UNIX, SOCK_STREAM, 0);
// ... bind, listen, accept, etc ...
// when a client connects, we send them the shared_memory_fd:
struct msghdr msg = {0};
struct cmsghdr *cmsg;
// reserve some space for the file descriptor:
char buf[CMSG_SPACE(sizeof(shared_memory_fd))];
memset(buf, 0, sizeof(buf));
// set some dummy data to send along with the control message:
struct iovec io = { .iov_base = "ABC", .iov_len = 3 };
msg.msg_iov = &io;
msg.msg_iovlen = 1;
// now, here goes the control message holding the file descriptor:
msg.msg_control = buf;
msg.msg_controllen = sizeof(buf);
cmsg = CMSG_FIRSTHDR(&msg);
// type and level for the control message:
cmsg->cmsg_level = SOL_SOCKET;
cmsg->cmsg_type = SCM_RIGHTS;
cmsg->cmsg_len = CMSG_LEN(sizeof(shared_memory_fd));
// set the 'body' of the control message to be the file descriptor:
memcpy(CMSG_DATA(cmsg), &shared_memory_fd, sizeof(shared_memory_fd));
// send!
if (sendmsg(socket_fd, &msg, 0) < 0) {
perror("failed to send shared memory fd");
exit(EXIT_FAILURE);
}
The client, on its end receives the control message and gets the file descriptor of the shared memory object:
// client side:
int socket_fd = socket(AF_UNIX, SOCK_STREAM, 0);
// ... connect to server, etc ...
// allocate some space where a message can be received:
struct msghdr msg = {0};
char m_buffer[256];
struct iovec io = { .iov_base = m_buffer, .iov_len = sizeof(m_buffer) };
msg.msg_iov = &io;
msg.msg_iovlen = 1;
// allocate some memory for the control message:
char c_buffer[256];
msg.msg_control = c_buffer;
msg.msg_controllen = sizeof(c_buffer);
// receive!
if (recvmsg(socket_fd, &msg, 0) < 0) {
perror("failed to receive shared memory fd");
exit(EXIT_FAILURE);
}
// get the file descriptor from the now populated buffers:
struct cmsghdr *cmsg = CMSG_FIRSTHDR(&msg);
int shared_memory_fd;
memcpy(&shared_memory_fd, CMSG_DATA(cmsg), sizeof(shared_memory_fd));
shared_memory_fd is now available on the “client-side” and can be mapped locally:
// "map" it the same was as in the server above:
void *shared_memory_ptr = mmap(NULL, size,
PROT_READ | PROT_WRITE,
MAP_SHARED,
shared_memory_fd, 0);
if (shared_memory_ptr == MAP_FAILED) {
perror("failed to map shared memory");
exit(EXIT_FAILURE);
}
Now both ends of the Unix domain socket have a pointer, shared_memory_ptr that
points to the same physical memory. Changes done on one side will be immediately
visible on the other.
Alright, so now lets move to implementing the same in Java, as it is the main
topic of this blog post.
Java traditionally focuses on platform independence and builds abstractions and APIs
that allow developers to leverage OS-dependent optimizations in an OS-agnostic way.
These usually take time to get standardized and become mainstream JDK features, hence the
need to also have a way of calling underlying OS APIs directly.
A good example of the above is the evolution of the Java-to-native interfaces. It began as
Java Native Interface (JNI) - a bunch of (mostly) C++ APIs that a developer could use to interact
with the JVM and a native keyword to mark a method’s implementation as being in C++.
The HDFS DataNode that we mentioned earlier actually uses a JNI library, called ‘libhadoop’ and
written in C++, to do its short-circuiting.
This JNI-interface later evolved into the java.lang.foreign package, which offers an OS-agnostic
way of interacting with the OS dynamic linker - loading libraries, looking up symbols, calling
functions, etc.
For the purpose of this post - java.lang.foreign is exactly the package we want to be using in
order to implement our memory sharing and short-circuit I/O. With it we can lookup all functions
we need to call - shm_open, ftruncate, mmap, etc and get invokeable method handles for each
of them. We can then allocate C-style strings and other memory pointers, pass them as arguments
to the syscalls, and get back the results.
So, for example, we can create a shared memory object like this:
import java.lang.foreign.*;
import java.lang.invoke.MethodHandle;
// create a C-style string
MemorySegment sharedObjectName = arena.allocateFrom("/my_shared_memory_object");
SymbolLookup symbols = SymbolLookup.loaderLookup();
// ask the OS dynamic linker for the address of `shm_open` symbol:
MemorySegment shmOpenAddr = symbols.lookup("shm_open").orElseThrow();
// obtain a handle to the `shm_open` function:
MethodHandle shm_open = Linker.nativeLinker().downcallHandle(
shmOpenAddr,
FunctionDescriptor.of(
(ValueLayout.OfInt) Linker.nativeLinker().canonicalLayouts().get("int"), // retrun type: int
((AddressLayout) Linker.nativeLinker().canonicalLayouts().get("void*")), // first argument: a pointer
(ValueLayout.OfInt) Linker.nativeLinker().canonicalLayouts().get("int"), // second argument: an int
(ValueLayout.OfInt) Linker.nativeLinker().canonicalLayouts().get("int") // third argument: an int
));
// now we can call `shm_open` just like in the C example above:
int sharedMemoryFd = (int) shm_open.invokeExact(
sharedObjectName, // the name of the shared memory object
0100 | 02, // O_CREAT | O_RDWR
0400 | 0200 // S_IRUSR | S_IWUSR
);
if (sharedMemoryFd == -1) {
// must check C's `errno` to know what went wrong
throw new RuntimeException("failed to create a shared memory object");
}
The code above initially looks simple, but you’re probably noticing that quite quickly
it is getting cumbersome - to both implement and maintain. First, we need to know the
exact signature of each syscall: correct return types, correct argument types.
Then, there’s all those C-macros that go as arguments - O_CREAT, O_RDWR, permissions
like S_IRUSR, S_IWUSR, memory protection modes like PROT_READ, PROT_WRITE, and so on.
All of these come from the C header files, and we have to go through quite a lot of them to
find out what all of them evaluate to, so we can write the same final value in Java.
For example, on Linux, /usr/include/sys/mman.h defines shm_open:
extern int shm_open (const char *__name, int __oflag, mode_t __mode);
So from this line we know in Java we should be looking for a function called shm_open that
returns an int and takes a pointer, an int, and a mode_t as arguments. What’s a mode_t
and how do we map it in Java? Well, we have to look at /usr/include/bits/types.h to find out:
// mode_t is actually a typedef to __mode_t...
typedef __mode_t mode_t;
/usr/include/bits/types.h then defines __mode_t:
__STD_TYPE __MODE_T_TYPE __mode_t; /* Type of file attribute bitmasks. */
In turn, /usr/include/bits/typesizes.h defines __MODE_T_TYPE:
#define __MODE_T_TYPE __U32_TYPE
… and so on.
Same with constants like O_CREAT, O_RDWR, S_IRUSR, S_IWUSR - we
have to look them up in /usr/include/bits/fcntl.h and /usr/include/bits/stat.h
to find out their actual values. Once we put the values (0100 for O_CREAT, 02 for O_RDWR,
0400 for S_IRUSR, and 0200 for S_IWUSR) our Java code quickly becomes unreadable and
unmaintainable.
That’s not an issue in C/C++ because the compiler will resolve all those macros for us, but
in Java we have to do it manually. Unless, of course, we have a tool that, just like the C compiler,
will read and process the same header files, find the definitions that we need, and generate
java code for us.
Enter jextract - a JDK tool that does
exactly that. Part of OpenJDK “Code Tools” project
and described as “Native library binding extraction tool”:
jextract is a tool which mechanically generates Java bindings from native
library headers. This tools leverages the clang C API in order to parse the
headers associated with a given native library, and the generated Java
bindings build upon the Foreign Function & Memory API.
it comes really handy when we want to have maintainable java.lang.foreign code.
Sticking to the example of shm_open above, we can run jextract on the Linux
headers to generate a Java class that captures the necessary details for us:
jextract \
# look for additional headers under /usr/include
-I /usr/include \
# generate a Java class named `ForeignSharedMemory`
--header-class-name ForeignSharedMemory \
# place the generated code in the `foreign.shm.package` package:
-t foreign.shm.package \
# extract some O_* mode constants:
--include-constant O_CREAT \
... \
# extract some S_* file permission constants:
--include-constant S_IRUSR \
... \
# extract some PROT_* memory protection constants:
--include-constant PROT_READ \
...
# extract some functions too:
--include-function shm_open \
--include-function shm_unlink \
...
# finally, start with these headers:
/usr/include/sys/stat.h \
/usr/include/bits/fcntl.h \
/usr/include/sys/mman.h
Running the above will create the java sources for a ForeignSharedMemory
class (and perhaps a few others) in the foreign.shm.package package. Lets
take a look at what’s inside:
private static final int O_RDWR = (int)2L;
/**
* {@snippet lang=c :
* #define O_RDWR 2
* }
*/
public static int O_RDWR() {
return O_RDWR;
}
private static final int O_CREAT = (int)64L;
/**
* {@snippet lang=c :
* #define O_CREAT 64
* }
*/
public static int O_CREAT() {
return O_CREAT;
}
Great! A bunch of static methods for our constants. And then a bit further:
private static class shm_open {
public static final FunctionDescriptor DESC = FunctionDescriptor.of(
ForeignSharedMemory.C_INT,
ForeignSharedMemory.C_POINTER,
ForeignSharedMemory.C_INT,
ForeignSharedMemory.C_INT
);
public static final MemorySegment ADDR = SYMBOL_LOOKUP.findOrThrow("shm_open");
public static final MethodHandle HANDLE = Linker.nativeLinker().downcallHandle(ADDR, DESC);
}
/**
* {@snippet lang=c :
* extern int shm_open(const char *__name, int __oflag, mode_t __mode)
* }
*/
public static int shm_open(MemorySegment __name, int __oflag, int __mode) {
var mh$ = shm_open.HANDLE;
try {
if (TRACE_DOWNCALLS) {
traceDowncall("shm_open", __name, __oflag, __mode);
}
return (int)mh$.invokeExact(__name, __oflag, __mode);
} catch (Error | RuntimeException ex) {
throw ex;
} catch (Throwable ex$) {
throw new AssertionError("should not reach here", ex$);
}
}
Using the generated static methods above makes our Java code much simpler and understandable:
import java.lang.foreign.*;
import static foreign.shm.package.ForeignSharedMemory.O_CREAT;
import static foreign.shm.package.ForeignSharedMemory.O_RDWR;
import static foreign.shm.package.ForeignSharedMemory.S_IRUSR;
import static foreign.shm.package.ForeignSharedMemory.S_IWUSR;
import static foreign.shm.package.ForeignSharedMemory.shm_open;
// create the same C-style string as before:
MemorySegment sharedObjectName = arena.allocateFrom("/my_shared_memory_object");
// Just call shm_open() using the much more meaningful literals:
int sharedMemoryFd = shm_open(
sharedObjectName,
O_CREAT() | O_RDWR(),
S_IRUSR() | S_IWUSR());
if (sharedMemoryFd == -1) {
throw new RuntimeException("failed to create a shared memory object");
}
We can, in fact, use the original Linux documentation for C developers and
follow it almost verbatim in our Java code. For example, the shm_open
documentation describes the function’s signature, how the system behaves
depending on the flags passed:
$ man shm_open
SYNOPSIS
#include <sys/mman.h>
#include <sys/stat.h> /* For mode constants */
#include <fcntl.h> /* For O_* constants */
int shm_open(const char *name, int oflag, mode_t mode);
int shm_unlink(const char *name);
DESCRIPTION
shm_open() creates and opens a new, or opens an existing, POSIX shared memory object. A POSIX shared memory object is in effect a handle which can be used by unrelated processes to mmap(2) the same region of
shared memory. The shm_unlink() function performs the converse operation, removing an object previously created by shm_open().
The operation of shm_open() is analogous to that of open(2). name specifies the shared memory object to be created or opened. For portable use, a shared memory object should be identified by a name of the
form /somename; that is, a null-terminated string of up to NAME_MAX (i.e., 255) characters consisting of an initial slash, followed by one or more characters, none of which are slashes.
oflag is a bit mask created by ORing together exactly one of O_RDONLY or O_RDWR and any of the other flags listed here:
O_RDONLY
Open the object for read access. A shared memory object opened in this way can be mmap(2)ed only for read (PROT_READ) access.
O_RDWR Open the object for read-write access.
O_CREAT
...
… and of course what errors might be encountered:
...
RETURN VALUE
On success, shm_open() returns a file descriptor (a nonnegative integer). On success, shm_unlink() returns 0. On failure, both functions return -1 and set errno to indicate the error.
ERRORS
EACCES Permission to shm_unlink() the shared memory object was denied.
EACCES Permission was denied to shm_open() name in the specified mode, or O_TRUNC was specified and the caller does not have write permission on the object.
EEXIST Both O_CREAT and O_EXCL were specified to shm_open() and the shared memory object specified by name already exists.
EINVAL The name argument to shm_open() was invalid.
EMFILE The per-process limit on the number of open file descriptors has been reached.
ENAMETOOLONG
The length of name exceeds PATH_MAX.
ENFILE The system-wide limit on the total number of open files has been reached.
ENOENT An attempt was made to shm_open() a name that did not exist, and O_CREAT was not specified.
ENOENT An attempt was to made to shm_unlink() a name that does not exist.
Another tricky bit with native C/C++ code is that members of a struct are
‘placed’ on memory addresses that align with the CPU architecture’s word size.
For example - int members will be aligned to 4-byte boundaries, long
members on 8-byte boundaries, and so on. In Java, when we need to pass a pointer
to a struct, or need to access members of a struct returned by a native
function - we need to follow the same alignment rules. Otherwise, we
might end up reading or writing the wrong memory locations within that struct.
The java.lang.foreign package takes this into account and provides abstractions
that allow developers to set the correct alignments. But again, just like with the
macros and constants above, the developer has to follow the exact struct definitions
in the C header files and workout the correct alignments from them. And this is
another scenario where jextract comes to the rescue:
$ jextract \
-I /usr/include \
--header-class-name SysSocketH \
-t foreign.socket.package \
...
# now, generate code for the `msghdr` and `cmsghdr` structs
# that we used in the C example above:
--include-struct msghdr \
--include-struct cmsghdr \
...
/usr/include/sys/socket.h
With the above command, two extra classes will be generated, one for each of the
msghdr and cmsghdr structs. In them, there will be the layout:
private static final GroupLayout $LAYOUT = MemoryLayout.structLayout(
SysSocketH.C_POINTER.withName("msg_name"),
SysSocketH.C_INT.withName("msg_namelen"),
MemoryLayout.paddingLayout(4),
SysSocketH.C_POINTER.withName("msg_iov"),
SysSocketH.C_LONG.withName("msg_iovlen"),
SysSocketH.C_POINTER.withName("msg_control"),
SysSocketH.C_LONG.withName("msg_controllen"),
SysSocketH.C_INT.withName("msg_flags"),
MemoryLayout.paddingLayout(4)
).withName("msghdr");
Additional information:Memory layout and alignment
Note the MemoryLayout.paddingLayout(4) calls - these are the alignment rules that
jextract has inferred from the C header files and added to the generated code.
Helper methods, getters and setters for the members of the struct are also generated.
These methods take into account the C-type and the correct offset of each member so
we don’t have to:
private static final OfInt msg_namelen$LAYOUT = (OfInt)$LAYOUT.select(groupElement("msg_namelen"));
private static final long msg_namelen$OFFSET = $LAYOUT.byteOffset(groupElement("msg_namelen"));
/**
* Getter for field:
* {@snippet lang=c :
* socklen_t msg_namelen
* }
*/
public static int msg_namelen(MemorySegment struct) {
return struct.get(msg_namelen$LAYOUT, msg_namelen$OFFSET);
}
/**
* Setter for field:
* {@snippet lang=c :
* socklen_t msg_namelen
* }
*/
public static void msg_namelen(MemorySegment struct, int fieldValue) {
struct.set(msg_namelen$LAYOUT, msg_namelen$OFFSET, fieldValue);
}
Now, server-side Java code will look something like this:
// allocate some dummy data to go with our control message:
MemorySegment data = arena.allocateFrom("ASDF");
MemorySegment iov = iovec.allocate(arena);
iovec.iov_base(iov, data);
iovec.iov_len(iov, data.byteSize());
// allocate one control message struct plus enough space for the file descriptor:
MemorySegment control = arena.allocate(
cmsghdr.sizeof() + ValueLayout.JAVA_INT.byteSize());
// configure control message level and type:
cmsghdr.cmsg_level(control, SysSocketH.SOL_SOCKET());
cmsghdr.cmsg_type(control, SysSocketH.SCM_RIGHTS());
// set the correct length of the control message:
cmsghdr.cmsg_len(control,
cmsghdr.sizeof() + ValueLayout.JAVA_INT.byteSize());
// set the file descriptor just after the control message struct:
control.set(ValueLayout.JAVA_INT, cmsghdr.sizeof(), sharedMemoryFd);
// allocate a msghdr struct and fill it in:
MemorySegment msgHdr = msghdr.allocate(arena);
msghdr.msg_name(msgHdr, MemorySegment.ofAddress(0));
msghdr.msg_namelen(msgHdr, 0);
msghdr.msg_iov(msgHdr, iov);
msghdr.msg_iovlen(msgHdr, 1);
msghdr.msg_control(msgHdr, control);
msghdr.msg_controllen(msgHdr, cmsghdr.sizeof() + ValueLayout.JAVA_INT.byteSize());
msghdr.msg_flags(msgHdr, 0);
// send!
long sentBytes = SysSocketH.sendmsg(peerSocket, msgHdr, 0);
if (sentBytes < 0) {
throw new RuntimeException("failed to send shared memory fd");
}
And a Java client looks something like this:
// a buffer that's large enough to hold the received dummy data:
MemorySegment clientData = arena.allocate(256);
MemorySegment clientIov = iovec.allocate(arena);
iovec.iov_base(clientIov, clientData);
iovec.iov_len(clientIov, 256);
// should be enough for a cmsghdr struct plus a bunch of file descriptors:
MemorySegment clientControl = arena.allocate(16 * cmsghdr.sizeof());
cmsghdr.cmsg_level(clientControl, SysSocketH.SOL_SOCKET());
cmsghdr.cmsg_type(clientControl, SysSocketH.SCM_RIGHTS());
cmsghdr.cmsg_len(clientControl, 16 * cmsghdr.sizeof());
// configure the msghdr struct with the buffers we want to be filled:
MemorySegment clientMsgHdr = msghdr.allocate(arena);
msghdr.msg_name(clientMsgHdr, MemorySegment.ofAddress(0));
msghdr.msg_namelen(clientMsgHdr, 0);
msghdr.msg_iov(clientMsgHdr, clientIov);
msghdr.msg_iovlen(clientMsgHdr, 1);
msghdr.msg_control(clientMsgHdr, clientControl);
msghdr.msg_controllen(clientMsgHdr, 16 * cmsghdr.sizeof());
msghdr.msg_flags(clientMsgHdr, 0);
long receivedBytes = SysSocketH.recvmsg(clientFD, clientMsgHdr, 0);
if (receivedBytes < 0) {
throw new RuntimeException("failed to receive shared memory fd");
}
// get the bytes of the received control message:
MemorySegment receivedCmsghdr = msghdr.msg_control(clientMsgHdr);
long recvCmsghdrSize = msghdr.msg_controllen(clientMsgHdr);
// get the file descriptor, it's located at an offset just after the cmsghdr struct:
int receivedFD = receivedCmsghdr.get(ValueLayout.JAVA_INT, cmsghdr.sizeof());
// now mmap() receivedFD...
MemorySegment sharedMemorySegment = mmap(...);
At this point - both Java processes have the same physical memory mapped into their
address space, under the MemorySegment Object returned by a successful call to mmap().
Both processes can then turn that memory in to a ByteBuffer (by calling MemorySegment.asByteBuffer()),
or into an array (with MemorySegment.toArray()), or use the MemorySegment getters and setters,
perhaps alongside a predefined GroupLayout to access the memory as a struct of some sort, or deserialize
Java POJOs, etc. MemorySegment offers a lot of flexibility and it’s really up to you to decide how best to use
that memory.
Additional information:Complete code examples available
We’ve implemented a complete Java lib that implements short-circuiting. It’s code is available on
our GitHub.
In another article of this series we’ll plug these shared memory optimizations into
Apache Arrow and share its buffers and vectors between apps (Java and/or Python).
Then, with the help of another native library, we’ll also add some GPU-processing power to the same
Apache Arrow vectors.
So, stay put! :)
