Add RegisterAll API to enable monitoring collections of key-values for refresh

src/Microsoft.Extensions.Configuration.AzureAppConfiguration/AzureAppConfigurationProvider.cs

@@ -32,6 +32,7 @@ internal class AzureAppConfigurationProvider : ConfigurationProvider, IConfigura
         private AzureAppConfigurationOptions _options;
         private Dictionary<string, ConfigurationSetting> _mappedData;
         private Dictionary<KeyValueIdentifier, ConfigurationSetting> _watchedIndividualKvs = new Dictionary<KeyValueIdentifier, ConfigurationSetting>();
+        private HashSet<string> _ffKeys = new HashSet<string>();
         private Dictionary<KeyValueSelector, IEnumerable<MatchConditions>> _kvEtags = new Dictionary<KeyValueSelector, IEnumerable<MatchConditions>>();
         private Dictionary<KeyValueSelector, IEnumerable<MatchConditions>> _ffEtags = new Dictionary<KeyValueSelector, IEnumerable<MatchConditions>>();
         private RequestTracingOptions _requestTracingOptions;
@@ -262,6 +263,7 @@ public async Task RefreshAsync(CancellationToken cancellationToken)
                     // Avoid instance state modification
                     Dictionary<KeyValueSelector, IEnumerable<MatchConditions>> ffEtags = null;
                     Dictionary<KeyValueSelector, IEnumerable<MatchConditions>> kvEtags = null;
+                    HashSet<string> ffKeys = null;
                     Dictionary<KeyValueIdentifier, ConfigurationSetting> watchedIndividualKvs = null;
                     List<KeyValueChange> keyValueChanges = null;
                     Dictionary<string, ConfigurationSetting> data = null;
