Free native allocations when ByteBuffer creation fails

saferalloc/src/main/java/org/ngengine/saferalloc/SaferAlloc.java

@@ -23,7 +23,7 @@ public static ByteBuffer malloc(int size) {
     ensureLoaded();
     requireNonNegativeSize(size);
     long addr = SaferAllocNative.malloc(size);
-    return wrapMemByteBuffer(addr, size);
+    return wrapNewAllocation(addr, size);
   }
 
   public static ByteBuffer calloc(int count, int size) {
@@ -32,7 +32,7 @@ public static ByteBuffer calloc(int count, int size) {
     requireNonNegativeSize(size);
     int total = requireBufferCapacity(count, size);
     long addr = SaferAllocNative.calloc(count, size);
-    return wrapMemByteBuffer(addr, total);
+    return wrapNewAllocation(addr, total);
   }
 
   public static ByteBuffer realloc(ByteBuffer buffer, int newSize) {
@@ -51,7 +51,7 @@ public static ByteBuffer mallocAligned(int size, int alignment) {
     requireNonNegativeSize(size);
     requireValidAlignment(alignment);
     long addr = SaferAllocNative.mallocAligned(size, alignment);
-    return wrapMemByteBuffer(addr, size);
+    return wrapNewAllocation(addr, size);
   }
 
   public static long address(Buffer buffer) {
@@ -122,6 +122,21 @@ private static int pointerSizeBytes() {
     return SaferAllocNative.pointerSizeBytes();
   }
 
+  private static ByteBuffer wrapNewAllocation(long addr, int size) {
+    try {
+      ByteBuffer buffer = wrapMemByteBuffer(addr, size);
+      if (buffer == null && addr != 0L) {
+        SaferAllocNative.free(addr);
+      }
+      return buffer;
+    } catch (RuntimeException | Error failure) {
+      if (addr != 0L) {
+        SaferAllocNative.free(addr);
+      }
+      throw failure;
+    }
+  }
+
   private static ByteBuffer wrapMemByteBuffer(long addr, int size) {
     requireNonNegativeSize(size);
     if (addr == 0L) {

saferalloc/src/test/java/org/ngengine/saferalloc/AllocationFailureMain.java

@@ -0,0 +1,44 @@
+package org.ngengine.saferalloc;
+
+import java.nio.ByteBuffer;
+
+/** Runs in a small, isolated heap so native allocation can outlive Java allocation. */
+public final class AllocationFailureMain {
+  private static Object[] retained;
+
+  public static void main(String[] args) {
+    String operation = args[0];
+    ByteBuffer warm = allocate(operation);
+    SaferAlloc.free(warm);
+    long before = SaferAlloc.currentAllocatedBytes();
+    retained = new Object[1000000];
+    int used = 0;
+    try {
+      while (used < retained.length) retained[used++] = new byte[16];
+    } catch (OutOfMemoryError expected) {
+      // Keep these objects reachable until the ByteBuffer allocation has failed.
+    }
+    boolean failed = false;
+    ByteBuffer result = null;
+    try {
+      result = allocate(operation);
+    } catch (OutOfMemoryError expected) {
+      failed = true;
+    }
+    long after = SaferAlloc.currentAllocatedBytes();
+    retained = null;
+    System.gc();
+    if (result != null) SaferAlloc.free(result);
+    if (!failed) throw new AssertionError("The test did not exhaust the Java heap");
+    if (after != before) {
+      throw new AssertionError(operation + " leaked " + (after - before) + " native bytes");
+    }
+    System.out.println(operation + ": wrapper allocation failed without leaking native memory");
+  }
+
+  private static ByteBuffer allocate(String operation) {
+    if ("calloc".equals(operation)) return SaferAlloc.calloc(64, 64);
+    if ("aligned".equals(operation)) return SaferAlloc.mallocAligned(4096, 64);
+    return SaferAlloc.malloc(4096);
+  }
+}

saferalloc/src/test/java/org/ngengine/saferalloc/AllocationFailureTest.java

@@ -0,0 +1,37 @@
+package org.ngengine.saferalloc;
+
+import java.io.File;
+import java.nio.file.Path;
+import java.nio.file.Files;
+import java.util.concurrent.TimeUnit;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+import static org.junit.jupiter.api.Assertions.*;
+
+class AllocationFailureTest {
+  @TempDir Path output;
+
+  @Test
+  void freesNewNativeAllocationsWhenJavaHeapIsExhausted() throws Exception {
+    String java = new File(System.getProperty("java.home"), "bin/java").getAbsolutePath();
+    String classpath = new File(AllocationFailureMain.class.getProtectionDomain()
+        .getCodeSource().getLocation().toURI()).getAbsolutePath() + File.pathSeparator
+        + new File(SaferAlloc.class.getProtectionDomain().getCodeSource()
+        .getLocation().toURI()).getAbsolutePath();
+    String nativePath = System.getProperty("saferalloc.native.override");
+    assertNotNull(nativePath, "the test needs the configured test native library");
+    for (String operation : new String[] {"malloc", "calloc", "aligned"}) {
+      Path log = output.resolve(operation + ".log");
+      Process child = new ProcessBuilder(java, "-Xmx16m", "-XX:-UseGCOverheadLimit",
+          "-Dsaferalloc.native.override=" + nativePath, "-cp", classpath,
+          AllocationFailureMain.class.getName(), operation)
+          .redirectErrorStream(true).redirectOutput(log.toFile()).start();
+      try {
+        assertTrue(child.waitFor(30, TimeUnit.SECONDS), "allocation test timed out");
+        assertEquals(0, child.exitValue(), new String(Files.readAllBytes(log), "UTF-8"));
+      } finally {
+        child.destroyForcibly();
+      }
+    }
+  }
+}