Export binary memory on ResponseMessage

Marfusios · Aug 8, 2023 · 3db795e · 3db795e
1 parent 638cf7f
commit 3db795e
Show file tree

Hide file tree

Showing 3 changed files with 56 additions and 2 deletions.
diff --git a/src/Websocket.Client/ResponseMessage.cs b/src/Websocket.Client/ResponseMessage.cs
@@ -1,4 +1,5 @@
-using System.Net.WebSockets;
+using System;
+using System.Net.WebSockets;
 
 namespace Websocket.Client
 {
@@ -7,22 +8,43 @@ namespace Websocket.Client
     /// </summary>
     public class ResponseMessage
     {
+#if NETSTANDARD2_1_OR_GREATER || NET5_0_OR_GREATER
+        private ResponseMessage(ref Memory<byte> binaryMemory, string text, WebSocketMessageType messageType)
+        {
+            BinaryMemory = binaryMemory;
+            Text = text;
+            MessageType = messageType;
+        }
+#else
         private ResponseMessage(byte[] binary, string text, WebSocketMessageType messageType)
         {
             Binary = binary;
             Text = text;
             MessageType = messageType;
         }
+#endif
 
         /// <summary>
         /// Received text message (only if type = WebSocketMessageType.Text)
         /// </summary>
         public string Text { get; }
 
+#if NETSTANDARD2_1_OR_GREATER || NET5_0_OR_GREATER
+        /// <summary>
+        /// Received text message (only if type = WebSocketMessageType.Binary)
+        /// </summary>
+        public byte[] Binary => BinaryMemory.ToArray();
+
+        /// <summary>
+        /// Memory representation of the binary message
+        /// </summary>
+        public ReadOnlyMemory<byte> BinaryMemory { get; }
+#else
         /// <summary>
         /// Received text message (only if type = WebSocketMessageType.Binary)
         /// </summary>
         public byte[] Binary { get; }
+#endif
 
         /// <summary>
         /// Current message type (Text or Binary)
@@ -42,6 +64,24 @@ public override string ToString()
             return $"Type binary, length: {Binary?.Length}";
         }
 
+#if NETSTANDARD2_1_OR_GREATER || NET5_0_OR_GREATER
+        /// <summary>
+        /// Create text response message
+        /// </summary>
+        public static ResponseMessage TextMessage(string data)
+        {
+            var emptyMemory = Memory<byte>.Empty;
+            return new ResponseMessage(ref emptyMemory, data, WebSocketMessageType.Text);
+        }
+
+        /// <summary>s
+        /// Create binary response message
+        /// </summary>
+        public static ResponseMessage BinaryMessage(ref Memory<byte> data)
+        {
+            return new ResponseMessage(ref data, null, WebSocketMessageType.Binary);
+        }
+#else
         /// <summary>
         /// Create text response message
         /// </summary>
@@ -57,5 +97,6 @@ public static ResponseMessage BinaryMessage(byte[] data)
         {
             return new ResponseMessage(data, null, WebSocketMessageType.Binary);
         }
+#endif
     }
 }
diff --git a/src/Websocket.Client/Websocket.Client.csproj b/src/Websocket.Client/Websocket.Client.csproj
@@ -31,6 +31,7 @@
 			<PrivateAssets>all</PrivateAssets>
 			<IncludeAssets>runtime; build; native; contentfiles; analyzers</IncludeAssets>
 		</PackageReference>
+		<PackageReference Include="Microsoft.IO.RecyclableMemoryStream" Version="2.3.2" />
 		<PackageReference Include="System.Reactive" Version="5.0.0" />
 		<PackageReference Include="System.Threading.Channels" Version="5.0.0" />
 	</ItemGroup>

diff --git a/src/Websocket.Client/WebsocketClient.cs b/src/Websocket.Client/WebsocketClient.cs
@@ -1,3 +1,4 @@
+using Microsoft.IO;
 using System;
 using System.Diagnostics;
 using System.IO;
@@ -24,6 +25,7 @@ public partial class WebsocketClient : IWebsocketClient
 
         private readonly WebsocketAsyncLock _locker = new WebsocketAsyncLock();
         private readonly Func<Uri, CancellationToken, Task<WebSocket>> _connectionFactory;
+        private static readonly RecyclableMemoryStreamManager _memoryStreamManager = new RecyclableMemoryStreamManager();
 
         private Uri _url;
         private Timer _lastChanceTimer;
@@ -450,7 +452,7 @@ private async Task Listen(WebSocket client, CancellationToken token)
                             if (ms == null)
                             {
                                 // create memory stream and insert first chunk
-                                ms = new MemoryStream();
+                                ms = _memoryStreamManager.GetStream();
                                 ms.Write(resultArrayWithTrailing, 0, resultArraySize);
                             }
 
@@ -522,12 +524,22 @@ await StopInternal(client, WebSocketCloseStatus.NormalClosure, "Closing",
                     {
                         if (ms != null)
                         {
+#if NETSTANDARD2_1_OR_GREATER || NET5_0_OR_GREATER
+                            var binaryMemory = ms.ToArray().AsMemory();
+                            message = ResponseMessage.BinaryMessage(ref binaryMemory);
+#else
                             message = ResponseMessage.BinaryMessage(ms.ToArray());
+#endif
                         }
                         else
                         {
                             Array.Resize(ref resultArrayWithTrailing, resultArraySize);
+#if NETSTANDARD2_1_OR_GREATER || NET5_0_OR_GREATER
+                            var binaryMemory = resultArrayWithTrailing.AsMemory();
+                            message = ResponseMessage.BinaryMessage(ref binaryMemory);
+#else
                             message = ResponseMessage.BinaryMessage(resultArrayWithTrailing);
+#endif
                         }
                     }