From 2ce8c7dc187f34e5479376f6a4c7e4c2bde60836 Mon Sep 17 00:00:00 2001
From: Chris Vest <christianvest_hansen@apple.com>
Date: Mon, 7 Dec 2020 17:35:21 +0100
Subject: [PATCH] Fix drop race with Cleaner

Motivation:
We were seeing rare test failures where a cleaner had raced to close a memory segment we were using or closing.
The cause was that a single MemorySegment ended up used in multiple Buf instances.
When the SizeClassedMemoryPool was closed, the memory segments could be disposed without closing the gate in the NativeMemoryCleanerDrop.
The gate is important because it prevents double-frees of the memory segment.

Modification:
The fix is to change how the SizeClassedMemoryPool is closed, such that it always releases memory by calling `close()` on its buffers, which in turn will close the gate. The program will then proceed through the SizeClassedMemoryPool.drop implementation, which in turn will observe that the allocator is closed, and *then* dispose of the memory.

Result:
We should hopefully not see any more random test failures, but if we do, they would at least indicate a different bug.
This particular one was mostly showing up inside the cleaner threads, which were ignoring the exception, but occasionally I think the race went the other way, causing a test failure.
---
 .../java/io/netty/buffer/api/NativeMemoryCleanerDrop.java | 6 +++---
 .../java/io/netty/buffer/api/SizeClassedMemoryPool.java   | 8 ++++++--
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/src/main/java/io/netty/buffer/api/NativeMemoryCleanerDrop.java b/src/main/java/io/netty/buffer/api/NativeMemoryCleanerDrop.java
index 5e448d2..68c4c1f 100644
--- a/src/main/java/io/netty/buffer/api/NativeMemoryCleanerDrop.java
+++ b/src/main/java/io/netty/buffer/api/NativeMemoryCleanerDrop.java
@@ -61,9 +61,9 @@ class NativeMemoryCleanerDrop implements Drop<Buf> {
         var mem = manager.unwrapRecoverableMemory(buf);
         var delegate = this.delegate;
         WeakReference<Buf> ref = new WeakReference<>(buf);
-        AtomicBoolean gate = new AtomicBoolean();
+        AtomicBoolean gate = new AtomicBoolean(true);
         cleanable = new GatedCleanable(gate, CLEANER.register(this, () -> {
-            if (gate.compareAndSet(false, true)) {
+            if (gate.getAndSet(false)) {
                 Buf b = ref.get();
                 if (b == null) {
                     pool.recoverMemory(mem);
@@ -84,7 +84,7 @@ class NativeMemoryCleanerDrop implements Drop<Buf> {
         }
 
         public void disable() {
-            gate.set(true);
+            gate.set(false);
         }
 
         @Override
diff --git a/src/main/java/io/netty/buffer/api/SizeClassedMemoryPool.java b/src/main/java/io/netty/buffer/api/SizeClassedMemoryPool.java
index b015064..dc547e1 100644
--- a/src/main/java/io/netty/buffer/api/SizeClassedMemoryPool.java
+++ b/src/main/java/io/netty/buffer/api/SizeClassedMemoryPool.java
@@ -70,7 +70,7 @@ class SizeClassedMemoryPool implements Allocator, AllocatorControl, Drop<Buf> {
                 Send<Buf> send;
                 while ((send = v.poll()) != null) {
                     try {
-                        dispose(send.receive());
+                        send.receive().close();
                     } catch (Exception e) {
                         capturedExceptions.add(e);
                     }
@@ -86,12 +86,16 @@ class SizeClassedMemoryPool implements Allocator, AllocatorControl, Drop<Buf> {
 
     @Override
     public void drop(Buf buf) {
+        if (closed) {
+            dispose(buf);
+            return;
+        }
         var sizeClassPool = getSizeClassPool(buf.capacity());
         sizeClassPool.offer(buf.send());
         if (closed) {
             Send<Buf> send;
             while ((send = sizeClassPool.poll()) != null) {
-                dispose(send.receive());
+                send.receive().close();
             }
         }
     }