Fix worker stopping in tests.

apache · Jul 30, 2024 · ac49b8e · ac49b8e
1 parent f9f9321
commit ac49b8e
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/...e/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestControllerContext.java b/...e/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/...s-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerClient.java b/...s-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();
+    });
   }
 }