dispatcher gracefully shutdown must be handlded for management connections

src/java/org/apache/cassandra/transport/Dispatcher.java

@@ -119,7 +119,16 @@ public class Dispatcher implements CQLMessageHandler.MessageConsumer<Message.Req
                                                                                  "Native-Transport-Management-Tasks");
 
     private static final ConcurrentMap<EventLoop, Flusher> flusherLookup = new ConcurrentHashMap<>();
+
+    /**
+     * Set while a management request is executing on the current thread. A management command may itself
+     * initiate the management server's stop (e.g. INVOKE COMMAND stopdaemon), in which case the drain-wait
+     * in {@link Server#close} runs on this very thread and must not wait for its own task to complete.
+     */
+    private static final ThreadLocal<Boolean> IN_MANAGEMENT_TASK = ThreadLocal.withInitial(() -> Boolean.FALSE);
+
     private final boolean useLegacyFlusher;
+    private final boolean isManagementDispatcher;
 
     /**
      * Takes a Channel, Request and the Response produced by processRequest and outputs a FlushItem
@@ -133,9 +142,10 @@ public class Dispatcher implements CQLMessageHandler.MessageConsumer<Message.Req
         FlushItem<?> toFlushItem(P param, Channel channel, Message.Request request, Message.Response response);
     }
 
-    public Dispatcher(boolean useLegacyFlusher)
+    public Dispatcher(boolean useLegacyFlusher, boolean isManagementDispatcher)
     {
         this.useLegacyFlusher = useLegacyFlusher;
+        this.isManagementDispatcher = isManagementDispatcher;
     }
 
     @Override
@@ -437,9 +447,14 @@ public class Dispatcher implements CQLMessageHandler.MessageConsumer<Message.Req
                 }
             }
 
-            // If validation passes, call the normal processRequest to execute
-            // This calls the instance method processRequest() which does all the work
-            processRequest(channel, request, forFlusher, flusherParam, backpressure, requestTime);
+            IN_MANAGEMENT_TASK.set(Boolean.TRUE);
+            try {
+                // If validation passes, call the normal processRequest to execute
+                // This calls the instance method processRequest() which does all the work
+                processRequest(channel, request, forFlusher, flusherParam, backpressure, requestTime);
+            } finally {
+                IN_MANAGEMENT_TASK.set(Boolean.FALSE);
+            }
         }
     }
 
@@ -670,9 +685,16 @@ public class Dispatcher implements CQLMessageHandler.MessageConsumer<Message.Req
         flusher.start();
     }
 
+    /**
+     * @return true once the executor serving this dispatcher's server has drained: {@link #managementExecutor}
+     * for the management transport, {@link #requestExecutor} otherwise. A management task running on the
+     * calling thread is excluded — it is the task that initiated the stop and cannot drain before it.
+     */
     public boolean isDone()
     {
-        return requestExecutor.getPendingTaskCount() == 0 && requestExecutor.getActiveTaskCount() == 0;
+        LocalAwareExecutorPlus executor = isManagementDispatcher ? managementExecutor : requestExecutor;
+        int self = isManagementDispatcher && IN_MANAGEMENT_TASK.get() ? 1 : 0;
+        return executor.getPendingTaskCount() == 0 && executor.getActiveTaskCount() - self <= 0;
     }
 
     public static void shutdown()

src/java/org/apache/cassandra/transport/PipelineConfigurator.java

@@ -142,7 +142,7 @@ public class PipelineConfigurator
         this.epoll               = epoll;
         this.keepAlive           = keepAlive;
         this.tlsEncryptionPolicy = encryptionPolicy;
-        this.dispatcher          = new Dispatcher(useLegacyFlusher);
+        this.dispatcher          = new Dispatcher(useLegacyFlusher, false);
         this.queueBackpressure   = QueueBackpressure.DEFAULT;
         this.isManagementConnection = false;
     }

src/java/org/apache/cassandra/transport/Server.java

@@ -112,7 +112,7 @@ public class Server implements CassandraDaemon.Server
                 workerGroup = new NioEventLoopGroup();
         }
 
-        dispatcher = new Dispatcher(DatabaseDescriptor.useNativeTransportLegacyFlusher());
+        dispatcher = new Dispatcher(DatabaseDescriptor.useNativeTransportLegacyFlusher(), builder.isManagementConnection);
         pipelineConfigurator = builder.pipelineConfigurator != null
                                ? builder.pipelineConfigurator
                                : new PipelineConfigurator(useEpoll,

test/unit/org/apache/cassandra/transport/MessageDispatcherTest.java

@@ -172,7 +172,7 @@ public class MessageDispatcherTest
     {
         public AuthTestDispatcher()
         {
-            super(false);
+            super(false, false);
         }
 
         @Override

test/unit/org/apache/cassandra/transport/MessageManagementDispatcherTest.java

@@ -19,7 +19,9 @@
 package org.apache.cassandra.transport;
 
 import java.util.concurrent.Callable;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import com.google.common.util.concurrent.Uninterruptibles;
 
@@ -332,6 +334,111 @@ public class MessageManagementDispatcherTest
                     Dispatcher.isManagementRequestAllowed(queryMessage("USE my_keyspace;")));
     }
 
