From ac49b8e5f355cc91895c9ef56f5440fc00d660ac Mon Sep 17 00:00:00 2001
From: Gian Merlino <gianmerlino@gmail.com>
Date: Tue, 30 Jul 2024 08:27:14 -0700
Subject: [PATCH] Fix worker stopping in tests.

---
 .../org/apache/druid/msq/test/MSQTestControllerContext.java  | 1 +
 .../java/org/apache/druid/msq/test/MSQTestWorkerClient.java  | 5 ++++-
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestControllerContext.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestControllerContext.java
index fdbb00ae1f5e..72a9cdfd70b2 100644
--- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestControllerContext.java
+++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestControllerContext.java
@@ -262,6 +262,7 @@ public ListenableFuture<Void> cancelTask(String workerId)
     {
       final Worker worker = inMemoryWorkers.remove(workerId);
       if (worker != null) {
+        worker.stop();
         worker.awaitStop();
       }
       return Futures.immediateFuture(null);
diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java
index f384397965b9..e6048fbf600c 100644
--- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java
+++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java
@@ -138,6 +138,9 @@ public ListenableFuture<Boolean> fetchChannelData(
   @Override
   public void close()
   {
-    inMemoryWorkers.forEach((k, v) -> v.awaitStop());
+    inMemoryWorkers.forEach((k, v) -> {
+      v.stop();
+      v.awaitStop();
+    });
   }
 }