Return Future instead of CompletableFuture

Sushisource · Sushisource · commit 87279003ece0 · 2025-03-20T10:06:13.000-07:00
diff --git a/temporal-sdk/src/main/java/io/temporal/internal/worker/ActivityPollTask.java b/temporal-sdk/src/main/java/io/temporal/internal/worker/ActivityPollTask.java
@@ -35,8 +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.concurrent.Future;
 import java.util.function.Supplier;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
@@ -98,7 +98,7 @@ public ActivityTask poll() {
     PollActivityTaskQueueResponse response;
     SlotPermit permit;
     boolean isSuccessful = false;
-    CompletableFuture<SlotPermit> future;
+    Future<SlotPermit> future;
     try {
       future =
           slotSupplier.reserveSlot(
diff --git a/temporal-sdk/src/main/java/io/temporal/internal/worker/LocalActivitySlotSupplierQueue.java b/temporal-sdk/src/main/java/io/temporal/internal/worker/LocalActivitySlotSupplierQueue.java
@@ -84,7 +84,7 @@ private void processQueue() {
       try {
         request = requestQueue.take();
 
-        CompletableFuture<SlotPermit> future = slotSupplier.reserveSlot(request.data);
+        Future<SlotPermit> future = slotSupplier.reserveSlot(request.data);
         try {
           slotPermit = future.get();
         } catch (InterruptedException e) {
diff --git a/temporal-sdk/src/main/java/io/temporal/internal/worker/NexusPollTask.java b/temporal-sdk/src/main/java/io/temporal/internal/worker/NexusPollTask.java
@@ -32,8 +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.concurrent.Future;
 import java.util.function.Supplier;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
@@ -88,7 +88,7 @@ public NexusTask poll() {
     PollNexusTaskQueueResponse response;
     SlotPermit permit;
     boolean isSuccessful = false;
-    CompletableFuture<SlotPermit> future =
+    Future<SlotPermit> future =
         slotSupplier.reserveSlot(
             new SlotReservationData(
                 pollRequest.getTaskQueue().getName(),
diff --git a/temporal-sdk/src/main/java/io/temporal/internal/worker/TrackingSlotSupplier.java b/temporal-sdk/src/main/java/io/temporal/internal/worker/TrackingSlotSupplier.java
@@ -26,8 +26,8 @@
 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.*;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 
 /**
@@ -49,15 +49,52 @@ public TrackingSlotSupplier(SlotSupplier<SI> inner, Scope metricsScope) {
     publishSlotsMetric();
   }
 
-  public CompletableFuture<SlotPermit> reserveSlot(SlotReservationData dat) {
-    CompletableFuture<SlotPermit> future = null;
+  public Future<SlotPermit> reserveSlot(SlotReservationData dat) {
+    final Future<SlotPermit> originalFuture;
     try {
-      future = inner.reserveSlot(createCtx(dat));
+      originalFuture = inner.reserveSlot(createCtx(dat));
     } catch (Exception e) {
       throw new RuntimeException(e);
     }
-    future.thenAccept(permit -> issuedSlots.incrementAndGet());
-    return future;
+
+    return new Future<SlotPermit>() {
+      private final AtomicBoolean callbackInvoked = new AtomicBoolean(false);
+
+      private SlotPermit executeCallbackIfNeeded(SlotPermit permit) {
+        if (callbackInvoked.compareAndSet(false, true)) {
+          issuedSlots.incrementAndGet();
+        }
+        return permit;
+      }
+
+      @Override
+      public boolean cancel(boolean mayInterruptIfRunning) {
+        return originalFuture.cancel(mayInterruptIfRunning);
+      }
+
+      @Override
+      public boolean isCancelled() {
+        return originalFuture.isCancelled();
+      }
+
+      @Override
+      public boolean isDone() {
+        return originalFuture.isDone();
+      }
+
+      @Override
+      public SlotPermit get() throws InterruptedException, ExecutionException {
+        SlotPermit permit = originalFuture.get();
+        return executeCallbackIfNeeded(permit);
+      }
+
+      @Override
+      public SlotPermit get(long timeout, TimeUnit unit)
+          throws InterruptedException, ExecutionException, TimeoutException {
+        SlotPermit permit = originalFuture.get(timeout, unit);
+        return executeCallbackIfNeeded(permit);
+      }
+    };
   }
 
   public Optional<SlotPermit> tryReserveSlot(SlotReservationData dat) {
diff --git a/temporal-sdk/src/main/java/io/temporal/internal/worker/WorkflowPollTask.java b/temporal-sdk/src/main/java/io/temporal/internal/worker/WorkflowPollTask.java
@@ -39,8 +39,8 @@
 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.concurrent.Future;
 import java.util.function.Supplier;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
@@ -127,7 +127,7 @@ public WorkflowPollTask(
   public WorkflowTask poll() {
     boolean isSuccessful = false;
     SlotPermit permit;
-    CompletableFuture<SlotPermit> future =
+    Future<SlotPermit> future =
         slotSupplier.reserveSlot(
             new SlotReservationData(
                 pollRequest.getTaskQueue().getName(),
diff --git a/temporal-sdk/src/main/java/io/temporal/worker/tuning/FixedSizeSlotSupplier.java b/temporal-sdk/src/main/java/io/temporal/worker/tuning/FixedSizeSlotSupplier.java
@@ -25,6 +25,7 @@
 import java.util.Optional;
 import java.util.Queue;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Future;
 import java.util.concurrent.locks.ReentrantLock;
 
 /**
@@ -110,7 +111,7 @@ public FixedSizeSlotSupplier(int numSlots) {
   }
 
   @Override
-  public CompletableFuture<SlotPermit> reserveSlot(SlotReserveContext<SI> ctx) throws Exception {
+  public Future<SlotPermit> reserveSlot(SlotReserveContext<SI> ctx) throws Exception {
     return executorSlotsSemaphore.acquire().thenApply(ignored -> new SlotPermit());
   }
 
diff --git a/temporal-sdk/src/main/java/io/temporal/worker/tuning/ResourceBasedSlotSupplier.java b/temporal-sdk/src/main/java/io/temporal/worker/tuning/ResourceBasedSlotSupplier.java
@@ -202,7 +202,7 @@ private ResourceBasedSlotSupplier(
   }
 
   @Override
-  public CompletableFuture<SlotPermit> reserveSlot(SlotReserveContext<SI> ctx) throws Exception {
+  public Future<SlotPermit> reserveSlot(SlotReserveContext<SI> ctx) throws Exception {
     if (ctx.getNumIssuedSlots() < options.getMinimumSlots()) {
       return CompletableFuture.completedFuture(new SlotPermit());
     }
diff --git a/temporal-sdk/src/main/java/io/temporal/worker/tuning/SlotSupplier.java b/temporal-sdk/src/main/java/io/temporal/worker/tuning/SlotSupplier.java
@@ -22,7 +22,7 @@
 
 import io.temporal.common.Experimental;
 import java.util.Optional;
-import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Future;
 
 /**
  * A SlotSupplier is responsible for managing the number of slots available for a given type of
@@ -49,7 +49,7 @@ public interface SlotSupplier<SI extends SlotInfo> {
    * @return A future that will be completed with a permit to use the slot when one becomes
    *     available. Never return null, or complete the future with null.
    */
-  CompletableFuture<SlotPermit> reserveSlot(SlotReserveContext<SI> ctx) throws Exception;
+  Future<SlotPermit> reserveSlot(SlotReserveContext<SI> ctx) throws Exception;
 
   /**
    * This function is called when trying to reserve slots for "eager" workflow and activity tasks.
diff --git a/temporal-sdk/src/test/java/io/temporal/testUtils/CountingSlotSupplier.java b/temporal-sdk/src/test/java/io/temporal/testUtils/CountingSlotSupplier.java
@@ -22,8 +22,8 @@
 
 import io.temporal.worker.tuning.*;
 import java.util.Optional;
-import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.*;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 
 public class CountingSlotSupplier<SI extends SlotInfo> extends FixedSizeSlotSupplier<SI> {
@@ -38,13 +38,47 @@ public CountingSlotSupplier(int numSlots) {
   }
 
   @Override
-  public CompletableFuture<SlotPermit> reserveSlot(SlotReserveContext<SI> ctx) throws Exception {
-    CompletableFuture<SlotPermit> p = super.reserveSlot(ctx);
-    return p.thenApply(
-        permit -> {
+  public Future<SlotPermit> reserveSlot(SlotReserveContext<SI> ctx) throws Exception {
+    Future<SlotPermit> originalFuture = super.reserveSlot(ctx);
+
+    return new Future<SlotPermit>() {
+      private final AtomicBoolean callbackInvoked = new AtomicBoolean(false);
+
+      private SlotPermit executeCallbackIfNeeded(SlotPermit permit) {
+        if (callbackInvoked.compareAndSet(false, true)) {
           reservedCount.incrementAndGet();
-          return permit;
-        });
+        }
+        return permit;
+      }
+
+      @Override
+      public boolean cancel(boolean mayInterruptIfRunning) {
+        return originalFuture.cancel(mayInterruptIfRunning);
+      }
+
+      @Override
+      public boolean isCancelled() {
+        return originalFuture.isCancelled();
+      }
+
+      @Override
+      public boolean isDone() {
+        return originalFuture.isDone();
+      }
+
+      @Override
+      public SlotPermit get() throws InterruptedException, ExecutionException {
+        SlotPermit permit = originalFuture.get();
+        return executeCallbackIfNeeded(permit);
+      }
+
+      @Override
+      public SlotPermit get(long timeout, TimeUnit unit)
+          throws InterruptedException, ExecutionException, TimeoutException {
+        SlotPermit permit = originalFuture.get(timeout, unit);
+        return executeCallbackIfNeeded(permit);
+      }
+    };
   }
 
   @Override

Original file line number	Diff line number	Diff line change
`@@ -25,6 +25,7 @@`
`25`	`25`	`import java.util.Optional;`
`26`	`26`	`import java.util.Queue;`
`27`	`27`	`import java.util.concurrent.CompletableFuture;`
	`28`	`+import java.util.concurrent.Future;`
`28`	`29`	`import java.util.concurrent.locks.ReentrantLock;`
`29`	`30`
`30`	`31`	`/**`
`@@ -110,7 +111,7 @@ public FixedSizeSlotSupplier(int numSlots) {`
`110`	`111`	`}`
`111`	`112`
`112`	`113`	`@Override`
`113`		`- public CompletableFuture<SlotPermit> reserveSlot(SlotReserveContext<SI> ctx) throws Exception {`
	`114`	`+ public Future<SlotPermit> reserveSlot(SlotReserveContext<SI> ctx) throws Exception {`
`114`	`115`	`return executorSlotsSemaphore.acquire().thenApply(ignored -> new SlotPermit());`
`115`	`116`	`}`
`116`	`117`
Original file line number	Diff line number	Diff line change
`@@ -202,7 +202,7 @@ private ResourceBasedSlotSupplier(`
`202`	`202`	`}`
`203`	`203`
`204`	`204`	`@Override`
`205`		`- public CompletableFuture<SlotPermit> reserveSlot(SlotReserveContext<SI> ctx) throws Exception {`
	`205`	`+ public Future<SlotPermit> reserveSlot(SlotReserveContext<SI> ctx) throws Exception {`
`206`	`206`	`if (ctx.getNumIssuedSlots() < options.getMinimumSlots()) {`
`207`	`207`	`return CompletableFuture.completedFuture(new SlotPermit());`
`208`	`208`	`}`