"feat: enable grpc configurator for client-side tracing"

Author: cindy-peng

google-cloud-datastore/pom.xml

@@ -147,6 +147,10 @@
       <artifactId>opentelemetry-context</artifactId>
       <version>${opentelemetry.version}</version>
     </dependency>
+    <dependency>
+      <groupId>io.opentelemetry.instrumentation</groupId>
+      <artifactId>opentelemetry-grpc-1.6</artifactId>
+    </dependency>
     <!-- END OpenTelemetry -->
 
     <!-- Test dependencies -->

google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreOptions.java

@@ -18,6 +18,7 @@
 
 import static com.google.cloud.datastore.Validator.validateNamespace;
 
+import com.google.api.core.ApiFunction;
 import com.google.api.core.BetaApi;
 import com.google.api.gax.grpc.ChannelPoolSettings;
 import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
@@ -35,6 +36,7 @@
 import com.google.cloud.http.HttpTransportOptions;
 import com.google.common.base.MoreObjects;
 import com.google.common.collect.ImmutableSet;
+import io.grpc.ManagedChannelBuilder;
 import java.io.IOException;
 import java.lang.reflect.Method;
 import java.util.Objects;
@@ -110,9 +112,11 @@ public static class Builder extends ServiceOptions.Builder<Datastore, DatastoreO
     private String host;
     private TransportOptions transportOptions;
 
-    @Nullable private DatastoreOpenTelemetryOptions openTelemetryOptions = null;
+    @Nullable
+    private DatastoreOpenTelemetryOptions openTelemetryOptions = null;
 
