xds: Check for env variable before doing custom LB config (#9165)

grpc · May 12, 2022 · c6bfce0 · c6bfce0
1 parent 36d1d5f
commit c6bfce0
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 21 deletions.
diff --git a/xds/src/main/java/io/grpc/xds/ClientXdsClient.java b/xds/src/main/java/io/grpc/xds/ClientXdsClient.java
@@ -162,6 +162,12 @@ final class ClientXdsClient extends XdsClient implements XdsResponseHandler, Res
       !Strings.isNullOrEmpty(System.getenv("GRPC_EXPERIMENTAL_ENABLE_LEAST_REQUEST"))
           ? Boolean.parseBoolean(System.getenv("GRPC_EXPERIMENTAL_ENABLE_LEAST_REQUEST"))
           : Boolean.parseBoolean(System.getProperty("io.grpc.xds.experimentalEnableLeastRequest"));
+  @VisibleForTesting
+  static boolean enableCustomLbConfig =
+      !Strings.isNullOrEmpty(System.getenv("GRPC_EXPERIMENTAL_ENABLE_CUSTOM_LB_CONFIG"))
+          ? Boolean.parseBoolean(System.getenv("GRPC_EXPERIMENTAL_ENABLE_CUSTOM_LB_CONFIG"))
+          : Boolean.parseBoolean(
+              System.getProperty("io.grpc.xds.experimentalEnableCustomLbConfig"));
   private static final String TYPE_URL_HTTP_CONNECTION_MANAGER_V2 =
       "type.googleapis.com/envoy.config.filter.network.http_connection_manager.v2"
           + ".HttpConnectionManager";
@@ -1636,7 +1642,7 @@ static CdsUpdate processCluster(Cluster cluster, Set<String> retainedEdsResource
     CdsUpdate.Builder updateBuilder = structOrError.getStruct();
 
     ImmutableMap<String, ?> lbPolicyConfig = LoadBalancerConfigFactory.newConfig(cluster,
-        enableLeastRequest);
+        enableLeastRequest, enableCustomLbConfig);
 
     // Validate the LB config by trying to parse it with the corresponding LB provider.
     LbConfig lbConfig = ServiceConfigUtil.unwrapLoadBalancingConfig(lbPolicyConfig);

diff --git a/xds/src/main/java/io/grpc/xds/LoadBalancerConfigFactory.java b/xds/src/main/java/io/grpc/xds/LoadBalancerConfigFactory.java
@@ -78,11 +78,12 @@ class LoadBalancerConfigFactory {
    *
    * @throws ResourceInvalidException If the {@link Cluster} has an invalid LB configuration.
    */
-  static ImmutableMap<String, ?> newConfig(Cluster cluster, boolean enableLeastRequest)
+  static ImmutableMap<String, ?> newConfig(Cluster cluster, boolean enableLeastRequest,
+      boolean enableCustomLbConfig)
       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()) {
+    if (cluster.hasLoadBalancingPolicy() && enableCustomLbConfig) {
       try {
         return LoadBalancingPolicyConverter.convertToServiceConfig(cluster.getLoadBalancingPolicy(),
             0);

diff --git a/xds/src/test/java/io/grpc/xds/LoadBalancerConfigFactoryTest.java b/xds/src/test/java/io/grpc/xds/LoadBalancerConfigFactoryTest.java
@@ -106,14 +106,14 @@ public void deregisterCustomProvider() {
   public void roundRobin() throws ResourceInvalidException {
     Cluster cluster = newCluster(buildWrrPolicy(ROUND_ROBIN_POLICY));
 
-    assertThat(newLbConfig(cluster, true)).isEqualTo(VALID_ROUND_ROBIN_CONFIG);
+    assertThat(newLbConfig(cluster, true, 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)).isEqualTo(VALID_ROUND_ROBIN_CONFIG);
+    assertThat(newLbConfig(cluster, true, true)).isEqualTo(VALID_ROUND_ROBIN_CONFIG);
   }
 
   @Test
@@ -122,7 +122,7 @@ public void ringHash() throws ResourceInvalidException {
         .setLoadBalancingPolicy(LoadBalancingPolicy.newBuilder().addPolicies(RING_HASH_POLICY))
         .build();
 
-    assertThat(newLbConfig(cluster, true)).isEqualTo(VALID_RING_HASH_CONFIG);
+    assertThat(newLbConfig(cluster, true, true)).isEqualTo(VALID_RING_HASH_CONFIG);
   }
 
   @Test
@@ -132,7 +132,7 @@ public void ringHash_legacy() throws ResourceInvalidException {
             .setMaximumRingSize(UInt64Value.of(RING_HASH_MAX_RING_SIZE))
             .setHashFunction(HashFunction.XX_HASH)).build();
 
-    assertThat(newLbConfig(cluster, true)).isEqualTo(VALID_RING_HASH_CONFIG);
+    assertThat(newLbConfig(cluster, true, true)).isEqualTo(VALID_RING_HASH_CONFIG);
   }
 
   @Test
@@ -144,15 +144,15 @@ public void ringHash_invalidHash() {
                     .setMaximumRingSize(UInt64Value.of(RING_HASH_MAX_RING_SIZE))
                     .setHashFunction(RingHash.HashFunction.MURMUR_HASH_2).build()))).build());
 
-    assertResourceInvalidExceptionThrown(cluster, true, "Invalid ring hash function");
+    assertResourceInvalidExceptionThrown(cluster, true, true, "Invalid ring hash function");
   }
 
   @Test
   public void ringHash_invalidHash_legacy() {
     Cluster cluster = Cluster.newBuilder().setLbPolicy(LbPolicy.RING_HASH).setRingHashLbConfig(
         RingHashLbConfig.newBuilder().setHashFunction(HashFunction.MURMUR_HASH_2)).build();
 
-    assertResourceInvalidExceptionThrown(cluster, true, "invalid ring hash function");
+    assertResourceInvalidExceptionThrown(cluster, true, true, "invalid ring hash function");
   }
 
   @Test
