feat(VariantBasedInjection): Keyed DI support

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

@@ -58,19 +58,19 @@ public static IFeatureManagementBuilder WithVariantService<TService>(this IFeatu
                 throw new InvalidOperationException($"A variant service of {typeof(TService).FullName} has already been added.");
             }
 
+            Func<IServiceProvider, IVariantServiceProvider<TService>> variantSpFactory = sp =>
+            {
+                var featureManager = sp.GetRequiredService<IVariantFeatureManager>();
+                return new VariantServiceProvider<TService>(featureName, featureManager, sp);
+            };
+
             if (builder.Services.Any(descriptor => descriptor.ServiceType == typeof(IFeatureManager) && descriptor.Lifetime == ServiceLifetime.Scoped))
             {
-                builder.Services.AddScoped<IVariantServiceProvider<TService>>(sp => new VariantServiceProvider<TService>(
-                    featureName,
-                    sp.GetRequiredService<IVariantFeatureManager>(),
-                    sp.GetRequiredService<IEnumerable<TService>>()));
+                builder.Services.AddScoped(variantSpFactory);
             }
             else
             {
-                builder.Services.AddSingleton<IVariantServiceProvider<TService>>(sp => new VariantServiceProvider<TService>(
-                    featureName,
-                    sp.GetRequiredService<IVariantFeatureManager>(),
-                    sp.GetRequiredService<IEnumerable<TService>>()));
+                builder.Services.AddSingleton(variantSpFactory);
             }
 
             return builder;

src/Microsoft.FeatureManagement/Microsoft.FeatureManagement.csproj

@@ -16,7 +16,7 @@
     <SignAssembly>true</SignAssembly>
     <DelaySign>false</DelaySign>
     <AssemblyOriginatorKeyFile>..\..\build\Microsoft.FeatureManagement.snk</AssemblyOriginatorKeyFile>
-    <!-- Microsoft.FeatureManagement uses the feature of async streams which is not supported in versions of C# earlier than 8.0. 
+    <!-- Microsoft.FeatureManagement uses the feature of async streams which is not supported in versions of C# earlier than 8.0.
         The library targets on netstandard 2.0. To ensure compatibility, the minimum language version requirement should be maintained. -->
     <LangVersion>8.0</LangVersion>
   </PropertyGroup>

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>();
         }
 
@@ -49,20 +50,36 @@ public async ValueTask<TService> GetServiceAsync(CancellationToken cancellationT
 
             Variant variant = await _featureManager.GetVariantAsync(_featureName, cancellationToken);
 
-            TService implementation = null;
+            return variant != null ? _variantServiceCache.GetOrAdd(variant.Name, ResolveVariantService) : null;
+        }
 
-            if (variant != null)
+        private TService ResolveVariantService(string variantName)
+        {
+            if (TryGetKeyedVariantService(variantName, out var keyedVariantService))
             {
-                implementation = _variantServiceCache.GetOrAdd(
-                    variant.Name,
-                    (_) => _services.FirstOrDefault(
-                        service => IsMatchingVariantName(
-                            service.GetType(),
-                            variant.Name))
-                );
+                return keyedVariantService;
             }
 
-            return implementation;
+            return GetVariantServiceFallback(variantName);
+        }
+
+        private bool TryGetKeyedVariantService(string variantName, out TService keyedService)
+        {
+            if (_serviceProvider is IKeyedServiceProvider keyedServiceProvider)
+            {
+                keyedService = keyedServiceProvider.GetKeyedService<TService>(variantName);
+                return keyedService != null;
+            }
+
+            keyedService = null;
+            return false;
+        }
+
+        private TService GetVariantServiceFallback(string variantName)
+        {
+            return _serviceProvider
+                .GetRequiredService<IEnumerable<TService>>()
+                .FirstOrDefault(service => IsMatchingVariantName(service.GetType(), variantName));
         }
 
         private bool IsMatchingVariantName(Type implementationType, string variantName)

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,237 @@ public async Task VariantBasedInjection()
             );
         }
 
+        [Fact]
+        public async Task LazyVariantBasedInjectionScoped()
+        {
+            IConfiguration configuration = new ConfigurationBuilder()
+                .AddJsonFile("appsettings.json")
+                .Build();
+
+            IServiceCollection services = new ServiceCollection();
+
+            services.AddKeyedScoped<IAlgorithm, AlgorithmBeta>(nameof(AlgorithmBeta));
+            services.AddKeyedScoped<IAlgorithm, AlgorithmSigma>(nameof(AlgorithmSigma));
+            services.AddKeyedScoped<IAlgorithm>("Omega", (sp, _) => new AlgorithmOmega("OMEGA"));
+
+            services.AddSingleton(configuration)
+                .AddScopedFeatureManagement()
+                .AddFeatureFilter<TargetingFilter>()
+                .WithVariantService<IAlgorithm>(Features.VariantImplementationFeature);
+
+            var targetingContextAccessor = new OnDemandTargetingContextAccessor();
+
+            services.AddSingleton<ITargetingContextAccessor>(targetingContextAccessor);
+
+            var serviceProvider = services.BuildServiceProvider().CreateScope().ServiceProvider;
+
+            IVariantFeatureManager featureManager = serviceProvider.GetRequiredService<IVariantFeatureManager>();
+
+            IVariantServiceProvider<IAlgorithm> featuredAlgorithm = serviceProvider.GetRequiredService<IVariantServiceProvider<IAlgorithm>>();
+
+            targetingContextAccessor.Current = new TargetingContext
+            {
+                UserId = "Guest"
+            };
+
+            IAlgorithm algorithm = await featuredAlgorithm.GetServiceAsync(CancellationToken.None);
+
+            Assert.Null(algorithm);
+
+            targetingContextAccessor.Current = new TargetingContext
+            {
+                UserId = "UserSigma"
+            };
+
+            algorithm = await featuredAlgorithm.GetServiceAsync(CancellationToken.None);
+
+            Assert.Null(algorithm);
+
+            targetingContextAccessor.Current = new TargetingContext
+            {
+                UserId = "UserBeta"
+            };
+
+            algorithm = await featuredAlgorithm.GetServiceAsync(CancellationToken.None);
+
+            Assert.NotNull(algorithm);
+            Assert.Equal("Beta", algorithm.Style);
+
+            targetingContextAccessor.Current = new TargetingContext
+            {
+                UserId = "UserOmega"
+            };
+
+            algorithm = await featuredAlgorithm.GetServiceAsync(CancellationToken.None);
+
+            Assert.NotNull(algorithm);
+            Assert.Equal("OMEGA", algorithm.Style);
+
+            services = new ServiceCollection();
+
+            Assert.Throws<InvalidOperationException>(() =>
+                {
+                    services.AddFeatureManagement()
+                        .WithVariantService<IAlgorithm>("DummyFeature1")
+                        .WithVariantService<IAlgorithm>("DummyFeature2");
+                }
+            );
+        }
+
+        [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()
         {