8319928: Exceptions thrown by cleanup actions should be handled correctly

src/java.base/share/classes/java/lang/foreign/Arena.java

@@ -28,7 +28,9 @@
 import jdk.internal.foreign.MemorySessionImpl;
 import jdk.internal.ref.CleanerFactory;
 
+import java.io.Serial;
 import java.lang.foreign.MemorySegment.Scope;
+import java.util.function.Consumer;
 
 /**
  * An arena controls the lifecycle of native memory segments, providing both flexible
@@ -317,8 +319,50 @@ static Arena ofShared() {
      * @throws WrongThreadException if this arena is confined, and this method is called
      *         from a thread other than the arena's owner thread
      * @throws UnsupportedOperationException if this arena cannot be closed explicitly
+     * @throws CleanupException if an exception is thrown while executing a custom cleanup action
+     *                          associated with this arena (e.g. as a result of calling
+     *                          {@link MemorySegment#reinterpret(long, Arena, Consumer)} or
+     *                          {@link MemorySegment#reinterpret(Arena, Consumer)}).
      */
     @Override
     void close();
 
+    /**
+     * Thrown when an arena is {@linkplain Arena#close() closed explicitly} to indicate that an exception has occurred
+     * while executing a custom cleanup action associated with the arena. Custom cleanup actions can be attached
+     * to an arena when calling {@link MemorySegment#reinterpret(long, Arena, Consumer)} or
+     * {@link MemorySegment#reinterpret(Arena, Consumer)}. If more than one cleanup action fails,
+     * then this exception models the first failure whereas subsequent failures are
+     * {@linkplain Throwable#addSuppressed(Throwable) suppressed}. The order in which cleanup actions
+     * are executed is unspecified.
+     *
+     * @see MemorySegment#reinterpret(Arena, Consumer)
+     * @see MemorySegment#reinterpret(long, Arena, Consumer)
+     * @since 22
+     */
+    final class CleanupException extends RuntimeException {
+        @Serial
+        private static final long serialVersionUID = -403985082231157866L;
+
+        /**
+         * Constructs a {@code CleanupException} with the given detail message and cause.
+         *
+         * @param  message the detail message, can be null
+         * @param  cause the cause, can be null
+         */
+        public CleanupException(String message, Throwable cause) {
+            super(message, cause);
+        }
+
+        /**
+         * Constructs a {@code CleanupException} with the given cause and a detail
+         * message of {@code (cause==null ? null : cause.toString())} (which
+         * typically contains the class and detail message of {@code cause}).
+         *
+         * @param  cause the cause, can be null
+         */
+        public CleanupException(Throwable cause) {
+            super(cause);
+        }
+    }
 }

src/java.base/share/classes/java/lang/foreign/MemorySegment.java

@@ -720,8 +720,9 @@ public sealed interface MemorySegment permits AbstractMemorySegmentImpl {
      * }
      * That is, the cleanup action receives a segment that is associated with the global
      * scope, and is accessible from any thread. The size of the segment accepted by the
