Concurrent/parallel ForEachAsync - proper handling of exceptions and cancellations

Question

I've created an asynchronous parallel ForEach mechanism so I could enumerate an enumerable by N degrees of concurrency and process an action for each item. Additionally, I have the option of stopping when first exception is encountered or processing the entire enumerable and returning all exceptions at the end. I'm trying to ensure that all my exception handling is being done correctly and I understand the flow correctly. My original code review can be found here (Original Answer). The code below incorporates suggestions from the original answer, but I'm not sure all the assumptions stated in that answer are happening.

My assumptions and concerns (some contrary to original answer) are posted with each test case. Please feel free to correct any incorrect assumptions.

Test Harness In All Test Cases

// This really would be a CreateLinkedTokenSource that could be cancelled
// from the caller of this code or internally within the body of the ForEachAsync body
var internalCancel = new CancellationTokenSource();
try
{
 // Simulate running something for 10 data batches...
 await Enumerable.Range(0, 10)
 .ForEachAsync(
 async jobData =>
 {
 // body of code that might throw exceptions or set internalCancel.Cancel()...
 internalCancel.Token.ThrowIfCancellationRequested();
 Console.WriteLine( $"Task {jobKey}: FooAsync - Finish DataChunk {jobData}" );
 },
 new AsyncParallelOptions { MaxDegreeOfParallelism = 4, CancellationToken = internalCancel.Token }
 );
}
catch (Exception ex)
{
 Console.WriteLine( $"Task {jobKey}: FooAsync - Exception: {ex.GetType().ToString()}, internalCancel.Token.IsCancellationRequested: {internalCancel.Token.IsCancellationRequested}" );
 throw;
}

Test Case 1 - Run body() to completion

This seems to work as expected.

Test Case 2 - Inside body() exception is thrown

Inside the body(), throw a NotSupportedException() (or any exception other than OperationCancelledException) and I get the following output.

ForEachAsync Extension #1: ContinueWith, t.Exception is null: False, t.IsCanceled: False, t.IsFaulted: True, cts.IsCancellationRequested: False
ForEachAsync Extension #2: Finished, 1 total, allDone.IsCanceled: False, allDone.IsFaulted: False, cts.IsCancellationRequested: True
ForEachAsync Extension #3: Throw Exceptions
Task 0: FooAsync - Exception: System.AggregateException, internalCancel.Token.IsCancellationRequested: False

My assumptions/concerns are:

1. Any exceptions thrown in body() are swallowed and presented inside ContinueWith() via t.Exception or t.IsFaulted.
2. Since I did not throw t.Exception and did not use cts.Token.ThrowIfCancellationRequested(), that is the reason why after await allDone, allDone.IsFaulted and allDone.IsCanceled is always false?
3. Original answer mentioned allDone.IsFaulted could be true if body() threw synchronously. I'm not sure what throwing synchronously means, but as my code stands I couldn't seem to ever have it set to true.
4. Original answer mentioned code after await allDone; would not run if thrown synchronously, but as my output displays, code is running after the await allDone. Is this because I didn't throw any exceptions from inside ContinueWith()?

Test Case 3 - Inside body() OperationCancelledException is thrown

Output:

ForEachAsync Extension #1: ContinueWith, t.Exception is null: True, t.IsCanceled: True, t.IsFaulted: False, cts.IsCancellationRequested: False
// Next three are dumps from the other 3 degrees of parallelism
ForEachAsync Extension #1: ContinueWith, t.Exception is null: True, t.IsCanceled: False, t.IsFaulted: False, cts.IsCancellationRequested: True
ForEachAsync Extension #1: ContinueWith, t.Exception is null: True, t.IsCanceled: False, t.IsFaulted: False, cts.IsCancellationRequested: False
ForEachAsync Extension #1: ContinueWith, t.Exception is null: True, t.IsCanceled: False, t.IsFaulted: False, cts.IsCancellationRequested: True
ForEachAsync Extension #2: Finished, 0 total, allDone.IsCanceled: False, allDone.IsFaulted: False, cts.IsCancellationRequested: True
ForEachAsync Extension #4: Throw OperationCanceledException
Task 0: FooAsync - Exception: System.OperationCanceledException, internalCancel.Token.IsCancellationRequested: False

Assumptions/Concerns:

1. Original answer stated that it was better practice if I used cts.Token.ThrowIfCancellationRequested() instead of checking cts.IsCancellationRequested however, I assume I'd place the ThrowIfCancellationRequested() as first line in while( partition.MoveNext() ) { } which means that if an exception or cancellation happened during the body, the next partition would be prepared for processing which could involve hitting web services and/or database calls. Checking cts.IsCancellationRequested seems to be a more efficient approach? Unless I'm over thinking things.
2. Original answer stated that I could just re throw t.Exception instead of adding to my collection then I could get my exception handling for free. I assume this means that await allDone; would then throw the exception and no code after that would be executed? This raises two questions.
   1. If I wanted entire enumerable processed and all exceptions thrown at end, this wouldn't allow that right? This would stop execution as soon as the first exception occurred?
   2. The suggestion to call t.Exception.Flatten() before adding to my exception would be lost, then requiring the caller of ForEachAsync to do that instead?

Test Case 4 - Inside body() internalCancel.Token.ThrowIfCancellationRequested() is triggered

Output:

ForEachAsync Extension #1: ContinueWith, t.Exception is null: True, t.IsCanceled: True, t.IsFaulted: False, cts.IsCancellationRequested: True
ForEachAsync Extension #1: ContinueWith, t.Exception is null: True, t.IsCanceled: True, t.IsFaulted: False, cts.IsCancellationRequested: True
ForEachAsync Extension #2: Finished, 0 total, allDone.IsCanceled: False, allDone.IsFaulted: False, cts.IsCancellationRequested: True
ForEachAsync Extension #4: Throw OperationCanceledException
Task 0: FooAsync - Exception: System.OperationCanceledException, internalCancel.Token.IsCancellationRequested: True

Assumptions/Concerns: 1. Similar to the above case, since I don't use cts.Token.ThrowIfCancellationRequested(); my while loop exits without exception. Then, exceptions.Count==0 so I simply throw a new OperationCancelledException and the flow seems to behave correctly (or at least how I expect it to).

If you've made it this far, I appreciate any and all comments.

