[feat]: support FLAT nesting type for map/parallel (#335)

* [feat]: support FLAT nesting type for map/parallel

Author: zhongkechen

examples/src/main/java/software/amazon/lambda/durable/examples/map/ComplexFlatMapExample.java

@@ -0,0 +1,70 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.lambda.durable.examples.map;
+
+import java.time.Duration;
+import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+import software.amazon.lambda.durable.DurableContext;
+import software.amazon.lambda.durable.DurableHandler;
+import software.amazon.lambda.durable.config.CompletionConfig;
+import software.amazon.lambda.durable.config.MapConfig;
+import software.amazon.lambda.durable.config.NestingType;
+
+/**
+ * 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 ComplexFlatMapExample extends DurableHandler<Integer, String> {
+
+    @Override
+    public String handleRequest(Integer input, DurableContext context) {
+        context.getLogger().info("Starting complex map example with {} items", input);
+
+        // Part 1: Concurrent map with step + wait inside each branch
+        var orderIds = IntStream.range(1, input + 1).mapToObj(x -> "order-" + x).collect(Collectors.toList());
+
+        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);
+
+                    // 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);
+                },
+                MapConfig.builder().nestingType(NestingType.FLAT).build());
+
+        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))
+                .nestingType(NestingType.FLAT)
+                .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/test/java/software/amazon/lambda/durable/examples/CloudBasedIntegrationTest.java

@@ -674,6 +674,25 @@ void testComplexMapExample() {
         assertTrue(output.contains("reason=MIN_SUCCESSFUL_REACHED"));
     }
 
+    @Test
+    void testComplexFlatMapExample() {
+        var runner = CloudDurableTestRunner.create(arn("complex-flat-map-example"), Integer.class, String.class);
+        var result = runner.run(100);
+
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+        var output = result.getResult();
+        assertNotNull(output);
+
+        // Part 1: Concurrent order processing with step + wait + step
+        assertTrue(output.contains("done:validated:order-1"));
+        assertTrue(output.contains("done:validated:order-2"));
+        assertTrue(output.contains("done:validated:order-100"));
+
+        // Part 2: Early termination — find 2 healthy servers then stop
+        assertTrue(output.contains("healthy"));
+        assertTrue(output.contains("reason=MIN_SUCCESSFUL_REACHED"));
+    }
+
     @Test
     void testWaitForConditionExample() {
         var runner = CloudDurableTestRunner.create(

examples/src/test/java/software/amazon/lambda/durable/examples/map/ComplexFlatMapExampleTest.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.map;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import org.junit.jupiter.api.Test;
+import software.amazon.lambda.durable.model.ExecutionStatus;
+import software.amazon.lambda.durable.testing.LocalDurableTestRunner;
+
+class ComplexFlatMapExampleTest {
+
+    @Test
+    void testComplexMapExample() {
+        var handler = new ComplexFlatMapExample();
+        var runner = LocalDurableTestRunner.create(Integer.class, handler);
+
+        var result = runner.runUntilComplete(50);
+
+        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"));
+        assertTrue(output.contains("done:validated:order-2"));
+        assertTrue(output.contains("done:validated:order-50"));
+
+        // 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 ComplexFlatMapExample();
+        var runner = LocalDurableTestRunner.create(Integer.class, handler);
+
+        var result1 = runner.runUntilComplete(50);
+        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(50);
+        assertEquals(ExecutionStatus.SUCCEEDED, result2.getStatus());
+        var output = result2.getResult(String.class);
+        assertTrue(output.contains("done:validated:order-1"));
+        assertTrue(output.contains("done:validated:order-2"));
+        assertTrue(output.contains("done:validated:order-50"));
+        assertTrue(output.contains("reason=MIN_SUCCESSFUL_REACHED"));
+        assertTrue(output.contains("healthy"));
+    }
+}

examples/template.yaml

@@ -364,6 +364,23 @@ Resources:
                 - lambda:GetDurableExecutionState
               Resource: !Sub "arn:aws:lambda:${AWS::Region}:${AWS::AccountId}:function:complex-map-example-${JavaVersion}-runtime"
 
+  ComplexFlatMapExampleFunction:
+    Type: AWS::Serverless::Function
+    Properties:
+      FunctionName: !Join
+        - '-'
+        - - 'complex-flat-map-example'
+          - !Ref JavaVersion
+          - runtime
+      Handler: "software.amazon.lambda.durable.examples.map.ComplexFlatMapExample"
+      Policies:
+        - Statement:
+            - Effect: Allow
+              Action:
+                - lambda:CheckpointDurableExecutions
+                - lambda:GetDurableExecutionState
+              Resource: !Sub "arn:aws:lambda:${AWS::Region}:${AWS::AccountId}:function:complex-flat-map-example-${JavaVersion}-runtime"
+
   WaitForConditionExampleFunction:
     Type: AWS::Serverless::Function
     Properties:
@@ -497,6 +514,10 @@ Outputs:
     Description: Complex Map Example Function ARN
     Value: !GetAtt ComplexMapExampleFunction.Arn
 
+  ComplexFlatMapExampleFunction:
+    Description: Complex Flat Map Example Function ARN
+    Value: !GetAtt ComplexFlatMapExampleFunction.Arn
+
   WaitForConditionExampleFunction:
     Description: Wait For Condition Example Function ARN
     Value: !GetAtt WaitForConditionExampleFunction.Arn