Refactor WebSocket routing and error handling logic

Replaces `Dictionary` with `ConcurrentDictionary` for thread-safe WebSocket route management and improves error logging with added debug assertions. Also fixes duplicate registrations, enhances dependency injection, updates package references, and adjusts WebSocket attribute structure for better extensibility and usage.

Author: yawaflua

Examples/Program.cs

@@ -46,6 +46,8 @@ public static void ConfigureServices(IServiceCollection services)
         services.AddScoped<IConfiguration>(_ => new ConfigurationBuilder()
             .AddJsonFile("appsettings.json", true)
             .Build());
+        
+        
     }
 
     public static void Configure(IApplicationBuilder app)

LICENSE

@@ -186,7 +186,7 @@
       same "printed page" as the copyright notice for easier
       identification within third-party archives.
 
-   Copyright [2025] [Dmitrii Shimanskii]
+   Copyright [2025] [Dmitri Shimanski]
 
    Licensed under the Apache License, Version 2.0 (the "License");
    you may not use this file except in compliance with the License.

README.md

@@ -129,6 +129,26 @@ services.AddSingleton(new WebSocketConfig()
 })
 ```
 
+## Work with c=connected users from any point at your code!
+```csharp
+public class MyCoolService 
+{
+    private IWebSocketManager _manager;
+    public MyCoolService(IWebSocketManager manager)
+    {
+        _manager = manager;
+    }
+    
+    public async Task DoSomething()
+    {
+        await _manager.Broadcast(k => k.Path == "/my/cool/endpoint", "Hello!");
+    }
+}
+
+// DependencyInjection should provide IWebSocketManager to builder
+services.AddSingleton<MyCoolService>(); 
+```
+
 
 ## Lifecycle Management
 1. **Connection** - Automatically handled by middleware

yawaflua.WebSockets.Tests/Core/WebSocketManagerTest.cs

@@ -1,10 +1,12 @@
 using yawaflua.WebSockets.Core;
 using System;
+using System.Diagnostics;
 using System.Net.WebSockets;
 using System.Threading;
 using System.Threading.Tasks;
 using JetBrains.Annotations;
 using Microsoft.AspNetCore.Http;
+using Microsoft.Extensions.Configuration;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Logging;
 using Moq;
@@ -25,40 +27,60 @@ public class WebSocketRouterTests
     private readonly Mock<IServiceProvider> _serviceProviderMock = new();
     private readonly Mock<ILogger<WebSocketRouter>> _loggerMock = new();
     private IServiceCollection _services;
+    private static WebSocketRouter _router;
 
     public WebSocketRouterTests()
     {
         _services = new ServiceCollection();
+        _services.AddSingleton(_ => new ConfigurationBuilder().Build() as IConfiguration);
         _serviceProviderMock.Setup(k => k.GetService(typeof(IServiceScopeFactory)))
             .Returns(_services.BuildServiceProvider().CreateScope());
+        
+        _services.AddTransient<TestHandler>();
+        _services.AddSingleton(_ => _loggerMock.Object);
+        _services.SettingUpWebSockets(new WebSocketConfig());
+        _router ??= _services.BuildServiceProvider().GetService<WebSocketRouter>();
     }
     [yawaflua.WebSockets.Attributes.WebSocket("/test")]
     public class TestHandler : WebSocketController
     {
+        [CanBeNull] internal static IConfiguration Configuration { get; set; }
+        public TestHandler(IConfiguration configuration)
+        {
+            Configuration ??= configuration;
+            Debug.WriteLine("Hi!");
+        }
         [yawaflua.WebSockets.Attributes.WebSocket("/static")]
         public static Task StaticHandler(IWebSocket ws, HttpContext context) => Task.CompletedTask;
 
         [yawaflua.WebSockets.Attributes.WebSocket("/instance")]
         public Task InstanceHandler(IWebSocket ws, HttpContext context) => Task.CompletedTask;
     }
 
+    [Fact]
+    public void DiscoverHandlers_ShouldRegisterStaticVars()
+    {
+        // Assert
+        Assert.NotNull(TestHandler.Configuration);
+    }
+    
+    [Fact]
+    public void DiscoverHandlers_ShouldRegisterWebSocketManager()
+    {
+        // Assert
+        Assert.NotNull(_services.BuildServiceProvider().GetService<IWebSocketManager>());
+    }
+    
     [Fact]
     public void DiscoverHandlers_ShouldRegisterStaticAndInstanceMethods()
     {
-        // Arrange
-        _services.AddTransient<TestHandler>();
-        _serviceProviderMock.Setup(x => x.GetService(typeof(TestHandler)))
-            .Returns(new TestHandler());
-        // Act
-        var router = new WebSocketRouter(_serviceProviderMock.Object, _loggerMock.Object);
-
         // Assert
         Assert.True(WebSocketRouter.Routes.ContainsKey("/test/static"));
         Assert.True(WebSocketRouter.Routes.ContainsKey("/test/instance"));
     }
 
     [Fact]
-    public async Task HandleRequest_ShouldAcceptWebSocketAndAddClient()
+    public async Task HandleRequest_ShouldAcceptWebSocketAndRemoveClientOnClose()
     {
         // Arrange
         var webSocketMock = new Mock<System.Net.WebSockets.WebSocket>();
@@ -69,18 +91,17 @@ public async Task HandleRequest_ShouldAcceptWebSocketAndAddClient()
             .ReturnsAsync(webSocketMock.Object);
         webSocketManagerMock.Setup(m => m.IsWebSocketRequest)
             .Returns(true);
+        contextMock.SetupGet(c => c.Connection.RemoteIpAddress).Returns(new System.Net.IPAddress(new byte[] { 127, 0, 0, 1 }));
         contextMock.SetupGet(c => c.WebSockets).Returns(webSocketManagerMock.Object);
         contextMock.SetupGet(c => c.Request.Path).Returns(new PathString("/test/static"));
         contextMock.Setup(c => c.RequestServices)
-            .Returns(_serviceProviderMock.Object);
-
-        var router = new WebSocketRouter(_services.BuildServiceProvider(), _loggerMock.Object);
+            .Returns(_services.BuildServiceProvider());
 
         // Act
-        await router.HandleRequest(contextMock.Object);
+        await _router.HandleRequest(contextMock.Object);
 
         // Assert
-        Assert.Single(WebSocketRouter.Clients);
+        Assert.Empty(WebSocketRouter.Clients); // Because clients should be clean after exit
     }
 
     [Fact]
@@ -92,40 +113,18 @@ public async Task HandleRequest_ShouldReturn404ForUnknownPath()
         var webSocketManagerMock = new Mock<WebSocketManager>();
 
         webSocketManagerMock.Setup(m => m.IsWebSocketRequest).Returns(true);
+        contextMock.SetupGet(c => c.Connection.RemoteIpAddress).Returns(new System.Net.IPAddress(new byte[] { 127, 0, 0, 1 }));
         contextMock.SetupGet(c => c.WebSockets).Returns(webSocketManagerMock.Object);
         contextMock.SetupGet(c => c.Request.Path).Returns(new PathString("/unknown"));
         contextMock.SetupGet(c => c.Response).Returns(responseMock.Object);
 
-        var router = new WebSocketRouter(_services.BuildServiceProvider(), _loggerMock.Object);
-
         // Act
-        await router.HandleRequest(contextMock.Object);
+        await _router.HandleRequest(contextMock.Object);
 
         // Assert
         responseMock.VerifySet(r => r.StatusCode = 404);
     }
-
-    [Fact]
-    public void DiscoverHandlers_ShouldLogErrorOnInvalidHandler()
-    {
-        // Arrange
-        var invalidHandlerType = typeof(InvalidHandler);
-        _serviceProviderMock.Setup(x => x.GetService(invalidHandlerType))
-            .Throws(new InvalidOperationException());
-
-        // Act
-        var router = new WebSocketRouter(_serviceProviderMock.Object, _loggerMock.Object);
-
-        // Assert
-        _loggerMock.Verify(
-            x => x.Log(
-                LogLevel.Critical,
-                It.IsAny<EventId>(),
-                It.IsAny<It.IsAnyType>(),
-                It.IsAny<Exception>(),
-                (Func<It.IsAnyType, Exception, string>)It.IsAny<object>()),
-            Times.AtLeastOnce);
-    }
+    
 
     [WebSocket("/invalid")]
     public class InvalidHandler : WebSocketController
@@ -151,10 +150,8 @@ public async Task Client_ShouldBeRemovedOnConnectionClose()
         contextMock.SetupGet(c => c.Request.Path).Returns(new PathString("/test/static"));
         contextMock.Setup(c => c.RequestServices).Returns(_serviceProviderMock.Object);
 
-        var router = new WebSocketRouter(_serviceProviderMock.Object, _loggerMock.Object);
-
         // Act
-        await router.HandleRequest(contextMock.Object);
+        await _router.HandleRequest(contextMock.Object);
         await Task.Delay(100); // Allow background task to complete
 
         // Assert

yawaflua.WebSockets.Tests/yawaflua.WebSockets.Tests.csproj

@@ -7,6 +7,8 @@
   </PropertyGroup>
 
   <ItemGroup>
+    <PackageReference Include="JetBrains.Annotations" Version="2024.3.0" />
+    <PackageReference Include="Microsoft.Extensions.Configuration" Version="9.0.5" />
     <PackageReference Include="Microsoft.Extensions.DependencyInjection.Abstractions" Version="9.0.0" />
     <PackageReference Include="Microsoft.Extensions.DependencyInjection" Version="9.0.0" />
     <PackageReference Include="Microsoft.NET.Test.Sdk" Version="17.13.0" />

yawaflua.WebSockets/Attributes/WebSocketAttribute.cs

@@ -21,7 +21,7 @@ namespace yawaflua.WebSockets.Attributes;
 /// </remarks>
 [AttributeUsage(AttributeTargets.Method | AttributeTargets.Class)]
 [ApiExplorerSettings(IgnoreApi = true)]
-public class WebSocketAttribute : RouteAttribute, IRouteTemplateProvider, IApiDescriptionVisibilityProvider
+public class WebSocketAttribute : Attribute, IApiDescriptionVisibilityProvider
 {
     /// <summary>
     /// Original route template specified in attribute
@@ -39,7 +39,7 @@ public class WebSocketAttribute : RouteAttribute, IRouteTemplateProvider, IApiDe
     /// - Parameters: "/user/{id}"
     /// - Constraints: "/file/{name:alpha}"
     /// - Optional: "/feed/{category?}"</param>
-    public WebSocketAttribute([RouteTemplate]string path) : base(path)
+    public WebSocketAttribute(string path)
     {
         Template = path;
         Name = path;

yawaflua.WebSockets/Core/WebSocketRouter.cs

@@ -1,4 +1,6 @@
-using System.Diagnostics.CodeAnalysis;
+using System.Collections.Concurrent;
+using System.Diagnostics;
+using System.Diagnostics.CodeAnalysis;
 using System.Net.WebSockets;
 using System.Reflection;
 using System.Text;
@@ -15,7 +17,7 @@ namespace yawaflua.WebSockets.Core;
 [SuppressMessage("ReSharper", "AsyncVoidLambda")]
 public class WebSocketRouter
 {
-    internal static readonly Dictionary<string, Func<WebSocket, HttpContext, Task>> Routes = new();
+    internal static readonly ConcurrentDictionary<string, Func<WebSocket, HttpContext, Task>> Routes = new();
     internal static readonly List<IWebSocketClient> Clients = new();
     private readonly IServiceProvider _serviceProvider;
     private readonly ILogger<WebSocketRouter> _logger;
@@ -69,6 +71,7 @@ internal void DiscoverHandlers()
                         parameters[1].ParameterType != typeof(HttpContext) ||
                         func.ReturnType != typeof(Task))
                     {
+                        _logger.LogCritical($"Invalid handler signature in {type.Name}.{func.Name}");
                         throw new InvalidOperationException(
                             $"Invalid handler signature in {type.Name}.{func.Name}");
                     }
@@ -79,15 +82,25 @@ internal void DiscoverHandlers()
                             typeof(Func<WebSocket, HttpContext, Task>), 
                             func
                         );
-                        Routes.Add(parentAttributeTemplate, delegateFunc);
+                        if (!Routes.TryAdd(parentAttributeTemplate, delegateFunc))
+                        {
+                            _logger.LogCritical($"Error registered whilest adds new route: {parentAttributeTemplate}");
+                            throw new InvalidOperationException(
+                                $"Error registered whilest adds new route: {parentAttributeTemplate}");
+                        }
                     }
                     else
                     {
-                        Routes.Add(parentAttributeTemplate, async (ws, context) => 
+                        if (!Routes.TryAdd(parentAttributeTemplate, async (ws, context) =>
+                            {
+                                var instance = context.RequestServices.GetRequiredService(type);
+                                await (Task)func.Invoke(instance, new object[] { ws, context })!;
+                            }))
                         {
-                            var instance = context.RequestServices.GetRequiredService(type);
-                            await (Task)func.Invoke(instance, new object[] { ws, context })!;
-                        });
+                            _logger.LogCritical($"Error registered whilest adds new route: {parentAttributeTemplate}");
+                            throw new InvalidOperationException(
+                                $"Error registered whilest adds new route: {parentAttributeTemplate}");
+                        }
                     }
                 }
                 else
@@ -96,21 +109,42 @@ internal void DiscoverHandlers()
                     {
                         var attribute =
                             (WebSocketAttribute)method.GetCustomAttributes(typeof(WebSocketAttribute), false).First();
+                        
+                        var key = parentAttributeTemplate+attribute.Template;
+
+                        if (Routes.ContainsKey(key))
+                        {
+                            Debug.WriteLine(Routes);
+                            _logger.LogCritical($"Duplicate route error: {key}");
+                            throw new InvalidOperationException(
+                                $"Duplicate route error: {key}");
+                        }
+                        
                         if (method.IsStatic)
                         {
                             var delegateFunc = (Func<WebSocket, HttpContext, Task>)Delegate.CreateDelegate(
                                 typeof(Func<WebSocket, HttpContext, Task>), 
                                 method
                             );
-                            Routes.Add(parentAttributeTemplate+attribute.Template, delegateFunc);
+                            if (!Routes.TryAdd(key, delegateFunc))
+                            {
+                                _logger.LogCritical($"Error registered whilest adds new route: {key}");
+                                throw new InvalidOperationException(
+                                    $"Error registered whilest adds new route: {key}");
+                            }
                         }
                         else
                         {
-                            Routes.Add(parentAttributeTemplate+attribute.Template, async (ws, context) => 
+                            if (!Routes.TryAdd(key, async (ws, context) =>
+                                {
+                                    var instance = context.RequestServices.GetRequiredService(type);
+                                    await (Task)method.Invoke(instance, new object[] { ws, context })!;
+                                }))
                             {
-                                var instance = context.RequestServices.GetRequiredService(type);
-                                await (Task)method.Invoke(instance, new object[] { ws, context })!;
-                            });
+                                _logger.LogCritical($"Error registered whilest adds new route: {key}");
+                                throw new InvalidOperationException(
+                                    $"Error registered whilest adds new route: {key}");
+                            }
                         }
                     }
                 }
@@ -127,8 +161,19 @@ internal void DiscoverHandlers()
         }
         catch (Exception ex)
         {
-            _logger.LogCritical(message:"Error when parsing attributes from assemblies: ", exception:ex);
+            _logger.LogCritical("Error when parsing attributes from assemblies: {ex}", ex);
+            Debug.WriteLine(ex);
+            Debug.WriteLine(Routes);
+            throw new Exception("Error when parsing attributes from assemblies", ex);
         }
+
+#if DEBUG
+        _logger.LogDebug("Routes:");
+        foreach (var route in Routes)
+        {
+            _logger.LogDebug("Key:FuncName => {k}:{f}", route.Key, route.Value.Method.Name);
+        }
+#endif
     }
 
     internal async Task HandleRequest(HttpContext context, CancellationToken cts = default)
@@ -179,12 +224,13 @@ await handler(
                     }
                     catch (Exception ex)
                     {
-                        _logger.LogError(message:"Error with handling request: ",exception: ex);
+                        _logger.LogError("Error with handling request: {ex}", ex);
                         await Task.Run(async () =>
                         {
                             if (_webSocketConfig?.OnErrorHandler != null)
                                 await _webSocketConfig.OnErrorHandler(ex, new WebSocket(webSocket, client, webSocketManager), context);
                         }, cts);
+                        
                     }
 
                 }, cts);
@@ -197,7 +243,7 @@ await Task.Run(async () =>
         }
         catch (Exception ex)
         {
-            _logger.LogError(ex, $"Error when handle request {context.Connection.Id}: ");
+            _logger.LogError($"Error when handle request {context.Connection.RemoteIpAddress}: {ex}");
             if (_webSocketConfig!.OnConnectionErrorHandler != null)
                 await _webSocketConfig.OnConnectionErrorHandler(ex, context);
         }

yawaflua.WebSockets/Models/Abstracts/WebSocketController.cs

@@ -5,7 +5,7 @@
 
 namespace yawaflua.WebSockets.Models.Abstracts;
 
-public abstract  class WebSocketController : IWebSocketController
+public abstract class WebSocketController : IWebSocketController
 {
     /// <summary>
     /// WebsocketManager provides work with all clients

yawaflua.WebSockets/ServiceBindings.cs

@@ -16,6 +16,8 @@ public static IServiceCollection SettingUpWebSockets(this IServiceCollection isc
         if (isc.All(k => k.ServiceType != typeof(WebSocketConfig)))
             isc.AddSingleton(new WebSocketConfig());
         isc.AddScoped<IWebSocketManager, WebSocketManager>();
+        isc.AddSingleton<IWebSocketManager, WebSocketManager>();
+        isc.AddTransient<IWebSocketManager, WebSocketManager>();
         isc.AddSingleton<WebSocketMiddleware>();
         return isc;
     }