Skip to content

Commit bb97593

Browse files
mrm9084mrm9084
authored
Fixing Load Balancing shuffle (#48274)
* Fixing Load Balancing shuffle * Update ConnectionManager.java * Update ConnectionManager.java
1 parent 6f2a3a5 commit bb97593

File tree

3 files changed

+96
-25
lines changed

3 files changed

+96
-25
lines changed

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

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import java.net.URL;
99
import java.time.Duration;
1010
import java.util.ArrayList;
11+
import java.util.Collections;
1112
import java.util.List;
1213
import java.util.function.Function;
1314
import java.util.regex.Matcher;
@@ -156,7 +157,8 @@ List<AppConfigurationReplicaClient> buildClients(ConfigStore configStore) {
156157
LOGGER.debug("Connecting to " + endpoint + " using Connecting String.");
157158
ConfigurationClientBuilder builder = createBuilderInstance().connectionString(connectionString);
158159

159-
clients.add(modifyAndBuildClient(builder, endpoint, configStore.getEndpoint(), connectionStrings.size() - 1));
160+
clients.add(
161+
modifyAndBuildClient(builder, endpoint, configStore.getEndpoint(), connectionStrings.size() - 1));
160162
}
161163
} else {
162164
DefaultAzureCredential defautAzureCredential = new DefaultAzureCredentialBuilder().build();
@@ -171,6 +173,11 @@ List<AppConfigurationReplicaClient> buildClients(ConfigStore configStore) {
171173
clients.add(modifyAndBuildClient(builder, endpoint, configStore.getEndpoint(), endpoints.size() - 1));
172174
}
173175
}
176+
177+
if (configStore.isLoadBalancingEnabled()) {
178+
Collections.shuffle(clients);
179+
}
180+
174181
return clients;
175182
}
176183

@@ -196,7 +203,8 @@ AppConfigurationReplicaClient buildClient(String failoverEndpoint, ConfigStore c
196203
}
197204
}
198205

