Abort the websocket connection a Close() leaves open

Brokerages/WebSocketClientWrapper.cs

@@ -223,6 +223,7 @@ private void HandleConnection()
             const int maximumWaitTimeOnError = 120 * 1000;
             const int minimumWaitTimeOnError = 2 * 1000;
             var waitTimeOnError = minimumWaitTimeOnError;
+            ClientWebSocket client = null;
             while (_cts is { IsCancellationRequested: false })
             {
                 Log.Trace($"WebSocketClientWrapper.HandleConnection({_url}): Connecting...");
@@ -233,7 +234,7 @@ private void HandleConnection()
                         lock (_locker)
                         {
                             _client.DisposeSafely();
-                            _client = new ClientWebSocket();
+                            _client = client = new ClientWebSocket();
                             if (_sessionToken != null)
                             {
                                 _client.Options.SetRequestHeader("x-session-token", _sessionToken);
@@ -280,6 +281,10 @@ private void HandleConnection()
                     }
                 }
             }
+
+            // Close() cancels the receive, which aborts the connection, but a cancellation that lands before the first receive,
+            // while the Open handlers run, leaves the connection open after Close() returns
+            client?.Abort();
         }
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]

Tests/Brokerages/WebSocketClientWrapperTests.cs

@@ -0,0 +1,77 @@
+/*
+ * QUANTCONNECT.COM - Democratizing Finance, Empowering Individuals.
+ * Lean Algorithmic Trading Engine v2.0. Copyright 2014 QuantConnect Corporation.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+*/
+
+using System;
+using System.Net;
+using System.Net.Sockets;
+using System.Net.WebSockets;
+using System.Threading;
+using System.Threading.Tasks;
+using NUnit.Framework;
+using QuantConnect.Brokerages;
+
+namespace QuantConnect.Tests.Brokerages
+{
+    [TestFixture]
+    public class WebSocketClientWrapperTests
+    {
+        [Test]
+        public void CloseWhileTheOpenHandlerRunsClosesTheConnection()
+        {
+            var tcpListener = new TcpListener(IPAddress.Loopback, 0);
+            tcpListener.Start();
+            var port = ((IPEndPoint)tcpListener.LocalEndpoint).Port;
+            tcpListener.Stop();
+
+            using var server = new HttpListener();
+            server.Prefixes.Add($"http://127.0.0.1:{port}/");
+            server.Start();
+            var serverConnectionEnded = Task.Run(async () =>
+            {
+                var context = await server.GetContextAsync();
+                var webSocket = (await context.AcceptWebSocketAsync(null)).WebSocket;
+                try
+                {
+                    await webSocket.ReceiveAsync(new byte[1024], CancellationToken.None);
+                }
+                catch (WebSocketException)
+                {
+                    // the client aborted the connection
+                }
+            });
+
+            var webSocketClient = new WebSocketClientWrapper();
+            webSocketClient.Initialize($"ws://127.0.0.1:{port}/");
+            using var opened = new ManualResetEventSlim();
+            using var closing = new ManualResetEventSlim();
+            webSocketClient.Open += (_, _) =>
+            {
+                opened.Set();
+                // hold the open handler until Close() has cancelled the connection loop, before the first receive
+                closing.Wait();
+                Thread.Sleep(TimeSpan.FromSeconds(1));
+            };
+
+            webSocketClient.Connect();
+            Assert.IsTrue(opened.Wait(TimeSpan.FromSeconds(10)));
+
+            closing.Set();
+            webSocketClient.Close();
+
+            Assert.IsFalse(webSocketClient.IsOpen);
+            Assert.IsTrue(serverConnectionEnded.Wait(TimeSpan.FromSeconds(5)), "The connection is still open on the server after Close()");
+        }
+    }
+}