xds: lazily and only parse headers with matchers matching the key (#8163)

In normal cases, we only have a few header matchers but the number of headers can be completely up to the application. Indexing headers eagerly parses all headers, even for those with no matcher matching the key. We should only parse header values for those with key matching the header matcher (aka, only call Metadata.get() with key that has some matcher looking for).

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

@@ -334,26 +334,12 @@ final class XdsNameResolver extends NameResolver {
   private final class ConfigSelector extends InternalConfigSelector {
     @Override
     public Result selectConfig(PickSubchannelArgs args) {
-      // Index ASCII headers by key, multi-value headers are concatenated for matching purposes.
-      Map<String, String> asciiHeaders = new HashMap<>();
-      Metadata headers = args.getHeaders();
-      for (String headerName : headers.keys()) {
-        if (headerName.endsWith(Metadata.BINARY_HEADER_SUFFIX)) {
-          continue;
-        }
-        Metadata.Key<String> key = Metadata.Key.of(headerName, Metadata.ASCII_STRING_MARSHALLER);
-        Iterable<String> values = headers.getAll(key);
-        if (values != null) {
-          asciiHeaders.put(headerName, Joiner.on(",").join(values));
-        }
-      }
-      // Special hack for exposing headers: "content-type".
-      asciiHeaders.put("content-type", "application/grpc");
       String cluster = null;
       Route selectedRoute = null;
       RoutingConfig routingCfg;
       Map<String, FilterConfig> selectedOverrideConfigs;
       List<ClientInterceptor> filterInterceptors = new ArrayList<>();
+      Metadata headers = args.getHeaders();
       do {
         routingCfg = routingConfig;
         selectedOverrideConfigs = new HashMap<>(routingCfg.virtualHostOverrideConfig);
@@ -363,7 +349,7 @@ final class XdsNameResolver extends NameResolver {
         }
         for (Route route : routingCfg.routes) {
           if (matchRoute(route.routeMatch(), "/" + args.getMethodDescriptor().getFullMethodName(),
-              asciiHeaders, random)) {
+              headers, random)) {
             selectedRoute = route;
             selectedOverrideConfigs.putAll(route.filterConfigOverrides());
             break;
@@ -442,7 +428,7 @@ final class XdsNameResolver extends NameResolver {
         }
       }
       final String finalCluster = cluster;
-      final long hash = generateHash(selectedRoute.routeAction().hashPolicies(), asciiHeaders);
+      final long hash = generateHash(selectedRoute.routeAction().hashPolicies(), headers);
       class ClusterSelectionInterceptor implements ClientInterceptor {
         @Override
         public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
@@ -517,13 +503,13 @@ final class XdsNameResolver extends NameResolver {
       }
     }
 
-    private long generateHash(List<HashPolicy> hashPolicies, Map<String, String> headers) {
+    private long generateHash(List<HashPolicy> hashPolicies, Metadata headers) {
       Long hash = null;
       for (HashPolicy policy : hashPolicies) {
         Long newHash = null;
         if (policy.type() == HashPolicy.Type.HEADER) {
-          if (headers.containsKey(policy.headerName())) {
+          String value = getHeaderValue(headers, policy.headerName());
-            String value = headers.get(policy.headerName());
+          if (value != null) {
             if (policy.regEx() != null && policy.regExSubstitution() != null) {
               value = policy.regEx().matcher(value).replaceAll(policy.regExSubstitution());
             }
@@ -565,19 +551,20 @@ final class XdsNameResolver extends NameResolver {
 
   @VisibleForTesting
   static boolean matchRoute(RouteMatch routeMatch, String fullMethodName,
-      Map<String, String> headers, ThreadSafeRandom random) {
+      Metadata headers, ThreadSafeRandom random) {
     if (!matchPath(routeMatch.pathMatcher(), fullMethodName)) {
       return false;
     }
-    if (!matchHeaders(routeMatch.headerMatchers(), headers)) {
+    for (HeaderMatcher headerMatcher : routeMatch.headerMatchers()) {
+      if (!matchHeader(headerMatcher, getHeaderValue(headers, headerMatcher.name()))) {
         return false;
       }
+    }
     FractionMatcher fraction = routeMatch.fractionMatcher();
     return fraction == null || random.nextInt(fraction.denominator()) < fraction.numerator();
   }
 
-  @VisibleForTesting
+  private static boolean matchPath(PathMatcher pathMatcher, String fullMethodName) {
-  static boolean matchPath(PathMatcher pathMatcher, String fullMethodName) {
     if (pathMatcher.path() != null) {
       return pathMatcher.caseSensitive()
           ? pathMatcher.path().equals(fullMethodName)
@@ -590,18 +577,7 @@ final class XdsNameResolver extends NameResolver {
     return pathMatcher.regEx().matches(fullMethodName);
   }
 
-  private static boolean matchHeaders(
+  private static boolean matchHeader(HeaderMatcher headerMatcher, @Nullable String value) {
-      List<HeaderMatcher> headerMatchers, Map<String, String> headers) {
-    for (HeaderMatcher headerMatcher : headerMatchers) {
-      if (!matchHeader(headerMatcher, headers.get(headerMatcher.name()))) {
-        return false;
-      }
-    }
-    return true;
-  }
-
-  @VisibleForTesting
-  static boolean matchHeader(HeaderMatcher headerMatcher, @Nullable String value) {
     if (headerMatcher.present() != null) {
       return (value == null) == headerMatcher.present().equals(headerMatcher.inverted());
     }
@@ -630,6 +606,24 @@ final class XdsNameResolver extends NameResolver {
     return baseMatch != headerMatcher.inverted();
   }
 
+  @Nullable
+  private static String getHeaderValue(Metadata headers, String headerName) {
+    if (headerName.endsWith(Metadata.BINARY_HEADER_SUFFIX)) {
+      return null;
+    }
+    if (headerName.equals("content-type")) {
+      return "application/grpc";
+    }
+    Metadata.Key<String> key;
+    try {
+      key = Metadata.Key.of(headerName, Metadata.ASCII_STRING_MARSHALLER);
+    } catch (IllegalArgumentException e) {
+      return null;
+    }
+    Iterable<String> values = headers.getAll(key);
+    return values == null ? null : Joiner.on(",").join(values);
+  }
+
   private class ResolveState implements LdsResourceWatcher {
     private final ConfigOrError emptyServiceConfig =
         serviceConfigParser.parseServiceConfig(Collections.<String, Object>emptyMap());

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

@@ -77,7 +77,6 @@ import io.grpc.xds.XdsNameResolverProvider.XdsClientPoolFactory;
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ScheduledExecutorService;
@@ -449,7 +448,38 @@ public class XdsNameResolverTest {
   }
 
   @Test
-  public void resolved_rpcHashingByHeader() {
+  public void resolved_rpcHashingByHeader_withoutSubstitution() {
+    resolver.start(mockListener);
+    FakeXdsClient xdsClient = (FakeXdsClient) resolver.getXdsClient();
+    xdsClient.deliverLdsUpdate(
+        Collections.singletonList(
+            Route.create(
+                RouteMatch.withPathExactOnly(
+                    "/" + TestMethodDescriptors.voidMethod().getFullMethodName()),
+                RouteAction.forCluster(cluster1, Collections.singletonList(HashPolicy.forHeader(
+                    false, "custom-key", null, null)),
+                    null),
+                ImmutableMap.<String, FilterConfig>of())));
+    verify(mockListener).onResult(resolutionResultCaptor.capture());
+    InternalConfigSelector configSelector =
+        resolutionResultCaptor.getValue().getAttributes().get(InternalConfigSelector.KEY);
+
+    // First call, with header "custom-key": "custom-value".
+    startNewCall(TestMethodDescriptors.voidMethod(), configSelector,
+        ImmutableMap.of("custom-key", "custom-value"), CallOptions.DEFAULT);
+    long hash1 = testCall.callOptions.getOption(XdsNameResolver.RPC_HASH_KEY);
+
+    // Second call, with header "custom-key": "custom-val".
+    startNewCall(TestMethodDescriptors.voidMethod(), configSelector,
+        ImmutableMap.of("custom-key", "custom-val"),
+        CallOptions.DEFAULT);
+    long hash2 = testCall.callOptions.getOption(XdsNameResolver.RPC_HASH_KEY);
+
+    assertThat(hash2).isNotEqualTo(hash1);
+  }
+
+  @Test
+  public void resolved_rpcHashingByHeader_withSubstitution() {
     resolver.start(mockListener);
     FakeXdsClient xdsClient = (FakeXdsClient) resolver.getXdsClient();
     xdsClient.deliverLdsUpdate(
@@ -1320,7 +1350,7 @@ public class XdsNameResolverTest {
 
   @Test
   public void routeMatching_pathOnly() {
-    Map<String, String> headers = Collections.emptyMap();
+    Metadata headers = new Metadata();
     ThreadSafeRandom random = mock(ThreadSafeRandom.class);
 
     RouteMatch routeMatch1 =
@@ -1351,14 +1381,36 @@ public class XdsNameResolverTest {
         .isTrue();
   }
 
+  @Test
+  public void routeMatching_pathOnly_caseInsensitive() {
+    Metadata headers = new Metadata();
+    ThreadSafeRandom random = mock(ThreadSafeRandom.class);
+
+    RouteMatch routeMatch1 =
+        RouteMatch.create(
+            PathMatcher.fromPath("/FooService/barMethod", false),
+            Collections.<HeaderMatcher>emptyList(), null);
+    assertThat(XdsNameResolver.matchRoute(routeMatch1, "/fooservice/barmethod", headers, random))
+        .isTrue();
+
+    RouteMatch routeMatch2 =
+        RouteMatch.create(
+            PathMatcher.fromPrefix("/FooService", false),
+            Collections.<HeaderMatcher>emptyList(), null);
+    assertThat(XdsNameResolver.matchRoute(routeMatch2, "/fooservice/barmethod", headers, random))
+        .isTrue();
+  }
+
   @Test
   public void routeMatching_withHeaders() {
-    Map<String, String> headers = new HashMap<>();
+    Metadata headers = new Metadata();
-    headers.put("authority", "foo.googleapis.com");
+    headers.put(Metadata.Key.of("authority", Metadata.ASCII_STRING_MARSHALLER),
-    headers.put("grpc-encoding", "gzip");
+        "foo.googleapis.com");
-    headers.put("user-agent", "gRPC-Java");
+    headers.put(Metadata.Key.of("grpc-encoding", Metadata.ASCII_STRING_MARSHALLER), "gzip");
-    headers.put("content-length", "1000");
+    headers.put(Metadata.Key.of("user-agent", Metadata.ASCII_STRING_MARSHALLER), "gRPC-Java");
-    headers.put("custom-key", "custom-value1,custom-value2");
+    headers.put(Metadata.Key.of("content-length", Metadata.ASCII_STRING_MARSHALLER), "1000");
+    headers.put(Metadata.Key.of("custom-key", Metadata.ASCII_STRING_MARSHALLER), "custom-value1");
+    headers.put(Metadata.Key.of("custom-key", Metadata.ASCII_STRING_MARSHALLER), "custom-value2");
     ThreadSafeRandom random = mock(ThreadSafeRandom.class);
 
     PathMatcher pathMatcher = PathMatcher.fromPath("/FooService/barMethod", true);
@@ -1419,15 +1471,22 @@ public class XdsNameResolverTest {
         null);
     assertThat(XdsNameResolver.matchRoute(routeMatch7, "/FooService/barMethod", headers, random))
         .isTrue();
-  }
 
-  @Test
+    RouteMatch routeMatch8 = RouteMatch.create(
-  public void pathMatching_caseInsensitive() {
+        pathMatcher,
-    PathMatcher pathMatcher1 = PathMatcher.fromPath("/FooService/barMethod", false);
+        Collections.singletonList(
-    assertThat(XdsNameResolver.matchPath(pathMatcher1, "/fooservice/barmethod")).isTrue();
+            HeaderMatcher.forExactValue("content-type", "application/grpc", false)),
+        null);
+    assertThat(XdsNameResolver.matchRoute(
+        routeMatch8, "/FooService/barMethod", new Metadata(), random)).isTrue();
 
-    PathMatcher pathMatcher2 = PathMatcher.fromPrefix("/FooService", false);
+    RouteMatch routeMatch9 = RouteMatch.create(
-    assertThat(XdsNameResolver.matchPath(pathMatcher2, "/fooservice/barmethod")).isTrue();
+        pathMatcher,
+        Collections.singletonList(
+            HeaderMatcher.forExactValue("custom-key!", "custom-value1,custom-value2", false)),
+        null);
+    assertThat(XdsNameResolver.matchRoute(routeMatch9, "/FooService/barMethod", headers, random))
+        .isFalse();
   }
 
   private final class FakeXdsClientPoolFactory implements XdsClientPoolFactory {