Allow slices to obtain ownership when parent is closed

Motivation:
It is kind of a weird internal and hidden state, that slices were special.
For instance, slices could not be sent, and they could never obtain ownership.
This means buffers from slices behaved differently from allocated buffers.
In doing so, they violated both the principle that magic should stay hidden, and the principle of consistent behaviour.

Modification:
- The special reference-counting drop implementation that was added to support bifurcation, has been renamed to ArcDrop (for atomic reference counting).
- The ArcDrop is then used throughout the MemSegBuffer implementation to account for every instance where multiple buffers reference the same memory, e.g. slices and the like.
- Borrows of a buffer is then the sum of borrows from the buffer itself, and its ArcDrop.
- Ownership is thus tied to both the buffer itself being owned, and the ArcDrop being in an owned state.
- SizeClassedMemoryPool is changed to pool recoverable memory instead of sends, because the sends could come from slices.
- We also take care to keep around a "base" memory segment, so that we don't return memory segment slices to the memory pool (doing so would leak the memory from the parent segment that is not part of the slice).
- CleanerPooledDrop now keeps a weak reference to itself, rather than the buffer, which is more correct anyway, but now also required because we cannot rely on the buffer reference the cleaner was created with.
- The CleanerPooledDrop now takes care to drop the buffer that is actually passed to it, rather than what it was referencing from some earlier point.
- MemoryManager can now disclose the size of recoverable memory, so that SizeClassedMemoryPool can pick the correct size pool to return memory to. It cannot rely on the passed down buffer instance for this, because that buffer might have been a slice.

Result:
It is now possible for slices to obtain ownership when their parent buffer is closed.

Author: chrisvest

src/main/java/io/netty/buffer/api/CleanerPooledDrop.java

@@ -45,6 +45,7 @@ public void drop(Buffer buf) {
         GatedCleanable c = (GatedCleanable) CLEANABLE.getAndSet(this, null);
         if (c != null) {
             c.clean();
+            delegate.drop(buf);
         }
     }
 
@@ -57,24 +58,38 @@ public void attach(Buffer buf) {
             c.clean();
         }
 
-        var pool = this.pool;
         var mem = manager.unwrapRecoverableMemory(buf);
-        var delegate = this.delegate;
-        WeakReference<Buffer> ref = new WeakReference<>(buf);
+        WeakReference<CleanerPooledDrop> ref = new WeakReference<>(this);
         AtomicBoolean gate = new AtomicBoolean(true);
-        cleanable = new GatedCleanable(gate, CLEANER.register(this, () -> {
+        cleanable = new GatedCleanable(gate, CLEANER.register(this, new CleanAction(pool, mem, ref, gate)));
+    }
+
+    private static final class CleanAction implements Runnable {
+        private final SizeClassedMemoryPool pool;
+        private final Object mem;
+        private final WeakReference<CleanerPooledDrop> ref;
+        private final AtomicBoolean gate;
+
+        private CleanAction(SizeClassedMemoryPool pool, Object mem, WeakReference<CleanerPooledDrop> ref,
+                            AtomicBoolean gate) {
+            this.pool = pool;
+            this.mem = mem;
+            this.ref = ref;
+            this.gate = gate;
+        }
+
+        @Override
+        public void run() {
             if (gate.getAndSet(false)) {
-                Buffer b = ref.get();
-                if (b == null) {
+                var monitored = ref.get();
+                if (monitored == null) {
                     pool.recoverMemory(mem);
-                } else {
-                    delegate.drop(b);
                 }
             }
-        }));
+        }
     }
 
