Skip to content

Commit 9f9e73d

Browse files
minborgminborg
committed
8349146: [REDO] Implement a better allocator for downcalls
Reviewed-by: mcimadamore, jvernee, liach
1 parent 995d541 commit 9f9e73d

File tree

11 files changed

+989
-21
lines changed

11 files changed

+989
-21
lines changed

‎src/java.base/share/classes/jdk/internal/foreign/BufferStack.java

Lines changed: 214 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,214 @@
1+
/*
2+
* Copyright (c) 2025, Oracle and/or its affiliates. All rights reserved.
3+
* DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
4+
*
5+
* This code is free software; you can redistribute it and/or modify it
6+
* under the terms of the GNU General Public License version 2 only, as
7+
* published by the Free Software Foundation. Oracle designates this
8+
* particular file as subject to the "Classpath" exception as provided
9+
* by Oracle in the LICENSE file that accompanied this code.
10+
*
11+
* This code is distributed in the hope that it will be useful, but WITHOUT
12+
* ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
13+
* FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
14+
* version 2 for more details (a copy is included in the LICENSE file that
15+
* accompanied this code).
16+
*
17+
* You should have received a copy of the GNU General Public License version
18+
* 2 along with this work; if not, write to the Free Software Foundation,
19+
* Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
20+
*
21+
* Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
22+
* or visit www.oracle.com if you need additional information or have any
23+
* questions.
24+
*/
25+
26+
package jdk.internal.foreign;
27+
28+
import jdk.internal.misc.CarrierThreadLocal;
29+
import jdk.internal.vm.annotation.ForceInline;
30+
31+
import java.lang.foreign.Arena;
32+
import java.lang.foreign.MemoryLayout;
33+
import java.lang.foreign.MemorySegment;
34+
import java.lang.foreign.SegmentAllocator;
35+
import java.lang.ref.Reference;
36+
import java.util.concurrent.locks.ReentrantLock;
37+
import java.util.function.Consumer;
38+
39+
/**
40+
* A buffer stack that allows efficient reuse of memory segments. This is useful in cases
41+
* where temporary memory is needed.
42+
* <p>
43+
* Use the factories {@code BufferStack.of(...)} to create new instances of this class.
44+
* <p>
45+
* Note: The reused segments are neither zeroed out before nor after re-use.
46+
*/
47+
public final class BufferStack {
48+
49+
private final long byteSize;
50+
private final long byteAlignment;
51+
private final CarrierThreadLocal<PerThread> tl;
52+
53+
private BufferStack(long byteSize, long byteAlignment) {
54+
this.byteSize = byteSize;
55+
this.byteAlignment = byteAlignment;
56+
this.tl = new CarrierThreadLocal<>() {
57+
@Override
58+
protected BufferStack.PerThread initialValue() {
59+
return BufferStack.PerThread.of(byteSize, byteAlignment);
60+
}
61+
};
62+
}
63+
64+
/**
65+
* {@return a new Arena that tries to provide {@code byteSize} and {@code byteAlignment}
66+
* allocations by recycling the BufferStack's internal memory}
67+
*
68+
* @param byteSize to be reserved from this BufferStack's internal memory
69+
* @param byteAlignment to be used for reservation
70+
*/
71+
@ForceInline
72+
public Arena pushFrame(long byteSize, long byteAlignment) {
73+
return tl.get().pushFrame(byteSize, byteAlignment);
74+
}
75+
76+
/**
77+
* {@return a new Arena that tries to provide {@code byteSize}
78+
* allocations by recycling the BufferStack's internal memory}
79+
*
80+
* @param byteSize to be reserved from this BufferStack's internal memory
81+
*/
82+
@ForceInline
83+
public Arena pushFrame(long byteSize) {
84+
return pushFrame(byteSize, 1);
85+
}
86+
87+
/**
88+
* {@return a new Arena that tries to provide {@code layout}
89+
* allocations by recycling the BufferStack's internal memory}
90+
*
91+
* @param layout for which to reserve internal memory
92+
*/
93+
@ForceInline
94+
public Arena pushFrame(MemoryLayout layout) {
95+
return pushFrame(layout.byteSize(), layout.byteAlignment());
96+
}
97+
98+
@Override
99+
public String toString() {
100+
return "BufferStack[byteSize=" + byteSize + ", byteAlignment=" + byteAlignment + "]";
101+
}
102+
103+
private record PerThread(ReentrantLock lock,
104+
Arena arena,
105+
SlicingAllocator stack,
106+
CleanupAction cleanupAction) {
107+
108+
@ForceInline
109+
public Arena pushFrame(long size, long byteAlignment) {
110+
boolean needsLock = Thread.currentThread().isVirtual() && !lock.isHeldByCurrentThread();
111+
if (needsLock && !lock.tryLock()) {
112+
// Rare: another virtual thread on the same carrier competed for acquisition.
113+
return Arena.ofConfined();
114+
}
115+
if (!stack.canAllocate(size, byteAlignment)) {
116+
if (needsLock) lock.unlock();
117+
return Arena.ofConfined();
118+
}
119+
return new Frame(needsLock, size, byteAlignment);
120+
}
121+
122+
static PerThread of(long byteSize, long byteAlignment) {
123+
final Arena arena = Arena.ofAuto();
124+
return new PerThread(new ReentrantLock(),
125+
arena,
126+
new SlicingAllocator(arena.allocate(byteSize, byteAlignment)),
127+
new CleanupAction(arena));
128+
}
129+
130+
private record CleanupAction(Arena arena) implements Consumer<MemorySegment> {
131+
@Override
132+
public void accept(MemorySegment memorySegment) {
133+
Reference.reachabilityFence(arena);
134+
}
135+
}
136+
137+
private final class Frame implements Arena {
138+
139+
private final boolean locked;
140+
private final long parentOffset;
141+
private final long topOfStack;
142+
private final Arena confinedArena;
143+
private final SegmentAllocator frame;
144+
145+
@SuppressWarnings("restricted")
146+
@ForceInline
147+
public Frame(boolean locked, long byteSize, long byteAlignment) {
148+
this.locked = locked;
149+
this.parentOffset = stack.currentOffset();
150+
final MemorySegment frameSegment = stack.allocate(byteSize, byteAlignment);
151+
this.topOfStack = stack.currentOffset();
152+
this.confinedArena = Arena.ofConfined();
153+
// The cleanup action will keep the original automatic `arena` (from which
154+
// the reusable segment is first allocated) alive even if this Frame
155+
// becomes unreachable but there are reachable segments still alive.
156+
this.frame = new SlicingAllocator(frameSegment.reinterpret(confinedArena, cleanupAction));
157+
}
158+
159+
@ForceInline
160+
private void assertOrder() {
161+
if (topOfStack != stack.currentOffset())
162+
throw new IllegalStateException("Out of order access: frame not top-of-stack");
163+
}
164+
165+
@ForceInline
166+
@Override
167+
@SuppressWarnings("restricted")
168+
public MemorySegment allocate(long byteSize, long byteAlignment) {
169+
// Make sure we are on the right thread and not closed
170+
MemorySessionImpl.toMemorySession(confinedArena).checkValidState();
171+
return frame.allocate(byteSize, byteAlignment);
172+
}
173+
174+
@ForceInline
175+
@Override
176+
public MemorySegment.Scope scope() {
177+
return confinedArena.scope();
178+
}
179+
180+
@ForceInline
181+
@Override
182+
public void close() {
183+
assertOrder();
184+
// the Arena::close method is called "early" as it checks thread
185+
// confinement and crucially before any mutation of the internal
186+
// state takes place.
187+
confinedArena.close();
188+
stack.resetTo(parentOffset);
189+
if (locked) {
190+
lock.unlock();
191+
}
192+
}
193+
}
194+
}
195+
196+
public static BufferStack of(long byteSize, long byteAlignment) {
197+
if (byteSize < 0) {
198+
throw new IllegalArgumentException("Negative byteSize: " + byteSize);
199+
}
200+
if (byteAlignment < 0) {
201+
throw new IllegalArgumentException("Negative byteAlignment: " + byteAlignment);
202+
}
203+
return new BufferStack(byteSize, byteAlignment);
204+
}
205+
206+
public static BufferStack of(long byteSize) {
207+
return new BufferStack(byteSize, 1);
208+
}
209+
210+
public static BufferStack of(MemoryLayout layout) {
211+
// Implicit null check
212+
return of(layout.byteSize(), layout.byteAlignment());
213+
}
214+
}

