Skip to content

Commit 1089f68

Browse files
thomhurstthomhurst
authored
Add cancellation-aware async selectors (#354)
* feat: add cancellation-aware selectors Add token-aware SelectAsync and ForEachAsync overloads for enumerable, async-enumerable, item, and execution-count builders. Adapt each delegate once to the builder's linked processor token so CancelAll, disposal, and external cancellation interrupt in-flight work without changing tokenless paths. * fix: await async selector cleanup External cancellation can end source enumeration before started selectors finish. Keep unbounded action and result execution alive through selector cleanup. Refs #336
1 parent 691ecf5 commit 1089f68

14 files changed

Lines changed: 443 additions & 44 deletions

‎EnumerableAsyncProcessor.UnitTests/CancellationAwareSelectorTests.cs‎

Lines changed: 229 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,229 @@
1+
using System;
2+
using System.Collections.Generic;
3+
using System.Runtime.CompilerServices;
4+
using System.Threading;
5+
using System.Threading.Tasks;
6+
using EnumerableAsyncProcessor.Builders;
7+
using EnumerableAsyncProcessor.Extensions;
8+
9+
namespace EnumerableAsyncProcessor.UnitTests;
10+
11+
public class CancellationAwareSelectorTests
12+
{
13+
[Test]
14+
public async Task CancelAll_Interrupts_InFlight_ForEachAsync_Selector(CancellationToken cancellationToken)
15+
{
16+
var started = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
17+
var interrupted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
18+
19+
await using var processor = new[] { 1 }
20+
.ForEachAsync((_, processorToken) => WaitUntilCanceledAsync(processorToken, started, interrupted))
21+
.ProcessInParallel();
22+
23+
await started.Task.WaitAsync(cancellationToken);
24+
processor.CancelAll();
25+
26+
await interrupted.Task.WaitAsync(cancellationToken);
27+
await Assert.ThrowsAsync<TaskCanceledException>(() => processor.WaitAsync());
28+
}
29+
30+
[Test]
31+
public async Task CancelAll_Interrupts_InFlight_SelectAsync_Selector(CancellationToken cancellationToken)
32+
{
33+
var started = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
34+
var interrupted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
35+
36+
await using var processor = new[] { 1 }
37+
.SelectAsync((item, processorToken) => WaitUntilCanceledAsync(item, processorToken, started, interrupted))
38+
.ProcessInParallel();
39+
40+
await started.Task.WaitAsync(cancellationToken);
41+
processor.CancelAll();
42+
43+
await interrupted.Task.WaitAsync(cancellationToken);
44+
await Assert.ThrowsAsync<TaskCanceledException>(() => processor.GetResultsAsync());
45+
}
46+
47+
[Test]
48+
public async Task CancelAll_Interrupts_ExecutionCount_Selector(CancellationToken cancellationToken)
49+
{
50+
var started = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
51+
var interrupted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
52+
53+
await using var processor = AsyncProcessorBuilder.WithExecutionCount(1)
54+
.ForEachAsync(processorToken => WaitUntilCanceledAsync(processorToken, started, interrupted))
55+
.ProcessInParallel();
56+
57+
await started.Task.WaitAsync(cancellationToken);
58+
processor.CancelAll();
59+
60+
await interrupted.Task.WaitAsync(cancellationToken);
61+
await Assert.ThrowsAsync<TaskCanceledException>(() => processor.WaitAsync());
62+
}
63+
64+
[Test]
65+
public async Task DisposeAsync_Interrupts_InFlight_Selector(CancellationToken cancellationToken)
66+
{
67+
var started = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
68+
var interrupted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
69+
70+
var processor = new[] { 1 }
71+
.ForEachAsync((_, processorToken) => WaitUntilCanceledAsync(processorToken, started, interrupted))
72+
.ProcessInParallel();
73+
74+
await started.Task.WaitAsync(cancellationToken);
75+
var disposeTask = processor.DisposeAsync().AsTask();
76+
77+
await interrupted.Task.WaitAsync(cancellationToken);
78+
await disposeTask.WaitAsync(cancellationToken);
79+
}
80+
81+
[Test]
82+
public async Task ExternalCancellation_Interrupts_AsyncEnumerable_Selector(CancellationToken cancellationToken)
83+
{
84+
using var cancellationTokenSource = new CancellationTokenSource();
85+
var started = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
86+
var interrupted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
87+
var allowCleanupToFinish = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
88+
89+
var processor = GetItemsAsync()
90+
.ForEachAsync(
91+
(_, processorToken) => WaitUntilCanceledAfterCleanupAsync(
92+
processorToken,
93+
started,
94+
interrupted,
95+
allowCleanupToFinish),
96+
cancellationTokenSource.Token)
97+
.ProcessInParallel();
98+
99+
var executeTask = processor.ExecuteAsync();
100+
await started.Task.WaitAsync(cancellationToken);
101+
await cancellationTokenSource.CancelAsync();
102+
103+
await AssertExecutionWaitsForCleanupAsync(
104+
executeTask,
105+
interrupted,
106+
allowCleanupToFinish,
107+
cancellationToken);
108+
}
109+
110+
[Test]
111+
public async Task ExternalCancellation_Waits_For_AsyncEnumerable_Result_Selector_Cleanup(CancellationToken cancellationToken)
112+
{
113+
using var cancellationTokenSource = new CancellationTokenSource();
114+
var started = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
115+
var interrupted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
116+
var allowCleanupToFinish = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
117+
118+
var processor = GetItemsAsync()
119+
.SelectAsync(
120+
(item, processorToken) => WaitUntilCanceledAfterCleanupAsync(
121+
item,
122+
processorToken,
123+
started,
124+
interrupted,
125+
allowCleanupToFinish),
126+
cancellationTokenSource.Token)
127+
.ProcessInParallel();
128+
129+
var executeTask = processor.ExecuteAsync().ToListAsync();
130+
await started.Task.WaitAsync(cancellationToken);
131+
await cancellationTokenSource.CancelAsync();
132+
133+
await AssertExecutionWaitsForCleanupAsync(
134+
executeTask,
135+
interrupted,
136+
allowCleanupToFinish,
137+
cancellationToken);
138+
}
139+
140+
private static async Task WaitUntilCanceledAsync(
141+
CancellationToken cancellationToken,
142+
TaskCompletionSource started,
143+
TaskCompletionSource interrupted)
144+
{
145+
started.TrySetResult();
146+
try
147+
{
148+
await Task.Delay(Timeout.InfiniteTimeSpan, cancellationToken);
149+
}
150+
catch (OperationCanceledException)
151+
{
152+
interrupted.TrySetResult();
153+
throw;
154+
}
155+
}
156+
157+
private static async Task<T> WaitUntilCanceledAsync<T>(
158+
T result,
159+
CancellationToken cancellationToken,
160+
TaskCompletionSource started,
161+
TaskCompletionSource interrupted)
162+
{
163+
await WaitUntilCanceledAsync(cancellationToken, started, interrupted);
164+
return result;
165+
}
166+
167+
private static async Task WaitUntilCanceledAfterCleanupAsync(
168+
CancellationToken cancellationToken,
169+
TaskCompletionSource started,
170+
TaskCompletionSource interrupted,
171+
TaskCompletionSource allowCleanupToFinish)
172+
{
173+
started.TrySetResult();
174+
try
175+
{
176+
await Task.Delay(Timeout.InfiniteTimeSpan, cancellationToken);
177+
}
178+
catch (OperationCanceledException)
179+
{
180+
interrupted.TrySetResult();
181+
await allowCleanupToFinish.Task;
182+
throw;
183+
}
184+
}
185+
186+
private static async Task<T> WaitUntilCanceledAfterCleanupAsync<T>(
187+
T result,
188+
CancellationToken cancellationToken,
189+
TaskCompletionSource started,
190+
TaskCompletionSource interrupted,
191+
TaskCompletionSource allowCleanupToFinish)
192+
{
193+
await WaitUntilCanceledAfterCleanupAsync(
194+
cancellationToken,
195+
started,
196+
interrupted,
197+
allowCleanupToFinish);
198+
199+
return result;
200+
}
201+
202+
private static async Task AssertExecutionWaitsForCleanupAsync(
203+
Task executeTask,
204+
TaskCompletionSource interrupted,
205+
TaskCompletionSource allowCleanupToFinish,
206+
CancellationToken cancellationToken)
207+
{
208+
await interrupted.Task.WaitAsync(cancellationToken);
209+
await Task.Delay(100, cancellationToken);
210+
211+
try
212+
{
213+
await Assert.That(executeTask.IsCompleted).IsFalse();
214+
}
215+
finally
216+
{
217+
allowCleanupToFinish.TrySetResult();
218+
}
219+
220+
await Assert.ThrowsAsync<OperationCanceledException>(() => executeTask);
221+
}
222+
223+
private static async IAsyncEnumerable<int> GetItemsAsync(
224+
[EnumeratorCancellation] CancellationToken cancellationToken = default)
225+
{
226+
yield return 1;
227+
await Task.Delay(Timeout.InfiniteTimeSpan, cancellationToken);
228+
}
229+
}

‎EnumerableAsyncProcessor/Builders/ActionAsyncProcessorBuilder.cs‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,13 @@ public ActionAsyncProcessorBuilder(int count, Func<Task> taskSelector, Cancellat
1717
_cancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
1818
}
1919

20+
public ActionAsyncProcessorBuilder(int count, Func<CancellationToken, Task> taskSelector, CancellationToken cancellationToken)
21+
{
22+
_count = count;
23+
_cancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
24+
_taskSelector = () => taskSelector(_cancellationTokenSource.Token);
25+
}
26+
2027
public IAsyncProcessor ProcessInBatches(int batchSize)
2128
{
2229
return new BatchAsyncProcessor(batchSize, _count, _taskSelector, _cancellationTokenSource).StartProcessing();
@@ -77,4 +84,4 @@ public IAsyncProcessor ProcessOneAtATime()
7784
return new OneAtATimeAsyncProcessor(_count, _taskSelector, _cancellationTokenSource).StartProcessing();
7885
}
7986

80-
}
87+
}

