Skip to content

Commit

Permalink
Use asynchronous lock in Parallel.ForEachAsync with synchronous enume…
Browse files Browse the repository at this point in the history 
…rable (#82501)

* Use asynchronous lock in Parallel.ForEachAsync with synchronous enumerable

Avoid blocking threads while waiting for access to the enumerator in the case of a slower MoveNext.

* Update src/libraries/System.Threading.Tasks.Parallel/src/System/Threading/Tasks/Parallel.ForEachAsync.cs

Co-authored-by: Tanner Gooding <tagoo@outlook.com>

---------

Co-authored-by: Tanner Gooding <tagoo@outlook.com>
  • Loading branch information
@stephentoub @tannergooding
stephentoub and tannergooding committed Feb 23, 2023
1 parent 72c9c36 commit 220437e
Showing 1 changed file with 31 additions and 17 deletions.
48 changes: 31 additions & 17 deletions ...aries/System.Threading.Tasks.Parallel/src/System/Threading/Tasks/Parallel.ForEachAsync.cs
View file
Expand Up @@ -93,17 +93,23 @@ private static Task ForEachAsync<TSource>(IEnumerable<TSource> source, int dop,
// Continue to loop while there are more elements to be processed.
while (!state.Cancellation.IsCancellationRequested)
{
// Get the next element from the enumerator. This requires asynchronously locking around MoveNextAsync/Current.
// Get the next element from the enumerator. This requires asynchronously locking around MoveNext/Current.
TSource element;
lock (state)
await state.AcquireLock();
try
{
if (!state.Enumerator.MoveNext())
if (state.Cancellation.IsCancellationRequested || // check now that the lock has been acquired
!state.Enumerator.MoveNext())
{
break;
}
element = state.Enumerator.Current;
}
finally
{
state.ReleaseLock();
}
// If the remaining dop allows it and we've not yet queued the next worker, do so now. We wait
// until after we've grabbed an item from the enumerator to a) avoid unnecessary contention on the
Expand Down Expand Up @@ -249,20 +255,11 @@ private static Task ForEachAsync<TSource>(IAsyncEnumerable<TSource> source, int
{
// Get the next element from the enumerator. This requires asynchronously locking around MoveNextAsync/Current.
TSource element;
await state.AcquireLock();
try
{
// TODO https://github.com/dotnet/runtime/issues/22144:
// Use a no-throwing await if/when one is available built-in.
await state.Lock.WaitAsync(state.Cancellation.Token);
}
catch (OperationCanceledException)
{
break;
}
try
{
if (!await state.Enumerator.MoveNextAsync())
if (state.Cancellation.IsCancellationRequested || // check now that the lock has been acquired
!await state.Enumerator.MoveNextAsync())
{
break;
}
Expand All @@ -271,7 +268,7 @@ private static Task ForEachAsync<TSource>(IAsyncEnumerable<TSource> source, int
}
finally
{
state.Lock.Release();
state.ReleaseLock();
}
// If the remaining dop allows it and we've not yet queued the next worker, do so now. We wait
Expand Down Expand Up @@ -354,6 +351,8 @@ private abstract class ForEachAsyncState<TSource> : TaskCompletionSource, IThrea
private readonly TaskScheduler _scheduler;
/// <summary>The <see cref="ExecutionContext"/> present at the time of the ForEachAsync invocation. This is only used if on the default scheduler.</summary>
private readonly ExecutionContext? _executionContext;
/// <summary>Semaphore used to provide exclusive access to the enumerator.</summary>
private readonly SemaphoreSlim _lock = new SemaphoreSlim(initialCount: 1, maxCount: 1);

/// <summary>The number of outstanding workers. When this hits 0, the operation has completed.</summary>
private int _completionRefCount;
Expand Down Expand Up @@ -417,6 +416,21 @@ public void QueueWorkerIfDopAvailable()
/// <returns>true if this is the last worker to complete iterating; otherwise, false.</returns>
public bool SignalWorkerCompletedIterating() => Interlocked.Decrement(ref _completionRefCount) == 0;

/// <summary>Asynchronously acquires exclusive access to the enumerator.</summary>
public Task AcquireLock() =>
// We explicitly don't pass this.Cancellation to WaitAsync. Doing so adds overhead, and it isn't actually
// necessary. All of the operations that monitor the lock are part of the same ForEachAsync operation, and the Task
// returned from ForEachAsync can't complete until all of the constituent operations have completed, including whoever
// holds the lock while this worker is waiting on the lock. Thus, the lock will need to be released for the overall
// operation to complete. Passing the token would allow the overall operation to potentially complete a bit faster in
// the face of cancellation, in exchange for making it a bit slower / more overhead in the common case of cancellation
// not being requested. We want to optimize for the latter. This also then avoids an exception throw / catch when
// cancellation is requested.
_lock.WaitAsync(CancellationToken.None);

/// <summary>Relinquishes exclusive access to the enumerator.</summary>
public void ReleaseLock() => _lock.Release();

/// <summary>Stores an exception and triggers cancellation in order to alert all workers to stop as soon as possible.</summary>
/// <param name="e">The exception.</param>
public void RecordException(Exception e)
Expand Down Expand Up @@ -444,6 +458,7 @@ public void Complete()
else if (_exceptions is null)
{
// Everything completed successfully.
Debug.Assert(!Cancellation.IsCancellationRequested);
taskSet = TrySetResult();
}
else
Expand Down Expand Up @@ -500,7 +515,6 @@ public void Dispose()
/// <typeparam name="TSource">Specifies the type of data being enumerated.</typeparam>
private sealed class AsyncForEachAsyncState<TSource> : ForEachAsyncState<TSource>, IAsyncDisposable
{
public readonly SemaphoreSlim Lock = new SemaphoreSlim(1, 1);
public readonly IAsyncEnumerator<TSource> Enumerator;

public AsyncForEachAsyncState(
Expand Down

0 comments on commit 220437e

Please sign in to comment.