Support keyed service for variant service provider (#606)

* support keyed service

* support variant service provider options

* update

* update

* revert variant service provider options

* add comment

Author: zhiyuanliang-ms

examples/VariantServiceDemo/Program.cs

@@ -19,10 +19,11 @@
 builder.Services.AddApplicationInsightsTelemetry();
 
 //
-// Add variant implementations of ICalculator
-builder.Services.AddSingleton<ICalculator, DefaultCalculator>();
+// Add variant implementations of ICalculator using keyed services so that only the
+// implementation matching the assigned variant is instantiated on demand.
+builder.Services.AddKeyedSingleton<ICalculator, DefaultCalculator>("DefaultCalculator");
 
-builder.Services.AddSingleton<ICalculator, RemoteCalculator>();
+builder.Services.AddKeyedSingleton<ICalculator, RemoteCalculator>("RemoteCalculator");
 
 //
 // Enter feature management

src/Microsoft.FeatureManagement/FeatureManagementBuilderExtensions.cs

@@ -63,14 +63,14 @@ public static IFeatureManagementBuilder WithVariantService<TService>(this IFeatu
                 builder.Services.AddScoped<IVariantServiceProvider<TService>>(sp => new VariantServiceProvider<TService>(
                     featureName,
                     sp.GetRequiredService<IVariantFeatureManager>(),
-                    sp.GetRequiredService<IEnumerable<TService>>()));
+                    sp));
             }
             else
             {
                 builder.Services.AddSingleton<IVariantServiceProvider<TService>>(sp => new VariantServiceProvider<TService>(
                     featureName,
                     sp.GetRequiredService<IVariantFeatureManager>(),
-                    sp.GetRequiredService<IEnumerable<TService>>()));
+                    sp));
             }
 
             return builder;

src/Microsoft.FeatureManagement/VariantServiceProvider.cs

@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT license.
 //
+using Microsoft.Extensions.DependencyInjection;
 using System;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
@@ -16,7 +17,7 @@ namespace Microsoft.FeatureManagement
     /// </summary>
     internal class VariantServiceProvider<TService> : IVariantServiceProvider<TService> where TService : class
     {
-        private readonly IEnumerable<TService> _services;
+        private readonly IServiceProvider _serviceProvider;
         private readonly IVariantFeatureManager _featureManager;
         private readonly string _featureName;
         private readonly ConcurrentDictionary<string, TService> _variantServiceCache;
@@ -26,15 +27,15 @@ internal class VariantServiceProvider<TService> : IVariantServiceProvider<TServi
         /// </summary>
         /// <param name="featureName">The feature flag that should be used to determine which variant of the service should be used.</param>
         /// <param name="featureManager">The feature manager to get the assigned variant of the feature flag.</param>
-        /// <param name="services">Implementation variants of TService.</param>
+        /// <param name="serviceProvider">The service provider used to resolve implementation variants of TService. If it implements <see cref="IKeyedServiceProvider"/>, keyed resolution is used to enable lazy instantiation; otherwise all registered implementations are enumerated.</param>
         /// <exception cref="ArgumentNullException">Thrown if <paramref name="featureName"/> is null.</exception>
         /// <exception cref="ArgumentNullException">Thrown if <paramref name="featureManager"/> is null.</exception>
-        /// <exception cref="ArgumentNullException">Thrown if <paramref name="services"/> is null.</exception>
-        public VariantServiceProvider(string featureName, IVariantFeatureManager featureManager, IEnumerable<TService> services)
+        /// <exception cref="ArgumentNullException">Thrown if <paramref name="serviceProvider"/> is null.</exception>
+        public VariantServiceProvider(string featureName, IVariantFeatureManager featureManager, IServiceProvider serviceProvider)
         {
             _featureName = featureName ?? throw new ArgumentNullException(nameof(featureName));
             _featureManager = featureManager ?? throw new ArgumentNullException(nameof(featureManager));
-            _services = services ?? throw new ArgumentNullException(nameof(services));
+            _serviceProvider = serviceProvider ?? throw new ArgumentNullException(nameof(serviceProvider));
             _variantServiceCache = new ConcurrentDictionary<string, TService>();
         }
 
@@ -55,16 +56,35 @@ public async ValueTask<TService> GetServiceAsync(CancellationToken cancellationT
             {
                 implementation = _variantServiceCache.GetOrAdd(
                     variant.Name,
-                    (_) => _services.FirstOrDefault(
-                        service => IsMatchingVariantName(
-                            service.GetType(),
-                            variant.Name))
-                );
+                    (variantName) => ResolveVariantService(variantName));
             }
 
             return implementation;
         }
 