‎EnumerableAsyncProcessor/Builders/ActionAsyncProcessorBuilder_1.cs‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,13 @@ internal ActionAsyncProcessorBuilder(int count, Func<Task<TOutput>> taskSelector
1717
_cancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
1818
}
1919

20+
internal ActionAsyncProcessorBuilder(int count, Func<CancellationToken, Task<TOutput>> taskSelector, CancellationToken cancellationToken)
21+
{
22+
_count = count;
23+
_cancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
24+
_taskSelector = () => taskSelector(_cancellationTokenSource.Token);
25+
}
26+
2027
public IAsyncProcessor<TOutput> ProcessInBatches(int batchSize)
2128
{
2229
return new ResultBatchAsyncProcessor<TOutput>(batchSize, _count, _taskSelector, _cancellationTokenSource).StartProcessing();
@@ -77,4 +84,4 @@ public IAsyncProcessor<TOutput> ProcessOneAtATime()
7784
return new ResultOneAtATimeAsyncProcessor<TOutput>(_count, _taskSelector, _cancellationTokenSource).StartProcessing();
7885
}
7986

80-
}
87+
}

‎EnumerableAsyncProcessor/Builders/AsyncEnumerableActionAsyncProcessorBuilder_1.cs‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,16 @@ public AsyncEnumerableActionAsyncProcessorBuilder(
1919
_cancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
2020
}
2121

