CORE-50: prototype native-memory tracking callbacks

src/main/java/net/openhft/chronicle/core/UnsafeMemory.java

@@ -5,6 +5,9 @@
 
 import net.openhft.chronicle.core.internal.Bootstrap;
 import net.openhft.chronicle.core.internal.util.DirectBufferUtil;
+import net.openhft.chronicle.core.alloc.AllocationListener;
+import net.openhft.chronicle.core.alloc.AllocationTrace;
+import net.openhft.chronicle.core.alloc.NoOpAllocationListener;
 import net.openhft.chronicle.core.util.MisAlignedAssertionError;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -54,6 +57,11 @@ public class UnsafeMemory implements Memory {
     public static final UnsafeMemory INSTANCE;
     public static final UnsafeMemory MEMORY;
 
+    private static volatile AllocationListener allocationListener = NoOpAllocationListener.INSTANCE;
+
+    private static final ThreadLocal<AllocationTrace> TRACE_POOL =
+        ThreadLocal.withInitial(AllocationTrace::new);
+
     // see java.nio.Bits.copyMemory
     // This number limits the number of bytes to copy per call to Unsafe's
     // copyMemory method. A limit is imposed to allow for safepoint polling
@@ -629,8 +637,17 @@ public void setMemory(Object o, long offset, long size, byte b) {
     @Override
     public void freeMemory(long address, long size) {
         assert SKIP_ASSERTIONS || size >= 0;
-        if (address != 0)
+        if (address != 0) {
+            AllocationListener listener = allocationListener;
+            if (listener != NoOpAllocationListener.INSTANCE) {
+                long timestampMs = System.currentTimeMillis();
+                AllocationTrace trace = TRACE_POOL.get();
+                trace.capture();
+                try { listener.onFree(timestampMs, address, size, trace); }
+                catch (Throwable t) { Jvm.warn().on(UnsafeMemory.class, "AllocationListener.onFree threw", t); }
+            }
             UNSAFE.freeMemory(address);
+        }
         nativeMemoryUsed.addAndGet(-size);
     }
 
@@ -651,7 +668,14 @@ public long allocate(long capacity) {
             throw new OutOfMemoryError("Not enough free native memory, capacity attempted: " + capacity / 1024 + " KiB");
 
         nativeMemoryUsed.addAndGet(capacity);
-
+        AllocationListener listener = allocationListener;
+        if (listener != NoOpAllocationListener.INSTANCE) {
+            long timestampMs = System.currentTimeMillis();
+            AllocationTrace trace = TRACE_POOL.get();
+            trace.capture();
+            try { listener.onAllocate(timestampMs, address, capacity, trace); }
+            catch (Throwable t) { Jvm.warn().on(UnsafeMemory.class, "AllocationListener.onAllocate threw", t); }
+        }
         return address;
     }
 
@@ -665,6 +689,20 @@ public long nativeMemoryUsed() {
         return nativeMemoryUsed.get();
     }
 
+    /**
+     * Installs an allocation listener. Pass {@code null} to reset to the no-op default.
+     */
+    public static void setAllocationListener(AllocationListener listener) {
+        allocationListener = (listener != null) ? listener : NoOpAllocationListener.INSTANCE;
+    }
+
+    /**
+     * Returns the currently installed allocation listener.
+     */
+    public static AllocationListener getAllocationListener() {
+        return allocationListener;
+    }
+
     /**
      * Writes a byte to the given memory address.
      *

src/main/java/net/openhft/chronicle/core/alloc/AllocationListener.java

@@ -0,0 +1,34 @@
+/*
+ * Copyright 2013-2025 chronicle.software; SPDX-License-Identifier: Apache-2.0
+ */
+package net.openhft.chronicle.core.alloc;
+
+/**
+ * Callback interface for native memory allocation events.
+ * <p>
+ * Implementations receive notifications when memory is allocated or freed
+ * via {@link net.openhft.chronicle.core.UnsafeMemory}. Also serves as the
+ * Chronicle Queue MethodWriter/MethodReader contract for the NMT agent.
+ */
+public interface AllocationListener {
+
+    /**
+     * Called after a native memory allocation succeeds.
+     *
+     * @param timestampMs millisecond wall-clock timestamp
+     * @param address     the native memory address returned by allocateMemory
+     * @param capacity    the number of bytes requested
+     * @param trace       pooled stack trace — must be read/copied before returning
+     */
+    void onAllocate(long timestampMs, long address, long capacity, AllocationTrace trace);
+
+    /**
+     * Called before native memory is freed.
+     *
+     * @param timestampMs millisecond wall-clock timestamp
+     * @param address     the native memory address about to be freed
+     * @param size        the number of bytes being freed
+     * @param trace       pooled stack trace — must be read/copied before returning
+     */
+    void onFree(long timestampMs, long address, long size, AllocationTrace trace);
+}

src/main/java/net/openhft/chronicle/core/alloc/AllocationTrace.java

@@ -0,0 +1,35 @@
+/*
+ * Copyright 2013-2025 chronicle.software; SPDX-License-Identifier: Apache-2.0
+ */
+package net.openhft.chronicle.core.alloc;
+
+/**
+ * Pooled, reusable stack-trace container. One instance per thread via ThreadLocal.
+ * <p>
+ * Listeners must read or copy data before returning — the object is recycled
+ * on the next allocation from the same thread.
+ * <p>
+ * Plain class with no Wire dependency; Chronicle Wire serialises it via
+ * reflection-based field access.
+ */
+public class AllocationTrace {
+    private StackTraceElement[] elements;
+    private int depth;
+
+    /**
+     * Captures the current stack trace, reusing the backing array where possible.
+     */
+    public void capture() {
+        Throwable t = new Throwable();
+        StackTraceElement[] raw = t.getStackTrace();
+        if (elements == null || elements.length < raw.length) {
+            elements = new StackTraceElement[raw.length];
+        }
+        System.arraycopy(raw, 0, elements, 0, raw.length);
+        depth = raw.length;
+    }
+
+    public StackTraceElement[] elements() { return elements; }
+
+    public int depth() { return depth; }
+}

src/main/java/net/openhft/chronicle/core/alloc/NoOpAllocationListener.java

@@ -0,0 +1,23 @@
+/*
+ * Copyright 2013-2025 chronicle.software; SPDX-License-Identifier: Apache-2.0
+ */
+package net.openhft.chronicle.core.alloc;
+
+/**
+ * No-op singleton used as the default listener.
+ * Identity-checked by {@link net.openhft.chronicle.core.UnsafeMemory} to skip
+ * ThreadLocal access and stack capture on the hot path.
+ */
+public enum NoOpAllocationListener implements AllocationListener {
+    INSTANCE;
+
+    @Override
+    public void onAllocate(long timestampMs, long address, long capacity, AllocationTrace trace) {
+        // no-op
+    }
+
+    @Override
+    public void onFree(long timestampMs, long address, long size, AllocationTrace trace) {
+        // no-op
+    }
+}

src/test/java/net/openhft/chronicle/core/alloc/AllocationListenerTest.java

@@ -0,0 +1,177 @@
+/*
+ * Copyright 2013-2025 chronicle.software; SPDX-License-Identifier: Apache-2.0
+ */
+package net.openhft.chronicle.core.alloc;
+
+import net.openhft.chronicle.core.UnsafeMemory;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+class AllocationListenerTest {
+
+    @AfterEach
+    void resetListener() {
+        UnsafeMemory.setAllocationListener(null);
+    }
+
+    @Test
+    void defaultListenerIsNoOp() {
+        assertSame(NoOpAllocationListener.INSTANCE, UnsafeMemory.getAllocationListener());
+    }
+
+    @Test
+    void setListenerInstallsAndRemoves() {
+        AllocationListener custom = new RecordingListener();
+        UnsafeMemory.setAllocationListener(custom);
+        assertSame(custom, UnsafeMemory.getAllocationListener());
+
+        UnsafeMemory.setAllocationListener(null);
+        assertSame(NoOpAllocationListener.INSTANCE, UnsafeMemory.getAllocationListener());
+    }
+
+    @Test
+    void onAllocateCalledWithCorrectParameters() {
+        RecordingListener listener = new RecordingListener();
+        UnsafeMemory.setAllocationListener(listener);
+
+        long beforeMs = System.currentTimeMillis();
+        long address = UnsafeMemory.INSTANCE.allocate(128);
+        long afterMs = System.currentTimeMillis();
+
+        try {
+            assertEquals(1, listener.allocations.size());
+            RecordedEvent event = listener.allocations.get(0);
+            assertTrue(event.timestampMs >= beforeMs && event.timestampMs <= afterMs,
+                    "timestampMs should be between before and after");
+            assertEquals(address, event.address);
+            assertEquals(128, event.size);
+            assertNotNull(event.traceElements);
+            assertTrue(event.traceDepth > 0, "trace depth should be > 0");
+        } finally {
+            UnsafeMemory.INSTANCE.freeMemory(address, 128);
+        }
+    }
+
+    @Test
+    void onFreeCalledWithCorrectParameters() {
+        RecordingListener listener = new RecordingListener();
+        long address = UnsafeMemory.INSTANCE.allocate(256);
+
+        // Install listener after allocate so we only capture the free
+        UnsafeMemory.setAllocationListener(listener);
+        long beforeMs = System.currentTimeMillis();
+        UnsafeMemory.INSTANCE.freeMemory(address, 256);
+        long afterMs = System.currentTimeMillis();
+
+        assertEquals(1, listener.frees.size());
+        RecordedEvent event = listener.frees.get(0);
+        assertTrue(event.timestampMs >= beforeMs && event.timestampMs <= afterMs,
+                "timestampMs should be between before and after");
+        assertEquals(address, event.address);
+        assertEquals(256, event.size);
+        assertNotNull(event.traceElements);
+        assertTrue(event.traceDepth > 0, "trace depth should be > 0");
+    }
+
+    @Test
+    void onFreeNotCalledWhenAddressIsZero() {
+        RecordingListener listener = new RecordingListener();
+        UnsafeMemory.setAllocationListener(listener);
+
+        UnsafeMemory.INSTANCE.freeMemory(0, 64);
+
+        assertTrue(listener.frees.isEmpty(), "onFree should not be called for address 0");
+    }
+
+    @Test
+    void listenerExceptionIsCaughtAndAllocationSucceeds() {
+        AllocationListener throwing = new AllocationListener() {
+            @Override
+            public void onAllocate(long timestampMs, long address, long capacity, AllocationTrace trace) {
+                throw new RuntimeException("test exception");
+            }
+
+            @Override
+            public void onFree(long timestampMs, long address, long size, AllocationTrace trace) {
+                throw new RuntimeException("test exception");
+            }
+        };
+        UnsafeMemory.setAllocationListener(throwing);
+
+        // Should not throw — exception is caught and logged
+        long address = UnsafeMemory.INSTANCE.allocate(64);
+        assertTrue(address != 0, "allocation should succeed despite listener exception");
+
+        // Free should also not throw
+        assertDoesNotThrow(() -> UnsafeMemory.INSTANCE.freeMemory(address, 64));
+    }
+
+    @Test
+    void allocationTraceIsReusedAcrossCallsOnSameThread() {
+        List<AllocationTrace> traces = new ArrayList<>();
+        AllocationListener capturing = new AllocationListener() {
+            @Override
+            public void onAllocate(long timestampMs, long address, long capacity, AllocationTrace trace) {
+                traces.add(trace);
+            }
+
+            @Override
+            public void onFree(long timestampMs, long address, long size, AllocationTrace trace) {
+                traces.add(trace);
+            }
+        };
+        UnsafeMemory.setAllocationListener(capturing);
+
+        long addr1 = UnsafeMemory.INSTANCE.allocate(32);
+        long addr2 = UnsafeMemory.INSTANCE.allocate(32);
+        UnsafeMemory.INSTANCE.freeMemory(addr1, 32);
+        UnsafeMemory.INSTANCE.freeMemory(addr2, 32);
+
+        assertEquals(4, traces.size());
+        // All traces should be the same object (ThreadLocal reuse)
+        AllocationTrace first = traces.get(0);
+        for (AllocationTrace t : traces) {
+            assertSame(first, t, "AllocationTrace should be reused on the same thread");
+        }
+    }
+
+    // ---- helpers ----
+
+    private static class RecordedEvent {
+        final long timestampMs;
+        final long address;
+        final long size;
+        final StackTraceElement[] traceElements;
+        final int traceDepth;
+
+        RecordedEvent(long timestampMs, long address, long size, AllocationTrace trace) {
+            this.timestampMs = timestampMs;
+            this.address = address;
+            this.size = size;
+            // Copy since trace is recycled
+            this.traceDepth = trace.depth();
+            this.traceElements = new StackTraceElement[trace.depth()];
+            System.arraycopy(trace.elements(), 0, this.traceElements, 0, trace.depth());
+        }
+    }
+
+    private static class RecordingListener implements AllocationListener {
+        final List<RecordedEvent> allocations = new ArrayList<>();
+        final List<RecordedEvent> frees = new ArrayList<>();
+
+        @Override
+        public void onAllocate(long timestampMs, long address, long capacity, AllocationTrace trace) {
+            allocations.add(new RecordedEvent(timestampMs, address, capacity, trace));
+        }
+
+        @Override
+        public void onFree(long timestampMs, long address, long size, AllocationTrace trace) {
+            frees.add(new RecordedEvent(timestampMs, address, size, trace));
+        }
+    }
+}