Keep the original allocation valid when realloc wrapping fails

README.md

@@ -31,8 +31,10 @@ On desktop, bundled JNI libraries are extracted into a fresh private directory f
 
 Notes:
 - `malloc/calloc/mallocAligned` return `null` on allocation failure.
-- `realloc(buffer, newSize)` throws `OutOfMemoryError` if growth fails and the old allocation is still valid.
-- `realloc(buffer, 0)` may return `null` and free the old allocation.
+- `realloc(buffer, newSize)` throws `OutOfMemoryError` if a non-null buffer cannot be resized, including when its replacement Java wrapper cannot be allocated. The old allocation remains valid on failure.
+- `realloc(buffer, 0)` returns `null` and frees the old allocation.
+- `realloc(null, newSize)` behaves like `malloc(newSize)` for positive sizes; `realloc(null, 0)` returns `null`.
+- High-level `realloc` allocates a separate native block and Java wrapper before copying `min(buffer.capacity(), newSize)` bytes and freeing the original. It ignores the original position and limit, and the returned buffer has position zero and limit `newSize`. This guarantees exception safety at the cost of copying and temporarily holding both allocations, even when shrinking or keeping the same size. Native reallocation function pointers retain their existing in-place-capable behavior.
 - invalid sizes throw `IllegalArgumentException`.
 - `mallocAligned(size, alignment)` requires `alignment` to be a power of two and a multiple of pointer size.
 - `calloc(count, size)` rejects capacities larger than `Integer.MAX_VALUE`.
@@ -41,5 +43,6 @@ Notes:
 ## Note
 
 - This API manages native memory manually. Always call `free(buffer)`.
-- After successful `realloc`, treat the old buffer as invalid and use the returned one only. If `realloc` throws, keep using the old buffer (it is still allocated).
+- Pass only live, owning buffers returned by SaferAlloc to `realloc` or `free`, not slices, duplicates, or buffers allocated elsewhere. These operations do not track ownership.
+- After successful `realloc`, treat the old buffer and any views as invalid and use the returned one only. If `realloc` throws, keep using the old buffer (it is still allocated).
 - `free(null)` is a no-op.

saferalloc/src/main/c/saferalloc.c

@@ -181,6 +181,72 @@ JNIEXPORT jlong JNICALL Java_org_ngengine_saferalloc_SaferAllocNative_realloc(JN
   return (jlong)(uintptr_t)safer_tracked_realloc((void*)(uintptr_t)addr, (size_t)newSize);
 }
 
