fix: misc fixes for bigtable-proxy (#9736)

* fix: misc fixes for bigtable-proxy

* fix dependecy conflict that shared-configuration created for the transitive dep google-cloud-core
* some docs
* typo in Verify & extract constants
* hard fail if call options credentials receives are missing a tracer
* rename the label apiclient -> api_client for readability
* remove stray imports
* format

* presence typo

* fix ping and warm

* update readme for verify

* add multiple response detection to ping and warm

* fix comment

* add a TODO for a debug metric

Author: igorbernstein2

bigtable/bigtable-proxy/README.md

@@ -64,14 +64,27 @@ in a project your choosing. The metrics will be published under the namespace
 # Build the binary
 mvn package
 
-# use the binary
+# unpack the binary on the proxy host
 unzip target/bigtable-proxy-0.0.1-SNAPSHOT-bin.zip
 cd bigtable-proxy-0.0.1-SNAPSHOT
+
+# Verify that the proxy has require permissions using an existing table. Please note that the table
+# data will not be modified, however a test metric will be written.
+./bigtable-verify.sh  \
+  --bigtable-project-id=$BIGTABLE_PROJECT_ID \
+  --bigtable-instance-id=$BIGTABLE_INSTANCE_ID \
+  --bigtable-table-id=$BIGTABLE_TABLE_ID \
+  --metrics-project-id=$METRICS_PROJECT_ID
+
+# Then start the proxy on the specified port. The proxy can forward requests for multiple
+# Bigtable projects/instances/tables. However it will export health metrics to a single project
+# specified by `metrics-project-id`. 
 ./bigtable-proxy.sh \
   --listen-port=1234 \
   --metrics-project-id=SOME_GCP_PROJECT
 
-export BIGTABLE_EMULATOR_HOST=1234
+# Start your application, and redirect the bigtable client to connect to the local proxy. 
+export BIGTABLE_EMULATOR_HOST="localhost:1234"
 path/to/application/with/bigtable/client
 ```
 

bigtable/bigtable-proxy/pom.xml

@@ -130,6 +130,11 @@
       <artifactId>exporter-metrics</artifactId>
       <version>${exporter-metrics.version}</version>
     </dependency>
+    <!-- Workaround shared-configuration incorrectly marking google-cloud-core as a test dep -->
+    <dependency>
+      <groupId>com.google.cloud</groupId>
+      <artifactId>google-cloud-core</artifactId>
+    </dependency>
     <dependency>
       <groupId>io.opentelemetry.contrib</groupId>
       <artifactId>opentelemetry-gcp-resources</artifactId>

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

@@ -18,8 +18,6 @@
 
 import com.google.cloud.bigtable.examples.proxy.commands.Serve;
 import com.google.cloud.bigtable.examples.proxy.commands.Verify;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 import org.slf4j.bridge.SLF4JBridgeHandler;
 import picocli.CommandLine;
 import picocli.CommandLine.Command;

bigtable/bigtable-proxy/src/main/java/com/google/cloud/bigtable/examples/proxy/channelpool/DataChannel.java

@@ -17,22 +17,28 @@
 package com.google.cloud.bigtable.examples.proxy.channelpool;
 
 import com.google.bigtable.v2.BigtableGrpc;
-import com.google.bigtable.v2.BigtableGrpc.BigtableFutureStub;
 import com.google.bigtable.v2.PingAndWarmRequest;
 import com.google.bigtable.v2.PingAndWarmResponse;
+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 com.google.common.util.concurrent.ListenableFuture;
+import com.google.common.util.concurrent.SettableFuture;
 import io.grpc.CallCredentials;
 import io.grpc.CallOptions;
 import io.grpc.ClientCall;
+import io.grpc.ClientCall.Listener;
 import io.grpc.ConnectivityState;
 import io.grpc.Deadline;
 import io.grpc.ExperimentalApi;
 import io.grpc.ManagedChannel;
 import io.grpc.ManagedChannelBuilder;
+import io.grpc.Metadata;
 import io.grpc.MethodDescriptor;
+import io.grpc.Status;
 import io.grpc.StatusRuntimeException;
+import java.net.URLEncoder;
+import java.nio.charset.StandardCharsets;
 import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.ExecutionException;
@@ -50,7 +56,7 @@ public class DataChannel extends ManagedChannel {
   private final ManagedChannel inner;
   private final Metrics metrics;
   private final ResourceCollector resourceCollector;
-  private final BigtableFutureStub warmingStub;
+  private final CallCredentials callCredentials;
   private final ScheduledFuture<?> antiIdleTask;
 
   private final AtomicBoolean closed = new AtomicBoolean();
@@ -65,6 +71,7 @@ public DataChannel(
       Metrics metrics) {
     this.resourceCollector = resourceCollector;
 
+    this.callCredentials = callCredentials;
     inner =
         ManagedChannelBuilder.forAddress(endpoint, port)
             .userAgent(userAgent)
@@ -76,8 +83,6 @@ public DataChannel(
     this.metrics = metrics;
 
     try {
-      warmingStub = BigtableGrpc.newFutureStub(inner).withCallCredentials(callCredentials);
-
       warm();
     } catch (RuntimeException e) {
       try {
@@ -107,10 +112,8 @@ private void warm() {
       return;
     }
 
-    BigtableFutureStub timedStub = warmingStub.withDeadline(Deadline.after(1, TimeUnit.MINUTES));
-
     List<ListenableFuture<PingAndWarmResponse>> futures =
-        requests.stream().map(timedStub::pingAndWarm).collect(Collectors.toList());
+        requests.stream().map(this::sendPingAndWarm).collect(Collectors.toList());
 
     int successCount = 0;
     int failures = 0;
@@ -148,6 +151,61 @@ private void warm() {
     }
   }
 
+  private ListenableFuture<PingAndWarmResponse> sendPingAndWarm(PingAndWarmRequest request) {
+    CallLabels callLabels =
+        CallLabels.create(
+            BigtableGrpc.getPingAndWarmMethod(),
+            Optional.of("bigtableproxy"),
+            Optional.of(request.getName()),
+            Optional.of(request.getAppProfileId()));
+    Tracer tracer = new Tracer(metrics, callLabels);
+
+    CallOptions callOptions =
+        CallOptions.DEFAULT
+            .withCallCredentials(callCredentials)
+            .withDeadline(Deadline.after(1, TimeUnit.MINUTES));
+    callOptions = tracer.injectIntoCallOptions(callOptions);
+
+    ClientCall<PingAndWarmRequest, PingAndWarmResponse> call =
+        inner.newCall(BigtableGrpc.getPingAndWarmMethod(), callOptions);
+
+    Metadata metadata = new Metadata();
+    metadata.put(
+        CallLabels.REQUEST_PARAMS,
+        String.format(
+            "name=projects/%s/instances/%s",
+            URLEncoder.encode(request.getName(), StandardCharsets.UTF_8),
+            URLEncoder.encode(request.getAppProfileId(), StandardCharsets.UTF_8)));
+
+    SettableFuture<PingAndWarmResponse> f = SettableFuture.create();
+    call.start(
+        new Listener<>() {
+          @Override
+          public void onMessage(PingAndWarmResponse response) {
+            if (!f.set(response)) {
+              // TODO: set a metric
+              LOGGER.warn("PingAndWarm returned multiple responses");
+            }
+          }
+
+          @Override
+          public void onClose(Status status, Metadata trailers) {
+            tracer.onCallFinished(status);
+
+            if (status.isOk()) {
+              f.setException(new IllegalStateException("PingAndWarm was missing a response"));
+            } else {
+              f.setException(status.asRuntimeException());
+            }
+          }
+        },
+        metadata);
+    call.sendMessage(request);
+    call.request(Integer.MAX_VALUE);
+
+    return f;
+  }
+
   @Override
   public ManagedChannel shutdown() {
     if (closed.compareAndSet(false, true)) {
@@ -208,9 +266,13 @@ public void enterIdle() {
   @Override
   public <RequestT, ResponseT> ClientCall<RequestT, ResponseT> newCall(
       MethodDescriptor<RequestT, ResponseT> methodDescriptor, CallOptions callOptions) {
-    Optional.ofNullable(Tracer.extractTracerFromCallOptions(callOptions))
-        .map(Tracer::getCallLabels)
-        .ifPresent(resourceCollector::collect);
+    Tracer tracer =
+        Optional.ofNullable(Tracer.extractTracerFromCallOptions(callOptions))
+            .orElseThrow(
+                () ->
+                    new IllegalStateException(
+                        "DataChannel failed to extract Tracer from CallOptions"));
+    resourceCollector.collect(tracer.getCallLabels());
 
     return inner.newCall(methodDescriptor, callOptions);
   }

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

@@ -106,6 +106,9 @@ void start() throws IOException {
         new InstrumentedCallCredentials(MoreCallCredentials.from(credentials));
 
     if (metrics == null) {
+      // InstrumentedCallCredentials expect to only be called when a Tracer is available in the
+      // CallOptions. This is only true for DataChannel pingAndWarm and things invoked by
+      // ProxyHandler. MetricsImpl does not do this, so it must get undecorated credentials.
       metrics = new MetricsImpl(credentials, metricsProjectId);
     }
 

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

@@ -29,10 +29,9 @@
 import com.google.bigtable.v2.RowFilter;
 import com.google.bigtable.v2.RowFilter.Chain;
 import com.google.bigtable.v2.RowSet;
+import com.google.cloud.bigtable.examples.proxy.metrics.MetricsImpl;
 import com.google.cloud.opentelemetry.metric.GoogleCloudMetricExporter;
 import com.google.cloud.opentelemetry.metric.MetricConfiguration;
-import com.google.cloud.opentelemetry.resource.GcpResource;
-import com.google.cloud.opentelemetry.resource.ResourceTranslator;
 import com.google.common.collect.ImmutableList;
 import com.google.common.net.PercentEscaper;
 import com.google.protobuf.ByteString;
@@ -53,7 +52,6 @@
 import io.opentelemetry.api.common.Attributes;
 import io.opentelemetry.contrib.gcp.resource.GCPResourceProvider;
 import io.opentelemetry.sdk.common.CompletableResultCode;
-import io.opentelemetry.sdk.common.InstrumentationScopeInfo;
 import io.opentelemetry.sdk.metrics.data.MetricData;
 import io.opentelemetry.sdk.metrics.export.MetricExporter;
 import io.opentelemetry.sdk.metrics.internal.data.ImmutableGaugeData;
@@ -102,7 +100,6 @@ public class Verify implements Callable<Void> {
       showDefaultValue = Visibility.ALWAYS)
   Endpoint dataEndpoint = Endpoint.create("bigtable.googleapis.com", 443);
 
-
   Credentials credentials = null;
 
   @Override
@@ -141,7 +138,7 @@ private void checkBigtable(CallCredentials callCredentials, String tableName) {
               .setTableName(
                   String.format(
                       "projects/%s/instances/%s/tables/%s",
-                      bigtableProjectId, bigtableTableId, bigtableTableId))
+                      bigtableProjectId, bigtableInstanceId, bigtableTableId))
               .setRowsLimit(1)
               .setRows(
                   RowSet.newBuilder().addRowKeys(ByteString.copyFromUtf8("some-nonexistent-row")))
@@ -194,7 +191,6 @@ void checkMetrics(Credentials creds) throws IOException {
 
     GCPResourceProvider resourceProvider = new GCPResourceProvider();
     Resource resource = Resource.create(resourceProvider.getAttributes());
-    GcpResource gcpResource = ResourceTranslator.mapResource(resource);
 
     MetricExporter exporter =
         GoogleCloudMetricExporter.createWithConfiguration(
@@ -208,10 +204,10 @@ void checkMetrics(Credentials creds) throws IOException {
         ImmutableList.of(
             ImmutableMetricData.createLongGauge(
                 resource,
-                InstrumentationScopeInfo.create("bigtable-proxy"),
-                "bigtableproxy.presence",
-                "Number of proxy processes",
-                "{process}",
+                MetricsImpl.INSTRUMENTATION_SCOPE_INFO,
+                MetricsImpl.METRIC_PRESENCE_NAME,
+                MetricsImpl.METRIC_PRESENCE_DESC,
+                MetricsImpl.METRIC_PRESENCE_UNIT,
                 ImmutableGaugeData.create(
                     ImmutableList.of(
                         ImmutableLongPointData.create(
@@ -222,6 +218,7 @@ void checkMetrics(Credentials creds) throws IOException {
     CompletableResultCode result = exporter.export(metricData);
     result.join(1, TimeUnit.MINUTES);
 
+    System.out.println("Metrics resource: " + resource);
     if (result.isSuccess()) {
       System.out.println("Metrics write: OK");
     } else {

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

@@ -29,16 +29,16 @@
  * 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>
+ * metrics & the upstream client. The primary headers consulted are:
  *
  * <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>
+ *   <li>{@code x-goog-request-params} - contains the resource and app profile id
+ *   <li>{@code x-goog-api-client} - contains the client info of the downstream client
  * </ul>
  */
 @AutoValue
 public abstract class CallLabels {
-  private static final Key<String> REQUEST_PARAMS =
+  public 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);

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

@@ -24,6 +24,7 @@
 import java.time.Duration;
 import java.util.concurrent.Executor;
 import java.util.concurrent.TimeUnit;
+import javax.annotation.Nullable;
 
 public class InstrumentedCallCredentials extends CallCredentials
     implements InternalMayRequireSpecificExecutor {
@@ -40,7 +41,13 @@ public InstrumentedCallCredentials(CallCredentials inner) {
   @Override
   public void applyRequestMetadata(
       RequestInfo requestInfo, Executor appExecutor, MetadataApplier applier) {
-    Tracer tracer = Tracer.extractTracerFromCallOptions(requestInfo.getCallOptions());
+    @Nullable Tracer tracer = Tracer.extractTracerFromCallOptions(requestInfo.getCallOptions());
+    if (tracer == null) {
+      applier.fail(
+          Status.INTERNAL.withDescription(
+              "InstrumentedCallCredentials failed to extract tracer from CallOptions"));
+      return;
+    }
     final Stopwatch stopwatch = Stopwatch.createStarted();
 
     inner.applyRequestMetadata(

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

@@ -28,6 +28,7 @@
 import io.opentelemetry.api.metrics.LongUpDownCounter;
 import io.opentelemetry.api.metrics.Meter;
 import io.opentelemetry.contrib.gcp.resource.GCPResourceProvider;
+import io.opentelemetry.sdk.common.InstrumentationScopeInfo;
 import io.opentelemetry.sdk.metrics.SdkMeterProvider;
 import io.opentelemetry.sdk.metrics.export.MetricExporter;
 import io.opentelemetry.sdk.metrics.export.PeriodicMetricReader;
@@ -39,14 +40,21 @@
 import java.util.function.Supplier;
 
 public class MetricsImpl implements Closeable, Metrics {
+  public static final InstrumentationScopeInfo INSTRUMENTATION_SCOPE_INFO =
+      InstrumentationScopeInfo.create("bigtable-proxy");
+
   public static final String METRIC_PREFIX = "bigtableproxy.";
 
-  static final AttributeKey<String> API_CLIENT_KEY = AttributeKey.stringKey("apiclient");
+  static final AttributeKey<String> API_CLIENT_KEY = AttributeKey.stringKey("api_client");
   static final AttributeKey<String> RESOURCE_KEY = AttributeKey.stringKey("resource");
   static final AttributeKey<String> APP_PROFILE_KEY = AttributeKey.stringKey("app_profile");
   static final AttributeKey<String> METHOD_KEY = AttributeKey.stringKey("method");
   static final AttributeKey<String> STATUS_KEY = AttributeKey.stringKey("status");
 
+  public static final String METRIC_PRESENCE_NAME = METRIC_PREFIX + "presence";
+  public static final String METRIC_PRESENCE_DESC = "Number of proxy processes";
+  public static final String METRIC_PRESENCE_UNIT = "{process}";
+
   private final SdkMeterProvider meterProvider;
 
   private final DoubleHistogram gfeLatency;
@@ -67,7 +75,11 @@ public class MetricsImpl implements Closeable, Metrics {
 
   public MetricsImpl(Credentials credentials, String projectId) throws IOException {
     meterProvider = createMeterProvider(credentials, projectId);
-    Meter meter = meterProvider.meterBuilder("bigtableproxy").build();
+    Meter meter =
+        meterProvider
+            .meterBuilder(INSTRUMENTATION_SCOPE_INFO.getName())
+            .setInstrumentationVersion(INSTRUMENTATION_SCOPE_INFO.getVersion())
+            .build();
 
     serverCallsStarted =
         meter
@@ -151,9 +163,9 @@ public MetricsImpl(Credentials credentials, String projectId) throws IOException
         .buildWithCallback(o -> o.record(maxSeen.getAndSet(0)));
 
     meter
-        .gaugeBuilder(METRIC_PREFIX + ".presence")
-        .setDescription("Number of proxy processes")
-        .setUnit("{process}")
+        .gaugeBuilder(METRIC_PRESENCE_NAME)
+        .setDescription(METRIC_PRESENCE_DESC)
+        .setUnit(METRIC_PRESENCE_UNIT)
         .ofLongs()
         .buildWithCallback(o -> o.record(1));
   }

bigtable/bigtable-proxy/src/test/java/com/google/cloud/bigtable/examples/proxy/commands/ServeMetricsTest.java

@@ -98,19 +98,13 @@ public void setUp() throws Exception {
 
     fakeServiceChannel =
         grpcCleanup.register(
-            ManagedChannelBuilder.forAddress("localhost", server.getPort())
-                .usePlaintext()
-                .build()
-        );
+            ManagedChannelBuilder.forAddress("localhost", server.getPort()).usePlaintext().build());
 
     serve = createAndStartCommand(fakeServiceChannel, fakeCredentials, mockMetrics);
 
     proxyChannel =
         grpcCleanup.register(
-            ManagedChannelBuilder.forAddress("localhost", serve.listenPort)
-                .usePlaintext()
-                .build()
-        );
+            ManagedChannelBuilder.forAddress("localhost", serve.listenPort).usePlaintext().build());
   }
 
   @After