session_arenas

Dissecting Memory Sessions

November 2022

For new users of the Foreign Function and Memory API (FFM API in short), the concept of memory session can seem jarring at first. Memory sessions force users to think explicitly about the lifetime of the memory segments they wish to create. Because of that, almost invariably, the first knee-jerk reaction is along the lines “but are memory sessions really needed?”

There are several forces pushing the FFM API towards a notion of a lifetime that can be shared across multiple segments:

• Slices of a segment share the same lifetime;
• Groups of segments are often allocated in the same lifetime, and freed all at once (e.g. inside a try-with-resources block in a method). We call this principle co-allocation - that is, if two or more segments are co-allocated, they can safely point to each other. Co-allocation shows up frequently when interacting with native code, and even in the FFM API itself (e.g. VaList::copy);
• Our secret sauce for preventing use-after-free in racy, shared cases is expensive, so guiding users towards amortizing these costs is good;
• Without a concept of lifetime, it would be hard (if not impossible) to reason about attaching a raw native pointers to an existing lifetime (e.g. MemorySegment::ofAddress).

The last point is especially important: in a way, the ability of creating an unsafe native memory segment, with given (raw) base address, size, session and cleanup action is the true primitive of the FFM API. Once we have that capability, we can easily express any kind of native segment: mapped segments, upcall segments, lookup symbols and variable-argument lists are nothing but simple wrappers around the unsafe native segment factory. For instance, a safe factory for native segments can be expressed as follows:

MemorySegment allocateNative(long size, MemorySession session) {

    long addr = Unsafe.allocateMemory(size);

    MemorySegment segment = MemorySegment.ofAddress(addr, size, session);

    session.addCleanupAction(() -> Unsafe.freeMemory(addr));

}

All this leads to an API where (perhaps surprisingly) segments are not individually allocated and/or released. Instead, the lifetime of multiple segments is managed in an atomic and centralized fashion, as follows:

try (MemorySession session = MemorySession.openConfined()) {

    MemorySegment segment1 = session.allocate(100);

    MemorySegment segment2 = session.allocate(100);

    MemorySegment segment3 = session.allocate(100);

    ...

} // segment1, segment2, segment3 freed together

Memory sessions especially shine in the context of native interop. When interacting with native code, memory sessions help developers managing lifetime of logically related allocations, with relative ease, as demonstrated in [1-8]. However, memory sessions are not simple to use, nor easy to understand, especially for users that are new to the FFM API.

This document has three goals: first, to highlight the issues that contribute to the perceived complexity of memory sessions; secondly, to discuss how to improve the status quo, primarily by teasing apart the various roles of memory sessions into separate abstractions; finally, to show how memory sessions could be integrated, with minor tweaks, with future enhancements, such as support for structured memory sessions.

Memory sessions: one lump too many?

There's no question that memory sessions are useful tools to manage lifetime of foreign memory. However, as defined in Java 19, the memory session API can look quite daunting at first:

sealed interface MemorySession extends SegmentAllocator, AutoCloseable {

    Thread ownerThread();

    boolean isAlive();

    void close();

    void addCloseAction(Runnable);

    boolean isCloseable();

    MemorySession asNonCloseable();

    void whileAlive(Runnable);

    boolean equals(Object);

    int hashCode();

    MemorySegment allocate(byte size, byte align) { ... }

    public static MemorySession openConfined() { ... }

    public static MemorySession openConfined(Cleaner) { ... }

    public static MemorySession openShared() { ... }

    public static MemorySession openShared(Cleaner) { ... }

    public static MemorySession openImplicit() { ... }

    public static MemorySession global() { ... }

}

There are many factors contributing to the complexity of memory sessions, some of which are necessary, but some which are more accidental. In the following sections we will discuss some problematic aspects of memory sessions.

Sessions and allocators

Memory sessions are not pure lifetime abstractions, as they also implement the SegmentAllocator interface. This was a result of a compromise the FFM API had to make, given that it is sometimes handy to be able to pass a MemorySession instance where a SegmentAllocator is expected. While this move is a pragmatic one, it creates some duplication in the API, as there's no difference between MemorySession::allocate and MemorySegment::allocateNative. And, the API does not provide solutions for clients who wish to associate a custom allocation strategy with a memory session. In such cases, clients would have to manage session and allocation abstractions independently, as shown below:

try (MemorySession session = MemorySession.openConfined()) {

    SegmentAllocator allocator = allocatorFor(session);

    ...

}

