fix: apply invocation deadlines to SDK waits and cleanup

docs/advanced/configuration.md

@@ -105,3 +105,34 @@ file retention.
 Custom serializers can override `serialize(value, SerDesContext)` and
 `deserialize(data, typeToken, SerDesContext)` to receive stable execution/entity
 identity. Existing implementations inherit defaults that call the original methods.
+
+## Invocation deadline cleanup
+
+The SDK captures the Lambda invocation's remaining time once and uses a monotonic
+budget for its waits and cleanup. Deadline exhaustion is an invocation interruption
+that can be retried; it is not checkpointed as a business-operation failure.
+Completed operation results and failures remain available for replay.
+
+Normal finalization drains invocation-owned user work before result preparation
+and invocation End. Cleanup first requests cooperative cancellation when work does
+not settle within its allowance, then rejects new durable admission. Required
+checkpoint waits and cleanup use the remaining budget. A terminal execution result
+already acknowledged by the service remains authoritative, while a JVM fatal is
+never replaced by a cleanup timeout or an acknowledged-result shortcut.
+
+Interruption does not skip the remaining cleanup attempts: the SDK records it,
+continues best-effort joins and checkpoint shutdown with the remaining budget, and
+restores the interrupt flag. If an uncooperative root exits after its caller has
+already timed out, its original worker still runs the matching invocation End once
+to restore Start-installed scopes. This eventual cleanup does not admit new durable
+work or prepare another result. An End hook that has already started can itself
+remain blocked; the SDK cannot forcibly stop arbitrary callbacks.
+
+This is not an unconditional bounded-return or worker-reclamation guarantee.
+Already-running application I/O, synchronous initial-state reads, custom clients,
+and plugin hooks cannot be forcibly terminated. In particular, an executor that
+runs the root inline, including a saturated `CallerRunsPolicy`, can block the runtime
+caller before it reaches its timeout-aware wait. Use finite downstream I/O timeouts
+and executor configurations that let the runtime caller make progress. Cancellation
+cannot undo a request already accepted by the service; interrupted side effects
+still require the documented retry and idempotency strategy.

otel-plugin/src/test/java/software/amazon/lambda/durable/otel/HandlerMdcIntegrationTest.java

@@ -22,6 +22,7 @@
 import org.slf4j.MDC;
 import software.amazon.lambda.durable.DurableConfig;
 import software.amazon.lambda.durable.TypeToken;
+import software.amazon.lambda.durable.execution.ExecutionManager;
 import software.amazon.lambda.durable.model.ExecutionStatus;
 import software.amazon.lambda.durable.plugin.DurableExecutionPlugin;
 import software.amazon.lambda.durable.plugin.InvocationEndInfo;
@@ -262,10 +263,19 @@ private static void awaitMdcLatch(CountDownLatch latch) {
     private static void awaitMdcCallerJoin(Thread caller) {
         var deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(3);
         while (System.nanoTime() < deadline) {
-            if (caller.getState() == Thread.State.WAITING
-                    && Arrays.stream(caller.getStackTrace())
+            var state = caller.getState();
+            var frames = caller.getStackTrace();
+            var legacyJoin = Arrays.stream(frames)
+                    .anyMatch(frame -> frame.getClassName().equals(CompletableFuture.class.getName())
+                            && frame.getMethodName().equals("join"));
+            var budgetedWait = Arrays.stream(frames)
+                            .anyMatch(frame -> frame.getClassName().equals(ExecutionManager.class.getName())
+                                    && frame.getMethodName().equals("awaitInvocationOutcome"))
+                    && Arrays.stream(frames)
                             .anyMatch(frame -> frame.getClassName().equals(CompletableFuture.class.getName())
-                                    && frame.getMethodName().equals("join"))) return;
+                                    && frame.getMethodName().equals("get"));
+            if ((state == Thread.State.WAITING || state == Thread.State.TIMED_WAITING) && (legacyJoin || budgetedWait))
+                return;
             LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(1));
         }
         fail("Caller did not attach its completion observer before releasing the handler");

otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationEndCompatibilityTest.java

@@ -21,6 +21,7 @@
 import software.amazon.lambda.durable.DurableConfig;
 import software.amazon.lambda.durable.TypeToken;
 import software.amazon.lambda.durable.execution.DurableExecutor;
