JDK-8244601: Cleanup support for upcall handles

Reviewed-by: sundar, jvernee

src/hotspot/share/prims/upcallStubs.cpp

@@ -25,7 +25,7 @@
 #include "runtime/jniHandles.inline.hpp"
 #include "runtime/interfaceSupport.inline.hpp"
 
-JVM_ENTRY(static jboolean, UH_FreeUpcallStub(JNIEnv *env, jobject _unused, jlong addr))
+JVM_ENTRY(static jboolean, UH_FreeUpcallStub0(JNIEnv *env, jobject _unused, jlong addr))
   //acquire code cache lock
   MutexLocker mu(CodeCache_lock, Mutex::_no_safepoint_check_flag);
   //find code blob
@@ -47,7 +47,7 @@ JVM_END
 
 // These are the native methods on jdk.internal.foreign.NativeInvoker.
 static JNINativeMethod UH_methods[] = {
-  {CC "freeUpcallStub",     CC "(J)Z",                FN_PTR(UH_FreeUpcallStub)}
+  {CC "freeUpcallStub0",     CC "(J)Z",                FN_PTR(UH_FreeUpcallStub0)}
 };
 
 /**

src/jdk.incubator.foreign/share/classes/jdk/incubator/foreign/SystemABI.java

@@ -74,25 +74,17 @@
     MethodHandle downcallHandle(MemoryAddress symbol, MethodType type, FunctionDescriptor function);
 
     /**
-     * Obtain the pointer to a native stub which
-     * can be used to upcall into a given method handle.
+     * Allocates a native stub segment which contains executable code to upcall into a given method handle.
+     * As such, the base address of the returned stub segment can be passed to other foreign functions
+     * (as a function pointer). The returned segment is <em>not</em> thread-confined, and it only features
+     * the {@link MemorySegment#CLOSE} access mode. When the returned segment is closed,
+     * the corresponding native stub will be deallocated.
      *
      * @param target the target method handle.
      * @param function the function descriptor.
-     * @return the upcall symbol.
+     * @return the native stub segment.
      */
