From 6b095cf533950e7fc93ebc475093bafac70058cb Mon Sep 17 00:00:00 2001
From: David Fowler <davidfowl@gmail.com>
Date: Wed, 31 Jan 2018 16:03:40 -0800
Subject: [PATCH] Allow overriding the hosting service provider  (#1325)

- Use the IServiceProviderFactory<IServiceCollection>
- Assert creation and disposal service providers
- Updated the tests to verify that service providers are created and disposed
- Called CreateBuilder even in the default case in case the service collection is modified as part of it.
---
 .../Internal/StartupLoader.cs                 |  5 +-
 .../WebHostBuilder.cs                         | 18 +++-
 .../WebHostBuilderTests.cs                    | 84 +++++++++++++++++++
 3 files changed, 103 insertions(+), 4 deletions(-)

diff --git a/src/Microsoft.AspNetCore.Hosting/Internal/StartupLoader.cs b/src/Microsoft.AspNetCore.Hosting/Internal/StartupLoader.cs
index 879f382f5c..d7211d39d9 100644
--- a/src/Microsoft.AspNetCore.Hosting/Internal/StartupLoader.cs
+++ b/src/Microsoft.AspNetCore.Hosting/Internal/StartupLoader.cs
@@ -143,9 +143,8 @@ namespace Microsoft.AspNetCore.Hosting.Internal
                     {
                         // Get the default factory
                         var serviceProviderFactory = HostingServiceProvider.GetRequiredService<IServiceProviderFactory<IServiceCollection>>();
-
-                        // Don't bother calling CreateBuilder since it just returns the default service collection
-                        applicationServiceProvider = serviceProviderFactory.CreateServiceProvider(services);
+                        var builder = serviceProviderFactory.CreateBuilder(services);
+                        applicationServiceProvider = serviceProviderFactory.CreateServiceProvider(builder);
                     }
 
                     return applicationServiceProvider ?? services.BuildServiceProvider();
diff --git a/src/Microsoft.AspNetCore.Hosting/WebHostBuilder.cs b/src/Microsoft.AspNetCore.Hosting/WebHostBuilder.cs
index 2d39c2cb86..74e18e628f 100644
--- a/src/Microsoft.AspNetCore.Hosting/WebHostBuilder.cs
+++ b/src/Microsoft.AspNetCore.Hosting/WebHostBuilder.cs
@@ -153,7 +153,7 @@ namespace Microsoft.AspNetCore.Hosting
 
             var hostingServices = BuildCommonServices(out var hostingStartupErrors);
             var applicationServices = hostingServices.Clone();
-            var hostingServiceProvider = hostingServices.BuildServiceProvider();
+            var hostingServiceProvider = GetProviderFromFactory(hostingServices);
 
             if (!_options.SuppressStatusMessages)
             {
@@ -202,6 +202,22 @@ namespace Microsoft.AspNetCore.Hosting
                 host.Dispose();
                 throw;
             }
+
+            IServiceProvider GetProviderFromFactory(IServiceCollection collection)
+            {
+                var provider = collection.BuildServiceProvider();
+                var factory = provider.GetService<IServiceProviderFactory<IServiceCollection>>();
+
+                if (factory != null)
+                {
+                    using (provider)
+                    {
+                        return factory.CreateServiceProvider(factory.CreateBuilder(collection));
+                    }
+                }
+                
+                return provider;                
+            }
         }
 
         private IServiceCollection BuildCommonServices(out AggregateException hostingStartupErrors)
diff --git a/test/Microsoft.AspNetCore.Hosting.Tests/WebHostBuilderTests.cs b/test/Microsoft.AspNetCore.Hosting.Tests/WebHostBuilderTests.cs
index af482b0e0b..c1244e5c8f 100644
--- a/test/Microsoft.AspNetCore.Hosting.Tests/WebHostBuilderTests.cs
+++ b/test/Microsoft.AspNetCore.Hosting.Tests/WebHostBuilderTests.cs
@@ -809,6 +809,45 @@ namespace Microsoft.AspNetCore.Hosting
             }
         }
 
+
+        [Fact]
+        public async Task ExternalContainerInstanceCanBeUsedForEverything()
+        {
+            var disposables = new List<DisposableService>();
+
+            var containerFactory = new ExternalContainerFactory(services =>
+            {
+                services.AddSingleton(sp =>
+                {
+                    var service = new DisposableService();
+                    disposables.Add(service);
+                    return service;
+                });
+            });
+
+            var host = new WebHostBuilder()
+                .UseStartup<StartupWithExternalServices>()
+                .UseServer(new TestServer())
+                .ConfigureServices(services =>
+                {
+                    services.AddSingleton<IServiceProviderFactory<IServiceCollection>>(containerFactory);
+                })
+                .Build();
+
+            using (host)
+            {
+                await host.StartAsync();
+            }
+
+            // We should create the hosting service provider and the application service provider
+            Assert.Equal(2, containerFactory.ServiceProviders.Count);
+            Assert.Equal(2, disposables.Count);
+
+            Assert.NotEqual(disposables[0], disposables[1]);
+            Assert.True(disposables[0].Disposed);
+            Assert.True(disposables[1].Disposed);
+        }
+
         [Fact]
         public void Build_HostingStartupAssemblyCanBeExcluded()
         {
@@ -1048,6 +1087,51 @@ namespace Microsoft.AspNetCore.Hosting
             public Task StopAsync(CancellationToken cancellationToken) => Task.CompletedTask;
         }
 
+        internal class ExternalContainerFactory : IServiceProviderFactory<IServiceCollection>
+        {
+            private readonly Action<IServiceCollection> _configureServices;
+            private readonly List<IServiceProvider> _serviceProviders = new List<IServiceProvider>();
+
+            public List<IServiceProvider> ServiceProviders => _serviceProviders;
+
+            public ExternalContainerFactory(Action<IServiceCollection> configureServices)
+            {
+                _configureServices = configureServices;
+            }
+
+            public IServiceCollection CreateBuilder(IServiceCollection services)
+            {
+                _configureServices(services);
+                return services;
+            }
+
+            public IServiceProvider CreateServiceProvider(IServiceCollection containerBuilder)
+            {
+                var provider = containerBuilder.BuildServiceProvider();
+                _serviceProviders.Add(provider);
+                return provider;
+            }
+        }
+
+        internal class StartupWithExternalServices
+        {
+            public DisposableService DisposableServiceCtor { get; set; }
+
+            public DisposableService DisposableServiceApp { get; set; }
+
+            public StartupWithExternalServices(DisposableService disposable)
+            {
+                DisposableServiceCtor = disposable;
+            }
+
+            public void ConfigureServices(IServiceCollection services) { }
+
+            public void Configure(IApplicationBuilder app, DisposableService disposable)
+            {
+                DisposableServiceApp = disposable;
+            }
+        }
+
         internal class StartupVerifyServiceA : IStartup
         {
             internal ServiceA ServiceA { get; set; }