Use worker pool for InParallelAsync

EnumerableAsyncProcessor.UnitTests/ParallelExtensionsTests.cs

@@ -0,0 +1,83 @@
+using System;
+using System.Collections.Concurrent;
+using System.Collections.Generic;
+using System.Linq;
+using System.Threading;
+using System.Threading.Tasks;
+using EnumerableAsyncProcessor.Extensions;
+
+namespace EnumerableAsyncProcessor.UnitTests;
+
+public class ParallelExtensionsTests
+{
+    [Test]
+    [NotInParallel]
+    public async Task CpuBoundActionRunsInParallel()
+    {
+        using var workersStarted = new CountdownEvent(4);
+        using var releaseWorkers = new ManualResetEventSlim();
+
+        var processingTask = Enumerable.Range(0, 4).InParallelAsync(
+            4,
+            (Action<int>)(_ =>
+            {
+                workersStarted.Signal();
+                releaseWorkers.Wait(TimeSpan.FromSeconds(10));
+            }));
+
+        var ranInParallel = workersStarted.Wait(TimeSpan.FromSeconds(10));
+        releaseWorkers.Set();
+        await processingTask;
+
+        await Assert.That(ranInParallel).IsTrue();
+    }
+
+    [Test]
+    public async Task WorkerPoolDrainsAllItemsAfterFailures()
+    {
+        var processedItems = new ConcurrentBag<int>();
+
+        var processingTask = Enumerable.Range(0, 20).InParallelAsync(
+            3,
+            (Func<int, Task>)(item =>
+            {
+                processedItems.Add(item);
+
+                return item % 4 == 0
+                    ? Task.FromException(new InvalidOperationException($"item {item}"))
+                    : Task.CompletedTask;
+            }));
+
+        try
+        {
+            await processingTask;
+        }
+        catch (InvalidOperationException)
+        {
+            // Expected: awaiting a multiply faulted Task surfaces its first exception.
+        }
+
+        await Assert.That(processedItems.Count).IsEqualTo(20);
+        await Assert.That(processingTask.Exception!.InnerExceptions.Count).IsEqualTo(5);
+    }
+
+    [Test]
+    public async Task SourceIsMaterializedOnceBeforeWorkersStart()
+    {
+        var enumerationCount = 0;
+
+        IEnumerable<int> Source()
+        {
+            Interlocked.Increment(ref enumerationCount);
+
+            for (var i = 0; i < 10; i++)
+            {
+                yield return i;
+            }
+        }
+
+        await Source().InParallelAsync(2, (Func<int, Task>)(_ => Task.CompletedTask));
+
+        await Assert.That(enumerationCount).IsEqualTo(1);
+    }
+}

EnumerableAsyncProcessor/Extensions/ParallelExtensions.cs

