fix: enforce invocation deadlines while awaiting owner cleanup

docs/advanced/configuration.md

@@ -95,3 +95,36 @@ 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 monotonic
+cutoffs for admission and its wait policy. Deadline exhaustion selects a retryable
+invocation interruption; it is not checkpointed as a business-operation failure.
+Catching an admission rejection cannot turn the stopped invocation into a
+successful fallback. Completed operation results and failures remain available for
+replay.
+
+Cleanup allowances determine when the SDK requests cooperative cancellation and
+stops ordinary admission. They do not make cancellation equivalent to actual task
+exit. Finalization still waits for admitted user work, continuations and in-flight
+checkpoint/poll transport, and runs the matching invocation End on its owner.
+The runtime caller waits for that owner and
+End to finish, including after a deadline or interruption. End is not moved to a
+completion-observer thread or deferred until after the SDK call returns.
+
+Interruption is recorded, remaining required cleanup still runs, and the interrupt
+flag is restored. A terminal execution result acknowledged by the
+service remains authoritative for a handled deadline cutoff. An accepted result
+does not suppress an otherwise-propagating delivery or metadata failure, and an
+available JVM fatal keeps its original identity and priority.
+
+These cutoffs do not guarantee bounded caller return or worker reclamation.
+Already-running application I/O, synchronous initial-state reads, custom clients,
+plugin hooks and other code that ignores interruption can keep the caller waiting
+for actual cleanup beyond the policy budget. An executor that runs the root inline,
+including a saturated `CallerRunsPolicy`, can also block the runtime caller before
+it reaches a timeout-aware wait. Use finite downstream I/O timeouts and executor
+configurations that allow progress. Cancellation cannot undo a request already
+accepted by the service; interrupted side effects still require the documented
+retry and idempotency strategy.

docs/design.md

