Skip to content

Commit e984839

Browse files
DaRosenbergDaRosenberg
authored
Add caller CancellationToken propagation for hedging and timeout (#3094)
Adds caller cancellation token propagation in hedging and timeout strategies.
1 parent 27eb2bb commit e984839

7 files changed

Lines changed: 368 additions & 3 deletions

File tree

‎src/Polly.Core/Hedging/HedgingResilienceStrategy.cs‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
using Polly.Hedging.Utils;
22
using Polly.Telemetry;
3+
using Polly.Utils;
34

45
namespace Polly.Hedging;
56

@@ -57,7 +58,7 @@ protected internal override async ValueTask<Outcome<T>> ExecuteCore<TState>(
5758

5859
if (loadedExecution.Outcome is Outcome<T> outcome)
5960
{
60-
return outcome;
61+
return outcome.WithCallerCancellationToken(cancellationToken);
6162
}
6263

6364
var delay = await GetHedgingDelayAsync(context, hedgingContext.LoadedTasks).ConfigureAwait(continueOnCapturedContext);
@@ -72,7 +73,7 @@ protected internal override async ValueTask<Outcome<T>> ExecuteCore<TState>(
7273
if (!execution.IsHandled)
7374
{
7475
execution.AcceptOutcome();
75-
return outcome;
76+
return outcome.WithCallerCancellationToken(cancellationToken);
7677
}
7778
}
7879
}

‎src/Polly.Core/Timeout/TimeoutResilienceStrategy.cs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
using Polly.Telemetry;
2+
using Polly.Utils;
23

34
namespace Polly.Timeout;
45

@@ -89,7 +90,7 @@ protected internal override async ValueTask<Outcome<TResult>> ExecuteCore<TResul
8990
return Outcome.FromException<TResult>(timeoutException.TrySetStackTrace());
9091
}
9192

92-
return outcome;
93+
return outcome.WithCallerCancellationToken(previousToken);
9394
}
9495

9596
private static CancellationTokenRegistration CreateRegistration(CancellationTokenSource cancellationSource, CancellationToken previousToken)

‎src/Polly.Core/Utils/OutcomeUtilities.cs‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
namespace Polly.Utils;
2+
3+
internal static class OutcomeUtilities
4+
{
5+
/// <summary>
6+
/// Ensures that an <see cref="OperationCanceledException"/> escaping a strategy that substituted the
7+
/// execution <see cref="CancellationToken"/> (e.g. timeout or hedging) carries the caller's token when
8+
/// the cancellation was caused by that token.
9+
/// </summary>
10+
/// <typeparam name="T">The result type of the outcome.</typeparam>
11+
/// <param name="outcome">The outcome produced by the strategy.</param>
12+
/// <param name="callerToken">The cancellation token that was associated with the execution before the strategy substituted it.</param>
13+
/// <returns>
14+
/// An outcome whose <see cref="OperationCanceledException"/> carries <paramref name="callerToken"/> when the caller
15+
/// requested cancellation, preserving the original exception as its <see cref="Exception.InnerException"/>;
16+
/// otherwise the original <paramref name="outcome"/> unchanged.
17+
/// </returns>
18+
/// <remarks>
19+
/// The rewrite happens only when <paramref name="callerToken"/> actually requested cancellation. This preserves
20+
/// the contract that a real timeout, or an unrelated <see cref="OperationCanceledException"/> thrown while the
21+
/// caller's token was not cancelled, is left untouched.
22+
/// </remarks>
23+
public static Outcome<T> WithCallerCancellationToken<T>(this Outcome<T> outcome, CancellationToken callerToken)
24+
{
25+
if (callerToken.IsCancellationRequested && outcome.Exception is OperationCanceledException oce && oce.CancellationToken != callerToken)
26+
{
27+
return Outcome.FromException<T>(new OperationCanceledException(oce.Message, oce, callerToken).TrySetStackTrace());
28+
}
29+
30+
return outcome;
31+
}
32+
}

‎test/Polly.Core.Tests/Hedging/HedgingResilienceStrategyTests.cs‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -935,6 +935,36 @@ public async Task ExecuteAsync_EnsureOnHedgingTelemetry()
935935
_events.Select(v => v.Event.EventName).Distinct().Count().ShouldBe(2);
936936
}
937937

938+
[Fact]
939+
public async Task ExecuteCore_CallerCancellation_EnsureExceptionCarriesCallerToken()
940+
{
941+
// arrange
942+
using var cancellationSource = new CancellationTokenSource();
943+
_options.MaxHedgedAttempts = 1;
944+
_options.Delay = LongDelay; // ensure no secondary attempt starts before the primary completes
945+
ConfigureHedging(); // ShouldHandle returns false, so the primary's outcome is accepted as-is
946+
947+
var strategy = (HedgingResilienceStrategy<string>)Create().GetPipelineDescriptor().FirstStrategy.StrategyInstance;
948+
949+
var context = ResilienceContextPool.Shared.Get(CancellationToken);
950+
context.CancellationToken = cancellationSource.Token;
951+
952+
// act
953+
var outcome = await strategy.ExecuteCore(
954+
(ctx, _) =>
955+
{
956+
cancellationSource.Cancel();
957+
ctx.CancellationToken.ThrowIfCancellationRequested();
958+
return Outcome.FromResultAsValueTask("unreachable");
959+
},
960+
context,
961+
"state");
962+
963+
// assert
964+
var exception = outcome.Exception.ShouldBeOfType<OperationCanceledException>();
965+
exception.CancellationToken.ShouldBe(cancellationSource.Token);
966+
}
967+
938968
private void ConfigureHedging() =>
939969
ConfigureHedging(_ => false, _actions.Generator);
940970

