Pool bounded async-enumerable workers (#358)

* perf: pool async-enumerable workers

Bounded async streams now use P persistent channel workers and capped ordered-result queues, keeping source read-ahead and coordination proportional to concurrency.

Refs #337

* fix: preserve async stream exception type (#358)

Author: thomhurst

EnumerableAsyncProcessor.UnitTests/AsyncEnumerableProcessorTests.cs

@@ -104,6 +104,71 @@ public async Task SelectAsync_ProcessInParallel_ReturnsAllTransformedItems()
         await Assert.That(results.OrderBy(x => x)).IsEquivalentTo(Enumerable.Range(1, 10).Select(x => x * 2));
     }
 
+    [Test]
+    public async Task SelectAsync_BoundedParallelism_PreservesInputOrder()
+    {
+        var results = await GenerateAsyncEnumerable(20)
+            .SelectAsync(async item =>
+            {
+                await Task.Delay((21 - item) * 2);
+                return item;
+            })
+            .ProcessInParallel(4)
+            .ExecuteAsync()
+            .ToListAsync();
+
+        await Assert.That(results.SequenceEqual(Enumerable.Range(1, 20))).IsTrue();
+    }
+
+    [Test, Timeout(10_000)]
+    public async Task BoundedParallelism_AppliesBackpressureToSource(CancellationToken cancellationToken)
+    {
+        const int maxConcurrency = 2;
+
+        var producedCount = 0;
+        var startedCount = 0;
+        var workersStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+        var releaseWorkers = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+
+        async IAsyncEnumerable<int> Source([EnumeratorCancellation] CancellationToken token = default)
+        {
+            for (var i = 0; i < 100; i++)
+            {
+                token.ThrowIfCancellationRequested();
+                Interlocked.Increment(ref producedCount);
+                yield return i;
+                await Task.Yield();
+            }
+        }
+
+        var processingTask = Source(cancellationToken)
+            .ForEachAsync(async _ =>
+            {
+                if (Interlocked.Increment(ref startedCount) == maxConcurrency)
+                {
+                    workersStarted.TrySetResult();
+                }
+
+                await releaseWorkers.Task;
+            }, cancellationToken)
+            .ProcessInParallel(maxConcurrency)
+            .ExecuteAsync();
+
+        try
+        {
+            await workersStarted.Task.WaitAsync(TimeSpan.FromSeconds(3), cancellationToken);
+            await Task.Delay(100, cancellationToken);
+
+            await Assert.That(producedCount).IsLessThanOrEqualTo((maxConcurrency * 2) + 1);
+        }
+        finally
+        {
+            releaseWorkers.TrySetResult();
+        }
+
+        await processingTask;
+    }
+
     [Test]
     public async Task ForEachAsync_ProcessInParallel_WithHighConcurrency_HandlesCorrectly()
     {
@@ -213,6 +278,34 @@ public async Task ForEachAsync_WithException_PropagatesException()
         var exception = await Assert.ThrowsAsync<InvalidOperationException>(async () => await task);
         await Assert.That(exception!.Message).IsEqualTo("Test exception");
     }
+
+    [Test, Timeout(10_000)]
+    public async Task ForEachAsync_WithConcurrentExceptions_PropagatesOriginalException(
+        CancellationToken cancellationToken)
+    {
+        var startedCount = 0;
+        var workersStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+        var releaseWorkers = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+
+        var task = GenerateAsyncEnumerable(2)
+            .ForEachAsync(async _ =>
+            {
+                if (Interlocked.Increment(ref startedCount) == 2)
+                {
+                    workersStarted.TrySetResult();
+                }
+
+                await releaseWorkers.Task.WaitAsync(cancellationToken);
+                throw new InvalidOperationException("Concurrent failure");
+            })
+            .ProcessInParallel(2)
+            .ExecuteAsync();
+
+        await workersStarted.Task.WaitAsync(cancellationToken);
+        releaseWorkers.TrySetResult();
+
+        await Assert.ThrowsAsync<InvalidOperationException>(() => task);
+    }
 
     [Test]
     public async Task ForEachAsync_ProcessInParallel_UnboundedConcurrency_ProcessesAllItems()

EnumerableAsyncProcessor/AsyncEnumerableWorkerPool.cs

@@ -0,0 +1,229 @@
+using System.Collections.Concurrent;
+using System.Runtime.CompilerServices;
+using System.Runtime.ExceptionServices;
+using System.Threading.Channels;
+
+namespace EnumerableAsyncProcessor;
+
+/// <summary>
+/// Processes asynchronous sources with a bounded channel and a fixed set of workers.
+/// Source read-ahead and queued results stay proportional to worker count.
+/// </summary>
+internal static class AsyncEnumerableWorkerPool
+{
+    internal static async Task ProcessAsync<TInput>(
+        IAsyncEnumerable<TInput> items,
+        Func<TInput, Task> taskSelector,
+        int workerCount,
+        CancellationToken cancellationToken)
+    {
+        using var pipelineCancellation = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+        var pipelineToken = pipelineCancellation.Token;
+        var channel = CreateChannel<TInput>(workerCount);
+        var exceptions = new ConcurrentQueue<Exception>();
+        var wasCanceled = 0;
+        var workers = StartWorkers(channel.Reader, taskSelector, workerCount, exceptions, () => Interlocked.Exchange(ref wasCanceled, 1), pipelineToken);
+
+        try
+        {
+            try
+            {
+                await foreach (var item in items.WithCancellation(pipelineToken).ConfigureAwait(false))
+                {
+                    await channel.Writer.WriteAsync(item, pipelineToken).ConfigureAwait(false);
+                }
+            }
+            catch (OperationCanceledException)
+            {
+                Interlocked.Exchange(ref wasCanceled, 1);
+            }
+            catch (Exception exception)
+            {
+                exceptions.Enqueue(exception);
+            }
+            finally
+            {
+                channel.Writer.TryComplete();
+            }
+
+            await Task.WhenAll(workers).ConfigureAwait(false);
+
+            ThrowIfFailed(exceptions, wasCanceled, cancellationToken);
+        }
+        finally
+        {
+            pipelineCancellation.Cancel();
+            channel.Writer.TryComplete();
+        }
+    }
+
+    internal static async IAsyncEnumerable<TOutput> ProcessResultsAsync<TInput, TOutput>(
+        IAsyncEnumerable<TInput> items,
+        Func<TInput, Task<TOutput>> taskSelector,
+        int workerCount,
+        [EnumeratorCancellation] CancellationToken cancellationToken)
+    {
+        using var pipelineCancellation = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+        var pipelineToken = pipelineCancellation.Token;
+        var channel = CreateChannel<ResultWorkItem<TInput, TOutput>>(workerCount);
+        var workers = StartResultWorkers(channel.Reader, taskSelector, workerCount, pipelineToken);
+        var pendingResults = new Queue<Task<TOutput>>(workerCount);
+
+        try
+        {
+            await foreach (var item in items.WithCancellation(pipelineToken).ConfigureAwait(false))
+            {
+                var completionSource = new TaskCompletionSource<TOutput>(TaskCreationOptions.RunContinuationsAsynchronously);
+                await channel.Writer.WriteAsync(new ResultWorkItem<TInput, TOutput>(item, completionSource), pipelineToken).ConfigureAwait(false);
+                pendingResults.Enqueue(completionSource.Task);
+
+                if (pendingResults.Count == workerCount)
+                {
+                    yield return await pendingResults.Dequeue().ConfigureAwait(false);
+                }
+            }
+
+            channel.Writer.TryComplete();
+
+            while (pendingResults.TryDequeue(out var resultTask))
+            {
+                yield return await resultTask.ConfigureAwait(false);
+            }
+
+            await Task.WhenAll(workers).ConfigureAwait(false);
+        }
+        finally
+        {
+            pipelineCancellation.Cancel();
+            channel.Writer.TryComplete();
+
+            try
+            {
+                await Task.WhenAll(workers).ConfigureAwait(false);
+            }
+            catch (OperationCanceledException) when (pipelineToken.IsCancellationRequested)
+            {
+                // Expected when enumeration is canceled or the consumer stops early.
+            }
+        }
+    }
+
+    private static Channel<T> CreateChannel<T>(int capacity)
+    {
+        return Channel.CreateBounded<T>(new BoundedChannelOptions(capacity)
+        {
+            SingleWriter = true,
+            SingleReader = false,
+            FullMode = BoundedChannelFullMode.Wait,
+            AllowSynchronousContinuations = false
+        });
+    }
+
+    private static Task[] StartWorkers<TInput>(
+        ChannelReader<TInput> reader,
+        Func<TInput, Task> taskSelector,
+        int workerCount,
+        ConcurrentQueue<Exception> exceptions,
+        Action recordCancellation,
+        CancellationToken cancellationToken)
+    {
+        var workers = new Task[workerCount];
+
+        for (var i = 0; i < workerCount; i++)
+        {
+            workers[i] = Task.Run(async () =>
+            {
+                await foreach (var item in reader.ReadAllAsync(cancellationToken).ConfigureAwait(false))
+                {
+                    Task? task = null;
+
+                    try
+                    {
+                        task = taskSelector(item);
+                        await task.ConfigureAwait(false);
+                    }
+                    catch (OperationCanceledException)
+                    {
+                        recordCancellation();
+                    }
+                    catch (Exception exception)
+                    {
+                        EnqueueExceptions(exceptions, task, exception);
+                    }
+                }
+            }, cancellationToken);
+        }
+
+        return workers;
+    }
+
+    private static Task[] StartResultWorkers<TInput, TOutput>(
+        ChannelReader<ResultWorkItem<TInput, TOutput>> reader,
+        Func<TInput, Task<TOutput>> taskSelector,
+        int workerCount,
+        CancellationToken cancellationToken)
+    {
+        var workers = new Task[workerCount];
+
+        for (var i = 0; i < workerCount; i++)
+        {
+            workers[i] = Task.Run(async () =>
+            {
+                await foreach (var workItem in reader.ReadAllAsync(cancellationToken).ConfigureAwait(false))
+                {
+                    Task<TOutput>? task = null;
+
+                    try
+                    {
+                        task = taskSelector(workItem.Input);
+                        workItem.CompletionSource.TrySetResult(await task.ConfigureAwait(false));
+                    }
+                    catch (Exception exception)
+                    {
+                        workItem.CompletionSource.TrySetFromFault(task, exception, cancellationToken);
+                    }
+                }
+            }, cancellationToken);
+        }
+
+        return workers;
+    }
+
+    private static void EnqueueExceptions(
+        ConcurrentQueue<Exception> exceptions,
+        Task? task,
+        Exception exception)
+    {
+        if (task is { IsFaulted: true })
+        {
+            foreach (var innerException in task.Exception!.InnerExceptions)
+            {
+                exceptions.Enqueue(innerException);
+            }
+
+            return;
+        }
+
+        exceptions.Enqueue(exception);
+    }
+
+    private static void ThrowIfFailed(
+        ConcurrentQueue<Exception> exceptions,
+        int wasCanceled,
+        CancellationToken cancellationToken)
+    {
+        if (exceptions.TryDequeue(out var firstException))
+        {
+            ExceptionDispatchInfo.Capture(firstException).Throw();
+        }
+
+        if (wasCanceled != 0)
+        {
+            throw new OperationCanceledException(cancellationToken);
+        }
+    }
+
+    private readonly record struct ResultWorkItem<TInput, TOutput>(
+        TInput Input,
+        TaskCompletionSource<TOutput> CompletionSource);
+}