rls: rls request factory (aka key builder map)

rls/src/main/java/io/grpc/rls/internal/RlsRequestFactory.java

@@ -0,0 +1,117 @@
+/*
+ * Copyright 2020 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.rls.internal;
+
+import static com.google.common.base.Preconditions.checkNotNull;
+
+import com.google.common.base.MoreObjects;
+import com.google.common.collect.HashBasedTable;
+import com.google.common.collect.Table;
+import io.grpc.Metadata;
+import io.grpc.Status;
+import io.grpc.StatusRuntimeException;
+import io.grpc.rls.internal.RlsProtoData.GrpcKeyBuilder;
+import io.grpc.rls.internal.RlsProtoData.GrpcKeyBuilder.Name;
+import io.grpc.rls.internal.RlsProtoData.NameMatcher;
+import io.grpc.rls.internal.RlsProtoData.RouteLookupConfig;
+import io.grpc.rls.internal.RlsProtoData.RouteLookupRequest;
+import java.util.HashMap;
+import java.util.Map;
+import javax.annotation.CheckReturnValue;
+
+/**
+ * A RlsRequestFactory creates {@link RouteLookupRequest} using key builder map from {@link
+ * RouteLookupConfig}.
+ */
+final class RlsRequestFactory {
+
+  private final String target;
+  /**
+   * schema: Path(serviceName/methodName or serviceName/*), rls request headerName, header fields.
+   */
+  private final Table<String, String, NameMatcher> keyBuilderTable;
+
+  /** Constructor. */
+  public RlsRequestFactory(RouteLookupConfig rlsConfig) {
+    checkNotNull(rlsConfig, "rlsConfig");
+    this.target = rlsConfig.getLookupService();
+    this.keyBuilderTable = createKeyBuilderTable(rlsConfig);
+  }
+
+  private static Table<String, String, NameMatcher> createKeyBuilderTable(
+      RouteLookupConfig config) {
+    Table<String, String, NameMatcher> table = HashBasedTable.create();
+    for (GrpcKeyBuilder grpcKeyBuilder : config.getGrpcKeyBuilders()) {
+      for (NameMatcher nameMatcher : grpcKeyBuilder.getHeaders()) {
+        for (Name name : grpcKeyBuilder.getNames()) {
+          String method =
+              name.getMethod() == null || name.getMethod().isEmpty()
+                  ? "*" : name.getMethod();
+          String path = name.getService() + "/" + method;
+          table.put(path, nameMatcher.getKey(), nameMatcher);
+        }
+      }
+    }
+    return table;
+  }
+
+  /** Creates a {@link RouteLookupRequest} for given request's metadata. */
+  @CheckReturnValue
+  public RouteLookupRequest create(String service, String method, Metadata metadata) {
+    checkNotNull(service, "service");
+    checkNotNull(method, "method");
+    String path = service + "/" + method;
+    Map<String, NameMatcher> keyBuilder = keyBuilderTable.row(path);
+    // if no matching keyBuilder found, fall back to wildcard match (ServiceName/*)
+    if (keyBuilder.isEmpty()) {
+      keyBuilder = keyBuilderTable.row(service + "/*");
+    }
+    Map<String, String> rlsRequestHeaders = createRequestHeaders(metadata, keyBuilder);
+    return new RouteLookupRequest(target, path, "grpc", rlsRequestHeaders);
+  }
+
+  private Map<String, String> createRequestHeaders(
+      Metadata metadata, Map<String, NameMatcher> keyBuilder) {
+    Map<String, String> rlsRequestHeaders = new HashMap<>();
+    for (Map.Entry<String, NameMatcher> entry : keyBuilder.entrySet()) {
+      NameMatcher nameMatcher = entry.getValue();
+      String value = null;
+      for (String requestHeaderName : nameMatcher.names()) {
+        value = metadata.get(Metadata.Key.of(requestHeaderName, Metadata.ASCII_STRING_MARSHALLER));
+        if (value != null) {
+          break;
+        }
+      }
+      if (value != null) {
+        rlsRequestHeaders.put(entry.getKey(), value);
+      } else if (!nameMatcher.isOptional()) {
+        throw new StatusRuntimeException(
+            Status.INVALID_ARGUMENT.withDescription(
+                String.format("Missing mandatory metadata(%s) not found", entry.getKey())));
+      }
+    }
+    return rlsRequestHeaders;
+  }
+
+  @Override
+  public String toString() {
+    return MoreObjects.toStringHelper(this)
+        .add("lookupService", target)
+        .add("keyBuilderTable", keyBuilderTable)
+        .toString();
+  }
+}