-     * cleanup action is {@link #byteSize()}. All exceptions thrown by the cleanup action are ignored,
-     * and do not affect other cleanup actions associated with the provided arena.
+     * cleanup action is {@link #byteSize()}. If the provided arena is {@linkplain Arena#close() closed explicitly},
+     * any exception thrown by the cleanup action is rethrown as a {@link Arena.CleanupException} when
+     * the arena is closed.
      *
      * @apiNote The cleanup action (if present) should take care not to leak the received
      *          segment to external clients that might access the segment after its
@@ -766,8 +767,9 @@ public sealed interface MemorySegment permits AbstractMemorySegmentImpl {
      * }
      * That is, the cleanup action receives a segment that is associated with the global
      * scope, and is accessible from any thread. The size of the segment accepted by the
-     * cleanup action is {@code newSize}. All exceptions thrown by the cleanup action are ignored,
-     * and do not affect other cleanup actions associated with the provided arena.
+     * cleanup action is {@code newSize}. If the provided arena is {@linkplain Arena#close() closed explicitly},
+     * any exception thrown by the cleanup action is rethrown as a {@link Arena.CleanupException} when
+     * the arena is closed.
      *
      * @apiNote The cleanup action (if present) should take care not to leak the received
      *          segment to external clients that might access the segment after its

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

@@ -26,6 +26,7 @@
 
 package jdk.internal.foreign;
 
+import java.lang.foreign.Arena.CleanupException;
 import java.lang.foreign.MemorySegment;
 import java.lang.foreign.Arena;
 import java.lang.foreign.MemorySegment.Scope;
@@ -254,11 +255,23 @@ public final void run() {
         }
 
         static void cleanup(ResourceCleanup first) {
+            CleanupException pendingException = null;
             ResourceCleanup current = first;
             while (current != null) {
-                current.cleanup();
+                try {
+                    current.cleanup();
+                } catch (Throwable ex) {
+                    if (pendingException == null) {
+                        pendingException = new CleanupException(ex);
+                    } else {
+                        pendingException.addSuppressed(ex);
+                    }
+                }
                 current = current.next;
             }
+            if (pendingException != null) {
+                throw pendingException;
+            }
         }
 
         public abstract static class ResourceCleanup {
@@ -277,11 +290,7 @@ static ResourceCleanup ofRunnable(Runnable cleanupAction) {
                 return new ResourceCleanup() {
                     @Override
                     public void cleanup() {
-                        try {
-                            cleanupAction.run();
-                        } catch (Throwable ex) {
-                            // swallow
-                        }
+                        cleanupAction.run();
                     }
                 };
             }

test/jdk/java/foreign/TestSegments.java

@@ -32,9 +32,12 @@
 import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
 
+import java.lang.foreign.Arena.CleanupException;
 import java.lang.invoke.VarHandle;
 import java.nio.ByteBuffer;
+import java.util.HashSet;
 import java.util.List;
+import java.util.Set;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.IntFunction;
@@ -384,12 +387,41 @@ void testReinterpret() {
     @Test
     void testThrowInCleanup() {
         AtomicInteger counter = new AtomicInteger();
-        try (Arena arena = Arena.ofConfined()){
-            MemorySegment.ofAddress(42).reinterpret(arena, seg -> { throw new AssertionError(); });
-            MemorySegment.ofAddress(42).reinterpret(100, arena, seg -> { throw new AssertionError(); });
+        CleanupException thrown = null;
+        Set<String> expected = new HashSet<>();
+        try (Arena arena = Arena.ofConfined()) {
+            for (int i = 0 ; i < 10 ; i++) {
+                String msg = "exception#" + i;
+                expected.add(msg);
+                MemorySegment.ofAddress(42).reinterpret(arena, seg -> {
+                    throw new AssertionError(msg);
+                });
+            }
+            for (int i = 10 ; i < 20 ; i++) {
+                String msg = "exception#" + i;
+                expected.add(msg);
+                MemorySegment.ofAddress(42).reinterpret(100, arena, seg -> {
+                    throw new AssertionError(msg);
+                });
+            }
             MemorySegment.ofAddress(42).reinterpret(arena, seg -> counter.incrementAndGet());
-        } // no exception should occur here!
+        } catch (CleanupException ex) {
+            thrown = ex;
+        }
+        assertNotNull(thrown);
         assertEquals(counter.get(), 1);
+        assertEquals(thrown.getSuppressed().length, 19);
+        Throwable[] errors = new AssertionError[20];
+        assertTrue(thrown.getCause() instanceof AssertionError);
+        errors[0] = thrown.getCause();
+        for (int i = 0 ; i < 19 ; i++) {
+            assertTrue(thrown.getSuppressed()[i] instanceof AssertionError);
+            errors[i + 1] = thrown.getSuppressed()[i];
+        }
+        for (Throwable t : errors) {
+            assertTrue(expected.remove(t.getMessage()));
+        }
+        assertTrue(expected.isEmpty());
     }
 
     @DataProvider(name = "badSizeAndAlignments")