Skip to content

Commit 691ecf5

Browse files
thomhurstthomhurst
authored
perf: use worker pool in InParallelAsync (#355)
1 parent 1f41055 commit 691ecf5

2 files changed

Lines changed: 227 additions & 60 deletions

File tree

‎EnumerableAsyncProcessor.UnitTests/ParallelExtensionsTests.cs‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
using System;
2+
using System.Collections.Concurrent;
3+
using System.Collections.Generic;
4+
using System.Linq;
5+
using System.Threading;
6+
using System.Threading.Tasks;
7+
using EnumerableAsyncProcessor.Extensions;
8+
9+
namespace EnumerableAsyncProcessor.UnitTests;
10+
11+
public class ParallelExtensionsTests
12+
{
13+
[Test]
14+
[NotInParallel]
15+
public async Task CpuBoundActionRunsInParallel()
16+
{
17+
using var workersStarted = new CountdownEvent(4);
18+
using var releaseWorkers = new ManualResetEventSlim();
19+
20+
var processingTask = Enumerable.Range(0, 4).InParallelAsync(
21+
4,
22+
(Action<int>)(_ =>
23+
{
24+
workersStarted.Signal();
25+
releaseWorkers.Wait(TimeSpan.FromSeconds(10));
26+
}));
27+
28+
var ranInParallel = workersStarted.Wait(TimeSpan.FromSeconds(10));
29+
releaseWorkers.Set();
30+
await processingTask;
31+
32+
await Assert.That(ranInParallel).IsTrue();
33+
}
34+
35+
[Test]
36+
public async Task WorkerPoolDrainsAllItemsAfterFailures()
37+
{
38+
var processedItems = new ConcurrentBag<int>();
39+
40+
var processingTask = Enumerable.Range(0, 20).InParallelAsync(
41+
3,
42+
(Func<int, Task>)(item =>
43+
{
44+
processedItems.Add(item);
45+
46+
return item % 4 == 0
47+
? Task.FromException(new InvalidOperationException($"item {item}"))
48+
: Task.CompletedTask;
49+
}));
50+
51+
try
52+
{
53+
await processingTask;
54+
}
55+
catch (InvalidOperationException)
56+
{
57+
// Expected: awaiting a multiply faulted Task surfaces its first exception.
58+
}
59+
60+
await Assert.That(processedItems.Count).IsEqualTo(20);
61+
await Assert.That(processingTask.Exception!.InnerExceptions.Count).IsEqualTo(5);
62+
}
63+
64+
[Test]
65+
public async Task SourceIsMaterializedOnceBeforeWorkersStart()
66+
{
67+
var enumerationCount = 0;
68+
69+
IEnumerable<int> Source()
70+
{
71+
Interlocked.Increment(ref enumerationCount);
72+
73+
for (var i = 0; i < 10; i++)
74+
{
75+
yield return i;
76+
}
77+
}
78+
79+
await Source().InParallelAsync(2, (Func<int, Task>)(_ => Task.CompletedTask));
80+
81+
await Assert.That(enumerationCount).IsEqualTo(1);
82+
}
83+
}

‎EnumerableAsyncProcessor/Extensions/ParallelExtensions.cs‎

Lines changed: 144 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -1,123 +1,207 @@
1-
namespace EnumerableAsyncProcessor.Extensions;
1+
using System.Collections.Concurrent;
2+
3+
namespace EnumerableAsyncProcessor.Extensions;
24

35
public static class ParallelExtensions
46
{
5-
public static async Task InParallelAsync<TSource, TResult>(
7+
public static Task InParallelAsync<TSource, TResult>(
68
this IEnumerable<TSource> source,
79
int levelOfParallelism,
810
Func<TSource, Task<TResult>> taskSelector)
911
{
10-
await InParallelAsync(source, levelOfParallelism, taskSelector, CancellationToken.None).ConfigureAwait(false);
12+
return InParallelAsync(source, levelOfParallelism, taskSelector, CancellationToken.None);
1113
}
1214

13-
public static async Task InParallelAsync<TSource, TResult>(
15+
public static Task InParallelAsync<TSource, TResult>(
1416
this IEnumerable<TSource> source,
1517
int levelOfParallelism,
1618
Func<TSource, Task<TResult>> taskSelector,
1719
CancellationToken cancellationToken)
1820
{
19-
if (levelOfParallelism <= 0)
20-
{
21-
levelOfParallelism = Environment.ProcessorCount;
22-
}
23-
24-
using var parallelLock = new SemaphoreSlim(initialCount:levelOfParallelism, maxCount:levelOfParallelism);
25-
26-
await Task.WhenAll(source.Select(item => ProcessAsync(item, taskSelector, parallelLock, cancellationToken))).ConfigureAwait(false);
21+
return StartWorkers(source, levelOfParallelism, taskSelector, cancellationToken);
2722
}
28-
29-
public static async Task InParallelAsync<TSource>(
23+
24+
public static Task InParallelAsync<TSource>(
3025
this IEnumerable<TSource> source,
3126
int levelOfParallelism,
3227
Func<TSource, Task> taskSelector)
3328
{
34-
await InParallelAsync(source, levelOfParallelism, taskSelector, CancellationToken.None).ConfigureAwait(false);
29+
return InParallelAsync(source, levelOfParallelism, taskSelector, CancellationToken.None);
3530
}
3631

37-
public static async Task InParallelAsync<TSource>(
32+
public static Task InParallelAsync<TSource>(
3833
this IEnumerable<TSource> source,
3934
int levelOfParallelism,
4035
Func<TSource, Task> taskSelector,
4136
CancellationToken cancellationToken)
4237
{
43-
if (levelOfParallelism <= 0)
44-
{
45-
levelOfParallelism = Environment.ProcessorCount;
46-
}
47-
48-
using var parallelLock = new SemaphoreSlim(initialCount:levelOfParallelism, maxCount:levelOfParallelism);
49-
50-
await Task.WhenAll(source.Select(item => ProcessAsync(item, taskSelector, parallelLock, cancellationToken))).ConfigureAwait(false);
38+
return StartWorkers(source, levelOfParallelism, taskSelector, cancellationToken);
5139
}
5240

5341
// Overloads for CPU-bound processing
54-
public static async Task InParallelAsync<TSource, TResult>(
42+
public static Task InParallelAsync<TSource, TResult>(
5543
this IEnumerable<TSource> source,
5644
int levelOfParallelism,
5745
Func<TSource, TResult> taskSelector,
5846
CancellationToken cancellationToken = default)
5947
{
60-
await InParallelAsync(source, levelOfParallelism, item => Task.FromResult(taskSelector(item)), cancellationToken).ConfigureAwait(false);
48+
return StartWorkers(
49+
source,
50+
levelOfParallelism,
51+
item =>
52+
{
53+
_ = taskSelector(item);
54+
return Task.CompletedTask;
55+
},
56+
cancellationToken);
6157
}
6258

63-
public static async Task InParallelAsync<TSource>(
59+
public static Task InParallelAsync<TSource>(
6460
this IEnumerable<TSource> source,
6561
int levelOfParallelism,
6662
Action<TSource> taskSelector,
6763
CancellationToken cancellationToken = default)
6864
{
69-
await InParallelAsync(source, levelOfParallelism, item => { taskSelector(item); return Task.CompletedTask; }, cancellationToken).ConfigureAwait(false);
65+
return StartWorkers(
66+
source,
67+
levelOfParallelism,
68+
item =>
69+
{
70+
taskSelector(item);
71+
return Task.CompletedTask;
72+
},
73+
cancellationToken);
7074
}
7175

72-
private static async Task<TResult> ProcessAsync<TSource, TResult>(
73-
TSource item,
74-
Func<TSource, Task<TResult>> taskSelector,
75-
SemaphoreSlim parallelLock,
76+
private static Task StartWorkers<TSource>(
77+
IEnumerable<TSource> source,
78+
int levelOfParallelism,
79+
Func<TSource, Task> taskSelector,
7680
CancellationToken cancellationToken)
7781
{
78-
var semaphoreAcquired = false;
79-
82+
TSource[] items;
83+
8084
try
8185
{
82-
await parallelLock.WaitAsync(cancellationToken).ConfigureAwait(false);
83-
semaphoreAcquired = true;
84-
85-
cancellationToken.ThrowIfCancellationRequested();
86+
items = source.ToArray();
87+
}
88+
catch (Exception exception)
89+
{
90+
return Task.FromException(exception);
91+
}
8692

87-
return await taskSelector(item).ConfigureAwait(false);
93+
if (items.Length == 0)
94+
{
95+
return Task.CompletedTask;
8896
}
89-
finally
97+
98+
if (levelOfParallelism <= 0)
9099
{
91-
if (semaphoreAcquired)
100+
levelOfParallelism = Environment.ProcessorCount;
101+
}
102+
103+
var workerCount = Math.Min(levelOfParallelism, items.Length);
104+
var completionSource = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
105+
var exceptions = new ConcurrentQueue<Exception>();
106+
var nextIndex = -1;
107+
var remainingWorkers = workerCount;
108+
var wasCanceled = 0;
109+
110+
for (var i = 0; i < workerCount; i++)
111+
{
112+
_ = Task.Run(async () =>
113+
{
114+
try
115+
{
116+
while (true)
117+
{
118+
if (cancellationToken.IsCancellationRequested)
119+
{
120+
Interlocked.Exchange(ref wasCanceled, 1);
121+
return;
122+
}
123+
124+
var index = Interlocked.Increment(ref nextIndex);
125+
126+
if (index >= items.Length)
127+
{
128+
return;
129+
}
130+
131+
Task? task = null;
132+
133+
try
134+
{
135+
cancellationToken.ThrowIfCancellationRequested();
136+
task = taskSelector(items[index]);
137+
await task.ConfigureAwait(false);
138+
}
139+
catch (OperationCanceledException)
140+
{
141+
Interlocked.Exchange(ref wasCanceled, 1);
142+
}
143+
catch (Exception exception)
144+
{
145+
EnqueueExceptions(exceptions, task, exception);
146+
}
147+
}
148+
}
149+
catch (Exception exception)
150+
{
151+
exceptions.Enqueue(exception);
152+
}
153+
finally
154+
{
155+
if (Interlocked.Decrement(ref remainingWorkers) == 0)
156+
{
157+
Complete(completionSource, exceptions, wasCanceled, cancellationToken);
158+
}
159+
}
160+
});
161+
}
162+
163+
return completionSource.Task;
164+
}
165+
166+
private static void EnqueueExceptions(
167+
ConcurrentQueue<Exception> exceptions,
168+
Task? task,
169+
Exception exception)
170+
{
171+
if (task is { IsFaulted: true })
172+
{
173+
foreach (var innerException in task.Exception!.InnerExceptions)
92174
{
93-
parallelLock.Release();
175+
exceptions.Enqueue(innerException);
94176
}
177+
178+
return;
95179
}
180+
181+
exceptions.Enqueue(exception);
96182
}
97-
98-
private static async Task ProcessAsync<TSource>(
99-
TSource item,
100-
Func<TSource, Task> taskSelector,
101-
SemaphoreSlim parallelLock,
183+
184+
private static void Complete(
185+
TaskCompletionSource completionSource,
186+
ConcurrentQueue<Exception> exceptions,
187+
int wasCanceled,
102188
CancellationToken cancellationToken)
103189
{
104-
var semaphoreAcquired = false;
105-
106-
try
190+
if (!exceptions.IsEmpty)
191+
{
192+
completionSource.TrySetException(exceptions);
193+
}
194+
else if (wasCanceled != 0)
107195
{
108-
await parallelLock.WaitAsync(cancellationToken).ConfigureAwait(false);
109-
semaphoreAcquired = true;
110-
111-
cancellationToken.ThrowIfCancellationRequested();
196+
var canceledToken = cancellationToken.IsCancellationRequested
197+
? cancellationToken
198+
: new CancellationToken(canceled: true);
112199

113-
await taskSelector(item).ConfigureAwait(false);
200+
completionSource.TrySetCanceled(canceledToken);
114201
}
115-
finally
202+
else
116203
{
117-
if (semaphoreAcquired)
118-
{
119-
parallelLock.Release();
120-
}
204+
completionSource.TrySetResult();
121205
}
122206
}
123207
}

0 commit comments

Comments
 (0)