rls/src/test/java/io/grpc/rls/internal/RlsRequestFactoryTest.java

@@ -0,0 +1,160 @@
+/*
+ * Copyright 2020 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.rls.internal;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.fail;
+
+import com.google.common.collect.ImmutableList;
+import io.grpc.Metadata;
+import io.grpc.Status.Code;
+import io.grpc.StatusRuntimeException;
+import io.grpc.rls.internal.RlsProtoData.GrpcKeyBuilder;
+import io.grpc.rls.internal.RlsProtoData.GrpcKeyBuilder.Name;
+import io.grpc.rls.internal.RlsProtoData.NameMatcher;
+import io.grpc.rls.internal.RlsProtoData.RequestProcessingStrategy;
+import io.grpc.rls.internal.RlsProtoData.RouteLookupConfig;
+import io.grpc.rls.internal.RlsProtoData.RouteLookupRequest;
+import java.util.concurrent.TimeUnit;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class RlsRequestFactoryTest {
+
+  private static final RouteLookupConfig RLS_CONFIG =
+      new RouteLookupConfig(
+          ImmutableList.of(
+              new GrpcKeyBuilder(
+                  ImmutableList.of(new Name("com.google.service1", "Create")),
+                  ImmutableList.of(
+                      new NameMatcher("user", ImmutableList.of("User", "Parent"), true),
+                      new NameMatcher("id", ImmutableList.of("X-Google-Id"), true))),
+              new GrpcKeyBuilder(
+                  ImmutableList.of(new Name("com.google.service1")),
+                  ImmutableList.of(
+                      new NameMatcher("user", ImmutableList.of("User", "Parent"), true),
+                      new NameMatcher("password", ImmutableList.of("Password"), true))),
+              new GrpcKeyBuilder(
+                  ImmutableList.of(new Name("com.google.service2")),
+                  ImmutableList.of(
+                      new NameMatcher("user", ImmutableList.of("User", "Parent"), false),
+                      new NameMatcher("password", ImmutableList.of("Password"), true))),
+              new GrpcKeyBuilder(
+                  ImmutableList.of(new Name("com.google.service3")),
+                  ImmutableList.of(
+                      new NameMatcher("user", ImmutableList.of("User", "Parent"), true)))),
+          /* lookupService= */ "foo.google.com",
+          /* lookupServiceTimeoutInMillis= */ TimeUnit.SECONDS.toMillis(2),
+          /* maxAgeInMillis= */ TimeUnit.SECONDS.toMillis(300),
+          /* staleAgeInMillis= */ TimeUnit.SECONDS.toMillis(240),
+          /* cacheSize= */ 1000,
+          /* validTargets= */ ImmutableList.of("a valid target"),
+          /* defaultTarget= */ "us_east_1.cloudbigtable.googleapis.com",
+          RequestProcessingStrategy.ASYNC_LOOKUP_DEFAULT_TARGET_ON_MISS);
+
+  private final RlsRequestFactory factory = new RlsRequestFactory(RLS_CONFIG);
+
+  @Test
+  public void create_pathMatches() {
+    Metadata metadata = new Metadata();
+    metadata.put(Metadata.Key.of("User", Metadata.ASCII_STRING_MARSHALLER), "test");
+    metadata.put(Metadata.Key.of("X-Google-Id", Metadata.ASCII_STRING_MARSHALLER), "123");
+    metadata.put(Metadata.Key.of("foo", Metadata.ASCII_STRING_MARSHALLER), "bar");
+
+    RouteLookupRequest request = factory.create("com.google.service1", "Create", metadata);
+    assertThat(request.getTargetType()).isEqualTo("grpc");
+    assertThat(request.getPath()).isEqualTo("com.google.service1/Create");
+    assertThat(request.getServer()).isEqualTo("foo.google.com");
+    assertThat(request.getKeyMap()).containsExactly("user", "test", "id", "123");
+  }
+
+  @Test
+  public void create_missingRequiredHeader() {
+    Metadata metadata = new Metadata();
+
+    try {
+      RouteLookupRequest unused = factory.create("com.google.service2", "Create", metadata);
+      fail();
+    } catch (StatusRuntimeException e) {
+      assertThat(e.getStatus().getCode()).isEqualTo(Code.INVALID_ARGUMENT);
+      assertThat(e.getStatus().getDescription())
+          .isEqualTo("Missing mandatory metadata(user) not found");
+    }
+  }
+
+  @Test
+  public void create_pathFallbackMatches() {
+    Metadata metadata = new Metadata();
+    metadata.put(Metadata.Key.of("Parent", Metadata.ASCII_STRING_MARSHALLER), "test");
+    metadata.put(Metadata.Key.of("Password", Metadata.ASCII_STRING_MARSHALLER), "hunter2");
+    metadata.put(Metadata.Key.of("foo", Metadata.ASCII_STRING_MARSHALLER), "bar");
+
+    RouteLookupRequest request = factory.create("com.google.service1" , "Update", metadata);
+
+    assertThat(request.getTargetType()).isEqualTo("grpc");
+    assertThat(request.getPath()).isEqualTo("com.google.service1/Update");
+    assertThat(request.getServer()).isEqualTo("foo.google.com");
+    assertThat(request.getKeyMap()).containsExactly("user", "test", "password", "hunter2");
+  }
+
+  @Test
+  public void create_pathFallbackMatches_optionalHeaderMissing() {
+    Metadata metadata = new Metadata();
+    metadata.put(Metadata.Key.of("User", Metadata.ASCII_STRING_MARSHALLER), "test");
+    metadata.put(Metadata.Key.of("X-Google-Id", Metadata.ASCII_STRING_MARSHALLER), "123");
+    metadata.put(Metadata.Key.of("foo", Metadata.ASCII_STRING_MARSHALLER), "bar");
+
+    RouteLookupRequest request = factory.create("com.google.service1", "Update", metadata);
+
+    assertThat(request.getTargetType()).isEqualTo("grpc");
+    assertThat(request.getPath()).isEqualTo("com.google.service1/Update");
+    assertThat(request.getServer()).isEqualTo("foo.google.com");
+    assertThat(request.getKeyMap()).containsExactly("user", "test");
+  }
+
+  @Test
+  public void create_unknownPath() {
+    Metadata metadata = new Metadata();
+    metadata.put(Metadata.Key.of("User", Metadata.ASCII_STRING_MARSHALLER), "test");
+    metadata.put(Metadata.Key.of("X-Google-Id", Metadata.ASCII_STRING_MARSHALLER), "123");
+    metadata.put(Metadata.Key.of("foo", Metadata.ASCII_STRING_MARSHALLER), "bar");
+
+    RouteLookupRequest request = factory.create("abc.def.service999", "Update", metadata);
+
+    assertThat(request.getTargetType()).isEqualTo("grpc");
+    assertThat(request.getPath()).isEqualTo("abc.def.service999/Update");
+    assertThat(request.getServer()).isEqualTo("foo.google.com");
+    assertThat(request.getKeyMap()).isEmpty();
+  }
+
+  @Test
+  public void create_noMethodInRlsConfig() {
+    Metadata metadata = new Metadata();
+    metadata.put(Metadata.Key.of("User", Metadata.ASCII_STRING_MARSHALLER), "test");
+    metadata.put(Metadata.Key.of("X-Google-Id", Metadata.ASCII_STRING_MARSHALLER), "123");
+    metadata.put(Metadata.Key.of("foo", Metadata.ASCII_STRING_MARSHALLER), "bar");
+
+    RouteLookupRequest request = factory.create("com.google.service3", "Update", metadata);
+
+    assertThat(request.getTargetType()).isEqualTo("grpc");
+    assertThat(request.getPath()).isEqualTo("com.google.service3/Update");
+    assertThat(request.getServer()).isEqualTo("foo.google.com");
+    assertThat(request.getKeyMap()).containsExactly("user", "test");
+  }
+}