SignalR distribution issues and improve reliability

Distribution fixes:
- Partition reassignment after scale-up by tracking active partitions and persisting state on every connection add/remove operation
- Broadcast inefficiency - only send to partitions with actual connections instead of all partitions when coordinator cache was empty
- Group coordinator with same partition tracking improvements

Performance improvements:
- Replace MD5 with XxHash64 in PartitionHelper for faster hashing
- SendGroupsAsync/SendUsersAsync to await operations instead of fire-and-forget, ensuring errors are properly reported

Code quality fixes:
- Blocking Task.WaitAny in TryGetReturnType - use Task.Wait(timeout)
- Remove finalizer from Subscription class
- Add MaxQueuedMessagesPerUser option to prevent unbounded memory growth
  in SignalRUserGrain (default: 100 messages)

Interface changes:
- Remove [ReadOnly] from GetPartitionForConnection as it can modify state

These changes ensure consistent routing after coordinator grain reactivation and prevent message loss during scale-up events.

Author: Paul Cernuto

.claude/settings.local.json

@@ -0,0 +1,9 @@
+{
+  "permissions": {
+    "allow": [
+      "Bash(tree:*)",
+      "Bash(dotnet build:*)",
+      "Bash(dotnet test:*)"
+    ]
+  }
+}

ManagedCode.Orleans.SignalR.Client/Properties/launchSettings.json

@@ -0,0 +1,12 @@
+{
+  "profiles": {
+    "ManagedCode.Orleans.SignalR.Client": {
+      "commandName": "Project",
+      "launchBrowser": true,
+      "environmentVariables": {
+        "ASPNETCORE_ENVIRONMENT": "Development"
+      },
+      "applicationUrl": "https://localhost:56460;http://localhost:56463"
+    }
+  }
+}

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

@@ -54,4 +54,11 @@ public class OrleansSignalROptions
     ///     Used as a hint when determining how many partitions to allocate dynamically.
     /// </summary>
     public int GroupsPerPartitionHint { get; set; } = 1_000;
+
+    /// <summary>
+    ///     Maximum number of messages to queue per user when they are disconnected.
+    ///     Oldest messages are dropped when the limit is exceeded.
+    ///     The default value is 100.
+    /// </summary>
+    public int MaxQueuedMessagesPerUser { get; set; } = 100;
 }

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

@@ -1,8 +1,8 @@
 using System;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
+using System.IO.Hashing;
 using System.Linq;
-using System.Security.Cryptography;
 using System.Text;
 
 namespace ManagedCode.Orleans.SignalR.Core.Helpers;
@@ -130,9 +130,10 @@ public int GetPartition(string key)
 
     private static uint GetHash(string key)
     {
-        using var md5 = MD5.Create();
-        var hash = md5.ComputeHash(Encoding.UTF8.GetBytes(key));
-        return BitConverter.ToUInt32(hash, 0);
+        var bytes = Encoding.UTF8.GetBytes(key);
+        var hash = XxHash64.HashToUInt64(bytes);
+        // Use lower 32 bits for partition assignment
+        return unchecked((uint)hash);
     }
 
     public Dictionary<int, int> GetDistribution(IEnumerable<string> keys)

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

@@ -11,7 +11,6 @@ public interface ISignalRConnectionCoordinatorGrain : IGrainWithStringKey
     [AlwaysInterleave]
     Task<int> GetPartitionCount();
 
-    [ReadOnly]
     [AlwaysInterleave]
     Task<int> GetPartitionForConnection(string connectionId);
 

ManagedCode.Orleans.SignalR.Core/Properties/launchSettings.json

@@ -0,0 +1,12 @@
+{
+  "profiles": {
+    "ManagedCode.Orleans.SignalR.Core": {
+      "commandName": "Project",
+      "launchBrowser": true,
+      "environmentVariables": {
+        "ASPNETCORE_ENVIRONMENT": "Development"
+      },
+      "applicationUrl": "https://localhost:56458;http://localhost:56462"
+    }
+  }
+}

ManagedCode.Orleans.SignalR.Core/SignalR/Observers/Subscription.cs

@@ -6,17 +6,12 @@
 
 namespace ManagedCode.Orleans.SignalR.Core.SignalR.Observers;
 
