Raise exception when called from nested steps and delete example.

Author: phipag

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

@@ -186,27 +186,6 @@ 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:
@@ -273,10 +252,3 @@ Outputs:
     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

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

@@ -0,0 +1,55 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package com.amazonaws.lambda.durable;
+
+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;
+
+/** Tests that nested step calling is properly rejected. */
+class NestedStepIntegrationTest {
+
+    @Test
+    void nestedStepCallingThrowsIllegalStateException() {
+        var runner = LocalDurableTestRunner.create(String.class, (input, context) -> {
+            // outer-step's supplier calls context.step() which internally calls stepAsync().get()
+            // The get() is called from the outer step's thread (named "1-step"), triggering the check
+            var future = context.stepAsync("outer-step", String.class, () -> {
+                return context.step("inner-step", String.class, () -> "inner-result");
+            });
+            return future.get();
+        });
+
+        var result = runner.run("test");
+
+        assertEquals(ExecutionStatus.FAILED, result.getStatus());
+        var errorMessage = result.getError().get().errorMessage();
+        assertTrue(
+                errorMessage.contains("Nested step calling is not supported"),
+                "Expected error about nested step calling, got: " + errorMessage);
+    }
+
+    @Test
+    void awaitingAsyncStepInsideSyncStepThrowsIllegalStateException() {
+        var runner = LocalDurableTestRunner.create(String.class, (input, context) -> {
+            // Start async step from handler thread
+            var asyncFuture = context.stepAsync("async-step", String.class, () -> "async-result");
+
+            // Sync step tries to await the async step's result inside its supplier
+            return context.step("sync-step", String.class, () -> {
+                // This get() is called from sync-step's thread ("2-step"), which is not allowed
+                return "combined: " + asyncFuture.get();
+            });
+        });
+
+        var result = runner.run("test");
+
+        assertEquals(ExecutionStatus.FAILED, result.getStatus());
+        var errorMessage = result.getError().get().errorMessage();
+        assertTrue(
+                errorMessage.contains("Nested step calling is not supported"),
+                "Expected error about nested step calling, got: " + errorMessage);
+    }
+}

examples/template.yaml

@@ -4,6 +4,7 @@
 
 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;
@@ -17,21 +18,22 @@
 import software.amazon.awssdk.services.lambda.model.OperationType;
 
 public class DurableContext {
+    private static final String HANDLER_CONTEXT_ID = "handler";
+
     private final ExecutionManager executionManager;
     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 context thread as active
-        executionManager.registerActiveThread(uniqueThreadName);
+        // Register handler context as active
+        executionManager.registerActiveThread(HANDLER_CONTEXT_ID, ThreadType.CONTEXT);
+        executionManager.enterContext(HANDLER_CONTEXT_ID, ThreadType.CONTEXT);
     }
 
     public <T> T step(String name, Class<T> resultType, Supplier<T> func) {

sdk-integration-tests/src/test/java/com/amazonaws/lambda/durable/NestedStepIntegrationTest.java

@@ -8,10 +8,8 @@
 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;
@@ -49,7 +47,9 @@ public class ExecutionManager {
     private final String durableExecutionArn;
 
     // ===== Thread Coordination =====
-    private final Set<String> activeThreads = Collections.synchronizedSet(new HashSet<>());
+    private final Map<String, ThreadType> activeThreads = Collections.synchronizedMap(new HashMap<>());
+    private static final ThreadLocal<String> currentContextId = new ThreadLocal<>();
+    private static final ThreadLocal<ThreadType> currentThreadType = new ThreadLocal<>();
     private final Map<String, Phaser> openPhasers = Collections.synchronizedMap(new HashMap<>());
     private final CompletableFuture<Void> suspendExecutionFuture = new CompletableFuture<>();
 
@@ -139,35 +139,70 @@ public Operation getExecutionOperation() {
 
     // ===== Thread Coordination =====
 
-    public void registerActiveThread(String threadId) {
-        if (activeThreads.contains(threadId)) {
-            logger.debug("Thread '{}' already registered as active", threadId);
+    public void registerActiveThread(String threadId, ThreadType threadType) {
+        if (activeThreads.containsKey(threadId)) {
+            logger.debug("Thread '{}' ({}) already registered as active", threadId, threadType);
             return;
         }
 
         synchronized (this) {
-            activeThreads.add(threadId);
-            logger.debug("Registered thread '{}' as active. Active threads: {}", threadId, activeThreads.size());
+            activeThreads.put(threadId, threadType);
+            logger.debug(
+                    "Registered thread '{}' ({}) as active. Active threads: {}",
+                    threadId,
+                    threadType,
+                    activeThreads.size());
         }
     }
 
+    /**
+     * Sets the current thread's context. Call this when entering a context (handler or step).
+     * This is separate from registerActiveThread to support cases where registration happens
+     * on a different thread than execution.
+     */
+    public void enterContext(String contextId, ThreadType threadType) {
+        currentContextId.set(contextId);
+        currentThreadType.set(threadType);
+    }
+
+    /**
+     * Clears the current thread's context. Call this when exiting a context.
+     */
+    public void exitContext() {
+        currentContextId.remove();
+        currentThreadType.remove();
+    }
+
+    /**
+     * Returns the ThreadType for the current thread, or null if not registered.
+     * This uses ThreadLocal to track context independent of thread naming.
+     */
+    public ThreadType getCurrentThreadType() {
+        return currentThreadType.get();
+    }
+
+    /**
+     * Returns the context ID for the current thread, or null if not registered.
+     * This uses ThreadLocal to track context independent of thread naming.
+     */
+    public String getCurrentContextId() {
+        return currentContextId.get();
+    }
+
     public void deregisterActiveThread(String threadId) {
         // Skip if already suspended
         if (suspendExecutionFuture.isDone()) {
             return;
         }
 
-        if (!activeThreads.contains(threadId)) {
-            logger.warn(
-                    "Thread '{}' not active, cannot deregister. Current thread: {}",
-                    threadId,
-                    Thread.currentThread().getName());
+        if (!activeThreads.containsKey(threadId)) {
+            logger.warn("Thread '{}' not active, cannot deregister", threadId);
             return;
         }
 
         synchronized (this) {
-            activeThreads.remove(threadId);
-            logger.debug("Deregistered thread '{}'. Active threads: {}", threadId, activeThreads.size());
+            ThreadType type = activeThreads.remove(threadId);
+            logger.debug("Deregistered thread '{}' ({}). Active threads: {}", threadId, type, activeThreads.size());
 
             if (activeThreads.isEmpty()) {
                 logger.info("No active threads remaining - suspending execution");

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

@@ -0,0 +1,24 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package com.amazonaws.lambda.durable.execution;
+
+/**
+ * Thread type enum for tracking conceptual threads in durable execution.
+ *
+ * <p>These are not physical OS threads, but logical threads representing different types of work in the execution.
+ */
+public enum ThreadType {
+    CONTEXT("Context"),
+    STEP("Step");
+
+    private final String displayName;
+
+    ThreadType(String displayName) {
+        this.displayName = displayName;
+    }
+
+    @Override
+    public String toString() {
+        return displayName;
+    }
+}