Don't wait for client methods to complete when stopping connection (#20104)

Author: BrennanConroy

src/SignalR/clients/csharp/Client.Core/src/HubConnection.cs

@@ -1128,7 +1128,9 @@ private async Task ReceiveLoop(ConnectionState connectionState)
             var uploadStreamSource = new CancellationTokenSource();
             connectionState.UploadStreamToken = uploadStreamSource.Token;
             var invocationMessageChannel = Channel.CreateUnbounded<InvocationMessage>(_receiveLoopOptions);
-            var invocationMessageReceiveTask = StartProcessingInvocationMessages(invocationMessageChannel.Reader);
+
+            // We can't safely wait for this task when closing without introducing deadlock potential when calling StopAsync in a .On method
+            connectionState.InvocationMessageReceiveTask = StartProcessingInvocationMessages(invocationMessageChannel.Reader);
 
             async Task StartProcessingInvocationMessages(ChannelReader<InvocationMessage> invocationMessageChannelReader)
             {
@@ -1223,9 +1225,6 @@ async Task StartProcessingInvocationMessages(ChannelReader<InvocationMessage> in
                 await timerTask;
                 uploadStreamSource.Cancel();
                 await HandleConnectionClose(connectionState);
-
-                // await after the connection has been closed, otherwise could deadlock on a user's .On callback(s)
-                await invocationMessageReceiveTask;
             }
         }
 
@@ -1658,6 +1657,9 @@ private class ConnectionState : IInvocationBinder
             public Exception CloseException { get; set; }
             public CancellationToken UploadStreamToken { get; set; }
 
+            // We store this task so we can view it in a dump file, but never await it
+            public Task InvocationMessageReceiveTask { get; set; }
+
             // Indicates the connection is stopping AND the client should NOT attempt to reconnect even if automatic reconnects are enabled.
             // This means either HubConnection.DisposeAsync/StopAsync was called OR a CloseMessage with AllowReconnects set to false was received.
             public bool Stopping

src/SignalR/clients/csharp/Client/test/UnitTests/HubConnectionTests.cs

@@ -78,6 +78,50 @@ public async Task ClosedEventRaisedWhenTheClientIsStopped()
             Assert.Null(await closedEventTcs.Task);
         }
 
+        [Fact]
+        public async Task StopAsyncCanBeCalledFromOnHandler()
+        {
+            var connection = new TestConnection();
+            var hubConnection = CreateHubConnection(connection, loggerFactory: LoggerFactory);
+
+            var tcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
+            hubConnection.On("method", async () =>
+            {
+                await hubConnection.StopAsync().OrTimeout();
+                tcs.SetResult(null);
+            });
+
+            await hubConnection.StartAsync().OrTimeout();
+
+            await connection.ReceiveJsonMessage(new { type = HubProtocolConstants.InvocationMessageType, target= "method", arguments = new object[] { } }).OrTimeout();
+
+            Assert.Null(await tcs.Task.OrTimeout());
+        }
+
+        [Fact]
+        public async Task StopAsyncDoesNotWaitForOnHandlers()
+        {
+            var connection = new TestConnection();
+            var hubConnection = CreateHubConnection(connection, loggerFactory: LoggerFactory);
+
+            var tcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
+            var methodCalledTcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
+            hubConnection.On("method", async () =>
+            {
+                methodCalledTcs.SetResult(null);
+                await tcs.Task;
+            });
+
+            await hubConnection.StartAsync().OrTimeout();
+
+            await connection.ReceiveJsonMessage(new { type = HubProtocolConstants.InvocationMessageType, target = "method", arguments = new object[] { } }).OrTimeout();
+
+            await methodCalledTcs.Task.OrTimeout();
+            await hubConnection.StopAsync().OrTimeout();
+
+            tcs.SetResult(null);
+        }
+
         [Fact]
         public async Task PendingInvocationsAreCanceledWhenConnectionClosesCleanly()
         {