JDK-8243669: Improve library loading for Panama libraries

src/java.base/share/classes/module-info.java

@@ -217,7 +217,8 @@
     exports jdk.internal.platform to
         jdk.management;
     exports jdk.internal.ref to
-        java.desktop;
+        java.desktop,
+        jdk.incubator.foreign;
     exports jdk.internal.reflect to
         java.logging,
         java.sql,

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

@@ -28,16 +28,29 @@
 import jdk.internal.foreign.LibrariesHelper;
 
 import java.io.File;
-import java.lang.invoke.MethodHandles;
+import java.lang.invoke.MethodType;
 import java.util.Objects;
 
 /**
  * A native library lookup. Exposes lookup operation for searching symbols, see {@link LibraryLookup#lookup(String)}.
+ * A given native library remains loaded as long as there is at least one <em>live</em> library lookup instance referring
+ * to it.
+ * <p>
+ * Memory address instances generated by a library lookup will contain a strong reference to the originating lookup object,
+ * therefore preventing library unloading; in turn method handle instances obtained from
+ * {@link SystemABI#downcallHandle(MemoryAddress, MethodType, FunctionDescriptor)}) also maintain a strong reference
+ * to the memory address parameter used for their construction. This means that there is always a strong reachability chain
+ * from a native method handle to a lookup object (the one that was used to lookup the native library symbol the method handle
+ * refers to); this is useful to prevent situations where a native library is unloaded in the middle of a native call.
+ * <p>
+ * To allow for a library to be unloaded, a client will have to discard any strong references it
+ * maintains, directly, or indirectly to a lookup object associated with given library.
  */
 public interface LibraryLookup {
 
     /**
-     * Lookups a symbol with given name in this library.
+     * Lookups a symbol with given name in this library. The returned address has the same temporal bounds
+     * as this lookup object.
      * @param name the symbol name.
      * @return the library symbol (if any)
      * @throws NoSuchMethodException if no symbol with given name could be found.
@@ -58,11 +71,10 @@ static LibraryLookup ofDefault() {
 
     /**
      * Obtain a library lookup object corresponding to a library identified by given path.
-     * @param lookup the contextual lookup object.
      * @param path the library path.
      * @return a library lookup object for given path.
      */
-    static LibraryLookup ofPath(MethodHandles.Lookup lookup, String path) {
+    static LibraryLookup ofPath(String path) {
         Objects.requireNonNull(path);
         SecurityManager security = System.getSecurityManager();
         if (security != null) {
@@ -72,16 +84,15 @@ static LibraryLookup ofPath(MethodHandles.Lookup lookup, String path) {
             throw new UnsatisfiedLinkError(
                     "Expecting an absolute path of the library: " + path);
         }
-        return LibrariesHelper.load(lookup, path);
+        return LibrariesHelper.load(path);
     }
 
     /**
      * Obtain a library lookup object corresponding to a library identified by given library name.
-     * @param lookup the contextual lookup object.
      * @param libName the library name.
      * @return a library lookup object for given library name.
      */
-    static LibraryLookup ofLibrary(MethodHandles.Lookup lookup, String libName) {
+    static LibraryLookup ofLibrary(String libName) {
         Objects.requireNonNull(libName);
         SecurityManager security = System.getSecurityManager();
         if (security != null) {
@@ -91,6 +102,6 @@ static LibraryLookup ofLibrary(MethodHandles.Lookup lookup, String libName) {
             throw new UnsatisfiedLinkError(
                     "Directory separator should not appear in library name: " + libName);
         }
-        return LibrariesHelper.loadLibrary(lookup, libName);
+        return LibrariesHelper.loadLibrary(libName);
     }
 }

src/jdk.incubator.foreign/share/classes/jdk/internal/foreign/LibrariesHelper.java

@@ -28,53 +28,51 @@
 import jdk.incubator.foreign.MemoryAddress;
 
 import java.io.File;
-import java.lang.invoke.MethodHandles.Lookup;
 import jdk.incubator.foreign.LibraryLookup;
 import jdk.internal.loader.NativeLibraries;
 import jdk.internal.loader.NativeLibrary;
+import jdk.internal.ref.CleanerFactory;
 
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.Arrays;
+import java.util.IdentityHashMap;
+import java.util.Map;
 import java.util.Optional;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Supplier;
 
 public final class LibrariesHelper {
     private LibrariesHelper() {}
 
-    // FIXME - revisit this (refcount for unload)
-    private final static ClassValue<NativeLibraries> nativeLibrary = new ClassValue<>() {
-        @Override
-        protected NativeLibraries computeValue(Class<?> type) {
-            return NativeLibraries.jniNativeLibraries(type.getClassLoader());
-        }
-    };
+    private final static NativeLibraries nativeLibraries =
+            NativeLibraries.rawNativeLibraries(LibrariesHelper.class, true);
 
+    private final static Map<NativeLibrary, AtomicInteger> loadedLibraries = new IdentityHashMap<>();
 
     /**
      * Load the specified shared library.
      *
-     * @param lookup Lookup object of the caller.
      * @param name Name of the shared library to load.
      */
-    public static LibraryLookup loadLibrary(Lookup lookup, String name) {
-        return new LibraryLookupImpl(nativeLibrary.get(lookup.lookupClass())
-                .loadLibrary(lookup.lookupClass(), name));
+    public static LibraryLookup loadLibrary(String name) {
+        return lookup(() -> nativeLibraries.loadLibrary(LibrariesHelper.class, name),
+                "Library not found: " + name);
     }
 
     /**
      * Load the specified shared library.
      *
-     * @param lookup Lookup object of the caller.
      * @param path Path of the shared library to load.
      */
-    public static LibraryLookup load(Lookup lookup, String path) {
+    public static LibraryLookup load(String path) {
         File file = new File(path);
         if (!file.isAbsolute()) {
             throw new UnsatisfiedLinkError(
                     "Expecting an absolute path of the library: " + path);
         }
-        return new LibraryLookupImpl(nativeLibrary.get(lookup.lookupClass())
-                .loadLibrary(lookup.lookupClass(), file));
+        return lookup(() -> nativeLibraries.loadLibrary(LibrariesHelper.class, file),
+                "Library not found: " + path);
     }
 
     // return the absolute path of the library of given name by searching
@@ -86,7 +84,27 @@ private static Optional<Path> findLibraryPath(Path[] paths, String libName) {
     }
 
     public static LibraryLookup getDefaultLibrary() {
-        return new LibraryLookupImpl(NativeLibraries.defaultLibrary);
+        return LibraryLookupImpl.DEFAULT_LOOKUP;
+    }
+
+    synchronized static LibraryLookupImpl lookup(Supplier<NativeLibrary> librarySupplier, String notFoundMsg) {
+        NativeLibrary library = librarySupplier.get();
+        if (library == null) {
+            throw new IllegalArgumentException(notFoundMsg);
+        }
+        AtomicInteger refCount = loadedLibraries.computeIfAbsent(library, lib -> new AtomicInteger());
+        refCount.incrementAndGet();
+        LibraryLookupImpl lookup = new LibraryLookupImpl(library);
+        CleanerFactory.cleaner().register(lookup, () -> tryUnload(library));
+        return lookup;
+    }
+
+    synchronized static void tryUnload(NativeLibrary library) {
+        AtomicInteger refCount = loadedLibraries.get(library);
+        if (refCount.decrementAndGet() == 0) {
+            loadedLibraries.remove(library);
+            nativeLibraries.unload(library);
+        }
     }
 
     static class LibraryLookupImpl implements LibraryLookup {
@@ -99,7 +117,16 @@ static class LibraryLookupImpl implements LibraryLookup {
         @Override
         public MemoryAddress lookup(String name) throws NoSuchMethodException {
             long addr = library.lookup(name);
-            return MemoryAddress.ofLong(addr);
+            return NativeMemorySegmentImpl.makeNativeSegmentUnchecked(MemoryAddress.ofLong(addr),
+                    0, null, null, this)
+                    .baseAddress();
         }
+
+        static LibraryLookup DEFAULT_LOOKUP = new LibraryLookupImpl(NativeLibraries.defaultLibrary);
+    }
+
+    /* used for testing */
+    public static int numLoadedLibraries() {
+        return loadedLibraries.size();
     }
 }

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

@@ -72,10 +72,10 @@ public class ProgrammableInvoker {
     private final FunctionDescriptor function;
     private final CallingSequence callingSequence;
 
-    private final long addr;
+    private final MemoryAddress addr;
     private final long stubAddress;
 
-    public ProgrammableInvoker(ABIDescriptor abi, long addr, CallingSequence callingSequence) {
+    public ProgrammableInvoker(ABIDescriptor abi, MemoryAddress addr, CallingSequence callingSequence) {
         this.abi = abi;
         this.layout = BufferLayout.of(abi);
         this.stubAddress = adapterStubs.computeIfAbsent(abi, key -> generateAdapter(key, layout));
@@ -109,7 +109,7 @@ Object invoke(Object[] args) {
                 stackArgs = MemoryAddressImpl.NULL;
             }
 
-            VH_LONG.set(argsPtr.addOffset(layout.arguments_next_pc), addr);
+            VH_LONG.set(argsPtr.addOffset(layout.arguments_next_pc), addr.toRawLongValue());
             VH_LONG.set(argsPtr.addOffset(layout.stack_args_bytes), stackArgsBytes);
             VH_LONG.set(argsPtr.addOffset(layout.stack_args), stackArgs.toRawLongValue());
 

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

@@ -55,7 +55,7 @@ public static AArch64ABI getInstance() {
 
     @Override
     public MethodHandle downcallHandle(MemoryAddress symbol, MethodType type, FunctionDescriptor function) {
-        return CallArranger.arrangeDowncall(symbol.toRawLongValue(), type, function);
+        return CallArranger.arrangeDowncall(symbol, type, function);
     }
 
     @Override

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

@@ -124,7 +124,7 @@ public static Bindings getBindings(MethodType mt, FunctionDescriptor cDesc, bool
         return new Bindings(csb.build(), returnInMemory);
     }
 
-    public static MethodHandle arrangeDowncall(long addr, MethodType mt, FunctionDescriptor cDesc) {
+    public static MethodHandle arrangeDowncall(MemoryAddress addr, MethodType mt, FunctionDescriptor cDesc) {
         Bindings bindings = getBindings(mt, cDesc, false);
 
         MethodHandle handle = new ProgrammableInvoker(C, addr, bindings.callingSequence).getBoundMethodHandle();

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

@@ -127,7 +127,7 @@ public static Bindings getBindings(MethodType mt, FunctionDescriptor cDesc, bool
         return new Bindings(csb.build(), returnInMemory, argCalc.storageCalculator.nVectorReg);
     }
 
-    public static MethodHandle arrangeDowncall(long addr, MethodType mt, FunctionDescriptor cDesc) {
+    public static MethodHandle arrangeDowncall(MemoryAddress addr, MethodType mt, FunctionDescriptor cDesc) {
         Bindings bindings = getBindings(mt, cDesc, false);
 
         MethodHandle handle = new ProgrammableInvoker(CSysV, addr, bindings.callingSequence).getBoundMethodHandle();

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

@@ -59,7 +59,7 @@ public static SysVx64ABI getInstance() {
 
     @Override
     public MethodHandle downcallHandle(MemoryAddress symbol, MethodType type, FunctionDescriptor function) {
-        return CallArranger.arrangeDowncall(symbol.toRawLongValue(), type, function);
+        return CallArranger.arrangeDowncall(symbol, type, function);
     }
 
     @Override

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

@@ -124,7 +124,7 @@ void setReturnBindings(Class<?> carrier, MemoryLayout layout) {
         return new Bindings(csb.csb.build(), returnInMemory);
     }
 
-    public static MethodHandle arrangeDowncall(long addr, MethodType mt, FunctionDescriptor cDesc) {
+    public static MethodHandle arrangeDowncall(MemoryAddress addr, MethodType mt, FunctionDescriptor cDesc) {
         Bindings bindings = getBindings(mt, cDesc, false);
 
         MethodHandle handle = new ProgrammableInvoker(CWindows, addr, bindings.callingSequence).getBoundMethodHandle();

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

@@ -64,7 +64,7 @@ public static Windowsx64ABI getInstance() {
 
     @Override
     public MethodHandle downcallHandle(MemoryAddress symbol, MethodType type, FunctionDescriptor function) {
-        return CallArranger.arrangeDowncall(symbol.toRawLongValue(), type, function);
+        return CallArranger.arrangeDowncall(symbol, type, function);
     }
 
     @Override

test/jdk/java/foreign/TestDowncall.java

@@ -51,7 +51,7 @@
 
 public class TestDowncall extends CallGeneratorHelper {
 
-    static LibraryLookup lib = LibraryLookup.ofLibrary(MethodHandles.lookup(), "TestDowncall");
+    static LibraryLookup lib = LibraryLookup.ofLibrary("TestDowncall");
     static SystemABI abi = SystemABI.getSystemABI();