Add support for instantiating the startup class (#24144)

* Add support for instantiating the startup class
- Adds an overload of UseStartup that takes a factory so users can control the instance creation. The factory is given the WebHostBuilderContext to expose access to configuration and IWebHostEnvironment.
- Make sure only one startup delegate runs, the last one registered.

* Update src/Hosting/Hosting/src/GenericHost/GenericWebHostBuilder.cs

Co-authored-by: Chris Ross <Tratcher@Outlook.com>

* PR feedback and bug fixes
- Use actual throw expressions...
- Added null checks

Co-authored-by: Chris Ross <Tratcher@Outlook.com>

src/Hosting/Hosting/src/GenericHost/GenericWebHostBuilder.cs

@@ -21,6 +21,7 @@ namespace Microsoft.AspNetCore.Hosting
     {
         private readonly IHostBuilder _builder;
         private readonly IConfiguration _config;
+        private object _startupObject;
         private readonly object _startupKey = new object();
 
         private AggregateException _hostingStartupErrors;
@@ -198,10 +199,12 @@ namespace Microsoft.AspNetCore.Hosting
         public IWebHostBuilder UseStartup(Type startupType)
         {
             // UseStartup can be called multiple times. Only run the last one.
-            _builder.Properties["UseStartup.StartupType"] = startupType;
+            _startupObject = startupType;
+
             _builder.ConfigureServices((context, services) =>
             {
-                if (_builder.Properties.TryGetValue("UseStartup.StartupType", out var cachedType) && (Type)cachedType == startupType)
+                // Run this delegate if the startup type matches
+                if (object.ReferenceEquals(_startupObject, startupType))
                 {
                     UseStartup(startupType, context, services);
                 }
@@ -210,13 +213,31 @@ namespace Microsoft.AspNetCore.Hosting
             return this;
         }
 
-        private void UseStartup(Type startupType, HostBuilderContext context, IServiceCollection services)
+        public IWebHostBuilder UseStartup(Func<WebHostBuilderContext, object> startupFactory)
+        {
+            // Clear the startup type
+            _startupObject = startupFactory;
+
+            _builder.ConfigureServices((context, services) =>
+            {
+                // UseStartup can be called multiple times. Only run the last one.
+                if (object.ReferenceEquals(_startupObject, startupFactory))
+                {
+                    var webHostBuilderContext = GetWebHostBuilderContext(context);
+                    var instance = startupFactory(webHostBuilderContext) ?? throw new InvalidOperationException("The specified factory returned null startup instance.");
+                    UseStartup(instance.GetType(), context, services, instance);
+                }
+            });
+
+            return this;
+        }
+
+        private void UseStartup(Type startupType, HostBuilderContext context, IServiceCollection services, object instance = null)
         {
             var webHostBuilderContext = GetWebHostBuilderContext(context);
             var webHostOptions = (WebHostOptions)context.Properties[typeof(WebHostOptions)];
 
             ExceptionDispatchInfo startupError = null;
-            object instance = null;
             ConfigureBuilder configureBuilder = null;
 
             try
@@ -231,7 +252,7 @@ namespace Microsoft.AspNetCore.Hosting
                     throw new NotSupportedException($"ConfigureServices returning an {typeof(IServiceProvider)} isn't supported.");
                 }
 
-                instance = ActivatorUtilities.CreateInstance(new HostServiceProvider(webHostBuilderContext), startupType);
+                instance ??= ActivatorUtilities.CreateInstance(new HostServiceProvider(webHostBuilderContext), startupType);
                 context.Properties[_startupKey] = instance;
 
                 // Startup.ConfigureServices
@@ -296,13 +317,19 @@ namespace Microsoft.AspNetCore.Hosting
 
         public IWebHostBuilder Configure(Action<WebHostBuilderContext, IApplicationBuilder> configure)
         {
+            // Clear the startup type
+            _startupObject = configure;
+
             _builder.ConfigureServices((context, services) =>
             {
-                services.Configure<GenericWebHostServiceOptions>(options =>
+                if (object.ReferenceEquals(_startupObject, configure))
                 {
-                    var webhostBuilderContext = GetWebHostBuilderContext(context);
-                    options.ConfigureApplication = app => configure(webhostBuilderContext, app);
-                });
+                    services.Configure<GenericWebHostServiceOptions>(options =>
+                    {
+                        var webhostBuilderContext = GetWebHostBuilderContext(context);
+                        options.ConfigureApplication = app => configure(webhostBuilderContext, app);
+                    });
+                }
             });
 
             return this;

src/Hosting/Hosting/src/GenericHost/HostingStartupWebHostBuilder.cs

@@ -75,5 +75,10 @@ namespace Microsoft.AspNetCore.Hosting
         {
             return _builder.UseStartup(startupType);
         }
+
+        public IWebHostBuilder UseStartup(Func<WebHostBuilderContext, object> startupFactory)
+        {
+            return _builder.UseStartup(startupFactory);
+        }
     }
 }

src/Hosting/Hosting/src/GenericHost/ISupportsStartup.cs

@@ -10,5 +10,6 @@ namespace Microsoft.AspNetCore.Hosting
     {
         IWebHostBuilder Configure(Action<WebHostBuilderContext, IApplicationBuilder> configure);
         IWebHostBuilder UseStartup(Type startupType);
+        IWebHostBuilder UseStartup(Func<WebHostBuilderContext, object> startupFactory);
     }
 }

src/Hosting/Hosting/src/Internal/StartupLoader.cs

@@ -37,15 +37,14 @@ namespace Microsoft.AspNetCore.Hosting
         // 
         // If the Startup class ConfigureServices returns an <see cref="IServiceProvider"/> and there is at least an <see cref="IStartupConfigureServicesFilter"/> registered we
         // throw as the filters can't be applied.
-        public static StartupMethods LoadMethods(IServiceProvider hostingServiceProvider, Type startupType, string environmentName)
+        public static StartupMethods LoadMethods(IServiceProvider hostingServiceProvider, Type startupType, string environmentName, object instance = null)
         {
             var configureMethod = FindConfigureDelegate(startupType, environmentName);
 
             var servicesMethod = FindConfigureServicesDelegate(startupType, environmentName);
             var configureContainerMethod = FindConfigureContainerDelegate(startupType, environmentName);
 
-            object instance = null;
-            if (!configureMethod.MethodInfo.IsStatic || (servicesMethod != null && !servicesMethod.MethodInfo.IsStatic))
+            if (instance == null && (!configureMethod.MethodInfo.IsStatic || (servicesMethod != null && !servicesMethod.MethodInfo.IsStatic)))
             {
                 instance = ActivatorUtilities.GetServiceOrCreateInstance(hostingServiceProvider, startupType);
             }
@@ -54,7 +53,7 @@ namespace Microsoft.AspNetCore.Hosting
             // going to be used for anything.
             var type = configureContainerMethod.MethodInfo != null ? configureContainerMethod.GetContainerType() : typeof(object);
 
-            var builder = (ConfigureServicesDelegateBuilder) Activator.CreateInstance(
+            var builder = (ConfigureServicesDelegateBuilder)Activator.CreateInstance(
                 typeof(ConfigureServicesDelegateBuilder<>).MakeGenericType(type),
                 hostingServiceProvider,
                 servicesMethod,
@@ -104,13 +103,13 @@ namespace Microsoft.AspNetCore.Hosting
 
                     // The ConfigureContainer pipeline needs an Action<TContainerBuilder> as source, so we just adapt the
                     // signature with this function.
-                    void Source(TContainerBuilder containerBuilder) => 
+                    void Source(TContainerBuilder containerBuilder) =>
                         action(containerBuilder);
 
                     // The ConfigureContainerBuilder.ConfigureContainerFilters expects an Action<object> as value, but our pipeline
                     // produces an Action<TContainerBuilder> given a source, so we wrap it on an Action<object> that internally casts
                     // the object containerBuilder to TContainerBuilder to match the expected signature of our ConfigureContainer pipeline.
-                    void Target(object containerBuilder) => 
+                    void Target(object containerBuilder) =>
                         BuildStartupConfigureContainerFiltersPipeline(Source)((TContainerBuilder)containerBuilder);
                 }
             }

src/Hosting/Hosting/src/WebHostBuilderExtensions.cs

@@ -61,6 +61,48 @@ namespace Microsoft.AspNetCore.Hosting
             });
         }
 
