add test cases for virtual thread pool

Author: zhongkechen

examples/pom.xml

@@ -79,6 +79,32 @@
         </dependency>
     </dependencies>
 
+    <profiles>
+        <profile>
+            <id>exclude-java21-code-on-lower-jdk</id>
+            <activation>
+                <!-- Activates if JDK is strictly lower than 21 -->
+                <jdk>(,21)</jdk>
+            </activation>
+            <build>
+                <plugins>
+                    <plugin>
+                        <groupId>org.apache.maven.plugins</groupId>
+                        <artifactId>maven-compiler-plugin</artifactId>
+                        <!-- Version 3.11.0+ recommended for best support -->
+                        <version>3.13.0</version>
+                        <configuration>
+                            <excludes>
+                                <!-- Path is relative to src/main/java -->
+                                <exclude>software/amazon/lambda/durable/examples/vt</exclude>
+                            </excludes>
+                        </configuration>
+                    </plugin>
+                </plugins>
+            </build>
+        </profile>
+    </profiles>
+
     <build>
         <plugins>
             <plugin>

examples/src/main/java/software/amazon/lambda/durable/examples/vt/ManyAsyncStepsVirtualThreadPoolExample.java

@@ -0,0 +1,78 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.lambda.durable.examples.vt;
+
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import software.amazon.lambda.durable.DurableConfig;
+import software.amazon.lambda.durable.DurableContext;
+import software.amazon.lambda.durable.DurableFuture;
+import software.amazon.lambda.durable.DurableHandler;
+
+/**
+ * Performance test example demonstrating concurrent async steps.
+ *
+ * <p>This example tests the SDK's ability to handle many concurrent operations:
+ *
+ * <ul>
+ *   <li>Creates async steps in a loop
+ *   <li>Each step performs a simple computation
+ *   <li>All results are collected using {@link DurableFuture#allOf}
+ * </ul>
+ */
+public class ManyAsyncStepsVirtualThreadPoolExample
+        extends DurableHandler<
+                ManyAsyncStepsVirtualThreadPoolExample.Input, ManyAsyncStepsVirtualThreadPoolExample.Output> {
+
+    public record Input(int multiplier, int steps) {}
+
+    public record Output(long result, long executionTimeMs, long replayTimeMs) {}
+
+    @Override
+    public Output handleRequest(Input input, DurableContext context) {
+        var startTime = System.nanoTime();
+        var multiplier = input.multiplier();
+        var steps = input.steps();
+        var logger = context.getLogger();
+
+        logger.info("Starting {} async steps with multiplier {}", steps, multiplier);
+
+        // Create async steps
+        var futures = new ArrayList<DurableFuture<Integer>>(steps);
+        for (var i = 0; i < steps; i++) {
+            var index = i;
+            var future = context.stepAsync("compute-" + i, Integer.class, stepCtx -> index * multiplier);
+            futures.add(future);
+        }
+
+        logger.info("All {} async steps created, collecting results", steps);
+
+        // Collect all results using allOf
+        var results = DurableFuture.allOf(futures);
+        var totalSum = results.stream().mapToInt(Integer::intValue).sum();
+
+        // checkpoint the executionTime so that we can have the same value when replay
+        var executionTimeMs = context.step(
+                "execution-time", Long.class, stepCtx -> TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startTime));
+        logger.info("Completed {} steps, total sum: {}, execution time: {}ms", steps, totalSum, executionTimeMs);
+
+        // Wait 2 seconds to test replay
+        context.wait("post-compute-wait", Duration.ofSeconds(2));
+
+        var replayTimeMs = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startTime);
+
+        return new Output(totalSum, executionTimeMs, replayTimeMs);
+    }
+
+    @Override
+    protected DurableConfig createConfiguration() {
+        // Add a small checkpoint delay to help batch the checkpoint requests and reduce the overall latencies
+        // when the function has many concurrent operations
+        return DurableConfig.builder()
+                .withCheckpointDelay(Duration.ofMillis(10))
+                .withExecutorService(Executors.newVirtualThreadPerTaskExecutor())
+                .build();
+    }
+}

examples/src/test/java/software/amazon/lambda/durable/examples/CloudBasedIntegrationTest.java

@@ -544,7 +544,7 @@ void testManyAsyncStepsExample(int steps, long maxExecutionTime, long maxReplayT
             assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
 
             var finalResult = result.getResult();
-            System.out.printf("ManyAsyncStepsExample result (%d steps): %s\n", steps, finalResult);
+            System.out.printf("ManyAsyncStepsVirtualThreadPoolExample result (%d steps): %s\n", steps, finalResult);
             assertNotNull(finalResult);
             assertEquals((long) steps * (steps - 1), finalResult.result()); // Sum of 0..steps * 2
 

examples/src/test/java/software/amazon/lambda/durable/examples/vt/ManyAsyncStepsVirtualTheadPoolExampleTest.java

@@ -0,0 +1,68 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.lambda.durable.examples.vt;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.condition.EnabledForJreRange;
+import org.junit.jupiter.api.condition.JRE;
+import software.amazon.lambda.durable.model.ExecutionStatus;
+import software.amazon.lambda.durable.testing.LocalDurableTestRunner;
+
+@EnabledForJreRange(min = JRE.JAVA_21)
+class ManyAsyncStepsVirtualTheadPoolExampleTest {
+
+    @Test
+    void testManyAsyncSteps() {
+        var handler = new ManyAsyncStepsVirtualThreadPoolExample();
+        var runner = LocalDurableTestRunner.create(ManyAsyncStepsVirtualThreadPoolExample.Input.class, handler);
+
+        var input = new ManyAsyncStepsVirtualThreadPoolExample.Input(2, 500);
+        var result = runner.runUntilComplete(input);
+
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+
+        var output = result.getResult(ManyAsyncStepsVirtualThreadPoolExample.Output.class);
+        assertNotNull(output);
+
+        // Sum of 0..499 * 2 = 499 * 500 / 2 * 2 = 249500
+        assertEquals(
+                249500,
+                result.getResult(ManyAsyncStepsVirtualThreadPoolExample.Output.class)
+                        .result());
+    }
+
+    @Test
+    void testManyAsyncStepsWithDefaultMultiplier() {
+        var handler = new ManyAsyncStepsVirtualThreadPoolExample();
+        var runner = LocalDurableTestRunner.create(ManyAsyncStepsVirtualThreadPoolExample.Input.class, handler);
+
+        var input = new ManyAsyncStepsVirtualThreadPoolExample.Input(1, 500);
+        var result = runner.runUntilComplete(input);
+
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+
+        // Sum of 0..499 = 499 * 500 / 2 = 124750
+        assertEquals(
+                124750,
+                result.getResult(ManyAsyncStepsVirtualThreadPoolExample.Output.class)
+                        .result());
+    }
+
+    @Test
+    void testOperationsAreTracked() {
+        var handler = new ManyAsyncStepsVirtualThreadPoolExample();
+        var runner = LocalDurableTestRunner.create(ManyAsyncStepsVirtualThreadPoolExample.Input.class, handler);
+
+        var result = runner.runUntilComplete(new ManyAsyncStepsVirtualThreadPoolExample.Input(1, 500));
+
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+
+        // Verify some operations are tracked
+        assertNotNull(result.getOperation("compute-0"));
+        assertNotNull(result.getOperation("compute-499"));
+        assertNotNull(result.getOperation("compute-250"));
+    }
+}