[release/8.0] SignalR performance: track groups per connection, remove on disconnect (#53862)

* SignalR: track groups per connection, remove on disconnect, fixes #48249 Instead of iterating over ALL the groups, which is slow and even introduces a DDoS vector, we remove from groups that are specific to this connection

* Removed unused method

* Apply suggestions from code review

Co-authored-by: Brennan <[email protected]>

* More code review suggestions with a minor fix

* Use lock instead of ToArray (the latter iterates anyway)

* Made the HashSet an internal property on HubConnectionContext

* Removed unneeded "using"

* Removed ignore-case from group tracking

* Addressed race condition from code review

---------

Co-authored-by: alex-jitbit <[email protected]>
Co-authored-by: Brennan <[email protected]>

Author: 3 people

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

@@ -42,7 +42,18 @@ public override Task AddToGroupAsync(string connectionId, string groupName, Canc
             return Task.CompletedTask;
         }
 
-        _groups.Add(connection, groupName);
+        // Track groups in the connection object
+        lock (connection.GroupNames)
+        {
+            if (!connection.GroupNames.Add(groupName))
+            {
+                // Connection already in group
+                return Task.CompletedTask;
+            }
+
+            _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,7 +75,17 @@ public override Task RemoveFromGroupAsync(string connectionId, string groupName,
             return Task.CompletedTask;
         }
 
-        _groups.Remove(connectionId, groupName);
+        // Remove from previously saved groups
+        lock (connection.GroupNames)
+        {
+            if (!connection.GroupNames.Remove(groupName))
+            {
+                // Connection not in group
+                return Task.CompletedTask;
+            }
+
+            _groups.Remove(connectionId, groupName);
+        }
 
         return Task.CompletedTask;
     }
@@ -277,8 +298,16 @@ public override Task OnConnectedAsync(HubConnectionContext connection)
     /// <inheritdoc />
     public override Task OnDisconnectedAsync(HubConnectionContext connection)
     {
+        lock (connection.GroupNames)
+        {
+            // Remove from tracked groups one by one
+            foreach (var groupName in connection.GroupNames)
+            {
+                _groups.Remove(connection.ConnectionId, groupName);
+            }
+        }
+
         _connections.Remove(connection);
-        _groups.RemoveDisconnectedConnection(connection.ConnectionId);
 
         return Task.CompletedTask;
     }

src/SignalR/server/Core/src/HubConnectionContext.cs

@@ -56,6 +56,9 @@ public partial class HubConnectionContext
     [MemberNotNullWhen(true, nameof(_messageBuffer))]
     internal bool UsingStatefulReconnect() => _useStatefulReconnect;
 
+    // Tracks groups that the connection has been added to
+    internal HashSet<string> GroupNames { get; } = new HashSet<string>();
+
     /// <summary>
     /// Initializes a new instance of the <see cref="HubConnectionContext"/> class.
     /// </summary>

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()