feat: add metrics

stack-info: PR: #9691, branch: igorbernstein2/stack/4

Author: igorbernstein2

bigtable/bigtable-proxy/pom.xml

@@ -23,6 +23,8 @@
     <!-- dep versions -->
     <libraries-bom.version>26.50.0</libraries-bom.version>
 
+    <otel.version>1.44.1</otel.version>
+    <exporter-metrics.version>0.33.0</exporter-metrics.version>
     <slf4j.version>2.0.16</slf4j.version>
     <logback.version>1.5.12</logback.version>
     <auto-value.version>1.11.0</auto-value.version>
@@ -40,6 +42,20 @@
         <type>pom</type>
         <scope>import</scope>
       </dependency>
+      <dependency>
+        <groupId>io.opentelemetry</groupId>
+        <artifactId>opentelemetry-bom</artifactId>
+        <version>${otel.version}</version>
+        <type>pom</type>
+        <scope>import</scope>
+      </dependency>
+      <dependency>
+        <groupId>org.mockito</groupId>
+        <artifactId>mockito-bom</artifactId>
+        <version>5.14.2</version>
+        <type>pom</type>
+        <scope>import</scope>
+      </dependency>
     </dependencies>
   </dependencyManagement>
 
@@ -93,6 +109,23 @@
       <artifactId>proto-google-common-protos</artifactId>
     </dependency>
 
+    <!-- Metrics -->
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-sdk</artifactId>
+      <!-- version managed by opentelemetry-bom -->
+    </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-sdk-metrics</artifactId>
+      <!-- version managed by opentelemetry-bom -->
+    </dependency>
+    <dependency>
+      <groupId>com.google.cloud.opentelemetry</groupId>
+      <artifactId>exporter-metrics</artifactId>
+      <version>${exporter-metrics.version}</version>
+    </dependency>
+
     <!-- Logging -->
     <dependency>
       <groupId>org.slf4j</groupId>
@@ -146,6 +179,12 @@
       <version>${truth.version}</version>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>org.mockito</groupId>
+      <artifactId>mockito-core</artifactId>
+      <!-- version managed by mockito-bom -->
+      <scope>test</scope>
+    </dependency>
   </dependencies>
 
   <build>

bigtable/bigtable-proxy/src/main/java/com/google/cloud/bigtable/examples/proxy/commands/Serve.java

@@ -23,6 +23,9 @@
 import com.google.bigtable.v2.BigtableGrpc;
 import com.google.cloud.bigtable.examples.proxy.core.ProxyHandler;
 import com.google.cloud.bigtable.examples.proxy.core.Registry;
+import com.google.cloud.bigtable.examples.proxy.metrics.InstrumentedCallCredentials;
+import com.google.cloud.bigtable.examples.proxy.metrics.Metrics;
+import com.google.cloud.bigtable.examples.proxy.metrics.MetricsImpl;
 import com.google.common.collect.ImmutableMap;
 import com.google.longrunning.OperationsGrpc;
 import io.grpc.CallCredentials;
@@ -69,10 +72,17 @@ public class Serve implements Callable<Void> {
       showDefaultValue = Visibility.ALWAYS)
   Endpoint adminEndpoint = Endpoint.create("bigtableadmin.googleapis.com", 443);
 
+  @Option(
+      names = "--metrics-project-id",
+      required = true,
+      description = "The project id where metrics should be exported")
+  String metricsProjectId = null;
+
   ManagedChannel adminChannel = null;
   ManagedChannel dataChannel = null;
   Credentials credentials = null;
   Server server;
