Add more validation

Author: Quinn-With-Two-Ns

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

@@ -11,6 +11,7 @@
 import io.temporal.worker.tuning.SlotSupplierFuture;
 import java.util.Collections;
 import java.util.List;
+import java.util.Objects;
 import java.util.concurrent.*;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -28,6 +29,8 @@ final class AsyncPoller<T extends ScalingTask> extends BasePoller<T> {
   private final PollerBehaviorAutoscaling pollerBehavior;
   private final Scope workerMetricsScope;
   private Throttler pollRateThrottler;
+  private final Thread.UncaughtExceptionHandler uncaughtExceptionHandler =
+      new PollerUncaughtExceptionHandler();
 
   AsyncPoller(
       TrackingSlotSupplier<?> slotSupplier,
@@ -53,18 +56,26 @@ final class AsyncPoller<T extends ScalingTask> extends BasePoller<T> {
       PollerOptions pollerOptions,
       Scope workerMetricsScope) {
     super(pollTaskExecutor);
+    Objects.requireNonNull(slotSupplier, "slot supplier cannot be null");
+    Objects.requireNonNull(slotReservationData, "slot reservation data should not be null");
+    Objects.requireNonNull(asyncTaskPollers, "async tasK pollers should not be null");
+    if (asyncTaskPollers.isEmpty()) {
+      throw new IllegalArgumentException("async task pollers must contain at least one poller");
+    }
+    Objects.requireNonNull(pollerOptions, "pollerOptions should not be null");
+    Objects.requireNonNull(workerMetricsScope, "workerMetricsScope should not be null");
     this.slotSupplier = slotSupplier;
     this.slotReservationData = slotReservationData;
     this.asyncTaskPollers = asyncTaskPollers;
-    this.pollerOptions = pollerOptions;
-    this.workerMetricsScope = workerMetricsScope;
     if (!(pollerOptions.getPollerBehavior() instanceof PollerBehaviorAutoscaling)) {
       throw new IllegalArgumentException(
           "PollerBehavior "
               + pollerOptions.getPollerBehavior()
               + " is not supported. Only PollerBehaviorSimpleMaximum is supported.");
     }
-    pollerBehavior = (PollerBehaviorAutoscaling) pollerOptions.getPollerBehavior();
+    this.pollerBehavior = (PollerBehaviorAutoscaling) pollerOptions.getPollerBehavior();
+    this.pollerOptions = pollerOptions;
+    this.workerMetricsScope = workerMetricsScope;
   }
 
   @Override
@@ -77,8 +88,14 @@ public boolean start() {
               pollerOptions.getMaximumPollRatePerSecond(),
               pollerOptions.getMaximumPollRateIntervalMilliseconds());
     }
-    // TODO Set thread factory
-    ScheduledExecutorService exec = Executors.newScheduledThreadPool(asyncTaskPollers.size() + 1);
+    // Each poller will have its own thread and one thread will be used to schedule the scale
+    // reporters
+    ScheduledExecutorService exec =
+        Executors.newScheduledThreadPool(
+            asyncTaskPollers.size() + 1,
+            new ExecutorThreadFactory(
+                pollerOptions.getPollThreadNamePrefix(),
+                pollerOptions.getUncaughtExceptionHandler()));
     for (PollTaskAsync<T> asyncTaskPoller : asyncTaskPollers) {
       AdjustableSemaphore pollerSemaphore = new AdjustableSemaphore();
       pollerSemaphore.setMaxPermits(pollerBehavior.getInitialMaxConcurrentTaskPollers());
@@ -190,7 +207,11 @@ public void run() {
                   return null;
                 });
       } catch (Throwable e) {
-
+        if (e instanceof InterruptedException) {
+          // we restore the flag here, so it can be checked and processed (with exit) in finally.
+          Thread.currentThread().interrupt();
+        }
+        uncaughtExceptionHandler.uncaughtException(Thread.currentThread(), e);
       } finally {
         if (!shouldTerminate()) {
           // Resubmit itself back to pollExecutor

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

@@ -1,5 +1,6 @@
 package io.temporal.internal.worker;
 
+import io.grpc.Status;
 import io.grpc.StatusRuntimeException;
 import io.temporal.internal.common.GrpcUtils;
 import io.temporal.worker.tuning.SlotPermit;
@@ -166,4 +167,39 @@ static boolean shouldIgnoreDuringShutdown(Throwable ex) {
         // javadoc.
         || ex.getCause() instanceof InterruptedException;
   }
+
+  protected final class PollerUncaughtExceptionHandler implements Thread.UncaughtExceptionHandler {
+
+    @Override
+    public void uncaughtException(Thread t, Throwable e) {
+      if (!pollExecutor.isShutdown() || !shouldIgnoreDuringShutdown(e)) {
+        logPollErrors(t, e);
+      } else {
+        logPollExceptionsSuppressedDuringShutdown(t, e);
+      }
+    }
+
+    private void logPollErrors(Thread t, Throwable e) {
+      if (e instanceof StatusRuntimeException) {
+        StatusRuntimeException te = (StatusRuntimeException) e;
+        if (te.getStatus().getCode() == Status.Code.DEADLINE_EXCEEDED) {
+          log.info("DEADLINE_EXCEEDED in poller thread {}", t.getName(), e);
+          return;
+        }
+      }
+      log.warn("Failure in poller thread {}", t.getName(), e);
+    }
+
+    /**
+     * Some exceptions are considered normal during shutdown {@link #shouldIgnoreDuringShutdown} and
+     * we log them in the most quiet manner.
+     *
+     * @param t thread where the exception happened
+     * @param e the exception itself
+     */
+    private void logPollExceptionsSuppressedDuringShutdown(Thread t, Throwable e) {
+      log.trace(
+          "Failure in thread {} is suppressed, considered normal during shutdown", t.getName(), e);
+    }
+  }
 }

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

@@ -210,39 +210,4 @@ public void run() throws Exception {
       }
     }
   }
