diff --git a/temporal-sdk/src/main/java/io/temporal/worker/WorkerFactory.java b/temporal-sdk/src/main/java/io/temporal/worker/WorkerFactory.java index dfb9b51326..b9abcec046 100644 --- a/temporal-sdk/src/main/java/io/temporal/worker/WorkerFactory.java +++ b/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 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 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); } diff --git a/temporal-sdk/src/test/java/io/temporal/worker/WorkerShutdownTest.java b/temporal-sdk/src/test/java/io/temporal/worker/WorkerShutdownTest.java index 07af0f5bec..6cb3ba3a13 100644 --- a/temporal-sdk/src/test/java/io/temporal/worker/WorkerShutdownTest.java +++ b/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);