Limit ring hash max size to 4K

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

@@ -26,6 +26,7 @@
 import io.grpc.Status;
 import io.grpc.internal.JsonUtil;
 import io.grpc.xds.RingHashLoadBalancer.RingHashConfig;
+import io.grpc.xds.RingHashOptions;
 import java.util.Map;
 
 /**
@@ -39,11 +40,7 @@ public final class RingHashLoadBalancerProvider extends LoadBalancerProvider {
   static final long DEFAULT_MIN_RING_SIZE = 1024L;
   // Same as ClientXdsClient.DEFAULT_RING_HASH_LB_POLICY_MAX_RING_SIZE
   @VisibleForTesting
-  static final long DEFAULT_MAX_RING_SIZE = 8 * 1024 * 1024L;
-  // Maximum number of ring entries allowed. Setting this too large can result in slow
-  // ring construction and OOM error.
-  // Same as ClientXdsClient.MAX_RING_HASH_LB_POLICY_RING_SIZE
-  static final long MAX_RING_SIZE = 8 * 1024 * 1024L;
+  static final long DEFAULT_MAX_RING_SIZE = 4 * 1024L;
 
   private static final boolean enableRingHash =
       Strings.isNullOrEmpty(System.getenv("GRPC_XDS_EXPERIMENTAL_ENABLE_RING_HASH"))
@@ -73,14 +70,20 @@ public String getPolicyName() {
   public ConfigOrError parseLoadBalancingPolicyConfig(Map<String, ?> rawLoadBalancingPolicyConfig) {
     Long minRingSize = JsonUtil.getNumberAsLong(rawLoadBalancingPolicyConfig, "minRingSize");
     Long maxRingSize = JsonUtil.getNumberAsLong(rawLoadBalancingPolicyConfig, "maxRingSize");
+    long maxRingSizeCap = RingHashOptions.getRingSizeCap();
     if (minRingSize == null) {
       minRingSize = DEFAULT_MIN_RING_SIZE;
     }
     if (maxRingSize == null) {
       maxRingSize = DEFAULT_MAX_RING_SIZE;
     }
-    if (minRingSize <= 0 || maxRingSize <= 0 || minRingSize > maxRingSize
-        || maxRingSize > MAX_RING_SIZE) {
+    if (minRingSize > maxRingSizeCap) {
+      minRingSize = maxRingSizeCap;
+    }
+    if (maxRingSize > maxRingSizeCap) {
+      maxRingSize = maxRingSizeCap;
+    }
+    if (minRingSize <= 0 || maxRingSize <= 0 || minRingSize > maxRingSize) {
       return ConfigOrError.fromError(Status.UNAVAILABLE.withDescription(
           "Invalid 'mingRingSize'/'maxRingSize'"));
     }

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

@@ -0,0 +1,62 @@
+/*
+ * Copyright 2021 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.xds;
+
+import com.google.common.annotations.VisibleForTesting;
+import io.grpc.ExperimentalApi;
+
+/**
+ * Utility class that provides a way to configure ring hash size limits. This is applicable
+ * for clients that use the ring hash load balancing policy. Note that size limits involve
+ * a tradeoff between client memory consumption and accuracy of load balancing weight
+ * representations. Also see https://github.com/grpc/proposal/pull/338.
+ */
+@ExperimentalApi("https://github.com/grpc/grpc-java/issues/9718")
+public final class RingHashOptions {
+  // Same as ClientXdsClient.DEFAULT_RING_HASH_LB_POLICY_MAX_RING_SIZE
+  @VisibleForTesting
+  static final long MAX_RING_SIZE_CAP = 8 * 1024 * 1024L;
+  @VisibleForTesting
+  // Same as RingHashLoadBalancerProvider.DEFAULT_MAX_RING_SIZE
+  static final long DEFAULT_RING_SIZE_CAP = 4 * 1024L;
+
+  // Limits ring hash sizes to restrict client memory usage.
+  private static volatile long ringSizeCap = DEFAULT_RING_SIZE_CAP;
+
+  private RingHashOptions() {} // Prevent instantiation
+
+  /**
+   * Set the global limit for the min and max number of ring hash entries per ring.
+   * Note that this limit is clamped between 1 entry and 8,388,608 entries, and new
+   * limits lying outside that range will be silently moved to the nearest number within
+   * that range. Defaults initially to 4096 entries.
+   */
+  @ExperimentalApi("https://github.com/grpc/grpc-java/issues/9718")
+  public static void setRingSizeCap(long ringSizeCap) {
+    ringSizeCap = Math.max(1, ringSizeCap);
+    ringSizeCap = Math.min(MAX_RING_SIZE_CAP, ringSizeCap);
+    RingHashOptions.ringSizeCap = ringSizeCap;
+  }
+
+  /**
+   * Get the global limit for min and max ring hash sizes.
+   */
+  @ExperimentalApi("https://github.com/grpc/grpc-java/issues/9718")
+  public static long getRingSizeCap() {
+    return RingHashOptions.ringSizeCap;
+  }
+}

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

@@ -29,6 +29,7 @@
 import io.grpc.SynchronizationContext;
 import io.grpc.internal.JsonParser;
 import io.grpc.xds.RingHashLoadBalancer.RingHashConfig;
+import io.grpc.xds.RingHashOptions;
 import java.io.IOException;
 import java.lang.Thread.UncaughtExceptionHandler;
 import java.util.Locale;
@@ -116,16 +117,84 @@ public void parseLoadBalancingConfig_invalid_minGreaterThanMax() throws IOExcept
   }
 
   @Test
