TryGetKeyedServiceKey<T>; allow [ServiceKey] injection during AnyKey resolution

src/Autofac/Features/Collections/CollectionRegistrationSource.cs

@@ -119,41 +119,12 @@ public IEnumerable<IComponentRegistration> RegistrationsFor(Service service, Fun
             limitType,
             (c, p) =>
             {
-                List<ServiceRegistration> itemRegistrations;
+                var registrationTuples = isAnyKeyQuery
+                    ? GetAllSpecificKeyedRegistrations(c.ComponentRegistry, elementType)
+                        .ConvertAll(static tuple => ((Service)tuple.KeyedService, tuple.Registration))
+                    : BuildStandardRegistrationList(c.ComponentRegistry, elementTypeService);
 
-                if (isAnyKeyQuery)
-                {
-                    // AnyKey queries for collections return _all_ keyed services.
-                    itemRegistrations = GetAllSpecificKeyedRegistrations(c.ComponentRegistry, elementType);
-                }
-                else
-                {
-                    itemRegistrations = c.ComponentRegistry
-                        .ServiceRegistrationsFor(elementTypeService)
-                        .Where(cr => !cr.Registration.Options.HasOption(RegistrationOptions.ExcludeFromCollections))
-                        .OrderBy(cr => cr.Registration.GetRegistrationOrder())
-                        .ToList();
-                }
-
-                var output = factory(itemRegistrations.Count);
-                var isFixedSize = output.IsFixedSize;
-
-                for (var i = 0; i < itemRegistrations.Count; i++)
-                {
-                    var itemRegistration = itemRegistrations[i];
-                    var resolveRequest = new ResolveRequest(elementTypeService, itemRegistration, p);
-                    var component = c.ResolveComponent(resolveRequest);
-                    if (isFixedSize)
-                    {
-                        output[i] = component;
-                    }
-                    else
-                    {
-                        output.Add(component);
-                    }
-                }
-
-                return output;
+                return BuildCollection(c, factory, registrationTuples, p);
             });
 
         var registration = new ComponentRegistration(
@@ -190,9 +161,27 @@ private static Func<int, IList> GenerateArrayFactory(Type elementType)
         return Expression.Lambda<Func<int, IList>>(newArray, parameter).Compile();
     }
 
-    private static List<ServiceRegistration> GetAllSpecificKeyedRegistrations(IComponentRegistry registry, Type elementType)
+    /// <summary>
+    /// When the query is for "any keyed" enumerable, we need to find all the
+    /// specific keyed registrations and return them.
+    /// </summary>
+    /// <param name="registry">
+    /// The registry to search for registrations. We need to search the entire
+    /// registry because "any keyed" could match any specific key.
+    /// </param>
+    /// <param name="elementType">
+    /// The element type of the enumerable being resolved. We need this to
+    /// filter the registry down to only the relevant registrations.
+    /// </param>
+    /// <returns>
+    /// A list of tuples containing the specific keyed service and the
+    /// registration for each matching registration. We return the specific
+    /// keyed service so that we can issue resolve requests that still know the
+    /// original key.
+    /// </returns>
+    private static List<(KeyedService KeyedService, ServiceRegistration Registration)> GetAllSpecificKeyedRegistrations(IComponentRegistry registry, Type elementType)
     {
-        var result = new List<ServiceRegistration>();
+        var result = new List<(KeyedService, ServiceRegistration)>();
         var processedServices = new HashSet<KeyedService>();
 
         foreach (var registration in registry.Registrations)
@@ -220,12 +209,55 @@ private static List<ServiceRegistration> GetAllSpecificKeyedRegistrations(ICompo
                         !cr.Registration.Options.HasOption(RegistrationOptions.ExcludeFromCollections) &&
                         !cr.Registration.Metadata.ContainsKey(MetadataKeys.AnyKeyAdapter));
 
-                result.AddRange(serviceRegistrations);
+                foreach (var serviceRegistration in serviceRegistrations)
+                {
+                    // Return both the keyed service and the registration so callers can issue
+                    // resolve requests that still know the original key.
+                    result.Add((keyed, serviceRegistration));
+                }
             }
         }
 
         return result
+            .OrderBy(tuple => tuple.Item2.Registration.GetRegistrationOrder())
+            .ToList();
+    }
+
+    private static List<(Service Service, ServiceRegistration Registration)> BuildStandardRegistrationList(IComponentRegistry registry, Service elementTypeService)
+    {
+        return registry
+            .ServiceRegistrationsFor(elementTypeService)
+            .Where(cr => !cr.Registration.Options.HasOption(RegistrationOptions.ExcludeFromCollections))
             .OrderBy(cr => cr.Registration.GetRegistrationOrder())
+            .Select(cr => ((Service)elementTypeService, cr))
             .ToList();
     }
