Add EnsureCompleted task extensions.

lukebakken · lukebakken · commit 48dac7bc58a6 · 2023-11-21T13:23:25.000-08:00
diff --git a/projects/RabbitMQ.Client/client/TaskExtensions.cs b/projects/RabbitMQ.Client/client/TaskExtensions.cs
@@ -87,40 +87,73 @@ static async Task DoTimeoutAfter(Task task, TimeSpan timeout)
 #endif
         }
 
-        public static async ValueTask TimeoutAfter(this ValueTask task, TimeSpan timeout)
+        public static async ValueTask TimeoutAfter(this ValueTask valueTask, TimeSpan timeout)
         {
-            if (task.IsCompletedSuccessfully)
+            if (valueTask.IsCompletedSuccessfully)
             {
                 return;
             }
 
 #if NET6_0_OR_GREATER
-            Task actualTask = task.AsTask();
-            await actualTask.WaitAsync(timeout)
+            Task task = valueTask.AsTask();
+            await task.WaitAsync(timeout)
                 .ConfigureAwait(false);
 #else
-            await DoTimeoutAfter(task, timeout)
+            await DoTimeoutAfter(valueTask, timeout)
                 .ConfigureAwait(false);
 
-            async static ValueTask DoTimeoutAfter(ValueTask task, TimeSpan timeout)
+            // https://github.com/davidfowl/AspNetCoreDiagnosticScenarios/blob/master/AsyncGuidance.md#using-a-timeout
+            static async ValueTask DoTimeoutAfter(ValueTask valueTask, TimeSpan timeout)
             {
-                Task actualTask = task.AsTask();
-                if (actualTask == await Task.WhenAny(actualTask, Task.Delay(timeout)).ConfigureAwait(false))
-                {
-                    await actualTask.ConfigureAwait(false);
-                }
-                else
+                Task task = valueTask.AsTask();
+                using (var cts = new CancellationTokenSource())
                 {
-                    Task supressErrorTask = actualTask.ContinueWith(
-                        continuationAction: continuation,
-                        state: null,
-                        cancellationToken: CancellationToken.None,
-                        continuationOptions: s_tco,
-                        scheduler: TaskScheduler.Default);
-                    throw new TimeoutException();
+                    Task delayTask = Task.Delay(timeout, cts.Token);
+                    Task resultTask = await Task.WhenAny(task, delayTask).ConfigureAwait(false);
+                    if (resultTask == delayTask)
+                    {
+                        Task supressErrorTask = task.ContinueWith(
+                            continuationAction: continuation,
+                            state: null,
+                            cancellationToken: CancellationToken.None,
+                            continuationOptions: s_tco,
+                            scheduler: TaskScheduler.Default);
+                        throw new TimeoutException();
+                    }
+                    else
+                    {
+                        cts.Cancel();
+                    }
+
+                    await valueTask.ConfigureAwait(false);
                 }
             }
 #endif
         }
+
+        /*
+         * https://devblogs.microsoft.com/dotnet/configureawait-faq/
+         * I'm using GetAwaiter().GetResult(). Do I need to use ConfigureAwait(false)?
+         * Answer: No
+         */
+        public static void EnsureCompleted(this Task task)
+        {
+            task.GetAwaiter().GetResult();
+        }
+
+        public static T EnsureCompleted<T>(this Task<T> task)
+        {
+            return task.GetAwaiter().GetResult();
+        }
+
+        public static T EnsureCompleted<T>(this ValueTask<T> task)
+        {
+            return task.GetAwaiter().GetResult();
+        }
+
+        public static void EnsureCompleted(this ValueTask task)
+        {
+            task.GetAwaiter().GetResult();
+        }
     }
 }
diff --git a/projects/RabbitMQ.Client/client/impl/ChannelBase.cs b/projects/RabbitMQ.Client/client/impl/ChannelBase.cs
@@ -287,6 +287,7 @@ await ConsumerDispatcher.WaitForShutdownAsync()
             }
         }
 
+        // TODO cancellation tokens
         internal async ValueTask ConnectionOpenAsync(string virtualHost)
         {
             var m = new ConnectionOpen(virtualHost);
diff --git a/projects/RabbitMQ.Client/client/impl/Connection.cs b/projects/RabbitMQ.Client/client/impl/Connection.cs
@@ -213,12 +213,7 @@ internal void TakeOver(Connection other)
 
         internal IConnection Open()
         {
-            /*
-             * https://devblogs.microsoft.com/dotnet/configureawait-faq/
-             * I'm using GetAwaiter().GetResult(). Do I need to use ConfigureAwait(false)?
-             * Answer: No
-             */
-            return OpenAsync().GetAwaiter().GetResult();
+            return OpenAsync().EnsureCompleted();
         }
 
         internal async ValueTask<IConnection> OpenAsync()
@@ -530,12 +525,7 @@ internal void Write(RentedMemory frames)
             ValueTask task = _frameHandler.WriteAsync(frames);
             if (!task.IsCompletedSuccessfully)
             {
-                /*
-                 * https://devblogs.microsoft.com/dotnet/configureawait-faq/
-                 * I'm using GetAwaiter().GetResult(). Do I need to use ConfigureAwait(false)?
-                 * Answer: No
-                 */
-                task.GetAwaiter().GetResult();
+                task.EnsureCompleted();
             }
         }
 
