[SignalR] Avoid deadlock with closing and user callbacks (#19612) (#19664)

* [SignalR] Avoid deadlock with closing and user callbacks (#19612)

* fb

Author: BrennanConroy

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

@@ -1207,11 +1207,13 @@ async Task StartProcessingInvocationMessages(ChannelReader<InvocationMessage> in
             finally
             {
                 invocationMessageChannel.Writer.TryComplete();
-                await invocationMessageReceiveTask;
                 timer.Stop();
                 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;
             }
         }
 

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

@@ -345,7 +345,7 @@ public async Task UploadStreamCancelationSendsStreamComplete()
                 var complete = await connection.ReadSentJsonAsync().OrTimeout();
                 Assert.Equal(HubProtocolConstants.CompletionMessageType, complete["type"]);
                 Assert.EndsWith("canceled by client.", ((string)complete["error"]));
-            } 
+            }
         }
 
         [Fact]
@@ -414,6 +414,52 @@ bool ExpectedErrors(WriteContext writeContext)
             }
         }
 
+        [Fact]
+        public async Task CanAwaitInvokeFromOnHandlerWithServerClosingConnection()
+        {
+            using (StartVerifiableLog())
+            {
+                var connection = new TestConnection();
+                var hubConnection = CreateHubConnection(connection, loggerFactory: LoggerFactory);
+                await hubConnection.StartAsync().OrTimeout();
+
+                var tcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
+                hubConnection.On<string>("Echo", async msg =>
+                {
+                    try
+                    {
+                        // This should be canceled when the connection is closed
+                        await hubConnection.InvokeAsync<string>("Echo", msg).OrTimeout();
+                    }
+                    catch (Exception ex)
+                    {
+                        tcs.SetException(ex);
+                        return;
+                    }
+
+                    tcs.SetResult(null);
+                });
+
+                var closedTcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
+                hubConnection.Closed += _ =>
+                {
+                    closedTcs.SetResult(null);
+
+                    return Task.CompletedTask;
+                };
+
+                await connection.ReceiveJsonMessage(new { type = HubProtocolConstants.InvocationMessageType, target = "Echo", arguments = new object[] { "42" } }).OrTimeout();
+
+                // Read sent message first to make sure invoke has been processed and is waiting for a response
+                await connection.ReadSentJsonAsync().OrTimeout();
+                await connection.ReceiveJsonMessage(new { type = HubProtocolConstants.CloseMessageType }).OrTimeout();
+
+                await closedTcs.Task.OrTimeout();
+
+                await Assert.ThrowsAsync<TaskCanceledException>(() => tcs.Task.OrTimeout());
+            }
+        }
+
         private class SampleObject
         {
             public SampleObject(string foo, int bar)