diff --git a/src/CodegenTests/Services/keyed_service_location.cs b/src/CodegenTests/Services/keyed_service_location.cs new file mode 100644 index 0000000..af92e29 --- /dev/null +++ b/src/CodegenTests/Services/keyed_service_location.cs @@ -0,0 +1,105 @@ +using JasperFx; +using JasperFx.CodeGeneration; +using JasperFx.CodeGeneration.Frames; +using JasperFx.CodeGeneration.Services; +using JasperFx.RuntimeCompiler; +using Microsoft.Extensions.DependencyInjection; +using Shouldly; +using Xunit.Abstractions; + +namespace CodegenTests.Services; + +// GH-2878 (JasperFx/wolverine): keyed services resolved through the scoped-IServiceProvider +// "service location" path lost their key — the frame always emitted GetRequiredService(provider) +// instead of GetRequiredKeyedService(provider, key). Service location kicks in whenever something +// in the same generated method forces it (a directly-injected IServiceProvider, or an opaque lambda +// registration like the ones MS Graph adds), which then drags every sibling dependency — including +// keyed ones — onto the service-location path where the key was being dropped. +public class keyed_service_location +{ + private readonly ITestOutputHelper _output; + private readonly ServiceCollection theServices = new(); + + public keyed_service_location(ITestOutputHelper output) + { + _output = output; + } + + // Generates a harness whose Build() calls KeyedHandler.Handle(...), exactly how Wolverine + // generates a handler: the keyed IWidget parameter is resolved via [FromKeyedServices] during + // code generation, and the opaque IScopedLambda sibling forces the method onto service location. + private (string Code, ServiceContainer Graph, GeneratedType Type) generateHandlerCall() + { + var graph = new ServiceContainer(theServices, theServices.BuildServiceProvider()); + + var assembly = new GeneratedAssembly(new GenerationRules()); + var type = assembly.AddType("KeyedHandlerHarness", typeof(ServiceHarness)); + var buildMethod = type.MethodFor("Build"); + + var call = new MethodCall(typeof(KeyedHandler), nameof(KeyedHandler.Handle)); + buildMethod.Frames.Add(call); + buildMethod.Frames.Code("return {0};", call.ReturnVariable!); + + var source = new ServiceCollectionServerVariableSource(graph); + source.StartNewType(); + source.StartNewMethod(); + + var code = assembly.GenerateCode(source); + _output.WriteLine(code); + return (code, graph, type); + } + + [Fact] + public void keyed_concrete_service_dragged_onto_location_keeps_its_key() + { + theServices.AddKeyedScoped("blue"); + theServices.AddScoped(_ => new ScopedLambda()); + + var (code, _, _) = generateHandlerCall(); + + code.ShouldContain("GetRequiredKeyedService"); + code.ShouldContain("\"blue\""); + code.ShouldNotContain("GetRequiredService"); + } + + [Fact] + public void keyed_opaque_registration_keeps_its_key() + { + theServices.AddKeyedScoped("blue", (_, _) => new BWidget()); + theServices.AddScoped(_ => new ScopedLambda()); + + var (code, _, _) = generateHandlerCall(); + + code.ShouldContain("GetRequiredKeyedService"); + code.ShouldContain("\"blue\""); + } + + [Fact] + public void keyed_service_resolves_correctly_at_runtime() + { + // End-to-end: compile the generated harness and prove the keyed service actually resolves + // (before the fix this threw "No service for type IWidget has been registered"). + theServices.AddKeyedScoped("blue"); + theServices.AddScoped(_ => new ScopedLambda()); + + var (code, graph, _) = generateHandlerCall(); + + var compiler = new AssemblyGenerator(); + compiler.ReferenceAssembly(GetType().Assembly); + var builtAssembly = compiler.Generate(code); + var builtType = builtAssembly.ExportedTypes.Single(); + + var result = ((ServiceHarness)graph.BuildFromType(builtType)).Build(); + result.Widget.ShouldBeOfType(); + } +} + +public class KeyedHandler +{ + public static WidgetResult Handle([FromKeyedServices("blue")] IWidget widget, IScopedLambda opaque) + { + return new WidgetResult(widget); + } +} + +public record WidgetResult(IWidget Widget); diff --git a/src/JasperFx/CodeGeneration/Services/GetServiceFromScopedContainerFrame.cs b/src/JasperFx/CodeGeneration/Services/GetServiceFromScopedContainerFrame.cs index b1362bb..001d734 100644 --- a/src/JasperFx/CodeGeneration/Services/GetServiceFromScopedContainerFrame.cs +++ b/src/JasperFx/CodeGeneration/Services/GetServiceFromScopedContainerFrame.cs @@ -8,17 +8,19 @@ namespace JasperFx.CodeGeneration.Services; public class GetServiceFromScopedContainerFrame : SyncFrame { private readonly Variable _scoped; + private readonly object? _serviceKey; - public GetServiceFromScopedContainerFrame(Variable scoped, Type serviceType) + public GetServiceFromScopedContainerFrame(Variable scoped, Type serviceType, object? serviceKey = null) { if (scoped.VariableType != typeof(IServiceProvider)) { throw new ArgumentOutOfRangeException(nameof(scoped), $"Wrong type for the variable. Expected {typeof(IServiceProvider).FullNameInCode()} but got {scoped.VariableType.FullNameInCode()}"); } - + _scoped = scoped; + _serviceKey = serviceKey; uses.Add(_scoped); Variable = new Variable(serviceType, this); @@ -60,8 +62,19 @@ public override void GenerateCode(GeneratedMethod method, ISourceWriter writer) Header.Write(writer); } - writer.Write( - $"var {Variable.Usage} = {typeof(ServiceProviderServiceExtensions).FullNameInCode()}.{nameof(ServiceProviderServiceExtensions.GetRequiredService)}<{Variable.VariableType.FullNameInCode()}>({_scoped.Usage});"); + if (_serviceKey != null) + { + // Keyed service resolved through service location. Without the key this would emit + // GetRequiredService and either throw or resolve the wrong registration. See GH-2878. + writer.Write( + $"var {Variable.Usage} = {typeof(ServiceProviderKeyedServiceExtensions).FullNameInCode()}.{nameof(ServiceProviderKeyedServiceExtensions.GetRequiredKeyedService)}<{Variable.VariableType.FullNameInCode()}>({_scoped.Usage}, {CodeFormatter.Write(_serviceKey)});"); + } + else + { + writer.Write( + $"var {Variable.Usage} = {typeof(ServiceProviderServiceExtensions).FullNameInCode()}.{nameof(ServiceProviderServiceExtensions.GetRequiredService)}<{Variable.VariableType.FullNameInCode()}>({_scoped.Usage});"); + } + Next?.GenerateCode(method, writer); } } \ No newline at end of file diff --git a/src/JasperFx/CodeGeneration/Services/ServiceCollectionServerVariableSource.cs b/src/JasperFx/CodeGeneration/Services/ServiceCollectionServerVariableSource.cs index 845b1a5..8c23a26 100644 --- a/src/JasperFx/CodeGeneration/Services/ServiceCollectionServerVariableSource.cs +++ b/src/JasperFx/CodeGeneration/Services/ServiceCollectionServerVariableSource.cs @@ -153,7 +153,11 @@ private void useServiceProvider(IMethodVariables method) var written = false; foreach (var standin in _standins) { - var frame = new GetServiceFromScopedContainerFrame(_scoped, standin.VariableType); + // Keyed services must keep their key when dragged onto the service-location path, + // otherwise the generated code emits GetRequiredService and loses the key. See GH-2878. + var descriptor = standin.Plan.Descriptor; + var serviceKey = descriptor is { IsKeyedService: true } ? descriptor.ServiceKey : null; + var frame = new GetServiceFromScopedContainerFrame(_scoped, standin.VariableType, serviceKey); var variable = frame.Variable; // Write description of why this had to use the nested container diff --git a/src/JasperFx/CodeGeneration/Services/ServiceLocationPlan.cs b/src/JasperFx/CodeGeneration/Services/ServiceLocationPlan.cs index cca94d6..faa0382 100644 --- a/src/JasperFx/CodeGeneration/Services/ServiceLocationPlan.cs +++ b/src/JasperFx/CodeGeneration/Services/ServiceLocationPlan.cs @@ -40,7 +40,8 @@ public override string WhyRequireServiceProvider(IMethodVariables method) public override Variable CreateVariable(ServiceVariables resolverVariables) { - return new GetServiceFromScopedContainerFrame(resolverVariables.ServiceProvider, Descriptor.ServiceType) + var serviceKey = Descriptor.IsKeyedService ? Descriptor.ServiceKey : null; + return new GetServiceFromScopedContainerFrame(resolverVariables.ServiceProvider, Descriptor.ServiceType, serviceKey) .Variable; } } \ No newline at end of file