Merge pull request #9 from aws/phipag/nested-steps

feat(nested-steps): Stop hardcoding Root context thread

Author: phipag

sdk-integration-tests/src/test/java/com/amazonaws/lambda/durable/NestedStepIntegrationTest.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);
+    }
+}

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

@@ -27,7 +27,12 @@ public class DurableContext {
     private final AtomicInteger operationCounter;
     private final DurableLogger logger;
 
-    DurableContext(ExecutionManager executionManager, SerDes serDes, Context lambdaContext, LoggerConfig loggerConfig) {
+    DurableContext(
+            ExecutionManager executionManager,
+            SerDes serDes,
+            Context lambdaContext,
+            LoggerConfig loggerConfig,
+            String contextId) {
         this.executionManager = executionManager;
         this.serDes = serDes;
         this.lambdaContext = lambdaContext;
@@ -41,9 +46,11 @@ public class DurableContext {
                 loggerConfig.suppressReplayLogs());
 
         // 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);
+        executionManager.registerActiveThreadWithContext(contextId, ThreadType.CONTEXT);
+    }
+
+    DurableContext(ExecutionManager executionManager, SerDes serDes, Context lambdaContext, LoggerConfig loggerConfig) {
+        this(executionManager, serDes, lambdaContext, loggerConfig, "Root");
     }
 
     public <T> T step(String name, Class<T> resultType, Supplier<T> func) {

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

@@ -70,11 +70,15 @@ 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, config.getLoggerConfig());
-
         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, config.getLoggerConfig());
+                        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

@@ -50,6 +50,7 @@ public class ExecutionManager {
 
     // ===== Thread Coordination =====
     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> suspendExecutionFuture = new CompletableFuture<>();
 
@@ -164,20 +165,51 @@ public Operation getExecutionOperation() {
 
     // ===== Thread Coordination =====
 
-    public void registerActiveThread(String threadId, ThreadType threadType) {
+    public void registerActiveThreadWithContext(String threadId, ThreadType threadType) {
         if (activeThreads.containsKey(threadId)) {
             logger.trace("Thread '{}' ({}) already registered as active", threadId, threadType);
             return;
         }
+        activeThreads.put(threadId, threadType);
+        currentContext.set(new OperationContext(threadId, threadType));
+        logger.trace(
+                "Registered thread '{}' ({}) as active. Active threads: {}",
+                threadId,
+                threadType,
+                activeThreads.size());
+    }
 
-        synchronized (this) {
-            activeThreads.put(threadId, threadType);
-            logger.trace(
-                    "Registered thread '{}' ({}) as active. Active threads: {}",
-                    threadId,
-                    threadType,
-                    activeThreads.size());
+    /**
+     * Registers a thread as active without setting the thread local OperationContext. Use this when registration must
+     * happen on a different thread than execution. Call setCurrentContext() on the execution thread to set the local
+     * OperationContext.
+     *
+     * @see OperationContext
+     */
+    public void registerActiveThread(String threadId, ThreadType threadType) {
+        if (activeThreads.containsKey(threadId)) {
+            logger.trace("Thread '{}' ({}) already registered as active", threadId, threadType);
+            return;
         }
+        activeThreads.put(threadId, threadType);
+        logger.trace(
+                "Registered thread '{}' ({}) as active (no context). Active threads: {}",
+                threadId,
+                threadType,
+                activeThreads.size());
+    }
+
+    /**
+     * Sets the current thread's context. Use after registerActiveThreadWithoutContext() when the execution thread is
+     * different from the registration thread.
+     */
+    public void setCurrentContext(String contextId, ThreadType threadType) {
+        currentContext.set(new OperationContext(contextId, threadType));
+    }
+
+    /** Returns the current thread's context, or null if not set. */
+    public OperationContext getCurrentContext() {
+        return currentContext.get();
     }
 
     public void deregisterActiveThread(String threadId) {
@@ -191,15 +223,14 @@ public void deregisterActiveThread(String threadId) {
             return;
         }
 
-        synchronized (this) {
-            ThreadType type = activeThreads.remove(threadId);
-            logger.trace("Deregistered thread '{}' ({}). Active threads: {}", threadId, type, activeThreads.size());
+        ThreadType type = activeThreads.remove(threadId);
+        currentContext.remove();
+        logger.trace("Deregistered thread '{}' ({}). Active threads: {}", threadId, type, activeThreads.size());
 
-            if (activeThreads.isEmpty()) {
-                logger.info("No active threads remaining - suspending execution");
-                suspendExecutionFuture.complete(null);
-                throw new SuspendExecutionException();
-            }
+        if (activeThreads.isEmpty()) {
+            logger.info("No active threads remaining - suspending execution");
+            suspendExecutionFuture.complete(null);
+            throw new SuspendExecutionException();
         }
     }
 

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

@@ -0,0 +1,6 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package com.amazonaws.lambda.durable.execution;
+
+/** Holds the current thread's execution context. */
+public record OperationContext(String contextId, ThreadType threadType) {}

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

@@ -163,12 +163,17 @@ public void execute() {
     }
 
     private void executeStepLogic(int attempt) {
-        // Register step thread as active
+        // TODO: Modify this logic when child contexts are introduced such that the child context id is in this key
         var stepThreadId = operationId + "-step";
+
+        // Register step thread as active BEFORE executor runs (prevents suspension when handler deregisters)
+        // thread local OperationContext is set inside the executor since that's where the step actually runs
         executionManager.registerActiveThread(stepThreadId, ThreadType.STEP);
 
         // Execute in managed executor
         executionManager.getManagedExecutor().execute(() -> {
+            // Set thread local OperationContext on the executor thread
+            executionManager.setCurrentContext(stepThreadId, ThreadType.STEP);
             // Set operation context for logging in this thread
             durableLogger.setOperationContext(operationId, name, attempt);
             try {
@@ -290,6 +295,15 @@ private void handleStepFailure(Throwable error, int attempt) {
 
     @Override
     public T get() {
+        // Get current context from ThreadLocal
+        var currentContext = executionManager.getCurrentContext();
+
+        // Nested steps are not supported
+        if (currentContext.threadType() == ThreadType.STEP) {
+            throw new IllegalStateException("Nested step calling is not supported. Cannot call get() on step '" + name
+                    + "' from within another step's execution.");
+        }
+
         // 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
@@ -298,17 +312,16 @@ public T get() {
             // Operation not done yet
             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");
+            // Deregister current context - allows suspension
+            logger.debug("StepOperation.get() attempting to deregister context: {}", currentContext.contextId());
+            executionManager.deregisterActiveThread(currentContext.contextId());
 
             // Block until operation completes
             logger.trace("Waiting for operation to finish {} (Phaser: {})", operationId, phaser);
             phaser.arriveAndAwaitAdvance(); // Wait for phase 0
 
-            // Reactivate current thread
-            executionManager.registerActiveThread("Root", ThreadType.CONTEXT);
+            // Reactivate current context
+            executionManager.registerActiveThreadWithContext(currentContext.contextId(), currentContext.threadType());
 
             // Complete phase 1
             phaser.arriveAndDeregister();

sdk/src/main/java/com/amazonaws/lambda/durable/operation/WaitOperation.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;
@@ -97,16 +96,19 @@ public Void get() {
         if (phaser.getPhase() == ExecutionPhase.RUNNING.getValue()) {
             phaser.register();
 
+            // Get current context from ThreadLocal
+            var currentContext = executionManager.getCurrentContext();
+
             // Deregister current thread - THIS is where suspension can happen!
             // If no other threads are active, this will throw SuspendExecutionException
-            executionManager.deregisterActiveThread("Root");
+            executionManager.deregisterActiveThread(currentContext.contextId());
 
             // 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.registerActiveThreadWithContext(currentContext.contextId(), currentContext.threadType());
 
             // Complete phase 1
             phaser.arriveAndDeregister();

sdk/src/test/java/com/amazonaws/lambda/durable/StepConfigTest.java

@@ -8,7 +8,6 @@
 
 import com.amazonaws.lambda.durable.retry.RetryStrategies;
 import com.amazonaws.lambda.durable.serde.JacksonSerDes;
-import com.amazonaws.lambda.durable.serde.SerDes;
 import org.junit.jupiter.api.Test;
 
 class StepConfigTest {
@@ -32,7 +31,7 @@ void testBuilderWithoutRetryStrategy() {
     @Test
     void testBuilderChaining() {
         var strategy = RetryStrategies.Presets.NO_RETRY;
-        SerDes customSerDes = new JacksonSerDes();
+        var customSerDes = new JacksonSerDes();
 
         var config = StepConfig.builder()
                 .retryStrategy(strategy)
@@ -61,7 +60,7 @@ void testSemanticsDefaultsToAtLeastOnce() {
 
     @Test
     void testBuilderWithCustomSerDes() {
-        SerDes customSerDes = new JacksonSerDes();
+        var customSerDes = new JacksonSerDes();
 
         var config = StepConfig.builder().serDes(customSerDes).build();
 
@@ -86,7 +85,7 @@ void testBuilderWithNullSerDes() {
     @Test
     void testBuilderWithAllOptions() {
         var strategy = RetryStrategies.Presets.DEFAULT;
-        SerDes customSerDes = new JacksonSerDes();
+        var customSerDes = new JacksonSerDes();
 
         var config = StepConfig.builder()
                 .retryStrategy(strategy)

sdk/src/test/java/com/amazonaws/lambda/durable/operation/StepOperationTest.java

@@ -0,0 +1,70 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package com.amazonaws.lambda.durable.operation;
+
+import static org.junit.jupiter.api.Assertions.*;
+import static org.mockito.Mockito.*;
+
+import com.amazonaws.lambda.durable.execution.ExecutionManager;
+import com.amazonaws.lambda.durable.execution.OperationContext;
+import com.amazonaws.lambda.durable.execution.ThreadType;
+import com.amazonaws.lambda.durable.logging.DurableLogger;
+import com.amazonaws.lambda.durable.serde.JacksonSerDes;
+import java.util.concurrent.Phaser;
+import org.junit.jupiter.api.Test;
+
+class StepOperationTest {
+
+    @Test
+    void getThrowsIllegalStateExceptionWhenCalledFromStepContext() {
+        var executionManager = mock(ExecutionManager.class);
+        var phaser = new Phaser(1);
+        when(executionManager.startPhaser(any())).thenReturn(phaser);
+        when(executionManager.getCurrentContext()).thenReturn(new OperationContext("1-step", ThreadType.STEP));
+
+        var operation = new StepOperation<>(
+                "1",
+                "test-step",
+                () -> "result",
+                String.class,
+                null,
+                executionManager,
+                mock(DurableLogger.class),
+                new JacksonSerDes());
+
+        var ex = assertThrows(IllegalStateException.class, operation::get);
+        assertTrue(ex.getMessage().contains("Nested step calling is not supported"));
+        assertTrue(ex.getMessage().contains("test-step"));
+    }
+
+    @Test
+    void getDoesNotThrowWhenCalledFromHandlerContext() {
+        var executionManager = mock(ExecutionManager.class);
+        var phaser = new Phaser(1);
+        phaser.arriveAndDeregister(); // Advance to phase 1 to skip blocking
+        when(executionManager.startPhaser(any())).thenReturn(phaser);
+        when(executionManager.getCurrentContext()).thenReturn(new OperationContext("handler", ThreadType.CONTEXT));
+        when(executionManager.getOperation("1"))
+                .thenReturn(software.amazon.awssdk.services.lambda.model.Operation.builder()
+                        .id("1")
+                        .name("test-step")
+                        .status(software.amazon.awssdk.services.lambda.model.OperationStatus.SUCCEEDED)
+                        .stepDetails(software.amazon.awssdk.services.lambda.model.StepDetails.builder()
+                                .result("\"cached-result\"")
+                                .build())
+                        .build());
+
+        var operation = new StepOperation<>(
+                "1",
+                "test-step",
+                () -> "result",
+                String.class,
+                null,
+                executionManager,
+                mock(DurableLogger.class),
+                new JacksonSerDes());
+
+        var result = operation.get();
+        assertEquals("cached-result", result);
+    }
+}