Refactor HeapArgInfo to Support Queue-Based Tracking

Aditi Srinivas · Aditi Srinivas · commit f11b808ac6f0 · 2025-08-22T18:34:56.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. Clears the most recent entry after completion. FIFO based management for heap-based arguments in Java 22+.
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
@@ -27,6 +27,8 @@
 /*[ENDIF] JAVA_SPEC_VERSION >= 22 */
 import java.util.HashMap;
 import java.util.List;
+import java.util.Deque;
+import java.util.ArrayDeque;
 /*[IF JAVA_SPEC_VERSION >= 21]*/
 import java.util.Objects;
 /*[ENDIF] JAVA_SPEC_VERSION >= 21 */
@@ -123,7 +125,8 @@ void clear() {
 		}
 	}
 
-	private final ThreadLocal<HeapArgInfo> heapArgInfo;
+	private static final ThreadLocal<Deque<HeapArgInfo>> heapArgInfo = ThreadLocal.withInitial(ArrayDeque::new);
+	private static final ThreadLocal<Boolean> flagCreateHeapArgInfo = ThreadLocal.withInitial(()->Boolean.TRUE);
 	/*[ENDIF] JAVA_SPEC_VERSION >= 22 */
 
 	/* The hashtables of sessions/scopes is intended for multithreading in which case
@@ -311,7 +314,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 +327,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(!flagCreateHeapArgInfo.get() && !infoStack.isEmpty()) {
+				infoStack.pop();
+				flagCreateHeapArgInfo.set(Boolean.TRUE);
 			}
 			/*[ENDIF] JAVA_SPEC_VERSION >= 22 */
 			throw e;
 		}
 
 		long address = argValue.address();
 		/*[IF JAVA_SPEC_VERSION >= 22]*/
-		if (info == null) {
+
+		if ((info == null) || flagCreateHeapArgInfo.get()) {
 			info = new HeapArgInfo(argLayoutArray.length);
-			heapArgInfo.set(info);
+			infoStack.push(info);
+			flagCreateHeapArgInfo.set(Boolean.FALSE);
 		}
 
 		if (!argValue.isNative() && options.allowsHeapAccess()) {
@@ -369,9 +376,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(!flagCreateHeapArgInfo.get() && !infoStack.isEmpty()) {
+				infoStack.pop();
+				flagCreateHeapArgInfo.set(Boolean.TRUE);
 			}
 			/*[ENDIF] JAVA_SPEC_VERSION >= 22 */
 			throw e;
@@ -506,10 +514,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 +886,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;
+		boolean pushed = true;
+		if (flagCreateHeapArgInfo.get()) {
+			info = null;
+			pushed = false;
+		}
+		else {
+			info = infoStack.peek();
+			flagCreateHeapArgInfo.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 +950,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 (!infoStack.isEmpty() && pushed) {
+				infoStack.pop();
 			}
 		}
 		/*[ENDIF] JAVA_SPEC_VERSION >= 22 */