@@ -1,123 +1,207 @@
-namespace EnumerableAsyncProcessor.Extensions;
+using System.Collections.Concurrent;
+
+namespace EnumerableAsyncProcessor.Extensions;
 
 public static class ParallelExtensions
 {
-    public static async Task InParallelAsync<TSource, TResult>(
+    public static Task InParallelAsync<TSource, TResult>(
         this IEnumerable<TSource> source,
         int levelOfParallelism,
         Func<TSource, Task<TResult>> taskSelector)
     {
-        await InParallelAsync(source, levelOfParallelism, taskSelector, CancellationToken.None).ConfigureAwait(false);
+        return InParallelAsync(source, levelOfParallelism, taskSelector, CancellationToken.None);
     }
 
-    public static async Task InParallelAsync<TSource, TResult>(
+    public static Task InParallelAsync<TSource, TResult>(
         this IEnumerable<TSource> source,
         int levelOfParallelism,
         Func<TSource, Task<TResult>> taskSelector,
         CancellationToken cancellationToken)
     {
-        if (levelOfParallelism <= 0)
-        {
-            levelOfParallelism = Environment.ProcessorCount;
-        }
-
-        using var parallelLock = new SemaphoreSlim(initialCount:levelOfParallelism, maxCount:levelOfParallelism);
-
-        await Task.WhenAll(source.Select(item => ProcessAsync(item, taskSelector, parallelLock, cancellationToken))).ConfigureAwait(false);
+        return StartWorkers(source, levelOfParallelism, taskSelector, cancellationToken);
     }
-    
-    public static async Task InParallelAsync<TSource>(
+
+    public static Task InParallelAsync<TSource>(
         this IEnumerable<TSource> source,
         int levelOfParallelism,
         Func<TSource, Task> taskSelector)
     {
-        await InParallelAsync(source, levelOfParallelism, taskSelector, CancellationToken.None).ConfigureAwait(false);
+        return InParallelAsync(source, levelOfParallelism, taskSelector, CancellationToken.None);
     }
 
-    public static async Task InParallelAsync<TSource>(
+    public static Task InParallelAsync<TSource>(
         this IEnumerable<TSource> source,
         int levelOfParallelism,
         Func<TSource, Task> taskSelector,
         CancellationToken cancellationToken)
     {
-        if (levelOfParallelism <= 0)
-        {
-            levelOfParallelism = Environment.ProcessorCount;
-        }
-
-        using var parallelLock = new SemaphoreSlim(initialCount:levelOfParallelism, maxCount:levelOfParallelism);
-
-        await Task.WhenAll(source.Select(item => ProcessAsync(item, taskSelector, parallelLock, cancellationToken))).ConfigureAwait(false);
+        return StartWorkers(source, levelOfParallelism, taskSelector, cancellationToken);
     }
 
     // Overloads for CPU-bound processing
-    public static async Task InParallelAsync<TSource, TResult>(
+    public static Task InParallelAsync<TSource, TResult>(
         this IEnumerable<TSource> source,
         int levelOfParallelism,
         Func<TSource, TResult> taskSelector,
         CancellationToken cancellationToken = default)
     {
-        await InParallelAsync(source, levelOfParallelism, item => Task.FromResult(taskSelector(item)), cancellationToken).ConfigureAwait(false);
+        return StartWorkers(
+            source,
+            levelOfParallelism,
+            item =>
+            {
+                _ = taskSelector(item);
+                return Task.CompletedTask;
+            },
+            cancellationToken);
     }
 
-    public static async Task InParallelAsync<TSource>(
+    public static Task InParallelAsync<TSource>(
         this IEnumerable<TSource> source,
         int levelOfParallelism,
         Action<TSource> taskSelector,
         CancellationToken cancellationToken = default)
     {
-        await InParallelAsync(source, levelOfParallelism, item => { taskSelector(item); return Task.CompletedTask; }, cancellationToken).ConfigureAwait(false);
+        return StartWorkers(
+            source,
+            levelOfParallelism,
+            item =>
+            {
+                taskSelector(item);
+                return Task.CompletedTask;
+            },
+            cancellationToken);
     }
 
-    private static async Task<TResult> ProcessAsync<TSource, TResult>(
-        TSource item,
-        Func<TSource, Task<TResult>> taskSelector,
-        SemaphoreSlim parallelLock,
+    private static Task StartWorkers<TSource>(
+        IEnumerable<TSource> source,
+        int levelOfParallelism,
+        Func<TSource, Task> taskSelector,
         CancellationToken cancellationToken)
     {
-        var semaphoreAcquired = false;
-        
+        TSource[] items;
+
         try
         {
-            await parallelLock.WaitAsync(cancellationToken).ConfigureAwait(false);
-            semaphoreAcquired = true;
-
-            cancellationToken.ThrowIfCancellationRequested();
+            items = source.ToArray();
+        }
+        catch (Exception exception)
+        {
+            return Task.FromException(exception);
+        }
 
-            return await taskSelector(item).ConfigureAwait(false);
+        if (items.Length == 0)
+        {
+            return Task.CompletedTask;
         }
-        finally
+
+        if (levelOfParallelism <= 0)
         {
-            if (semaphoreAcquired)
+            levelOfParallelism = Environment.ProcessorCount;
+        }
+
+        var workerCount = Math.Min(levelOfParallelism, items.Length);
+        var completionSource = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+        var exceptions = new ConcurrentQueue<Exception>();
+        var nextIndex = -1;
+        var remainingWorkers = workerCount;
+        var wasCanceled = 0;
+
+        for (var i = 0; i < workerCount; i++)
+        {
+            _ = Task.Run(async () =>
+            {
+                try
+                {
+                    while (true)
+                    {
+                        if (cancellationToken.IsCancellationRequested)
+                        {
+                            Interlocked.Exchange(ref wasCanceled, 1);
+                            return;
+                        }
+
+                        var index = Interlocked.Increment(ref nextIndex);
+
+                        if (index >= items.Length)
+                        {
+                            return;
+                        }
+
+                        Task? task = null;
+
+                        try
+                        {
+                            cancellationToken.ThrowIfCancellationRequested();
+                            task = taskSelector(items[index]);
+                            await task.ConfigureAwait(false);
+                        }
+                        catch (OperationCanceledException)
+                        {
+                            Interlocked.Exchange(ref wasCanceled, 1);
+                        }
+                        catch (Exception exception)
+                        {
+                            EnqueueExceptions(exceptions, task, exception);
+                        }
+                    }
+                }
+                catch (Exception exception)
+                {
+                    exceptions.Enqueue(exception);
+                }
+                finally
+                {
+                    if (Interlocked.Decrement(ref remainingWorkers) == 0)
+                    {
+                        Complete(completionSource, exceptions, wasCanceled, cancellationToken);
+                    }
+                }
+            });
+        }
+
+        return completionSource.Task;
+    }
+
+    private static void EnqueueExceptions(
+        ConcurrentQueue<Exception> exceptions,
+        Task? task,
+        Exception exception)
+    {
+        if (task is { IsFaulted: true })
+        {
+            foreach (var innerException in task.Exception!.InnerExceptions)
             {
-                parallelLock.Release();
+                exceptions.Enqueue(innerException);
             }
+
+            return;
         }
+
+        exceptions.Enqueue(exception);
     }
-    
-    private static async Task ProcessAsync<TSource>(
-        TSource item,
-        Func<TSource, Task> taskSelector,
-        SemaphoreSlim parallelLock,
+
+    private static void Complete(
+        TaskCompletionSource completionSource,
+        ConcurrentQueue<Exception> exceptions,
+        int wasCanceled,
         CancellationToken cancellationToken)
     {
-        var semaphoreAcquired = false;
-
-        try
+        if (!exceptions.IsEmpty)
+        {
+            completionSource.TrySetException(exceptions);
+        }
+        else if (wasCanceled != 0)
         {
-            await parallelLock.WaitAsync(cancellationToken).ConfigureAwait(false);
-            semaphoreAcquired = true;
-
-            cancellationToken.ThrowIfCancellationRequested();
+            var canceledToken = cancellationToken.IsCancellationRequested
+                ? cancellationToken
+                : new CancellationToken(canceled: true);
 
-            await taskSelector(item).ConfigureAwait(false);
+            completionSource.TrySetCanceled(canceledToken);
         }
-        finally
+        else
         {
-            if (semaphoreAcquired)
-            {
-                parallelLock.Release();
-            }
+            completionSource.TrySetResult();
         }
     }
 }