Asyncify slot suppliers

Author: Sushisource

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

@@ -35,6 +35,8 @@
 import io.temporal.worker.MetricsType;
 import io.temporal.worker.tuning.*;
 import java.util.Objects;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
 import java.util.function.Supplier;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
@@ -96,18 +98,19 @@ public ActivityTask poll() {
     PollActivityTaskQueueResponse response;
     SlotPermit permit;
     boolean isSuccessful = false;
-
+    CompletableFuture<SlotPermit> future =
+        slotSupplier.reserveSlot(
+            new SlotReservationData(
+                pollRequest.getTaskQueue().getName(),
+                pollRequest.getIdentity(),
+                pollRequest.getWorkerVersionCapabilities().getBuildId()));
     try {
-      permit =
-          slotSupplier.reserveSlot(
-              new SlotReservationData(
-                  pollRequest.getTaskQueue().getName(),
-                  pollRequest.getIdentity(),
-                  pollRequest.getWorkerVersionCapabilities().getBuildId()));
+      permit = future.get();
     } catch (InterruptedException e) {
+      future.cancel(true);
       Thread.currentThread().interrupt();
       return null;
-    } catch (Exception e) {
+    } catch (ExecutionException e) {
       log.warn("Error while trying to reserve a slot for an activity", e.getCause());
       return null;
     }

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

@@ -83,18 +83,22 @@ private void processQueue() {
       QueuedLARequest request = null;
       try {
         request = requestQueue.take();
+
+        CompletableFuture<SlotPermit> future = slotSupplier.reserveSlot(request.data);
         try {
-          slotPermit = slotSupplier.reserveSlot(request.data);
+          slotPermit = future.get();
         } catch (InterruptedException e) {
+          future.cancel(true);
           Thread.currentThread().interrupt();
           return;
-        } catch (Exception e) {
+        } catch (ExecutionException e) {
           log.error(
               "Error reserving local activity slot, dropped activity id {}",
               request.task.getActivityId(),
               e);
           continue;
         }
+
         request.task.getExecutionContext().setPermit(slotPermit);
         afterReservedCallback.apply(request.task);
       } catch (InterruptedException e) {

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

@@ -32,6 +32,8 @@
 import io.temporal.worker.MetricsType;
 import io.temporal.worker.tuning.*;
 import java.util.Objects;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
 import java.util.function.Supplier;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
@@ -86,18 +88,19 @@ public NexusTask poll() {
     PollNexusTaskQueueResponse response;
     SlotPermit permit;
     boolean isSuccessful = false;
-
+    CompletableFuture<SlotPermit> future =
+        slotSupplier.reserveSlot(
+            new SlotReservationData(
+                pollRequest.getTaskQueue().getName(),
+                pollRequest.getIdentity(),
+                pollRequest.getWorkerVersionCapabilities().getBuildId()));
     try {
-      permit =
-          slotSupplier.reserveSlot(
-              new SlotReservationData(
-                  pollRequest.getTaskQueue().getName(),
-                  pollRequest.getIdentity(),
-                  pollRequest.getWorkerVersionCapabilities().getBuildId()));
+      permit = future.get();
     } catch (InterruptedException e) {
+      future.cancel(true);
       Thread.currentThread().interrupt();
       return null;
-    } catch (Exception e) {
+    } catch (ExecutionException e) {
       log.warn("Error while trying to reserve a slot for a nexus task", e.getCause());
       return null;
     }

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

@@ -26,6 +26,7 @@
 import java.util.Collections;
 import java.util.Map;
 import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicInteger;
 
@@ -48,10 +49,10 @@ public TrackingSlotSupplier(SlotSupplier<SI> inner, Scope metricsScope) {
     publishSlotsMetric();
   }
 
-  public SlotPermit reserveSlot(SlotReservationData dat) throws InterruptedException {
-    SlotPermit p = inner.reserveSlot(createCtx(dat));
-    issuedSlots.incrementAndGet();
-    return p;
+  public CompletableFuture<SlotPermit> reserveSlot(SlotReservationData dat) {
+    CompletableFuture<SlotPermit> future = inner.reserveSlot(createCtx(dat));
+    future.thenAccept(permit -> issuedSlots.incrementAndGet());
+    return future;
   }
 
   public Optional<SlotPermit> tryReserveSlot(SlotReservationData dat) {

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

@@ -35,8 +35,12 @@
 import io.temporal.serviceclient.MetricsTag;
 import io.temporal.serviceclient.WorkflowServiceStubs;
 import io.temporal.worker.MetricsType;
-import io.temporal.worker.tuning.*;
+import io.temporal.worker.tuning.SlotPermit;
+import io.temporal.worker.tuning.SlotReleaseReason;
+import io.temporal.worker.tuning.WorkflowSlotInfo;
 import java.util.Objects;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
 import java.util.function.Supplier;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
@@ -123,17 +127,19 @@ public WorkflowPollTask(
   public WorkflowTask poll() {
     boolean isSuccessful = false;
     SlotPermit permit;
+    CompletableFuture<SlotPermit> future =
+        slotSupplier.reserveSlot(
+            new SlotReservationData(
+                pollRequest.getTaskQueue().getName(),
+                pollRequest.getIdentity(),
+                pollRequest.getWorkerVersionCapabilities().getBuildId()));
     try {
-      permit =
-          slotSupplier.reserveSlot(
-              new SlotReservationData(
-                  pollRequest.getTaskQueue().getName(),
-                  pollRequest.getIdentity(),
-                  pollRequest.getWorkerVersionCapabilities().getBuildId()));
+      permit = future.get();
     } catch (InterruptedException e) {
+      future.cancel(true);
       Thread.currentThread().interrupt();
       return null;
-    } catch (Exception e) {
+    } catch (ExecutionException e) {
       log.warn("Error while trying to reserve a slot for workflow task", e.getCause());
       return null;
     }

temporal-sdk/src/main/java/io/temporal/worker/tuning/FixedSizeSlotSupplier.java

@@ -21,8 +21,11 @@
 package io.temporal.worker.tuning;
 
 import com.google.common.base.Preconditions;
+import java.util.ArrayDeque;
 import java.util.Optional;
-import java.util.concurrent.*;
+import java.util.Queue;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.locks.ReentrantLock;
 
 /**
  * This implementation of {@link SlotSupplier} provides a fixed number of slots backed by a
@@ -32,18 +35,83 @@
  */
 public class FixedSizeSlotSupplier<SI extends SlotInfo> implements SlotSupplier<SI> {
   private final int numSlots;
-  private final Semaphore executorSlotsSemaphore;
+  private final AsyncSemaphore executorSlotsSemaphore;
+
+  /**
+   * A simple version of an async semaphore. Unfortunately there's not any readily available
+   * properly licensed library I could find for this which is a bit shocking, but this
+   * implementation should be suitable for our needs
+   */
+  static class AsyncSemaphore {
+    private final ReentrantLock lock = new ReentrantLock();
+    private final Queue<CompletableFuture<Void>> waiters = new ArrayDeque<>();
+    private int permits;
+
+    AsyncSemaphore(int initialPermits) {
+      this.permits = initialPermits;
+    }
+
+    /**
+     * Acquire a permit asynchronously. If a permit is available, returns a completed future,
+     * otherwise returns a future that will be completed when a permit is released.
+     */
+    public CompletableFuture<Void> acquire() {
+      lock.lock();
+      try {
+        if (permits > 0) {
+          permits--;
+          return CompletableFuture.completedFuture(null);
+        } else {
+          CompletableFuture<Void> waiter = new CompletableFuture<>();
+          waiters.add(waiter);
+          return waiter;
+        }
+      } finally {
+        lock.unlock();
+      }
+    }
+
+    public boolean tryAcquire() {
+      lock.lock();
+      try {
+        if (permits > 0) {
+          permits--;
+          return true;
+        }
+        return false;
+      } finally {
+        lock.unlock();
+      }
+    }
+
+    /**
+     * Release a permit. If there are waiting futures, completes the next one instead of
+     * incrementing the permit count.
+     */
+    public void release() {
+      lock.lock();
+      try {
+        CompletableFuture<Void> waiter = waiters.poll();
+        if (waiter != null) {
+          waiter.complete(null);
+        } else {
+          permits++;
+        }
+      } finally {
+        lock.unlock();
+      }
+    }
+  }
 
   public FixedSizeSlotSupplier(int numSlots) {
     Preconditions.checkArgument(numSlots > 0, "FixedSizeSlotSupplier must have at least one slot");
     this.numSlots = numSlots;
-    executorSlotsSemaphore = new Semaphore(numSlots);
+    executorSlotsSemaphore = new AsyncSemaphore(numSlots);
   }
 
   @Override
-  public SlotPermit reserveSlot(SlotReserveContext<SI> ctx) throws InterruptedException {
-    executorSlotsSemaphore.acquire();
-    return new SlotPermit();
+  public CompletableFuture<SlotPermit> reserveSlot(SlotReserveContext<SI> ctx) {
+    return executorSlotsSemaphore.acquire().thenApply(ignored -> new SlotPermit());
   }
 
   @Override

temporal-sdk/src/main/java/io/temporal/worker/tuning/ResourceBasedSlotSupplier.java

@@ -24,6 +24,7 @@
 import java.time.Duration;
 import java.time.Instant;
 import java.util.Optional;
+import java.util.concurrent.*;
 
 /** Implements a {@link SlotSupplier} based on resource usage for a particular slot type. */
 @Experimental
@@ -32,6 +33,18 @@ public class ResourceBasedSlotSupplier<SI extends SlotInfo> implements SlotSuppl
   private final ResourceBasedController resourceController;
   private final ResourceBasedSlotOptions options;
   private Instant lastSlotIssuedAt = Instant.EPOCH;
+  // For slot reservations that are waiting to re-check resource usage
+  private static final ScheduledExecutorService scheduler =
+      Executors.newScheduledThreadPool(
+          // Two threads seem needed here, so that reading PID decisions doesn't interfere overly
+          // with firing off scheduled tasks or one another.
+          2,
+          r -> {
+            Thread t = new Thread(r);
+            t.setName("ResourceBasedSlotSupplier.scheduler");
+            t.setDaemon(true);
+            return t;
+          });
 
   /**
    * Construct a slot supplier for workflow tasks with the given resource controller and options.
@@ -139,29 +152,43 @@ private ResourceBasedSlotSupplier(
   }
 
   @Override
-  public SlotPermit reserveSlot(SlotReserveContext<SI> ctx) throws InterruptedException {
-    while (true) {
-      if (ctx.getNumIssuedSlots() < options.getMinimumSlots()) {
-        return new SlotPermit();
-      } else {
-        Duration mustWaitFor;
-        try {
-          mustWaitFor = options.getRampThrottle().minus(timeSinceLastSlotIssued());
-        } catch (ArithmeticException e) {
-          mustWaitFor = Duration.ZERO;
-        }
-        if (mustWaitFor.compareTo(Duration.ZERO) > 0) {
-          Thread.sleep(mustWaitFor.toMillis());
-        }
-
-        Optional<SlotPermit> permit = tryReserveSlot(ctx);
-        if (permit.isPresent()) {
-          return permit.get();
-        } else {
-          Thread.sleep(10);
-        }
-      }
+  public CompletableFuture<SlotPermit> reserveSlot(SlotReserveContext<SI> ctx) {
+    if (ctx.getNumIssuedSlots() < options.getMinimumSlots()) {
+      return CompletableFuture.completedFuture(new SlotPermit());
     }
+    return tryReserveSlot(ctx)
+        .map(CompletableFuture::completedFuture)
+        .orElseGet(() -> scheduleSlotAcquisition(ctx));
+  }
+
+  private CompletableFuture<SlotPermit> scheduleSlotAcquisition(SlotReserveContext<SI> ctx) {
+    Duration mustWaitFor;
+    try {
+      mustWaitFor = options.getRampThrottle().minus(timeSinceLastSlotIssued());
+    } catch (ArithmeticException e) {
+      mustWaitFor = Duration.ZERO;
+    }
+
+    CompletableFuture<Void> permitFuture;
+    if (mustWaitFor.compareTo(Duration.ZERO) > 0) {
+      permitFuture =
+          CompletableFuture.supplyAsync(() -> null, delayedExecutor(mustWaitFor.toMillis()));
+    } else {
+      permitFuture = CompletableFuture.completedFuture(null);
+    }
+
+    // After the delay, try to reserve the slot
+    return permitFuture.thenCompose(
+        ignored -> {
+          Optional<SlotPermit> permit = tryReserveSlot(ctx);
+          // If we couldn't get a slot this time, delay for a short period and try again
+          return permit
+              .map(CompletableFuture::completedFuture)
+              .orElseGet(
+                  () ->
+                      CompletableFuture.supplyAsync(() -> null, delayedExecutor(10))
+                          .thenCompose(ig -> scheduleSlotAcquisition(ctx)));
+        });
   }
 
   @Override
@@ -190,4 +217,9 @@ public ResourceBasedController getResourceController() {
   private Duration timeSinceLastSlotIssued() {
     return Duration.between(lastSlotIssuedAt, Instant.now());
   }
+
+  // Polyfill for Java 9 delayedExecutor
+  private static Executor delayedExecutor(long delay) {
+    return r -> scheduler.schedule(() -> scheduler.execute(r), delay, TimeUnit.MILLISECONDS);
+  }
 }