xds: explicitly set request hash key for the ring hash LB policy

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

@@ -34,6 +34,7 @@
 import io.grpc.EquivalentAddressGroup;
 import io.grpc.InternalLogId;
 import io.grpc.LoadBalancer;
+import io.grpc.Metadata;
 import io.grpc.Status;
 import io.grpc.SynchronizationContext;
 import io.grpc.util.MultiChildLoadBalancer;
@@ -47,6 +48,7 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Random;
 import java.util.Set;
 import java.util.stream.Collectors;
 import javax.annotation.Nullable;
@@ -70,6 +72,7 @@
   private final XdsLogger logger;
   private final SynchronizationContext syncContext;
   private List<RingEntry> ring;
+  private String requestHashHeader = "";
 
   RingHashLoadBalancer(Helper helper) {
     super(helper);
@@ -99,6 +102,7 @@
       if (config == null) {
         throw new IllegalArgumentException("Missing RingHash configuration");
       }
+      requestHashHeader = config.requestHashHeader;
       Map<EquivalentAddressGroup, Long> serverWeights = new HashMap<>();
       long totalWeight = 0L;
       for (EquivalentAddressGroup eag : addrList) {
@@ -213,7 +217,8 @@
       overallState = TRANSIENT_FAILURE;
     }
 
-    RingHashPicker picker = new RingHashPicker(syncContext, ring, getChildLbStates());
+    RingHashPicker picker =
+        new RingHashPicker(syncContext, ring, getChildLbStates(), requestHashHeader);
     getHelper().updateBalancingState(overallState, picker);
     this.currentConnectivityState = overallState;
   }