ForEachAsync Extension Method

 public static async Task ForEachAsync<T>( this IEnumerable<T> source, Func<T, Task> body, AsyncParallelOptions parallelOptions )
 {
 ConcurrentBag<Exception> exceptions = new ConcurrentBag<Exception>();
 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 != CancellationToken.None
 ? CancellationTokenSource.CreateLinkedTokenSource( parallelOptions.CancellationToken )
 : new CancellationTokenSource();
 var allDone = Task.WhenAll(
 from partition in Partitioner.Create( source ).GetPartitions( maxDegreeOfConcurrency )
 select Task.Run( async delegate
 {
 using ( partition )
 {
 while ( !cts.IsCancellationRequested /* either from caller or failImmediately */ && partition.MoveNext() )
 {
 await body( partition.Current ).ContinueWith( t =>
 {
 Console.WriteLine( $"ForEachAsync Extension #1: ContinueWith, t.Exception is null: {t.Exception == null}, t.IsCanceled: {t.IsCanceled}, t.IsFaulted: {t.IsFaulted}, cts.IsCancellationRequested: {cts.IsCancellationRequested}" );
 // If body() threw an error, cancel if a CancellationTokenSource is present.
 if ( t.Exception != null )
 {
 if ( parallelOptions.FailImmediately )
 {
 cts.Cancel();
 }
 // Always gather the exception to throw at the end
 foreach ( var ex in t.Exception.Flatten().InnerExceptions )
 {
 exceptions.Add( ex );
 }
 }
 else if (t.IsCanceled)
 {
 // Needed in case OperationCanceledException() is thrown manually without calling 
 // .Cancel() on any linked token sources
 cts.Cancel(); 
 } 
 }
 );
 }
 }
 }, 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 #2: Finished, {exceptions?.Count ?? 0} total, allDone.IsCanceled: {allDone.IsCanceled}, allDone.IsFaulted: {allDone.IsFaulted}, cts.IsCancellationRequested: {cts.IsCancellationRequested}" );
 if ( exceptions.Count > 0 )
 {
 Console.WriteLine( $"ForEachAsync Extension #3: Throw Exceptions" );
 throw new AggregateException( exceptions );
 }
 if ( cts.IsCancellationRequested )
 {
 Console.WriteLine( $"ForEachAsync Extension #4: Throw OperationCanceledException" );
 throw new OperationCanceledException();
 }
 }

Answer 1

Test Case 2

1. Correct, but only if those exceptions are thrown asynchronously. See (3) below.
2. Correct. When you do await task.ContinueWith(...), you are no longer awaiting the original task, rather its continuation. The fact that the original task threw will only reflect in the task object passed to the continuation (e.g. t.Exception). You are allowed to observe the exception, but unless you throw anything from the continuation it will not be considered faulted or canceled.
3. If you throw synchronously (e.g. body = i => throw new Exception("foo")), you'll never get back a task and evaluation will stop before ContinueWith is ever called, let alone awaited. However, since you're using an async delegate in your test harness ( async jobData => {...}) that would never happen, as async delegates make sure a Task is always returned (and when you throw an exception, it is attached to that task).
4. It's because you don't throw from either body (synchronously) or from ContinueWith so basically in the eyes of the TPL everything is completing successfully.

Test Case 3

1. If that scenario concerns you, you could always do something like while (true) {cts.Token.ThrowIfCancellationRequested(); if (!partition.MoveNext()) {break;} }
2. Correct. 2(1). If you rethrow the exceptions from the continuation (and don't catch them in the outer await task.ContinueWith) then it won't be guaranteed either way. This is because execution would stop only for the offending partition's specific task. And since the partitions are dynamic, the enumerable will keep getting processed in the other tasks, and that will continue until all elements are processed, because Task.WhenAll waits for all tasks to complete (even if some failed along the way). Of course if you throw as many exceptions as your degree of parallelism, all tasks will fault and execution will stop. But that's a pretty random condition, so I'd say for this approach to be reasonable you'd have to always trigger cancellation as well (not just when FailImmediately is true). 2(2). Correct.

Test Case 4

There is a race condition here. If one of the tasks hits the cancellation code quickly enough, one of the Task.Run calls will complete as canceled before ever running its task (once it started running a task, it can't magically cancel it, so the token is only relevant before it actually started it). In that case, cancellation will propagate to allDone, and a TaskCanceledException will be thrown at await allDone. To see this in action (in high probability), try increasing your range to say 1000 and your degree of parallelism to say 100. The other option of the race condition is that the tasks are launched so quickly that nothing gets actually cancelled (since, again, true cancellation is only possibe before task launch in your code, and everything was launched by the time cancellation was triggered), leading to the scenario you've described.

More Comments

1. You should still be able to do cts = CancellationTokenSource.CreateLinkedTokenSource(parallelOptions.CancellationToken), don't see the need for the separation.
2. I don't like your manual check of t.IsCanceled, the user should specify cancelation via his token. Maybe he wanted to cancel a specific task but not the entire operation.
3. It turns out ConcurrentQueue (which is generally more lightweight) performs better in this use case, both in terms of run-time performance and memory. For proof: http://pastebin.com/ig47x6VV.

• Terry: Awesome information. I've applied your suggestions for Test case #1/Assumption 1, and your 'More Comments' section. Two questions. Test case #2/Assumption 3 - is there a way to force caller to use an async lamda? Just curious. Test case #4 - Race condition - I tried 1000 range and 100 degrees of parallelism and my code always seemed to 'exit' correctly. I don't seem to get a TaskCanceledException ever thrown. More importantly, any chance you could put a suggestion/code fix in your answer?
• Ohad Schneider: @Terry Test case #2/Assumption 3 - no, because async methods are just syntactic sugar for "plain" Task returning methods (the compiler transforms your async method to a regular Task returning method) . So there is no way to differentiate between them in your signature. Test case #4 it doesn't make sense that you can't repro this. Are you sure you have internalCancel.Cancel() inside the test body? Add something like Console.WriteLine(Interlocked.Increment(i) before using (partition) and also add a trace right after you cancel the token. Then share the output at pastebin.com
• Ohad Schneider: As for a fix (assuming that by fix you mean more consistent behavior that [almost] always throws when canceled) , like I said using ThrowIfCancellationRequested should do it as that would trigger a true TPL cancellation from within the task. Of course if you cancel it at the last minute, while the tasks are still running but are about to finish and are just after the last ThrowIfCancellationRequested, no exception will be thrown. Such is the nature of async calcellation. you could add another ThrowIfCancellationRequested right after await allDone for consistent behavior in that case.
• Terry: I've posted the final chapter of review (again with most of your suggestions) at codereview.stackexchange.com/questions/146129/…