+import software.amazon.lambda.durable.execution.ExecutionManager;
 import software.amazon.lambda.durable.model.DurableExecutionInput;
 import software.amazon.lambda.durable.model.ExecutionStatus;
 import software.amazon.lambda.durable.plugin.DurableExecutionPlugin;
@@ -262,10 +263,19 @@ private record Observation(
     private static void awaitCallerJoin(Thread caller) {
         var end = System.nanoTime() + TimeUnit.SECONDS.toNanos(3);
         while (System.nanoTime() < end) {
-            if (caller.getState() == Thread.State.WAITING
-                    && Arrays.stream(caller.getStackTrace())
+            var state = caller.getState();
+            var frames = caller.getStackTrace();
+            var legacyJoin = Arrays.stream(frames)
+                    .anyMatch(frame -> frame.getClassName().equals(CompletableFuture.class.getName())
+                            && frame.getMethodName().equals("join"));
+            var budgetedWait = Arrays.stream(frames)
+                            .anyMatch(frame -> frame.getClassName().equals(ExecutionManager.class.getName())
+                                    && frame.getMethodName().equals("awaitInvocationOutcome"))
+                    && Arrays.stream(frames)
                             .anyMatch(frame -> frame.getClassName().equals(CompletableFuture.class.getName())
-                                    && frame.getMethodName().equals("join"))) return;
+                                    && frame.getMethodName().equals("get"));
+            if ((state == Thread.State.WAITING || state == Thread.State.TIMED_WAITING) && (legacyJoin || budgetedWait))
+                return;
             LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(1));
         }
         fail("invocation caller did not reach its future join before releasing the handler");

otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationPreparationAdmissionProbeTest.java

@@ -21,6 +21,7 @@
 import software.amazon.lambda.durable.config.*;
 import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException;
 import software.amazon.lambda.durable.execution.DurableExecutor;
+import software.amazon.lambda.durable.execution.ExecutionManager;
 import software.amazon.lambda.durable.model.DurableExecutionInput;
 import software.amazon.lambda.durable.model.ExecutionStatus;
 import software.amazon.lambda.durable.model.WaitForConditionResult;
@@ -273,7 +274,9 @@ private static void awaitDrain(Thread root) {
         var deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(3);
         while (System.nanoTime() < deadline) {
             if (Arrays.stream(root.getStackTrace())
-                    .anyMatch(frame -> frame.getMethodName().equals("waitForRunningUserHandlers"))) return;
+                    .anyMatch(frame -> frame.getClassName().equals(ExecutionManager.class.getName())
+                            && (frame.getMethodName().equals("waitForRunningUserHandlers")
+                                    || frame.getMethodName().equals("waitForRunningUserHandlersUntil")))) return;
             LockSupport.parkNanos(100_000);
         }
         fail("Root did not begin draining the pending child");

sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ContinuationEndCleanupOrderTest.java

