[SDK Sentinel] Fix staged workflow worker shutdown race

temporal-sdk/src/main/java/io/temporal/internal/worker/WorkflowWorker.java

@@ -42,8 +42,10 @@
 import java.util.*;
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.RejectedExecutionException;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 import org.slf4j.Logger;
@@ -81,6 +83,7 @@ final class WorkflowWorker implements SuspendableWorker {
   private final NamespaceCapabilities namespaceCapabilities;
 
   private PollTaskExecutor<WorkflowTask> pollTaskExecutor;
+  private volatile CompletableFuture<Void> shutdownFuture = CompletableFuture.completedFuture(null);
 
   // TODO this ideally should be volatile or final (and NoopWorker should go away)
   //  Currently the implementation looks safe without volatile, but it's brittle.
@@ -259,30 +262,46 @@ public CompletableFuture<Void> shutdown(ShutdownManager shutdownManager, boolean
                             stickyQueueBalancer, options.getDrainStickyTaskQueueTimeout())
                         : CompletableFuture.completedFuture(null))
             .thenCompose(ignore -> poller.shutdown(shutdownManager, interruptTasks));
-    return pollerShutdown
-        .thenCompose(
-            ignore ->
-                !interruptTasks
-                    ? shutdownManager.waitForSupplierPermitsReleasedUnlimited(
-                        slotSupplier, supplierName)
-                    : CompletableFuture.completedFuture(null))
-        .thenCompose(
-            ignore ->
-                pollTaskExecutor != null
-                    ? pollTaskExecutor.shutdown(shutdownManager, interruptTasks)
-                    : CompletableFuture.completedFuture(null))
-        .exceptionally(
-            e -> {
-              log.error("Unexpected exception during shutdown", e);
-              return null;
-            });
+    CompletableFuture<Void> result =
+        pollerShutdown
+            .thenCompose(
+                ignore ->
+                    !interruptTasks
+                        ? shutdownManager.waitForSupplierPermitsReleasedUnlimited(
+                            slotSupplier, supplierName)
+                        : CompletableFuture.completedFuture(null))
+            .thenCompose(
+                ignore ->
+                    pollTaskExecutor != null
+                        ? pollTaskExecutor.shutdown(shutdownManager, interruptTasks)
+                        : CompletableFuture.completedFuture(null))
+            .exceptionally(
+                e -> {
+                  log.error("Unexpected exception during shutdown", e);
+                  return null;
+                });
+    shutdownFuture = result;
+    return result;
   }
 
   @Override
   public void awaitTermination(long timeout, TimeUnit unit) {
-    long timeoutMillis = ShutdownManager.awaitTermination(poller, unit.toMillis(timeout));
-    // relies on the fact that the pollTaskExecutor is the last one to be shutdown, no need to
-    // wait separately for intermediate steps
+    long timeoutMillis = unit.toMillis(timeout);
+    long shutdownTimeoutMillis = timeoutMillis;
+    CompletableFuture<Void> currentShutdownFuture = shutdownFuture;
+    timeoutMillis =
+        ShutdownManager.runAndGetRemainingTimeoutMs(
+            shutdownTimeoutMillis,
+            () -> {
+              try {
+                currentShutdownFuture.get(shutdownTimeoutMillis, TimeUnit.MILLISECONDS);
+              } catch (InterruptedException e) {
+                Thread.currentThread().interrupt();
+              } catch (ExecutionException | TimeoutException e) {
+                // Shutdown failures are logged by shutdown(), and a timeout is expected here.
+              }
+            });
+    timeoutMillis = ShutdownManager.awaitTermination(poller, timeoutMillis);
     ShutdownManager.awaitTermination(pollTaskExecutor, timeoutMillis);
   }
 

temporal-sdk/src/test/java/io/temporal/internal/worker/WorkflowWorkerTest.java

@@ -60,6 +60,7 @@
 import io.temporal.worker.tuning.PollerBehaviorSimpleMaximum;
 import io.temporal.worker.tuning.SlotSupplier;
 import io.temporal.worker.tuning.WorkflowSlotInfo;
+import java.lang.reflect.Field;
 import java.time.Duration;
 import java.util.UUID;
 import java.util.concurrent.*;
@@ -84,6 +85,62 @@ public class WorkflowWorkerTest {
   private static final String RUN_ID = "test-run-id";
   private static final String WORKFLOW_TYPE = "test-workflow-type";
 
+  @Test
+  public void awaitTerminationWaitsForStagedShutdown() throws Exception {
+    WorkflowServiceStubs service = mock(WorkflowServiceStubs.class);
+    WorkflowRunLockManager runLockManager = new WorkflowRunLockManager();
+    Scope metricsScope = new NoopScope();
+    WorkflowWorker worker =
+        new WorkflowWorker(
+            service,
+            "default",
+            "task_queue",
+            null,
+            SingleWorkerOptions.newBuilder()
+                .setIdentity("test_identity")
+                .setBuildId("test_build_id")
+                .setWorkerInstanceKey("test_worker_instance_key")
+                .setMetricsScope(metricsScope)
+                .build(),
+            runLockManager,
+            new WorkflowExecutorCache(10, runLockManager, metricsScope),
+            mock(WorkflowTaskHandler.class),
+            mock(EagerActivityDispatcher.class),
+            3,
+            new FixedSizeSlotSupplier<>(10),
+            new NamespaceCapabilities());
+
+    SuspendableWorker poller = mock(SuspendableWorker.class);
+    CompletableFuture<Void> pollerShutdown = new CompletableFuture<>();
+    when(poller.shutdown(any(ShutdownManager.class), eq(true))).thenReturn(pollerShutdown);
+    Field pollerField = WorkflowWorker.class.getDeclaredField("poller");
+    pollerField.setAccessible(true);
+    pollerField.set(worker, poller);
+
+    ShutdownManager shutdownManager = new ShutdownManager();
+    worker.shutdown(shutdownManager, true);
+    CountDownLatch awaitStarted = new CountDownLatch(1);
+    CompletableFuture<Void> await =
+        CompletableFuture.runAsync(
+            () -> {
+              awaitStarted.countDown();
+              worker.awaitTermination(5, TimeUnit.SECONDS);
+            });
+    assertTrue(awaitStarted.await(5, TimeUnit.SECONDS));
+    try {
+      await.get(100, TimeUnit.MILLISECONDS);
+      fail("awaitTermination returned before staged shutdown completed");
+    } catch (TimeoutException expected) {
+      // Expected until the first shutdown stage completes.
+    }
+    verify(poller, never()).awaitTermination(anyLong(), any(TimeUnit.class));
+
+    pollerShutdown.complete(null);
+    await.get(5, TimeUnit.SECONDS);
+    verify(poller).awaitTermination(anyLong(), eq(TimeUnit.MILLISECONDS));
+    shutdownManager.close();
+  }
+
   @Test
   public void concurrentPollRequestLockTest() throws Exception {
     // Test that if the server sends multiple concurrent workflow tasks for the same workflow the