22+
public AsyncEnumerableActionAsyncProcessorBuilder(
23+
IAsyncEnumerable<TInput> items,
24+
Func<TInput, CancellationToken, Task> taskSelector,
25+
CancellationToken cancellationToken)
26+
{
27+
_items = items;
28+
_cancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
29+
_taskSelector = item => taskSelector(item, _cancellationTokenSource.Token);
30+
}
31+
2232
/// <summary>
2333
/// Process items in parallel without concurrency limits.
2434
/// </summary>
@@ -70,7 +80,6 @@ public IAsyncEnumerableProcessor ProcessInParallel(int? maxConcurrency, bool sch
7080
_items, _taskSelector, maxConcurrency, scheduleOnThreadPool, _cancellationTokenSource);
7181
}
7282

73-
7483
/// <summary>
7584
/// Process items one at a time (sequential processing).
7685
/// </summary>
@@ -91,4 +100,4 @@ public IAsyncEnumerableProcessor ProcessInBatches(int batchSize)
91100
_items, _taskSelector, batchSize, _cancellationTokenSource);
92101
}
93102

94-
}
103+
}

‎EnumerableAsyncProcessor/Builders/AsyncEnumerableActionAsyncProcessorBuilder_2.cs‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,16 @@ public AsyncEnumerableActionAsyncProcessorBuilder(
1919
_cancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
2020
}
2121

