8253243: Investigate ways to make MemorySegment::ofNativeRestricted more composable

src/jdk.incubator.foreign/share/classes/jdk/incubator/foreign/MemoryAddress.java

@@ -27,6 +27,10 @@
 package jdk.incubator.foreign;
 
 import jdk.internal.foreign.MemoryAddressImpl;
+import jdk.internal.foreign.NativeMemorySegmentImpl;
+import jdk.internal.foreign.Utils;
+
+import java.util.Objects;
 
 /**
  * A memory address models a reference into a memory location. Memory addresses are typically obtained using the
@@ -83,9 +87,40 @@ default MemoryAddress address() {
     long segmentOffset(MemorySegment segment);
 
     /**
-     * Returns the raw long value associated to this memory address.
-     * @return The raw long value associated to this memory address.
-     * @throws UnsupportedOperationException if this memory address is associated with an heap segment.
+     * Returns a new confined native memory segment with given size, and whose base address is this address; the returned segment has its own temporal
+     * bounds, and can therefore be closed. This method can be very useful when interacting with custom native memory sources (e.g. custom allocators),
+     * where an address to some underlying memory region is typically obtained from native code
+     * (often as a plain {@code long} value). The returned segment will feature all <a href="#access-modes">access modes</a>
+     * (see {@link MemorySegment#ALL_ACCESS}), and its confinement thread is the current thread (see {@link Thread#currentThread()}).
+     * <p>
+     * Calling {@link MemorySegment#close()} on the returned segment will <em>not</em> result in releasing any
+     * memory resources which might implicitly be associated with the segment. If the client wants to specify
+     * a cleanup action to be executed when the returned segment is closed, the {@link MemorySegment#withCleanupAction(Runnable)}
+     * method should be used.
+     * <p>
+     * This method is <em>restricted</em>. Restricted methods are unsafe, and, if used incorrectly, their use might crash
+     * the JVM or, worse, silently result in memory corruption. Thus, clients should refrain from depending on
+     * restricted methods, and use safe and supported functionalities, where possible.
+     *
+     * @param bytesSize the desired size.
+     * @return a new confined native memory segment with given base address and size.
+     * @throws IllegalArgumentException if {@code bytesSize <= 0}.
+     * @throws UnsupportedOperationException if this address is an heap address.
+     * @throws IllegalAccessError if the runtime property {@code foreign.restricted} is not set to either
+     * {@code permit}, {@code warn} or {@code debug} (the default value is set to {@code deny}).
+     */
+    default MemorySegment asSegmentRestricted(long bytesSize) {
+        Utils.checkRestrictedAccess("MemoryAddress.asSegmentRestricted");
+        if (bytesSize <= 0) {
+            throw new IllegalArgumentException("Invalid size : " + bytesSize);
+        }
+        return NativeMemorySegmentImpl.makeNativeSegmentUnchecked(this, bytesSize);
+    }
+
+    /**
+     * Returns the raw long value associated with this memory address.
+     * @return The raw long value associated with this memory address.
+     * @throws UnsupportedOperationException if this memory address is an heap address.
      */
     long toRawLongValue();
 

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

@@ -295,6 +295,12 @@ public MemorySegment withOwnerThread(Thread newOwner) {
         }
     }
 
