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

google-cloud-datastore/pom.xml

@@ -111,6 +111,14 @@
       <artifactId>jsr305</artifactId>
     </dependency>
 
+    <!-- OpenTelemetry -->
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-api</artifactId>
+      <version>1.37.0</version>
+    </dependency>
+    <!-- END OpenTelemetry -->
+
     <!-- Test dependencies -->
     <dependency>
       <groupId>${project.groupId}</groupId>
@@ -160,6 +168,46 @@
       <version>1.4.2</version>
       <scope>test</scope>
     </dependency>
+    <!-- OpenTelemetry -->
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-sdk</artifactId>
+      <version>1.37.0</version>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-sdk-testing</artifactId>
+      <version>1.37.0</version>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-sdk-trace</artifactId>
+      <version>1.37.0</version>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-sdk-common</artifactId>
+      <version>1.37.0</version>
+      <scope>test</scope>
+    </dependency>
+    <!-- END OpenTelemetry -->
+    <!-- Cloud Ops -->
+    <dependency>
+      <groupId>com.google.api.grpc</groupId>
+      <artifactId>proto-google-cloud-trace-v1</artifactId>
+      <version>1.3.0</version>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>com.google.cloud.opentelemetry</groupId>
+      <artifactId>exporter-trace</artifactId>
+      <version>0.15.0</version>
+      <scope>test</scope>
+    </dependency>
+    <!-- END OpenTelemetry -->
   </dependencies>
 
   <build>

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

@@ -0,0 +1,98 @@
+/*
+ * 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 getEnabled() {
+      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 Firestore instance. If telemetry collection
+       * is enabled, but an `OpenTelemetry` is not provided, the Firestore SDK will attempt to use
+       * the `GlobalOpenTelemetry`.
+       *
+       * @param openTelemetry The OpenTelemetry that should be used by this Firestore 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 @Nonnull DatastoreOpenTelemetryOptions openTelemetryOptions;
+  private final @Nonnull com.google.cloud.datastore.telemetry.TraceUtil traceUtil;
+
   public static class DefaultDatastoreFactory implements DatastoreFactory {
 
     private static final DatastoreFactory INSTANCE = new DefaultDatastoreFactory();
@@ -63,17 +69,32 @@ 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
@@ -87,6 +108,9 @@ public Builder setTransportOptions(TransportOptions transportOptions) {
 
     @Override
     public DatastoreOptions build() {
+      if (this.openTelemetryOptions == null) {
+        this.setOpenTelemetryOptions(DatastoreOpenTelemetryOptions.newBuilder().build());
+      }
       return new DatastoreOptions(this);
     }
 
@@ -100,10 +124,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,103 @@
+/*
+ * 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 io.grpc.ManagedChannelBuilder;
+import java.util.Map;
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
+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();
+  }
+}