added cancellation token for websocket send async method

kerryjiang · Apr 15, 2024 · 4cea48c · 4cea48c
1 parent 6e33402
commit 4cea48c
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 19 deletions.
diff --git a/src/SuperSocket.WebSocket.Server/WebSocketPackageHandler.cs b/src/SuperSocket.WebSocket.Server/WebSocketPackageHandler.cs
@@ -129,7 +129,7 @@ public async ValueTask Handle(IAppSession session, WebSocketPackage package, Can
 
                     try
                     {
-                        await websocketSession.SendAsync(package);
+                        await websocketSession.SendAsync(package, cancellationToken);
                     }
                     catch (InvalidOperationException)
                     {
@@ -146,7 +146,7 @@ public async ValueTask Handle(IAppSession session, WebSocketPackage package, Can
             else if (package.OpCode == OpCode.Ping)
             {
                 package.OpCode = OpCode.Pong;
-                await websocketSession.SendAsync(package);
+                await websocketSession.SendAsync(package, cancellationToken);
                 return;
             }
             else if (package.OpCode == OpCode.Pong)

diff --git a/src/SuperSocket.WebSocket.Server/WebSocketSession.cs b/src/SuperSocket.WebSocket.Server/WebSocketSession.cs
@@ -2,6 +2,7 @@
 using System.Buffers;
 using System.Collections.Specialized;
 using System.Text;
+using System.Threading;
 using System.Threading.Tasks;
 using SuperSocket.ProtoBase;
 using SuperSocket.Server;
@@ -33,30 +34,32 @@ public string Path
 
         internal IPackageEncoder<WebSocketPackage> MessageEncoder { get; set; }
 
-        public virtual ValueTask SendAsync(WebSocketPackage message)
+        public virtual ValueTask SendAsync(WebSocketPackage message, CancellationToken cancellationToken = default)
         {
-            return this.Connection.SendAsync(MessageEncoder, message);
+            return this.Connection.SendAsync(MessageEncoder, message, cancellationToken);
         }
 
-        public virtual ValueTask SendAsync(string message)
+        public virtual ValueTask SendAsync(string message, CancellationToken cancellationToken = default)
         {
             return SendAsync(new WebSocketPackage
-            {
-                OpCode = OpCode.Text,
-                Message = message,
-            });
+                {
+                    OpCode = OpCode.Text,
+                    Message = message,
+                },
+                cancellationToken);
         }
 
-        public virtual ValueTask SendAsync(ReadOnlyMemory<byte> data)
+        public virtual ValueTask SendAsync(ReadOnlyMemory<byte> data, CancellationToken cancellationToken = default)
         {
             return SendAsync(new WebSocketPackage
-            {
-                OpCode = OpCode.Binary,
-                Data = new ReadOnlySequence<byte>(data),
-            });
+                {
+                    OpCode = OpCode.Binary,
+                    Data = new ReadOnlySequence<byte>(data),
+                },
+                cancellationToken);
         }
 
-        public ValueTask CloseAsync(CloseReason reason, string reasonText = null)
+        public ValueTask CloseAsync(CloseReason reason, string reasonText = null, CancellationToken cancellationToken = default)
         {
             var closeReasonCode = (short)reason;
 
@@ -90,10 +93,11 @@ public ValueTask CloseAsync(CloseReason reason, string reasonText = null)
             OnCloseHandshakeStarted();
 
             return SendAsync(new WebSocketPackage
-            {
-                OpCode = OpCode.Close,
-                Data = new ReadOnlySequence<byte>(buffer, 0, length)
-            });
+                {
+                    OpCode = OpCode.Close,
+                    Data = new ReadOnlySequence<byte>(buffer, 0, length)
+                },
+                cancellationToken);
         }
 
         private void OnCloseHandshakeStarted()