8246050: Improve scalability of MemoryScope

Reiplement memory scope using StampedLock

Reviewed-by: psandoz

Author: plevart

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

@@ -26,6 +26,8 @@
 package jdk.internal.foreign;
 
 import jdk.incubator.foreign.MemoryAddress;
+import jdk.incubator.foreign.MemoryLayout;
+import jdk.incubator.foreign.MemoryLayouts;
 import jdk.incubator.foreign.MemorySegment;
 import jdk.incubator.foreign.SequenceLayout;
 import jdk.internal.access.JavaNioAccess;
@@ -68,22 +70,20 @@ public abstract class AbstractMemorySegmentImpl implements MemorySegment, Memory
 
     final long length;
     final int mask;
-    final Thread owner;
     final MemoryScope scope;
 
     @ForceInline
-    AbstractMemorySegmentImpl(long length, int mask, Thread owner, MemoryScope scope) {
+    AbstractMemorySegmentImpl(long length, int mask, MemoryScope scope) {
         this.length = length;
         this.mask = mask;
-        this.owner = owner;
         this.scope = scope;
     }
 
     abstract long min();
 
     abstract Object base();
 
-    abstract AbstractMemorySegmentImpl dup(long offset, long size, int mask, Thread owner, MemoryScope scope);
+    abstract AbstractMemorySegmentImpl dup(long offset, long size, int mask, MemoryScope scope);
 
     abstract ByteBuffer makeByteBuffer();
 
@@ -100,7 +100,7 @@ public AbstractMemorySegmentImpl asSlice(long offset, long newSize) {
     }
 
     private AbstractMemorySegmentImpl asSliceNoCheck(long offset, long newSize) {
-        return dup(offset, newSize, mask, owner, scope);
+        return dup(offset, newSize, mask, scope);
     }
 
     @SuppressWarnings("unchecked")
@@ -145,12 +145,12 @@ public final long byteSize() {
 
     @Override
     public final boolean isAlive() {
-        return scope.isAliveThreadSafe();
+        return scope.isAlive();
     }
 
     @Override
     public Thread ownerThread() {
-        return owner;
+        return scope.ownerThread();
     }
 
     @Override
@@ -159,7 +159,7 @@ public AbstractMemorySegmentImpl withAccessModes(int accessModes) {
         if ((~accessModes() & accessModes) != 0) {
             throw new IllegalArgumentException("Cannot acquire more access modes");
         }
-        return dup(0, length, (mask & ~ACCESS_MASK) | accessModes, owner, scope);
+        return dup(0, length, (mask & ~ACCESS_MASK) | accessModes, scope);
     }
 
     @Override
@@ -177,17 +177,16 @@ private void checkAccessModes(int accessModes) {
     @Override
     public MemorySegment withOwnerThread(Thread newOwner) {
         Objects.requireNonNull(newOwner);
-        checkValidState();
         if (!isSet(HANDOFF)) {
             throw unsupportedAccessMode(HANDOFF);
         }
-        if (owner == newOwner) {
+        if (scope.ownerThread() == newOwner) {
             throw new IllegalArgumentException("Segment already owned by thread: " + newOwner);
         } else {
             try {
-                return dup(0L, length, mask, newOwner, scope.dup());
+                return dup(0L, length, mask, scope.dup(newOwner));
             } finally {
-                //flush read/writes to memory before returning the new segment
+                //flush read/writes to segment memory before returning the new segment
                 VarHandle.fullFence();
             }
         }
@@ -198,19 +197,18 @@ public final void close() {
         if (!isSet(CLOSE)) {
             throw unsupportedAccessMode(CLOSE);
         }
-        checkValidState();
         closeNoCheck();
     }
 
     private final void closeNoCheck() {
-        scope.close(true);
+        scope.close();
     }
 
     final AbstractMemorySegmentImpl acquire() {
         if (Thread.currentThread() != ownerThread() && !isSet(ACQUIRE)) {
             throw unsupportedAccessMode(ACQUIRE);
         }
-        return dup(0, length, mask, Thread.currentThread(), scope.acquire());
+        return dup(0, length, mask, scope.acquire());
     }
 
     @Override
@@ -227,7 +225,7 @@ boolean isSmall() {
     }
 
     void checkRange(long offset, long length, boolean writeAccess) {
-        checkValidState();
+        scope.checkValidState();
         if (writeAccess && !isSet(WRITE)) {
             throw unsupportedAccessMode(WRITE);
         } else if (!writeAccess && !isSet(READ)) {
@@ -238,10 +236,7 @@ void checkRange(long offset, long length, boolean writeAccess) {
 
     @Override
     public final void checkValidState() {
-        if (owner != null && owner != Thread.currentThread()) {
-            throw new IllegalStateException("Attempt to access segment outside owning thread");
-        }
-        scope.checkAliveConfined();
+        scope.checkValidState();
     }
 
     // Helper methods
@@ -415,29 +410,28 @@ public static AbstractMemorySegmentImpl ofBuffer(ByteBuffer bb) {
         AbstractMemorySegmentImpl bufferSegment = (AbstractMemorySegmentImpl)nioAccess.bufferSegment(bb);
         final MemoryScope bufferScope;
         int modes;
-        final Thread owner;
         if (bufferSegment != null) {
             bufferScope = bufferSegment.scope;
             modes = bufferSegment.mask;
-            owner = bufferSegment.owner;
         } else {
-            bufferScope = new MemoryScope(bb, null);
+            bufferScope = MemoryScope.create(bb, null);
             modes = defaultAccessModes(size);
-            owner = Thread.currentThread();
         }
         if (bb.isReadOnly()) {
             modes &= ~WRITE;
         }
         if (base != null) {
-            return new HeapMemorySegmentImpl<>(bbAddress + pos, () -> (byte[])base, size, modes, owner, bufferScope);
+            return new HeapMemorySegmentImpl<>(bbAddress + pos, () -> (byte[])base, size, modes, bufferScope);
         } else if (unmapper == null) {
-            return new NativeMemorySegmentImpl(bbAddress + pos, size, modes, owner, bufferScope);
+            return new NativeMemorySegmentImpl(bbAddress + pos, size, modes, bufferScope);
         } else {
-            return new MappedMemorySegmentImpl(bbAddress + pos, unmapper, size, modes, owner, bufferScope);
+            return new MappedMemorySegmentImpl(bbAddress + pos, unmapper, size, modes, bufferScope);
         }
     }
 
-    public static AbstractMemorySegmentImpl NOTHING = new AbstractMemorySegmentImpl(0, 0, null, MemoryScope.GLOBAL) {
+    public static final AbstractMemorySegmentImpl NOTHING = new AbstractMemorySegmentImpl(
+        0, 0, MemoryScope.createUnchecked(null, null, null)
+    ) {
         @Override
         ByteBuffer makeByteBuffer() {
             throw new UnsupportedOperationException();
@@ -454,7 +448,7 @@ Object base() {
         }
 
         @Override
-        AbstractMemorySegmentImpl dup(long offset, long size, int mask, Thread owner, MemoryScope scope) {
+        AbstractMemorySegmentImpl dup(long offset, long size, int mask, MemoryScope scope) {
             throw new UnsupportedOperationException();
         }
     };

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

@@ -52,8 +52,8 @@ public class HeapMemorySegmentImpl<H> extends AbstractMemorySegmentImpl {
     final Supplier<H> baseProvider;
 
     @ForceInline
-    HeapMemorySegmentImpl(long offset, Supplier<H> baseProvider, long length, int mask, Thread owner, MemoryScope scope) {
-        super(length, mask, owner, scope);
+    HeapMemorySegmentImpl(long offset, Supplier<H> baseProvider, long length, int mask, MemoryScope scope) {
+        super(length, mask, scope);
         this.offset = offset;
         this.baseProvider = baseProvider;
     }
@@ -69,8 +69,8 @@ long min() {
     }
 
     @Override
-    HeapMemorySegmentImpl<H> dup(long offset, long size, int mask, Thread owner, MemoryScope scope) {
-        return new HeapMemorySegmentImpl<H>(this.offset + offset, baseProvider, size, mask, owner, scope);
+    HeapMemorySegmentImpl<H> dup(long offset, long size, int mask, MemoryScope scope) {
+        return new HeapMemorySegmentImpl<>(this.offset + offset, baseProvider, size, mask, scope);
     }
 
     @Override
@@ -121,7 +121,7 @@ public static MemorySegment makeArraySegment(double[] arr) {
 
     static <Z> HeapMemorySegmentImpl<Z> makeHeapSegment(Supplier<Z> obj, int length, int base, int scale) {
         int byteSize = length * scale;
-        MemoryScope scope = new MemoryScope(null, null);
-        return new HeapMemorySegmentImpl<>(base, obj, byteSize, defaultAccessModes(byteSize), Thread.currentThread(), scope);
+        MemoryScope scope = MemoryScope.create(null, null);
+        return new HeapMemorySegmentImpl<>(base, obj, byteSize, defaultAccessModes(byteSize), scope);
     }
 }

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

@@ -48,8 +48,8 @@ public class MappedMemorySegmentImpl extends NativeMemorySegmentImpl implements
 
     private final UnmapperProxy unmapper;
 
-    MappedMemorySegmentImpl(long min, UnmapperProxy unmapper, long length, int mask, Thread owner, MemoryScope scope) {
-        super(min, length, mask, owner, scope);
+    MappedMemorySegmentImpl(long min, UnmapperProxy unmapper, long length, int mask, MemoryScope scope) {
+        super(min, length, mask, scope);
         this.unmapper = unmapper;
     }
 
@@ -60,8 +60,8 @@ ByteBuffer makeByteBuffer() {
     }
 
     @Override
-    MappedMemorySegmentImpl dup(long offset, long size, int mask, Thread owner, MemoryScope scope) {
-        return new MappedMemorySegmentImpl(min + offset, unmapper, size, mask, owner, scope);
+    MappedMemorySegmentImpl dup(long offset, long size, int mask, MemoryScope scope) {
+        return new MappedMemorySegmentImpl(min + offset, unmapper, size, mask, scope);
     }
 
     // mapped segment methods
@@ -103,13 +103,13 @@ public static MappedMemorySegment makeMappedSegment(Path path, long bytesSize, F
         if (bytesSize <= 0) throw new IllegalArgumentException("Requested bytes size must be > 0.");
         try (FileChannelImpl channelImpl = (FileChannelImpl)FileChannel.open(path, openOptions(mapMode))) {
             UnmapperProxy unmapperProxy = channelImpl.mapInternal(mapMode, 0L, bytesSize);
-            MemoryScope scope = new MemoryScope(null, unmapperProxy::unmap);
+            MemoryScope scope = MemoryScope.create(null, unmapperProxy::unmap);
             int modes = defaultAccessModes(bytesSize);
             if (mapMode == FileChannel.MapMode.READ_ONLY) {
                 modes &= ~WRITE;
             }
             return new MappedMemorySegmentImpl(unmapperProxy.address(), unmapperProxy, bytesSize,
-                    modes, Thread.currentThread(), scope);
+                    modes, scope);
         }
     }