xds: pick_first LB configuration (#10181)

Configure pick_first load balancing if the xDS PickFirst LB config
message is used in the Cluster load balancing config.

repositories.bzl

@@ -136,10 +136,10 @@ def grpc_java_repositories():
     if not native.existing_rule("envoy_api"):
         http_archive(
             name = "envoy_api",
-            sha256 = "74156c0d8738d0469f23047f0fd0f8846fdd0d59d7b55c76cd8cb9ebf2fa3a01",
+            sha256 = "b426904abf51ba21dd8947a05694bb3c861d6f5e436e4673e74d7d7bfb6d3188",
-            strip_prefix = "data-plane-api-b1d2e441133c00bfe8412dfd6e93ea85e66da9bb",
+            strip_prefix = "data-plane-api-268824e4eee3d7770a347a5dc5aaddc0b1b14e24",
             urls = [
-                "https://github.com/envoyproxy/data-plane-api/archive/b1d2e441133c00bfe8412dfd6e93ea85e66da9bb.tar.gz",
+                "https://github.com/envoyproxy/data-plane-api/archive/268824e4eee3d7770a347a5dc5aaddc0b1b14e24.tar.gz",
             ],
         )
 

xds/BUILD.bazel

@@ -88,6 +88,7 @@ java_proto_library(
         "@envoy_api//envoy/extensions/filters/network/http_connection_manager/v3:pkg",
         "@envoy_api//envoy/extensions/load_balancing_policies/client_side_weighted_round_robin/v3:pkg",
         "@envoy_api//envoy/extensions/load_balancing_policies/least_request/v3:pkg",
+        "@envoy_api//envoy/extensions/load_balancing_policies/pick_first/v3:pkg",
         "@envoy_api//envoy/extensions/load_balancing_policies/ring_hash/v3:pkg",
         "@envoy_api//envoy/extensions/load_balancing_policies/round_robin/v3:pkg",
         "@envoy_api//envoy/extensions/load_balancing_policies/wrr_locality/v3:pkg",

xds/src/main/java/io/grpc/xds/LoadBalancerConfigFactory.java

@@ -31,6 +31,7 @@ import io.envoyproxy.envoy.config.cluster.v3.LoadBalancingPolicy;
 import io.envoyproxy.envoy.config.cluster.v3.LoadBalancingPolicy.Policy;
 import io.envoyproxy.envoy.extensions.load_balancing_policies.client_side_weighted_round_robin.v3.ClientSideWeightedRoundRobin;
 import io.envoyproxy.envoy.extensions.load_balancing_policies.least_request.v3.LeastRequest;
+import io.envoyproxy.envoy.extensions.load_balancing_policies.pick_first.v3.PickFirst;
 import io.envoyproxy.envoy.extensions.load_balancing_policies.ring_hash.v3.RingHash;
 import io.envoyproxy.envoy.extensions.load_balancing_policies.round_robin.v3.RoundRobin;
 import io.envoyproxy.envoy.extensions.load_balancing_policies.wrr_locality.v3.WrrLocality;
@@ -85,6 +86,9 @@ class LoadBalancerConfigFactory {
 
   static final String WEIGHT_UPDATE_PERIOD = "weightUpdatePeriod";
 
+  static final String PICK_FIRST_FIELD_NAME = "pick_first";
+  static final String SHUFFLE_ADDRESS_LIST_FIELD_NAME = "shuffleAddressList";
+
   /**
    * Factory method for creating a new {link LoadBalancerConfigConverter} for a given xDS {@link
    * Cluster}.
@@ -92,14 +96,14 @@ class LoadBalancerConfigFactory {
    * @throws ResourceInvalidException If the {@link Cluster} has an invalid LB configuration.
    */
   static ImmutableMap<String, ?> newConfig(Cluster cluster, boolean enableLeastRequest,
-      boolean enableWrr)
+      boolean enableWrr, boolean enablePickFirst)
       throws ResourceInvalidException {
     // The new load_balancing_policy will always be used if it is set, but for backward
     // compatibility we will fall back to using the old lb_policy field if the new field is not set.
     if (cluster.hasLoadBalancingPolicy()) {
       try {
         return LoadBalancingPolicyConverter.convertToServiceConfig(cluster.getLoadBalancingPolicy(),
-            0, enableWrr);
+            0, enableWrr, enablePickFirst);
       } catch (MaxRecursionReachedException e) {
         throw new ResourceInvalidException("Maximum LB config recursion depth reached", e);
       }
@@ -154,7 +158,7 @@ class LoadBalancerConfigFactory {
 
   /**
    * Builds a service config JSON object for the least_request load balancer config based on the
-   * given config values..
+   * given config values.
    */
   private static ImmutableMap<String, ?> buildLeastRequestConfig(Integer choiceCount) {
     ImmutableMap.Builder<String, Object> configBuilder = ImmutableMap.builder();
@@ -180,6 +184,16 @@ class LoadBalancerConfigFactory {
     return ImmutableMap.of(ROUND_ROBIN_FIELD_NAME, ImmutableMap.of());
   }
 
+  /**
+   * Builds a service config JSON object for the pick_first load balancer config based on the
+   * given config values.
+   */
+  private static ImmutableMap<String, ?> buildPickFirstConfig(boolean shuffleAddressList) {
+    ImmutableMap.Builder<String, Object> configBuilder = ImmutableMap.builder();
+    configBuilder.put(SHUFFLE_ADDRESS_LIST_FIELD_NAME, shuffleAddressList);
+    return ImmutableMap.of(PICK_FIRST_FIELD_NAME, configBuilder.buildOrThrow());
+  }
+
   /**
    * Responsible for converting from a {@code envoy.config.cluster.v3.LoadBalancingPolicy} proto
    * message to a gRPC service config format.
@@ -192,7 +206,8 @@ class LoadBalancerConfigFactory {
      * Converts a {@link LoadBalancingPolicy} object to a service config JSON object.
      */
     private static ImmutableMap<String, ?> convertToServiceConfig(
-        LoadBalancingPolicy loadBalancingPolicy, int recursionDepth, boolean enableWrr)
+        LoadBalancingPolicy loadBalancingPolicy, int recursionDepth, boolean enableWrr,
+        boolean enablePickFirst)
         throws ResourceInvalidException, MaxRecursionReachedException {
       if (recursionDepth > MAX_RECURSION) {
         throw new MaxRecursionReachedException();
@@ -206,7 +221,7 @@ class LoadBalancerConfigFactory {
             serviceConfig = convertRingHashConfig(typedConfig.unpack(RingHash.class));
           } else if (typedConfig.is(WrrLocality.class)) {
             serviceConfig = convertWrrLocalityConfig(typedConfig.unpack(WrrLocality.class),
-                recursionDepth, enableWrr);
+                recursionDepth, enableWrr, enablePickFirst);
           } else if (typedConfig.is(RoundRobin.class)) {
             serviceConfig = convertRoundRobinConfig();
           } else if (typedConfig.is(LeastRequest.class)) {
@@ -216,6 +231,10 @@ class LoadBalancerConfigFactory {
               serviceConfig = convertWeightedRoundRobinConfig(
                   typedConfig.unpack(ClientSideWeightedRoundRobin.class));
             }
+          } else if (typedConfig.is(PickFirst.class)) {
+            if (enablePickFirst) {
+              serviceConfig = convertPickFirstConfig(typedConfig.unpack(PickFirst.class));
+            }
           } else if (typedConfig.is(com.github.xds.type.v3.TypedStruct.class)) {
             serviceConfig = convertCustomConfig(
                 typedConfig.unpack(com.github.xds.type.v3.TypedStruct.class));
@@ -283,11 +302,12 @@ class LoadBalancerConfigFactory {
      * Converts a wrr_locality {@link Any} configuration to service config format.
      */
     private static ImmutableMap<String, ?> convertWrrLocalityConfig(WrrLocality wrrLocality,
-        int recursionDepth, boolean enableWrr) throws ResourceInvalidException,
+        int recursionDepth, boolean enableWrr, boolean enablePickFirst)
+        throws ResourceInvalidException,
         MaxRecursionReachedException {
       return buildWrrLocalityConfig(
           convertToServiceConfig(wrrLocality.getEndpointPickingPolicy(),
-              recursionDepth + 1, enableWrr));
+              recursionDepth + 1, enableWrr, enablePickFirst));
     }
 
     /**
@@ -297,6 +317,13 @@ class LoadBalancerConfigFactory {
       return buildRoundRobinConfig();
     }
 
+    /**
+     * "Converts" a pick_first configuration to service config format.
+     */
+    private static ImmutableMap<String, ?> convertPickFirstConfig(PickFirst pickFirst) {
+      return buildPickFirstConfig(pickFirst.getShuffleAddressList());
+    }
+
     /**
      * Converts a least_request {@link Any} configuration to service config format.
      */

xds/src/main/java/io/grpc/xds/XdsClusterResource.java

@@ -133,7 +133,7 @@ class XdsClusterResource extends XdsResourceType<CdsUpdate> {
     CdsUpdate.Builder updateBuilder = structOrError.getStruct();
 
     ImmutableMap<String, ?> lbPolicyConfig = LoadBalancerConfigFactory.newConfig(cluster,
-        enableLeastRequest, enableWrr);
+        enableLeastRequest, enableWrr, enablePickFirst);
 
     // Validate the LB config by trying to parse it with the corresponding LB provider.
     LbConfig lbConfig = ServiceConfigUtil.unwrapLoadBalancingConfig(lbPolicyConfig);

xds/src/main/java/io/grpc/xds/XdsResourceType.java

@@ -56,6 +56,10 @@ abstract class XdsResourceType<T extends ResourceUpdate> {
 
   @VisibleForTesting
   static boolean enableWrr = getFlag("GRPC_EXPERIMENTAL_XDS_WRR_LB", true);
+
+  @VisibleForTesting
+  static boolean enablePickFirst = getFlag("GRPC_EXPERIMENTAL_PICKFIRST_LB_CONFIG", false);
+
   static final String TYPE_URL_CLUSTER_CONFIG =
       "type.googleapis.com/envoy.extensions.clusters.aggregate.v3.ClusterConfig";
   static final String TYPE_URL_TYPED_STRUCT_UDPA =

xds/src/test/java/io/grpc/xds/LoadBalancerConfigFactoryTest.java

@@ -40,6 +40,7 @@ import io.envoyproxy.envoy.config.cluster.v3.LoadBalancingPolicy.Policy;
 import io.envoyproxy.envoy.config.core.v3.TypedExtensionConfig;
 import io.envoyproxy.envoy.extensions.load_balancing_policies.client_side_weighted_round_robin.v3.ClientSideWeightedRoundRobin;
 import io.envoyproxy.envoy.extensions.load_balancing_policies.least_request.v3.LeastRequest;
+import io.envoyproxy.envoy.extensions.load_balancing_policies.pick_first.v3.PickFirst;
 import io.envoyproxy.envoy.extensions.load_balancing_policies.ring_hash.v3.RingHash;
 import io.envoyproxy.envoy.extensions.load_balancing_policies.round_robin.v3.RoundRobin;
 import io.envoyproxy.envoy.extensions.load_balancing_policies.wrr_locality.v3.WrrLocality;
@@ -81,6 +82,11 @@ public class LoadBalancerConfigFactoryTest {
           LeastRequest.newBuilder().setChoiceCount(UInt32Value.of(LEAST_REQUEST_CHOICE_COUNT))
               .build()))).build();
 
+  private static final Policy PICK_FIRST_POLICY = Policy.newBuilder().setTypedExtensionConfig(
+      TypedExtensionConfig.newBuilder().setTypedConfig(Any.pack(
+          PickFirst.newBuilder().setShuffleAddressList(true)
+              .build()))).build();
+
   private static final Policy WRR_POLICY = Policy.newBuilder()
               .setTypedExtensionConfig(TypedExtensionConfig.newBuilder()
                   .setName("backend")
@@ -132,6 +138,9 @@ public class LoadBalancerConfigFactoryTest {
   private static final LbConfig VALID_LEAST_REQUEST_CONFIG = new LbConfig(
       "least_request_experimental",
       ImmutableMap.of("choiceCount", (double) LEAST_REQUEST_CHOICE_COUNT));
+  private static final LbConfig VALID_PICK_FIRST_CONFIG = new LbConfig(
+      "pick_first",
+      ImmutableMap.of("shuffleAddressList", true));
 
   @After
   public void deregisterCustomProvider() {
@@ -142,14 +151,14 @@ public class LoadBalancerConfigFactoryTest {
   public void roundRobin() throws ResourceInvalidException {
     Cluster cluster = newCluster(buildWrrPolicy(ROUND_ROBIN_POLICY));
 
-    assertThat(newLbConfig(cluster, true, true)).isEqualTo(VALID_ROUND_ROBIN_CONFIG);
+    assertThat(newLbConfig(cluster, true, true, true)).isEqualTo(VALID_ROUND_ROBIN_CONFIG);
   }
 
   @Test
   public void weightedRoundRobin() throws ResourceInvalidException {
     Cluster cluster = newCluster(buildWrrPolicy(WRR_POLICY));
 
-    assertThat(newLbConfig(cluster, true, true)).isEqualTo(VALID_WRR_CONFIG);
+    assertThat(newLbConfig(cluster, true, true, true)).isEqualTo(VALID_WRR_CONFIG);
   }
 
   @Test
@@ -166,7 +175,7 @@ public class LoadBalancerConfigFactoryTest {
             .build())
         .build()));
 
-    assertResourceInvalidExceptionThrown(cluster, true, true,
+    assertResourceInvalidExceptionThrown(cluster, true, true, true,
         "Invalid duration in weighted round robin config");
   }
 
@@ -174,14 +183,14 @@ public class LoadBalancerConfigFactoryTest {
   public void weightedRoundRobin_fallback_roundrobin() throws ResourceInvalidException {
     Cluster cluster = newCluster(buildWrrPolicy(WRR_POLICY, ROUND_ROBIN_POLICY));
 
-    assertThat(newLbConfig(cluster, true, false)).isEqualTo(VALID_ROUND_ROBIN_CONFIG);
+    assertThat(newLbConfig(cluster, true, false, true)).isEqualTo(VALID_ROUND_ROBIN_CONFIG);
   }
 
   @Test
   public void roundRobin_legacy() throws ResourceInvalidException {
     Cluster cluster = Cluster.newBuilder().setLbPolicy(LbPolicy.ROUND_ROBIN).build();
 
-    assertThat(newLbConfig(cluster, true, true)).isEqualTo(VALID_ROUND_ROBIN_CONFIG);
+    assertThat(newLbConfig(cluster, true, true, true)).isEqualTo(VALID_ROUND_ROBIN_CONFIG);
   }
 
   @Test
@@ -190,7 +199,7 @@ public class LoadBalancerConfigFactoryTest {
         .setLoadBalancingPolicy(LoadBalancingPolicy.newBuilder().addPolicies(RING_HASH_POLICY))
         .build();
 
-    assertThat(newLbConfig(cluster, true, true)).isEqualTo(VALID_RING_HASH_CONFIG);
+    assertThat(newLbConfig(cluster, true, true, true)).isEqualTo(VALID_RING_HASH_CONFIG);
   }
 
   @Test
@@ -200,7 +209,7 @@ public class LoadBalancerConfigFactoryTest {
             .setMaximumRingSize(UInt64Value.of(RING_HASH_MAX_RING_SIZE))
             .setHashFunction(HashFunction.XX_HASH)).build();
 
-    assertThat(newLbConfig(cluster, true, true)).isEqualTo(VALID_RING_HASH_CONFIG);
+    assertThat(newLbConfig(cluster, true, true, true)).isEqualTo(VALID_RING_HASH_CONFIG);
   }
 
   @Test
@@ -212,7 +221,7 @@ public class LoadBalancerConfigFactoryTest {
                     .setMaximumRingSize(UInt64Value.of(RING_HASH_MAX_RING_SIZE))
                     .setHashFunction(RingHash.HashFunction.MURMUR_HASH_2).build()))).build());
 
-    assertResourceInvalidExceptionThrown(cluster, true, true, "Invalid ring hash function");
+    assertResourceInvalidExceptionThrown(cluster, true, true, true, "Invalid ring hash function");
   }
 
   @Test
@@ -220,7 +229,7 @@ public class LoadBalancerConfigFactoryTest {
     Cluster cluster = Cluster.newBuilder().setLbPolicy(LbPolicy.RING_HASH).setRingHashLbConfig(
         RingHashLbConfig.newBuilder().setHashFunction(HashFunction.MURMUR_HASH_2)).build();
 
-    assertResourceInvalidExceptionThrown(cluster, true, true, "invalid ring hash function");
+    assertResourceInvalidExceptionThrown(cluster, true, true, true, "invalid ring hash function");
   }
 
   @Test
@@ -229,7 +238,7 @@ public class LoadBalancerConfigFactoryTest {
         .setLoadBalancingPolicy(LoadBalancingPolicy.newBuilder().addPolicies(LEAST_REQUEST_POLICY))
         .build();
 
-    assertThat(newLbConfig(cluster, true, true)).isEqualTo(VALID_LEAST_REQUEST_CONFIG);
+    assertThat(newLbConfig(cluster, true, true, true)).isEqualTo(VALID_LEAST_REQUEST_CONFIG);
   }
 
   @Test
@@ -241,7 +250,7 @@ public class LoadBalancerConfigFactoryTest {
             LeastRequestLbConfig.newBuilder()
                 .setChoiceCount(UInt32Value.of(LEAST_REQUEST_CHOICE_COUNT))).build();
 
-    LbConfig lbConfig = newLbConfig(cluster, true, true);
+    LbConfig lbConfig = newLbConfig(cluster, true, true, true);
     assertThat(lbConfig.getPolicyName()).isEqualTo("wrr_locality_experimental");
 
     List<LbConfig> childConfigs = ServiceConfigUtil.unwrapLoadBalancingConfigList(
@@ -254,11 +263,27 @@ public class LoadBalancerConfigFactoryTest {
 
   @Test
   public void leastRequest_notEnabled() {
-    System.setProperty("io.grpc.xds.experimentalEnableLeastRequest", "false");
-
     Cluster cluster = Cluster.newBuilder().setLbPolicy(LbPolicy.LEAST_REQUEST).build();
 
-    assertResourceInvalidExceptionThrown(cluster, false, true, "unsupported lb policy");
+    assertResourceInvalidExceptionThrown(cluster, false, true, true, "unsupported lb policy");
+  }
+
+  @Test
+  public void pickFirst() throws ResourceInvalidException {
+    Cluster cluster = Cluster.newBuilder()
+        .setLoadBalancingPolicy(LoadBalancingPolicy.newBuilder().addPolicies(PICK_FIRST_POLICY))
+        .build();
+
+    assertThat(newLbConfig(cluster, true, true, true)).isEqualTo(VALID_PICK_FIRST_CONFIG);
+  }
+
+  @Test
+  public void pickFirst_notEnabled() throws ResourceInvalidException {
+    Cluster cluster = Cluster.newBuilder()
+        .setLoadBalancingPolicy(LoadBalancingPolicy.newBuilder().addPolicies(PICK_FIRST_POLICY))
+        .build();
+
+    assertResourceInvalidExceptionThrown(cluster, true, true, false, "Invalid LoadBalancingPolicy");
   }
 
   @Test
@@ -266,7 +291,7 @@ public class LoadBalancerConfigFactoryTest {
     LoadBalancerRegistry.getDefaultRegistry().register(CUSTOM_POLICY_PROVIDER);
 
     assertThat(newLbConfig(newCluster(CUSTOM_POLICY), false,
-        true)).isEqualTo(VALID_CUSTOM_CONFIG);
+        true, true)).isEqualTo(VALID_CUSTOM_CONFIG);
   }
 
   @Test
@@ -275,7 +300,7 @@ public class LoadBalancerConfigFactoryTest {
         .setLoadBalancingPolicy(LoadBalancingPolicy.newBuilder().addPolicies(CUSTOM_POLICY))
         .build();
 
-    assertResourceInvalidExceptionThrown(cluster, false, true,"Invalid LoadBalancingPolicy");
+    assertResourceInvalidExceptionThrown(cluster, false, true, true, "Invalid LoadBalancingPolicy");
   }
 
   // When a provider for the endpoint picking custom policy is available, the configuration should
@@ -287,7 +312,7 @@ public class LoadBalancerConfigFactoryTest {
     Cluster cluster = Cluster.newBuilder().setLoadBalancingPolicy(LoadBalancingPolicy.newBuilder()
         .addPolicies(buildWrrPolicy(CUSTOM_POLICY, ROUND_ROBIN_POLICY))).build();
 
-    assertThat(newLbConfig(cluster, false, true)).isEqualTo(VALID_CUSTOM_CONFIG_IN_WRR);
+    assertThat(newLbConfig(cluster, false, true, true)).isEqualTo(VALID_CUSTOM_CONFIG_IN_WRR);
   }
 
   // When a provider for the endpoint picking custom policy is available, the configuration should
@@ -299,7 +324,7 @@ public class LoadBalancerConfigFactoryTest {
     Cluster cluster = Cluster.newBuilder().setLoadBalancingPolicy(LoadBalancingPolicy.newBuilder()
         .addPolicies(buildWrrPolicy(CUSTOM_POLICY_UDPA, ROUND_ROBIN_POLICY))).build();
 
-    assertThat(newLbConfig(cluster, false, true)).isEqualTo(VALID_CUSTOM_CONFIG_IN_WRR);
+    assertThat(newLbConfig(cluster, false, true, true)).isEqualTo(VALID_CUSTOM_CONFIG_IN_WRR);
   }
 
   // When a provider for the custom wrr_locality child policy is NOT available, we should fall back
@@ -309,7 +334,7 @@ public class LoadBalancerConfigFactoryTest {
     Cluster cluster = Cluster.newBuilder().setLoadBalancingPolicy(LoadBalancingPolicy.newBuilder()
         .addPolicies(buildWrrPolicy(CUSTOM_POLICY, ROUND_ROBIN_POLICY))).build();
 
-    assertThat(newLbConfig(cluster, false, true)).isEqualTo(VALID_ROUND_ROBIN_CONFIG);
+    assertThat(newLbConfig(cluster, false, true, true)).isEqualTo(VALID_ROUND_ROBIN_CONFIG);
   }
 
   // When a provider for the custom wrr_locality child policy is NOT available and no alternative
@@ -319,7 +344,7 @@ public class LoadBalancerConfigFactoryTest {
     Cluster cluster = Cluster.newBuilder().setLoadBalancingPolicy(
         LoadBalancingPolicy.newBuilder().addPolicies(buildWrrPolicy(CUSTOM_POLICY))).build();
 
-    assertResourceInvalidExceptionThrown(cluster, false, true, "Invalid LoadBalancingPolicy");
+    assertResourceInvalidExceptionThrown(cluster, false, true, true, "Invalid LoadBalancingPolicy");
   }
 
   @Test
@@ -346,7 +371,7 @@ public class LoadBalancerConfigFactoryTest {
                                               buildWrrPolicy(
                                                 ROUND_ROBIN_POLICY))))))))))))))))))).build();
 
