Check CancellationToken before TimeSpan.Zero in TimeProviderTaskExten…

…sions

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 #106996

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

@@ -316,12 +316,26 @@ public static void RunDelayTests(TimeProvider provider, ITestTaskFactory taskFac
             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));
+            try
+            {
+                var e = Assert.Throws<AggregateException>(task3.Wait);
+                Assert.IsAssignableFrom<OperationCanceledException>(e.InnerException);
+            }
+            catch (Exception e)
+            {
+                Assert.Fail(string.Format("RunDelayTests:    > FAILED.  Unexpected exception on canceled task Wait(): {0}", e));
+            }
+
+            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 +373,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);
@@ -161,14 +161,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);