This is clearly suboptimal: now the client has two allocators: session (since MemorySession implements SegmentAllocator) and allocator , and it is up to developers to use the correct one. Ideally, what the code would like to do is to create a custom allocator, which can be wrapped in a try-with-resources block, and use that. But doing that is not trivial, primarily because an allocator cannot be used where a session is expected, and exposing a session accessor from the custom allocator can be problematic (as we shall see in a later section).

Closing sessions

Some sessions can be closed explicitly, such as those created via MemorySession::openConfined or MemorySession::openShared. But this is not captured in the API, which does not differentiate between sessions that can be closed explicitly and sessions that can never be closed (MemorySession::global) or sessions that are only closed implicitly, by the garbage collector (MemorySession::openImplicit). Moreover, having all memory sessions implement AutoCloseable is also problematic: code such as MemorySession.global().allocate(100) can easily lead to warnings, as an AutoCloseable resource (the session) is used without a try-with-resources block.

Sharing segments

When MemorySession was first added to the FFM API in Java 17 (back then it was called ResourceScope), our aim has been to move the close operation away from memory segments [9]. However, that problem was never fully addressed: since clients might need to co-allocate a segment into an existing session (e.g. this is useful if a new segment should contain pointers to the existing segment), all memory segments expose a session accessor:

void accept(MemorySegment segment) {

    MemorySegment newSegment = MemorySegment.allocateNative(100, segment.session());

    newSegment.setAtIndex(ValueLayout.ADDRESS, 0, segment); // ok, co-allocation

    segment.session().close() // whoops!

}

Accessing a segment session is crucial to allow co-allocation. But having session accessors in segments is a risky move: since a session can be closed, clients of a segment can first access the segment session and then close it, thus potentially breaking encapsulation. This is a questionable API choice, especially when compared with other capability-based APIs: a MethodHandle doesn't expose the MethodHandles.Lookup instance with which it was created; nor does a Future exposes the ExecutorService which created it.

Our attempts at dropping session accessors [11] have not been successful. First, while we could simply remove the accessors and add some extra methods (e.g. liveness predicates) on MemorySegment, doing so would come at the expense of the important co-allocation use case. Second, memory segments are not the only entities which might expose session accessors, as custom allocator abstractions (as discussed above) might need those too: consider a SegmentAllocator which implements AutoCloseable, whose allocated segments are associated with the same lifetime.

For these reasons, the Java 19 API ended up with another compromise: all sessions are closeable, and segment clients can always access the segment session. But API writers can obtain a non-closeable view of an existing session: that is, a session that denotes the same lifetime as the original session, but whose close method always throws an exception:

private MemorySession privateSession = MemorySession.openShared();

...

public MemorySegment getSegment() {

    MemorySegment segment = MemorySegment.allocateNative(100, privateSession.asNonCloseable());

    return segment;

}

The above code creates a session (privateSession) and then allocates a segment in a non-closeable view of privateSession. As a result, clients will not be able to close the session associated with the returned segment.

Equals, hashCode, isCloseable

Once we buy the idea of having multiple views associated with a given lifetime, we need some ways to compare them. E.g. a client might wish to compare two non-closeable session views, and ask whether they are associated with the same lifetime. In other words, memory sessions should implement equals/hashCode. And, since now we have some sessions views that cannot be closed, the MemorySession::isCloseable predicate is also required. Unfortunately, the isCloseable predicate is itself a source of confusion. For instance, it might be odd that a session whose isCloseable predicate returns false can be still closed, either explicitly (from another session, for which the predicate yields true), or implicitly (by the garbage collector). That is, all sessions are closeable (except for the global session), in some way or another.

Owner thread

Some memory sessions (e.g. confined sessions) can have an owner thread. The owner thread, when set, determines which thread can close the session and access segments associated with that session. Exposing the owner thread directly, as FFM does, is problematic, for a couple of reasons. First, the MemorySession::ownerThread method is leaking too much information - the only sensible thing clients can do with the returned thread is to compare it either against null, or Thread.currentThread(). Secondly, the method conflates session ownership and access capabilities. As we shall see, other kinds of sessions are possible, structured sessions, which can be closed by one thread (the owner thread), but can be accessed by multiple threads.

Keeping sessions alive

In a world where segments can be invalidated at any time (even racily, by threads running concurrently), it might be sometimes desirable to keep the session associated with a memory segment alive for a period of time, e.g. while performing a certain critical operation. The memory session API provides a method to do this, namely MemorySession::whileAlive, which runs a user-defined action while keeping a session alive. This method works, but can only be used, effectively, in a stack-confined situation, e.g. where the computation can be expressed using a lambda expression. As we have seen in [10], other, more unstructured cases, are also possible (e.g. async I/O operations).