+
+    private static IList BuildCollection(
+        IComponentContext context,
+        Func<int, IList> factory,
+        List<(Service Service, ServiceRegistration Registration)> registrations,
+        IEnumerable<Parameter> parameters)
+    {
+        var output = factory(registrations.Count);
+        var isFixedSize = output.IsFixedSize;
+
+        for (var i = 0; i < registrations.Count; i++)
+        {
+            var (service, registration) = registrations[i];
+            var resolveRequest = new ResolveRequest(service, registration, parameters);
+            var component = context.ResolveComponent(resolveRequest);
+
+            if (isFixedSize)
+            {
+                output[i] = component;
+            }
+            else
+            {
+                output.Add(component);
+            }
+        }
+
+        return output;
+    }
 }

src/Autofac/ParameterExtensions.cs

@@ -95,8 +95,24 @@ public static T TypedAs<T>(this IEnumerable<Parameter> parameters)
     /// <typeparam name="T">The type to which the returned value will be cast.</typeparam>
     /// <param name="parameters">The available parameters to choose from.</param>
     /// <returns>The value of the keyed service key.</returns>
-    /// <seealso cref="KeyedServiceParameterInjector"/>
     public static T KeyedServiceKey<T>(this IEnumerable<Parameter> parameters)
+    {
+        if (!TryGetKeyedServiceKey(parameters, out T value))
+        {
+            throw new InvalidOperationException(ResolutionExtensionsResources.KeyedServiceKeyUnavailable);
+        }
+
+        return value;
+    }
+
+    /// <summary>
+    /// Attempts to retrieve the keyed service key value associated with the current resolve operation.
+    /// </summary>
+    /// <typeparam name="T">The type to which the returned value will be cast.</typeparam>
+    /// <param name="parameters">The available parameters to choose from.</param>
+    /// <param name="value">The value of the keyed service key.</param>
+    /// <returns><see langword="true"/> if a keyed service key is available; otherwise, <see langword="false"/>.</returns>
+    public static bool TryGetKeyedServiceKey<T>(this IEnumerable<Parameter> parameters, [NotNullWhen(true)] out T value)
     {
         if (parameters == null)
         {
@@ -107,11 +123,13 @@ public static T KeyedServiceKey<T>(this IEnumerable<Parameter> parameters)
         {
             if (parameter is KeyedServiceKeyParameter keyParameter)
             {
-                return (T)keyParameter.ServiceKey;
+                value = (T)keyParameter.ServiceKey;
+                return true;
             }
         }
 
-        throw new InvalidOperationException(ResolutionExtensionsResources.KeyedServiceKeyUnavailable);
+        value = default!;
+        return false;
     }
 
     private static TValue ConstantValue<TParameter, TValue>(IEnumerable<Parameter> parameters, Func<TParameter, bool> predicate)

test/Autofac.Specification.Test/Features/KeyedServiceTests.cs

@@ -901,6 +901,20 @@ void VerifyException()
         }
     }
 
+    [Fact]
+    public void ResolveAnyKeyWithInjectedKeyedParameter()
+    {
+        var builder = new ContainerBuilder();
+        builder.RegisterType<Service>().Keyed<IService>("a");
+        builder.RegisterType<Service>().Keyed<IService>("b");
+        var provider = builder.Build();
+
+        var services = provider.ResolveKeyed<IEnumerable<IService>>(KeyedService.AnyKey).ToList();
+        Assert.Equal(2, services.Count);
+        Assert.Equal("a", services[0].ToString());
+        Assert.Equal("b", services[1].ToString());
+    }
+
     private interface IService
     {
     }

test/Autofac.Test/Core/ParameterExtensionsTests.cs

@@ -8,7 +8,7 @@ namespace Autofac.Test.Core;
 public class ParameterExtensionsTests
 {
     [Fact]
-    public void KeyedServiceKey_ReturnsValue()
+    public void KeyedServiceKey_Found()
     {
         var result = new Parameter[] { new KeyedServiceKeyParameter("expected") }
             .KeyedServiceKey<string>();
@@ -22,4 +22,26 @@ public void KeyedServiceKey_NotFound()
         Assert.Throws<InvalidOperationException>(
             () => Array.Empty<Parameter>().KeyedServiceKey<string>());
     }
+
+    [Fact]
+    public void TryGetKeyedServiceKey_Found()
+    {
+        var parameters = new Parameter[] { new KeyedServiceKeyParameter("expected") };
+
+        var result = parameters.TryGetKeyedServiceKey(out string value);
+
+        Assert.True(result);
+        Assert.Equal("expected", value);
+    }
+
+    [Fact]
+    public void TryGetKeyedServiceKey_NotFound()
+    {
+        var parameters = Array.Empty<Parameter>();
+
+        var result = parameters.TryGetKeyedServiceKey(out string value);
+
+        Assert.False(result);
+        Assert.Null(value);
+    }
 }