-    private static class GatedCleanable implements Cleanable {
+    private static final class GatedCleanable implements Cleanable {
         private final AtomicBoolean gate;
         private final Cleanable cleanable;
 

src/main/java/io/netty/buffer/api/MemoryManager.java

@@ -34,5 +34,6 @@ static MemoryManager getNativeMemoryManager() {
     Buffer allocateShared(AllocatorControl allo, long size, Drop<Buffer> drop, Cleaner cleaner);
     Drop<Buffer> drop();
     Object unwrapRecoverableMemory(Buffer buf);
+    int capacityOfRecoverableMemory(Object memory);
     Buffer recoverMemory(Object recoverableMemory, Drop<Buffer> drop);
 }

src/main/java/io/netty/buffer/api/RcSupport.java

@@ -37,7 +37,7 @@ protected RcSupport(Drop<T> drop) {
     @Override
     public final I acquire() {
         if (acquires < 0) {
-            throw attachTrace(new IllegalStateException("Resource is closed."));
+            throw attachTrace(new IllegalStateException("This resource is closed: " + this + '.'));
         }
         if (acquires == Integer.MAX_VALUE) {
             throw new IllegalStateException("Cannot acquire more references; counter would overflow.");

src/main/java/io/netty/buffer/api/SizeClassedMemoryPool.java

@@ -28,7 +28,7 @@ class SizeClassedMemoryPool implements BufferAllocator, AllocatorControl, Drop<B
     private static final VarHandle CLOSE = Statics.findVarHandle(
             lookup(), SizeClassedMemoryPool.class, "closed", boolean.class);
     private final MemoryManager manager;
-    private final ConcurrentHashMap<Integer, ConcurrentLinkedQueue<Send<Buffer>>> pool;
+    private final ConcurrentHashMap<Integer, ConcurrentLinkedQueue<Object>> pool;
     @SuppressWarnings("unused")
     private volatile boolean closed;
 
@@ -41,9 +41,9 @@ protected SizeClassedMemoryPool(MemoryManager manager) {
     public Buffer allocate(int size) {
         BufferAllocator.checkSize(size);
         var sizeClassPool = getSizeClassPool(size);
-        Send<Buffer> send = sizeClassPool.poll();
-        if (send != null) {
-            return send.receive()
+        Object memory = sizeClassPool.poll();
+        if (memory != null) {
+            return recoverMemoryIntoBuffer(memory)
                        .reset()
                        .readOnly(false)
                        .fill((byte) 0)
@@ -71,10 +71,10 @@ public void close() {
         if (CLOSE.compareAndSet(this, false, true)) {
             var capturedExceptions = new ArrayList<Exception>(4);
             pool.forEach((k, v) -> {
-                Send<Buffer> send;
-                while ((send = v.poll()) != null) {
+                Object memory;
+                while ((memory = v.poll()) != null) {
                     try {
-                        send.receive().close();
+                        dispose(recoverMemoryIntoBuffer(memory));
                     } catch (Exception e) {
                         capturedExceptions.add(e);
                     }
@@ -94,40 +94,42 @@ public void drop(Buffer buf) {
             dispose(buf);
             return;
         }
-        var sizeClassPool = getSizeClassPool(buf.capacity());
-        sizeClassPool.offer(buf.send());
+        Object mem = manager.unwrapRecoverableMemory(buf);
+        var sizeClassPool = getSizeClassPool(manager.capacityOfRecoverableMemory(mem));
+        sizeClassPool.offer(mem);
         if (closed) {
-            Send<Buffer> send;
-            while ((send = sizeClassPool.poll()) != null) {
-                send.receive().close();
+            Object memory;
+            while ((memory = sizeClassPool.poll()) != null) {
+                dispose(recoverMemoryIntoBuffer(memory));
             }
         }
     }
 
     @Override
     public Object allocateUntethered(Buffer originator, int size) {
         var sizeClassPool = getSizeClassPool(size);
-        Send<Buffer> send = sizeClassPool.poll();
-        Buffer untetheredBuf;
-        if (send != null) {
-            var transfer = (TransferSend<Buffer, Buffer>) send;
-            var owned = transfer.unsafeUnwrapOwned();
-            untetheredBuf = owned.transferOwnership(NO_OP_DROP);
-        } else {
-            untetheredBuf = createBuf(size, NO_OP_DROP);
+        Object memory = sizeClassPool.poll();
+        if (memory == null) {
+            Buffer untetheredBuf = createBuf(size, NO_OP_DROP);
+            memory = manager.unwrapRecoverableMemory(untetheredBuf);
         }
-        return manager.unwrapRecoverableMemory(untetheredBuf);
+        return memory;
     }
 
     @Override
     public void recoverMemory(Object memory) {
+        Buffer buf = recoverMemoryIntoBuffer(memory);
+        buf.close();
+    }
+
+    private Buffer recoverMemoryIntoBuffer(Object memory) {
         var drop = getDrop();
         var buf = manager.recoverMemory(memory, drop);
         drop.attach(buf);
-        buf.close();
+        return buf;
     }
 
-    private ConcurrentLinkedQueue<Send<Buffer>> getSizeClassPool(int size) {
+    private ConcurrentLinkedQueue<Object> getSizeClassPool(int size) {
         return pool.computeIfAbsent(size, k -> new ConcurrentLinkedQueue<>());
     }
 

src/main/java/io/netty/buffer/api/memseg/AbstractMemorySegmentManager.java

@@ -34,7 +34,7 @@ public Buffer allocateConfined(AllocatorControl alloc, long size, Drop<Buffer> d
         if (cleaner != null) {
             segment = segment.registerCleaner(cleaner);
         }
-        return new MemSegBuffer(segment, convert(drop), alloc);
+        return new MemSegBuffer(segment, segment, convert(drop), alloc);
     }
 
     @Override
@@ -43,7 +43,7 @@ public Buffer allocateShared(AllocatorControl alloc, long size, Drop<Buffer> dro
         if (cleaner != null) {
             segment = segment.registerCleaner(cleaner);
         }
-        return new MemSegBuffer(segment, convert(drop), alloc);
+        return new MemSegBuffer(segment, segment, convert(drop), alloc);
     }
 
     protected abstract MemorySegment createSegment(long size);
@@ -59,6 +59,11 @@ public Object unwrapRecoverableMemory(Buffer buf) {
         return b.recoverableMemory();
     }
 
+    @Override
+    public int capacityOfRecoverableMemory(Object memory) {
+        return ((RecoverableMemory) memory).capacity();
+    }
+
     @Override
     public Buffer recoverMemory(Object recoverableMemory, Drop<Buffer> drop) {
         var recovery = (RecoverableMemory) recoverableMemory;

src/main/java/io/netty/buffer/api/memseg/BifurcatedDrop.java renamed to src/main/java/io/netty/buffer/api/memseg/ArcDrop.java

@@ -20,33 +20,47 @@
 import java.lang.invoke.MethodHandles;
 import java.lang.invoke.VarHandle;
 
-class BifurcatedDrop implements Drop<MemSegBuffer> {
+final class ArcDrop implements Drop<MemSegBuffer> {
     private static final VarHandle COUNT;
     static {
         try {
-            COUNT = MethodHandles.lookup().findVarHandle(BifurcatedDrop.class, "count", int.class);
+            COUNT = MethodHandles.lookup().findVarHandle(ArcDrop.class, "count", int.class);
         } catch (Exception e) {
             throw new ExceptionInInitializerError(e);
         }
     }
 
-    private final MemSegBuffer originalBuf;
     private final Drop<MemSegBuffer> delegate;
     @SuppressWarnings("FieldMayBeFinal")
     private volatile int count;
 
-    BifurcatedDrop(MemSegBuffer originalBuf, Drop<MemSegBuffer> delegate) {
-        this.originalBuf = originalBuf;
+    ArcDrop(Drop<MemSegBuffer> delegate) {
         this.delegate = delegate;
-        count = 2; // These are created by buffer bifurcation, so we initially have 2 references to this drop.
+        count = 1;
     }
 
-    void increment() {
+    static Drop<MemSegBuffer> wrap(Drop<MemSegBuffer> drop) {
+        if (drop.getClass() == ArcDrop.class) {
+            return drop;
+        }
+        return new ArcDrop(drop);
+    }
+
+    static Drop<MemSegBuffer> acquire(Drop<MemSegBuffer> drop) {
+        if (drop.getClass() == ArcDrop.class) {
+            ((ArcDrop) drop).increment();
+            return drop;
+        }
+        return new ArcDrop(drop);
+    }
+
+    ArcDrop increment() {
         int c;
         do {
             c = count;
             checkValidState(c);
         } while (!COUNT.compareAndSet(this, c, c + 1));
+        return this;
     }
 
     @Override
@@ -59,17 +73,23 @@ public void drop(MemSegBuffer buf) {
             checkValidState(c);
         } while (!COUNT.compareAndSet(this, c, n));
         if (n == 0) {
-            delegate.attach(originalBuf);
-            delegate.drop(originalBuf);
+            delegate.drop(buf);
         }
-        buf.makeInaccessible();
     }
 
     @Override
     public void attach(MemSegBuffer obj) {
         delegate.attach(obj);
     }
 
+    boolean isOwned() {
+        return count <= 1;
+    }
+
+    int countBorrows() {
+        return count - 1;
+    }
+
     Drop<MemSegBuffer> unwrap() {
         return delegate;
     }