Wait for worker slots to be fully released in the graceful worker shutdown (#1679)

Author: Spikhalskiy

temporal-sdk/src/main/java/io/temporal/internal/sync/POJOWorkflowImplementationFactory.java

@@ -62,6 +62,9 @@
 import org.slf4j.LoggerFactory;
 
 public final class POJOWorkflowImplementationFactory implements ReplayWorkflowFactory {
+  private static final Logger log =
+      LoggerFactory.getLogger(POJOWorkflowImplementationFactory.class);
+
   public static final ImmutableSet<String> WORKFLOW_HANDLER_STACKTRACE_CUTOFF =
       ImmutableSet.<String>builder()
           // POJO
@@ -71,13 +74,8 @@ public final class POJOWorkflowImplementationFactory implements ReplayWorkflowFa
           // Dynamic
           .add(
               ReflectionUtils.getMethodNameForStackTraceCutoff(
-                  DynamicSyncWorkflowDefinition.RootWorkflowInboundCallsInterceptor.class,
-                  "execute",
-                  WorkflowInboundCallsInterceptor.WorkflowInput.class))
+                  DynamicSyncWorkflowDefinition.class, "execute", Header.class, Optional.class))
           .build();
-
-  private static final Logger log =
-      LoggerFactory.getLogger(POJOWorkflowImplementationFactory.class);
   private final WorkerInterceptor[] workerInterceptors;
 
   private final DataConverter dataConverter;

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