@@ -163,7 +163,7 @@ public void leastRequest_legacy() throws ResourceInvalidException {
         .setLeastRequestLbConfig(
             LeastRequestLbConfig.newBuilder().setChoiceCount(UInt32Value.of(10))).build();
 
-    LbConfig lbConfig = newLbConfig(cluster, true);
+    LbConfig lbConfig = newLbConfig(cluster, true, true);
     assertThat(lbConfig.getPolicyName()).isEqualTo("wrr_locality_experimental");
 
     List<LbConfig> childConfigs = ServiceConfigUtil.unwrapLoadBalancingConfigList(
@@ -180,14 +180,14 @@ public void leastRequest_notEnabled() {
 
     Cluster cluster = Cluster.newBuilder().setLbPolicy(LbPolicy.LEAST_REQUEST).build();
 
-    assertResourceInvalidExceptionThrown(cluster, false, "unsupported lb policy");
+    assertResourceInvalidExceptionThrown(cluster, false, true, "unsupported lb policy");
   }
 
   @Test
   public void customRootLb_providerRegistered() throws ResourceInvalidException {
     LoadBalancerRegistry.getDefaultRegistry().register(CUSTOM_POLICY_PROVIDER);
 
-    assertThat(newLbConfig(newCluster(CUSTOM_POLICY), false)).isEqualTo(VALID_CUSTOM_CONFIG);
+    assertThat(newLbConfig(newCluster(CUSTOM_POLICY), false, true)).isEqualTo(VALID_CUSTOM_CONFIG);
   }
 
   @Test
@@ -196,7 +196,7 @@ public void customRootLb_providerNotRegistered() throws ResourceInvalidException
         .setLoadBalancingPolicy(LoadBalancingPolicy.newBuilder().addPolicies(CUSTOM_POLICY))
         .build();
 
-    assertResourceInvalidExceptionThrown(cluster, false, "Invalid LoadBalancingPolicy");
+    assertResourceInvalidExceptionThrown(cluster, false, true, "Invalid LoadBalancingPolicy");
   }
 
   // When a provider for the endpoint picking custom policy is available, the configuration should
@@ -208,7 +208,7 @@ public void customLbInWrr_providerRegistered() throws ResourceInvalidException {
     Cluster cluster = Cluster.newBuilder().setLoadBalancingPolicy(LoadBalancingPolicy.newBuilder()
         .addPolicies(buildWrrPolicy(CUSTOM_POLICY, ROUND_ROBIN_POLICY))).build();
 
-    assertThat(newLbConfig(cluster, false)).isEqualTo(VALID_CUSTOM_CONFIG_IN_WRR);
+    assertThat(newLbConfig(cluster, false, true)).isEqualTo(VALID_CUSTOM_CONFIG_IN_WRR);
   }
 
   // When a provider for the custom wrr_locality child policy is NOT available, we should fall back
@@ -218,7 +218,7 @@ public void customLbInWrr_providerNotRegistered() throws ResourceInvalidExceptio
     Cluster cluster = Cluster.newBuilder().setLoadBalancingPolicy(LoadBalancingPolicy.newBuilder()
         .addPolicies(buildWrrPolicy(CUSTOM_POLICY, ROUND_ROBIN_POLICY))).build();
 
-    assertThat(newLbConfig(cluster, false)).isEqualTo(VALID_ROUND_ROBIN_CONFIG);
+    assertThat(newLbConfig(cluster, false, true)).isEqualTo(VALID_ROUND_ROBIN_CONFIG);
   }
 
   // When a provider for the custom wrr_locality child policy is NOT available and no alternative