@@ -151,10 +151,17 @@ public class MyHandler extends DurableHandler<Input, Output> {
 
 ### Thread Pool Architecture
 
+Each invocation tracks logical task outcomes separately from actual task exit.
+Deadline and cleanup cutoffs fence admission and request cooperative cancellation,
+while finalization still waits for admitted work and in-flight transport to finish.
+The matching invocation End runs on its owner, and the runtime caller waits for
+that owner and End before returning. Code that ignores interruption can therefore
+outlast the policy budget; see [invocation deadline cleanup](advanced/configuration.md#invocation-deadline-cleanup).
+
 The SDK uses two separate thread pools with distinct responsibilities:
 
 **User Executor (`DurableConfig.executorService`):**
-- Runs user-defined operations (the code passed to `ctx.step()` and `ctx.stepAsync()`)
+- Runs the root handler, steps, child-context user code, and map/parallel coordination
 - Configurable via `DurableConfig.builder().withExecutorService()`
 - Default: cached daemon thread pool
 

lmi-tests/README.md

@@ -19,6 +19,15 @@ those original snapshots and retain the existing zero-work and +16-thread limits
 Newly discovered request identities invalidate the warm cohort instead of moving
 its original baseline or final measurement to another request.
 
+The current deadline controls fence admission and request cooperative cancellation,
+but the SDK still awaits actual owner cleanup, admitted transport and the matching
+invocation End before returning. Non-cooperative work can therefore keep the
+runtime caller blocked beyond its deadline. This remains a partial change for
+[#726](https://github.com/aws/aws-durable-execution-sdk-java/issues/726), not proof
+that every desired LMI lifecycle bound is satisfied. The acceptance limits below
+remain unchanged; local tests do not replace a deployed cloud acceptance run, and
+eventual cleanup cannot turn an earlier lifecycle violation into a pass.
+
 ## Test design
 
 * Each fixture is a published durable function with LMI invocation concurrency
@@ -219,7 +228,7 @@ barrier diagnostics never authorize a retry. Every attempted cohort remains in
 the escape ledger, so a later successful cohort cannot hide an earlier pool escape.
 
 `timeouts/*.json` distinguishes server timeout logs from an SDK deadline
-cancellation that returns early with an invocation error. A client HTTP timeout
+cancellation reported as an invocation error after owner cleanup. A client HTTP timeout
 is a collection error. Deadline victims use asynchronous service invocation so a
 longer durable retry does not consume the driver HTTP timeout; invocation
 outcomes come from request-correlated runtime logs. Invocation request failures

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

@@ -24,6 +24,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.DurableExecutionPluginFactory;
@@ -325,10 +326,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/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");

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

@@ -19,6 +19,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;
@@ -147,10 +148,19 @@ public void onInvocationEnd(InvocationEndInfo info) {
     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/SerializationCancellationOtelIntegrationTest.java

@@ -0,0 +1,144 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.lambda.durable.otel;
+
+import static org.junit.jupiter.api.Assertions.*;
+import static org.mockito.Mockito.*;
+import static software.amazon.lambda.durable.otel.SpanAttributes.DURABLE_INVOCATION_STATUS;
+
+import com.amazonaws.services.lambda.runtime.Context;
+import io.opentelemetry.api.GlobalOpenTelemetry;
+import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter;
+import io.opentelemetry.sdk.trace.SdkTracerProvider;
+import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor;
+import java.time.Instant;
+import java.util.concurrent.*;
+import java.util.concurrent.atomic.*;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Timeout;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+import software.amazon.awssdk.services.lambda.model.*;
+import software.amazon.lambda.durable.DurableConfig;
+import software.amazon.lambda.durable.TypeToken;
+import software.amazon.lambda.durable.client.DurableExecutionClient;
+import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException;
+import software.amazon.lambda.durable.execution.DurableExecutor;
+import software.amazon.lambda.durable.model.DurableExecutionInput;
+import software.amazon.lambda.durable.serde.JacksonSerDes;
+
+@Timeout(20)
+class SerializationCancellationOtelIntegrationTest {
+    @BeforeEach
+    @AfterEach
+    void clearState() {
+        GlobalOpenTelemetry.resetForTest();
+        DeterministicIdGenerator.clearSharedStateForTest();
+        DurableSamplingDecision.clearSharedStateForTest();
+        OtelPluginAutoConfigurationState.resetInstalledForTest();
+    }
+
+    @ParameterizedTest
+    @ValueSource(booleans = {false, true})
+    void cancellationDuringSmallResultSerializationDoesNotCompleteWorkflow(boolean executionView) throws Exception {
+        var exporter = InMemorySpanExporter.create();
+        var provider = SdkTracerProvider.builder().addSpanProcessor(SimpleSpanProcessor.create(exporter));
+        var pluginConfig = OtelPluginConfig.builder()
+                .contextExtractor(() -> null)
+                .enableMdc(false)
+                .build();
+        var factory = executionView
+                ? ExecutionOtelPlugin.factory(provider, pluginConfig)
+                : InvocationOtelPlugin.factory(provider, pluginConfig);
+        var serializationEntered = new CountDownLatch(1);
+        var serializationInterrupted = new CountDownLatch(1);
+        var releaseSerialization = new CountDownLatch(1);
+        var caller = new AtomicReference<Thread>();
+        var callerInterrupted = new AtomicBoolean();
+        var workers = Executors.newSingleThreadExecutor();
+        var callers = Executors.newSingleThreadExecutor();
+        var serializer = new JacksonSerDes() {
+            @Override
+            public String serialize(Object value) {
+                if ("result".equals(value)) {
+                    serializationEntered.countDown();
+                    var deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5);
+                    while (true) {
+                        try {
+                            assertTrue(releaseSerialization.await(
+                                    Math.max(0, deadline - System.nanoTime()), TimeUnit.NANOSECONDS));
+                            break;
+                        } catch (InterruptedException ignored) {
+                            serializationInterrupted.countDown();
+                        }
+                    }
+                }
+                return super.serialize(value);
+            }
+        };
+        var config = DurableConfig.builder()
+                .withExecutorService(workers)
+                .withDurableExecutionClient(mock(DurableExecutionClient.class))
+                .withSerDes(serializer)
+                .withPlugins(factory)
+                .build();
+        var runtime = mock(Context.class);
+        when(runtime.getRemainingTimeInMillis()).thenReturn(10_000);
+        when(runtime.getAwsRequestId()).thenReturn("serialization-cancellation");
+        var input = new DurableExecutionInput(
+                "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/otel/exec",
+                "token",
+                CheckpointUpdatedExecutionState.builder()
+                        .operations(Operation.builder()
+                                .id("exec")
+                                .type(OperationType.EXECUTION)
+                                .status(OperationStatus.STARTED)
+                                .startTimestamp(Instant.now())
+                                .executionDetails(ExecutionDetails.builder()
+                                        .inputPayload("\"input\"")
+                                        .build())
+                                .build())
+                        .build());
+        try {
+            var response = callers.submit(() -> {
+                caller.set(Thread.currentThread());
+                try {
+                    return DurableExecutor.execute(
+                            input, runtime, TypeToken.get(String.class), (value, context) -> "result", config);
+                } finally {
+                    callerInterrupted.set(Thread.interrupted());
+                }
+            });
+            assertTrue(serializationEntered.await(5, TimeUnit.SECONDS));
+            caller.get().interrupt();
+            assertTrue(serializationInterrupted.await(5, TimeUnit.SECONDS));
+            assertThrows(TimeoutException.class, () -> response.get(100, TimeUnit.MILLISECONDS));
+            releaseSerialization.countDown();
+            var failure = assertInstanceOf(
+                    UnrecoverableDurableExecutionException.class,
+                    assertThrows(ExecutionException.class, () -> response.get(5, TimeUnit.SECONDS))
+                            .getCause());
+            assertTrue(failure.isRetryable());
+            assertTrue(callerInterrupted.get());
+            var spans = exporter.getFinishedSpanItems();
+            var invocations = spans.stream()
+                    .filter(span -> span.getName().equals("Invocation"))
+                    .toList();
+            assertEquals(1, invocations.size(), "The real plugin must still export the completed invocation");
+            assertAll(
+                    () -> assertEquals(
+                            "RETRYING", invocations.get(0).getAttributes().get(DURABLE_INVOCATION_STATUS)),
+                    () -> assertTrue(
+                            spans.stream().noneMatch(span -> span.getName().equals("Workflow")),
+                            "A cancelled preparation must not emit a terminal successful Workflow span"));
+        } finally {
+            releaseSerialization.countDown();
+            workers.shutdownNow();
+            callers.shutdownNow();
+            assertTrue(workers.awaitTermination(5, TimeUnit.SECONDS));
+            assertTrue(callers.awaitTermination(5, TimeUnit.SECONDS));
+            exporter.close();
+        }
+    }
+}

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

@@ -11,7 +11,9 @@
 import java.util.concurrent.CompletionException;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 import java.util.function.Function;
@@ -28,7 +30,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;
@@ -134,6 +136,12 @@ void shutdown() {
         if (fatal != null) throw new CompletionException(fatal);
     }
 
+    /** 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) {
@@ -145,8 +153,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<>();