+        /// <summary>
+        /// Specify a factory that creates the startup instance to be used by the web host.
+        /// </summary>
+        /// <param name="hostBuilder">The <see cref="IWebHostBuilder"/> to configure.</param>
+        /// <param name="startupFactory">A delegate that specifies a factory for the startup class.</param>
+        /// <returns>The <see cref="IWebHostBuilder"/>.</returns>
+        public static IWebHostBuilder UseStartup(this IWebHostBuilder hostBuilder, Func<WebHostBuilderContext, object> startupFactory)
+        {
+            if (startupFactory == null)
+            {
+                throw new ArgumentNullException(nameof(startupFactory));
+            }
+
+            var startupAssemblyName = startupFactory.GetMethodInfo().DeclaringType.GetTypeInfo().Assembly.GetName().Name;
+
+            hostBuilder.UseSetting(WebHostDefaults.ApplicationKey, startupAssemblyName);
+
+            // Light up the GenericWebHostBuilder implementation
+            if (hostBuilder is ISupportsStartup supportsStartup)
+            {
+                return supportsStartup.UseStartup(startupFactory);
+            }
+
+            return hostBuilder
+                .ConfigureServices((context, services) =>
+                {
+                    services.AddSingleton(typeof(IStartup), sp =>
+                    {
+                        var instance = startupFactory(context) ?? throw new InvalidOperationException("The specified factory returned null startup instance.");
+
+                        var hostingEnvironment = sp.GetRequiredService<IHostEnvironment>();
+
+                        // Check if the instance implements IStartup before wrapping
+                        if (instance is IStartup startup)
+                        {
+                            return startup;
+                        }
+
+                        return new ConventionBasedStartup(StartupLoader.LoadMethods(sp, instance.GetType(), hostingEnvironment.EnvironmentName, instance));
+                    });
+                });
+        }
 
         /// <summary>
         /// Specify the startup type to be used by the web host.
