Repository: netty/netty
File: buffer/src/test/java/io/netty5/buffer/api/tests/SensitiveBufferTest.java (Java)

Code before:
        try (Buffer buffer = MemoryManager.using(stubManager, () -> sensitiveOffHeapAllocator().allocate(8))) {
            buffer.writeLong(0x0102030405060708L).makeReadOnly();
            try (Buffer split = buffer.readSplit(4)) {
                assertTrue(split.readOnly());
                split.copy(true).close();
            }
            final Send<Buffer> send;
            try (Buffer copy = buffer.copy(true)) {
                send = buffer.send();
                copy.readSplit(2).close();
            }
            assertEquals(0, stubManager.getBytesCleared());
            send.close();
            assertEquals(8, stubManager.getBytesCleared());
        }
    }

    private static final class StubManager implements MemoryManager {
        private final MemoryManager baseMemoryManager;
        private final AtomicInteger bytesCleared = new AtomicInteger();

        StubManager(MemoryManager baseMemoryManager) {
            this.baseMemoryManager = baseMemoryManager;
        }

        @Override
        public Buffer allocateShared(AllocatorControl control, long size,
                                     Function<Drop<Buffer>, Drop<Buffer>> dropDecorator,
                                     AllocationType allocationType) {
            return baseMemoryManager.allocateShared(
                    control,
                    size,
                    drop -> dropDecorator.apply(new CheckingDrop(bytesCleared, drop)),
                    allocationType);
        }

        @Override
        public Buffer allocateConstChild(Buffer readOnlyConstParent) {
            throw new UnsupportedOperationException();
        }

        @Override
        public Object unwrapRecoverableMemory(Buffer buf) {
            return baseMemoryManager.unwrapRecoverableMemory(buf);
        }

        @Override
        public Buffer recoverMemory(AllocatorControl control, Object recoverableMemory, Drop<Buffer> drop) {
            return baseMemoryManager.recoverMemory(control, recoverableMemory, drop);
        }


Review comment (suggestion):
Shouldn’t we close the allocator once done (consider using try with resources) ? Same comment everywhere

Code after:
    public void closingBuffersWithComplicatedStructuralSharingMustNotZeroBytesUntilAllPartsAreClosed() {
        MemoryManager baseMemoryManager = MemoryManager.instance();
        StubManager stubManager = new StubManager(baseMemoryManager);
        try (Buffer buffer = MemoryManager.using(stubManager, () -> sensitiveOffHeapAllocator().allocate(8))) {
            buffer.writeLong(0x0102030405060708L).makeReadOnly();
            try (Buffer split = buffer.readSplit(4)) {
                assertTrue(split.readOnly());
                split.copy(true).close();
            }
            final Send<Buffer> send;
            try (Buffer copy = buffer.copy(true)) {
                send = buffer.send();
                copy.readSplit(2).close();
            }
            assertEquals(0, stubManager.getBytesCleared());
            send.close();
            assertEquals(8, stubManager.getBytesCleared());
        }
    }

    @Test
    public void sensitiveAllocatorMustCaptureMemoryManager() {
        MemoryManager baseMemoryManager = MemoryManager.instance();
        StubManager stubManager = new StubManager(baseMemoryManager);

        try (BufferAllocator allocator = MemoryManager.using(stubManager, () -> sensitiveOffHeapAllocator())) {
            allocator.allocate(8).close();
            assertThat(stubManager.getBytesCleared()).isEqualTo(8);
        }

        try (BufferAllocator allocator = sensitiveOffHeapAllocator()) {
            allocator.allocate(8).close();
            assertThat(stubManager.getBytesCleared()).isEqualTo(8); // No change, stub not captured here.
        }

        try (BufferAllocator allocator = sensitiveOffHeapAllocator()) {
            MemoryManager.using(stubManager, () -> allocator.allocate(8)).close();
            assertThat(stubManager.getBytesCleared()).isEqualTo(16); // Stub captured on allocation, as override.
        }
    }

    private static final class StubManager implements MemoryManager {
        private final MemoryManager baseMemoryManager;
        private final AtomicInteger bytesCleared = new AtomicInteger();

        StubManager(MemoryManager baseMemoryManager) {
            this.baseMemoryManager = baseMemoryManager;
        }

        @Override
        public Buffer allocateShared(AllocatorControl control, long size,