From b49e62146808070e75695dd4ce6580b01ca489e5 Mon Sep 17 00:00:00 2001 From: Norm Johanson Date: Tue, 3 Nov 2026 02:20:31 +0000 Subject: [PATCH 1/6] Support Lambda Managed Instances in Amazon.Lambda.DurableExecution - Cancel the workflow token synchronously at teardown, before disposal. Termination delivered the cancel via a threadpool continuation that could run after disposal and be silently dropped. - Clean up abandoned work: keep the workflow token readable after dispose, reject new durable operations once the invocation ends, and observe the abandoned workflow task's fault. - Lambda Managed Instances only: apply the invocation deadline from RemainingTime. Cancel workflow work and state hydration at the deadline, bound checkpoint cleanup, and throw TimeoutException so the invocation is retried rather than recorded as a durable failure. On-demand behavior is unchanged. - Add multi-concurrency unit tests for isolation, peer timeouts, replay alongside a live peer, and teardown ordering. Refs #2604 --- .../46bb3a38-e1f3-48d1-ab9b-10a36475e901.json | 12 + .../DurableContext.cs | 9 + .../DurableExecutionHandler.cs | 71 +- .../DurableFunction.cs | 106 ++- .../Internal/CheckpointBatcher.cs | 54 +- .../Internal/DurableInvocationOptions.cs | 26 + .../Internal/WorkflowCancellation.cs | 197 +++++- .../docs/core/cancellation.md | 13 + .../LambdaManagedInstancesTests.cs | 644 ++++++++++++++++++ 9 files changed, 1099 insertions(+), 33 deletions(-) create mode 100644 .autover/changes/46bb3a38-e1f3-48d1-ab9b-10a36475e901.json create mode 100644 Libraries/src/Amazon.Lambda.DurableExecution/Internal/DurableInvocationOptions.cs create mode 100644 Libraries/test/Amazon.Lambda.DurableExecution.Tests/LambdaManagedInstancesTests.cs diff --git a/.autover/changes/46bb3a38-e1f3-48d1-ab9b-10a36475e901.json b/.autover/changes/46bb3a38-e1f3-48d1-ab9b-10a36475e901.json new file mode 100644 index 000000000..4eee32762 --- /dev/null +++ b/.autover/changes/46bb3a38-e1f3-48d1-ab9b-10a36475e901.json @@ -0,0 +1,12 @@ +{ + "Projects": [ + { + "Name": "Amazon.Lambda.DurableExecution", + "Type": "Patch", + "ChangelogMessages": [ + "Cancel user work that is still running when an invocation ends, and reject new durable operations started after that point.", + "On Lambda Managed Instances, stop durable work at the invocation deadline and throw TimeoutException so the invocation is retried instead of left running." + ] + } + ] +} diff --git a/Libraries/src/Amazon.Lambda.DurableExecution/DurableContext.cs b/Libraries/src/Amazon.Lambda.DurableExecution/DurableContext.cs index 11530af7d..d2444c3a8 100644 --- a/Libraries/src/Amazon.Lambda.DurableExecution/DurableContext.cs +++ b/Libraries/src/Amazon.Lambda.DurableExecution/DurableContext.cs @@ -84,6 +84,7 @@ private Task RunStep( StepConfig? config, CancellationToken cancellationToken) { + _workflowCancellation.ThrowIfTornDown(); var defaultSerializer = LambdaSerializerHelper.GetRequired(LambdaContext); var serializer = LambdaSerializerHelper.WithDefaultInner(config?.Serializer ?? defaultSerializer, defaultSerializer); @@ -99,6 +100,7 @@ public Task WaitAsync( string? name = null, CancellationToken cancellationToken = default) { + _workflowCancellation.ThrowIfTornDown(); // Service timer granularity is 1 second; sub-second waits would round to 0. // WaitOptions.WaitSeconds is integer in [1, 31_622_400] (1 second to ~1 year). if (duration < TimeSpan.FromSeconds(1)) @@ -144,6 +146,7 @@ public Task WaitForConditionAsync( string? name = null, CancellationToken cancellationToken = default) { + _workflowCancellation.ThrowIfTornDown(); ArgumentNullException.ThrowIfNull(check); ArgumentNullException.ThrowIfNull(config); ArgumentNullException.ThrowIfNull(config.WaitStrategy); @@ -163,6 +166,7 @@ private Task RunChildContext( ChildContextConfig? config, CancellationToken cancellationToken) { + _workflowCancellation.ThrowIfTornDown(); var defaultSerializer = LambdaSerializerHelper.GetRequired(LambdaContext); var serializer = LambdaSerializerHelper.WithDefaultInner(config?.Serializer ?? defaultSerializer, defaultSerializer); @@ -186,6 +190,7 @@ private Task> RunCallback( CallbackConfig? config, CancellationToken cancellationToken) { + _workflowCancellation.ThrowIfTornDown(); var serializer = config?.Serializer ?? LambdaSerializerHelper.GetRequired(LambdaContext); var operationId = _idGenerator.NextId(); @@ -199,6 +204,7 @@ public IDurableParallel CreateParallel( string? name = null, ParallelConfig? config = null) { + _workflowCancellation.ThrowIfTornDown(); var effectiveConfig = config ?? new ParallelConfig(); // Operation-level default for per-branch result serialization: the config's // ItemSerializer if set, else the globally-registered serializer. Individual @@ -266,6 +272,7 @@ private Task> RunParallel( ParallelConfig? config, CancellationToken cancellationToken) { + _workflowCancellation.ThrowIfTornDown(); if (branches == null) throw new ArgumentNullException(nameof(branches)); for (var i = 0; i < branches.Count; i++) { @@ -306,6 +313,7 @@ private Task> RunMap( MapConfig? config, CancellationToken cancellationToken) { + _workflowCancellation.ThrowIfTornDown(); if (items == null) throw new ArgumentNullException(nameof(items)); if (func == null) throw new ArgumentNullException(nameof(func)); @@ -528,6 +536,7 @@ private Task RunInvoke( InvokeConfig? config, CancellationToken cancellationToken) { + _workflowCancellation.ThrowIfTornDown(); // Argument validation runs synchronously at the call site (matches the // .NET convention of failing fast for misuse). Match Python/JS/Java // parity: only check for null/empty here; the durable execution service diff --git a/Libraries/src/Amazon.Lambda.DurableExecution/DurableExecutionHandler.cs b/Libraries/src/Amazon.Lambda.DurableExecution/DurableExecutionHandler.cs index 971ac6f64..c205a8462 100644 --- a/Libraries/src/Amazon.Lambda.DurableExecution/DurableExecutionHandler.cs +++ b/Libraries/src/Amazon.Lambda.DurableExecution/DurableExecutionHandler.cs @@ -14,6 +14,13 @@ internal sealed class HandlerResult public TResult? Result { get; init; } public string? Message { get; init; } public Exception? Exception { get; init; } + + /// + /// True when the invocation deadline passed before the workflow completed. + /// is meaningless in that case: the wrapper fails the + /// invocation so Lambda retries it, rather than recording a durable outcome. + /// + public bool DeadlineExceeded { get; init; } } /// @@ -23,6 +30,9 @@ internal sealed class HandlerResult /// internal static class DurableExecutionHandler { + // Stands in for the deadline when the caller passes no WorkflowCancellation. + private static readonly Task NeverCompletes = new TaskCompletionSource().Task; + /// /// Runs the user's workflow function within the durable execution engine. /// @@ -49,8 +59,10 @@ internal static class DurableExecutionHandler /// Key insight: WaitAsync never returns a completed Task — it hands back /// a TaskCompletionSource that is never resolved. The user's await blocks /// indefinitely. The escape signal is terminationManager.Terminate(), - /// which Task.WhenAny picks up. We return Pending; the dangling user - /// Task is GC'd. The service flushes checkpoints, fires the wait timer, then + /// which Task.WhenAny picks up. We return Pending and observe the + /// abandoned user Task's eventual fault; then + /// cancels the workflow token so any user work still running unwinds. The + /// service flushes checkpoints, fires the wait timer, then /// re-invokes Lambda — on replay, WaitAsync sees the matching SUCCEED /// checkpoint and returns Task.CompletedTask normally. /// @@ -63,11 +75,16 @@ internal static class DurableExecutionHandler /// Hydrated execution state from prior invocations. /// Manages the suspension signal. /// The user's workflow function receiving a DurableContext. + /// + /// Optional workflow cancellation for this invocation. When it carries a + /// deadline (Lambda Managed Instances only), the deadline joins the race. + /// /// The handler result indicating SUCCEEDED, FAILED, or PENDING. internal static async Task> RunAsync( ExecutionState executionState, TerminationManager terminationManager, - Func> userHandler) + Func> userHandler, + WorkflowCancellation? workflowCancellation = null) { // Run user code on a threadpool thread so it executes independently of // the termination signal. When TerminationManager fires (e.g., WaitAsync), @@ -75,9 +92,32 @@ internal static async Task> RunAsync( // for the user task to reach an await point. var userTask = Task.Run(userHandler); - // Race: user code completing vs. termination signal (wait/callback/retry). - // If termination wins, we return PENDING and the abandoned userTask is never awaited. - var winner = await Task.WhenAny(userTask, terminationManager.TerminationTask); + // Race: user code completing vs. termination signal (wait/callback/retry) + // vs. the invocation deadline. DeadlineTask never completes when there is + // no deadline, so on-demand Lambda races only the first two. + var deadlineTask = workflowCancellation?.DeadlineTask ?? NeverCompletes; + var winner = await Task.WhenAny(userTask, terminationManager.TerminationTask, deadlineTask); + + if (winner != userTask) + { + // userTask is abandoned. Observe its eventual fault so it never + // surfaces through TaskScheduler.UnobservedTaskException, which in a + // shared process (Lambda Managed Instances) would be reported against + // whichever invocation happens to be running when the GC finalizes it. + ObserveAbandoned(userTask); + } + + // The deadline wins over a termination that raced it: the invocation has + // run out of time either way, and the service must re-invoke it. + if (winner == deadlineTask + || (winner == terminationManager.TerminationTask && workflowCancellation?.IsDeadlineExceeded == true)) + { + return new HandlerResult + { + Status = InvocationStatus.Failed, + DeadlineExceeded = true, + }; + } if (winner == terminationManager.TerminationTask) { @@ -111,6 +151,18 @@ internal static async Task> RunAsync( } catch (Exception ex) { + // A workflow that faulted because the deadline cancelled it did not + // really fail. Report the deadline instead, so the invocation is + // retried rather than the execution being failed permanently. + if (workflowCancellation?.IsDeadlineExceeded == true) + { + return new HandlerResult + { + Status = InvocationStatus.Failed, + DeadlineExceeded = true, + }; + } + return new HandlerResult { Status = InvocationStatus.Failed, @@ -119,4 +171,11 @@ internal static async Task> RunAsync( }; } } + + private static void ObserveAbandoned(Task task) => + task.ContinueWith( + static t => _ = t.Exception, + CancellationToken.None, + TaskContinuationOptions.OnlyOnFaulted | TaskContinuationOptions.ExecuteSynchronously, + TaskScheduler.Default); } diff --git a/Libraries/src/Amazon.Lambda.DurableExecution/DurableFunction.cs b/Libraries/src/Amazon.Lambda.DurableExecution/DurableFunction.cs index 895e96631..d9d95b3ec 100644 --- a/Libraries/src/Amazon.Lambda.DurableExecution/DurableFunction.cs +++ b/Libraries/src/Amazon.Lambda.DurableExecution/DurableFunction.cs @@ -1,6 +1,7 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 +using System.Diagnostics; using System.IO; using System.Text; using System.Threading; @@ -83,33 +84,64 @@ internal static Task WrapAsync WrapAsyncCore(workflow, invocationInput, lambdaContext, serviceClient); + => WrapAsyncCore(workflow, invocationInput, lambdaContext, serviceClient, options: null); + + /// + /// Internal overload for tests that override + /// (Lambda Managed Instances detection and the cleanup budget). + /// + internal static Task WrapAsync( + Func> workflow, + DurableExecutionInvocationInput invocationInput, + ILambdaContext lambdaContext, + IDurableServiceClient serviceClient, + DurableInvocationOptions options) + => WrapAsyncCore(workflow, invocationInput, lambdaContext, serviceClient, options); private static async Task WrapAsyncCore( Func> workflow, DurableExecutionInvocationInput invocationInput, ILambdaContext lambdaContext, - IDurableServiceClient serviceClient) + IDurableServiceClient serviceClient, + DurableInvocationOptions? options = null) { var serializer = LambdaSerializerHelper.GetRequired(lambdaContext); + // Lambda Managed Instances only: run the invocation against its deadline. + // On-demand Lambda gets no deadline, so every wait below is unbounded as + // before and the environment's own timeout stops the invocation. + var startTimestamp = Stopwatch.GetTimestamp(); + var deadline = InvocationDeadline.Resolve( + lambdaContext.RemainingTime, + options?.IsMultiConcurrency ?? InvocationDeadline.IsMultiConcurrency()); + var cleanupTimeout = options?.CleanupTimeout ?? DurableInvocationOptions.DefaultCleanupTimeout; + + var terminationManager = new TerminationManager(); + using var workflowCancellation = new WorkflowCancellation(terminationManager, deadline); + var state = new ExecutionState(); state.LoadFromCheckpoint(invocationInput.InitialExecutionState); var checkpointToken = invocationInput.CheckpointToken; var nextMarker = invocationInput.InitialExecutionState?.NextMarker; - while (!string.IsNullOrEmpty(nextMarker)) + try + { + while (!string.IsNullOrEmpty(nextMarker)) + { + var (operations, marker) = await serviceClient.GetExecutionStateAsync( + invocationInput.DurableExecutionArn, checkpointToken, nextMarker, + workflowCancellation.DeadlineToken); + state.AddOperations(operations); + nextMarker = marker; + } + } + catch (OperationCanceledException ex) when (workflowCancellation.IsDeadlineExceeded) { - var (operations, marker) = await serviceClient.GetExecutionStateAsync( - invocationInput.DurableExecutionArn, checkpointToken, nextMarker); - state.AddOperations(operations); - nextMarker = marker; + throw CreateDeadlineException(invocationInput.DurableExecutionArn, ex); } var userPayload = ExtractUserPayload(invocationInput, serializer); - var terminationManager = new TerminationManager(); - using var workflowCancellation = new WorkflowCancellation(terminationManager); var idGenerator = new OperationIdGenerator(); await using var batcher = new CheckpointBatcher( @@ -145,15 +177,44 @@ private static async Task WrapAsyncCore( state, terminationManager, - async () => await workflow(userPayload, context)); + async () => await workflow(userPayload, context), + workflowCancellation); + } + + // The invocation's outcome is decided. Cancel the workflow token now, + // synchronously, so user work abandoned by a suspension (a losing + // parallel branch, a step body still running) unwinds before the + // drain below. Waiting for TerminationManager's threadpool + // continuation instead could let it run after disposal, losing the + // cancel. From here on, new durable operations are rejected. + workflowCancellation.BeginTeardown(); + + if (deadline is null) + { + await batcher.DrainAsync(); + } + else + { + // Lambda Managed Instances: a stalled checkpoint call must not hold + // this invocation (and one of the runtime's processing loops) open + // indefinitely. Allow time up to the deadline plus a fixed cleanup + // budget, then abandon the flush. + var untilDeadline = deadline.Value - Stopwatch.GetElapsedTime(startTimestamp); + var drainTimeout = (untilDeadline > TimeSpan.Zero ? untilDeadline : TimeSpan.Zero) + cleanupTimeout; + if (!await batcher.DrainAsync(drainTimeout)) + throw CreateDeadlineException(invocationInput.DurableExecutionArn, innerException: null); } - await batcher.DrainAsync(); + if (result.DeadlineExceeded) + throw CreateDeadlineException(invocationInput.DurableExecutionArn, innerException: null); if (batcher.IsCheckpointTokenMissing) result = new HandlerResult { Status = InvocationStatus.Pending }; } - catch (DurableExecutionException ex) when (ex.InnerException is AmazonServiceException sdkEx && IsTerminalCheckpointError(sdkEx)) + catch (DurableExecutionException ex) when ( + !workflowCancellation.IsDeadlineExceeded + && ex.InnerException is AmazonServiceException sdkEx + && IsTerminalCheckpointError(sdkEx)) { return new DurableExecutionInvocationOutput { @@ -161,10 +222,29 @@ private static async Task WrapAsyncCore + /// The exception thrown when a Lambda Managed Instances invocation runs out + /// of time. It escapes to the Lambda runtime as an invocation error, the same + /// outcome as an on-demand timeout, so the durable execution is re-invoked and + /// replays from its last checkpoint instead of being marked failed. + /// + private static TimeoutException CreateDeadlineException(string durableExecutionArn, Exception? innerException) => + new($"The invocation for durable execution '{durableExecutionArn}' reached its deadline before it completed. " + + "Durable work for this invocation was cancelled; work that was not checkpointed runs again on replay.", + innerException); + /// /// Returns true for checkpoint-flush SDK errors that should fail the workflow /// (Failed envelope) instead of escaping to the host (Lambda retry). The catch @@ -187,7 +267,7 @@ private static async Task WrapAsyncCoreStepAsync call /// (the user awaits EnqueueAsync → batch flush → SDK throws → service client /// wraps). - /// 2. The final after the workflow returns. + /// 2. The final after the workflow returns. /// /// State-hydration errors (GetExecutionStateAsync) propagate as /// too, but they are NOT caught here — they diff --git a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/CheckpointBatcher.cs b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/CheckpointBatcher.cs index 588263096..f08101fda 100644 --- a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/CheckpointBatcher.cs +++ b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/CheckpointBatcher.cs @@ -16,7 +16,7 @@ namespace Amazon.Lambda.DurableExecution.Internal; /// /// Fire-and-forget semantics are achieved by simply not awaiting the returned /// Task. Errors still surface deterministically via _terminalError: the -/// next sync or rethrows. +/// next sync or rethrows. /// Callers using fire-and-forget should observe the discarded Task's exception /// (see StepOperation.FireAndForget) so it doesn't trip the runtime's /// UnobservedTaskException event. @@ -34,6 +34,7 @@ internal sealed class CheckpointBatcher : IAsyncDisposable private string? _checkpointToken; private Exception? _terminalError; private int _disposed; + private int _abandoned; private int _tokenMissing; // Per-update wire-footprint estimate constants. Deliberate over-estimates: @@ -139,11 +140,62 @@ public async Task DrainAsync() if (terminal != null) ExceptionDispatchInfo.Throw(terminal); } + /// + /// Like , but gives up after . + /// Used for Lambda Managed Instances, where a stalled checkpoint call must not + /// hold the invocation open past its deadline. + /// + /// + /// true when the worker finished. false when the timeout elapsed: + /// the in-flight flush is cancelled and the worker is abandoned, so any + /// checkpoint it had not yet sent is not persisted (replay re-runs that work). + /// + public async Task DrainAsync(TimeSpan timeout) + { + _channel.Writer.TryComplete(); + try + { + await _worker.WaitAsync(timeout).ConfigureAwait(false); + } + catch (TimeoutException) + { + Abandon(); + return false; + } + catch + { + // Surfaced via _terminalError below. + } + + var terminal = Volatile.Read(ref _terminalError); + if (terminal != null) ExceptionDispatchInfo.Throw(terminal); + return true; + } + + private void Abandon() + { + if (Interlocked.Exchange(ref _abandoned, 1) != 0) return; + + // Ask the in-flight flush to stop. A flush that ignores the token keeps + // running until it returns; its outcome is no longer awaited. + try { _shutdownCts.Cancel(); } catch (ObjectDisposedException) { } + _worker.ContinueWith( + static t => _ = t.Exception, + CancellationToken.None, + TaskContinuationOptions.OnlyOnFaulted | TaskContinuationOptions.ExecuteSynchronously, + TaskScheduler.Default); + } + public async ValueTask DisposeAsync() { if (Interlocked.Exchange(ref _disposed, 1) != 0) return; _channel.Writer.TryComplete(); + + // An abandoned worker may never finish. Don't wait for it, and leave its + // cancellation source undisposed because the worker may still read the token. + if (Volatile.Read(ref _abandoned) == 1) return; + _shutdownCts.Cancel(); try { await _worker.ConfigureAwait(false); } catch { /* swallow on dispose */ } diff --git a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/DurableInvocationOptions.cs b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/DurableInvocationOptions.cs new file mode 100644 index 000000000..cf21405bf --- /dev/null +++ b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/DurableInvocationOptions.cs @@ -0,0 +1,26 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 + +namespace Amazon.Lambda.DurableExecution.Internal; + +/// +/// Per-invocation settings for that tests override. +/// Production code uses the defaults. +/// +internal sealed class DurableInvocationOptions +{ + /// + /// How long checkpoint cleanup may run after the invocation deadline before + /// the SDK gives up on it. Only applies to Lambda Managed Instances. + /// + public static readonly TimeSpan DefaultCleanupTimeout = TimeSpan.FromSeconds(2); + + /// + /// Overrides detection of Lambda Managed Instances. null reads the + /// AWS_LAMBDA_MAX_CONCURRENCY environment variable. + /// + public bool? IsMultiConcurrency { get; init; } + + /// See . + public TimeSpan CleanupTimeout { get; init; } = DefaultCleanupTimeout; +} diff --git a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/WorkflowCancellation.cs b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/WorkflowCancellation.cs index 115dd585e..9d248d5e2 100644 --- a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/WorkflowCancellation.cs +++ b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/WorkflowCancellation.cs @@ -3,12 +3,31 @@ namespace Amazon.Lambda.DurableExecution.Internal; +/// +/// Why the workflow-scoped token was cancelled. The first reason wins. +/// +internal enum WorkflowCancellationReason +{ + None = 0, + + /// The resolved (suspend or checkpoint failure). + Terminated = 1, + + /// The invocation deadline passed (Lambda Managed Instances only). + DeadlineExceeded = 2, + + /// The wrapper finished the invocation and is tearing it down. + Teardown = 3, +} + /// /// Workflow-scoped cancellation source. Cancels when the -/// resolves so abandoned user-Func bodies -/// (the WhenAny loser in ) unwind via +/// resolves, when the optional invocation +/// deadline passes, or when tears the invocation +/// down. Abandoned user-Func bodies (the WhenAny loser in +/// ) unwind via /// instead of running to completion on -/// the threadpool while Lambda is mid-response. +/// the threadpool. /// /// /// One instance per durable function invocation, constructed and disposed by @@ -20,27 +39,179 @@ namespace Amazon.Lambda.DurableExecution.Internal; /// observe this token: successful work must persist even when the workflow is /// being torn down. /// +/// +/// Lifetime rules that matter when several invocations share one process +/// (Lambda Managed Instances): +/// +/// +/// and cancel the token +/// synchronously before releasing anything. Termination delivers its cancel +/// through a threadpool continuation, which could otherwise run after +/// disposal and be lost. +/// is captured at construction, so abandoned user code +/// that reads it after disposal gets a cancelled token instead of an +/// . +/// /// internal sealed class WorkflowCancellation : IDisposable { private readonly CancellationTokenSource _cts = new(); + private readonly CancellationToken _token; + private readonly CancellationTokenSource? _deadlineCts; + private readonly CancellationTokenRegistration _deadlineRegistration; + private readonly TaskCompletionSource _deadlineTcs = new(TaskCreationOptions.RunContinuationsAsynchronously); + private int _reason; + private int _tornDown; + private int _disposed; + + public CancellationToken Token => _token; - public CancellationToken Token => _cts.Token; + /// + /// The reason the token was cancelled, or . + /// + public WorkflowCancellationReason Reason => (WorkflowCancellationReason)Volatile.Read(ref _reason); - public WorkflowCancellation(TerminationManager terminationManager) + /// True once the invocation deadline passed. + public bool IsDeadlineExceeded => Reason == WorkflowCancellationReason.DeadlineExceeded; + + /// + /// True once the invocation is being torn down (the wrapper finished or the + /// deadline passed). New durable operations are rejected from this point. + /// + public bool IsTornDown => Volatile.Read(ref _tornDown) == 1; + + /// + /// Cancels only when the invocation deadline passes. Passed to SDK-owned + /// service calls (state hydration). when + /// there is no deadline, so on-demand Lambda calls are unchanged. + /// + public CancellationToken DeadlineToken { get; } + + /// + /// Completes when the invocation deadline passes. Never completes when there + /// is no deadline. + /// + public Task DeadlineTask => _deadlineTcs.Task; + + /// Suspension signal for this invocation. + /// + /// Time until the invocation deadline, or null for no deadline. Only + /// set for Lambda Managed Instances; see . + /// + public WorkflowCancellation(TerminationManager terminationManager, TimeSpan? deadline = null) { + _token = _cts.Token; + terminationManager.TerminationTask.ContinueWith( - static (_, state) => - { - var cts = (CancellationTokenSource)state!; - try { cts.Cancel(); } - catch (ObjectDisposedException) { } - }, - _cts, + static (_, state) => ((WorkflowCancellation)state!).Cancel(WorkflowCancellationReason.Terminated), + this, CancellationToken.None, TaskContinuationOptions.ExecuteSynchronously, TaskScheduler.Default); + + if (deadline is { } remaining) + { + _deadlineCts = new CancellationTokenSource(remaining); + DeadlineToken = _deadlineCts.Token; + _deadlineRegistration = _deadlineCts.Token.Register( + static state => ((WorkflowCancellation)state!).OnDeadline(), + this); + } + } + + /// + /// Marks the invocation as torn down and cancels the token synchronously, + /// so the cancel is delivered before the caller releases any resources. + /// + public void BeginTeardown() + { + Volatile.Write(ref _tornDown, 1); + Cancel(WorkflowCancellationReason.Teardown); + } + + /// + /// Throws when the invocation is + /// torn down, so abandoned user code cannot start new durable operations. + /// + public void ThrowIfTornDown() + { + if (IsTornDown) + throw new OperationCanceledException( + "The durable invocation has ended; no new durable operations can be started.", _token); + } + + private void OnDeadline() + { + // Record the reason and stop admitting new operations before cancelling + // the token, so code that observes the cancellation also sees why. + Interlocked.CompareExchange(ref _reason, (int)WorkflowCancellationReason.DeadlineExceeded, 0); + Volatile.Write(ref _tornDown, 1); + Cancel(WorkflowCancellationReason.DeadlineExceeded); + _deadlineTcs.TrySetResult(); + } + + private void Cancel(WorkflowCancellationReason reason) + { + Interlocked.CompareExchange(ref _reason, (int)reason, 0); + try { _cts.Cancel(); } + catch (ObjectDisposedException) + { + // Only reachable from the late termination continuation or deadline + // timer. Dispose cancels before disposing, so the token is already + // cancelled. + } + } + + public void Dispose() + { + if (Interlocked.Exchange(ref _disposed, 1) != 0) return; + + BeginTeardown(); + _deadlineRegistration.Dispose(); + _deadlineCts?.Dispose(); + _cts.Dispose(); } +} - public void Dispose() => _cts.Dispose(); +/// +/// Decides whether an invocation gets a deadline that cancels durable work. +/// +/// +/// Only Lambda Managed Instances get a deadline. There, several invocations +/// share one process and Lambda does not stop code when an invocation times out, +/// so the SDK has to stop its own work. In on-demand Lambda the execution +/// environment is stopped at the timeout, so the SDK keeps its existing +/// behavior and never cancels on a deadline. +/// +internal static class InvocationDeadline +{ + /// + /// Environment variable Lambda sets for functions running on Lambda Managed + /// Instances (multi-concurrency mode). + /// + internal const string MaxConcurrencyEnvironmentVariable = "AWS_LAMBDA_MAX_CONCURRENCY"; + + /// True when the function runs on Lambda Managed Instances. + public static bool IsMultiConcurrency() => + !string.IsNullOrEmpty(Environment.GetEnvironmentVariable(MaxConcurrencyEnvironmentVariable)); + + /// + /// Returns the time until the deadline, or null for no deadline. + /// + /// The invocation's ILambdaContext.RemainingTime. + /// Whether the function runs on Lambda Managed Instances. + /// + /// A zero or negative remaining time means no deadline. That is what an + /// unconfigured test context reports, and it must not cancel immediately. + /// + public static TimeSpan? Resolve(TimeSpan remainingTime, bool isMultiConcurrency) + { + if (!isMultiConcurrency) return null; + if (remainingTime <= TimeSpan.Zero) return null; + + // CancellationTokenSource rejects delays above int.MaxValue milliseconds. + // Lambda timeouts are far below that; treat anything larger as no deadline. + if (remainingTime.TotalMilliseconds >= int.MaxValue) return null; + return remainingTime; + } } diff --git a/Libraries/src/Amazon.Lambda.DurableExecution/docs/core/cancellation.md b/Libraries/src/Amazon.Lambda.DurableExecution/docs/core/cancellation.md index 1abc75858..e8104f37f 100644 --- a/Libraries/src/Amazon.Lambda.DurableExecution/docs/core/cancellation.md +++ b/Libraries/src/Amazon.Lambda.DurableExecution/docs/core/cancellation.md @@ -22,6 +22,19 @@ When either trigger fires, the token transitions to `IsCancellationRequested = t - **`OperationCanceledException` thrown out of a step body via the token** (i.e. `linked.IsCancellationRequested` is true) is treated as cancellation: no FAIL checkpoint is written, no retry is consulted. The exception propagates up. - **`OperationCanceledException` thrown by user code for unrelated reasons** (token never fired) is treated as a normal step failure: FAIL checkpoint, retry per the configured `RetryStrategy`. - **The SDK's own writes (checkpoints, batcher flush, the runtime API response)** never observe the workflow-shutdown signal. Successful work is never lost to teardown. +- **When the invocation ends**, the token is cancelled before `DurableFunction.WrapAsync` returns. Work the workflow started but did not await (for example a step still running when a sibling suspended) observes the cancellation. Starting a new durable operation after this point throws `OperationCanceledException`. + +## Lambda Managed Instances + +On Lambda Managed Instances several invocations run in the same process, and Lambda does not stop an invocation that times out. The SDK therefore applies the invocation deadline (`ILambdaContext.RemainingTime`) itself: + +- At the deadline the token is cancelled, new durable operations are rejected, and state loading is cancelled. +- Checkpoint cleanup gets a short budget past the deadline, then is abandoned. +- `WrapAsync` throws `TimeoutException`. Lambda treats this as an invocation error, the same outcome as a timeout in on-demand Lambda: the execution is not marked failed, and work that was not checkpointed runs again on replay. + +Only cooperative code can be stopped. Code that ignores the token, or blocks synchronously, keeps running until it returns, and keeps a thread busy that other invocations could use. + +On-demand Lambda is unchanged: the SDK applies no deadline, and the environment's timeout ends the invocation. ## Guidance diff --git a/Libraries/test/Amazon.Lambda.DurableExecution.Tests/LambdaManagedInstancesTests.cs b/Libraries/test/Amazon.Lambda.DurableExecution.Tests/LambdaManagedInstancesTests.cs new file mode 100644 index 000000000..7c608c382 --- /dev/null +++ b/Libraries/test/Amazon.Lambda.DurableExecution.Tests/LambdaManagedInstancesTests.cs @@ -0,0 +1,644 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 + +using System.Collections.Concurrent; +using Amazon.Lambda.DurableExecution.Internal; +using Amazon.Lambda.DurableExecution.Services; +using Amazon.Lambda.Serialization.SystemTextJson; +using Amazon.Lambda.TestUtilities; +using Xunit; +using SdkOperationUpdate = Amazon.Lambda.Model.OperationUpdate; + +namespace Amazon.Lambda.DurableExecution.Tests; + +/// +/// Tests for running durable functions on Lambda Managed Instances (LMI), where +/// several invocations run in one process at the same time and Lambda does not +/// stop an invocation that times out. +/// +/// +/// These tests drive directly with an in-memory +/// service client, running invocations concurrently the way RuntimeSupport does +/// in multi-concurrency mode. LMI mode is selected through +/// rather than the +/// process-wide environment variable, so these tests can run in parallel. +/// Ordering is controlled with signals, never with sleeps; the short deadlines +/// used below only need to expire eventually, and no test asserts on elapsed time. +/// +public class LambdaManagedInstancesTests +{ + // Guards against a hang turning into a stuck test run. Not a timing assertion. + private static readonly TimeSpan HangGuard = TimeSpan.FromSeconds(30); + + private static readonly DurableInvocationOptions LmiOptions = new() { IsMultiConcurrency = true }; + private static readonly DurableInvocationOptions OnDemandOptions = new() { IsMultiConcurrency = false }; + + private static string IdAt(int position) => OperationIdGenerator.HashOperationId(position.ToString()); + + // ── Teardown ordering (all modes) ────────────────────────────────── + + [Fact] + public void WorkflowCancellation_TerminateThenDispose_TokenIsAlwaysCancelled() + { + // Termination delivers its cancel through a threadpool continuation. + // Disposing right after Terminate used to win that race and drop the + // cancel. Repeat to make the old race show up reliably. + for (var i = 0; i < 500; i++) + { + var tm = new TerminationManager(); + var wfc = new WorkflowCancellation(tm); + var token = wfc.Token; + + tm.Terminate(TerminationReason.WaitScheduled); + wfc.Dispose(); + + Assert.True(token.IsCancellationRequested, $"Cancellation lost on iteration {i}."); + } + } + + [Fact] + public void WorkflowCancellation_TokenIsUsableAfterDispose() + { + var wfc = new WorkflowCancellation(new TerminationManager()); + wfc.Dispose(); + + // Abandoned user code reads the token after the invocation ends. It must + // get a cancelled token, not an ObjectDisposedException. + var token = wfc.Token; + Assert.True(token.IsCancellationRequested); + using var linked = CancellationTokenSource.CreateLinkedTokenSource(CancellationToken.None, token); + Assert.True(linked.IsCancellationRequested); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task Suspended_AbandonedStepTokenIsCancelledBeforeWrapperReturns(bool isMultiConcurrency) + { + // Repeat to cover warm, back-to-back invocations of the same handler. + for (var i = 0; i < 20; i++) + { + var client = new FakeServiceClient(); + var entered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + CancellationToken stepToken = default; + + var output = await DurableFunction.WrapAsync( + async (input, ctx) => + { + // Not awaited: this step is still running when the workflow suspends. + _ = ctx.StepAsync(async (_, ct) => + { + stepToken = ct; + entered.TrySetResult(); + await Task.Delay(Timeout.Infinite, ct); + return 0; + }, name: "slow"); + await entered.Task; + + await ctx.WaitAsync(TimeSpan.FromSeconds(30), name: "wait"); + return "unreachable"; + }, + Input($"arn:abandoned-{i}"), + Context(TimeSpan.FromMinutes(5)), + client, + new DurableInvocationOptions { IsMultiConcurrency = isMultiConcurrency }); + + Assert.Equal(InvocationStatus.Pending, output.Status); + // Checked immediately: cancellation must already be delivered. + Assert.True(stepToken.IsCancellationRequested, $"Abandoned step not cancelled on iteration {i}."); + } + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task AfterTeardown_AbandonedCodeCannotStartNewOperations(bool isMultiConcurrency) + { + var client = new FakeServiceClient(); + var release = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var bodyRan = false; + Task? leaked = null; + + var output = await DurableFunction.WrapAsync( + async (_, ctx) => + { + // Work the workflow started but never awaited; it resumes after the + // invocation has ended. + leaked = Task.Run(async () => + { + await release.Task; + return await ctx.StepAsync(async (_, _) => + { + bodyRan = true; + await Task.CompletedTask; + return 1; + }, name: "late"); + }); + + await ctx.WaitAsync(TimeSpan.FromSeconds(30), name: "wait"); + return "unreachable"; + }, + Input("arn:leaked"), + Context(TimeSpan.FromMinutes(5)), + client, + new DurableInvocationOptions { IsMultiConcurrency = isMultiConcurrency }); + + Assert.Equal(InvocationStatus.Pending, output.Status); + var checkpointsAtReturn = client.Calls.Count; + + release.SetResult(); + await Assert.ThrowsAnyAsync(() => leaked!.WaitAsync(HangGuard)); + + Assert.False(bodyRan); + Assert.Equal(checkpointsAtReturn, client.Calls.Count); + } + + // ── Deadline: on-demand is unchanged ─────────────────────────────── + + [Theory] + [InlineData(false, 300_000, null)] // on-demand: never a deadline + [InlineData(false, 1, null)] + [InlineData(true, 0, null)] // unset test context: no deadline + [InlineData(true, -5, null)] + [InlineData(true, 300_000, 300_000)] + public void InvocationDeadline_Resolve(bool isMultiConcurrency, long remainingMs, long? expectedMs) + { + var deadline = InvocationDeadline.Resolve(TimeSpan.FromMilliseconds(remainingMs), isMultiConcurrency); + Assert.Equal(expectedMs, deadline is { } d ? (long)d.TotalMilliseconds : null); + } + + [Fact] + public void InvocationDeadline_Resolve_IgnoresValuesBeyondTimerRange() + { + Assert.Null(InvocationDeadline.Resolve(TimeSpan.FromDays(60), isMultiConcurrency: true)); + } + + [Fact] + public async Task OnDemand_HydrationIsNotCancellable() + { + var client = new FakeServiceClient(); + + var output = await DurableFunction.WrapAsync( + (input, _) => Task.FromResult(input.Id!), + Input("arn:hydrate-od", nextMarker: "page-2"), + Context(TimeSpan.FromMilliseconds(1)), + client, + OnDemandOptions); + + Assert.Equal(InvocationStatus.Succeeded, output.Status); + Assert.False(Assert.Single(client.HydrationTokens).CanBeCanceled); + } + + [Fact] + public async Task Lmi_HydrationReceivesDeadlineToken() + { + var client = new FakeServiceClient(); + + var output = await DurableFunction.WrapAsync( + (input, _) => Task.FromResult(input.Id!), + Input("arn:hydrate-lmi", nextMarker: "page-2"), + Context(TimeSpan.FromMinutes(5)), + client, + LmiOptions); + + Assert.Equal(InvocationStatus.Succeeded, output.Status); + Assert.True(Assert.Single(client.HydrationTokens).CanBeCanceled); + } + + // ── Deadline: LMI ────────────────────────────────────────────────── + + [Fact] + public async Task Lmi_DeadlinePassesDuringStep_ThrowsTimeout_WithoutFailCheckpoint() + { + var client = new FakeServiceClient(); + CancellationToken stepToken = default; + + var ex = await Assert.ThrowsAsync(() => + DurableFunction.WrapAsync( + async (_, ctx) => await ctx.StepAsync(async (_, ct) => + { + stepToken = ct; + await Task.Delay(Timeout.Infinite, ct); + return "unreachable"; + }, name: "slow"), + Input("arn:deadline-step"), + Context(TimeSpan.FromMilliseconds(100)), + client, + LmiOptions).WaitAsync(HangGuard)); + + Assert.Contains("arn:deadline-step", ex.Message); + Assert.True(stepToken.IsCancellationRequested); + // A timeout must not be recorded as a durable failure; replay retries it. + Assert.DoesNotContain(client.AllUpdates, u => u.Action == OperationAction.FAIL); + } + + [Fact] + public async Task Lmi_WorkflowConvertsCancellationToOwnError_StillReportsTimeout() + { + // User code that catches the cancellation and throws its own exception + // must not turn a timeout into a permanently failed execution. + var client = new FakeServiceClient(); + + await Assert.ThrowsAsync(() => + DurableFunction.WrapAsync( + async (_, ctx) => + { + try + { + return await ctx.StepAsync(async (_, ct) => + { + await Task.Delay(Timeout.Infinite, ct); + return "unreachable"; + }, name: "slow"); + } + catch (OperationCanceledException) + { + throw new InvalidOperationException("user code gave up"); + } + }, + Input("arn:deadline-user-catch"), + Context(TimeSpan.FromMilliseconds(100)), + client, + LmiOptions).WaitAsync(HangGuard)); + } + + [Fact] + public async Task RunAsync_WorkflowFaultsFromDeadlineBeforeDeadlineTaskCompletes_ReportsDeadline() + { + // Forces the race where the workflow task faults before DeadlineTask + // completes: the fault is raised synchronously inside the token's + // cancellation callback, which runs before the deadline task is + // signalled. The fault must still be reported as a deadline, not as a + // durable failure. + var tm = new TerminationManager(); + using var wfc = new WorkflowCancellation(tm, TimeSpan.FromMilliseconds(50)); + + var result = await DurableExecutionHandler.RunAsync( + new ExecutionState(), + tm, + async () => + { + // No RunContinuationsAsynchronously: the await below resumes + // inline inside the cancellation callback. + var faulted = new TaskCompletionSource(); + wfc.Token.Register(() => faulted.TrySetException(new InvalidOperationException("user gave up"))); + return await faulted.Task; + }, + wfc).WaitAsync(HangGuard); + + Assert.True(result.DeadlineExceeded); + Assert.Null(result.Exception); + } + + [Fact] + public async Task Lmi_DeadlinePassesDuringHydration_ThrowsTimeout() + { + var client = new FakeServiceClient + { + OnGetExecutionState = ct => Task.Delay(Timeout.Infinite, ct), + }; + var workflowRan = false; + + var ex = await Assert.ThrowsAsync(() => + DurableFunction.WrapAsync( + (input, _) => { workflowRan = true; return Task.FromResult(input.Id!); }, + Input("arn:deadline-hydrate", nextMarker: "page-2"), + Context(TimeSpan.FromMilliseconds(100)), + client, + LmiOptions).WaitAsync(HangGuard)); + + Assert.IsAssignableFrom(ex.InnerException); + Assert.False(workflowRan); + } + + [Fact] + public async Task Lmi_CheckpointCallStallsPastDeadline_CleanupIsBounded() + { + // The checkpoint call ignores its cancellation token and never returns. + // The invocation must still end instead of holding a processing loop. + var neverCompletes = new TaskCompletionSource(); + var client = new FakeServiceClient { OnCheckpoint = _ => neverCompletes.Task }; + + await Assert.ThrowsAsync(() => + DurableFunction.WrapAsync( + async (_, ctx) => await ctx.StepAsync( + async (_, _) => { await Task.CompletedTask; return "done"; }, + name: "step"), + Input("arn:stalled-checkpoint"), + Context(TimeSpan.FromMilliseconds(100)), + client, + new DurableInvocationOptions { IsMultiConcurrency = true, CleanupTimeout = TimeSpan.FromMilliseconds(100) }) + .WaitAsync(HangGuard)); + } + + [Fact] + public async Task Lmi_CompletesBeforeDeadline_Succeeds() + { + var client = new FakeServiceClient(); + + var output = await DurableFunction.WrapAsync( + async (input, ctx) => await ctx.StepAsync( + async (_, _) => { await Task.CompletedTask; return input.Id!; }, name: "step"), + Input("arn:in-time", id: "abc"), + Context(TimeSpan.FromMinutes(5)), + client, + LmiOptions); + + Assert.Equal(InvocationStatus.Succeeded, output.Status); + Assert.Equal("\"abc\"", output.Result); + } + + // ── Concurrent invocations ───────────────────────────────────────── + + [Fact] + public async Task Lmi_ConcurrentInvocations_AreIsolated() + { + const int count = 8; + var client = new FakeServiceClient(); + var handler = new SharedHandler(count); + + // All invocations go through one handler instance, as they would in one + // LMI execution environment. Even-numbered ones suspend, odd ones finish. + var tasks = Enumerable.Range(0, count).Select(i => Task.Run(() => + DurableFunction.WrapAsync( + handler.RunAsync, + Input($"arn:exec-{i}", id: $"id-{i}", checkpointToken: $"initial-{i}"), + Context(TimeSpan.FromMinutes(5), requestId: $"req-{i}"), + client, + LmiOptions))).ToArray(); + + var outputs = await Task.WhenAll(tasks).WaitAsync(HangGuard); + + for (var i = 0; i < count; i++) + { + var arn = $"arn:exec-{i}"; + + if (i % 2 == 0) + { + Assert.Equal(InvocationStatus.Pending, outputs[i].Status); + } + else + { + Assert.Equal(InvocationStatus.Succeeded, outputs[i].Status); + Assert.Equal($"\"id-{i}:id-{i}\"", outputs[i].Result); + } + + // Each step saw only its own invocation's context. + Assert.Equal(arn, handler.ObservedArns[$"id-{i}"]); + Assert.Equal($"req-{i}", handler.ObservedRequestIds[$"id-{i}"]); + + // Each execution's checkpoints carry only its own ids, payloads and + // tokens, and its token chain starts from its own initial token. + var calls = client.CallsFor(arn); + Assert.NotEmpty(calls); + Assert.Equal($"initial-{i}", calls[0].TokenIn); + for (var c = 1; c < calls.Count; c++) + Assert.Equal(calls[c - 1].TokenOut, calls[c].TokenIn); + + var updates = calls.SelectMany(c => c.Updates).ToList(); + Assert.All(updates, u => Assert.Contains(u.Id, new[] { IdAt(1), IdAt(2), IdAt(3) })); + Assert.All(updates.Where(u => u.Payload != null), u => Assert.Contains($"id-{i}", u.Payload)); + Assert.Contains(updates, u => u.Id == IdAt(1) && u.Action == OperationAction.SUCCEED); + } + } + + [Fact] + public async Task Lmi_PeerTimesOut_HealthyPeerIsUnaffected() + { + var client = new FakeServiceClient(); + var peerFinished = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var healthyEntered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + CancellationToken healthyToken = default; + + var healthy = Task.Run(() => DurableFunction.WrapAsync( + async (input, ctx) => await ctx.StepAsync(async (_, ct) => + { + healthyToken = ct; + healthyEntered.TrySetResult(); + // Keep running until the peer has timed out and been cleaned up. + await peerFinished.Task; + return input.Id!; + }, name: "healthy"), + Input("arn:healthy", id: "ok"), + Context(TimeSpan.FromMinutes(5)), + client, + LmiOptions)); + + await healthyEntered.Task.WaitAsync(HangGuard); + + var timingOut = DurableFunction.WrapAsync( + async (_, ctx) => await ctx.StepAsync(async (_, ct) => + { + await Task.Delay(Timeout.Infinite, ct); + return "unreachable"; + }, name: "slow"), + Input("arn:timing-out"), + Context(TimeSpan.FromMilliseconds(100)), + client, + LmiOptions); + + await Assert.ThrowsAsync(() => timingOut.WaitAsync(HangGuard)); + Assert.False(healthyToken.IsCancellationRequested); + + peerFinished.SetResult(); + var output = await healthy.WaitAsync(HangGuard); + + Assert.Equal(InvocationStatus.Succeeded, output.Status); + Assert.Equal("\"ok\"", output.Result); + Assert.Contains(client.CallsFor("arn:healthy").SelectMany(c => c.Updates), + u => u.Action == OperationAction.SUCCEED); + } + + [Fact] + public async Task Lmi_ReplayWhilePeerIsLive_SkipsCompletedSteps() + { + var client = new FakeServiceClient(); + var bothInside = new AsyncBarrier(2); + var replayBodyRuns = 0; + var freshBodyRuns = 0; + + // The replaying execution already completed step 1; only step 2 runs. + var replay = Task.Run(() => DurableFunction.WrapAsync( + async (_, ctx) => + { + var first = await ctx.StepAsync(async (_, _) => + { + Interlocked.Increment(ref replayBodyRuns); + await Task.CompletedTask; + return "recomputed"; + }, name: "first"); + var second = await ctx.StepAsync(async (_, _) => + { + await bothInside.SignalAndWaitAsync(); + return "second"; + }, name: "second"); + return $"{first}+{second}"; + }, + Input("arn:replay", extraOperations: new[] + { + new Operation + { + Id = IdAt(1), Type = OperationTypes.Step, Status = OperationStatuses.Succeeded, + Name = "first", StepDetails = new StepDetails { Result = "\"cached\"" }, + }, + }), + Context(TimeSpan.FromMinutes(5)), + client, + LmiOptions)); + + var fresh = Task.Run(() => DurableFunction.WrapAsync( + async (_, ctx) => await ctx.StepAsync(async (_, _) => + { + Interlocked.Increment(ref freshBodyRuns); + await bothInside.SignalAndWaitAsync(); + return "fresh"; + }, name: "first"), + Input("arn:fresh"), + Context(TimeSpan.FromMinutes(5)), + client, + LmiOptions)); + + var outputs = await Task.WhenAll(replay, fresh).WaitAsync(HangGuard); + + Assert.Equal("\"cached+second\"", outputs[0].Result); + Assert.Equal("\"fresh\"", outputs[1].Result); + Assert.Equal(0, replayBodyRuns); + Assert.Equal(1, freshBodyRuns); + Assert.DoesNotContain(client.CallsFor("arn:replay").SelectMany(c => c.Updates), u => u.Id == IdAt(1)); + } + + // ── Helpers ───────────────────────────────────────────────────────── + + /// + /// One handler instance shared by every concurrent invocation. Every step + /// waits until all invocations are inside a step, so they truly overlap. + /// + private sealed class SharedHandler + { + private readonly AsyncBarrier _allInside; + public ConcurrentDictionary ObservedArns { get; } = new(); + public ConcurrentDictionary ObservedRequestIds { get; } = new(); + + public SharedHandler(int count) => _allInside = new AsyncBarrier(count); + + public async Task RunAsync(Payload input, IDurableContext ctx) + { + var first = await ctx.StepAsync(async (_, _) => + { + ObservedArns[input.Id!] = ctx.ExecutionContext.DurableExecutionArn; + ObservedRequestIds[input.Id!] = ctx.LambdaContext.AwsRequestId; + await _allInside.SignalAndWaitAsync(); + return input.Id!; + }, name: "first"); + + var second = await ctx.StepAsync( + async (_, _) => { await Task.Yield(); return input.Id!; }, name: "second"); + + var index = int.Parse(input.Id!.Substring("id-".Length)); + if (index % 2 == 0) + await ctx.WaitAsync(TimeSpan.FromSeconds(30), name: "wait"); + + return $"{first}:{second}"; + } + } + + private sealed class AsyncBarrier + { + private int _remaining; + private readonly TaskCompletionSource _released = new(TaskCreationOptions.RunContinuationsAsynchronously); + + public AsyncBarrier(int participants) => _remaining = participants; + + public Task SignalAndWaitAsync() + { + if (Interlocked.Decrement(ref _remaining) == 0) + _released.TrySetResult(); + return _released.Task; + } + } + + public class Payload + { + public string? Id { get; set; } + } + + private static TestLambdaContext Context(TimeSpan remainingTime, string requestId = "req") => new() + { + Serializer = new DefaultLambdaJsonSerializer(), + RemainingTime = remainingTime, + AwsRequestId = requestId, + }; + + private static DurableExecutionInvocationInput Input( + string arn, + string id = "id", + string checkpointToken = "initial", + string? nextMarker = null, + IEnumerable? extraOperations = null) + { + var operations = new List + { + new() + { + Id = "exec-0", + Type = OperationTypes.Execution, + Status = OperationStatuses.Started, + ExecutionDetails = new ExecutionDetails { InputPayload = $"{{\"Id\":\"{id}\"}}" }, + }, + }; + if (extraOperations != null) operations.AddRange(extraOperations); + + return new DurableExecutionInvocationInput + { + DurableExecutionArn = arn, + CheckpointToken = checkpointToken, + InitialExecutionState = new InitialExecutionState { Operations = operations, NextMarker = nextMarker }, + }; + } + + private sealed record CheckpointCall(string Arn, string? TokenIn, string TokenOut, IReadOnlyList Updates); + + /// + /// In-memory durable service shared by concurrent invocations. Issues a new + /// checkpoint token per call, scoped to the execution ARN. + /// + private sealed class FakeServiceClient : IDurableServiceClient + { + private readonly ConcurrentQueue _calls = new(); + private readonly ConcurrentQueue _hydrationTokens = new(); + private int _tokenCounter; + + public Func? OnCheckpoint { get; init; } + public Func? OnGetExecutionState { get; init; } + + public IReadOnlyList Calls => _calls.ToArray(); + public IReadOnlyList HydrationTokens => _hydrationTokens.ToArray(); + public IEnumerable AllUpdates => Calls.SelectMany(c => c.Updates); + public List CallsFor(string arn) => Calls.Where(c => c.Arn == arn).ToList(); + + public async Task CheckpointAsync( + string durableExecutionArn, + string? checkpointToken, + IReadOnlyList pendingOperations, + Action>? onNewOperations = null, + CancellationToken cancellationToken = default) + { + if (OnCheckpoint != null) await OnCheckpoint(cancellationToken); + + var tokenOut = $"{durableExecutionArn}#{Interlocked.Increment(ref _tokenCounter)}"; + _calls.Enqueue(new CheckpointCall(durableExecutionArn, checkpointToken, tokenOut, pendingOperations.ToArray())); + return tokenOut; + } + + public async Task<(List Operations, string? NextMarker)> GetExecutionStateAsync( + string durableExecutionArn, + string? checkpointToken, + string marker, + CancellationToken cancellationToken = default) + { + _hydrationTokens.Enqueue(cancellationToken); + if (OnGetExecutionState != null) await OnGetExecutionState(cancellationToken); + return (new List(), null); + } + } +} From 1fafa7dc839597835e0b021f17acfe97818af43e Mon Sep 17 00:00:00 2001 From: Norm Johanson Date: Wed, 4 Nov 2026 00:59:43 +0000 Subject: [PATCH 2/6] Address Copilot review findings on LMI deadline and teardown handling - Contain exceptions thrown by user cancellation callbacks so they cannot crash the shared process from the deadline timer or replace the outcome during teardown. - Signal the deadline before cancelling, so a blocking user callback cannot stop the handler from seeing it; unregister rather than dispose the deadline registration. - Read IsDeadlineExceeded from the deadline token, independent of which cancellation reason came first. - Let the deadline decide the outcome whichever task wins the race. - Treat an already-expired LMI invocation as an immediate timeout. - Do not checkpoint FAIL or RETRY for exceptions raised after the deadline in steps, child contexts and condition checks. - Re-check cancellation before running a step body, child function or condition check, and when the queued workflow actually starts. - Re-check the deadline after the final drain, bound hydration even when the client ignores cancellation, and stop the batcher from sending further batches once it is shut down. - Tests: add a regression test per fix and a hang guard that does not throw TimeoutException, so a hang can no longer pass a deadline test. --- .../DurableExecutionHandler.cs | 32 +- .../DurableFunction.cs | 38 +- .../Internal/CheckpointBatcher.cs | 7 + .../Internal/ChildContextOperation.cs | 10 + .../Internal/StepOperation.cs | 12 + .../Internal/WaitForConditionOperation.cs | 9 + .../Internal/WorkflowCancellation.cs | 54 ++- .../docs/core/cancellation.md | 2 + .../LambdaManagedInstancesTests.cs | 393 +++++++++++++++++- 9 files changed, 519 insertions(+), 38 deletions(-) diff --git a/Libraries/src/Amazon.Lambda.DurableExecution/DurableExecutionHandler.cs b/Libraries/src/Amazon.Lambda.DurableExecution/DurableExecutionHandler.cs index c205a8462..395a61cec 100644 --- a/Libraries/src/Amazon.Lambda.DurableExecution/DurableExecutionHandler.cs +++ b/Libraries/src/Amazon.Lambda.DurableExecution/DurableExecutionHandler.cs @@ -90,7 +90,15 @@ internal static async Task> RunAsync( // the termination signal. When TerminationManager fires (e.g., WaitAsync), // we need the WhenAny race below to resolve immediately without waiting // for the user task to reach an await point. - var userTask = Task.Run(userHandler); + // + // The workflow is only queued here; under load it may start after the + // deadline passed or after this invocation already returned. Check again + // when it actually starts so a late start runs no workflow code. + var userTask = Task.Run(() => + { + workflowCancellation?.ThrowIfTornDown(); + return userHandler(); + }); // Race: user code completing vs. termination signal (wait/callback/retry) // vs. the invocation deadline. DeadlineTask never completes when there is @@ -98,19 +106,23 @@ internal static async Task> RunAsync( var deadlineTask = workflowCancellation?.DeadlineTask ?? NeverCompletes; var winner = await Task.WhenAny(userTask, terminationManager.TerminationTask, deadlineTask); - if (winner != userTask) + // Once the deadline has passed, it decides the outcome whichever task won. + // A workflow can catch the cancellation and return a fallback value, or a + // termination can race the deadline; neither is a real result, and the + // service must re-invoke the execution. + var deadlineExceeded = workflowCancellation?.IsDeadlineExceeded == true; + + if (winner != userTask || deadlineExceeded) { - // userTask is abandoned. Observe its eventual fault so it never - // surfaces through TaskScheduler.UnobservedTaskException, which in a - // shared process (Lambda Managed Instances) would be reported against - // whichever invocation happens to be running when the GC finalizes it. + // userTask is abandoned or its outcome is ignored. Observe its fault so + // it never surfaces through TaskScheduler.UnobservedTaskException, which + // in a shared process (Lambda Managed Instances) would be reported + // against whichever invocation happens to be running when the GC + // finalizes it. ObserveAbandoned(userTask); } - // The deadline wins over a termination that raced it: the invocation has - // run out of time either way, and the service must re-invoke it. - if (winner == deadlineTask - || (winner == terminationManager.TerminationTask && workflowCancellation?.IsDeadlineExceeded == true)) + if (deadlineExceeded) { return new HandlerResult { diff --git a/Libraries/src/Amazon.Lambda.DurableExecution/DurableFunction.cs b/Libraries/src/Amazon.Lambda.DurableExecution/DurableFunction.cs index d9d95b3ec..b9005d0ef 100644 --- a/Libraries/src/Amazon.Lambda.DurableExecution/DurableFunction.cs +++ b/Libraries/src/Amazon.Lambda.DurableExecution/DurableFunction.cs @@ -124,14 +124,21 @@ private static async Task WrapAsyncCore WrapAsyncCore WrapAsyncCore + /// Awaits , but stops waiting when + /// is cancelled. A task that is abandoned this + /// way has its eventual fault observed. With no deadline this is a plain await. + /// + private static async Task WaitWithDeadlineAsync(Task task, CancellationToken deadlineToken) + { + if (!deadlineToken.CanBeCanceled) + return await task.ConfigureAwait(false); + + try + { + return await task.WaitAsync(deadlineToken).ConfigureAwait(false); + } + catch (OperationCanceledException) when (!task.IsCompleted) + { + _ = task.ContinueWith( + static t => _ = t.Exception, + CancellationToken.None, + TaskContinuationOptions.OnlyOnFaulted | TaskContinuationOptions.ExecuteSynchronously, + TaskScheduler.Default); + throw; + } + } + /// /// The exception thrown when a Lambda Managed Instances invocation runs out /// of time. It escapes to the Lambda runtime as an invocation error, the same diff --git a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/CheckpointBatcher.cs b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/CheckpointBatcher.cs index f08101fda..824cf2358 100644 --- a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/CheckpointBatcher.cs +++ b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/CheckpointBatcher.cs @@ -306,6 +306,13 @@ private async Task FlushBatchAsync(IReadOnlyList batch, CancellationT try { + // Never start a new request once the worker is shut down or abandoned. + // If an abandoned flush ignored cancellation and later returned, the + // worker would otherwise go on to send queued batches with this + // invocation's checkpoint token, racing the invocation that replays it. + // Throwing here routes through the catch below, which fails the batch. + cancellationToken.ThrowIfCancellationRequested(); + var newToken = await _flushAsync(_checkpointToken, updates, cancellationToken).ConfigureAwait(false); Volatile.Write(ref _checkpointToken, newToken); diff --git a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/ChildContextOperation.cs b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/ChildContextOperation.cs index 55c09d05f..ec375548d 100644 --- a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/ChildContextOperation.cs +++ b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/ChildContextOperation.cs @@ -187,6 +187,9 @@ private async Task ExecuteFunc(CancellationToken cancellationToken) T result; try { + // The START checkpoint may have been awaited; the workflow can have + // been torn down meanwhile. Never start the child function then. + _workflowCancellation.Token.ThrowIfCancellationRequested(); result = await _func(childContext, linked.Token); } catch (OperationCanceledException) when (linked.IsCancellationRequested) @@ -196,6 +199,13 @@ private async Task ExecuteFunc(CancellationToken cancellationToken) // (or upstream cancel) owns the outcome. throw; } + catch (Exception) when (_workflowCancellation.IsDeadlineExceeded) + { + // The invocation deadline passed (Lambda Managed Instances only). An + // exception raised because work was cancelled is not a real failure: + // do NOT checkpoint CONTEXT FAIL. The invocation is retried. + throw; + } catch (NonDeterministicExecutionException) { // Replay-mismatch from an inner operation means the entire execution diff --git a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/StepOperation.cs b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/StepOperation.cs index 462b146bf..5a02b17e1 100644 --- a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/StepOperation.cs +++ b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/StepOperation.cs @@ -233,6 +233,9 @@ private async Task ExecuteFunc(int attemptNumber, CancellationToken cancellat ["attempt"] = attemptNumber, })) { + // The START checkpoint above may have been awaited; the workflow + // can have been torn down meanwhile. Never start the body then. + _workflowCancellation.Token.ThrowIfCancellationRequested(); result = await _func(stepContext, linked.Token); } } @@ -245,6 +248,15 @@ private async Task ExecuteFunc(int attemptNumber, CancellationToken cancellat // as a fault on the workflow user task. throw; } + catch (Exception) when (_workflowCancellation.IsDeadlineExceeded) + { + // The invocation deadline passed (Lambda Managed Instances only). A + // body that converts its cancellation into another exception did not + // really fail: do NOT checkpoint FAIL or RETRY, or replay would see a + // permanent failure caused by a timeout. The invocation is retried + // and the step runs again. + throw; + } catch (Exception ex) { // The user step body itself threw. Funnel into the retry/fail decision diff --git a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/WaitForConditionOperation.cs b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/WaitForConditionOperation.cs index 2345eda9d..9649ec441 100644 --- a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/WaitForConditionOperation.cs +++ b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/WaitForConditionOperation.cs @@ -191,6 +191,8 @@ private async Task ExecuteIteration(TState currentState, int attemptNumb try { var checkContext = new ConditionCheckContext(attemptNumber, _logger); + // Never start the check once the workflow has been torn down. + _workflowCancellation.Token.ThrowIfCancellationRequested(); newState = await _check(currentState, checkContext, linked.Token); } catch (OperationCanceledException) when (linked.IsCancellationRequested) @@ -200,6 +202,13 @@ private async Task ExecuteIteration(TState currentState, int attemptNumb // upstream cancel) owns the outcome. throw; } + catch (Exception) when (_workflowCancellation.IsDeadlineExceeded) + { + // The invocation deadline passed (Lambda Managed Instances only). An + // exception raised because the check was cancelled is not a real + // failure: do NOT checkpoint FAIL. The invocation is retried. + throw; + } catch (Exception ex) { // The check threw. WaitForCondition has no per-exception retry diff --git a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/WorkflowCancellation.cs b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/WorkflowCancellation.cs index 9d248d5e2..9d22e9707 100644 --- a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/WorkflowCancellation.cs +++ b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/WorkflowCancellation.cs @@ -1,6 +1,8 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 +using Amazon.Lambda.Core; + namespace Amazon.Lambda.DurableExecution.Internal; /// @@ -71,8 +73,12 @@ internal sealed class WorkflowCancellation : IDisposable /// public WorkflowCancellationReason Reason => (WorkflowCancellationReason)Volatile.Read(ref _reason); - /// True once the invocation deadline passed. - public bool IsDeadlineExceeded => Reason == WorkflowCancellationReason.DeadlineExceeded; + /// + /// True once the invocation deadline passed. Read from the deadline token + /// itself, so it is accurate even when a suspension or teardown cancelled the + /// workflow first, and becomes true before any cancellation callback runs. + /// + public bool IsDeadlineExceeded => DeadlineToken.IsCancellationRequested; /// /// True once the invocation is being torn down (the wrapper finished or the @@ -111,7 +117,12 @@ public WorkflowCancellation(TerminationManager terminationManager, TimeSpan? dea if (deadline is { } remaining) { - _deadlineCts = new CancellationTokenSource(remaining); + // A zero deadline means the invocation was already out of time when + // it started: cancel at once rather than through the timer. + _deadlineCts = remaining > TimeSpan.Zero + ? new CancellationTokenSource(remaining) + : new CancellationTokenSource(); + if (remaining <= TimeSpan.Zero) _deadlineCts.Cancel(); DeadlineToken = _deadlineCts.Token; _deadlineRegistration = _deadlineCts.Token.Register( static state => ((WorkflowCancellation)state!).OnDeadline(), @@ -142,12 +153,15 @@ public void ThrowIfTornDown() private void OnDeadline() { - // Record the reason and stop admitting new operations before cancelling - // the token, so code that observes the cancellation also sees why. + // Stop admitting new operations and signal the deadline BEFORE cancelling + // the token. Cancel runs user cancellation callbacks synchronously; a + // callback that blocks must not stop DurableExecutionHandler from seeing + // the deadline. _deadlineTcs runs its continuations asynchronously, so + // signalling it never runs the handler on this timer thread. Interlocked.CompareExchange(ref _reason, (int)WorkflowCancellationReason.DeadlineExceeded, 0); Volatile.Write(ref _tornDown, 1); - Cancel(WorkflowCancellationReason.DeadlineExceeded); _deadlineTcs.TrySetResult(); + Cancel(WorkflowCancellationReason.DeadlineExceeded); } private void Cancel(WorkflowCancellationReason reason) @@ -160,6 +174,17 @@ private void Cancel(WorkflowCancellationReason reason) // timer. Dispose cancels before disposing, so the token is already // cancelled. } + catch (AggregateException ex) + { + // A user cancellation callback threw. Cancel still ran every other + // callback before throwing, so cancellation is fully delivered. Never + // let the exception escape: on the deadline timer it would crash the + // process (and every invocation sharing it), and during teardown it + // would replace the invocation's real outcome. + LambdaLogger.Log( + "Warning", ex, + "A cancellation callback registered by the durable workflow threw while the workflow was being cancelled."); + } } public void Dispose() @@ -167,7 +192,11 @@ public void Dispose() if (Interlocked.Exchange(ref _disposed, 1) != 0) return; BeginTeardown(); - _deadlineRegistration.Dispose(); + + // Unregister rather than Dispose: Dispose waits for a deadline callback + // that is still running, and a user cancellation callback that blocks + // would then hang the invocation here. + _deadlineRegistration.Unregister(); _deadlineCts?.Dispose(); _cts.Dispose(); } @@ -196,18 +225,21 @@ public static bool IsMultiConcurrency() => !string.IsNullOrEmpty(Environment.GetEnvironmentVariable(MaxConcurrencyEnvironmentVariable)); /// - /// Returns the time until the deadline, or null for no deadline. + /// Returns the time until the deadline, when the + /// deadline has already passed, or null for no deadline. /// /// The invocation's ILambdaContext.RemainingTime. /// Whether the function runs on Lambda Managed Instances. /// - /// A zero or negative remaining time means no deadline. That is what an - /// unconfigured test context reports, and it must not cancel immediately. + /// On Lambda Managed Instances a zero or negative remaining time means the + /// invocation was already out of time when it started, so it times out + /// immediately. Tests that exercise this mode must set + /// TestLambdaContext.RemainingTime. /// public static TimeSpan? Resolve(TimeSpan remainingTime, bool isMultiConcurrency) { if (!isMultiConcurrency) return null; - if (remainingTime <= TimeSpan.Zero) return null; + if (remainingTime <= TimeSpan.Zero) return TimeSpan.Zero; // CancellationTokenSource rejects delays above int.MaxValue milliseconds. // Lambda timeouts are far below that; treat anything larger as no deadline. diff --git a/Libraries/src/Amazon.Lambda.DurableExecution/docs/core/cancellation.md b/Libraries/src/Amazon.Lambda.DurableExecution/docs/core/cancellation.md index e8104f37f..c9d4b3eaf 100644 --- a/Libraries/src/Amazon.Lambda.DurableExecution/docs/core/cancellation.md +++ b/Libraries/src/Amazon.Lambda.DurableExecution/docs/core/cancellation.md @@ -29,6 +29,8 @@ When either trigger fires, the token transitions to `IsCancellationRequested = t On Lambda Managed Instances several invocations run in the same process, and Lambda does not stop an invocation that times out. The SDK therefore applies the invocation deadline (`ILambdaContext.RemainingTime`) itself: - At the deadline the token is cancelled, new durable operations are rejected, and state loading is cancelled. +- An exception your code throws after the deadline (for example, after catching the cancellation) is not recorded as a step failure, and does not consume a retry. +- An invocation that starts with its deadline already passed times out immediately. When testing in this mode, set `TestLambdaContext.RemainingTime`. - Checkpoint cleanup gets a short budget past the deadline, then is abandoned. - `WrapAsync` throws `TimeoutException`. Lambda treats this as an invocation error, the same outcome as a timeout in on-demand Lambda: the execution is not marked failed, and work that was not checkpointed runs again on replay. diff --git a/Libraries/test/Amazon.Lambda.DurableExecution.Tests/LambdaManagedInstancesTests.cs b/Libraries/test/Amazon.Lambda.DurableExecution.Tests/LambdaManagedInstancesTests.cs index 7c608c382..53236877a 100644 --- a/Libraries/test/Amazon.Lambda.DurableExecution.Tests/LambdaManagedInstancesTests.cs +++ b/Libraries/test/Amazon.Lambda.DurableExecution.Tests/LambdaManagedInstancesTests.cs @@ -86,13 +86,13 @@ public async Task Suspended_AbandonedStepTokenIsCancelledBeforeWrapperReturns(bo async (input, ctx) => { // Not awaited: this step is still running when the workflow suspends. - _ = ctx.StepAsync(async (_, ct) => + ctx.StepAsync(async (_, ct) => { stepToken = ct; entered.TrySetResult(); await Task.Delay(Timeout.Infinite, ct); return 0; - }, name: "slow"); + }, name: "slow").Observe(); await entered.Task; await ctx.WaitAsync(TimeSpan.FromSeconds(30), name: "wait"); @@ -147,7 +147,7 @@ public async Task AfterTeardown_AbandonedCodeCannotStartNewOperations(bool isMul var checkpointsAtReturn = client.Calls.Count; release.SetResult(); - await Assert.ThrowsAnyAsync(() => leaked!.WaitAsync(HangGuard)); + await Assert.ThrowsAnyAsync(() => leaked!.NoHang()); Assert.False(bodyRan); Assert.Equal(checkpointsAtReturn, client.Calls.Count); @@ -158,8 +158,8 @@ public async Task AfterTeardown_AbandonedCodeCannotStartNewOperations(bool isMul [Theory] [InlineData(false, 300_000, null)] // on-demand: never a deadline [InlineData(false, 1, null)] - [InlineData(true, 0, null)] // unset test context: no deadline - [InlineData(true, -5, null)] + [InlineData(true, 0, 0)] // LMI and already expired: time out at once + [InlineData(true, -5, 0)] [InlineData(true, 300_000, 300_000)] public void InvocationDeadline_Resolve(bool isMultiConcurrency, long remainingMs, long? expectedMs) { @@ -224,7 +224,7 @@ public async Task Lmi_DeadlinePassesDuringStep_ThrowsTimeout_WithoutFailCheckpoi Input("arn:deadline-step"), Context(TimeSpan.FromMilliseconds(100)), client, - LmiOptions).WaitAsync(HangGuard)); + LmiOptions).NoHang()); Assert.Contains("arn:deadline-step", ex.Message); Assert.True(stepToken.IsCancellationRequested); @@ -259,7 +259,7 @@ await Assert.ThrowsAsync(() => Input("arn:deadline-user-catch"), Context(TimeSpan.FromMilliseconds(100)), client, - LmiOptions).WaitAsync(HangGuard)); + LmiOptions).NoHang()); } [Fact] @@ -284,12 +284,341 @@ public async Task RunAsync_WorkflowFaultsFromDeadlineBeforeDeadlineTaskCompletes wfc.Token.Register(() => faulted.TrySetException(new InvalidOperationException("user gave up"))); return await faulted.Task; }, - wfc).WaitAsync(HangGuard); + wfc).NoHang(); Assert.True(result.DeadlineExceeded); Assert.Null(result.Exception); } + // ── Copilot review round 1 ───────────────────────────────────────── + + [Fact] + public async Task Lmi_AlreadyExpiredInvocation_TimesOutWithoutRunningWorkflow() + { + var client = new FakeServiceClient(); + var ran = false; + + await Assert.ThrowsAsync(() => + DurableFunction.WrapAsync( + (input, _) => { ran = true; return Task.FromResult(input.Id!); }, + Input("arn:expired"), + Context(TimeSpan.FromMilliseconds(-1)), + client, + LmiOptions).NoHang()); + + Assert.False(ran); + Assert.Empty(client.Calls); + } + + [Fact] + public void WorkflowCancellation_DeadlineIsReported_EvenWhenTerminationCancelledFirst() + { + var tm = new TerminationManager(); + using var wfc = new WorkflowCancellation(tm, TimeSpan.FromMilliseconds(50)); + var deadlineHit = new ManualResetEventSlim(); + wfc.DeadlineToken.Register(deadlineHit.Set); + + // A suspension sets the first cancellation reason before the deadline. + wfc.BeginTeardown(); + Assert.True(deadlineHit.Wait(HangGuard)); + + Assert.True(wfc.IsDeadlineExceeded); + } + + [Fact] + public async Task WorkflowCancellation_ThrowingCallback_DoesNotEscapeDeadlineOrTeardown() + { + var tm = new TerminationManager(); + var wfc = new WorkflowCancellation(tm, TimeSpan.FromMilliseconds(50)); + // Callbacks run in reverse registration order: the throwing one runs + // first, then this one. It still running proves Cancel delivered to + // every callback despite the throw. + var otherCallbackRan = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + wfc.Token.Register(() => otherCallbackRan.TrySetResult()); + wfc.Token.Register(() => throw new InvalidOperationException("user callback threw")); + + // The deadline fires on a timer thread: the throwing callback must not + // escape there (that would crash the test host) or block the signal. + await wfc.DeadlineTask.NoHang(); + await otherCallbackRan.Task.NoHang(); + Assert.True(wfc.Token.IsCancellationRequested); + + // Teardown and dispose must not throw either. + wfc.BeginTeardown(); + wfc.Dispose(); + } + + [Fact] + public async Task WorkflowCancellation_BlockingCallback_DoesNotDelayDeadlineSignal() + { + var tm = new TerminationManager(); + using var release = new ManualResetEventSlim(); + var wfc = new WorkflowCancellation(tm, TimeSpan.FromMilliseconds(50)); + wfc.Token.Register(() => release.Wait()); + + try + { + // DeadlineTask completes even though a cancellation callback blocks + // the thread that is delivering the cancel. + await wfc.DeadlineTask.NoHang(); + Assert.True(wfc.IsDeadlineExceeded); + } + finally + { + release.Set(); + wfc.Dispose(); + } + } + + [Fact] + public async Task RunAsync_WorkflowReturnsFallbackAfterDeadline_ReportsDeadline() + { + // The workflow observes cancellation and returns a value. It can win the + // race against DeadlineTask; the deadline must still decide the outcome. + var tm = new TerminationManager(); + using var wfc = new WorkflowCancellation(tm, TimeSpan.FromMilliseconds(50)); + + var result = await DurableExecutionHandler.RunAsync( + new ExecutionState(), + tm, + async () => + { + // Completes inline in a deadline-token callback registered after + // WorkflowCancellation's own, so it runs first (callbacks run in + // reverse order): the workflow finishes while the deadline has + // passed but DeadlineTask is not yet signalled. The workflow wins + // the race; only the deadline check can catch it. + var done = new TaskCompletionSource(); + wfc.DeadlineToken.Register(() => done.TrySetResult("fallback")); + return await done.Task; + }, + wfc).NoHang(); + + Assert.True(result.DeadlineExceeded); + Assert.Null(result.Result); + } + + [Fact] + public async Task RunAsync_WorkflowStartedAfterTeardown_RunsNoWorkflowCode() + { + var tm = new TerminationManager(); + using var wfc = new WorkflowCancellation(tm); + wfc.BeginTeardown(); + var ran = false; + + await DurableExecutionHandler.RunAsync( + new ExecutionState(), tm, + () => { ran = true; return Task.FromResult("x"); }, + wfc).NoHang(); + + Assert.False(ran); + } + + [Theory] + [InlineData("step")] + [InlineData("child")] + [InlineData("condition")] + public async Task Lmi_BodyConvertsDeadlineCancellation_NoFailOrRetryCheckpoint(string kind) + { + // The user body catches its cancelled token and throws a different + // exception. That must not be recorded as a real failure. + var client = new FakeServiceClient(); + var retry = new StepConfig { RetryStrategy = RetryStrategy.Exponential(maxAttempts: 3) }; + + async Task Body(CancellationToken ct) + { + try { await Task.Delay(Timeout.Infinite, ct); } + catch (OperationCanceledException) { throw new InvalidOperationException("gave up"); } + return "unreachable"; + } + + Func> workflow = kind switch + { + "step" => async (_, ctx) => await ctx.StepAsync((_, ct) => Body(ct), name: "s", config: retry), + "child" => async (_, ctx) => await ctx.RunInChildContextAsync((_, ct) => Body(ct), name: "c"), + _ => async (_, ctx) => (await ctx.WaitForConditionAsync( + async (state, _, ct) => { await Body(ct); return state; }, + new WaitForConditionConfig + { + InitialState = 0, + WaitStrategy = WaitStrategy.Fixed(TimeSpan.FromSeconds(1)), + }, + name: "w")).ToString(), + }; + + await Assert.ThrowsAsync(() => + DurableFunction.WrapAsync(workflow, Input($"arn:convert-{kind}"), + Context(TimeSpan.FromMilliseconds(100)), client, LmiOptions).NoHang()); + + Assert.DoesNotContain(client.AllUpdates, + u => u.Action == OperationAction.FAIL || u.Action == OperationAction.RETRY); + } + + [Theory] + [InlineData("step")] + [InlineData("child")] + [InlineData("condition")] + public async Task Operation_BodyConvertsDeadlineCancellation_EnqueuesNoFailOrRetry(string kind) + { + // Drives the operation directly, with nothing tearing the batcher down, + // so a FAIL/RETRY the operation enqueued would definitely be flushed. + var tm = new TerminationManager(); + using var wfc = new WorkflowCancellation(tm, TimeSpan.FromMilliseconds(50)); + var recorder = new RecordingBatcher(); + var ctx = new DurableContext(new ExecutionState(), tm, wfc, new OperationIdGenerator(), + "arn:test", Context(TimeSpan.FromMinutes(5)), recorder.Batcher); + var retry = new StepConfig { RetryStrategy = RetryStrategy.Exponential(maxAttempts: 3) }; + + async Task Body(CancellationToken ct) + { + try { await Task.Delay(Timeout.Infinite, ct); } + catch (OperationCanceledException) { throw new InvalidOperationException("gave up"); } + return 0; + } + + Task op = kind switch + { + "step" => ctx.StepAsync((_, ct) => Body(ct), name: "s", config: retry), + "child" => ctx.RunInChildContextAsync((_, ct) => Body(ct), name: "c"), + _ => ctx.WaitForConditionAsync( + async (state, _, ct) => { await Body(ct); return state; }, + new WaitForConditionConfig + { + InitialState = 0, + WaitStrategy = WaitStrategy.Fixed(TimeSpan.FromSeconds(1)), + }, + name: "w"), + }; + + // The operation must surface the error without recording a failure. A + // RETRY would suspend instead, so only wait for completion or suspension. + await Task.WhenAny(op, tm.TerminationTask).NoHang(); + await recorder.Batcher.DrainAsync(); + + Assert.True(op.IsFaulted, "The operation should rethrow the body's exception."); + // Observe the fault: other tests count process-wide unobserved exceptions. + Assert.NotNull(op.Exception); + Assert.DoesNotContain(recorder.Flushed, + u => u.Action == OperationAction.FAIL || u.Action == OperationAction.RETRY); + } + + [Fact] + public async Task StepAsync_TornDownWhileStartCheckpointIsAwaited_BodyDoesNotRun() + { + // AtMostOncePerRetry awaits its START checkpoint before the body. The + // workflow is torn down while that flush is in progress. + var flushEntered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var releaseFlush = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var tm = new TerminationManager(); + using var wfc = new WorkflowCancellation(tm); + await using var batcher = new CheckpointBatcher("t0", (_, _, _) => + { + flushEntered.TrySetResult(); + return releaseFlush.Task; + }); + var ctx = new DurableContext(new ExecutionState(), tm, wfc, new OperationIdGenerator(), + "arn:test", Context(TimeSpan.FromMinutes(5)), batcher); + var bodyRan = false; + + var step = ctx.StepAsync( + async (_, _) => { bodyRan = true; await Task.CompletedTask; return 1; }, + name: "s", + config: new StepConfig { Semantics = StepSemantics.AtMostOncePerRetry }); + + await flushEntered.Task.NoHang(); + wfc.BeginTeardown(); + releaseFlush.SetResult("t1"); + + await Assert.ThrowsAnyAsync(() => step.NoHang()); + Assert.False(bodyRan); + } + + [Fact] + public async Task Lmi_DeadlinePassesDuringFinalDrain_ThrowsTimeout() + { + // The workflow finishes in time; its final checkpoint is still in flight + // when the deadline passes and completes inside the cleanup budget. + var client = new FakeServiceClient + { + OnCheckpoint = _ => Task.Delay(400), + }; + + await Assert.ThrowsAsync(() => + DurableFunction.WrapAsync( + async (input, ctx) => + { + // Fire-and-forget START plus a SUCCEED queued at the very end. + ctx.StepAsync(async (_, _) => { await Task.CompletedTask; return "a"; }, name: "a").Observe(); + await Task.Yield(); + return "done"; + }, + Input("arn:late-drain"), + Context(TimeSpan.FromMilliseconds(100)), + client, + new DurableInvocationOptions { IsMultiConcurrency = true, CleanupTimeout = TimeSpan.FromSeconds(5) }) + .NoHang()); + } + + [Fact] + public async Task Lmi_HydrationIgnoringCancellation_IsStillBounded() + { + var never = new TaskCompletionSource(); + var client = new FakeServiceClient { OnGetExecutionState = _ => never.Task }; + + await Assert.ThrowsAsync(() => + DurableFunction.WrapAsync( + (input, _) => Task.FromResult(input.Id!), + Input("arn:hydrate-ignores-token", nextMarker: "page-2"), + Context(TimeSpan.FromMilliseconds(100)), + client, + LmiOptions).NoHang()); + } + + [Fact] + public async Task CheckpointBatcher_AfterAbandon_SendsNoFurtherBatches() + { + // The worker must send two batches in one pass: with one operation per + // batch, reading "a" and "b" together flushes [a] mid-pass, then [b]. + // The [a] flush ignores cancellation and stalls; the drain gives up and + // abandons the worker. Releasing [a] afterwards must not send [b]. + var calls = 0; + var warmupEntered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var releaseWarmup = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var stalledEntered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var releaseStalled = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var batcher = new CheckpointBatcher("t0", (_, _, _) => + { + switch (Interlocked.Increment(ref calls)) + { + case 1: warmupEntered.TrySetResult(); return releaseWarmup.Task; + case 2: stalledEntered.TrySetResult(); return releaseStalled.Task; + default: return Task.FromResult("t-next"); + } + }, new CheckpointBatcherConfig { MaxBatchOperations = 1 }); + + // Hold the worker in a warm-up flush so "a" and "b" queue up together. + batcher.EnqueueAsync(Update("w")).Observe(); + await warmupEntered.Task.NoHang(); + batcher.EnqueueAsync(Update("a")).Observe(); + var second = batcher.EnqueueAsync(Update("b")); + releaseWarmup.SetResult("t1"); + + await stalledEntered.Task.NoHang(); + Assert.False(await batcher.DrainAsync(TimeSpan.FromMilliseconds(50))); + releaseStalled.SetResult("t2"); + + await Assert.ThrowsAnyAsync(() => second.NoHang()); + Assert.Equal(2, Volatile.Read(ref calls)); + await batcher.DisposeAsync(); + } + + private static SdkOperationUpdate Update(string id) => new() + { + Id = id, + Type = OperationTypes.Step, + Action = OperationAction.START, + }; + [Fact] public async Task Lmi_DeadlinePassesDuringHydration_ThrowsTimeout() { @@ -305,7 +634,7 @@ public async Task Lmi_DeadlinePassesDuringHydration_ThrowsTimeout() Input("arn:deadline-hydrate", nextMarker: "page-2"), Context(TimeSpan.FromMilliseconds(100)), client, - LmiOptions).WaitAsync(HangGuard)); + LmiOptions).NoHang()); Assert.IsAssignableFrom(ex.InnerException); Assert.False(workflowRan); @@ -328,7 +657,7 @@ await Assert.ThrowsAsync(() => Context(TimeSpan.FromMilliseconds(100)), client, new DurableInvocationOptions { IsMultiConcurrency = true, CleanupTimeout = TimeSpan.FromMilliseconds(100) }) - .WaitAsync(HangGuard)); + .NoHang()); } [Fact] @@ -367,7 +696,7 @@ public async Task Lmi_ConcurrentInvocations_AreIsolated() client, LmiOptions))).ToArray(); - var outputs = await Task.WhenAll(tasks).WaitAsync(HangGuard); + var outputs = await Task.WhenAll(tasks).NoHang(); for (var i = 0; i < count; i++) { @@ -424,7 +753,7 @@ public async Task Lmi_PeerTimesOut_HealthyPeerIsUnaffected() client, LmiOptions)); - await healthyEntered.Task.WaitAsync(HangGuard); + await healthyEntered.Task.NoHang(); var timingOut = DurableFunction.WrapAsync( async (_, ctx) => await ctx.StepAsync(async (_, ct) => @@ -437,11 +766,11 @@ public async Task Lmi_PeerTimesOut_HealthyPeerIsUnaffected() client, LmiOptions); - await Assert.ThrowsAsync(() => timingOut.WaitAsync(HangGuard)); + await Assert.ThrowsAsync(() => timingOut.NoHang()); Assert.False(healthyToken.IsCancellationRequested); peerFinished.SetResult(); - var output = await healthy.WaitAsync(HangGuard); + var output = await healthy.NoHang(); Assert.Equal(InvocationStatus.Succeeded, output.Status); Assert.Equal("\"ok\"", output.Result); @@ -498,7 +827,7 @@ public async Task Lmi_ReplayWhilePeerIsLive_SkipsCompletedSteps() client, LmiOptions)); - var outputs = await Task.WhenAll(replay, fresh).WaitAsync(HangGuard); + var outputs = await Task.WhenAll(replay, fresh).NoHang(); Assert.Equal("\"cached+second\"", outputs[0].Result); Assert.Equal("\"fresh\"", outputs[1].Result); @@ -509,6 +838,9 @@ public async Task Lmi_ReplayWhilePeerIsLive_SkipsCompletedSteps() // ── Helpers ───────────────────────────────────────────────────────── + /// The configured hang guard, for . + internal static TimeSpan HangGuardTimeout => HangGuard; + /// /// One handler instance shared by every concurrent invocation. Every step /// waits until all invocations are inside a step, so they truly overlap. @@ -642,3 +974,34 @@ private sealed class FakeServiceClient : IDurableServiceClient } } } + +/// +/// Guards a test against hanging. Fails with an xUnit exception rather than +/// : the deadline tests assert on +/// , so Task.WaitAsync(TimeSpan) would turn +/// a hang into a false pass. +/// +internal static class HangGuardExtensions +{ + /// + /// Observes a task the test deliberately does not await, so its fault never + /// reaches TaskScheduler.UnobservedTaskException. Other tests in this + /// assembly count that process-wide event. + /// + public static void Observe(this Task task) => + task.ContinueWith(static t => _ = t.Exception, CancellationToken.None, + TaskContinuationOptions.OnlyOnFaulted | TaskContinuationOptions.ExecuteSynchronously, TaskScheduler.Default); + + public static async Task NoHang(this Task task) + { + if (await Task.WhenAny(task, Task.Delay(LambdaManagedInstancesTests.HangGuardTimeout)) != task) + throw new Xunit.Sdk.XunitException("The operation did not complete within the hang guard; it appears to hang."); + await task; + } + + public static async Task NoHang(this Task task) + { + await ((Task)task).NoHang(); + return await task; + } +} From 36bfc3b3240ac7b21bf27721517cfc3b4a447c3f Mon Sep 17 00:00:00 2001 From: Norm Johanson Date: Wed, 4 Nov 2026 01:46:27 +0000 Subject: [PATCH 3/6] Address second round of Copilot review findings - Treat an OCE as cancellation when the workflow token is cancelled even if the linked token has not been notified yet (step, child context, condition check), so a body that never ran is never recorded as failed. - Map/parallel: check the workflow token directly before recording the parent outcome, matching IncrementalParallelOperation. - Condition checks: do not record RETRY/FAIL/SUCCEED for a check that returned after the LMI deadline. - Own the deadline timer so a throwing deadline-token callback (for example from the service client) cannot crash the process. - Checkpoint batcher: cancel an abandoned flush with CancelAsync so blocking client callbacks cannot stall the drain, and release pending EnqueueAsync awaiters on abandonment. - Use one combined deadline + cleanup budget for the final drain. - Report a timeout when output serialization crosses the deadline. - Always observe an abandoned hydration request's fault. - Tests: regression test per fix; the deadline fault-race test registers on the deadline token so it exercises the race again; add a teardown-path test for throwing cancellation callbacks. --- .../DurableFunction.cs | 39 ++- .../Internal/CheckpointBatcher.cs | 49 ++- .../Internal/ChildContextOperation.cs | 3 +- .../Internal/ConcurrentOperation.cs | 6 + .../Internal/StepOperation.cs | 3 +- .../Internal/WaitForConditionOperation.cs | 26 +- .../Internal/WorkflowCancellation.cs | 38 ++- .../LambdaManagedInstancesTests.cs | 287 +++++++++++++++++- 8 files changed, 424 insertions(+), 27 deletions(-) diff --git a/Libraries/src/Amazon.Lambda.DurableExecution/DurableFunction.cs b/Libraries/src/Amazon.Lambda.DurableExecution/DurableFunction.cs index b9005d0ef..2d566ba76 100644 --- a/Libraries/src/Amazon.Lambda.DurableExecution/DurableFunction.cs +++ b/Libraries/src/Amazon.Lambda.DurableExecution/DurableFunction.cs @@ -206,8 +206,7 @@ private static async Task WrapAsyncCore TimeSpan.Zero ? untilDeadline : TimeSpan.Zero) + cleanupTimeout; + var drainTimeout = ComputeDrainTimeout(deadline.Value, cleanupTimeout, Stopwatch.GetElapsedTime(startTimestamp)); if (!await batcher.DrainAsync(drainTimeout)) throw CreateDeadlineException(invocationInput.DurableExecutionArn, innerException: null); } @@ -240,7 +239,35 @@ private static async Task WrapAsyncCore + /// How long the final checkpoint drain may run on Lambda Managed Instances: + /// what is left of one combined budget (the deadline plus the cleanup + /// budget), measured from invocation start. Time already spent past the + /// deadline comes out of the cleanup budget rather than restarting it. + /// + internal static TimeSpan ComputeDrainTimeout(TimeSpan deadline, TimeSpan cleanupTimeout, TimeSpan elapsed) + { + var remaining = deadline + cleanupTimeout - elapsed; + return remaining > TimeSpan.Zero ? remaining : TimeSpan.Zero; } /// @@ -257,8 +284,12 @@ private static async Task WaitWithDeadlineAsync(Task task, Cancellation { return await task.WaitAsync(deadlineToken).ConfigureAwait(false); } - catch (OperationCanceledException) when (!task.IsCompleted) + catch (OperationCanceledException) { + // Observe the request in every case: it may fault after WaitAsync chose + // cancellation but before this handler runs, in which case nothing else + // ever awaits it. Only a fault is observed; a request that honored the + // token completes cancelled and needs nothing. _ = task.ContinueWith( static t => _ = t.Exception, CancellationToken.None, diff --git a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/CheckpointBatcher.cs b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/CheckpointBatcher.cs index 824cf2358..6a42f564e 100644 --- a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/CheckpointBatcher.cs +++ b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/CheckpointBatcher.cs @@ -31,6 +31,11 @@ internal sealed class CheckpointBatcher : IAsyncDisposable private readonly Task _worker; private readonly CancellationTokenSource _shutdownCts = new(); + // Cancelled only when the worker is abandoned (Lambda Managed Instances + // cleanup budget exceeded). Never disposed: it is cheap, and abandoned + // awaiters may still be registered on it. + private readonly CancellationTokenSource _abandonCts = new(); + private string? _checkpointToken; private Exception? _terminalError; private int _disposed; @@ -117,9 +122,31 @@ public async Task EnqueueAsync(SdkOperationUpdate update, CancellationToken canc } // After a missing token this never completes (the invocation is suspended). - await tcs.Task.WaitAsync(cancellationToken).ConfigureAwait(false); + // It is also released when the worker is abandoned: a stalled flush may + // never return, and the waiting workflow must still be able to unwind + // rather than being kept alive by that request. + var completion = tcs.Task.WaitAsync(cancellationToken); + try + { + await completion.WaitAsync(_abandonCts.Token).ConfigureAwait(false); + } + catch (OperationCanceledException) when (_abandonCts.IsCancellationRequested && !completion.IsCompleted) + { + // The worker may still fail this item later; observe that so it never + // reaches TaskScheduler.UnobservedTaskException. + ObserveFault(completion); + throw new OperationCanceledException( + "The checkpoint was not confirmed before the invocation's cleanup budget ran out.", _abandonCts.Token); + } } + private static void ObserveFault(Task task) => + task.ContinueWith( + static t => _ = t.Exception, + CancellationToken.None, + TaskContinuationOptions.OnlyOnFaulted | TaskContinuationOptions.ExecuteSynchronously, + TaskScheduler.Default); + /// /// Closes the channel and awaits the worker. Any items already enqueued are /// flushed; any subsequent call throws. @@ -176,14 +203,18 @@ private void Abandon() { if (Interlocked.Exchange(ref _abandoned, 1) != 0) return; - // Ask the in-flight flush to stop. A flush that ignores the token keeps - // running until it returns; its outcome is no longer awaited. - try { _shutdownCts.Cancel(); } catch (ObjectDisposedException) { } - _worker.ContinueWith( - static t => _ = t.Exception, - CancellationToken.None, - TaskContinuationOptions.OnlyOnFaulted | TaskContinuationOptions.ExecuteSynchronously, - TaskScheduler.Default); + // Release every EnqueueAsync awaiter, in flight or queued. Only the SDK's + // own WaitAsync promises are registered on this token, so cancelling it + // synchronously runs no foreign code. + _abandonCts.Cancel(); + + // Ask the in-flight flush to stop. CancelAsync runs the callbacks the + // service client registered on the threadpool, so a callback that blocks + // or throws cannot stall the drain past its budget. A flush that ignores + // the token keeps running until it returns; its outcome is not awaited. + // The source is never disposed once abandoned: the worker may still read it. + ObserveFault(_shutdownCts.CancelAsync()); + ObserveFault(_worker); } public async ValueTask DisposeAsync() diff --git a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/ChildContextOperation.cs b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/ChildContextOperation.cs index ec375548d..1a4bf47e4 100644 --- a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/ChildContextOperation.cs +++ b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/ChildContextOperation.cs @@ -192,7 +192,8 @@ private async Task ExecuteFunc(CancellationToken cancellationToken) _workflowCancellation.Token.ThrowIfCancellationRequested(); result = await _func(childContext, linked.Token); } - catch (OperationCanceledException) when (linked.IsCancellationRequested) + catch (OperationCanceledException) when ( + linked.IsCancellationRequested || _workflowCancellation.Token.IsCancellationRequested) { // Cancellation owned by the linked source — caller cancel or workflow // shutdown. Do NOT checkpoint CONTEXT FAIL: the termination signal diff --git a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/ConcurrentOperation.cs b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/ConcurrentOperation.cs index 01f8c9547..e208ff007 100644 --- a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/ConcurrentOperation.cs +++ b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/ConcurrentOperation.cs @@ -340,6 +340,12 @@ void SignalShortCircuit() // Surfacing it here means a torn-down operation propagates an // OperationCanceledException instead of synthesizing a spurious // FailureToleranceExceeded verdict from units that merely unwound. + // + // Check the workflow token itself as well (as IncrementalParallelOperation + // does): controlToken is linked to it through a callback, so while + // cancellation is still being delivered it can lag behind. A unit that + // failed because of the cancellation must never reach the parent checkpoint. + _workflowCancellation.Token.ThrowIfCancellationRequested(); controlToken.ThrowIfCancellationRequested(); // Build BatchItems ONLY for units that were dispatched — matching the JS diff --git a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/StepOperation.cs b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/StepOperation.cs index 5a02b17e1..a0f22a7a9 100644 --- a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/StepOperation.cs +++ b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/StepOperation.cs @@ -239,7 +239,8 @@ private async Task ExecuteFunc(int attemptNumber, CancellationToken cancellat result = await _func(stepContext, linked.Token); } } - catch (OperationCanceledException) when (linked.IsCancellationRequested) + catch (OperationCanceledException) when ( + linked.IsCancellationRequested || _workflowCancellation.Token.IsCancellationRequested) { // Cancellation owned by the linked source (caller-cancel or workflow // shutdown). Do NOT checkpoint FAIL and do NOT consult the retry diff --git a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/WaitForConditionOperation.cs b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/WaitForConditionOperation.cs index 9649ec441..e7156069a 100644 --- a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/WaitForConditionOperation.cs +++ b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/WaitForConditionOperation.cs @@ -148,6 +148,14 @@ private Task ReplayReady(Operation ready, CancellationToken cancellation return ExecuteIteration(priorState, attemptNumber, existingStatus: ready.Status, cancellationToken); } + private void ThrowIfDeadlineExceeded() + { + if (_workflowCancellation.IsDeadlineExceeded) + throw new OperationCanceledException( + "The invocation deadline passed during the condition check; its outcome is not recorded.", + _workflowCancellation.Token); + } + private async Task ExecuteIteration(TState currentState, int attemptNumber, string? existingStatus, CancellationToken cancellationToken) { cancellationToken.ThrowIfCancellationRequested(); @@ -195,7 +203,8 @@ private async Task ExecuteIteration(TState currentState, int attemptNumb _workflowCancellation.Token.ThrowIfCancellationRequested(); newState = await _check(currentState, checkContext, linked.Token); } - catch (OperationCanceledException) when (linked.IsCancellationRequested) + catch (OperationCanceledException) when ( + linked.IsCancellationRequested || _workflowCancellation.Token.IsCancellationRequested) { // Cancellation owned by the linked source — caller cancel or workflow // shutdown. Do NOT checkpoint FAIL: the termination signal (or @@ -232,11 +241,23 @@ await EnqueueAsync(new SdkOperationUpdate }; } + // The invocation deadline passed (Lambda Managed Instances only) while the + // check ran. A check that swallowed the cancellation and returned a state + // did not complete a real poll: record no RETRY (which would consume an + // attempt) and no FAIL. The poll runs again when the invocation is retried. + ThrowIfDeadlineExceeded(); + WaitDecision decision; try { decision = _config.WaitStrategy.Decide(newState, attemptNumber); } + catch (WaitForConditionException) when (_workflowCancellation.IsDeadlineExceeded) + { + // The deadline passed while the strategy ran: never record FAIL for it. + ThrowIfDeadlineExceeded(); + throw; + } catch (WaitForConditionException maxEx) { // Strategy is signaling max-attempts reached. The strategy @@ -273,6 +294,9 @@ await EnqueueAsync(new SdkOperationUpdate throw enriched; } + // Decide is user code too; the deadline can pass while it runs. + ThrowIfDeadlineExceeded(); + if (!decision.ShouldContinue) { // Stop() means the condition has been met. Persist the final diff --git a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/WorkflowCancellation.cs b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/WorkflowCancellation.cs index 9d22e9707..a20846755 100644 --- a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/WorkflowCancellation.cs +++ b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/WorkflowCancellation.cs @@ -61,6 +61,7 @@ internal sealed class WorkflowCancellation : IDisposable private readonly CancellationToken _token; private readonly CancellationTokenSource? _deadlineCts; private readonly CancellationTokenRegistration _deadlineRegistration; + private readonly Timer? _deadlineTimer; private readonly TaskCompletionSource _deadlineTcs = new(TaskCreationOptions.RunContinuationsAsynchronously); private int _reason; private int _tornDown; @@ -117,16 +118,22 @@ public WorkflowCancellation(TerminationManager terminationManager, TimeSpan? dea if (deadline is { } remaining) { - // A zero deadline means the invocation was already out of time when - // it started: cancel at once rather than through the timer. - _deadlineCts = remaining > TimeSpan.Zero - ? new CancellationTokenSource(remaining) - : new CancellationTokenSource(); - if (remaining <= TimeSpan.Zero) _deadlineCts.Cancel(); + _deadlineCts = new CancellationTokenSource(); DeadlineToken = _deadlineCts.Token; _deadlineRegistration = _deadlineCts.Token.Register( static state => ((WorkflowCancellation)state!).OnDeadline(), this); + + // Own the timer rather than using CancellationTokenSource(TimeSpan): + // the deadline token is also handed to the service client, and if one + // of its cancellation callbacks threw on the CTS's internal timer the + // exception would crash the process. FireDeadline contains it. + // A zero deadline means the invocation was already out of time when it + // started, so it fires at once. + if (remaining > TimeSpan.Zero) + _deadlineTimer = new Timer(static state => ((WorkflowCancellation)state!).FireDeadline(), this, remaining, Timeout.InfiniteTimeSpan); + else + FireDeadline(); } } @@ -151,6 +158,24 @@ public void ThrowIfTornDown() "The durable invocation has ended; no new durable operations can be started.", _token); } + private void FireDeadline() + { + try { _deadlineCts!.Cancel(); } + catch (ObjectDisposedException) + { + // The timer fired while the invocation was being disposed. + } + catch (AggregateException ex) + { + // A callback registered on the deadline token (for example by the + // service client) threw. Cancel ran every other callback, including + // OnDeadline, before throwing. + LambdaLogger.Log( + "Warning", ex, + "A cancellation callback registered on the durable invocation deadline threw."); + } + } + private void OnDeadline() { // Stop admitting new operations and signal the deadline BEFORE cancelling @@ -197,6 +222,7 @@ public void Dispose() // that is still running, and a user cancellation callback that blocks // would then hang the invocation here. _deadlineRegistration.Unregister(); + _deadlineTimer?.Dispose(); _deadlineCts?.Dispose(); _cts.Dispose(); } diff --git a/Libraries/test/Amazon.Lambda.DurableExecution.Tests/LambdaManagedInstancesTests.cs b/Libraries/test/Amazon.Lambda.DurableExecution.Tests/LambdaManagedInstancesTests.cs index 53236877a..ae6d52383 100644 --- a/Libraries/test/Amazon.Lambda.DurableExecution.Tests/LambdaManagedInstancesTests.cs +++ b/Libraries/test/Amazon.Lambda.DurableExecution.Tests/LambdaManagedInstancesTests.cs @@ -266,26 +266,30 @@ await Assert.ThrowsAsync(() => public async Task RunAsync_WorkflowFaultsFromDeadlineBeforeDeadlineTaskCompletes_ReportsDeadline() { // Forces the race where the workflow task faults before DeadlineTask - // completes: the fault is raised synchronously inside the token's - // cancellation callback, which runs before the deadline task is - // signalled. The fault must still be reported as a deadline, not as a - // durable failure. + // completes. The workflow registers on the deadline token before the + // deadline fires; that callback runs before WorkflowCancellation's own + // (callbacks run in reverse order), so the workflow faults while the + // deadline has passed but DeadlineTask is not yet signalled. The fault + // must still be reported as a deadline, not as a durable failure. var tm = new TerminationManager(); using var wfc = new WorkflowCancellation(tm, TimeSpan.FromMilliseconds(50)); + var workflowRan = false; var result = await DurableExecutionHandler.RunAsync( new ExecutionState(), tm, async () => { + workflowRan = true; // No RunContinuationsAsynchronously: the await below resumes // inline inside the cancellation callback. var faulted = new TaskCompletionSource(); - wfc.Token.Register(() => faulted.TrySetException(new InvalidOperationException("user gave up"))); + wfc.DeadlineToken.Register(() => faulted.TrySetException(new InvalidOperationException("user gave up"))); return await faulted.Task; }, wfc).NoHang(); + Assert.True(workflowRan); Assert.True(result.DeadlineExceeded); Assert.Null(result.Exception); } @@ -348,6 +352,24 @@ public async Task WorkflowCancellation_ThrowingCallback_DoesNotEscapeDeadlineOrT wfc.Dispose(); } + [Fact] + public void WorkflowCancellation_ThrowingCallback_DoesNotEscapeTeardown() + { + // No deadline: the first cancel comes from the wrapper's teardown, as on + // every invocation. A throwing user callback must not replace the + // invocation's outcome. + var wfc = new WorkflowCancellation(new TerminationManager()); + var otherCallbackRan = false; + wfc.Token.Register(() => otherCallbackRan = true); + wfc.Token.Register(() => throw new InvalidOperationException("user callback threw")); + + wfc.BeginTeardown(); + wfc.Dispose(); + + Assert.True(wfc.Token.IsCancellationRequested); + Assert.True(otherCallbackRan); + } + [Fact] public async Task WorkflowCancellation_BlockingCallback_DoesNotDelayDeadlineSignal() { @@ -612,6 +634,261 @@ public async Task CheckpointBatcher_AfterAbandon_SendsNoFurtherBatches() await batcher.DisposeAsync(); } + // ── Copilot review round 2 ───────────────────────────────────────── + + [Fact] + public async Task CheckpointBatcher_Abandon_DoesNotWaitForBlockingClientCallback() + { + // The in-flight flush registers a cancellation callback that blocks. The + // drain must still give up once its budget is spent. + using var releaseCallback = new ManualResetEventSlim(); + var flushEntered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var never = new TaskCompletionSource(); + var batcher = new CheckpointBatcher("t0", (_, _, ct) => + { + ct.Register(() => releaseCallback.Wait()); + flushEntered.TrySetResult(); + return never.Task; + }); + + try + { + batcher.EnqueueAsync(Update("a")).Observe(); + await flushEntered.Task.NoHang(); + + Assert.False(await batcher.DrainAsync(TimeSpan.FromMilliseconds(50)).NoHang()); + } + finally + { + releaseCallback.Set(); + } + } + + [Fact] + public async Task CheckpointBatcher_Abandon_ReleasesPendingAwaiters() + { + // A flush that never returns would otherwise keep every waiting caller + // (in flight or queued) suspended forever. + var flushEntered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var never = new TaskCompletionSource(); + var batcher = new CheckpointBatcher("t0", (_, _, _) => + { + flushEntered.TrySetResult(); + return never.Task; + }); + + var inFlight = batcher.EnqueueAsync(Update("a")); + await flushEntered.Task.NoHang(); + var queued = batcher.EnqueueAsync(Update("b")); + + Assert.False(await batcher.DrainAsync(TimeSpan.FromMilliseconds(50))); + + await Assert.ThrowsAnyAsync(() => inFlight.NoHang()); + await Assert.ThrowsAnyAsync(() => queued.NoHang()); + } + + [Fact] + public async Task StepAsync_WorkflowCancelledButLinkedTokenLagging_NoFailCheckpoint() + { + // Cancellation is delivered on another thread and stalls in a callback + // registered after the step's linked token, so the workflow token is + // cancelled while the linked one is not yet. The pre-body check then + // throws; that must count as cancellation, never as a step failure. + // A custom logger's BeginScope runs after the linked token is created and + // before the check, which is where the test arranges the lag. + var tm = new TerminationManager(); + using var wfc = new WorkflowCancellation(tm); + var recorder = new RecordingBatcher(); + var ctx = new DurableContext(new ExecutionState(), tm, wfc, new OperationIdGenerator(), + "arn:test", Context(TimeSpan.FromMinutes(5)), recorder.Batcher); + using var releaseDelivery = new ManualResetEventSlim(); + using var deliveryStalled = new ManualResetEventSlim(); + ctx.ConfigureLogger(new LoggerConfig + { + CustomLogger = new ScopeHookLogger(() => + { + wfc.Token.Register(() => { deliveryStalled.Set(); releaseDelivery.Wait(); }); + // Suspension: TerminationManager delivers the cancel on the threadpool. + tm.Terminate(TerminationReason.WaitScheduled); + Assert.True(deliveryStalled.Wait(LambdaManagedInstancesTests.HangGuardTimeout)); + }), + }); + var bodyRan = false; + + try + { + var step = ctx.StepAsync( + async (_, _) => { bodyRan = true; await Task.CompletedTask; return 1; }, name: "s"); + await Assert.ThrowsAnyAsync(() => step.NoHang()); + } + finally + { + releaseDelivery.Set(); + } + + await recorder.Batcher.DrainAsync(); + Assert.False(bodyRan); + Assert.DoesNotContain(recorder.Flushed, + u => u.Action == OperationAction.FAIL || u.Action == OperationAction.RETRY); + } + + [Fact] + public async Task ParallelAsync_BranchFailsWhileDeadlineDeliveryLags_NoParentCheckpoint() + { + // Deadline delivery stalls in a callback registered last, so it runs + // first and holds back every other callback, including the parallel's + // control token. The branches fail (converting the cancellation) before + // the control token is notified. The parent must not checkpoint an + // outcome that contains those failures. + var tm = new TerminationManager(); + using var wfc = new WorkflowCancellation(tm, TimeSpan.FromMilliseconds(50)); + var recorder = new RecordingBatcher(); + var ctx = new DurableContext(new ExecutionState(), tm, wfc, new OperationIdGenerator(), + "arn:test", Context(TimeSpan.FromMinutes(5)), recorder.Batcher); + using var releaseDelivery = new ManualResetEventSlim(); + var deliveryStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var entered = new CountdownEvent(2); + + var branches = new Func>[2]; + for (var i = 0; i < 2; i++) + { + branches[i] = async (_, _) => + { + entered.Signal(); + await deliveryStarted.Task; + throw new InvalidOperationException("gave up"); + }; + } + + try + { + var run = ctx.ParallelAsync(branches, name: "fanout", + config: new ParallelConfig { CompletionConfig = CompletionConfig.AllCompleted() }); + Assert.True(entered.Wait(LambdaManagedInstancesTests.HangGuardTimeout)); + // Registered after the parallel's control token, so it runs before it. + wfc.Token.Register(() => { deliveryStarted.TrySetResult(); releaseDelivery.Wait(); }); + + await Assert.ThrowsAnyAsync(() => run.NoHang()); + } + finally + { + releaseDelivery.Set(); + } + + await recorder.Batcher.DrainAsync(); + Assert.DoesNotContain(recorder.Flushed, + u => u.Type == OperationTypes.Context && u.SubType == OperationSubTypes.Parallel + && (u.Action == OperationAction.SUCCEED || u.Action == OperationAction.FAIL)); + } + + [Fact] + public async Task WaitForCondition_CheckReturnsStateAfterDeadline_RecordsNoRetryOrFail() + { + var tm = new TerminationManager(); + using var wfc = new WorkflowCancellation(tm, TimeSpan.FromMilliseconds(50)); + var recorder = new RecordingBatcher(); + var ctx = new DurableContext(new ExecutionState(), tm, wfc, new OperationIdGenerator(), + "arn:test", Context(TimeSpan.FromMinutes(5)), recorder.Batcher); + + // The check swallows the cancellation and returns a state. + var op = ctx.WaitForConditionAsync( + async (state, _, ct) => + { + try { await Task.Delay(Timeout.Infinite, ct); } + catch (OperationCanceledException) { } + return state + 1; + }, + new WaitForConditionConfig + { + InitialState = 0, + WaitStrategy = WaitStrategy.Fixed(TimeSpan.FromSeconds(1)), + }, + name: "w"); + + await Task.WhenAny(op, tm.TerminationTask).NoHang(); + await recorder.Batcher.DrainAsync(); + + Assert.True(op.IsCanceled || op.IsFaulted, "The check outcome must not be accepted after the deadline."); + if (op.IsFaulted) Assert.NotNull(op.Exception); + Assert.DoesNotContain(recorder.Flushed, + u => u.Action == OperationAction.FAIL || u.Action == OperationAction.RETRY || u.Action == OperationAction.SUCCEED); + } + + [Fact] + public async Task WorkflowCancellation_ThrowingDeadlineTokenCallback_IsContained() + { + // The deadline token is also given to the service client. A callback it + // registers that throws must not escape the deadline timer (which would + // crash the test host process, as it would an LMI environment). + var tm = new TerminationManager(); + using var wfc = new WorkflowCancellation(tm, TimeSpan.FromMilliseconds(50)); + wfc.DeadlineToken.Register(() => throw new InvalidOperationException("client callback threw")); + + await wfc.DeadlineTask.NoHang(); + Assert.True(wfc.Token.IsCancellationRequested); + } + + [Theory] + [InlineData(10_000, 2_000, 4_000, 8_000)] // before the deadline: rest of it plus cleanup + [InlineData(10_000, 2_000, 11_500, 500)] // 1.5s past the deadline: 0.5s of cleanup left + [InlineData(10_000, 2_000, 13_000, 0)] // budget spent + public void DrainTimeout_IsTheRemainingCombinedBudget(int deadlineMs, int cleanupMs, int elapsedMs, int expectedMs) + { + var timeout = DurableFunction.ComputeDrainTimeout( + TimeSpan.FromMilliseconds(deadlineMs), TimeSpan.FromMilliseconds(cleanupMs), TimeSpan.FromMilliseconds(elapsedMs)); + Assert.Equal(TimeSpan.FromMilliseconds(expectedMs), timeout); + } + + [Fact] + public async Task Lmi_OutputSerializationCrossesDeadline_ThrowsTimeout() + { + var client = new FakeServiceClient(); + var context = Context(TimeSpan.FromMilliseconds(100)); + // Serializing the output takes longer than the time left. A sleep is used + // only to make the deadline pass; nothing asserts on elapsed time. + context.Serializer = new SlowSerializer(TimeSpan.FromMilliseconds(400)); + + await Assert.ThrowsAsync(() => + DurableFunction.WrapAsync( + (_, _) => Task.FromResult("done"), + Input("arn:slow-output"), + context, + client, + LmiOptions).NoHang()); + } + + /// Delegates to the default serializer; Serialize sleeps first. + private sealed class SlowSerializer : Amazon.Lambda.Core.ILambdaSerializer + { + private readonly DefaultLambdaJsonSerializer _inner = new(); + private readonly TimeSpan _delay; + public SlowSerializer(TimeSpan delay) => _delay = delay; + public T Deserialize(Stream requestStream) => _inner.Deserialize(requestStream); + public void Serialize(T response, Stream responseStream) + { + Thread.Sleep(_delay); + _inner.Serialize(response, responseStream); + } + } + + /// A logger whose first BeginScope call runs a hook. + private sealed class ScopeHookLogger : Microsoft.Extensions.Logging.ILogger + { + private Action? _hook; + public ScopeHookLogger(Action hook) => _hook = hook; + + public IDisposable? BeginScope(TState state) where TState : notnull + { + Interlocked.Exchange(ref _hook, null)?.Invoke(); + return null; + } + + public bool IsEnabled(Microsoft.Extensions.Logging.LogLevel logLevel) => false; + + public void Log(Microsoft.Extensions.Logging.LogLevel logLevel, Microsoft.Extensions.Logging.EventId eventId, + TState state, Exception? exception, Func formatter) { } + } + private static SdkOperationUpdate Update(string id) => new() { Id = id, From 904f209e59c4fe0f9b0e258f9671b90f3d3fc6dd Mon Sep 17 00:00:00 2001 From: Norm Johanson Date: Wed, 4 Nov 2026 02:43:14 +0000 Subject: [PATCH 4/6] Never run foreign cancellation callbacks on SDK threads Third round of Copilot review findings, addressed with one rule instead of per-site patches: the SDK never runs cancellation callbacks it does not own (user code, the service client) on its own threads, and it publishes its own deadline and teardown state before any foreign code runs. - Every cancel the SDK starts uses CancellationTokenSource.CancelAsync: tokens report cancellation immediately, callbacks run on the threadpool, and their exceptions are logged. BeginTeardown can no longer block in a user callback, in any mode. - A short lock around starting the cancel ensures every caller returns only after the token reports cancellation. - Sources are disposed only after callback delivery finishes, so no callback is skipped. - The deadline is anchored to a monotonic timestamp taken when RemainingTime was read. IsDeadlineExceeded and admission compare the clock, and fire the deadline themselves if the timer is late. - Firing the deadline records the reason, closes admission and signals DeadlineTask before cancelling the workflow token, then the client's deadline token. Hydration races DeadlineTask, not the client token. - Map/parallel parents check the deadline before recording an outcome. - Condition checks re-check the deadline after serializing state, on the RETRY, SUCCEED and max-attempts FAIL paths. - Checkpoint batcher always observes an abandoned item's fault. - Tests: deadline-dependent tests fire the deadline themselves after setup (TriggerDeadline) instead of racing a real 50 ms timer, and a test-only option leaves the timer unarmed to model a late timer. --- .../DurableFunction.cs | 48 +-- .../Internal/CheckpointBatcher.cs | 7 +- .../Internal/ConcurrentOperation.cs | 4 + .../Internal/IncrementalParallelOperation.cs | 4 + .../Internal/WaitForConditionOperation.cs | 27 +- .../Internal/WorkflowCancellation.cs | 276 ++++++++++----- .../docs/core/cancellation.md | 1 + .../LambdaManagedInstancesTests.cs | 326 +++++++++++++++--- 8 files changed, 536 insertions(+), 157 deletions(-) diff --git a/Libraries/src/Amazon.Lambda.DurableExecution/DurableFunction.cs b/Libraries/src/Amazon.Lambda.DurableExecution/DurableFunction.cs index 2d566ba76..1fd784073 100644 --- a/Libraries/src/Amazon.Lambda.DurableExecution/DurableFunction.cs +++ b/Libraries/src/Amazon.Lambda.DurableExecution/DurableFunction.cs @@ -117,7 +117,8 @@ private static async Task WrapAsyncCore WrapAsyncCore - /// Awaits , but stops waiting when - /// is cancelled. A task that is abandoned this - /// way has its eventual fault observed. With no deadline this is a plain await. + /// Awaits , but stops waiting when the invocation + /// deadline passes, and then throws . + /// Races , which the SDK + /// signals before any foreign callback runs, rather than the deadline token: + /// a blocking callback the client registered on that token could otherwise + /// delay the wait. A task that is abandoned this way has its eventual fault + /// observed. With no deadline this is a plain await. /// - private static async Task WaitWithDeadlineAsync(Task task, CancellationToken deadlineToken) + internal static async Task WaitWithDeadlineAsync(Task task, WorkflowCancellation workflowCancellation) { - if (!deadlineToken.CanBeCanceled) + if (!workflowCancellation.DeadlineToken.CanBeCanceled) return await task.ConfigureAwait(false); - try - { - return await task.WaitAsync(deadlineToken).ConfigureAwait(false); - } - catch (OperationCanceledException) - { - // Observe the request in every case: it may fault after WaitAsync chose - // cancellation but before this handler runs, in which case nothing else - // ever awaits it. Only a fault is observed; a request that honored the - // token completes cancelled and needs nothing. - _ = task.ContinueWith( - static t => _ = t.Exception, - CancellationToken.None, - TaskContinuationOptions.OnlyOnFaulted | TaskContinuationOptions.ExecuteSynchronously, - TaskScheduler.Default); - throw; - } + if (await Task.WhenAny(task, workflowCancellation.DeadlineTask).ConfigureAwait(false) == task) + return await task.ConfigureAwait(false); + + // Nothing else ever awaits the abandoned request; observe its fault. + // A request that honored the token completes cancelled and needs nothing. + _ = task.ContinueWith( + static t => _ = t.Exception, + CancellationToken.None, + TaskContinuationOptions.OnlyOnFaulted | TaskContinuationOptions.ExecuteSynchronously, + TaskScheduler.Default); + throw new OperationCanceledException("The invocation deadline passed while loading execution state.", + workflowCancellation.DeadlineToken); } /// diff --git a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/CheckpointBatcher.cs b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/CheckpointBatcher.cs index 6a42f564e..90c8c0cbe 100644 --- a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/CheckpointBatcher.cs +++ b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/CheckpointBatcher.cs @@ -130,10 +130,11 @@ public async Task EnqueueAsync(SdkOperationUpdate update, CancellationToken canc { await completion.WaitAsync(_abandonCts.Token).ConfigureAwait(false); } - catch (OperationCanceledException) when (_abandonCts.IsCancellationRequested && !completion.IsCompleted) + catch (OperationCanceledException) when (_abandonCts.IsCancellationRequested) { - // The worker may still fail this item later; observe that so it never - // reaches TaskScheduler.UnobservedTaskException. + // The worker may fail this item later, or may have just failed it after + // the abandonment won the wait. Either way nothing else awaits it, so + // always observe it; it never reaches TaskScheduler.UnobservedTaskException. ObserveFault(completion); throw new OperationCanceledException( "The checkpoint was not confirmed before the invocation's cleanup budget ran out.", _abandonCts.Token); diff --git a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/ConcurrentOperation.cs b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/ConcurrentOperation.cs index e208ff007..0e023ebe0 100644 --- a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/ConcurrentOperation.cs +++ b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/ConcurrentOperation.cs @@ -345,6 +345,10 @@ void SignalShortCircuit() // does): controlToken is linked to it through a callback, so while // cancellation is still being delivered it can lag behind. A unit that // failed because of the cancellation must never reach the parent checkpoint. + // The deadline is checked against the clock too, in case the deadline + // timer has not run yet when units failed because time ran out. + _workflowCancellation.ThrowIfDeadlineExceeded( + "The invocation deadline passed before the operation completed; its outcome is not recorded."); _workflowCancellation.Token.ThrowIfCancellationRequested(); controlToken.ThrowIfCancellationRequested(); diff --git a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/IncrementalParallelOperation.cs b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/IncrementalParallelOperation.cs index 295d23e8f..72cbe74a8 100644 --- a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/IncrementalParallelOperation.cs +++ b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/IncrementalParallelOperation.cs @@ -556,6 +556,10 @@ private async Task CompleteCoreAsync(CancellationToken cancellatio } // A torn-down operation propagates cancellation rather than a synthesized verdict. + // The deadline is checked against the clock too, in case the deadline timer + // has not run yet when branches failed because time ran out. + _workflowCancellation.ThrowIfDeadlineExceeded( + "The invocation deadline passed before the parallel operation completed; its outcome is not recorded."); _workflowCancellation.Token.ThrowIfCancellationRequested(); cancellationToken.ThrowIfCancellationRequested(); diff --git a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/WaitForConditionOperation.cs b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/WaitForConditionOperation.cs index e7156069a..4c5bff500 100644 --- a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/WaitForConditionOperation.cs +++ b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/WaitForConditionOperation.cs @@ -148,12 +148,20 @@ private Task ReplayReady(Operation ready, CancellationToken cancellation return ExecuteIteration(priorState, attemptNumber, existingStatus: ready.Status, cancellationToken); } - private void ThrowIfDeadlineExceeded() + private void ThrowIfDeadlineExceeded() => + _workflowCancellation.ThrowIfDeadlineExceeded( + "The invocation deadline passed during the condition check; its outcome is not recorded."); + + /// + /// Serializes for a checkpoint, then re-checks the + /// deadline: serialization runs user serializer code, which can take long + /// enough to cross it. Used on every path that records an outcome. + /// + private string SerializeStateForCheckpoint(TState state) { - if (_workflowCancellation.IsDeadlineExceeded) - throw new OperationCanceledException( - "The invocation deadline passed during the condition check; its outcome is not recorded.", - _workflowCancellation.Token); + var serialized = SerializeState(state); + ThrowIfDeadlineExceeded(); + return serialized; } private async Task ExecuteIteration(TState currentState, int attemptNumber, string? existingStatus, CancellationToken cancellationToken) @@ -275,6 +283,7 @@ await EnqueueAsync(new SdkOperationUpdate // that hits this cached FAIL can reconstruct LastState identically // to the live throw. The wire protocol forbids a Payload on FAIL // updates, so ErrorData is the only field that survives replay. + var lastStateData = SerializeStateForCheckpoint(newState); await EnqueueAsync(new SdkOperationUpdate { Id = OperationId, @@ -287,7 +296,7 @@ await EnqueueAsync(new SdkOperationUpdate { ErrorType = typeof(WaitForConditionException).FullName, ErrorMessage = enriched.Message, - ErrorData = SerializeState(newState) + ErrorData = lastStateData } }, cancellationToken); @@ -301,6 +310,7 @@ await EnqueueAsync(new SdkOperationUpdate { // Stop() means the condition has been met. Persist the final // state and return it to the caller. + var finalState = SerializeStateForCheckpoint(newState); await EnqueueAsync(new SdkOperationUpdate { Id = OperationId, @@ -309,7 +319,7 @@ await EnqueueAsync(new SdkOperationUpdate Action = OperationAction.SUCCEED, SubType = OperationSubTypes.WaitForCondition, Name = Name, - Payload = SerializeState(newState) + Payload = finalState }, cancellationToken); return newState; @@ -320,6 +330,7 @@ await EnqueueAsync(new SdkOperationUpdate // service definitely has the new state and timer scheduled before // we suspend. var delaySeconds = (int)Math.Max(1, Math.Ceiling(decision.Delay.TotalSeconds)); + var nextState = SerializeStateForCheckpoint(newState); await EnqueueAsync(new SdkOperationUpdate { Id = OperationId, @@ -328,7 +339,7 @@ await EnqueueAsync(new SdkOperationUpdate Action = OperationAction.RETRY, SubType = OperationSubTypes.WaitForCondition, Name = Name, - Payload = SerializeState(newState), + Payload = nextState, StepOptions = new SdkStepOptions { NextAttemptDelaySeconds = delaySeconds } }, cancellationToken); diff --git a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/WorkflowCancellation.cs b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/WorkflowCancellation.cs index a20846755..a1cfa40e0 100644 --- a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/WorkflowCancellation.cs +++ b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/WorkflowCancellation.cs @@ -1,6 +1,7 @@ // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // SPDX-License-Identifier: Apache-2.0 +using System.Diagnostics; using Amazon.Lambda.Core; namespace Amazon.Lambda.DurableExecution.Internal; @@ -42,16 +43,28 @@ internal enum WorkflowCancellationReason /// being torn down. /// /// -/// Lifetime rules that matter when several invocations share one process -/// (Lambda Managed Instances): +/// Rules that matter when several invocations share one process (Lambda +/// Managed Instances), where a stuck SDK thread blocks a runtime processing loop: /// /// -/// and cancel the token -/// synchronously before releasing anything. Termination delivers its cancel -/// through a threadpool continuation, which could otherwise run after -/// disposal and be lost. -/// is captured at construction, so abandoned user code -/// that reads it after disposal gets a cancelled token instead of an +/// The SDK never runs cancellation callbacks it does not own on its own +/// threads. User code and the service client register callbacks on these +/// tokens (directly, or through linked tokens), and a callback can block or +/// throw. Every cancel the SDK starts uses +/// : the token reports +/// IsCancellationRequested immediately, and callbacks run on the +/// threadpool. Their exceptions are observed and logged. +/// The SDK's own decisions are published before any foreign code runs: +/// admission is closed, the deadline is recorded and +/// is signalled first. +/// The deadline is a monotonic timestamp, not just a timer. A timer +/// callback can run late when the threadpool is busy, so +/// also compares the clock, and fires the +/// deadline itself when the timer has not yet. +/// Cancellation sources are disposed only after callback delivery has +/// finished, so no registered callback is ever skipped. +/// is captured at construction, so abandoned user code that reads it after +/// disposal gets a cancelled token instead of an /// . /// /// @@ -60,11 +73,15 @@ internal sealed class WorkflowCancellation : IDisposable private readonly CancellationTokenSource _cts = new(); private readonly CancellationToken _token; private readonly CancellationTokenSource? _deadlineCts; - private readonly CancellationTokenRegistration _deadlineRegistration; + private readonly long _deadlineTimestamp; private readonly Timer? _deadlineTimer; private readonly TaskCompletionSource _deadlineTcs = new(TaskCreationOptions.RunContinuationsAsynchronously); + private readonly object _cancelLock = new(); + private Task? _cancelDelivery; + private Task? _deadlineDelivery; private int _reason; private int _tornDown; + private int _deadlineFired; private int _disposed; public CancellationToken Token => _token; @@ -75,37 +92,70 @@ internal sealed class WorkflowCancellation : IDisposable public WorkflowCancellationReason Reason => (WorkflowCancellationReason)Volatile.Read(ref _reason); /// - /// True once the invocation deadline passed. Read from the deadline token - /// itself, so it is accurate even when a suspension or teardown cancelled the - /// workflow first, and becomes true before any cancellation callback runs. + /// True once the invocation deadline passed. Compares the clock as well as + /// checking whether the deadline fired, so it is accurate even when the + /// threadpool delays the timer or a suspension cancelled the workflow first. + /// When the clock shows the deadline has passed but the timer has not run, + /// this fires the deadline, so every observer sees one consistent outcome. /// - public bool IsDeadlineExceeded => DeadlineToken.IsCancellationRequested; + public bool IsDeadlineExceeded + { + get + { + if (_deadlineCts is null) return false; + if (Volatile.Read(ref _deadlineFired) == 1) return true; + if (Stopwatch.GetTimestamp() < _deadlineTimestamp) return false; + TriggerDeadline(); + return true; + } + } /// /// True once the invocation is being torn down (the wrapper finished or the /// deadline passed). New durable operations are rejected from this point. /// - public bool IsTornDown => Volatile.Read(ref _tornDown) == 1; + public bool IsTornDown => Volatile.Read(ref _tornDown) == 1 || IsDeadlineExceeded; /// - /// Cancels only when the invocation deadline passes. Passed to SDK-owned - /// service calls (state hydration). when - /// there is no deadline, so on-demand Lambda calls are unchanged. + /// Cancels when the invocation deadline passes. Passed to SDK-owned service + /// calls (state hydration). when there is + /// no deadline, so on-demand Lambda calls are unchanged. Do not wait on this + /// token to learn the deadline passed: its callbacks belong to the client and + /// may be delayed. Use or . /// public CancellationToken DeadlineToken { get; } /// - /// Completes when the invocation deadline passes. Never completes when there - /// is no deadline. + /// Completes when the invocation deadline passes, before any foreign callback + /// runs. Never completes when there is no deadline. /// public Task DeadlineTask => _deadlineTcs.Task; + /// + /// Completes once every callback registered on has run. + /// Already complete before the token is cancelled. Exposed for tests and + /// for disposal ordering. + /// + internal Task CancellationDelivered => Volatile.Read(ref _cancelDelivery) ?? Task.CompletedTask; + /// Suspension signal for this invocation. /// /// Time until the invocation deadline, or null for no deadline. Only /// set for Lambda Managed Instances; see . /// - public WorkflowCancellation(TerminationManager terminationManager, TimeSpan? deadline = null) + /// + /// The at which + /// was measured. Time elapsed since then is subtracted, so the deadline does + /// not drift by however long it took to get here. Defaults to now. + /// + /// + /// Test seam. false leaves the deadline to fire only through + /// or the clock check in + /// , which models a timer the threadpool has + /// not run yet. Production code always arms the timer. + /// + public WorkflowCancellation(TerminationManager terminationManager, TimeSpan? deadline = null, long? anchorTimestamp = null, + bool armTimer = true) { _token = _cts.Token; @@ -120,26 +170,28 @@ public WorkflowCancellation(TerminationManager terminationManager, TimeSpan? dea { _deadlineCts = new CancellationTokenSource(); DeadlineToken = _deadlineCts.Token; - _deadlineRegistration = _deadlineCts.Token.Register( - static state => ((WorkflowCancellation)state!).OnDeadline(), - this); - - // Own the timer rather than using CancellationTokenSource(TimeSpan): - // the deadline token is also handed to the service client, and if one - // of its cancellation callbacks threw on the CTS's internal timer the - // exception would crash the process. FireDeadline contains it. - // A zero deadline means the invocation was already out of time when it - // started, so it fires at once. - if (remaining > TimeSpan.Zero) - _deadlineTimer = new Timer(static state => ((WorkflowCancellation)state!).FireDeadline(), this, remaining, Timeout.InfiniteTimeSpan); + + var anchor = anchorTimestamp ?? Stopwatch.GetTimestamp(); + _deadlineTimestamp = anchor + (long)(remaining.TotalSeconds * Stopwatch.Frequency); + var untilDeadline = remaining - Stopwatch.GetElapsedTime(anchor); + + // A deadline already reached (including one that ran out between + // reading RemainingTime and getting here) fires at once. + if (!armTimer) + { + // Test seam: nothing fires the deadline on its own. + } + else if (untilDeadline > TimeSpan.Zero) + _deadlineTimer = new Timer(static state => ((WorkflowCancellation)state!).TriggerDeadline(), this, untilDeadline, Timeout.InfiniteTimeSpan); else - FireDeadline(); + TriggerDeadline(); } } /// - /// Marks the invocation as torn down and cancels the token synchronously, - /// so the cancel is delivered before the caller releases any resources. + /// Marks the invocation as torn down and cancels the token. Returns without + /// running user cancellation callbacks: they run on the threadpool, so one + /// that blocks cannot stall the wrapper before its bounded drain. /// public void BeginTeardown() { @@ -158,58 +210,96 @@ public void ThrowIfTornDown() "The durable invocation has ended; no new durable operations can be started.", _token); } - private void FireDeadline() + /// + /// Throws when the invocation + /// deadline has passed, so a timed-out invocation records no outcome. + /// + public void ThrowIfDeadlineExceeded(string message) { - try { _deadlineCts!.Cancel(); } - catch (ObjectDisposedException) - { - // The timer fired while the invocation was being disposed. - } - catch (AggregateException ex) - { - // A callback registered on the deadline token (for example by the - // service client) threw. Cancel ran every other callback, including - // OnDeadline, before throwing. - LambdaLogger.Log( - "Warning", ex, - "A cancellation callback registered on the durable invocation deadline threw."); - } + if (IsDeadlineExceeded) + throw new OperationCanceledException(message, _token); } - private void OnDeadline() + /// + /// Fires the deadline now: by the timer, by an observer that saw the clock + /// pass it, or by a test. Idempotent. Runs no foreign code on the caller's + /// thread, so it is safe to call from anywhere. + /// + internal void TriggerDeadline() { - // Stop admitting new operations and signal the deadline BEFORE cancelling - // the token. Cancel runs user cancellation callbacks synchronously; a - // callback that blocks must not stop DurableExecutionHandler from seeing - // the deadline. _deadlineTcs runs its continuations asynchronously, so - // signalling it never runs the handler on this timer thread. + if (_deadlineCts is null) return; + if (Interlocked.Exchange(ref _deadlineFired, 1) != 0) return; + + // The SDK's own state first: record why, close admission, and signal the + // handler. _deadlineTcs runs its continuations asynchronously, so this + // never runs the handler on the calling thread. Interlocked.CompareExchange(ref _reason, (int)WorkflowCancellationReason.DeadlineExceeded, 0); Volatile.Write(ref _tornDown, 1); _deadlineTcs.TrySetResult(); + + // Then the workflow token, then the client's deadline token. Both only + // queue their callbacks. Cancel(WorkflowCancellationReason.DeadlineExceeded); + StartCancelOnce(ref _deadlineDelivery, _deadlineCts, + "A cancellation callback registered on the durable invocation deadline threw."); } private void Cancel(WorkflowCancellationReason reason) { Interlocked.CompareExchange(ref _reason, (int)reason, 0); - try { _cts.Cancel(); } - catch (ObjectDisposedException) + StartCancelOnce(ref _cancelDelivery, _cts, + "A cancellation callback registered by the durable workflow threw while the workflow was being cancelled."); + } + + /// + /// Starts cancelling exactly once, and records the + /// delivery so can wait for it. + /// + /// + /// The lock makes every caller return only after the source reports + /// IsCancellationRequested: a caller that loses the race waits for the + /// winner to start cancelling instead of returning early. It is held only + /// while flips the state and + /// queues the callbacks, so no foreign code ever runs under it. + /// + private void StartCancelOnce(ref Task? deliverySlot, CancellationTokenSource source, string callbackFailureMessage) + { + lock (_cancelLock) { - // Only reachable from the late termination continuation or deadline - // timer. Dispose cancels before disposing, so the token is already - // cancelled. + if (deliverySlot != null) return; + Volatile.Write(ref deliverySlot, StartCancel(source, callbackFailureMessage)); } - catch (AggregateException ex) + } + + /// + /// Cancels without running its callbacks on this + /// thread. The returned task completes when every callback has run; callback + /// exceptions are logged, never rethrown. + /// + private static Task StartCancel(CancellationTokenSource source, string callbackFailureMessage) + { + Task delivery; + try + { + delivery = source.CancelAsync(); + } + catch (ObjectDisposedException) { - // A user cancellation callback threw. Cancel still ran every other - // callback before throwing, so cancellation is fully delivered. Never - // let the exception escape: on the deadline timer it would crash the - // process (and every invocation sharing it), and during teardown it - // would replace the invocation's real outcome. - LambdaLogger.Log( - "Warning", ex, - "A cancellation callback registered by the durable workflow threw while the workflow was being cancelled."); + // Only reachable after disposal, which happens once delivery finished, + // so the source was already cancelled. + return Task.CompletedTask; } + + return delivery.ContinueWith( + static (t, message) => + { + if (t.Exception is { } ex) + LambdaLogger.Log("Warning", ex.Flatten(), (string)message!); + }, + callbackFailureMessage, + CancellationToken.None, + TaskContinuationOptions.ExecuteSynchronously, + TaskScheduler.Default); } public void Dispose() @@ -217,14 +307,42 @@ public void Dispose() if (Interlocked.Exchange(ref _disposed, 1) != 0) return; BeginTeardown(); - - // Unregister rather than Dispose: Dispose waits for a deadline callback - // that is still running, and a user cancellation callback that blocks - // would then hang the invocation here. - _deadlineRegistration.Unregister(); _deadlineTimer?.Dispose(); - _deadlineCts?.Dispose(); - _cts.Dispose(); + + // Dispose each source only after its callbacks have all run. Disposing + // earlier could skip callbacks that have not run yet, so linked tokens + // (and the user code behind them) would never be cancelled. A callback + // that never returns just leaves its source undisposed. Read the delivery + // handles under the lock, so a cancel that is starting is never missed. + Task workflowDelivery; + Task? deadlineDelivery; + lock (_cancelLock) + { + workflowDelivery = _cancelDelivery ?? Task.CompletedTask; + deadlineDelivery = _deadlineDelivery; + } + + DisposeAfter(_cts, workflowDelivery); + if (_deadlineCts != null) + { + // A deadline that never fired was never cancelled: dispose it now and + // the timer can no longer fire it (CancelAsync on a disposed source is + // tolerated by StartCancel). + DisposeAfter(_deadlineCts, deadlineDelivery ?? Task.CompletedTask); + } + } + + private static void DisposeAfter(CancellationTokenSource source, Task delivery) + { + if (delivery.IsCompleted) + source.Dispose(); + else + delivery.ContinueWith( + static (_, s) => ((CancellationTokenSource)s!).Dispose(), + source, + CancellationToken.None, + TaskContinuationOptions.ExecuteSynchronously, + TaskScheduler.Default); } } @@ -267,8 +385,8 @@ public static bool IsMultiConcurrency() => if (!isMultiConcurrency) return null; if (remainingTime <= TimeSpan.Zero) return TimeSpan.Zero; - // CancellationTokenSource rejects delays above int.MaxValue milliseconds. - // Lambda timeouts are far below that; treat anything larger as no deadline. + // Timer rejects delays above int.MaxValue milliseconds. Lambda timeouts + // are far below that; treat anything larger as no deadline. if (remainingTime.TotalMilliseconds >= int.MaxValue) return null; return remainingTime; } diff --git a/Libraries/src/Amazon.Lambda.DurableExecution/docs/core/cancellation.md b/Libraries/src/Amazon.Lambda.DurableExecution/docs/core/cancellation.md index c9d4b3eaf..ef3da6beb 100644 --- a/Libraries/src/Amazon.Lambda.DurableExecution/docs/core/cancellation.md +++ b/Libraries/src/Amazon.Lambda.DurableExecution/docs/core/cancellation.md @@ -22,6 +22,7 @@ When either trigger fires, the token transitions to `IsCancellationRequested = t - **`OperationCanceledException` thrown out of a step body via the token** (i.e. `linked.IsCancellationRequested` is true) is treated as cancellation: no FAIL checkpoint is written, no retry is consulted. The exception propagates up. - **`OperationCanceledException` thrown by user code for unrelated reasons** (token never fired) is treated as a normal step failure: FAIL checkpoint, retry per the configured `RetryStrategy`. - **The SDK's own writes (checkpoints, batcher flush, the runtime API response)** never observe the workflow-shutdown signal. Successful work is never lost to teardown. +- **Callbacks you register on the token** (`ct.Register(...)`) run on the thread pool, not on the thread that cancelled. The token itself reports `IsCancellationRequested` immediately. A callback that throws is logged, not rethrown. - **When the invocation ends**, the token is cancelled before `DurableFunction.WrapAsync` returns. Work the workflow started but did not await (for example a step still running when a sibling suspended) observes the cancellation. Starting a new durable operation after this point throws `OperationCanceledException`. ## Lambda Managed Instances diff --git a/Libraries/test/Amazon.Lambda.DurableExecution.Tests/LambdaManagedInstancesTests.cs b/Libraries/test/Amazon.Lambda.DurableExecution.Tests/LambdaManagedInstancesTests.cs index ae6d52383..1fba24f44 100644 --- a/Libraries/test/Amazon.Lambda.DurableExecution.Tests/LambdaManagedInstancesTests.cs +++ b/Libraries/test/Amazon.Lambda.DurableExecution.Tests/LambdaManagedInstancesTests.cs @@ -104,8 +104,13 @@ public async Task Suspended_AbandonedStepTokenIsCancelledBeforeWrapperReturns(bo new DurableInvocationOptions { IsMultiConcurrency = isMultiConcurrency }); Assert.Equal(InvocationStatus.Pending, output.Status); - // Checked immediately: cancellation must already be delivered. - Assert.True(stepToken.IsCancellationRequested, $"Abandoned step not cancelled on iteration {i}."); + // Cancellation is requested before the wrapper returns, and its + // callbacks are delivered on the threadpool (never on the wrapper's + // thread, where a blocking user callback could stall it). It must + // arrive, never be lost to disposal. + var cancelled = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + stepToken.Register(() => cancelled.TrySetResult()); + await cancelled.Task.NoHang(); } } @@ -265,14 +270,14 @@ await Assert.ThrowsAsync(() => [Fact] public async Task RunAsync_WorkflowFaultsFromDeadlineBeforeDeadlineTaskCompletes_ReportsDeadline() { - // Forces the race where the workflow task faults before DeadlineTask - // completes. The workflow registers on the deadline token before the - // deadline fires; that callback runs before WorkflowCancellation's own - // (callbacks run in reverse order), so the workflow faults while the - // deadline has passed but DeadlineTask is not yet signalled. The fault - // must still be reported as a deadline, not as a durable failure. + // The workflow faults while the deadline has passed but DeadlineTask is + // not yet signalled: the deadline is already behind the clock and its + // timer has not run (as when the threadpool delays it). The workflow wins + // the race; the fault must still be reported as a deadline, not as a + // durable failure. var tm = new TerminationManager(); - using var wfc = new WorkflowCancellation(tm, TimeSpan.FromMilliseconds(50)); + var (wfc, deadlinePassed) = DeadlineWithTimerNotRun(tm); + using var _ = wfc; var workflowRan = false; var result = await DurableExecutionHandler.RunAsync( @@ -281,11 +286,8 @@ public async Task RunAsync_WorkflowFaultsFromDeadlineBeforeDeadlineTaskCompletes async () => { workflowRan = true; - // No RunContinuationsAsynchronously: the await below resumes - // inline inside the cancellation callback. - var faulted = new TaskCompletionSource(); - wfc.DeadlineToken.Register(() => faulted.TrySetException(new InvalidOperationException("user gave up"))); - return await faulted.Task; + await deadlinePassed(); + throw new InvalidOperationException("user gave up"); }, wfc).NoHang(); @@ -315,16 +317,17 @@ await Assert.ThrowsAsync(() => } [Fact] - public void WorkflowCancellation_DeadlineIsReported_EvenWhenTerminationCancelledFirst() + public async Task WorkflowCancellation_DeadlineIsReported_EvenWhenTerminationCancelledFirst() { var tm = new TerminationManager(); - using var wfc = new WorkflowCancellation(tm, TimeSpan.FromMilliseconds(50)); - var deadlineHit = new ManualResetEventSlim(); - wfc.DeadlineToken.Register(deadlineHit.Set); + using var wfc = new WorkflowCancellation(tm, ControlledDeadline); + var deadlineHit = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + wfc.DeadlineToken.Register(() => deadlineHit.TrySetResult()); // A suspension sets the first cancellation reason before the deadline. wfc.BeginTeardown(); - Assert.True(deadlineHit.Wait(HangGuard)); + wfc.TriggerDeadline(); + await deadlineHit.Task.NoHang(); Assert.True(wfc.IsDeadlineExceeded); } @@ -333,18 +336,20 @@ public void WorkflowCancellation_DeadlineIsReported_EvenWhenTerminationCancelled public async Task WorkflowCancellation_ThrowingCallback_DoesNotEscapeDeadlineOrTeardown() { var tm = new TerminationManager(); - var wfc = new WorkflowCancellation(tm, TimeSpan.FromMilliseconds(50)); + var wfc = new WorkflowCancellation(tm, ControlledDeadline); // Callbacks run in reverse registration order: the throwing one runs - // first, then this one. It still running proves Cancel delivered to - // every callback despite the throw. + // first, then this one. It still running proves every callback was + // delivered despite the throw. var otherCallbackRan = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); wfc.Token.Register(() => otherCallbackRan.TrySetResult()); wfc.Token.Register(() => throw new InvalidOperationException("user callback threw")); - // The deadline fires on a timer thread: the throwing callback must not - // escape there (that would crash the test host) or block the signal. + // Firing the deadline must not throw on the caller's thread (on the timer + // thread that would crash the process) or block the signal. + wfc.TriggerDeadline(); await wfc.DeadlineTask.NoHang(); await otherCallbackRan.Task.NoHang(); + await wfc.CancellationDelivered.NoHang(); Assert.True(wfc.Token.IsCancellationRequested); // Teardown and dispose must not throw either. @@ -353,7 +358,7 @@ public async Task WorkflowCancellation_ThrowingCallback_DoesNotEscapeDeadlineOrT } [Fact] - public void WorkflowCancellation_ThrowingCallback_DoesNotEscapeTeardown() + public async Task WorkflowCancellation_ThrowingCallback_DoesNotEscapeTeardown() { // No deadline: the first cancel comes from the wrapper's teardown, as on // every invocation. A throwing user callback must not replace the @@ -364,6 +369,7 @@ public void WorkflowCancellation_ThrowingCallback_DoesNotEscapeTeardown() wfc.Token.Register(() => throw new InvalidOperationException("user callback threw")); wfc.BeginTeardown(); + await wfc.CancellationDelivered.NoHang(); wfc.Dispose(); Assert.True(wfc.Token.IsCancellationRequested); @@ -375,13 +381,14 @@ public async Task WorkflowCancellation_BlockingCallback_DoesNotDelayDeadlineSign { var tm = new TerminationManager(); using var release = new ManualResetEventSlim(); - var wfc = new WorkflowCancellation(tm, TimeSpan.FromMilliseconds(50)); + var wfc = new WorkflowCancellation(tm, ControlledDeadline); wfc.Token.Register(() => release.Wait()); try { - // DeadlineTask completes even though a cancellation callback blocks - // the thread that is delivering the cancel. + // Firing returns at once and DeadlineTask completes, even though a + // cancellation callback blocks the thread delivering the cancel. + await Task.Run(wfc.TriggerDeadline).NoHang(); await wfc.DeadlineTask.NoHang(); Assert.True(wfc.IsDeadlineExceeded); } @@ -392,30 +399,213 @@ public async Task WorkflowCancellation_BlockingCallback_DoesNotDelayDeadlineSign } } + // ── Copilot review round 3 ───────────────────────────────────────── + + [Fact] + public async Task WorkflowCancellation_BeginTeardown_DoesNotRunUserCallbacksOnCaller() + { + // A step's linked token carries a user callback that blocks. Teardown is + // the first cancel; it must return without waiting for that callback. + var wfc = new WorkflowCancellation(new TerminationManager()); + using var release = new ManualResetEventSlim(); + using var stepToken = CancellationTokenSource.CreateLinkedTokenSource(CancellationToken.None, wfc.Token); + stepToken.Token.Register(() => release.Wait()); + + try + { + await Task.Run(wfc.BeginTeardown).NoHang(); + // The token reports cancellation immediately; delivery is still pending. + Assert.True(wfc.Token.IsCancellationRequested); + Assert.False(wfc.CancellationDelivered.IsCompleted); + } + finally + { + release.Set(); + } + + await wfc.CancellationDelivered.NoHang(); + Assert.True(stepToken.IsCancellationRequested); + wfc.Dispose(); + } + + [Fact] + public async Task OnDemand_UnawaitedStepWithBlockingCallback_WrapperStillReturns() + { + // End to end, in on-demand mode: the workflow returns while an unawaited + // step has a blocking callback on its token. Teardown runs on every + // invocation and must not stall in that callback. + var client = new FakeServiceClient(); + using var release = new ManualResetEventSlim(); + var registered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + + try + { + var output = await DurableFunction.WrapAsync( + async (_, ctx) => + { + ctx.StepAsync(async (_, ct) => + { + ct.Register(() => release.Wait()); + registered.TrySetResult(); + await Task.Delay(Timeout.Infinite, ct); + return 0; + }, name: "slow").Observe(); + await registered.Task; + return "done"; + }, + Input("arn:blocking-teardown"), + Context(TimeSpan.FromMinutes(5)), + client, + OnDemandOptions).NoHang(); + + Assert.Equal(InvocationStatus.Succeeded, output.Status); + } + finally + { + release.Set(); + } + } + + [Fact] + public void WorkflowCancellation_DeadlinePassedButTimerNotRun_IsReportedAndClosesAdmission() + { + var tm = new TerminationManager(); + using var wfc = ExpiredWithTimerNotRun(tm); + + // Nothing has fired the deadline, but the clock is past it. + Assert.False(wfc.DeadlineTask.IsCompleted); + Assert.True(wfc.IsTornDown); + Assert.Throws(wfc.ThrowIfTornDown); + Assert.True(wfc.IsDeadlineExceeded); + + // Observing it fires the deadline, so every later observer agrees. + Assert.True(wfc.DeadlineTask.IsCompleted); + Assert.True(wfc.Token.IsCancellationRequested); + } + + [Fact] + public async Task WorkflowCancellation_BlockingDeadlineTokenCallback_DoesNotBlockDeadline() + { + // The service client registers a callback on the deadline token that + // blocks. The SDK's own deadline work must still complete. + var tm = new TerminationManager(); + using var release = new ManualResetEventSlim(); + var wfc = new WorkflowCancellation(tm, ControlledDeadline); + wfc.DeadlineToken.Register(() => release.Wait()); + + try + { + await Task.Run(wfc.TriggerDeadline).NoHang(); + await wfc.DeadlineTask.NoHang(); + Assert.True(wfc.Token.IsCancellationRequested); + Assert.True(wfc.IsTornDown); + } + finally + { + release.Set(); + wfc.Dispose(); + } + } + + [Fact] + public async Task Hydration_BlockingClientDeadlineCallback_DoesNotBlockTheWait() + { + // The client registers a blocking callback on the deadline token and + // never completes its request. The wait must still end at the deadline. + var tm = new TerminationManager(); + using var release = new ManualResetEventSlim(); + var wfc = new WorkflowCancellation(tm, ControlledDeadline); + var never = new TaskCompletionSource(); + + try + { + var wait = DurableFunction.WaitWithDeadlineAsync(never.Task, wfc); + // Registered after the wait started (as a client that registers once + // its request is under way), so it runs before anything the wait + // itself registered on this token. + wfc.DeadlineToken.Register(() => release.Wait()); + wfc.TriggerDeadline(); + await Assert.ThrowsAnyAsync(() => wait.NoHang()); + } + finally + { + release.Set(); + wfc.Dispose(); + } + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task WaitForCondition_DeadlinePassesDuringStateSerialization_RecordsNoOutcome(bool conditionMet) + { + // The state serializer runs after the deadline check that follows the + // condition check; the deadline passes while it runs (the serializer + // fires it, deterministically). No RETRY, SUCCEED or FAIL may be recorded. + var tm = new TerminationManager(); + using var wfc = new WorkflowCancellation(tm, ControlledDeadline); + var recorder = new RecordingBatcher(); + var ctx = new DurableContext(new ExecutionState(), tm, wfc, new OperationIdGenerator(), + "arn:test", Context(TimeSpan.FromMinutes(5)), recorder.Batcher); + + var op = ctx.WaitForConditionAsync( + (state, _, _) => Task.FromResult(state + 1), + new WaitForConditionConfig + { + InitialState = 0, + WaitStrategy = conditionMet + ? WaitStrategy.FromDelegate((_, _) => WaitDecision.Stop()) + : WaitStrategy.Fixed(TimeSpan.FromSeconds(1)), + Serializer = new DeadlineTriggeringSerializer(wfc), + }, + name: "w"); + + await Task.WhenAny(op, tm.TerminationTask).NoHang(); + await recorder.Batcher.DrainAsync(); + + Assert.True(op.IsCanceled || op.IsFaulted, "The outcome must not be accepted after the deadline."); + if (op.IsFaulted) Assert.NotNull(op.Exception); + Assert.DoesNotContain(recorder.Flushed, + u => u.Action == OperationAction.FAIL || u.Action == OperationAction.RETRY || u.Action == OperationAction.SUCCEED); + } + + /// Default serializer whose Serialize fires the deadline first. + private sealed class DeadlineTriggeringSerializer : Amazon.Lambda.Core.ILambdaSerializer + { + private readonly DefaultLambdaJsonSerializer _inner = new(); + private readonly WorkflowCancellation _wfc; + public DeadlineTriggeringSerializer(WorkflowCancellation wfc) => _wfc = wfc; + public T Deserialize(Stream requestStream) => _inner.Deserialize(requestStream); + public void Serialize(T response, Stream responseStream) + { + _wfc.TriggerDeadline(); + _inner.Serialize(response, responseStream); + } + } + [Fact] public async Task RunAsync_WorkflowReturnsFallbackAfterDeadline_ReportsDeadline() { - // The workflow observes cancellation and returns a value. It can win the - // race against DeadlineTask; the deadline must still decide the outcome. + // The workflow returns a value after the deadline passed, before its timer + // ran. It wins the race against DeadlineTask; the deadline must still + // decide the outcome. var tm = new TerminationManager(); - using var wfc = new WorkflowCancellation(tm, TimeSpan.FromMilliseconds(50)); + var (wfc, deadlinePassed) = DeadlineWithTimerNotRun(tm); + using var _ = wfc; + var workflowRan = false; var result = await DurableExecutionHandler.RunAsync( new ExecutionState(), tm, async () => { - // Completes inline in a deadline-token callback registered after - // WorkflowCancellation's own, so it runs first (callbacks run in - // reverse order): the workflow finishes while the deadline has - // passed but DeadlineTask is not yet signalled. The workflow wins - // the race; only the deadline check can catch it. - var done = new TaskCompletionSource(); - wfc.DeadlineToken.Register(() => done.TrySetResult("fallback")); - return await done.Task; + workflowRan = true; + await deadlinePassed(); + return "fallback"; }, wfc).NoHang(); + Assert.True(workflowRan); Assert.True(result.DeadlineExceeded); Assert.Null(result.Result); } @@ -485,14 +675,16 @@ public async Task Operation_BodyConvertsDeadlineCancellation_EnqueuesNoFailOrRet // Drives the operation directly, with nothing tearing the batcher down, // so a FAIL/RETRY the operation enqueued would definitely be flushed. var tm = new TerminationManager(); - using var wfc = new WorkflowCancellation(tm, TimeSpan.FromMilliseconds(50)); + using var wfc = new WorkflowCancellation(tm, ControlledDeadline); var recorder = new RecordingBatcher(); var ctx = new DurableContext(new ExecutionState(), tm, wfc, new OperationIdGenerator(), "arn:test", Context(TimeSpan.FromMinutes(5)), recorder.Batcher); var retry = new StepConfig { RetryStrategy = RetryStrategy.Exponential(maxAttempts: 3) }; + var entered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); async Task Body(CancellationToken ct) { + entered.TrySetResult(); try { await Task.Delay(Timeout.Infinite, ct); } catch (OperationCanceledException) { throw new InvalidOperationException("gave up"); } return 0; @@ -512,6 +704,9 @@ async Task Body(CancellationToken ct) name: "w"), }; + await entered.Task.NoHang(); + wfc.TriggerDeadline(); + // The operation must surface the error without recording a failure. A // RETRY would suspend instead, so only wait for completion or suspension. await Task.WhenAny(op, tm.TerminationTask).NoHang(); @@ -741,7 +936,7 @@ public async Task ParallelAsync_BranchFailsWhileDeadlineDeliveryLags_NoParentChe // the control token is notified. The parent must not checkpoint an // outcome that contains those failures. var tm = new TerminationManager(); - using var wfc = new WorkflowCancellation(tm, TimeSpan.FromMilliseconds(50)); + using var wfc = new WorkflowCancellation(tm, ControlledDeadline); var recorder = new RecordingBatcher(); var ctx = new DurableContext(new ExecutionState(), tm, wfc, new OperationIdGenerator(), "arn:test", Context(TimeSpan.FromMinutes(5)), recorder.Batcher); @@ -767,6 +962,7 @@ public async Task ParallelAsync_BranchFailsWhileDeadlineDeliveryLags_NoParentChe Assert.True(entered.Wait(LambdaManagedInstancesTests.HangGuardTimeout)); // Registered after the parallel's control token, so it runs before it. wfc.Token.Register(() => { deliveryStarted.TrySetResult(); releaseDelivery.Wait(); }); + wfc.TriggerDeadline(); await Assert.ThrowsAnyAsync(() => run.NoHang()); } @@ -785,15 +981,17 @@ public async Task ParallelAsync_BranchFailsWhileDeadlineDeliveryLags_NoParentChe public async Task WaitForCondition_CheckReturnsStateAfterDeadline_RecordsNoRetryOrFail() { var tm = new TerminationManager(); - using var wfc = new WorkflowCancellation(tm, TimeSpan.FromMilliseconds(50)); + using var wfc = new WorkflowCancellation(tm, ControlledDeadline); var recorder = new RecordingBatcher(); var ctx = new DurableContext(new ExecutionState(), tm, wfc, new OperationIdGenerator(), "arn:test", Context(TimeSpan.FromMinutes(5)), recorder.Batcher); + var entered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); // The check swallows the cancellation and returns a state. var op = ctx.WaitForConditionAsync( async (state, _, ct) => { + entered.TrySetResult(); try { await Task.Delay(Timeout.Infinite, ct); } catch (OperationCanceledException) { } return state + 1; @@ -805,6 +1003,8 @@ public async Task WaitForCondition_CheckReturnsStateAfterDeadline_RecordsNoRetry }, name: "w"); + await entered.Task.NoHang(); + wfc.TriggerDeadline(); await Task.WhenAny(op, tm.TerminationTask).NoHang(); await recorder.Batcher.DrainAsync(); @@ -821,10 +1021,15 @@ public async Task WorkflowCancellation_ThrowingDeadlineTokenCallback_IsContained // registers that throws must not escape the deadline timer (which would // crash the test host process, as it would an LMI environment). var tm = new TerminationManager(); - using var wfc = new WorkflowCancellation(tm, TimeSpan.FromMilliseconds(50)); + using var wfc = new WorkflowCancellation(tm, ControlledDeadline); + var laterCallbackRan = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + wfc.DeadlineToken.Register(() => laterCallbackRan.TrySetResult()); wfc.DeadlineToken.Register(() => throw new InvalidOperationException("client callback threw")); + // Firing must not throw here, on the caller's thread, either. + wfc.TriggerDeadline(); await wfc.DeadlineTask.NoHang(); + await laterCallbackRan.Task.NoHang(); Assert.True(wfc.Token.IsCancellationRequested); } @@ -1118,6 +1323,41 @@ public async Task Lmi_ReplayWhilePeerIsLive_SkipsCompletedSteps() /// The configured hang guard, for . internal static TimeSpan HangGuardTimeout => HangGuard; + /// + /// A deadline far enough away that its timer never fires during a test. + /// Tests fire it themselves with + /// once their setup is in place, so setup never races a real timer. + /// + private static readonly TimeSpan ControlledDeadline = TimeSpan.FromMinutes(10); + + /// + /// A short deadline whose timer is never armed, plus a function that waits + /// until the clock has passed it: lets a workflow start, then finish after + /// the deadline, before anything fired it. Models the threadpool delaying + /// the timer callback. Only the outcome is asserted, never elapsed time. + /// + private static (WorkflowCancellation Wfc, Func DeadlinePassed) DeadlineWithTimerNotRun(TerminationManager tm) + { + var anchor = System.Diagnostics.Stopwatch.GetTimestamp(); + var deadline = TimeSpan.FromMilliseconds(100); + var wfc = new WorkflowCancellation(tm, deadline, anchor, armTimer: false); + async Task DeadlinePassed() + { + while (System.Diagnostics.Stopwatch.GetElapsedTime(anchor) <= deadline) + await Task.Delay(10); + } + return (wfc, DeadlinePassed); + } + + /// + /// A deadline the clock has already passed, whose timer has not run: models + /// the threadpool delaying the timer callback. + /// + private static WorkflowCancellation ExpiredWithTimerNotRun(TerminationManager tm) => + new(tm, TimeSpan.FromSeconds(1), + anchorTimestamp: System.Diagnostics.Stopwatch.GetTimestamp() - 5 * System.Diagnostics.Stopwatch.Frequency, + armTimer: false); + /// /// One handler instance shared by every concurrent invocation. Every step /// waits until all invocations are inside a step, so they truly overlap. From 4d7422b19cffb26f96aaedc32059c2aafcb59a55 Mon Sep 17 00:00:00 2001 From: Norm Johanson Date: Wed, 4 Nov 2026 21:19:21 +0000 Subject: [PATCH 5/6] Address fourth round of Copilot review findings - Start the deadline-token and workflow-token cancels before signalling DeadlineTask, so a waiter that ends the invocation and disposes at once always finds the deadline cancel under way. Dispose and starting a cancel now coordinate under one lock: a cancel is either already started (and waited for) or never starts on a released source. - Check the clock-based deadline immediately before running a step body, a child context function or a condition check, alongside the token check, so an awaited START or a replay state restore that crosses the deadline before its timer runs cannot start user code. - Tests: barrier test disposing right after DeadlineTask is signalled; unarmed-timer tests for step and child START waits and condition replay; the step-deadline wrapper test now fires a controllable deadline after the step is entered instead of racing a real 100 ms deadline. --- .../DurableFunction.cs | 1 + .../Internal/ChildContextOperation.cs | 4 + .../Internal/DurableInvocationOptions.cs | 8 + .../Internal/StepOperation.cs | 4 + .../Internal/WaitForConditionOperation.cs | 3 + .../Internal/WorkflowCancellation.cs | 43 +++-- .../LambdaManagedInstancesTests.cs | 147 ++++++++++++++++-- 7 files changed, 186 insertions(+), 24 deletions(-) diff --git a/Libraries/src/Amazon.Lambda.DurableExecution/DurableFunction.cs b/Libraries/src/Amazon.Lambda.DurableExecution/DurableFunction.cs index 1fd784073..8374a4453 100644 --- a/Libraries/src/Amazon.Lambda.DurableExecution/DurableFunction.cs +++ b/Libraries/src/Amazon.Lambda.DurableExecution/DurableFunction.cs @@ -119,6 +119,7 @@ private static async Task WrapAsyncCore ExecuteFunc(CancellationToken cancellationToken) // The START checkpoint may have been awaited; the workflow can have // been torn down meanwhile. Never start the child function then. _workflowCancellation.Token.ThrowIfCancellationRequested(); + // Also against the clock: the awaited CONTEXT START can finish after + // the deadline while its timer has not run yet. + _workflowCancellation.ThrowIfDeadlineExceeded( + "The invocation deadline passed before the child context started."); result = await _func(childContext, linked.Token); } catch (OperationCanceledException) when ( diff --git a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/DurableInvocationOptions.cs b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/DurableInvocationOptions.cs index cf21405bf..5f14140eb 100644 --- a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/DurableInvocationOptions.cs +++ b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/DurableInvocationOptions.cs @@ -23,4 +23,12 @@ internal sealed class DurableInvocationOptions /// See . public TimeSpan CleanupTimeout { get; init; } = DefaultCleanupTimeout; + + /// + /// Test seam: receives the invocation's + /// as soon as it is created, so a test can fire the deadline itself + /// () at a chosen point + /// instead of racing a real timer. + /// + public Action? OnWorkflowCancellationCreated { get; init; } } diff --git a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/StepOperation.cs b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/StepOperation.cs index a0f22a7a9..2dee83f69 100644 --- a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/StepOperation.cs +++ b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/StepOperation.cs @@ -236,6 +236,10 @@ private async Task ExecuteFunc(int attemptNumber, CancellationToken cancellat // The START checkpoint above may have been awaited; the workflow // can have been torn down meanwhile. Never start the body then. _workflowCancellation.Token.ThrowIfCancellationRequested(); + // Also against the clock: the awaited START can finish after the + // deadline while its timer has not run yet. + _workflowCancellation.ThrowIfDeadlineExceeded( + "The invocation deadline passed before the step body started."); result = await _func(stepContext, linked.Token); } } diff --git a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/WaitForConditionOperation.cs b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/WaitForConditionOperation.cs index 4c5bff500..f6768f082 100644 --- a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/WaitForConditionOperation.cs +++ b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/WaitForConditionOperation.cs @@ -209,6 +209,9 @@ private async Task ExecuteIteration(TState currentState, int attemptNumb var checkContext = new ConditionCheckContext(attemptNumber, _logger); // Never start the check once the workflow has been torn down. _workflowCancellation.Token.ThrowIfCancellationRequested(); + // Also against the clock: restoring the prior state on replay can + // cross the deadline while its timer has not run yet. + ThrowIfDeadlineExceeded(); newState = await _check(currentState, checkContext, linked.Token); } catch (OperationCanceledException) when ( diff --git a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/WorkflowCancellation.cs b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/WorkflowCancellation.cs index a1cfa40e0..1d6e61d65 100644 --- a/Libraries/src/Amazon.Lambda.DurableExecution/Internal/WorkflowCancellation.cs +++ b/Libraries/src/Amazon.Lambda.DurableExecution/Internal/WorkflowCancellation.cs @@ -77,6 +77,7 @@ internal sealed class WorkflowCancellation : IDisposable private readonly Timer? _deadlineTimer; private readonly TaskCompletionSource _deadlineTcs = new(TaskCreationOptions.RunContinuationsAsynchronously); private readonly object _cancelLock = new(); + private bool _sourcesReleased; private Task? _cancelDelivery; private Task? _deadlineDelivery; private int _reason; @@ -230,20 +231,33 @@ internal void TriggerDeadline() if (_deadlineCts is null) return; if (Interlocked.Exchange(ref _deadlineFired, 1) != 0) return; - // The SDK's own state first: record why, close admission, and signal the - // handler. _deadlineTcs runs its continuations asynchronously, so this - // never runs the handler on the calling thread. + // The SDK's own state first: record why and close admission. Interlocked.CompareExchange(ref _reason, (int)WorkflowCancellationReason.DeadlineExceeded, 0); Volatile.Write(ref _tornDown, 1); - _deadlineTcs.TrySetResult(); - // Then the workflow token, then the client's deadline token. Both only - // queue their callbacks. + // Then start cancelling the workflow token and the client's deadline + // token. Both only queue their callbacks, so no foreign code runs here. + // This must happen BEFORE DeadlineTask is signalled: a waiter on + // DeadlineTask (the hydration wait) can end the invocation and reach + // Dispose at once, and Dispose must find the deadline cancel already + // under way, not dispose a source whose cancel has not started. Cancel(WorkflowCancellationReason.DeadlineExceeded); StartCancelOnce(ref _deadlineDelivery, _deadlineCts, "A cancellation callback registered on the durable invocation deadline threw."); + + // Last, signal the handler. _deadlineTcs runs its continuations + // asynchronously, so this never runs the handler on the calling thread. + _deadlineTcs.TrySetResult(); + AfterDeadlineSignalledForTest?.Invoke(); } + /// + /// Test seam: runs right after is signalled, on the + /// thread firing the deadline. Lets a test act (for example dispose) at the + /// earliest point a waiter on could. + /// + internal Action? AfterDeadlineSignalledForTest { get; set; } + private void Cancel(WorkflowCancellationReason reason) { Interlocked.CompareExchange(ref _reason, (int)reason, 0); @@ -266,7 +280,9 @@ private void StartCancelOnce(ref Task? deliverySlot, CancellationTokenSource sou { lock (_cancelLock) { - if (deliverySlot != null) return; + // Once Dispose has taken the delivery handles, the sources may be + // disposed; a cancel that had not started by then never starts. + if (deliverySlot != null || _sourcesReleased) return; Volatile.Write(ref deliverySlot, StartCancel(source, callbackFailureMessage)); } } @@ -312,12 +328,15 @@ public void Dispose() // Dispose each source only after its callbacks have all run. Disposing // earlier could skip callbacks that have not run yet, so linked tokens // (and the user code behind them) would never be cancelled. A callback - // that never returns just leaves its source undisposed. Read the delivery - // handles under the lock, so a cancel that is starting is never missed. + // that never returns just leaves its source undisposed. Taking the + // delivery handles and closing out new cancels happen together under the + // lock: a cancel is either already under way (and is waited for) or never + // starts, so a source is never disposed while its cancel is starting. Task workflowDelivery; Task? deadlineDelivery; lock (_cancelLock) { + _sourcesReleased = true; workflowDelivery = _cancelDelivery ?? Task.CompletedTask; deadlineDelivery = _deadlineDelivery; } @@ -325,9 +344,9 @@ public void Dispose() DisposeAfter(_cts, workflowDelivery); if (_deadlineCts != null) { - // A deadline that never fired was never cancelled: dispose it now and - // the timer can no longer fire it (CancelAsync on a disposed source is - // tolerated by StartCancel). + // A deadline that never fired was never cancelled: dispose it now. A + // timer callback still in flight finds _sourcesReleased set and does + // not touch the source. DisposeAfter(_deadlineCts, deadlineDelivery ?? Task.CompletedTask); } } diff --git a/Libraries/test/Amazon.Lambda.DurableExecution.Tests/LambdaManagedInstancesTests.cs b/Libraries/test/Amazon.Lambda.DurableExecution.Tests/LambdaManagedInstancesTests.cs index 1fba24f44..44c234a6b 100644 --- a/Libraries/test/Amazon.Lambda.DurableExecution.Tests/LambdaManagedInstancesTests.cs +++ b/Libraries/test/Amazon.Lambda.DurableExecution.Tests/LambdaManagedInstancesTests.cs @@ -215,21 +215,36 @@ public async Task Lmi_HydrationReceivesDeadlineToken() [Fact] public async Task Lmi_DeadlinePassesDuringStep_ThrowsTimeout_WithoutFailCheckpoint() { + // A controllable deadline, fired only once the step body is running, so + // the test always exercises a deadline during the step. A real short + // deadline could expire before the queued workflow even starts, which the + // wrapper validly rejects without running the body. var client = new FakeServiceClient(); CancellationToken stepToken = default; + var stepEntered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + WorkflowCancellation? wfc = null; + var options = new DurableInvocationOptions + { + IsMultiConcurrency = true, + OnWorkflowCancellationCreated = created => wfc = created, + }; - var ex = await Assert.ThrowsAsync(() => - DurableFunction.WrapAsync( - async (_, ctx) => await ctx.StepAsync(async (_, ct) => - { - stepToken = ct; - await Task.Delay(Timeout.Infinite, ct); - return "unreachable"; - }, name: "slow"), - Input("arn:deadline-step"), - Context(TimeSpan.FromMilliseconds(100)), - client, - LmiOptions).NoHang()); + var run = DurableFunction.WrapAsync( + async (_, ctx) => await ctx.StepAsync(async (_, ct) => + { + stepToken = ct; + stepEntered.TrySetResult(); + await Task.Delay(Timeout.Infinite, ct); + return "unreachable"; + }, name: "slow"), + Input("arn:deadline-step"), + Context(ControlledDeadline), + client, + options); + + await stepEntered.Task.NoHang(); + wfc!.TriggerDeadline(); + var ex = await Assert.ThrowsAsync(() => run.NoHang()); Assert.Contains("arn:deadline-step", ex.Message); Assert.True(stepToken.IsCancellationRequested); @@ -569,6 +584,114 @@ public async Task WaitForCondition_DeadlinePassesDuringStateSerialization_Record u => u.Action == OperationAction.FAIL || u.Action == OperationAction.RETRY || u.Action == OperationAction.SUCCEED); } + // ── Copilot review round 4 ───────────────────────────────────────── + + [Fact] + public async Task TriggerDeadline_DisposedAsSoonAsDeadlineTaskIsSignalled_StillCancelsDeadlineToken() + { + // A hydration wait ends on DeadlineTask, and the wrapper then disposes at + // once. Dispose at exactly that point (on the firing thread, right after + // DeadlineTask is signalled): the client's deadline token must still be + // cancelled, so a cooperative hydration request stops. + var wfc = new WorkflowCancellation(new TerminationManager(), ControlledDeadline); + var clientCancelled = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + wfc.DeadlineToken.Register(() => clientCancelled.TrySetResult()); + wfc.AfterDeadlineSignalledForTest = wfc.Dispose; + + wfc.TriggerDeadline(); + + await clientCancelled.Task.NoHang(); + } + + [Theory] + [InlineData("step")] + [InlineData("child")] + public async Task StartCheckpointFinishesAfterDeadlineBeforeTimerRuns_UserCodeDoesNotStart(string kind) + { + // The awaited START checkpoint completes only after the clock has passed + // the deadline; the deadline timer has not run (the threadpool delayed + // it), so the workflow token is not cancelled yet. + var tm = new TerminationManager(); + var (wfc, deadlinePassed) = DeadlineWithTimerNotRun(tm); + using var _ = wfc; + await using var batcher = new CheckpointBatcher("t0", async (_, _, _) => + { + await deadlinePassed(); + return "t1"; + }); + var ctx = new DurableContext(new ExecutionState(), tm, wfc, new OperationIdGenerator(), + "arn:test", Context(TimeSpan.FromMinutes(5)), batcher); + var userCodeRan = false; + + Task op = kind == "step" + ? ctx.StepAsync( + async (_, _) => { userCodeRan = true; await Task.CompletedTask; return 1; }, + name: "s", + config: new StepConfig { Semantics = StepSemantics.AtMostOncePerRetry }) + : ctx.RunInChildContextAsync( + (_, _) => { userCodeRan = true; return Task.FromResult(1); }, + name: "c"); + + await Assert.ThrowsAnyAsync(() => op.NoHang()); + Assert.False(userCodeRan); + } + + [Fact] + public async Task WaitForConditionReplay_StateRestoreCrossesDeadline_CheckDoesNotStart() + { + // READY replay: restoring the prior state runs the user serializer, which + // takes long enough to cross the deadline before its timer has run. + var tm = new TerminationManager(); + var (wfc, deadlinePassed) = DeadlineWithTimerNotRun(tm); + using var _ = wfc; + var state = new ExecutionState(); + state.LoadFromCheckpoint(new InitialExecutionState + { + Operations = new List + { + new() + { + Id = OperationIdGenerator.HashOperationId("1"), + Type = OperationTypes.Step, + SubType = OperationSubTypes.WaitForCondition, + Status = OperationStatuses.Ready, + Name = "w", + StepDetails = new StepDetails { Result = "5", Attempt = 1 }, + }, + }, + }); + var ctx = new DurableContext(state, tm, wfc, new OperationIdGenerator(), + "arn:test", Context(TimeSpan.FromMinutes(5))); + var checkRan = false; + + var op = ctx.WaitForConditionAsync( + (s, _, _) => { checkRan = true; return Task.FromResult(s + 1); }, + new WaitForConditionConfig + { + InitialState = 0, + WaitStrategy = WaitStrategy.Fixed(TimeSpan.FromSeconds(1)), + Serializer = new SlowDeserializer(() => deadlinePassed().GetAwaiter().GetResult()), + }, + name: "w"); + + await Assert.ThrowsAnyAsync(() => op.NoHang()); + Assert.False(checkRan); + } + + /// Default serializer whose Deserialize first runs a blocking wait. + private sealed class SlowDeserializer : Amazon.Lambda.Core.ILambdaSerializer + { + private readonly DefaultLambdaJsonSerializer _inner = new(); + private readonly Action _beforeDeserialize; + public SlowDeserializer(Action beforeDeserialize) => _beforeDeserialize = beforeDeserialize; + public T Deserialize(Stream requestStream) + { + _beforeDeserialize(); + return _inner.Deserialize(requestStream); + } + public void Serialize(T response, Stream responseStream) => _inner.Serialize(response, responseStream); + } + /// Default serializer whose Serialize fires the deadline first. private sealed class DeadlineTriggeringSerializer : Amazon.Lambda.Core.ILambdaSerializer { From fc87fb06ea985faad3a2acb71a561630fece8b55 Mon Sep 17 00:00:00 2001 From: Norm Johanson Date: Fri, 6 Nov 2026 01:57:15 +0000 Subject: [PATCH 6/6] Report any hydration failure after the deadline as a timeout A hydration request that faults after the LMI deadline (for example as it is cut off) is now reported as TimeoutException with the original as the inner exception, matching the rest of the wrapper. --- .../DurableFunction.cs | 5 ++- .../LambdaManagedInstancesTests.cs | 33 +++++++++++++++++++ 2 files changed, 37 insertions(+), 1 deletion(-) diff --git a/Libraries/src/Amazon.Lambda.DurableExecution/DurableFunction.cs b/Libraries/src/Amazon.Lambda.DurableExecution/DurableFunction.cs index 8374a4453..d195b18d9 100644 --- a/Libraries/src/Amazon.Lambda.DurableExecution/DurableFunction.cs +++ b/Libraries/src/Amazon.Lambda.DurableExecution/DurableFunction.cs @@ -145,8 +145,11 @@ private static async Task WrapAsyncCore + { + wfc!.TriggerDeadline(); + throw new DurableExecutionException("request failed as it was cut off"); + }, + }; + var options = new DurableInvocationOptions + { + IsMultiConcurrency = true, + OnWorkflowCancellationCreated = created => wfc = created, + }; + + var ex = await Assert.ThrowsAsync(() => + DurableFunction.WrapAsync( + (input, _) => Task.FromResult(input.Id!), + Input("arn:hydrate-fault-after-deadline", nextMarker: "page-2"), + Context(ControlledDeadline), + client, + options).NoHang()); + + Assert.IsType(ex.InnerException); + } + /// Default serializer whose Deserialize first runs a blocking wait. private sealed class SlowDeserializer : Amazon.Lambda.Core.ILambdaSerializer {