Stop the buffer pool pruner thread when the pool is empty

PROBLEM

The driver starts a thread that removes idle buffers from
PowerOfTwoBufferPool.DEFAULT. That thread never stops. A thread that
runs forever keeps the class loader of all driver classes in memory. The
static data of those classes also stays in memory. Then an application
server cannot unload an application, and the memory of that application
stays in use. Users report this behavior in GitHub issue 2029 and in
JAVA-5643.

CAUSE

DEFAULT is a static field, and it calls enablePruning() during class
initialization. That method schedules a periodic task. The executor
starts the worker thread for the first task, but the pool is empty at
that time. Therefore the thread has no work, and it continues to wake up
forever.

The thread keeps the class loader in memory because the JVM captures the
class loader when it constructs the thread. The data that the thread
holds is not the cause. For this reason, a change to the references of
the thread cannot release the class loader. The thread must stop.

SOLUTION

enablePruning() now sets a flag. It does not schedule a task. The
release() method schedules one prune when it puts a buffer into the
pool. Each prune schedules the next prune, but only if the pool still
holds a buffer. The pruner does not schedule the next prune when the
pool becomes empty.

The pruner also uses these settings on its executor:

  - allowCoreThreadTimeOut(true), so that the worker thread can stop
  - a keep-alive time of maxIdleTime / 2
  - setRemoveOnCancelPolicy(true)

The work queue becomes empty after the last prune. Then the keep-alive
time expires, and the worker thread stops. A later call to release()
schedules a new prune, and the executor starts a new thread.

Two threads must not schedule a prune at the same time. A prune must
also not stop while a different thread adds a buffer to the pool. The
AtomicBoolean pruningScheduled prevents both conditions. The prune
clears the flag, and then it examines the pool one more time before it
stops.

DRAWBACKS

The pool releases the class loader about 90 seconds after the last
buffer release. The default value of maxIdleTime is one minute. A buffer
is old enough to remove only after two prunes, and the keep-alive time
adds 30 seconds. The class loader stays in memory during that period. A
tool that examines threads at the moment of an undeployment can still
find a live thread.

A pool that becomes idle and then busy starts a new thread. This adds a
small cost. A test measures this cost. A busy pool keeps one thread,
because new work arrives before the keep-alive time expires.

prune() keeps its current behavior after an error. It writes a log
message and throws the error again, and the pruner does not start again.
This behavior is the same as before this change.

PRIOR ART

Netty has the same problem and uses the same solution.
GlobalEventExecutor is a single-thread singleton. It starts its thread
when work arrives, and it stops the thread when the task queue stays
empty for a quiet period. The deprecated ThreadDeathWatcher class uses
the same pattern. The steps that this change uses to clear and then
examine the flag follow GlobalEventExecutor.TaskRunner.

Netty also sets the context class loader of a new thread to null. See
netty#7290 and JDK-7008595. That change corrects a different problem,
which is a driver thread that keeps an application class loader in
memory. This commit does not include that change.

TESTS

New tests in PowerOfTwoBufferPoolTest show three results. An empty pool
starts no thread. The thread stops after the pruner empties the pool.
The pruner starts again after the thread stops.

A separate test harness measures class loader retention. That harness
loads the driver into a child class loader, opens a MongoClient, closes
it, and then waits for the class loader to become unreachable. Before
this change, the class loader stayed in memory. After this change, the
JVM collects it. The harness is a local development tool, and it is not
part of this commit.

JAVA-6279

Author: rozza

driver-core/src/main/com/mongodb/internal/connection/PowerOfTwoBufferPool.java

@@ -28,9 +28,10 @@
 import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.ConcurrentLinkedDeque;
