Synchronize step ids with thread names and de-register / register caller threads in .get().

Author: phipag

examples/src/main/java/com/amazonaws/lambda/durable/examples/NestedStepExample.java

@@ -0,0 +1,50 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package com.amazonaws.lambda.durable.examples;
+
+import com.amazonaws.lambda.durable.DurableContext;
+import com.amazonaws.lambda.durable.DurableHandler;
+import com.amazonaws.lambda.durable.StepConfig;
+import com.amazonaws.lambda.durable.retry.RetryStrategies;
+
+/**
+ * Example demonstrating nested step calling with stepAsync.
+ *
+ * <p>This example shows how to:
+ *
+ * <ol>
+ *   <li>Create an async step that performs long-running work
+ *   <li>Use the result of that async step within another step
+ *   <li>Properly coordinate execution between multiple steps
+ * </ol>
+ */
+public class NestedStepExample extends DurableHandler<Object, String> {
+
+    @Override
+    public String handleRequest(Object input, DurableContext context) {
+        // Step 1: Create an async step that performs long-running work
+        var durableFuture1 = context.stepAsync(
+                "async-step",
+                String.class,
+                () -> {
+                    try {
+                        Thread.sleep(10000); // Simulate 10 seconds of work
+                    } catch (InterruptedException e) {
+                        Thread.currentThread().interrupt();
+                        throw new RuntimeException("Interrupted", e);
+                    }
+                    return "async-result";
+                },
+                StepConfig.builder()
+                        .retryStrategy(RetryStrategies.Presets.DEFAULT)
+                        .build());
+
+        // Step 2: Process the result of the async step
+        var result = context.step("process-result", String.class, () -> {
+            var asyncResult = durableFuture1.get(); // Wait for async step to complete
+            return asyncResult + "-processed";
+        });
+
+        return result;
+    }
+}

examples/src/test/java/com/amazonaws/lambda/durable/examples/NestedStepExampleTest.java

@@ -0,0 +1,42 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package com.amazonaws.lambda.durable.examples;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+import com.amazonaws.lambda.durable.model.ExecutionStatus;
+import com.amazonaws.lambda.durable.testing.LocalDurableTestRunner;
+import org.junit.jupiter.api.Test;
+
+class NestedStepExampleTest {
+
+    @Test
+    void testNestedStepExample() {
+        // Test nested step calling with stepAsync
+        var handler = new NestedStepExample();
+        var runner = LocalDurableTestRunner.create(Object.class, handler);
+
+        var result = runner.run("test-input");
+
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+        assertEquals("async-result-processed", result.getResult(String.class));
+    }
+
+    @Test
+    void testReplay() {
+        // Test replay behavior with nested steps
+        var handler = new NestedStepExample();
+        var runner = LocalDurableTestRunner.create(Object.class, handler);
+
+        var input = "replay-test";
+
+        // First execution
+        var result1 = runner.run(input);
+
+        // Second execution (replay)
+        var result2 = runner.run(input);
+
+        assertEquals(result1.getStatus(), result2.getStatus());
+        assertEquals(result1.getResult(String.class), result2.getResult(String.class));
+    }
+}

examples/template.yaml

@@ -186,6 +186,28 @@ Resources:
       DockerContext: ../
       DockerTag: durable-examples
 
+  NestedStepExampleFunction:
+    Type: AWS::Serverless::Function
+    Properties:
+      PackageType: Image
+      FunctionName: nested-step-example
+      ImageConfig:
+        Command: ["com.amazonaws.lambda.durable.examples.NestedStepExample::handleRequest"]
+      DurableConfig:
+        ExecutionTimeout: 300
+        RetentionPeriodInDays: 7
+      Policies:
+        - Statement:
+            - Effect: Allow
+              Action:
+                - lambda:CheckpointDurableExecutions
+                - lambda:GetDurableExecutionState
+              Resource: !Sub "arn:aws:lambda:${AWS::Region}:${AWS::AccountId}:function:nested-step-example"
+    Metadata:
+      Dockerfile: examples/Dockerfile
+      DockerContext: ../
+      DockerTag: durable-examples
+
 Outputs:
   SimpleStepExampleFunction:
     Description: Simple Step Example Function ARN
@@ -250,3 +272,11 @@ Outputs:
   CustomConfigExampleFunctionName:
     Description: Custom Config Example Function Name
     Value: !Ref CustomConfigExampleFunction
+
+  NestedStepExampleFunction:
+    Description: Nested Step Example Function ARN
+    Value: !GetAtt NestedStepExampleFunction.Arn
+
+  NestedStepExampleFunctionName:
+    Description: Nested Step Example Function Name
+    Value: !Ref NestedStepExampleFunction

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