-  public void parseLoadBalancingConfig_invalid_ringTooLarge() throws IOException {
-    long ringSize = RingHashLoadBalancerProvider.MAX_RING_SIZE + 1;
+  public void parseLoadBalancingConfig_ringTooLargeUsesCap() throws IOException {
+    long ringSize = RingHashOptions.MAX_RING_SIZE_CAP + 1;
     String lbConfig =
         String.format(Locale.US, "{\"minRingSize\" : 10, \"maxRingSize\" : %d}", ringSize);
     ConfigOrError configOrError =
         provider.parseLoadBalancingPolicyConfig(parseJsonObject(lbConfig));
-    assertThat(configOrError.getError()).isNotNull();
-    assertThat(configOrError.getError().getCode()).isEqualTo(Code.UNAVAILABLE);
-    assertThat(configOrError.getError().getDescription())
-        .isEqualTo("Invalid 'mingRingSize'/'maxRingSize'");
+    assertThat(configOrError.getConfig()).isNotNull();
+    RingHashConfig config = (RingHashConfig) configOrError.getConfig();
+    assertThat(config.minRingSize).isEqualTo(10);
+    assertThat(config.maxRingSize).isEqualTo(RingHashOptions.DEFAULT_RING_SIZE_CAP);
+  }
+
+  @Test
+  public void parseLoadBalancingConfig_ringCapCanBeRaised() throws IOException {
+    RingHashOptions.setRingSizeCap(RingHashOptions.MAX_RING_SIZE_CAP);
+    long ringSize = RingHashOptions.MAX_RING_SIZE_CAP;
+    String lbConfig =
+        String.format(
+            Locale.US, "{\"minRingSize\" : %d, \"maxRingSize\" : %d}", ringSize, ringSize);
+    ConfigOrError configOrError =
+        provider.parseLoadBalancingPolicyConfig(parseJsonObject(lbConfig));
+    assertThat(configOrError.getConfig()).isNotNull();
+    RingHashConfig config = (RingHashConfig) configOrError.getConfig();
+    assertThat(config.minRingSize).isEqualTo(RingHashOptions.MAX_RING_SIZE_CAP);
+    assertThat(config.maxRingSize).isEqualTo(RingHashOptions.MAX_RING_SIZE_CAP);
+    // Reset to avoid affecting subsequent test cases
+    RingHashOptions.setRingSizeCap(RingHashOptions.DEFAULT_RING_SIZE_CAP);
+  }
+
+  @Test
+  public void parseLoadBalancingConfig_ringCapIsClampedTo8M() throws IOException {
+    RingHashOptions.setRingSizeCap(RingHashOptions.MAX_RING_SIZE_CAP + 1);
+    long ringSize = RingHashOptions.MAX_RING_SIZE_CAP + 1;
+    String lbConfig =
+        String.format(
+            Locale.US, "{\"minRingSize\" : %d, \"maxRingSize\" : %d}", ringSize, ringSize);
+    ConfigOrError configOrError =
+        provider.parseLoadBalancingPolicyConfig(parseJsonObject(lbConfig));
+    assertThat(configOrError.getConfig()).isNotNull();
+    RingHashConfig config = (RingHashConfig) configOrError.getConfig();
+    assertThat(config.minRingSize).isEqualTo(RingHashOptions.MAX_RING_SIZE_CAP);
+    assertThat(config.maxRingSize).isEqualTo(RingHashOptions.MAX_RING_SIZE_CAP);
+    // Reset to avoid affecting subsequent test cases
+    RingHashOptions.setRingSizeCap(RingHashOptions.DEFAULT_RING_SIZE_CAP);
+  }
+
+  @Test
+  public void parseLoadBalancingConfig_ringCapCanBeLowered() throws IOException {
+    RingHashOptions.setRingSizeCap(1);
+    long ringSize = 2;
+    String lbConfig =
+        String.format(
+            Locale.US, "{\"minRingSize\" : %d, \"maxRingSize\" : %d}", ringSize, ringSize);
+    ConfigOrError configOrError =
+        provider.parseLoadBalancingPolicyConfig(parseJsonObject(lbConfig));
+    assertThat(configOrError.getConfig()).isNotNull();
+    RingHashConfig config = (RingHashConfig) configOrError.getConfig();
+    assertThat(config.minRingSize).isEqualTo(1);
+    assertThat(config.maxRingSize).isEqualTo(1);
+    // Reset to avoid affecting subsequent test cases
+    RingHashOptions.setRingSizeCap(RingHashOptions.DEFAULT_RING_SIZE_CAP);
+  }
+
+  @Test
+  public void parseLoadBalancingConfig_ringCapLowerLimitIs1() throws IOException {
+    RingHashOptions.setRingSizeCap(0);
+    long ringSize = 2;
+    String lbConfig =
+        String.format(
+            Locale.US, "{\"minRingSize\" : %d, \"maxRingSize\" : %d}", ringSize, ringSize);
+    ConfigOrError configOrError =
+        provider.parseLoadBalancingPolicyConfig(parseJsonObject(lbConfig));
+    assertThat(configOrError.getConfig()).isNotNull();
+    RingHashConfig config = (RingHashConfig) configOrError.getConfig();
+    assertThat(config.minRingSize).isEqualTo(1);
+    assertThat(config.maxRingSize).isEqualTo(1);
+    // Reset to avoid affecting subsequent test cases
+    RingHashOptions.setRingSizeCap(RingHashOptions.DEFAULT_RING_SIZE_CAP);
   }
 
   @Test