8275054: Racy updates when releasing confined scopes

Reviewed-by: jvernee

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

@@ -27,6 +27,8 @@
 
 import jdk.internal.vm.annotation.ForceInline;
 
+import java.lang.invoke.MethodHandles;
+import java.lang.invoke.VarHandle;
 import java.lang.ref.Cleaner;
 
 /**
@@ -39,8 +41,19 @@ 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;
+
+    static {
+        try {
+            ASYNC_RELEASE_COUNT = MethodHandles.lookup().findVarHandle(ConfinedScope.class, "asyncReleaseCount", int.class);
+        } catch (Throwable ex) {
+            throw new ExceptionInInitializerError(ex);
+        }
+    }
+
     public ConfinedScope(Thread owner, Cleaner cleaner) {
         super(new ConfinedResourceList(), cleaner);
         this.owner = owner;
@@ -74,12 +87,23 @@ public void acquire0() {
     @Override
     @ForceInline
     public void release0() {
-        lockCount--;
+        if (Thread.currentThread() == owner) {
+            lockCount--;
+        } 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,
+            // this scope might be kept alive by a shared scope, which means the release call can come from any
+            // thread.
+            int value;
+            do {
+                value = (int)ASYNC_RELEASE_COUNT.getVolatile(this);
+            } while (!ASYNC_RELEASE_COUNT.compareAndSet(this, value, value + 1));
+        }
     }
 
     void justClose() {
         this.checkValidState();
-        if (lockCount == 0) {
+        if (lockCount == 0 || lockCount - ((int)ASYNC_RELEASE_COUNT.getVolatile(this)) == 0) {
             closed = true;
         } else {
             throw new IllegalStateException("Scope is kept alive by " + lockCount + " scopes");

test/jdk/java/foreign/TestResourceScope.java

@@ -39,6 +39,7 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Set;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Supplier;
@@ -270,6 +271,63 @@ private void acquireRecursive(ResourceScope scope, int acquireCount) {
         }
     }
 
+    @Test
+    public void testConfinedScopeWithImplicitDependency() {
+        ResourceScope root = ResourceScope.newConfinedScope();
+        // Create many implicit scopes which depend on 'root', and let them become unreachable.
+        for (int i = 0; i < N_THREADS; i++) {
+            ResourceScope.newConfinedScope(Cleaner.create()).keepAlive(root);
+        }
+        // Now let's keep trying to close 'root' until we succeed. This is trickier than it seems: cleanup action
+        // might be called from another thread (the Cleaner thread), so that the confined scope lock count is updated racily.
+        // If that happens, the loop below never terminates.
+        while (true) {
+            try {
+                root.close();
+                break; // success!
+            } catch (IllegalStateException ex) {
+                kickGC();
+                for (int i = 0 ; i < N_THREADS ; i++) {  // add more races from current thread
+                    try (ResourceScope scope = ResourceScope.newConfinedScope()) {
+                        scope.keepAlive(root);
+                        // dummy
+                    }
+                }
+                // try again
+            }
+        }
+    }
+
+    @Test
+    public void testConfinedScopeWithSharedDependency() {
+        ResourceScope root = ResourceScope.newConfinedScope();
+        List<Thread> threads = new ArrayList<>();
+        // Create many implicit scopes which depend on 'root', and let them become unreachable.
+        for (int i = 0; i < N_THREADS; i++) {
+            ResourceScope scope = ResourceScope.newSharedScope(); // create scope inside same thread!
+            scope.keepAlive(root);
+            Thread t = new Thread(scope::close); // close from another thread!
+            threads.add(t);
+            t.start();
+        }
+        for (int i = 0 ; i < N_THREADS ; i++) { // add more races from current thread
+            try (ResourceScope scope = ResourceScope.newConfinedScope()) {
+                scope.keepAlive(root);
+                // dummy
+            }
+        }
+        threads.forEach(t -> {
+            try {
+                t.join();
+            } catch (InterruptedException ex) {
+                // ok
+            }
+        });
+        // Now let's close 'root'. This is trickier than it seems: releases of the confined scope happen in different
+        // threads, so that the confined scope lock count is updated racily. If that happens, the following close will blow up.
+        root.close();
+    }
+
     private void waitSomeTime() {
         try {
             Thread.sleep(10);