+  Metrics metrics;
 
   @Override
   public Void call() throws Exception {
@@ -103,18 +113,23 @@ void start() throws IOException {
     if (credentials == null) {
       credentials = GoogleCredentials.getApplicationDefault();
     }
-    CallCredentials callCredentials = MoreCallCredentials.from(credentials);
+    CallCredentials callCredentials =
+        new InstrumentedCallCredentials(MoreCallCredentials.from(credentials));
+
+    if (metrics == null) {
+      metrics = new MetricsImpl(credentials, metricsProjectId);
+    }
 
     Map<String, ServerCallHandler<byte[], byte[]>> serviceMap =
         ImmutableMap.of(
             BigtableGrpc.SERVICE_NAME,
-            new ProxyHandler<>(dataChannel, callCredentials),
+            new ProxyHandler<>(metrics, dataChannel, callCredentials),
             BigtableInstanceAdminGrpc.SERVICE_NAME,
-            new ProxyHandler<>(adminChannel, callCredentials),
+            new ProxyHandler<>(metrics, adminChannel, callCredentials),
             BigtableTableAdminGrpc.SERVICE_NAME,
-            new ProxyHandler<>(adminChannel, callCredentials),
+            new ProxyHandler<>(metrics, adminChannel, callCredentials),
             OperationsGrpc.SERVICE_NAME,
-            new ProxyHandler<>(adminChannel, callCredentials));
+            new ProxyHandler<>(metrics, adminChannel, callCredentials));
 
     server =
         NettyServerBuilder.forAddress(

bigtable/bigtable-proxy/src/main/java/com/google/cloud/bigtable/examples/proxy/core/CallProxy.java

@@ -16,6 +16,7 @@
 
 package com.google.cloud.bigtable.examples.proxy.core;
 
+import com.google.cloud.bigtable.examples.proxy.metrics.Tracer;
 import io.grpc.ClientCall;
 import io.grpc.Metadata;
 import io.grpc.ServerCall;
@@ -24,15 +25,20 @@
 
 /** A per gppc RPC proxy. */
 class CallProxy<ReqT, RespT> {
+
+  private final Tracer tracer;
   final RequestProxy serverCallListener;
   final ResponseProxy clientCallListener;
 
   /**
+   * @param tracer a lifecycle observer to publish metrics.
    * @param serverCall the incoming server call. This will be triggered a customer client.
    * @param clientCall the outgoing call to Bigtable service. This will be created by {@link
    *     ProxyHandler}
    */
-  public CallProxy(ServerCall<ReqT, RespT> serverCall, ClientCall<ReqT, RespT> clientCall) {
+  public CallProxy(
+      Tracer tracer, ServerCall<ReqT, RespT> serverCall, ClientCall<ReqT, RespT> clientCall) {
+    this.tracer = tracer;
     // Listen for incoming request messages and send them to the upstream ClientCall
     // The RequestProxy will respect back pressure from the ClientCall and only request a new
     // message from the incoming rpc when the upstream client call is ready,
@@ -131,6 +137,8 @@ public ResponseProxy(ServerCall<?, RespT> serverCall) {
 
     @Override
     public void onClose(Status status, Metadata trailers) {
+      tracer.onCallFinished(status);
+
       serverCall.close(status, trailers);
     }
 

bigtable/bigtable-proxy/src/main/java/com/google/cloud/bigtable/examples/proxy/core/ProxyHandler.java

@@ -16,6 +16,9 @@
 
 package com.google.cloud.bigtable.examples.proxy.core;
 
+import com.google.cloud.bigtable.examples.proxy.metrics.CallLabels;
+import com.google.cloud.bigtable.examples.proxy.metrics.Metrics;
+import com.google.cloud.bigtable.examples.proxy.metrics.Tracer;
 import io.grpc.CallCredentials;
 import io.grpc.CallOptions;
 import io.grpc.Channel;
@@ -29,25 +32,32 @@ public final class ProxyHandler<ReqT, RespT> implements ServerCallHandler<ReqT,
   private static final Metadata.Key<String> AUTHORIZATION_KEY =
       Metadata.Key.of("Authorization", Metadata.ASCII_STRING_MARSHALLER);
 
+  private final Metrics metrics;
   private final Channel channel;
   private final CallCredentials callCredentials;
 
-  public ProxyHandler(Channel channel, CallCredentials callCredentials) {
+  public ProxyHandler(Metrics metrics, Channel channel, CallCredentials callCredentials) {
+    this.metrics = metrics;
     this.channel = channel;
     this.callCredentials = callCredentials;
   }
 
   @Override
   public ServerCall.Listener<ReqT> startCall(ServerCall<ReqT, RespT> serverCall, Metadata headers) {
-    // Strip incoming credentials
-    headers.removeAll(AUTHORIZATION_KEY);
+    CallLabels callLabels = CallLabels.create(serverCall.getMethodDescriptor(), headers);
+    Tracer tracer = new Tracer(metrics, callLabels);
+
     // Inject proxy credentials
     CallOptions callOptions = CallOptions.DEFAULT.withCallCredentials(callCredentials);
+    callOptions = tracer.injectIntoCallOptions(callOptions);
+
+    // Strip incoming credentials
+    headers.removeAll(AUTHORIZATION_KEY);
 
     ClientCall<ReqT, RespT> clientCall =
         channel.newCall(serverCall.getMethodDescriptor(), callOptions);
 
-    CallProxy<ReqT, RespT> proxy = new CallProxy<>(serverCall, clientCall);
+    CallProxy<ReqT, RespT> proxy = new CallProxy<>(tracer, serverCall, clientCall);
     clientCall.start(proxy.clientCallListener, headers);
     serverCall.request(1);
     clientCall.request(1);

bigtable/bigtable-proxy/src/main/java/com/google/cloud/bigtable/examples/proxy/metrics/CallLabels.java

@@ -0,0 +1,168 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * 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 com.google.cloud.bigtable.examples.proxy.metrics;
+
+import com.google.auto.value.AutoValue;
+import io.grpc.Metadata;
+import io.grpc.Metadata.Key;
+import io.grpc.MethodDescriptor;
+import io.opentelemetry.api.common.Attributes;
+import java.net.URLDecoder;
+import java.nio.charset.StandardCharsets;
+import java.util.Optional;
+
+/**
+ * A value class to encapsulate call identity.
+ *
+ * <p>This call extracts relevant information from request headers and makes it accessible to
+ * metrics & the upstream client. The primary headers consulted are:</p>
+ *
+ * <ul>
+ *   <li>{@code x-goog-request-params} - contains the resource and app profile id</li>
+ *   <li>{@code x-goog-api-client} - contains the client info of the downstream client</li>
+ * </ul>
+ */
+
+@AutoValue
+public abstract class CallLabels {
+  private static final Key<String> REQUEST_PARAMS =
+      Key.of("x-goog-request-params", Metadata.ASCII_STRING_MARSHALLER);
+  private static final Key<String> API_CLIENT =
+      Key.of("x-goog-api-client", Metadata.ASCII_STRING_MARSHALLER);
+
+  enum ResourceNameType {
+    Parent("parent", 0),
+    Name("name", 1),
+    TableName("table_name", 2);
+
+    private final String name;
+    private final int priority;
+
+    ResourceNameType(String name, int priority) {
+      this.name = name;
+      this.priority = priority;
+    }
+  }
+
+  @AutoValue
+  abstract static class ResourceName {
+
+    abstract ResourceNameType getType();
+
+    abstract String getValue();
+
+    static ResourceName create(ResourceNameType type, String value) {
+      return new AutoValue_CallLabels_ResourceName(type, value);
+    }
+  }
+
+  abstract Optional<String> getApiClient();
+
+  public abstract Optional<String> getResourceName();
+
+  public abstract Optional<String> getAppProfileId();
+
+  abstract String getMethodName();
+
+  public abstract Attributes getOtelAttributes();
+
+  public static CallLabels create(MethodDescriptor<?, ?> method, Metadata headers) {
+    Optional<String> apiClient = Optional.ofNullable(headers.get(API_CLIENT));
+
+    String requestParams = Optional.ofNullable(headers.get(REQUEST_PARAMS)).orElse("");
+    String[] encodedKvPairs = requestParams.split("&");
+    Optional<String> resourceName = extractResourceName(encodedKvPairs).map(ResourceName::getValue);
+    Optional<String> appProfile = extractAppProfileId(encodedKvPairs);
+
+    return create(method, apiClient, resourceName, appProfile);
+  }
+
+  public static CallLabels create(
+      MethodDescriptor<?, ?> method,
+      Optional<String> apiClient,
+      Optional<String> resourceName,
+      Optional<String> appProfile) {
+    Attributes otelAttrs =
+        Attributes.builder()
+            .put(MetricsImpl.API_CLIENT_KEY, apiClient.orElse("<missing>"))
+            .put(MetricsImpl.RESOURCE_KEY, resourceName.orElse("<missing>"))
+            .put(MetricsImpl.APP_PROFILE_KEY, appProfile.orElse("<missing>"))
+            .put(MetricsImpl.METHOD_KEY, method.getFullMethodName())
+            .build();
+    return new AutoValue_CallLabels(
+        apiClient, resourceName, appProfile, method.getFullMethodName(), otelAttrs);
+  }
+
+  private static Optional<ResourceName> extractResourceName(String[] encodedKvPairs) {
+    Optional<ResourceName> resourceName = Optional.empty();
+
+    for (String encodedKv : encodedKvPairs) {
+      String[] split = encodedKv.split("=", 2);
+      if (split.length != 2) {
+        continue;
+      }
+      String encodedKey = split[0];
+      String encodedValue = split[1];
+      if (encodedKey.isEmpty() || encodedValue.isEmpty()) {
+        continue;
+      }
+
+      Optional<ResourceNameType> newType = findType(encodedKey);
+
+      if (newType.isEmpty()) {
+        continue;
+      }
+      // Skip if we previously found a resource name and the new resource name type has a lower
+      // priority
+      if (resourceName.isPresent()
+          && newType.get().priority <= resourceName.get().getType().priority) {
+        continue;
+      }
+      String decodedValue = percentDecode(encodedValue);
+
+      resourceName = Optional.of(ResourceName.create(newType.get(), decodedValue));
+    }
+    return resourceName;
+  }
+
+  private static Optional<ResourceNameType> findType(String encodedKey) {
+    String decodedKey = percentDecode(encodedKey);
+
+    for (ResourceNameType type : ResourceNameType.values()) {
+      if (type.name.equals(decodedKey)) {
+        return Optional.of(type);
+      }
+    }
+    return Optional.empty();
+  }
+
+  private static Optional<String> extractAppProfileId(String[] encodedKvPairs) {
+    for (String encodedPair : encodedKvPairs) {
+      if (!encodedPair.startsWith("app_profile_id=")) {
+        continue;
+      }
+      String[] parts = encodedPair.split("=", 2);
+      String encodedValue = parts.length > 1 ? parts[1] : "";
+      return Optional.of(percentDecode(encodedValue));
+    }
+    return Optional.empty();
+  }
+
+  private static String percentDecode(String s) {
+    return URLDecoder.decode(s, StandardCharsets.UTF_8);
+  }
+}