‎test/Polly.Core.Tests/Issues/IssuesTests.CancellationTokenPropagation_3086.cs‎

Lines changed: 203 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,203 @@
1+
using Polly.Hedging;
2+
using Polly.Timeout;
3+
4+
namespace Polly.Core.Tests.Issues;
5+
6+
public partial class IssuesTests
7+
{
8+
// https://github.com/App-vNext/Polly/issues/3086
9+
// When a strategy substitutes the caller's CancellationToken with an internal one (timeout, hedging),
10+
// a caller-initiated cancellation must still surface an OperationCanceledException carrying the
11+
// caller's token, so that callers can distinguish caller cancellation from other failures.
12+
[Fact]
13+
public async Task Timeout_CallerCancellation_ExceptionCarriesCallerToken_3086()
14+
{
15+
var pipeline = new ResiliencePipelineBuilder()
16+
.AddTimeout(TimeSpan.FromMinutes(1))
17+
.Build();
18+
19+
using var cts = new CancellationTokenSource();
20+
21+
var exception = await Should.ThrowAsync<OperationCanceledException>(() =>
22+
pipeline.ExecuteAsync(
23+
async token =>
24+
{
25+
cts.Cancel(); // simulate cancellation request from upstream caller
26+
token.ThrowIfCancellationRequested(); // simulate cancellation response from downstream code
27+
},
28+
cts.Token).AsTask());
29+
30+
exception.CancellationToken.ShouldBe(cts.Token);
31+
}
32+
33+
[Fact]
34+
public async Task TimeoutThenRetry_CallerCancellation_ExceptionCarriesCallerToken_3086()
35+
{
36+
var pipeline = new ResiliencePipelineBuilder()
37+
.AddTimeout(TimeSpan.FromMinutes(1))
38+
.AddRetry(new() { MaxRetryAttempts = 3, Delay = TimeSpan.Zero })
39+
.Build();
40+
41+
using var cts = new CancellationTokenSource();
42+
43+
var exception = await Should.ThrowAsync<OperationCanceledException>(() =>
44+
pipeline.ExecuteAsync(
45+
async token =>
46+
{
47+
cts.Cancel();
48+
token.ThrowIfCancellationRequested();
49+
},
50+
cts.Token).AsTask());
51+
52+
exception.CancellationToken.ShouldBe(cts.Token);
53+
}
54+
55+
[Fact]
56+
public async Task RetryThenTimeout_CallerCancellation_ExceptionCarriesCallerToken_3086()
57+
{
58+
var pipeline = new ResiliencePipelineBuilder()
59+
.AddRetry(new() { MaxRetryAttempts = 3, Delay = TimeSpan.Zero })
60+
.AddTimeout(TimeSpan.FromMinutes(1))
61+
.Build();
62+
63+
using var cts = new CancellationTokenSource();
64+
65+
var exception = await Should.ThrowAsync<OperationCanceledException>(() =>
66+
pipeline.ExecuteAsync(
67+
async token =>
68+
{
69+
cts.Cancel();
70+
token.ThrowIfCancellationRequested();
71+
},
72+
cts.Token).AsTask());
73+
74+
exception.CancellationToken.ShouldBe(cts.Token);
75+
}
76+
77+
[Fact]
78+
public async Task NestedTimeouts_CallerCancellation_ExceptionCarriesCallerToken_3086()
79+
{
80+
var innermost = new ResiliencePipelineBuilder()
81+
.AddTimeout(TimeSpan.FromMinutes(1))
82+
.Build();
83+
84+
var inner = new ResiliencePipelineBuilder()
85+
.AddTimeout(TimeSpan.FromMinutes(2))
86+
.AddPipeline(innermost)
87+
.Build();
88+
89+
var pipeline = new ResiliencePipelineBuilder()
90+
.AddTimeout(TimeSpan.FromMinutes(3))
91+
.AddPipeline(inner)
92+
.Build();
93+
94+
using var cts = new CancellationTokenSource();
95+
96+
var exception = await Should.ThrowAsync<OperationCanceledException>(() =>
97+
pipeline.ExecuteAsync(
98+
async token =>
99+
{
100+
cts.Cancel();
101+
token.ThrowIfCancellationRequested();
102+
},
103+
cts.Token).AsTask());
104+
105+
exception.CancellationToken.ShouldBe(cts.Token);
106+
}
107+
108+
[Fact]
109+
public async Task Hedging_CallerCancellation_ExceptionCarriesCallerToken_3086()
110+
{
111+
var pipeline = new ResiliencePipelineBuilder<string>()
112+
.AddHedging(new HedgingStrategyOptions<string>
113+
{
114+
MaxHedgedAttempts = 1,
115+
Delay = TimeSpan.FromMinutes(1), // ensure only the primary attempt runs before cancellation
116+
ShouldHandle = _ => PredicateResult.False(), // accept the primary outcome as-is
117+
})
118+
.Build();
119+
120+
using var cts = new CancellationTokenSource();
121+
122+
var exception = await Should.ThrowAsync<OperationCanceledException>(() =>
123+
pipeline.ExecuteAsync(
124+
async token =>
125+
{
126+
cts.Cancel();
127+
token.ThrowIfCancellationRequested();
128+
return "unreachable";
129+
},
130+
cts.Token).AsTask());
131+
132+
exception.CancellationToken.ShouldBe(cts.Token);
133+
}
134+
135+
[Fact]
136+
public async Task WithoutTokenSubstitution_CallerCancellation_ExceptionCarriesCallerToken_3086()
137+
{
138+
// A pipeline that does not substitute the token has always surfaced the caller's token
139+
// and this case simply guards against regressions for the common path.
140+
var pipeline = new ResiliencePipelineBuilder()
141+
.AddRetry(new() { MaxRetryAttempts = 3, Delay = TimeSpan.Zero })
142+
.Build();
143+
144+
using var cts = new CancellationTokenSource();
145+
146+
var exception = await Should.ThrowAsync<OperationCanceledException>(() =>
147+
pipeline.ExecuteAsync(
148+
async token =>
149+
{
150+
cts.Cancel();
151+
token.ThrowIfCancellationRequested();
152+
},
153+
cts.Token).AsTask());
154+
155+
exception.CancellationToken.ShouldBe(cts.Token);
156+
}
157+
158+
[Fact]
159+
public async Task Timeout_RealTimeout_StillThrowsTimeoutRejectedException_3086()
160+
{
161+
// "only if": a real timeout (caller token not cancelled) must remain a TimeoutRejectedException,
162+
// and must not be rewritten into a caller-cancellation.
163+
var pipeline = new ResiliencePipelineBuilder { TimeProvider = TimeProvider }
164+
.AddTimeout(TimeSpan.FromSeconds(1))
165+
.Build();
166+
167+
await Should.ThrowAsync<TimeoutRejectedException>(() =>
168+
pipeline.ExecuteAsync(
169+
async token =>
170+
{
171+
var delay = TimeProvider.Delay(TimeSpan.FromSeconds(10), token);
172+
TimeProvider.Advance(TimeSpan.FromSeconds(2));
173+
await delay;
174+
},
175+
CancellationToken.None).AsTask());
176+
}
177+
178+
[Fact]
179+
public async Task Timeout_UnrelatedCancellation_ExceptionPreserved_3086()
180+
{
181+
// "only if": when the caller token is not cancelled, an unrelated OperationCanceledException
182+
// thrown by user code must propagate unchanged (its own token must be preserved).
183+
var pipeline = new ResiliencePipelineBuilder()
184+
.AddTimeout(TimeSpan.FromMinutes(1))
185+
.Build();
186+
187+
using var callerCts = new CancellationTokenSource();
188+
using var unrelatedCts = new CancellationTokenSource();
189+
190+
unrelatedCts.Cancel();
191+
192+
var exception = await Should.ThrowAsync<OperationCanceledException>(() =>
193+
pipeline.ExecuteAsync(
194+
async token =>
195+
{
196+
await Task.Yield();
197+
throw new OperationCanceledException(unrelatedCts.Token);
198+
},
199+
callerCts.Token).AsTask());
200+
201+
exception.CancellationToken.ShouldBe(unrelatedCts.Token);
202+
}
203+
}

‎test/Polly.Core.Tests/Timeout/TimeoutResilienceStrategyTests.cs‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -225,6 +225,28 @@ await Should.ThrowAsync<OperationCanceledException>(
225225
_args.ShouldBeEmpty();
226226
}
227227

228+
[Fact]
229+
public async Task Execute_CallerCancellation_EnsureExceptionCarriesCallerToken()
230+
{
231+
var delay = TimeSpan.FromSeconds(10);
232+
233+
using var cts = new CancellationTokenSource();
234+
SetTimeout(TimeSpan.FromSeconds(10));
235+
236+
var sut = CreateSut();
237+
238+
var exception = await Should.ThrowAsync<OperationCanceledException>(
239+
() => sut.ExecuteAsync(async token =>
240+
{
241+
var task = _timeProvider.Delay(delay, token);
242+
cts.Cancel();
243+
await task;
244+
},
245+
cts.Token).AsTask());
246+
247+
exception.CancellationToken.ShouldBe(cts.Token);
248+
}
249+
228250
[Fact]
229251
public async Task Execute_NoTimeoutOrCancellation_EnsureCancellationTokenRestored()
230252
{

0 commit comments

Comments
 (0)
⚡