From b7d3b3aa13f0f0e18532307f886fac5127885411 Mon Sep 17 00:00:00 2001 From: David Fowler Date: Thu, 7 Feb 2019 19:49:41 -0800 Subject: [PATCH] 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/Internal/AsyncDisposableExtensions.cs | 24 +++++++++++++ .../Core/src/Internal/DefaultHubDispatcher.cs | 35 ++++++++++++++----- .../HubConnectionHandlerTestUtils/Hubs.cs | 15 ++++++++ .../HubConnectionHandlerTestUtils/Utils.cs | 17 +++++++++ .../SignalR/test/HubConnectionHandlerTests.cs | 31 ++++++++++++++++ 5 files changed, 113 insertions(+), 9 deletions(-) create mode 100644 src/SignalR/server/Core/src/Internal/AsyncDisposableExtensions.cs diff --git a/src/SignalR/server/Core/src/Internal/AsyncDisposableExtensions.cs b/src/SignalR/server/Core/src/Internal/AsyncDisposableExtensions.cs new file mode 100644 index 0000000000..f92fc3e00a --- /dev/null +++ b/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; + } + } +} diff --git a/src/SignalR/server/Core/src/Internal/DefaultHubDispatcher.cs b/src/SignalR/server/Core/src/Internal/DefaultHubDispatcher.cs index fe65bf13c6..e137cc5998 100644 --- a/src/SignalR/server/Core/src/Internal/DefaultHubDispatcher.cs +++ b/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>(); 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>(); 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 hubActivator, + private ValueTask CleanupInvocation(HubConnectionContext connection, HubMethodInvocationMessage hubMessage, IHubActivator 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 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(); diff --git a/src/SignalR/server/SignalR/test/HubConnectionHandlerTestUtils/Hubs.cs b/src/SignalR/server/SignalR/test/HubConnectionHandlerTestUtils/Hubs.cs index c4988c378d..681c4759fd 100644 --- a/src/SignalR/server/SignalR/test/HubConnectionHandlerTestUtils/Hubs.cs +++ b/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() diff --git a/src/SignalR/server/SignalR/test/HubConnectionHandlerTestUtils/Utils.cs b/src/SignalR/server/SignalR/test/HubConnectionHandlerTestUtils/Utils.cs index fed151dd6d..6f9055d4c9 100644 --- a/src/SignalR/server/SignalR/test/HubConnectionHandlerTestUtils/Utils.cs +++ b/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; + } + } } diff --git a/src/SignalR/server/SignalR/test/HubConnectionHandlerTests.cs b/src/SignalR/server/SignalR/test/HubConnectionHandlerTests.cs index 14b2583314..03c6438810 100644 --- a/src/SignalR/server/SignalR/test/HubConnectionHandlerTests.cs +++ b/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(); + s.AddSingleton(trackDispose); + }, + LoggerFactory); + var connectionHandler = serviceProvider.GetService>(); + + 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() {