Expose eager activity concurrency limit

Author: Sushisource

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

@@ -218,8 +218,9 @@ public WorkerLifecycleState getLifecycleState() {
     return poller.getLifecycleState();
   }
 
-  public EagerActivityDispatcher getEagerActivityDispatcher() {
-    return new EagerActivityDispatcherImpl();
+  public EagerActivityDispatcher getEagerActivityDispatcher(
+      int maxConcurrentEagerActivityExecutionSize) {
+    return new EagerActivityDispatcherImpl(maxConcurrentEagerActivityExecutionSize);
   }
 
   private PollerOptions getPollerOptions(SingleWorkerOptions options) {
@@ -485,6 +486,13 @@ private void logExceptionDuringResultReporting(
   }
 
   private final class EagerActivityDispatcherImpl implements EagerActivityDispatcher {
+    private final EagerActivitySlotLimiter eagerActivitySlotLimiter;
+
+    private EagerActivityDispatcherImpl(int maxConcurrentEagerActivityExecutionSize) {
+      this.eagerActivitySlotLimiter =
+          new EagerActivitySlotLimiter(maxConcurrentEagerActivityExecutionSize);
+    }
+
     @Override
     public Optional<SlotPermit> tryReserveActivitySlot(
         ScheduleActivityTaskCommandAttributesOrBuilder commandAttributes) {
@@ -493,15 +501,31 @@ public Optional<SlotPermit> tryReserveActivitySlot(
               commandAttributes.getTaskQueue().getName(), ActivityWorker.this.taskQueue)) {
         return Optional.empty();
       }
-      return ActivityWorker.this.slotSupplier.tryReserveSlot(
-          new SlotReservationData(
-              ActivityWorker.this.taskQueue, options.getIdentity(), options.getBuildId()));
+      if (!eagerActivitySlotLimiter.tryReserve()) {
+        return Optional.empty();
+      }
+      Optional<SlotPermit> permit = Optional.empty();
+      try {
+        permit =
+            ActivityWorker.this.slotSupplier.tryReserveSlot(
+                new SlotReservationData(
+                    ActivityWorker.this.taskQueue, options.getIdentity(), options.getBuildId()));
+        return permit;
+      } finally {
+        if (!permit.isPresent()) {
+          eagerActivitySlotLimiter.release();
+        }
+      }
     }
 
     @Override
     public void releaseActivitySlotReservations(Iterable<SlotPermit> permits) {
       for (SlotPermit permit : permits) {
-        ActivityWorker.this.slotSupplier.releaseSlot(SlotReleaseReason.neverUsed(), permit);
+        try {
+          ActivityWorker.this.slotSupplier.releaseSlot(SlotReleaseReason.neverUsed(), permit);
+        } finally {
+          eagerActivitySlotLimiter.release();
+        }
       }
     }
 
@@ -511,9 +535,39 @@ public void dispatchActivity(PollActivityTaskQueueResponse activity, SlotPermit
           new ActivityTask(
               activity,
               permit,
-              () ->
+              () -> {
+                try {
                   ActivityWorker.this.slotSupplier.releaseSlot(
-                      SlotReleaseReason.taskComplete(), permit)));
+                      SlotReleaseReason.taskComplete(), permit);
+                } finally {
+                  eagerActivitySlotLimiter.release();
+                }
+              }));
+    }
+  }
+
+  static final class EagerActivitySlotLimiter {
+    private final int maxConcurrent;
+    private int heldSlotCount;
+
+    EagerActivitySlotLimiter(int maxConcurrent) {
+      this.maxConcurrent = maxConcurrent;
+    }
+
+    synchronized boolean tryReserve() {
+      if (maxConcurrent > 0 && heldSlotCount >= maxConcurrent) {
+        return false;
+      }
+      heldSlotCount++;
+      return true;
+    }
+
+    synchronized void release() {
+      if (heldSlotCount <= 0) {
+        throw new IllegalStateException(
+            "Trying to release an unreserved eager activity slot. This is an SDK bug.");
+      }
+      heldSlotCount--;
     }
   }
 }

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

@@ -163,8 +163,9 @@ public WorkerLifecycleState getLifecycleState() {
     }
   }
 
