Fix session manager bug

src/Microsoft.FeatureManagement/FeatureManager.cs

@@ -355,14 +355,6 @@ private async ValueTask<EvaluationEvent> EvaluateFeature<TContext>(string featur
                     }
                 }
 
-                if (_sessionManagers != null)
-                {
-                    foreach (ISessionManager sessionManager in _sessionManagers)
-                    {
-                        await sessionManager.SetAsync(evaluationEvent.FeatureDefinition.Name, evaluationEvent.Enabled).ConfigureAwait(false);
-                    }
-                }
-
                 // Only add an activity event if telemetry is enabled for the feature and the activity is valid
                 if (telemetryEnabled &&
                     Activity.Current != null &&
@@ -371,6 +363,28 @@ private async ValueTask<EvaluationEvent> EvaluateFeature<TContext>(string featur
                     FeatureEvaluationTelemetry.Publish(evaluationEvent, Logger);
                 }
             }
+            else if (_sessionManagers != null)
+            {
+                foreach (ISessionManager sessionManager in _sessionManagers)
+                {
+                    bool? readSessionResult = await sessionManager.GetAsync(feature).ConfigureAwait(false);
+
+                    if (readSessionResult.HasValue)
+                    {
+                        evaluationEvent.Enabled = readSessionResult.Value;
+
+                        break;
+                    }
+                }
+            }
+
+            if (_sessionManagers != null)
+            {
+                foreach (ISessionManager sessionManager in _sessionManagers)
+                {
+                    await sessionManager.SetAsync(feature, evaluationEvent.Enabled).ConfigureAwait(false);
+                }
+            }
 
             return evaluationEvent;
         }

tests/Tests.FeatureManagement/FeatureManagementTest.cs

@@ -280,6 +280,61 @@ public async Task ThrowsForMissingFeatures()
                 featureManager.IsEnabledAsync("NonExistentFeature"));
         }
 
+        [Fact]
+        public async Task SessionManagerQueriedWhenFeatureDefinitionIsNull()
+        {
+            IConfiguration config = new ConfigurationBuilder().AddJsonFile("appsettings.json").Build();
+
+            var services = new ServiceCollection();
+
+            ISessionManager sessionManager = new TestSessionManager();
+
+            await sessionManager.SetAsync("UnexistedFeature", true);
+
+            services
+                .AddSingleton(config)
+                .AddSingleton<ISessionManager>(sessionManager)
+                .AddFeatureManagement();
+
+            ServiceProvider serviceProvider = services.BuildServiceProvider();
+
+            IFeatureManager featureManager = serviceProvider.GetRequiredService<IFeatureManager>();
+
+            // Feature doesn't exist in configuration, but should return true from session
+            Assert.True(await featureManager.IsEnabledAsync("UnexistedFeature"));
+
+            // Set the feature to false in session
+            await sessionManager.SetAsync("UnexistedFeature", false);
+
+            // Should return false from session
+            Assert.False(await featureManager.IsEnabledAsync("UnexistedFeature"));
+        }
+
+        [Fact]
+        public async Task SetResultInSessionManagerWhenFeatureDefinitionIsNull()
+        {
+            IConfiguration config = new ConfigurationBuilder().AddJsonFile("appsettings.json").Build();
+
+            var services = new ServiceCollection();
+
+            ISessionManager sessionManager = new TestSessionManager();
+
+            services
+                .AddSingleton(config)
+                .AddSingleton<ISessionManager>(sessionManager)
+                .AddFeatureManagement();
+
+            ServiceProvider serviceProvider = services.BuildServiceProvider();
+
+            IFeatureManager featureManager = serviceProvider.GetRequiredService<IFeatureManager>();
+
+            // session manager is clean here
+            var result = await featureManager.IsEnabledAsync("UnexistedFeature");
+
+            // session manager should store the result after evaluation
+            Assert.Equal(result, await sessionManager.GetAsync("UnexistedFeature"));
+        }
+
         [Fact]
         public async Task ThreadSafeSnapshot()
         {

tests/Tests.FeatureManagement/TestSessionManager.cs

@@ -0,0 +1,31 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+//
+using Microsoft.FeatureManagement;
+using System.Collections.Concurrent;
+using System.Threading.Tasks;
+
+namespace Tests.FeatureManagement
+{
+    class TestSessionManager : ISessionManager
+    {
+        private readonly ConcurrentDictionary<string, bool> _session = new ConcurrentDictionary<string, bool>();
+
+        public Task SetAsync(string featureName, bool enabled)
+        {
+            _session[featureName] = enabled;
+
+            return Task.CompletedTask;
+        }
+
+        public Task<bool?> GetAsync(string featureName)
+        {
+            if (_session.TryGetValue(featureName, out bool enabled))
+            {
+                return Task.FromResult<bool?>(enabled);
+            }
+
+            return Task.FromResult<bool?>(null);
+        }
+    }
+}