From cb55973c27ca4bf2333c509edadab62c909e3c52 Mon Sep 17 00:00:00 2001
From: David Fowler <davidfowl@gmail.com>
Date: Mon, 29 Jan 2018 02:56:19 -0800
Subject: [PATCH] Dispose the host if Initialize fails (#1318)

---
 .../WebHostBuilder.cs                         | 17 +++++--
 .../WebHostBuilderTests.cs                    | 45 +++++++++++++++++++
 2 files changed, 58 insertions(+), 4 deletions(-)

diff --git a/src/Microsoft.AspNetCore.Hosting/WebHostBuilder.cs b/src/Microsoft.AspNetCore.Hosting/WebHostBuilder.cs
index 31f07c6d71..2d39c2cb86 100644
--- a/src/Microsoft.AspNetCore.Hosting/WebHostBuilder.cs
+++ b/src/Microsoft.AspNetCore.Hosting/WebHostBuilder.cs
@@ -100,7 +100,7 @@ namespace Microsoft.AspNetCore.Hosting
                 throw new ArgumentNullException(nameof(configureServices));
             }
 
-            return ConfigureServices((_ , services) => configureServices(services));
+            return ConfigureServices((_, services) => configureServices(services));
         }
 
         /// <summary>
@@ -189,10 +189,19 @@ namespace Microsoft.AspNetCore.Hosting
                 _options,
                 _config,
                 hostingStartupErrors);
+            try
+            {
+                host.Initialize();
 
-            host.Initialize();
-
-            return host;
+                return host;
+            }
+            catch
+            {
+                // Dispose the host if there's a failure to initialize, this should clean up
+                // will dispose services that were constructed until the exception was thrown
+                host.Dispose();
+                throw;
+            }
         }
 
         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 a5765a52b2..af482b0e0b 100644
--- a/test/Microsoft.AspNetCore.Hosting.Tests/WebHostBuilderTests.cs
+++ b/test/Microsoft.AspNetCore.Hosting.Tests/WebHostBuilderTests.cs
@@ -369,6 +369,23 @@ namespace Microsoft.AspNetCore.Hosting
             Assert.Equal("A public method named 'ConfigureProduction' or 'Configure' could not be found in the 'Microsoft.AspNetCore.Hosting.Fakes.StartupBoom' type.", exception.Message);
         }
 
+        [Fact]
+        public void ServiceProviderDisposedOnBuildException()
+        {
+            var service = new DisposableService();
+            var hostBuilder = new WebHostBuilder()
+                .UseServer(new TestServer())
+                .ConfigureServices(services =>
+                {
+                    // Added as a factory since instances are never disposed by the container
+                    services.AddSingleton(sp => service);
+                })
+                .UseStartup<StartupWithResolvedDisposableThatThrows>();
+
+            Assert.Throws<InvalidOperationException>(() => hostBuilder.Build());
+            Assert.True(service.Disposed);
+        }
+
         [Fact]
         public void CaptureStartupErrorsHonored()
         {
@@ -1050,6 +1067,16 @@ namespace Microsoft.AspNetCore.Hosting
             }
         }
 
+        public class DisposableService : IDisposable
+        {
+            public bool Disposed { get; private set; }
+
+            public void Dispose()
+            {
+                Disposed = true;
+            }
+        }
+
         public class TestHostingStartup : IHostingStartup
         {
             public void Configure(IWebHostBuilder builder)
@@ -1068,6 +1095,24 @@ namespace Microsoft.AspNetCore.Hosting
             }
         }
 
+        public class StartupWithResolvedDisposableThatThrows
+        {
+            public StartupWithResolvedDisposableThatThrows(DisposableService service)
+            {
+
+            }
+
+            public void ConfigureServices(IServiceCollection services)
+            {
+                throw new InvalidOperationException();
+            }
+
+            public void Configure(IApplicationBuilder app)
+            {
+
+            }
+        }
+
         public class TestLoggerProvider : ILoggerProvider
         {
             public TestSink Sink { get; set; } = new TestSink();