WIPdavidfowl/fix-service-ordering

4 files changed, 95 insertions, 35 deletions

diff --git a/src/DefaultBuilder/src/ConfigureHostBuilder.cs b/src/DefaultBuilder/src/ConfigureHostBuilder.cs
index 8f66f3b585..8426ce23d0 100644
--- a/src/DefaultBuilder/src/ConfigureHostBuilder.cs
+++ b/src/DefaultBuilder/src/ConfigureHostBuilder.cs
@@ -22,11 +22,11 @@ namespace Microsoft.AspNetCore.Builder
 
         private readonly WebHostEnvironment _environment;
         private readonly ConfigurationManager _configuration;
-        private readonly IServiceCollection _services;
+        private readonly WebApplicationServiceCollection _services;
 
         private readonly HostBuilderContext _context;
 
-        internal ConfigureHostBuilder(ConfigurationManager configuration, WebHostEnvironment environment, IServiceCollection services)
+        internal ConfigureHostBuilder(ConfigurationManager configuration, WebHostEnvironment environment, WebApplicationServiceCollection services)
         {
             _configuration = configuration;
             _environment = environment;
@@ -87,13 +87,36 @@ namespace Microsoft.AspNetCore.Builder
         /// <inheritdoc />
         public IHostBuilder ConfigureServices(Action<HostBuilderContext, IServiceCollection> configureDelegate)
         {
+            AddServiceOptions();
+
+            _services.StopTracking();
+
             // Run these immediately so that they are observable by the imperative code
             configureDelegate(_context, _services);
 
+            _services.StartTracking();
+
+            // Still defer because we want to replay these after we have the final state
             _operations.Add(b => b.ConfigureServices(configureDelegate));
             return this;
         }
 
+        private void AddServiceOptions()
+        {
+            var operations = _services.GetOperations();
+
+            if (operations.Count > 0)
+            {
+                _operations.Add(b => b.ConfigureServices(services =>
+                {
+                    foreach (var operation in operations)
+                    {
+                        operation(services);
+                    }
+                }));
+            }
+        }
+
         /// <inheritdoc />
         public IHostBuilder UseServiceProviderFactory<TContainerBuilder>(IServiceProviderFactory<TContainerBuilder> factory) where TContainerBuilder : notnull
         {
@@ -120,6 +143,8 @@ namespace Microsoft.AspNetCore.Builder
 
         internal void RunDeferredCallbacks(IHostBuilder hostBuilder)
         {
+            AddServiceOptions();
+
             foreach (var operation in _operations)
             {
                 operation(hostBuilder);
diff --git a/src/DefaultBuilder/src/ConfigureWebHostBuilder.cs b/src/DefaultBuilder/src/ConfigureWebHostBuilder.cs
index 0ac2ec7045..4f4b0396df 100644
--- a/src/DefaultBuilder/src/ConfigureWebHostBuilder.cs
+++ b/src/DefaultBuilder/src/ConfigureWebHostBuilder.cs
@@ -18,15 +18,15 @@ namespace Microsoft.AspNetCore.Builder
         private readonly WebHostEnvironment _environment;
         private readonly ConfigurationManager _configuration;
         private readonly Dictionary<string, string?> _settings = new(StringComparer.OrdinalIgnoreCase);
-        private readonly IServiceCollection _services;
+        private readonly ConfigureHostBuilder _hostBuilder;
 
         private readonly WebHostBuilderContext _context;
 
-        internal ConfigureWebHostBuilder(ConfigurationManager configuration, WebHostEnvironment environment, IServiceCollection services)
+        internal ConfigureWebHostBuilder(ConfigurationManager configuration, WebHostEnvironment environment, ConfigureHostBuilder hostBuilder)
         {
             _configuration = configuration;
             _environment = environment;
-            _services = services;
+            _hostBuilder = hostBuilder;
 
             _context = new WebHostBuilderContext
             {
@@ -52,8 +52,12 @@ namespace Microsoft.AspNetCore.Builder
         /// <inheritdoc />
         public IWebHostBuilder ConfigureServices(Action<WebHostBuilderContext, IServiceCollection> configureServices)
         {
-            // Run these immediately so that they are observable by the imperative code
-            configureServices(_context, _services);
+            _hostBuilder.ConfigureServices((context, services) =>
+            {
+                configureServices(_context, services);
+            });
+            //// Run these immediately so that they are observable by the imperative code
+            //configureServices(_context, _services);
             return this;
         }
 
diff --git a/src/DefaultBuilder/src/WebApplicationBuilder.cs b/src/DefaultBuilder/src/WebApplicationBuilder.cs
index 1427618401..6df2cbad3d 100644
--- a/src/DefaultBuilder/src/WebApplicationBuilder.cs
+++ b/src/DefaultBuilder/src/WebApplicationBuilder.cs
@@ -22,9 +22,12 @@ namespace Microsoft.AspNetCore.Builder
         private readonly ConfigureWebHostBuilder _deferredWebHostBuilder;
         private readonly WebHostEnvironment _environment;
         private WebApplication? _builtApplication;
+        private readonly WebApplicationServiceCollection _services = new();
 
         internal WebApplicationBuilder(Assembly? callingAssembly, string[]? args = null)
         {
+            Services = _services;
+
             // HACK: MVC and Identity do this horrible thing to get the hosting environment as an instance
             // from the service collection before it is built. That needs to be fixed...
             Environment = _environment = new WebHostEnvironment(callingAssembly);
@@ -41,8 +44,8 @@ namespace Microsoft.AspNetCore.Builder
             bootstrapBuilder.RunConfigurationCallbacks();
 
             Logging = new LoggingBuilder(Services);
-            WebHost = _deferredWebHostBuilder = new ConfigureWebHostBuilder(Configuration, _environment, Services);
-            Host = _deferredHostBuilder = new ConfigureHostBuilder(Configuration, _environment, Services);
+            Host = _deferredHostBuilder = new ConfigureHostBuilder(Configuration, _environment, _services);
+            WebHost = _deferredWebHostBuilder = new ConfigureWebHostBuilder(Configuration, _environment, _deferredHostBuilder);
 
             // Register Configuration as IConfiguration so updates can be observed even after the WebApplication is built.
             Services.AddSingleton<IConfiguration>(Configuration);
@@ -72,7 +75,7 @@ namespace Microsoft.AspNetCore.Builder
         /// <summary>
         /// A collection of services for the application to compose. This is useful for adding user provided or framework provided services.
         /// </summary>
-        public IServiceCollection Services { get; } = new ServiceCollection();
+        public IServiceCollection Services { get; }
 
         /// <summary>
         /// A collection of configuration providers for the application to compose. This is useful for adding new configuration sources and providers.
@@ -102,10 +105,7 @@ namespace Microsoft.AspNetCore.Builder
         /// <returns>A configured <see cref="WebApplication"/>.</returns>
         public WebApplication Build()
         {
-            // We call ConfigureWebHostDefaults AGAIN because config might be added like "ForwardedHeaders_Enabled"
-            // which can add even more services. If not for that, we probably call _hostBuilder.ConfigureWebHost(ConfigureWebHost)
-            // instead in order to avoid duplicate service registration.
-            _hostBuilder.ConfigureWebHostDefaults(ConfigureWebHost);
+            _hostBuilder.ConfigureWebHost(ConfigureWebHost);
             return _builtApplication = new WebApplication(_hostBuilder.Build());
         }
 
@@ -189,24 +189,6 @@ namespace Microsoft.AspNetCore.Builder
                 builder.AddConfiguration(Configuration, shouldDisposeConfiguration: true);
             });
 
-            genericWebHostBuilder.ConfigureServices((context, services) =>
-            {
-                // We've only added services configured by the GenericWebHostBuilder and WebHost.ConfigureWebDefaults
-                // at this point. HostBuilder news up a new ServiceCollection in HostBuilder.Build() we haven't seen
-                // until now, so we cannot clear these services even though some are redundant because
-                // we called ConfigureWebHostDefaults on both the _deferredHostBuilder and _hostBuilder.
-
-                // Ideally, we'd only call _hostBuilder.ConfigureWebHost(ConfigureWebHost) instead of
-                // _hostBuilder.ConfigureWebHostDefaults(ConfigureWebHost) to avoid some duplicate service descriptors,
-                // but we want to add services in the WebApplicationBuilder constructor so code can inspect
-                // WebApplicationBuilder.Services. At the same time, we want to be able which services are loaded
-                // to react to config changes (e.g. ForwardedHeadersStartupFilter).
-                foreach (var s in Services)
-                {
-                    services.Add(s);
-                }
-            });
-
             genericWebHostBuilder.Configure(ConfigureApplication);
 
             _deferredHostBuilder.RunDeferredCallbacks(_hostBuilder);
diff --git a/src/DefaultBuilder/src/WebApplicationServiceCollection.cs b/src/DefaultBuilder/src/WebApplicationServiceCollection.cs
index 4d507a564c..26bba0e417 100644
--- a/src/DefaultBuilder/src/WebApplicationServiceCollection.cs
+++ b/src/DefaultBuilder/src/WebApplicationServiceCollection.cs
@@ -14,10 +14,13 @@ namespace Microsoft.AspNetCore
     internal class WebApplicationServiceCollection : IServiceCollection
     {
         private readonly IServiceCollection _serviceCollection;
+        private readonly List<Action<IServiceCollection>> _operations = new();
+        private bool _trackOperations;
 
-        public WebApplicationServiceCollection(ServiceCollection serviceCollection)
+        public WebApplicationServiceCollection()
         {
-            _serviceCollection = serviceCollection;
+            _serviceCollection = new ServiceCollection();
+            _trackOperations = true;
         }
 
         public ServiceDescriptor this[int index] { get => _serviceCollection[index]; set => _serviceCollection[index] = value; }
@@ -29,11 +32,21 @@ namespace Microsoft.AspNetCore
         public void Add(ServiceDescriptor item)
         {
             _serviceCollection.Add(item);
+
+            if (_trackOperations)
+            {
+                _operations.Add(sc => sc.Add(item));
+            }
         }
 
         public void Clear()
         {
             _serviceCollection.Clear();
+
+            if (_trackOperations)
+            {
+                _operations.Add(sc => sc.Clear());
+            }
         }
 
         public bool Contains(ServiceDescriptor item)
@@ -59,21 +72,57 @@ namespace Microsoft.AspNetCore
         public void Insert(int index, ServiceDescriptor item)
         {
             _serviceCollection.Insert(index, item);
+
+            if (_trackOperations)
+            {
+                _operations.Add(sc => sc.Insert(index, item));
+            }
         }
 
         public bool Remove(ServiceDescriptor item)
         {
-            return _serviceCollection.Remove(item);
+            var removed = _serviceCollection.Remove(item);
+            if (_trackOperations)
+            {
+                _operations.Add(sc => sc.Remove(item));
+            }
+            return removed;
         }
 
         public void RemoveAt(int index)
         {
             _serviceCollection.RemoveAt(index);
+            if (_trackOperations)
+            {
+                _operations.Add(sc => sc.RemoveAt(index));
+            }
         }
 
         IEnumerator IEnumerable.GetEnumerator()
         {
             return GetEnumerator();
         }
+
+        public List<Action<IServiceCollection>> GetOperations()
+        {
+            if (_operations.Count == 0)
+            {
+                return _operations;
+            }
+            // Copy the list
+            var operations = _operations.ToList();
+            _operations.Clear();
+            return operations;
+        }
+
+        public void StartTracking()
+        {
+            _trackOperations = true;
+        }
+
+        public void StopTracking()
+        {
+            _trackOperations = false;
+        }
     }
 }