diff --git a/EnumerableAsyncProcessor.UnitTests/CancellationRegressionTests.cs b/EnumerableAsyncProcessor.UnitTests/CancellationRegressionTests.cs new file mode 100644 index 0000000..e3518a4 --- /dev/null +++ b/EnumerableAsyncProcessor.UnitTests/CancellationRegressionTests.cs @@ -0,0 +1,70 @@ +using System; +using System.Collections.Generic; +using System.Threading; +using System.Threading.Tasks; +using EnumerableAsyncProcessor.Extensions; + +namespace EnumerableAsyncProcessor.UnitTests; + +/// +/// Guards against cancellation hanging the bounded async-enumerable result pipeline: +/// a worker that observed cancellation between an item being written to the channel +/// and that item being claimed used to abandon the item's completion source, leaving +/// the consumer awaiting a task that never completes. +/// +public class CancellationRegressionTests +{ + [Test, Timeout(120_000)] + public async Task Cancelling_Bounded_Result_Streaming_Never_Hangs(CancellationToken cancellationToken) + { + for (var iteration = 0; iteration < 100; iteration++) + { + using var cts = new CancellationTokenSource(); + var cancelAfter = 1 + iteration % 10; + var consumed = 0; + + var stream = InfiniteSource() + .SelectAsync(async item => + { + await Task.Yield(); + return item; + }, cts.Token) + .ProcessInParallel(maxConcurrency: 2); + + async Task ConsumeAsync() + { + await foreach (var _ in stream.ExecuteAsync()) + { + if (Interlocked.Increment(ref consumed) == cancelAfter) + { + cts.Cancel(); + } + } + } + + Exception? caught = null; + try + { + await ConsumeAsync().WaitAsync(TimeSpan.FromSeconds(10), cancellationToken); + } + catch (Exception exception) + { + caught = exception; + } + + // A TimeoutException here means the pipeline hung instead of observing cancellation. + await Assert.That(caught is OperationCanceledException).IsTrue(); + } + } + + // Deliberately ignores cancellation so the pipeline's own guards are what stop it. + private static async IAsyncEnumerable InfiniteSource() + { + var i = 0; + while (true) + { + yield return i++; + await Task.Yield(); + } + } +} diff --git a/EnumerableAsyncProcessor.UnitTests/ExceptionFidelityTests.cs b/EnumerableAsyncProcessor.UnitTests/ExceptionFidelityTests.cs index 14de4b7..26ea6b8 100644 --- a/EnumerableAsyncProcessor.UnitTests/ExceptionFidelityTests.cs +++ b/EnumerableAsyncProcessor.UnitTests/ExceptionFidelityTests.cs @@ -100,6 +100,57 @@ public async Task Failed_Items_Do_Not_Prevent_Remaining_Items_From_Processing() await Assert.That(processor.GetEnumerableTasks().Count(x => x.IsFaulted)).IsEqualTo(10); } + [Test] + public async Task Streaming_Results_Surface_The_Original_Exception_Unwrapped() + { + await using var processor = Enumerable.Range(0, 5).ToList() + .SelectAsync(i => i == 2 + ? Task.FromException(new InvalidOperationException("item 2 failed")) + : Task.FromResult(i)) + .ProcessInParallel(2); + + Exception? caught = null; + try + { + await foreach (var _ in processor.GetResultsAsyncEnumerable()) + { + } + } + catch (Exception exception) + { + caught = exception; + } + + await Assert.That(caught).IsNotNull(); + await Assert.That(caught!.GetType()).IsEqualTo(typeof(InvalidOperationException)); + await Assert.That(caught.Message).IsEqualTo("item 2 failed"); + } + + [Test] + public async Task Streaming_Results_Surface_Cancellation_As_OperationCanceledException_Not_AggregateException() + { + using var itemCancellation = new CancellationTokenSource(); + itemCancellation.Cancel(); + + await using var processor = new[] { 1 } + .SelectAsync(_ => Task.FromCanceled(itemCancellation.Token)) + .ProcessInParallel(2); + + Exception? caught = null; + try + { + await foreach (var _ in processor.GetResultsAsyncEnumerable()) + { + } + } + catch (Exception exception) + { + caught = exception; + } + + await Assert.That(caught is OperationCanceledException).IsTrue(); + } + [Test] public async Task Item_Task_Continuations_Are_Not_Forced_Inline() { diff --git a/EnumerableAsyncProcessor/AsyncEnumerableWorkerPool.cs b/EnumerableAsyncProcessor/AsyncEnumerableWorkerPool.cs index 6ae9bc2..d1163cd 100644 --- a/EnumerableAsyncProcessor/AsyncEnumerableWorkerPool.cs +++ b/EnumerableAsyncProcessor/AsyncEnumerableWorkerPool.cs @@ -79,15 +79,23 @@ internal static async IAsyncEnumerable ProcessResultsAsync 0) { - yield return await resultTask.ConfigureAwait(false); + var value = await pendingResults.Peek().WaitAsync(pipelineToken).ConfigureAwait(false); + pendingResults.Dequeue(); + yield return value; } await Task.WhenAll(workers).ConfigureAwait(false); @@ -105,6 +113,17 @@ internal static async IAsyncEnumerable ProcessResultsAsync _ = t.Exception, + CancellationToken.None, + TaskContinuationOptions.ExecuteSynchronously, + TaskScheduler.Default); + } } } @@ -169,19 +188,33 @@ private static Task[] StartResultWorkers( { workers[i] = Task.Run(async () => { - await foreach (var workItem in reader.ReadAllAsync(cancellationToken).ConfigureAwait(false)) + try { - Task? task = null; - - try + await foreach (var workItem in reader.ReadAllAsync(cancellationToken).ConfigureAwait(false)) { - task = taskSelector(workItem.Input); - workItem.CompletionSource.TrySetResult(await task.ConfigureAwait(false)); + Task? task = null; + + try + { + task = taskSelector(workItem.Input); + workItem.CompletionSource.TrySetResult(await task.ConfigureAwait(false)); + } + catch (Exception exception) + { + workItem.CompletionSource.TrySetFromFault(task, exception, cancellationToken); + } } - catch (Exception exception) + } + catch (OperationCanceledException) + { + // Cancellation can strand items that were written but never claimed; + // complete them so nothing awaiting their results hangs. + while (reader.TryRead(out var abandonedItem)) { - workItem.CompletionSource.TrySetFromFault(task, exception, cancellationToken); + abandonedItem.CompletionSource.TrySetCanceled(cancellationToken); } + + throw; } }, cancellationToken); } diff --git a/EnumerableAsyncProcessor/Extensions/EnumerableExtensions.cs b/EnumerableAsyncProcessor/Extensions/EnumerableExtensions.cs index 72bbe0f..422e24d 100644 --- a/EnumerableAsyncProcessor/Extensions/EnumerableExtensions.cs +++ b/EnumerableAsyncProcessor/Extensions/EnumerableExtensions.cs @@ -155,7 +155,10 @@ internal static async IAsyncEnumerable ToIAsyncEnumerable(this IEnumerable #if NET9_0_OR_GREATER await foreach (var task in Task.WhenEach(tasks).ConfigureAwait(false).WithCancellation(cancellationToken)) { - yield return task.Result; + // Await rather than .Result: the task is already complete, but .Result wraps + // failures in AggregateException while the net8 fallback path rethrows the + // original exception. Both paths must surface identical exceptions. + yield return await task.ConfigureAwait(false); } #else // Interleaving via completion-order buckets: each task's continuation claims the next