Send CloseMessage in more cases (#25693)

* Send CloseMessage in more cases

* fb

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

@@ -1349,7 +1349,11 @@ namespace Microsoft.AspNetCore.SignalR.Client.FunctionalTests
             };
 
             var protocol = HubProtocols[protocolName];
-            await using (var server = await StartServer<Startup>(write => write.EventId.Name == "FailedWritingMessage"))
+            await using (var server = await StartServer<Startup>(write =>
+            {
+                return write.EventId.Name == "FailedWritingMessage" || write.EventId.Name == "ReceivedCloseWithError"
+                    || write.EventId.Name == "ShutdownWithError";
+            }))
             {
                 var connection = CreateHubConnection(server.Url, "/default", HttpTransportType.WebSockets, protocol, LoggerFactory);
                 var closedTcs = new TaskCompletionSource<Exception>(TaskCreationOptions.RunContinuationsAsynchronously);
@@ -1361,9 +1365,12 @@ namespace Microsoft.AspNetCore.SignalR.Client.FunctionalTests
                     var result = connection.InvokeAsync<string>(nameof(TestHub.CallWithUnserializableObject));
 
                     // The connection should close.
-                    Assert.Null(await closedTcs.Task.OrTimeout());
+                    var exception = await closedTcs.Task.OrTimeout();
+                    Assert.Contains("Connection closed with an error.", exception.Message);
 
-                    await Assert.ThrowsAsync<TaskCanceledException>(() => result).OrTimeout();
+                    var hubException = await Assert.ThrowsAsync<HubException>(() => result).OrTimeout();
+                    Assert.Contains("Connection closed with an error.", hubException.Message);
+                    Assert.Contains(exceptionSubstring, hubException.Message);
                 }
                 catch (Exception ex)
                 {
@@ -1396,7 +1403,11 @@ namespace Microsoft.AspNetCore.SignalR.Client.FunctionalTests
             };
 
             var protocol = HubProtocols[protocolName];
