add an optional stepContext parameter

Author: zhongkechen

examples/src/main/java/software/amazon/lambda/durable/examples/RetryExample.java

@@ -31,7 +31,7 @@ public class RetryExample extends DurableHandler<Object, String> {
     @Override
     public String handleRequest(Object input, DurableContext context) {
         // Step 1: Record start time
-        startTime = context.step("record-start-time", Instant.class, Instant::now);
+        startTime = context.step("record-start-time", Instant.class, () -> Instant.now());
         logger.info("Recorded start time: {}", startTime);
 
         // Step 2: Call that never retries (fails immediately)

sdk/src/main/java/software/amazon/lambda/durable/BaseContext.java

@@ -0,0 +1,79 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.lambda.durable;
+
+import com.amazonaws.services.lambda.runtime.Context;
+import software.amazon.lambda.durable.execution.ExecutionManager;
+import software.amazon.lambda.durable.logging.DurableLogger;
+
+public abstract class BaseContext {
+    protected final ExecutionManager executionManager;
+    private final DurableConfig durableConfig;
+    private final Context lambdaContext;
+    private final ExecutionContext executionContext;
+    private final String parentContextId;
+    private boolean isReplaying;
+
+    /** Shared initialization — sets all fields but performs no thread registration. */
+    protected BaseContext(
+            ExecutionManager executionManager, DurableConfig durableConfig, Context lambdaContext, String contextId) {
+        this.executionManager = executionManager;
+        this.durableConfig = durableConfig;
+        this.lambdaContext = lambdaContext;
+        this.parentContextId = contextId;
+        this.executionContext = new ExecutionContext(executionManager.getDurableExecutionArn());
+        this.isReplaying = executionManager.hasOperationsForContext(contextId);
+    }
+
+    // =============== accessors ================
+    public abstract DurableLogger getLogger();
+
+    public Context getLambdaContext() {
+        return lambdaContext;
+    }
+
+    /**
+     * Returns metadata about the current durable execution.
+     *
+     * <p>The execution context provides information that remains constant throughout the execution lifecycle, such as
+     * the durable execution ARN. This is useful for tracking execution progress, correlating logs, and referencing this
+     * execution in external systems.
+     *
+     * @return the execution context
+     */
+    public ExecutionContext getExecutionContext() {
+        return executionContext;
+    }
+
+    /**
+     * Returns the configuration for durable execution behavior.
+     *
+     * @return the durable configuration
+     */
+    public DurableConfig getDurableConfig() {
+        return durableConfig;
+    }
+
+    // ============= internal utilities ===============
+
+    /** Gets the context ID for this context. Null for root context, set for child contexts. */
+    public String getParentContextId() {
+        return parentContextId;
+    }
+
+    public ExecutionManager getExecutionManager() {
+        return executionManager;
+    }
+
+    /** Returns whether this context is currently in replay mode. */
+    boolean isReplaying() {
+        return isReplaying;
+    }
+
+    /**
+     * Transitions this context from replay to execution mode. Called when the first un-cached operation is encountered.
+     */
+    void setExecutionMode() {
+        this.isReplaying = false;
+    }
+}

sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java

@@ -20,28 +20,17 @@
 import software.amazon.lambda.durable.operation.WaitOperation;
 import software.amazon.lambda.durable.validation.ParameterValidator;
 
-public class DurableContext {
-    private static final String ROOT_CONTEXT = "Root";
-
-    private final ExecutionManager executionManager;
-    private final DurableConfig durableConfig;
-    private final Context lambdaContext;
+public class DurableContext extends BaseContext {
     private final AtomicInteger operationCounter;
     private final DurableLogger logger;
     private final ExecutionContext executionContext;
-    private final String contextId;
-    private boolean isReplaying;
 
     /** Shared initialization — sets all fields but performs no thread registration. */
     private DurableContext(
             ExecutionManager executionManager, DurableConfig durableConfig, Context lambdaContext, String contextId) {
-        this.executionManager = executionManager;
-        this.durableConfig = durableConfig;
-        this.lambdaContext = lambdaContext;
-        this.contextId = contextId;
+        super(executionManager, durableConfig, lambdaContext, contextId);
         this.operationCounter = new AtomicInteger(0);
         this.executionContext = new ExecutionContext(executionManager.getDurableExecutionArn());
-        this.isReplaying = executionManager.hasOperationsForContext(contextId);
 
         var requestId = lambdaContext != null ? lambdaContext.getAwsRequestId() : null;
         this.logger = new DurableLogger(
@@ -54,19 +43,18 @@ private DurableContext(
     /**
      * Creates a root context and registers the current thread for execution coordination.
      *
-     * <p>The context itself always has a null contextId (making it a root context). The thread is registered with the
-     * ExecutionManager using the default {@link #ROOT_CONTEXT} identifier.
+     * <p>The context itself always has a null contextId (making it a root context).
      *
      * @param executionManager the execution manager
      * @param durableConfig the durable configuration
      * @param lambdaContext the Lambda context
      * @return a new root DurableContext
      */
-    static DurableContext createRootContext(
+    public static DurableContext createRootContext(
             ExecutionManager executionManager, DurableConfig durableConfig, Context lambdaContext) {
         var ctx = new DurableContext(executionManager, durableConfig, lambdaContext, null);
-        executionManager.registerActiveThread(ROOT_CONTEXT);
-        executionManager.setCurrentThreadContext(new ThreadContext(ROOT_CONTEXT, ThreadType.CONTEXT));
+        executionManager.registerActiveThread(null);
+        executionManager.setCurrentThreadContext(new ThreadContext(null, ThreadType.CONTEXT));
         return ctx;
     }
 
@@ -75,67 +63,108 @@ static DurableContext createRootContext(
      * ChildContextOperation, which registers on the parent thread before the executor runs and sets the context on the
      * child thread inside the executor.
      *
-     * @param executionManager the execution manager
-     * @param durableConfig the durable configuration
-     * @param lambdaContext the Lambda context
-     * @param contextId the child context's ID (the CONTEXT operation's operation ID)
+     * @param childContextId the child context's ID (the CONTEXT operation's operation ID)
      * @return a new DurableContext for the child context
      */
-    public static DurableContext createChildContext(
-            ExecutionManager executionManager, DurableConfig durableConfig, Context lambdaContext, String contextId) {
-        return new DurableContext(executionManager, durableConfig, lambdaContext, contextId);
+    public DurableContext createChildContext(String childContextId) {
+        return new DurableContext(executionManager, getDurableConfig(), getLambdaContext(), childContextId);
+    }
+
+    /**
+     * Creates a step context for executing step operations.
+     *
+     * @return a new StepContext instance
+     */
+    public StepContext createStepContext(String stepOperationId) {
+        return new StepContext(executionManager, getDurableConfig(), getLambdaContext(), stepOperationId);
     }
 
     // ========== step methods ==========
 
-    public <T> T step(String name, Class<T> resultType, Supplier<T> func) {
+    public <T> T step(String name, Class<T> resultType, Function<StepContext, T> func) {
         return step(name, TypeToken.get(resultType), func, StepConfig.builder().build());
     }
 
-    public <T> T step(String name, Class<T> resultType, Supplier<T> func, StepConfig config) {
+    public <T> T step(String name, Class<T> resultType, Function<StepContext, T> func, StepConfig config) {
         // Simply delegate to stepAsync and block on the result
         return stepAsync(name, resultType, func, config).get();
     }
 
-    public <T> T step(String name, TypeToken<T> typeToken, Supplier<T> func) {
+    public <T> T step(String name, TypeToken<T> typeToken, Function<StepContext, T> func) {
         return step(name, typeToken, func, StepConfig.builder().build());
     }
 
-    public <T> T step(String name, TypeToken<T> typeToken, Supplier<T> func, StepConfig config) {
+    public <T> T step(String name, TypeToken<T> typeToken, Function<StepContext, T> func, StepConfig config) {
         // Simply delegate to stepAsync and block on the result
         return stepAsync(name, typeToken, func, config).get();
     }
 
-    public <T> DurableFuture<T> stepAsync(String name, Class<T> resultType, Supplier<T> func) {
+    public <T> DurableFuture<T> stepAsync(String name, Class<T> resultType, Function<StepContext, T> func) {
         return stepAsync(
                 name, TypeToken.get(resultType), func, StepConfig.builder().build());
     }
 
-    public <T> DurableFuture<T> stepAsync(String name, Class<T> resultType, Supplier<T> func, StepConfig config) {
+    public <T> DurableFuture<T> stepAsync(
+            String name, Class<T> resultType, Function<StepContext, T> func, StepConfig config) {
         return stepAsync(name, TypeToken.get(resultType), func, config);
     }
 
-    public <T> DurableFuture<T> stepAsync(String name, TypeToken<T> typeToken, Supplier<T> func) {
+    public <T> DurableFuture<T> stepAsync(String name, TypeToken<T> typeToken, Function<StepContext, T> func) {
         return stepAsync(name, typeToken, func, StepConfig.builder().build());
     }
 
-    public <T> DurableFuture<T> stepAsync(String name, TypeToken<T> typeToken, Supplier<T> func, StepConfig config) {
+    public <T> DurableFuture<T> stepAsync(
+            String name, TypeToken<T> typeToken, Function<StepContext, T> func, StepConfig config) {
         Objects.requireNonNull(config, "config cannot be null");
         Objects.requireNonNull(typeToken, "typeToken cannot be null");
         if (config.serDes() == null) {
-            config = config.toBuilder().serDes(durableConfig.getSerDes()).build();
+            config = config.toBuilder().serDes(getDurableConfig().getSerDes()).build();
         }
         var operationId = nextOperationId();
 
         // Create and start step operation with TypeToken
-        var operation = new StepOperation<>(
-                operationId, name, func, typeToken, config, executionManager, logger, durableConfig, contextId);
+        var operation = new StepOperation<>(operationId, name, func, typeToken, config, this);
 
         operation.execute(); // Start the step (returns immediately)
 
         return operation;
     }
 
+    public <T> T step(String name, Class<T> resultType, Supplier<T> func) {
+        return step(name, TypeToken.get(resultType), func, StepConfig.builder().build());
+    }
+
+    public <T> T step(String name, Class<T> resultType, Supplier<T> func, StepConfig config) {
+        // Simply delegate to stepAsync and block on the result
+        return stepAsync(name, resultType, func, config).get();
+    }
+
+    public <T> T step(String name, TypeToken<T> typeToken, Supplier<T> func) {
+        return step(name, typeToken, func, StepConfig.builder().build());
+    }
+
+    public <T> T step(String name, TypeToken<T> typeToken, Supplier<T> func, StepConfig config) {
+        // Simply delegate to stepAsync and block on the result
+        return stepAsync(name, typeToken, func, config).get();
+    }
+
+    public <T> DurableFuture<T> stepAsync(String name, Class<T> resultType, Supplier<T> func) {
+        return stepAsync(
+                name, TypeToken.get(resultType), func, StepConfig.builder().build());
+    }
+
+    public <T> DurableFuture<T> stepAsync(String name, Class<T> resultType, Supplier<T> func, StepConfig config) {
+        return stepAsync(name, TypeToken.get(resultType), func, config);
+    }
+
+    public <T> DurableFuture<T> stepAsync(String name, TypeToken<T> typeToken, Supplier<T> func) {
+        return stepAsync(name, typeToken, func, StepConfig.builder().build());
+    }
+
+    public <T> DurableFuture<T> stepAsync(String name, TypeToken<T> typeToken, Supplier<T> func, StepConfig config) {
+        return stepAsync(name, typeToken, stepContext -> func.get(), config);
+    }
+
     // ========== wait methods ==========
 
     public Void wait(Duration duration) {
@@ -147,7 +176,7 @@ public Void wait(String waitName, Duration duration) {
         var operationId = nextOperationId();
 
         // Create and start wait operation
-        var operation = new WaitOperation(operationId, waitName, duration, executionManager, contextId);
+        var operation = new WaitOperation(operationId, waitName, duration, this);
 
         operation.execute(); // Checkpoint the wait
         return operation.get(); // Block (will throw SuspendExecutionException if needed)
@@ -208,16 +237,17 @@ public <T, U> DurableFuture<T> invokeAsync(
         Objects.requireNonNull(config, "config cannot be null");
         Objects.requireNonNull(typeToken, "typeToken cannot be null");
         if (config.serDes() == null) {
-            config = config.toBuilder().serDes(durableConfig.getSerDes()).build();
+            config = config.toBuilder().serDes(getDurableConfig().getSerDes()).build();
         }
         if (config.payloadSerDes() == null) {
-            config = config.toBuilder().payloadSerDes(durableConfig.getSerDes()).build();
+            config = config.toBuilder()
+                    .payloadSerDes(getDurableConfig().getSerDes())
+                    .build();
         }
         var operationId = nextOperationId();
 
         // Create and start invoke operation
-        var operation = new InvokeOperation<>(
-                operationId, name, functionName, payload, typeToken, config, executionManager, contextId);
+        var operation = new InvokeOperation<>(operationId, name, functionName, payload, typeToken, config, this);
 
         operation.execute(); // checkpoint the invoke operation
         return operation; // Block (will throw SuspendExecutionException if needed)
@@ -239,11 +269,11 @@ public <T> DurableCallbackFuture<T> createCallback(String name, Class<T> resultT
 
     public <T> DurableCallbackFuture<T> createCallback(String name, TypeToken<T> typeToken, CallbackConfig config) {
         if (config.serDes() == null) {
-            config = config.toBuilder().serDes(durableConfig.getSerDes()).build();
+            config = config.toBuilder().serDes(getDurableConfig().getSerDes()).build();
         }
         var operationId = nextOperationId();
 
-        var operation = new CallbackOperation<>(operationId, name, typeToken, config, executionManager, contextId);
+        var operation = new CallbackOperation<>(operationId, name, typeToken, config, this);
         operation.execute();
 
         return operation;
@@ -270,69 +300,25 @@ public <T> DurableFuture<T> runInChildContextAsync(
         var operationId = nextOperationId();
 
         var operation = new ChildContextOperation<>(
-                operationId,
-                name,
-                func,
-                typeToken,
-                durableConfig.getSerDes(),
-                executionManager,
-                durableConfig,
-                lambdaContext,
-                contextId);
+                operationId, name, func, typeToken, getDurableConfig().getSerDes(), this);
 
         operation.execute();
         return operation;
     }
 
     // =============== accessors ================
-
-    public Context getLambdaContext() {
-        return lambdaContext;
-    }
-
     public DurableLogger getLogger() {
         return logger;
     }
 
-    /**
-     * Returns metadata about the current durable execution.
-     *
-     * <p>The execution context provides information that remains constant throughout the execution lifecycle, such as
-     * the durable execution ARN. This is useful for tracking execution progress, correlating logs, and referencing this
-     * execution in external systems.
-     *
-     * @return the execution context
-     */
-    public ExecutionContext getExecutionContext() {
-        return executionContext;
-    }
-
-    // ============= internal utilities ===============
-
-    /** Gets the context ID for this context. Null for root context, set for child contexts. */
-    String getContextId() {
-        return contextId;
-    }
-
-    /** Returns whether this context is currently in replay mode. */
-    boolean isReplaying() {
-        return isReplaying;
-    }
-
-    /**
-     * Transitions this context from replay to execution mode. Called when the first un-cached operation is encountered.
-     */
-    void setExecutionMode() {
-        this.isReplaying = false;
-    }
-
     /**
      * Get the next operationId. For root contexts, returns sequential IDs like "1", "2", "3". For child contexts,
-     * prefixes with the contextId to ensure global uniqueness, e.g. "1-1", "1-2" for operations inside child context
-     * "1". This matches the JavaScript SDK's stepPrefix convention and prevents ID collisions in checkpoint batches.
+     * prefixes with the getContextId() to ensure global uniqueness, e.g. "1-1", "1-2" for operations inside child
+     * context "1". This matches the JavaScript SDK's stepPrefix convention and prevents ID collisions in checkpoint
+     * batches.
      */
     private String nextOperationId() {
         var counter = String.valueOf(operationCounter.incrementAndGet());
-        return contextId != null ? contextId + "-" + counter : counter;
+        return getParentContextId() != null ? getParentContextId() + "-" + counter : counter;
     }
 }