Add DisposeAsync support to generic host (#1010)

Author: pakrym

src/DependencyInjection/DI/src/ServiceLookup/ServiceProviderEngineScope.cs

@@ -10,6 +10,9 @@
 namespace Microsoft.Extensions.DependencyInjection.ServiceLookup
 {
     internal class ServiceProviderEngineScope : IServiceScope, IServiceProvider
+#if DISPOSE_ASYNC
+        , IAsyncDisposable
+#endif
     {
         // For testing only
         internal Action<object> _captureDisposableCallback;

src/DependencyInjection/DI/test/ServiceProviderContainerTests.AsyncDisposable.cs

@@ -68,6 +68,67 @@ public void ProviderDisposeThrowsWhenOnlyDisposeAsyncImplemented()
                 exception.Message);
         }
 
+        [Fact]
+        public async Task ProviderScopeDisposeAsyncCallsDisposeAsyncOnServices()
+        {
+            var serviceCollection = new ServiceCollection();
+            serviceCollection.AddTransient<AsyncDisposable>();
+
+            var serviceProvider = CreateServiceProvider(serviceCollection);
+            var scope = serviceProvider.CreateScope();
+            var disposable = scope.ServiceProvider.GetService<AsyncDisposable>();
+
+            await (scope as IAsyncDisposable).DisposeAsync();
+
+            Assert.True(disposable.DisposeAsyncCalled);
+        }
+
+        [Fact]
+        public async Task ProviderScopeDisposeAsyncPrefersDisposeAsyncOnServices()
+        {
+            var serviceCollection = new ServiceCollection();
+            serviceCollection.AddTransient<SyncAsyncDisposable>();
+
+            var serviceProvider = CreateServiceProvider(serviceCollection);
+            var scope = serviceProvider.CreateScope();
+            var disposable = scope.ServiceProvider.GetService<SyncAsyncDisposable>();
+
+            await (scope as IAsyncDisposable).DisposeAsync();
+
+            Assert.True(disposable.DisposeAsyncCalled);
+        }
+
+        [Fact]
+        public void ProviderScopeDisposePrefersServiceDispose()
+        {
+            var serviceCollection = new ServiceCollection();
+            serviceCollection.AddTransient<SyncAsyncDisposable>();
+
+            var serviceProvider = CreateServiceProvider(serviceCollection);
+            var scope = serviceProvider.CreateScope();
+            var disposable = scope.ServiceProvider.GetService<SyncAsyncDisposable>();
+
+            (scope as IDisposable).Dispose();
+
+            Assert.True(disposable.DisposeCalled);
+        }
+
+        [Fact]
+        public void ProviderScopeDisposeThrowsWhenOnlyDisposeAsyncImplemented()
+        {
+            var serviceCollection = new ServiceCollection();
+            serviceCollection.AddTransient<AsyncDisposable>();
+
+            var serviceProvider = CreateServiceProvider(serviceCollection);
+            var scope = serviceProvider.CreateScope();
+            var disposable = scope.ServiceProvider.GetService<AsyncDisposable>();
+
+            var exception = Assert.Throws<InvalidOperationException>(() => (scope as IDisposable).Dispose());
+            Assert.Equal(
+                "'Microsoft.Extensions.DependencyInjection.Tests.ServiceProviderContainerTests+AsyncDisposable' type only implements IAsyncDisposable. Use DisposeAsync to dispose the container.",
+                exception.Message);
+        }
+
         private class AsyncDisposable: IFakeService, IAsyncDisposable
         {
             public bool DisposeAsyncCalled { get; private set; }

src/Hosting/Abstractions/src/HostingAbstractionsHostExtensions.cs

@@ -56,12 +56,26 @@ public static void Run(this IHost host)
         /// <param name="token">The token to trigger shutdown.</param>
         public static async Task RunAsync(this IHost host, CancellationToken token = default)
         {
-            using (host)
+            try
             {
                 await host.StartAsync(token);
 
                 await host.WaitForShutdownAsync(token);
             }
+            finally
+            {
+#if DISPOSE_ASYNC
+                if (host is IAsyncDisposable asyncDisposable)
+                {
+                    await asyncDisposable.DisposeAsync();
+                }
+                else
+#endif
+                {
+                    host.Dispose();
+                }
+
+            }
         }
 
         /// <summary>
@@ -92,4 +106,4 @@ public static async Task WaitForShutdownAsync(this IHost host, CancellationToken
             await host.StopAsync();
         }
     }
-}
+}

src/Hosting/Abstractions/src/Microsoft.Extensions.Hosting.Abstractions.csproj

@@ -2,12 +2,12 @@
 
   <PropertyGroup>
     <Description>.NET Core hosting and startup abstractions for applications.</Description>
-    <TargetFramework>netstandard2.0</TargetFramework>
+    <TargetFrameworks>netstandard2.0;netcoreapp3.0</TargetFrameworks>
     <NoWarn>$(NoWarn);CS1591</NoWarn>
     <GenerateDocumentationFile>true</GenerateDocumentationFile>
     <PackageTags>hosting</PackageTags>
     <RootNamespace>Microsoft.Extensions.Hosting</RootNamespace>
