Ensure shutdown of LA slot queue isn't swallowed (#2161)

temporalio · Jul 30, 2024 · 1acafa3 · 1acafa3
1 parent e5c08a1
commit 1acafa3
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 4 deletions.
diff --git a/temporal-sdk/src/main/java/io/temporal/internal/worker/LocalActivitySlotSupplierQueue.java b/temporal-sdk/src/main/java/io/temporal/internal/worker/LocalActivitySlotSupplierQueue.java
@@ -84,6 +84,9 @@ private void processQueue() {
         SlotPermit slotPermit;
         try {
           slotPermit = slotSupplier.reserveSlot(request.data);
+        } catch (InterruptedException e) {
+          Thread.currentThread().interrupt();
+          return;
         } catch (Exception e) {
           log.error(
               "Error reserving local activity slot, dropped activity id {}",
@@ -131,8 +134,4 @@ void submitAttempt(SlotReservationData data, boolean isRetry, LocalActivityAttem
       newExecutionsBackpressureSemaphore.release();
     }
   }
-
-  TrackingSlotSupplier<LocalActivitySlotInfo> getSlotSupplier() {
-    return slotSupplier;
-  }
 }
diff --git a/temporal-sdk/src/main/java/io/temporal/internal/worker/LocalActivityWorker.java b/temporal-sdk/src/main/java/io/temporal/internal/worker/LocalActivityWorker.java
@@ -717,6 +717,7 @@ public CompletableFuture<Void> shutdown(ShutdownManager shutdownManager, boolean
 
   @Override
   public void awaitTermination(long timeout, TimeUnit unit) {
+    slotQueue.shutdown();
     long timeoutMillis = unit.toMillis(timeout);
     ShutdownManager.awaitTermination(scheduledExecutor, timeoutMillis);
   }

diff --git a/temporal-sdk/src/test/java/io/temporal/internal/worker/WorkflowSlotsSmallSizeTests.java b/temporal-sdk/src/test/java/io/temporal/internal/worker/WorkflowSlotsSmallSizeTests.java
@@ -47,6 +47,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.Semaphore;
+import java.util.concurrent.TimeUnit;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Rule;
@@ -249,6 +250,29 @@ public void TestLocalActivitySlotAtLimit() throws InterruptedException {
         MAX_CONCURRENT_LOCAL_ACTIVITY_EXECUTION_SIZE);
   }
 
+  @Test
+  public void TestLocalActivityShutdownWhileWaitingOnSlot() throws InterruptedException {
+    testWorkflowRule.getTestEnvironment().start();
+    WorkflowClient client = testWorkflowRule.getWorkflowClient();
+    TestWorkflow workflow =
+        client.newWorkflowStub(
+            TestWorkflow.class,
+            WorkflowOptions.newBuilder()
+                .setTaskQueue(testWorkflowRule.getTaskQueue())
+                .validateBuildWithDefaults());
+    WorkflowClient.start(workflow::workflow, activitiesAreLocal);
+    workflow.unblock();
+    parallelSemRunning.acquire(2);
+    testWorkflowRule.getTestEnvironment().getWorkerFactory().shutdownNow();
+    parallelSemBlocked.release(2);
+    testWorkflowRule.getTestEnvironment().getWorkerFactory().awaitTermination(3, TimeUnit.SECONDS);
+    // All slots should be available
+    assertWorkerSlotCount(
+        MAX_CONCURRENT_WORKFLOW_TASK_EXECUTION_SIZE,
+        MAX_CONCURRENT_ACTIVITY_EXECUTION_SIZE,
+        MAX_CONCURRENT_LOCAL_ACTIVITY_EXECUTION_SIZE);
+  }
+
   @Test
   public void TestLocalActivitySlotHitsCapacity() throws InterruptedException {
     testWorkflowRule.getTestEnvironment().start();