Deterministically dispose instances created by WebHostBuilder (#868)

src/Microsoft.AspNetCore.Hosting.Abstractions/HostingAbstractionsWebHostBuilderExtensions.cs

@@ -77,8 +77,7 @@ public static IWebHostBuilder UseServer(this IWebHostBuilder hostBuilder, IServe
             {
                 // It would be nicer if this was transient but we need to pass in the
                 // factory instance directly
-                // Registering as factory so server gets disposed along with a WebHost
-                services.AddSingleton(provider => server);
+                services.AddSingleton(server);
             });
         }
 

src/Microsoft.AspNetCore.Hosting/Internal/ServiceCollectionExtensions.cs

@@ -0,0 +1,20 @@
+// 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 Microsoft.Extensions.DependencyInjection;
+
+namespace Microsoft.AspNetCore.Hosting.Internal
+{
+    internal static class ServiceCollectionExtensions
+    {
+        public static IServiceCollection Clone(this IServiceCollection serviceCollection)
+        {
+            IServiceCollection clone = new ServiceCollection();
+            foreach (var service in serviceCollection)
+            {
+                clone.Add(service);
+            }
+            return clone;
+        }
+    }
+}

src/Microsoft.AspNetCore.Hosting/Internal/WebHost.cs

@@ -245,23 +245,9 @@ public void Dispose()
         {
             _logger?.Shutdown();
             _applicationLifetime.StopApplication();
+            (_hostingServiceProvider as IDisposable)?.Dispose();
             (_applicationServices as IDisposable)?.Dispose();
             _applicationLifetime.NotifyStopped();
         }
-
-        private class Disposable : IDisposable
-        {
-            private Action _dispose;
-
-            public Disposable(Action dispose)
-            {
-                _dispose = dispose;
-            }
-
-            public void Dispose()
-            {
-                Interlocked.Exchange(ref _dispose, () => { }).Invoke();
-            }
-        }
     }
 }

src/Microsoft.AspNetCore.Hosting/Resources.resx

@@ -126,4 +126,7 @@
   <data name="ErrorPageHtml_UnknownLocation" xml:space="preserve">
     <value>Unknown location</value>
   </data>
+  <data name="WebHostBuilder_SingleInstance" xml:space="preserve">
+    <value>WebHostBuilder allows creation only of a single instance of WebHost</value>
+  </data>
 </root>

src/Microsoft.AspNetCore.Hosting/WebHostBuilder.cs

@@ -31,6 +31,7 @@ public class WebHostBuilder : IWebHostBuilder
         private IConfiguration _config;
         private ILoggerFactory _loggerFactory;
         private WebHostOptions _options;
+        private bool _webHostBuilt;
 
         /// <summary>
         /// Initializes a new instance of the <see cref="WebHostBuilder"/> class.
