Add DisposeAsync support to WebHost and RequestServices (#7091)

src/Hosting/Hosting/src/Internal/WebHost.cs

@@ -24,7 +24,7 @@ using Microsoft.Extensions.StackTrace.Sources;
 
 namespace Microsoft.AspNetCore.Hosting.Internal
 {
-    internal class WebHost : IWebHost
+    internal class WebHost : IWebHost, IAsyncDisposable
     {
         private static readonly string DeprecatedServerUrlsKey = "server.urls";
 
@@ -342,12 +342,17 @@ namespace Microsoft.AspNetCore.Hosting.Internal
         }
 
         public void Dispose()
+        {
+            DisposeAsync().ConfigureAwait(false).GetAwaiter().GetResult();
+        }
+
+        public async ValueTask DisposeAsync()
         {
             if (!_stopped)
             {
                 try
                 {
-                    StopAsync().GetAwaiter().GetResult();
+                    await StopAsync().ConfigureAwait(false);
                 }
                 catch (Exception ex)
                 {
@@ -355,8 +360,21 @@ namespace Microsoft.AspNetCore.Hosting.Internal
                 }
             }
 
-            (_applicationServices as IDisposable)?.Dispose();
-            (_hostingServiceProvider as IDisposable)?.Dispose();
+            await DisposeServiceProviderAsync(_applicationServices).ConfigureAwait(false);
+            await DisposeServiceProviderAsync(_hostingServiceProvider).ConfigureAwait(false);
+        }
+
+        private async ValueTask DisposeServiceProviderAsync(IServiceProvider serviceProvider)
+        {
+            switch (serviceProvider)
+            {
+                case IAsyncDisposable asyncDisposable:
+                    await asyncDisposable.DisposeAsync();
+                    break;
+                case IDisposable disposable:
+                    disposable.Dispose();
+                    break;
+            }
         }
     }
 }

src/Hosting/Hosting/src/WebHostExtensions.cs

@@ -101,7 +101,7 @@ namespace Microsoft.AspNetCore.Hosting
 
         private static async Task RunAsync(this IWebHost host, CancellationToken token, string shutdownMessage)
         {
-            using (host)
+            try
             {
                 await host.StartAsync(token);
 
@@ -131,6 +131,17 @@ namespace Microsoft.AspNetCore.Hosting
 
                 await host.WaitForTokenShutdownAsync(token);
             }
+            finally
+            {
+                if (host is IAsyncDisposable asyncDisposable)
+                {
+                    await asyncDisposable.DisposeAsync().ConfigureAwait(false);
+                }
+                else
+                {
+                    host.Dispose();
+                }
+            }
         }
 
         private static async Task WaitForTokenShutdownAsync(this IWebHost host, CancellationToken token)

src/Hosting/Hosting/test/WebHostTests.AsyncDisposable.cs

@@ -0,0 +1,93 @@
+// 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.Threading.Tasks;
+using Microsoft.Extensions.DependencyInjection;
+using Xunit;
+
+namespace Microsoft.AspNetCore.Hosting
+{
+    public partial class WebHostTests
+    {
+        [Fact]
+        public async Task DisposingHostCallsDisposeAsyncOnProvider()
+        {
+            var providerFactory = new AsyncServiceProviderFactory();
+            using (var host = CreateBuilder()
+                .UseFakeServer()
+                .ConfigureServices((context, services) =>
+                    services.Add(ServiceDescriptor.Singleton<IServiceProviderFactory<IServiceCollection>>(providerFactory)
+                    ))
+                .UseStartup("Microsoft.AspNetCore.Hosting.Tests")
+                .Build())
+            {
+                await host.StartAsync();
+
+                Assert.Equal(2, providerFactory.Providers.Count);
+
+                await host.StopAsync();
+
+                Assert.All(providerFactory.Providers, provider => {
+                    Assert.False(provider.DisposeCalled);
+                    Assert.False(provider.DisposeAsyncCalled);
+                });
+
+                host.Dispose();
+
+                Assert.All(providerFactory.Providers, provider => {
+                    Assert.False(provider.DisposeCalled);
+                    Assert.True(provider.DisposeAsyncCalled);
+                });
+            }
+        }
+
+        private class AsyncServiceProviderFactory : IServiceProviderFactory<IServiceCollection>
+        {
+            public List<AsyncDisposableServiceProvider> Providers { get; } = new List<AsyncDisposableServiceProvider>();
+
+            public IServiceCollection CreateBuilder(IServiceCollection services)
+            {
+                return services;
+            }
+
+            public IServiceProvider CreateServiceProvider(IServiceCollection containerBuilder)
+            {
+                var provider = new AsyncDisposableServiceProvider(containerBuilder.BuildServiceProvider());
+                Providers.Add(provider);
+                return provider;
+            }
+        }
+
+        private class AsyncDisposableServiceProvider : IServiceProvider, IDisposable, IAsyncDisposable
+        {
+            private readonly ServiceProvider _serviceProvider;
+
+            public AsyncDisposableServiceProvider(ServiceProvider serviceProvider)
+            {
+                _serviceProvider = serviceProvider;
+            }
+
+            public bool DisposeCalled { get; set; }
+
+            public bool DisposeAsyncCalled { get; set; }
+
+            public object GetService(Type serviceType) => _serviceProvider.GetService(serviceType);
+
+            public void Dispose()
+            {
+                DisposeCalled = true;
+                _serviceProvider.Dispose();
+            }
+
+            public ValueTask DisposeAsync()
+            {
+                DisposeAsyncCalled = true;
+                _serviceProvider.Dispose();
+                return default;
+            }
+        }
+    }
+
+}

