Add CancellationTokenSource.TryReset (#50346)

* Add CancellationTokenSource.TryReset

* Update src/libraries/System.Private.CoreLib/src/System/Threading/CancellationTokenSource.cs

src/libraries/System.Private.CoreLib/src/System/Threading/CancellationTokenSource.cs

@@ -66,7 +66,7 @@ private static void TimerCallback(object? state) => // separated out into a name
         /// canceled concurrently.
         /// </para>
         /// </remarks>
-        public bool IsCancellationRequested => _state >= NotifyingState;
+        public bool IsCancellationRequested => _state != NotCanceledState;
 
         /// <summary>A simple helper to determine whether cancellation has finished.</summary>
         internal bool IsCancellationCompleted => _state == NotifyingCompleteState;
@@ -365,6 +365,54 @@ private void CancelAfter(uint millisecondsDelay)
             }
         }
 
+        /// <summary>
+        /// Attempts to reset the <see cref="CancellationTokenSource"/> to be used for an unrelated operation.
+        /// </summary>
+        /// <returns>
+        /// true if the <see cref="CancellationTokenSource"/> has not had cancellation requested and could
+        /// have its state reset to be reused for a subsequent operation; otherwise, false.
+        /// </returns>
+        /// <remarks>
+        /// <see cref="TryReset"/> is intended to be used by the sole owner of the <see cref="CancellationTokenSource"/>
+        /// when it is known that the operation with which the <see cref="CancellationTokenSource"/> was used has
+        /// completed, no one else will be attempting to cancel it, and any registrations still remaining are erroneous.
+        /// Upon a successful reset, such registrations will no longer be notified for any subsequent cancellation of the
+        /// <see cref="CancellationTokenSource"/>; however, if any component still holds a reference to this
+        /// <see cref="CancellationTokenSource"/> either directly or indirectly via a <see cref="CancellationToken"/>
+        /// handed out from it, polling via their reference will show the current state any time after the reset as
+        /// it's the same instance.  Usage of <see cref="TryReset"/> concurrently with requesting cancellation is not
+        /// thread-safe and may result in TryReset returning true even if cancellation was already requested and may result
+        /// in registrations not being invoked as part of the concurrent cancellation request.
+        /// </remarks>
+        public bool TryReset()
+        {
+            ThrowIfDisposed();
+
+            // We can only reset if cancellation has not yet been requested: we never want to allow a CancellationToken
+            // to transition from canceled to non-canceled.
+            if (_state == NotCanceledState)
+            {
+                // If there is no timer, then we're free to reset.  If there is a timer, then we need to first try
+                // to reset it to be infinite so that it won't fire, and then recognize that it could have already
+                // fired by the time we successfully changed it, and so check to see whether that's possibly the case.
+                // If we successfully reset it and it never fired, then we can be sure it won't trigger cancellation.
+                bool reset =
+                    _timer is not TimerQueueTimer timer ||
+                    (timer.Change(Timeout.UnsignedInfinite, Timeout.UnsignedInfinite) && !timer._everQueued);
+
+                if (reset)
+                {
+                    // We're not canceled and no timer will run to cancel us.
+                    // Clear out all the registrations, and return that we've successfully reset.
+                    Volatile.Read(ref _registrations)?.UnregisterAll();
+                    return true;
+                }
+            }
+
+            // Failed to reset.
+            return false;
+        }
+
         /// <summary>Releases the resources used by this <see cref="CancellationTokenSource" />.</summary>
         /// <remarks>This method is not thread-safe for any other concurrent calls.</remarks>
         public void Dispose()
@@ -434,10 +482,7 @@ private void ThrowIfDisposed()
         {
             if (_disposed)
             {
-                ThrowObjectDisposedException();
-
-                [DoesNotReturn]
-                static void ThrowObjectDisposedException() => throw new ObjectDisposedException(null, SR.CancellationTokenSource_Disposed);
+                ThrowHelper.ThrowObjectDisposedException(ExceptionResource.CancellationTokenSource_Disposed);
             }
         }
 
@@ -876,6 +921,25 @@ internal sealed class Registrations
             /// <param name="source">The associated source.</param>
             public Registrations(CancellationTokenSource source) => Source = source;
 
+            [MethodImpl(MethodImplOptions.AggressiveInlining)] // used in only two places, one of which is a hot path
+            private void Recycle(CallbackNode node)
+            {
+                Debug.Assert(_lock == 1);
+
+                // Clear out the unused node and put it on the singly-linked free list.
+                // The only field we don't clear out is the associated Registrations, as that's fixed
+                // throughout the node's lifetime.
+                node.Id = 0;
+                node.Callback = null;
+                node.CallbackState = null;
+                node.ExecutionContext = null;
+                node.SynchronizationContext = null;
+
+                node.Prev = null;
+                node.Next = FreeNodeList;
+                FreeNodeList = node;
+            }
+
             /// <summary>Unregisters a callback.</summary>
             /// <param name="id">The expected id of the registration.</param>
             /// <param name="node">The callback node.</param>
@@ -925,17 +989,7 @@ public bool Unregister(long id, CallbackNode node)
                         node.Next.Prev = node.Prev;
                     }
 
-                    // Clear out the now unused node and put it on the singly-linked free list.
-                    // The only field we don't clear out is the associated Source, as that's fixed
-                    // throughout the nodes lifetime.
-                    node.Id = 0;
-                    node.Callback = null;
-                    node.CallbackState = null;
-                    node.ExecutionContext = null;
-                    node.SynchronizationContext = null;
-                    node.Prev = null;
-                    node.Next = FreeNodeList;
-                    FreeNodeList = node;
+                    Recycle(node);
 
                     return true;
                 }