+// The Java-facing reallocation must not release the original allocation until
+// the replacement's Java wrapper exists. Raw realloc entry points retain their
+// usual native semantics and may resize in place.
+JNIEXPORT jobject JNICALL Java_org_ngengine_saferalloc_SaferAllocNative_reallocBuffer(JNIEnv* env, jclass cls, jobject buffer, jint newSize) {
+  (void)cls;
+  if (newSize < 0) {
+    safer_throw_illegal_argument(env, "size must be >= 0");
+    return NULL;
+  }
+
+  void* old_ptr = NULL;
+  jlong old_capacity = 0;
+  if (buffer != NULL) {
+    old_ptr = (*env)->GetDirectBufferAddress(env, buffer);
+    if ((*env)->ExceptionCheck(env)) {
+      return NULL;
+    }
+    old_capacity = (*env)->GetDirectBufferCapacity(env, buffer);
+    if ((*env)->ExceptionCheck(env)) {
+      return NULL;
+    }
+    if (old_ptr == NULL || old_capacity < 0) {
+      safer_throw_illegal_argument(env, "buffer must be a direct ByteBuffer");
+      return NULL;
+    }
+  }
+
+  if (newSize == 0) {
+    safer_tracked_free(old_ptr);
+    return NULL;
+  }
+
+  void* new_ptr = safer_tracked_malloc((size_t)newSize);
+  if (new_ptr == NULL) {
+    // Preserve realloc(NULL, size)'s allocation-failure behavior.
+    if (old_ptr != NULL) {
+      jclass oom = (*env)->FindClass(env, "java/lang/OutOfMemoryError");
+      if (oom != NULL) {
+        (*env)->ThrowNew(env, oom, "realloc failed; original buffer is still valid");
+      }
+    }
+    return NULL;
+  }
+
+  jobject replacement = safer_new_direct_byte_buffer(env, new_ptr, (jlong)newSize);
+  if (replacement == NULL || (*env)->ExceptionCheck(env)) {
+    // In particular, NewDirectByteBuffer may throw while Java heap is exhausted.
+    safer_tracked_free(new_ptr);
+    if (!(*env)->ExceptionCheck(env)) {
+      jclass oom = (*env)->FindClass(env, "java/lang/OutOfMemoryError");
+      if (oom != NULL) {
+        (*env)->ThrowNew(env, oom, "could not create replacement ByteBuffer");
+      }
+    }
+    return NULL;
+  }
+
+  // Copy the allocation's contents, regardless of the old position and limit.
+  size_t copy_size = (size_t)(old_capacity < (jlong)newSize ? old_capacity : (jlong)newSize);
+  if (copy_size > 0) {
+    memcpy(new_ptr, old_ptr, copy_size);
+  }
+  safer_tracked_free(old_ptr);
+  return replacement;
+}
+
 JNIEXPORT void JNICALL Java_org_ngengine_saferalloc_SaferAllocNative_free(JNIEnv* env, jclass cls, jlong addr) {
   (void)env;
   (void)cls;

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

@@ -38,12 +38,7 @@ public static ByteBuffer calloc(int count, int size) {
   public static ByteBuffer realloc(ByteBuffer buffer, int newSize) {
     ensureLoaded();
     requireNonNegativeSize(newSize);
-    long addr = address(buffer);
-    long newAddr = SaferAllocNative.realloc(addr, newSize);
-    if (newAddr == 0L && addr != 0L && newSize > 0) {
-      throw new OutOfMemoryError("realloc failed; original buffer is still valid");
-    }
-    return wrapMemByteBuffer(newAddr, newSize);
+    return SaferAllocNative.reallocBuffer(buffer, newSize);
   }
 
   public static ByteBuffer mallocAligned(int size, int alignment) {

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

@@ -10,6 +10,7 @@ private SaferAllocNative() {}
   public static native long malloc(long size);
   public static native long calloc(long count, long size);
   public static native long realloc(long addr, long newSize);
+  static native ByteBuffer reallocBuffer(ByteBuffer buffer, int newSize);
   public static native void free(long addr);
   public static native long mallocAligned(long size, long alignment);
   public static  native long currentAllocatedBytes();

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

@@ -0,0 +1,88 @@
+package org.ngengine.saferalloc;
+
+import java.nio.ByteBuffer;
+
+/** Exhausts only a small child JVM's Java heap, while native allocation still succeeds. */
+public final class ReallocationFailureMain {
+  private static Object[] retained;
+  private static ByteBuffer original;
+
+  public static void main(String[] args) {
+    // Load and warm the complete path before exhausting the heap, so the failure
+    // occurs when JNI creates the replacement wrapper rather than loading classes.
+    ByteBuffer warm = SaferAlloc.malloc(16);
+    warm = SaferAlloc.realloc(warm, 4096);
+    SaferAlloc.free(warm);
+
+    long baseline = SaferAlloc.currentAllocatedBytes();
+    original = SaferAlloc.malloc(16);
+    if (original == null) throw new AssertionError("initial native allocation failed");
+    for (int i = 0; i < original.capacity(); i++) {
+      original.put(i, (byte) (i + 1));
+    }
+    original.position(3);
+    original.limit(8);
+    long address = SaferAlloc.address(original);
+    long before = SaferAlloc.currentAllocatedBytes();
+
+    retained = new Object[1000000];
+    int used = 0;
+    // A first OOM can release JVM-internal emergency resources. Refill after
+    // that, keeping every allocated object strongly reachable through the call.
+    for (int pass = 0; pass < 4; pass++) {
+      try {
+        while (used < retained.length) {
+          retained[used] = new byte[16];
+          used++;
+        }
+      } catch (OutOfMemoryError expected) {
+        // The error and its stack trace also occupy heap. Retain them so GC
+        // cannot reclaim that space to satisfy the subsequent wrapper allocation.
+        if (used < retained.length) retained[used++] = expected;
+      }
+    }
+
+    boolean failed = false;
+    ByteBuffer result = null;
+    try {
+      result = SaferAlloc.realloc(original, 4096);
+    } catch (OutOfMemoryError expected) {
+      failed = true;
+    }
+    long after = SaferAlloc.currentAllocatedBytes();
+    retained = null;
+    System.gc();
+
+    if (!failed) {
+      if (result != null) SaferAlloc.free(result);
+      throw new AssertionError("replacement wrapper did not fail under Java heap exhaustion");
+    }
+    // Check accounting BEFORE touching or freeing the original. On the buggy
+    // implementation it has already been resized/freed, so dereferencing it here
+    // would make this regression unsafe. Growing 16 to 4096 makes that observable.
+    if (after != before) {
+      throw new AssertionError("native accounting changed after wrapper failure: " + before + " -> " + after);
+    }
+    if (SaferAlloc.address(original) != address) {
+      throw new AssertionError("original buffer address changed");
+    }
+    if (original.capacity() != 16 || original.position() != 3 || original.limit() != 8) {
+      throw new AssertionError("original buffer state changed");
+    }
+    original.clear();
+    for (int i = 0; i < original.capacity(); i++) {
+      if (original.get(i) != (byte) (i + 1)) {
+        throw new AssertionError("original content changed at index " + i);
+      }
+    }
+    // Check that the retained allocation is still writable and may be freed once.
+    original.put(0, (byte) 99);
+    if (original.get(0) != (byte) 99) throw new AssertionError("original is not writable");
+    SaferAlloc.free(original);
+    original = null;
+    if (SaferAlloc.currentAllocatedBytes() != baseline) {
+      throw new AssertionError("native accounting did not return to baseline");
+    }
+    System.out.println("REALLOC_WRAPPER_FAILURE_PRESERVED_ORIGINAL");
+  }
+}

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

@@ -0,0 +1,52 @@
+package org.ngengine.saferalloc;
+
+import java.io.File;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+import org.junit.jupiter.api.Test;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+class ReallocationFailureTest {
+  @Test
+  void wrapperAllocationFailurePreservesOriginal() throws Exception {
+    List<String> command = new ArrayList<>();
+    command.add(System.getProperty("java.home") + "/bin/java");
+    command.add("-Xms16m");
+    command.add("-Xmx16m");
+    command.add("-XX:+UseSerialGC");
+    command.add("-XX:-UseGCOverheadLimit");
+    String nativeOverride = System.getProperty("saferalloc.native.override");
+    assertNotNull(nativeOverride, "the test needs the configured test native library");
+    command.add("-Dsaferalloc.native.override=" + nativeOverride);
+    // Gradle workers and the JUnit console need not expose test classes through
+    // java.class.path. These code sources also handle separate main/test outputs.
+    String classpath = new File(ReallocationFailureMain.class.getProtectionDomain()
+        .getCodeSource().getLocation().toURI()).getAbsolutePath() + File.pathSeparator
+        + new File(SaferAlloc.class.getProtectionDomain().getCodeSource()
+        .getLocation().toURI()).getAbsolutePath();
+    command.add("-cp");
+    command.add(classpath);
+    command.add(ReallocationFailureMain.class.getName());
+
+    Path output = Files.createTempFile("saferalloc-realloc-failure-", ".log");
+    Process child = null;
+    try {
+      child = new ProcessBuilder(command).redirectErrorStream(true).redirectOutput(output.toFile()).start();
+      assertTrue(child.waitFor(60, TimeUnit.SECONDS), "heap-exhaustion child JVM timed out");
+      String log = new String(Files.readAllBytes(output), StandardCharsets.UTF_8);
+      assertEquals(0, child.exitValue(), log);
+      assertTrue(log.contains("REALLOC_WRAPPER_FAILURE_PRESERVED_ORIGINAL"), log);
+    } finally {
+      if (child != null && child.isAlive()) {
+        child.destroyForcibly();
+        child.waitFor();
+      }
+      Files.deleteIfExists(output);
+    }
+  }
+}