@@ -70,6 +112,11 @@ namespace Microsoft.AspNetCore.Hosting
         /// <returns>The <see cref="IWebHostBuilder"/>.</returns>
         public static IWebHostBuilder UseStartup(this IWebHostBuilder hostBuilder, Type startupType)
         {
+            if (startupType == null)
+            {
+                throw new ArgumentNullException(nameof(startupType));
+            }
+
             var startupAssemblyName = startupType.GetTypeInfo().Assembly.GetName().Name;
 
             hostBuilder.UseSetting(WebHostDefaults.ApplicationKey, startupAssemblyName);

src/Hosting/Hosting/test/Fakes/GenericWebHostBuilderWrapper.cs

@@ -73,5 +73,11 @@ namespace Microsoft.AspNetCore.Hosting.Tests.Fakes
             _builder.UseStartup(startupType);
             return this;
         }
+
+        public IWebHostBuilder UseStartup(Func<WebHostBuilderContext, object> startupFactory)
+        {
+            _builder.UseStartup(startupFactory);
+            return this;
+        }
     }
 }

src/Hosting/Hosting/test/WebHostBuilderTests.cs

@@ -6,15 +6,19 @@ using System.Collections.Generic;
 using System.IO;
 using System.Linq;
 using System.Reflection;
+using System.Runtime.ExceptionServices;
 using System.Threading;
 using System.Threading.Tasks;
