8319928: Exceptions thrown by cleanup actions should be handled correctly

Reviewed-by: jvernee

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

@@ -29,6 +29,7 @@ import jdk.internal.foreign.MemorySessionImpl;
 import jdk.internal.ref.CleanerFactory;
 
 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 +318,11 @@ public interface Arena extends SegmentAllocator, AutoCloseable {
      * @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 RuntimeException 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();
-
 }

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

@@ -254,11 +254,24 @@ public abstract sealed class MemorySessionImpl
         }
 
         static void cleanup(ResourceCleanup first) {
+            RuntimeException pendingException = null;
             ResourceCleanup current = first;
             while (current != null) {
-                current.cleanup();
+                try {
+                    current.cleanup();
+                } catch (RuntimeException ex) {
+                    if (pendingException == null) {
+                        pendingException = ex;
+                    } else if (ex != pendingException) {
+                        // note: self-suppression is not supported
+                        pendingException.addSuppressed(ex);
+                    }
+                }
                 current = current.next;
             }
+            if (pendingException != null) {
+                throw pendingException;
+            }
         }
 
         public abstract static class ResourceCleanup {

test/jdk/java/foreign/TestSegments.java

@@ -34,7 +34,9 @@ import org.testng.annotations.Test;
 
 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;
@@ -381,6 +383,71 @@ public class TestSegments {
         assertEquals(counter.get(), 2);
     }
 
+    @Test
+    void testThrowInCleanup() {
+        AtomicInteger counter = new AtomicInteger();
+        RuntimeException 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 IllegalArgumentException(msg);
+                });
+            }
+            for (int i = 10 ; i < 20 ; i++) {
+                String msg = "exception#" + i;
+                expected.add(msg);
+                MemorySegment.ofAddress(42).reinterpret(100, arena, seg -> {
+                    throw new IllegalArgumentException(msg);
+                });
+            }
+            MemorySegment.ofAddress(42).reinterpret(arena, seg -> counter.incrementAndGet());
+        } catch (RuntimeException ex) {
+            thrown = ex;
+        }
+        assertNotNull(thrown);
+        assertEquals(counter.get(), 1);
+        assertEquals(thrown.getSuppressed().length, 19);
+        Throwable[] errors = new IllegalArgumentException[20];
+        assertTrue(thrown instanceof IllegalArgumentException);
+        errors[0] = thrown;
+        for (int i = 0 ; i < 19 ; i++) {
+            assertTrue(thrown.getSuppressed()[i] instanceof IllegalArgumentException);
+            errors[i + 1] = thrown.getSuppressed()[i];
+        }
+        for (Throwable t : errors) {
+            assertTrue(expected.remove(t.getMessage()));
+        }
+        assertTrue(expected.isEmpty());
+    }
+
+    @Test
+    void testThrowInCleanupSame() {
+        AtomicInteger counter = new AtomicInteger();
+        Throwable thrown = null;
+        IllegalArgumentException iae = new IllegalArgumentException();
+        try (Arena arena = Arena.ofConfined()) {
+            for (int i = 0 ; i < 10 ; i++) {
+                MemorySegment.ofAddress(42).reinterpret(arena, seg -> {
+                    throw iae;
+                });
+            }
+            for (int i = 10 ; i < 20 ; i++) {
+                MemorySegment.ofAddress(42).reinterpret(100, arena, seg -> {
+                    throw iae;
+                });
+            }
+            MemorySegment.ofAddress(42).reinterpret(arena, seg -> counter.incrementAndGet());
+        } catch (RuntimeException ex) {
+            thrown = ex;
+        }
+        assertEquals(thrown, iae);
+        assertEquals(counter.get(), 1);
+        assertEquals(thrown.getSuppressed().length, 0);
+    }
+
     @DataProvider(name = "badSizeAndAlignments")
     public Object[][] sizesAndAlignments() {
         return new Object[][] {