replace phaser with future

Author: zhongkechen

sdk/src/main/java/com/amazonaws/lambda/durable/DurableFuture.java

@@ -9,8 +9,8 @@ public interface DurableFuture<T> {
     /**
      * Blocks until the operation completes and returns the result.
      *
-     * <p>This delegates to operation.get() which handles: - Phaser blocking (arriveAndAwaitAdvance) - Thread
-     * deregistration (allows suspension) - Thread reactivation (resumes execution) - Result retrieval
+     * <p>This delegates to operation.get() which handles: - Thread deregistration (allows suspension) - Thread
+     * reactivation (resumes execution) - Result retrieval
      *
      * @return the operation result
      */

sdk/src/main/java/com/amazonaws/lambda/durable/execution/CheckpointBatcher.java

@@ -139,10 +139,6 @@ private void checkpointBatch(List<OperationUpdate> updates) {
             logger.debug("Durable checkpoint API called (latency={}ns): {}.", System.nanoTime() - startTime, response);
 
             // Notify callback of completion
-            // TODO: sam local backend returns no new execution state when called with zero
-            // updates. WHY?
-            // This means the polling will never receive an operation update and complete
-            // the Phaser.
             checkpointToken = response.checkpointToken();
             if (response.newExecutionState() != null) {
                 // fetch all pages of operations

sdk/src/main/java/com/amazonaws/lambda/durable/execution/ExecutionManager.java

@@ -4,13 +4,13 @@
 
 import com.amazonaws.lambda.durable.DurableConfig;
 import com.amazonaws.lambda.durable.exception.UnrecoverableDurableExecutionException;
+import com.amazonaws.lambda.durable.operation.BaseDurableOperation;
 import java.time.Duration;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.Phaser;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
 import org.slf4j.Logger;
@@ -28,8 +28,8 @@
  * <ul>
  *   <li>Execution state (operations, checkpoint token)
  *   <li>Thread lifecycle (registration/deregistration)
- *   <li>Phaser management (coordination)
- *   <li>Checkpoint batching (via CheckpointManager)
+ *   <li>Checkpoint batching (via CheckpointBatcher)
+ *   <li>Checkpoint result handling (CheckpointBatcher callback)
  *   <li>Polling (for waits and retries)
  * </ul>
  *
@@ -43,15 +43,16 @@ public class ExecutionManager {
     private static final Logger logger = LoggerFactory.getLogger(ExecutionManager.class);
 
     // ===== Execution State =====
-    private final Map<String, Operation> operations;
+    private final Map<String, Operation> operationStorage;
     private final String executionOperationId;
     private final String durableExecutionArn;
     private final AtomicReference<ExecutionMode> executionMode;
 
     // ===== Thread Coordination =====
+    private final Map<String, BaseDurableOperation<?>> registeredOperations =
+            Collections.synchronizedMap(new HashMap<>());
     private final Map<String, ThreadType> activeThreads = Collections.synchronizedMap(new HashMap<>());
     private static final ThreadLocal<OperationContext> currentContext = new ThreadLocal<>();
-    private final Map<String, Phaser> openPhasers = Collections.synchronizedMap(new HashMap<>());
     private final CompletableFuture<Void> executionExceptionFuture = new CompletableFuture<>();
 
     // ===== Checkpoint Batching =====
@@ -69,12 +70,12 @@ public ExecutionManager(
         this.checkpointBatcher =
                 new CheckpointBatcher(config, durableExecutionArn, checkpointToken, this::onCheckpointComplete);
 
-        this.operations = checkpointBatcher.fetchAllPages(initialExecutionState).stream()
+        this.operationStorage = checkpointBatcher.fetchAllPages(initialExecutionState).stream()
                 .collect(Collectors.toConcurrentMap(Operation::id, op -> op));
 
         // Start in REPLAY mode if we have more than just the initial EXECUTION operation
         this.executionMode =
-                new AtomicReference<>(operations.size() > 1 ? ExecutionMode.REPLAY : ExecutionMode.EXECUTION);
+                new AtomicReference<>(operationStorage.size() > 1 ? ExecutionMode.REPLAY : ExecutionMode.EXECUTION);
     }
 
     // ===== State Management =====
@@ -87,24 +88,22 @@ public boolean isReplaying() {
         return executionMode.get() == ExecutionMode.REPLAY;
     }
 
+    public void registerOperation(BaseDurableOperation<?> operation) {
+        registeredOperations.put(operation.getOperationId(), operation);
+    }
+
     // ===== Checkpoint Completion Handler =====
-    /** Called by CheckpointManager when a checkpoint completes. Updates state and advances phasers. */
+    /** Called by CheckpointManager when a checkpoint completes. Updates operationStorage and notify operations . */
     private void onCheckpointComplete(List<Operation> newOperations) {
-        // Update operation storage
-        newOperations.forEach(op -> operations.put(op.id(), op));
-
-        // Advance phasers for completed operations
-        for (Operation operation : newOperations) {
-            if (openPhasers.containsKey(operation.id()) && isTerminalStatus(operation.status())) {
-                var phaser = openPhasers.get(operation.id());
-
-                // Two-phase completion
-                logger.trace("Advancing phaser 0 -> 1: {}", phaser);
-                phaser.arriveAndAwaitAdvance();
-                logger.trace("Advancing phaser 1 -> 2: {}", phaser);
-                phaser.arriveAndAwaitAdvance();
-            }
-        }
+        newOperations.forEach(op -> {
+            // Update operation storage
+            operationStorage.put(op.id(), op);
+            // call registered operation's onCheckpointComplete method for completed operations
+            registeredOperations.computeIfPresent(op.id(), (id, operation) -> {
+                operation.onCheckpointComplete(op);
+                return operation;
+            });
+        });
     }
 
     /**
@@ -115,7 +114,7 @@ private void onCheckpointComplete(List<Operation> newOperations) {
      * @return the existing operation, or null if not found (first execution)
      */
     public Operation getOperationAndUpdateReplayState(String operationId) {
-        var existing = operations.get(operationId);
+        var existing = operationStorage.get(operationId);
         if (executionMode.get() == ExecutionMode.REPLAY) {
             if (existing == null || !isTerminalStatus(existing.status())) {
                 if (executionMode.compareAndSet(ExecutionMode.REPLAY, ExecutionMode.EXECUTION)) {
@@ -127,7 +126,7 @@ public Operation getOperationAndUpdateReplayState(String operationId) {
     }
 
     public Operation getExecutionOperation() {
-        return operations.get(executionOperationId);
+        return operationStorage.get(executionOperationId);
     }
 
     // ===== Thread Coordination =====
@@ -185,15 +184,6 @@ public void deregisterActiveThreadAndUnsetCurrentContext(String threadId) {
         }
     }
 
-    // ===== Phaser Management =====
-
-    public Phaser startPhaser(String operationId) {
-        var phaser = new Phaser(1);
-        openPhasers.put(operationId, phaser);
-        logger.trace("Started phaser for operation '{}'", operationId);
-        return phaser;
-    }
-
     // ===== Checkpointing =====
 
     // This method will checkpoint the operation updates to the durable backend and return a future which completes

sdk/src/main/java/com/amazonaws/lambda/durable/execution/ExecutionPhase.java

@@ -9,9 +9,8 @@
 /**
  * Shared executor for internal SDK coordination tasks.
  *
- * <p>This executor is used for SDK-internal operations such as checkpoint batching, polling for wait completion, and
- * phaser coordination. It is separate from the user-configured executor (via {@code DurableConfig}) which runs
- * user-defined operations.
+ * <p>This executor is used for SDK-internal operations such as checkpoint batching. It is separate from the
+ * user-configured executor (via {@code DurableConfig}) which runs user-defined operations.
  *
  * <p>Using a dedicated thread pool ensures SDK coordination tasks are isolated from user code.
  */

sdk/src/main/java/com/amazonaws/lambda/durable/execution/InternalExecutor.java

@@ -9,18 +9,17 @@
 import com.amazonaws.lambda.durable.exception.SerDesException;
 import com.amazonaws.lambda.durable.exception.UnrecoverableDurableExecutionException;
 import com.amazonaws.lambda.durable.execution.ExecutionManager;
-import com.amazonaws.lambda.durable.execution.ExecutionPhase;
 import com.amazonaws.lambda.durable.execution.ThreadType;
 import com.amazonaws.lambda.durable.serde.SerDes;
 import com.amazonaws.lambda.durable.util.ExceptionHelper;
 import java.time.Duration;
 import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.Phaser;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import software.amazon.awssdk.services.lambda.model.ErrorObject;
 import software.amazon.awssdk.services.lambda.model.Operation;
+import software.amazon.awssdk.services.lambda.model.OperationStatus;
 import software.amazon.awssdk.services.lambda.model.OperationType;
 import software.amazon.awssdk.services.lambda.model.OperationUpdate;
 
@@ -39,7 +38,7 @@
  * <ul>
  *   <li>Starting multiple async operations quickly
  *   <li>Blocking on results later when needed
- *   <li>Proper thread coordination via Phasers
+ *   <li>Proper thread coordination via future
  * </ul>
  */
 public abstract class BaseDurableOperation<T> implements DurableFuture<T> {
@@ -51,7 +50,7 @@ public abstract class BaseDurableOperation<T> implements DurableFuture<T> {
     private final ExecutionManager executionManager;
     private final TypeToken<T> resultTypeToken;
     private final SerDes resultSerDes;
-    private final Phaser phaser;
+    private final CompletableFuture<Void> completionFuture;
 
     public BaseDurableOperation(
             String operationId,
@@ -67,8 +66,10 @@ public BaseDurableOperation(
         this.resultTypeToken = resultTypeToken;
         this.resultSerDes = resultSerDes;
 
-        // todo: phaser could be used only in ExecutionManager and invisible from operations.
-        this.phaser = executionManager.startPhaser(operationId);
+        this.completionFuture = new CompletableFuture<>();
+
+        // register this operation in ExecutionManager so that the operation can receive updates from ExecutionManager
+        executionManager.registerOperation(this);
     }
 
     /** Gets the unique identifier for this operation. */
@@ -96,7 +97,7 @@ public OperationType getType() {
      *
      * <ul>
      *   <li>Thread deregistration (allows suspension)
-     *   <li>Phaser blocking (waits for operation to complete)
+     *   <li>Blocking (waits for operation to complete)
      *   <li>Thread reactivation (resumes execution)
      *   <li>Result retrieval
      * </ul>
@@ -116,7 +117,7 @@ protected Operation getOperation() {
     }
 
     /**
-     * check if it's called from a Step.
+     * Checks if it's called from a Step.
      *
      * @throws IllegalDurableOperationException if it's in a step
      */
@@ -131,43 +132,43 @@ private void validateCurrentThreadType() {
         }
     }
 
-    // phase control utilities
+    /** Checks if this operation is completed */
+    protected boolean isOperationCompleted() {
+        return completionFuture.isDone();
+    }
+
+    /** Waits for the operation to complete and suspends the execution if no active thread is running */
     protected Operation waitForOperationCompletion() {
 
         validateCurrentThreadType();
 
-        // register to prevent the state from advancing
-        phaser.register();
-
-        // If we are in a replay where the operation is already complete (SUCCEEDED /
-        // FAILED), the Phaser will be
-        // advanced in .execute() already and we don't block but return the result
-        // immediately.
-        if (phaser.getPhase() == ExecutionPhase.RUNNING.getValue()) {
-            // Operation not done yet
-            var context = executionManager.getCurrentContext();
-            // Deregister current context - allows suspension
-            logger.debug(
-                    "get() on {} attempting to deregister context: {}",
-                    getType(),
-                    executionManager.getCurrentContext().contextId());
-            deregisterActiveThreadAndUnsetCurrentContext(context.contextId());
-
-            // Block until operation completes
-            logger.trace("Waiting for operation to finish {} (Phaser: {})", getOperationId(), phaser);
-            phaser.arriveAndAwaitAdvance();
-
-            // Reactivate current context
-            registerActiveThread(context.contextId(), context.threadType());
-            setCurrentContext(context.contextId(), context.threadType());
-
-            // Complete phase 1
-            phaser.arriveAndDeregister();
-        } else {
-            // The phaser is already completed. Deregister now.
-            phaser.arriveAndDeregister();
+        var context = executionManager.getCurrentContext();
+
+        // Use a synchronized block here to prevent the completionFuture from being completed by the execution thread
+        // (a step or child context thread) when it's inside the `if` block where the completion check is done (not
+        // completed) while the callback isn't added to the completionFuture or the current thread isn't deregistered.
+        synchronized (completionFuture) {
+            if (!isOperationCompleted()) {
+                // Operation not done yet
+                logger.debug("get() on {} attempting to deregister context: {}", getType(), context.contextId());
+
+                // Add a callback to completionFuture so that when the completionFuture is completed,
+                // it will register the current Context thread synchronously to make sure it is always registered
+                // before the execution thread (Step or child context) is deregistered.
+                completionFuture.thenRun(() -> registerActiveThread(context.contextId(), context.threadType()));
+
+                // Deregister the current thread to allow suspension
+                deregisterActiveThreadAndUnsetCurrentContext(context.contextId());
+            }
         }
 
+        // Block until operation completes. No-op if the future is already completed.
+        logger.trace("Waiting for operation to finish {} ({})", getOperationId(), completionFuture);
+        completionFuture.join();
+
+        // Reactivate current context. No-op if this is called twice.
+        setCurrentContext(context.contextId(), context.threadType());
+
         // Get result based on status
         var op = getOperation();
         if (op == null) {
@@ -177,13 +178,27 @@ protected Operation waitForOperationCompletion() {
         return op;
     }
 
+    /** Receives operation updates from ExecutionManager and updates the internal state of the operation */
+    public void onCheckpointComplete(Operation operation) {
+        if (isTerminalStatus(operation.status())) {
+            synchronized (completionFuture) {
+                logger.trace("In onCheckpointComplete, completing operation {} ({})", operationId, completionFuture);
+                completionFuture.complete(null);
+            }
+        }
+    }
+
+    /** Marks the operation as already completed (in replay). */
     protected void markAlreadyCompleted() {
-        // Operation is already completed in a relay. We advance and deregister from the Phaser
-        // so that get method doesn't block and returns the result immediately.
-        logger.trace("Detected terminal status during replay. Advancing phaser 0 -> 1 {}.", phaser);
-        phaser.arriveAndDeregister(); // Phase 0 -> 1
+        // When the operation is already completed in a replay, we complete completionFuture immediately
+        // so that the `get` method will be unblocked and the context thread will be registered
+        logger.trace("In markAlreadyCompleted, completing operation: {} ({}).", operationId, completionFuture);
+        synchronized (completionFuture) {
+            completionFuture.complete(null);
+        }
     }
 
+    // terminate the execution
     protected T terminateExecution(UnrecoverableDurableExecutionException exception) {
         executionManager.terminateExecution(exception);
         // Exception is already thrown from above. Keep the throw statement below to make tests happy
@@ -194,7 +209,7 @@ protected T terminateExecutionWithIllegalDurableOperationException(String messag
         return terminateExecution(new IllegalDurableOperationException(message));
     }
 
-    // advanced phase control used by Step only
+    // advanced thread and context control
     protected void deregisterActiveThreadAndUnsetCurrentContext(String threadId) {
         executionManager.deregisterActiveThreadAndUnsetCurrentContext(threadId);
     }
@@ -298,4 +313,12 @@ protected void validateReplay(Operation checkpointed) {
                     operationId, checkpointed.name(), getName())));
         }
     }
+
+    private boolean isTerminalStatus(OperationStatus status) {
+        return status == OperationStatus.SUCCEEDED
+                || status == OperationStatus.FAILED
+                || status == OperationStatus.CANCELLED
+                || status == OperationStatus.TIMED_OUT
+                || status == OperationStatus.STOPPED;
+    }
 }

sdk/src/main/java/com/amazonaws/lambda/durable/operation/BaseDurableOperation.java

@@ -52,7 +52,7 @@ public void execute() {
 
             switch (existing.status()) {
                 case SUCCEEDED, FAILED, TIMED_OUT -> {
-                    // Terminal state - complete phaser immediately
+                    // Terminal state - complete the operation immediately
                     markAlreadyCompleted();
                     return;
                 }