+    @Override
+    public MemorySegment withCleanupAction(Runnable action) {
+        checkValidState();
+        return dup(0L, length, mask, scope.wrapAction(action));
+    }
+
     @Override
     public void registerCleaner(Cleaner cleaner) {
         checkAccessModes(CLOSE);

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

@@ -63,17 +63,7 @@ private MemoryScope(Object ref, CleanupAction cleanupAction) {
      * @return a confined memory scope
      */
     static MemoryScope createConfined(Object ref, CleanupAction cleanupAction) {
-        return createConfined(Thread.currentThread(), ref, cleanupAction);
-    }
-
-    /**
-     * Creates a confined memory scope with given attachment, cleanup action and owner thread.
-     * @param ref           an optional reference to an instance that needs to be kept reachable
-     * @param cleanupAction a cleanup action to be executed when returned scope is closed
-     * @return a confined memory scope
-     */
-    static MemoryScope createConfined(Thread owner, Object ref, CleanupAction cleanupAction) {
-        return new ConfinedScope(owner, ref, cleanupAction);
+        return new ConfinedScope(Thread.currentThread(), ref, cleanupAction);
     }
 
     /**
@@ -147,6 +137,17 @@ MemoryScope share() {
         }
     }
 
+    final MemoryScope wrapAction(Runnable runnable) {
+        try {
+            justClose();
+            return ownerThread() == null ?
+                    new SharedScope(ref, cleanupAction.wrap(runnable)) :
+                    new ConfinedScope(ownerThread(), ref, cleanupAction.wrap(runnable));
+        } finally {
+            Reference.reachabilityFence(this);
+        }
+    }
+
     /**
      * Returns "owner" thread of this scope.
      * @return owner thread (or null for a shared scope)
@@ -274,6 +275,7 @@ void justClose() {
     interface CleanupAction extends Runnable {
         void cleanup();
         CleanupAction dup();
+        CleanupAction wrap(Runnable runnable);
 
         @Override
         default void run() {
@@ -291,6 +293,11 @@ public void cleanup() {
             public CleanupAction dup() {
                 return this;
             }
+
+            @Override
+            public CleanupAction wrap(Runnable runnable) {
+                return AtMostOnceOnly.of(runnable);
+            }
         };
 
         /**
@@ -325,6 +332,20 @@ public CleanupAction dup() {
                 return new DupAction(this);
             }
 
+            @Override
+            public CleanupAction wrap(Runnable runnable) {
+                disable();
+                return AtMostOnceOnly.of(() -> {
+                    try {
+                        runnable.run();
+                    } catch (Throwable t) {
+                        // ignore
+                    } finally {
+                        doCleanup();
+                    }
+                });
+            }
+
             //where
             static class DupAction extends AtMostOnceOnly {
                 final AtMostOnceOnly root;

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

@@ -41,8 +41,9 @@
  */
 public class NativeMemorySegmentImpl extends AbstractMemorySegmentImpl {
 
-    public static final MemorySegment EVERYTHING = NativeMemorySegmentImpl.makeNativeSegmentUnchecked(MemoryAddress.NULL,
-            Long.MAX_VALUE, null, null, null).withAccessModes(READ | WRITE);
+    public static final MemorySegment EVERYTHING = makeNativeSegmentUnchecked(MemoryAddress.NULL, Long.MAX_VALUE)
+            .withOwnerThread(null)
+            .withAccessModes(READ | WRITE);
 
     private static final Unsafe unsafe = Unsafe.getUnsafe();
 
@@ -113,12 +114,8 @@ void doCleanup() {
         return segment;
     }
 
-    public static MemorySegment makeNativeSegmentUnchecked(MemoryAddress min, long bytesSize, Thread owner, Runnable cleanup, Object attachment) {
-        MemoryScope.CleanupAction cleanupAction = cleanup != null ?
-                MemoryScope.CleanupAction.AtMostOnceOnly.of(cleanup) : MemoryScope.CleanupAction.DUMMY;
-        MemoryScope scope = owner == null ?
-                MemoryScope.createShared(attachment, cleanupAction) :
-                MemoryScope.createConfined(owner, attachment, cleanupAction);
-        return new NativeMemorySegmentImpl(min.toRawLongValue(), bytesSize, defaultAccessModes(bytesSize), scope);
+    public static MemorySegment makeNativeSegmentUnchecked(MemoryAddress min, long bytesSize) {
+        return new NativeMemorySegmentImpl(min.toRawLongValue(), bytesSize, defaultAccessModes(bytesSize),
+                MemoryScope.createConfined(null, MemoryScope.CleanupAction.DUMMY));
     }
 }

test/jdk/java/foreign/TestArrays.java

@@ -39,7 +39,6 @@
 import java.util.function.BiFunction;
 import java.util.function.Consumer;
 import java.util.function.Function;
-import java.util.function.ToIntFunction;
 
 import org.testng.annotations.*;
 
@@ -113,7 +112,7 @@ public void testArrays(Consumer<MemorySegment> init, Consumer<MemorySegment> che
     public void testTooBigForArray(MemoryLayout layout, Function<MemorySegment, Object> arrayFactory) {
         MemoryLayout seq = MemoryLayout.ofSequence((Integer.MAX_VALUE * layout.byteSize()) + 1, layout);
         //do not really allocate here, as it's way too much memory
-        try (MemorySegment segment = MemorySegment.ofNativeRestricted(MemoryAddress.NULL, seq.byteSize(), null, null, null)) {
+        try (MemorySegment segment = MemoryAddress.NULL.asSegmentRestricted(seq.byteSize())) {
             arrayFactory.apply(segment);
         }
     }

test/jdk/java/foreign/TestCleaner.java

@@ -27,10 +27,9 @@
  * @test
  * @modules java.base/jdk.internal.ref
  *          jdk.incubator.foreign/jdk.incubator.foreign
- * @run testng/othervm -Dforeign.restricted=permit TestCleaner
+ * @run testng TestCleaner
  */
 
-import jdk.incubator.foreign.MemoryAccess;
 import jdk.incubator.foreign.MemoryAddress;
 import jdk.incubator.foreign.MemorySegment;
 import java.lang.ref.Cleaner;
@@ -61,7 +60,8 @@ int cleanupCalls() {
     @Test(dataProvider = "cleaners")
     public void test(int n, Supplier<Cleaner> cleanerFactory, SegmentFunction segmentFunction) {
         SegmentState segmentState = new SegmentState();
-        MemorySegment segment = makeSegment(segmentState);
+        MemorySegment segment = MemorySegment.allocateNative(10)
+                .withCleanupAction(segmentState::cleanup);
         // register cleaners before
         for (int i = 0 ; i < n ; i++) {
             segment.registerCleaner(cleanerFactory.get());
@@ -88,10 +88,6 @@ public void test(int n, Supplier<Cleaner> cleanerFactory, SegmentFunction segmen
         assertEquals(segmentState.cleanupCalls(), 1);
     }
 
-    MemorySegment makeSegment(SegmentState segmentState) {
-        return MemorySegment.ofNativeRestricted(MemoryAddress.NULL, 10, Thread.currentThread(), segmentState::cleanup, null);
-    }
-
     enum SegmentFunction implements Function<MemorySegment, MemorySegment> {
         IDENTITY(Function.identity()),
         CLOSE(s -> { s.close(); return s; }),

test/jdk/java/foreign/TestNative.java

@@ -167,8 +167,8 @@ public void testNativeCapacity(Function<ByteBuffer, Buffer> bufferFunction, int
     @Test
     public void testDefaultAccessModes() {
         MemoryAddress addr = MemoryAddress.ofLong(allocate(12));
-        MemorySegment mallocSegment = MemorySegment.ofNativeRestricted(addr, 12, null,
-                () -> free(addr.toRawLongValue()), null);
+        MemorySegment mallocSegment = addr.asSegmentRestricted(12)
+                .withCleanupAction(() -> free(addr.toRawLongValue()));
         try (MemorySegment segment = mallocSegment) {
             assertTrue(segment.hasAccessModes(ALL_ACCESS));
             assertEquals(segment.accessModes(), ALL_ACCESS);
@@ -185,8 +185,8 @@ public void testDefaultAccessModesEverthing() {
     @Test
     public void testMallocSegment() {
         MemoryAddress addr = MemoryAddress.ofLong(allocate(12));
-        MemorySegment mallocSegment = MemorySegment.ofNativeRestricted(addr, 12, null,
-                () -> free(addr.toRawLongValue()), null);
+        MemorySegment mallocSegment = addr.asSegmentRestricted(12)
+                .withCleanupAction(() -> free(addr.toRawLongValue()));
         assertEquals(mallocSegment.byteSize(), 12);
         mallocSegment.close(); //free here
         assertTrue(!mallocSegment.isAlive());
@@ -204,15 +204,10 @@ public void testEverythingSegment() {
     @Test(expectedExceptions = IllegalArgumentException.class)
     public void testBadResize() {
         try (MemorySegment segment = MemorySegment.allocateNative(4)) {
-            MemorySegment.ofNativeRestricted(segment.address(), 0, null, null, null);
+            segment.address().asSegmentRestricted(0);
         }
     }
 
-    @Test(expectedExceptions = NullPointerException.class)
-    public void testNullUnsafeSegment() {
-        MemorySegment.ofNativeRestricted(null, 10, null, null, null);
-    }
-
     static {
         System.loadLibrary("NativeAccess");
     }

test/jdk/java/foreign/TestNoForeignUnsafeOverride.java

@@ -29,8 +29,8 @@
  */
 
 import jdk.incubator.foreign.MemoryAddress;
-import jdk.incubator.foreign.MemorySegment;
 
+import jdk.incubator.foreign.MemorySegment;
 import org.testng.annotations.Test;
 
 public class TestNoForeignUnsafeOverride {
@@ -40,6 +40,6 @@ public class TestNoForeignUnsafeOverride {
 
     @Test(expectedExceptions = IllegalAccessError.class)
     public void testUnsafeAccess() {
-        MemorySegment.ofNativeRestricted(MemoryAddress.ofLong(42), 10, null, null, null);
+        MemorySegment.ofNativeRestricted();
     }
 }

test/jdk/java/foreign/TestSegments.java

@@ -336,7 +336,8 @@ static class SegmentMember {
                 "toLongArray",
                 "toDoubleArray",
                 "withOwnerThread",
-                "registerCleaner"
+                "registerCleaner",
+                "withCleanupAction"
         );
 
         public SegmentMember(Method method, Object[] params) {

test/jdk/java/foreign/TestSharedAccess.java

@@ -27,10 +27,7 @@
  * @run testng/othervm -Dforeign.restricted=permit TestSharedAccess
  */
 
-import jdk.incubator.foreign.MemoryLayout;
-import jdk.incubator.foreign.MemorySegment;
-import jdk.incubator.foreign.MemoryLayouts;
-import jdk.incubator.foreign.SequenceLayout;
+import jdk.incubator.foreign.*;
 import org.testng.annotations.*;
 
 import java.lang.invoke.VarHandle;
@@ -127,8 +124,8 @@ public void testSharedUnsafe() throws Throwable {
             setInt(s, 42);
             assertEquals(getInt(s), 42);
             List<Thread> threads = new ArrayList<>();
-            MemorySegment sharedSegment = MemorySegment.ofNativeRestricted(
-                    s.address(), s.byteSize(), null, null, null);
+            MemorySegment sharedSegment = s.address().asSegmentRestricted(s.byteSize())
+                    .withOwnerThread(null);
             for (int i = 0 ; i < 1000 ; i++) {
                 threads.add(new Thread(() -> {
                     assertEquals(getInt(sharedSegment), 42);