+using System.Web;
 using Microsoft.AspNetCore.Builder;
 using Microsoft.AspNetCore.Hosting;
 using Microsoft.AspNetCore.Hosting.Fakes;
 using Microsoft.AspNetCore.Hosting.Server;
 using Microsoft.AspNetCore.Hosting.Tests.Fakes;
 using Microsoft.AspNetCore.Http;
+using Microsoft.AspNetCore.Http.Extensions;
 using Microsoft.AspNetCore.Http.Features;
+using Microsoft.AspNetCore.WebUtilities;
 using Microsoft.Extensions.Configuration;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Hosting;
@@ -68,6 +72,71 @@ namespace Microsoft.AspNetCore.Hosting
             }
         }
 
+        [Theory]
+        [MemberData(nameof(DefaultWebHostBuildersWithConfig))]
+        public void UseStartupThrowsWhenFactoryIsNull(IWebHostBuilder builder)
+        {
+            var server = new TestServer();
+            Assert.Throws<ArgumentNullException>(() => builder.UseServer(server).UseStartup((Func<WebHostBuilderContext, object>)null));
+        }
+
+        [Theory]
+        [MemberData(nameof(DefaultWebHostBuilders))]
+        public void UseStartupThrowsWhenFactoryReturnsNull(IWebHostBuilder builder)
+        {
+            var server = new TestServer();
+            var ex = Assert.Throws<InvalidOperationException>(() => builder.UseServer(server).UseStartup(context => null).Build());
+            Assert.Equal("The specified factory returned null startup instance.", ex.Message);
+        }
+
+        [Theory]
+        [MemberData(nameof(DefaultWebHostBuildersWithConfig))]
+        public async Task MultipleUseStartupCallsLastWins(IWebHostBuilder builder)
+        {
+            var server = new TestServer();
+            var host = builder.UseServer(server)
+                              .UseStartup<StartupCtorThrows>()
+                              .UseStartup(context => throw new InvalidOperationException("This doesn't run"))
+                              .Configure(app =>
+                              {
+                                  throw new InvalidOperationException("This doesn't run");
+                              })
+                              .Configure(app =>
+                              {
+                                  app.Run(context =>
+                                  {
+                                      return context.Response.WriteAsync("This wins");
+                                  });
+                              })
+                              .Build();
+            using (host)
+            {
+                await host.StartAsync();
+                await AssertResponseContains(server.RequestDelegate, "This wins");
+            }
+        }
+
+        [Theory]
+        [MemberData(nameof(DefaultWebHostBuildersWithConfig))]
+        public async Task UseStartupFactoryWorks(IWebHostBuilder builder)
+        {
+            void ConfigureServices(IServiceCollection services) { }
+            void Configure(IApplicationBuilder app)
+            {
+                app.Run(context => context.Response.WriteAsync("UseStartupFactoryWorks"));
+            }
+
+            var server = new TestServer();
+            var host = builder.UseServer(server)
+                              .UseStartup(context => new DelegatingStartup(ConfigureServices, Configure))
+                              .Build();
+            using (host)
+            {
+                await host.StartAsync();
+                await AssertResponseContains(server.RequestDelegate, "UseStartupFactoryWorks");
+            }
+        }
+
         [Theory]
         [MemberData(nameof(DefaultWebHostBuildersWithConfig))]
         public async Task StartupCtorThrows_Fallback(IWebHostBuilder builder)
@@ -199,7 +268,7 @@ namespace Microsoft.AspNetCore.Hosting
                     options.ValidateScopes = true;
                 });
 
-            using  var host = hostBuilder.Build();
+            using var host = hostBuilder.Build();
             Assert.Throws<InvalidOperationException>(() => host.Start());
             Assert.True(configurationCallbackCalled);
         }
