Check CancellationToken before TimeSpan in TimeProviderTaskExtensions (…

…dotnet#107009)

* Check CancellationToken before TimeSpan.Zero in TimeProviderTaskExtensions

This makes the down-level behavior for Delay and WaitAsync match the
behavior in .NET 8 and later, prefering to return a Canceled task over
a Completed task or TimeoutException when the TimeSpan is Zero and the
CancellationToken is already Canceled.

Fixes dotnet#106996

* Test improvements based on feedback

---------

Co-authored-by: Tarek Mahmoud Sayed <10833894+tarekgh@users.noreply.github.com>

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

@@ -310,18 +310,25 @@ public static void RunDelayTests(TimeProvider provider, ITestTaskFactory taskFac
             }
             catch (Exception e)
             {
-                Assert.Fail(string.Format("RunDelayTests:    > FAILED.  Unexpected exception on WaitAll(simple tasks): {0}", e));
+                Assert.Fail($"RunDelayTests:    > FAILED.  Unexpected exception on WaitAll(simple tasks): {e}");
             }
 
             Assert.True(task1.Status == TaskStatus.RanToCompletion, "    > FAILED.  Expected Delay(TimeSpan(0), timeProvider) to run to completion");
             Assert.True(task2.Status == TaskStatus.RanToCompletion, "    > FAILED.  Expected Delay(TimeSpan(0), timeProvider, uncanceledToken) to run to completion");
 
+            // This should complete quickly with a CANCELED status.
+            Task task3 = taskFactory.Delay(provider, new TimeSpan(0), new CancellationToken(true));
+            var canceledException = Record.Exception(task3.Wait);
+            Assert.True(canceledException is AggregateException { InnerException: OperationCanceledException },
+                $"RunDelayTests:    > FAILED.  Unexpected exception on canceled task Wait(): {canceledException}");
+            Assert.True(task3.Status == TaskStatus.Canceled, "    > FAILED.  Expected Delay(timeProvider, TimeSpan(0), canceledToken) to be canceled");
+
             // This should take some time
-            Task task3 = taskFactory.Delay(provider, TimeSpan.FromMilliseconds(20000));
+            Task task4 = taskFactory.Delay(provider, TimeSpan.FromMilliseconds(20000));
 
-            Assert.False(task3.IsCompleted, "RunDelayTests:    > FAILED.  Delay(20000) appears to have completed too soon(1).");
+            Assert.False(task4.IsCompleted, "RunDelayTests:    > FAILED.  Delay(20000) appears to have completed too soon(1).");
             Task t2 = Task.Delay(TimeSpan.FromMilliseconds(10));
-            Assert.False(task3.IsCompleted, "RunDelayTests:    > FAILED.  Delay(10000) appears to have completed too soon(2).");
+            Assert.False(task4.IsCompleted, "RunDelayTests:    > FAILED.  Delay(10000) appears to have completed too soon(2).");
         }
 
         [ConditionalTheory(typeof(PlatformDetection), nameof(PlatformDetection.IsThreadingSupported))]
@@ -359,6 +366,7 @@ public static async Task RunWaitAsyncTests(TimeProvider provider, ITestTaskFacto
 
             cts.Cancel();
             await Assert.ThrowsAsync<TaskCanceledException>(() => taskFactory.WaitAsync(tcs5.Task, TimeSpan.FromMilliseconds(10), provider, cts.Token));
+            await Assert.ThrowsAsync<TaskCanceledException>(() => taskFactory.WaitAsync(tcs5.Task, TimeSpan.Zero, provider, cts.Token));
 
             using CancellationTokenSource cts1 = new CancellationTokenSource();
             Task task5 = Task.Run(() => { while (!cts1.Token.IsCancellationRequested) { Thread.Sleep(10); } });

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

@@ -66,14 +66,14 @@ public static Task Delay(this TimeProvider timeProvider, TimeSpan delay, Cancell
                 throw new ArgumentOutOfRangeException(nameof(delay));
             }
 
-            if (delay == TimeSpan.Zero)
+            if (cancellationToken.IsCancellationRequested)
             {
-                return Task.CompletedTask;
+                return Task.FromCanceled(cancellationToken);
             }
 
-            if (cancellationToken.IsCancellationRequested)
+            if (delay == TimeSpan.Zero)
             {
-                return Task.FromCanceled(cancellationToken);
+                return Task.CompletedTask;
             }
 
             DelayState state = new(cancellationToken);
@@ -160,14 +160,14 @@ public static Task WaitAsync(this Task task, TimeSpan timeout, TimeProvider time
                 return task;
             }
 
-            if (timeout == TimeSpan.Zero)
+            if (cancellationToken.IsCancellationRequested)
             {
-                return Task.FromException(new TimeoutException());
+                return Task.FromCanceled(cancellationToken);
             }
 
-            if (cancellationToken.IsCancellationRequested)
+            if (timeout == TimeSpan.Zero)
             {
-                return Task.FromCanceled(cancellationToken);
+                return Task.FromException(new TimeoutException());
             }
 
             WaitAsyncState state = new(cancellationToken);