dotnet · stephentoub · Mar 11, 2021 · Mar 9, 2021
diff --git a/src/libraries/Common/src/System/Threading/Tasks/TaskTimeoutExtensions.cs b/src/libraries/Common/src/System/Threading/Tasks/TaskTimeoutExtensions.cs
diff --git a/src/libraries/Common/tests/StreamConformanceTests/System/IO/StreamConformanceTests.cs b/src/libraries/Common/tests/StreamConformanceTests/System/IO/StreamConformanceTests.cs
@@ -558,9 +558,12 @@ protected async Task WhenAllOrAnyFailed(Task task1, Task task2)
             }
             else
             {
-                var cts = new CancellationTokenSource();
-                await Task.WhenAny(incomplete, Task.Delay(500, cts.Token)); // give second task a chance to complete
-                cts.Cancel();
+                try
+                {
+                    await incomplete.WaitAsync(TimeSpan.FromMilliseconds(500)); // give second task a chance to complete
+                }
+                catch (TimeoutException) { }
+
                 await (incomplete.IsCompleted ? Task.WhenAll(completed, incomplete) : completed);
             }
         }

diff --git a/src/libraries/Common/tests/System/Net/Http/GenericLoopbackServer.cs b/src/libraries/Common/tests/System/Net/Http/GenericLoopbackServer.cs
@@ -34,7 +34,7 @@ public Task CreateClientAndServerAsync(Func<Uri, Task> clientFunc, Func<GenericL
                 Task serverTask = serverFunc(server);
 
                 await new Task[] { clientTask, serverTask }.WhenAllOrAnyFailed().ConfigureAwait(false);
-            }, options: options).TimeoutAfter(millisecondsTimeout);
+            }, options: options).WaitAsync(TimeSpan.FromMilliseconds(millisecondsTimeout));
         }
     }
 

diff --git a/src/libraries/Common/tests/System/Net/Http/Http2LoopbackConnection.cs b/src/libraries/Common/tests/System/Net/Http/Http2LoopbackConnection.cs
@@ -241,8 +241,7 @@ public async Task ExpectSettingsAckAsync(int timeoutMs = 5000)
             Task currentTask = _ignoredSettingsAckPromise?.Task;
             if (currentTask != null)
             {
-                var timeout = TimeSpan.FromMilliseconds(timeoutMs);
-                await currentTask.TimeoutAfter(timeout);
+                await currentTask.WaitAsync(TimeSpan.FromMilliseconds(timeoutMs));
             }
 
             _ignoredSettingsAckPromise = new TaskCompletionSource<bool>();
