core: Optional address shuffle in PickFirstLoadBalancer (#10110)

If provided with the new PickFirstLoadBalancerConfig,
PickFirstLoadBalancer will shuffle the list of addresses it receives
from the name resolver.

PickFirstLoadBalancerProvider will now support the new config
if enabled by an env variable.

Author: temawi

core/src/main/java/io/grpc/internal/PickFirstLoadBalancer.java

@@ -28,8 +28,12 @@
 import io.grpc.EquivalentAddressGroup;
 import io.grpc.LoadBalancer;
 import io.grpc.Status;
+import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
+import java.util.Random;
 import java.util.concurrent.atomic.AtomicBoolean;
+import javax.annotation.Nullable;
 
 /**
  * A {@link LoadBalancer} that provides no load-balancing over the addresses from the {@link
@@ -55,6 +59,18 @@ public boolean acceptResolvedAddresses(ResolvedAddresses resolvedAddresses) {
       return false;
     }
 
+    // We can optionally be configured to shuffle the address list. This can help better distribute
+    // the load.
+    if (resolvedAddresses.getLoadBalancingPolicyConfig() instanceof PickFirstLoadBalancerConfig) {
+      PickFirstLoadBalancerConfig config
+          = (PickFirstLoadBalancerConfig) resolvedAddresses.getLoadBalancingPolicyConfig();
+      if (config.shuffleAddressList != null && config.shuffleAddressList) {
+        servers = new ArrayList<EquivalentAddressGroup>(servers);
+        Collections.shuffle(servers,
+            config.randomSeed != null ? new Random(config.randomSeed) : new Random());
+      }
+    }
+
     if (subchannel == null) {
       final Subchannel subchannel = helper.createSubchannel(
           CreateSubchannelArgs.newBuilder()
@@ -199,4 +215,22 @@ public void run() {
       return PickResult.withNoResult();
     }
   }
+
+  public static final class PickFirstLoadBalancerConfig {
+
+    @Nullable
+    public final Boolean shuffleAddressList;
+
+    // For testing purposes only, not meant to be parsed from a real config.
+    @Nullable final Long randomSeed;
+
+    public PickFirstLoadBalancerConfig(@Nullable Boolean shuffleAddressList) {
+      this(shuffleAddressList, null);
+    }
+
+    PickFirstLoadBalancerConfig(@Nullable Boolean shuffleAddressList, @Nullable Long randomSeed) {
+      this.shuffleAddressList = shuffleAddressList;
+      this.randomSeed = randomSeed;
+    }
+  }
 }

core/src/main/java/io/grpc/internal/PickFirstLoadBalancerProvider.java

@@ -16,10 +16,13 @@
 
 package io.grpc.internal;
 
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Strings;
 import io.grpc.LoadBalancer;
 import io.grpc.LoadBalancerProvider;
 import io.grpc.NameResolver;
 import io.grpc.NameResolver.ConfigOrError;
+import io.grpc.internal.PickFirstLoadBalancer.PickFirstLoadBalancerConfig;
 import java.util.Map;
 
 /**
@@ -30,6 +33,10 @@
  */
 public final class PickFirstLoadBalancerProvider extends LoadBalancerProvider {
   private static final String NO_CONFIG = "no service config";
+  private static final String SHUFFLE_ADDRESS_LIST_KEY = "shuffleAddressList";
+  private static final String CONFIG_FLAG_NAME = "GRPC_EXPERIMENTAL_PICKFIRST_LB_CONFIG";
+  @VisibleForTesting
+  static boolean enablePickFirstConfig = !Strings.isNullOrEmpty(System.getenv(CONFIG_FLAG_NAME));
 
   @Override
   public boolean isAvailable() {
@@ -54,6 +61,12 @@ public LoadBalancer newLoadBalancer(LoadBalancer.Helper helper) {
   @Override
   public ConfigOrError parseLoadBalancingPolicyConfig(
       Map<String, ?> rawLoadBalancingPolicyConfig) {
-    return ConfigOrError.fromConfig(NO_CONFIG);
+    if (enablePickFirstConfig) {
+      return ConfigOrError.fromConfig(
+          new PickFirstLoadBalancerConfig(JsonUtil.getBoolean(rawLoadBalancingPolicyConfig,
+              SHUFFLE_ADDRESS_LIST_KEY)));
+    } else {
+      return ConfigOrError.fromConfig(NO_CONFIG);
+    }
   }
 }

core/src/test/java/io/grpc/internal/PickFirstLoadBalancerProviderTest.java

@@ -0,0 +1,61 @@
+/*
+ * Copyright 2023 The gRPC Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.grpc.internal;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import io.grpc.NameResolver.ConfigOrError;
+import io.grpc.internal.PickFirstLoadBalancer.PickFirstLoadBalancerConfig;
+import java.util.HashMap;
+import java.util.Map;
+import org.junit.After;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class PickFirstLoadBalancerProviderTest {
+
+  @After
+  public void resetConfigFlag() {
+    PickFirstLoadBalancerProvider.enablePickFirstConfig = false;
+  }
+
+  @Test
+  public void parseWithConfigEnabled() {
+    PickFirstLoadBalancerProvider.enablePickFirstConfig = true;
+    Map<String, Object> rawConfig = new HashMap<>();
+    rawConfig.put("shuffleAddressList", true);
+    ConfigOrError parsedConfig = new PickFirstLoadBalancerProvider().parseLoadBalancingPolicyConfig(
+        rawConfig);
+    PickFirstLoadBalancerConfig config = (PickFirstLoadBalancerConfig) parsedConfig.getConfig();
+
+    assertThat(config.shuffleAddressList).isTrue();
+  }
+
+  @Test
+  public void parseWithConfigDisabled() {
+    PickFirstLoadBalancerProvider.enablePickFirstConfig = false;
+    Map<String, Object> rawConfig = new HashMap<>();
+    rawConfig.put("shuffleAddressList", true);
+    ConfigOrError parsedConfig = new PickFirstLoadBalancerProvider().parseLoadBalancingPolicyConfig(
+        rawConfig);
+    String config = (String) parsedConfig.getConfig();
+
+    assertThat(config).isEqualTo("no service config");
+  }
+}

core/src/test/java/io/grpc/internal/PickFirstLoadBalancerTest.java

@@ -52,6 +52,7 @@
 import io.grpc.Status;
 import io.grpc.Status.Code;
 import io.grpc.SynchronizationContext;
+import io.grpc.internal.PickFirstLoadBalancer.PickFirstLoadBalancerConfig;
 import java.net.SocketAddress;
 import java.util.List;
 import org.junit.After;
@@ -140,6 +141,49 @@ public void pickAfterResolved() throws Exception {
     verifyNoMoreInteractions(mockHelper);
   }
 
+  @Test
+  public void pickAfterResolved_shuffle() throws Exception {
+    loadBalancer.acceptResolvedAddresses(
+        ResolvedAddresses.newBuilder().setAddresses(servers).setAttributes(affinity)
+            .setLoadBalancingPolicyConfig(new PickFirstLoadBalancerConfig(true, 123L)).build());
+
+    verify(mockHelper).createSubchannel(createArgsCaptor.capture());
+    CreateSubchannelArgs args = createArgsCaptor.getValue();
+    // We should still see the same set of addresses.
+    assertThat(args.getAddresses()).containsExactlyElementsIn(servers);
+    // Because we use a fixed seed, the addresses should always be shuffled in this order.
+    assertThat(args.getAddresses().get(0)).isEqualTo(servers.get(1));
+    assertThat(args.getAddresses().get(1)).isEqualTo(servers.get(0));
+    assertThat(args.getAddresses().get(2)).isEqualTo(servers.get(2));
+    verify(mockHelper).updateBalancingState(eq(CONNECTING), pickerCaptor.capture());
+    verify(mockSubchannel).requestConnection();
+
+    // Calling pickSubchannel() twice gave the same result
+    assertEquals(pickerCaptor.getValue().pickSubchannel(mockArgs),
+        pickerCaptor.getValue().pickSubchannel(mockArgs));
+
+    verifyNoMoreInteractions(mockHelper);
+  }
+
+  @Test
+  public void pickAfterResolved_noShuffle() throws Exception {
+    loadBalancer.acceptResolvedAddresses(
+        ResolvedAddresses.newBuilder().setAddresses(servers).setAttributes(affinity)
+            .setLoadBalancingPolicyConfig(new PickFirstLoadBalancerConfig(false)).build());
+
+    verify(mockHelper).createSubchannel(createArgsCaptor.capture());
+    CreateSubchannelArgs args = createArgsCaptor.getValue();
+    assertThat(args.getAddresses()).isEqualTo(servers);
+    verify(mockHelper).updateBalancingState(eq(CONNECTING), pickerCaptor.capture());
+    verify(mockSubchannel).requestConnection();
+
+    // Calling pickSubchannel() twice gave the same result
+    assertEquals(pickerCaptor.getValue().pickSubchannel(mockArgs),
+        pickerCaptor.getValue().pickSubchannel(mockArgs));
+
+    verifyNoMoreInteractions(mockHelper);
+  }
+
   @Test
   public void requestConnectionPicker() throws Exception {
     loadBalancer.acceptResolvedAddresses(