@@ -183,10 +183,11 @@ public static final class Builder {
         private Builder() {}
 
         /**
-         * Sets a custom LambdaClient for production use. Use this method to customize the AWS SDK client
-         * with specific regions, credentials, timeouts, or retry policies.
+         * Sets a custom LambdaClient for production use. Use this method to customize the AWS SDK client with specific
+         * regions, credentials, timeouts, or retry policies.
          *
          * <p>Example:
+         *
          * <pre>{@code
          * LambdaClient lambdaClient = LambdaClient.builder()
          *     .region(Region.US_WEST_2)
@@ -211,9 +212,9 @@ public Builder withLambdaClient(LambdaClient lambdaClient) {
         /**
          * Sets a custom DurableExecutionClient.
          *
-         * <p><b>Note:</b> This method is primarily intended for testing with mock clients
-         * (e.g., {@code LocalMemoryExecutionClient}). For production use with a custom AWS SDK client,
-         * prefer {@link #withLambdaClient(LambdaClient)}.
+         * <p><b>Note:</b> This method is primarily intended for testing with mock clients (e.g.,
+         * {@code LocalMemoryExecutionClient}). For production use with a custom AWS SDK client, prefer
+         * {@link #withLambdaClient(LambdaClient)}.
          *
          * @param durableExecutionClient Custom DurableExecutionClient instance
          * @return This builder

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

@@ -4,7 +4,6 @@
 
 import com.amazonaws.lambda.durable.exception.NonDeterministicExecutionException;
 import com.amazonaws.lambda.durable.execution.ExecutionManager;
-import com.amazonaws.lambda.durable.execution.ThreadType;
 import com.amazonaws.lambda.durable.operation.StepOperation;
 import com.amazonaws.lambda.durable.operation.WaitOperation;
 import com.amazonaws.lambda.durable.retry.RetryStrategies;
@@ -22,17 +21,17 @@ public class DurableContext {
     private final SerDes serDes;
     private final Context lambdaContext;
     private final AtomicInteger operationCounter;
+    private final String uniqueThreadName;
 
     DurableContext(ExecutionManager executionManager, SerDes serDes, Context lambdaContext) {
         this.executionManager = executionManager;
         this.serDes = serDes;
         this.lambdaContext = lambdaContext;
         this.operationCounter = new AtomicInteger(0);
+        this.uniqueThreadName = Thread.currentThread().getName(); // Auto-detect handler thread
 
-        // Register root context thread as active
-        // TODO: Once we implement child contexts, the threadId needs to be the ID of
-        // the child context
-        executionManager.registerActiveThread("Root", ThreadType.CONTEXT);
+        // Register context thread as active
+        executionManager.registerActiveThread(uniqueThreadName);
     }
 
     public <T> T step(String name, Class<T> resultType, Supplier<T> func) {

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

@@ -71,11 +71,14 @@ public static <I, O> DurableExecutionOutput execute(
         var serDes = config.getSerDes();
         var userInput = extractUserInput(executionOp, serDes, inputType);
 
-        // Create context
-        var context = new DurableContext(executionManager, serDes, lambdaContext);
-
         try {
-            var handlerFuture = CompletableFuture.supplyAsync(() -> handler.apply(userInput, context), executor);
+            var handlerFuture = CompletableFuture.supplyAsync(
+                    () -> {
+                        // Create context in the executor thread so it detects the correct thread name
+                        var context = new DurableContext(executionManager, serDes, lambdaContext);
+                        return handler.apply(userInput, context);
+                    },
+                    executor);
 
             // Get suspend future from ExecutionManager. If this future completes, it
             // indicates

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

@@ -8,8 +8,10 @@
 import java.time.Instant;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.Executor;
@@ -47,7 +49,7 @@ public class ExecutionManager {
     private final String durableExecutionArn;
 
     // ===== Thread Coordination =====
-    private final Map<String, ThreadType> activeThreads = Collections.synchronizedMap(new HashMap<>());
+    private final Set<String> activeThreads = Collections.synchronizedSet(new HashSet<>());
     private final Map<String, Phaser> openPhasers = Collections.synchronizedMap(new HashMap<>());
     private final CompletableFuture<Void> suspendExecutionFuture = new CompletableFuture<>();
 
@@ -137,19 +139,15 @@ public Operation getExecutionOperation() {
 
     // ===== Thread Coordination =====
 
-    public void registerActiveThread(String threadId, ThreadType threadType) {
-        if (activeThreads.containsKey(threadId)) {
-            logger.debug("Thread '{}' ({}) already registered as active", threadId, threadType);
+    public void registerActiveThread(String threadId) {
+        if (activeThreads.contains(threadId)) {
+            logger.debug("Thread '{}' already registered as active", threadId);
             return;
         }
 
         synchronized (this) {
-            activeThreads.put(threadId, threadType);
-            logger.debug(
-                    "Registered thread '{}' ({}) as active. Active threads: {}",
-                    threadId,
-                    threadType,
-                    activeThreads.size());
+            activeThreads.add(threadId);
+            logger.debug("Registered thread '{}' as active. Active threads: {}", threadId, activeThreads.size());
         }
     }
 
@@ -159,14 +157,17 @@ public void deregisterActiveThread(String threadId) {
             return;
         }
 
-        if (!activeThreads.containsKey(threadId)) {
-            logger.warn("Thread '{}' not active, cannot deregister", threadId);
+        if (!activeThreads.contains(threadId)) {
+            logger.warn(
+                    "Thread '{}' not active, cannot deregister. Current thread: {}",
+                    threadId,
+                    Thread.currentThread().getName());
             return;
         }
 
         synchronized (this) {
-            ThreadType type = activeThreads.remove(threadId);
-            logger.debug("Deregistered thread '{}' ({}). Active threads: {}", threadId, type, activeThreads.size());
+            activeThreads.remove(threadId);
+            logger.debug("Deregistered thread '{}'. Active threads: {}", threadId, activeThreads.size());
 
             if (activeThreads.isEmpty()) {
                 logger.info("No active threads remaining - suspending execution");

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

@@ -9,7 +9,6 @@
 import com.amazonaws.lambda.durable.exception.StepInterruptedException;
 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 java.time.Duration;
 import java.time.Instant;
@@ -159,10 +158,12 @@ public void execute() {
     private void executeStepLogic(int attempt) {
         // Register step thread as active
         var stepThreadId = operationId + "-step";
-        executionManager.registerActiveThread(stepThreadId, ThreadType.STEP);
+        executionManager.registerActiveThread(stepThreadId);
 
         // Execute in managed executor
         executionManager.getManagedExecutor().execute(() -> {
+            // Set thread name to match the step thread ID for consistent logging and identification
+            Thread.currentThread().setName(stepThreadId);
             try {
                 // Check if we need to send START
                 var existing = executionManager.getOperation(operationId);
@@ -290,16 +291,16 @@ public T get() {
             phaser.register();
 
             // Deregister current thread - allows suspension
-            // TODO: The threadId here should be the (potential childContext) thread id that
-            // is calling .get()
-            executionManager.deregisterActiveThread("Root");
+            String currentThreadName = Thread.currentThread().getName();
+            logger.debug("StepOperation.get() attempting to deregister thread: {}", currentThreadName);
+            executionManager.deregisterActiveThread(currentThreadName);
 
             // Block until operation completes
             logger.debug("Waiting for operation to finish {} (Phaser: {})", operationId, phaser);
             phaser.arriveAndAwaitAdvance(); // Wait for phase 0
 
             // Reactivate current thread
-            executionManager.registerActiveThread("Root", ThreadType.CONTEXT);
+            executionManager.registerActiveThread(currentThreadName);
 
             // Complete phase 1
             phaser.arriveAndDeregister();
@@ -346,8 +347,7 @@ public T get() {
             throw new StepFailedException(
                     String.format(
                             "Step failed with error of type %s. Message: %s",
-                            errorType,
-                            op.stepDetails().error().errorMessage()),
+                            errorType, op.stepDetails().error().errorMessage()),
                     null,
                     // Preserve original stack trace
                     StepFailedException.deserializeStackTrace(

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

@@ -4,7 +4,6 @@
 
 import com.amazonaws.lambda.durable.execution.ExecutionManager;
 import com.amazonaws.lambda.durable.execution.ExecutionPhase;
-import com.amazonaws.lambda.durable.execution.ThreadType;
 import java.time.Duration;
 import java.time.Instant;
 import java.util.concurrent.Phaser;
@@ -100,14 +99,15 @@ public Void get() {
 
             // Deregister current thread - THIS is where suspension can happen!
             // If no other threads are active, this will throw SuspendExecutionException
-            executionManager.deregisterActiveThread("Root");
+            String callingThreadName = Thread.currentThread().getName();
+            executionManager.deregisterActiveThread(callingThreadName);
 
             // Complete the wait phaser immediately (we don't actually wait in Lambda)
             // The backend handles the wait duration
             phaser.arriveAndAwaitAdvance(); // Phase 0 -> 1
 
             // Reactivate current thread
-            executionManager.registerActiveThread("Root", ThreadType.CONTEXT);
+            executionManager.registerActiveThread(callingThreadName);
 
             // Complete phase 1
             phaser.arriveAndDeregister();