SignalR performance: track groups per connection, remove on disconnect

src/SignalR/server/Core/src/DefaultHubLifetimeManager.cs

@@ -7,6 +7,7 @@
 using Microsoft.AspNetCore.SignalR.Internal;
 using Microsoft.AspNetCore.SignalR.Protocol;
 using Microsoft.Extensions.Logging;
+using Microsoft.AspNetCore.Http.Features;
 
 namespace Microsoft.AspNetCore.SignalR;
 
@@ -42,6 +43,16 @@ public override Task AddToGroupAsync(string connectionId, string groupName, Canc
             return Task.CompletedTask;
         }
 
+        //track groups in the connection object
+        var groupNames = connection.Features.GetRequiredFeature<GroupTrackerFeature>().Groups;
+        lock (groupNames)
+        {
+            if (!groupNames.Add(groupName))
+            {
+                return Task.CompletedTask; // Connection already in group
+            }
+        }
+
         _groups.Add(connection, groupName);
         // Connection disconnected while adding to group, remove it in case the Add was called after OnDisconnectedAsync removed items from the group
         if (connection.ConnectionAborted.IsCancellationRequested)
@@ -64,6 +75,16 @@ public override Task RemoveFromGroupAsync(string connectionId, string groupName,
             return Task.CompletedTask;
         }
 
+        //remove from previouslyy saved groups
+        var groupNames = connection.Features.GetRequiredFeature<GroupTrackerFeature>().Groups;
+        lock (groupNames)
+        {
+            if (!groupNames.Remove(groupName))
+            {
+                return Task.CompletedTask; // Connection not in group
+            }
+        }
+
         _groups.Remove(connectionId, groupName);
 
         return Task.CompletedTask;
@@ -271,14 +292,22 @@ public override Task SendUserAsync(string userId, string methodName, object?[] a
     public override Task OnConnectedAsync(HubConnectionContext connection)
     {
         _connections.Add(connection);
+        connection.Features.Set(new GroupTrackerFeature()); //add a group tracker to every concection
         return Task.CompletedTask;
     }
 
     /// <inheritdoc />
     public override Task OnDisconnectedAsync(HubConnectionContext connection)
     {
+        //now remove from tracked groups one by one
+        //this is faster than calling _groups.RemoveDisconnectedConnection
+        //because that method iteratas through ALL the groups
+        foreach (var grpName in connection.Features.GetRequiredFeature<GroupTrackerFeature>().Groups.ToArray()) //copy to array because groups can be modified in other methods, prevent "collection was modified"
+        {
+            _groups.Remove(connection.ConnectionId, grpName);
+        }
+
         _connections.Remove(connection);
-        _groups.RemoveDisconnectedConnection(connection.ConnectionId);
 
         return Task.CompletedTask;
     }
@@ -351,6 +380,11 @@ public override async Task<T> InvokeConnectionAsync<T>(string connectionId, stri
         }
     }
 
+    private class GroupTrackerFeature
+    {
+        public HashSet<string> Groups { get; } = new HashSet<string>(StringComparer.OrdinalIgnoreCase);
+    }
+
     /// <inheritdoc/>
     public override Task SetConnectionResultAsync(string connectionId, CompletionMessage result)
     {

src/SignalR/server/Core/src/Internal/HubGroupList.cs

@@ -3,7 +3,6 @@
 
 using System.Collections;
 using System.Collections.Concurrent;
-using System.Linq;
 
 namespace Microsoft.AspNetCore.SignalR.Internal;
 
@@ -43,15 +42,6 @@ public void Remove(string connectionId, string groupName)
         }
     }
 
-    public void RemoveDisconnectedConnection(string connectionId)
-    {
-        var groupNames = _groups.Where(x => x.Value.ContainsKey(connectionId)).Select(x => x.Key);
-        foreach (var groupName in groupNames)
-        {
-            Remove(connectionId, groupName);
-        }
-    }
-
     public int Count => _groups.Count;
 
     public IEnumerator<ConcurrentDictionary<string, HubConnectionContext>> GetEnumerator()