Adjusted snapshot to use a shared IsEnabled cache

microsoft · Sep 25, 2024 · 3e07522 · 3e07522
1 parent ea8fea1
commit 3e07522
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 9 deletions.
diff --git a/src/Microsoft.FeatureManagement/FeatureManagerSnapshot.cs b/src/Microsoft.FeatureManagement/FeatureManagerSnapshot.cs
@@ -18,8 +18,7 @@ class FeatureManagerSnapshot : IFeatureManagerSnapshot, IVariantFeatureManagerSn
     {
         private readonly IFeatureManager _featureManager;
         private readonly IVariantFeatureManager _variantFeatureManager;
-        private readonly ConcurrentDictionary<string, Task<bool>> _flagCache = new ConcurrentDictionary<string, Task<bool>>();
-        private readonly ConcurrentDictionary<string, ValueTask<bool>> _variantFlagCache = new ConcurrentDictionary<string, ValueTask<bool>>();
+        private readonly ConcurrentDictionary<string, ValueTask<bool>> _flagCache = new ConcurrentDictionary<string, ValueTask<bool>>();
         private readonly ConcurrentDictionary<string, Variant> _variantCache = new ConcurrentDictionary<string, Variant>();
         private IEnumerable<string> _featureNames;
 
@@ -74,26 +73,26 @@ public Task<bool> IsEnabledAsync(string feature)
         {
             return _flagCache.GetOrAdd(
                 feature,
-                (key) => _featureManager.IsEnabledAsync(key));
+                (key) => new ValueTask<bool>(_featureManager.IsEnabledAsync(key))).AsTask();
         }
 
         public Task<bool> IsEnabledAsync<TContext>(string feature, TContext context)
         {
             return _flagCache.GetOrAdd(
                 feature,
-                (key) => _featureManager.IsEnabledAsync(key, context));
+                (key) => new ValueTask<bool>(_featureManager.IsEnabledAsync(key, context))).AsTask();
         }
 
         public ValueTask<bool> IsEnabledAsync(string feature, CancellationToken cancellationToken)
         {
-            return _variantFlagCache.GetOrAdd(
+            return _flagCache.GetOrAdd(
                 feature,
                 (key) => _variantFeatureManager.IsEnabledAsync(key, cancellationToken));
         }
 
         public ValueTask<bool> IsEnabledAsync<TContext>(string feature, TContext context, CancellationToken cancellationToken)
         {
-            return _variantFlagCache.GetOrAdd(
+            return _flagCache.GetOrAdd(
                 feature,
                 (key) => _variantFeatureManager.IsEnabledAsync(key, context, cancellationToken));
         }

diff --git a/tests/Tests.FeatureManagement/FeatureManagementTest.cs b/tests/Tests.FeatureManagement/FeatureManagementTest.cs
@@ -1917,10 +1917,10 @@ public async Task CustomIFeatureManagerTest()
             Assert.False(await featureManagerSnapshot.IsEnabledAsync("NotTest"));
             Assert.False(await featureManagerSnapshot.IsEnabledAsync("OnTestFeature"));
 
-            // But use IVariantFeatureManager when using new interface
-            Assert.False(await featureManagerSnapshot.IsEnabledAsync("Test", CancellationToken.None));
+            // Use snapshot results even though IVariantFeatureManager would be called here
+            Assert.True(await featureManagerSnapshot.IsEnabledAsync("Test", CancellationToken.None));
             Assert.False(await featureManagerSnapshot.IsEnabledAsync("NotTest", CancellationToken.None));
-            Assert.True(await featureManagerSnapshot.IsEnabledAsync("OnTestFeature", CancellationToken.None));
+            Assert.False(await featureManagerSnapshot.IsEnabledAsync("OnTestFeature", CancellationToken.None));
         }
 
         [Fact]