-
-  private final class PollerUncaughtExceptionHandler implements Thread.UncaughtExceptionHandler {
-
-    @Override
-    public void uncaughtException(Thread t, Throwable e) {
-      if (!pollExecutor.isShutdown() || !shouldIgnoreDuringShutdown(e)) {
-        logPollErrors(t, e);
-      } else {
-        logPollExceptionsSuppressedDuringShutdown(t, e);
-      }
-    }
-
-    private void logPollErrors(Thread t, Throwable e) {
-      if (e instanceof StatusRuntimeException) {
-        StatusRuntimeException te = (StatusRuntimeException) e;
-        if (te.getStatus().getCode() == Status.Code.DEADLINE_EXCEEDED) {
-          log.info("DEADLINE_EXCEEDED in poller thread {}", t.getName(), e);
-          return;
-        }
-      }
-      log.warn("Failure in poller thread {}", t.getName(), e);
-    }
-
-    /**
-     * Some exceptions are considered normal during shutdown {@link #shouldIgnoreDuringShutdown} and
-     * we log them in the most quiet manner.
-     *
-     * @param t thread where the exception happened
-     * @param e the exception itself
-     */
-    private void logPollExceptionsSuppressedDuringShutdown(Thread t, Throwable e) {
-      log.trace(
-          "Failure in thread {} is suppressed, considered normal during shutdown", t.getName(), e);
-    }
-  }
 }

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

