Lifetimemanager, message holder

Author: gedem

src/NetCoreStack.WebSockets/ConnectionManager.cs

@@ -24,19 +24,22 @@ public class ConnectionManager : IConnectionManager
         private readonly IHandshakeStateTransport _initState;
         private readonly ILoggerFactory _loggerFactory;
         private readonly IStreamCompressor _compressor;
+        private readonly TransportLifetimeManager _lifetimeManager;
         public ConcurrentDictionary<string, WebSocketTransport> Connections { get; }
 
         public ConnectionManager(IStreamCompressor compressor,
+            TransportLifetimeManager lifetimeManager,
             InvocatorRegistry invocatorRegistry,
             IOptions<ServerSocketsOptions> options,
             IHandshakeStateTransport initState,
             ILoggerFactory loggerFactory)
         {
+            _compressor = compressor;
+            _lifetimeManager = lifetimeManager;
             _invocatorRegistry = invocatorRegistry;
             _options = options.Value;
             _initState = initState;
-            _loggerFactory = loggerFactory;
-            _compressor = compressor;
+            _loggerFactory = loggerFactory;            
             Connections = new ConcurrentDictionary<string, WebSocketTransport>(StringComparer.OrdinalIgnoreCase);
         }
 
@@ -89,10 +92,18 @@ private async Task SendAsync(WebSocketTransport transport, WebSocketMessageDescr
                 throw new ArgumentNullException(nameof(descriptor.Segments));
             }
 
-            await transport.WebSocket.SendAsync(descriptor.Segments, 
-                descriptor.MessageType, 
-                descriptor.EndOfMessage, 
-                CancellationToken.None);
+            if (!transport.WebSocket.CloseStatus.HasValue)
+            {
+                await transport.WebSocket.SendAsync(descriptor.Segments,
+                   descriptor.MessageType,
+                   descriptor.EndOfMessage,
+                   CancellationToken.None);
+            }
+            else
+            {
+                // Only text messages
+                _lifetimeManager.AddQueue(new MessageHolder { ConnectionId = transport.ConnectionId, Segments = descriptor.Segments });
+            }
         }
 
         private async Task SendBinaryAsync(WebSocketTransport transport, byte[] chunkedBytes, bool endOfMessage, CancellationToken token)
@@ -104,10 +115,13 @@ private async Task SendBinaryAsync(WebSocketTransport transport, byte[] chunkedB
 
             var segments = new ArraySegment<byte>(chunkedBytes);
 
-            await transport.WebSocket.SendAsync(segments,
-                           WebSocketMessageType.Binary,
-                           endOfMessage,
-                           token);
+            if (!transport.WebSocket.CloseStatus.HasValue)
+            {
+                await transport.WebSocket.SendAsync(segments,
+                   WebSocketMessageType.Binary,
+                   endOfMessage,
+                   token);
+            }
         }
 
         public async Task ConnectAsync(WebSocket webSocket, string connectorName = "")

src/NetCoreStack.WebSockets/Extensions/SocketServiceCollectionExtensions.cs

@@ -21,6 +21,7 @@ public static void AddNativeWebSockets(this IServiceCollection services, Action<
                 throw new ArgumentNullException(nameof(setup));
             }
 
+            services.AddSingleton<TransportLifetimeManager>();
             services.TryAdd(ServiceDescriptor.Singleton<ILoggerFactory, LoggerFactory>());
             services.TryAdd(ServiceDescriptor.Singleton<IStreamCompressor, GZipStreamCompressor>());
             services.TryAdd(ServiceDescriptor.Transient<IHandshakeStateTransport, DefaultHandshakeStateTransport>());

src/NetCoreStack.WebSockets/Internal/MessageHolder.cs

@@ -0,0 +1,10 @@
+using System;
+
+namespace NetCoreStack.WebSockets.Internal
+{
+    public class MessageHolder
+    {
+        public string ConnectionId { get; set; }
+        public ArraySegment<byte> Segments { get; set; }
+    }
+}

src/NetCoreStack.WebSockets/Internal/TransportLifetimeManager.cs

@@ -0,0 +1,19 @@
+using System.Collections.Concurrent;
+
+namespace NetCoreStack.WebSockets.Internal
+{
+    public class TransportLifetimeManager
+    {
+        public ConcurrentQueue<MessageHolder> Queue { get; }
+
+        public TransportLifetimeManager()
+        {
+            Queue = new ConcurrentQueue<MessageHolder>();
+        }
+
+        public void AddQueue(MessageHolder holder)
+        {
+            Queue.Enqueue(holder);
+        }
+    }
+}

src/NetCoreStack.WebSockets/Internal/WebSocketMessageDescriptor.cs

@@ -16,4 +16,4 @@ public WebSocketMessageDescriptor()
             CancellationToken = CancellationToken.None;
         }
     }
-}
+}