Review comments

Author: Sushisource

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

@@ -65,7 +65,8 @@ public CompletableFuture<Void> shutdown(ShutdownManager shutdownManager, boolean
       // 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.shutdownExecutorUntimed(pollExecutor, this + "#pollExecutor");
+          shutdownManager.shutdownExecutor(
+              pollExecutor, this + "#pollExecutor", Duration.ofSeconds(80));
     } else {
       // Old behaviour forcibly stops outstanding polls.
       pollExecutorShutdown =

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;
 
 /**
@@ -13,12 +14,17 @@ public final class NamespaceCapabilities {
   private final AtomicBoolean workerHeartbeats = new AtomicBoolean(false);
 
 
-  public boolean isPollerAutoscaling() {
-    return pollerAutoscaling.get();
+  public void setFromCapabilities(Capabilities capabilities) {
+    if (capabilities.getPollerAutoscaling()) {
+      pollerAutoscaling.set(true);
+    }
+    if (capabilities.getWorkerPollCompleteOnShutdown()) {
+      gracefulPollShutdown.set(true);
+    }
   }
 
-  public void setPollerAutoscaling(boolean value) {
-    pollerAutoscaling.set(value);
+  public boolean isPollerAutoscaling() {
+    return pollerAutoscaling.get();
   }
 
   public boolean isGracefulPollShutdown() {

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

@@ -268,23 +268,8 @@ public synchronized void start() {
                 DescribeNamespaceRequest.newBuilder()
                     .setNamespace(workflowClient.getOptions().getNamespace())
                     .build());
-    if (describeNamespaceResponse.getNamespaceInfo().getCapabilities().getWorkerHeartbeats()) {
-      namespaceCapabilities.setWorkerHeartbeats(true);
-    } else {
-      log.debug(
-          "Server does not support worker heartbeats for namespace {}",
-          workflowClient.getOptions().getNamespace());
-    }
-
-    if (describeNamespaceResponse.getNamespaceInfo().getCapabilities().getPollerAutoscaling()) {
-      namespaceCapabilities.setPollerAutoscaling(true);
-    }
-    if (describeNamespaceResponse
-        .getNamespaceInfo()
-        .getCapabilities()
-        .getWorkerPollCompleteOnShutdown()) {
-      namespaceCapabilities.setGracefulPollShutdown(true);
-    }
+    namespaceCapabilities.setFromCapabilities(
+        describeNamespaceResponse.getNamespaceInfo().getCapabilities());
 
     // Build plugin execution chain (reverse order for proper nesting)
     Consumer<WorkerFactory> startChain = WorkerFactory::doStart;

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

@@ -3,6 +3,7 @@
 import static org.junit.Assert.*;
 
 import com.uber.m3.tally.NoopScope;
+import io.temporal.api.namespace.v1.NamespaceInfo.Capabilities;
 import io.temporal.worker.tuning.PollerBehaviorSimpleMaximum;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CountDownLatch;
@@ -30,7 +31,8 @@ public static Object[] data() {
   @Test(timeout = 10_000)
   public void inflightPollSurvivesShutdownOnlyWhenGraceful() throws Exception {
     NamespaceCapabilities capabilities = new NamespaceCapabilities();
-    capabilities.setGracefulPollShutdown(graceful);
+    capabilities.setFromCapabilities(
+        Capabilities.newBuilder().setWorkerPollCompleteOnShutdown(true).build());
 
     AtomicReference<String> processedTask = new AtomicReference<>();
     CountDownLatch taskProcessedLatch = new CountDownLatch(1);