8343394: Make MemorySessionImpl.state a stable field

Co-authored-by: Maurizio Cimadamore <[email protected]>
Reviewed-by: mcimadamore, jvernee

src/java.base/share/classes/jdk/internal/foreign/ConfinedSession.java

@@ -41,8 +41,7 @@ final class ConfinedSession extends MemorySessionImpl {
 
     private int asyncReleaseCount = 0;
 
-    static final VarHandle ASYNC_RELEASE_COUNT= MhUtil.findVarHandle(
-            MethodHandles.lookup(), "asyncReleaseCount", int.class);
+    static final VarHandle ASYNC_RELEASE_COUNT= MhUtil.findVarHandle(MethodHandles.lookup(), "asyncReleaseCount", int.class);
 
     public ConfinedSession(Thread owner) {
         super(owner, new ConfinedResourceList());
@@ -52,17 +51,17 @@ public ConfinedSession(Thread owner) {
     @ForceInline
     public void acquire0() {
         checkValidState();
-        if (state == MAX_FORKS) {
+        if (acquireCount == MAX_FORKS) {
             throw tooManyAcquires();
         }
-        state++;
+        acquireCount++;
     }
 
     @Override
     @ForceInline
     public void release0() {
         if (Thread.currentThread() == owner) {
-            state--;
+            acquireCount--;
         } else {
             // It is possible to end up here in two cases: this session was kept alive by some other confined session
             // which is implicitly released (in which case the release call comes from the cleaner thread). Or,
@@ -75,11 +74,11 @@ public void release0() {
     void justClose() {
         checkValidState();
         int asyncCount = (int)ASYNC_RELEASE_COUNT.getVolatile(this);
-        if ((state == 0 && asyncCount == 0)
-                || ((state - asyncCount) == 0)) {
+        int acquire = acquireCount - asyncCount;
+        if (acquire == 0) {
             state = CLOSED;
         } else {
-            throw alreadyAcquired(state - asyncCount);
+            throw alreadyAcquired(acquire);
         }
     }
 

src/java.base/share/classes/jdk/internal/foreign/GlobalSession.java

@@ -42,6 +42,7 @@ non-sealed class GlobalSession extends MemorySessionImpl {
 
     public GlobalSession() {
         super(null, null);
+        this.state = NONCLOSEABLE;
     }
 
     @Override
@@ -50,11 +51,6 @@ public void release0() {
         // do nothing
     }
 
-    @Override
-    public boolean isCloseable() {
-        return false;
-    }
-
     @Override
     @ForceInline
     public void acquire0() {

src/java.base/share/classes/jdk/internal/foreign/ImplicitSession.java

@@ -42,6 +42,7 @@ final class ImplicitSession extends SharedSession {
 
     public ImplicitSession(Cleaner cleaner) {
         super();
+        this.state = NONCLOSEABLE;
         cleaner.register(this, resourceList);
     }
 
@@ -55,11 +56,6 @@ public void acquire0() {
         // do nothing
     }
 
-    @Override
-    public boolean isCloseable() {
-        return false;
-    }
-
     @Override
     public void justClose() {
         throw nonCloseable();

src/java.base/share/classes/jdk/internal/foreign/MemorySessionImpl.java

@@ -38,6 +38,7 @@
 import jdk.internal.misc.ScopedMemoryAccess;
 import jdk.internal.invoke.MhUtil;
 import jdk.internal.vm.annotation.ForceInline;
+import jdk.internal.vm.annotation.Stable;
 
 /**
  * This class manages the temporal bounds associated with a memory segment as well
@@ -55,11 +56,19 @@
 public abstract sealed class MemorySessionImpl
         implements Scope
         permits ConfinedSession, GlobalSession, SharedSession {
+
+    /**
+     * The value of the {@code state} of a {@code MemorySessionImpl}. The only possible transition
+     * is OPEN -> CLOSED. As a result, the states CLOSED and NONCLOSEABLE are stable. This allows
+     * us to annotate {@code state} with {@link Stable} and elide liveness check on non-closeable
+     * constant scopes, such as {@code GLOBAL_SESSION}.
+     */
     static final int OPEN = 0;
     static final int CLOSED = -1;
+    static final int NONCLOSEABLE = 1;
 
-    static final VarHandle STATE = MhUtil.findVarHandle(
-            MethodHandles.lookup(), "state", int.class);
+    static final VarHandle STATE = MhUtil.findVarHandle(MethodHandles.lookup(), "state", int.class);
+    static final VarHandle ACQUIRE_COUNT = MhUtil.findVarHandle(MethodHandles.lookup(), "acquireCount", int.class);
 
     static final int MAX_FORKS = Integer.MAX_VALUE;
 
@@ -70,7 +79,11 @@ public abstract sealed class MemorySessionImpl
 
     final ResourceList resourceList;
     final Thread owner;
-    int state = OPEN;
+
+    @Stable
+    int state;
+
+    int acquireCount;
 
     public Arena asArena() {
         return new ArenaImpl(this);
@@ -214,8 +227,8 @@ protected Object clone() throws CloneNotSupportedException {
         throw new CloneNotSupportedException();
     }
 
-    public boolean isCloseable() {
-        return true;
+    public final boolean isCloseable() {
+        return state <= OPEN;
     }
 
     /**

src/java.base/share/classes/jdk/internal/foreign/SharedSession.java

@@ -44,6 +44,8 @@ sealed class SharedSession extends MemorySessionImpl permits ImplicitSession {
 
     private static final ScopedMemoryAccess SCOPED_MEMORY_ACCESS = ScopedMemoryAccess.getScopedMemoryAccess();
 
+    private static final int CLOSED_ACQUIRE_COUNT = -1;
+
     SharedSession() {
         super(null, new SharedResourceList());
     }
@@ -53,40 +55,51 @@ sealed class SharedSession extends MemorySessionImpl permits ImplicitSession {
     public void acquire0() {
         int value;
         do {
-            value = (int) STATE.getVolatile(this);
-            if (value < OPEN) {
+            value = (int) ACQUIRE_COUNT.getVolatile(this);
+            if (value < 0) {
                 //segment is not open!
-                throw alreadyClosed();
+                throw sharedSessionAlreadyClosed();
             } else if (value == MAX_FORKS) {
                 //overflow
                 throw tooManyAcquires();
             }
-        } while (!STATE.compareAndSet(this, value, value + 1));
+        } while (!ACQUIRE_COUNT.compareAndSet(this, value, value + 1));
     }
 
     @Override
     @ForceInline
     public void release0() {
         int value;
         do {
-            value = (int) STATE.getVolatile(this);
-            if (value <= OPEN) {
+            value = (int) ACQUIRE_COUNT.getVolatile(this);
+            if (value <= 0) {
                 //cannot get here - we can't close segment twice
-                throw alreadyClosed();
+                throw sharedSessionAlreadyClosed();
             }
-        } while (!STATE.compareAndSet(this, value, value - 1));
+        } while (!ACQUIRE_COUNT.compareAndSet(this, value, value - 1));
     }
 
     void justClose() {
-        int prevState = (int) STATE.compareAndExchange(this, OPEN, CLOSED);
-        if (prevState < 0) {
-            throw alreadyClosed();
-        } else if (prevState != OPEN) {
-            throw alreadyAcquired(prevState);
+        int acquireCount = (int) ACQUIRE_COUNT.compareAndExchange(this, 0, CLOSED_ACQUIRE_COUNT);
+        if (acquireCount < 0) {
+            throw sharedSessionAlreadyClosed();
+        } else if (acquireCount > 0) {
+            throw alreadyAcquired(acquireCount);
         }
+
+        STATE.setVolatile(this, CLOSED);
         SCOPED_MEMORY_ACCESS.closeScope(this, ALREADY_CLOSED);
     }
 
+    private IllegalStateException sharedSessionAlreadyClosed() {
+        // To avoid the situation where a scope fails to be acquired or closed but still reports as
+        // alive afterward, we wait for the state to change before throwing the exception
+        while ((int) STATE.getVolatile(this) == OPEN) {
+            Thread.onSpinWait();
+        }
+        return alreadyClosed();
+    }
+
     /**
      * A shared resource list; this implementation has to handle add vs. add races, as well as add vs. cleanup races.
      */

test/jdk/java/foreign/TestMemorySession.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021, 2023, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2021, 2024, Oracle and/or its affiliates. All rights reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -28,19 +28,18 @@
  */
 
 import java.lang.foreign.Arena;
-
-import jdk.internal.foreign.MemorySessionImpl;
-import org.testng.annotations.DataProvider;
-import org.testng.annotations.Test;
-import static org.testng.Assert.*;
-
 import java.lang.reflect.Method;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Supplier;
 import java.util.stream.IntStream;
+import jdk.internal.foreign.MemorySessionImpl;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.*;
 
 public class TestMemorySession {
 
@@ -319,6 +318,70 @@ public void testIsCloseableBy(ArenaSupplier arenaSupplier) {
         assertEquals(sessionImpl.isCloseableBy(otherThread), isCloseableByOther);
     }
 
+    /**
+     * Test that a thread failing to acquire a scope will not observe it as alive afterwards.
+     */
+    @Test
+    public void testAcquireCloseRace() throws InterruptedException {
+        int iteration = 1000;
+        AtomicInteger lock = new AtomicInteger();
+        boolean[] result = new boolean[1];
+        lock.set(-2);
+        MemorySessionImpl[] scopes = new MemorySessionImpl[iteration];
+        for (int i = 0; i < iteration; i++) {
+            scopes[i] = MemorySessionImpl.toMemorySession(Arena.ofShared());
+        }
+
+        // This thread tries to close the scopes
+        Thread t1 = new Thread(() -> {
+            for (int i = 0; i < iteration; i++) {
+                MemorySessionImpl scope = scopes[i];
+                while (true) {
+                    try {
+                        scope.close();
+                        break;
+                    } catch (IllegalStateException e) {}
+                }
+                // Keep the 2 threads operating on the same scope
+                int k = lock.getAndAdd(1) + 1;
+                while (k != i * 2) {
+                    Thread.onSpinWait();
+                    k = lock.get();
+                }
+            }
+        });
+
+        // This thread tries to acquire the scopes, then check if it is alive after an acquire failure
+        Thread t2 = new Thread(() -> {
+            for (int i = 0; i < iteration; i++) {
+                MemorySessionImpl scope = scopes[i];
+                while (true) {
+                    try {
+                        scope.acquire0();
+                    } catch (IllegalStateException e) {
+                        if (scope.isAlive()) {
+                            result[0] = true;
+                        }
+                        break;
+                    }
+                    scope.release0();
+                }
+                // Keep the 2 threads operating on the same scope
+                int k = lock.getAndAdd(1) + 1;
+                while (k != i * 2) {
+                    Thread.onSpinWait();
+                    k = lock.get();
+                }
+            }
+        });
+
+        t1.start();
+        t2.start();
+        t1.join();
+        t2.join();
+        assertFalse(result[0]);
+    }
+
     private void waitSomeTime() {
         try {
             Thread.sleep(10);