+        private TService ResolveVariantService(string variantName)
+        {
+            //
+            // If the service provider supports keyed services, try to resolve the variant by its name as the key first.
+            // This allows lazy instantiation of the variant service.
+            if (_serviceProvider is IKeyedServiceProvider)
+            {
+                TService keyedService = _serviceProvider.GetKeyedService<TService>(variantName);
+
+                if (keyedService != null)
+                {
+                    return keyedService;
+                }
+            }
+
+            //
+            // Fall back to enumerating all non-keyed registrations of TService and matching by VariantServiceAliasAttribute or the implementation type name.
+            IEnumerable<TService> services = _serviceProvider.GetRequiredService<IEnumerable<TService>>();
+
+            return services.FirstOrDefault(
+                service => IsMatchingVariantName(service.GetType(), variantName));
+        }
+
         private bool IsMatchingVariantName(Type implementationType, string variantName)
         {
             string implementationName = ((VariantServiceAliasAttribute)Attribute.GetCustomAttribute(implementationType, typeof(VariantServiceAliasAttribute)))?.Alias;

tests/Tests.FeatureManagement/FeatureManagementTest.cs

@@ -524,12 +524,12 @@ public async Task MergesFeatureFlagsFromDifferentConfigurationSources()
              * Feature1: true
              * Feature2: true
              * FeatureA: true
-             * 
+             *
              * appsettings2.json
              * Feature1: true
              * Feature2: false
              * FeatureB: true
-             * 
+             *
              * appsettings3.json
              * Feature1: false
              * Feature2: false
@@ -2234,6 +2234,157 @@ public async Task VariantBasedInjection()
             );
         }
 
