add retryable exception

Author: zhongkechen

sdk/src/main/java/software/amazon/lambda/durable/exception/UnrecoverableDurableExecutionException.java

@@ -7,14 +7,25 @@
 /** Exception thrown when the execution is not recoverable. The durable execution will be immediately terminated. */
 public class UnrecoverableDurableExecutionException extends DurableExecutionException {
     private final ErrorObject errorObject;
+    private final boolean retryable;
 
-    public UnrecoverableDurableExecutionException(ErrorObject errorObject) {
+    public UnrecoverableDurableExecutionException(ErrorObject errorObject, boolean retryable) {
         super(errorObject.errorMessage());
         this.errorObject = errorObject;
+        this.retryable = retryable;
+    }
+
+    public UnrecoverableDurableExecutionException(ErrorObject errorObject) {
+        this(errorObject, false);
     }
 
     /** Returns the error details for this unrecoverable exception. */
     public ErrorObject getErrorObject() {
         return errorObject;
     }
+
+    /** Returns true if the execution can be retried. */
+    public boolean isRetryable() {
+        return retryable;
+    }
 }

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

@@ -6,6 +6,7 @@
 import com.amazonaws.services.lambda.runtime.RequestHandler;
 import java.nio.charset.StandardCharsets;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
 import java.util.function.BiFunction;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -65,26 +66,43 @@ public static <I, O> DurableExecutionOutput execute(
             // 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();
+            try {
+                return executionManager
+                        .runUntilCompleteOrSuspend(handlerFuture)
+                        .handle((result, ex) -> {
+                            if (ex != null) {
+                                // an exception thrown from handlerFuture or suspension/termination occurred
+                                Throwable cause = ExceptionHelper.unwrapCompletableFuture(ex);
+
+                                // return PENDING if it's SuspendExecutionException
+                                if (cause instanceof SuspendExecutionException) {
+                                    return DurableExecutionOutput.pending();
+                                }
+
+                                // let the backend retry the invocation if the exception is retryable
+                                if (cause
+                                                instanceof
+                                                UnrecoverableDurableExecutionException
+                                                        unrecoverableDurableExecutionException
+                                        && unrecoverableDurableExecutionException.isRetryable()) {
+                                    throw unrecoverableDurableExecutionException;
+                                }
+
+                                // fail the execution otherwise
+                                logger.debug("Execution failed: {}", cause.getMessage());
+                                return DurableExecutionOutput.failure(buildErrorObject(cause, config.getSerDes()));
                             }
-
-                            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));
-                    })
-                    .join();
+                            // user handler complete successfully
+                            logger.debug("Execution completed");
+                            var outputPayload = config.getSerDes().serialize(result);
+                            return DurableExecutionOutput.success(handleLargePayload(executionManager, outputPayload));
+                        })
+                        .join();
+            } catch (CompletionException e) {
+                // unwrap the CompletionException and rethrow the wrapped exception
+                ExceptionHelper.sneakyThrow(ExceptionHelper.unwrapCompletableFuture(e));
+                return null;
+            }
         }
     }
 

sdk/src/test/java/software/amazon/lambda/durable/execution/DurableExecutionTest.java

@@ -14,13 +14,15 @@
 import java.util.concurrent.ExecutorService;
 import org.junit.jupiter.api.Test;
 import software.amazon.awssdk.services.lambda.model.CheckpointUpdatedExecutionState;
+import software.amazon.awssdk.services.lambda.model.ErrorObject;
 import software.amazon.awssdk.services.lambda.model.ExecutionDetails;
 import software.amazon.awssdk.services.lambda.model.Operation;
 import software.amazon.awssdk.services.lambda.model.OperationStatus;
 import software.amazon.awssdk.services.lambda.model.OperationType;
 import software.amazon.awssdk.services.lambda.model.StepDetails;
 import software.amazon.lambda.durable.DurableConfig;
 import software.amazon.lambda.durable.TestUtils;
+import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException;
 import software.amazon.lambda.durable.model.DurableExecutionInput;
 import software.amazon.lambda.durable.model.ExecutionStatus;
 
@@ -134,6 +136,42 @@ void testExecuteFailure() {
         assertEquals("Test error", output.error().errorMessage());
     }
 
+    @Test
+    void testRetryableExceptions() {
+        var executionOp = Operation.builder()
+                .id(EXECUTION_OP_ID)
+                .type(OperationType.EXECUTION)
+                .status(OperationStatus.STARTED)
+                .executionDetails(ExecutionDetails.builder()
+                        .inputPayload("\"test-input\"")
+                        .build())
+                .build();
+
+        var input = new DurableExecutionInput(
+                EXECUTION_ARN,
+                "token1",
+                CheckpointUpdatedExecutionState.builder()
+                        .operations(List.of(executionOp))
+                        .build());
+
+        UnrecoverableDurableExecutionException ex = assertThrows(
+                UnrecoverableDurableExecutionException.class,
+                () -> DurableExecutor.execute(
+                        input,
+                        null,
+                        get(String.class),
+                        (userInput, ctx) -> {
+                            throw new UnrecoverableDurableExecutionException(
+                                    ErrorObject.builder()
+                                            .errorMessage("Test error")
+                                            .build(),
+                                    true);
+                        },
+                        configWithMockClient()));
+
+        assertTrue(ex.isRetryable());
+    }
+
     @Test
     void testExecuteReplay() {
         var executionOp = Operation.builder()