feat: Adding TraceUtil interface and its implementation to enable Tracing controls via DatastoreOptions

google-cloud-datastore/pom.xml

@@ -16,6 +16,7 @@
   </parent>
   <properties>
     <site.installationModule>google-cloud-datastore</site.installationModule>
+    <opentelemetry.version>1.37.0</opentelemetry.version>
   </properties>
   <dependencies>
     <dependency>
@@ -38,6 +39,10 @@
       <groupId>com.google.cloud.datastore</groupId>
       <artifactId>datastore-v1-proto-client</artifactId>
     </dependency>
+    <dependency>
+      <groupId>com.google.auth</groupId>
+      <artifactId>google-auth-library-credentials</artifactId>
+    </dependency>
     <dependency>
       <groupId>io.grpc</groupId>
       <artifactId>grpc-api</artifactId>
@@ -111,6 +116,19 @@
       <artifactId>jsr305</artifactId>
     </dependency>
 
+    <!-- OpenTelemetry -->
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-api</artifactId>
+      <version>${opentelemetry.version}</version>
+    </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-context</artifactId>
+      <version>${opentelemetry.version}</version>
+    </dependency>
+    <!-- END OpenTelemetry -->
+
     <!-- Test dependencies -->
     <dependency>
       <groupId>${project.groupId}</groupId>
@@ -160,6 +178,12 @@
       <version>1.4.2</version>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-sdk</artifactId>
+      <version>${opentelemetry.version}</version>
+      <scope>test</scope>
+    </dependency>
   </dependencies>
 
   <build>

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

@@ -0,0 +1,97 @@
+/*
+ * 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.datastore;
+
+import io.opentelemetry.api.OpenTelemetry;
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
+public class DatastoreOpenTelemetryOptions {
+  private final boolean enabled;
+  private final @Nullable OpenTelemetry openTelemetry;
+
+  DatastoreOpenTelemetryOptions(Builder builder) {
+    this.enabled = builder.enabled;
+    this.openTelemetry = builder.openTelemetry;
+  }
+
+  public boolean isEnabled() {
+    return enabled;
+  }
+
+  @Nullable
+  public OpenTelemetry getOpenTelemetry() {
+    return openTelemetry;
+  }
+
+  @Nonnull
+  public DatastoreOpenTelemetryOptions.Builder toBuilder() {
+    return new DatastoreOpenTelemetryOptions.Builder(this);
+  }
+
+  @Nonnull
+  public static DatastoreOpenTelemetryOptions.Builder newBuilder() {
+    return new DatastoreOpenTelemetryOptions.Builder();
+  }
+
+  public static class Builder {
+
+    private boolean enabled;
+
+    @Nullable private OpenTelemetry openTelemetry;
+
+    private Builder() {
+      enabled = false;
+      openTelemetry = null;
+    }
+
+    private Builder(DatastoreOpenTelemetryOptions options) {
+      this.enabled = options.enabled;
+      this.openTelemetry = options.openTelemetry;
+    }
+
+    @Nonnull
+    public DatastoreOpenTelemetryOptions build() {
+      return new DatastoreOpenTelemetryOptions(this);
+    }
+
+    /**
+     * Sets whether tracing should be enabled.
+     *
+     * @param enable Whether tracing should be enabled.
+     */
+    @Nonnull
+    public DatastoreOpenTelemetryOptions.Builder setTracingEnabled(boolean enable) {
+      this.enabled = enable;
+      return this;
+    }
+
+    /**
+     * Sets the {@link OpenTelemetry} to use with this Datastore instance. If telemetry collection
+     * is enabled, but an `OpenTelemetry` is not provided, the Datastore SDK will attempt to use the
+     * `GlobalOpenTelemetry`.
+     *
+     * @param openTelemetry The OpenTelemetry that should be used by this Datastore instance.
+     */
+    @Nonnull
+    public DatastoreOpenTelemetryOptions.Builder setOpenTelemetry(
+        @Nonnull OpenTelemetry openTelemetry) {
+      this.openTelemetry = openTelemetry;
+      return this;
+    }
+  }
+}

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.BetaApi;
 import com.google.cloud.ServiceDefaults;
 import com.google.cloud.ServiceOptions;
 import com.google.cloud.ServiceRpc;
