RFR: 8349146: [REDO] Implement a better allocator for downcalls [v7]

[Shaojin Wen]

On Fri, 2 May 2025 10:40:04 GMT, Per Minborg <pminborg at openjdk.org> wrote:

>> This PR is based on the work of @mernst-github and aims to implement an _internal_ thread-local 'stack' allocator, which works like a dynamically sized arena, but with reset functionality to reset the allocated size back to a certain level. The underlying memory could stay around between calls, which could improve performance.
>> 
>> Re-allocated segments are not zeroed between allocations.
>
> Per Minborg has updated the pull request incrementally with one additional commit since the last revision:
> 
>   Address comments

src/java.base/share/classes/jdk/internal/foreign/BufferStack.java line 186:

> 184:             }
> 185:         }
> 186:     }

Suggestion:

    private record PerThread(ReentrantLock lock, Arena arena, SlicingAllocator stack) {

        @ForceInline
        @SuppressWarnings("restricted")
        public Arena pushFrame(long size, long byteAlignment) {
            boolean needsLock = Thread.currentThread().isVirtual() && !lock.isHeldByCurrentThread();
            if (needsLock && !lock.tryLock()) {
                // Rare: another virtual thread on the same carrier competed for acquisition.
                return Arena.ofConfined();
            }
            if (!stack.canAllocate(size, byteAlignment)) {
                if (needsLock) lock.unlock();
                return Arena.ofConfined();
            }

            Arena confinedArena = Arena.ofConfined();
            long parentOffset = stack.currentOffset();
            MemorySegment frameSegment = stack.allocate(size, byteAlignment);
            long topOfStack = stack.currentOffset();
            // The cleanup action will keep the original automatic `arena` (from which
            // the reusable segment is first allocated) alive even if this Frame
            // becomes unreachable but there are reachable segments still alive.
            SegmentAllocator frame = new SlicingAllocator(frameSegment.reinterpret(confinedArena, new Frame.CleanupAction(arena)));

            return new Frame(this, needsLock, parentOffset, topOfStack, confinedArena, frame);
        }

        static PerThread of(long byteSize, long byteAlignment) {
            final Arena arena = Arena.ofAuto();
            return new PerThread(new ReentrantLock(),
                    arena,
                    new SlicingAllocator(arena.allocate(byteSize, byteAlignment)));
        }
    }

    private record Frame(PerThread thread, boolean locked, long parentOffset, long topOfStack, Arena confinedArena, SegmentAllocator frame) implements Arena {
        record CleanupAction(Arena arena) implements Consumer<MemorySegment> {
            @Override
            public void accept(MemorySegment memorySegment) {
                Reference.reachabilityFence(arena);
            }
        }

        @ForceInline
        private void assertOrder() {
            if (topOfStack != thread.stack.currentOffset())
                throw new IllegalStateException("Out of order access: frame not top-of-stack");
        }

        @ForceInline
        @Override
        @SuppressWarnings("restricted")
        public MemorySegment allocate(long byteSize, long byteAlignment) {
            // Make sure we are on the right thread and not closed
            MemorySessionImpl.toMemorySession(confinedArena).checkValidState();
            return frame.allocate(byteSize, byteAlignment);
        }

        @ForceInline
        @Override
        public MemorySegment.Scope scope() {
            return confinedArena.scope();
        }

        @ForceInline
        @Override
        public void close() {
            assertOrder();
            // the Arena::close method is called "early" as it checks thread
            // confinement and crucially before any mutation of the internal
            // state takes place.
            confinedArena.close();
            thread.stack.resetTo(parentOffset);
            if (locked) {
                thread.lock.unlock();
            }
        }
    }

I find it a bit strange to nest a class inside a record. What do you think if I change it to two records?

-------------

PR Review Comment: https://git.openjdk.org/jdk/pull/24829#discussion_r2071471740