feat: Add OpenTelemetry Traces to GRPC (#2783)

sydney-munro · web-flow · commit 5e6eb84a7720 · 2024-10-22T09:34:10.000-07:00
* feat: Add OpenTelemetry to GRPC

* adding rpc.system value

* refactor test for both transports

* lint

* add branch protection rules to otel feature branch

* pr comment
diff --git a/.github/sync-repo-settings.yaml b/.github/sync-repo-settings.yaml
@@ -162,6 +162,22 @@ branchProtectionRules:
       - 'Kokoro - Test: Java GraalVM Native Image'
       - 'Kokoro - Test: Java 17 GraalVM Native Image'
       - javadoc
+  - pattern: otel-v1-branch
+    isAdminEnforced: true
+    requiredApprovingReviewCount: 1
+    requiresCodeOwnerReviews: true
+    requiresStrictStatusChecks: false
+    requiredStatusCheckContexts:
+    - dependencies (17)
+    - lint
+    - clirr
+    - units (8)
+    - units (11)
+    - 'Kokoro - Test: Integration'
+    - cla/google
+    - 'Kokoro - Test: Java GraalVM Native Image'
+    - 'Kokoro - Test: Java 17 GraalVM Native Image'
+    - javadoc
 permissionRules:
   - team: yoshi-admins
     permission: admin
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/GrpcStorageImpl.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/GrpcStorageImpl.java
@@ -68,6 +68,7 @@
 import com.google.cloud.storage.UnifiedOpts.Opts;
 import com.google.cloud.storage.UnifiedOpts.ProjectId;
 import com.google.cloud.storage.UnifiedOpts.UserProject;
+import com.google.cloud.storage.otel.OpenTelemetryTraceUtil;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
@@ -168,6 +169,7 @@ final class GrpcStorageImpl extends BaseService<StorageOptions>
   // workaround for https://github.com/googleapis/java-storage/issues/1736
   private final Opts<UserProject> defaultOpts;
   @Deprecated private final ProjectId defaultProjectId;
+  private final OpenTelemetryTraceUtil openTelemetryTraceUtil;
 
   GrpcStorageImpl(
       GrpcStorageOptions options,
@@ -184,6 +186,7 @@ final class GrpcStorageImpl extends BaseService<StorageOptions>
     this.retryAlgorithmManager = options.getRetryAlgorithmManager();
     this.syntaxDecoders = new SyntaxDecoders();
     this.defaultProjectId = UnifiedOpts.projectId(options.getProjectId());
+    this.openTelemetryTraceUtil = OpenTelemetryTraceUtil.getInstance(options);
   }
 
   @Override
@@ -198,6 +201,8 @@ public void close() throws Exception {
 
   @Override
   public Bucket create(BucketInfo bucketInfo, BucketTargetOption... options) {
+    OpenTelemetryTraceUtil.Span otelSpan =
+        openTelemetryTraceUtil.startSpan("create(Bucket, BucketTargetOption");
     Opts<BucketTargetOpt> opts = Opts.unwrap(options).resolveFrom(bucketInfo).prepend(defaultOpts);
     GrpcCallContext grpcCallContext =
         opts.grpcMetadataMapper().apply(GrpcCallContext.createDefault());
@@ -212,11 +217,20 @@ public Bucket create(BucketInfo bucketInfo, BucketTargetOption... options) {
             .setParent("projects/_");
     CreateBucketRequest req = opts.createBucketsRequest().apply(builder).build();
     GrpcCallContext merge = Utils.merge(grpcCallContext, Retrying.newCallContext());
-    return Retrying.run(
-        getOptions(),
-        retryAlgorithmManager.getFor(req),
-        () -> storageClient.createBucketCallable().call(req, merge),
-        syntaxDecoders.bucket);
+    try (OpenTelemetryTraceUtil.Scope unused = otelSpan.makeCurrent()) {
+      return Retrying.run(
+          getOptions(),
+          retryAlgorithmManager.getFor(req),
+          () -> storageClient.createBucketCallable().call(req, merge),
+          syntaxDecoders.bucket);
+    } catch (Exception ex) {
+      otelSpan.recordException(ex);
+      otelSpan.setStatus(
+          io.opentelemetry.api.trace.StatusCode.ERROR, ex.getClass().getSimpleName());
+      throw StorageException.coalesce(ex);
+    } finally {
+      otelSpan.end();
+    }
   }
 
   @Override
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/GrpcStorageOptions.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/GrpcStorageOptions.java
@@ -119,6 +119,7 @@ public final class GrpcStorageOptions extends StorageOptions
   private final boolean grpcClientMetricsManuallyEnabled;
   private final GrpcInterceptorProvider grpcInterceptorProvider;
   private final BlobWriteSessionConfig blobWriteSessionConfig;
+  private OpenTelemetrySdk openTelemetrySdk;
 
   private GrpcStorageOptions(Builder builder, GrpcStorageDefaults serviceDefaults) {
     super(builder, serviceDefaults);
@@ -134,6 +135,7 @@ private GrpcStorageOptions(Builder builder, GrpcStorageDefaults serviceDefaults)
     this.grpcClientMetricsManuallyEnabled = builder.grpcMetricsManuallyEnabled;
     this.grpcInterceptorProvider = builder.grpcInterceptorProvider;
     this.blobWriteSessionConfig = builder.blobWriteSessionConfig;
+    this.openTelemetrySdk = builder.openTelemetrySdk;
   }
 
   @Override
@@ -350,7 +352,7 @@ private Tuple<StorageSettings, Opts<UserProject>> resolveSettingsAndOpts() throw
 
   @Override
   public OpenTelemetrySdk getOpenTelemetrySdk() {
-    return null;
+    return openTelemetrySdk;
   }
 
   /** @since 2.14.0 This new api is in preview and is subject to breaking changes. */
@@ -435,6 +437,8 @@ public static final class Builder extends StorageOptions.Builder {
 
     private boolean grpcMetricsManuallyEnabled = false;
 
+    private OpenTelemetrySdk openTelemetrySdk;
+
     Builder() {}
 
     Builder(StorageOptions options) {
@@ -446,6 +450,7 @@ public static final class Builder extends StorageOptions.Builder {
       this.enableGrpcClientMetrics = gso.enableGrpcClientMetrics;
       this.grpcInterceptorProvider = gso.grpcInterceptorProvider;
       this.blobWriteSessionConfig = gso.blobWriteSessionConfig;
+      this.openTelemetrySdk = gso.openTelemetrySdk;
     }
 
     /**
@@ -634,9 +639,15 @@ public GrpcStorageOptions.Builder setBlobWriteSessionConfig(
       return this;
     }
 
-    @Override
-    public StorageOptions.Builder setOpenTelemetrySdk(@NonNull OpenTelemetrySdk openTelemetrySdk) {
-      return null;
+    /**
+     * Enable OpenTelemetry Tracing and provide an instance for the client to use.
+     *
+     * @param openTelemetrySdk User defined instance of OpenTelemetry SDK to be used by the library
+     */
+    public GrpcStorageOptions.Builder setOpenTelemetrySdk(OpenTelemetrySdk openTelemetrySdk) {
+      requireNonNull(openTelemetrySdk, "openTelemetry must be non null");
+      this.openTelemetrySdk = openTelemetrySdk;
+      return this;
     }
 
     /** @since 2.14.0 This new api is in preview and is subject to breaking changes. */
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/otel/OpenTelemetryInstance.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/otel/OpenTelemetryInstance.java
@@ -18,10 +18,12 @@
 
 import com.google.api.core.ApiFuture;
 import com.google.api.gax.core.GaxProperties;
+import com.google.cloud.storage.GrpcStorageOptions;
 import com.google.cloud.storage.StorageOptions;
 import com.google.cloud.storage.otel.OpenTelemetryTraceUtil.Context;
 import com.google.cloud.storage.otel.OpenTelemetryTraceUtil.Span;
 import io.opentelemetry.api.OpenTelemetry;
+import io.opentelemetry.api.common.AttributeKey;
 import io.opentelemetry.api.common.Attributes;
 import io.opentelemetry.api.common.AttributesBuilder;
 import io.opentelemetry.api.trace.SpanBuilder;
@@ -38,11 +40,14 @@ class OpenTelemetryInstance implements OpenTelemetryTraceUtil {
 
   private static final String LIBRARY_NAME = "cloud.google.com/java/storage";
 
+  private final String transport;
+
   public OpenTelemetryInstance(StorageOptions storageOptions) {
     this.storageOptions = storageOptions;
     this.openTelemetry = storageOptions.getOpenTelemetrySdk();
     this.tracer =
         openTelemetry.getTracer(LIBRARY_NAME, GaxProperties.getLibraryVersion(this.getClass()));
+    this.transport = storageOptions instanceof GrpcStorageOptions ? "grpc" : "http";
   }
 
   static class Span implements OpenTelemetryTraceUtil.Span {
@@ -56,7 +61,12 @@ static class Span implements OpenTelemetryTraceUtil.Span {
 
     @Override
     public OpenTelemetryTraceUtil.Span recordException(Throwable error) {
-      span.recordException(error);
+      span.recordException(
+          error,
+          Attributes.of(
+              AttributeKey.stringKey("exception.message"), error.getMessage(),
+              AttributeKey.stringKey("exception.type"), error.getClass().getName(),
+              AttributeKey.stringKey("exception.stacktrace"), error.getStackTrace().toString()));
       return this;
     }
 
@@ -146,6 +156,7 @@ public Scope makeCurrent() {
   public OpenTelemetryTraceUtil.Span startSpan(String methodName) {
     String formatSpanName = String.format("%s.%s/%s", "storage", "client", methodName);
     SpanBuilder spanBuilder = tracer.spanBuilder(formatSpanName).setSpanKind(SpanKind.CLIENT);
+    spanBuilder.setAttribute("rpc.system", transport);
     io.opentelemetry.api.trace.Span span =
         addSettingsAttributesToCurrentSpan(spanBuilder).startSpan();
     return new Span(span, formatSpanName);
diff --git a/google-cloud-storage/src/test/java/com/google/cloud/storage/ITOpenTelemetryTest.java b/google-cloud-storage/src/test/java/com/google/cloud/storage/ITOpenTelemetryTest.java
@@ -16,12 +16,7 @@
 
 package com.google.cloud.storage;
 
-import com.google.cloud.storage.TransportCompatibility.Transport;
-import com.google.cloud.storage.it.runner.StorageITRunner;
-import com.google.cloud.storage.it.runner.annotations.Backend;
-import com.google.cloud.storage.it.runner.annotations.CrossRun;
-import com.google.cloud.storage.it.runner.annotations.Inject;
-import com.google.cloud.storage.it.runner.registry.Generator;
+import com.google.cloud.storage.testing.RemoteStorageHelper;
 import io.opentelemetry.api.common.AttributeKey;
 import io.opentelemetry.sdk.OpenTelemetrySdk;
 import io.opentelemetry.sdk.common.CompletableResultCode;
@@ -33,16 +28,11 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
+import java.util.UUID;
 import org.junit.Assert;
 import org.junit.Test;
-import org.junit.runner.RunWith;
 
-@RunWith(StorageITRunner.class)
-@CrossRun(
-    backends = {Backend.PROD},
-    transports = {Transport.HTTP})
 public final class ITOpenTelemetryTest {
-  @Inject public Generator generator;
 
   @Test
   public void checkInstrumentation() {
@@ -58,29 +48,80 @@ public void checkInstrumentation() {
     StorageOptions storageOptions =
         StorageOptions.http().setOpenTelemetrySdk(openTelemetrySdk).build();
     Storage storage = storageOptions.getService();
-    storage.create(BucketInfo.of(generator.randomBucketName()));
+    String bucket = randomBucketName();
+    storage.create(BucketInfo.of(bucket));
     TestExporter testExported = (TestExporter) exporter;
     SpanData spanData = testExported.getExportedSpans().get(0);
     Assert.assertEquals("Storage", getAttributeValue(spanData, "gcp.client.service"));
     Assert.assertEquals("googleapis/java-storage", getAttributeValue(spanData, "gcp.client.repo"));
     Assert.assertEquals(
         "com.google.cloud.google-cloud-storage",
         getAttributeValue(spanData, "gcp.client.artifact"));
+    Assert.assertEquals("http", getAttributeValue(spanData, "rpc.system"));
+
+    // Cleanup
+    RemoteStorageHelper.forceDelete(storage, bucket);
   }
 
   @Test
-  public void noOpDoesNothing() {
-    StorageOptions storageOptions = StorageOptions.http().build();
+  public void checkInstrumentationGrpc() {
+    SpanExporter exporter = new TestExporter();
+
+    OpenTelemetrySdk openTelemetrySdk =
+        OpenTelemetrySdk.builder()
+            .setTracerProvider(
+                SdkTracerProvider.builder()
+                    .addSpanProcessor(SimpleSpanProcessor.create(exporter))
+                    .build())
+            .build();
+    StorageOptions storageOptions =
+        StorageOptions.grpc().setOpenTelemetrySdk(openTelemetrySdk).build();
     Storage storage = storageOptions.getService();
-    storage.create(BucketInfo.of(generator.randomBucketName()));
+    String bucket = randomBucketName();
+    storage.create(BucketInfo.of(bucket));
+    TestExporter testExported = (TestExporter) exporter;
+    SpanData spanData = testExported.getExportedSpans().get(0);
+    Assert.assertEquals("Storage", getAttributeValue(spanData, "gcp.client.service"));
+    Assert.assertEquals("googleapis/java-storage", getAttributeValue(spanData, "gcp.client.repo"));
+    Assert.assertEquals(
+        "com.google.cloud.google-cloud-storage",
+        getAttributeValue(spanData, "gcp.client.artifact"));
+    Assert.assertEquals("grpc", getAttributeValue(spanData, "rpc.system"));
+
+    // Cleanup
+    RemoteStorageHelper.forceDelete(storage, bucket);
+  }
+
+  @Test
+  public void noOpDoesNothing() {
+    String httpBucket = randomBucketName();
+    String grpcBucket = randomBucketName();
+    // NoOp for HTTP
+    StorageOptions storageOptionsHttp = StorageOptions.http().build();
+    Storage storageHttp = storageOptionsHttp.getService();
+    storageHttp.create(BucketInfo.of(httpBucket));
+
+    // NoOp for Grpc
+    StorageOptions storageOptionsGrpc = StorageOptions.grpc().build();
+    Storage storageGrpc = storageOptionsGrpc.getService();
+    storageGrpc.create(BucketInfo.of(grpcBucket));
+
+    // cleanup
+    RemoteStorageHelper.forceDelete(storageHttp, httpBucket);
+    RemoteStorageHelper.forceDelete(storageGrpc, grpcBucket);
   }
 
   private String getAttributeValue(SpanData spanData, String key) {
     return spanData.getAttributes().get(AttributeKey.stringKey(key)).toString();
   }
+
+  public String randomBucketName() {
+    return "java-storage-grpc-rand-" + UUID.randomUUID();
+  }
 }
 
 class TestExporter implements SpanExporter {
+
   public final List<SpanData> exportedSpans = Collections.synchronizedList(new ArrayList<>());
 
   @Override
