Add ability to reuse an external RecyclableMemoryStreamManager

Marfusios · Apr 9, 2024 · 1a01b60 · 1a01b60
1 parent c002823
commit 1a01b60
Showing 1 changed file with 15 additions and 7 deletions.
diff --git a/src/Websocket.Client/WebsocketClient.cs b/src/Websocket.Client/WebsocketClient.cs
@@ -22,8 +22,12 @@ public partial class WebsocketClient : IWebsocketClient
         private readonly ILogger<WebsocketClient> _logger;
         private readonly WebsocketAsyncLock _locker = new WebsocketAsyncLock();
         private readonly Func<Uri, CancellationToken, Task<WebSocket>> _connectionFactory;
-        private static readonly RecyclableMemoryStreamManager _memoryStreamManager = new RecyclableMemoryStreamManager();
+        private readonly RecyclableMemoryStreamManager? _memoryStreamManager;
+        private static readonly Lazy<RecyclableMemoryStreamManager> _globalMemoryStreamManager = new Lazy<RecyclableMemoryStreamManager>(new RecyclableMemoryStreamManager());
 
+        private RecyclableMemoryStreamManager MemoryStreamManager =>
+            _memoryStreamManager ?? _globalMemoryStreamManager.Value;
+
         private Uri _url;
         private Timer? _lastChanceTimer;
         private DateTime _lastReceivedMsg = DateTime.UtcNow;
@@ -47,8 +51,9 @@ public partial class WebsocketClient : IWebsocketClient
         /// </summary>
         /// <param name="url">Target websocket url (wss://)</param>
         /// <param name="clientFactory">Optional factory for native ClientWebSocket, use it whenever you need some custom features (proxy, settings, etc)</param>
-        public WebsocketClient(Uri url, Func<ClientWebSocket>? clientFactory = null)
-            : this(url, null, GetClientFactory(clientFactory))
+        /// <param name="memoryStreamManager">Optional memory stream manager that can be reused on multiple occasions</param>
+        public WebsocketClient(Uri url, Func<ClientWebSocket>? clientFactory = null, RecyclableMemoryStreamManager? memoryStreamManager = null)
+            : this(url, null, GetClientFactory(clientFactory), memoryStreamManager)
         {
         }
 
@@ -58,8 +63,9 @@ public WebsocketClient(Uri url, Func<ClientWebSocket>? clientFactory = null)
         /// <param name="url">Target websocket url (wss://)</param>
         /// <param name="logger">Logger instance, can be null</param>
         /// <param name="clientFactory">Optional factory for native ClientWebSocket, use it whenever you need some custom features (proxy, settings, etc)</param>
-        public WebsocketClient(Uri url, ILogger<WebsocketClient>? logger, Func<ClientWebSocket>? clientFactory = null)
-            : this(url, logger, GetClientFactory(clientFactory))
+        /// <param name="memoryStreamManager">Optional memory stream manager that can be reused on multiple occasions</param>
+        public WebsocketClient(Uri url, ILogger<WebsocketClient>? logger, Func<ClientWebSocket>? clientFactory = null, RecyclableMemoryStreamManager? memoryStreamManager = null)
+            : this(url, logger, GetClientFactory(clientFactory), memoryStreamManager)
         {
         }
 
@@ -69,7 +75,8 @@ public WebsocketClient(Uri url, ILogger<WebsocketClient>? logger, Func<ClientWeb
         /// <param name="url">Target websocket url (wss://)</param>
         /// <param name="logger">Logger instance, can be null</param>
         /// <param name="connectionFactory">Optional factory for native creating and connecting to a websocket. The method should return a <see cref="WebSocket"/> which is connected. Use it whenever you need some custom features (proxy, settings, etc)</param>
-        public WebsocketClient(Uri url, ILogger<WebsocketClient>? logger, Func<Uri, CancellationToken, Task<WebSocket>>? connectionFactory)
+        /// <param name="memoryStreamManager">Optional memory stream manager that can be reused on multiple occasions</param>
+        public WebsocketClient(Uri url, ILogger<WebsocketClient>? logger, Func<Uri, CancellationToken, Task<WebSocket>>? connectionFactory, RecyclableMemoryStreamManager? memoryStreamManager = null)
         {
             Validations.Validations.ValidateInput(url, nameof(url));
 
@@ -85,6 +92,7 @@ public WebsocketClient(Uri url, ILogger<WebsocketClient>? logger, Func<Uri, Canc
                 await client.ConnectAsync(uri, token).ConfigureAwait(false);
                 return client;
             });
+            _memoryStreamManager = memoryStreamManager;
         }
 
         /// <inheritdoc />
@@ -458,7 +466,7 @@ private async Task Listen(WebSocket client, CancellationToken token)
                 do
                 {
                     ValueWebSocketReceiveResult result;
-                    var ms = (RecyclableMemoryStream)_memoryStreamManager.GetStream();
+                    var ms = MemoryStreamManager.GetStream();
 
                     while (true)
                     {