Merge remote-tracking branch 'origin/main' into exception

Author: zhongkechen

sdk/src/main/java/com/amazonaws/lambda/durable/DurableExecutor.java

@@ -4,6 +4,7 @@
 
 import com.amazonaws.lambda.durable.exception.DurableOperationException;
 import com.amazonaws.lambda.durable.execution.ExecutionManager;
+import com.amazonaws.lambda.durable.execution.SuspendExecutionException;
 import com.amazonaws.lambda.durable.model.DurableExecutionInput;
 import com.amazonaws.lambda.durable.model.DurableExecutionOutput;
 import com.amazonaws.lambda.durable.serde.SerDes;
@@ -81,11 +82,7 @@ public static <I, O> DurableExecutionOutput execute(
                     userExecutor);
 
             // Get suspend future from ExecutionManager. If this future completes, it
-            // indicates
-            // that no threads are active and we can safely suspend. This is useful for
-            // async scenarios where multiple operations are scheduled concurrently and
-            // awaited
-            // at a later point.
+            // indicates that no threads are active and we can safely suspend.
             var suspendFuture = executionManager.getSuspendExecutionFuture();
 
             // Wait for either handler to complete or suspension to occur
@@ -96,16 +93,6 @@ public static <I, O> DurableExecutionOutput execute(
                 return DurableExecutionOutput.pending();
             }
 
-            if (handlerFuture.isCompletedExceptionally()) {
-                try {
-                    handlerFuture.join(); // Will throw the exception
-                } catch (Exception e) {
-                    Throwable cause = ExceptionHelper.unwrapCompletableFuture(e);
-                    logger.debug("Execution failed: {}", cause.getMessage());
-                    return DurableExecutionOutput.failure(buildErrorObject(cause, serDes));
-                }
-            }
-
             var result = handlerFuture.get();
             var outputPayload = serDes.serialize(result);
 
@@ -137,7 +124,13 @@ public static <I, O> DurableExecutionOutput execute(
             logger.debug("Execution completed");
             return DurableExecutionOutput.success(outputPayload);
         } catch (Exception e) {
-            return DurableExecutionOutput.failure(buildErrorObject(ExceptionHelper.unwrapCompletableFuture(e), serDes));
+            Throwable cause = ExceptionHelper.unwrapCompletableFuture(e);
+            if (cause instanceof SuspendExecutionException) {
+                logger.debug("Execution suspended");
+                return DurableExecutionOutput.pending();
+            }
+            logger.debug("Execution failed: {}", cause.getMessage());
+            return DurableExecutionOutput.failure(buildErrorObject(cause, serDes));
         } finally {
             // We shutdown the execution to make sure remaining checkpoint calls in the queue are drained
             executionManager.shutdown();

sdk/src/main/java/com/amazonaws/lambda/durable/operation/StepOperation.java

@@ -10,6 +10,7 @@
 import com.amazonaws.lambda.durable.exception.StepInterruptedException;
 import com.amazonaws.lambda.durable.execution.ExecutionManager;
 import com.amazonaws.lambda.durable.execution.ExecutionPhase;
+import com.amazonaws.lambda.durable.execution.SuspendExecutionException;
 import com.amazonaws.lambda.durable.execution.ThreadType;
 import com.amazonaws.lambda.durable.logging.DurableLogger;
 import com.amazonaws.lambda.durable.serde.SerDes;
@@ -190,7 +191,14 @@ private void executeStepLogic(int attempt) {
             } catch (Throwable e) {
                 handleStepError(e, attempt);
             } finally {
-                executionManager.deregisterActiveThread(stepThreadId);
+                try {
+                    executionManager.deregisterActiveThread(stepThreadId);
+                } catch (SuspendExecutionException e) {
+                    // Expected when this is the last active thread. Must catch here because:
+                    // 1/ This runs in a worker thread detached from handlerFuture
+                    // 2/ Uncaught exception would prevent phaser from advancing, blocking stepAsync().get()
+                    // Suspension is already signaled via suspendExecutionFuture before the throw.
+                }
                 durableLogger.clearOperationContext();
             }
         });