From cb5bba36fcd9233f499bec429c7ce6ccafb98fcb Mon Sep 17 00:00:00 2001
From: David Fowler <davidfowl@gmail.com>
Date: Wed, 11 Apr 2018 10:39:09 -0700
Subject: [PATCH] Get off the event loop before dispatch (#1952)

---
 .../AwaitableThreadPool.cs                    |  2 +-
 ...crosoft.AspNetCore.Http.Connections.csproj |  1 +
 .../HubConnection.cs                          | 43 +++++++++++--------
 ...soft.AspNetCore.SignalR.Client.Core.csproj |  1 +
 .../HubConnectionTests.Protocol.cs            |  4 +-
 5 files changed, 28 insertions(+), 23 deletions(-)
 rename src/{Microsoft.AspNetCore.Http.Connections/Internal => Common}/AwaitableThreadPool.cs (94%)

diff --git a/src/Microsoft.AspNetCore.Http.Connections/Internal/AwaitableThreadPool.cs b/src/Common/AwaitableThreadPool.cs
similarity index 94%
rename from src/Microsoft.AspNetCore.Http.Connections/Internal/AwaitableThreadPool.cs
rename to src/Common/AwaitableThreadPool.cs
index d4b14df268..e27c8cbbbf 100644
--- a/src/Microsoft.AspNetCore.Http.Connections/Internal/AwaitableThreadPool.cs
+++ b/src/Common/AwaitableThreadPool.cs
@@ -7,7 +7,7 @@ using System.Linq;
 using System.Runtime.CompilerServices;
 using System.Threading.Tasks;
 
-namespace Microsoft.AspNetCore.Http.Connections.Internal
+namespace Microsoft.AspNetCore.Internal
 {
     public static class AwaitableThreadPool
     {
diff --git a/src/Microsoft.AspNetCore.Http.Connections/Microsoft.AspNetCore.Http.Connections.csproj b/src/Microsoft.AspNetCore.Http.Connections/Microsoft.AspNetCore.Http.Connections.csproj
index cd36dcd2e0..0b0a149c3a 100644
--- a/src/Microsoft.AspNetCore.Http.Connections/Microsoft.AspNetCore.Http.Connections.csproj
+++ b/src/Microsoft.AspNetCore.Http.Connections/Microsoft.AspNetCore.Http.Connections.csproj
@@ -6,6 +6,7 @@
   </PropertyGroup>
 
   <ItemGroup>
+    <Compile Include="..\Common\AwaitableThreadPool.cs" Link="AwaitableThreadPool.cs" />
     <Compile Include="..\Common\MemoryBufferWriter.cs" Link="MemoryBufferWriter.cs" />
     <Compile Include="..\Common\PipeWriterStream.cs" Link="PipeWriterStream.cs" />
     <Compile Include="..\Common\WebSocketExtensions.cs" Link="WebSocketExtensions.cs" />
diff --git a/src/Microsoft.AspNetCore.SignalR.Client.Core/HubConnection.cs b/src/Microsoft.AspNetCore.SignalR.Client.Core/HubConnection.cs
index fa9e33ccf8..cbdae52a48 100644
--- a/src/Microsoft.AspNetCore.SignalR.Client.Core/HubConnection.cs
+++ b/src/Microsoft.AspNetCore.SignalR.Client.Core/HubConnection.cs
@@ -12,6 +12,7 @@ using System.Threading.Channels;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Connections;
 using Microsoft.AspNetCore.Connections.Features;
+using Microsoft.AspNetCore.Internal;
 using Microsoft.AspNetCore.SignalR.Internal;
 using Microsoft.AspNetCore.SignalR.Internal.Protocol;
 using Microsoft.Extensions.DependencyInjection;
@@ -433,6 +434,9 @@ namespace Microsoft.AspNetCore.SignalR.Client
 
         private async Task DispatchInvocationAsync(InvocationMessage invocation)
         {
+            // Make sure we get off the main event loop before we dispatch into user code
+            await AwaitableThreadPool.Yield();
+
             // Find the handler
             if (!_handlers.TryGetValue(invocation.Target, out var handlers))
             {
@@ -678,29 +682,30 @@ namespace Microsoft.AspNetCore.SignalR.Client
                 Log.ShutdownConnection(_logger);
             }
 
-            // Fire-and-forget the closed event
-            RunClosedEvent(connectionState.CloseException);
-        }
-
-        private void RunClosedEvent(Exception closeException)
-        {
             var closed = Closed;
 
             // There is no need to start a new task if there is no Closed event registered
             if (closed != null)
             {
-                _ = Task.Run(() =>
-                {
-                    try
-                    {
-                        Log.InvokingClosedEventHandler(_logger);
-                        closed.Invoke(closeException);
-                    }
-                    catch (Exception ex)
-                    {
-                        Log.ErrorDuringClosedEvent(_logger, ex);
-                    }
-                });
+
+                // Fire-and-forget the closed event
+                _ = RunClosedEvent(closed, connectionState.CloseException);
+            }
+        }
+
+        private async Task RunClosedEvent(Action<Exception> closed, Exception closeException)
+        {
+            // Dispatch to the thread pool before we invoke the user callback
+            await AwaitableThreadPool.Yield();
+
+            try
+            {
+                Log.InvokingClosedEventHandler(_logger);
+                closed.Invoke(closeException);
+            }
+            catch (Exception ex)
+            {
+                Log.ErrorDuringClosedEvent(_logger, ex);
             }
         }
 
@@ -925,7 +930,7 @@ namespace Microsoft.AspNetCore.SignalR.Client
                     }
                     else
                     {
-                        _stopTcs = new TaskCompletionSource<object>();
+                        _stopTcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
                         return StopAsyncCore(timeout);
                     }
                 }
diff --git a/src/Microsoft.AspNetCore.SignalR.Client.Core/Microsoft.AspNetCore.SignalR.Client.Core.csproj b/src/Microsoft.AspNetCore.SignalR.Client.Core/Microsoft.AspNetCore.SignalR.Client.Core.csproj
index 8cd58928fc..1d91bf64cc 100644
--- a/src/Microsoft.AspNetCore.SignalR.Client.Core/Microsoft.AspNetCore.SignalR.Client.Core.csproj
+++ b/src/Microsoft.AspNetCore.SignalR.Client.Core/Microsoft.AspNetCore.SignalR.Client.Core.csproj
@@ -7,6 +7,7 @@
   </PropertyGroup>
 
   <ItemGroup>
+    <Compile Include="..\Common\AwaitableThreadPool.cs" Link="AwaitableThreadPool.cs" />
     <Compile Include="..\Common\ForceAsyncAwaiter.cs" Link="ForceAsyncAwaiter.cs" />
     <Compile Include="..\Common\PipeWriterStream.cs" Link="PipeWriterStream.cs" />
   </ItemGroup>
diff --git a/test/Microsoft.AspNetCore.SignalR.Client.Tests/HubConnectionTests.Protocol.cs b/test/Microsoft.AspNetCore.SignalR.Client.Tests/HubConnectionTests.Protocol.cs
index 39019b8aaa..87a0682edd 100644
--- a/test/Microsoft.AspNetCore.SignalR.Client.Tests/HubConnectionTests.Protocol.cs
+++ b/test/Microsoft.AspNetCore.SignalR.Client.Tests/HubConnectionTests.Protocol.cs
@@ -485,9 +485,7 @@ namespace Microsoft.AspNetCore.SignalR.Client.Tests
 
                     await connection.ReceiveTextAsync(":[\"hello\"]}\u001e");
 
-                    Assert.True(tcs.Task.IsCompleted);
-
-                    var response = await tcs.Task;
+                    var response = await tcs.Task.OrTimeout();
 
                     Assert.Equal("hello", response);
                 }