@@ -728,6 +797,22 @@ namespace Microsoft.AspNetCore.Hosting
             }
         }
 
+        [Theory]
+        [MemberData(nameof(DefaultWebHostBuilders))]
+        public void DefaultApplicationNameWithUseStartupFactory(IWebHostBuilder builder)
+        {
+            using (var host = builder
+                .UseServer(new TestServer())
+                .UseStartup(context => new DelegatingStartup(s => { }, app => { }))
+                .Build())
+            {
+                var hostingEnv = host.Services.GetService<IHostEnvironment>();
+
+                // Should be the assembly containing this test, because that's where the delegate comes from
+                Assert.Equal(typeof(WebHostBuilderTests).Assembly.GetName().Name, hostingEnv.ApplicationName);
+            }
+        }
+
         [Theory]
         [MemberData(nameof(DefaultWebHostBuilders))]
         public void Configure_SupportsNonStaticMethodDelegate(IWebHostBuilder builder)
@@ -770,6 +855,27 @@ namespace Microsoft.AspNetCore.Hosting
             }
         }
 
+        [Fact]
+        public async Task UseStartupImplementingIStartupWorks()
+        {
+            void Configure(IApplicationBuilder app)
+            {
+                app.Run(context => context.Response.WriteAsync("Configure"));
+            }
+
+            IServiceProvider ConfigureServices(IServiceCollection services) => services.BuildServiceProvider();
+
+            var builder = CreateWebHostBuilder();
+            var server = new TestServer();
+            using (var host = builder.UseServer(server)
+                .UseStartup(context => new DelegatingStartupWithIStartup(ConfigureServices, Configure))
+                .Build())
+            {
+                await host.StartAsync();
+                await AssertResponseContains(server.RequestDelegate, "Configure");
+            }
+        }
+
         [Theory]
         [MemberData(nameof(DefaultWebHostBuildersWithConfig))]
         public void Build_DoesNotOverrideILoggerFactorySetByConfigureServices(IWebHostBuilder builder)
@@ -1218,7 +1324,7 @@ namespace Microsoft.AspNetCore.Hosting
 
             Assert.Equal("nestedvalue", builder.GetSetting("key"));
 
-            using  var host = builder.Build();
+            using var host = builder.Build();
             var appConfig = host.Services.GetRequiredService<IConfiguration>();
             Assert.Equal("nestedvalue", appConfig["key"]);
         }
@@ -1574,6 +1680,37 @@ namespace Microsoft.AspNetCore.Hosting
             }
         }
 
+        private class DelegatingStartupWithIStartup : IStartup
+        {
+            private readonly Func<IServiceCollection, IServiceProvider> _configureServices;
+            private readonly Action<IApplicationBuilder> _configure;
+
+            public DelegatingStartupWithIStartup(Func<IServiceCollection, IServiceProvider> configureServices, Action<IApplicationBuilder> configure)
+            {
+                _configureServices = configureServices;
+                _configure = configure;
+            }
+
+            // These are explicitly implemented to verify they don't get called via reflection
+            IServiceProvider IStartup.ConfigureServices(IServiceCollection services) => _configureServices(services);
+            void IStartup.Configure(IApplicationBuilder app) => _configure(app);
+        }
+
+        public class DelegatingStartup
+        {
+            private readonly Action<IServiceCollection> _configureServices;
+            private readonly Action<IApplicationBuilder> _configure;
+
+            public DelegatingStartup(Action<IServiceCollection> configureServices, Action<IApplicationBuilder> configure)
+            {
+                _configureServices = configureServices;
+                _configure = configure;
+            }
+
+            public void ConfigureServices(IServiceCollection services) => _configureServices(services);
+            public void Configure(IApplicationBuilder app) => _configure(app);
+        }
+
         public class StartupWithResolvedDisposableThatThrows
         {
             public StartupWithResolvedDisposableThatThrows(DisposableService service)