Refactor HeapArgInfo to Support Recursive Downcalls

Aditi Srinivas · Aditi Srinivas · commit 22869c23567d · 2025-08-26T05:47:08.000-07:00
Fix for Issue #21944 Fix for failing java/foreign/TestBufferStack.java test. Replaced single HeapArgInfo with a thread-local Deque. To track multiple argument sets for recursions. Each downcall pushes a new HeapArgInfo into Deque. Clears the most recent entry after downcall completion.
diff --git a/jcl/src/java.base/share/classes/openj9/internal/foreign/abi/InternalDowncallHandler.java b/jcl/src/java.base/share/classes/openj9/internal/foreign/abi/InternalDowncallHandler.java
@@ -23,7 +23,9 @@
 package openj9.internal.foreign.abi;
 
 /*[IF JAVA_SPEC_VERSION >= 22]*/
+import java.util.ArrayDeque;
 import java.util.Arrays;
+import java.util.Deque;
 /*[ENDIF] JAVA_SPEC_VERSION >= 22 */
 import java.util.HashMap;
 import java.util.List;
@@ -123,7 +125,18 @@ void clear() {
 		}
 	}
 
-	private final ThreadLocal<HeapArgInfo> heapArgInfo;
+	/* heapArgInfo holds a per-thread stack of HeapArgInfo structures.
+	 * Each entry tracks Java references to heap-allocated memory segments
+	 * Required for recursive downcalls.
+	 * Each downcall maintains its own set of references until completion.
+	 *
+	 * createHeapArgInfo controls when a new entry should be created.
+	 * If true, a new HeapArgInfo is pushed for the current downcall.
+	 * A single downcall may have multiple pointer arguments,
+	 * and all of them share the same entry.
+	 */
+	private static final ThreadLocal<Deque<HeapArgInfo>> heapArgInfo = ThreadLocal.withInitial(ArrayDeque::new);
+	private static final ThreadLocal<Boolean> createHeapArgInfo = ThreadLocal.withInitial(()->Boolean.TRUE);
 	/*[ENDIF] JAVA_SPEC_VERSION >= 22 */
 
 	/* The hashtables of sessions/scopes is intended for multithreading in which case
@@ -311,7 +324,8 @@ private void validateMemScope(ResourceScope memScope) throws IllegalStateExcepti
 	 */
 	private final long memSegmtOfPtrToLongArg(MemorySegment argValue, LinkerOptions options) throws IllegalStateException {
 		/*[IF JAVA_SPEC_VERSION >= 22]*/
-		HeapArgInfo info = heapArgInfo.get();
+		Deque<HeapArgInfo> infoStack = heapArgInfo.get();
+		HeapArgInfo info = infoStack.peek();
 		/*[ENDIF] JAVA_SPEC_VERSION >= 22 */
 
 		try {
@@ -323,18 +337,21 @@ private final long memSegmtOfPtrToLongArg(MemorySegment argValue, LinkerOptions
 			 * is captured so as to reset the internal index and avoid retaining the references
 			 * to the unreachable objects.
 			 */
-			if (info != null) {
-				info.clear();
+			if (!createHeapArgInfo.get()) {
+				infoStack.pop();
+				createHeapArgInfo.set(Boolean.TRUE);
 			}
 			/*[ENDIF] JAVA_SPEC_VERSION >= 22 */
 			throw e;
 		}
 
 		long address = argValue.address();
 		/*[IF JAVA_SPEC_VERSION >= 22]*/
-		if (info == null) {
+
+		if (createHeapArgInfo.get()) {
 			info = new HeapArgInfo(argLayoutArray.length);
-			heapArgInfo.set(info);
+			infoStack.push(info);
+			createHeapArgInfo.set(Boolean.FALSE);
 		}
 
 		if (!argValue.isNative() && options.allowsHeapAccess()) {
@@ -369,9 +386,10 @@ private final long memSegmtToLongArg(MemorySegment argValue) throws IllegalState
 			 * is captured so as to reset the internal index and avoid retaining the references
 			 * to the unreachable objects.
 			 */
-			HeapArgInfo info = heapArgInfo.get();
-			if (info != null) {
-				info.clear();
+			Deque<HeapArgInfo> infoStack = heapArgInfo.get();
+			if (!createHeapArgInfo.get()) {
+				infoStack.pop();
+				createHeapArgInfo.set(Boolean.TRUE);
 			}
 			/*[ENDIF] JAVA_SPEC_VERSION >= 22 */
 			throw e;
@@ -506,10 +524,6 @@ public InternalDowncallHandler(MethodType functionMethodType, FunctionDescriptor
 		scopeHandleMap = new ConcurrentHashMap<>();
 		/*[ENDIF] JAVA_SPEC_VERSION == 17 */
 
-		/*[IF JAVA_SPEC_VERSION >= 22]*/
-		heapArgInfo = new ThreadLocal<>();
-		/*[ENDIF] JAVA_SPEC_VERSION >= 22 */
-
 		try {
 			/*[IF JAVA_SPEC_VERSION >= 21]*/
 			longObjToMemSegmtRetFilter = lookup.bind(this, "longObjToMemSegmtRet", methodType(MemorySegment.class, Object.class));
@@ -882,7 +896,17 @@ Object runNativeMethod(Addressable downcallAddr, SegmentAllocator segmtAllocator
 
 		long returnVal;
 		/*[IF JAVA_SPEC_VERSION >= 22]*/
-		HeapArgInfo info = heapArgInfo.get();
+		Deque<HeapArgInfo> infoStack = heapArgInfo.get();
+		HeapArgInfo info;
+		/* Check if current downcall has pointer arguments. */
+		boolean pointerArgs = true;
+		if (createHeapArgInfo.get()) {
+			info = null;
+			pointerArgs = false;
+		} else {
+			info = infoStack.peek();
+			createHeapArgInfo.set(Boolean.TRUE);
+		}
 		/*[ENDIF] JAVA_SPEC_VERSION >= 22 */
 		/* The scope associated with memory specific arguments must be kept alive
 		 * during the downcall since JDK17, including the downcall adddress.
@@ -936,8 +960,8 @@ Object runNativeMethod(Addressable downcallAddr, SegmentAllocator segmtAllocator
 			 * so as to reset the internal index and avoid retaining the references to the
 			 * unreachable objects.
 			 */
-			if (info != null) {
-				info.clear();
+			if (pointerArgs) {
+				infoStack.pop();
 			}
 		}
 		/*[ENDIF] JAVA_SPEC_VERSION >= 22 */