src/Hosting/Hosting/test/WebHostTests.cs

@@ -25,7 +25,7 @@ using Xunit;
 
 namespace Microsoft.AspNetCore.Hosting
 {
-    public class WebHostTests
+    public partial class WebHostTests
     {
         [Fact]
         public async Task WebHostThrowsWithNoServer()
@@ -1325,4 +1325,4 @@ namespace Microsoft.AspNetCore.Hosting
             return builder.ConfigureServices(services => services.AddSingleton<IServer, WebHostTests.FakeServer>());
         }
     }
-}
+}

src/Http/Http.Abstractions/src/HttpResponse.cs

@@ -22,6 +22,8 @@ namespace Microsoft.AspNetCore.Http
             return Task.CompletedTask;
         };
 
+        private static readonly Func<object, Task> _disposeAsyncDelegate = disposable => ((IAsyncDisposable)disposable).DisposeAsync().AsTask();
+
         /// <summary>
         /// Gets the <see cref="HttpContext"/> for this response.
         /// </summary>
@@ -93,6 +95,12 @@ namespace Microsoft.AspNetCore.Http
         /// <param name="disposable">The object to be disposed.</param>
         public virtual void RegisterForDispose(IDisposable disposable) => OnCompleted(_disposeDelegate, disposable);
 
+        /// <summary>
+        /// Registers an object for asynchronous disposal by the host once the request has finished processing.
+        /// </summary>
+        /// <param name="disposable">The object to be disposed asynchronously.</param>
+        public virtual void RegisterForDisposeAsync(IAsyncDisposable disposable) => OnCompleted(_disposeAsyncDelegate, disposable);
+
         /// <summary>
         /// Adds a delegate to be invoked after the response has finished being sent to the client.
         /// </summary>

src/Http/Http/src/Features/RequestServicesFeature.cs

