add ExecutionOperation as a special durable operation

Author: zhongkechen

sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java

@@ -6,13 +6,16 @@
 import java.time.Duration;
 import java.util.Objects;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.BiFunction;
 import java.util.function.Function;
 import java.util.function.Supplier;
 import org.slf4j.LoggerFactory;
 import software.amazon.lambda.durable.execution.ExecutionManager;
 import software.amazon.lambda.durable.logging.DurableLogger;
+import software.amazon.lambda.durable.model.DurableExecutionOutput;
 import software.amazon.lambda.durable.operation.CallbackOperation;
 import software.amazon.lambda.durable.operation.ChildContextOperation;
+import software.amazon.lambda.durable.operation.ExecutionOperation;
 import software.amazon.lambda.durable.operation.InvokeOperation;
 import software.amazon.lambda.durable.operation.StepOperation;
 import software.amazon.lambda.durable.operation.WaitOperation;
@@ -298,6 +301,28 @@ public <T> DurableFuture<T> runInChildContextAsync(
         return operation;
     }
 
+    // =============== execute ================
+
+    /**
+     * Executes a durable function with automatic checkpointing and replay support.
+     *
+     * <p>This is the entry point for durable execution. The function will be executed with the ability to checkpoint
+     * state and resume from the last checkpoint on subsequent invocations. The execution automatically handles
+     * serialization, deserialization, and replay of operations.
+     *
+     * @param <I> the input type for the durable function
+     * @param <O> the output type returned by the durable function
+     * @param inputType the Class object representing the input type, used for deserialization
+     * @param func the durable function to execute, receiving the deserialized input and this DurableContext
+     * @return a DurableExecutionOutput containing the execution result and metadata
+     */
+    <I, O> DurableExecutionOutput execute(Class<I> inputType, BiFunction<I, DurableContext, O> func) {
+        var operation = new ExecutionOperation<>(
+                getExecutionManager().getExecutionOperation().id(), inputType, func, this);
+        operation.execute();
+        return operation.get();
+    }
+
     // =============== accessors ================
     /**
      * Returns a logger with execution context information for replay-aware logging.

sdk/src/main/java/software/amazon/lambda/durable/DurableExecutor.java

@@ -4,136 +4,27 @@
 
 import com.amazonaws.services.lambda.runtime.Context;
 import com.amazonaws.services.lambda.runtime.RequestHandler;
-import java.nio.charset.StandardCharsets;
-import java.util.concurrent.CompletableFuture;
 import java.util.function.BiFunction;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-import software.amazon.awssdk.services.lambda.model.ErrorObject;
-import software.amazon.awssdk.services.lambda.model.Operation;
-import software.amazon.awssdk.services.lambda.model.OperationAction;
-import software.amazon.awssdk.services.lambda.model.OperationType;
-import software.amazon.awssdk.services.lambda.model.OperationUpdate;
-import software.amazon.lambda.durable.exception.DurableOperationException;
-import software.amazon.lambda.durable.exception.IllegalDurableOperationException;
-import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException;
 import software.amazon.lambda.durable.execution.ExecutionManager;
-import software.amazon.lambda.durable.execution.SuspendExecutionException;
-import software.amazon.lambda.durable.execution.ThreadContext;
-import software.amazon.lambda.durable.execution.ThreadType;
 import software.amazon.lambda.durable.model.DurableExecutionInput;
 import software.amazon.lambda.durable.model.DurableExecutionOutput;
-import software.amazon.lambda.durable.serde.SerDes;
-import software.amazon.lambda.durable.util.ExceptionHelper;
 
 public class DurableExecutor {
-    private static final Logger logger = LoggerFactory.getLogger(DurableExecutor.class);
-
-    // Lambda response size limit is 6MB minus small epsilon for envelope
-    private static final int LAMBDA_RESPONSE_SIZE_LIMIT = 6 * 1024 * 1024 - 50;
 
     public static <I, O> DurableExecutionOutput execute(
             DurableExecutionInput input,
             Context lambdaContext,
             Class<I> inputType,
             BiFunction<I, DurableContext, O> handler,
             DurableConfig config) {
-        var executionManager = new ExecutionManager(
-                input.durableExecutionArn(), input.checkpointToken(), input.initialExecutionState(), config);
-
-        var handlerFuture = CompletableFuture.supplyAsync(
-                () -> {
-                    var userInput =
-                            extractUserInput(executionManager.getExecutionOperation(), config.getSerDes(), inputType);
-                    // Create context in the executor thread so it detects the correct thread name
-                    var context = DurableContext.createRootContext(executionManager, config, lambdaContext);
-                    executionManager.registerActiveThread(null);
-                    executionManager.setCurrentThreadContext(new ThreadContext(null, ThreadType.CONTEXT));
-                    return handler.apply(userInput, context);
-                },
-                config.getExecutorService()); // Get executor from config for running user code
-
-        // Execute the handlerFuture in ExecutionManager. If it completes successfully, the output of user function
-        // will be returned. Otherwise, it will complete exceptionally with a SuspendExecutionException or a failure.
-        return executionManager
-                .runUntilCompleteOrSuspend(handlerFuture)
-                .handle((result, ex) -> {
-                    if (ex != null) {
-                        // an exception thrown from handlerFuture or suspension/termination occurred
-                        Throwable cause = ExceptionHelper.unwrapCompletableFuture(ex);
-                        if (cause instanceof SuspendExecutionException) {
-                            return DurableExecutionOutput.pending();
-                        }
-
-                        logger.debug("Execution failed: {}", cause.getMessage());
-                        return DurableExecutionOutput.failure(buildErrorObject(cause, config.getSerDes()));
-                    }
-                    // user handler complete successfully
-                    var outputPayload = config.getSerDes().serialize(result);
-
-                    logger.debug("Execution completed");
-                    return DurableExecutionOutput.success(handleLargePayload(executionManager, outputPayload));
-                })
-                .whenComplete((v, ex) -> {
-                    // We shutdown the execution to make sure remaining checkpoint calls in the queue are drained
-                    // We DO NOT shutdown the executor since it should stay warm for re-invokes against a warm Lambda
-                    // runtime.
-                    // For example, a re-invoke after a wait should re-use the same executor instance from
-                    // DurableConfig.
-                    // userExecutor.shutdown();
-                    executionManager.shutdown();
-                })
-                .join();
-    }
-
-    private static String handleLargePayload(ExecutionManager executionManager, String outputPayload) {
-        // Check if the serialized payload exceeds Lambda response size limit
-        var payloadSize = outputPayload != null ? outputPayload.getBytes(StandardCharsets.UTF_8).length : 0;
-
-        if (payloadSize > LAMBDA_RESPONSE_SIZE_LIMIT) {
-            logger.debug(
-                    "Response size ({} bytes) exceeds Lambda limit ({} bytes). Checkpointing result.",
-                    payloadSize,
-                    LAMBDA_RESPONSE_SIZE_LIMIT);
-
-            // Checkpoint the large result and wait for it to complete
-            executionManager
-                    .sendOperationUpdate(OperationUpdate.builder()
-                            .type(OperationType.EXECUTION)
-                            .id(executionManager.getExecutionOperation().id())
-                            .action(OperationAction.SUCCEED)
-                            .payload(outputPayload)
-                            .build())
-                    .join();
 
-            // Return empty result, we checkpointed the data manually
-            logger.debug("Execution completed (large response checkpointed)");
-            return "";
+        // We shut down the executionManager to make sure remaining checkpoint calls in the queue are drained
+        // We DO NOT shut down the executor since it should stay warm for re-invokes against a warm Lambda runtime.
+        // For example, a re-invoke after a wait should re-use the same executor instance from DurableConfig.
+        try (var executionManager = new ExecutionManager(input, config)) {
+            var context = DurableContext.createRootContext(executionManager, config, lambdaContext);
+            return context.execute(inputType, handler);
         }
-
-        // If response size is acceptable, return the result directly
-        return outputPayload;
-    }
-
-    private static ErrorObject buildErrorObject(Throwable e, SerDes serDes) {
-        // exceptions thrown from operations, e.g. Step
-        if (e instanceof DurableOperationException) {
-            return ((DurableOperationException) e).getErrorObject();
-        }
-        if (e instanceof UnrecoverableDurableExecutionException) {
-            return ((UnrecoverableDurableExecutionException) e).getErrorObject();
-        }
-        // exceptions thrown from non-operation code
-        return ExceptionHelper.buildErrorObject(e, serDes);
-    }
-
-    private static <I> I extractUserInput(Operation executionOp, SerDes serDes, Class<I> inputType) {
-        if (executionOp.executionDetails() == null) {
-            throw new IllegalDurableOperationException("EXECUTION operation missing executionDetails");
-        }
-
-        var inputPayload = executionOp.executionDetails().inputPayload();
-        return serDes.deserialize(inputPayload, TypeToken.get(inputType));
     }
 
     public static <I, O> RequestHandler<DurableExecutionInput, DurableExecutionOutput> wrap(

sdk/src/main/java/software/amazon/lambda/durable/execution/ExecutionManager.java

@@ -22,6 +22,7 @@
 import software.amazon.awssdk.services.lambda.model.OperationUpdate;
 import software.amazon.lambda.durable.DurableConfig;
 import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException;
+import software.amazon.lambda.durable.model.DurableExecutionInput;
 import software.amazon.lambda.durable.operation.BaseDurableOperation;
 
 /**
@@ -45,7 +46,7 @@
  *
  * @see InternalExecutor
  */
