Skip to content

Commit 03785b4

Browse files
thomhurstthomhurst
committed
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 e1cb4b7 commit 03785b4

3 files changed

Lines changed: 155 additions & 37 deletions

File tree

‎EnumerableAsyncProcessor.UnitTests/CancellationAwareSelectorTests.cs‎

Lines changed: 101 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
using System;
22
using System.Collections.Generic;
3+
using System.Runtime.CompilerServices;
34
using System.Threading;
45
using System.Threading.Tasks;
56
using EnumerableAsyncProcessor.Builders;
@@ -83,19 +84,57 @@ public async Task ExternalCancellation_Interrupts_AsyncEnumerable_Selector(Cance
8384
using var cancellationTokenSource = new CancellationTokenSource();
8485
var started = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
8586
var interrupted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
87+
var allowCleanupToFinish = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
8688

8789
var processor = GetItemsAsync()
8890
.ForEachAsync(
89-
(_, processorToken) => WaitUntilCanceledAsync(processorToken, started, interrupted),
91+
(_, processorToken) => WaitUntilCanceledAfterCleanupAsync(
92+
processorToken,
93+
started,
94+
interrupted,
95+
allowCleanupToFinish),
9096
cancellationTokenSource.Token)
9197
.ProcessInParallel();
9298

9399
var executeTask = processor.ExecuteAsync();
94100
await started.Task.WaitAsync(cancellationToken);
95101
await cancellationTokenSource.CancelAsync();
96102

97-
await interrupted.Task.WaitAsync(cancellationToken);
98-
await Assert.ThrowsAsync<OperationCanceledException>(() => executeTask);
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);
99138
}
100139

101140
private static async Task WaitUntilCanceledAsync(
@@ -125,9 +164,66 @@ private static async Task<T> WaitUntilCanceledAsync<T>(
125164
return result;
126165
}
127166

128-
private static async IAsyncEnumerable<int> GetItemsAsync()
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)
129225
{
130226
yield return 1;
131-
await Task.CompletedTask;
227+
await Task.Delay(Timeout.InfiniteTimeSpan, cancellationToken);
132228
}
133229
}

‎EnumerableAsyncProcessor/RunnableProcessors/AsyncEnumerable/AsyncEnumerableParallelProcessor.cs‎

Lines changed: 22 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -69,28 +69,33 @@ public async Task ExecuteAsync()
6969
{
7070
// Unbounded parallel processing
7171
var tasks = new List<Task>();
72-
73-
await foreach (var item in _items.WithCancellation(cancellationToken).ConfigureAwait(false))
72+
73+
try
7474
{
75-
var capturedItem = item;
76-
77-
Task task;
78-
if (_scheduleOnThreadPool)
79-
{
80-
task = Task.Run(async () => await _taskSelector(capturedItem).ConfigureAwait(false), cancellationToken);
81-
}
82-
else
75+
await foreach (var item in _items.WithCancellation(cancellationToken).ConfigureAwait(false))
8376
{
84-
task = _taskSelector(capturedItem);
77+
var capturedItem = item;
78+
79+
Task task;
80+
if (_scheduleOnThreadPool)
81+
{
82+
task = Task.Run(() => _taskSelector(capturedItem), cancellationToken);
83+
}
84+
else
85+
{
86+
task = _taskSelector(capturedItem);
87+
}
88+
89+
tasks.Add(task);
8590
}
86-
87-
tasks.Add(task);
8891
}
89-
90-
if (tasks.Count > 0)
92+
finally
9193
{
92-
await Task.WhenAll(tasks).ConfigureAwait(false);
94+
if (tasks.Count > 0)
95+
{
96+
await Task.WhenAll(tasks).ConfigureAwait(false);
97+
}
9398
}
9499
}
95100
}
96-
}
101+
}

‎EnumerableAsyncProcessor/RunnableProcessors/AsyncEnumerable/ResultProcessors/ResultAsyncEnumerableParallelProcessor.cs‎

Lines changed: 32 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -91,28 +91,45 @@ public async IAsyncEnumerable<TOutput> ExecuteAsync()
9191
else
9292
{
9393
// Unbounded parallel processing
94-
await foreach (var item in _items.WithCancellation(cancellationToken).ConfigureAwait(false))
94+
try
9595
{
96-
var capturedItem = item;
97-
98-
Task<TOutput> task;
99-
if (_scheduleOnThreadPool)
96+
await foreach (var item in _items.WithCancellation(cancellationToken).ConfigureAwait(false))
10097
{
101-
task = Task.Run(async () => await _taskSelector(capturedItem).ConfigureAwait(false), cancellationToken);
98+
var capturedItem = item;
99+
100+
Task<TOutput> task;
101+
if (_scheduleOnThreadPool)
102+
{
103+
task = Task.Run(() => _taskSelector(capturedItem), cancellationToken);
104+
}
105+
else
106+
{
107+
task = _taskSelector(capturedItem);
108+
}
109+
110+
tasks.Add(task);
102111
}
103-
else
112+
113+
// Yield all results as they complete
114+
await foreach (var result in tasks.ToIAsyncEnumerable(cancellationToken).ConfigureAwait(false))
104115
{
105-
task = _taskSelector(capturedItem);
116+
yield return result;
106117
}
107-
108-
tasks.Add(task);
109118
}
110-
111-
// Yield all results as they complete
112-
await foreach (var result in tasks.ToIAsyncEnumerable(cancellationToken).ConfigureAwait(false))
119+
finally
113120
{
114-
yield return result;
121+
if (tasks.Count > 0)
122+
{
123+
try
124+
{
125+
await Task.WhenAll(tasks).ConfigureAwait(false);
126+
}
127+
catch
128+
{
129+
// Preserve the exception already propagating from enumeration or result consumption.
130+
}
131+
}
115132
}
116133
}
117134
}
118-
}
135+
}

0 commit comments

Comments
 (0)