@@ -48,7 +49,7 @@ public WebHostBuilder()
             if (string.IsNullOrEmpty(GetSetting(WebHostDefaults.EnvironmentKey)))
             {
                 // Try adding legacy environment keys, never remove these.
-                UseSetting(WebHostDefaults.EnvironmentKey, Environment.GetEnvironmentVariable("Hosting:Environment") 
+                UseSetting(WebHostDefaults.EnvironmentKey, Environment.GetEnvironmentVariable("Hosting:Environment")
                     ?? Environment.GetEnvironmentVariable("ASPNET_ENV"));
             }
 
@@ -135,6 +136,12 @@ public IWebHostBuilder ConfigureLogging(Action<ILoggerFactory> configureLogging)
         /// </summary>
         public IWebHost Build()
         {
+            if (_webHostBuilt)
+            {
+                throw new InvalidOperationException(Resources.WebHostBuilder_SingleInstance);
+            }
+            _webHostBuilt = true;
+
             // Warn about deprecated environment variables
             if (Environment.GetEnvironmentVariable("Hosting:Environment") != null)
             {
@@ -151,17 +158,24 @@ public IWebHost Build()
                 Console.WriteLine("The environment variable 'ASPNETCORE_SERVER.URLS' is obsolete and has been replaced with 'ASPNETCORE_URLS'");
             }
 
-            var hostingServices = BuildHostingServices();
-            var hostingContainer = hostingServices.BuildServiceProvider();
+            var hostingServices = BuildCommonServices();
+            var applicationServices = hostingServices.Clone();
+            var hostingServiceProvider = hostingServices.BuildServiceProvider();
 
-            var host = new WebHost(hostingServices, hostingContainer, _options, _config);
+            AddApplicationServices(applicationServices, hostingServiceProvider);
+
+            var host = new WebHost(
+                applicationServices,
+                hostingServiceProvider,
+                _options,
+                _config);
 
             host.Initialize();
 
             return host;
         }
 
-        private IServiceCollection BuildHostingServices()
+        private IServiceCollection BuildCommonServices()
         {
             _options = new WebHostOptions(_config);
 
@@ -175,30 +189,33 @@ private IServiceCollection BuildHostingServices()
             var services = new ServiceCollection();
             services.AddSingleton(_hostingEnvironment);
 
+            // The configured ILoggerFactory is added as a singleton here. AddLogging below will not add an additional one.
             if (_loggerFactory == null)
             {
                 _loggerFactory = new LoggerFactory();
+                services.AddSingleton(provider => _loggerFactory);
+            }
+            else
+            {
+                services.AddSingleton(_loggerFactory);
             }
 
             foreach (var configureLogging in _configureLoggingDelegates)
             {
                 configureLogging(_loggerFactory);
             }
 
-            //The configured ILoggerFactory is added as a singleton here. AddLogging below will not add an additional one.
-            services.AddSingleton(_loggerFactory);
-
             //This is required to add ILogger of T.
             services.AddLogging();
 
+            var listener = new DiagnosticListener("Microsoft.AspNetCore");
+            services.AddSingleton<DiagnosticListener>(listener);
+            services.AddSingleton<DiagnosticSource>(listener);
+
             services.AddTransient<IApplicationBuilderFactory, ApplicationBuilderFactory>();
             services.AddTransient<IHttpContextFactory, HttpContextFactory>();
             services.AddOptions();
 
-            var diagnosticSource = new DiagnosticListener("Microsoft.AspNetCore");
-            services.AddSingleton<DiagnosticSource>(diagnosticSource);
-            services.AddSingleton<DiagnosticListener>(diagnosticSource);
-
             // Conjure up a RequestServices
             services.AddTransient<IStartupFilter, AutoRequestServicesStartupFilter>();
             services.AddTransient<IServiceProviderFactory<IServiceCollection>, DefaultServiceProviderFactory>();
@@ -245,6 +262,20 @@ private IServiceCollection BuildHostingServices()
             return services;
         }
 
+        private void AddApplicationServices(IServiceCollection services, IServiceProvider hostingServiceProvider)
+        {
+            // We are forwarding services from hosting contrainer so hosting container
+            // can still manage their lifetime (disposal) shared instances with application services.
+            // NOTE: This code overrides original services lifetime. Instances would always be singleton in
+            // application container.
+            var loggerFactory = hostingServiceProvider.GetService<ILoggerFactory>();
+            services.Replace(ServiceDescriptor.Singleton(typeof(ILoggerFactory), loggerFactory));
+
+            var listener = hostingServiceProvider.GetService<DiagnosticListener>();
+            services.Replace(ServiceDescriptor.Singleton(typeof(DiagnosticListener), listener));
+            services.Replace(ServiceDescriptor.Singleton(typeof(DiagnosticSource), listener));
+        }
+
         private string ResolveContentRootPath(string contentRootPath, string basePath)
         {
             if (string.IsNullOrEmpty(contentRootPath))

src/Microsoft.AspNetCore.Hosting/WebHostBuilderExtensions.cs

@@ -35,7 +35,7 @@ public static IWebHostBuilder Configure(this IWebHostBuilder hostBuilder, Action
 
 
         /// <summary>
-        /// Specify the startup type to be used by the web host. 
+        /// Specify the startup type to be used by the web host.
         /// </summary>
         /// <param name="hostBuilder">The <see cref="IWebHostBuilder"/> to configure.</param>
         /// <param name="startupType">The <see cref="Type"/> to be used.</param>

test/Microsoft.AspNetCore.Hosting.Tests/Fakes/StartupWithILoggerFactory.cs

@@ -0,0 +1,31 @@
+// 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 Microsoft.AspNetCore.Builder;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Logging;
+
+namespace Microsoft.AspNetCore.Hosting.Fakes
+{
+    public class StartupWithILoggerFactory
+    {
+        public ILoggerFactory ConstructorLoggerFactory { get; set; }
+
+        public ILoggerFactory ConfigureLoggerFactory { get; set; }
+
+        public StartupWithILoggerFactory(ILoggerFactory constructorLoggerFactory)
+        {
+             ConstructorLoggerFactory = constructorLoggerFactory;
+        }
+
+        public void ConfigureServices(IServiceCollection collection)
+        {
+            collection.AddSingleton(this);
+        }
+
+        public void Configure(IApplicationBuilder builder, ILoggerFactory loggerFactory)
+        {
+            ConfigureLoggerFactory = loggerFactory;
+        }
+    }
+}

test/Microsoft.AspNetCore.Hosting.Tests/WebHostBuilderTests.cs

@@ -15,6 +15,7 @@
 using Microsoft.Extensions.Configuration;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Logging;
+using Microsoft.Extensions.Logging.Abstractions;
 using Microsoft.Extensions.ObjectPool;
 using Microsoft.Extensions.PlatformAbstractions;
 using Xunit;
@@ -494,6 +495,84 @@ public void Configure_SupportsStaticMethodDelegate()
             Assert.Equal("Microsoft.AspNetCore.Hosting.Tests", hostingEnv.ApplicationName);
         }
 
+        [Fact]
+        public void Build_DoesNotAllowBuildingMuiltipleTimes()
+        {
+            var builder = CreateWebHostBuilder();
+            var server = new TestServer();
+            builder.UseServer(server)
+                .UseStartup<StartupNoServices>()
+                .Build();
+
+            var ex = Assert.Throws<InvalidOperationException>(() => builder.Build());
+
+            Assert.Equal("WebHostBuilder allows creation only of a single instance of WebHost", ex.Message);
+        }
+
+        [Fact]
+        public void Build_PassesSameAutoCreatedILoggerFactoryEverywhere()
+        {
+            var builder = CreateWebHostBuilder();
+            var server = new TestServer();
+            var host = builder.UseServer(server)
+                .UseStartup<StartupWithILoggerFactory>()
+                .Build();
+
+            var startup = host.Services.GetService<StartupWithILoggerFactory>();
+
+            Assert.Equal(startup.ConfigureLoggerFactory, startup.ConstructorLoggerFactory);
+        }
+
+        [Fact]
+        public void Build_PassesSamePassedILoggerFactoryEverywhere()
+        {
+            var factory = new LoggerFactory();
+            var builder = CreateWebHostBuilder();
+            var server = new TestServer();
+            var host = builder.UseServer(server)
+                .UseLoggerFactory(factory)
+                .UseStartup<StartupWithILoggerFactory>()
+                .Build();
+
+            var startup = host.Services.GetService<StartupWithILoggerFactory>();
+
+            Assert.Equal(factory, startup.ConfigureLoggerFactory);
+            Assert.Equal(factory, startup.ConstructorLoggerFactory);
+        }
+
+        [Fact]
+        public void Build_PassedILoggerFactoryNotDisposed()
+        {
+            var factory = new DisposableLoggerFactory();
+            var builder = CreateWebHostBuilder();
+            var server = new TestServer();
+
+            var host = builder.UseServer(server)
+                .UseLoggerFactory(factory)
+                .UseStartup<StartupWithILoggerFactory>()
+                .Build();
+
+            host.Dispose();
+
+            Assert.Equal(false, factory.Disposed);
+        }
+
+        [Fact]
+        public void Build_DoesNotOverrideILoggerFactorySetByConfigureServices()
+        {
+            var factory = new DisposableLoggerFactory();
+            var builder = CreateWebHostBuilder();
+            var server = new TestServer();
+
+            var host = builder.UseServer(server)
+                .ConfigureServices(collection => collection.AddSingleton<ILoggerFactory>(factory))
+                .UseStartup<StartupWithILoggerFactory>()
+                .Build();
+
+            var factoryFromHost = host.Services.GetService<ILoggerFactory>();
+            Assert.Equal(factory, factoryFromHost);
+        }
+
         private static void StaticConfigureMethod(IApplicationBuilder app)
         { }
 
@@ -558,5 +637,24 @@ private class ServiceB
         {
 
         }
+
+        private class DisposableLoggerFactory : ILoggerFactory
+        {
+            public void Dispose()
+            {
+                Disposed = true;
+            }
+
+            public bool Disposed { get; set; }
+
+            public ILogger CreateLogger(string categoryName)
+            {
+                return NullLogger.Instance;
+            }
+
+            public void AddProvider(ILoggerProvider provider)
+            {
+            }
+        }
     }
 }

test/Microsoft.AspNetCore.Hosting.Tests/WebHostTests.cs

@@ -135,7 +135,7 @@ public void WebHostCanBeStarted()
 
             host.Dispose();
 
-            Assert.Equal(1, _startInstances[0].DisposeCalls);
+            Assert.Equal(0, _startInstances[0].DisposeCalls);
         }
 
         [Fact]
@@ -163,7 +163,7 @@ public void WebHostShutsDownWhenTokenTriggers()
             // Wait on the host to shutdown
             lifetime.ApplicationStopped.WaitHandle.WaitOne();
 
-            Assert.Equal(1, _startInstances[0].DisposeCalls);
+            Assert.Equal(0, _startInstances[0].DisposeCalls);
         }
 
         [Fact]