-public class ExecutionManager {
+public class ExecutionManager implements AutoCloseable {
 
     private static final Logger logger = LoggerFactory.getLogger(ExecutionManager.class);
 
@@ -65,25 +66,21 @@ public class ExecutionManager {
     // ===== Checkpoint Batching =====
     private final CheckpointBatcher checkpointBatcher;
 
-    public ExecutionManager(
-            String durableExecutionArn,
-            String checkpointToken,
-            CheckpointUpdatedExecutionState initialExecutionState,
-            DurableConfig config) {
-        this.durableExecutionArn = durableExecutionArn;
+    public ExecutionManager(DurableExecutionInput input, DurableConfig config) {
+        this.durableExecutionArn = input.durableExecutionArn();
 
         // Create checkpoint batcher for internal coordination
         this.checkpointBatcher =
-                new CheckpointBatcher(config, durableExecutionArn, checkpointToken, this::onCheckpointComplete);
+                new CheckpointBatcher(config, durableExecutionArn, input.checkpointToken(), this::onCheckpointComplete);
 
-        this.operationStorage = checkpointBatcher.fetchAllPages(initialExecutionState).stream()
+        this.operationStorage = checkpointBatcher.fetchAllPages(input.initialExecutionState()).stream()
                 .collect(Collectors.toConcurrentMap(Operation::id, op -> op));
 
         // Start in REPLAY mode if we have more than just the initial EXECUTION operation
         this.executionMode =
                 new AtomicReference<>(operationStorage.size() > 1 ? ExecutionMode.REPLAY : ExecutionMode.EXECUTION);
 
-        executionOp = findExecutionOp(initialExecutionState);
+        executionOp = findExecutionOp(input.initialExecutionState());
 
         // Validate initial operation is an EXECUTION operation
         if (executionOp == null) {
@@ -248,7 +245,17 @@ public CompletableFuture<Operation> pollForOperationUpdates(String operationId,
     }
 
     // ===== Utilities =====
-    public void shutdown() {
+
+    /**
+     * Closes this stream and releases any system resources associated with it. If the stream is already closed then
+     * invoking this method has no effect.
+     *
+     * <p>As noted in {@link AutoCloseable#close()}, cases where the close may fail require careful attention. It is
+     * strongly advised to relinquish the underlying resources and to internally <em>mark</em> the {@code Closeable} as
+     * closed, prior to throwing the {@code IOException}.
+     */
+    @Override
+    public void close() {
         checkpointBatcher.shutdown();
     }
 

sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java

@@ -150,8 +150,8 @@ private void checkpointSuccess(T result) {
     private void handleChildContextFailure(Throwable exception) {
         exception = ExceptionHelper.unwrapCompletableFuture(exception);
         if (exception instanceof SuspendExecutionException) {
-            // Rethrow Error immediately — do not checkpoint
-            ExceptionHelper.sneakyThrow(exception);
+            // The execution is going to suspend. do nothing
+            return;
         }
         if (exception instanceof UnrecoverableDurableExecutionException) {
             terminateExecution((UnrecoverableDurableExecutionException) exception);