@@ -7,6 +7,10 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+/**
+ * PollScaleReportHandle is responsible for managing the scaling of pollers based on the scaling
+ * feedback attached to the task by the server.
+ */
 @ThreadSafe
 public class PollScaleReportHandle<T extends ScalingTask> implements Runnable {
   private static final Logger logger = LoggerFactory.getLogger(PollScaleReportHandle.class);

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

@@ -493,37 +493,52 @@ public Builder setDeploymentOptions(WorkerDeploymentOptions deploymentOptions) {
       return this;
     }
 
-    /** Set the poller behavior for workflow task pollers. */
+    /**
+     * Set the poller behavior for workflow task pollers.
+     *
+     * <p>If the sticky queue is enabled, the poller behavior will be used for the sticky queue as
+     * well.
+     */
     @Experimental
     public Builder setWorkflowTaskPollersBehaviour(PollerBehaviorAutoscaling pollerBehavior) {
       this.workflowTaskPollersBehaviour = pollerBehavior;
       return this;
     }
 
+    /**
+     * Set the poller behavior for workflow task pollers.
+     *
+     * <p>If the sticky queue is enabled, the poller behavior will be used for the sticky queue as
+     * well.
+     */
     @Experimental
     public Builder setWorkflowTaskPollersBehaviour(PollerBehaviorSimpleMaximum pollerBehavior) {
       this.workflowTaskPollersBehaviour = pollerBehavior;
       return this;
     }
 
+    /** Set the poller behavior for activity task pollers. */
     @Experimental
     public Builder setActivityTaskPollersBehaviour(PollerBehaviorAutoscaling pollerBehavior) {
       this.activityTaskPollersBehaviour = pollerBehavior;
       return this;
     }
 
+    /** Set the poller behavior for activity task pollers. */
     @Experimental
     public Builder setActivityTaskPollersBehaviour(PollerBehaviorSimpleMaximum pollerBehavior) {
       this.activityTaskPollersBehaviour = pollerBehavior;
       return this;
     }
 
+    /** Set the poller behavior for nexus task pollers. */
     @Experimental
     public Builder setNexusTaskPollersBehaviour(PollerBehaviorAutoscaling pollerBehavior) {
       this.nexusTaskPollersBehaviour = pollerBehavior;
       return this;
     }
 
+    /** Set the poller behavior for nexus task pollers. */
     @Experimental
     public Builder setNexusTaskPollersBehaviour(PollerBehaviorSimpleMaximum pollerBehavior) {
       this.nexusTaskPollersBehaviour = pollerBehavior;

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

@@ -0,0 +1,59 @@
+package io.temporal.internal.worker;
+
+import io.grpc.Status;
+import io.grpc.StatusRuntimeException;
+import io.temporal.workflow.Functions;
+import org.junit.Test;
+import org.mockito.Mockito;
+
+public class PollScaleReportHandleTest {
+
+  @Test
+  public void handleResourceExhaustedError() {
+    // Mock dependencies
+    Functions.Proc1<Integer> mockScaleCallback = Mockito.mock(Functions.Proc1.class);
+    PollScaleReportHandle<ScalingTask> handle =
+        new PollScaleReportHandle<>(1, 10, 8, mockScaleCallback);
+
+    // Simulate RESOURCE_EXHAUSTED error
+    StatusRuntimeException exception = new StatusRuntimeException(Status.RESOURCE_EXHAUSTED);
+    handle.report(null, exception);
+
+    // Verify target poller count is halved and callback is invoked
+    Mockito.verify(mockScaleCallback).apply(4);
+  }
+
+  @Test
+  public void handleGenericError() {
+    // Mock dependencies
+    Functions.Proc1<Integer> mockScaleCallback = Mockito.mock(Functions.Proc1.class);
+    PollScaleReportHandle<ScalingTask> handle =
+        new PollScaleReportHandle<>(1, 10, 5, mockScaleCallback);
+
+    // Simulate a generic error
+    handle.report(null, new RuntimeException("Generic error"));
+
+    // Verify target poller count is decremented and callback is invoked
+    Mockito.verify(mockScaleCallback).apply(4);
+  }
+
+  @Test
+  public void applyScalingDecisionDeltaWhenAllowed() {
+    // Mock dependencies
+    Functions.Proc1<Integer> mockScaleCallback = Mockito.mock(Functions.Proc1.class);
+    ScalingTask mockTask = Mockito.mock(ScalingTask.class);
+    ScalingTask.ScalingDecision mockDecision = Mockito.mock(ScalingTask.ScalingDecision.class);
+    Mockito.when(mockTask.getScalingDecision()).thenReturn(mockDecision);
+    Mockito.when(mockDecision.getPollRequestDeltaSuggestion()).thenReturn(3);
+
+    PollScaleReportHandle<ScalingTask> handle =
+        new PollScaleReportHandle<>(1, 10, 5, mockScaleCallback);
+    handle.run(); // Enable scale-up
+
+    // Report a task with a scaling decision
+    handle.report(mockTask, null);
+
+    // Verify target poller count is updated and callback is invoked
+    Mockito.verify(mockScaleCallback).apply(8);
+  }
+}