修复WebSocket连接断开后没有释放资源的BUG

Author: aaasoft

Quick.Protocol.WebSocket.Server.AspNetCore/QpWebSocketServer.cs

@@ -1,7 +1,6 @@
 using Microsoft.AspNetCore.Http;
 using Quick.Protocol.Utils;
 using System;
-using System.Collections;
 using System.Collections.Generic;
 using System.Threading;
 using System.Threading.Tasks;
@@ -17,29 +16,38 @@ private class WebSocketContext
         {
             public string ConnectionInfo { get; set; }
             public System.Net.WebSockets.WebSocket WebSocket { get; set; }
+            public CancellationTokenSource Cts { get; set; }
 
             public WebSocketContext(string connectionInfo, System.Net.WebSockets.WebSocket webSocket, CancellationTokenSource cts)
             {
                 ConnectionInfo = connectionInfo;
                 WebSocket = webSocket;
+                Cts = cts;
             }
         }
 
         public QpWebSocketServer(QpWebSocketServerOptions options) : base(options) { }
 
         public override void Start()
         {
-            isStarted=true;
+            isStarted = true;
             lock (webSocketContextQueue)
+            {
+                foreach (var webSocket in webSocketContextQueue)
+                    webSocket.Cts.Cancel();
                 webSocketContextQueue.Clear();
+            }
             base.Start();
         }
 
-        public Task OnNewConnection(System.Net.WebSockets.WebSocket webSocket, ConnectionInfo connectionInfo)
+        public async Task OnNewConnection(System.Net.WebSockets.WebSocket webSocket, ConnectionInfo connectionInfo)
         {
             //如果还没有开始接收，则直接关闭
             if (!isStarted)
-                return webSocket.CloseAsync(System.Net.WebSockets.WebSocketCloseStatus.NormalClosure, "", CancellationToken.None);
+            {
+                await webSocket.CloseAsync(System.Net.WebSockets.WebSocketCloseStatus.NormalClosure, "", CancellationToken.None);
+                return;
+            }
 
             var connectionInfoStr = $"WebSocket:{connectionInfo.RemoteIpAddress}:{connectionInfo.RemotePort}";
             var cts = new CancellationTokenSource();
@@ -49,7 +57,7 @@ public Task OnNewConnection(System.Net.WebSockets.WebSocket webSocket, Connectio
                         connectionInfoStr,
                         webSocket,
                         cts));
-            return Task.Delay(-1, cts.Token).ContinueWith(t =>
+            await Task.Delay(-1, cts.Token).ContinueWith(t =>
              {
                  if (LogUtils.LogConnection)
                      LogUtils.Log("[Connection]{0} disconnected.", connectionInfoStr);
@@ -58,9 +66,13 @@ public Task OnNewConnection(System.Net.WebSockets.WebSocket webSocket, Connectio
 
         public override void Stop()
         {
-            isStarted=false;
+            isStarted = false;
             lock (webSocketContextQueue)
+            {
+                foreach (var webSocket in webSocketContextQueue)
+                    webSocket.Cts.Cancel();
                 webSocketContextQueue.Clear();
+            }
             base.Stop();
         }
 
@@ -73,7 +85,7 @@ protected override async Task InnerAcceptAsync(CancellationToken token)
                 webSocketContextQueue.Clear();
             }
             //如果当前没有WebSocket连接，则等待0.1秒后再返回
-            if (webSocketContexts == null || webSocketContexts.Length==0)
+            if (webSocketContexts == null || webSocketContexts.Length == 0)
             {
                 await Task.Delay(100).ConfigureAwait(false);
                 return;
@@ -84,15 +96,19 @@ protected override async Task InnerAcceptAsync(CancellationToken token)
                 {
                     if (LogUtils.LogConnection)
                         LogUtils.Log("[Connection]{0} connected.", context.ConnectionInfo);
-                    OnNewChannelConnected(new WebSocketServerStream(context.WebSocket, token), context.ConnectionInfo, token);
+                    OnNewChannelConnected(new WebSocketServerStream(context.WebSocket, context.Cts), context.ConnectionInfo, token);
                 }
                 catch (Exception ex)
                 {
+                    context.Cts.Cancel();
                     if (LogUtils.LogConnection)
                         LogUtils.Log("[Connection]Init&Start Channel error,reason:{0}", ex.ToString());
-                    try { await context.WebSocket
+                    try
+                    {
+                        await context.WebSocket
                             .CloseAsync(System.Net.WebSockets.WebSocketCloseStatus.InternalServerError, ex.Message, CancellationToken.None)
-                            .ConfigureAwait(false); }
+                            .ConfigureAwait(false);
+                    }
                     catch { }
                 }
             }

Quick.Protocol.WebSocket.Server.AspNetCore/WebSocketServerStream.cs

@@ -1,8 +1,5 @@
 using System;
-using System.Collections.Generic;
 using System.IO;
-using System.Linq;
-using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
 
@@ -11,12 +8,12 @@ namespace Quick.Protocol.WebSocket.Server.AspNetCore
     internal class WebSocketServerStream : Stream
     {
         private System.Net.WebSockets.WebSocket webSocket;
-        private CancellationToken cancellationToken;
+        private CancellationTokenSource cts;
 
-        public WebSocketServerStream(System.Net.WebSockets.WebSocket webSocket, CancellationToken cancellationToken)
+        public WebSocketServerStream(System.Net.WebSockets.WebSocket webSocket, CancellationTokenSource cts)
         {
             this.webSocket = webSocket;
-            this.cancellationToken = cancellationToken;
+            this.cts = cts;
         }
 
         public override bool CanSeek => throw new NotImplementedException();
@@ -33,7 +30,7 @@ public override void Flush() { }
 
         public override int Read(byte[] buffer, int offset, int count)
         {
-            var result = webSocket.ReceiveAsync(new ArraySegment<byte>(buffer, offset, count), cancellationToken).Result;
+            var result = webSocket.ReceiveAsync(new ArraySegment<byte>(buffer, offset, count), CancellationToken.None).Result;
             return result.Count;
         }
 
@@ -50,7 +47,7 @@ public override void Write(byte[] buffer, int offset, int count)
                 new ArraySegment<byte>(buffer, offset, count),
                 System.Net.WebSockets.WebSocketMessageType.Binary,
                 true,
-                cancellationToken)
+                CancellationToken.None)
                 .Wait();
         }
 
@@ -66,6 +63,7 @@ public override Task WriteAsync(byte[] buffer, int offset, int count, Cancellati
 
         protected override void Dispose(bool disposing)
         {
+            cts.Dispose();
             webSocket.Dispose();
             base.Dispose(disposing);
         }