diff --git a/src/libraries/Common/tests/System/TimeProviderTests.cs b/src/libraries/Common/tests/System/TimeProviderTests.cs index a4209e83f5236c..45148afbcbb4fc 100644 --- a/src/libraries/Common/tests/System/TimeProviderTests.cs +++ b/src/libraries/Common/tests/System/TimeProviderTests.cs @@ -638,6 +638,69 @@ public override ITimer CreateTimer(TimerCallback callback, object? state, TimeSp } } + public static IEnumerable TaskFactoryData() + { + yield return new object[] { taskFactory }; + +#if TESTEXTENSIONS + yield return new object[] { extensionsTaskFactory }; +#endif // TESTEXTENSIONS + } + + [ConditionalTheory(typeof(PlatformDetection), nameof(PlatformDetection.IsThreadingSupported))] + [MemberData(nameof(TaskFactoryData))] + public async Task TestDelayTaskContinuation(ITestTaskFactory taskFactory) + { + // + // Test time expiration and validate continuation is called synchronously. + // + + ManualTimeProvider manualTimeProvider = new ManualTimeProvider(); + int callbackCount = 0; + + _ = Continuation(manualTimeProvider, default, () => callbackCount++); + + Assert.NotNull(manualTimeProvider.Timer); + manualTimeProvider.Timer.Fire(); + + // Delay should be completed and the continuation should be called synchronously. + Assert.Equal(1, callbackCount); + + // + // Test cancellation and validate continuation is called asynchronously. + // + + manualTimeProvider = new ManualTimeProvider(); + CancellationTokenSource cts = new CancellationTokenSource(); + + var tl = new ThreadLocal(); + tl.Value = 10; + int t1Value = 0; + + Task task = Continuation(manualTimeProvider, cts.Token, () => t1Value = tl.Value); + cts.Cancel(); + + // reset the thread local value as the continuation callback could end up running on this thread pool thread. + tl.Value = 0; + await task; + + Assert.NotEqual(10, t1Value); + + async Task Continuation(TimeProvider timeProvider, CancellationToken token, Action callback) + { + try + { + await taskFactory.Delay(timeProvider, TimeSpan.FromSeconds(10), token); + } + catch (OperationCanceledException) + { + // Ignore + } + + callback(); + } + } + [Fact] // 1- Creates the CTS with a delay that we control via the time provider. // 2- Disposes the CTS. diff --git a/src/libraries/Microsoft.Bcl.TimeProvider/src/System/Threading/Tasks/TimeProviderTaskExtensions.cs b/src/libraries/Microsoft.Bcl.TimeProvider/src/System/Threading/Tasks/TimeProviderTaskExtensions.cs index e410f4a5c9d32b..1497b84c4aa3a6 100644 --- a/src/libraries/Microsoft.Bcl.TimeProvider/src/System/Threading/Tasks/TimeProviderTaskExtensions.cs +++ b/src/libraries/Microsoft.Bcl.TimeProvider/src/System/Threading/Tasks/TimeProviderTaskExtensions.cs @@ -15,7 +15,7 @@ public static class TimeProviderTaskExtensions #if !NET8_0_OR_GREATER private sealed class DelayState : TaskCompletionSource { - public DelayState(CancellationToken cancellationToken) : base(TaskCreationOptions.RunContinuationsAsynchronously) + public DelayState(CancellationToken cancellationToken) { CancellationToken = cancellationToken; } @@ -89,7 +89,15 @@ public static Task Delay(this TimeProvider timeProvider, TimeSpan delay, Cancell state.Registration = cancellationToken.Register(static delayState => { DelayState s = (DelayState)delayState!; - s.TrySetCanceled(s.CancellationToken); + + // When cancellation is requested, we need to force the task continuation to run asynchronously + // to avoid doing arbitrary amounts of work as part of a call to CancellationTokenSource.Cancel. + ThreadPool.UnsafeQueueUserWorkItem(static state => + { + DelayState theState = (DelayState)state; + theState.TrySetCanceled(theState.CancellationToken); + }, s); + s.Registration.Dispose(); s.Timer?.Dispose(); }, state);