Add plumbing to support decorating keyed services

khellang · Nov 15, 2023 · 4913f3a · 4913f3a
1 parent 1a99cfb
commit 4913f3a
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 18 deletions.
diff --git a/src/Scrutor/ClosedTypeDecorationStrategy.cs b/src/Scrutor/ClosedTypeDecorationStrategy.cs
@@ -4,7 +4,7 @@ namespace Scrutor;
 
 internal sealed class ClosedTypeDecorationStrategy : DecorationStrategy
 {
-    public ClosedTypeDecorationStrategy(Type serviceType, Type? decoratorType, Func<object, IServiceProvider, object>? decoratorFactory) : base(serviceType)
+    public ClosedTypeDecorationStrategy(Type serviceType, string? serviceKey, Type? decoratorType, Func<object, IServiceProvider, object>? decoratorFactory) : base(serviceType, serviceKey)
     {
         DecoratorType = decoratorType;
         DecoratorFactory = decoratorFactory;
@@ -14,7 +14,7 @@ public ClosedTypeDecorationStrategy(Type serviceType, Type? decoratorType, Func<
 
     private Func<object, IServiceProvider, object>? DecoratorFactory { get; }
 
-    public override bool CanDecorate(Type serviceType) => ServiceType == serviceType;
+    protected override bool CanDecorate(Type serviceType) => ServiceType == serviceType;
 
     public override Func<IServiceProvider, object?, object> CreateDecorator(Type serviceType, string serviceKey)
     {

diff --git a/src/Scrutor/DecorationStrategy.cs b/src/Scrutor/DecorationStrategy.cs
@@ -5,22 +5,28 @@ namespace Scrutor;
 
 public abstract class DecorationStrategy
 {
-    protected DecorationStrategy(Type serviceType)
+    protected DecorationStrategy(Type serviceType, string? serviceKey)
     {
         ServiceType = serviceType;
+        ServiceKey = serviceKey;
     }
 
     public Type ServiceType { get; }
 
-    public abstract bool CanDecorate(Type serviceType);
+    public string? ServiceKey { get; }
+
+    public virtual bool CanDecorate(ServiceDescriptor descriptor) =>
+        string.Equals(ServiceKey, descriptor.ServiceKey) && CanDecorate(descriptor.ServiceType);
+
+    protected abstract bool CanDecorate(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 WithType(Type serviceType, string? serviceKey, Type decoratorType) =>
+        Create(serviceType, serviceKey, decoratorType, decoratorFactory: null);
 
-    internal static DecorationStrategy WithFactory(Type serviceType, Func<object, IServiceProvider, object> decoratorFactory) =>
-        Create(serviceType, decoratorType: null, decoratorFactory);
+    internal static DecorationStrategy WithFactory(Type serviceType, string? serviceKey, Func<object, IServiceProvider, object> decoratorFactory) =>
+        Create(serviceType, serviceKey, decoratorType: null, decoratorFactory);
 
     protected static Func<IServiceProvider, object?, object> TypeDecorator(Type serviceType, string serviceKey, Type decoratorType)
     {
@@ -38,13 +44,13 @@ internal static DecorationStrategy WithFactory(Type serviceType, Func<object, IS
         return decoratorFactory(instanceToDecorate, serviceProvider);
     };
 
-    private static DecorationStrategy Create(Type serviceType, Type? decoratorType, Func<object, IServiceProvider, object>? decoratorFactory)
+    private static DecorationStrategy Create(Type serviceType, string? serviceKey, Type? decoratorType, Func<object, IServiceProvider, object>? decoratorFactory)
     {
         if (serviceType.IsOpenGeneric())
         {
-            return new OpenGenericDecorationStrategy(serviceType, decoratorType, decoratorFactory);
+            return new OpenGenericDecorationStrategy(serviceType, serviceKey, decoratorType, decoratorFactory);
         }
 
-        return new ClosedTypeDecorationStrategy(serviceType, decoratorType, decoratorFactory);
+        return new ClosedTypeDecorationStrategy(serviceType, serviceKey, decoratorType, decoratorFactory);
     }
 }
diff --git a/src/Scrutor/OpenGenericDecorationStrategy.cs b/src/Scrutor/OpenGenericDecorationStrategy.cs
@@ -4,7 +4,7 @@ namespace Scrutor;
 
 public class OpenGenericDecorationStrategy : DecorationStrategy
 {
-    public OpenGenericDecorationStrategy(Type serviceType, Type? decoratorType, Func<object, IServiceProvider, object>? decoratorFactory) : base(serviceType)
+    public OpenGenericDecorationStrategy(Type serviceType, string? serviceKey, Type? decoratorType, Func<object, IServiceProvider, object>? decoratorFactory) : base(serviceType, serviceKey)
     {
         DecoratorType = decoratorType;
         DecoratorFactory = decoratorFactory;
@@ -14,7 +14,7 @@ public OpenGenericDecorationStrategy(Type serviceType, Type? decoratorType, Func
 
     private Func<object, IServiceProvider, object>? DecoratorFactory { get; }
 
-    public override bool CanDecorate(Type serviceType) =>
+    protected override bool CanDecorate(Type serviceType) =>
         serviceType.IsGenericType
             && !serviceType.IsGenericTypeDefinition
             && serviceType.GetGenericTypeDefinition() == ServiceType.GetGenericTypeDefinition()

diff --git a/src/Scrutor/ServiceCollectionExtensions.Decoration.cs b/src/Scrutor/ServiceCollectionExtensions.Decoration.cs
@@ -55,7 +55,7 @@ public static IServiceCollection Decorate(this IServiceCollection services, Type
         Preconditions.NotNull(serviceType, nameof(serviceType));
         Preconditions.NotNull(decoratorType, nameof(decoratorType));
 
-        return services.Decorate(DecorationStrategy.WithType(serviceType, decoratorType));
+        return services.Decorate(DecorationStrategy.WithType(serviceType, serviceKey: null, decoratorType));
     }
 
     /// <summary>
@@ -73,7 +73,7 @@ public static bool TryDecorate(this IServiceCollection services, Type serviceTyp
         Preconditions.NotNull(serviceType, nameof(serviceType));
         Preconditions.NotNull(decoratorType, nameof(decoratorType));
 
-        return services.TryDecorate(DecorationStrategy.WithType(serviceType, decoratorType));
+        return services.TryDecorate(DecorationStrategy.WithType(serviceType, serviceKey: null, decoratorType));
     }
 
     /// <summary>
@@ -199,7 +199,7 @@ public static IServiceCollection Decorate(this IServiceCollection services, Type
         Preconditions.NotNull(serviceType, nameof(serviceType));
         Preconditions.NotNull(decorator, nameof(decorator));
 
-        return services.Decorate(DecorationStrategy.WithFactory(serviceType, decorator));
+        return services.Decorate(DecorationStrategy.WithFactory(serviceType, serviceKey: null, decorator));
     }
 
     /// <summary>
@@ -217,7 +217,7 @@ public static bool TryDecorate(this IServiceCollection services, Type serviceTyp
         Preconditions.NotNull(serviceType, nameof(serviceType));
         Preconditions.NotNull(decorator, nameof(decorator));
 
-        return services.TryDecorate(DecorationStrategy.WithFactory(serviceType, decorator));
+        return services.TryDecorate(DecorationStrategy.WithFactory(serviceType, serviceKey: null, decorator));
     }
 
     /// <summary>
@@ -252,7 +252,7 @@ public static bool TryDecorate(this IServiceCollection services, DecorationStrat
         {
             var serviceDescriptor = services[i];
 
-            if (IsDecorated(serviceDescriptor) || !strategy.CanDecorate(serviceDescriptor.ServiceType))
+            if (IsDecorated(serviceDescriptor) || !strategy.CanDecorate(serviceDescriptor))
             {
                 continue;
             }