App Config Spring - Refresh Refactor

sdk/spring/spring-cloud-azure-appconfiguration-config/src/main/java/com/azure/spring/cloud/appconfiguration/config/AppConfigurationWatchAutoConfiguration.java

@@ -2,11 +2,11 @@
 // Licensed under the MIT License.
 package com.azure.spring.cloud.appconfiguration.config;
 
-import org.springframework.boot.bootstrap.BootstrapContext;
 import org.springframework.boot.autoconfigure.AutoConfiguration;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
+import org.springframework.boot.bootstrap.BootstrapContext;
 import org.springframework.boot.context.properties.EnableConfigurationProperties;
 import org.springframework.cloud.endpoint.RefreshEndpoint;
 import org.springframework.context.annotation.Bean;
@@ -15,6 +15,7 @@
 import com.azure.spring.cloud.appconfiguration.config.implementation.AppConfigurationPullRefresh;
 import com.azure.spring.cloud.appconfiguration.config.implementation.AppConfigurationRefreshUtil;
 import com.azure.spring.cloud.appconfiguration.config.implementation.AppConfigurationReplicaClientFactory;
+import com.azure.spring.cloud.appconfiguration.config.implementation.StateHolder;
 import com.azure.spring.cloud.appconfiguration.config.implementation.autofailover.ReplicaLookUp;
 import com.azure.spring.cloud.appconfiguration.config.implementation.properties.AppConfigurationProperties;
 
@@ -37,15 +38,16 @@ public AppConfigurationWatchAutoConfiguration() {
     @Bean
     @ConditionalOnMissingBean
     AppConfigurationRefresh appConfigurationRefresh(AppConfigurationProperties properties, BootstrapContext context) {
-        AppConfigurationReplicaClientFactory clientFactory = context
-            .getOrElse(AppConfigurationReplicaClientFactory.class, null);
+        AppConfigurationReplicaClientFactory clientFactory = context.getOrElse(AppConfigurationReplicaClientFactory.class, null);
         ReplicaLookUp replicaLookUp = context.getOrElse(ReplicaLookUp.class, null);
 
+        StateHolder stateHolder = context.get(StateHolder.class);
+
         if (clientFactory == null || replicaLookUp == null) {
             return null;
         }
 
         return new AppConfigurationPullRefresh(clientFactory, properties.getRefreshInterval(), replicaLookUp,
-            new AppConfigurationRefreshUtil());
+            stateHolder, new AppConfigurationRefreshUtil(stateHolder));
     }
 }

sdk/spring/spring-cloud-azure-appconfiguration-config/src/main/java/com/azure/spring/cloud/appconfiguration/config/implementation/AppConfigurationApplicationSettingPropertySource.java

@@ -65,6 +65,7 @@ class AppConfigurationApplicationSettingPropertySource extends AppConfigurationP
      * @param keyPrefixTrimValues prefixs to trim from key values
      * @throws InvalidConfigurationPropertyValueException thrown if fails to parse Json content type
      */
