Stop the buffer pool pruner thread when the pool is empty

.gitignore

@@ -52,6 +52,7 @@ local.properties
 *.sh
 !.evergreen/*.sh
 !scripts/*.sh
+!testing/java-6279-poc/*.sh
 
 # security-sensitive files
 *.gpg

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();
+    }
 }