-    assertResourceInvalidExceptionThrown(cluster, false, true,
+    assertResourceInvalidExceptionThrown(cluster, false, true, true,
         "Maximum LB config recursion depth reached");
   }
 
@@ -362,17 +387,18 @@ public class LoadBalancerConfigFactoryTest {
             .build()))).build();
   }
 
-  private LbConfig newLbConfig(Cluster cluster, boolean enableLeastRequest, boolean enableWrr)
+  private LbConfig newLbConfig(Cluster cluster, boolean enableLeastRequest, boolean enableWrr,
+      boolean enablePickFirst)
       throws ResourceInvalidException {
     return ServiceConfigUtil.unwrapLoadBalancingConfig(
         LoadBalancerConfigFactory.newConfig(cluster, enableLeastRequest,
-            enableWrr));
+            enableWrr, enablePickFirst));
   }
 
   private void assertResourceInvalidExceptionThrown(Cluster cluster, boolean enableLeastRequest,
-      boolean enableWrr, String expectedMessage) {
+      boolean enableWrr, boolean enablePickFirst, String expectedMessage) {
     try {
-      newLbConfig(cluster, enableLeastRequest, enableWrr);
+      newLbConfig(cluster, enableLeastRequest, enableWrr, enablePickFirst);
     } catch (ResourceInvalidException e) {
       assertThat(e).hasMessageThat().contains(expectedMessage);
       return;