-    MemoryAddress upcallStub(MethodHandle target, FunctionDescriptor function);
-
-    /**
-     * Frees an upcall stub given it's memory address.
-     *
-     * @param address the memory address of the upcall stub, returned from
-     *                {@link SystemABI#upcallStub(MethodHandle, FunctionDescriptor)}.
-     * @throws IllegalArgumentException if the given address is not a valid upcall stub address.
-     */
-    default void freeUpcallStub(MemoryAddress address) {
-        UpcallStubs.freeUpcallStub(address);
-    }
+    MemorySegment upcallStub(MethodHandle target, FunctionDescriptor function);
 
     /**
      * Returns the name of this ABI.

src/jdk.incubator.foreign/share/classes/jdk/internal/foreign/abi/UpcallStubs.java

@@ -25,26 +25,29 @@
 package jdk.internal.foreign.abi;
 
 import jdk.incubator.foreign.MemoryAddress;
+import jdk.incubator.foreign.MemorySegment;
 import jdk.internal.foreign.MemoryAddressImpl;
+import jdk.internal.foreign.NativeMemorySegmentImpl;
 
 public class UpcallStubs {
 
-    public static MemoryAddress upcallAddress(UpcallHandler handler) {
-        long addr = handler.entryPoint();
-        return MemoryAddress.ofLong(addr);
-    }
+    public static MemorySegment upcallAddress(UpcallHandler handler) {
+        long stubAddress = handler.entryPoint();
+        return NativeMemorySegmentImpl.makeNativeSegmentUnchecked(MemoryAddress.ofLong(stubAddress),
+                0, null, () -> freeUpcallStub(stubAddress), null)
+                .withAccessModes(MemorySegment.CLOSE);
+    };
 
-    public static void freeUpcallStub(MemoryAddress address) {
-        MemoryAddressImpl ma = (MemoryAddressImpl) address;
-        if (ma.unsafeGetBase() != null || !freeUpcallStub(ma.unsafeGetOffset())) {
-            throw new IllegalArgumentException("Not a stub address: " + address);
+    private static void freeUpcallStub(long stubAddress) {
+        if (!freeUpcallStub0(stubAddress)) {
+            throw new IllegalStateException("Not a stub address: " + stubAddress);
         }
     }
 
     // natives
 
     // returns true if the stub was found (and freed)
-    private static native boolean freeUpcallStub(long addr);
+    private static native boolean freeUpcallStub0(long addr);
 
     private static native void registerNatives();
     static {

src/jdk.incubator.foreign/share/classes/jdk/internal/foreign/abi/aarch64/AArch64ABI.java

@@ -28,6 +28,7 @@
 import jdk.incubator.foreign.FunctionDescriptor;
 import jdk.incubator.foreign.MemoryAddress;
 import jdk.incubator.foreign.MemoryLayout;
+import jdk.incubator.foreign.MemorySegment;
 import jdk.incubator.foreign.SystemABI;
 import jdk.internal.foreign.abi.*;
 
@@ -54,7 +55,7 @@ public MethodHandle downcallHandle(MemoryAddress symbol, MethodType type, Functi
     }
 
     @Override
-    public MemoryAddress upcallStub(MethodHandle target, FunctionDescriptor function) {
+    public MemorySegment upcallStub(MethodHandle target, FunctionDescriptor function) {
         return UpcallStubs.upcallAddress(CallArranger.arrangeUpcall(target, target.type(), function));
     }
 

src/jdk.incubator.foreign/share/classes/jdk/internal/foreign/abi/x64/sysv/SysVx64ABI.java

@@ -27,6 +27,7 @@
 import jdk.incubator.foreign.FunctionDescriptor;
 import jdk.incubator.foreign.MemoryAddress;
 import jdk.incubator.foreign.MemoryLayout;
+import jdk.incubator.foreign.MemorySegment;
 import jdk.incubator.foreign.SystemABI;
 import jdk.internal.foreign.abi.*;
 
@@ -59,7 +60,7 @@ public MethodHandle downcallHandle(MemoryAddress symbol, MethodType type, Functi
     }
 
     @Override
-    public MemoryAddress upcallStub(MethodHandle target, FunctionDescriptor function) {
+    public MemorySegment upcallStub(MethodHandle target, FunctionDescriptor function) {
         return UpcallStubs.upcallAddress(CallArranger.arrangeUpcall(target, target.type(), function));
     }
 

src/jdk.incubator.foreign/share/classes/jdk/internal/foreign/abi/x64/windows/Windowsx64ABI.java

@@ -27,6 +27,7 @@
 import jdk.incubator.foreign.FunctionDescriptor;
 import jdk.incubator.foreign.MemoryAddress;
 import jdk.incubator.foreign.MemoryLayout;
+import jdk.incubator.foreign.MemorySegment;
 import jdk.incubator.foreign.SystemABI;
 import jdk.internal.foreign.abi.x64.sysv.ArgumentClassImpl;
 import jdk.internal.foreign.abi.*;
@@ -61,7 +62,7 @@ public MethodHandle downcallHandle(MemoryAddress symbol, MethodType type, Functi
     }
 
     @Override
-    public MemoryAddress upcallStub(MethodHandle target, FunctionDescriptor function) {
+    public MemorySegment upcallStub(MethodHandle target, FunctionDescriptor function) {
         return UpcallStubs.upcallAddress(CallArranger.arrangeUpcall(target, target.type(), function));
     }
 

test/jdk/java/foreign/CallGeneratorHelper.java

@@ -333,12 +333,7 @@ static void cleanup(Object arg) {
         if (arg instanceof MemoryAddress) {
             cleanup(((MemoryAddress)arg).segment());
         } else if (arg instanceof MemorySegment) {
-            try {
-                ((MemorySegment) arg).close();
-            } catch (UnsupportedOperationException e) {
-                assertTrue(e.getMessage().contains("Required access mode"));
-                // fine, NOTHING segment for upcall stubs
-            }
+            ((MemorySegment) arg).close();
         }
     }
 

test/jdk/java/foreign/StdLibTest.java

@@ -308,9 +308,9 @@ boolean isdst() {
                         .forEach(i -> intArrHandle.set(nativeArr.baseAddress(), i, arr[i]));
 
                 //call qsort
-                MemoryAddress qsortUpcallAddr = abi.upcallStub(qsortCompar.bindTo(nativeArr), qsortComparFunction);
-                qsort.invokeExact(nativeArr.baseAddress(), seq.elementCount().getAsLong(), C_INT.byteSize(), qsortUpcallAddr);
-                abi.freeUpcallStub(qsortUpcallAddr);
+                try (MemorySegment qsortUpcallStub = abi.upcallStub(qsortCompar.bindTo(nativeArr), qsortComparFunction)) {
+                    qsort.invokeExact(nativeArr.baseAddress(), seq.elementCount().getAsLong(), C_INT.byteSize(), qsortUpcallStub.baseAddress());
+                }
 
                 //convert back to Java array
                 return LongStream.range(0, arr.length)

test/jdk/java/foreign/TestUpcall.java

@@ -41,6 +41,8 @@
 import jdk.incubator.foreign.MemorySegment;
 import jdk.incubator.foreign.SystemABI;
 import jdk.incubator.foreign.ValueLayout;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
 
 import java.lang.invoke.MethodHandle;
@@ -63,8 +65,6 @@
 
     static LibraryLookup lib = LibraryLookup.ofLibrary("TestUpcall");
     static SystemABI abi = SystemABI.getSystemABI();
-    static final MemoryAddress dummyAddress;
-    static final Cleaner cleaner = Cleaner.create();
 
     static MethodHandle DUMMY;
     static MethodHandle PASS_AND_SAVE;
@@ -73,14 +73,22 @@
         try {
             DUMMY = MethodHandles.lookup().findStatic(TestUpcall.class, "dummy", MethodType.methodType(void.class));
             PASS_AND_SAVE = MethodHandles.lookup().findStatic(TestUpcall.class, "passAndSave", MethodType.methodType(Object.class, Object[].class, AtomicReference.class));
-
-            dummyAddress = abi.upcallStub(DUMMY, FunctionDescriptor.ofVoid());
-            cleaner.register(dummyAddress, () -> abi.freeUpcallStub(dummyAddress));
         } catch (Throwable ex) {
             throw new IllegalStateException(ex);
         }
     }
 
+    static MemoryAddress dummyAddress;
+
+    @BeforeClass
+    void setup() {
+        dummyAddress = abi.upcallStub(DUMMY, FunctionDescriptor.ofVoid()).baseAddress();
+    }
+
+    @AfterClass
+    void teardown() {
+        dummyAddress.segment().close();
+    }
 
     @Test(dataProvider="functions", dataProviderClass=CallGeneratorHelper.class)
     public void testUpcalls(String fName, Ret ret, List<ParamType> paramTypes, List<StructFieldType> fields) throws Throwable {
@@ -96,7 +104,9 @@ public void testUpcalls(String fName, Ret ret, List<ParamType> paramTypes, List<
             returnChecks.forEach(c -> c.accept(res));
         }
         for (Object arg : args) {
-            cleanup(arg);
+            if (arg != dummyAddress) {
+                cleanup(arg);
+            }
         }
     }
 
@@ -168,8 +178,7 @@ static MemoryAddress makeCallback(Ret ret, List<ParamType> params, List<StructFi
         FunctionDescriptor func = ret != Ret.VOID
                 ? FunctionDescriptor.of(firstlayout, paramLayouts)
                 : FunctionDescriptor.ofVoid(paramLayouts);
-        MemoryAddress stub = abi.upcallStub(mh, func);
-        cleaner.register(stub, () -> abi.freeUpcallStub(stub));
+        MemoryAddress stub = abi.upcallStub(mh, func).baseAddress();
         return stub;
     }
 

test/jdk/java/foreign/TestUpcallStubs.java

@@ -39,6 +39,7 @@
 import java.lang.invoke.VarHandle;
 
 import static jdk.incubator.foreign.MemoryLayouts.JAVA_INT;
+import static org.testng.Assert.assertFalse;
 
 public class TestUpcallStubs {
 
@@ -54,40 +55,30 @@
         }
     }
 
-    private static MemoryAddress getStub() {
+    private static MemorySegment getStub() {
         return abi.upcallStub(MH_dummy, FunctionDescriptor.ofVoid());
     }
 
     @Test(expectedExceptions = UnsupportedOperationException.class)
     public void testNoAccess() {
-        MemoryAddress stub = getStub();
-        try {
+        try (MemorySegment stub = getStub()) {
             VarHandle vh = JAVA_INT.varHandle(int.class);
-            vh.set(stub, 10);
-        } finally {
-            abi.freeUpcallStub(stub);
+            vh.set(stub.baseAddress(), 10);
         }
     }
 
     @Test
     public void testFree() {
-        MemoryAddress stub = getStub();
-        abi.freeUpcallStub(stub);
+        MemorySegment stub = getStub();
+        stub.close();
+        assertFalse(stub.isAlive());
     }
 
-    @Test(expectedExceptions = IllegalArgumentException.class,
-          expectedExceptionsMessageRegExp = ".*Not a stub address: .*")
+    @Test(expectedExceptions = IllegalStateException.class)
     public void testAlreadyFreed() {
-        MemoryAddress stub = getStub();
-        abi.freeUpcallStub(stub);
-        abi.freeUpcallStub(stub); // should fail
-    }
-
-    @Test(expectedExceptions = IllegalArgumentException.class,
-          expectedExceptionsMessageRegExp = ".*Not a stub address: .*",
-          dataProvider = "badAddresses")
-    public void testCanNotFree(MemoryAddress ma) {
-        abi.freeUpcallStub(ma);
+        MemorySegment stub = getStub();
+        stub.close();
+        stub.close(); // should fail
     }
 
     @DataProvider