Leverage keyed services for decoration instead of DecoratedType hack

khellang · Aug 18, 2023 · 096a14f · 096a14f
1 parent 6955471
commit 096a14f
Show file tree

Hide file tree

Showing 7 changed files with 83 additions and 147 deletions.
diff --git a/src/Scrutor/ClosedTypeDecorationStrategy.cs b/src/Scrutor/ClosedTypeDecorationStrategy.cs
@@ -16,16 +16,16 @@ public ClosedTypeDecorationStrategy(Type serviceType, Type? decoratorType, Func<
 
     public override bool CanDecorate(Type serviceType) => ServiceType == serviceType;
 
-    public override Func<IServiceProvider, object> CreateDecorator(Type serviceType)
+    public override Func<IServiceProvider, object?, object> CreateDecorator(Type serviceType, string serviceKey)
     {
         if (DecoratorType is not null)
         {
-            return TypeDecorator(serviceType, DecoratorType);
+            return TypeDecorator(serviceType, serviceKey, DecoratorType);
         }
 
         if (DecoratorFactory is not null)
         {
-            return FactoryDecorator(serviceType, DecoratorFactory);
+            return FactoryDecorator(serviceType, serviceKey, DecoratorFactory);
         }
 
         throw new InvalidOperationException($"Both serviceType and decoratorFactory can not be null.");

diff --git a/src/Scrutor/DecoratedType.cs b/src/Scrutor/DecoratedType.cs
diff --git a/src/Scrutor/DecorationStrategy.cs b/src/Scrutor/DecorationStrategy.cs
@@ -14,23 +14,23 @@ protected DecorationStrategy(Type serviceType)
 
     public abstract bool CanDecorate(Type serviceType);
 
-    public abstract Func<IServiceProvider, object> CreateDecorator(Type serviceType);
+    public abstract Func<IServiceProvider, object?, object> CreateDecorator(Type serviceType, string serviceKey);
 
     internal static DecorationStrategy WithType(Type serviceType, Type decoratorType) =>
         Create(serviceType, decoratorType, decoratorFactory: null);
 
     internal static DecorationStrategy WithFactory(Type serviceType, Func<object, IServiceProvider, object> decoratorFactory) =>
         Create(serviceType, decoratorType: null, decoratorFactory);
 
-    protected static Func<IServiceProvider, object> TypeDecorator(Type serviceType, Type decoratorType) => serviceProvider =>
+    protected static Func<IServiceProvider, object?, object> TypeDecorator(Type serviceType, string serviceKey, Type decoratorType) => (serviceProvider, _) =>
     {
-        var instanceToDecorate = serviceProvider.GetRequiredService(serviceType);
+        var instanceToDecorate = serviceProvider.GetRequiredKeyedService(serviceType, serviceKey);
         return ActivatorUtilities.CreateInstance(serviceProvider, decoratorType, instanceToDecorate);
     };
 
-    protected static Func<IServiceProvider, object> FactoryDecorator(Type decorated, Func<object, IServiceProvider, object> decoratorFactory) => serviceProvider =>
+    protected static Func<IServiceProvider, object?, object> FactoryDecorator(Type serviceType, string serviceKey, Func<object, IServiceProvider, object> decoratorFactory) => (serviceProvider, _) =>
     {
-        var instanceToDecorate = serviceProvider.GetRequiredService(decorated);
+        var instanceToDecorate = serviceProvider.GetRequiredKeyedService(serviceType, serviceKey);
         return decoratorFactory(instanceToDecorate, serviceProvider);
     };
 

diff --git a/src/Scrutor/OpenGenericDecorationStrategy.cs b/src/Scrutor/OpenGenericDecorationStrategy.cs
@@ -20,19 +20,19 @@ public override bool CanDecorate(Type serviceType) =>
             && serviceType.GetGenericTypeDefinition() == ServiceType.GetGenericTypeDefinition()
             && (DecoratorType is null || serviceType.HasCompatibleGenericArguments(DecoratorType));
 
-    public override Func<IServiceProvider, object> CreateDecorator(Type serviceType)
+    public override Func<IServiceProvider, object?, object> CreateDecorator(Type serviceType, string serviceKey)
     {
         if (DecoratorType is not null)
         {
             var genericArguments = serviceType.GetGenericArguments();
             var closedDecorator = DecoratorType.MakeGenericType(genericArguments);
 
-            return TypeDecorator(serviceType, closedDecorator);
+            return TypeDecorator(serviceType, serviceKey, closedDecorator);
         }
 
         if (DecoratorFactory is not null)
         {
-            return FactoryDecorator(serviceType, DecoratorFactory);
+            return FactoryDecorator(serviceType, serviceKey, DecoratorFactory);
         }
 
         throw new InvalidOperationException($"Both serviceType and decoratorFactory can not be null.");

diff --git a/src/Scrutor/ServiceCollectionExtensions.Decoration.cs b/src/Scrutor/ServiceCollectionExtensions.Decoration.cs
@@ -8,6 +8,8 @@ namespace Microsoft.Extensions.DependencyInjection;
 [PublicAPI]
 public static partial class ServiceCollectionExtensions
 {
+    private const string DecoratedServiceKeySuffix = "+Decorated";
+
     /// <summary>
     /// Decorates all registered services of type <typeparamref name="TService"/>
     /// using the specified type <typeparamref name="TDecorator"/>.
@@ -250,27 +252,48 @@ public static bool TryDecorate(this IServiceCollection services, DecorationStrat
         {
             var serviceDescriptor = services[i];
 
-            if (serviceDescriptor.ServiceType is DecoratedType)
+            if (IsDecorated(serviceDescriptor) || !strategy.CanDecorate(serviceDescriptor.ServiceType))
             {
-                continue; // Service has already been decorated.
+                continue;
             }
 
-            if (!strategy.CanDecorate(serviceDescriptor.ServiceType))
+            var serviceKey = GetDecoratorKey(serviceDescriptor);
+            if (serviceKey is null)
             {
-                continue; // Unable to decorate using the specified strategy.
+                return false;
             }
 
-            var decoratedType = new DecoratedType(serviceDescriptor.ServiceType);
-
             // Insert decorated
-            services.Add(serviceDescriptor.WithServiceType(decoratedType));
+            services.Add(serviceDescriptor.WithServiceKey(serviceKey));
 
             // Replace decorator
-            services[i] = serviceDescriptor.WithImplementationFactory(strategy.CreateDecorator(decoratedType));
+            services[i] = serviceDescriptor.WithImplementationFactory(strategy.CreateDecorator(serviceDescriptor.ServiceType, serviceKey));
 
             decorated = true;
         }
 
         return decorated;
     }
+
+    private static string? GetDecoratorKey(ServiceDescriptor descriptor)
+    {
+        var uniqueId = Guid.NewGuid().ToString("n");
+
+        if (descriptor.ServiceKey is null)
+        {
+            return $"{descriptor.ServiceType.Name}+{uniqueId}{DecoratedServiceKeySuffix}";
+        }
+
+        if (descriptor.ServiceKey is string stringKey)
+        {
+            return $"{stringKey}+{uniqueId}{DecoratedServiceKeySuffix}";
+        }
+
+        return null;
+    }
+
+    private static bool IsDecorated(ServiceDescriptor descriptor)
+    {
+        return descriptor.ServiceKey is string stringKey && stringKey.EndsWith(DecoratedServiceKeySuffix);
+    }
 }
diff --git a/src/Scrutor/ServiceDescriptorExtensions.cs b/src/Scrutor/ServiceDescriptorExtensions.cs
@@ -5,14 +5,46 @@ namespace Scrutor;
 
 internal static class ServiceDescriptorExtensions
 {
-    public static ServiceDescriptor WithImplementationFactory(this ServiceDescriptor descriptor, Func<IServiceProvider, object> implementationFactory) =>
-        new(descriptor.ServiceType, implementationFactory, descriptor.Lifetime);
+    public static ServiceDescriptor WithImplementationFactory(this ServiceDescriptor descriptor, Func<IServiceProvider, object?, object> implementationFactory) =>
+        new(descriptor.ServiceType, descriptor.ServiceKey, implementationFactory, descriptor.Lifetime);
 
-    public static ServiceDescriptor WithServiceType(this ServiceDescriptor descriptor, Type serviceType) => descriptor switch
+    public static ServiceDescriptor WithServiceKey(this ServiceDescriptor descriptor, string serviceKey)
     {
-        { ImplementationType: not null } => new ServiceDescriptor(serviceType, descriptor.ImplementationType, descriptor.Lifetime),
-        { ImplementationFactory: not null } => new ServiceDescriptor(serviceType, descriptor.ImplementationFactory, descriptor.Lifetime),
-        { ImplementationInstance: not null } => new ServiceDescriptor(serviceType, descriptor.ImplementationInstance),
-        _ => throw new ArgumentException($"No implementation factory or instance or type found for {descriptor.ServiceType}.", nameof(descriptor))
-    };
+        if (descriptor.IsKeyedService)
+        {
+            if (descriptor.KeyedImplementationType is not null)
+            {
+                return new ServiceDescriptor(descriptor.ServiceType, serviceKey, descriptor.KeyedImplementationType, descriptor.Lifetime);
+            }
+
+            if (descriptor.KeyedImplementationInstance is not null)
+            {
+                return new ServiceDescriptor(descriptor.ServiceType, serviceKey, descriptor.KeyedImplementationInstance);
+            }
+
+            if (descriptor.KeyedImplementationFactory is not null)
+            {
+                return new ServiceDescriptor(descriptor.ServiceType, serviceKey, descriptor.KeyedImplementationFactory, descriptor.Lifetime);
+            }
+
+            throw new InvalidOperationException($"One of the following properties must be set: {nameof(ServiceDescriptor.KeyedImplementationType)}, {nameof(ServiceDescriptor.KeyedImplementationInstance)} or {nameof(ServiceDescriptor.KeyedImplementationFactory)}");
+        }
+
+        if (descriptor.ImplementationType is not null)
+        {
+            return new ServiceDescriptor(descriptor.ServiceType, serviceKey, descriptor.ImplementationType, descriptor.Lifetime);
+        }
+
+        if (descriptor.ImplementationInstance is not null)
+        {
+            return new ServiceDescriptor(descriptor.ServiceType, serviceKey, descriptor.ImplementationInstance);
+        }
+
+        if (descriptor.ImplementationFactory is not null)
+        {
+            return new ServiceDescriptor(descriptor.ServiceType, serviceKey, (sp, key) => descriptor.ImplementationFactory(sp), descriptor.Lifetime);
+        }
+
+        throw new InvalidOperationException($"One of the following properties must be set: {nameof(ServiceDescriptor.ImplementationType)}, {nameof(ServiceDescriptor.ImplementationInstance)} or {nameof(ServiceDescriptor.ImplementationFactory)}");
+    }
 }
diff --git a/test/Scrutor.Tests/ServiceCollectionExtensions.cs b/test/Scrutor.Tests/ServiceCollectionExtensions.cs
@@ -18,6 +18,6 @@ public static ServiceDescriptor[] GetDescriptors<T>(this IServiceCollection serv
 
     public static ServiceDescriptor[] GetDescriptors(this IServiceCollection services, Type serviceType)
     {
-        return services.Where(x => x.ServiceType == serviceType).ToArray();
+        return services.Where(x => x.ServiceType == serviceType && x.ServiceKey is null).ToArray();
     }
 }