Load balance bug

sdk/spring/spring-cloud-azure-appconfiguration-config/CHANGELOG.md

@@ -10,6 +10,7 @@
 
 ### Bugs Fixed
 
+- Fixed an issue where feature flag–based refresh did not work when load balancing was enabled with a single configuration store. Feature flag refresh now uses the same load-balanced client selection as configuration refresh, including the single-store scenario.
 - Fixed YAML configuration binding for `label-filter` by adding standard no-arg getter methods to `AppConfigurationKeyValueSelector` and `FeatureFlagKeyValueSelector`, enabling proper type resolution by Spring Boot's `@ConfigurationProperties` binder.
 
 ### Other Changes

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

@@ -69,8 +69,6 @@ RefreshEventData refreshStoresCheck(AppConfigurationReplicaClientFactory clientF
             for (Entry<String, ConnectionManager> entry : clientFactory.getConnections().entrySet()) {
                 String originEndpoint = entry.getKey();
                 ConnectionManager connection = entry.getValue();
-                // For safety reset current used replica.
-                clientFactory.setCurrentConfigStoreClient(originEndpoint, originEndpoint);
 
                 AppConfigurationStoreMonitoring monitor = connection.getMonitoring();
 
@@ -94,7 +92,8 @@ RefreshEventData refreshStoresCheck(AppConfigurationReplicaClientFactory clientF
                             monitor.getRefreshInterval(), data, replicaLookUp, ctx),
                         eventData,
                         context,
-                        "configuration refresh check");
+                        "configuration refresh check",
+                        false);
                     if (result != null) {
                         return result;
                     }
@@ -113,7 +112,8 @@ RefreshEventData refreshStoresCheck(AppConfigurationReplicaClientFactory clientF
                             monitor.getFeatureFlagRefreshInterval(), data, replicaLookUp, ctx),
                         eventData,
                         context,
-                        "feature flag refresh check");
+                        "feature flag refresh check",
+                        true);
                     if (result != null) {
                         return result;
                     }
