Merge pull request #119 from autofac/bug/keyed-middleware-performance

Add optimisation for reflection-activated components to avoid adding a parameter in the resolve path if possible.

src/Autofac.Extensions.DependencyInjection/AutofacRegistration.cs

@@ -2,6 +2,10 @@
 // Licensed under the MIT License. See LICENSE in the project root for license information.
 using System.Reflection;
 using Autofac.Builder;
+using Autofac.Core;
+using Autofac.Core.Activators;
+using Autofac.Core.Activators.Delegate;
+using Autofac.Core.Activators.Reflection;
 using Autofac.Core.Resolving.Pipeline;
 using Microsoft.Extensions.DependencyInjection;
 
@@ -63,11 +67,16 @@ public static void Populate(
         IEnumerable<ServiceDescriptor> descriptors,
         object? lifetimeScopeTagForSingletons)
     {
-        if (descriptors == null)
+        if (descriptors is null)
         {
             throw new ArgumentNullException(nameof(descriptors));
         }
 
+        if (builder is null)
+        {
+            throw new ArgumentNullException(nameof(builder));
+        }
+
         builder.RegisterType<AutofacServiceProvider>()
                .As<IServiceProvider>()
                .As<IServiceProviderIsService>()
@@ -82,12 +91,80 @@ public static void Populate(
             .SingleInstance();
 
         // Shims for keyed service compatibility.
-        builder.RegisterServiceMiddlewareSource(new KeyedServiceMiddlewareSource());
         builder.RegisterSource<AnyKeyRegistrationSource>();
+        builder.ComponentRegistryBuilder.Registered += AddFromKeyedServiceParameterMiddleware;
 
         Register(builder, descriptors, lifetimeScopeTagForSingletons);
     }
 
+    /// <summary>
+    /// Inspect each component registration, and determine whether or not we can avoid adding the
+    /// <see cref="FromKeyedServicesAttribute"/> parameter to the resolve pipeline.
+    /// </summary>
+    private static void AddFromKeyedServiceParameterMiddleware(object? sender, ComponentRegisteredEventArgs e)
+    {
+        var needFromKeyedServiceParameter = false;
+
+        // We can optimise quite significantly in the case where we are using the reflection activator.
+        // In that state we can inspect the constructors ahead of time and determine whether the parameter will even need to be added.
+        if (e.ComponentRegistration.Activator is ReflectionActivator reflectionActivator)
+        {
+            var constructors = reflectionActivator.ConstructorFinder.FindConstructors(reflectionActivator.LimitType);
+
+            // Go through all the constructors; if any have a FromKeyedServices, then we must add our component middleware to
+            // the pipeline.
+            foreach (var constructor in constructors)
+            {
+                foreach (var constructorParameter in constructor.GetParameters())
+                {
+                    if (constructorParameter.GetCustomAttribute<FromKeyedServicesAttribute>() is not null)
+                    {
+                        // One or more of the constructors we will use to activate has a FromKeyedServicesAttribute,
+                        // we must add our middleware.
+                        needFromKeyedServiceParameter = true;
+                        break;
+                    }
+                }
+
+                if (needFromKeyedServiceParameter)
+                {
+                    break;
+                }
+            }
+        }
+        else if (e.ComponentRegistration.Activator is DelegateActivator)
+        {
+            // For delegate activation there are very few paths that would let the FromKeyedServicesAttribute
+            // actually work, and none that MSDI supports directly.
+            // We're explicitly choosing here not to support [FromKeyedServices] on the Autofac-specific generic
+            // delegate resolve methods, to improve performance for the 99% case of other delegates that only
+            // receive an IComponentContext or an IServiceProvider.
+            needFromKeyedServiceParameter = false;
+        }
+        else if (e.ComponentRegistration.Activator is InstanceActivator)
+        {
+            // Instance activators don't use parameters.
+            needFromKeyedServiceParameter = false;
+        }
+        else
+        {
+            // Unknown activator, assume we need the parameter.
+            needFromKeyedServiceParameter = true;
+        }
+
+        e.ComponentRegistration.PipelineBuilding += (sender, pipeline) =>
+        {
+            if (needFromKeyedServiceParameter)
+            {
+                pipeline.Use(KeyedServiceMiddleware.InstanceWithFromKeyedServicesParameter, MiddlewareInsertionMode.StartOfPhase);
+            }
+            else
+            {
+                pipeline.Use(KeyedServiceMiddleware.InstanceWithoutFromKeyedServicesParameter, MiddlewareInsertionMode.StartOfPhase);
+            }
+        };
+    }
+
     /// <summary>
     /// Configures the exposed service type on a service registration.
     /// </summary>

src/Autofac.Extensions.DependencyInjection/KeyedServiceMiddleware.cs

@@ -13,16 +13,54 @@ namespace Autofac.Extensions.DependencyInjection;
 /// </summary>
 internal class KeyedServiceMiddleware : IResolveMiddleware
 {
+    // [FromKeyedServices("key")] - Specifies a keyed service
+    // for injection into a constructor. This is similar to the
+    // Autofac [KeyFilter] attribute.
+    private static readonly Parameter FromKeyedServicesParameter = new ResolvedParameter(
+            (p, c) =>
+            {
+                var filter = p.GetCustomAttributes<FromKeyedServicesAttribute>(true).FirstOrDefault();
+                return filter is not null && filter.CanResolveParameter(p, c);
+            },
+            (p, c) =>
+            {
+                var filter = p.GetCustomAttributes<FromKeyedServicesAttribute>(true).First();
+                return filter.ResolveParameter(p, c);
+            });
+
+    /// <summary>
+    /// Gets a single instance of this middleware that does not add the keyed services parameter.
+    /// </summary>
+    public static KeyedServiceMiddleware InstanceWithoutFromKeyedServicesParameter { get; } = new(addFromKeyedServiceParameter: false);
+
+    /// <summary>
+    /// Gets a single instance of this middleware that adds the keyed services parameter.
+    /// </summary>
+    public static KeyedServiceMiddleware InstanceWithFromKeyedServicesParameter { get; } = new(addFromKeyedServiceParameter: true);
+
+    private readonly bool _addFromKeyedServiceParameter;
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="KeyedServiceMiddleware"/> class.
+    /// </summary>
+    /// <param name="addFromKeyedServiceParameter">Whether or not the from-keyed-service parameter should be added.</param>
+    public KeyedServiceMiddleware(bool addFromKeyedServiceParameter)
+    {
+        _addFromKeyedServiceParameter = addFromKeyedServiceParameter;
+    }
+
     /// <inheritdoc />
-    public PipelinePhase Phase => PipelinePhase.ResolveRequestStart;
+    public PipelinePhase Phase => PipelinePhase.Activation;
 
     /// <inheritdoc />
     public void Execute(ResolveRequestContext context, Action<ResolveRequestContext> next)
     {
-        var newParameters = new List<Parameter>(context.Parameters);
+        List<Parameter>? newParameters = null;
 
         if (context.Service is Autofac.Core.KeyedService keyedService)
         {
+            newParameters = new List<Parameter>(context.Parameters);
+
             var key = keyedService.ServiceKey;
 
             // [ServiceKey] - indicates that the parameter value should
@@ -41,22 +79,20 @@ public void Execute(ResolveRequestContext context, Action<ResolveRequestContext>
                 }));
         }
 
-        // [FromKeyedServices("key")] - Specifies a keyed service
-        // for injection into a constructor. This is similar to the
-        // Autofac [KeyFilter] attribute.
-        newParameters.Add(new ResolvedParameter(
-            (p, c) =>
-            {
-                var filter = p.GetCustomAttributes<FromKeyedServicesAttribute>(true).FirstOrDefault();
-                return filter is not null && filter.CanResolveParameter(p, c);
-            },
-            (p, c) =>
-            {
-                var filter = p.GetCustomAttributes<FromKeyedServicesAttribute>(true).First();
-                return filter.ResolveParameter(p, c);
-            }));
+        if (_addFromKeyedServiceParameter)
+        {
+            newParameters ??= new List<Parameter>(context.Parameters);
 
-        context.ChangeParameters(newParameters);
+            // [FromKeyedServices("key")] - Specifies a keyed service
+            // for injection into a constructor. This is similar to the
+            // Autofac [KeyFilter] attribute.
+            newParameters.Add(FromKeyedServicesParameter);
+        }
+
+        if (newParameters is not null)
+        {
+            context.ChangeParameters(newParameters);
+        }
 
         next(context);
     }

test/Autofac.Extensions.DependencyInjection.Test/AutofacRegistrationTests.cs

@@ -21,6 +21,18 @@ public void PopulateRegistersServiceProvider()
         container.AssertRegistered<IServiceProvider>();
     }
 
+    [Fact]
+    public void PopulateThrowsForNullBuilder()
+    {
+        Assert.Throws<ArgumentNullException>(() => AutofacRegistration.Populate(null, Enumerable.Empty<ServiceDescriptor>()));
+    }
+
+    [Fact]
+    public void PopulateThrowsForNullDescriptors()
+    {
+        Assert.Throws<ArgumentNullException>(() => new ContainerBuilder().Populate(null));
+    }
+
     [Fact]
     public void CorrectServiceProviderIsRegistered()
     {