FIx PersistentTTLNode thread leak (#1264)

chevaris · web-flow · commit 57cf5135aa7c · 2025-04-21T18:24:30.000+08:00
Fixes #1263.
diff --git a/curator-client/src/main/java/org/apache/curator/utils/CloseableScheduledExecutorService.java b/curator-client/src/main/java/org/apache/curator/utils/CloseableScheduledExecutorService.java
@@ -95,4 +95,28 @@ public Future<?> scheduleWithFixedDelay(Runnable task, long initialDelay, long d
                 scheduledExecutorService.scheduleWithFixedDelay(task, initialDelay, delay, unit);
         return new InternalScheduledFutureTask(scheduledFuture);
     }
+
+    /**
+     * Creates and executes a periodic action that becomes enabled first after the given initial
+     * delay, and subsequently with the given period; that is executions will commence after {@code
+     * initialDelay} then {@code initialDelay+period}, then {@code initialDelay + 2 * period}, and so
+     * on. If any execution of the task encounters an exception, subsequent executions are suppressed.
+     * Otherwise, the task will only terminate via cancellation or termination of the executor. If any
+     * execution of this task takes longer than its period, then subsequent executions may start late,
+     * but will not concurrently execute.
+     *
+     * @param task the task to execute
+     * @param initialDelay the time to delay first execution
+     * @param period the period between successive executions
+     * @param unit the time unit of the initialDelay and delay parameters
+     * @return a Future representing pending completion of the task, and whose <tt>get()</tt> method
+     *     will throw an exception upon cancellation
+     */
+    public Future<?> scheduleAtFixedRate(Runnable task, long initialDelay, long period, TimeUnit unit) {
+        Preconditions.checkState(isOpen.get(), "CloseableExecutorService is closed");
+
+        ScheduledFuture<?> scheduledFuture =
+                scheduledExecutorService.scheduleAtFixedRate(task, initialDelay, period, unit);
+        return new InternalScheduledFutureTask(scheduledFuture);
+    }
 }
diff --git a/curator-client/src/test/java/org/apache/curator/utils/TestCloseableScheduledExecutorService.java b/curator-client/src/test/java/org/apache/curator/utils/TestCloseableScheduledExecutorService.java
@@ -20,9 +20,12 @@
 package org.apache.curator.utils;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.Executors;
+import java.util.concurrent.RejectedExecutionException;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -110,4 +113,84 @@ public void run() {
         assertTrue(newValue > value);
         assertEquals(innerValue, innerCounter.get());
     }
+
+    @Test
+    public void testCloseableScheduleAtFixedRate() throws InterruptedException {
+        CloseableScheduledExecutorService service = new CloseableScheduledExecutorService(executorService);
+
+        final CountDownLatch latch = new CountDownLatch(QTY);
+        final AtomicInteger fixedDelayCounter = new AtomicInteger();
+        service.scheduleAtFixedRate(
+                new Runnable() {
+                    @Override
+                    public void run() {
+                        latch.countDown();
+                        // This delay is almost NOT impacting when using scheduleAtFixedRate
+                        try {
+                            Thread.sleep(DELAY_MS);
+                        } catch (InterruptedException e) {
+                            // Do nothing
+                        }
+                    }
+                },
+                DELAY_MS,
+                DELAY_MS,
+                TimeUnit.MILLISECONDS);
+        service.scheduleWithFixedDelay(
+                new Runnable() {
+                    @Override
+                    public void run() {
+                        fixedDelayCounter.incrementAndGet();
+                        // This delay is impacting when using scheduleWithFixedDelay
+                        try {
+                            Thread.sleep(DELAY_MS);
+                        } catch (InterruptedException e) {
+                            // Do nothing
+                        }
+                    }
+                },
+                DELAY_MS,
+                DELAY_MS,
+                TimeUnit.MILLISECONDS);
+        assertTrue(latch.await((QTY * 2) * DELAY_MS, TimeUnit.MILLISECONDS));
+        assertTrue(fixedDelayCounter.get() <= (QTY / 2 + 1));
+    }
+
+    @Test
+    public void testCloseWithoutShutdown() throws InterruptedException {
+        final CountDownLatch latch = new CountDownLatch(1);
+        try (CloseableScheduledExecutorService service = new CloseableScheduledExecutorService(executorService)) {
+            service.submit(latch::countDown);
+            assertTrue(latch.await(1, TimeUnit.SECONDS));
+        }
+        assertFalse(executorService.isShutdown());
+    }
+
+    @Test
+    public void testCloseWithShutdown() throws InterruptedException {
+        final CountDownLatch latch = new CountDownLatch(1);
+        try (CloseableScheduledExecutorService service = new CloseableScheduledExecutorService(executorService, true)) {
+            service.submit(latch::countDown);
+            assertTrue(latch.await(1, TimeUnit.SECONDS));
+        }
+        assertTrue(executorService.isShutdown());
+        assertThrows(RejectedExecutionException.class, () -> executorService.submit(() -> System.out.println("Hello")));
+    }
+
+    @Test
+    public void testCloseMultipleTimes() throws InterruptedException {
+        final CountDownLatch latch = new CountDownLatch(1);
+        CloseableScheduledExecutorService service = null;
+        try {
+            service = new CloseableScheduledExecutorService(executorService, true);
+            service.submit(latch::countDown);
+            assertTrue(latch.await(1, TimeUnit.SECONDS));
+        } finally {
+            if (service != null) {
+                service.close();
+                service.close();
+                service.close();
+            }
+        }
+    }
 }
