Implement IAsyncDisposable on CancellationTokenRegistration (dotnet/coreclr#20645)

Commit migrated from https://github.com/dotnet/coreclr/commit/d66952478bd8f2edec3ea97b0d57072102513ad8

diff --git a/src/coreclr/src/System.Private.CoreLib/src/System/Threading/CancellationTokenRegistration.cs b/src/coreclr/src/System.Private.CoreLib/src/System/Threading/CancellationTokenRegistration.cs
index 815c9cc..4261b89 100644 (file)
--- a/src/coreclr/src/System.Private.CoreLib/src/System/Threading/CancellationTokenRegistration.cs
+++ b/src/coreclr/src/System.Private.CoreLib/src/System/Threading/CancellationTokenRegistration.cs
@@ -2,6 +2,8 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
+using System.Threading.Tasks;
+
 namespace System.Threading
 {
     /// <summary>
@@ -10,7 +12,7 @@ namespace System.Threading
     /// <remarks>
     /// To unregister a callback, dispose the corresponding Registration instance.
     /// </remarks>
-    public readonly struct CancellationTokenRegistration : IEquatable<CancellationTokenRegistration>, IDisposable
+    public readonly struct CancellationTokenRegistration : IEquatable<CancellationTokenRegistration>, IDisposable, IAsyncDisposable
     {
         private readonly long _id;
         private readonly CancellationTokenSource.CallbackNode _node;
@@ -37,6 +39,21 @@ namespace System.Threading
         }
 
         /// <summary>
+        /// Disposes of the registration and unregisters the target callback from the associated 
+        /// <see cref="T:System.Threading.CancellationToken">CancellationToken</see>.
+        /// The returned <see cref="ValueTask"/> will complete once the associated callback
+        /// is unregistered without having executed or once it's finished executing, except
+        /// in the degenerate case where the callback itself is unregistering itself.
+        /// </summary>
+        public ValueTask DisposeAsync()
+        {
+            CancellationTokenSource.CallbackNode node = _node;
+            return node != null && !node.Partition.Unregister(_id, node) ?
+                WaitForCallbackIfNecessaryAsync() :
+                default;
+        }
+
+        /// <summary>
         /// Gets the <see cref="CancellationToken"/> with which this registration is associated.  If the
         /// registration isn't associated with a token (such as after the registration has been disposed),
         /// this will return a default token.
@@ -69,12 +86,30 @@ namespace System.Threading
                 !source.IsCancellationCompleted && // Running callbacks hasn't finished.
                 source.ThreadIDExecutingCallbacks != Thread.CurrentThread.ManagedThreadId) // The executing thread ID is not this thread's ID.
             {
-                // Callback execution is in progress, the executing thread is different to us and has taken the callback for execution
+                // Callback execution is in progress, the executing thread is different from this thread and has taken the callback for execution
                 // so observe and wait until this target callback is no longer the executing callback.
                 source.WaitForCallbackToComplete(_id);
             }
         }
 
+        private ValueTask WaitForCallbackIfNecessaryAsync()
+        {
+            // Same as WaitForCallbackIfNecessary, except returning a task that'll be completed when callbacks complete.
+
+            CancellationTokenSource source = _node.Partition.Source;
+            if (source.IsCancellationRequested && // Running callbacks has commenced.
+                !source.IsCancellationCompleted && // Running callbacks hasn't finished.
+                source.ThreadIDExecutingCallbacks != Thread.CurrentThread.ManagedThreadId) // The executing thread ID is not this thread's ID.
+            {
+                // Callback execution is in progress, the executing thread is different from this thread and has taken the callback for execution
+                // so get a task that'll complete when this target callback is no longer the executing callback.
+                return source.WaitForCallbackToCompleteAsync(_id);
+            }
+
+            // Callback is either already completed, won't execute, or the callback itself is calling this.
+            return default;
+        }
+
         /// <summary>
         /// Determines whether two <see
         /// cref="T:System.Threading.CancellationTokenRegistration">CancellationTokenRegistration</see>

diff --git a/src/coreclr/src/System.Private.CoreLib/src/System/Threading/CancellationTokenSource.cs b/src/coreclr/src/System.Private.CoreLib/src/System/Threading/CancellationTokenSource.cs
index 8fd6220..d0656c3 100644 (file)
--- a/src/coreclr/src/System.Private.CoreLib/src/System/Threading/CancellationTokenSource.cs
+++ b/src/coreclr/src/System.Private.CoreLib/src/System/Threading/CancellationTokenSource.cs
@@ -4,6 +4,7 @@
 
 using System.Collections.Generic;
 using System.Diagnostics;
+using System.Threading.Tasks;
 
 namespace System.Threading
 {
@@ -786,6 +787,33 @@ namespace System.Threading
             }
         }
 
+        /// <summary>
+        /// Asynchronously wait for a single callback to complete (or, more specifically, to not be running).
+        /// It is ok to call this method if the callback has already finished.
+        /// Calling this method before the target callback has been selected for execution would be an error.
+        /// </summary>
+        internal ValueTask WaitForCallbackToCompleteAsync(long id)
+        {
+            // If the currently executing callback is not the target one, then the target one has already
+            // completed and we can simply return.  This should be the most common case, as the caller
+            // calls if we're currently canceling but doesn't know what callback is running, if any.
+            if (ExecutingCallback != id)
+            {
+                return default;
+            }
+
+            // The specified callback is actually running: queue a task that'll poll for the currently
+            // executing callback to complete. In general scheduling such a work item that polls is a really
+            // unfortunate thing to do.  However, we expect this to be a rare case (disposing while the associated
+            // callback is running), and brief when it happens (so the polling will be minimal), and making
+            // this work with a callback mechanism will add additional cost to other more common cases.
+            return new ValueTask(Task.Factory.StartNew(s =>
+            {
+                var state = (Tuple<CancellationTokenSource, long>)s;
+                state.Item1.WaitForCallbackToComplete(state.Item2);
+            }, Tuple.Create(this, id), CancellationToken.None, TaskCreationOptions.None, TaskScheduler.Default));
+        }
+
         private sealed class Linked1CancellationTokenSource : CancellationTokenSource
         {
             private readonly CancellationTokenRegistration _reg1;