8318072: DowncallLinker does not acquire/release segments in interpreter

Reviewed-by: mcimadamore

src/java.base/share/classes/jdk/internal/foreign/abi/Binding.java

@@ -464,7 +464,7 @@ public sealed interface Binding {
         @Override
         public void interpret(Deque<Object> stack, StoreFunc storeFunc,
                               LoadFunc loadFunc, SegmentAllocator allocator) {
-            storeFunc.store(storage(), type(), stack.pop());
+            storeFunc.store(storage(), stack.pop());
         }
     }
 

src/java.base/share/classes/jdk/internal/foreign/abi/BindingInterpreter.java

@@ -50,7 +50,7 @@ public class BindingInterpreter {
 
     @FunctionalInterface
     public interface StoreFunc {
-        void store(VMStorage storage, Class<?> type, Object o);
+        void store(VMStorage storage, Object o);
     }
 
     @FunctionalInterface

src/java.base/share/classes/jdk/internal/foreign/abi/DowncallLinker.java

@@ -28,16 +28,22 @@ import jdk.internal.access.JavaLangInvokeAccess;
 import jdk.internal.access.SharedSecrets;
 import sun.security.action.GetPropertyAction;
 
+import java.lang.foreign.AddressLayout;
 import java.lang.foreign.Arena;
 import java.lang.foreign.MemorySegment;
 import java.lang.foreign.SegmentAllocator;
 import java.lang.invoke.MethodHandle;
 import java.lang.invoke.MethodHandles;
 import java.lang.invoke.MethodType;
+import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.List;
 import java.util.Map;
 import java.util.stream.Stream;
 
+import jdk.internal.foreign.AbstractMemorySegmentImpl;
+import jdk.internal.foreign.MemorySessionImpl;
+
 import static java.lang.invoke.MethodHandles.collectArguments;
 import static java.lang.invoke.MethodHandles.foldArguments;
 import static java.lang.invoke.MethodHandles.identity;
@@ -93,9 +99,8 @@ public class DowncallLinker {
             handle = BindingSpecializer.specializeDowncall(handle, callingSequence, abi);
          } else {
             Map<VMStorage, Integer> argIndexMap = SharedUtils.indexMap(argMoves);
-            Map<VMStorage, Integer> retIndexMap = SharedUtils.indexMap(retMoves);
 
-            InvocationData invData = new InvocationData(handle, argIndexMap, retIndexMap);
+            InvocationData invData = new InvocationData(handle, callingSequence, argIndexMap);
             handle = insertArguments(MH_INVOKE_INTERP_BINDINGS.bindTo(this), 2, invData);
             MethodType interpType = callingSequence.callerMethodType();
             if (callingSequence.needsReturnBuffer()) {
@@ -146,17 +151,17 @@ public class DowncallLinker {
         return Arrays.stream(moves).map(Binding.Move::storage).toArray(VMStorage[]::new);
     }
 
-    private record InvocationData(MethodHandle leaf, Map<VMStorage, Integer> argIndexMap, Map<VMStorage, Integer> retIndexMap) {}
+    private record InvocationData(MethodHandle leaf, CallingSequence callingSequence, Map<VMStorage, Integer> argIndexMap) {}
 
     Object invokeInterpBindings(SegmentAllocator allocator, Object[] args, InvocationData invData) throws Throwable {
         Arena unboxArena = callingSequence.allocationSize() != 0
                 ? SharedUtils.newBoundedArena(callingSequence.allocationSize())
                 : SharedUtils.DUMMY_ARENA;
+        List<MemorySessionImpl> acquiredScopes = new ArrayList<>();
         try (unboxArena) {
             MemorySegment returnBuffer = null;
 
             // do argument processing, get Object[] as result
-            Object[] leafArgs = new Object[invData.leaf.type().parameterCount()];
             if (callingSequence.needsReturnBuffer()) {
                 // we supply the return buffer (argument array does not contain it)
                 Object[] prefixedArgs = new Object[args.length + 1];
@@ -165,10 +170,21 @@ public class DowncallLinker {
                 System.arraycopy(args, 0, prefixedArgs, 1, args.length);
                 args = prefixedArgs;
             }
+
+            Object[] leafArgs = new Object[invData.leaf.type().parameterCount()];
             for (int i = 0; i < args.length; i++) {
                 Object arg = args[i];
+                if (callingSequence.functionDesc().argumentLayouts().get(i) instanceof AddressLayout) {
+                    MemorySessionImpl sessionImpl = ((AbstractMemorySegmentImpl) arg).sessionImpl();
+                    if (!(callingSequence.needsReturnBuffer() && i == 0)) { // don't acquire unboxArena's scope
+                        sessionImpl.acquire0();
+                        // add this scope _after_ we acquire, so we only release scopes we actually acquired
+                        // in case an exception occurs
+                        acquiredScopes.add(sessionImpl);
+                    }
+                }
                 BindingInterpreter.unbox(arg, callingSequence.argumentBindings(i),
-                        (storage, type, value) -> leafArgs[invData.argIndexMap.get(storage)] = value, unboxArena);
+                    (storage, value) -> leafArgs[invData.argIndexMap.get(storage)] = value, unboxArena);
             }
 
             // call leaf
@@ -194,6 +210,10 @@ public class DowncallLinker {
                 return BindingInterpreter.box(callingSequence.returnBindings(), (storage, type) -> o,
                         allocator);
             }
+        } finally {
+            for (MemorySessionImpl sessionImpl : acquiredScopes) {
+                sessionImpl.release0();
+            }
         }
     }
 }

src/java.base/share/classes/jdk/internal/foreign/abi/UpcallLinker.java

@@ -177,7 +177,7 @@ public class UpcallLinker {
             Object[] returnValues = new Object[invData.retIndexMap.size()];
             if (leaf.type().returnType() != void.class) {
                 BindingInterpreter.unbox(o, invData.callingSequence.returnBindings(),
-                        (storage, type, value) -> returnValues[invData.retIndexMap.get(storage)] = value, null);
+                        (storage, value) -> returnValues[invData.retIndexMap.get(storage)] = value, null);
             }
 
             if (returnValues.length == 0) {
@@ -187,15 +187,10 @@ public class UpcallLinker {
             } else {
                 assert invData.callingSequence.needsReturnBuffer();
 
-                Binding.VMStore[] retMoves = invData.callingSequence.returnBindings().stream()
+                assert returnValues.length == invData.retMoves().length;
-                        .filter(Binding.VMStore.class::isInstance)
-                        .map(Binding.VMStore.class::cast)
-                        .toArray(Binding.VMStore[]::new);
-
-                assert returnValues.length == retMoves.length;
                 int retBufWriteOffset = 0;
-                for (int i = 0; i < retMoves.length; i++) {
+                for (int i = 0; i < invData.retMoves().length; i++) {
-                    Binding.VMStore store = retMoves[i];
+                    Binding.VMStore store = invData.retMoves()[i];
                     Object value = returnValues[i];
                     SharedUtils.writeOverSized(returnBuffer.asSlice(retBufWriteOffset), store.type(), value);
                     retBufWriteOffset += invData.abi.arch.typeSize(store.storage().type());

test/jdk/java/foreign/LibraryLookupTest.java

@@ -34,8 +34,19 @@ import java.util.concurrent.TimeUnit;
 import static org.testng.Assert.*;
 
 /*
- * @test
+ * @test id=specialized
- * @run testng/othervm --enable-native-access=ALL-UNNAMED LibraryLookupTest
+ * @run testng/othervm
+ *  -Djdk.internal.foreign.DowncallLinker.USE_SPEC=true
+ *  --enable-native-access=ALL-UNNAMED
+ *  LibraryLookupTest
+ */
+
+/*
+ * @test id=interpreted
+ * @run testng/othervm
+ *   -Djdk.internal.foreign.DowncallLinker.USE_SPEC=false
+ *   --enable-native-access=ALL-UNNAMED
+ *   LibraryLookupTest
  */
 public class LibraryLookupTest {
 

test/jdk/java/foreign/SafeFunctionAccessTest.java

@@ -22,8 +22,19 @@
  */
 
 /*
- * @test
+ * @test id=specialized
- * @run testng/othervm --enable-native-access=ALL-UNNAMED SafeFunctionAccessTest
+ * @run testng/othervm
+ *  -Djdk.internal.foreign.DowncallLinker.USE_SPEC=true
+ *  --enable-native-access=ALL-UNNAMED
+ *  SafeFunctionAccessTest
+ */
+
+/*
+ * @test id=interpreted
+ * @run testng/othervm
+ *   -Djdk.internal.foreign.DowncallLinker.USE_SPEC=false
+ *   --enable-native-access=ALL-UNNAMED
+ *   SafeFunctionAccessTest
  */
 
 import java.lang.foreign.Arena;