@@ -31,6 +32,8 @@
 import java.lang.reflect.Method;
 import java.util.Objects;
 import java.util.Set;
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 
 public class DatastoreOptions extends ServiceOptions<Datastore, DatastoreOptions> {
 
@@ -43,6 +46,9 @@ public class DatastoreOptions extends ServiceOptions<Datastore, DatastoreOptions
   private final String namespace;
   private final String databaseId;
 
+  private final transient @Nonnull DatastoreOpenTelemetryOptions openTelemetryOptions;
+  private final transient @Nonnull com.google.cloud.datastore.telemetry.TraceUtil traceUtil;
+
   public static class DefaultDatastoreFactory implements DatastoreFactory {
 
     private static final DatastoreFactory INSTANCE = new DefaultDatastoreFactory();
@@ -63,17 +69,31 @@ public ServiceRpc create(DatastoreOptions options) {
     }
   }
 
+  @Nonnull
+  com.google.cloud.datastore.telemetry.TraceUtil getTraceUtil() {
+    return traceUtil;
+  }
+
+  @BetaApi
+  @Nonnull
+  public DatastoreOpenTelemetryOptions getOpenTelemetryOptions() {
+    return openTelemetryOptions;
+  }
+
   public static class Builder extends ServiceOptions.Builder<Datastore, DatastoreOptions, Builder> {
 
     private String namespace;
     private String databaseId;
 
+    @Nullable private DatastoreOpenTelemetryOptions openTelemetryOptions = null;
+
     private Builder() {}
 
     private Builder(DatastoreOptions options) {
       super(options);
       namespace = options.namespace;
       databaseId = options.databaseId;
+      this.openTelemetryOptions = options.openTelemetryOptions;
     }
 
     @Override
@@ -100,10 +120,30 @@ public Builder setDatabaseId(String databaseId) {
       this.databaseId = databaseId;
       return this;
     }
+
+    /**
+     * Sets the {@link DatastoreOpenTelemetryOptions} to be used for this Firestore instance.
+     *
+     * @param openTelemetryOptions The `DatastoreOpenTelemetryOptions` to use.
+     */
+    @BetaApi
+    @Nonnull
+    public Builder setOpenTelemetryOptions(
+        @Nonnull DatastoreOpenTelemetryOptions openTelemetryOptions) {
+      this.openTelemetryOptions = openTelemetryOptions;
+      return this;
+    }
   }
 
   private DatastoreOptions(Builder builder) {
     super(DatastoreFactory.class, DatastoreRpcFactory.class, builder, new DatastoreDefaults());
+
+    this.openTelemetryOptions =
+        builder.openTelemetryOptions != null
+            ? builder.openTelemetryOptions
+            : DatastoreOpenTelemetryOptions.newBuilder().build();
+    this.traceUtil = com.google.cloud.datastore.telemetry.TraceUtil.getInstance(this);
+
     namespace = MoreObjects.firstNonNull(builder.namespace, defaultNamespace());
     databaseId = MoreObjects.firstNonNull(builder.databaseId, DEFAULT_DATABASE_ID);
   }

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

@@ -0,0 +1,109 @@
+/*
+ * 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.datastore.telemetry;
+
+import com.google.api.core.ApiFunction;
+import com.google.api.core.ApiFuture;
+import com.google.api.core.InternalApi;
+import io.grpc.ManagedChannelBuilder;
+import java.util.Map;
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
+/**
+ * Tracing utility implementation, used to stub out tracing instrumentation when tracing is
+ * disabled.
+ */
+@InternalApi
+public class DisabledTraceUtil implements TraceUtil {
+
+  static class Span implements TraceUtil.Span {
+    @Override
+    public void end() {}
+
+    @Override
+    public void end(Throwable error) {}
+
+    @Override
+    public <T> void endAtFuture(ApiFuture<T> futureValue) {}
+
+    @Override
+    public TraceUtil.Span addEvent(String name) {
+      return this;
+    }
+
+    @Override
+    public TraceUtil.Span addEvent(String name, Map<String, Object> attributes) {
+      return this;
+    }
+
+    @Override
+    public TraceUtil.Span setAttribute(String key, int value) {
+      return this;
+    }
+
+    @Override
+    public TraceUtil.Span setAttribute(String key, String value) {
+      return this;
+    }
+
+    @Override
+    public Scope makeCurrent() {
+      return new Scope();
+    }
+  }
+
+  static class Context implements TraceUtil.Context {
+    @Override
+    public Scope makeCurrent() {
+      return new Scope();
+    }
+  }
+
+  static class Scope implements TraceUtil.Scope {
+    @Override
+    public void close() {}
+  }
+
+  @Nullable
+  @Override
+  public ApiFunction<ManagedChannelBuilder, ManagedChannelBuilder> getChannelConfigurator() {
+    return null;
+  }
+
+  @Override
+  public Span startSpan(String spanName) {
+    return new Span();
+  }
+
+  @Override
+  public TraceUtil.Span startSpan(String spanName, TraceUtil.Context parent) {
+    return new Span();
+  }
+
+  @Nonnull
+  @Override
+  public TraceUtil.Span currentSpan() {
+    return new Span();
+  }
+
+  @Nonnull
+  @Override
+  public TraceUtil.Context currentContext() {
+    return new Context();
+  }
+}