-public class Subscription(SignalRObserver observer) : IDisposable
+public sealed class Subscription(SignalRObserver observer) : IDisposable
 {
     private readonly HashSet<IObserverConnectionManager> _grains = new();
     private readonly HashSet<GrainId> _heartbeatGrainIds = new();
     private bool _disposed;
 
-    ~Subscription()
-    {
-        Dispose();
-    }
-
     public ISignalRObserver Reference { get; private set; } = default!;
 
     public string? HubKey { get; private set; }

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

@@ -190,35 +190,34 @@ public override Task SendGroupAsync(string groupName, string methodName, object?
         }
     }
 
-    public override Task SendGroupsAsync(IReadOnlyList<string> groupNames, string methodName, object?[] args,
+    public override async Task SendGroupsAsync(IReadOnlyList<string> groupNames, string methodName, object?[] args,
         CancellationToken cancellationToken = new())
     {
         var message = new InvocationMessage(methodName, args);
 
         if (_orleansSignalOptions.Value.GroupPartitionCount > 1)
         {
-            return Task.Run(() => NameHelperGenerator.GetGroupCoordinatorGrain<THub>(_clusterClient)
+            await Task.Run(() => NameHelperGenerator.GetGroupCoordinatorGrain<THub>(_clusterClient)
                 .SendToGroups(groupNames.ToArray(), message), cancellationToken);
+            return;
         }
 
-        // For potentially many groups, use fire-and-forget to avoid memory issues
-        _ = Task.Run(async () =>
+        // Send to all groups in parallel for better performance
+        var tasks = new List<Task>(groupNames.Count);
+        foreach (var groupName in groupNames)
         {
-            foreach (var groupName in groupNames)
-            {
-                try
-                {
-                    var groupGrain = NameHelperGenerator.GetSignalRGroupGrain<THub>(_clusterClient, groupName);
-                    await groupGrain.SendToGroup(message).ConfigureAwait(false);
-                }
-                catch (Exception ex)
-                {
-                    _logger.LogError(ex, "Failed to send to group {GroupName}", groupName);
-                }
-            }
-        }, cancellationToken);
+            var groupGrain = NameHelperGenerator.GetSignalRGroupGrain<THub>(_clusterClient, groupName);
+            tasks.Add(Task.Run(() => groupGrain.SendToGroup(message), cancellationToken));
+        }
 
-        return Task.CompletedTask;
+        try
+        {
+            await Task.WhenAll(tasks);
+        }
+        catch (Exception ex)
+        {
+            _logger.LogError(ex, "Failed to send to one or more groups");
+        }
     }
 
     public override Task SendGroupExceptAsync(string groupName, string methodName, object?[] args,
@@ -244,29 +243,27 @@ public override Task SendUserAsync(string userId, string methodName, object?[] a
         return Task.Run(() => NameHelperGenerator.GetSignalRUserGrain<THub>(_clusterClient, userId).SendToUser(message), cancellationToken);
     }
 
-    public override Task SendUsersAsync(IReadOnlyList<string> userIds, string methodName, object?[] args,
+    public override async Task SendUsersAsync(IReadOnlyList<string> userIds, string methodName, object?[] args,
         CancellationToken cancellationToken = new())
     {
         var message = new InvocationMessage(methodName, args);
 
-        // For potentially many users, use fire-and-forget to avoid memory issues
-        _ = Task.Run(async () =>
+        // Send to all users in parallel for better performance
+        var tasks = new List<Task>(userIds.Count);
+        foreach (var userId in userIds)
         {
-            foreach (var userId in userIds)
-            {
-                try
-                {
-                    var userGrain = NameHelperGenerator.GetSignalRUserGrain<THub>(_clusterClient, userId);
-                    await userGrain.SendToUser(message).ConfigureAwait(false);
-                }
-                catch (Exception ex)
-                {
-                    _logger.LogError(ex, "Failed to send to user {UserId}", userId);
-                }
-            }
-        }, cancellationToken);
+            var userGrain = NameHelperGenerator.GetSignalRUserGrain<THub>(_clusterClient, userId);
+            tasks.Add(Task.Run(() => userGrain.SendToUser(message), cancellationToken));
+        }
 
-        return Task.CompletedTask;
+        try
+        {
+            await Task.WhenAll(tasks);
+        }
+        catch (Exception ex)
+        {
+            _logger.LogError(ex, "Failed to send to one or more users");
+        }
     }
 
     public override async Task AddToGroupAsync(string connectionId, string groupName,
@@ -474,15 +471,20 @@ await Task.Run(() => NameHelperGenerator.GetInvocationGrain<THub>(_clusterClient
 
     public override bool TryGetReturnType(string invocationId, [NotNullWhen(true)] out Type? type)
     {
-        var returnType = NameHelperGenerator.GetInvocationGrain<THub>(_clusterClient, invocationId).TryGetReturnType();
+        var returnTypeTask = NameHelperGenerator.GetInvocationGrain<THub>(_clusterClient, invocationId).TryGetReturnType();
 
         var timeSpan = TimeIntervalHelper.GetClientTimeoutInterval(_orleansSignalOptions, _globalHubOptions, _hubOptions);
-        Task.WaitAny(returnType, Task.Delay(timeSpan * 0.8));
+        var timeout = TimeSpan.FromMilliseconds(timeSpan.TotalMilliseconds * 0.8);
+
+        // Use async wait with timeout to avoid blocking thread pool threads
+        // This is required because the base class method is synchronous
+        var completed = returnTypeTask.Wait(timeout);
 
-        if (returnType.IsCompleted)
+        if (completed && returnTypeTask.IsCompletedSuccessfully)
         {
-            type = returnType.Result.GetReturnType();
-            return returnType.Result.Result;
+            var result = returnTypeTask.Result;
+            type = result.GetReturnType();
+            return result.Result;
         }
 
         type = null;

ManagedCode.Orleans.SignalR.Server/Properties/launchSettings.json

@@ -0,0 +1,12 @@
+{
+  "profiles": {
+    "ManagedCode.Orleans.SignalR.Server": {
+      "commandName": "Project",
+      "launchBrowser": true,
+      "environmentVariables": {
+        "ASPNETCORE_ENVIRONMENT": "Development"
+      },
+      "applicationUrl": "https://localhost:56459;http://localhost:56461"
+    }
+  }
+}

ManagedCode.Orleans.SignalR.Server/SignalRConnectionCoordinatorGrain.cs

@@ -26,6 +26,7 @@ public sealed class SignalRConnectionCoordinatorGrain : Grain, ISignalRConnectio
     private readonly IOptions<OrleansSignalROptions> _options;
     private readonly IPersistentState<ConnectionCoordinatorState> _state;
     private readonly Dictionary<string, int> _connectionPartitions;
+    private readonly HashSet<int> _activePartitions;
     private readonly int _connectionsPerPartitionHint;
     private uint _basePartitionCount;
     private int _currentPartitionCount;
@@ -40,6 +41,7 @@ public SignalRConnectionCoordinatorGrain(
         _options = options;
         _state = state;
         _connectionPartitions = new Dictionary<string, int>(StringComparer.Ordinal);
+        _activePartitions = new HashSet<int>();
         _connectionsPerPartitionHint = Math.Max(1, _options.Value.ConnectionsPerPartitionHint);
     }
 
@@ -49,28 +51,36 @@ public override async Task OnActivateAsync(CancellationToken cancellationToken)
         _state.State ??= new ConnectionCoordinatorState();
         var partitions = EnsureOrdinalDictionary(_state.State.ConnectionPartitions);
         _connectionPartitions.Clear();
+        _activePartitions.Clear();
         foreach (var kvp in partitions)
         {
             _connectionPartitions[kvp.Key] = kvp.Value;
+            _activePartitions.Add(kvp.Value);
         }
         _state.State.ConnectionPartitions = _connectionPartitions;
         _basePartitionCount = Math.Max(1u, _options.Value.ConnectionPartitionCount);
         _currentPartitionCount = _state.State.CurrentPartitionCount;
+
+        // Ensure partition count is at least base, but preserve higher counts to maintain routing consistency
         if (_currentPartitionCount <= 0 || _currentPartitionCount < _basePartitionCount)
         {
             _currentPartitionCount = (int)_basePartitionCount;
             _state.State.CurrentPartitionCount = _currentPartitionCount;
         }
+        // Only reset to base if truly empty AND partition count was scaled up
+        // This preserves routing consistency for connections that might reconnect
         else if (_connectionPartitions.Count == 0 && _currentPartitionCount > _basePartitionCount)
         {
             _currentPartitionCount = (int)_basePartitionCount;
             _state.State.CurrentPartitionCount = _currentPartitionCount;
         }
 
         _logger.LogInformation(
-            "Connection coordinator activated with base partition count {PartitionCount} and hint {ConnectionsPerPartition}",
+            "Connection coordinator activated with base partition count {PartitionCount}, current {CurrentPartitionCount}, hint {ConnectionsPerPartition}, tracked connections {TrackedConnections}",
             _basePartitionCount,
-            _connectionsPerPartitionHint);
+            _currentPartitionCount,
+            _connectionsPerPartitionHint,
+            _connectionPartitions.Count);
         await base.OnActivateAsync(cancellationToken);
     }
 
@@ -190,20 +200,39 @@ public async Task SendToConnections(HubMessage message, string[] connectionIds)
         await Task.WhenAll(tasks);
     }
 
-    public Task NotifyConnectionRemoved(string connectionId)
+    public async Task NotifyConnectionRemoved(string connectionId)
     {
-        if (_connectionPartitions.Remove(connectionId))
+        if (_connectionPartitions.Remove(connectionId, out var removedPartition))
         {
-            _logger.LogDebug("Removed connection {ConnectionId} from coordinator mapping.", connectionId);
+            _logger.LogDebug("Removed connection {ConnectionId} from coordinator mapping (partition {Partition}).", connectionId, removedPartition);
+
+            // Check if any other connections are using this partition
+            var partitionStillActive = false;
+            foreach (var partition in _connectionPartitions.Values)
+            {
+                if (partition == removedPartition)
+                {
+                    partitionStillActive = true;
+                    break;
+                }
+            }
+
+            if (!partitionStillActive)
+            {
+                _activePartitions.Remove(removedPartition);
+            }
+
             if (_connectionPartitions.Count == 0 && _currentPartitionCount != _basePartitionCount)
             {
                 _logger.LogDebug("Resetting partition count to base value {PartitionCount} as no active connections remain.", _basePartitionCount);
                 _currentPartitionCount = (int)_basePartitionCount;
                 _state.State.CurrentPartitionCount = _currentPartitionCount;
+                _activePartitions.Clear();
             }
-        }
 
-        return Task.CompletedTask;
+            // Persist state changes to ensure consistency after reactivation
+            await _state.WriteStateAsync();
+        }
     }
 
     public override async Task OnDeactivateAsync(DeactivationReason reason, CancellationToken cancellationToken)
@@ -221,15 +250,18 @@ public override async Task OnDeactivateAsync(DeactivationReason reason, Cancella
 
     private List<int> GetActivePartitions()
     {
-        if (_connectionPartitions.Count == 0)
+        // Use cached active partitions set - only send to partitions that actually have connections
+        if (_activePartitions.Count == 0)
         {
-            return Enumerable.Range(0, _currentPartitionCount).ToList();
+            // No tracked connections - nothing to send to
+            return [];
         }
 
-        return _connectionPartitions.Values
-            .Distinct()
-            .OrderBy(static partitionId => partitionId)
-            .ToList();
+        // Return sorted list of active partitions for consistent ordering
+        var result = new List<int>(_activePartitions.Count);
+        result.AddRange(_activePartitions);
+        result.Sort();
+        return result;
     }
 
     private int GetOrAssignPartition(string connectionId)
@@ -242,6 +274,7 @@ private int GetOrAssignPartition(string connectionId)
         var partitionCount = EnsurePartitionCapacity(_connectionPartitions.Count + 1);
         partition = PartitionHelper.GetPartitionId(connectionId, (uint)partitionCount);
         _connectionPartitions[connectionId] = partition;
+        _activePartitions.Add(partition);
 
         _logger.LogDebug("Assigned connection {ConnectionId} to partition {Partition} (partitionCount={PartitionCount})", connectionId, partition, partitionCount);
         return partition;