‎src/java.base/share/classes/jdk/internal/foreign/SlicingAllocator.java

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright (c) 2021, 2023, Oracle and/or its affiliates. All rights reserved.
2+
* Copyright (c) 2021, 2025, Oracle and/or its affiliates. All rights reserved.
33
* DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
44
*
55
* This code is free software; you can redistribute it and/or modify it
@@ -38,6 +38,22 @@ public SlicingAllocator(MemorySegment segment) {
3838
this.segment = segment;
3939
}
4040

41+
public long currentOffset() {
42+
return sp;
43+
}
44+
45+
public void resetTo(long offset) {
46+
if (offset < 0 || offset > sp)
47+
throw new IllegalArgumentException(String.format("offset %d should be in [0, %d] ", offset, sp));
48+
this.sp = offset;
49+
}
50+
51+
public boolean canAllocate(long byteSize, long byteAlignment) {
52+
long min = segment.address();
53+
long start = Utils.alignUp(min + sp, byteAlignment) - min;
54+
return start + byteSize <= segment.byteSize();
55+
}
56+
4157
MemorySegment trySlice(long byteSize, long byteAlignment) {
4258
long min = segment.address();
4359
long start = Utils.alignUp(min + sp, byteAlignment) - min;

‎src/java.base/share/classes/jdk/internal/foreign/abi/SharedUtils.java

Lines changed: 6 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -24,9 +24,9 @@
2424
*/
2525
package jdk.internal.foreign.abi;
2626

27-
import jdk.internal.access.JavaLangAccess;
2827
import jdk.internal.access.JavaLangInvokeAccess;
2928
import jdk.internal.access.SharedSecrets;
29+
import jdk.internal.foreign.BufferStack;
3030
import jdk.internal.foreign.CABI;
3131
import jdk.internal.foreign.abi.AbstractLinker.UpcallStubFactory;
3232
import jdk.internal.foreign.abi.aarch64.linux.LinuxAArch64Linker;
@@ -390,26 +390,12 @@ static long pickChunkOffset(long chunkOffset, long byteWidth, int chunkWidth) {
390390
: chunkOffset;
391391
}
392392

393-
public static Arena newBoundedArena(long size) {
394-
return new Arena() {
395-
final Arena arena = Arena.ofConfined();
396-
final SegmentAllocator slicingAllocator = SegmentAllocator.slicingAllocator(arena.allocate(size));
397-
398-
@Override
399-
public Scope scope() {
400-
return arena.scope();
401-
}
393+
private static final int LINKER_STACK_SIZE = Integer.getInteger("jdk.internal.foreign.LINKER_STACK_SIZE", 256);
394+
private static final BufferStack LINKER_STACK = BufferStack.of(LINKER_STACK_SIZE, 1);
402395

403-
@Override
404-
public void close() {
405-
arena.close();
406-
}
407-
408-
@Override
409-
public MemorySegment allocate(long byteSize, long byteAlignment) {
410-
return slicingAllocator.allocate(byteSize, byteAlignment);
411-
}
412-
};
396+
@ForceInline
397+
public static Arena newBoundedArena(long size) {
398+
return LINKER_STACK.pushFrame(size, 8);
413399
}
414400

415401
public static Arena newEmptyArena() {

‎test/jdk/ProblemList.txt

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -784,6 +784,8 @@ jdk/jfr/jvm/TestWaste.java 8282427 generic-
784784

785785
# jdk_foreign
786786

787+
java/foreign/TestBufferStackStress.java 8350455 macosx-all
788+
787789
############################################################################
788790
# Client manual tests
789791

0 commit comments

Comments
 (0)