diff --git a/sdk/spring/spring-cloud-azure-appconfiguration-config/src/main/java/com/azure/spring/cloud/appconfiguration/config/implementation/AppConfigurationReplicaClientsBuilder.java b/sdk/spring/spring-cloud-azure-appconfiguration-config/src/main/java/com/azure/spring/cloud/appconfiguration/config/implementation/AppConfigurationReplicaClientsBuilder.java index 6b4d42827626..8832196d5dc2 100644 --- a/sdk/spring/spring-cloud-azure-appconfiguration-config/src/main/java/com/azure/spring/cloud/appconfiguration/config/implementation/AppConfigurationReplicaClientsBuilder.java +++ b/sdk/spring/spring-cloud-azure-appconfiguration-config/src/main/java/com/azure/spring/cloud/appconfiguration/config/implementation/AppConfigurationReplicaClientsBuilder.java @@ -8,6 +8,7 @@ import java.net.URL; import java.time.Duration; import java.util.ArrayList; +import java.util.Collections; import java.util.List; import java.util.function.Function; import java.util.regex.Matcher; @@ -156,7 +157,8 @@ List buildClients(ConfigStore configStore) { LOGGER.debug("Connecting to " + endpoint + " using Connecting String."); ConfigurationClientBuilder builder = createBuilderInstance().connectionString(connectionString); - clients.add(modifyAndBuildClient(builder, endpoint, configStore.getEndpoint(), connectionStrings.size() - 1)); + clients.add( + modifyAndBuildClient(builder, endpoint, configStore.getEndpoint(), connectionStrings.size() - 1)); } } else { DefaultAzureCredential defautAzureCredential = new DefaultAzureCredentialBuilder().build(); @@ -171,6 +173,11 @@ List buildClients(ConfigStore configStore) { clients.add(modifyAndBuildClient(builder, endpoint, configStore.getEndpoint(), endpoints.size() - 1)); } } + + if (configStore.isLoadBalancingEnabled()) { + Collections.shuffle(clients); + } + return clients; } @@ -196,7 +203,8 @@ AppConfigurationReplicaClient buildClient(String failoverEndpoint, ConfigStore c } } - private AppConfigurationReplicaClient modifyAndBuildClient(ConfigurationClientBuilder builder, String endpoint, String originEndpoint, + private AppConfigurationReplicaClient modifyAndBuildClient(ConfigurationClientBuilder builder, String endpoint, + String originEndpoint, Integer replicaCount) { TracingInfo tracingInfo = new TracingInfo(isKeyVaultConfigured, replicaCount, Configuration.getGlobalConfiguration()); diff --git a/sdk/spring/spring-cloud-azure-appconfiguration-config/src/main/java/com/azure/spring/cloud/appconfiguration/config/implementation/ConnectionManager.java b/sdk/spring/spring-cloud-azure-appconfiguration-config/src/main/java/com/azure/spring/cloud/appconfiguration/config/implementation/ConnectionManager.java index 4b7f6f2c1ca1..be0581a20f74 100644 --- a/sdk/spring/spring-cloud-azure-appconfiguration-config/src/main/java/com/azure/spring/cloud/appconfiguration/config/implementation/ConnectionManager.java +++ b/sdk/spring/spring-cloud-azure-appconfiguration-config/src/main/java/com/azure/spring/cloud/appconfiguration/config/implementation/ConnectionManager.java @@ -97,7 +97,7 @@ String getMainEndpoint() { AppConfigurationReplicaClient getNextActiveClient(boolean useLastActive) { if (useLastActive) { List clients = getAvailableClients(); - for (AppConfigurationReplicaClient client: clients) { + for (AppConfigurationReplicaClient client : clients) { if (client.getEndpoint().equals(lastActiveClient)) { return client; } @@ -106,7 +106,7 @@ AppConfigurationReplicaClient getNextActiveClient(boolean useLastActive) { if (activeClients.isEmpty()) { lastActiveClient = ""; return null; - } + } if (!configStore.isLoadBalancingEnabled()) { if (!activeClients.isEmpty()) { @@ -115,7 +115,8 @@ 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(). + // 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; @@ -125,27 +126,17 @@ AppConfigurationReplicaClient getNextActiveClient(boolean useLastActive) { * Finds the currently active clients for a given origin endpoint. */ void findActiveClients() { - activeClients = getAvailableClients(); - - if (!configStore.isLoadBalancingEnabled() || lastActiveClient.isEmpty()) { + // Load balancing enabled: only refresh if no active clients (rotation happens in getNextActiveClient) + if (configStore.isLoadBalancingEnabled()) { + if (activeClients.isEmpty()) { + activeClients = getAvailableClients(); + } return; } - - for (int i = 0; i < activeClients.size(); i++) { - AppConfigurationReplicaClient client = activeClients.get(i); - if (client.getEndpoint().equals(lastActiveClient)) { - int swapPoint = (i + 1) % activeClients.size(); - List rotatedClients = new ArrayList<>(); - - // Add elements from swapPoint to end - rotatedClients.addAll(activeClients.subList(swapPoint, activeClients.size())); - - // Add elements from beginning to swapPoint - rotatedClients.addAll(activeClients.subList(0, swapPoint)); - - activeClients = rotatedClients; - return; - } + // Load balancing disabled: always refresh to use the most preferred available client + List availableClients = getAvailableClients(); + if (!availableClients.isEmpty()) { + activeClients = availableClients; } } @@ -177,7 +168,7 @@ public List getAvailableClients() { } } } else if (configStore.isLoadBalancingEnabled()) { - for (AppConfigurationReplicaClient client: clients) { + for (AppConfigurationReplicaClient client : clients) { if (client.getBackoffEndTime().isBefore(Instant.now())) { availableClients.add(client); } diff --git a/sdk/spring/spring-cloud-azure-appconfiguration-config/src/test/java/com/azure/spring/cloud/appconfiguration/config/implementation/AppConfigurationReplicaClientBuilderTest.java b/sdk/spring/spring-cloud-azure-appconfiguration-config/src/test/java/com/azure/spring/cloud/appconfiguration/config/implementation/AppConfigurationReplicaClientBuilderTest.java index 10e10c33787b..3b95e77724f2 100644 --- a/sdk/spring/spring-cloud-azure-appconfiguration-config/src/test/java/com/azure/spring/cloud/appconfiguration/config/implementation/AppConfigurationReplicaClientBuilderTest.java +++ b/sdk/spring/spring-cloud-azure-appconfiguration-config/src/test/java/com/azure/spring/cloud/appconfiguration/config/implementation/AppConfigurationReplicaClientBuilderTest.java @@ -298,4 +298,76 @@ public void buildClientConnectionStringInvalid3Test() { exception.getMessage()); } + @Test + public void buildClientsWithLoadBalancingEnabledTest() { + // Test that load balancing shuffles the order of clients + configStore = new ConfigStore(); + List endpoints = new ArrayList<>(); + + // Add multiple endpoints to shuffle + endpoints.add(TEST_ENDPOINT); + endpoints.add(TEST_ENDPOINT_GEO); + endpoints.add("https://third.test.config.io"); + endpoints.add("https://fourth.test.config.io"); + + configStore.setEndpoints(endpoints); + configStore.setLoadBalancingEnabled(true); + configStore.validateAndInit(); + + clientBuilder = new AppConfigurationReplicaClientsBuilder(clientFactoryMock, null, false, false); + AppConfigurationReplicaClientsBuilder spy = Mockito.spy(clientBuilder); + + ConfigurationClientBuilder builder = new ConfigurationClientBuilder(); + when(builderMock.endpoint(Mockito.anyString())).thenReturn(builder); + when(builderMock.addPolicy(Mockito.any())).thenReturn(builderMock); + when(clientFactoryMock.build()).thenReturn(builderMock); + + // Build clients multiple times and verify that order changes (due to shuffle) + List clients1 = spy.buildClients(configStore); + List clients2 = spy.buildClients(configStore); + + assertEquals(4, clients1.size()); + assertEquals(4, clients2.size()); + + // The order should potentially differ due to random shuffling + // We can't guarantee they're different, but we can verify the same endpoints exist + List endpoints1 = clients1.stream().map(AppConfigurationReplicaClient::getEndpoint).toList(); + List endpoints2 = clients2.stream().map(AppConfigurationReplicaClient::getEndpoint).toList(); + + // All endpoints should be present in both lists (just potentially in different order) + assertTrue(endpoints1.containsAll(endpoints)); + assertTrue(endpoints2.containsAll(endpoints)); + } + + @Test + public void buildClientsWithLoadBalancingDisabledTest() { + // Test that without load balancing, order is preserved + configStore = new ConfigStore(); + List endpoints = new ArrayList<>(); + + endpoints.add(TEST_ENDPOINT); + endpoints.add(TEST_ENDPOINT_GEO); + endpoints.add("https://third.test.config.io"); + + configStore.setEndpoints(endpoints); + configStore.setLoadBalancingEnabled(false); + configStore.validateAndInit(); + + clientBuilder = new AppConfigurationReplicaClientsBuilder(clientFactoryMock, null, false, false); + AppConfigurationReplicaClientsBuilder spy = Mockito.spy(clientBuilder); + + ConfigurationClientBuilder builder = new ConfigurationClientBuilder(); + when(builderMock.endpoint(Mockito.anyString())).thenReturn(builder); + when(builderMock.addPolicy(Mockito.any())).thenReturn(builderMock); + when(clientFactoryMock.build()).thenReturn(builderMock); + + List clients = spy.buildClients(configStore); + + assertEquals(3, clients.size()); + // When load balancing is disabled, order should match input order + assertEquals(TEST_ENDPOINT, clients.get(0).getEndpoint()); + assertEquals(TEST_ENDPOINT_GEO, clients.get(1).getEndpoint()); + assertEquals("https://third.test.config.io", clients.get(2).getEndpoint()); + } + }