Merge pull request #19008 from ChengJin01/enable_heap_arg_downcall_v0.44.0

[FFI/v0.44] Enable heap related code in downcall

Author: pshipton

jcl/src/java.base/share/classes/openj9/internal/foreign/abi/InternalDowncallHandler.java

@@ -49,10 +49,11 @@
 import java.lang.foreign.ValueLayout;
 /*[IF JAVA_SPEC_VERSION >= 22]*/
 import jdk.internal.access.SharedSecrets;
+import jdk.internal.foreign.AbstractMemorySegmentImpl;
 /*[ENDIF] JAVA_SPEC_VERSION >= 22 */
+import jdk.internal.foreign.MemorySessionImpl;
 import jdk.internal.foreign.Utils;
 import jdk.internal.foreign.abi.LinkerOptions;
-import jdk.internal.foreign.MemorySessionImpl;
 /*[ELSE] JAVA_SPEC_VERSION >= 21 */
 import jdk.incubator.foreign.Addressable;
 import jdk.incubator.foreign.FunctionDescriptor;
@@ -89,6 +90,39 @@ public class InternalDowncallHandler {
 	private MemoryLayout realReturnLayout;
 	private MethodHandle boundMH;
 
+	/*[IF JAVA_SPEC_VERSION >= 22]*/
+	/* The identifier denotes the passed-in argument is a heap segment which helps
+	 * extract the address from the heap argument's base object in native.
+	 */
+	private static final long DOWNCALL_HEAP_ARGUMENT_ID = 0xFFFFFFFFFFFFFFFFL;
+
+	/* The ThreadLocal holds the information of the heap argument which includes
+	 * the base object array, the offset array plus the current heap argument index
+	 * for the current thread in multithreading. The heap base array stores the base
+	 * object (which contains the heap address as defined in OpenJDK) of the heap
+	 * arguments to extract the heap address from the base object in native.
+	 */
+	private static final class HeapArgInfo {
+		final Object[] bases;
+		final long[] offsets;
+		int index;
+
+		HeapArgInfo(int size) {
+			bases = new Object[size];
+			offsets = new long[size];
+			index = 0;
+		}
+
+		void append(Object base, long offset) {
+			bases[index] = base;
+			offsets[index] = offset;
+			index += 1;
+		}
+	}
+
+	private final ThreadLocal<HeapArgInfo> heapArgInfo;
+	/*[ENDIF] JAVA_SPEC_VERSION >= 22 */
+
 	/* The hashtables of sessions/scopes is intended for multithreading in which case
 	 * the same downcall handler might hold various sessions/scopes being used by
 	 * different threads in downcall.
@@ -151,11 +185,16 @@ public class InternalDowncallHandler {
 	/*[ENDIF] JAVA_SPEC_VERSION >= 22 */
 	private static synchronized native void resolveRequiredFields();
 	private native void initCifNativeThunkData(String[] argLayouts, String retLayout, boolean newArgTypes, int varArgIndex);
-	/*[IF JAVA_SPEC_VERSION >= 21]*/
-	private native long invokeNative(boolean isInTrivialDownCall, long returnStateMemAddr, long returnStructMemAddr, long functionAddress, long calloutThunk, long[] argValues);
-	/*[ELSE] JAVA_SPEC_VERSION >= 21 */
-	private native long invokeNative(long returnStructMemAddr, long functionAddress, long calloutThunk, long[] argValues);
-	/*[ENDIF] JAVA_SPEC_VERSION >= 21 */
+	private native long invokeNative(
+			/*[IF JAVA_SPEC_VERSION >= 22]*/
+			Object[] bases,
+			long[] offsets,
+			/*[ENDIF] JAVA_SPEC_VERSION >= 22 */
+			/*[IF JAVA_SPEC_VERSION >= 21]*/
+			boolean isInCriticalDownCall, long returnStateMemAddr,
+			/*[ENDIF] JAVA_SPEC_VERSION >= 21 */
+			long returnStructMemAddr, long functionAddress, long calloutThunk, long[] argValues
+			);
 
 	private static final class PrivateClassLock {
 		PrivateClassLock() {}
@@ -266,10 +305,25 @@ private void validateMemScope(ResourceScope memScope) throws IllegalStateExcepti
 	/* Intended for memSegmtOfPtrToLongArgFilter that converts the memory segment
 	 * of the passed-in pointer argument to long.
 	 */
-	private final long memSegmtOfPtrToLongArg(MemorySegment argValue) throws IllegalStateException {
-		UpcallMHMetaData.validateNativeArgRetSegmentOfPtr(argValue);
+	private final long memSegmtOfPtrToLongArg(MemorySegment argValue, LinkerOptions options) throws IllegalStateException {
+		UpcallMHMetaData.validateNativeArgRetSegmentOfPtr(argValue, options);
 		addMemArgScope(argValue.scope());
-		return argValue.address();
+
+		long address = argValue.address();
+		/*[IF JAVA_SPEC_VERSION >= 22]*/
+		if (!argValue.isNative() && options.allowsHeapAccess()) {
+			HeapArgInfo info = heapArgInfo.get();
+			if (info == null) {
+				info = new HeapArgInfo(argLayoutArray.length);
+				heapArgInfo.set(info);
+			}
+			/* Store the heap argument's base object and offset. */
+			AbstractMemorySegmentImpl segment = (AbstractMemorySegmentImpl)argValue;
+			info.append(segment.unsafeGetBase(), segment.unsafeGetOffset());
+			address = DOWNCALL_HEAP_ARGUMENT_ID;
+		}
+		/*[ENDIF] JAVA_SPEC_VERSION >= 22 */
+		return address;
 	}
 	/*[ELSE] JAVA_SPEC_VERSION >= 21 */
 	/* Intended for memAddrToLongArgFilter that converts the memory address to long. */
@@ -414,10 +468,14 @@ 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));
-			memSegmtOfPtrToLongArgFilter = lookup.bind(this, "memSegmtOfPtrToLongArg", methodType(long.class, MemorySegment.class));
+			memSegmtOfPtrToLongArgFilter = lookup.bind(this, "memSegmtOfPtrToLongArg", methodType(long.class, MemorySegment.class, LinkerOptions.class));
 			/*[ELSE] JAVA_SPEC_VERSION >= 21 */
 			memAddrToLongArgFilter = lookup.bind(this, "memAddrToLongArg", methodType(long.class, MemoryAddress.class));
 			/*[ENDIF] JAVA_SPEC_VERSION >= 21 */
@@ -603,7 +661,7 @@ private MethodHandle getArgumentFilter(Class<?> argTypeClass, MemoryLayout argLa
 			 * Note: AddressLayout is introduced in JDK21 to replace OfAddress.
 			 */
 			if (argLayout instanceof AddressLayout) {
-				filterMH = memSegmtOfPtrToLongArgFilter;
+				filterMH = MethodHandles.insertArguments(memSegmtOfPtrToLongArgFilter, 1, linkerOpts);
 			} else
 			/*[ENDIF] JAVA_SPEC_VERSION >= 21 */
 			{
@@ -785,6 +843,9 @@ Object runNativeMethod(Addressable downcallAddr, SegmentAllocator segmtAllocator
 		/*[ENDIF] JAVA_SPEC_VERSION >= 21 */
 
 		long returnVal;
+		/*[IF JAVA_SPEC_VERSION >= 22]*/
+		HeapArgInfo info = heapArgInfo.get();
+		/*[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.
 		 *
@@ -795,6 +856,8 @@ Object runNativeMethod(Addressable downcallAddr, SegmentAllocator segmtAllocator
 			SetDependency(arena.scope());
 			returnVal = invokeNative(
 					/*[IF JAVA_SPEC_VERSION >= 22]*/
+					(info != null) ? info.bases : null,
+					(info != null) ? info.offsets : null,
 					linkerOpts.isCritical(),
 					/*[ELSE] JAVA_SPEC_VERSION >= 22 */
 					linkerOpts.isTrivial(),
@@ -811,6 +874,16 @@ Object runNativeMethod(Addressable downcallAddr, SegmentAllocator segmtAllocator
 		releaseScope();
 		/*[ENDIF] JAVA_SPEC_VERSION >= 21 */
 
+		/*[IF JAVA_SPEC_VERSION >= 22]*/
+		/* Reset the heap argument information for the current thread so as to
+		 * update the base object of different heap arguments passed by the same
+		 * downcall handler within the current thread.
+		 */
+		if (info != null) {
+			heapArgInfo.remove();
+		}
+		/*[ENDIF] JAVA_SPEC_VERSION >= 22 */
+
 		/* This struct specific MemorySegment object returns to the current thread in the multithreading environment,
 		 * in which case the native invocations from threads end up with distinct returned structs.
 		 */

jcl/src/java.base/share/classes/openj9/internal/foreign/abi/InternalUpcallHandler.java

@@ -62,6 +62,9 @@ public final class InternalUpcallHandler {
 	private final MemoryLayout realReturnLayout;
 	private final long thunkAddr;
 	private UpcallMHMetaData metaData;
+	/*[IF JAVA_SPEC_VERSION >= 21]*/
+	private final LinkerOptions linkerOpts;
+	/*[ENDIF] JAVA_SPEC_VERSION >= 21 */
 
 	static final Lookup lookup = MethodHandles.lookup();
 
@@ -76,7 +79,7 @@ public final class InternalUpcallHandler {
 	static {
 		try {
 			/*[IF JAVA_SPEC_VERSION >= 21]*/
-			argRetSegmtOfPtrFilter = lookup.findStatic(InternalUpcallHandler.class, "argRetSegmtOfPtr", methodType(MemorySegment.class, MemorySegment.class, MemoryLayout.class));
+			argRetSegmtOfPtrFilter = lookup.findStatic(InternalUpcallHandler.class, "argRetSegmtOfPtr", methodType(MemorySegment.class, MemorySegment.class, MemoryLayout.class, LinkerOptions.class));
 			/*[ELSE] JAVA_SPEC_VERSION >= 21 */
 			argRetAddrOfPtrFilter = lookup.findStatic(InternalUpcallHandler.class, "argRetAddrOfPtr", methodType(MemoryAddress.class, MemoryAddress.class));
 			/*[ENDIF] JAVA_SPEC_VERSION >= 21 */
@@ -90,8 +93,8 @@ public final class InternalUpcallHandler {
 	 * of the passed-in pointer argument.
 	 */
 	/*[IF JAVA_SPEC_VERSION >= 21]*/
-	private static MemorySegment argRetSegmtOfPtr(MemorySegment argValue, MemoryLayout layout) throws IllegalStateException {
-		UpcallMHMetaData.validateNativeArgRetSegmentOfPtr(argValue);
+	private static MemorySegment argRetSegmtOfPtr(MemorySegment argValue, MemoryLayout layout, LinkerOptions options) throws IllegalStateException {
+		UpcallMHMetaData.validateNativeArgRetSegmentOfPtr(argValue, options);
 		return UpcallMHMetaData.getArgRetAlignedSegmentOfPtr(argValue.address(), layout);
 	}
 	/*[ELSE] JAVA_SPEC_VERSION >= 21 */
@@ -136,6 +139,9 @@ public InternalUpcallHandler(MethodHandle target, MethodType mt, FunctionDescrip
 		List<MemoryLayout> argLayouts = cDesc.argumentLayouts();
 		argLayoutArray = argLayouts.toArray(new MemoryLayout[argLayouts.size()]);
 		realReturnLayout = cDesc.returnLayout().orElse(null); // Set to null for void
+		/*[IF JAVA_SPEC_VERSION >= 21]*/
+		linkerOpts = options;
+		/*[ENDIF] JAVA_SPEC_VERSION >= 21 */
 
 		/*[IF JAVA_SPEC_VERSION == 17]*/
 		/* The layout check against the method type is still required for Java 17 in that both
@@ -244,7 +250,7 @@ else if (argTypeClass == MemoryAddress.class) {
 		else if (argLayout instanceof ValueLayout) {
 			/*[IF JAVA_SPEC_VERSION >= 21]*/
 			if (argLayout instanceof AddressLayout) {
-				filterMH = MethodHandles.insertArguments(argRetSegmtOfPtrFilter, 1, argLayout);
+				filterMH = MethodHandles.insertArguments(argRetSegmtOfPtrFilter, 1, argLayout, linkerOpts);
 			} else
 			/*[ENDIF] JAVA_SPEC_VERSION >= 21 */
 			{

jcl/src/java.base/share/classes/openj9/internal/foreign/abi/UpcallMHMetaData.java

@@ -34,16 +34,20 @@
 import java.lang.foreign.MemorySegment;
 import java.lang.foreign.MemoryLayout;
 import java.lang.foreign.MemorySegment.Scope;
+import jdk.internal.foreign.MemorySessionImpl;
 import jdk.internal.foreign.Utils;
 import jdk.internal.foreign.abi.LinkerOptions;
-import jdk.internal.foreign.MemorySessionImpl;
 /*[ELSE] JAVA_SPEC_VERSION >= 21 */
 import jdk.incubator.foreign.Addressable;
 import jdk.incubator.foreign.MemoryAddress;
 import jdk.incubator.foreign.MemorySegment;
 import jdk.incubator.foreign.ResourceScope;
 /*[ENDIF] JAVA_SPEC_VERSION >= 21 */
 
+/*[IF JAVA_SPEC_VERSION >= 22]*/
+import static java.lang.foreign.ValueLayout.*;
+/*[ENDIF] JAVA_SPEC_VERSION >= 22 */
+
 /**
  * The meta data consists of the callee MH and a cache of 2 elements for MH resolution,
  * which are used to generate an upcall handler to the requested java method.
@@ -127,11 +131,15 @@ static MemorySegment getArgRetAlignedSegmentOfPtr(long addrValue, MemoryLayout l
 	 * The method is shared in downcall and upcall.
 	 */
 	/*[IF JAVA_SPEC_VERSION >= 21]*/
-	static void validateNativeArgRetSegmentOfPtr(MemorySegment argRetSegmentOfPtr) {
+	static void validateNativeArgRetSegmentOfPtr(MemorySegment argRetSegmentOfPtr, LinkerOptions options) {
 		if (argRetSegmentOfPtr == null) {
 			throw new NullPointerException("A null pointer is not allowed.");
 		}
-		if (!argRetSegmentOfPtr.isNative()) {
+		if (!argRetSegmentOfPtr.isNative()
+		/*[IF JAVA_SPEC_VERSION >= 22]*/
+			&& !options.allowsHeapAccess()
+		/*[ENDIF] JAVA_SPEC_VERSION >= 22 */
+		) {
 			throw new IllegalArgumentException("Heap segment not allowed: " + argRetSegmentOfPtr);
 		}
 	}

runtime/oti/j9nonbuilder.h

@@ -5437,7 +5437,7 @@ typedef struct J9VMThread {
 	J9VMContinuation **continuationT1Cache;
 #endif /* JAVA_SPEC_VERSION >= 19 */
 #if JAVA_SPEC_VERSION >= 21
-	BOOLEAN isInTrivialDownCall;
+	BOOLEAN isInCriticalDownCall;
 #endif /* JAVA_SPEC_VERSION >= 21 */
 } J9VMThread;
 
@@ -6062,6 +6062,9 @@ typedef struct J9JavaVM {
 #define J9JAVAVM_DISCONTIGUOUS_INDEXABLE_HEADER_SIZE(vm) ((vm)->discontiguousIndexableHeaderSize)
 
 #if JAVA_SPEC_VERSION >= 16
+#if JAVA_SPEC_VERSION >= 22
+#define J9_FFI_DOWNCALL_HEAP_ARGUMENT_ID J9CONST_U64(0xFFFFFFFFFFFFFFFF)
+#endif /* JAVA_SPEC_VERSION >= 22 */
 
 /* The mask for the signature type identifier */
 #define J9_FFI_UPCALL_SIG_TYPE_MASK 0xF

runtime/tests/clinkerffi/CMakeLists.txt

@@ -332,6 +332,14 @@ omr_add_exports(clinkerffitests
 	addShortAndBytesFromUnion_Nested4ByteStruct
 	addIntAndIntShortFromUnion_Nested2ShortStruct
 	addIntAndShortsFromUnion_Nested4ShortStruct
+	addBoolsFromMultipleStructPtrs_returnStruct
+	addBytesFromMultipleStructPtrs_returnStruct
+	addCharsFromMultipleStructPtrs_returnStruct
+	addShortsFromMultipleStructPtrs_returnStruct
+	addIntsFromMultipleStructPtrs_returnStruct
+	addLongsFromMultipleStructPtrs_returnStruct
+	addFloatsFromMultipleStructPtrs_returnStruct
+	addDoublesFromMultipleStructPtrs_returnStruct
 	add2BoolsWithOrByUpcallMH
 	addBoolAndBoolFromPointerWithOrByUpcallMH
 	addBoolAndBoolFromNativePtrWithOrByUpcallMH