more tests

Author: KSemenenko

ManagedCode.Orleans.SignalR.Core/Config/OrleansSignalROptions.cs

@@ -34,11 +34,24 @@ public class OrleansSignalROptions
     /// </summary>
     public uint ConnectionPartitionCount { get; set; } = 4;
 
+    /// <summary>
+    ///     Target number of concurrent connections per partition.
+    ///     Used as a hint when determining how many partitions to allocate dynamically.
+    ///     Lower values result in more partitions.
+    /// </summary>
+    public int ConnectionsPerPartitionHint { get; set; } = 10_000;
+    
     /// <summary>
     ///     Number of partitions to use for group distribution.
     ///     Set to 1 to disable partitioning.
     ///     Increase this value for better scalability with millions of groups.
     ///     The default value is 4.
     /// </summary>
     public uint GroupPartitionCount { get; set; } = 4;
+
+    /// <summary>
+    ///     Target number of groups per partition.
+    ///     Used as a hint when determining how many partitions to allocate dynamically.
+    /// </summary>
+    public int GroupsPerPartitionHint { get; set; } = 1_000;
 }

ManagedCode.Orleans.SignalR.Core/Helpers/PartitionHelper.cs

@@ -28,22 +28,37 @@ public static int GetPartitionId(string connectionId, uint partitionCount)
 
     public static int GetOptimalPartitionCount(int expectedConnections)
     {
-        // Rule of thumb: ~10,000 connections per partition
-        const int connectionsPerPartition = 10000;
-        var partitions = Math.Max(1, (expectedConnections + connectionsPerPartition - 1) / connectionsPerPartition);
-        
-        // Round to nearest power of 2 for better hash distribution
-        return (int)Math.Pow(2, Math.Ceiling(Math.Log(partitions, 2)));
+        return GetOptimalPartitionCount(expectedConnections, 10_000);
+    }
+    
+    public static int GetOptimalPartitionCount(int expectedConnections, int connectionsPerPartition)
+    {
+        var perPartition = Math.Max(1, connectionsPerPartition);
+        var partitions = Math.Max(1, (expectedConnections + perPartition - 1) / perPartition);
+        return ToPowerOfTwo(partitions);
     }
 
     public static int GetOptimalGroupPartitionCount(int expectedGroups)
     {
-        // Rule of thumb: ~1,000 groups per partition
-        const int groupsPerPartition = 1000;
-        var partitions = Math.Max(1, (expectedGroups + groupsPerPartition - 1) / groupsPerPartition);
-        
-        // Round to nearest power of 2 for better hash distribution
-        return (int)Math.Pow(2, Math.Ceiling(Math.Log(partitions, 2)));
+        return GetOptimalGroupPartitionCount(expectedGroups, 1_000);
+    }
+
+    public static int GetOptimalGroupPartitionCount(int expectedGroups, int groupsPerPartition)
+    {
+        var perPartition = Math.Max(1, groupsPerPartition);
+        var partitions = Math.Max(1, (expectedGroups + perPartition - 1) / perPartition);
+        return ToPowerOfTwo(partitions);
+    }
+
+    private static int ToPowerOfTwo(int value)
+    {
+        if (value <= 1)
+        {
+            return 1;
+        }
+
+        var power = (int)Math.Ceiling(Math.Log(value, 2));
+        return (int)Math.Pow(2, power);
     }
 
     private readonly record struct RingCacheKey(int PartitionCount, int VirtualNodes);

ManagedCode.Orleans.SignalR.Core/Interfaces/ISignalRConnectionCoordinatorGrain.cs

@@ -23,4 +23,6 @@ public interface ISignalRConnectionCoordinatorGrain : IGrainWithStringKey
 
     [OneWay]
     Task SendToConnections(HubMessage message, string[] connectionIds);
+
+    Task NotifyConnectionRemoved(string connectionId);
 }

ManagedCode.Orleans.SignalR.Core/Interfaces/ISignalRGroupCoordinatorGrain.cs

@@ -25,4 +25,6 @@ public interface ISignalRGroupCoordinatorGrain : IGrainWithStringKey
     Task AddConnectionToGroup(string groupName, string connectionId, ISignalRObserver observer);
 
     Task RemoveConnectionFromGroup(string groupName, string connectionId, ISignalRObserver observer);