-  public EagerActivityDispatcher getEagerActivityDispatcher() {
-    return this.worker.getEagerActivityDispatcher();
+  public EagerActivityDispatcher getEagerActivityDispatcher(
+      int maxConcurrentEagerActivityExecutionSize) {
+    return this.worker.getEagerActivityDispatcher(maxConcurrentEagerActivityExecutionSize);
   }
 
   public boolean isAnyTypeSupported() {

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

@@ -173,7 +173,8 @@ private static final class TaskSnapshot {
 
     EagerActivityDispatcher eagerActivityDispatcher =
         (activityWorker != null && !this.options.isEagerExecutionDisabled())
-            ? activityWorker.getEagerActivityDispatcher()
+            ? activityWorker.getEagerActivityDispatcher(
+                this.options.getMaxConcurrentEagerActivityExecutionSize())
             : new EagerActivityDispatcher.NoopEagerActivityDispatcher();
 
     SingleWorkerOptions nexusOptions =

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

@@ -64,6 +64,7 @@ public static final class Builder {
     private Duration defaultHeartbeatThrottleInterval;
     private Duration stickyQueueScheduleToStartTimeout;
     private boolean disableEagerExecution;
+    private int maxConcurrentEagerActivityExecutionSize;
     private String buildId;
     private boolean useBuildIdForVersioning;
     private Duration stickyTaskQueueDrainTimeout;
@@ -109,6 +110,7 @@ private Builder(WorkerOptions o) {
       this.defaultHeartbeatThrottleInterval = o.defaultHeartbeatThrottleInterval;
       this.stickyQueueScheduleToStartTimeout = o.stickyQueueScheduleToStartTimeout;
       this.disableEagerExecution = o.disableEagerExecution;
+      this.maxConcurrentEagerActivityExecutionSize = o.maxConcurrentEagerActivityExecutionSize;
       this.useBuildIdForVersioning = o.useBuildIdForVersioning;
       this.buildId = o.buildId;
       this.stickyTaskQueueDrainTimeout = o.stickyTaskQueueDrainTimeout;
@@ -400,6 +402,19 @@ public Builder setDisableEagerExecution(boolean disableEagerExecution) {
       return this;
     }
 
+    /**
+     * Sets the maximum number of eager activities that can be running concurrently.
+     *
+     * <p>When nonzero, eager activity execution will not be requested if it would cause the number
+     * of running eager activities to exceed this value. The default of zero means unlimited and
+     * therefore only bound by the activity slot supplier.
+     */
+    public Builder setMaxConcurrentEagerActivityExecutionSize(
+        int maxConcurrentEagerActivityExecutionSize) {
+      this.maxConcurrentEagerActivityExecutionSize = maxConcurrentEagerActivityExecutionSize;
+      return this;
+    }
+
     /**
      * Opts the worker in to the Build-ID-based versioning feature. This ensures that the worker
      * will only receive tasks which it is compatible with.
@@ -623,6 +638,7 @@ public WorkerOptions build() {
           defaultHeartbeatThrottleInterval,
           stickyQueueScheduleToStartTimeout,
           disableEagerExecution,
+          maxConcurrentEagerActivityExecutionSize,
           useBuildIdForVersioning,
           buildId,
           stickyTaskQueueDrainTimeout,
@@ -647,6 +663,9 @@ public WorkerOptions validateAndBuildWithDefaults() {
           maxWorkerActivitiesPerSecond >= 0, "negative maxActivitiesPerSecond");
       Preconditions.checkState(
           maxConcurrentActivityExecutionSize >= 0, "negative maxConcurrentActivityExecutionSize");
+      Preconditions.checkState(
+          maxConcurrentEagerActivityExecutionSize >= 0,
+          "negative maxConcurrentEagerActivityExecutionSize");
       Preconditions.checkState(
           maxConcurrentWorkflowTaskExecutionSize >= 0,
           "negative maxConcurrentWorkflowTaskExecutionSize");
@@ -758,6 +777,7 @@ public WorkerOptions validateAndBuildWithDefaults() {
               ? DEFAULT_STICKY_SCHEDULE_TO_START_TIMEOUT
               : stickyQueueScheduleToStartTimeout,
           disableEagerExecution,
+          maxConcurrentEagerActivityExecutionSize,
           useBuildIdForVersioning,
           buildId,
           stickyTaskQueueDrainTimeout == null
@@ -796,6 +816,7 @@ public WorkerOptions validateAndBuildWithDefaults() {
   private final Duration defaultHeartbeatThrottleInterval;
   private final @Nonnull Duration stickyQueueScheduleToStartTimeout;
   private final boolean disableEagerExecution;
+  private final int maxConcurrentEagerActivityExecutionSize;
   private final boolean useBuildIdForVersioning;
   private final String buildId;
   private final Duration stickyTaskQueueDrainTimeout;
@@ -831,6 +852,7 @@ private WorkerOptions(
       Duration defaultHeartbeatThrottleInterval,
       @Nonnull Duration stickyQueueScheduleToStartTimeout,
       boolean disableEagerExecution,
+      int maxConcurrentEagerActivityExecutionSize,
       boolean useBuildIdForVersioning,
       String buildId,
       Duration stickyTaskQueueDrainTimeout,
@@ -864,6 +886,7 @@ private WorkerOptions(
     this.defaultHeartbeatThrottleInterval = defaultHeartbeatThrottleInterval;
     this.stickyQueueScheduleToStartTimeout = stickyQueueScheduleToStartTimeout;
     this.disableEagerExecution = maxTaskQueueActivitiesPerSecond > 0 ? true : disableEagerExecution;
+    this.maxConcurrentEagerActivityExecutionSize = maxConcurrentEagerActivityExecutionSize;
     this.useBuildIdForVersioning = useBuildIdForVersioning;
     this.buildId = buildId;
     this.stickyTaskQueueDrainTimeout = stickyTaskQueueDrainTimeout;
@@ -989,6 +1012,10 @@ public boolean isEagerExecutionDisabled() {
     return disableEagerExecution;
   }
 
+  public int getMaxConcurrentEagerActivityExecutionSize() {
+    return maxConcurrentEagerActivityExecutionSize;
+  }
+
   public boolean isUsingBuildIdForVersioning() {
     return useBuildIdForVersioning;
   }
@@ -1070,6 +1097,7 @@ && compare(maxTaskQueueActivitiesPerSecond, that.maxTaskQueueActivitiesPerSecond
         && localActivityWorkerOnly == that.localActivityWorkerOnly
         && defaultDeadlockDetectionTimeout == that.defaultDeadlockDetectionTimeout
         && disableEagerExecution == that.disableEagerExecution
+        && maxConcurrentEagerActivityExecutionSize == that.maxConcurrentEagerActivityExecutionSize
         && useBuildIdForVersioning == that.useBuildIdForVersioning
         && Objects.equals(workerTuner, that.workerTuner)
         && Objects.equals(maxHeartbeatThrottleInterval, that.maxHeartbeatThrottleInterval)
@@ -1109,6 +1137,7 @@ public int hashCode() {
         defaultHeartbeatThrottleInterval,
         stickyQueueScheduleToStartTimeout,
         disableEagerExecution,
+        maxConcurrentEagerActivityExecutionSize,
         useBuildIdForVersioning,
         buildId,
         stickyTaskQueueDrainTimeout,
@@ -1160,6 +1189,8 @@ public String toString() {
         + stickyQueueScheduleToStartTimeout
         + ", disableEagerExecution="
         + disableEagerExecution
+        + ", maxConcurrentEagerActivityExecutionSize="
+        + maxConcurrentEagerActivityExecutionSize
         + ", useBuildIdForVersioning="
         + useBuildIdForVersioning
         + ", buildId='"

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

@@ -0,0 +1,31 @@
+package io.temporal.internal.worker;
+
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+import org.junit.Test;
+
+public class EagerActivitySlotLimiterTest {
+  @Test
+  public void enforcesLimitUntilReservationIsReleased() {
+    ActivityWorker.EagerActivitySlotLimiter limiter =
+        new ActivityWorker.EagerActivitySlotLimiter(2);
+
+    assertTrue(limiter.tryReserve());
+    assertTrue(limiter.tryReserve());
+    assertFalse(limiter.tryReserve());
+
+    limiter.release();
+    assertTrue(limiter.tryReserve());
+  }
+
+  @Test
+  public void zeroAllowsUnlimitedReservations() {
+    ActivityWorker.EagerActivitySlotLimiter limiter =
+        new ActivityWorker.EagerActivitySlotLimiter(0);
+
+    for (int i = 0; i < 1000; i++) {
+      assertTrue(limiter.tryReserve());
+    }
+  }
+}

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

@@ -55,6 +55,7 @@ public void verifyNewBuilderFromExistingWorkerOptions() {
             .setDefaultHeartbeatThrottleInterval(Duration.ofSeconds(7))
             .setStickyQueueScheduleToStartTimeout(Duration.ofSeconds(60))
             .setDisableEagerExecution(false)
+            .setMaxConcurrentEagerActivityExecutionSize(17)
             .setUseBuildIdForVersioning(false)
             .setBuildId("build-id")
             .setStickyTaskQueueDrainTimeout(Duration.ofSeconds(15))
@@ -90,6 +91,9 @@ public void verifyNewBuilderFromExistingWorkerOptions() {
     assertEquals(
         w1.getStickyQueueScheduleToStartTimeout(), w2.getStickyQueueScheduleToStartTimeout());
     assertEquals(w1.isEagerExecutionDisabled(), w2.isEagerExecutionDisabled());
+    assertEquals(
+        w1.getMaxConcurrentEagerActivityExecutionSize(),
+        w2.getMaxConcurrentEagerActivityExecutionSize());
     assertEquals(w1.isUsingBuildIdForVersioning(), w2.isUsingBuildIdForVersioning());
     assertEquals(w1.getBuildId(), w2.getBuildId());
     assertEquals(w1.getStickyTaskQueueDrainTimeout(), w2.getStickyTaskQueueDrainTimeout());