diff --git a/google-cloud-datastore/pom.xml b/google-cloud-datastore/pom.xml
index b06df4c0f..f11be9f56 100644
--- a/google-cloud-datastore/pom.xml
+++ b/google-cloud-datastore/pom.xml
@@ -16,6 +16,7 @@
google-cloud-datastore
+ 1.37.0
@@ -38,6 +39,10 @@
com.google.cloud.datastore
datastore-v1-proto-client
+
+ com.google.auth
+ google-auth-library-credentials
+
io.grpc
grpc-api
@@ -111,6 +116,19 @@
jsr305
+
+
+ io.opentelemetry
+ opentelemetry-api
+ ${opentelemetry.version}
+
+
+ io.opentelemetry
+ opentelemetry-context
+ ${opentelemetry.version}
+
+
+
${project.groupId}
@@ -160,6 +178,12 @@
1.4.3
test
+
+ io.opentelemetry
+ opentelemetry-sdk
+ ${opentelemetry.version}
+ test
+
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreOpenTelemetryOptions.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreOpenTelemetryOptions.java
new file mode 100644
index 000000000..ac266562e
--- /dev/null
+++ b/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 enabled Whether tracing should be enabled.
+ */
+ @Nonnull
+ public DatastoreOpenTelemetryOptions.Builder setTracingEnabled(boolean enabled) {
+ this.enabled = enabled;
+ 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;
+ }
+ }
+}
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreOptions.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreOptions.java
index 8437c3e22..cef40eedd 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreOptions.java
+++ b/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 {
@@ -43,6 +46,9 @@ public class DatastoreOptions extends ServiceOptions {
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);
}
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/DisabledTraceUtil.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/DisabledTraceUtil.java
new file mode 100644
index 000000000..21321897d
--- /dev/null
+++ b/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 void endAtFuture(ApiFuture futureValue) {}
+
+ @Override
+ public TraceUtil.Span addEvent(String name) {
+ return this;
+ }
+
+ @Override
+ public TraceUtil.Span addEvent(String name, Map 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 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();
+ }
+}
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/EnabledTraceUtil.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/EnabledTraceUtil.java
new file mode 100644
index 000000000..8a5b0b36e
--- /dev/null
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/EnabledTraceUtil.java
@@ -0,0 +1,306 @@
+/*
+ * 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.ApiFutureCallback;
+import com.google.api.core.ApiFutures;
+import com.google.api.core.InternalApi;
+import com.google.cloud.datastore.DatastoreOptions;
+import com.google.common.base.Throwables;
+import io.grpc.ManagedChannelBuilder;
+import io.opentelemetry.api.GlobalOpenTelemetry;
+import io.opentelemetry.api.OpenTelemetry;
+import io.opentelemetry.api.common.Attributes;
+import io.opentelemetry.api.common.AttributesBuilder;
+import io.opentelemetry.api.trace.SpanBuilder;
+import io.opentelemetry.api.trace.SpanKind;
+import io.opentelemetry.api.trace.StatusCode;
+import io.opentelemetry.api.trace.Tracer;
+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.
+ */
+@InternalApi
+public class EnabledTraceUtil implements TraceUtil {
+ private final Tracer tracer;
+ private final OpenTelemetry openTelemetry;
+ private final DatastoreOptions datastoreOptions;
+
+ EnabledTraceUtil(DatastoreOptions datastoreOptions) {
+ OpenTelemetry openTelemetry = datastoreOptions.getOpenTelemetryOptions().getOpenTelemetry();
+
+ // If tracing is enabled, but an OpenTelemetry instance is not provided, fall back
+ // to using GlobalOpenTelemetry.
+ if (openTelemetry == null) {
+ openTelemetry = GlobalOpenTelemetry.get();
+ }
+
+ this.datastoreOptions = datastoreOptions;
+ this.openTelemetry = openTelemetry;
+ this.tracer = openTelemetry.getTracer(LIBRARY_NAME);
+ }
+
+ public OpenTelemetry getOpenTelemetry() {
+ return openTelemetry;
+ }
+
+ @Override
+ @Nullable
+ public ApiFunction getChannelConfigurator() {
+ // TODO(jimit) Update this to return a gRPC Channel Configurator after gRPC upgrade.
+ return null;
+ }
+
+ static class Span implements TraceUtil.Span {
+ private final io.opentelemetry.api.trace.Span span;
+ private final String spanName;
+
+ public Span(io.opentelemetry.api.trace.Span span, String spanName) {
+ this.span = span;
+ this.spanName = spanName;
+ }
+
+ /** Ends this span. */
+ @Override
+ public void end() {
+ span.end();
+ }
+
+ /** Ends this span in an error. */
+ @Override
+ public void end(Throwable error) {
+ span.setStatus(StatusCode.ERROR, error.getMessage());
+ span.recordException(
+ error,
+ Attributes.builder()
+ .put("exception.message", error.getMessage())
+ .put("exception.type", error.getClass().getName())
+ .put("exception.stacktrace", Throwables.getStackTraceAsString(error))
+ .build());
+ span.end();
+ }
+
+ /**
+ * If an operation ends in the future, its relevant span should end _after_ the future has been
+ * completed. This method "appends" the span completion code at the completion of the given
+ * future. In order for telemetry info to be recorded, the future returned by this method should
+ * be completed.
+ */
+ @Override
+ public void endAtFuture(ApiFuture futureValue) {
+ io.opentelemetry.context.Context asyncContext = io.opentelemetry.context.Context.current();
+ ApiFutures.addCallback(
+ futureValue,
+ new ApiFutureCallback() {
+ @Override
+ public void onFailure(Throwable t) {
+ try (io.opentelemetry.context.Scope scope = asyncContext.makeCurrent()) {
+ span.addEvent(spanName + " failed.");
+ end(t);
+ }
+ }
+
+ @Override
+ public void onSuccess(T result) {
+ try (io.opentelemetry.context.Scope scope = asyncContext.makeCurrent()) {
+ span.addEvent(spanName + " succeeded.");
+ end();
+ }
+ }
+ });
+ }
+
+ /** Adds the given event to this span. */
+ @Override
+ public TraceUtil.Span addEvent(String name) {
+ span.addEvent(name);
+ return this;
+ }
+
+ @Override
+ public TraceUtil.Span addEvent(String name, Map attributes) {
+ AttributesBuilder attributesBuilder = Attributes.builder();
+ attributes.forEach(
+ (key, value) -> {
+ if (value instanceof Integer) {
+ attributesBuilder.put(key, (int) value);
+ } else if (value instanceof Long) {
+ attributesBuilder.put(key, (long) value);
+ } else if (value instanceof Double) {
+ attributesBuilder.put(key, (double) value);
+ } else if (value instanceof Float) {
+ attributesBuilder.put(key, (float) value);
+ } else if (value instanceof Boolean) {
+ attributesBuilder.put(key, (boolean) value);
+ } else if (value instanceof String) {
+ attributesBuilder.put(key, (String) value);
+ } else {
+ // OpenTelemetry APIs do not support any other type.
+ throw new IllegalArgumentException(
+ "Unknown attribute type:" + value.getClass().getSimpleName());
+ }
+ });
+ span.addEvent(name, attributesBuilder.build());
+ return this;
+ }
+
+ @Override
+ public TraceUtil.Span setAttribute(String key, int value) {
+ span.setAttribute(ATTRIBUTE_SERVICE_PREFIX + key, value);
+ return this;
+ }
+
+ @Override
+ public TraceUtil.Span setAttribute(String key, String value) {
+ span.setAttribute(ATTRIBUTE_SERVICE_PREFIX + key, value);
+ return this;
+ }
+
+ @Override
+ public Scope makeCurrent() {
+ try (io.opentelemetry.context.Scope scope = span.makeCurrent()) {
+ return new Scope(scope);
+ }
+ }
+ }
+
+ static class Scope implements TraceUtil.Scope {
+ private final io.opentelemetry.context.Scope scope;
+
+ Scope(io.opentelemetry.context.Scope scope) {
+ this.scope = scope;
+ }
+
+ @Override
+ public void close() {
+ scope.close();
+ }
+ }
+
+ static class Context implements TraceUtil.Context {
+ private final io.opentelemetry.context.Context context;
+
+ Context(io.opentelemetry.context.Context context) {
+ this.context = context;
+ }
+
+ @Override
+ public Scope makeCurrent() {
+ try (io.opentelemetry.context.Scope scope = context.makeCurrent()) {
+ return new Scope(scope);
+ }
+ }
+ }
+
+ /** Applies the current Datastore instance settings as attributes to the current Span */
+ private SpanBuilder addSettingsAttributesToCurrentSpan(SpanBuilder spanBuilder) {
+ spanBuilder =
+ spanBuilder.setAllAttributes(
+ Attributes.builder()
+ .put(
+ ATTRIBUTE_SERVICE_PREFIX + "settings.databaseId",
+ datastoreOptions.getDatabaseId())
+ .put(ATTRIBUTE_SERVICE_PREFIX + "settings.host", datastoreOptions.getHost())
+ .build());
+
+ if (datastoreOptions.getCredentials() != null) {
+ spanBuilder =
+ spanBuilder.setAttribute(
+ ATTRIBUTE_SERVICE_PREFIX + "settings.credentials.authenticationType",
+ datastoreOptions.getCredentials().getAuthenticationType());
+ }
+
+ if (datastoreOptions.getRetrySettings() != null) {
+ spanBuilder =
+ spanBuilder.setAllAttributes(
+ Attributes.builder()
+ .put(
+ ATTRIBUTE_SERVICE_PREFIX + "settings.retrySettings.initialRetryDelay",
+ datastoreOptions.getRetrySettings().getInitialRetryDelay().toString())
+ .put(
+ ATTRIBUTE_SERVICE_PREFIX + "settings.retrySettings.maxRetryDelay",
+ datastoreOptions.getRetrySettings().getMaxRetryDelay().toString())
+ .put(
+ ATTRIBUTE_SERVICE_PREFIX + "settings.retrySettings.retryDelayMultiplier",
+ String.valueOf(datastoreOptions.getRetrySettings().getRetryDelayMultiplier()))
+ .put(
+ ATTRIBUTE_SERVICE_PREFIX + "settings.retrySettings.maxAttempts",
+ String.valueOf(datastoreOptions.getRetrySettings().getMaxAttempts()))
+ .put(
+ ATTRIBUTE_SERVICE_PREFIX + "settings.retrySettings.initialRpcTimeout",
+ datastoreOptions.getRetrySettings().getInitialRpcTimeout().toString())
+ .put(
+ ATTRIBUTE_SERVICE_PREFIX + "settings.retrySettings.maxRpcTimeout",
+ datastoreOptions.getRetrySettings().getMaxRpcTimeout().toString())
+ .put(
+ ATTRIBUTE_SERVICE_PREFIX + "settings.retrySettings.rpcTimeoutMultiplier",
+ String.valueOf(datastoreOptions.getRetrySettings().getRpcTimeoutMultiplier()))
+ .put(
+ ATTRIBUTE_SERVICE_PREFIX + "settings.retrySettings.totalTimeout",
+ datastoreOptions.getRetrySettings().getTotalTimeout().toString())
+ .build());
+ }
+
+ // Add the memory utilization of the client at the time this trace was collected.
+ long totalMemory = Runtime.getRuntime().totalMemory();
+ long freeMemory = Runtime.getRuntime().freeMemory();
+ double memoryUtilization = ((double) (totalMemory - freeMemory)) / totalMemory;
+ spanBuilder.setAttribute(
+ ATTRIBUTE_SERVICE_PREFIX + "memoryUtilization",
+ String.format("%.2f", memoryUtilization * 100) + "%");
+
+ return spanBuilder;
+ }
+
+ @Override
+ public Span startSpan(String spanName) {
+ SpanBuilder spanBuilder = tracer.spanBuilder(spanName).setSpanKind(SpanKind.PRODUCER);
+ io.opentelemetry.api.trace.Span span =
+ addSettingsAttributesToCurrentSpan(spanBuilder).startSpan();
+ return new Span(span, spanName);
+ }
+
+ @Override
+ public TraceUtil.Span startSpan(String spanName, TraceUtil.Context parent) {
+ assert (parent instanceof Context);
+ SpanBuilder spanBuilder =
+ tracer
+ .spanBuilder(spanName)
+ .setSpanKind(SpanKind.PRODUCER)
+ .setParent(((Context) parent).context);
+ io.opentelemetry.api.trace.Span span =
+ addSettingsAttributesToCurrentSpan(spanBuilder).startSpan();
+ return new Span(span, spanName);
+ }
+
+ @Nonnull
+ @Override
+ public TraceUtil.Span currentSpan() {
+ return new Span(io.opentelemetry.api.trace.Span.current(), "");
+ }
+
+ @Nonnull
+ @Override
+ public TraceUtil.Context currentContext() {
+ return new Context(io.opentelemetry.context.Context.current());
+ }
+}
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java
new file mode 100644
index 000000000..4e55a1ff6
--- /dev/null
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java
@@ -0,0 +1,129 @@
+/*
+ * 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.InternalExtensionOnly;
+import com.google.cloud.datastore.DatastoreOptions;
+import io.grpc.ManagedChannelBuilder;
+import java.util.Map;
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
+/** Utility interface to manage OpenTelemetry tracing instrumentation based on the configuration. */
+@InternalExtensionOnly
+public interface TraceUtil {
+ static final String ATTRIBUTE_SERVICE_PREFIX = "gcp.datastore.";
+ static final String ENABLE_TRACING_ENV_VAR = "DATASTORE_ENABLE_TRACING";
+ static final String LIBRARY_NAME = "com.google.cloud.datastore";
+
+ /**
+ * Creates and returns an instance of the TraceUtil class.
+ *
+ * @param datastoreOptions The DatastoreOptions object that is requesting an instance of
+ * TraceUtil.
+ * @return An instance of the TraceUtil class.
+ */
+ static TraceUtil getInstance(@Nonnull DatastoreOptions datastoreOptions) {
+ boolean createEnabledInstance = datastoreOptions.getOpenTelemetryOptions().isEnabled();
+
+ // The environment variable can override options to enable/disable telemetry collection.
+ String enableTracingEnvVar = System.getenv(ENABLE_TRACING_ENV_VAR);
+ if (enableTracingEnvVar != null) {
+ if (enableTracingEnvVar.equalsIgnoreCase("true")
+ || enableTracingEnvVar.equalsIgnoreCase("on")) {
+ createEnabledInstance = true;
+ }
+ if (enableTracingEnvVar.equalsIgnoreCase("false")
+ || enableTracingEnvVar.equalsIgnoreCase("off")) {
+ createEnabledInstance = false;
+ }
+ }
+
+ if (createEnabledInstance) {
+ return new EnabledTraceUtil(datastoreOptions);
+ } else {
+ return new DisabledTraceUtil();
+ }
+ }
+
+ /** Returns a channel configurator for gRPC, or {@code null} if tracing is disabled. */
+ @Nullable
+ ApiFunction getChannelConfigurator();
+
+ /** Represents a trace span. */
+ interface Span {
+ /** Adds the given event to this span. */
+ Span addEvent(String name);
+
+ /** Adds the given event with the given attributes to this span. */
+ Span addEvent(String name, Map attributes);
+
+ /** Adds the given attribute to this span. */
+ Span setAttribute(String key, int value);
+
+ /** Adds the given attribute to this span. */
+ Span setAttribute(String key, String value);
+
+ /** Marks this span as the current span. */
+ Scope makeCurrent();
+
+ /** Ends this span. */
+ void end();
+
+ /** Ends this span in an error. */
+ void end(Throwable error);
+
+ /**
+ * If an operation ends in the future, its relevant span should end _after_ the future has been
+ * completed. This method "appends" the span completion code at the completion of the given
+ * future. In order for telemetry info to be recorded, the future returned by this method should
+ * be completed.
+ */
+ void endAtFuture(ApiFuture futureValue);
+ }
+
+ /** Represents a trace context. */
+ interface Context {
+ /** Makes this context the current context. */
+ Scope makeCurrent();
+ }
+
+ /** Represents a trace scope. */
+ interface Scope extends AutoCloseable {
+ /** Closes the current scope. */
+ void close();
+ }
+
+ /** Starts a new span with the given name, sets it as the current span, and returns it. */
+ Span startSpan(String spanName);
+
+ /**
+ * Starts a new span with the given name and the given context as its parent, sets it as the
+ * current span, and returns it.
+ */
+ Span startSpan(String spanName, Context parent);
+
+ /** Returns the current span. */
+ @Nonnull
+ Span currentSpan();
+
+ /** Returns the current Context. */
+ @Nonnull
+ Context currentContext();
+}
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreOptionsTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreOptionsTest.java
index a545580e2..85703f739 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreOptionsTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreOptionsTest.java
@@ -70,6 +70,20 @@ public void testHost() {
assertEquals("http://localhost:" + PORT, options.build().getHost());
}
+ @Test
+ public void testOpenTelemetryOptionsEnabled() {
+ options.setOpenTelemetryOptions(
+ DatastoreOpenTelemetryOptions.newBuilder().setTracingEnabled(true).build());
+ assertTrue(options.build().getOpenTelemetryOptions().isEnabled());
+ }
+
+ @Test
+ public void testOpenTelemetryOptionsDisabled() {
+ options.setOpenTelemetryOptions(
+ DatastoreOpenTelemetryOptions.newBuilder().setTracingEnabled(false).build());
+ assertTrue(!options.build().getOpenTelemetryOptions().isEnabled());
+ }
+
@Test
public void testNamespace() {
assertTrue(options.build().getNamespace().isEmpty());
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/DisabledTraceUtilTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/DisabledTraceUtilTest.java
new file mode 100644
index 000000000..0f3f183cd
--- /dev/null
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/DisabledTraceUtilTest.java
@@ -0,0 +1,53 @@
+/*
+ * 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 static com.google.common.truth.Truth.assertThat;
+
+import org.junit.Test;
+
+public class DisabledTraceUtilTest {
+ @Test
+ public void disabledTraceUtilDoesNotProvideChannelConfigurator() {
+ DisabledTraceUtil traceUtil = new DisabledTraceUtil();
+ assertThat(traceUtil.getChannelConfigurator()).isNull();
+ }
+
+ @Test
+ public void usesDisabledContext() {
+ DisabledTraceUtil traceUtil = new DisabledTraceUtil();
+ assertThat(traceUtil.currentContext() instanceof DisabledTraceUtil.Context).isTrue();
+ }
+
+ @Test
+ public void usesDisabledSpan() {
+ DisabledTraceUtil traceUtil = new DisabledTraceUtil();
+ assertThat(traceUtil.currentSpan() instanceof DisabledTraceUtil.Span).isTrue();
+ assertThat(traceUtil.startSpan("foo") instanceof DisabledTraceUtil.Span).isTrue();
+ assertThat(
+ traceUtil.startSpan("foo", traceUtil.currentContext())
+ instanceof DisabledTraceUtil.Span)
+ .isTrue();
+ }
+
+ @Test
+ public void usesDisabledScope() {
+ DisabledTraceUtil traceUtil = new DisabledTraceUtil();
+ assertThat(traceUtil.currentContext().makeCurrent() instanceof DisabledTraceUtil.Scope)
+ .isTrue();
+ assertThat(traceUtil.currentSpan().makeCurrent() instanceof DisabledTraceUtil.Scope).isTrue();
+ }
+}
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/EnabledTraceUtilTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/EnabledTraceUtilTest.java
new file mode 100644
index 000000000..e88e1a849
--- /dev/null
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/EnabledTraceUtilTest.java
@@ -0,0 +1,104 @@
+/*
+ * 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 static com.google.common.truth.Truth.assertThat;
+
+import com.google.cloud.NoCredentials;
+import com.google.cloud.datastore.DatastoreOpenTelemetryOptions;
+import com.google.cloud.datastore.DatastoreOptions;
+import io.opentelemetry.api.GlobalOpenTelemetry;
+import io.opentelemetry.sdk.OpenTelemetrySdk;
+import org.junit.Before;
+import org.junit.Test;
+
+public class EnabledTraceUtilTest {
+ @Before
+ public void setUp() {
+ GlobalOpenTelemetry.resetForTest();
+ }
+
+ DatastoreOptions.Builder getBaseOptions() {
+ return DatastoreOptions.newBuilder()
+ .setProjectId("test-project")
+ .setCredentials(NoCredentials.getInstance());
+ }
+
+ DatastoreOptions getTracingEnabledOptions() {
+ return getBaseOptions()
+ .setOpenTelemetryOptions(
+ DatastoreOpenTelemetryOptions.newBuilder().setTracingEnabled(true).build())
+ .build();
+ }
+
+ EnabledTraceUtil newEnabledTraceUtil() {
+ return new EnabledTraceUtil(getTracingEnabledOptions());
+ }
+
+ @Test
+ public void usesOpenTelemetryFromOptions() {
+ OpenTelemetrySdk myOpenTelemetrySdk = OpenTelemetrySdk.builder().build();
+ DatastoreOptions firestoreOptions =
+ getBaseOptions()
+ .setOpenTelemetryOptions(
+ DatastoreOpenTelemetryOptions.newBuilder()
+ .setTracingEnabled(true)
+ .setOpenTelemetry(myOpenTelemetrySdk)
+ .build())
+ .build();
+ EnabledTraceUtil traceUtil = new EnabledTraceUtil(firestoreOptions);
+ assertThat(traceUtil.getOpenTelemetry()).isEqualTo(myOpenTelemetrySdk);
+ }
+
+ @Test
+ public void usesGlobalOpenTelemetryIfOpenTelemetryInstanceNotProvided() {
+ OpenTelemetrySdk globalOpenTelemetrySdk = OpenTelemetrySdk.builder().buildAndRegisterGlobal();
+ DatastoreOptions firestoreOptions =
+ getBaseOptions()
+ .setOpenTelemetryOptions(
+ DatastoreOpenTelemetryOptions.newBuilder().setTracingEnabled(true).build())
+ .build();
+ EnabledTraceUtil traceUtil = new EnabledTraceUtil(firestoreOptions);
+ assertThat(traceUtil.getOpenTelemetry()).isEqualTo(GlobalOpenTelemetry.get());
+ }
+
+ @Test
+ public void enabledTraceUtilProvidesChannelConfigurator() {
+ assertThat(newEnabledTraceUtil().getChannelConfigurator()).isNull();
+ }
+
+ @Test
+ public void usesEnabledContext() {
+ assertThat(newEnabledTraceUtil().currentContext() instanceof EnabledTraceUtil.Context).isTrue();
+ }
+
+ @Test
+ public void usesEnabledSpan() {
+ EnabledTraceUtil traceUtil = newEnabledTraceUtil();
+ assertThat(traceUtil.currentSpan() instanceof EnabledTraceUtil.Span).isTrue();
+ assertThat(traceUtil.startSpan("foo") != null).isTrue();
+ assertThat(
+ traceUtil.startSpan("foo", traceUtil.currentContext()) instanceof EnabledTraceUtil.Span)
+ .isTrue();
+ }
+
+ @Test
+ public void usesEnabledScope() {
+ EnabledTraceUtil traceUtil = newEnabledTraceUtil();
+ assertThat(traceUtil.currentContext().makeCurrent() instanceof EnabledTraceUtil.Scope).isTrue();
+ assertThat(traceUtil.currentSpan().makeCurrent() instanceof EnabledTraceUtil.Scope).isTrue();
+ }
+}
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/TraceUtilTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/TraceUtilTest.java
new file mode 100644
index 000000000..f1cce8006
--- /dev/null
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/TraceUtilTest.java
@@ -0,0 +1,62 @@
+/*
+ * 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 static com.google.common.truth.Truth.assertThat;
+
+import com.google.cloud.NoCredentials;
+import com.google.cloud.datastore.DatastoreOpenTelemetryOptions;
+import com.google.cloud.datastore.DatastoreOptions;
+import org.junit.Test;
+
+public class TraceUtilTest {
+ @Test
+ public void defaultOptionsUseDisabledTraceUtil() {
+ TraceUtil traceUtil =
+ TraceUtil.getInstance(
+ DatastoreOptions.newBuilder()
+ .setProjectId("test-project")
+ .setCredentials(NoCredentials.getInstance())
+ .build());
+ assertThat(traceUtil instanceof DisabledTraceUtil).isTrue();
+ }
+
+ @Test
+ public void tracingDisabledOptionsUseDisabledTraceUtil() {
+ TraceUtil traceUtil =
+ TraceUtil.getInstance(
+ DatastoreOptions.newBuilder()
+ .setProjectId("test-project")
+ .setCredentials(NoCredentials.getInstance())
+ .setOpenTelemetryOptions(
+ DatastoreOpenTelemetryOptions.newBuilder().setTracingEnabled(false).build())
+ .build());
+ assertThat(traceUtil instanceof DisabledTraceUtil).isTrue();
+ }
+
+ @Test
+ public void tracingEnabledOptionsUseEnabledTraceUtil() {
+ TraceUtil traceUtil =
+ TraceUtil.getInstance(
+ DatastoreOptions.newBuilder()
+ .setProjectId("test-project")
+ .setCredentials(NoCredentials.getInstance())
+ .setOpenTelemetryOptions(
+ DatastoreOpenTelemetryOptions.newBuilder().setTracingEnabled(true).build())
+ .build());
+ assertThat(traceUtil instanceof EnabledTraceUtil).isTrue();
+ }
+}