+
+    Task NotifyGroupRemoved(string groupName);
 }

ManagedCode.Orleans.SignalR.Core/Interfaces/ISignalRGroupPartitionGrain.cs

@@ -19,4 +19,6 @@ public interface ISignalRGroupPartitionGrain : IGrainWithIntegerKey, IObserverCo
 
     [ReadOnly]
     Task<bool> HasConnection(string connectionId);
+
+    Task EnsureInitialized(string hubKey);
 }

ManagedCode.Orleans.SignalR.Core/Models/GroupPartitionState.cs

@@ -16,6 +16,9 @@ public class GroupPartitionState
     [Id(2)]
     public Dictionary<string, string> ConnectionObservers { get; set; } = new();
 
+    [Id(3)]
+    public string? HubKey { get; set; }
+
     public bool IsEmpty =>
         Groups.Count == 0 &&
         ConnectionGroups.Count == 0 &&

ManagedCode.Orleans.SignalR.Core/SignalR/NameHelperGenerator.cs

@@ -60,6 +60,11 @@ public static ISignalRGroupCoordinatorGrain GetGroupCoordinatorGrain<THub>(IGrai
         return grainFactory.GetGrain<ISignalRGroupCoordinatorGrain>(CleanString(typeof(THub).FullName!));
     }
 
+    public static ISignalRGroupCoordinatorGrain GetGroupCoordinatorGrain(IGrainFactory grainFactory, string hubKey)
+    {
+        return grainFactory.GetGrain<ISignalRGroupCoordinatorGrain>(CleanString(hubKey));
+    }
+    
     public static ISignalRGroupPartitionGrain GetGroupPartitionGrain<THub>(IGrainFactory grainFactory, int partitionId)
     {
         var key = GetPartitionGrainKey(typeof(THub).FullName!, partitionId, alreadyCleaned: false);

ManagedCode.Orleans.SignalR.Core/SignalR/OrleansHubLifetimeManager.cs

@@ -91,6 +91,9 @@ public override async Task OnDisconnectedAsync(HubConnectionContext connection)
         // For small number of grains (typical case), await all tasks
         // This ensures proper cleanup on disconnect
         await Task.Run(() => Task.WhenAll(tasks));
+
+        await Task.Run(() => NameHelperGenerator.GetConnectionCoordinatorGrain<THub>(_clusterClient)
+            .NotifyConnectionRemoved(connection.ConnectionId));
     }
 
     public override Task SendAllAsync(string methodName, object?[] args, CancellationToken cancellationToken = new())
@@ -231,6 +234,9 @@ public override async Task AddToGroupAsync(string connectionId, string groupName
             var coordinatorGrain = NameHelperGenerator.GetGroupCoordinatorGrain<THub>(_clusterClient);
             var partitionId = await Task.Run(() => coordinatorGrain.GetPartitionForGroup(groupName), cancellationToken);
             var partitionGrain = NameHelperGenerator.GetGroupPartitionGrain<THub>(_clusterClient, partitionId);
+            var hubKey = NameHelperGenerator.CleanString(typeof(THub).FullName!);
+
+            await Task.Run(() => partitionGrain.EnsureInitialized(hubKey), cancellationToken);
 
             subscription.AddGrain(partitionGrain);
             await Task.Run(() => partitionGrain.AddConnection(connectionId, subscription.Reference), cancellationToken);
@@ -256,6 +262,9 @@ public override async Task RemoveFromGroupAsync(string connectionId, string grou
             var coordinatorGrain = NameHelperGenerator.GetGroupCoordinatorGrain<THub>(_clusterClient);
             var partitionId = await Task.Run(() => coordinatorGrain.GetPartitionForGroup(groupName), cancellationToken);
             var partitionGrain = NameHelperGenerator.GetGroupPartitionGrain<THub>(_clusterClient, partitionId);
+            var hubKey = NameHelperGenerator.CleanString(typeof(THub).FullName!);
+
+            await Task.Run(() => partitionGrain.EnsureInitialized(hubKey), cancellationToken);
 
             await Task.Run(() => coordinatorGrain.RemoveConnectionFromGroup(groupName, connectionId, subscription.Reference), cancellationToken);