@@ -64,6 +64,7 @@ final class ActivityWorker implements SuspendableWorker {
   private final Scope workerMetricsScope;
   private final GrpcRetryer grpcRetryer;
   private final GrpcRetryer.GrpcRetryerOptions replyGrpcRetryerOptions;
+  private final int executorSlots;
   private final Semaphore executorSlotsSemaphore;
 
   public ActivityWorker(
@@ -86,7 +87,8 @@ public ActivityWorker(
     this.replyGrpcRetryerOptions =
         new GrpcRetryer.GrpcRetryerOptions(
             DefaultStubServiceOperationRpcRetryOptions.INSTANCE, null);
-    this.executorSlotsSemaphore = new Semaphore(options.getTaskExecutorThreadPoolSize());
+    this.executorSlots = options.getTaskExecutorThreadPoolSize();
+    this.executorSlotsSemaphore = new Semaphore(executorSlots);
   }
 
   @Override
@@ -126,12 +128,33 @@ public boolean start() {
 
   @Override
   public CompletableFuture<Void> shutdown(ShutdownManager shutdownManager, boolean interruptTasks) {
-    return poller.shutdown(shutdownManager, interruptTasks);
+    String semaphoreName = this + "#executorSlotsSemaphore";
+    return poller
+        .shutdown(shutdownManager, interruptTasks)
+        .thenCompose(
+            ignore ->
+                !interruptTasks
+                    ? shutdownManager.waitForSemaphorePermitsReleaseUntimed(
+                        executorSlotsSemaphore, executorSlots, semaphoreName)
+                    : 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;
+            });
   }
 
   @Override
   public void awaitTermination(long timeout, TimeUnit unit) {
-    poller.awaitTermination(timeout, 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
+    ShutdownManager.awaitTermination(pollTaskExecutor, timeoutMillis);
   }
 
   @Override
@@ -151,7 +174,7 @@ public boolean isShutdown() {
 
   @Override
   public boolean isTerminated() {
-    return poller.isTerminated();
+    return poller.isTerminated() && (pollTaskExecutor == null || pollTaskExecutor.isTerminated());
   }
 
   @Override
@@ -180,6 +203,13 @@ private PollerOptions getPollerOptions(SingleWorkerOptions options) {
     return pollerOptions;
   }
 
+  @Override
+  public String toString() {
+    return String.format(
+        "ActivityWorker{identity=%s, namespace=%s, taskQueue=%s}",
+        options.getIdentity(), namespace, taskQueue);
+  }
+
   private class TaskHandlerImpl implements PollTaskExecutor.TaskHandler<ActivityTask> {
 
     final ActivityTaskHandler handler;

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

@@ -134,9 +134,6 @@ public CompletableFuture<Void> shutdown(ShutdownManager shutdownManager, boolean
         // it's ok to forcefully shutdown pollers, especially because they stuck in a long poll call
         // we don't lose any progress doing that
         .shutdownExecutorNow(pollExecutor, this + "#pollExecutor", Duration.ofSeconds(1))
-        // TODO Poller shouldn't shutdown taskExecutor, because it gets it already created
-        //  externally. Creator of taskExecutor should be responsible for it's shutdown.
-        .thenCompose(ignore -> taskExecutor.shutdown(shutdownManager, interruptTasks))
         .exceptionally(
             e -> {
               log.error("Unexpected exception during shutdown", e);
@@ -154,8 +151,7 @@ public void awaitTermination(long timeout, TimeUnit unit) {
     }
 
     long timeoutMillis = unit.toMillis(timeout);
-    timeoutMillis = ShutdownManager.awaitTermination(pollExecutor, timeoutMillis);
-    ShutdownManager.awaitTermination(taskExecutor, timeoutMillis);
+    ShutdownManager.awaitTermination(pollExecutor, timeoutMillis);
   }
 
   @Override

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

@@ -68,6 +68,15 @@ public CompletableFuture<Void> shutdownExecutorUntimed(
     return untimedWait(executorToShutdown, executorName);
   }
 
+  public CompletableFuture<Void> waitForSemaphorePermitsReleaseUntimed(
+      Semaphore semaphore, int initialSemaphorePermits, String semaphoreName) {
+    CompletableFuture<Void> future = new CompletableFuture<>();
+    scheduledExecutorService.submit(
+        new SemaphoreReportingDelayShutdown(
+            semaphore, initialSemaphorePermits, semaphoreName, future));
+    return future;
+  }
+
   /**
    * Wait for {@code executorToShutdown} to terminate. Only completes the returned CompletableFuture
    * when the executor is terminated.
@@ -76,7 +85,7 @@ private CompletableFuture<Void> untimedWait(
       ExecutorService executorToShutdown, String executorName) {
     CompletableFuture<Void> future = new CompletableFuture<>();
     scheduledExecutorService.submit(
-        new ReportingDelayShutdown(executorToShutdown, executorName, future));
+        new ExecutorReportingDelayShutdown(executorToShutdown, executorName, future));
     return future;
   }
 
@@ -91,7 +100,7 @@ private CompletableFuture<Void> limitedWait(
 
     CompletableFuture<Void> future = new CompletableFuture<>();
     scheduledExecutorService.submit(
-        new LimitedWaitShutdown(executorToShutdown, attempts, executorName, future));
+        new ExecutorLimitedWaitShutdown(executorToShutdown, attempts, executorName, future));
     return future;
   }
 
@@ -100,67 +109,90 @@ public void close() {
     scheduledExecutorService.shutdownNow();
   }
 
-  private class LimitedWaitShutdown implements Runnable {
-    private final ExecutorService executorToShutdown;
+  private abstract class LimitedWaitShutdown implements Runnable {
     private final CompletableFuture<Void> promise;
     private final int maxAttempts;
-    private final String executorName;
     private int attempt;
 
-    public LimitedWaitShutdown(
-        ExecutorService executorToShutdown,
-        int maxAttempts,
-        String executorName,
-        CompletableFuture<Void> promise) {
-      this.executorToShutdown = executorToShutdown;
+    public LimitedWaitShutdown(int maxAttempts, CompletableFuture<Void> promise) {
       this.promise = promise;
       this.maxAttempts = maxAttempts;
-      this.executorName = executorName;
     }
 
     @Override
     public void run() {
-      if (executorToShutdown.isTerminated()) {
+      if (isTerminated()) {
+        onSuccessfulTermination();
         promise.complete(null);
         return;
       }
       attempt++;
       if (attempt > maxAttempts) {
-        log.warn(
-            "Wait for a graceful shutdown of {} timed out, fallback to shutdownNow()",
-            executorName);
-        executorToShutdown.shutdownNow();
+        onAttemptExhaustion();
         // we don't want to complicate shutdown with dealing of exceptions and errors of all sorts,
         // so just log and complete the promise
         promise.complete(null);
         return;
       }
       scheduledExecutorService.schedule(this, CHECK_PERIOD_MS, TimeUnit.MILLISECONDS);
     }
+
+    abstract boolean isTerminated();
+
+    abstract void onAttemptExhaustion();
+
+    abstract void onSuccessfulTermination();
   }
 
-  private class ReportingDelayShutdown implements Runnable {
+  private class ExecutorLimitedWaitShutdown extends LimitedWaitShutdown {
+    private final ExecutorService executorToShutdown;
+    private final String executorName;
+
+    public ExecutorLimitedWaitShutdown(
+        ExecutorService executorToShutdown,
+        int maxAttempts,
+        String executorName,
+        CompletableFuture<Void> promise) {
+      super(maxAttempts, promise);
+      this.executorToShutdown = executorToShutdown;
+      this.executorName = executorName;
+    }
+
+    @Override
+    boolean isTerminated() {
+      return executorToShutdown.isTerminated();
+    }
+
+    @Override
+    void onAttemptExhaustion() {
+      log.warn(
+          "Wait for a graceful shutdown of {} timed out, fallback to shutdownNow()", executorName);
+      executorToShutdown.shutdownNow();
+    }
+
+    @Override
+    void onSuccessfulTermination() {}
+  }
+
+  private abstract class ReportingDelayShutdown implements Runnable {
+    // measures in attempts count, not in ms
     private static final int BLOCKED_REPORTING_THRESHOLD = 60;
     private static final int BLOCKED_REPORTING_PERIOD = 20;
 
-    private final ExecutorService executorToShutdown;
     private final CompletableFuture<Void> promise;
-    private final String executorName;
     private int attempt;
 
-    public ReportingDelayShutdown(
-        ExecutorService executorToShutdown, String executorName, CompletableFuture<Void> promise) {
-      this.executorToShutdown = executorToShutdown;
+    public ReportingDelayShutdown(CompletableFuture<Void> promise) {
       this.promise = promise;
-      this.executorName = executorName;
     }
 
     @Override
     public void run() {
-      if (executorToShutdown.isTerminated()) {
+      if (isTerminated()) {
         if (attempt > BLOCKED_REPORTING_THRESHOLD) {
-          // log warn only if we already logged a shutdown being delayed
-          log.warn("{} successfully terminated", executorName);
+          onSlowSuccessfulTermination();
+        } else {
+          onSuccessfulTermination();
         }
         promise.complete(null);
         return;
@@ -170,13 +202,87 @@ public void run() {
       if (attempt >= BLOCKED_REPORTING_THRESHOLD) {
         // and repeat every BLOCKED_REPORTING_PERIOD attempts
         if (((float) (attempt - BLOCKED_REPORTING_THRESHOLD) % BLOCKED_REPORTING_PERIOD) < 0.001) {
-          log.warn(
-              "Graceful shutdown of {} is blocked by one of the long currently processing tasks",
-              executorName);
+          onSlowTermination();
         }
       }
       scheduledExecutorService.schedule(this, CHECK_PERIOD_MS, TimeUnit.MILLISECONDS);
     }
+
+    abstract boolean isTerminated();
+
+    abstract void onSlowTermination();
+
+    abstract void onSuccessfulTermination();
+
+    /** Called only if {@link #onSlowTermination()} was called before */
+    abstract void onSlowSuccessfulTermination();
+  }
+
+  private class ExecutorReportingDelayShutdown extends ReportingDelayShutdown {
+    private final ExecutorService executorToShutdown;
+    private final String executorName;
+
+    public ExecutorReportingDelayShutdown(
+        ExecutorService executorToShutdown, String executorName, CompletableFuture<Void> promise) {
+      super(promise);
+      this.executorToShutdown = executorToShutdown;
+      this.executorName = executorName;
+    }
+
+    @Override
+    boolean isTerminated() {
+      return executorToShutdown.isTerminated();
+    }
+
+    @Override
+    void onSlowTermination() {
+      log.warn(
+          "Graceful shutdown of {} is blocked by one of the long currently processing tasks",
+          executorName);
+    }
+
+    @Override
+    void onSuccessfulTermination() {}
+
+    @Override
+    void onSlowSuccessfulTermination() {
+      log.warn("{} successfully terminated", executorName);
+    }
+  }
+
+  private class SemaphoreReportingDelayShutdown extends ReportingDelayShutdown {
+    private final Semaphore semaphore;
+    private final int initialSemaphorePermits;
+    private final String semaphoreName;
+
+    public SemaphoreReportingDelayShutdown(
+        Semaphore semaphore,
+        int initialSemaphorePermits,
+        String semaphoreName,
+        CompletableFuture<Void> promise) {
+      super(promise);
+      this.semaphore = semaphore;
+      this.initialSemaphorePermits = initialSemaphorePermits;
+      this.semaphoreName = semaphoreName;
+    }
+
+    @Override
+    boolean isTerminated() {
+      return semaphore.availablePermits() == initialSemaphorePermits;
+    }
+
+    @Override
+    void onSlowTermination() {
+      log.warn("Wait for release of slots of {} takes a long time", semaphoreName);
+    }
+
+    @Override
+    void onSuccessfulTermination() {}
+
+    @Override
+    void onSlowSuccessfulTermination() {
+      log.warn("All slots of {} were successfully released", semaphoreName);
+    }
   }
 
   public static long awaitTermination(@Nullable ExecutorService s, long timeoutMillis) {
@@ -187,7 +293,7 @@ public static long awaitTermination(@Nullable ExecutorService s, long timeoutMil
         timeoutMillis,
         () -> {
           try {
-            s.awaitTermination(timeoutMillis, TimeUnit.MILLISECONDS);
+            boolean ignored = s.awaitTermination(timeoutMillis, TimeUnit.MILLISECONDS);
           } catch (InterruptedException e) {
             Thread.currentThread().interrupt();
           }