+    @Test
+    public void testIsDoneTracksManagementExecutor() throws Exception
+    {
+        Dispatcher managementDispatcher = new ManagementTestDispatcher(true);
+        CountDownLatch started = new CountDownLatch(1);
+        CountDownLatch release = new CountDownLatch(1);
+        Dispatcher.managementExecutor.submit(() -> {
+            started.countDown();
+            Uninterruptibles.awaitUninterruptibly(release);
+        });
+        try
+        {
+            assertTrue(started.await(10, TimeUnit.SECONDS));
+            assertFalse("An in-flight management task should block the management drain",
+                        managementDispatcher.isDone());
+            awaitTrue("The regular dispatcher drain should ignore the management executor",
+                      dispatch::isDone);
+        }
+        finally
+        {
+            release.countDown();
+        }
+        awaitTrue("The management drain should complete once its task finishes",
+                  managementDispatcher::isDone);
+    }
+
+    @Test
+    public void testIsDoneIgnoresRequestExecutorBacklog() throws Exception
+    {
+        Dispatcher managementDispatcher = new ManagementTestDispatcher(true);
+        CountDownLatch started = new CountDownLatch(1);
+        CountDownLatch release = new CountDownLatch(1);
+        Dispatcher.requestExecutor.submit(() -> {
+            started.countDown();
+            Uninterruptibles.awaitUninterruptibly(release);
+        });
+        try
+        {
+            assertTrue(started.await(10, TimeUnit.SECONDS));
+            assertFalse("An in-flight regular task should block the regular drain",
+                        dispatch.isDone());
+            assertTrue("The management drain should ignore the regular request executor",
+                       managementDispatcher.isDone());
+        }
+        finally
+        {
+            release.countDown();
+        }
+        awaitTrue("The regular drain should complete once its task finishes",
+                  dispatch::isDone);
+    }
+
+    /**
+     * A management command may itself initiate the server stop (e.g. INVOKE COMMAND stopdaemon), in which
+     * case the drain-wait runs on the very thread executing the command and must not wait for its own task.
+     */
+    @Test
+    public void testIsDoneExcludesStopInitiatingManagementTask() throws Exception
+    {
+        CountDownLatch entered = new CountDownLatch(1);
+        CountDownLatch release = new CountDownLatch(1);
+        AtomicBoolean isDoneInsideTask = new AtomicBoolean();
+
+        Dispatcher managementDispatcher = new ManagementTestDispatcher(true)
+        {
+            @Override
+            void processRequest(Channel channel,
+                                Message.Request request,
+                                FlushItemConverter forFlusher,
+                                ClientResourceLimits.Overload backpressure,
+                                RequestTime requestTime)
+            {
+                isDoneInsideTask.set(isDone());
+                entered.countDown();
+                Uninterruptibles.awaitUninterruptibly(release);
+            }
+        };
+
+        Message.Request request = createManagementRequest(Message.Type.STARTUP);
+        managementDispatcher.dispatch(request.connection().channel(), request, (channel, req, response) -> null,
+                                      ClientResourceLimits.Overload.NONE);
+        try
+        {
+            assertTrue(entered.await(10, TimeUnit.SECONDS));
+            assertTrue("The drain must not wait for the management task that initiated it",
+                       isDoneInsideTask.get());
+            assertFalse("Other threads should still see the in-flight management task",
+                        managementDispatcher.isDone());
+        }
+        finally
+        {
+            release.countDown();
+        }
+        awaitTrue("The management drain should complete once its task finishes",
+                  managementDispatcher::isDone);
+    }
+
+    private static void awaitTrue(String message, Callable<Boolean> condition) throws Exception
+    {
+        long timeout = Clock.Global.currentTimeMillis();
+        while (!condition.call() && Clock.Global.currentTimeMillis() - timeout < 10_000)
+            Uninterruptibles.sleepUninterruptibly(10, TimeUnit.MILLISECONDS);
+        assertTrue(message, condition.call());
+    }
+
     private long completedRequests()
     {
         return Dispatcher.requestExecutor.getCompletedTaskCount();
@@ -429,7 +536,12 @@ public class MessageManagementDispatcherTest
     {
         public ManagementTestDispatcher()
         {
-            super(false);
+            this(false);
+        }
+
+        public ManagementTestDispatcher(boolean isManagementDispatcher)
+        {
+            super(false, isManagementDispatcher);
         }
 
         @Override