@@ -341,9 +341,17 @@ public void onInvocationEnd(InvocationEndInfo info) {
     private static void awaitCallerOutcomeJoin(Thread caller) {
         var deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(3);
         while (System.nanoTime() < deadline) {
-            if (Arrays.stream(caller.getStackTrace())
+            var frames = caller.getStackTrace();
+            var legacyJoin = Arrays.stream(frames)
                     .anyMatch(frame -> frame.getClassName().equals(CompletableFuture.class.getName())
-                            && frame.getMethodName().equals("join"))) return;
+                            && frame.getMethodName().equals("join"));
+            var budgetedWait = Arrays.stream(frames)
+                            .anyMatch(frame -> frame.getClassName().equals(ExecutionManager.class.getName())
+                                    && frame.getMethodName().equals("awaitInvocationOutcome"))
+                    && Arrays.stream(frames)
+                            .anyMatch(frame -> frame.getClassName().equals(CompletableFuture.class.getName())
+                                    && frame.getMethodName().equals("get"));
+            if (legacyJoin || budgetedWait) return;
             Thread.yield();
         }
         fail("Invocation caller did not reach its installed outcome observer");

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

@@ -8,7 +8,9 @@
 import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import java.util.function.Consumer;
 import java.util.function.Function;
 
@@ -23,7 +25,7 @@
  * @param <T> Request type
  */
 public class ApiRequestDelayedBatcher<T> {
-    private static final Duration MAX_DELAY = Duration.ofMinutes(60);
+    private static final long NO_FLUSH_DEADLINE = Long.MAX_VALUE;
 
     /** Maximum items allowed in a single batch */
     private final int maxItemCount;
@@ -108,6 +110,12 @@ void shutdown() {
         flushingQueueFuture.join();
     }
 
+    /** Flushes pending requests without waiting beyond the caller's remaining cleanup budget. */
+    void shutdown(Duration timeout) throws InterruptedException, ExecutionException, TimeoutException {
+        flush();
+        flushingQueueFuture.get(Math.max(0, timeout.toNanos()), TimeUnit.NANOSECONDS);
+    }
+
     /** Schedules pending requests immediately without waiting for the current batch. */
     void flush() {
         synchronized (delayedBatch) {
@@ -121,8 +129,8 @@ void flush() {
     /** clear the current batch and creates a new batch */
     private void initializeDelayedBatch() {
         this.delayedBatch.clear();
-        // MAX_DELAY is longer than a single Lambda invocation
-        this.delayedBatchFlushTime = System.nanoTime() + MAX_DELAY.toNanos();
+        // No timer is scheduled until an item supplies its flush time.
+        this.delayedBatchFlushTime = NO_FLUSH_DEADLINE;
 
         // the timer future is created initially without a timeout until an item is added to the batch
         this.delayedBatchFlushTimer = new CompletableFuture<>();

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

@@ -11,16 +11,20 @@
 import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeoutException;
 import java.util.function.BooleanSupplier;
 import java.util.function.Consumer;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import software.amazon.awssdk.awscore.exception.AwsServiceException;
 import software.amazon.awssdk.services.lambda.model.CheckpointUpdatedExecutionState;
 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.DurableConfig;
+import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException;
 import software.amazon.lambda.durable.retry.PollingStrategies;
 import software.amazon.lambda.durable.retry.PollingStrategy;
 import software.amazon.lambda.durable.util.DurableApiErrorClassifier;
@@ -47,6 +51,9 @@ class CheckpointManager {
     private final BooleanSupplier tryStartCheckpointProcessing;
     private final Runnable onCheckpointProcessingFinish;
     private final Runnable onTokenRevoked;
+    private final Runnable requestAdmission;
+    private final Object checkpointLock = new Object();
+    private volatile boolean terminalExecutionAcknowledged;
     private volatile String checkpointToken;
     private boolean pollingStopped; // Guarded by pollingFutures.
     /**
@@ -71,6 +78,27 @@ class CheckpointManager {
             BooleanSupplier tryStartCheckpointProcessing,
             Runnable onCheckpointProcessingFinish,
             Runnable onTokenRevoked) {
+        this(
+                config,
+                durableExecutionArn,
+                checkpointToken,
+                callback,
+                tryStartCheckpointProcessing,
+                onCheckpointProcessingFinish,
+                onTokenRevoked,
+                () -> {});
+    }
+
+    CheckpointManager(
+            DurableConfig config,
+            String durableExecutionArn,
+            String checkpointToken,
+            Consumer<List<Operation>> callback,
+            BooleanSupplier tryStartCheckpointProcessing,
+            Runnable onCheckpointProcessingFinish,
+            Runnable onTokenRevoked,
+            Runnable requestAdmission) {
+        this.requestAdmission = requestAdmission;
         this.config = config;
         this.durableExecutionArn = durableExecutionArn;
         this.callback = callback;
@@ -186,6 +214,15 @@ void shutdown() {
         checkpointApiRequestDelayedBatcher.shutdown();
     }
 
+    void shutdown(Duration timeout) throws InterruptedException, ExecutionException, TimeoutException {
+        stopPolling(new IllegalStateException("CheckpointManager shutdown"));
+        checkpointApiRequestDelayedBatcher.shutdown(timeout);
+    }
+
+    boolean isTerminalExecutionAcknowledged() {
+        return terminalExecutionAcknowledged;
+    }
+
     private void stopPolling(Throwable cause) {
         List<List<CompletableFuture<Operation>>> allFutures;
         synchronized (pollingFutures) {
@@ -204,6 +241,11 @@ private void stopPolling(Throwable cause) {
      * nextMarker)
      */
     List<Operation> fetchAllPages(CheckpointUpdatedExecutionState checkpointUpdatedExecutionState) {
+        return fetchAllPages(checkpointUpdatedExecutionState, false);
+    }
+
+    private List<Operation> fetchAllPages(
+            CheckpointUpdatedExecutionState checkpointUpdatedExecutionState, boolean terminal) {
         List<Operation> operations = new ArrayList<>();
         if (checkpointUpdatedExecutionState == null) {
             return operations;
@@ -213,6 +255,12 @@ List<Operation> fetchAllPages(CheckpointUpdatedExecutionState checkpointUpdatedE
         }
         var nextMarker = checkpointUpdatedExecutionState.nextMarker();
         while (nextMarker != null && !nextMarker.isEmpty()) {
+            try {
+                requestAdmission.run();
+            } catch (UnrecoverableDurableExecutionException failure) {
+                if (!terminal || !ExecutionManager.isLifecycleFailure(failure)) throw failure;
+                break; // Preserve inline/already-read metadata without issuing a new request after terminal ACK.
+            }
             var startTime = System.nanoTime();
             try {
                 var response = config.getDurableExecutionClient()
@@ -231,12 +279,15 @@ List<Operation> fetchAllPages(CheckpointUpdatedExecutionState checkpointUpdatedE
     }
 
     private void checkpointBatch(List<OperationUpdate> updates) {
-        synchronized (pollingFutures) {
-            // A READY recheck can complete a poll before its scheduled batch. Keep only consumers still waiting.
-            pollingFutures.entrySet().removeIf(entry -> {
-                entry.getValue().removeIf(CompletableFuture::isDone);
-                return entry.getValue().isEmpty();
-            });
+        synchronized (checkpointLock) {
+            synchronized (pollingFutures) {
+                // A READY recheck can complete a poll before its scheduled batch. Keep only consumers still waiting.
+                pollingFutures.entrySet().removeIf(entry -> {
+                    entry.getValue().removeIf(CompletableFuture::isDone);
+                    return entry.getValue().isEmpty();
+                });
+            }
+
             // filter the null values from pollers
             var request = updates.stream().filter(Objects::nonNull).toList();
 
@@ -277,7 +328,14 @@ private void checkpointBatch(List<OperationUpdate> updates) {
     private void checkpointRequest(List<OperationUpdate> request) {
         var startTime = System.nanoTime();
         logger.debug("Calling durable checkpoint API with {} updates: {}", request.size(), request);
+        requestAdmission.run();
         var response = config.getDurableExecutionClient().checkpoint(durableExecutionArn, checkpointToken, request);
+        if (request.stream()
+                .anyMatch(update ->
+                        update.type() == OperationType.EXECUTION && update.action() == OperationAction.SUCCEED))
+            terminalExecutionAcknowledged = true;
+        // A late response is durable at the service, but must not resume old SDK work after invocation cleanup.
+        if (!terminalExecutionAcknowledged) requestAdmission.run();
         logger.debug("Durable checkpoint API called (latency={}ns): {}.", System.nanoTime() - startTime, response);
 
         checkpointToken = response.checkpointToken();
@@ -288,7 +346,7 @@ private void checkpointRequest(List<OperationUpdate> request) {
 
         var state = response.newExecutionState();
         if (state != null) {
-            processUpdatedOperations(fetchAllPages(state));
+            processUpdatedOperations(fetchAllPages(state, terminalExecutionAcknowledged));
         }
     }
 
@@ -322,7 +380,10 @@ private void processUpdatedOperations(List<Operation> operations) {
         callback.accept(operations);
 
         for (var operation : operations) {
-            var pollers = pollingFutures.remove(operation.id());
+            List<CompletableFuture<Operation>> pollers;
+            synchronized (pollingFutures) {
+                pollers = pollingFutures.remove(operation.id());
+            }
             if (pollers != null) {
                 completedFutures += pollers.size();
                 pollers.forEach(poller -> poller.complete(operation));