22+
public AsyncEnumerableActionAsyncProcessorBuilder(
23+
IAsyncEnumerable<TInput> items,
24+
Func<TInput, CancellationToken, Task<TOutput>> taskSelector,
25+
CancellationToken cancellationToken)
26+
{
27+
_items = items;
28+
_cancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
29+
_taskSelector = item => taskSelector(item, _cancellationTokenSource.Token);
30+
}
31+
2232
/// <summary>
2333
/// Process items in parallel without concurrency limits and return results.
2434
/// </summary>
@@ -70,7 +80,6 @@ public IAsyncEnumerableProcessor<TOutput> ProcessInParallel(int? maxConcurrency,
7080
_items, _taskSelector, maxConcurrency, scheduleOnThreadPool, _cancellationTokenSource);
7181
}
7282

73-
7483
/// <summary>
7584
/// Process items one at a time and return results in order.
7685
/// </summary>
@@ -91,4 +100,4 @@ public IAsyncEnumerableProcessor<TOutput> ProcessInBatches(int batchSize)
91100
_items, _taskSelector, batchSize, _cancellationTokenSource);
92101
}
93102

94-
}
103+
}

‎EnumerableAsyncProcessor/Builders/AsyncEnumerableAsyncProcessorBuilder.cs‎

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,19 @@ public AsyncEnumerableActionAsyncProcessorBuilder<TInput, TOutput> SelectAsync<T
2222
return new AsyncEnumerableActionAsyncProcessorBuilder<TInput, TOutput>(_items, taskSelector, cancellationToken);
2323
}
2424

25+
public AsyncEnumerableActionAsyncProcessorBuilder<TInput, TOutput> SelectAsync<TOutput>(
26+
Func<TInput, CancellationToken, Task<TOutput>> taskSelector)
27+
{
28+
return SelectAsync(taskSelector, CancellationToken.None);
29+
}
30+
31+
public AsyncEnumerableActionAsyncProcessorBuilder<TInput, TOutput> SelectAsync<TOutput>(
32+
Func<TInput, CancellationToken, Task<TOutput>> taskSelector,
33+
CancellationToken cancellationToken)
34+
{
35+
return new AsyncEnumerableActionAsyncProcessorBuilder<TInput, TOutput>(_items, taskSelector, cancellationToken);
36+
}
37+
2538
public AsyncEnumerableActionAsyncProcessorBuilder<TInput> ForEachAsync(
2639
Func<TInput, Task> taskSelector)
2740
{
@@ -34,4 +47,17 @@ public AsyncEnumerableActionAsyncProcessorBuilder<TInput> ForEachAsync(
3447
{
3548
return new AsyncEnumerableActionAsyncProcessorBuilder<TInput>(_items, taskSelector, cancellationToken);
3649
}
37-
}
50+
51+
public AsyncEnumerableActionAsyncProcessorBuilder<TInput> ForEachAsync(
52+
Func<TInput, CancellationToken, Task> taskSelector)
53+
{
54+
return ForEachAsync(taskSelector, CancellationToken.None);
55+
}
56+
57+
public AsyncEnumerableActionAsyncProcessorBuilder<TInput> ForEachAsync(
58+
Func<TInput, CancellationToken, Task> taskSelector,
59+
CancellationToken cancellationToken)
60+
{
61+
return new AsyncEnumerableActionAsyncProcessorBuilder<TInput>(_items, taskSelector, cancellationToken);
62+
}
63+
}

0 commit comments

Comments
 (0)