Fix cancellation unregistration in DataflowBlock.OutputAvailableAsync

OutputAvailableAsync is not unregistering from the supplied CancellationToken. If a cancelable token is supplied and is long lived, each call with that token to OutputAvailableAsync will add another callback into that token, and that will continue to grow until either the token is dropped or has been cancellation requested.  For a long-lived cancellation token, this is akin to a leak.

Author: stephentoub

src/libraries/System.Threading.Tasks.Dataflow/src/Base/DataflowBlock.cs

@@ -1476,6 +1476,14 @@ public static Task<bool> OutputAvailableAsync<TOutput>(
                         static state => OutputAvailableAsyncTarget<TOutput>.CancelAndUnlink(state, default),
 #endif
                         target);
+
+                    if (target.Task.IsCompleted)
+                    {
+                        // If the source calls to the target to complete it before we register for cancellation,
+                        // we need to ensure that we unregister from the cancellation token. If the task hasn't
+                        // yet completed, then its completion will handle the unregistration.
+                        target._ctr.Dispose();
+                    }
                 }
 
                 return target.Task;
@@ -1504,18 +1512,6 @@ public OutputAvailableAsyncTarget() :
             {
             }
 
-            /// <summary>
-            /// Cached continuation delegate that unregisters from cancellation and
-            /// marshals the antecedent's result to the return value.
-            /// </summary>
-            internal static readonly Func<Task<bool>, object?, bool> s_handleCompletion = (antecedent, state) =>
-            {
-                var target = state as OutputAvailableAsyncTarget<T>;
-                Debug.Assert(target != null, "Expected non-null target");
-                target._ctr.Dispose();
-                return antecedent.GetAwaiter().GetResult();
-            };
-
             /// <summary>Cancels the target and unlinks the target from the source.</summary>
             /// <param name="state">An OutputAvailableAsyncTarget.</param>
             /// <param name="cancellationToken">The token that triggered cancellation</param>
@@ -1551,13 +1547,16 @@ DataflowMessageStatus ITargetBlock<T>.OfferMessage(DataflowMessageHeader message
                 if (source == null) throw new ArgumentNullException(nameof(source));
 
                 TrySetResult(true);
+                _ctr.Dispose();
+
                 return DataflowMessageStatus.DecliningPermanently;
             }
 
             /// <include file='XmlDocs/CommonXmlDocComments.xml' path='CommonXmlDocComments/Blocks/Member[@name="Complete"]/*' />
             void IDataflowBlock.Complete()
             {
                 TrySetResult(false);
+                _ctr.Dispose();
             }
 
             /// <include file='XmlDocs/CommonXmlDocComments.xml' path='CommonXmlDocComments/Blocks/Member[@name="Fault"]/*' />
@@ -1569,6 +1568,7 @@ void IDataflowBlock.Fault(Exception exception)
                 }
 
                 TrySetResult(false);
+                _ctr.Dispose();
             }
 
             /// <include file='XmlDocs/CommonXmlDocComments.xml' path='CommonXmlDocComments/Blocks/Member[@name="Completion"]/*' />