diff --git a/packages/netty/src/main/java/dev/elide/bemo/transport/NativeDirectByteBufAllocator.java b/packages/netty/src/main/java/dev/elide/bemo/transport/NativeDirectByteBufAllocator.java index 310fb0a..663dc94 100644 --- a/packages/netty/src/main/java/dev/elide/bemo/transport/NativeDirectByteBufAllocator.java +++ b/packages/netty/src/main/java/dev/elide/bemo/transport/NativeDirectByteBufAllocator.java @@ -15,6 +15,7 @@ import java.nio.ByteBuffer; import java.nio.ByteOrder; import java.util.concurrent.atomic.AtomicLong; +import org.jspecify.annotations.Nullable; /** * Direct buffers whose memory comes from Bemo's native allocator, the same allocator the transport @@ -129,24 +130,93 @@ public void clean() {} } private static final class UnsafeBuffer extends UnpooledUnsafeDirectByteBuf { + // The superclass treats the initial externally supplied buffer as borrowed and frees it here. + private @Nullable CleanableDirectBuffer initial; + UnsafeBuffer(NativeDirectByteBufAllocator allocator, int initialCapacity, int maxCapacity) { - super(allocator, initialCapacity, maxCapacity); + this(allocator, allocator.allocate(initialCapacity), maxCapacity); + } + + private UnsafeBuffer( + NativeDirectByteBufAllocator allocator, CleanableDirectBuffer allocation, int maxCapacity) { + super(allocator, allocation.buffer(), maxCapacity); + this.initial = allocation; + setIndex(0, 0); } @Override protected CleanableDirectBuffer allocateDirectBuffer(int capacity) { return ((NativeDirectByteBufAllocator) alloc()).allocate(capacity); } + + @Override + public ByteBuf capacity(int capacity) { + ensureAccessible(); + if (capacity == capacity()) return this; + super.capacity(capacity); + releaseInitial(); + return this; + } + + @Override + protected void deallocate() { + try { + super.deallocate(); + } finally { + releaseInitial(); + } + } + + private void releaseInitial() { + if (initial != null) { + initial.clean(); + initial = null; + } + } } private static final class SafeBuffer extends UnpooledDirectByteBuf { + private @Nullable CleanableDirectBuffer initial; + SafeBuffer(NativeDirectByteBufAllocator allocator, int initialCapacity, int maxCapacity) { - super(allocator, initialCapacity, maxCapacity); + this(allocator, allocator.allocate(initialCapacity), maxCapacity); + } + + private SafeBuffer( + NativeDirectByteBufAllocator allocator, CleanableDirectBuffer allocation, int maxCapacity) { + super(allocator, allocation.buffer(), maxCapacity); + this.initial = allocation; + setIndex(0, 0); } @Override protected CleanableDirectBuffer allocateDirectBuffer(int capacity) { return ((NativeDirectByteBufAllocator) alloc()).allocate(capacity); } + + @Override + public ByteBuf capacity(int capacity) { + ensureAccessible(); + if (capacity == capacity()) return this; + super.capacity(capacity); + releaseInitial(); + return this; + } + + @Override + protected void deallocate() { + try { + super.deallocate(); + } finally { + releaseInitial(); + } + } + + private void releaseInitial() { + if (initial != null) { + initial.clean(); + initial = null; + } + } } } diff --git a/tests/transport/java/NativeDirectAllocatorTest.java b/tests/transport/java/NativeDirectAllocatorTest.java index 50f3778..72326ad 100644 --- a/tests/transport/java/NativeDirectAllocatorTest.java +++ b/tests/transport/java/NativeDirectAllocatorTest.java @@ -40,6 +40,10 @@ public static void verify(TransportNative api) throws Exception { ByteBuf buffer = allocator.directBuffer(8, 1 << 20); if (!buffer.isDirect()) throw new AssertionError("expected a direct buffer"); String kind = buffer.getClass().getSimpleName() + (buffer.hasMemoryAddress() ? "/address" : ""); + // A freshly allocated buffer (before any growth) must already be backed by Bemo's native + // allocator; the initial allocation may not silently fall back to Netty's direct path. + if (allocator.nativeAllocatedBytes() < baseline + buffer.capacity()) + throw new AssertionError("fresh direct buffer is not backed by the Bemo native allocator"); // Write across several reallocations, checking the whole prefix after each one. int capacity = buffer.capacity(); for (int index = 0; index < 4096; index++) { @@ -71,6 +75,11 @@ public static void verify(TransportNative api) throws Exception { List buffers = new ArrayList<>(); for (int index = 0; index < 256; index++) buffers.add(allocator.directBuffer(1024).writeLong(index)); + // These buffers are never grown, so their storage must come from Bemo's native allocator for + // the whole lifetime (the common request/response case); no part of them is Netty-owned. + if (allocator.nativeAllocatedBytes() < baseline + (long) buffers.size() * 1024) + throw new AssertionError( + "never-grown direct buffers are not backed by the Bemo native allocator"); var failure = new AtomicReference(); var done = new CountDownLatch(1); Thread release =