diff --git a/curator-recipes/src/main/java/org/apache/curator/framework/recipes/nodes/PersistentTtlNode.java b/curator-recipes/src/main/java/org/apache/curator/framework/recipes/nodes/PersistentTtlNode.java
@@ -24,11 +24,10 @@
 import java.io.IOException;
 import java.util.Objects;
 import java.util.concurrent.Executors;
-import java.util.concurrent.Future;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicReference;
 import org.apache.curator.framework.CuratorFramework;
+import org.apache.curator.utils.CloseableScheduledExecutorService;
 import org.apache.curator.utils.ThreadUtils;
 import org.apache.curator.utils.ZKPaths;
 import org.apache.zookeeper.CreateMode;
@@ -61,13 +60,15 @@ public class PersistentTtlNode implements Closeable {
     public static final int DEFAULT_TOUCH_SCHEDULE_FACTOR = 2;
     public static final boolean DEFAULT_USE_PARENT_CREATION = true;
 
+    @VisibleForTesting
+    static final String TOUCH_THREAD_NAME = "PersistentTtlNode";
+
     private final Logger log = LoggerFactory.getLogger(getClass());
     private final PersistentNode node;
     private final CuratorFramework client;
     private final long ttlMs;
     private final int touchScheduleFactor;
-    private final ScheduledExecutorService executorService;
-    private final AtomicReference<Future<?>> futureRef = new AtomicReference<>();
+    private final CloseableScheduledExecutorService closeableExecutorService;
     private final String childPath;
 
     /**
@@ -79,7 +80,9 @@ public class PersistentTtlNode implements Closeable {
     public PersistentTtlNode(CuratorFramework client, String path, long ttlMs, byte[] initData) {
         this(
                 client,
-                Executors.newSingleThreadScheduledExecutor(ThreadUtils.newThreadFactory("PersistentTtlNode")),
+                new CloseableScheduledExecutorService(
+                        Executors.newSingleThreadScheduledExecutor(ThreadUtils.newThreadFactory(TOUCH_THREAD_NAME)),
+                        true),
                 path,
                 ttlMs,
                 initData,
@@ -99,7 +102,9 @@ public PersistentTtlNode(
             CuratorFramework client, String path, long ttlMs, byte[] initData, boolean useParentCreation) {
         this(
                 client,
-                Executors.newSingleThreadScheduledExecutor(ThreadUtils.newThreadFactory("PersistentTtlNode")),
+                new CloseableScheduledExecutorService(
+                        Executors.newSingleThreadScheduledExecutor(ThreadUtils.newThreadFactory(TOUCH_THREAD_NAME)),
+                        true),
                 path,
                 ttlMs,
                 initData,
@@ -128,7 +133,7 @@ public PersistentTtlNode(
             int touchScheduleFactor) {
         this(
                 client,
-                executorService,
+                new CloseableScheduledExecutorService(executorService),
                 path,
                 ttlMs,
                 initData,
@@ -157,6 +162,26 @@ public PersistentTtlNode(
             String childNodeName,
             int touchScheduleFactor,
             boolean useParentCreation) {
+        this(
+                client,
+                new CloseableScheduledExecutorService(executorService, false),
+                path,
+                ttlMs,
+                initData,
+                childNodeName,
+                touchScheduleFactor,
+                useParentCreation);
+    }
+
+    private PersistentTtlNode(
+            CuratorFramework client,
+            CloseableScheduledExecutorService closeableExecutorService,
+            String path,
+            long ttlMs,
+            byte[] initData,
+            String childNodeName,
+            int touchScheduleFactor,
+            boolean useParentCreation) {
         this.client = Objects.requireNonNull(client, "client cannot be null");
         this.ttlMs = ttlMs;
         this.touchScheduleFactor = touchScheduleFactor;
@@ -168,7 +193,7 @@ protected void deleteNode() {
                         // NOP
                     }
                 };
-        this.executorService = Objects.requireNonNull(executorService, "executorService cannot be null");
+        this.closeableExecutorService = closeableExecutorService;
         childPath = ZKPaths.makePath(Objects.requireNonNull(path, "path cannot be null"), childNodeName);
     }
 
@@ -198,9 +223,8 @@ void touch() {
      */
     public void start() {
         node.start();
-        Future<?> future = executorService.scheduleAtFixedRate(
+        closeableExecutorService.scheduleAtFixedRate(
                 this::touch, ttlMs / touchScheduleFactor, ttlMs / touchScheduleFactor, TimeUnit.MILLISECONDS);
-        futureRef.set(future);
     }
 
     /**
@@ -238,17 +262,19 @@ public byte[] getData() {
         return node.getData();
     }
 
+    @VisibleForTesting
+    CloseableScheduledExecutorService getCloseableScheduledExecutorService() {
+        return closeableExecutorService;
+    }
+
     /**
      * Call when you are done with the PersistentTtlNode. Note: the ZNode is <em>not</em> immediately
      * deleted. However, if no other PersistentTtlNode with the same path is running the node will get deleted
      * based on the ttl.
      */
     @Override
     public void close() {
-        Future<?> future = futureRef.getAndSet(null);
-        if (future != null) {
-            future.cancel(true);
-        }
+        closeableExecutorService.close();
         try {
             node.close();
         } catch (IOException e) {
diff --git a/curator-recipes/src/test/java/org/apache/curator/framework/recipes/nodes/TestPersistentTtlNode.java b/curator-recipes/src/test/java/org/apache/curator/framework/recipes/nodes/TestPersistentTtlNode.java
