feat(map): Fix map replay, wait-inside-map, and concurrency race conditions (#229)

* feat(map): Fix map replay, wait-inside-map, and concurrency race conditions

Author: ayushiahjolia

AGENTS.md

@@ -208,7 +208,7 @@ void testAgainstRealLambda() {
 | `StepOperation` | Executes steps with retry logic |
 | `WaitOperation` | Handles wait checkpointing |
 | `MapOperation` | Applies a function across items concurrently via child contexts |
-| `BaseConcurrentOperation` | Shared base for map/parallel: concurrency limiting, completion evaluation |
+| `ConcurrencyOperation` | Shared base for map/parallel: concurrency limiting, completion evaluation |
 
 ## Common Tasks
 

docs/core/map.md

@@ -41,14 +41,14 @@ MapResult<OrderResult> result = future.get();
 | Method | Description |
 |--------|-------------|
 | `getResult(i)` | Result at index `i`, or `null` if that item failed |
-| `getError(i)` | Error at index `i`, or `null` if that item succeeded |
+| `getError(i)` | `ErrorObject` at index `i`, or `null` if that item succeeded |
 | `getItem(i)` | The `MapResultItem` at index `i` with status, result, and error |
 | `allSucceeded()` | `true` if every item succeeded |
 | `size()` | Number of items in the result |
 | `items()` | All result items as an unmodifiable list |
 | `results()` | All results as an unmodifiable list (nulls for failed items) |
 | `succeeded()` | Only the non-null (successful) results |
-| `failed()` | Only the non-null errors |
+| `failed()` | Only the non-null `ErrorObject`s |
 | `completionReason()` | Why the operation completed (`ALL_COMPLETED`, `MIN_SUCCESSFUL_REACHED`, `FAILURE_TOLERANCE_EXCEEDED`) |
 
 ### MapResultItem
@@ -57,9 +57,9 @@ Each `MapResultItem<T>` contains:
 
 | Field | Description |
 |-------|-------------|
-| `status()` | `SUCCEEDED`, `FAILED`, or `null` (not started) |
+| `status()` | `SUCCEEDED`, `FAILED`, or `NOT_STARTED` |
 | `result()` | The result value, or `null` if failed/not started |
-| `error()` | The error, or `null` if succeeded/not started |
+| `error()` | The error details as `ErrorObject`, or `null` if succeeded/not started |
 
 ### Error Isolation
 
@@ -126,7 +126,7 @@ var result = ctx.map("find-two", items, String.class, fn, config);
 assertEquals(CompletionReason.MIN_SUCCESSFUL_REACHED, result.completionReason());
 ```
 
-When early termination triggers, items that were never started have `null` for both result and error in the `MapResult`.
+When early termination triggers, items that were never started have `NOT_STARTED` status with `null` for both result and error in the `MapResult`.
 
 ### Checkpoint-and-Replay
 

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

@@ -0,0 +1,63 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.lambda.durable.examples;
+
+import java.time.Duration;
+import java.util.List;
+import java.util.stream.Collectors;
+import software.amazon.lambda.durable.CompletionConfig;
+import software.amazon.lambda.durable.DurableContext;
+import software.amazon.lambda.durable.DurableHandler;
+import software.amazon.lambda.durable.MapConfig;
+
+/**
+ * Example demonstrating advanced map features: wait operations inside branches, error handling, and early termination.
+ *
+ * <ol>
+ *   <li>Concurrent map with step + wait + step inside each branch — simulates multi-stage order processing with a
+ *       cooldown between stages
+ *   <li>Early termination with {@code minSuccessful(2)} — finds 2 healthy servers then stops
+ * </ol>
+ */
+public class ComplexMapExample extends DurableHandler<GreetingRequest, String> {
+
+    @Override
+    public String handleRequest(GreetingRequest input, DurableContext context) {
+        var name = input.getName();
+        context.getLogger().info("Starting complex map example for {}", name);
+
+        // Part 1: Concurrent map with step + wait inside each branch
+        var orderIds = List.of("order-1", "order-2", "order-3");
+
+        var orderResult = context.map("process-orders", orderIds, String.class, (orderId, index, ctx) -> {
+            // Step 1: validate the order
+            var validated = ctx.step("validate-" + index, String.class, stepCtx -> "validated:" + orderId + ":" + name);
+
+            // Wait between stages (simulates a cooldown or external dependency)
+            ctx.wait("cooldown-" + index, Duration.ofSeconds(1));
+
+            // Step 2: finalize the order
+            return ctx.step("finalize-" + index, String.class, stepCtx -> "done:" + validated);
+        });
+
+        var orderSummary = String.join(", ", orderResult.results());
+
+        // Part 2: Early termination — find 2 healthy servers then stop
+        var servers = List.of("server-1", "server-2", "server-3", "server-4", "server-5");
+        var earlyTermConfig = MapConfig.builder()
+                .completionConfig(CompletionConfig.minSuccessful(2))
+                .build();
+
+        var serverResult = context.map(
+                "find-healthy-servers",
+                servers,
+                String.class,
+                (server, index, ctx) -> ctx.step("health-check-" + index, String.class, stepCtx -> server + ":healthy"),
+                earlyTermConfig);
+
+        var healthyServers = serverResult.succeeded().stream().collect(Collectors.joining(", "));
+
+        return String.format(
+                "orders=[%s] | servers=[%s] reason=%s", orderSummary, healthyServers, serverResult.completionReason());
+    }
+}

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

@@ -571,39 +571,21 @@ void testSimpleMapExample() {
     }
 
     @Test
-    void testMapErrorHandlingExample() {
-        var runner = CloudDurableTestRunner.create(
-                arn("map-error-handling-example"), GreetingRequest.class, String.class, lambdaClient);
-        var result = runner.run(new GreetingRequest("Alice"));
-
-        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
-        var output = result.getResult(String.class);
-        assertNotNull(output);
-
-        // 3 of 5 orders succeed, 2 fail
-        assertTrue(output.contains("succeeded=3"));
-        assertTrue(output.contains("failed=2"));
-        assertTrue(output.contains("Processed order-1 for Alice"));
-        assertTrue(output.contains("Processed order-3 for Alice"));
-        assertTrue(output.contains("Processed order-5 for Alice"));
-    }
-
-    @Test
-    void testMapConfigExample() {
-        var runner = CloudDurableTestRunner.create(
-                arn("map-config-example"), GreetingRequest.class, String.class, lambdaClient);
+    void testComplexMapExample() {
+        var runner = CloudDurableTestRunner.create(arn("complex-map-example"), GreetingRequest.class, String.class);
         var result = runner.run(new GreetingRequest("Alice"));
 
         assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
         var output = result.getResult(String.class);
         assertNotNull(output);
 
-        // Sequential part: all 3 items processed
-        assertTrue(output.contains("ALPHA-Alice"));
-        assertTrue(output.contains("BETA-Alice"));
-        assertTrue(output.contains("GAMMA-Alice"));
+        // Part 1: Concurrent order processing with step + wait + step
+        assertTrue(output.contains("done:validated:order-1:Alice"));
+        assertTrue(output.contains("done:validated:order-2:Alice"));
+        assertTrue(output.contains("done:validated:order-3:Alice"));
 
-        // Early termination part
+        // Part 2: Early termination — find 2 healthy servers then stop
+        assertTrue(output.contains("healthy"));
         assertTrue(output.contains("reason=MIN_SUCCESSFUL_REACHED"));
     }
 }

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

@@ -0,0 +1,54 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.lambda.durable.examples;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+import org.junit.jupiter.api.Test;
+import software.amazon.lambda.durable.model.ExecutionStatus;
+import software.amazon.lambda.durable.testing.LocalDurableTestRunner;
+
+class ComplexMapExampleTest {
+
+    @Test
+    void testComplexMapExample() {
+        var handler = new ComplexMapExample();
+        var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler);
+
+        var result = runner.runUntilComplete(new GreetingRequest("Alice"));
+
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+        var output = result.getResult(String.class);
+
+        // Part 1: all 3 orders processed with step + wait + step
+        assertTrue(output.contains("done:validated:order-1:Alice"));
+        assertTrue(output.contains("done:validated:order-2:Alice"));
+        assertTrue(output.contains("done:validated:order-3:Alice"));
+
+        // Part 2: early termination after 2 healthy servers
+        assertTrue(output.contains("reason=MIN_SUCCESSFUL_REACHED"));
+        assertTrue(output.contains("healthy"));
+    }
+
+    @Test
+    void testReplay() {
+        var handler = new ComplexMapExample();
+        var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler);
+
+        var input = new GreetingRequest("Bob");
+        var result1 = runner.runUntilComplete(input);
+        assertEquals(ExecutionStatus.SUCCEEDED, result1.getStatus());
+
+        // Replay — should use cached results.
+        // Structural assertion because the first map has wait() inside branches with unlimited
+        // concurrency, which can cause non-deterministic thread scheduling across invocations.
+        var result2 = runner.runUntilComplete(input);
+        assertEquals(ExecutionStatus.SUCCEEDED, result2.getStatus());
+        var output = result2.getResult(String.class);
+        assertTrue(output.contains("done:validated:order-1:Bob"));
+        assertTrue(output.contains("done:validated:order-2:Bob"));
+        assertTrue(output.contains("done:validated:order-3:Bob"));
+        assertTrue(output.contains("reason=MIN_SUCCESSFUL_REACHED"));
+        assertTrue(output.contains("healthy"));
+    }
+}