Adding TimeProvider CreateCancellationTokenSource extension method

src/libraries/Common/tests/System/TimeProviderTests.cs

@@ -213,33 +213,34 @@ public static void CancellationTokenSourceWithTimer(TimeProvider provider)
             //
             // Test out some int-based timeout logic
             //
-#if NETFRAMEWORK
-            CancellationTokenSource cts = new CancellationTokenSource(Timeout.InfiniteTimeSpan); // should be an infinite timeout
+#if TESTEXTENSIONS
+            CancellationTokenSource cts = provider.CreateCancellationTokenSource(Timeout.InfiniteTimeSpan); // should be an infinite timeout
 #else
             CancellationTokenSource cts = new CancellationTokenSource(Timeout.InfiniteTimeSpan, provider); // should be an infinite timeout
-#endif // NETFRAMEWORK
-            CancellationToken token = cts.Token;
+#endif // TESTEXTENSIONS
             ManualResetEventSlim mres = new ManualResetEventSlim(false);
-            CancellationTokenRegistration ctr = token.Register(() => mres.Set());
 
-            Assert.False(token.IsCancellationRequested,
+            Assert.False(cts.Token.IsCancellationRequested,
                "CancellationTokenSourceWithTimer:  Cancellation signaled on infinite timeout (int)!");
 
-#if NETFRAMEWORK
-            CancelAfter(provider, cts, TimeSpan.FromMilliseconds(1000000));
+#if TESTEXTENSIONS
+            cts.Dispose();
+            cts = provider.CreateCancellationTokenSource(TimeSpan.FromMilliseconds(1000000));
 #else
             cts.CancelAfter(1000000);
-#endif // NETFRAMEWORK
+#endif // TESTEXTENSIONS
 
-            Assert.False(token.IsCancellationRequested,
+            Assert.False(cts.Token.IsCancellationRequested,
                "CancellationTokenSourceWithTimer:  Cancellation signaled on super-long timeout (int) !");
 
-#if NETFRAMEWORK
-            CancelAfter(provider, cts, TimeSpan.FromMilliseconds(1));
+#if TESTEXTENSIONS
+            cts.Dispose();
+            cts = provider.CreateCancellationTokenSource(TimeSpan.FromMilliseconds(1));
 #else
             cts.CancelAfter(1);
-#endif // NETFRAMEWORK
+#endif // TESTEXTENSIONS
 
+            CancellationTokenRegistration ctr = cts.Token.Register(() => mres.Set());
             Debug.WriteLine("CancellationTokenSourceWithTimer: > About to wait on cancellation that should occur soon (int)... if we hang, something bad happened");
 
             mres.Wait();
@@ -250,33 +251,34 @@ public static void CancellationTokenSourceWithTimer(TimeProvider provider)
             // Test out some TimeSpan-based timeout logic
             //
             TimeSpan prettyLong = new TimeSpan(1, 0, 0);
-#if NETFRAMEWORK
-            cts = new CancellationTokenSource(prettyLong);
+#if TESTEXTENSIONS
+            cts = provider.CreateCancellationTokenSource(prettyLong);
 #else
             cts = new CancellationTokenSource(prettyLong, provider);
-#endif // NETFRAMEWORK
+#endif // TESTEXTENSIONS
 
-            token = cts.Token;
             mres = new ManualResetEventSlim(false);
-            ctr = token.Register(() => mres.Set());
 
-            Assert.False(token.IsCancellationRequested,
+            Assert.False(cts.Token.IsCancellationRequested,
                "CancellationTokenSourceWithTimer:  Cancellation signaled on super-long timeout (TimeSpan,1)!");
 
-#if NETFRAMEWORK
-            CancelAfter(provider, cts, prettyLong);
+#if TESTEXTENSIONS
+            cts.Dispose();
+            cts = provider.CreateCancellationTokenSource(prettyLong);
 #else
             cts.CancelAfter(prettyLong);
-#endif // NETFRAMEWORK
+#endif // TESTEXTENSIONS
 
-            Assert.False(token.IsCancellationRequested,
+            Assert.False(cts.Token.IsCancellationRequested,
                "CancellationTokenSourceWithTimer:  Cancellation signaled on super-long timeout (TimeSpan,2) !");
 
-#if NETFRAMEWORK
-            CancelAfter(provider, cts, TimeSpan.FromMilliseconds(1000));
+#if TESTEXTENSIONS
+            cts.Dispose();
+            cts = provider.CreateCancellationTokenSource(TimeSpan.FromMilliseconds(1000));
 #else
             cts.CancelAfter(new TimeSpan(1000));
-#endif // NETFRAMEWORK
+#endif // TESTEXTENSIONS
+            ctr = cts.Token.Register(() => mres.Set());
 
             Debug.WriteLine("CancellationTokenSourceWithTimer: > About to wait on cancellation that should occur soon (TimeSpan)... if we hang, something bad happened");
 
@@ -555,5 +557,65 @@ public Task<TResult> WaitAsync<TResult>(Task<TResult> task, TimeSpan timeout, Ti
 
         private static TestExtensionsTaskFactory extensionsTaskFactory = new();
 #endif // TESTEXTENSIONS
+
+        // A timer that get fired on demand
+        private class ManualTimer : ITimer
+        {
+            TimerCallback _callback;
+            object? _state;
+
+            public ManualTimer(TimerCallback callback, object? state, TimeSpan dueTime, TimeSpan period)
+            {
+                _callback = callback;
+                _state = state;
+            }
+
+            public bool Change(TimeSpan dueTime, TimeSpan period) => true;
+
+            public void Fire()
+            {
+                _callback?.Invoke(_state);
+                IsFired = true;
+            }
+
+            public bool IsFired { get; set; }
+
+            public void Dispose() { }
+            public ValueTask DisposeAsync () { return default; }
+        }
+
+        private class ManualTimeProvider : TimeProvider
+        {
+            public ManualTimer Timer { get; set; }
+
+            public ManualTimeProvider() { }
+            public override ITimer CreateTimer(TimerCallback callback, object? state, TimeSpan dueTime, TimeSpan period)
+            {
+                Timer = new ManualTimer(callback, state, dueTime, period);
+                return Timer;
+            }
+        }
+
+        [Fact]
+        // 1- Creates the CTS with a delay that we control via the time provider.
+        // 2- Disposes the CTS.
+        // 3- Then fires the timer. We want to validate the process doesn't crash.
+        public static void TestCTSWithDelayFiringAfterCancellation()
+        {
+            ManualTimeProvider manualTimer = new ManualTimeProvider();
+#if TESTEXTENSIONS
+            CancellationTokenSource cts = manualTimer.CreateCancellationTokenSource(TimeSpan.FromSeconds(60));
+#else
+            CancellationTokenSource cts = new CancellationTokenSource(TimeSpan.FromSeconds(60), manualTimer);
+#endif // TESTEXTENSIONS
+
+            Assert.NotNull(manualTimer.Timer);
+            Assert.False(manualTimer.Timer.IsFired);
+
+            cts.Dispose();
+
+            manualTimer.Timer.Fire();
+            Assert.True(manualTimer.Timer.IsFired);
+        }
     }
 }

src/libraries/Microsoft.Bcl.TimeProvider/ref/Microsoft.Bcl.TimeProvider.Common.cs

@@ -8,5 +8,6 @@ public static class TimeProviderTaskExtensions
         public static System.Threading.Tasks.Task Delay(this System.TimeProvider timeProvider, System.TimeSpan delay, System.Threading.CancellationToken cancellationToken = default) { throw null; }
         public static System.Threading.Tasks.Task<TResult> WaitAsync<TResult>(this System.Threading.Tasks.Task<TResult> task, System.TimeSpan timeout, System.TimeProvider timeProvider, System.Threading.CancellationToken cancellationToken = default) { throw null; }
         public static System.Threading.Tasks.Task WaitAsync(this System.Threading.Tasks.Task task, System.TimeSpan timeout, System.TimeProvider timeProvider, System.Threading.CancellationToken cancellationToken = default) { throw null; }
+        public static System.Threading.CancellationTokenSource CreateCancellationTokenSource(this System.TimeProvider timeProvider, System.TimeSpan delay) { throw null; }
     }
 }

src/libraries/Microsoft.Bcl.TimeProvider/src/System/Threading/Tasks/TimeProviderTaskExtensions.cs

@@ -6,6 +6,10 @@ namespace System.Threading.Tasks
     /// <summary>
     /// Provide extensions methods for <see cref="Task"/> operations with <see cref="TimeProvider"/>.
     /// </summary>
+    /// <remarks>
+    /// The Microsoft.Bcl.TimeProvider library interfaces are intended solely for use in building against pre-.NET 8 surface area.
+    /// If your code is being built against .NET 8 or higher, then this library should not be utilized.
+    /// </remarks>
     public static class TimeProviderTaskExtensions
     {
 #if !NET8_0_OR_GREATER
@@ -218,5 +222,57 @@ public static async Task<TResult> WaitAsync<TResult>(this Task<TResult> task, Ti
             return task.Result;
         }
 #endif // NET8_0_OR_GREATER
+
+        /// <summary>Initializes a new instance of the <see cref="CancellationTokenSource"/> class that will be canceled after the specified <see cref="TimeSpan"/>. </summary>
+        /// <param name="timeProvider">The <see cref="TimeProvider"/> with which to interpret the <paramref name="delay"/>. </param>
+        /// <param name="delay">The time interval to wait before canceling this <see cref="CancellationTokenSource"/>. </param>
+        /// <exception cref="ArgumentOutOfRangeException"> The <paramref name="delay"/> is negative and not equal to <see cref="Timeout.InfiniteTimeSpan" /> or greater than maximum allowed timer duration.</exception>
+        /// <returns><see cref="CancellationTokenSource"/> that will be canceled after the specified <paramref name="delay"/>.</returns>
+        /// <remarks>
+        /// <para>
+        /// The countdown for the delay starts during the call to the constructor. When the delay expires,
+        /// the constructed <see cref="CancellationTokenSource"/> is canceled if it has
+        /// not been canceled already.
+        /// </para>
+        /// <para>
+        /// If running on .NET versions earlier than .NET 8.0, there is a constraint when invoking <see cref="CancellationTokenSource.CancelAfter(TimeSpan)"/> on the resultant object.
+        /// This action will not terminate the initial timer indicated by <paramref name="delay"/>. However, this restriction does not apply on .NET 8.0 and later versions.
+        /// </para>
+        /// </remarks>
+        public static CancellationTokenSource CreateCancellationTokenSource(this TimeProvider timeProvider, TimeSpan delay)
+        {
+#if NET8_0_OR_GREATER
+            return new CancellationTokenSource(delay, timeProvider);
+#else
+            if (timeProvider is null)
+            {
+                throw new ArgumentNullException(nameof(timeProvider));
+            }
+
+            if (delay != Timeout.InfiniteTimeSpan && delay < TimeSpan.Zero)
+            {
+                throw new ArgumentOutOfRangeException(nameof(delay));
+            }
+
+            if (timeProvider == TimeProvider.System)
+            {
+                return new CancellationTokenSource(delay);
+            }
+
+            var cts = new CancellationTokenSource();
+
+            ITimer timer = timeProvider.CreateTimer(s =>
+            {
+                try
+                {
+                    ((CancellationTokenSource)s).Cancel();
+                }
+                catch (ObjectDisposedException) { }
+            }, cts, delay, Timeout.InfiniteTimeSpan);
+
+            cts.Token.Register(t => ((ITimer)t).Dispose(), timer);
+            return cts;
+#endif // NET8_0_OR_GREATER
+        }
     }
 }