@@ -945,6 +999,30 @@ public bool Unregister(long id, CallbackNode node)
                 }
             }
 
+            /// <summary>Moves all registrations to the free list.</summary>
+            public void UnregisterAll()
+            {
+                EnterLock();
+                try
+                {
+                    // Null out all callbacks.
+                    CallbackNode? node = Callbacks;
+                    Callbacks = null;
+
+                    // Reset and move each node that was in the callbacks list to the free list.
+                    while (node != null)
+                    {
+                        CallbackNode? next = node.Next;
+                        Recycle(node);
+                        node = next;
+                    }
+                }
+                finally
+                {
+                    ExitLock();
+                }
+            }
+
             /// <summary>
             /// 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.

src/libraries/System.Private.CoreLib/src/System/Threading/Timer.cs

@@ -206,6 +206,7 @@ private void FireNextTimers()
                         if (remaining <= 0)
                         {
                             // Timer is ready to fire.
+                            timer._everQueued = true;
 
                             if (timer._period != Timeout.UnsignedInfinite)
                             {
@@ -476,6 +477,7 @@ internal sealed partial class TimerQueueTimer : IThreadPoolWorkItem
         // instead of with a provided WaitHandle.
         private int _callbacksRunning;
         private bool _canceled;
+        internal bool _everQueued;
         private object? _notifyWhenNoCallbacksRunning; // may be either WaitHandle or Task<bool>
 
         internal TimerQueueTimer(TimerCallback timerCallback, object? state, uint dueTime, uint period, bool flowExecutionContext)

src/libraries/System.Private.CoreLib/src/System/ThrowHelper.cs

@@ -911,6 +911,8 @@ private static string GetResourceString(ExceptionResource resource)
                     return SR.Argument_SpansMustHaveSameLength;
                 case ExceptionResource.Argument_InvalidFlag:
                     return SR.Argument_InvalidFlag;
+                case ExceptionResource.CancellationTokenSource_Disposed:
+                    return SR.CancellationTokenSource_Disposed;
                 default:
                     Debug.Fail("The enum value is not defined, please check the ExceptionResource Enum.");
                     return "";
@@ -1090,5 +1092,6 @@ internal enum ExceptionResource
         Arg_TypeNotSupported,
         Argument_SpansMustHaveSameLength,
         Argument_InvalidFlag,
+        CancellationTokenSource_Disposed,
     }
 }

src/libraries/System.Runtime/ref/System.Runtime.cs

@@ -11020,6 +11020,7 @@ public void CancelAfter(System.TimeSpan delay) { }
         public static System.Threading.CancellationTokenSource CreateLinkedTokenSource(params System.Threading.CancellationToken[] tokens) { throw null; }
         public void Dispose() { }
         protected virtual void Dispose(bool disposing) { }
+        public bool TryReset() { throw null; }
     }
     public enum LazyThreadSafetyMode
     {

src/libraries/System.Threading.Tasks/tests/CancellationTokenTests.cs

@@ -1049,6 +1049,7 @@ public static void CancellationTokenSourceWithTimer()
 
             cts.Dispose();
         }
+
         [Fact]
         public static void CancellationTokenSourceWithTimer_Negative()
         {
@@ -1076,6 +1077,54 @@ public static void CancellationTokenSourceWithTimer_Negative()
             Assert.Throws<ObjectDisposedException>(() => { cts.CancelAfter(reasonableTimeSpan); });
         }
 
+        [Fact]
+        public static void CancellationTokenSource_TryReset_ReturnsFalseIfAlreadyCanceled()
+        {
+            var cts = new CancellationTokenSource();
+            cts.Cancel();
+            Assert.False(cts.TryReset());
+            Assert.True(cts.IsCancellationRequested);
+        }
+
+        [Fact]
+        public static void CancellationTokenSource_TryReset_ReturnsTrueIfNotCanceledAndNoTimer()
+        {
+            var cts = new CancellationTokenSource();
+            Assert.True(cts.TryReset());
+            Assert.True(cts.TryReset());
+        }
+
+        [Fact]
+        public static void CancellationTokenSource_TryReset_ReturnsTrueIfNotCanceledAndTimerHasntFired()
+        {
+            var cts = new CancellationTokenSource();
+            cts.CancelAfter(TimeSpan.FromDays(1));
+            Assert.True(cts.TryReset());
+        }
+
+        [Fact]
+        public static void CancellationTokenSource_TryReset_UnregistersAll()
+        {
+            bool registration1Invoked = false;
+            bool registration2Invoked = false;
+
+            var cts = new CancellationTokenSource();
+            CancellationTokenRegistration ctr1 = cts.Token.Register(() => registration1Invoked = true);
+            Assert.True(cts.TryReset());
+            CancellationTokenRegistration ctr2 = cts.Token.Register(() => registration2Invoked = true);
+
+            cts.Cancel();
+
+            Assert.False(registration1Invoked);
+            Assert.True(registration2Invoked);
+
+            Assert.False(ctr1.Unregister());
+            Assert.False(ctr2.Unregister());
+
+            Assert.Equal(cts.Token, ctr1.Token);
+            Assert.Equal(cts.Token, ctr2.Token);
+        }
+
         [ConditionalFact(typeof(PlatformDetection), nameof(PlatformDetection.IsThreadingSupported))]
         public static void EnlistWithSyncContext_BeforeCancel()
         {