move nextOperationId from DurableContext to ExecutionManager

Author: zhongkechen

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

@@ -12,7 +12,6 @@
 import com.amazonaws.services.lambda.runtime.Context;
 import java.time.Duration;
 import java.util.Objects;
-import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Supplier;
 import org.slf4j.LoggerFactory;
 
@@ -22,7 +21,6 @@ public class DurableContext {
     private final ExecutionManager executionManager;
     private final DurableConfig durableConfig;
     private final Context lambdaContext;
-    private final AtomicInteger operationCounter;
     private final DurableLogger logger;
     private final ExecutionContext executionContext;
 
@@ -31,7 +29,6 @@ public class DurableContext {
         this.executionManager = executionManager;
         this.durableConfig = durableConfig;
         this.lambdaContext = lambdaContext;
-        this.operationCounter = new AtomicInteger(0);
         this.executionContext = new ExecutionContext(executionManager.getDurableExecutionArn());
 
         var requestId = lambdaContext != null ? lambdaContext.getAwsRequestId() : null;
@@ -89,11 +86,8 @@ public <T> DurableFuture<T> stepAsync(String name, TypeToken<T> typeToken, Suppl
         if (config.serDes() == null) {
             config = config.toBuilder().serDes(durableConfig.getSerDes()).build();
         }
-        var operationId = nextOperationId();
-
         // Create and start step operation with TypeToken
-        var operation = new StepOperation<>(
-                operationId, name, func, typeToken, config, executionManager, logger, durableConfig);
+        var operation = new StepOperation<>(name, func, typeToken, config, executionManager, logger, durableConfig);
 
         operation.execute(); // Start the step (returns immediately)
 
@@ -107,10 +101,8 @@ public Void wait(Duration duration) {
     }
 
     public Void wait(String waitName, Duration duration) {
-        var operationId = nextOperationId();
-
         // Create and start wait operation
-        var operation = new WaitOperation(operationId, waitName, duration, executionManager);
+        var operation = new WaitOperation(waitName, duration, executionManager);
 
         operation.execute(); // Checkpoint the wait
         return operation.get(); // Block (will throw SuspendExecutionException if needed)
@@ -176,11 +168,8 @@ public <T, U> DurableFuture<T> invokeAsync(
         if (config.payloadSerDes() == null) {
             config = config.toBuilder().payloadSerDes(durableConfig.getSerDes()).build();
         }
-        var operationId = nextOperationId();
-
         // Create and start invoke operation
-        var operation =
-                new InvokeOperation<>(operationId, name, functionName, payload, typeToken, config, executionManager);
+        var operation = new InvokeOperation<>(name, functionName, payload, typeToken, config, executionManager);
 
         operation.execute(); // checkpoint the invoke operation
         return operation; // Block (will throw SuspendExecutionException if needed)
@@ -204,9 +193,7 @@ public <T> DurableCallbackFuture<T> createCallback(String name, TypeToken<T> typ
         if (config.serDes() == null) {
             config = config.toBuilder().serDes(durableConfig.getSerDes()).build();
         }
-        var operationId = nextOperationId();
-
-        var operation = new CallbackOperation<>(operationId, name, typeToken, config, executionManager);
+        var operation = new CallbackOperation<>(name, typeToken, config, executionManager);
         operation.execute();
 
         return operation;
@@ -234,11 +221,4 @@ public DurableLogger getLogger() {
     public ExecutionContext getExecutionContext() {
         return executionContext;
     }
-
-    // ============= internal utilities ===============
-
-    /** Get the next operationId (latest operationId + 1) */
-    private String nextOperationId() {
-        return String.valueOf(operationCounter.incrementAndGet());
-    }
 }

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

@@ -49,7 +49,7 @@ class CheckpointBatcher {
                 MAX_ITEM_COUNT, MAX_BATCH_SIZE_BYTES, CheckpointBatcher::estimateSize, this::checkpointBatch);
     }
 
-    /** Queues a checkpoint request for batched execution */
+    /** Queues a operation update for batched checkpoint */
     CompletableFuture<Void> checkpoint(OperationUpdate update) {
         logger.debug("Checkpoint request received: Action {}", update.action());
         return checkpointApiRequestBatcher.submit(update, config.getCheckpointDelay());

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

@@ -11,6 +11,7 @@
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Phaser;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
 import org.slf4j.Logger;
@@ -47,6 +48,7 @@ public class ExecutionManager {
     private final String executionOperationId;
     private final String durableExecutionArn;
     private final AtomicReference<ExecutionMode> executionMode;
+    private final AtomicInteger operationCounter;
 
     // ===== Thread Coordination =====
     private final Map<String, ThreadType> activeThreads = Collections.synchronizedMap(new HashMap<>());
@@ -62,6 +64,7 @@ public ExecutionManager(
             String checkpointToken,
             CheckpointUpdatedExecutionState initialExecutionState,
             DurableConfig config) {
+        this.operationCounter = new AtomicInteger(0);
         this.durableExecutionArn = durableExecutionArn;
         this.executionOperationId = initialExecutionState.operations().get(0).id();
 
@@ -217,7 +220,15 @@ public CompletableFuture<Operation> pollForOperationUpdates(String operationId,
         return checkpointBatcher.pollForUpdate(operationId, delay);
     }
 
-    // ===== Utilities =====
+    // ============= internal utilities ===============
+
+    /** Get the next operationId (latest operationId + 1) */
+    public String nextOperationId() {
+        return String.valueOf(operationCounter.incrementAndGet());
+    }
+
+    // ============= lifecycle management =================
+
     public void shutdown() {
         checkpointBatcher.shutdown();
     }

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

@@ -54,19 +54,20 @@ public abstract class BaseDurableOperation<T> implements DurableFuture<T> {
     private final Phaser phaser;
 
     public BaseDurableOperation(
-            String operationId,
             String name,
             OperationType operationType,
             TypeToken<T> resultTypeToken,
             SerDes resultSerDes,
             ExecutionManager executionManager) {
-        this.operationId = operationId;
         this.name = name;
         this.operationType = operationType;
         this.executionManager = executionManager;
         this.resultTypeToken = resultTypeToken;
         this.resultSerDes = resultSerDes;
 
+        // get the next operation id from executionManager
+        this.operationId = executionManager.nextOperationId();
+
         // todo: phaser could be used only in ExecutionManager and invisible from operations.
         this.phaser = executionManager.startPhaser(operationId);
     }

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

@@ -25,12 +25,8 @@ public class CallbackOperation<T> extends BaseDurableOperation<T> implements Dur
     private String callbackId;
 
     public CallbackOperation(
-            String operationId,
-            String name,
-            TypeToken<T> resultTypeToken,
-            CallbackConfig config,
-            ExecutionManager executionManager) {
-        super(operationId, name, OperationType.CALLBACK, resultTypeToken, config.serDes(), executionManager);
+            String name, TypeToken<T> resultTypeToken, CallbackConfig config, ExecutionManager executionManager) {
+        super(name, OperationType.CALLBACK, resultTypeToken, config.serDes(), executionManager);
         this.config = config;
     }
 

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

@@ -22,14 +22,13 @@ public class InvokeOperation<T, U> extends BaseDurableOperation<T> {
     private final SerDes payloadSerDes;
 
     public InvokeOperation(
-            String operationId,
             String name,
             String functionName,
             U payload,
             TypeToken<T> resultTypeToken,
             InvokeConfig config,
             ExecutionManager executionManager) {
-        super(operationId, name, OperationType.CHAINED_INVOKE, resultTypeToken, config.serDes(), executionManager);
+        super(name, OperationType.CHAINED_INVOKE, resultTypeToken, config.serDes(), executionManager);
 
         this.functionName = functionName;
         this.payload = payload;

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

@@ -37,15 +37,14 @@ public class StepOperation<T> extends BaseDurableOperation<T> {
     private final ExecutorService userExecutor;
 
     public StepOperation(
-            String operationId,
             String name,
             Supplier<T> function,
             TypeToken<T> resultTypeToken,
             StepConfig config,
             ExecutionManager executionManager,
             DurableLogger durableLogger,
             DurableConfig durableConfig) {
-        super(operationId, name, OperationType.STEP, resultTypeToken, config.serDes(), executionManager);
+        super(name, OperationType.STEP, resultTypeToken, config.serDes(), executionManager);
 
         this.function = function;
         this.config = config;

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

@@ -23,8 +23,8 @@ public class WaitOperation extends BaseDurableOperation<Void> {
 
     private final Duration duration;
 
-    public WaitOperation(String operationId, String name, Duration duration, ExecutionManager executionManager) {
-        super(operationId, name, OperationType.WAIT, TypeToken.get(Void.class), NOOP_SER_DES, executionManager);
+    public WaitOperation(String name, Duration duration, ExecutionManager executionManager) {
+        super(name, OperationType.WAIT, TypeToken.get(Void.class), NOOP_SER_DES, executionManager);
         this.duration = duration;
     }