@@ -2,17 +2,18 @@
 // 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.Extensions.DependencyInjection;
 
 namespace Microsoft.AspNetCore.Http.Features
 {
-    public class RequestServicesFeature : IServiceProvidersFeature, IDisposable
+    public class RequestServicesFeature : IServiceProvidersFeature, IDisposable, IAsyncDisposable
     {
         private readonly IServiceScopeFactory _scopeFactory;
         private IServiceProvider _requestServices;
         private IServiceScope _scope;
         private bool _requestServicesSet;
-        private HttpContext _context;
+        private readonly HttpContext _context;
 
         public RequestServicesFeature(HttpContext context, IServiceScopeFactory scopeFactory)
         {
@@ -26,7 +27,7 @@ namespace Microsoft.AspNetCore.Http.Features
             {
                 if (!_requestServicesSet && _scopeFactory != null)
                 {
-                    _context.Response.RegisterForDispose(this);
+                    _context.Response.RegisterForDisposeAsync(this);
                     _scope = _scopeFactory.CreateScope();
                     _requestServices = _scope.ServiceProvider;
                     _requestServicesSet = true;
@@ -41,11 +42,25 @@ namespace Microsoft.AspNetCore.Http.Features
             }
         }
 
-        public void Dispose()
+        public async ValueTask DisposeAsync()
         {
-            _scope?.Dispose();
+            switch (_scope)
+            {
+                case IAsyncDisposable asyncDisposable:
+                    await asyncDisposable.DisposeAsync();
+                    break;
+                case IDisposable disposable:
+                    disposable.Dispose();
+                    break;
+            }
+
             _scope = null;
             _requestServices = null;
         }
+
+        public void Dispose()
+        {
+            DisposeAsync().ConfigureAwait(false).GetAwaiter().GetResult();
+        }
     }
 }

src/Http/Http/test/DefaultHttpContextTests.cs

@@ -199,7 +199,7 @@ namespace Microsoft.AspNetCore.Http
                         .BuildServiceProvider();
 
             var scopeFactory = serviceProvider.GetRequiredService<IServiceScopeFactory>();
-            
+
             var context = new DefaultHttpContext();
             context.ServiceScopeFactory = scopeFactory;
             context.RequestServices = serviceProvider;
@@ -211,8 +211,8 @@ namespace Microsoft.AspNetCore.Http
         public async Task RequestServicesAreDisposedOnCompleted()
         {
             var serviceProvider = new ServiceCollection()
-                        .AddTransient<DisposableThing>()
-                        .BuildServiceProvider();
+                .AddTransient<DisposableThing>()
+                .BuildServiceProvider();
 
             var scopeFactory = serviceProvider.GetRequiredService<IServiceScopeFactory>();
             DisposableThing instance = null;
@@ -234,6 +234,36 @@ namespace Microsoft.AspNetCore.Http
             Assert.True(instance.Disposed);
         }
 
+        [Fact]
+        public async Task RequestServicesAreDisposedAsynOnCompleted()
+        {
+            var serviceProvider = new AsyncDisposableServiceProvider(new ServiceCollection()
+                .AddTransient<DisposableThing>()
+                .BuildServiceProvider());
+
+            var scopeFactory = serviceProvider.GetRequiredService<IServiceScopeFactory>();
+            DisposableThing instance = null;
+
+            var context = new DefaultHttpContext();
+            context.ServiceScopeFactory = scopeFactory;
+            var responseFeature = new TestHttpResponseFeature();
+            context.Features.Set<IHttpResponseFeature>(responseFeature);
+
+            Assert.NotNull(context.RequestServices);
+            Assert.Single(responseFeature.CompletedCallbacks);
+
+            instance = context.RequestServices.GetRequiredService<DisposableThing>();
+
+            var callback = responseFeature.CompletedCallbacks[0];
+            await callback.callback(callback.state);
+
+            Assert.Null(context.RequestServices);
+            Assert.True(instance.Disposed);
+            var scope = Assert.Single(serviceProvider.Scopes);
+            Assert.True(scope.DisposeAsyncCalled);
+            Assert.False(scope.DisposeCalled);
+        }
+
         void TestAllCachedFeaturesAreNull(HttpContext context, IFeatureCollection features)
         {
             TestCachedFeaturesAreNull(context, features);
@@ -427,5 +457,68 @@ namespace Microsoft.AspNetCore.Http
                 throw new NotImplementedException();
             }
         }
+
+        private class AsyncDisposableServiceProvider : IServiceProvider, IDisposable, IServiceScopeFactory
+        {
+            private readonly ServiceProvider _serviceProvider;
+
+            public AsyncDisposableServiceProvider(ServiceProvider serviceProvider)
+            {
+                _serviceProvider = serviceProvider;
+            }
+
+            public List<AsyncServiceScope> Scopes { get; } = new List<AsyncServiceScope>();
+
+            public object GetService(Type serviceType)
+            {
+                if (serviceType == typeof(IServiceScopeFactory))
+                {
+                    return this;
+                }
+
+                return _serviceProvider.GetService(serviceType);
+            }
+
+            public void Dispose()
+            {
+                _serviceProvider.Dispose();
+            }
+
+            public IServiceScope CreateScope()
+            {
+                var scope = new AsyncServiceScope(_serviceProvider.GetService<IServiceScopeFactory>().CreateScope());
+                Scopes.Add(scope);
+                return scope;
+            }
+
+            internal class AsyncServiceScope : IServiceScope, IAsyncDisposable
+            {
+                private readonly IServiceScope _scope;
+
+                public AsyncServiceScope(IServiceScope scope)
+                {
+                    _scope = scope;
+                }
+
+                public bool DisposeCalled { get; set; }
+
+                public bool DisposeAsyncCalled { get; set; }
+
+                public void Dispose()
+                {
+                    DisposeCalled = true;
+                    _scope.Dispose();
+                }
+
+                public ValueTask DisposeAsync()
+                {
+                    DisposeAsyncCalled = true;
+                    _scope.Dispose();
+                    return default;
+                }
+
+                public IServiceProvider ServiceProvider => _scope.ServiceProvider;
+            }
+        }
     }
 }