Fix AddToGroupAsync race with OnDisconnectedAsync (#42117)

Author: BrennanConroy

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

@@ -51,6 +51,11 @@ public override Task AddToGroupAsync(string connectionId, string groupName, Canc
         }
 
         _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)
+        {
+            _groups.Remove(connection.ConnectionId, groupName);
+        }
 
         return Task.CompletedTask;
     }

src/SignalR/server/StackExchangeRedis/src/Internal/RedisSubscriptionManager.cs

@@ -1,4 +1,4 @@
-// Licensed to the .NET Foundation under one or more agreements.
+// Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System.Collections.Concurrent;
@@ -16,6 +16,13 @@ public async Task AddSubscriptionAsync(string id, HubConnectionContext connectio
 
         try
         {
+            // Avoid adding subscription if connection is closing/closed
+            // We're in a lock and ConnectionAborted is triggered before OnDisconnectedAsync is called so this is guaranteed to be safe when adding while connection is closing and removing items
+            if (connection.ConnectionAborted.IsCancellationRequested)
+            {
+                return;
+            }
+
             var subscription = _subscriptions.GetOrAdd(id, _ => new HubConnectionStore());
 
             subscription.Add(connection);
@@ -32,7 +39,7 @@ public async Task AddSubscriptionAsync(string id, HubConnectionContext connectio
         }
     }
 
-    public async Task RemoveSubscriptionAsync(string id, HubConnectionContext connection, Func<string, Task> unsubscribeMethod)
+    public async Task RemoveSubscriptionAsync(string id, HubConnectionContext connection, object state, Func<object, string, Task> unsubscribeMethod)
     {
         await _lock.WaitAsync();
 
@@ -48,7 +55,7 @@ public async Task RemoveSubscriptionAsync(string id, HubConnectionContext connec
             if (subscription.Count == 0)
             {
                 _subscriptions.TryRemove(id, out _);
-                await unsubscribeMethod(id);
+                await unsubscribeMethod(state, id);
             }
         }
         finally

src/SignalR/server/StackExchangeRedis/src/RedisHubLifetimeManager.cs

@@ -354,10 +354,11 @@ private async Task RemoveGroupAsyncCore(HubConnectionContext connection, string
     {
         var groupChannel = _channels.Group(groupName);
 
-        await _groups.RemoveSubscriptionAsync(groupChannel, connection, channelName =>
+        await _groups.RemoveSubscriptionAsync(groupChannel, connection, this, static (state, channelName) =>
         {
-            RedisLog.Unsubscribe(_logger, channelName);
-            return _bus!.UnsubscribeAsync(channelName);
+            var lifetimeManager = (RedisHubLifetimeManager<THub>)state;
+            RedisLog.Unsubscribe(lifetimeManager._logger, channelName);
+            return lifetimeManager._bus!.UnsubscribeAsync(channelName);
         });
 
         var feature = connection.Features.GetRequiredFeature<IRedisFeature>();
@@ -386,10 +387,11 @@ private Task RemoveUserAsync(HubConnectionContext connection)
     {
         var userChannel = _channels.User(connection.UserIdentifier!);
 
-        return _users.RemoveSubscriptionAsync(userChannel, connection, channelName =>
+        return _users.RemoveSubscriptionAsync(userChannel, connection, this, static (state, channelName) =>
         {
-            RedisLog.Unsubscribe(_logger, channelName);
-            return _bus!.UnsubscribeAsync(channelName);
+            var lifetimeManager = (RedisHubLifetimeManager<THub>)state;
+            RedisLog.Unsubscribe(lifetimeManager._logger, channelName);
+            return lifetimeManager._bus!.UnsubscribeAsync(channelName);
         });
     }