@@ -228,7 +228,18 @@ public void customLbInWrr_providerNotRegistered_noFallback() throws ResourceInva
     Cluster cluster = Cluster.newBuilder().setLoadBalancingPolicy(
         LoadBalancingPolicy.newBuilder().addPolicies(buildWrrPolicy(CUSTOM_POLICY))).build();
 
-    assertResourceInvalidExceptionThrown(cluster, false, "Invalid LoadBalancingPolicy");
+    assertResourceInvalidExceptionThrown(cluster, false, true, "Invalid LoadBalancingPolicy");
+  }
+
+  @Test
+  public void customConfig_notEnabled() throws ResourceInvalidException {
+    Cluster cluster = Cluster.newBuilder()
+        .setLoadBalancingPolicy(
+            LoadBalancingPolicy.newBuilder().addPolicies(RING_HASH_POLICY))
+        .build();
+
+    // Custom LB flag not set, so we use old logic that will default to round_robin.
+    assertThat(newLbConfig(cluster, true, false)).isEqualTo(VALID_ROUND_ROBIN_CONFIG);
   }
 
   @Test
@@ -255,7 +266,7 @@ public void maxRecursion() {
                                               buildWrrPolicy(
                                                 ROUND_ROBIN_POLICY))))))))))))))))))).build();
 
-    assertResourceInvalidExceptionThrown(cluster, false,
+    assertResourceInvalidExceptionThrown(cluster, false, true,
         "Maximum LB config recursion depth reached");
   }
 
@@ -271,16 +282,17 @@ private static Policy buildWrrPolicy(Policy... childPolicy) {
             .build()))).build();
   }
 
-  private LbConfig newLbConfig(Cluster cluster, boolean enableLeastRequest)
+  private LbConfig newLbConfig(Cluster cluster, boolean enableLeastRequest,
+      boolean enableCustomConfig)
       throws ResourceInvalidException {
     return ServiceConfigUtil.unwrapLoadBalancingConfig(
-        LoadBalancerConfigFactory.newConfig(cluster, enableLeastRequest));
+        LoadBalancerConfigFactory.newConfig(cluster, enableLeastRequest, enableCustomConfig));
   }
 
   private void assertResourceInvalidExceptionThrown(Cluster cluster, boolean enableLeastRequest,
-      String expectedMessage) {
+      boolean enableCustomConfig, String expectedMessage) {
     try {
-      newLbConfig(cluster, enableLeastRequest);
+      newLbConfig(cluster, enableLeastRequest, enableCustomConfig);
     } catch (ResourceInvalidException e) {
       assertThat(e).hasMessageThat().contains(expectedMessage);
       return;