From fa3babd8ec89bfbd2edb3124b17109aa45bf68e6 Mon Sep 17 00:00:00 2001
From: Stephen Toub <stoub@microsoft.com>
Date: Wed, 18 Oct 2023 17:36:37 -0400
Subject: [PATCH] Ensure Parallel.ForAsync unregisters from CancellationToken
 (#93672)

---
 .../src/System/Threading/Tasks/Parallel.ForEachAsync.cs      | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

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 b3c6d50cafe2e..48d76dff065ea 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
@@ -186,6 +186,7 @@ static bool CompareExchange(ref T location, T value, T comparand) =>
                     // If we're the last worker to complete, complete the operation.
                     if (state.SignalWorkerCompletedIterating())
                     {
+                        state.Dispose();
                         state.Complete();
                     }
                 }
@@ -745,7 +746,7 @@ public ValueTask DisposeAsync()
 
         /// <summary>Stores the state associated with an IAsyncEnumerable ForEachAsync operation, shared between all its workers.</summary>
         /// <typeparam name="T">Specifies the type of data being enumerated.</typeparam>
-        private sealed class ForEachState<T> : ForEachAsyncState<T>
+        private sealed class ForEachState<T> : ForEachAsyncState<T>, IDisposable
         {
             public T NextAvailable;
             public readonly T ToExclusive;
@@ -759,6 +760,8 @@ public ForEachState(
                 NextAvailable = fromExclusive;
                 ToExclusive = toExclusive;
             }
+
+            public void Dispose() => _registration.Dispose();
         }
     }
 }