199-
private AppConfigurationReplicaClient modifyAndBuildClient(ConfigurationClientBuilder builder, String endpoint, String originEndpoint,
206+
private AppConfigurationReplicaClient modifyAndBuildClient(ConfigurationClientBuilder builder, String endpoint,
207+
String originEndpoint,
200208
Integer replicaCount) {
201209
TracingInfo tracingInfo = new TracingInfo(isKeyVaultConfigured, replicaCount,
202210
Configuration.getGlobalConfiguration());

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

Lines changed: 14 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -97,7 +97,7 @@ String getMainEndpoint() {
9797
AppConfigurationReplicaClient getNextActiveClient(boolean useLastActive) {
9898
if (useLastActive) {
9999
List<AppConfigurationReplicaClient> clients = getAvailableClients();
100-
for (AppConfigurationReplicaClient client: clients) {
100+
for (AppConfigurationReplicaClient client : clients) {
101101
if (client.getEndpoint().equals(lastActiveClient)) {
102102
return client;
103103
}
@@ -106,7 +106,7 @@ AppConfigurationReplicaClient getNextActiveClient(boolean useLastActive) {
106106
if (activeClients.isEmpty()) {
107107
lastActiveClient = "";
108108
return null;
109-
}
109+
}
110110

111111
if (!configStore.isLoadBalancingEnabled()) {
112112
if (!activeClients.isEmpty()) {
@@ -115,7 +115,8 @@ AppConfigurationReplicaClient getNextActiveClient(boolean useLastActive) {
115115
return null;
116116
}
117117

118-
// Remove the current client from the list. The list will be rebuilt and rotated on the next refresh cycle by findActiveClients().
118+
// Remove the current client from the list. The list will be rebuilt and rotated on the next refresh cycle by
119+
// findActiveClients().
119120
AppConfigurationReplicaClient nextClient = activeClients.remove(0);
120121
lastActiveClient = nextClient.getEndpoint();
121122
return nextClient;
@@ -125,27 +126,17 @@ AppConfigurationReplicaClient getNextActiveClient(boolean useLastActive) {
125126
* Finds the currently active clients for a given origin endpoint.
126127
*/
127128
void findActiveClients() {
128-
activeClients = getAvailableClients();
129-
130-
if (!configStore.isLoadBalancingEnabled() || lastActiveClient.isEmpty()) {
129+
// Load balancing enabled: only refresh if no active clients (rotation happens in getNextActiveClient)
130+
if (configStore.isLoadBalancingEnabled()) {
131+
if (activeClients.isEmpty()) {
132+
activeClients = getAvailableClients();
133+
}
131134
return;
132135
}
133-
134-
for (int i = 0; i < activeClients.size(); i++) {
135-
AppConfigurationReplicaClient client = activeClients.get(i);
136-
if (client.getEndpoint().equals(lastActiveClient)) {
137-
int swapPoint = (i + 1) % activeClients.size();
138-
List<AppConfigurationReplicaClient> rotatedClients = new ArrayList<>();
139-
140-
// Add elements from swapPoint to end
141-
rotatedClients.addAll(activeClients.subList(swapPoint, activeClients.size()));
142-
143-
// Add elements from beginning to swapPoint
144-
rotatedClients.addAll(activeClients.subList(0, swapPoint));
145-
146-
activeClients = rotatedClients;
147-
return;
148-
}
136+
// Load balancing disabled: always refresh to use the most preferred available client
137+
List<AppConfigurationReplicaClient> availableClients = getAvailableClients();
138+
if (!availableClients.isEmpty()) {
139+
activeClients = availableClients;
149140
}
150141
}
151142

@@ -177,7 +168,7 @@ public List<AppConfigurationReplicaClient> getAvailableClients() {
177168
}
178169
}
179170
} else if (configStore.isLoadBalancingEnabled()) {
180-
for (AppConfigurationReplicaClient client: clients) {
171+
for (AppConfigurationReplicaClient client : clients) {
181172
if (client.getBackoffEndTime().isBefore(Instant.now())) {
182173
availableClients.add(client);
183174
}

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

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -298,4 +298,76 @@ public void buildClientConnectionStringInvalid3Test() {
298298
exception.getMessage());
299299
}
300300

301+
@Test
302+
public void buildClientsWithLoadBalancingEnabledTest() {
303+
// Test that load balancing shuffles the order of clients
304+
configStore = new ConfigStore();
305+
List<String> endpoints = new ArrayList<>();
306+
307+
// Add multiple endpoints to shuffle
308+
endpoints.add(TEST_ENDPOINT);
309+
endpoints.add(TEST_ENDPOINT_GEO);
310+
endpoints.add("https://third.test.config.io");
311+
endpoints.add("https://fourth.test.config.io");
312+
313+
configStore.setEndpoints(endpoints);
314+
configStore.setLoadBalancingEnabled(true);
315+
configStore.validateAndInit();
316+
317+
clientBuilder = new AppConfigurationReplicaClientsBuilder(clientFactoryMock, null, false, false);
318+
AppConfigurationReplicaClientsBuilder spy = Mockito.spy(clientBuilder);
319+
320+
ConfigurationClientBuilder builder = new ConfigurationClientBuilder();
321+
when(builderMock.endpoint(Mockito.anyString())).thenReturn(builder);
322+
when(builderMock.addPolicy(Mockito.any())).thenReturn(builderMock);
323+
when(clientFactoryMock.build()).thenReturn(builderMock);
324+
325+
// Build clients multiple times and verify that order changes (due to shuffle)
326+
List<AppConfigurationReplicaClient> clients1 = spy.buildClients(configStore);
327+
List<AppConfigurationReplicaClient> clients2 = spy.buildClients(configStore);
328+
329+
assertEquals(4, clients1.size());
330+
assertEquals(4, clients2.size());
331+
332+
// The order should potentially differ due to random shuffling
333+
// We can't guarantee they're different, but we can verify the same endpoints exist
334+
List<String> endpoints1 = clients1.stream().map(AppConfigurationReplicaClient::getEndpoint).toList();
335+
List<String> endpoints2 = clients2.stream().map(AppConfigurationReplicaClient::getEndpoint).toList();
336+
337+
// All endpoints should be present in both lists (just potentially in different order)
338+
assertTrue(endpoints1.containsAll(endpoints));
339+
assertTrue(endpoints2.containsAll(endpoints));
340+
}
341+
342+
@Test
343+
public void buildClientsWithLoadBalancingDisabledTest() {
344+
// Test that without load balancing, order is preserved
345+
configStore = new ConfigStore();
346+
List<String> endpoints = new ArrayList<>();
347+
348+
endpoints.add(TEST_ENDPOINT);
349+
endpoints.add(TEST_ENDPOINT_GEO);
350+
endpoints.add("https://third.test.config.io");
351+
352+
configStore.setEndpoints(endpoints);
353+
configStore.setLoadBalancingEnabled(false);
354+
configStore.validateAndInit();
355+
356+
clientBuilder = new AppConfigurationReplicaClientsBuilder(clientFactoryMock, null, false, false);
357+
AppConfigurationReplicaClientsBuilder spy = Mockito.spy(clientBuilder);
358+
359+
ConfigurationClientBuilder builder = new ConfigurationClientBuilder();
360+
when(builderMock.endpoint(Mockito.anyString())).thenReturn(builder);
361+
when(builderMock.addPolicy(Mockito.any())).thenReturn(builderMock);
362+
when(clientFactoryMock.build()).thenReturn(builderMock);
363+
364+
List<AppConfigurationReplicaClient> clients = spy.buildClients(configStore);
365+
366+
assertEquals(3, clients.size());
367+
// When load balancing is disabled, order should match input order
368+
assertEquals(TEST_ENDPOINT, clients.get(0).getEndpoint());
369+
assertEquals(TEST_ENDPOINT_GEO, clients.get(1).getEndpoint());
370+
assertEquals("https://third.test.config.io", clients.get(2).getEndpoint());
371+
}
372+
301373
}

0 commit comments

Comments
 (0)