Handle IAsyncDisposable scoped objects (#7343)

- We make a scope today around hub invocations, with IAsyncDisposable now implemented in the DI container, we need to support IServiceScope being IAsyncDisposable and IDisposable

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

@@ -0,0 +1,24 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System;
+using System.Collections.Generic;
+using System.Text;
+using System.Threading.Tasks;
+
+namespace Microsoft.AspNetCore.SignalR.Internal
+{
+    internal static class AsyncDisposableExtensions
+    {
+        // Does a light up check to see if a type is IAsyncDisposable and calls DisposeAsync if it is
+        public static ValueTask DisposeAsync(this IDisposable disposable)
+        {
+            if (disposable is IAsyncDisposable asyncDisposable)
+            {
+                return asyncDisposable.DisposeAsync();
+            }
+            disposable.Dispose();
+            return default;
+        }
+    }
+}

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

@@ -39,8 +39,12 @@ namespace Microsoft.AspNetCore.SignalR.Internal
 
         public override async Task OnConnectedAsync(HubConnectionContext connection)
         {
-            using (var scope = _serviceScopeFactory.CreateScope())
+            IServiceScope scope = null;
+
+            try
             {
+                scope = _serviceScopeFactory.CreateScope();
+
                 var hubActivator = scope.ServiceProvider.GetRequiredService<IHubActivator<THub>>();
                 var hub = hubActivator.Create();
                 try
@@ -53,12 +57,20 @@ namespace Microsoft.AspNetCore.SignalR.Internal
                     hubActivator.Release(hub);
                 }
             }
+            finally
+            {
+                await scope.DisposeAsync();
+            }
         }
 
         public override async Task OnDisconnectedAsync(HubConnectionContext connection, Exception exception)
         {
-            using (var scope = _serviceScopeFactory.CreateScope())
+            IServiceScope scope = null;
+
+            try
             {
+                scope = _serviceScopeFactory.CreateScope();
+
                 var hubActivator = scope.ServiceProvider.GetRequiredService<IHubActivator<THub>>();
                 var hub = hubActivator.Create();
                 try
@@ -71,6 +83,10 @@ namespace Microsoft.AspNetCore.SignalR.Internal
                     hubActivator.Release(hub);
                 }
             }
+            finally
+            {
+                await scope.DisposeAsync();
+            }
         }
 
         public override Task DispatchMessageAsync(HubConnectionContext connection, HubMessage hubMessage)
@@ -304,7 +320,7 @@ namespace Microsoft.AspNetCore.SignalR.Internal
                                 // And normal invocations handle cleanup below in the finally
                                 if (isStreamCall)
                                 {
-                                    CleanupInvocation(connection, hubMethodInvocationMessage, hubActivator, hub, scope);
+                                    await CleanupInvocation(connection, hubMethodInvocationMessage, hubActivator, hub, scope);
                                 }
                             }
 
@@ -342,17 +358,14 @@ namespace Microsoft.AspNetCore.SignalR.Internal
             {
                 if (disposeScope)
                 {
-                    CleanupInvocation(connection, hubMethodInvocationMessage, hubActivator, hub, scope);
+                    await CleanupInvocation(connection, hubMethodInvocationMessage, hubActivator, hub, scope);
                 }
             }
         }
 
-        private void CleanupInvocation(HubConnectionContext connection, HubMethodInvocationMessage hubMessage, IHubActivator<THub> hubActivator,
+        private ValueTask CleanupInvocation(HubConnectionContext connection, HubMethodInvocationMessage hubMessage, IHubActivator<THub> hubActivator,
             THub hub, IServiceScope scope)
         {
-            hubActivator?.Release(hub);
-            scope.Dispose();
-
             if (hubMessage.StreamIds != null)
             {
                 foreach (var stream in hubMessage.StreamIds)
@@ -365,6 +378,10 @@ namespace Microsoft.AspNetCore.SignalR.Internal
                     catch (KeyNotFoundException) { }
                 }
             }
+
+            hubActivator?.Release(hub);
+
+            return scope.DisposeAsync();
         }
 
         private async Task StreamResultsAsync(string invocationId, HubConnectionContext connection, IAsyncEnumerator<object> enumerator, IServiceScope scope,
@@ -398,7 +415,7 @@ namespace Microsoft.AspNetCore.SignalR.Internal
             {
                 (enumerator as IDisposable)?.Dispose();
 
-                CleanupInvocation(connection, hubMethodInvocationMessage, hubActivator, hub, scope);
+                await CleanupInvocation(connection, hubMethodInvocationMessage, hubActivator, hub, scope);
 
                 // Dispose the linked CTS for the stream.
                 streamCts.Dispose();

src/SignalR/server/SignalR/test/HubConnectionHandlerTestUtils/Hubs.cs

@@ -538,6 +538,21 @@ namespace Microsoft.AspNetCore.SignalR.Tests
         }
     }
 
+    public class HubWithAsyncDisposable : TestHub
+    {
+        private AsyncDisposable _disposable;
+
+        public HubWithAsyncDisposable(AsyncDisposable disposable)
+        {
+            _disposable = disposable;
+        }
+
+        public void Test()
+        {
+
+        }
+    }
+
     public class AbortHub : Hub
     {
         public void Kill()

src/SignalR/server/SignalR/test/HubConnectionHandlerTestUtils/Utils.cs

@@ -2,6 +2,7 @@
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System;
+using System.Threading.Tasks;
 using Microsoft.AspNetCore.SignalR.Protocol;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Logging;
@@ -91,4 +92,20 @@ namespace Microsoft.AspNetCore.SignalR.Tests
     {
         public int DisposeCount = 0;
     }
+
+    public class AsyncDisposable : IAsyncDisposable
+    {
+        private readonly TrackDispose _trackDispose;
+
+        public AsyncDisposable(TrackDispose trackDispose)
+        {
+            _trackDispose = trackDispose;
+        }
+
+        public ValueTask DisposeAsync()
+        {
+            _trackDispose.DisposeCount++;
+            return default;
+        }
+    }
 }

src/SignalR/server/SignalR/test/HubConnectionHandlerTests.cs

@@ -54,6 +54,37 @@ namespace Microsoft.AspNetCore.SignalR.Tests
             }
         }
 
+        [Fact]
+        public async Task AsyncDisposablesInHubsAreSupported()
+        {
+            using (StartVerifiableLog())
+            {
+                var trackDispose = new TrackDispose();
+                var serviceProvider = HubConnectionHandlerTestUtils.CreateServiceProvider(s =>
+                {
+                    s.AddScoped<AsyncDisposable>();
+                    s.AddSingleton(trackDispose);
+                },
+                LoggerFactory);
+                var connectionHandler = serviceProvider.GetService<HubConnectionHandler<HubWithAsyncDisposable>>();
+
+                using (var client = new TestClient())
+                {
+                    var connectionHandlerTask = await client.ConnectAsync(connectionHandler);
+
+                    var result = (await client.InvokeAsync(nameof(HubWithAsyncDisposable.Test)).OrTimeout());
+                    Assert.NotNull(result);
+
+                    // kill the connection
+                    client.Dispose();
+
+                    await connectionHandlerTask;
+
+                    Assert.Equal(3, trackDispose.DisposeCount);
+                }
+            }
+        }
+
         [Fact]
         public async Task ConnectionAbortedTokenTriggers()
         {