-    private Builder() {}
+    private Builder() {
+    }
 
     private Builder(DatastoreOptions options) {
       super(options);
@@ -155,8 +159,8 @@ public Builder setHost(String host) {
      *
      * <p>This functionality is experimental and subject to change.
      *
-     * @param channelProvider A InstantiatingGrpcChannelProvider object that defines the transport
-     *     provider for this client.
+     * @param channelProvider A InstantiatingGrpcChannelProvider object that defines the
+     *     transport provider for this client.
      */
     @BetaApi
     public Builder setChannelProvider(TransportChannelProvider channelProvider) {
@@ -172,7 +176,9 @@ public DatastoreOptions build() {
       return new DatastoreOptions(this);
     }
 
-    /** Sets the default namespace to be used by the datastore service. */
+    /**
+     * Sets the default namespace to be used by the datastore service.
+     */
     public Builder setNamespace(String namespace) {
       this.namespace = validateNamespace(namespace);
       return this;
@@ -222,20 +228,35 @@ private DatastoreOptions(Builder builder) {
       throw new IllegalArgumentException(
           "Only gRPC transport allows setting of channel provider or credentials provider");
     } else if (getTransportOptions() instanceof GrpcTransportOptions) {
-      // For grpc transport options, configure default gRPC Connection pool with minChannelCount = 1
-      // and maxChannelCount = 4
-      this.channelProvider =
-          builder.channelProvider != null
-              ? builder.channelProvider
-              : GrpcTransportOptions.setUpChannelProvider(
+      if (builder.channelProvider == null) {
+        /*
+         The default gRPC connection pool is configured with a minimum of 1 channel.
+         The maximum channel count automatically defaults to 200 (Defined in gax-grpc).
+        */
+        ChannelPoolSettings datastoreChannelPoolSettings = ChannelPoolSettings.builder()
+            .setInitialChannelCount(INIT_CHANNEL_COUNT)
+            .setMinChannelCount(MIN_CHANNEL_COUNT)
+            .build();
+
+        ApiFunction<ManagedChannelBuilder, ManagedChannelBuilder> channelConfigurator =
+            this.traceUtil.getChannelConfigurator();
+        if (channelConfigurator == null) {
+          this.channelProvider =
+              GrpcTransportOptions.setUpChannelProvider(
+                  DatastoreSettings.defaultGrpcTransportProviderBuilder()
+                      .setChannelPoolSettings(datastoreChannelPoolSettings), this);
+        } else {
+          // Intercept the grpc channel calls to add telemetry info.
+          this.channelProvider =
+              GrpcTransportOptions.setUpChannelProvider(
                   DatastoreSettings.defaultGrpcTransportProviderBuilder()
-                      .setChannelPoolSettings(
-                          ChannelPoolSettings.builder()
-                              .setInitialChannelCount(INIT_CHANNEL_COUNT)
-                              .setMinChannelCount(MIN_CHANNEL_COUNT)
-                              .setMaxChannelCount(MAX_CHANNEL_COUNT)
-                              .build()),
+                      .setChannelPoolSettings(datastoreChannelPoolSettings)
+                      .setChannelConfigurator(channelConfigurator),
                   this);
+        }
+      } else {
+        this.channelProvider = builder.channelProvider;
+      }
     }
   }
 
@@ -256,6 +277,7 @@ protected String getDefaultProject() {
   }
 
   private static class DatastoreDefaults implements ServiceDefaults<Datastore, DatastoreOptions> {
+
     private final TransportOptions TRANSPORT_OPTIONS = getDefaultTransportOptionsBuilder().build();
 
     @Override
@@ -286,7 +308,9 @@ public static GrpcTransportOptions getDefaultGrpcTransportOptions() {
     return GrpcTransportOptions.newBuilder().build();
   }
 
-  /** Returns the default namespace to be used by the datastore service. */
+  /**
+   * Returns the default namespace to be used by the datastore service.
+   */
   public String getNamespace() {
     return namespace;
   }
@@ -295,7 +319,9 @@ public String getDatabaseId() {
     return this.databaseId;
   }
 
-  /** Returns a default {@code DatastoreOptions} instance. */
+  /**
+   * Returns a default {@code DatastoreOptions} instance.
+   */
   public static DatastoreOptions getDefaultInstance() {
     return newBuilder().build();
   }

google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/EnabledTraceUtil.java

@@ -35,12 +35,15 @@
 import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.api.trace.StatusCode;
 import io.opentelemetry.api.trace.Tracer;
+import io.opentelemetry.api.trace.TracerProvider;
+import io.opentelemetry.instrumentation.grpc.v1_6.GrpcTelemetry;
 import java.util.Map;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 
 /**
- * Tracing utility implementation, used to stub out tracing instrumentation when tracing is enabled.
+ * Tracing utility implementation, used to stub out tracing instrumentation when tracing is
+ * enabled.
  */
 @InternalApi
 public class EnabledTraceUtil implements TraceUtil {
@@ -66,14 +69,30 @@ public OpenTelemetry getOpenTelemetry() {
     return openTelemetry;
   }
 
+  // The gRPC channel configurator that intercepts gRPC calls for tracing purposes.
+  public class OpenTelemetryGrpcChannelConfigurator
+      implements ApiFunction<ManagedChannelBuilder, ManagedChannelBuilder> {
+
+    @Override
+    public ManagedChannelBuilder apply(ManagedChannelBuilder managedChannelBuilder) {
+      GrpcTelemetry grpcTelemetry = GrpcTelemetry.create(getOpenTelemetry());
+      return managedChannelBuilder.intercept(grpcTelemetry.newClientInterceptor());
+    }
+  }
+
   @Override
   @Nullable
   public ApiFunction<ManagedChannelBuilder, ManagedChannelBuilder> getChannelConfigurator() {
-    // TODO(jimit) Update this to return a gRPC Channel Configurator after gRPC upgrade.
-    return null;
+    // Note: using `==` rather than `.equals` since OpenTelemetry has only 1 static instance of
+    // `TracerProvider.noop`.
+    if (openTelemetry.getTracerProvider() == TracerProvider.noop()) {
+      return null;
+    }
+    return new OpenTelemetryGrpcChannelConfigurator();
   }
 
   static class Span implements TraceUtil.Span {
+
     private final io.opentelemetry.api.trace.Span span;
     private final String spanName;
 
@@ -87,13 +106,17 @@ public io.opentelemetry.api.trace.Span getSpan() {
       return this.span;
     }
 
-    /** Ends this span. */
+    /**
+     * Ends this span.
+     */
     @Override
     public void end() {
       span.end();
     }
 
-    /** Ends this span in an error. */
+    /**
+     * Ends this span in an error.
+     */
     @Override
     public void end(Throwable error) {
       span.setStatus(StatusCode.ERROR, error.getMessage());
@@ -137,7 +160,9 @@ public void onSuccess(T result) {
           });
     }
 
-    /** Adds the given event to this span. */
+    /**
+     * Adds the given event to this span.
+     */
     @Override
     public TraceUtil.Span addEvent(String name) {
       span.addEvent(name);
@@ -198,6 +223,7 @@ public Scope makeCurrent() {
   }
 
   static class Scope implements TraceUtil.Scope {
+
     private final io.opentelemetry.context.Scope scope;
 
     Scope(io.opentelemetry.context.Scope scope) {
@@ -211,6 +237,7 @@ public void close() {
   }
 
   static class Context implements TraceUtil.Context {
+
     private final io.opentelemetry.context.Context context;
 
     Context(io.opentelemetry.context.Context context) {
@@ -225,7 +252,9 @@ public Scope makeCurrent() {
     }
   }
 
-  /** Applies the current Datastore instance settings as attributes to the current Span */
+  /**
+   * Applies the current Datastore instance settings as attributes to the current Span
+   */
   @Override
   public SpanBuilder addSettingsAttributesToCurrentSpan(SpanBuilder spanBuilder) {
     spanBuilder =

google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/EnabledTraceUtilTest.java

@@ -66,12 +66,12 @@ public void usesOpenTelemetryFromOptions() {
   @Test
   public void usesGlobalOpenTelemetryIfOpenTelemetryInstanceNotProvided() {
     OpenTelemetrySdk ignored = OpenTelemetrySdk.builder().buildAndRegisterGlobal();
-    DatastoreOptions firestoreOptions =
+    DatastoreOptions datastoreOptions =
         getBaseOptions()
             .setOpenTelemetryOptions(
                 DatastoreOpenTelemetryOptions.newBuilder().setTracingEnabled(true).build())
             .build();
-    EnabledTraceUtil traceUtil = new EnabledTraceUtil(firestoreOptions);
+    EnabledTraceUtil traceUtil = new EnabledTraceUtil(datastoreOptions);
     assertThat(traceUtil.getOpenTelemetry()).isEqualTo(GlobalOpenTelemetry.get());
   }
 
@@ -80,6 +80,34 @@ public void enabledTraceUtilProvidesChannelConfigurator() {
     assertThat(newEnabledTraceUtil().getChannelConfigurator()).isNull();
   }
 
+  @Test
+  public void openTelemetryInstanceRegistersGrpcChannelConfigurator() {
+    OpenTelemetrySdk myOpenTelemetrySdk = OpenTelemetrySdk.builder().build();
+    DatastoreOptions firestoreOptions =
+        getBaseOptions()
+            .setOpenTelemetryOptions(
+                DatastoreOpenTelemetryOptions.newBuilder()
+                    .setTracingEnabled(true)
+                    .setOpenTelemetry(myOpenTelemetrySdk)
+                    .build())
+            .build();
+    EnabledTraceUtil traceUtil = new EnabledTraceUtil(firestoreOptions);
+    assertThat(traceUtil.getChannelConfigurator()).isNotNull();
+  }
+
+  @Test
+  public void globalOpenTelemetryRegistersGrpcChannelConfigurator() {
+
+    OpenTelemetrySdk.builder().buildAndRegisterGlobal();
+    DatastoreOptions datastoreOptions =
+        getBaseOptions()
+            .setOpenTelemetryOptions(
+                DatastoreOpenTelemetryOptions.newBuilder().setTracingEnabled(true).build())
+            .build();
+    EnabledTraceUtil traceUtil = new EnabledTraceUtil(datastoreOptions);
+    assertThat(traceUtil.getChannelConfigurator()).isNotNull();
+  }
+
   @Test
   public void usesEnabledContext() {
     assertThat(newEnabledTraceUtil().getCurrentContext() instanceof EnabledTraceUtil.Context)
@@ -92,7 +120,7 @@ public void usesEnabledSpan() {
     assertThat(traceUtil.getCurrentSpan() instanceof EnabledTraceUtil.Span).isTrue();
     assertThat(traceUtil.startSpan("foo") != null).isTrue();
     assertThat(
-            traceUtil.startSpan("foo", traceUtil.getCurrentSpan()) instanceof EnabledTraceUtil.Span)
+        traceUtil.startSpan("foo", traceUtil.getCurrentSpan()) instanceof EnabledTraceUtil.Span)
         .isTrue();
   }