[Backport] [SignalR] Fix WebSocket client close when network disappears (#43576)

Author: BrennanConroy

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

@@ -0,0 +1,79 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Net.WebSockets;
+using Microsoft.AspNetCore.Connections;
+using Microsoft.AspNetCore.Http.Connections.Client;
+using Microsoft.AspNetCore.Http.Connections.Client.Internal;
+using Microsoft.AspNetCore.SignalR.Tests;
+using Microsoft.AspNetCore.Testing;
+
+namespace Microsoft.AspNetCore.SignalR.Client.Tests;
+
+public class WebSocketsTransportTests : VerifiableLoggedTest
+{
+    // Tests that the transport can still be stopped if SendAsync and ReceiveAsync are hanging (ethernet unplugged for example)
+    [Fact]
+    public async Task StopCancelsSendAndReceive()
+    {
+        var options = new HttpConnectionOptions()
+        {
+            WebSocketFactory = (context, token) =>
+            {
+                return ValueTask.FromResult((WebSocket)new TestWebSocket());
+            },
+            CloseTimeout = TimeSpan.FromMilliseconds(1),
+        };
+
+        using (StartVerifiableLog())
+        {
+            var webSocketsTransport = new WebSocketsTransport(options, loggerFactory: LoggerFactory, () => Task.FromResult<string>(null));
+
+            await webSocketsTransport.StartAsync(
+                new Uri("http://fakeuri.org"), TransferFormat.Text).DefaultTimeout();
+
+            await webSocketsTransport.StopAsync().DefaultTimeout();
+
+            await webSocketsTransport.Running.DefaultTimeout();
+        }
+    }
+
+    internal class TestWebSocket : WebSocket
+    {
+        public Task ConnectAsync(Uri uri, CancellationToken cancellationToken) => Task.CompletedTask;
+
+        public override WebSocketCloseStatus? CloseStatus => null;
+
+        public override string CloseStatusDescription => string.Empty;
+
+        public override WebSocketState State => WebSocketState.Open;
+
+        public override string SubProtocol => string.Empty;
+
+        public override void Abort() { }
+
+        public override Task CloseAsync(WebSocketCloseStatus closeStatus, string statusDescription, CancellationToken cancellationToken)
+            => Task.CompletedTask;
+
+        public override async Task CloseOutputAsync(WebSocketCloseStatus closeStatus, string statusDescription, CancellationToken cancellationToken)
+        {
+            await cancellationToken.WaitForCancellationAsync();
+            cancellationToken.ThrowIfCancellationRequested();
+        }
+
+        public override void Dispose() { }
+
+        public override async Task<WebSocketReceiveResult> ReceiveAsync(ArraySegment<byte> buffer, CancellationToken cancellationToken)
+        {
+            await cancellationToken.WaitForCancellationAsync();
+            cancellationToken.ThrowIfCancellationRequested();
+            return new WebSocketReceiveResult(0, WebSocketMessageType.Text, true);
+        }
+
+        public override async Task SendAsync(ArraySegment<byte> buffer, WebSocketMessageType messageType, bool endOfMessage, CancellationToken cancellationToken)
+        {
+            await cancellationToken.WaitForCancellationAsync();
+            cancellationToken.ThrowIfCancellationRequested();
+        }
+    }
+}

src/SignalR/clients/csharp/Http.Connections.Client/src/Internal/WebSocketsTransport.cs

@@ -24,6 +24,7 @@ internal partial class WebSocketsTransport : ITransport
         private readonly TimeSpan _closeTimeout;
         private volatile bool _aborted;
         private readonly HttpConnectionOptions _httpConnectionOptions;
+        private readonly CancellationTokenSource _stopCts = new CancellationTokenSource();
 
         private IDuplexPipe? _transport;
 
@@ -204,6 +205,8 @@ private async Task ProcessSocketAsync(WebSocket socket)
                 // Wait for send or receive to complete
                 var trigger = await Task.WhenAny(receiving, sending);
 
+                _stopCts.CancelAfter(_closeTimeout);
+
                 if (trigger == receiving)
                 {
                     // We're waiting for the application to finish and there are 2 things it could be doing
@@ -213,22 +216,14 @@ private async Task ProcessSocketAsync(WebSocket socket)
                     // Cancel the application so that ReadAsync yields
                     _application.Input.CancelPendingRead();
 
-                    using (var delayCts = new CancellationTokenSource())
-                    {
-                        var resultTask = await Task.WhenAny(sending, Task.Delay(_closeTimeout, delayCts.Token));
+                    var resultTask = await Task.WhenAny(sending, Task.Delay(_closeTimeout, _stopCts.Token));
 
-                        if (resultTask != sending)
-                        {
-                            _aborted = true;
+                    if (resultTask != sending)
+                    {
+                        _aborted = true;
 
-                            // Abort the websocket if we're stuck in a pending send to the client
-                            socket.Abort();
-                        }
-                        else
-                        {
-                            // Cancel the timeout
-                            delayCts.Cancel();
-                        }
+                        // Abort the websocket if we're stuck in a pending send to the client
+                        socket.Abort();
                     }
                 }
                 else
@@ -258,7 +253,7 @@ private async Task StartReceiving(WebSocket socket)
                 {
 #if NETSTANDARD2_1 || NETCOREAPP
                     // Do a 0 byte read so that idle connections don't allocate a buffer when waiting for a read
-                    var result = await socket.ReceiveAsync(Memory<byte>.Empty, CancellationToken.None);
+                    var result = await socket.ReceiveAsync(Memory<byte>.Empty, _stopCts.Token);
 
                     if (result.MessageType == WebSocketMessageType.Close)
                     {
@@ -275,13 +270,13 @@ private async Task StartReceiving(WebSocket socket)
                     var memory = _application.Output.GetMemory();
 #if NETSTANDARD2_1 || NETCOREAPP
                     // Because we checked the CloseStatus from the 0 byte read above, we don't need to check again after reading
-                    var receiveResult = await socket.ReceiveAsync(memory, CancellationToken.None);
+                    var receiveResult = await socket.ReceiveAsync(memory, _stopCts.Token);
 #elif NETSTANDARD2_0 || NET461
                     var isArray = MemoryMarshal.TryGetArray<byte>(memory, out var arraySegment);
                     Debug.Assert(isArray);
 
                     // Exceptions are handled above where the send and receive tasks are being run.
-                    var receiveResult = await socket.ReceiveAsync(arraySegment, CancellationToken.None);
+                    var receiveResult = await socket.ReceiveAsync(arraySegment, _stopCts.Token);
 #else
 #error TFMs need to be updated
 #endif
@@ -400,7 +395,7 @@ private async Task StartSending(WebSocket socket)
                     try
                     {
                         // We're done sending, send the close frame to the client if the websocket is still open
-                        await socket.CloseOutputAsync(error != null ? WebSocketCloseStatus.InternalServerError : WebSocketCloseStatus.NormalClosure, "", CancellationToken.None);
+                        await socket.CloseOutputAsync(error != null ? WebSocketCloseStatus.InternalServerError : WebSocketCloseStatus.NormalClosure, "", _stopCts.Token);
                     }
                     catch (Exception ex)
                     {
@@ -452,6 +447,9 @@ public async Task StopAsync()
             // Cancel any pending reads from the application, this should start the entire shutdown process
             _application.Input.CancelPendingRead();
 
+            // Start ungraceful close timer
+            _stopCts.CancelAfter(_closeTimeout);
+
             try
             {
                 await Running;
@@ -465,6 +463,7 @@ public async Task StopAsync()
             finally
             {
                 _webSocket?.Dispose();
+                _stopCts.Dispose();
             }
 
             Log.TransportStopped(_logger, null);