8279527: Dereferencing segments backed by different scopes leads to pollution

Reviewed-by: psandoz, jvernee

Author: mcimadamore

src/java.base/share/classes/java/nio/Buffer.java

@@ -767,7 +767,11 @@ final ScopedMemoryAccess.Scope scope() {
     final void checkScope() {
         ScopedMemoryAccess.Scope scope = scope();
         if (scope != null) {
-            scope.checkValidState();
+            try {
+                scope.checkValidState();
+            } catch (ScopedMemoryAccess.Scope.ScopedAccessError e) {
+                throw new IllegalStateException("This segment is already closed");
+            }
         }
     }
 

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

@@ -364,11 +364,7 @@ public void checkAccess(long offset, long length, boolean readOnly) {
     }
 
     void checkValidState() {
-        try {
-            scope.checkValidState();
-        } catch (ScopedMemoryAccess.Scope.ScopedAccessError ex) {
-            throw new IllegalStateException("This segment is already closed");
-        }
+        scope.checkValidStateSlow();
     }
 
     @Override

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

@@ -39,10 +39,7 @@
  */
 final class ConfinedScope extends ResourceScopeImpl {
 
-    private boolean closed; // = false
-    private int lockCount = 0;
     private int asyncReleaseCount = 0;
-    private final Thread owner;
 
     static final VarHandle ASYNC_RELEASE_COUNT;
 
@@ -55,40 +52,29 @@ final class ConfinedScope extends ResourceScopeImpl {
     }
 
     public ConfinedScope(Thread owner, Cleaner cleaner) {
-        super(new ConfinedResourceList(), cleaner);
-        this.owner = owner;
-    }
-
-    @ForceInline
-    public final void checkValidState() {
-        if (owner != Thread.currentThread()) {
-            throw new IllegalStateException("Attempted access outside owning thread");
-        }
-        if (closed) {
-            throw new IllegalStateException("Already closed");
-        }
+        super(owner, new ConfinedResourceList(), cleaner);
     }
 
     @Override
     public boolean isAlive() {
-        return !closed;
+        return state != CLOSED;
     }
 
     @Override
     @ForceInline
     public void acquire0() {
-        checkValidState();
-        if (lockCount == MAX_FORKS) {
+        checkValidStateSlow();
+        if (state == MAX_FORKS) {
             throw new IllegalStateException("Scope keep alive limit exceeded");
         }
-        lockCount++;
+        state++;
     }
 
     @Override
     @ForceInline
     public void release0() {
         if (Thread.currentThread() == owner) {
-            lockCount--;
+            state--;
         } else {
             // It is possible to end up here in two cases: this scope was kept alive by some other confined scope
             // which is implicitly released (in which case the release call comes from the cleaner thread). Or,
@@ -99,19 +85,14 @@ public void release0() {
     }
 
     void justClose() {
-        this.checkValidState();
-        if (lockCount == 0 || lockCount - ((int)ASYNC_RELEASE_COUNT.getVolatile(this)) == 0) {
-            closed = true;
+        checkValidStateSlow();
+        if (state == 0 || state - ((int)ASYNC_RELEASE_COUNT.getVolatile(this)) == 0) {
+            state = CLOSED;
         } else {
-            throw new IllegalStateException("Scope is kept alive by " + lockCount + " scopes");
+            throw new IllegalStateException("Scope is kept alive by " + state + " scopes");
         }
     }
 
-    @Override
-    public Thread ownerThread() {
-        return owner;
-    }
-
     /**
      * A confined resource list; no races are possible here.
      */

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

@@ -28,11 +28,12 @@
 import jdk.incubator.foreign.MemoryAddress;
 import jdk.incubator.foreign.NativeSymbol;
 import jdk.incubator.foreign.ResourceScope;
+import jdk.internal.misc.ScopedMemoryAccess;
 
 public record NativeSymbolImpl(String name, MemoryAddress address, ResourceScope scope) implements NativeSymbol, Scoped {
     @Override
     public MemoryAddress address() {
-        ((ResourceScopeImpl)scope).checkValidState();
+        ((ResourceScopeImpl)scope).checkValidStateSlow();
         return address;
     }
 }

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

@@ -32,6 +32,8 @@
 import jdk.internal.misc.ScopedMemoryAccess;
 import jdk.internal.vm.annotation.ForceInline;
 
+import java.lang.invoke.MethodHandles;
+import java.lang.invoke.VarHandle;
 import java.lang.ref.Cleaner;
 import java.lang.ref.Reference;
 import java.util.Objects;
@@ -53,6 +55,23 @@ public abstract non-sealed class ResourceScopeImpl implements ResourceScope, Seg
 
     final ResourceList resourceList;
     final Cleaner.Cleanable cleanable;
+    final Thread owner;
+
+    static final int ALIVE = 0;
+    static final int CLOSING = -1;
+    static final int CLOSED = -2;
+
+    int state = ALIVE;
+
+    static final VarHandle STATE;
+
+    static {
+        try {
+            STATE = MethodHandles.lookup().findVarHandle(ResourceScopeImpl.class, "state", int.class);
+        } catch (Throwable ex) {
+            throw new ExceptionInInitializerError(ex);
+        }
+    }
 
     static final int MAX_FORKS = Integer.MAX_VALUE;
 
@@ -89,7 +108,8 @@ void addInternal(ResourceList.ResourceCleanup resource) {
         }
     }
 
-    protected ResourceScopeImpl(ResourceList resourceList, Cleaner cleaner) {
+    protected ResourceScopeImpl(Thread owner, ResourceList resourceList, Cleaner cleaner) {
+        this.owner = owner;
         this.resourceList = resourceList;
         cleanable = (cleaner != null) ?
             cleaner.register(this, resourceList) : null;
@@ -147,30 +167,41 @@ public void close() {
      * Returns "owner" thread of this scope.
      * @return owner thread (or null for a shared scope)
      */
-    public abstract Thread ownerThread();
+    public final Thread ownerThread() {
+        return owner;
+    }
 
     /**
      * Returns true, if this scope is still alive. This method may be called in any thread.
      * @return {@code true} if this scope is not closed yet.
      */
     public abstract boolean isAlive();
 
-
     /**
      * This is a faster version of {@link #checkValidStateSlow()}, which is called upon memory access, and which
-     * relies on invariants associated with the memory scope implementations (typically, volatile access
-     * to the closed state bit is replaced with plain access, and ownership check is removed where not needed.
-     * Should be used with care.
+     * relies on invariants associated with the memory scope implementations (volatile access
+     * to the closed state bit is replaced with plain access). This method should be monomorphic,
+     * to avoid virtual calls in the memory access hot path. This method is not intended as general purpose method
+     * and should only be used in the memory access handle hot path; for liveness checks triggered by other API methods,
+     * please use {@link #checkValidStateSlow()}.
      */
-    public abstract void checkValidState();
+    @ForceInline
+    public final void checkValidState() {
+        if (owner != null && owner != Thread.currentThread()) {
+            throw new IllegalStateException("Attempted access outside owning thread");
+        }
+        if (state < ALIVE) {
+            throw ScopedAccessError.INSTANCE;
+        }
+    }
 
     /**
      * Checks that this scope is still alive (see {@link #isAlive()}).
      * @throws IllegalStateException if this scope is already closed or if this is
      * a confined scope and this method is called outside of the owner thread.
      */
     public final void checkValidStateSlow() {
-        if (ownerThread() != null && Thread.currentThread() != ownerThread()) {
+        if (owner != null && Thread.currentThread() != owner) {
             throw new IllegalStateException("Attempted access outside owning thread");
         } else if (!isAlive()) {
             throw new IllegalStateException("Already closed");

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

@@ -45,36 +45,8 @@ class SharedScope extends ResourceScopeImpl {
 
     private static final ScopedMemoryAccess SCOPED_MEMORY_ACCESS = ScopedMemoryAccess.getScopedMemoryAccess();
 
-    private static final int ALIVE = 0;
-    private static final int CLOSING = -1;
-    private static final int CLOSED = -2;
-
-    private int state = ALIVE;
-
-    private static final VarHandle STATE;
-
-    static {
-        try {
-            STATE = MethodHandles.lookup().findVarHandle(jdk.internal.foreign.SharedScope.class, "state", int.class);
-        } catch (Throwable ex) {
-            throw new ExceptionInInitializerError(ex);
-        }
-    }
-
     SharedScope(Cleaner cleaner) {
-        super(new SharedResourceList(), cleaner);
-    }
-
-    @Override
-    public Thread ownerThread() {
-        return null;
-    }
-
-    @Override
-    public void checkValidState() {
-        if (state < ALIVE) {
-            throw ScopedAccessError.INSTANCE;
-        }
+        super(null, new SharedResourceList(), cleaner);
     }
 
     @Override

test/jdk/java/foreign/TestByteBuffer.java

@@ -369,7 +369,7 @@ public void testScopedBuffer(Function<ByteBuffer, Buffer> bufferFactory, @NoInje
             Throwable cause = ex.getCause();
             if (cause instanceof IllegalStateException) {
                 //all get/set buffer operation should fail because of the scope check
-                assertTrue(ex.getCause().getMessage().contains("Already closed"));
+                assertTrue(ex.getCause().getMessage().contains("already closed"));
             } else {
                 //all other exceptions were unexpected - fail
                 fail("Unexpected exception", cause);
@@ -406,7 +406,7 @@ public void testScopedBufferAndVarHandle(VarHandle bufferHandle) {
                 handle.invoke(e.getValue());
                 fail();
             } catch (IllegalStateException ex) {
-                assertTrue(ex.getMessage().contains("Already closed"));
+                assertTrue(ex.getMessage().contains("already closed"));
             } catch (UnsupportedOperationException ex) {
                 //skip
             } catch (Throwable ex) {

test/micro/org/openjdk/bench/jdk/incubator/foreign/LoopOverPollutedSegments.java

@@ -58,7 +58,7 @@ public class LoopOverPollutedSegments {
 
     static final Unsafe unsafe = Utils.unsafe;
 
-    MemorySegment nativeSegment, heapSegmentBytes, heapSegmentFloats;
+    MemorySegment nativeSegment, nativeSharedSegment, heapSegmentBytes, heapSegmentFloats;
     byte[] arr;
     long addr;
 
@@ -73,6 +73,7 @@ public void setup() {
         }
         arr = new byte[ALLOC_SIZE];
         nativeSegment = MemorySegment.allocateNative(ALLOC_SIZE, 4, ResourceScope.newConfinedScope());
+        nativeSharedSegment = MemorySegment.allocateNative(ALLOC_SIZE, 4, ResourceScope.newSharedScope());
         heapSegmentBytes = MemorySegment.ofArray(new byte[ALLOC_SIZE]);
         heapSegmentFloats = MemorySegment.ofArray(new float[ELEM_SIZE]);
 
@@ -81,6 +82,8 @@ public void setup() {
                 unsafe.putInt(arr, Unsafe.ARRAY_BYTE_BASE_OFFSET + (i * 4), i);
                 nativeSegment.setAtIndex(JAVA_INT, i, i);
                 nativeSegment.setAtIndex(JAVA_FLOAT, i, i);
+                nativeSharedSegment.setAtIndex(JAVA_INT, i, i);
+                nativeSharedSegment.setAtIndex(JAVA_FLOAT, i, i);
                 intHandle.set(nativeSegment, (long)i, i);
                 heapSegmentBytes.setAtIndex(JAVA_INT, i, i);
                 heapSegmentBytes.setAtIndex(JAVA_FLOAT, i, i);