-    <IsShipping>true</IsShipping>
+    <DefineConstants Condition="'$(TargetFramework)' == 'netcoreapp3.0'">$(DefineConstants);DISPOSE_ASYNC</DefineConstants>
   </PropertyGroup>
 
   <ItemGroup>

src/Hosting/Hosting/src/Internal/Host.cs

@@ -1,4 +1,4 @@
-// Copyright (c) .NET Foundation. All rights reserved.
+// 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;
@@ -13,6 +13,9 @@
 namespace Microsoft.Extensions.Hosting.Internal
 {
     internal class Host : IHost
+#if DISPOSE_ASYNC
+      , IAsyncDisposable
+#endif
     {
         private readonly ILogger<Host> _logger;
         private readonly IHostLifetime _hostLifetime;
@@ -98,9 +101,29 @@ public async Task StopAsync(CancellationToken cancellationToken = default)
             _logger.Stopped();
         }
 
+#if DISPOSE_ASYNC
+        public void Dispose()
+        {
+            DisposeAsync().GetAwaiter().GetResult();
+        }
+
+        public async ValueTask DisposeAsync()
+        {
+            switch (Services)
+            {
+                case IAsyncDisposable asyncDisposable:
+                    await asyncDisposable.DisposeAsync();
+                    break;
+                case IDisposable disposable:
+                    disposable.Dispose();
+                    break;
+            }
+        }
+#else
         public void Dispose()
         {
             (Services as IDisposable)?.Dispose();
         }
+#endif
     }
 }

src/Hosting/Hosting/src/Microsoft.Extensions.Hosting.csproj

@@ -2,11 +2,12 @@
 
   <PropertyGroup>
     <Description>.NET Core hosting and startup infrastructures for applications.</Description>
-    <TargetFramework>netstandard2.0</TargetFramework>
+    <TargetFrameworks>netstandard2.0;netcoreapp3.0</TargetFrameworks>
     <NoWarn>$(NoWarn);CS1591</NoWarn>
     <GenerateDocumentationFile>true</GenerateDocumentationFile>
     <PackageTags>hosting</PackageTags>
     <IsShipping>true</IsShipping>
+    <DefineConstants Condition="'$(TargetFramework)' == 'netcoreapp3.0'">$(DefineConstants);DISPOSE_ASYNC</DefineConstants>
   </PropertyGroup>
 
   <ItemGroup>

src/Hosting/Hosting/test/HostTests.cs

@@ -14,7 +14,7 @@
 
 namespace Microsoft.Extensions.Hosting
 {
-    public class HostTests
+    public partial class HostTests
     {
         [Fact]
         public void CreateDefaultBuilder_IncludesContentRootByDefault()

src/Hosting/Hosting/test/Internal/HostTests.AsyncDisposable.cs

@@ -0,0 +1,80 @@
+// 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.
+
+#if NETCOREAPP
+
+using System;
+using System.Collections.Generic;
+using System.Threading.Tasks;
+using Microsoft.Extensions.DependencyInjection;
+using Xunit;
+
+namespace Microsoft.Extensions.Hosting.Internal
+{
+    public partial class HostTests
+    {
+        [Fact]
+        public async Task HostCallsDisposeAsyncOnServiceProvider()
+        {
+            using (var host = CreateBuilder()
+                .ConfigureServices((hostContext, services) =>
+                {
+                    services.AddSingleton<AsyncDisposableService>();
+                })
+                .Build())
+            {
+                await host.StartAsync();
+
+                var asyncDisposableService = host.Services.GetService<AsyncDisposableService>();
+
+                Assert.False(asyncDisposableService.DisposeAsyncCalled);
+
+                await host.StopAsync();
+
+                Assert.False(asyncDisposableService.DisposeAsyncCalled);
+
+                host.Dispose();
+
+                Assert.True(asyncDisposableService.DisposeAsyncCalled);
+            }
+        }
+
+        [Fact]
+        public async Task HostCallsDisposeAsyncOnServiceProviderWhenDisposeAsyncCalled()
+        {
+            using (var host = CreateBuilder()
+                .ConfigureServices((hostContext, services) =>
+                {
+                    services.AddSingleton<AsyncDisposableService>();
+                })
+                .Build())
+            {
+                await host.StartAsync();
+
+                var asyncDisposableService = host.Services.GetService<AsyncDisposableService>();
+
+                Assert.False(asyncDisposableService.DisposeAsyncCalled);
+
+                await host.StopAsync();
+
+                Assert.False(asyncDisposableService.DisposeAsyncCalled);
+
+                await ((IAsyncDisposable)host).DisposeAsync();
+
+                Assert.True(asyncDisposableService.DisposeAsyncCalled);
+            }
+        }
+
+        private class AsyncDisposableService: IAsyncDisposable
+        {
+            public bool DisposeAsyncCalled { get; set; }
+
+            public ValueTask DisposeAsync()
+            {
+                DisposeAsyncCalled = true;
+                return default;
+            }
+        }
+    }
+}
+#endif

src/Hosting/Hosting/test/Internal/HostTests.cs

@@ -16,7 +16,7 @@
 
 namespace Microsoft.Extensions.Hosting.Internal
 {
-    public class HostTests
+    public partial class HostTests
     {
         [Fact]
         public async Task HostInjectsHostingEnvironment()