[SDK Sentinel] Fix WorkerFactory shutdown await race

temporal-sdk/src/main/java/io/temporal/worker/WorkerFactory.java

@@ -31,10 +31,12 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.SynchronousQueue;
 import java.util.concurrent.ThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.BiConsumer;
 import java.util.function.Consumer;
@@ -68,6 +70,7 @@ public final class WorkerFactory {
   private final NamespaceCapabilities namespaceCapabilities = new NamespaceCapabilities();
 
   private SuspendableWorker workerCommandWorker;
+  private volatile CompletableFuture<Void> shutdownFuture;
 
   private State state = State.Initial;
 
@@ -476,7 +479,8 @@ private void doShutdown(boolean interruptUserTasks) {
       shutdownFutures.add(workerCommandWorker.shutdown(shutdownManager, true));
     }
 
-    CompletableFuture.allOf(shutdownFutures.toArray(new CompletableFuture[0]))
+    shutdownFuture = CompletableFuture.allOf(shutdownFutures.toArray(new CompletableFuture[0]));
+    shutdownFuture
         .thenApply(
             r -> {
               // Unregister workers from heartbeat manager only after full shutdown,
@@ -520,9 +524,27 @@ public void awaitTermination(long timeout, TimeUnit unit) {
               t, () -> worker.awaitTermination(t, TimeUnit.MILLISECONDS));
     }
     if (workerCommandWorker != null) {
+      long t = timeoutMillis;
+      timeoutMillis =
+          ShutdownManager.runAndGetRemainingTimeoutMs(
+              t, () -> workerCommandWorker.awaitTermination(t, TimeUnit.MILLISECONDS));
+    }
+    CompletableFuture<Void> currentShutdownFuture = shutdownFuture;
+    if (currentShutdownFuture != null) {
       long t = timeoutMillis;
       ShutdownManager.runAndGetRemainingTimeoutMs(
-          t, () -> workerCommandWorker.awaitTermination(t, TimeUnit.MILLISECONDS));
+          t,
+          () -> {
+            try {
+              currentShutdownFuture.get(t, TimeUnit.MILLISECONDS);
+            } catch (InterruptedException e) {
+              Thread.currentThread().interrupt();
+            } catch (ExecutionException e) {
+              log.warn("Exception while waiting for worker factory termination", e.getCause());
+            } catch (TimeoutException e) {
+              // The caller supplied the timeout, so returning is the expected behavior.
+            }
+          });
     }
     log.debug("awaitTermination done: {}", this);
   }

temporal-sdk/src/test/java/io/temporal/worker/WorkerShutdownTest.java

@@ -75,6 +75,11 @@ public void awaitTerminationDoesNotRaceWithWorkerCommandWorkerCleanup() throws E
                 });
 
         factory.awaitTermination(1, TimeUnit.SECONDS);
+        // One wait covers the command worker and the other covers the asynchronous factory
+        // shutdown chain.
+        shutdownManager.verify(
+            () -> ShutdownManager.runAndGetRemainingTimeoutMs(anyLong(), any(Runnable.class)),
+            times(2));
       }
 
       verify(workerCommandWorker).awaitTermination(1_000, TimeUnit.MILLISECONDS);