Support for [FromKeyedServices] in both handlers and HTTP endpoints

src/Http/Wolverine.Http.Tests/using_keyed_services.cs

@@ -0,0 +1,20 @@
+namespace Wolverine.Http.Tests;
+
+public class using_keyed_services : IntegrationContext
+{
+    public using_keyed_services(AppFixture fixture) : base(fixture)
+    {
+    }
+
+    [Fact]
+    public async Task using_singleton_service()
+    {
+        await Scenario(x => x.Get.Url("/thing/red"));
+    }
+
+    [Fact]
+    public async Task using_scoped_service()
+    {
+        await Scenario(x => x.Get.Url("/thing/blue"));
+    }
+}

src/Http/WolverineWebApi/Program.cs

@@ -63,6 +63,9 @@
 builder.Services.AddDbContextWithWolverineIntegration<ItemsDbContext>(
     x => x.UseNpgsql(Servers.PostgresConnectionString));
 
+builder.Services.AddKeyedSingleton<IThing, RedThing>("Red");
+builder.Services.AddKeyedScoped<IThing, BlueThing>("Blue");
+builder.Services.AddKeyedTransient<IThing, GreenThing>("Green");
 
 builder.Services.AddMarten(opts =>
 {

src/Testing/CoreTests/Acceptance/using_with_keyed_services.cs

@@ -0,0 +1,161 @@
+using System.ComponentModel.Design;
+using System.Diagnostics;
+using CoreTests.Codegen;
+using Lamar.Microsoft.DependencyInjection;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Hosting;
+using Wolverine.Attributes;
+using Xunit;
+
+namespace CoreTests.Acceptance;
+
+public class using_with_keyed_services : IAsyncLifetime
+{
+    private IHost _host;
+
+    public async Task InitializeAsync()
+    {
+        _host = await Host.CreateDefaultBuilder()
+            .UseWolverine(opts =>
+            {
+                opts.Discovery.DisableConventionalDiscovery().IncludeType<ThingHandler>();
+
+                opts.Services.AddKeyedSingleton<IThing, RedThing>("Red");
+                opts.Services.AddKeyedScoped<IThing, BlueThing>("Blue");
+                opts.Services.AddKeyedSingleton<IThing, GreenThing>("Green");
+
+                opts.Services.AddTransient<ThingUser>();
+                opts.Services.AddTransient<ThingUserHolder>();
+            }).StartAsync();
+    }
+
+    public Task DisposeAsync()
+    {
+        return _host.StopAsync();
+    }
+
+    [Fact]
+    public async Task use_inside_of_deep_dependency_chain()
+    {
+        var container = _host.Services.GetRequiredService<Wolverine.Runtime.IServiceContainer>();
+        var holder = container.QuickBuild<ThingUserHolder>();
+        holder.ThingUser.Thing.ShouldBeOfType<RedThing>();
+
+        await _host.InvokeAsync(new UseThingHolder());
+    }
+
+    [Fact]
+    public async Task use_as_single_parameter_on_handler()
+    {
+        await _host.InvokeAsync(new UseThingDirectly());
+    }
+
+    [Fact]
+    public async Task use_as_singleton_parameter_on_handler()
+    {
+        await _host.InvokeAsync(new UseSingletonThingDirectly());
+    }
+
+    [Fact]
+    public async Task use_multiple_parameters_on_handler()
+    {
+        await _host.InvokeAsync(new UseMultipleThings());
+    }
+}
+
+public class using_with_keyed_services_and_lamar : IAsyncLifetime
+{
+    private IHost _host;
+
+    public async Task InitializeAsync()
+    {
+        _host = await Host.CreateDefaultBuilder()
+            .UseLamar()
+            .UseWolverine(opts =>
+            {
+                opts.Discovery.DisableConventionalDiscovery().IncludeType<ThingHandler>();
+
+                opts.Services.AddKeyedSingleton<IThing, RedThing>("Red");
+                opts.Services.AddKeyedScoped<IThing, BlueThing>("Blue");
+                opts.Services.AddKeyedSingleton<IThing, GreenThing>("Green");
+
+                opts.Services.AddTransient<ThingUser>();
+                opts.Services.AddTransient<ThingUserHolder>();
+            }).StartAsync();
+    }
+
+    public Task DisposeAsync()
+    {
+        return _host.StopAsync();
+    }
+
+    [Fact]
+    public async Task use_inside_of_deep_dependency_chain()
+    {
+        var container = _host.Services.GetRequiredService<Wolverine.Runtime.IServiceContainer>();
+        var holder = container.QuickBuild<ThingUserHolder>();
+        holder.ThingUser.Thing.ShouldBeOfType<RedThing>();
+
+        await _host.InvokeAsync(new UseThingHolder());
+    }
+
+    [Fact]
+    public async Task use_as_single_parameter_on_handler()
+    {
+        await _host.InvokeAsync(new UseThingDirectly());
+    }
+
+    [Fact]
+    public async Task use_as_singleton_parameter_on_handler()
+    {
+        await _host.InvokeAsync(new UseSingletonThingDirectly());
+    }
+
+    [Fact]
+    public async Task use_multiple_parameters_on_handler()
+    {
+        await _host.InvokeAsync(new UseMultipleThings());
+    }
+}
+
+public interface IThing;
+public class RedThing : IThing;
+public class BlueThing : IThing;
+public class GreenThing : IThing;
+
+public record ThingUser([FromKeyedServices("Red")] IThing Thing);
+
+public record ThingUserHolder(ThingUser ThingUser);
+
+public record UseThingDirectly;
+public record UseSingletonThingDirectly;
+
+public record UseMultipleThings;
+
+public record UseThingHolder;
+
+[WolverineIgnore]
+public class ThingHandler
+{
+    public static void Handle(UseThingHolder command, ThingUserHolder holder)
+    {
+        holder.ThingUser.Thing.ShouldBeOfType<RedThing>();
+    }
+
+    public static void Handle(UseThingDirectly command, [FromKeyedServices("Blue")] IThing thing)
+    {
+        thing.ShouldBeOfType<BlueThing>();
+    }
+
+    public static void Handle(UseSingletonThingDirectly command, [FromKeyedServices("Red")] IThing thing)
+    {
+        thing.ShouldBeOfType<RedThing>();
+    }
+
+    public static void Handle(UseMultipleThings command, [FromKeyedServices("Green")] IThing green,
+        [FromKeyedServices("Red")] IThing red)
+    {
+        green.ShouldBeOfType<GreenThing>();
+        red.ShouldBeOfType<RedThing>();
+    }
+}

src/Testing/CoreTests/CoreTests.csproj

@@ -5,7 +5,7 @@
     </PropertyGroup>
 
     <ItemGroup>
-        <PackageReference Include="Lamar.Microsoft.DependencyInjection" Version="15.0.0-alpha-4" />
+        <PackageReference Include="Lamar.Microsoft.DependencyInjection" Version="15.0.0-rc-1" />
         <PackageReference Include="Microsoft.FeatureManagement" Version="3.2.0"/>
         <PackageReference Include="Microsoft.NET.Test.Sdk" Version="15.9.0"/>
         <PackageReference Include="OpenTelemetry.Exporter.Console" Version="1.6.0"/>

src/Wolverine/Codegen/ConstructorPlan.cs

@@ -47,10 +47,13 @@ public static bool TryBuildPlan(List<ServiceDescriptor> trail, ServiceDescriptor
         }
 
         trail.Add(descriptor);
+
+        var implementationType = descriptor.IsKeyedService ? descriptor.KeyedImplementationType : descriptor.ImplementationType;
 
         try
         {
-            var constructors = FindPublicConstructorCandidates(descriptor.ImplementationType);
+
+            var constructors = FindPublicConstructorCandidates(implementationType);
 
             // If no public constructors, get out of here
             if (!constructors.Any())
@@ -65,13 +68,43 @@ public static bool TryBuildPlan(List<ServiceDescriptor> trail, ServiceDescriptor
                 return true;
             }
 
+            Func<ParameterInfo, bool> hasPlan = parameter =>
+            {
+                if (parameter.TryGetAttribute<FromKeyedServicesAttribute>(out var att))
+                {
+                    var descriptor = graph
+                        .RegistrationsFor(parameter.ParameterType)
+                        .Where(x => x.IsKeyedService)
+                        .FirstOrDefault(x => x.ServiceKey.ToString() == att.Key.ToString());
+
+                    if (descriptor == null) return false;
+
+                    return graph.PlanFor(descriptor, trail) is not InvalidPlan;
+                }
+
+                return graph.FindDefault(parameter.ParameterType, trail) is not InvalidPlan;
+            };
+
             var constructor = constructors
-                .Where(x => x.GetParameters().All(p => graph.FindDefault(p.ParameterType, trail) is not InvalidPlan))
+                .Where(x => x.GetParameters().All(x => hasPlan(x)))
                 .MinBy(x => x.GetParameters().Length);
 
             if (constructor != null)
             {
-                var dependencies = constructor.GetParameters().Select(x => graph.FindDefault(x.ParameterType, trail)).ToArray();
+                var dependencies = constructor.GetParameters().Select(x =>
+                {
+                    if (x.TryGetAttribute<FromKeyedServicesAttribute>(out var att))
+                    {
+                        var descriptor = graph
+                            .RegistrationsFor(x.ParameterType)
+                            .Where(x => x.IsKeyedService)
+                            .FirstOrDefault(x => x.ServiceKey.ToString() == att.Key.ToString());
+
+                        return graph.PlanFor(descriptor, trail);
+                    }
+
+                    return graph.FindDefault(x.ParameterType, trail);
+                }).ToArray();
 
                 if (dependencies.OfType<InvalidPlan>().Any() || dependencies.Any(x => x == null))
                 {
@@ -99,10 +132,6 @@ public ConstructorPlan(ServiceDescriptor descriptor, ConstructorInfo constructor
         if (descriptor.Lifetime == ServiceLifetime.Singleton)
             throw new ArgumentOutOfRangeException(nameof(descriptor),
                 $"Only {ServiceLifetime.Scoped} or {ServiceLifetime.Transient} lifecycles are valid");
-#if NET8_0_OR_GREATER
-
-        if (descriptor.IsKeyedService) throw new ArgumentOutOfRangeException(nameof(descriptor), "Cannot yet support keyed services");
-#endif
 
         if (constructor.GetParameters().Length != dependencies.Length)
         {
@@ -159,8 +188,9 @@ public override string WhyRequireServiceProvider(IMethodVariables method)
 
     public override Variable CreateVariable(ServiceVariables resolverVariables)
     {
-        var frame = new ConstructorFrame(Descriptor.ImplementationType, Constructor);
-        if (Descriptor.ImplementationType.CanBeCastTo<IDisposable>() || Descriptor.ImplementationType.CanBeCastTo<IAsyncDisposable>())
+        var implementationType = Descriptor.IsKeyedService ? Descriptor.KeyedImplementationType : Descriptor.ImplementationType;
+        var frame = new ConstructorFrame(implementationType, Constructor);
+        if (implementationType.CanBeCastTo<IDisposable>() || implementationType.CanBeCastTo<IAsyncDisposable>())
         {
             frame.Mode = ConstructorCallMode.UsingNestedVariable;
         }

src/Wolverine/Codegen/InjectedSingleton.cs

@@ -1,4 +1,5 @@
 using JasperFx.CodeGeneration.Model;
+using JasperFx.Core.Reflection;
 using Microsoft.Extensions.DependencyInjection;
 
 namespace Wolverine.Codegen;
@@ -12,7 +13,18 @@ public InjectedSingleton(ServiceDescriptor descriptor) : base(descriptor.Service
     {
         Descriptor = descriptor;
     }
-
+
+    public override string CtorArgDeclaration()
+    {
+        if (Descriptor.IsKeyedService)
+        {
+            return $"[{typeof(FromKeyedServicesAttribute).FullNameInCode().Replace("Attribute", "")}(\"{Descriptor.ServiceKey}\")] " +
+                   base.CtorArgDeclaration();
+        }
+
+        return base.CtorArgDeclaration();
+    }
+
     public bool IsOnlyOne
     {
         private get => _isOnlyOne;
@@ -27,4 +39,34 @@ public bool IsOnlyOne
             }
         }
     }
+
+    protected bool Equals(InjectedSingleton other)
+    {
+        return base.Equals(other) && Descriptor.Equals(other.Descriptor);
+    }
+
+    public override bool Equals(object? obj)
+    {
+        if (obj is null)
+        {
+            return false;
+        }
+
+        if (ReferenceEquals(this, obj))
+        {
+            return true;
+        }
+
+        if (obj.GetType() != GetType())
+        {
+            return false;
+        }
+
+        return Equals((InjectedSingleton)obj);
+    }
+
+    public override int GetHashCode()
+    {
+        return HashCode.Combine(base.GetHashCode(), Descriptor);
+    }
 }