Creating a concurrent/parallel ForEachAsync that can fail immediately or wait until enumeration is complete

Question

I needed an asynchronous parallel ForEach mechanism so I could enumerate an enumerable by N degrees of concurrency and process an action for each item. Searching around, my code was modeled/merged from a few posts I found.

The tests I have been running are:

1. Run body() to completion
2. Inside the body(), do a CancellationTokenSource.Cancel()
3. Inside the body(), throw a OperationCancelledException()
4. Inside the body(), throw a NotSupportedException()

First, this points to a Stephen Toub post here. I came up with combination of Toub's methods to have maxDegreeOfConcurrency and exception options (following bits from here)..

public static async Task ForEachAsync<T>( this IEnumerable<T> source, Func<T, Task> body, AsyncParallelOptions parallelOptions )
{
    ConcurrentBag<Exception> exceptions = null;

    var maxDegreeOfConcurrency = parallelOptions.MaxDegreeOfParallelism;

    // If they pass in a CancellationToken from caller of ForEachAsync need to create linked token source in case caller cancels, I want
    // ForEachAsync to cancel as well.  If they want to failImmediately, make a new CancellationTokenSource so I can stop processing partitions
    var cts = parallelOptions.CancellationToken != null ? CancellationTokenSource.CreateLinkedTokenSource( parallelOptions.CancellationToken ) :
              parallelOptions.FailImmediately ? new CancellationTokenSource() : null;

    var allDone = Task.WhenAll(
        from partition in Partitioner.Create( source ).GetPartitions( maxDegreeOfConcurrency )
        select Task.Run( async delegate
        {
            using ( partition )
            {
                while ( ( cts == null || !cts.IsCancellationRequested /* either from caller or failImmediately */ ) && partition.MoveNext() )
                {
                    await body( partition.Current )
                    .ContinueWith( t =>
                        {
                            // If body() threw an error, cancel if a CancellationTokenSource is present.
                            if ( t.IsFaulted )
                            {
                                if ( parallelOptions.FailImmediately )
                                {
                                    cts.Cancel();
                                }

                                // Always gather the exception to throw at the end
                                if ( exceptions == null ) exceptions = new ConcurrentBag<Exception>();

                                foreach ( var ex in t.Exception.InnerExceptions )
                                {
                                    exceptions.Add( ex );
                                }
                            }
                        }
                );
                }
            }
        }, cts.Token ) );

    // Wait until all finished (or errored out) and then return exceptions
    await allDone;

    // Question: allDone is never going to have IsCanceled or IsFaulted correct?  because await body() will swallow all exceptions?
    Console.WriteLine($"ForEachAsync Extension: Check Exceptions, {exceptions?.Count ?? 0} total, allDone.IsCanceled: {allDone.IsCanceled}, allDone.IsFaulted: {allDone.IsFaulted}, cts.IsCancellationRequested: {cts.IsCancellationRequested}" );

    if ( exceptions != null )
    {
        throw new AggregateException( exceptions );
    }
    else if ( cts.IsCancellationRequested )
    {
        throw new OperationCanceledException();
    }
}

Answer 1

1. CancaellationToken is a struct, not a class, so it can never be null. When people don't want to cancel stuff they just pass CancellationToken.None. 
2. However note that If you change your cts assignment logic to something like parallelOptions.CancellationToken != CancellationToken.None you could hit a NullReferenceException when the cancellation token is None and FailImmediately is false. One simple way to make things work would be to just always assign cts = CancellationTokenSource.CreateLinkedTokenSource(parallelOptions.CancellationToken). Then you won't need the cts nullity check in the loop as well.
3. Instead of checking for IsCancellationRequested directly, it's a better practice to call cts.Token.ThrowIfCancellationRequested() - that way you will actually mark the task as canceled.
4. Instead of checking t.IsFaulted, it's a better practice to check if t.Exception != null.
5. Instead of iterating over t.Exception.InnerExceptions, you should iterate over t.Exception.Flatten().InnerExceptions`.
6. Creating the ConcurrentBag inside the task like you're doing now is not thread safe. I suggest just creating it upon declaration (instead of assigning null like you do now).
7. allDone issues:
   1. allDone could be faulted if one of the body methods throw synchronously. 
   2. allDone would have IsCanceled set if cancellation was initiated either by the caller or by the FailImmediately logic
   3. In either case, you won't get to the following line as await allDone would throw.
8. You don't need else after you throw new AggregateException.
9. If all you do is re-throw the exceptions at the end of the method, you could just await the Task.WhenAll and get the same behavior for free (without the Console.WriteLine though). To make this work, you'll have to re-throw the exception instead of adding its inner exceptions to the concurrent bag.