Fix test failures

Author: Sushisource

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

@@ -134,7 +134,7 @@ public CompletableFuture<Void> shutdown(ShutdownManager shutdownManager, boolean
     return super.shutdown(shutdownManager, interruptTasks)
         .thenApply(
             (f) -> {
-              if (!namespaceCapabilities.isGracefulPollShutdown()) {
+              if (interruptTasks || !namespaceCapabilities.isGracefulPollShutdown()) {
                 for (PollTaskAsync<T> asyncTaskPoller : asyncTaskPollers) {
                   try {
                     log.debug("Shutting down async poller: {}", asyncTaskPoller.getLabel());

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

@@ -61,14 +61,14 @@ public CompletableFuture<Void> shutdown(ShutdownManager shutdownManager, boolean
     }
 
     CompletableFuture<Void> pollExecutorShutdown;
-    if (namespaceCapabilities.isGracefulPollShutdown()) {
+    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 {
-      // Old behaviour forcibly stops outstanding polls.
+      // ShutdownNow and old servers forcibly stop outstanding polls.
       pollExecutorShutdown =
           shutdownManager.shutdownExecutorNow(
               pollExecutor, this + "#pollExecutor", Duration.ofSeconds(1));

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

@@ -149,10 +149,18 @@ public CompletableFuture<Void> waitOnWorkerShutdownRequest(
             future.complete(null);
           }
         },
-        scheduledExecutorService);
+        this::executeOrRunDirect);
     return future;
   }
 
+  private void executeOrRunDirect(Runnable command) {
+    try {
+      scheduledExecutorService.execute(command);
+    } catch (RejectedExecutionException e) {
+      command.run();
+    }
+  }
+
   @Override
   public void close() {
     scheduledExecutorService.shutdownNow();

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

@@ -214,7 +214,8 @@ public CompletableFuture<Void> shutdown(ShutdownManager shutdownManager, boolean
         !interruptTasks
             && !options.getDrainStickyTaskQueueTimeout().isZero()
             && stickyTaskQueueName != null
-            && stickyQueueBalancer != null;
+            && stickyQueueBalancer != null
+            && !namespaceCapabilities.isGracefulPollShutdown();
 
     CompletableFuture<Void> pollerShutdown =
         CompletableFuture.completedFuture(null)

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

@@ -147,4 +147,100 @@ public synchronized String poll() {
 
     shutdownManager.close();
   }
+
+  @Test(timeout = 10_000)
+  public void shutdownNowInterruptsInflightPollWhenGraceful() throws Exception {
+    NamespaceCapabilities capabilities = new NamespaceCapabilities();
+    capabilities.setFromCapabilities(
+        Capabilities.newBuilder().setWorkerPollCompleteOnShutdown(true).build());
+
+    AtomicReference<String> processedTask = new AtomicReference<>();
+    CountDownLatch taskProcessedLatch = new CountDownLatch(1);
+    ShutdownableTaskExecutor<String> taskExecutor =
+        new ShutdownableTaskExecutor<String>() {
+          @Override
+          public void process(@Nonnull String task) {
+            processedTask.set(task);
+            taskProcessedLatch.countDown();
+          }
+
+          @Override
+          public boolean isShutdown() {
+            return false;
+          }
+
+          @Override
+          public boolean isTerminated() {
+            return false;
+          }
+
+          @Override
+          public CompletableFuture<Void> shutdown(
+              ShutdownManager shutdownManager, boolean interruptTasks) {
+            return CompletableFuture.completedFuture(null);
+          }
+
+          @Override
+          public void awaitTermination(long timeout, TimeUnit unit) {}
+        };
+
+    CountDownLatch secondPollStarted = new CountDownLatch(1);
+    CountDownLatch releasePoll = new CountDownLatch(1);
+
+    MultiThreadedPoller.PollTask<String> pollTask =
+        new MultiThreadedPoller.PollTask<String>() {
+          private int callCount = 0;
+
+          @Override
+          public synchronized String poll() {
+            callCount++;
+            if (callCount == 1) {
+              return "task-1";
+            } else if (callCount == 2) {
+              secondPollStarted.countDown();
+              try {
+                releasePoll.await();
+              } catch (InterruptedException e) {
+                Thread.currentThread().interrupt();
+                return null;
+              }
+              return "task-2";
+            }
+            try {
+              Thread.sleep(Long.MAX_VALUE);
+            } catch (InterruptedException e) {
+              Thread.currentThread().interrupt();
+            }
+            return null;
+          }
+        };
+
+    MultiThreadedPoller<String> poller =
+        new MultiThreadedPoller<>(
+            "test-identity",
+            pollTask,
+            taskExecutor,
+            PollerOptions.newBuilder()
+                .setPollerBehavior(new PollerBehaviorSimpleMaximum(1))
+                .setPollThreadNamePrefix("test-poller")
+                .build(),
+            new NoopScope(),
+            capabilities);
+
+    poller.start();
+
+    assertTrue("first task should be processed", taskProcessedLatch.await(5, TimeUnit.SECONDS));
+    assertEquals("task-1", processedTask.get());
+    assertTrue("second poll should start", secondPollStarted.await(5, TimeUnit.SECONDS));
+
+    ShutdownManager shutdownManager = new ShutdownManager();
+    try {
+      poller.shutdown(shutdownManager, true).get(5, TimeUnit.SECONDS);
+      assertNotEquals(
+          "task-2 should not be processed by shutdownNow", "task-2", processedTask.get());
+    } finally {
+      releasePoll.countDown();
+      shutdownManager.close();
+    }
+  }
 }

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

@@ -97,6 +97,7 @@ public void activeTaskQueueTypesEvaluatedAtShutdownTime() throws Exception {
         .thenReturn(
             WorkflowClientOptions.newBuilder()
                 .setNamespace("test-ns")
+                .setIdentity("test-worker")
                 .validateAndBuildWithDefaults());
 
     Scope metricsScope = new NoopScope();
@@ -113,7 +114,7 @@ public void activeTaskQueueTypesEvaluatedAtShutdownTime() throws Exception {
             metricsScope,
             runLocks,
             cache,
-            false,
+            true,
             wfThreadExecutor,
             Collections.emptyList(),
             Collections.emptyList(),
@@ -147,5 +148,8 @@ public void activeTaskQueueTypesEvaluatedAtShutdownTime() throws Exception {
         "ShutdownWorkerRequest heartbeat should report SHUTTING_DOWN",
         WorkerStatus.WORKER_STATUS_SHUTTING_DOWN,
         captor.getValue().getWorkerHeartbeat().getStatus());
+    assertTrue(
+        "ShutdownWorkerRequest sticky task queue should be derived from worker identity",
+        captor.getValue().getStickyTaskQueue().startsWith("test-worker:"));
   }
 }

temporal-sdk/src/test/java/io/temporal/worker/shutdown/StickyWorkflowDrainShutdownTest.java

@@ -3,6 +3,8 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 
+import io.temporal.api.workflowservice.v1.DescribeNamespaceRequest;
+import io.temporal.api.workflowservice.v1.DescribeNamespaceResponse;
 import io.temporal.client.WorkflowClient;
 import io.temporal.client.WorkflowStub;
 import io.temporal.serviceclient.WorkflowServiceStubsOptions;
@@ -53,6 +55,7 @@ public StickyWorkflowDrainShutdownTest(PollerBehavior pollerBehaviorAutoscaling)
 
   @Test
   public void testShutdown() throws InterruptedException {
+    boolean gracefulPollShutdownSupported = isGracefulPollShutdownSupported();
     TestWorkflow1 workflow = testWorkflowRule.newWorkflowStub(TestWorkflow1.class);
     WorkflowClient.start(workflow::execute, null);
     testWorkflowRule.getTestEnvironment().shutdown();
@@ -62,10 +65,16 @@ public void testShutdown() throws InterruptedException {
     assertTrue(testWorkflowRule.getTestEnvironment().getWorkerFactory().isTerminated());
     System.out.println("Shutdown completed");
     long endTime = System.currentTimeMillis();
-    assertTrue("Drain time should be respected", endTime - startTime > DRAIN_TIME.toMillis());
-    // Workflow should complete successfully since the drain time is longer than the workflow
-    // execution time
-    assertEquals("Success", workflow.execute(null));
+    WorkflowStub untyped = WorkflowStub.fromTyped(workflow);
+    if (gracefulPollShutdownSupported) {
+      assertTrue("Drain time should be skipped", endTime - startTime < DRAIN_TIME.toMillis());
+      untyped.terminate("test cleanup");
+    } else {
+      assertTrue("Drain time should be respected", endTime - startTime > DRAIN_TIME.toMillis());
+      // Workflow should complete successfully since the drain time is longer than the workflow
+      // execution time.
+      assertEquals("Success", untyped.getResult(String.class));
+    }
   }
 
   @Test
@@ -94,4 +103,17 @@ public String execute(String now) {
       return "Success";
     }
   }
+
+  private boolean isGracefulPollShutdownSupported() {
+    DescribeNamespaceResponse response =
+        testWorkflowRule
+            .getWorkflowClient()
+            .getWorkflowServiceStubs()
+            .blockingStub()
+            .describeNamespace(
+                DescribeNamespaceRequest.newBuilder()
+                    .setNamespace(testWorkflowRule.getWorkflowClient().getOptions().getNamespace())
+                    .build());
+    return response.getNamespaceInfo().getCapabilities().getWorkerPollCompleteOnShutdown();
+  }
 }