Shutdown task loss prevention (#2820)

Author: Sushisource

.github/workflows/ci.yml

@@ -60,7 +60,7 @@ jobs:
           report_paths: "**/build/test-results/test/TEST-*.xml"
 
   unit_test_jdk8:
-    name: Unit test with docker service [JDK8]
+    name: Unit test with CLI
     runs-on: ubuntu-latest-16-cores
     timeout-minutes: 30
     steps:
@@ -82,7 +82,7 @@ jobs:
       - name: Set up Gradle
         uses: gradle/actions/setup-gradle@ac396bf1a80af16236baf54bd7330ae21dc6ece5 # v6
 
-      - name: Start containerized server and dependencies
+      - name: Start CLI server
         env:
           TEMPORAL_CLI_VERSION: 1.7.0
         run: |
@@ -110,6 +110,7 @@ jobs:
               --dynamic-config-value history.enableRequestIdRefLinks=true \
               --dynamic-config-value frontend.WorkerHeartbeatsEnabled=true \
               --dynamic-config-value frontend.ListWorkersEnabled=true \
+              --dynamic-config-value frontend.enableCancelWorkerPollsOnShutdown=true \
               --dynamic-config-value 'component.callbacks.allowedAddresses=[{"Pattern":"localhost:7243","AllowInsecure":true}]' \
               --dynamic-config-value frontend.activityAPIsEnabled=true \
               --dynamic-config-value activity.enableStandalone=true \

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

@@ -37,6 +37,7 @@ public ActivityPollTask(
       @Nonnull String namespace,
       @Nonnull String taskQueue,
       @Nonnull String identity,
+      @Nonnull String workerInstanceKey,
       @Nonnull WorkerVersioningOptions versioningOptions,
       double activitiesPerSecond,
       @Nonnull TrackingSlotSupplier<ActivitySlotInfo> slotSupplier,
@@ -53,6 +54,7 @@ public ActivityPollTask(
             .setNamespace(namespace)
             .setIdentity(identity)
             .setTaskQueue(TaskQueue.newBuilder().setName(taskQueue));
+    pollRequest.setWorkerInstanceKey(workerInstanceKey);
     if (activitiesPerSecond > 0) {
       pollRequest.setTaskQueueMetadata(
           TaskQueueMetadata.newBuilder()

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

@@ -105,6 +105,7 @@ public boolean start() {
                     namespace,
                     taskQueue,
                     options.getIdentity(),
+                    options.getWorkerInstanceKey(),
                     options.getWorkerVersioningOptions(),
                     taskQueueActivitiesPerSecond,
                     this.slotSupplier,
@@ -113,7 +114,7 @@ public boolean start() {
                     pollerTracker),
                 this.pollTaskExecutor,
                 pollerOptions,
-                namespaceCapabilities.isPollerAutoscaling(),
+                namespaceCapabilities,
                 workerMetricsScope);
 
       } else {
@@ -125,6 +126,7 @@ public boolean start() {
                     namespace,
                     taskQueue,
                     options.getIdentity(),
+                    options.getWorkerInstanceKey(),
                     options.getWorkerVersioningOptions(),
                     taskQueueActivitiesPerSecond,
                     this.slotSupplier,
@@ -133,7 +135,8 @@ public boolean start() {
                     pollerTracker),
                 this.pollTaskExecutor,
                 pollerOptions,
-                workerMetricsScope);
+                workerMetricsScope,
+                namespaceCapabilities);
       }
       poller.start();
       workerMetricsScope.counter(MetricsType.WORKER_START_COUNTER).inc(1);

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

@@ -43,6 +43,7 @@ public AsyncActivityPollTask(
       @Nonnull String namespace,
       @Nonnull String taskQueue,
       @Nonnull String identity,
+      @Nonnull String workerInstanceKey,
       @Nonnull WorkerVersioningOptions versioningOptions,
       double activitiesPerSecond,
       @Nonnull TrackingSlotSupplier<ActivitySlotInfo> slotSupplier,
@@ -59,6 +60,7 @@ public AsyncActivityPollTask(
             .setNamespace(namespace)
             .setIdentity(identity)
             .setTaskQueue(TaskQueue.newBuilder().setName(taskQueue));
+    pollRequest.setWorkerInstanceKey(workerInstanceKey);
     if (activitiesPerSecond > 0) {
       pollRequest.setTaskQueueMetadata(
           TaskQueueMetadata.newBuilder()

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

@@ -41,6 +41,7 @@ public AsyncNexusPollTask(
       @Nonnull String namespace,
       @Nonnull String taskQueue,
       @Nonnull String identity,
+      @Nonnull String workerInstanceKey,
       @Nonnull WorkerVersioningOptions versioningOptions,
       @Nonnull Scope metricsScope,
       @Nonnull Supplier<GetSystemInfoResponse.Capabilities> serverCapabilities,
@@ -57,6 +58,8 @@ public AsyncNexusPollTask(
             .setIdentity(identity)
             .setTaskQueue(TaskQueue.newBuilder().setName(taskQueue));
 
+    pollRequest.setWorkerInstanceKey(workerInstanceKey);
+
     if (versioningOptions.getWorkerDeploymentOptions() != null) {
       pollRequest.setDeploymentOptions(
           WorkerVersioningProtoUtils.deploymentOptionsToProto(

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

@@ -29,7 +29,6 @@ final class AsyncPoller<T extends ScalingTask> extends BasePoller<T> {
   private final List<PollTaskAsync<T>> asyncTaskPollers;
   private final PollerOptions pollerOptions;
   private final PollerBehaviorAutoscaling pollerBehavior;
-  private final boolean serverSupportsAutoscaling;
   private final Scope workerMetricsScope;
   private Throttler pollRateThrottler;
   private final Thread.UncaughtExceptionHandler uncaughtExceptionHandler =
@@ -43,15 +42,15 @@ final class AsyncPoller<T extends ScalingTask> extends BasePoller<T> {
       PollTaskAsync<T> asyncTaskPoller,
       ShutdownableTaskExecutor<T> taskExecutor,
       PollerOptions pollerOptions,
-      boolean serverSupportsAutoscaling,
+      NamespaceCapabilities namespaceCapabilities,
       Scope workerMetricsScope) {
     this(
         slotSupplier,
         slotReservationData,
         Collections.singletonList(asyncTaskPoller),
         taskExecutor,
         pollerOptions,
-        serverSupportsAutoscaling,
+        namespaceCapabilities,
         workerMetricsScope);
   }
 
@@ -61,9 +60,9 @@ final class AsyncPoller<T extends ScalingTask> extends BasePoller<T> {
       List<PollTaskAsync<T>> asyncTaskPollers,
       ShutdownableTaskExecutor<T> taskExecutor,
       PollerOptions pollerOptions,
-      boolean serverSupportsAutoscaling,
+      NamespaceCapabilities namespaceCapabilities,
       Scope workerMetricsScope) {
-    super(taskExecutor);
+    super(taskExecutor, namespaceCapabilities);
     Objects.requireNonNull(slotSupplier, "slotSupplier cannot be null");
     Objects.requireNonNull(slotReservationData, "slotReservation data should not be null");
     Objects.requireNonNull(asyncTaskPollers, "asyncTaskPollers should not be null");
@@ -82,7 +81,6 @@ final class AsyncPoller<T extends ScalingTask> extends BasePoller<T> {
               + " is not supported for AsyncPoller. Only PollerBehaviorAutoscaling is supported.");
     }
     this.pollerBehavior = (PollerBehaviorAutoscaling) pollerOptions.getPollerBehavior();
-    this.serverSupportsAutoscaling = serverSupportsAutoscaling;
     this.pollerOptions = pollerOptions;
     this.workerMetricsScope = workerMetricsScope;
   }
@@ -114,7 +112,7 @@ public boolean start() {
               pollerBehavior.getMinConcurrentTaskPollers(),
               pollerBehavior.getMaxConcurrentTaskPollers(),
               pollerBehavior.getInitialConcurrentTaskPollers(),
-              serverSupportsAutoscaling,
+              namespaceCapabilities.isPollerAutoscaling(),
               (newTarget) -> {
                 log.debug(
                     "Updating maximum number of pollers for {} to: {}",
@@ -136,12 +134,14 @@ public CompletableFuture<Void> shutdown(ShutdownManager shutdownManager, boolean
     return super.shutdown(shutdownManager, interruptTasks)
         .thenApply(
             (f) -> {
-              for (PollTaskAsync<T> asyncTaskPoller : asyncTaskPollers) {
-                try {
-                  log.debug("Shutting down async poller: {}", asyncTaskPoller.getLabel());
-                  asyncTaskPoller.cancel(new RuntimeException("Shutting down poller"));
-                } catch (Throwable e) {
-                  log.error("Error while cancelling poll task", e);
+              if (interruptTasks || !namespaceCapabilities.isGracefulPollShutdown()) {
+                for (PollTaskAsync<T> asyncTaskPoller : asyncTaskPollers) {
+                  try {
+                    log.debug("Shutting down async poller: {}", asyncTaskPoller.getLabel());
+                    asyncTaskPoller.cancel(new RuntimeException("Shutting down poller"));
+                  } catch (Throwable e) {
+                    log.error("Error while cancelling poll task", e);
+                  }
                 }
               }
               return null;

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

@@ -52,6 +52,7 @@ public AsyncWorkflowPollTask(
       @Nonnull String taskQueue,
       @Nullable String stickyTaskQueue,
       @Nonnull String identity,
+      @Nonnull String workerInstanceKey,
       @Nonnull WorkerVersioningOptions versioningOptions,
       @Nonnull TrackingSlotSupplier<WorkflowSlotInfo> slotSupplier,
       @Nonnull Scope metricsScope,
@@ -67,6 +68,8 @@ public AsyncWorkflowPollTask(
             .setNamespace(Objects.requireNonNull(namespace))
             .setIdentity(Objects.requireNonNull(identity));
 
+    pollRequestBuilder.setWorkerInstanceKey(workerInstanceKey);
+
     if (versioningOptions.getWorkerDeploymentOptions() != null) {
       pollRequestBuilder.setDeploymentOptions(
           WorkerVersioningProtoUtils.deploymentOptionsToProto(

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

@@ -27,9 +27,14 @@ abstract class BasePoller<T> implements SuspendableWorker {
 
   protected ExecutorService pollExecutor;
 
-  protected BasePoller(ShutdownableTaskExecutor<T> taskExecutor) {
+  protected final NamespaceCapabilities namespaceCapabilities;
+
+  protected BasePoller(
+      ShutdownableTaskExecutor<T> taskExecutor, NamespaceCapabilities namespaceCapabilities) {
     Objects.requireNonNull(taskExecutor, "taskExecutor should not be null");
     this.taskExecutor = taskExecutor;
+    this.namespaceCapabilities =
+        Objects.requireNonNull(namespaceCapabilities, "namespaceCapabilities should not be null");
   }
 
   @Override
@@ -55,15 +60,24 @@ public CompletableFuture<Void> shutdown(ShutdownManager shutdownManager, boolean
         return CompletableFuture.completedFuture(null);
     }
 
-    return shutdownManager
-        // it's ok to forcefully shutdown pollers, because they are stuck in a long poll call
-        // so we don't risk loosing any progress doing that.
-        .shutdownExecutorNow(pollExecutor, this + "#pollExecutor", Duration.ofSeconds(1))
-        .exceptionally(
-            e -> {
-              log.error("Unexpected exception during shutdown", e);
-              return null;
-            });
+    CompletableFuture<Void> pollExecutorShutdown;
+    if (namespaceCapabilities.isGracefulPollShutdown() && !interruptTasks) {
+      // When graceful poll shutdown is enabled, the server will complete outstanding polls with
+      // empty responses after ShutdownWorker is called. We simply wait for polls to return.
+      pollExecutorShutdown =
+          shutdownManager.shutdownExecutor(
+              pollExecutor, this + "#pollExecutor", Duration.ofSeconds(80));
+    } else {
+      // ShutdownNow and old servers forcibly stop outstanding polls.
+      pollExecutorShutdown =
+          shutdownManager.shutdownExecutorNow(
+              pollExecutor, this + "#pollExecutor", Duration.ofSeconds(1));
+    }
+    return pollExecutorShutdown.exceptionally(
+        e -> {
+          log.error("Unexpected exception during shutdown", e);
+          return null;
+        });
   }
 
   @Override

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

@@ -52,8 +52,9 @@ public MultiThreadedPoller(
       PollTask<T> pollTask,
       ShutdownableTaskExecutor<T> taskExecutor,
       PollerOptions pollerOptions,
-      Scope workerMetricsScope) {
-    super(taskExecutor);
+      Scope workerMetricsScope,
+      NamespaceCapabilities namespaceCapabilities) {
+    super(taskExecutor, namespaceCapabilities);
     Objects.requireNonNull(identity, "identity cannot be null");
     Objects.requireNonNull(pollTask, "poll service should not be null");
     Objects.requireNonNull(pollerOptions, "pollerOptions should not be null");

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

@@ -1,5 +1,6 @@
 package io.temporal.internal.worker;
 
+import io.temporal.api.namespace.v1.NamespaceInfo.Capabilities;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 /**
@@ -9,14 +10,31 @@
  */
 public final class NamespaceCapabilities {
   private final AtomicBoolean pollerAutoscaling = new AtomicBoolean(false);
+  private final AtomicBoolean gracefulPollShutdown = new AtomicBoolean(false);
   private final AtomicBoolean workerHeartbeats = new AtomicBoolean(false);
 
+  public void setFromCapabilities(Capabilities capabilities) {
+    if (capabilities.getPollerAutoscaling()) {
+      pollerAutoscaling.set(true);
+    }
+    if (capabilities.getWorkerPollCompleteOnShutdown()) {
+      gracefulPollShutdown.set(true);
+    }
+    if (capabilities.getWorkerHeartbeats()) {
+      workerHeartbeats.set(true);
+    }
+  }
+
   public boolean isPollerAutoscaling() {
     return pollerAutoscaling.get();
   }
 
-  public void setPollerAutoscaling(boolean value) {
-    pollerAutoscaling.set(value);
+  public boolean isGracefulPollShutdown() {
+    return gracefulPollShutdown.get();
+  }
+
+  public void setGracefulPollShutdown(boolean value) {
+    gracefulPollShutdown.set(value);
   }
 
   public boolean isWorkerHeartbeats() {