@@ -909,7 +908,7 @@ public override async Task WaitForCancellationAsync(bool ignoreIncomingData = tr
             Frame frame;
             do
             {
-                frame = await ReadFrameAsync(TimeSpan.FromMilliseconds(TestHelper.PassingTestTimeoutMilliseconds));
+                frame = await ReadFrameAsync(TestHelper.PassingTestTimeout);
                 Assert.NotNull(frame); // We should get Rst before closing connection.
                 Assert.Equal(0, (int)(frame.Flags & FrameFlags.EndStream));
                 if (ignoreIncomingData)

diff --git a/src/libraries/Common/tests/System/Net/Http/Http2LoopbackServer.cs b/src/libraries/Common/tests/System/Net/Http/Http2LoopbackServer.cs
@@ -192,7 +192,7 @@ public static async Task CreateServerAsync(Func<Http2LoopbackServer, Uri, Task>
         {
             using (var server = Http2LoopbackServer.CreateServer())
             {
-                await funcAsync(server, server.Address).TimeoutAfter(millisecondsTimeout).ConfigureAwait(false);
+                await funcAsync(server, server.Address).WaitAsync(TimeSpan.FromMilliseconds(millisecondsTimeout));
             }
         }
 
@@ -223,7 +223,7 @@ public override async Task CreateServerAsync(Func<GenericLoopbackServer, Uri, Ta
         {
             using (var server = CreateServer(options))
             {
-                await funcAsync(server, server.Address).TimeoutAfter(millisecondsTimeout).ConfigureAwait(false);
+                await funcAsync(server, server.Address).WaitAsync(TimeSpan.FromMilliseconds(millisecondsTimeout));
             }
         }
 

diff --git a/src/libraries/Common/tests/System/Net/Http/Http3LoopbackServer.cs b/src/libraries/Common/tests/System/Net/Http/Http3LoopbackServer.cs
@@ -89,7 +89,7 @@ public override GenericLoopbackServer CreateServer(GenericLoopbackOptions option
         public override async Task CreateServerAsync(Func<GenericLoopbackServer, Uri, Task> funcAsync, int millisecondsTimeout = 60000, GenericLoopbackOptions options = null)
         {
             using GenericLoopbackServer server = CreateServer(options);
-            await funcAsync(server, server.Address).TimeoutAfter(millisecondsTimeout).ConfigureAwait(false);
+            await funcAsync(server, server.Address).WaitAsync(TimeSpan.FromMilliseconds(millisecondsTimeout));
         }
 
         public override Task<GenericLoopbackConnection> CreateConnectionAsync(Socket socket, Stream stream, GenericLoopbackOptions options = null)

diff --git a/src/libraries/Common/tests/System/Net/Http/HttpAgnosticLoopbackServer.cs b/src/libraries/Common/tests/System/Net/Http/HttpAgnosticLoopbackServer.cs
@@ -208,7 +208,7 @@ public static async Task CreateServerAsync(Func<HttpAgnosticLoopbackServer, Uri,
         {
             using (var server = HttpAgnosticLoopbackServer.CreateServer())
             {
-                await funcAsync(server, server.Address).TimeoutAfter(millisecondsTimeout).ConfigureAwait(false);
+                await funcAsync(server, server.Address).WaitAsync(TimeSpan.FromMilliseconds(millisecondsTimeout));
             }
         }
 
@@ -240,7 +240,7 @@ public override async Task CreateServerAsync(Func<GenericLoopbackServer, Uri, Ta
         {
             using (var server = CreateServer(options))
             {
-                await funcAsync(server, server.Address).TimeoutAfter(millisecondsTimeout).ConfigureAwait(false);
+                await funcAsync(server, server.Address).WaitAsync(TimeSpan.FromMilliseconds(millisecondsTimeout));
             }
         }
 

diff --git a/src/libraries/Common/tests/System/Net/Http/HttpClientHandlerTest.Proxy.cs b/src/libraries/Common/tests/System/Net/Http/HttpClientHandlerTest.Proxy.cs
@@ -361,7 +361,7 @@ await LoopbackServer.CreateServerAsync(async (proxyServer, proxyUrl) =>
 
                     //  Send Digest challenge.
                     Task<List<string>> serverTask = proxyServer.AcceptConnectionSendResponseAndCloseAsync(HttpStatusCode.ProxyAuthenticationRequired, authHeader);
-                    if (clientTask == await Task.WhenAny(clientTask, serverTask).TimeoutAfter(TestHelper.PassingTestTimeoutMilliseconds))
+                    if (clientTask == await Task.WhenAny(clientTask, serverTask).WaitAsync(TestHelper.PassingTestTimeout))
                     {
                         // Client task shouldn't have completed successfully; propagate failure.
                         Assert.NotEqual(TaskStatus.RanToCompletion, clientTask.Status);

diff --git a/src/libraries/Common/tests/System/Net/Http/HttpClientHandlerTest.cs b/src/libraries/Common/tests/System/Net/Http/HttpClientHandlerTest.cs
@@ -1513,7 +1513,7 @@ await server.AcceptConnectionAsync(async connection =>
                         await connection.ReadRequestDataAsync(readBody: true);
                     }
                     catch { } // Eat errors from client disconnect.
-                    await clientFinished.Task.TimeoutAfter(TimeSpan.FromMinutes(2));
+                    await clientFinished.Task.WaitAsync(TimeSpan.FromMinutes(2));
                 });
             });
         }

diff --git a/src/libraries/Common/tests/System/Net/Http/TestHelper.cs b/src/libraries/Common/tests/System/Net/Http/TestHelper.cs
@@ -17,7 +17,9 @@ namespace System.Net.Http.Functional.Tests
 {
     public static class TestHelper
     {
+        public static TimeSpan PassingTestTimeout => TimeSpan.FromMilliseconds(PassingTestTimeoutMilliseconds);
         public static int PassingTestTimeoutMilliseconds => 60 * 1000;
+
         public static bool JsonMessageContainsKeyValue(string message, string key, string value)
         {
             // Deal with JSON encoding of '\' and '"' in value

diff --git a/src/libraries/Common/tests/System/Threading/Tasks/TaskTimeoutExtensions.cs b/src/libraries/Common/tests/System/Threading/Tasks/TaskTimeoutExtensions.cs
@@ -3,91 +3,50 @@
 
 using System.Collections.Generic;
 using System.Diagnostics;
+using System.Runtime.CompilerServices;
 
-/// <summary>
-/// Task timeout helper based on https://devblogs.microsoft.com/pfxteam/crafting-a-task-timeoutafter-method/
-/// </summary>
 namespace System.Threading.Tasks
 {
     public static class TaskTimeoutExtensions
     {
-        public static async Task WithCancellation(this Task task, CancellationToken cancellationToken)
-        {
-            var tcs = new TaskCompletionSource<bool>();
-            using (cancellationToken.Register(s => ((TaskCompletionSource<bool>)s).TrySetResult(true), tcs))
-            {
-                if (task != await Task.WhenAny(task, tcs.Task).ConfigureAwait(false))
-                {
-                    throw new OperationCanceledException(cancellationToken);
-                }
-                await task; // already completed; propagate any exception
-            }
-        }
-
-        public static Task TimeoutAfter(this Task task, int millisecondsTimeout)
-            => task.TimeoutAfter(TimeSpan.FromMilliseconds(millisecondsTimeout));
-
-        public static async Task TimeoutAfter(this Task task, TimeSpan timeout)
-        {
-            var cts = new CancellationTokenSource();
+        #region WaitAsync polyfills
+        // Test polyfills when targeting a platform that doesn't have these ConfigureAwait overloads on Task
 
-            if (task == await Task.WhenAny(task, Task.Delay(timeout, cts.Token)).ConfigureAwait(false))
-            {
-                cts.Cancel();
-                await task.ConfigureAwait(false);
-            }
-            else
-            {
-                throw new TimeoutException($"Task timed out after {timeout}");
-            }
-        }
+        public static Task WaitAsync(this Task task, TimeSpan timeout) =>
+            WaitAsync(task, timeout, default);
 
-        public static Task<TResult> TimeoutAfter<TResult>(this Task<TResult> task, int millisecondsTimeout)
-            => task.TimeoutAfter(TimeSpan.FromMilliseconds(millisecondsTimeout));
+        public static Task WaitAsync(this Task task, CancellationToken cancellationToken) =>
+            WaitAsync(task, Timeout.InfiniteTimeSpan, cancellationToken);
 
-        public static async Task<TResult> TimeoutAfter<TResult>(this Task<TResult> task, TimeSpan timeout)
+        public async static Task WaitAsync(this Task task, TimeSpan timeout, CancellationToken cancellationToken)
         {
-            var cts = new CancellationTokenSource();
-
-            if (task == await Task<TResult>.WhenAny(task, Task<TResult>.Delay(timeout, cts.Token)).ConfigureAwait(false))
-            {
-                cts.Cancel();
-                return await task.ConfigureAwait(false);
-            }
-            else
+            var tcs = new TaskCompletionSource<bool>();
+            using (new Timer(s => ((TaskCompletionSource<bool>)s).TrySetException(new TimeoutException()), tcs, timeout, Timeout.InfiniteTimeSpan))
+            using (cancellationToken.Register(s => ((TaskCompletionSource<bool>)s).TrySetCanceled(), tcs))
             {
-                throw new TimeoutException($"Task timed out after {timeout}");
+                await(await Task.WhenAny(task, tcs.Task).ConfigureAwait(false)).ConfigureAwait(false);
             }
         }
 
-#if !NETFRAMEWORK
-        public static Task TimeoutAfter(this ValueTask task, int millisecondsTimeout)
-            => task.AsTask().TimeoutAfter(TimeSpan.FromMilliseconds(millisecondsTimeout));
-
-        public static Task TimeoutAfter(this ValueTask task, TimeSpan timeout)
-            => task.AsTask().TimeoutAfter(timeout);
-
-        public static Task<TResult> TimeoutAfter<TResult>(this ValueTask<TResult> task, int millisecondsTimeout)
-            => task.AsTask().TimeoutAfter(TimeSpan.FromMilliseconds(millisecondsTimeout));
+        public static Task<TResult> WaitAsync<TResult>(this Task<TResult> task, TimeSpan timeout) =>
+            WaitAsync(task, timeout, default);
 
-        public static Task<TResult> TimeoutAfter<TResult>(this ValueTask<TResult> task, TimeSpan timeout)
-            => task.AsTask().TimeoutAfter(timeout);
-#endif
+        public static Task<TResult> WaitAsync<TResult>(this Task<TResult> task, CancellationToken cancellationToken) =>
+            WaitAsync(task, Timeout.InfiniteTimeSpan, cancellationToken);
 
-        public static async Task WhenAllOrAnyFailed(this Task[] tasks, int millisecondsTimeout)
+        public static async Task<TResult> WaitAsync<TResult>(this Task<TResult> task, TimeSpan timeout, CancellationToken cancellationToken)
         {
-            var cts = new CancellationTokenSource();
-            Task task = tasks.WhenAllOrAnyFailed();
-            if (task == await Task.WhenAny(task, Task.Delay(millisecondsTimeout, cts.Token)).ConfigureAwait(false))
+            var tcs = new TaskCompletionSource<TResult>();
+            using (new Timer(s => ((TaskCompletionSource<TResult>)s).TrySetException(new TimeoutException()), tcs, timeout, Timeout.InfiniteTimeSpan))
+            using (cancellationToken.Register(s => ((TaskCompletionSource<TResult>)s).TrySetCanceled(), tcs))
             {
-                cts.Cancel();
-                await task.ConfigureAwait(false);
-            }
-            else
-            {
-                throw new TimeoutException($"{nameof(WhenAllOrAnyFailed)} timed out after {millisecondsTimeout}ms");
+                return await (await Task.WhenAny(task, tcs.Task).ConfigureAwait(false)).ConfigureAwait(false);
             }
         }
+        #endregion
+
+        public static async Task WhenAllOrAnyFailed(this Task[] tasks, int millisecondsTimeout) =>
+            await tasks.WhenAllOrAnyFailed().WaitAsync(TimeSpan.FromMilliseconds(millisecondsTimeout));
 
         public static async Task WhenAllOrAnyFailed(this Task[] tasks)
         {
@@ -99,12 +58,11 @@ public static async Task WhenAllOrAnyFailed(this Task[] tasks)
             {
                 // Wait a bit to allow other tasks to complete so we can include their exceptions
                 // in the error we throw.
-                using (var cts = new CancellationTokenSource())
+                try
                 {
-                    await Task.WhenAny(
-                        Task.WhenAll(tasks),
-                        Task.Delay(3_000, cts.Token)).ConfigureAwait(false); // arbitrary delay; can be dialed up or down in the future
+                    await Task.WhenAll(tasks).WaitAsync(TimeSpan.FromSeconds(3)); // arbitrary delay; can be dialed up or down in the future
                 }
+                catch { }
 
                 var exceptions = new List<Exception>();
                 foreach (Task t in tasks)

diff --git a/src/libraries/Microsoft.Extensions.FileProviders.Physical/tests/PhysicalFileProviderTests.cs b/src/libraries/Microsoft.Extensions.FileProviders.Physical/tests/PhysicalFileProviderTests.cs
@@ -1286,7 +1286,7 @@ void Fail(object state)
 
                 fileSystemWatcher.CallOnRenamed(new RenamedEventArgs(WatcherChangeTypes.Renamed, root.RootPath, newDirectoryName, oldDirectoryName));
 
-                await Task.WhenAll(oldDirectoryTcs.Task, newDirectoryTcs.Task, newSubDirectoryTcs.Task, newFileTcs.Task).TimeoutAfter(TimeSpan.FromSeconds(30));
+                await Task.WhenAll(oldDirectoryTcs.Task, newDirectoryTcs.Task, newSubDirectoryTcs.Task, newFileTcs.Task).WaitAsync(TimeSpan.FromSeconds(30));
 
                 Assert.False(oldSubDirectoryToken.HasChanged, "Old subdirectory token should not have changed");
                 Assert.False(oldFileToken.HasChanged, "Old file token should not have changed");

diff --git a/...nsions.Hosting/tests/FunctionalTests/IntegrationTesting/src/Deployers/SelfHostDeployer.cs b/...nsions.Hosting/tests/FunctionalTests/IntegrationTesting/src/Deployers/SelfHostDeployer.cs
@@ -164,7 +164,7 @@ protected async Task<CancellationToken> StartSelfHostAsync()
                     // The timeout here is large, because we don't know how long the test could need
                     // We cover a lot of error cases above, but I want to make sure we eventually give up and don't hang the build
                     // just in case we missed one -anurse
-                    await started.Task.TimeoutAfter(TimeSpan.FromMinutes(10));
+                    await started.Task.WaitAsync(TimeSpan.FromMinutes(10));
                 }
 
                 return hostExitTokenSource.Token;

diff --git a/src/libraries/Microsoft.Extensions.Hosting/tests/FunctionalTests/ShutdownTests.cs b/src/libraries/Microsoft.Extensions.Hosting/tests/FunctionalTests/ShutdownTests.cs
@@ -96,11 +96,11 @@ private async Task ExecuteShutdownTest(string testName, string shutdownMechanic)
                     }
                 };
 
-                await started.Task.TimeoutAfter(TimeSpan.FromSeconds(60));
+                await started.Task.WaitAsync(TimeSpan.FromSeconds(60));
 
                 SendShutdownSignal(deployer.HostProcess);
 
-                await completed.Task.TimeoutAfter(TimeSpan.FromSeconds(60));
+                await completed.Task.WaitAsync(TimeSpan.FromSeconds(60));
 
                 WaitForExitOrKill(deployer.HostProcess);
 

diff --git a/...ollections/tests/Generic/Dictionary/Dictionary.Generic.Tests.ConcurrentAccessDetection.cs b/...ollections/tests/Generic/Dictionary/Dictionary.Generic.Tests.ConcurrentAccessDetection.cs
@@ -34,7 +34,7 @@ private async Task DictionaryConcurrentAccessDetection<TKey, TValue>(Dictionary<
             }, TaskCreationOptions.LongRunning);
 
             // If Dictionary regresses, we do not want to hang here indefinitely
-            Assert.True((await Task.WhenAny(task, Task.Delay(TimeSpan.FromSeconds(60))) == task) && task.IsCompletedSuccessfully);
+            await task.WaitAsync(TimeSpan.FromSeconds(60));
         }
 
         [ConditionalTheory(typeof(PlatformDetection), nameof(PlatformDetection.IsThreadingSupported))]

diff --git a/src/libraries/System.ComponentModel.EventBasedAsync/tests/BackgroundWorkerTests.cs b/src/libraries/System.ComponentModel.EventBasedAsync/tests/BackgroundWorkerTests.cs
@@ -84,8 +84,7 @@ public async Task RunWorkerAsync_NoOnWorkHandler_SetsResultToNull()
 
             backgroundWorker.RunWorkerAsync();
 
-            await Task.WhenAny(tcs.Task, Task.Delay(TimeSpan.FromSeconds(10))); // Usually takes 100th of a sec
-            Assert.True(tcs.Task.IsCompleted);
+            await tcs.Task.WaitAsync(TimeSpan.FromSeconds(10)); // Usually takes 100th of a sec
         }
 
         #region TestCancelAsync

diff --git a/src/libraries/System.Diagnostics.Process/src/System.Diagnostics.Process.csproj b/src/libraries/System.Diagnostics.Process/src/System.Diagnostics.Process.csproj
@@ -37,10 +37,6 @@
              Link="System\PasteArguments.cs" />
     <Compile Include="$(CommonPath)Interop\Windows\Interop.Errors.cs"
              Link="Common\Interop\Windows\Interop.Errors.cs" />
-    <Compile Include="$(CommonPath)System\Threading\Tasks\TaskCompletionSourceWithCancellation.cs"
-             Link="Common\System\Threading\Tasks\TaskCompletionSourceWithCancellation.cs" />
-    <Compile Include="$(CommonPath)System\Threading\Tasks\TaskTimeoutExtensions.cs"
-             Link="Common\System\Threading\Tasks\TaskTimeoutExtensions.cs" />
     <Compile Include="$(CommonPath)System\Text\ValueStringBuilder.cs"
              Link="Common\System\Text\ValueStringBuilder.cs" />
   </ItemGroup>

diff --git a/src/libraries/System.Diagnostics.Process/src/System/Diagnostics/AsyncStreamReader.cs b/src/libraries/System.Diagnostics.Process/src/System/Diagnostics/AsyncStreamReader.cs
@@ -251,25 +251,7 @@ private bool FlushMessageQueue(bool rethrowInNewThread)
             }
         }
 
-        // Wait until we hit EOF. This is called from Process.WaitForExit
-        // We will lose some information if we don't do this.
-        internal void WaitUntilEOF()
-        {
-            if (_readToBufferTask is Task task)
-            {
-                task.GetAwaiter().GetResult();
-            }
-        }
-
-        internal Task WaitUntilEOFAsync(CancellationToken cancellationToken)
-        {
-            if (_readToBufferTask is Task task)
-            {
-                return task.WithCancellation(cancellationToken);
-            }
-
-            return Task.CompletedTask;
-        }
+        internal Task EOF => _readToBufferTask ?? Task.CompletedTask;
 
         public void Dispose()
         {