@@ -139,6 +139,7 @@ RefreshEventData refreshStoresCheck(AppConfigurationReplicaClientFactory clientF
      * @param eventData the refresh event data to update
      * @param context the operation context
      * @param checkType description of the check type for logging (e.g., "configuration refresh check")
+     * @param useLastActive whether to reuse the last active client
      * @return the eventData if refresh is needed, null otherwise
      */
     private RefreshEventData executeRefreshWithRetry(
@@ -147,14 +148,14 @@ private RefreshEventData executeRefreshWithRetry(
         RefreshOperation operation,
         RefreshEventData eventData,
         Context context,
-        String checkType) {
-        AppConfigurationReplicaClient client = clientFactory.getNextActiveClient(originEndpoint, false);
+        String checkType,
+        boolean useLastActive) {
+        AppConfigurationReplicaClient client = clientFactory.getNextActiveClient(originEndpoint, useLastActive);
 
         while (client != null) {
             try {
                 operation.execute(client, eventData, context);
                 if (eventData.getDoRefresh()) {
-                    clientFactory.setCurrentConfigStoreClient(originEndpoint, client.getEndpoint());
                     return eventData;
                 }
                 // If check didn't throw an error, other clients don't need to be checked.

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

@@ -107,16 +107,6 @@ String findOriginForEndpoint(String endpoint) {
         return endpoint;
     }
 
-    /**
-     * Sets the current active replica for a configuration store.
-     * 
-     * @param originEndpoint the origin configuration store endpoint
-     * @param replicaEndpoint the replica endpoint that was successfully connected to
-     */
-    void setCurrentConfigStoreClient(String originEndpoint, String replicaEndpoint) {
-        CONNECTIONS.get(originEndpoint).setCurrentClient(replicaEndpoint);
-    }
-
     /**
      * Updates the sync token for a specific replica endpoint.
      * 

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

@@ -34,9 +34,6 @@ class ConnectionManager {
     /** Map of auto-discovered failover clients, keyed by endpoint URL. */
     private final Map<String, AppConfigurationReplicaClient> autoFailoverClients;
 
-    /** Currently active replica endpoint being used for requests. */
-    private String currentReplica;
-
     /** Current health status of the App Configuration store connection. */
     private AppConfigurationStoreHealth health;
 
@@ -67,7 +64,6 @@ class ConnectionManager {
         this.configStore = configStore;
         this.originEndpoint = configStore.getEndpoint();
         this.health = AppConfigurationStoreHealth.NOT_LOADED;
-        this.currentReplica = configStore.getEndpoint();
         this.autoFailoverClients = new HashMap<>();
         this.replicaLookUp = replicaLookUp;
         this.activeClients = new ArrayList<>();
@@ -83,15 +79,6 @@ AppConfigurationStoreHealth getHealth() {
         return this.health;
     }
 
-    /**
-     * Sets the current active replica endpoint for client routing.
-     * 
-     * @param replicaEndpoint the endpoint URL to set as current; may be null to reset to primary endpoint
-     */
-    void setCurrentClient(String replicaEndpoint) {
-        this.currentReplica = replicaEndpoint;
-    }
-
     /**
      * Retrieves the primary (origin) endpoint URL for the App Configuration store.
      * 
@@ -108,17 +95,18 @@ String getMainEndpoint() {
      * @return the next active AppConfigurationReplicaClient
      */
     AppConfigurationReplicaClient getNextActiveClient(boolean useLastActive) {
-        if (activeClients.isEmpty()) {
-            lastActiveClient = "";
-            return null;
-        } else if (useLastActive) {
+        if (useLastActive) {
             List<AppConfigurationReplicaClient> clients = getAvailableClients();
             for (AppConfigurationReplicaClient client: clients) {
                 if (client.getEndpoint().equals(lastActiveClient)) {
                     return client;
                 }
             }
         }
+        if (activeClients.isEmpty()) {
+            lastActiveClient = "";
+            return null;
+        } 
 
         if (!configStore.isLoadBalancingEnabled()) {
             if (!activeClients.isEmpty()) {
@@ -127,6 +115,7 @@ AppConfigurationReplicaClient getNextActiveClient(boolean useLastActive) {
             return null;
         }
 
+        // Remove the current client from the list. The list will be rebuilt and rotated on the next refresh cycle by findActiveClients().
         AppConfigurationReplicaClient nextClient = activeClients.remove(0);
         lastActiveClient = nextClient.getEndpoint();
         return nextClient;

sdk/spring/spring-cloud-azure-appconfiguration-config/src/test/java/com/azure/spring/cloud/appconfiguration/config/implementation/AppConfigurationRefreshUtilTest.java

@@ -248,7 +248,6 @@ public void refreshStoresCheckSettingsTestNotEnabled(TestInfo testInfo) {
                 Duration.ofMinutes(10),
                 (long) 60, replicaLookUpMock);
             assertFalse(eventData.getDoRefresh());
-            verify(clientFactoryMock, times(1)).setCurrentConfigStoreClient(Mockito.eq(endpoint), Mockito.eq(endpoint));
             verify(clientOriginMock, times(0)).getWatchKey(Mockito.anyString(), Mockito.anyString(),
                 Mockito.any(Context.class));
         }
@@ -271,7 +270,6 @@ public void refreshStoresCheckSettingsTestNotLoaded(TestInfo testInfo) {
             RefreshEventData eventData = new AppConfigurationRefreshUtil().refreshStoresCheck(clientFactoryMock,
                 Duration.ofMinutes(10), (long) 60, replicaLookUpMock);
             assertFalse(eventData.getDoRefresh());
-            verify(clientFactoryMock, times(1)).setCurrentConfigStoreClient(Mockito.eq(endpoint), Mockito.eq(endpoint));
             verify(clientOriginMock, times(0)).getWatchKey(Mockito.anyString(), Mockito.anyString(),
                 Mockito.any(Context.class));
         }
@@ -292,7 +290,6 @@ public void refreshStoresCheckSettingsTestNotRefreshTime(TestInfo testInfo) {
                 Duration.ofMinutes(10),
                 (long) 60, replicaLookUpMock);
             assertFalse(eventData.getDoRefresh());
-            verify(clientFactoryMock, times(1)).setCurrentConfigStoreClient(Mockito.eq(endpoint), Mockito.eq(endpoint));
             verify(clientOriginMock, times(0)).getWatchKey(Mockito.anyString(), Mockito.anyString(),
                 Mockito.any(Context.class));
         }
@@ -316,7 +313,6 @@ public void refreshStoresCheckSettingsTestFailedRequest(TestInfo testInfo) {
                 (long) 60, replicaLookUpMock);
             assertFalse(eventData.getDoRefresh());
             ArgumentCaptor<Context> captorParam = ArgumentCaptor.forClass(Context.class);
-            verify(clientFactoryMock, times(1)).setCurrentConfigStoreClient(Mockito.eq(endpoint), Mockito.eq(endpoint));
             verify(clientOriginMock, times(1)).getWatchKey(Mockito.anyString(), Mockito.anyString(),
                 captorParam.capture());
             assertEquals(newState, StateHolder.getState(endpoint));
@@ -346,7 +342,6 @@ public void refreshStoresCheckSettingsTestRefreshTimeNoChange(TestInfo testInfo)
                 Duration.ofMinutes(10), (long) 60, replicaLookUpMock);
             assertEquals(newState, StateHolder.getState(endpoint));
             assertFalse(eventData.getDoRefresh());
-            verify(clientFactoryMock, times(1)).setCurrentConfigStoreClient(Mockito.eq(endpoint), Mockito.eq(endpoint));
             verify(clientOriginMock, times(1)).getWatchKey(Mockito.anyString(), Mockito.anyString(),
                 Mockito.any(Context.class));
         }
@@ -380,7 +375,6 @@ public void refreshStoresPushRefreshEnabledPrimary(TestInfo testInfo) {
             assertEquals(newState, StateHolder.getState(endpoint));
             assertFalse(eventData.getDoRefresh());
             ArgumentCaptor<Context> captorParam = ArgumentCaptor.forClass(Context.class);
-            verify(clientFactoryMock, times(1)).setCurrentConfigStoreClient(Mockito.eq(endpoint), Mockito.eq(endpoint));
             verify(clientOriginMock, times(1)).getWatchKey(Mockito.anyString(), Mockito.anyString(),
                 captorParam.capture());
             Context testContext = captorParam.getValue();
@@ -417,7 +411,6 @@ public void refreshStoresPushRefreshEnabledSecondary(TestInfo testInfo) {
             assertEquals(newState, StateHolder.getState(endpoint));
             assertFalse(eventData.getDoRefresh());
             ArgumentCaptor<Context> captorParam = ArgumentCaptor.forClass(Context.class);
-            verify(clientFactoryMock, times(1)).setCurrentConfigStoreClient(Mockito.eq(endpoint), Mockito.eq(endpoint));
             verify(clientOriginMock, times(1)).getWatchKey(Mockito.anyString(), Mockito.anyString(),
                 captorParam.capture());
             Context testContext = captorParam.getValue();
@@ -452,7 +445,6 @@ public void refreshStoresCheckSettingsTestTriggerRefresh(TestInfo testInfo) {
                 Duration.ofMinutes(10),
                 (long) 60, replicaLookUpMock);
             assertTrue(eventData.getDoRefresh());
-            verify(clientFactoryMock, times(1)).setCurrentConfigStoreClient(Mockito.eq(endpoint), Mockito.eq(endpoint));
             verify(clientOriginMock, times(1)).getWatchKey(Mockito.anyString(), Mockito.anyString(),
                 Mockito.any(Context.class));
             verify(currentStateMock, times(1)).updateStateRefresh(Mockito.any(), Mockito.any());
@@ -476,7 +468,6 @@ public void refreshStoresCheckFeatureFlagTestNotLoaded(TestInfo testInfo) {
                 Duration.ofMinutes(10),
                 (long) 60, replicaLookUpMock);
             assertFalse(eventData.getDoRefresh());
-            verify(clientFactoryMock, times(1)).setCurrentConfigStoreClient(Mockito.eq(endpoint), Mockito.eq(endpoint));
             verify(clientOriginMock, times(0)).getWatchKey(Mockito.anyString(), Mockito.anyString(),
                 Mockito.any(Context.class));
         }
@@ -499,7 +490,6 @@ public void refreshStoresCheckFeatureFlagTestNotRefreshTime(TestInfo testInfo) {
                 Duration.ofMinutes(10),
                 (long) 60, replicaLookUpMock);
             assertFalse(eventData.getDoRefresh());
-            verify(clientFactoryMock, times(1)).setCurrentConfigStoreClient(Mockito.eq(endpoint), Mockito.eq(endpoint));
             verify(clientOriginMock, times(0)).getWatchKey(Mockito.anyString(), Mockito.anyString(),
                 Mockito.any(Context.class));
         }
@@ -529,7 +519,6 @@ public void refreshStoresCheckFeatureFlagTestNoChange(TestInfo testInfo) {
             RefreshEventData eventData = new AppConfigurationRefreshUtil().refreshStoresCheck(clientFactoryMock,
                 Duration.ofMinutes(10), (long) 60, replicaLookUpMock);
             assertFalse(eventData.getDoRefresh());
-            verify(clientFactoryMock, times(1)).setCurrentConfigStoreClient(Mockito.eq(endpoint), Mockito.eq(endpoint));
             verify(clientOriginMock, times(0)).getWatchKey(Mockito.anyString(), Mockito.anyString(),
                 Mockito.any(Context.class));
             verify(currentStateMock, times(1)).updateFeatureFlagStateRefresh(Mockito.any(), Mockito.any());
@@ -558,7 +547,6 @@ public void refreshStoresCheckFeatureFlagTestTriggerRefresh(TestInfo testInfo) {
             RefreshEventData eventData = new AppConfigurationRefreshUtil().refreshStoresCheck(clientFactoryMock,
                 Duration.ofMinutes(10), (long) 60, replicaLookUpMock);
             assertTrue(eventData.getDoRefresh());
-            verify(clientFactoryMock, times(1)).setCurrentConfigStoreClient(Mockito.eq(endpoint), Mockito.eq(endpoint));
             verify(clientOriginMock, times(0)).getWatchKey(Mockito.anyString(), Mockito.anyString(),
                 Mockito.any(Context.class));
         }
@@ -631,7 +619,6 @@ public void refreshAllWithWatchedConfigurationSettingsTest(TestInfo testInfo) {
                 clientFactoryMock, Duration.ofMinutes(10), (long) 60, replicaLookUpMock);
 
             assertTrue(eventData.getDoRefresh());
-            verify(clientFactoryMock, times(1)).setCurrentConfigStoreClient(Mockito.eq(endpoint), Mockito.eq(endpoint));
             // Verify checkWatchKeys is called (watched configuration settings path)
             verify(clientOriginMock, times(1)).checkWatchKeys(Mockito.any(SettingSelector.class),
                 Mockito.any(Context.class));

sdk/spring/spring-cloud-azure-appconfiguration-config/src/test/java/com/azure/spring/cloud/appconfiguration/config/implementation/ConnectionManagerTest.java

@@ -277,7 +277,7 @@ public void getNextActiveClientWithLoadBalancingTest() {
             throw new RuntimeException("Test verification failed", e);
         }
 
-        // activeClients list should now only have the second client
+        // activeClients list should now contain only the second client
         try {
             java.lang.reflect.Field activeClientsField = ConnectionManager.class.getDeclaredField("activeClients");
             activeClientsField.setAccessible(true);