diff --git a/projects/RabbitMQ.Client/client/impl/SocketFrameHandler.cs b/projects/RabbitMQ.Client/client/impl/SocketFrameHandler.cs
@@ -183,12 +183,7 @@ public TimeSpan WriteTimeout
 
         public void Close()
         {
-            /*
-             * https://devblogs.microsoft.com/dotnet/configureawait-faq/
-             * I'm using GetAwaiter().GetResult(). Do I need to use ConfigureAwait(false)?
-             * Answer: No
-             */
-            CloseAsync().GetAwaiter().GetResult();
+            CloseAsync().EnsureCompleted();
         }
 
         public async ValueTask CloseAsync()
@@ -339,11 +334,9 @@ private void ConnectOrFail(ITcpClient socket, IPEndPoint endpoint, TimeSpan time
         {
             try
             {
+                // this ensures exceptions aren't wrapped in an AggregateException
                 socket.ConnectAsync(endpoint.Address, endpoint.Port)
-                     .TimeoutAfter(timeout)
-                     // this ensures exceptions aren't wrapped in an AggregateException
-                     .GetAwaiter()
-                     .GetResult();
+                     .TimeoutAfter(timeout).EnsureCompleted();
             }
             catch (ArgumentException e)
             {
diff --git a/projects/RabbitMQ.Client/client/impl/SslHelper.cs b/projects/RabbitMQ.Client/client/impl/SslHelper.cs
@@ -53,6 +53,7 @@ private SslHelper(SslOption sslOption)
         /// <summary>
         /// Upgrade a Tcp stream to an Ssl stream using the TLS options provided.
         /// </summary>
+        // TODO async
         public static Stream TcpUpgrade(Stream tcpStream, SslOption options)
         {
             var helper = new SslHelper(options);
@@ -67,8 +68,9 @@ public static Stream TcpUpgrade(Stream tcpStream, SslOption options)
             Action<SslOption> TryAuthenticating = (SslOption opts) =>
             {
                 sslStream.AuthenticateAsClientAsync(opts.ServerName, opts.Certs, opts.Version,
-                                                    opts.CheckCertificateRevocation).GetAwaiter().GetResult();
+                    opts.CheckCertificateRevocation).EnsureCompleted();
             };
+
             try
             {
                 // TODO async
diff --git a/projects/Test/SequentialIntegration/TestConnectionRecoveryBase.cs b/projects/Test/SequentialIntegration/TestConnectionRecoveryBase.cs
@@ -255,12 +255,7 @@ internal void RestartServerAndWaitForRecovery(AutorecoveringConnection conn)
         protected bool WaitForConfirms(IChannel m)
         {
             using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(4));
-            /*
-             * https://devblogs.microsoft.com/dotnet/configureawait-faq/
-             * I'm using GetAwaiter().GetResult(). Do I need to use ConfigureAwait(false)?
-             * Answer: No
-             */
-            return m.WaitForConfirmsAsync(cts.Token).GetAwaiter().GetResult();
+            return m.WaitForConfirmsAsync(cts.Token).EnsureCompleted();
         }
 
         protected void WaitForRecovery()

Original file line number	Diff line number	Diff line change
`@@ -287,6 +287,7 @@ await ConsumerDispatcher.WaitForShutdownAsync()`
`287`	`287`	`}`
`288`	`288`	`}`
`289`	`289`
	`290`	`+ // TODO cancellation tokens`
`290`	`291`	`internal async ValueTask ConnectionOpenAsync(string virtualHost)`
`291`	`292`	`{`
`292`	`293`	`var m = new ConnectionOpen(virtualHost);`
Original file line number	Diff line number	Diff line change
`@@ -213,12 +213,7 @@ internal void TakeOver(Connection other)`
`213`	`213`
`214`	`214`	`internal IConnection Open()`
`215`	`215`	`{`
`216`		`- /*`
`217`		`- * https://devblogs.microsoft.com/dotnet/configureawait-faq/`
`218`		`- * I'm using GetAwaiter().GetResult(). Do I need to use ConfigureAwait(false)?`
`219`		`- * Answer: No`
`220`		`- */`
`221`		`- return OpenAsync().GetAwaiter().GetResult();`
	`216`	`+ return OpenAsync().EnsureCompleted();`
`222`	`217`	`}`
`223`	`218`
`224`	`219`	`internal async ValueTask<IConnection> OpenAsync()`
`@@ -530,12 +525,7 @@ internal void Write(RentedMemory frames)`
`530`	`525`	`ValueTask task = _frameHandler.WriteAsync(frames);`
`531`	`526`	`if (!task.IsCompletedSuccessfully)`
`532`	`527`	`{`
`533`		`- /*`
`534`		`- * https://devblogs.microsoft.com/dotnet/configureawait-faq/`
`535`		`- * I'm using GetAwaiter().GetResult(). Do I need to use ConfigureAwait(false)?`
`536`		`- * Answer: No`
`537`		`- */`
`538`		`- task.GetAwaiter().GetResult();`
	`528`	`+ task.EnsureCompleted();`
`539`	`529`	`}`
`540`	`530`	`}`
`541`	`531`
Original file line number	Diff line number	Diff line change
`@@ -183,12 +183,7 @@ public TimeSpan WriteTimeout`
`183`	`183`
`184`	`184`	`public void Close()`
`185`	`185`	`{`
`186`		`- /*`
`187`		`- * https://devblogs.microsoft.com/dotnet/configureawait-faq/`
`188`		`- * I'm using GetAwaiter().GetResult(). Do I need to use ConfigureAwait(false)?`
`189`		`- * Answer: No`
`190`		`- */`
`191`		`- CloseAsync().GetAwaiter().GetResult();`
	`186`	`+ CloseAsync().EnsureCompleted();`
`192`	`187`	`}`
`193`	`188`
`194`	`189`	`public async ValueTask CloseAsync()`
`@@ -339,11 +334,9 @@ private void ConnectOrFail(ITcpClient socket, IPEndPoint endpoint, TimeSpan time`
`339`	`334`	`{`
`340`	`335`	`try`
`341`	`336`	`{`
	`337`	`+ // this ensures exceptions aren't wrapped in an AggregateException`
`342`	`338`	`socket.ConnectAsync(endpoint.Address, endpoint.Port)`
`343`		`- .TimeoutAfter(timeout)`
`344`		`- // this ensures exceptions aren't wrapped in an AggregateException`
`345`		`- .GetAwaiter()`
`346`		`- .GetResult();`
	`339`	`+ .TimeoutAfter(timeout).EnsureCompleted();`
`347`	`340`	`}`
`348`	`341`	`catch (ArgumentException e)`
`349`	`342`	`{`
Original file line number	Diff line number	Diff line change
`@@ -53,6 +53,7 @@ private SslHelper(SslOption sslOption)`
`53`	`53`	`/// <summary>`
`54`	`54`	`/// Upgrade a Tcp stream to an Ssl stream using the TLS options provided.`
`55`	`55`	`/// </summary>`
	`56`	`+ // TODO async`
`56`	`57`	`public static Stream TcpUpgrade(Stream tcpStream, SslOption options)`
`57`	`58`	`{`
`58`	`59`	`var helper = new SslHelper(options);`
`@@ -67,8 +68,9 @@ public static Stream TcpUpgrade(Stream tcpStream, SslOption options)`
`67`	`68`	`Action<SslOption> TryAuthenticating = (SslOption opts) =>`
`68`	`69`	`{`
`69`	`70`	`sslStream.AuthenticateAsClientAsync(opts.ServerName, opts.Certs, opts.Version,`
`70`		`- opts.CheckCertificateRevocation).GetAwaiter().GetResult();`
	`71`	`+ opts.CheckCertificateRevocation).EnsureCompleted();`
`71`	`72`	`};`
	`73`	`+`
`72`	`74`	`try`
`73`	`75`	`{`
`74`	`76`	`// TODO async`
Original file line number	Diff line number	Diff line change
`@@ -255,12 +255,7 @@ internal void RestartServerAndWaitForRecovery(AutorecoveringConnection conn)`
`255`	`255`	`protected bool WaitForConfirms(IChannel m)`
`256`	`256`	`{`
`257`	`257`	`using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(4));`
`258`		`- /*`
`259`		`- * https://devblogs.microsoft.com/dotnet/configureawait-faq/`
`260`		`- * I'm using GetAwaiter().GetResult(). Do I need to use ConfigureAwait(false)?`
`261`		`- * Answer: No`
`262`		`- */`
`263`		`- return m.WaitForConfirmsAsync(cts.Token).GetAwaiter().GetResult();`
	`258`	`+ return m.WaitForConfirmsAsync(cts.Token).EnsureCompleted();`
`264`	`259`	`}`
`265`	`260`
`266`	`261`	`protected void WaitForRecovery()`