@@ -334,23 +339,30 @@
   }
 
   private static final class RingHashPicker extends SubchannelPicker {
+    private final Random random = new Random();
     private final SynchronizationContext syncContext;
     private final List<RingEntry> ring;
     // Avoid synchronization between pickSubchannel and subchannel's connectivity state change,
     // freeze picker's view of subchannel's connectivity state.
     // TODO(chengyuanzhang): can be more performance-friendly with
     //  IdentityHashMap<Subchannel, ConnectivityStateInfo> and RingEntry contains Subchannel.
     private final Map<Endpoint, SubchannelView> pickableSubchannels;  // read-only
+    private final String requestHashHeader;
+    private boolean hasEndpointInConnectingState = false;
 
     private RingHashPicker(
         SynchronizationContext syncContext, List<RingEntry> ring,
-        Collection<ChildLbState> children) {
+        Collection<ChildLbState> children, String requestHashHeader) {
       this.syncContext = syncContext;
       this.ring = ring;
+      this.requestHashHeader = requestHashHeader;
       pickableSubchannels = new HashMap<>(children.size());
       for (ChildLbState childLbState : children) {
         pickableSubchannels.put((Endpoint)childLbState.getKey(),
             new SubchannelView(childLbState, childLbState.getCurrentState()));
+        if (childLbState.getCurrentState() == CONNECTING) {
+          hasEndpointInConnectingState = true;
+        }
       }
     }
 
@@ -381,38 +393,78 @@
 
     @Override
     public PickResult pickSubchannel(PickSubchannelArgs args) {
-      Long requestHash = args.getCallOptions().getOption(XdsNameResolver.RPC_HASH_KEY);
-      if (requestHash == null) {
-        return PickResult.withError(RPC_HASH_NOT_FOUND);
+      // Determine request hash.
+      boolean usingRandomHash = false;
+      Long requestHash;
+      if (requestHashHeader.isEmpty()) {
+        // Set by the xDS config selector.
+        requestHash = args.getCallOptions().getOption(XdsNameResolver.RPC_HASH_KEY);
+        if (requestHash == null) {
+          return PickResult.withError(RPC_HASH_NOT_FOUND);
+        }
+      } else {
+        String headerValue =
+            args.getHeaders()
+                .get(Metadata.Key.of(requestHashHeader, Metadata.ASCII_STRING_MARSHALLER));
+        if (headerValue != null) {
+          requestHash = hashFunc.hashAsciiString(headerValue);
+        } else {
+          requestHash = random.nextLong();
+          usingRandomHash = true;
+        }
       }
 
       int targetIndex = getTargetIndex(requestHash);
 
-      // Per gRFC A61, because of sticky-TF with PickFirst's auto reconnect on TF, we ignore
-      // all TF subchannels and find the first ring entry in READY, CONNECTING or IDLE.  If
-      // CONNECTING or IDLE we return a pick with no results.  Additionally, if that entry is in
-      // IDLE, we initiate a connection.
-      for (int i = 0; i < ring.size(); i++) {
-        int index = (targetIndex + i) % ring.size();
-        SubchannelView subchannelView = pickableSubchannels.get(ring.get(index).addrKey);
-        ChildLbState childLbState = subchannelView.childLbState;
-
-        if (subchannelView.connectivityState  == READY) {
-          return childLbState.getCurrentPicker().pickSubchannel(args);
+      if (!usingRandomHash) {
+        // Per gRFC A61, because of sticky-TF with PickFirst's auto reconnect on TF, we ignore
+        // all TF subchannels and find the first ring entry in READY, CONNECTING or IDLE.  If
+        // CONNECTING or IDLE we return a pick with no results.  Additionally, if that entry is in
+        // IDLE, we initiate a connection.
+        for (int i = 0; i < ring.size(); i++) {
+          int index = (targetIndex + i) % ring.size();
+          SubchannelView subchannelView = pickableSubchannels.get(ring.get(index).addrKey);
+          ChildLbState childLbState = subchannelView.childLbState;
+
+          if (subchannelView.connectivityState  == READY) {
+            return childLbState.getCurrentPicker().pickSubchannel(args);
+          }
+
+          // RPCs can be buffered if the next subchannel is pending (per A62). Otherwise, RPCs
+          // are failed unless there is a READY connection.
+          if (subchannelView.connectivityState == CONNECTING) {
+            return PickResult.withNoResult();
+          }
+
+          if (subchannelView.connectivityState == IDLE) {
+            syncContext.execute(() -> {
+              childLbState.getLb().requestConnection();
+            });
+
+            return PickResult.withNoResult(); // Indicates that this should be retried after backoff
+          }
         }
-
-        // RPCs can be buffered if the next subchannel is pending (per A62). Otherwise, RPCs
-        // are failed unless there is a READY connection.
-        if (subchannelView.connectivityState == CONNECTING) {
-          return PickResult.withNoResult();
+      } else {
+        // Using a random hash. Find and use the first READY ring entry, triggering at most one
+        // entry to attempt connection.
+        boolean requestedConnection = hasEndpointInConnectingState;
+        for (int i = 0; i < ring.size(); i++) {
+          int index = (targetIndex + i) % ring.size();
+          SubchannelView subchannelView = pickableSubchannels.get(ring.get(index).addrKey);
+          ChildLbState childLbState = subchannelView.childLbState;
+          if (subchannelView.connectivityState == READY) {
+            return childLbState.getCurrentPicker().pickSubchannel(args);
+          }
+          if (!requestedConnection && subchannelView.connectivityState == IDLE) {
+            syncContext.execute(
+                () -> {
+                  childLbState.getLb().requestConnection();
+                });
+            requestedConnection = true;
+          }
         }
-
-        if (subchannelView.connectivityState == IDLE) {
-          syncContext.execute(() -> {
-            childLbState.getLb().requestConnection();
-          });
-
-          return PickResult.withNoResult(); // Indicates that this should be retried after backoff
+        if (requestedConnection) {
+          return PickResult.withNoResult();
         }
       }
 
@@ -460,20 +512,24 @@
   static final class RingHashConfig {
     final long minRingSize;
     final long maxRingSize;
+    final String requestHashHeader;
 
-    RingHashConfig(long minRingSize, long maxRingSize) {
+    RingHashConfig(long minRingSize, long maxRingSize, String requestHashHeader) {
       checkArgument(minRingSize > 0, "minRingSize <= 0");
       checkArgument(maxRingSize > 0, "maxRingSize <= 0");
       checkArgument(minRingSize <= maxRingSize, "minRingSize > maxRingSize");
+      checkNotNull(requestHashHeader);
       this.minRingSize = minRingSize;
       this.maxRingSize = maxRingSize;
+      this.requestHashHeader = requestHashHeader;
     }
 
     @Override
     public String toString() {
       return MoreObjects.toStringHelper(this)
           .add("minRingSize", minRingSize)
           .add("maxRingSize", maxRingSize)
+          .add("requestHashHeader", requestHashHeader)
           .toString();
     }
   }

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

@@ -24,6 +24,7 @@
 import io.grpc.LoadBalancerProvider;
 import io.grpc.NameResolver.ConfigOrError;
 import io.grpc.Status;
+import io.grpc.internal.GrpcUtil;
 import io.grpc.internal.JsonUtil;
 import io.grpc.xds.RingHashLoadBalancer.RingHashConfig;
 import io.grpc.xds.RingHashOptions;
@@ -81,13 +82,20 @@ private ConfigOrError parseLoadBalancingPolicyConfigInternal(
       Map<String, ?> rawLoadBalancingPolicyConfig) {
     Long minRingSize = JsonUtil.getNumberAsLong(rawLoadBalancingPolicyConfig, "minRingSize");
     Long maxRingSize = JsonUtil.getNumberAsLong(rawLoadBalancingPolicyConfig, "maxRingSize");
+    String requestHashHeader = "";
+    if (GrpcUtil.getFlag("GRPC_EXPERIMENTAL_RING_HASH_SET_REQUEST_HASH_KEY", false)) {
+      requestHashHeader = JsonUtil.getString(rawLoadBalancingPolicyConfig, "requestHashHeader");
+    }
     long maxRingSizeCap = RingHashOptions.getRingSizeCap();
     if (minRingSize == null) {
       minRingSize = DEFAULT_MIN_RING_SIZE;
     }
     if (maxRingSize == null) {
       maxRingSize = DEFAULT_MAX_RING_SIZE;
     }
+    if (requestHashHeader == null) {
+      requestHashHeader = "";
+    }
     if (minRingSize > maxRingSizeCap) {
       minRingSize = maxRingSizeCap;
     }
@@ -98,6 +106,7 @@ private ConfigOrError parseLoadBalancingPolicyConfigInternal(
       return ConfigOrError.fromError(Status.UNAVAILABLE.withDescription(
           "Invalid 'mingRingSize'/'maxRingSize'"));
     }
-    return ConfigOrError.fromConfig(new RingHashConfig(minRingSize, maxRingSize));
+    return ConfigOrError.fromConfig(
+        new RingHashConfig(minRingSize, maxRingSize, requestHashHeader));
   }
 }

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

@@ -163,7 +163,7 @@ public void uncaughtException(Thread t, Throwable e) {
           GracefulSwitchLoadBalancer.createLoadBalancingPolicyConfig(
               new FakeLoadBalancerProvider("round_robin"), null)));
   private final Object ringHash = GracefulSwitchLoadBalancer.createLoadBalancingPolicyConfig(
-      new FakeLoadBalancerProvider("ring_hash_experimental"), new RingHashConfig(10L, 100L));
+      new FakeLoadBalancerProvider("ring_hash_experimental"), new RingHashConfig(10L, 100L, ""));
   private final Object leastRequest = GracefulSwitchLoadBalancer.createLoadBalancingPolicyConfig(
       new FakeLoadBalancerProvider("wrr_locality_experimental"), new WrrLocalityConfig(
           GracefulSwitchLoadBalancer.createLoadBalancingPolicyConfig(

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

@@ -42,6 +42,8 @@
 @RunWith(JUnit4.class)
 public class RingHashLoadBalancerProviderTest {
   private static final String AUTHORITY = "foo.googleapis.com";
+  private static final String GRPC_EXPERIMENTAL_RING_HASH_SET_REQUEST_HASH_KEY =
+      "GRPC_EXPERIMENTAL_RING_HASH_SET_REQUEST_HASH_KEY";
 
   private final SynchronizationContext syncContext = new SynchronizationContext(
       new UncaughtExceptionHandler() {
@@ -81,6 +83,7 @@ public void parseLoadBalancingConfig_valid() throws IOException {
     RingHashConfig config = (RingHashConfig) configOrError.getConfig();
     assertThat(config.minRingSize).isEqualTo(10L);
     assertThat(config.maxRingSize).isEqualTo(100L);
+    assertThat(config.requestHashHeader).isEmpty();
   }
 
   @Test
@@ -92,6 +95,7 @@ public void parseLoadBalancingConfig_missingRingSize_useDefaults() throws IOExce
     RingHashConfig config = (RingHashConfig) configOrError.getConfig();
     assertThat(config.minRingSize).isEqualTo(RingHashLoadBalancerProvider.DEFAULT_MIN_RING_SIZE);
     assertThat(config.maxRingSize).isEqualTo(RingHashLoadBalancerProvider.DEFAULT_MAX_RING_SIZE);
+    assertThat(config.requestHashHeader).isEmpty();
   }
 
   @Test
@@ -127,6 +131,7 @@ public void parseLoadBalancingConfig_ringTooLargeUsesCap() throws IOException {
     RingHashConfig config = (RingHashConfig) configOrError.getConfig();
     assertThat(config.minRingSize).isEqualTo(10);
     assertThat(config.maxRingSize).isEqualTo(RingHashOptions.DEFAULT_RING_SIZE_CAP);
+    assertThat(config.requestHashHeader).isEmpty();
   }
 
   @Test
@@ -142,6 +147,7 @@ public void parseLoadBalancingConfig_ringCapCanBeRaised() throws IOException {
     RingHashConfig config = (RingHashConfig) configOrError.getConfig();
     assertThat(config.minRingSize).isEqualTo(RingHashOptions.MAX_RING_SIZE_CAP);
     assertThat(config.maxRingSize).isEqualTo(RingHashOptions.MAX_RING_SIZE_CAP);
+    assertThat(config.requestHashHeader).isEmpty();
     // Reset to avoid affecting subsequent test cases
     RingHashOptions.setRingSizeCap(RingHashOptions.DEFAULT_RING_SIZE_CAP);
   }
@@ -159,6 +165,7 @@ public void parseLoadBalancingConfig_ringCapIsClampedTo8M() throws IOException {
     RingHashConfig config = (RingHashConfig) configOrError.getConfig();
     assertThat(config.minRingSize).isEqualTo(RingHashOptions.MAX_RING_SIZE_CAP);
     assertThat(config.maxRingSize).isEqualTo(RingHashOptions.MAX_RING_SIZE_CAP);
+    assertThat(config.requestHashHeader).isEmpty();
     // Reset to avoid affecting subsequent test cases
     RingHashOptions.setRingSizeCap(RingHashOptions.DEFAULT_RING_SIZE_CAP);
   }
@@ -176,6 +183,7 @@ public void parseLoadBalancingConfig_ringCapCanBeLowered() throws IOException {
     RingHashConfig config = (RingHashConfig) configOrError.getConfig();
     assertThat(config.minRingSize).isEqualTo(1);
     assertThat(config.maxRingSize).isEqualTo(1);
+    assertThat(config.requestHashHeader).isEmpty();
     // Reset to avoid affecting subsequent test cases
     RingHashOptions.setRingSizeCap(RingHashOptions.DEFAULT_RING_SIZE_CAP);
   }
@@ -193,6 +201,7 @@ public void parseLoadBalancingConfig_ringCapLowerLimitIs1() throws IOException {
     RingHashConfig config = (RingHashConfig) configOrError.getConfig();
     assertThat(config.minRingSize).isEqualTo(1);
     assertThat(config.maxRingSize).isEqualTo(1);
+    assertThat(config.requestHashHeader).isEmpty();
     // Reset to avoid affecting subsequent test cases
     RingHashOptions.setRingSizeCap(RingHashOptions.DEFAULT_RING_SIZE_CAP);
   }
@@ -219,6 +228,59 @@ public void parseLoadBalancingConfig_minRingSizeGreaterThanMaxRingSize() throws
         .isEqualTo("Invalid 'mingRingSize'/'maxRingSize'");
   }
 
+  @Test
+  public void parseLoadBalancingConfig_requestHashHeaderIgnoredWhenEnvVarNotSet()
+      throws IOException {
+    String lbConfig =
+        "{\"minRingSize\" : 10, \"maxRingSize\" : 100, \"requestHashHeader\" : \"dummy-hash\"}";
+    ConfigOrError configOrError =
+        provider.parseLoadBalancingPolicyConfig(parseJsonObject(lbConfig));
+    assertThat(configOrError.getConfig()).isNotNull();
+    RingHashConfig config = (RingHashConfig) configOrError.getConfig();
+    assertThat(config.minRingSize).isEqualTo(10L);
+    assertThat(config.maxRingSize).isEqualTo(100L);
+    assertThat(config.requestHashHeader).isEmpty();
+  }
+
+  @Test
+  public void parseLoadBalancingConfig_requestHashHeaderSetWhenEnvVarSet() throws IOException {
+    System.setProperty(GRPC_EXPERIMENTAL_RING_HASH_SET_REQUEST_HASH_KEY, "true");
+    try {
+      String lbConfig =
+          "{\"minRingSize\" : 10, \"maxRingSize\" : 100, \"requestHashHeader\" : \"dummy-hash\"}";
+      ConfigOrError configOrError =
+          provider.parseLoadBalancingPolicyConfig(parseJsonObject(lbConfig));
+      assertThat(configOrError.getConfig()).isNotNull();
+      RingHashConfig config = (RingHashConfig) configOrError.getConfig();
+      assertThat(config.minRingSize).isEqualTo(10L);
+      assertThat(config.maxRingSize).isEqualTo(100L);
+      assertThat(config.requestHashHeader).isEqualTo("dummy-hash");
+      assertThat(config.toString()).contains("minRingSize=10");
+      assertThat(config.toString()).contains("maxRingSize=100");
+      assertThat(config.toString()).contains("requestHashHeader=dummy-hash");
+    } finally {
+      System.clearProperty(GRPC_EXPERIMENTAL_RING_HASH_SET_REQUEST_HASH_KEY);
+    }
+  }
+
+  @Test
+  public void parseLoadBalancingConfig_requestHashHeaderUnsetWhenEnvVarSet_useDefaults()
+      throws IOException {
+    System.setProperty(GRPC_EXPERIMENTAL_RING_HASH_SET_REQUEST_HASH_KEY, "true");
+    try {
+      String lbConfig = "{\"minRingSize\" : 10, \"maxRingSize\" : 100}";
+      ConfigOrError configOrError =
+          provider.parseLoadBalancingPolicyConfig(parseJsonObject(lbConfig));
+      assertThat(configOrError.getConfig()).isNotNull();
+      RingHashConfig config = (RingHashConfig) configOrError.getConfig();
+      assertThat(config.minRingSize).isEqualTo(10L);
+      assertThat(config.maxRingSize).isEqualTo(100L);
+      assertThat(config.requestHashHeader).isEmpty();
+    } finally {
+      System.clearProperty(GRPC_EXPERIMENTAL_RING_HASH_SET_REQUEST_HASH_KEY);
+    }
+  }
+
   @SuppressWarnings("unchecked")
   private static Map<String, ?> parseJsonObject(String json) throws IOException {
     return (Map<String, ?>) JsonParser.parse(json);