Shutdown ScheduledExecutors as part of node drainage

patch by Stefan Miklosovic; reviewed by Jon Meredith for CASSANDRA-17493

CHANGES.txt

@@ -1,4 +1,5 @@
 4.1
+ * Shutdown ScheduledExecutors as part of node drainage (CASSANDRA-17493)
  * Provide JMX endpoint to allow transient logging of blocking read repairs (CASSANDRA-17471)
  * Add guardrail for GROUP BY queries (CASSANDRA-17509)
  * make pylib PEP and pylint compliant (CASSANDRA-17546)

src/java/org/apache/cassandra/service/StorageService.java

@@ -810,7 +810,18 @@ public class StorageService extends NotificationBroadcasterSupport implements IE
             public void runMayThrow() throws InterruptedException, ExecutionException, IOException
             {
                 drain(true);
-                LoggingSupportFactory.getLoggingSupport().onShutdown();
+                try
+                {
+                    ExecutorUtils.shutdownNowAndWait(1, MINUTES, ScheduledExecutors.scheduledFastTasks);
+                }
+                catch (Throwable t)
+                {
+                    logger.warn("Unable to terminate fast tasks within 1 minute.", t);
+                }
+                finally
+                {
+                    LoggingSupportFactory.getLoggingSupport().onShutdown();
+                }
             }
         }, "StorageServiceShutdownHook");
         Runtime.getRuntime().addShutdownHook(drainOnShutdown);
@@ -5287,12 +5298,21 @@ public class StorageService extends NotificationBroadcasterSupport implements IE
             CommitLog.instance.shutdownBlocking();
 
             // wait for miscellaneous tasks like sstable and commitlog segment deletion
-            ScheduledExecutors.nonPeriodicTasks.shutdown();
-            if (!ScheduledExecutors.nonPeriodicTasks.awaitTermination(1, MINUTES))
-                logger.warn("Unable to terminate non-periodic tasks within 1 minute.");
-
             ColumnFamilyStore.shutdownPostFlushExecutor();
-            setMode(Mode.DRAINED, !isFinalShutdown);
+
+            try
+            {
+                // we are not shutting down ScheduledExecutors#scheduledFastTasks to be still able to progress time
+                // fast-tasks executor is shut down in StorageService's shutdown hook added to Runtime
+                ExecutorUtils.shutdownNowAndWait(1, MINUTES,
+                                                 ScheduledExecutors.nonPeriodicTasks,
+                                                 ScheduledExecutors.scheduledTasks,
+                                                 ScheduledExecutors.optionalTasks);
+            }
+            finally
+            {
+                setMode(Mode.DRAINED, !isFinalShutdown);
+            }
         }
         catch (Throwable t)
         {

test/unit/org/apache/cassandra/service/StorageServiceTest.java

@@ -18,25 +18,35 @@
 
 package org.apache.cassandra.service;
 
-import org.apache.cassandra.locator.EndpointsByReplica;
-import org.apache.cassandra.locator.ReplicaCollection;
+import java.util.concurrent.ScheduledFuture;
+import java.util.concurrent.atomic.AtomicInteger;
+
 import org.junit.Before;
 import org.junit.BeforeClass;
 import org.junit.Test;
 
+import org.apache.cassandra.concurrent.ScheduledExecutors;
 import org.apache.cassandra.config.DatabaseDescriptor;
+import org.apache.cassandra.db.commitlog.CommitLog;
 import org.apache.cassandra.dht.RandomPartitioner;
 import org.apache.cassandra.dht.Range;
 import org.apache.cassandra.dht.Token;
 import org.apache.cassandra.locator.AbstractEndpointSnitch;
 import org.apache.cassandra.locator.AbstractReplicationStrategy;
+import org.apache.cassandra.locator.EndpointsByReplica;
 import org.apache.cassandra.locator.IEndpointSnitch;
 import org.apache.cassandra.locator.InetAddressAndPort;
 import org.apache.cassandra.locator.Replica;
+import org.apache.cassandra.locator.ReplicaCollection;
 import org.apache.cassandra.locator.ReplicaMultimap;
 import org.apache.cassandra.locator.SimpleStrategy;
 import org.apache.cassandra.locator.TokenMetadata;
 
+import static java.util.concurrent.TimeUnit.MINUTES;
+import static java.util.concurrent.TimeUnit.SECONDS;
+import static org.awaitility.Awaitility.await;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
 public class StorageServiceTest
@@ -93,6 +103,7 @@ public class StorageServiceTest
         };
 
         DatabaseDescriptor.setEndpointSnitch(snitch);
+        CommitLog.instance.start();
     }
 
     private AbstractReplicationStrategy simpleStrategy(TokenMetadata tmd)
@@ -156,4 +167,27 @@ public class StorageServiceTest
         expectedResult.put(new Replica(aAddress, dRange, false), new Replica(bAddress, dRange, false));
         assertMultimapEqualsIgnoreOrder(result, expectedResult.build());
     }
+
+    @Test
+    public void testScheduledExecutorsShutdownOnDrain() throws Throwable
+    {
+        final AtomicInteger numberOfRuns = new AtomicInteger(0);
+
+        ScheduledFuture<?> f = ScheduledExecutors.scheduledTasks.scheduleAtFixedRate(numberOfRuns::incrementAndGet,
+                                                                                     0, 1, SECONDS);
+
+        // Prove the task was scheduled more than once before checking cancelled.
+        await("first run").atMost(1, MINUTES).until(() -> numberOfRuns.get() > 1);
+
+        assertFalse(f.isCancelled());
+        StorageService.instance.drain();
+        assertTrue(f.isCancelled());
+
+        assertTrue(ScheduledExecutors.scheduledTasks.isTerminated());
+        assertTrue(ScheduledExecutors.nonPeriodicTasks.isTerminated());
+        assertTrue(ScheduledExecutors.optionalTasks.isTerminated());
+
+        // fast tasks are shut down as part of the Runtime shutdown hook.
+        assertFalse(ScheduledExecutors.scheduledFastTasks.isTerminated());
+    }
 }