-import java.util.concurrent.Executors;
-import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.RejectedExecutionException;
+import java.util.concurrent.ScheduledThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 /**
  * <p>This class is not part of the public API and may be removed or changed at any time</p>
@@ -40,6 +41,13 @@ public class PowerOfTwoBufferPool implements BufferProvider {
 
     /**
      * The global default pool.  Pruning is enabled on this pool. Idle buffers are pruned after one minute.
+     *
+     * <p>The pruner thread does not run all the time. It starts when the pool holds a buffer. It stops when the pool
+     * becomes empty.</p>
+     *
+     * <p>The pruner thread must stop. A thread that runs forever keeps the class loader of all driver classes in
+     * memory. The static data of those classes also stays in memory. Then an application server cannot unload the
+     * application. See <a href="https://jira.mongodb.org/browse/JAVA-6279">JAVA-6279</a>.</p>
      */
     public static final PowerOfTwoBufferPool DEFAULT = new PowerOfTwoBufferPool().enablePruning();
 
@@ -63,7 +71,13 @@ public ByteBuffer getBuffer() {
 
     private final Map<Integer, BufferPool> powerOfTwoToPoolMap = new HashMap<>();
     private final long maxIdleTimeNanos;
-    private final ScheduledExecutorService pruner;
+    private final ScheduledThreadPoolExecutor pruner;
+    /**
+     * True if the pruner has a scheduled prune. Two threads must not schedule a prune at the same time, and this flag
+     * prevents that. The method {@link #pruneAndRescheduleIfNeeded()} also uses this flag when it stops the pruner.
+     */
+    private final AtomicBoolean pruningScheduled = new AtomicBoolean();
+    private volatile boolean pruningEnabled;
 
     /**
      * Construct an instance with a highest power of two of 24.
@@ -96,21 +110,51 @@ public ByteBuffer getBuffer() {
             powerOfTwo = powerOfTwo << 1;
         }
         maxIdleTimeNanos = timeUnit.toNanos(maxIdleTime);
-        pruner = Executors.newSingleThreadScheduledExecutor(new DaemonThreadFactory("BufferPoolPruner"));
+        pruner = new ScheduledThreadPoolExecutor(1, new DaemonThreadFactory("BufferPoolPruner"));
+        // The worker thread must stop when it has no more work. Then an idle pool holds no thread.
+        //
+        // These three settings are sufficient only because this class schedules one prune at a time. It schedules the
+        // next prune only if the pool is not empty. Then the work queue becomes empty and the keep-alive time expires.
+        // A periodic task stays in the work queue forever. Then the worker thread always has a task to wait for, and
+        // the keep-alive time never expires.
+        //
+        // The keep-alive time applies only after the last prune. While a prune is in the work queue, the worker thread
+        // waits for that prune. Because of this, a short keep-alive time does not change the interval between prunes.
+        // A short keep-alive time also decreases the time that an idle pool keeps our class loader in memory.
+        pruner.setKeepAliveTime(Math.max(1, maxIdleTimeNanos / 2), TimeUnit.NANOSECONDS);
+        pruner.allowCoreThreadTimeOut(true);
+        pruner.setRemoveOnCancelPolicy(true);
     }
 
     /**
-     * Call this method at most once to enable a background thread that prunes idle buffers from the pool
+     * Call this method one time only. It permits the pool to prune idle buffers.
+     *
+     * <p>This method does not start a thread. An empty pool has no buffers to prune. The pruner starts when you
+     * {@linkplain #release(ByteBuffer) release} a buffer. The pruner stops when the pool becomes empty.</p>
      */
     PowerOfTwoBufferPool enablePruning() {
-        pruner.scheduleAtFixedRate(this::prune, maxIdleTimeNanos, maxIdleTimeNanos / 2, TimeUnit.NANOSECONDS);
+        pruningEnabled = true;
+        if (!allPoolsEmpty()) {
+            // The pool can hold buffers from before this call, and those buffers also need a prune. An empty pool
+            // must not start a thread.
+            startPruningIfNeeded();
+        }
         return this;
     }
 
     void disablePruning() {
+        pruningEnabled = false;
         pruner.shutdownNow();
     }
 
+    /**
+     * @return The number of threads that the pruner uses. This method is package-private because the tests must show
+     * that no thread runs when the pool has no buffers to prune. JAVA-6279 is about that behavior.
+     */
+    int prunerThreadCount() {
+        return pruner.getPoolSize();
+    }
+
     @Override
     public ByteBuf getBuffer(final int size) {
         return new PooledByteBufNIO(getByteBuffer(size));
@@ -136,7 +180,59 @@ public void release(final ByteBuffer buffer) {
                 powerOfTwoToPoolMap.get(log2(roundUpToNextHighestPowerOfTwo(buffer.capacity())));
         if (pool != null) {
             pool.release(new IdleTrackingByteBuffer(buffer));
+            startPruningIfNeeded();
+        }
+    }
+
+    private void startPruningIfNeeded() {
+        if (pruningEnabled && pruningScheduled.compareAndSet(false, true)) {
+            schedulePrune();
+        }
+    }
+
+    private void schedulePrune() {
+        try {
+            pruner.schedule(this::pruneAndRescheduleIfNeeded, maxIdleTimeNanos / 2, TimeUnit.NANOSECONDS);
+        } catch (RejectedExecutionException e) {
+            // Another thread called `disablePruning` and stopped the executor. A release of a buffer must not fail
+            // because of this.
+            pruningScheduled.set(false);
+        }
+    }
+
+    /**
+     * Prunes the pool. Then schedules the next prune, but only if the pool is not empty.
+     *
+     * <p>This method does not cancel a task to stop the pruner. It stops the pruner when it does not schedule the next
+     * prune. Then the work queue becomes empty and the pruner thread stops.</p>
+     *
+     * <p>The steps below prevent a lost pruner. A thread that releases a buffer reads {@link #pruningScheduled}. If
+     * that flag is true, the thread does not schedule a prune, because it relies on this method to schedule the next
+     * prune. For this reason, this method clears the flag and then examines the pool one more time. If the pool is not
+     * empty, this method takes the next prune. If it cannot take the next prune, the other thread has taken it. The
+     * class {@code io.netty.util.concurrent.GlobalEventExecutor.TaskRunner} uses the same steps.</p>
+     */
+    private void pruneAndRescheduleIfNeeded() {
+        prune();
+        if (allPoolsEmpty()) {
+            pruningScheduled.set(false);
+            if (allPoolsEmpty()) {
+                return;
+            }
+            if (!pruningScheduled.compareAndSet(false, true)) {
+                return;
+            }
+        }
+        schedulePrune();
+    }
+
+    private boolean allPoolsEmpty() {
+        for (BufferPool pool : powerOfTwoToPoolMap.values()) {
+            if (!pool.isEmpty()) {
+                return false;
+            }
         }
+        return true;
     }
 
     private void prune() {
@@ -204,5 +300,9 @@ void prune() {
             long now = System.nanoTime();
             available.removeIf(cur -> now - cur.getLastUsedNanos() >= maxIdleTimeNanos);
         }
+
+        boolean isEmpty() {
+            return available.isEmpty();
+        }
     }
 }

driver-core/src/test/unit/com/mongodb/internal/connection/PowerOfTwoBufferPoolTest.java

@@ -22,10 +22,12 @@
 
 import java.nio.ByteBuffer;
 import java.util.concurrent.TimeUnit;
+import java.util.function.BooleanSupplier;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotSame;
 import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
 
 public class PowerOfTwoBufferPoolTest {
     private PowerOfTwoBufferPool pool;
@@ -75,7 +77,6 @@ public void testHugeBufferRequest() {
         assertNotSame(buf, pool.getBuffer((int) Math.pow(2, 10) + 1));
     }
 
-    // Racy test
     @Test
     public void testPruning() throws InterruptedException {
         PowerOfTwoBufferPool pool = new PowerOfTwoBufferPool(10, 5, TimeUnit.MILLISECONDS)
@@ -84,11 +85,88 @@ public void testPruning() throws InterruptedException {
             ByteBuf byteBuf = pool.getBuffer(256);
             ByteBuffer wrappedByteBuf = byteBuf.asNIO();
             byteBuf.release();
-            Thread.sleep(50);
+            // The pruner stops only after it empties the pool. Therefore a thread count of zero shows that the pruner
+            // removed the buffer. A wait for a fixed period would make this test racy.
+            assertTrue("the pruner must empty the pool", await(() -> pool.prunerThreadCount() == 0));
             ByteBuf newByteBuf = pool.getBuffer(256);
             assertNotSame(wrappedByteBuf, newByteBuf.asNIO());
         } finally {
             pool.disablePruning();
         }
     }
+
+    /**
+     * The pruner removes idle buffers, and an empty pool has no idle buffers. Therefore {@code enablePruning} must not
+     * start a thread. A thread that runs keeps the class loader of all driver classes in memory. See JAVA-6279.
+     */
+    @Test
+    public void testEnablePruningStartsNoThreadWhileThePoolIsEmpty() {
+        PowerOfTwoBufferPool pool = new PowerOfTwoBufferPool(10, 5, TimeUnit.MILLISECONDS).enablePruning();
+        try {
+            assertEquals(0, pool.prunerThreadCount());
+        } finally {
+            pool.disablePruning();
+        }
+    }
+
+    /**
+     * The pruner empties the pool. Then it has no more work, and the thread must stop. The thread must not continue to
+     * wake up. This behavior is the correction for JAVA-6279.
+     */
+    @Test
+    public void testPrunerThreadTerminatesOnceThePoolIsDrained() throws InterruptedException {
+        PowerOfTwoBufferPool pool = new PowerOfTwoBufferPool(10, 5, TimeUnit.MILLISECONDS).enablePruning();
+        try {
+            pool.getBuffer(256).release();
+            assertTrue("the pruner thread should terminate once the pool is drained",
+                    await(() -> pool.prunerThreadCount() == 0));
+        } finally {
+            pool.disablePruning();
+        }
+    }
+
+    /**
+     * The pruner must start again. A pool can become idle and then busy. If the pruner does not start again, the pool
+     * keeps the buffers that you release after the idle period.
+     */
+    @Test
+    public void testPruningResumesAfterTheThreadHasTerminated() throws InterruptedException {
+        PowerOfTwoBufferPool pool = new PowerOfTwoBufferPool(10, 5, TimeUnit.MILLISECONDS).enablePruning();
+        try {
+            pool.getBuffer(256).release();
+            assertTrue("precondition: the pruner thread terminates once drained",
+                    await(() -> pool.prunerThreadCount() == 0));
+
+            ByteBuf byteBuf = pool.getBuffer(256);
+            ByteBuffer wrapped = byteBuf.asNIO();
+            byteBuf.release();
+            assertTrue("a buffer released after termination should still be pruned",
+                    await(() -> pool.getBuffer(256).asNIO() != wrapped));
+        } finally {
+            pool.disablePruning();
+        }
+    }
+
+    /** A pool without pruning must not start a pruner thread. The number of buffers does not change this behavior. */
+    @Test
+    public void testPruningDisabledPoolNeverStartsAThread() {
+        ByteBuf byteBuf = pool.getBuffer(256);
+        ByteBuffer wrapped = byteBuf.asNIO();
+        byteBuf.release();
+        // This assertion needs no wait. The executor creates its worker thread when it accepts a task, and not when it
+        // runs that task. Therefore a pool that schedules a prune has a thread before `release` returns.
+        assertEquals(0, pool.prunerThreadCount());
+        assertSame("the pool must keep the buffer because it does not prune", wrapped, pool.getBuffer(256).asNIO());
+    }
+
+    private static boolean await(final BooleanSupplier condition) throws InterruptedException {
+        long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5);
+        while (System.nanoTime() < deadline) {
+            if (condition.getAsBoolean()) {
+                return true;
+            }
+            Thread.sleep(5);
+        }
+        return condition.getAsBoolean();
+    }
 }