@@ -273,17 +275,18 @@ public async Task RefreshAsync(CancellationToken cancellationToken)
 
                     await ExecuteWithFailOverPolicyAsync(clients, async (client) =>
                     {
-                        data = null;
-                        ffCollectionData = null;
-                        watchedIndividualKvs = null;
                         kvEtags = null;
                         ffEtags = null;
+                        ffKeys = null;
+                        watchedIndividualKvs = null;
                         keyValueChanges = new List<KeyValueChange>();
+                        data = null;
+                        ffCollectionData = null;
                         ffCollectionUpdated = false;
                         refreshAll = false;
-                        Uri endpoint = _configClientManager.GetEndpointForClient(client);
                         logDebugBuilder.Clear();
                         logInfoBuilder.Clear();
+                        Uri endpoint = _configClientManager.GetEndpointForClient(client);
 
                         if (_options.RegisterAllEnabled)
                         {
@@ -315,8 +318,9 @@ await ExecuteWithFailOverPolicyAsync(clients, async (client) =>
                             // or if any key-value collection change was detected.
                             kvEtags = new Dictionary<KeyValueSelector, IEnumerable<MatchConditions>>();
                             ffEtags = new Dictionary<KeyValueSelector, IEnumerable<MatchConditions>>();
+                            ffKeys = new HashSet<string>();
 
-                            data = await LoadSelected(client, kvEtags, ffEtags, _options.Selectors, cancellationToken).ConfigureAwait(false);
+                            data = await LoadSelected(client, kvEtags, ffEtags, _options.Selectors, ffKeys, cancellationToken).ConfigureAwait(false);
                             watchedIndividualKvs = await LoadKeyValuesRegisteredForRefresh(client, data, cancellationToken).ConfigureAwait(false);
                             logInfoBuilder.AppendLine(LogHelper.BuildConfigurationUpdatedMessage());
                             return;
@@ -332,12 +336,14 @@ await ExecuteWithFailOverPolicyAsync(clients, async (client) =>
                         if (ffCollectionUpdated)
                         {
                             ffEtags = new Dictionary<KeyValueSelector, IEnumerable<MatchConditions>>();
+                            ffKeys = new HashSet<string>();
 
                             ffCollectionData = await LoadSelected(
                                 client,
                                 new Dictionary<KeyValueSelector, IEnumerable<MatchConditions>>(),
                                 ffEtags,
                                 _options.Selectors.Where(selector => selector.IsFeatureFlagSelector),
+                                ffKeys,
                                 cancellationToken).ConfigureAwait(false);
 
                             logInfoBuilder.Append(LogHelper.BuildFeatureFlagsUpdatedMessage());
@@ -374,6 +380,12 @@ await ExecuteWithFailOverPolicyAsync(clients, async (client) =>
 
                         if (ffCollectionUpdated)
                         {
+                            // Remove all feature flag keys that are not present in the latest loading of feature flags, but were loaded previously
+                            foreach (string key in _ffKeys.Except(ffKeys))
+                            {
+                                _mappedData.Remove(key);
+                            }
+
                             Dictionary<string, ConfigurationSetting> mappedFfData = await MapConfigurationSettings(ffCollectionData).ConfigureAwait(false);
 
                             foreach (KeyValuePair<string, ConfigurationSetting> kvp in mappedFfData)
@@ -397,11 +409,13 @@ await ExecuteWithFailOverPolicyAsync(clients, async (client) =>
 
                     if (_options.Adapters.Any(adapter => adapter.NeedsRefresh()) || keyValueChanges.Any() || refreshAll || ffCollectionUpdated)
                     {
-                        _watchedIndividualKvs = watchedIndividualKvs;
+                        _watchedIndividualKvs = watchedIndividualKvs ?? _watchedIndividualKvs;
+
+                        _ffEtags = ffEtags ?? _ffEtags;
 
-                        _ffEtags = ffEtags;
+                        _kvEtags = kvEtags ?? _kvEtags;
 
-                        _kvEtags = kvEtags;
+                        _ffKeys = ffKeys ?? _ffKeys;
 
                         if (logDebugBuilder.Length > 0)
                         {
@@ -698,6 +712,7 @@ private async Task InitializeAsync(IEnumerable<ConfigurationClient> clients, Can
             Dictionary<KeyValueSelector, IEnumerable<MatchConditions>> kvEtags = new Dictionary<KeyValueSelector, IEnumerable<MatchConditions>>();
             Dictionary<KeyValueSelector, IEnumerable<MatchConditions>> ffEtags = new Dictionary<KeyValueSelector, IEnumerable<MatchConditions>>();
             Dictionary<KeyValueIdentifier, ConfigurationSetting> watchedIndividualKvs = null;
+            HashSet<string> ffKeys = new HashSet<string>();
 
             await ExecuteWithFailOverPolicyAsync(
                 clients,
@@ -708,6 +723,7 @@ await ExecuteWithFailOverPolicyAsync(
                         kvEtags,
                         ffEtags,
                         _options.Selectors,
+                        ffKeys,
                         cancellationToken)
                         .ConfigureAwait(false);
 
@@ -755,6 +771,7 @@ private async Task<Dictionary<string, ConfigurationSetting>> LoadSelected(
             Dictionary<KeyValueSelector, IEnumerable<MatchConditions>> kvEtags,
             Dictionary<KeyValueSelector, IEnumerable<MatchConditions>> ffEtags,
             IEnumerable<KeyValueSelector> selectors,
+            HashSet<string> ffKeys,
             CancellationToken cancellationToken)
         {
             Dictionary<string, ConfigurationSetting> data = new Dictionary<string, ConfigurationSetting>();
@@ -784,6 +801,11 @@ await CallWithRequestTracing(async () =>
                             foreach (ConfigurationSetting setting in page.Values)
                             {
                                 data[setting.Key] = setting;
+
+                                if (loadOption.IsFeatureFlagSelector)
+                                {
+                                    ffKeys.Add(setting.Key);
+                                }
                             }
 
                             matchConditions.Add(new MatchConditions { IfNoneMatch = serverEtag });

tests/Tests.AzureAppConfiguration/RefreshTests.cs

@@ -7,6 +7,7 @@
 using Azure.Identity;
 using Microsoft.Extensions.Configuration;
 using Microsoft.Extensions.Configuration.AzureAppConfiguration;
+using Microsoft.Extensions.Configuration.AzureAppConfiguration.FeatureManagement;
 using Microsoft.Extensions.Logging.Abstractions;
 using Moq;
 using System;
@@ -1108,6 +1109,118 @@ public async Task RefreshTests_SelectedKeysRefreshWithRegisterAll()
             Assert.Null(config["TestKey3"]);
         }
 
+        [Fact]
+        public async Task RefreshTests_RegisterAllRefreshesFeatureFlags()
+        {
+            IConfigurationRefresher refresher = null;
+            var mockClient = GetMockConfigurationClient();
+
+            var featureFlags = new List<ConfigurationSetting> {
+                ConfigurationModelFactory.ConfigurationSetting(
+                    key: FeatureManagementConstants.FeatureFlagMarker + "myFeature",
+                    value: @"
+                            {
+                                ""id"": ""MyFeature"",
+                                ""description"": ""The new beta version of our web site."",
+                                ""display_name"": ""Beta Feature"",
+                                ""enabled"": true,
+                                ""conditions"": {
+                                ""client_filters"": [
+                                    {
+                                    ""name"": ""SuperUsers""
+                                    }
+                                ]
+                                }
+                            }
+                            ",
+                    label: default,
+                    contentType: FeatureManagementConstants.ContentType + ";charset=utf-8",
+                    eTag: new ETag("c3c231fd-39a0-4cb6-3237-4614474b92c1"))
+            };
+
+            var mockAsyncPageableKv = new MockAsyncPageable(_kvCollection);
+
+            var mockAsyncPageableFf = new MockAsyncPageable(featureFlags);
+
+            MockAsyncPageable GetTestKeys(SettingSelector selector, CancellationToken ct)
+            {
+                if (selector.KeyFilter.StartsWith(FeatureManagementConstants.FeatureFlagMarker))
+                {
+                    mockAsyncPageableFf.UpdateCollection(featureFlags);
+
+                    return mockAsyncPageableFf;
+                }
+
+                mockAsyncPageableKv.UpdateCollection(_kvCollection);
+
+                return mockAsyncPageableKv;
+            }
+
+            mockClient.Setup(c => c.GetConfigurationSettingsAsync(It.IsAny<SettingSelector>(), It.IsAny<CancellationToken>()))
+                .Returns((Func<SettingSelector, CancellationToken, MockAsyncPageable>)GetTestKeys);
+
+            var config = new ConfigurationBuilder()
+                .AddAzureAppConfiguration(options =>
+                {
+                    options.ClientManager = TestHelpers.CreateMockedConfigurationClientManager(mockClient.Object);
+                    options.Select("TestKey*");
+                    options.ConfigurationSettingPageIterator = new MockConfigurationSettingPageIterator();
+                    options.ConfigureRefresh(refreshOptions =>
+                    {
+                        refreshOptions.RegisterAll()
+                            .SetRefreshInterval(TimeSpan.FromSeconds(1));
+                    });
+                    options.UseFeatureFlags();
+
+                    refresher = options.GetRefresher();
+                })
+                .Build();
+
+            Assert.Equal("TestValue1", config["TestKey1"]);
+            Assert.Equal("SuperUsers", config["FeatureManagement:MyFeature:EnabledFor:0:Name"]);
+
+            FirstKeyValue.Value = "newValue1";
+            featureFlags[0] = ConfigurationModelFactory.ConfigurationSetting(
+                key: FeatureManagementConstants.FeatureFlagMarker + "myFeature",
+                value: @"
+                        {
+                            ""id"": ""MyFeature"",
+                            ""description"": ""The new beta version of our web site."",
+                            ""display_name"": ""Beta Feature"",
+                            ""enabled"": true,
+                            ""conditions"": {
+                            ""client_filters"": [
+                                {
+                                ""name"": ""AllUsers""
+                                }
+                            ]
+                            }
+                        }
+                        ",
+                label: default,
+                contentType: FeatureManagementConstants.ContentType + ";charset=utf-8",
+                eTag: new ETag("c3c231fd-39a0-4cb6-3237-4614474b92c1"));
+
+            // Wait for the cache to expire
+            Thread.Sleep(1500);
+
+            await refresher.RefreshAsync();
+
+            Assert.Equal("newValue1", config["TestKey1"]);
+            Assert.Equal("AllUsers", config["FeatureManagement:MyFeature:EnabledFor:0:Name"]);
+
+            FirstKeyValue.Value = "newerValue1";
+            featureFlags.RemoveAt(0);
+
+            // Wait for the cache to expire
+            Thread.Sleep(1500);
+
+            await refresher.RefreshAsync();
+
+            Assert.Equal("newerValue1", config["TestKey1"]);
+            Assert.Null(config["FeatureManagement:MyFeature"]);
+        }
+
 #if NET8_0
         [Fact]
         public void RefreshTests_ChainedConfigurationProviderUsedAsRootForRefresherProvider()