From c547a874fc9a039a2db01932e71cf169cb073c17 Mon Sep 17 00:00:00 2001 From: Maxim Muzafarov Date: Sun, 12 Jul 2026 19:31:54 +0200 Subject: [PATCH] dispatcher gracefully shutdown must be handlded for management connections --- .../cassandra/transport/Dispatcher.java | 32 ++++- .../transport/PipelineConfigurator.java | 2 +- .../apache/cassandra/transport/Server.java | 2 +- .../transport/MessageDispatcherTest.java | 2 +- .../MessageManagementDispatcherTest.java | 114 +++++++++++++++++- 5 files changed, 143 insertions(+), 9 deletions(-) diff --git a/src/java/org/apache/cassandra/transport/Dispatcher.java b/src/java/org/apache/cassandra/transport/Dispatcher.java index fa0692a281..762b653caa 100644 --- a/src/java/org/apache/cassandra/transport/Dispatcher.java +++ b/src/java/org/apache/cassandra/transport/Dispatcher.java @@ -119,7 +119,16 @@ public class Dispatcher implements CQLMessageHandler.MessageConsumer 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 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 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 { + 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 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