+    @Override
     public void initProperties(List<String> keyPrefixTrimValues, Context context) throws InvalidConfigurationPropertyValueException {
 
         List<String> labels = Arrays.asList(labelFilters);
@@ -136,7 +137,6 @@ private void handleKeyVaultReference(String key, SecretReferenceConfigurationSet
     void handleFeatureFlag(String key, FeatureFlagConfigurationSetting setting, List<String> trimStrings)
         throws InvalidConfigurationPropertyValueException {
         // Feature Flags aren't loaded as configuration, but are loaded as feature flags when loading a snapshot.
-        return;
     }
 
     private void handleJson(ConfigurationSetting setting, List<String> keyPrefixTrimValues)

sdk/spring/spring-cloud-azure-appconfiguration-config/src/main/java/com/azure/spring/cloud/appconfiguration/config/implementation/AppConfigurationPullRefresh.java

@@ -4,6 +4,7 @@
 
 import java.time.Duration;
 import java.util.Map;
+import java.util.Objects;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.slf4j.Logger;
@@ -36,7 +37,6 @@ public class AppConfigurationPullRefresh implements AppConfigurationRefresh {
      * Publisher for Spring refresh events.
      */
     private ApplicationEventPublisher publisher;
-    private final Long defaultMinBackoff = (long) 30;
 
     /**
      * Default minimum backoff duration in seconds when refresh operations fail.
@@ -63,19 +63,30 @@ public class AppConfigurationPullRefresh implements AppConfigurationRefresh {
      */
     private final AppConfigurationRefreshUtil refreshUtils;
 
+    /**
+     * Holds configuration state between refreshes.
+     */
+    private final StateHolder stateHolder;
+
     /**
      * Creates a new AppConfigurationPullRefresh component.
      *
      * @param clientFactory factory for creating App Configuration clients to connect to stores
      * @param refreshInterval time duration between refresh interval checks
      * @param replicaLookUp component for handling replica lookup and failover
+     * @param stateHolder holds configuration state between refreshes
      * @param refreshUtils utility component for refresh operations
      */
     public AppConfigurationPullRefresh(AppConfigurationReplicaClientFactory clientFactory, Duration refreshInterval,
-        ReplicaLookUp replicaLookUp, AppConfigurationRefreshUtil refreshUtils) {
+        ReplicaLookUp replicaLookUp, StateHolder stateHolder, AppConfigurationRefreshUtil refreshUtils) {
         this.refreshInterval = refreshInterval;
         this.clientFactory = clientFactory;
         this.replicaLookUp = replicaLookUp;
+        if (Objects.isNull(stateHolder)) {
+            // StateHolder is null if all stores are disabled.
+            stateHolder = new StateHolder();
+        }
+        this.stateHolder = stateHolder;
         this.refreshUtils = refreshUtils;
     }
 
@@ -96,6 +107,7 @@ public void setApplicationEventPublisher(ApplicationEventPublisher applicationEv
      * @return a Mono containing a boolean indicating if a RefreshEvent was published. Returns {@code false} if
      * refreshConfigurations is currently being executed elsewhere.
      */
+    @Override
     public Mono<Boolean> refreshConfigurations() {
         return Mono.just(refreshStores());
     }
@@ -107,14 +119,15 @@ public Mono<Boolean> refreshConfigurations() {
      * @param endpoint the Config Store endpoint to expire refresh interval on
      * @param syncToken the syncToken to verify the latest changes are available on pull
      */
+    @Override
     public void expireRefreshInterval(String endpoint, String syncToken) {
         LOGGER.debug("Expiring refresh interval for " + endpoint);
 
         String originEndpoint = clientFactory.findOriginForEndpoint(endpoint);
 
         clientFactory.updateSyncToken(originEndpoint, endpoint, syncToken);
 
-        StateHolder.getCurrentState().expireState(originEndpoint);
+        stateHolder.expireState(originEndpoint);
     }
 
     /**
@@ -135,7 +148,7 @@ private boolean refreshStores() {
             } catch (Exception e) {
                 LOGGER.warn("Error occurred during configuration refresh, will retry at next interval", e);
                 // The next refresh will happen sooner if refresh interval is expired.
-                StateHolder.getCurrentState().updateNextRefreshTime(refreshInterval, DEFAULT_MIN_BACKOFF_SECONDS);
+                stateHolder.updateNextRefreshTime(refreshInterval, DEFAULT_MIN_BACKOFF_SECONDS);
                 throw e;
             } finally {
                 running.set(false);

sdk/spring/spring-cloud-azure-appconfiguration-config/src/main/java/com/azure/spring/cloud/appconfiguration/config/implementation/AppConfigurationRefreshUtil.java

@@ -32,6 +32,21 @@ public class AppConfigurationRefreshUtil {
 
     private static final String FEATURE_FLAG_PREFIX = ".appconfig.featureflag/*";
 
+    private final StateHolder stateHolder;
+
+    /**
+     * Creates a new AppConfigurationRefreshUtil with the specified state holder.
+     *
+     * @param stateHolder the state holder for managing configuration and feature flag states
+     */
+    public AppConfigurationRefreshUtil(StateHolder stateHolder) {
+        if (stateHolder == null) {
+            // This is a fallback if all stores are disabled.
+            stateHolder = new StateHolder();
+        }
+        this.stateHolder = stateHolder;
+    }
+
     /**
      * Functional interface for refresh operations that can throw AppConfigurationStatusException.
      */
@@ -56,8 +71,8 @@ RefreshEventData refreshStoresCheck(AppConfigurationReplicaClientFactory clientF
         RefreshEventData eventData = new RefreshEventData();
 
         try {
-            if (refreshInterval != null && StateHolder.getNextForcedRefresh() != null
-                && Instant.now().isAfter(StateHolder.getNextForcedRefresh())) {
+            if (refreshInterval != null && stateHolder.getNextForcedRefresh() != null
+                && Instant.now().isAfter(stateHolder.getNextForcedRefresh())) {
                 String eventDataInfo = "Minimum refresh period reached. Refreshing configurations.";
 
                 LOGGER.info(eventDataInfo);
@@ -84,12 +99,20 @@ RefreshEventData refreshStoresCheck(AppConfigurationReplicaClientFactory clientF
 
                 clientFactory.findActiveClients(originEndpoint);
 
-                if (monitor.isEnabled() && StateHolder.getLoadState(originEndpoint)) {
+                if (monitor.isEnabled() && stateHolder.getLoadState(originEndpoint)) {
                     RefreshEventData result = executeRefreshWithRetry(
                         clientFactory,
                         originEndpoint,
-                        (client, data, ctx) -> refreshWithTime(client, StateHolder.getState(originEndpoint),
-                            monitor.getRefreshInterval(), data, replicaLookUp, ctx),
+                        (client, data, ctx) -> {
+                            if (stateHolder.getState(originEndpoint) == null) {
+                                LOGGER.debug(
+                                    "Skipping configuration refresh check for {} because monitoring state is not initialized.",
+                                    originEndpoint);
+                                return;
+                            }
+                            refreshWithTime(client, stateHolder.getState(originEndpoint),
+                                monitor.getRefreshInterval(), data, replicaLookUp, ctx);
+                        },
                         eventData,
                         context,
                         "configuration refresh check",
@@ -103,12 +126,12 @@ RefreshEventData refreshStoresCheck(AppConfigurationReplicaClientFactory clientF
 
                 FeatureFlagStore featureStore = connection.getFeatureFlagStore();
 
-                if (featureStore.getEnabled() && StateHolder.getStateFeatureFlag(originEndpoint) != null) {
+                if (featureStore.getEnabled() && stateHolder.getStateFeatureFlag(originEndpoint) != null) {
                     RefreshEventData result = executeRefreshWithRetry(
                         clientFactory,
                         originEndpoint,
                         (client, data, ctx) -> refreshWithTimeFeatureFlags(client,
-                            StateHolder.getStateFeatureFlag(originEndpoint),
+                            stateHolder.getStateFeatureFlag(originEndpoint),
                             monitor.getFeatureFlagRefreshInterval(), data, replicaLookUp, ctx),
                         eventData,
                         context,
@@ -124,7 +147,7 @@ RefreshEventData refreshStoresCheck(AppConfigurationReplicaClientFactory clientF
             }
         } catch (Exception e) {
             // The next refresh will happen sooner if refresh interval is expired.
-            StateHolder.getCurrentState().updateNextRefreshTime(refreshInterval, defaultMinBackoff);
+            stateHolder.updateNextRefreshTime(refreshInterval, defaultMinBackoff);
             throw e;
         }
         return eventData;
@@ -179,12 +202,20 @@ private RefreshEventData executeRefreshWithRetry(
      * @param client the client for checking refresh status
      * @param originEndpoint the original config store endpoint
      * @param context the operation context
+     * @param stateHolder the state holder instance
      * @return true if a refresh should be triggered, false otherwise
      */
-    static boolean refreshStoreCheck(AppConfigurationReplicaClient client, String originEndpoint, Context context) {
+    static boolean refreshStoreCheck(AppConfigurationReplicaClient client, String originEndpoint, Context context,
+        StateHolder stateHolder) {
         RefreshEventData eventData = new RefreshEventData();
-        if (StateHolder.getLoadState(originEndpoint)) {
-            refreshWithoutTime(client, StateHolder.getState(originEndpoint).getWatchKeys(), eventData, context);
+        if (stateHolder.getLoadState(originEndpoint)) {
+            State state = stateHolder.getState(originEndpoint);
+            if (state != null) {
+                refreshWithoutTime(client, state.getWatchKeys(), eventData, context);
+            } else {
+                LOGGER.debug("Skipping configuration refresh check for {} as no watched state is available",
+                    originEndpoint);
+            }
         }
         return eventData.getDoRefresh();
     }
@@ -198,13 +229,13 @@ static boolean refreshStoreCheck(AppConfigurationReplicaClient client, String or
      * @param context the operation context
      * @return true if a refresh should be triggered, false otherwise
      */
-    static boolean refreshStoreFeatureFlagCheck(Boolean featureStoreEnabled,
+    boolean refreshStoreFeatureFlagCheck(Boolean featureStoreEnabled,
         AppConfigurationReplicaClient client, Context context) {
         RefreshEventData eventData = new RefreshEventData();
         String endpoint = client.getEndpoint();
 
-        if (featureStoreEnabled && StateHolder.getStateFeatureFlag(endpoint) != null) {
-            refreshWithoutTimeFeatureFlags(client, StateHolder.getStateFeatureFlag(endpoint), eventData, context);
+        if (featureStoreEnabled && stateHolder.getStateFeatureFlag(endpoint) != null) {
+            refreshWithoutTimeFeatureFlags(client, stateHolder.getStateFeatureFlag(endpoint), eventData, context);
         } else {
             LOGGER.debug("Skipping feature flag refresh check for {}", endpoint);
         }
@@ -223,7 +254,7 @@ static boolean refreshStoreFeatureFlagCheck(Boolean featureStoreEnabled,
      * @param context the operation context
      * @throws AppConfigurationStatusException if there's an error during the refresh check
      */
-    private static void refreshWithTime(AppConfigurationReplicaClient client, State state, Duration refreshInterval,
+    private void refreshWithTime(AppConfigurationReplicaClient client, State state, Duration refreshInterval,
         RefreshEventData eventData, ReplicaLookUp replicaLookUp, Context context)
         throws AppConfigurationStatusException {
         if (Instant.now().isAfter(state.getNextRefreshCheck())) {
@@ -238,7 +269,7 @@ private static void refreshWithTime(AppConfigurationReplicaClient client, State
                 refreshWithoutTime(client, state.getWatchKeys(), eventData, context);
             }
 
-            StateHolder.getCurrentState().updateStateRefresh(state, refreshInterval);
+            stateHolder.updateStateRefresh(state, refreshInterval);
         }
     }
 
@@ -308,7 +339,7 @@ private static void refreshWithoutTimeWatchedConfigurationSettings(AppConfigurat
      * @param context the operation context
      * @throws AppConfigurationStatusException if there's an error during the refresh check
      */
-    private static void refreshWithTimeFeatureFlags(AppConfigurationReplicaClient client, FeatureFlagState state,
+    private void refreshWithTimeFeatureFlags(AppConfigurationReplicaClient client, FeatureFlagState state,
         Duration refreshInterval, RefreshEventData eventData, ReplicaLookUp replicaLookUp, Context context)
         throws AppConfigurationStatusException {
         Instant date = Instant.now();
@@ -327,7 +358,7 @@ private static void refreshWithTimeFeatureFlags(AppConfigurationReplicaClient cl
 
             }
 
-            StateHolder.getCurrentState().updateFeatureFlagStateRefresh(state, refreshInterval);
+            stateHolder.updateFeatureFlagStateRefresh(state, refreshInterval);
         }
     }
 

sdk/spring/spring-cloud-azure-appconfiguration-config/src/main/java/com/azure/spring/cloud/appconfiguration/config/implementation/AppConfigurationReplicaClient.java

@@ -249,7 +249,7 @@ boolean checkWatchKeys(SettingSelector settingSelector, Context context) {
         List<PagedResponse<ConfigurationSetting>> results = client
             .listConfigurationSettings(settingSelector, context)
             .streamByPage().filter(pagedResponse -> pagedResponse.getStatusCode() != HTTP_NOT_MODIFIED).toList();
-        return results.size() > 0;
+        return !results.isEmpty();
     }
 
     /**

sdk/spring/spring-cloud-azure-appconfiguration-config/src/main/java/com/azure/spring/cloud/appconfiguration/config/implementation/AzureAppConfigBoostrapRegistrar.java → sdk/spring/spring-cloud-azure-appconfiguration-config/src/main/java/com/azure/spring/cloud/appconfiguration/config/implementation/AzureAppConfigBootstrapRegistrar.java

@@ -6,6 +6,7 @@
 import org.springframework.boot.context.config.ConfigDataLocationResolverContext;
 import org.springframework.boot.context.properties.bind.Bindable;
 import org.springframework.boot.context.properties.bind.Binder;
+import org.springframework.context.ConfigurableApplicationContext;
 import org.springframework.util.StringUtils;
 
 import com.azure.data.appconfiguration.ConfigurationClientBuilder;
@@ -51,6 +52,17 @@ static void register(ConfigDataLocationResolverContext context, Binder binder,
             InstanceSupplier.from(() -> keyVaultClientFactory));
         context.getBootstrapContext().registerIfAbsent(AppConfigurationReplicaClientFactory.class,
             InstanceSupplier.from(() -> buildClientFactory(replicaClientsBuilder, properties, replicaLookup)));
+
+        // Register StateHolder and promote it to ApplicationContext on close
+        context.getBootstrapContext().registerIfAbsent(StateHolder.class,
+            InstanceSupplier.from(StateHolder::new));
+        context.getBootstrapContext().addCloseListener(event -> {
+            StateHolder stateHolder = event.getBootstrapContext().get(StateHolder.class);
+            ConfigurableApplicationContext applicationContext = event.getApplicationContext();
+            if (!applicationContext.getBeanFactory().containsBean("appConfigurationStateHolder")) {
+                applicationContext.getBeanFactory().registerSingleton("appConfigurationStateHolder", stateHolder);
+            }
+        });
     }
 
     private static AppConfigurationKeyVaultClientFactory appConfigurationKeyVaultClientFactory(