Close actions

The MemorySession::addCloseAction method was originally intended for attaching cleanup actions to unsafe native segments, i.e. created using MemorySession::ofAddress. Looking at existing code using the FFM API [1-8], it is clear that this method has become an attractive nuisance, where many uses of this method could be replaced by a finally block. Even worse, having a way to be notified when a session becomes invalid brings up uncomfortable questions: should a client be notified before, or after a session is closed? The only thing the FFM API can do safely is the latter, but any custom cleanup logic the client might wish to run on some segments allocated in a session would require the former.

Session factories

There are several memory session factories. First, there is a singleton global session, a session that is always alive and cannot be closed (see MemorySession::global). Then there's a factory to create a new implicit memory session, which cannot be closed explicitly, but is closed implicitly, by the garbage collector (see MemorySession::openImplicit). Then we have factories to create new confined and shared sessions (see MemorySession::openConfined and MemorySession::openShared, respectively). These latter factories come with overloads, as clients can also create sessions that are both explicitly and implicitly closeable. This functionality is never used in practice [1-8]. Supporting confined sessions that can also be closed implicitly is problematic: confined sessions are meant to be accessed by a single thread, but cleaners typically run on a different thread (the Java 19 implementation needs to handle these cases with extreme care). And, even in the shared case, there is evident overlapping with MemorySession::openImplicit - and clients could still (manually) take a shared closeable session and have its close method called by a cleaner (this is done in [3]).

Towards more principled sessions

The memory session API seems to be doing too many things at once. As we have seen, a session acts as a lifetime abstraction for memory segments, which is critical to ensure that access operations on memory segments are temporally safe. Furthermore, memory sessions enable the important co-allocation use case (see above).

But a memory session is also a capability: it has a close method which can be used by clients to atomically and deterministically invalidate all segments associated with that session. The lack of separation between the lifetime and capability roles of memory sessions is a red herring: as we have seen, it creates the need, for API developers, to protect against untrusted close (MemorySession::asNonCloseableView), which in turn leads to even more API noise (e..g MemorySession::equals, MemorySession::hashCode and MemorySession::isCloseable). And, sessions also act as allocators, but, as we have seen it is not very easy for clients to create custom temporally-bounded allocators.

What if we had a simpler memory session abstraction? Something like this:

sealed interface MemorySession {

    boolean isAlive();

    Thread ownerThread();

    ...

    static MemorySession openImplicit() { ... }

    static MemorySession global() { ... }

}

Here, MemorySession does not have a close method, nor does it implement AutoCloseable/SegmentAllocator. That is, MemorySession is now a pure lifetime abstraction, with some thread-confinement properties coming along for the ride. It is easy to see that, with a MemorySession defined this way, there is no issue with exposing a session accessor from memory segments (or from custom allocators) - after all, a session cannot be closed. Another advantage is that the above definition makes it clearer that implicit sessions and the global session cannot be closed explicitly. Overall, this seems a good step forward.

Arenas

But what if clients want to take advantage of deterministic deallocation? We could introduce a separate abstraction, called arena:

interface Arena extends SegmentAllocator, AutoCloseable {

    public MemorySegment allocate(long byteSize, long byteAlign);

    public MemorySession session();    

    public void close();

    static Arena openConfined() { ... }

    static Arena openShared() { ... }

}

An arena is a closeable allocator that is associated with a session - the arena session. All the segments allocated by the arena are associated with the arena session. When the arena is closed, the underlying session (which can be either shared or confined) is closed, thus invalidating all the segments associated with the arena session. Since arenas implement AutoCloseable, client code can be expressed succinctly, as follows:

try (Arena arena = Arena.open()) {

    MemorySegment nativeArray = arena.allocateArray(ValueLayout.JAVA_INT, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9);

    MemorySegment nativeString = arena.allocateUtf8String("Hello!");

    MemorySegment upcallStub = linker.upcallStub(handle, desc, arena.session());

    ...

} // memory released here

In other words, client code can use Arena pretty much in the same way they would use MemorySession in the Java 19 API.

Custom arenas

The attentive reader might have noted the relationship between arenas and sessions. An arena has a session. This means there's nothing too special about an arena: an arena is just a temporally-bounded abstraction, associated with some memory session (in the same way a memory segment is). This makes it easy to define custom temporally-bounded abstractions, based on Arena, as shown below:

class CustomArena implements Arena {

    private final Arena arena;

    CustomArena(Arena arena) {

        this.arena = arena;

    }

    public MemorySegment allocate(long byteSize, long byteAlign) {

        // insert custom allocation logic here

    }

    public MemorySession session() {

        return arena.session();

    }

    public void close() {

        // insert pre-close actions here

        arena.close();

        // insert post-close actions here

    }

}

A custom arena is just a class wrapping an Arena and adding some custom behavior, like pre- or post-close actions, and/or support for custom allocation strategies. Since exposing session accessors is now completely safe, the custom arena can just return the session associated with the wrapped arena, without the need to take extra precautions (such as resorting to non-closeable session views, as in the Java 19 API).

Sharper close actions

The ability to add close actions to an entire session is handy, but also overly general, and prone to both misunderstanding and abuse. The method MemorySession::addCloseAction could be easily replaced with an additional overload of MemorySegment::ofAddress which allows clients to attach a Runnable to an unsafely created native segments. This would still cover the vast majority of use cases, without compromising readability of the code using the FFM API. If a client wants to define an arena with custom pre- or post-close actions, they can do so by defining a custom arena, as shown in the above section.

Structured arenas (optional)

All the arenas we have seen so far are completely unstructured. That is, they can be closed at any time, and, in case of a shared session, accessed (and closed) by any thread. Conversely, structured arenas build on the same principle behind Loom's StructuredTaskScope abstraction [12]: the session associated with a structured arena can only be closed by the thread that created it, but it can be safely accessed not just by a single thread, but by all the threads forked inside that arena (e.g. using a StructuredTaskScope) - these threads keep the structured session alive:

try (Arena arena = Arena.openStructured()) {

    MemorySegment segment = arena.allocate(100);

    try (var scope = new StructuredTaskScope.ShutdownOnFailure()) {

        Future<Integer> first = scope.fork(() -> segment.getAtIndex(ValueLayout.JAVA_INT, 0);

        Future<Integer> second = scope.fork(() -> segment.getAtIndex(ValueLayout.JAVA_INT, 0);

        scope.join();

        int result = first.get() + second.get();

    } // close task scope    

} // release memory

As shown above, a structured session is something that sits in between confined sessions and shared sessions. It can only be closed by the thread that created it (as a confined session), but it can be accessed by multiple threads (as a shared session), all the threads that have been forked from a StructuredTaskScope nested inside the structured arena's try-with-resources block. Any attempt to close a structured arena when any of the threads accessing it have not yet terminated, will result in a StructureViolationException.

More temporal predicates (optional)

With the addition of structured session, clients might be interested to ask whether a session is nested inside another session. In other words, comparing two sessions for equality might not be enough in the world of structured sessions. To accommodate this, we might add a method on sessions, namely MemorySession::isAliveIn(MemorySession) which returns true if a session is always alive in the context of the provided session. As a bonus, we could make this method always return true if the provided session is the global session (since all sessions can be seen as nested inside the global session, which is always alive).

Unstructured dependencies (optional)

Structured arenas are handy, not only because they provide the flexibility of shared sessions with none of the costs, but also because they allow programs to reason about lifetime containment (e.g. the MemorySession::isAliveIn). Consider an hypothetical factory for two-dimensional matrices, such as the one below:

Matrix2d createMatrix(MemorySession matrixSession, MemorySegment matrixBuffer);

This factory takes a session, namely the lifetime associated with the returned matrix instance, and an external buffer, presumably containing the data associated with the matrix to be created. Now, for this API to work safely, and w/o the possibilities of JVM crashes, the factory would need to ensure that some temporal dependency exists between matrixSession and matrixBuffer. In other words:

assertTrue(matrixBuffer.session().isAliveIn(matrixSession));

Now, for structured sessions, we can imagine non-trivial cases where this check would succeed. But for plain unstructured sessions, we quickly run out of gas: the only case we can safely support is if either the buffer is associated with matrixSession, or if the buffer is associated with the global session.

It would be interesting to allow for ad-hoc dependencies between unstructured sessions to be expressed in the FFM API, following the approach described in [10]. If we had dependencies between unstructured sessions, we then could write code like this:

MemorySession bufferSession = ...

MemorySegment buffer = MemorySegment.allocate(1024, bufferSession);

...

try (Arena arena = Arena.openConfined(bufferSession)) {

    Matrix2d matrix = Matrix2d.createMatrix(arena.session(), buffer); // ok!

    ...

}

The above code creates a new confined arena, which declares an explicit dependency on the buffer session. This means that the buffer session cannot be closed until the arena is also closed. Since there's a clear lifetime dependency between the matrix session and the buffer session, MemorySession::isAliveIn returns true, and the matrix creation succeeds.

Finer-grained ownership predicates

We have seen how MemorySession::ownerThread cannot handle cases where a session can be accessed by more thread than just the owner (which is the case for structured session). We can replace ownerThread with a better set of predicates:

boolean isOwnedBy(Thread thread);

boolean isAccessibleBy(Thread thread);

The new predicates are, crucially, more granular, allowing clients to distinguish between ownership (who can call close) and access. Under this new light, our various session kinds can be clearly classified as follows:

• implicit and shared session have no owner, and can be accessed by any thread;
• a session confined on thread T is owned by T and can only be accessed by T;
• a structured session created on thread T is owned by T, and can be accessed by all threads forked inside the session try-with-resource block (e.g. using a StructuredTaskScope)

As a nice bonus, the new predicates also fully encapsulate the session's owner thread.

Conclusions

By dropping AutoCloseable and SegmentAllocator, we can turnMemorySession into a pure lifetime abstraction: since there's no close method, sessions can always be shared safely, and there's no need for MemorySession::asNonCloseable . This, in turn, allows us to also drop isCloseable, equals and hashCode from the API. We can also leave out MemorySession::addCloseAction - as that can be replaced with a more focused API (e.g. overload of MemorySegment::ofAddress, see above). In fact, this new Runnable-accepting unsafe segment factory is the true primitive on top of which mapped segments, upcall segments and library symbol segments are all built.

Moreover, a session is no longer a SegmentAllocator. The addition of the separate Arena API allows us to simplify user code in the vastly common case [1-8] where a bunch of logically related allocations occurs within a single code block. Moreover, arenas provide a solid foundation for building custom temporal abstractions, such as custom allocators.

Since structured sessions have a well-behaved close semantics, clients using them do not have to worry about premature close: methods like MemorySession::whileAlive are no-op on structured sessions. And, as virtual threads will be used more and more to turn asynchronous code into synchronous one, whileAlive might just be good enough, especially if we also add dependencies between unstructured sessions.

By providing a clear separation between the lifetime and capability roles of the memory session API, not only we have arrived at an API that looks conceptually simpler than the API in Java 19, but also one on top of which clients can build custom temporal abstractions in a more clean and predictable fashion.

Appendix: Full API

We report below, for completeness, the full API of MemorySession and Arena which includes all the tweaks discussed in the previous sections.

sealed interface MemorySession {

    boolean isOwnedBy(Thread thread);

    boolean isAccessibleBy(Thread thread); // 1

    boolean isAlive();

    boolean isAliveIn(MemorySession other); // 1, 2

    void whileAlive(Runnable action);

    public static MemorySession implicit() { ... }

    public static MemorySession implicit(MemorySession... parents) { ... } // 2

    public static MemorySession global() { ... }

}

interface Arena extends SegmentAllocator, AutoCloseable {

    public MemorySegment allocate(long byteSize, long byteAlign);

    public MemorySession session();

    public void close();

    static Arena openConfined() { ... }

    static Arena openConfined(MemorySession... parents) { ... } // 2

    static Arena openShared() { ... }

    static Arena openShared(MemorySession... parents) { ... } // 2

    static Arena openStructured() { ... } // 1

}

The methods in (1) and (2) are associated with optional extensions. More specifically, methods in (1) add support for structured sessions, whereas methods in (2) add support for dependencies between unstructured sessions.

References

[1] - https://github.com/manuelbl/JavaDoesUSB [2] - https://github.com/carldea/panama4newbies [3] - https://github.com/rmaucher/openssl-panama-foreign [4] - https://github.com/openjdk/jextract [5] - https://github.com/netty/netty-incubator-buffer-api [6] - https://github.com/apache/lucene [7] - https://github.com/cryptomator/jfuse [8] - https://github.com/jzy3d/panama-gl [9] - https://inside.java/2021/01/25/memory-access-pulling-all-the-threads/ [10] - https://inside.java/2021/10/12/panama-scope-dependencies/ [11] - https://mail.openjdk.org/pipermail/panama-dev/2022-February/016152.html [12] - https://docs.oracle.com/en/java/javase/19/docs/api/jdk.incubator.concurrent/jdk/incubator/concurrent/StructuredTaskScope.html