From 004eb7556063cd3f5c508287089266e2571fb437 Mon Sep 17 00:00:00 2001
From: Stephen Toub <stoub@microsoft.com>
Date: Mon, 25 Jan 2021 12:00:57 -0500
Subject: [PATCH] Fix cancellation race condition in Parallel.ForEachAsync
 (#47396)

If cancellation is requested at just the "wrong" time, we may end up canceling the operation but reporting it as a successful completion, due to our relying on _registration.Token being the caller-provided token, but that not being the case if cancellation was requested prior to UnsafeRegister being called, in which case it'll return a default registration.
---
 .../src/System/Threading/Tasks/Parallel.ForEachAsync.cs    | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/src/libraries/System.Threading.Tasks.Parallel/src/System/Threading/Tasks/Parallel.ForEachAsync.cs b/src/libraries/System.Threading.Tasks.Parallel/src/System/Threading/Tasks/Parallel.ForEachAsync.cs
index ee3b794e09dfa..ee56e0d1a18d1 100644
--- a/src/libraries/System.Threading.Tasks.Parallel/src/System/Threading/Tasks/Parallel.ForEachAsync.cs
+++ b/src/libraries/System.Threading.Tasks.Parallel/src/System/Threading/Tasks/Parallel.ForEachAsync.cs
@@ -370,6 +370,8 @@ private static Task ForEachAsync<TSource>(IAsyncEnumerable<TSource> source, int
         /// <typeparam name="TSource">Specifies the type of data being enumerated.</typeparam>
         private abstract class ForEachAsyncState<TSource> : TaskCompletionSource, IThreadPoolWorkItem
         {
+            /// <summary>The caller-provided cancellation token.</summary>
+            private readonly CancellationToken _externalCancellationToken;
             /// <summary>Registration with caller-provided cancellation token.</summary>
             protected readonly CancellationTokenRegistration _registration;
             /// <summary>
@@ -409,6 +411,7 @@ protected ForEachAsyncState(Func<object, Task> taskBody, int dop, TaskScheduler
                     _executionContext = ExecutionContext.Capture();
                 }
 
+                _externalCancellationToken = cancellationToken;
                 _registration = cancellationToken.UnsafeRegister(static o => ((ForEachAsyncState<TSource>)o!).Cancellation.Cancel(), this);
             }
 
@@ -464,11 +467,11 @@ public void Complete()
                 Debug.Assert(_completionRefCount == 0, $"Expected {nameof(_completionRefCount)} == 0, got {_completionRefCount}");
 
                 bool taskSet;
-                if (_registration.Token.IsCancellationRequested)
+                if (_externalCancellationToken.IsCancellationRequested)
                 {
                     // The externally provided token had cancellation requested. Assume that any exceptions
                     // then are due to that, and just cancel the resulting task.
-                    taskSet = TrySetCanceled(_registration.Token);
+                    taskSet = TrySetCanceled(_externalCancellationToken);
                 }
                 else if (_exceptions is null)
                 {