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.

Examples/Program.cs

@@ -46,6 +46,8 @@ internal class Startup
         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,16 +27,29 @@ 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;
 
@@ -42,23 +57,30 @@ public class WebSocketRouterTests
         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 class WebSocketRouterTests
             .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 class WebSocketRouterTests
         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 class WebSocketRouterTests
         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 @@ public class WebSocketRouter
                         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 @@ public class WebSocketRouter
                             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 @@ public class WebSocketRouter
                     {
                         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 @@ public class WebSocketRouter
         }
         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 @@ public class WebSocketRouter
                     }
                     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 @@ public class WebSocketRouter
         }
         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 @@ using WebSocketManager = yawaflua.WebSockets.Core.WebSocketManager;
 
 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 class ServiceBindings
         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;
     }

yawaflua.WebSockets/yawaflua.WebSockets.csproj

@@ -4,7 +4,7 @@
         <TargetFrameworks>net6.0;net7.0;net8.0;net9.0</TargetFrameworks>
         <ImplicitUsings>enable</ImplicitUsings>
         <Nullable>enable</Nullable>
-        <Version>1.0.1</Version>
+        <Version>1.0.2</Version>
         <Title>yawaflua.WebSockets</Title>
         <Description>New AspNet controllers looks like websocket manager </Description>
         <Copyright>Dmitrii Shimanskii</Copyright>
@@ -20,6 +20,7 @@
 
     <ItemGroup>
         <PackageReference Include="JetBrains.Annotations" Version=">= (2023.3.0,)"/>
+        <PackageReference Include="Microsoft.AspNetCore.Http.Abstractions" Version="2.3.0" />
         <PackageReference Include="Microsoft.AspNetCore.WebSockets" Version=">= (2.1.7,)" />
         <PackageReference Include="Microsoft.Extensions.Hosting.Abstractions" Version=">= (6.0.0,)" />
         <PackageReference Include="System.Net.WebSockets" Version=">= (4.0.0,)"/>