-            await using (var server = await StartServer<Startup>(write => write.EventId.Name == "FailedWritingMessage"))
+            await using (var server = await StartServer<Startup>(write =>
+            {
+                return write.EventId.Name == "FailedWritingMessage" || write.EventId.Name == "ReceivedCloseWithError"
+                    || write.EventId.Name == "ShutdownWithError";
+            }))
             {
                 var connection = CreateHubConnection(server.Url, "/default", HttpTransportType.LongPolling, protocol, LoggerFactory);
                 var closedTcs = new TaskCompletionSource<Exception>(TaskCreationOptions.RunContinuationsAsynchronously);
@@ -1408,9 +1419,12 @@ namespace Microsoft.AspNetCore.SignalR.Client.FunctionalTests
                     var result = connection.InvokeAsync<string>(nameof(TestHub.GetUnserializableObject)).OrTimeout();
 
                     // The connection should close.
-                    Assert.Null(await closedTcs.Task.OrTimeout());
+                    var exception = await closedTcs.Task.OrTimeout();
+                    Assert.Contains("Connection closed with an error.", exception.Message);
 
-                    await Assert.ThrowsAsync<TaskCanceledException>(() => result).OrTimeout();
+                    var hubException = await Assert.ThrowsAsync<HubException>(() => result).OrTimeout();
+                    Assert.Contains("Connection closed with an error.", hubException.Message);
+                    Assert.Contains(exceptionSubstring, hubException.Message);
                 }
                 catch (Exception ex)
                 {

src/SignalR/server/Core/src/HubConnectionContext.cs

@@ -149,14 +149,19 @@ namespace Microsoft.AspNetCore.SignalR
 
         [SuppressMessage("ApiDesign", "RS0026:Do not add multiple overloads with optional parameters", Justification = "Required to maintain compatibility")]
         public virtual ValueTask WriteAsync(HubMessage message, CancellationToken cancellationToken = default)
+        {
+            return WriteAsync(message, ignoreAbort: false, cancellationToken);
+        }
+
+        internal ValueTask WriteAsync(HubMessage message, bool ignoreAbort, CancellationToken cancellationToken = default)
         {
             // Try to grab the lock synchronously, if we fail, go to the slower path
             if (!_writeLock.Wait(0))
             {
-                return new ValueTask(WriteSlowAsync(message, cancellationToken));
+                return new ValueTask(WriteSlowAsync(message, ignoreAbort, cancellationToken));
             }
 
-            if (_connectionAborted)
+            if (_connectionAborted && !ignoreAbort)
             {
                 _writeLock.Release();
                 return default;
@@ -275,14 +280,14 @@ namespace Microsoft.AspNetCore.SignalR
             }
         }
 
-        private async Task WriteSlowAsync(HubMessage message, CancellationToken cancellationToken)
+        private async Task WriteSlowAsync(HubMessage message, bool ignoreAbort, CancellationToken cancellationToken)
         {
             // Failed to get the lock immediately when entering WriteAsync so await until it is available
             await _writeLock.WaitAsync(cancellationToken);
 
             try
             {
-                if (_connectionAborted)
+                if (_connectionAborted && !ignoreAbort)
                 {
                     return;
                 }
@@ -304,7 +309,7 @@ namespace Microsoft.AspNetCore.SignalR
         private async Task WriteSlowAsync(SerializedHubMessage message, CancellationToken cancellationToken)
         {
             // Failed to get the lock immediately when entering WriteAsync so await until it is available
-            await _writeLock.WaitAsync();
+            await _writeLock.WaitAsync(cancellationToken);
 
             try
             {

src/SignalR/server/Core/src/HubConnectionHandler.cs

@@ -226,7 +226,7 @@ namespace Microsoft.AspNetCore.SignalR
 
             try
             {
-                await connection.WriteAsync(closeMessage);
+                await connection.WriteAsync(closeMessage, ignoreAbort: true);
             }
             catch (Exception ex)
             {

src/SignalR/server/SignalR/test/HubConnectionHandlerTests.cs

@@ -155,6 +155,9 @@ namespace Microsoft.AspNetCore.SignalR.Tests
 
                     await client.SendInvocationAsync(nameof(AbortHub.Kill)).OrTimeout();
 
+                    var close = Assert.IsType<CloseMessage>(await client.ReadAsync().OrTimeout());
+                    Assert.False(close.AllowReconnect);
+
                     await connectionHandlerTask.OrTimeout();
 
                     Assert.Null(client.TryRead());
@@ -955,15 +958,18 @@ namespace Microsoft.AspNetCore.SignalR.Tests
                 {
                     var connectionHandlerTask = await client.ConnectAsync(connectionHandler);
 
-                    var invokeTask = client.InvokeAsync(nameof(MethodHub.BlockingMethod));
+                    await client.SendInvocationAsync(nameof(MethodHub.BlockingMethod)).OrTimeout();
 
                     client.Connection.Abort();
 
+                    var closeMessage = Assert.IsType<CloseMessage>(await client.ReadAsync().OrTimeout());
+                    Assert.False(closeMessage.AllowReconnect);
+
                     // If this completes then the server has completed the connection
                     await connectionHandlerTask.OrTimeout();
 
                     // Nothing written to connection because it was closed
-                    Assert.False(invokeTask.IsCompleted);
+                    Assert.Null(client.TryRead());
                 }
             }
         }
@@ -1019,16 +1025,11 @@ namespace Microsoft.AspNetCore.SignalR.Tests
                     // kill the connection
                     client.Dispose();
 
+                    var message = Assert.IsType<CloseMessage>(client.TryRead());
+                    Assert.True(message.AllowReconnect);
+
                     // Ensure the client channel is empty
-                    var message = client.TryRead();
+                    Assert.Null(client.TryRead());
-                    switch (message)
-                    {
-                        case CloseMessage close:
-                            break;
-                        default:
-                            Assert.Null(message);
-                            break;
-                    }
 
                     await connectionHandlerTask.OrTimeout();
                 }