+        [Fact]
+        public async Task VariantServiceProviderResolvesKeyedService()
+        {
+            IConfiguration configuration = new ConfigurationBuilder()
+                .AddJsonFile("appsettings.json")
+                .Build();
+
+            IServiceCollection services = new ServiceCollection();
+
+            services.AddKeyedSingleton<IAlgorithm, AlgorithmBeta>("AlgorithmBeta");
+            services.AddKeyedSingleton<IAlgorithm, AlgorithmSigma>("Sigma");
+            services.AddKeyedSingleton<IAlgorithm>("Omega", (sp, _) => new AlgorithmOmega("OMEGA"));
+
+            services.AddSingleton(configuration)
+                .AddFeatureManagement()
+                .AddFeatureFilter<TargetingFilter>()
+                .WithVariantService<IAlgorithm>(Features.VariantImplementationFeature);
+
+            var targetingContextAccessor = new OnDemandTargetingContextAccessor();
+
+            services.AddSingleton<ITargetingContextAccessor>(targetingContextAccessor);
+
+            ServiceProvider serviceProvider = services.BuildServiceProvider();
+
+            IVariantServiceProvider<IAlgorithm> featuredAlgorithm = serviceProvider.GetRequiredService<IVariantServiceProvider<IAlgorithm>>();
+
+            targetingContextAccessor.Current = new TargetingContext { UserId = "UserBeta" };
+            IAlgorithm algorithm = await featuredAlgorithm.GetServiceAsync(CancellationToken.None);
+            Assert.NotNull(algorithm);
+            Assert.Equal("Beta", algorithm.Style);
+
+            targetingContextAccessor.Current = new TargetingContext { UserId = "UserSigma" };
+            algorithm = await featuredAlgorithm.GetServiceAsync(CancellationToken.None);
+            Assert.NotNull(algorithm);
+            Assert.Equal("Sigma", algorithm.Style);
+
+            targetingContextAccessor.Current = new TargetingContext { UserId = "UserOmega" };
+            algorithm = await featuredAlgorithm.GetServiceAsync(CancellationToken.None);
+            Assert.NotNull(algorithm);
+            Assert.Equal("OMEGA", algorithm.Style);
+        }
+
+        [Fact]
+        public async Task VariantServiceProviderKeyedServiceIsLazilyInstantiated()
+        {
+            IConfiguration configuration = new ConfigurationBuilder()
+                .AddJsonFile("appsettings.json")
+                .Build();
+
+            IServiceCollection services = new ServiceCollection();
+
+            int betaInstantiationCount = 0;
+            int sigmaInstantiationCount = 0;
+            int omegaInstantiationCount = 0;
+
+            services.AddKeyedSingleton<IAlgorithm>("AlgorithmBeta", (sp, _) =>
+            {
+                betaInstantiationCount++;
+                return new AlgorithmBeta();
+            });
+            services.AddKeyedSingleton<IAlgorithm>("Sigma", (sp, _) =>
+            {
+                sigmaInstantiationCount++;
+                return new AlgorithmSigma();
+            });
+            services.AddKeyedSingleton<IAlgorithm>("Omega", (sp, _) =>
+            {
+                omegaInstantiationCount++;
+                return new AlgorithmOmega("OMEGA");
+            });
+
+            services.AddSingleton(configuration)
+                .AddFeatureManagement()
+                .AddFeatureFilter<TargetingFilter>()
+                .WithVariantService<IAlgorithm>(Features.VariantImplementationFeature);
+
+            var targetingContextAccessor = new OnDemandTargetingContextAccessor();
+
+            services.AddSingleton<ITargetingContextAccessor>(targetingContextAccessor);
+
+            ServiceProvider serviceProvider = services.BuildServiceProvider();
+
+            IVariantServiceProvider<IAlgorithm> featuredAlgorithm = serviceProvider.GetRequiredService<IVariantServiceProvider<IAlgorithm>>();
+
+            //
+            // No variant resolved yet - nothing should be instantiated.
+            Assert.Equal(0, betaInstantiationCount);
+            Assert.Equal(0, sigmaInstantiationCount);
+            Assert.Equal(0, omegaInstantiationCount);
+
+            //
+            // Resolve the Beta variant. Only AlgorithmBeta should be instantiated.
+            targetingContextAccessor.Current = new TargetingContext { UserId = "UserBeta" };
+            IAlgorithm algorithm = await featuredAlgorithm.GetServiceAsync(CancellationToken.None);
+            Assert.Equal("Beta", algorithm.Style);
+            Assert.Equal(1, betaInstantiationCount);
+            Assert.Equal(0, sigmaInstantiationCount);
+            Assert.Equal(0, omegaInstantiationCount);
+
+            //
+            // Resolving Beta again should reuse the cached instance - no new instantiation.
+            algorithm = await featuredAlgorithm.GetServiceAsync(CancellationToken.None);
+            Assert.Equal("Beta", algorithm.Style);
+            Assert.Equal(1, betaInstantiationCount);
+            Assert.Equal(0, sigmaInstantiationCount);
+            Assert.Equal(0, omegaInstantiationCount);
+
+            //
+            // Resolve the Sigma variant. Only AlgorithmSigma should be instantiated additionally.
+            targetingContextAccessor.Current = new TargetingContext { UserId = "UserSigma" };
+            algorithm = await featuredAlgorithm.GetServiceAsync(CancellationToken.None);
+            Assert.Equal("Sigma", algorithm.Style);
+            Assert.Equal(1, betaInstantiationCount);
+            Assert.Equal(1, sigmaInstantiationCount);
+            Assert.Equal(0, omegaInstantiationCount);
+        }
+
+        [Fact]
+        public async Task VariantServiceProviderPrefersKeyedOverNonKeyed()
+        {
+            IConfiguration configuration = new ConfigurationBuilder()
+                .AddJsonFile("appsettings.json")
+                .Build();
+
+            IServiceCollection services = new ServiceCollection();
+
+            //
+            // Register both keyed and non-keyed implementations matching the same variant name.
+            // The keyed registration should take precedence.
+            services.AddSingleton<IAlgorithm, AlgorithmBeta>();
+            services.AddKeyedSingleton<IAlgorithm>("AlgorithmBeta", (sp, _) => new AlgorithmOmega("KeyedBeta"));
+
+            services.AddSingleton(configuration)
+                .AddFeatureManagement()
+                .AddFeatureFilter<TargetingFilter>()
+                .WithVariantService<IAlgorithm>(Features.VariantImplementationFeature);
+
+            var targetingContextAccessor = new OnDemandTargetingContextAccessor();
+
+            services.AddSingleton<ITargetingContextAccessor>(targetingContextAccessor);
+
+            ServiceProvider serviceProvider = services.BuildServiceProvider();
+
+            IVariantServiceProvider<IAlgorithm> featuredAlgorithm = serviceProvider.GetRequiredService<IVariantServiceProvider<IAlgorithm>>();
+
+            targetingContextAccessor.Current = new TargetingContext { UserId = "UserBeta" };
+            IAlgorithm algorithm = await featuredAlgorithm.GetServiceAsync(CancellationToken.None);
+            Assert.NotNull(algorithm);
+            Assert.Equal("KeyedBeta", algorithm.Style);
+        }
+
         [Fact]
         public async Task VariantFeatureFlagWithContextualFeatureFilter()
         {