() {
@Override
@@ -758,10 +757,10 @@ public Void call() throws DatastoreException {
EXCEPTION_HANDLER,
getOptions().getClock());
} catch (RetryHelperException e) {
- span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
+ span.end(e);
throw DatastoreException.translateAndThrow(e);
} finally {
- span.end(TraceUtil.END_SPAN_OPTIONS);
+ span.end();
}
}
}
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecorator.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecorator.java
index 7abf1d360..e1ea6e8dd 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecorator.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecorator.java
@@ -22,6 +22,7 @@
import com.google.cloud.RetryHelper;
import com.google.cloud.RetryHelper.RetryHelperException;
import com.google.cloud.datastore.spi.v1.DatastoreRpc;
+import com.google.cloud.datastore.telemetry.TraceUtil;
import com.google.datastore.v1.AllocateIdsRequest;
import com.google.datastore.v1.AllocateIdsResponse;
import com.google.datastore.v1.BeginTransactionRequest;
@@ -55,13 +56,13 @@ public class RetryAndTraceDatastoreRpcDecorator implements DatastoreRpc {
public RetryAndTraceDatastoreRpcDecorator(
DatastoreRpc datastoreRpc,
- TraceUtil traceUtil,
+ TraceUtil otelTraceUtil,
RetrySettings retrySettings,
DatastoreOptions datastoreOptions) {
this.datastoreRpc = datastoreRpc;
this.retrySettings = retrySettings;
this.datastoreOptions = datastoreOptions;
- this.otelTraceUtil = datastoreOptions.getTraceUtil();
+ this.otelTraceUtil = otelTraceUtil;
}
@Override
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TraceUtil.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TraceUtil.java
deleted file mode 100644
index 6b51d6a87..000000000
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TraceUtil.java
+++ /dev/null
@@ -1,75 +0,0 @@
-/*
- * Copyright 2020 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 com.google.cloud.datastore.spi.v1.HttpDatastoreRpc;
-import io.opencensus.trace.EndSpanOptions;
-import io.opencensus.trace.Span;
-import io.opencensus.trace.Tracer;
-import io.opencensus.trace.Tracing;
-
-/**
- * Helper class for tracing utility. It is used for instrumenting {@link HttpDatastoreRpc} with
- * OpenCensus APIs.
- *
- * TraceUtil instances are created by the {@link TraceUtil#getInstance()} method.
- */
-public class TraceUtil {
- private final Tracer tracer = Tracing.getTracer();
- private static final TraceUtil traceUtil = new TraceUtil();
- static final String SPAN_NAME_TRANSACTION = "CloudDatastoreOperation.readWriteTransaction";
- static final String SPAN_NAME_BEGINTRANSACTION = "CloudDatastoreOperation.beginTransaction";
- static final String SPAN_NAME_COMMIT = "CloudDatastoreOperation.commit";
- static final String SPAN_NAME_LOOKUP = "CloudDatastoreOperation.lookup";
- static final String SPAN_NAME_RESERVEIDS = "CloudDatastoreOperation.reserveIds";
- static final String SPAN_NAME_ROLLBACK = "CloudDatastoreOperation.rollback";
- static final String SPAN_NAME_RUNQUERY = "CloudDatastoreOperation.runQuery";
- static final String SPAN_NAME_RUN_AGGREGATION_QUERY =
- "CloudDatastoreOperation.runAggregationQuery";
- static final EndSpanOptions END_SPAN_OPTIONS =
- EndSpanOptions.builder().setSampleToLocalSpanStore(true).build();
-
- /**
- * Starts a new span.
- *
- * @param spanName The name of the returned Span.
- * @return The newly created {@link Span}.
- */
- protected Span startSpan(String spanName) {
- return tracer.spanBuilder(spanName).startSpan();
- }
-
- /**
- * Return the global {@link Tracer}.
- *
- * @return The global {@link Tracer}.
- */
- public Tracer getTracer() {
- return tracer;
- }
-
- /**
- * Return TraceUtil Object.
- *
- * @return An instance of {@link TraceUtil}
- */
- public static TraceUtil getInstance() {
- return traceUtil;
- }
-
- private TraceUtil() {}
-}
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/HttpDatastoreRpc.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/HttpDatastoreRpc.java
index fd3cdc658..d927a2d7f 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/HttpDatastoreRpc.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/HttpDatastoreRpc.java
@@ -21,7 +21,6 @@
import com.google.api.client.http.HttpTransport;
import com.google.cloud.datastore.DatastoreException;
import com.google.cloud.datastore.DatastoreOptions;
-import com.google.cloud.datastore.TraceUtil;
import com.google.cloud.http.CensusHttpModule;
import com.google.cloud.http.HttpTransportOptions;
import com.google.datastore.v1.AllocateIdsRequest;
@@ -40,6 +39,7 @@
import com.google.datastore.v1.RunAggregationQueryResponse;
import com.google.datastore.v1.RunQueryRequest;
import com.google.datastore.v1.RunQueryResponse;
+import io.opencensus.trace.Tracing;
import java.io.IOException;
import java.net.InetAddress;
import java.net.URL;
@@ -80,8 +80,7 @@ public HttpDatastoreRpc(DatastoreOptions options) {
private HttpRequestInitializer getHttpRequestInitializer(
final DatastoreOptions options, HttpTransportOptions httpTransportOptions) {
// Open Census initialization
- CensusHttpModule censusHttpModule =
- new CensusHttpModule(TraceUtil.getInstance().getTracer(), true);
+ CensusHttpModule censusHttpModule = new CensusHttpModule(Tracing.getTracer(), true);
final HttpRequestInitializer censusHttpModuleHttpRequestInitializer =
censusHttpModule.getHttpRequestInitializer(
httpTransportOptions.getHttpRequestInitializer(options));
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java
index 34c7ea858..b8174184e 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java
@@ -22,6 +22,7 @@
import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_COMMIT;
import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_LOOKUP;
import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RESERVE_IDS;
+import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_ROLLBACK;
import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RUN_AGGREGATION_QUERY;
import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RUN_QUERY;
import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_COMMIT;
@@ -374,6 +375,7 @@ public void after() throws Exception {
tracer = null;
retrievedTrace = null;
customSpanContext = null;
+ openTelemetrySdk = null;
}
@AfterClass
@@ -793,7 +795,7 @@ public void runAggregationQueryTraceTest() throws Exception {
}
@Test
- public void transactionalLookupTest() throws Exception {
+ public void newTransactionReadTest() throws Exception {
assertNotNull(customSpanContext);
Span rootSpan = getNewRootSpanWithContext();
@@ -817,7 +819,7 @@ public void transactionalLookupTest() throws Exception {
}
@Test
- public void transactionQueryTest() throws Exception {
+ public void newTransactionQueryTest() throws Exception {
// Set up
Entity entity1 = Entity.newBuilder(KEY1).set("test_field", "test_value1").build();
Entity entity2 = Entity.newBuilder(KEY2).set("test_field", "test_value2").build();
@@ -856,6 +858,116 @@ public void transactionQueryTest() throws Exception {
Collections.singletonList(SPAN_NAME_TRANSACTION_COMMIT)));
}
+ @Test
+ public void newTransactionReadWriteTraceTest() throws Exception {
+ // Set up
+ Entity entity1 = Entity.newBuilder(KEY1).set("pepper_type", "jalapeno").build();
+ Entity entity2 = Entity.newBuilder(KEY2).set("pepper_type", "habanero").build();
+ List entityList = new ArrayList<>();
+ entityList.add(entity1);
+ entityList.add(entity2);
+
+ List response = datastore.add(entity1, entity2);
+ assertEquals(entityList, response);
+
+ String simplified_spice_level = "not_spicy";
+ Entity entity1update =
+ Entity.newBuilder(entity1).set("spice_level", simplified_spice_level).build();
+
+ assertNotNull(customSpanContext);
+
+ // Test
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ Transaction transaction = datastore.newTransaction();
+ entity1 = transaction.get(KEY1);
+ switch (entity1.getString("pepper_type")) {
+ case "jalapeno":
+ simplified_spice_level = "mild";
+ break;
+
+ case "habanero":
+ simplified_spice_level = "hot";
+ break;
+ }
+ transaction.update(entity1update);
+ transaction.delete(KEY2);
+ transaction.commit();
+ assertFalse(transaction.isActive());
+ } finally {
+ rootSpan.end();
+ }
+
+ waitForTracesToComplete();
+
+ List list = datastore.fetch(KEY1, KEY2);
+ assertEquals(list.get(0), entity1update);
+ assertNull(list.get(1));
+
+ fetchAndValidateTrace(
+ customSpanContext.getTraceId(),
+ /*numExpectedSpans=*/ 3,
+ Arrays.asList(
+ Collections.singletonList(SPAN_NAME_BEGIN_TRANSACTION),
+ Collections.singletonList(SPAN_NAME_TRANSACTION_LOOKUP),
+ Collections.singletonList(SPAN_NAME_TRANSACTION_COMMIT)));
+ }
+
+ @Test
+ public void newTransactionRollbackTest() throws Exception {
+ // Set up
+ Entity entity1 = Entity.newBuilder(KEY1).set("pepper_type", "jalapeno").build();
+ Entity entity2 = Entity.newBuilder(KEY2).set("pepper_type", "habanero").build();
+ List entityList = new ArrayList<>();
+ entityList.add(entity1);
+ entityList.add(entity2);
+
+ List response = datastore.add(entity1, entity2);
+ assertEquals(entityList, response);
+
+ String simplified_spice_level = "not_spicy";
+ Entity entity1update =
+ Entity.newBuilder(entity1).set("spice_level", simplified_spice_level).build();
+
+ assertNotNull(customSpanContext);
+
+ // Test
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ Transaction transaction = datastore.newTransaction();
+ entity1 = transaction.get(KEY1);
+ switch (entity1.getString("pepper_type")) {
+ case "jalapeno":
+ simplified_spice_level = "mild";
+ break;
+
+ case "habanero":
+ simplified_spice_level = "hot";
+ break;
+ }
+ transaction.update(entity1update);
+ transaction.delete(KEY2);
+ transaction.rollback();
+ assertFalse(transaction.isActive());
+ } finally {
+ rootSpan.end();
+ }
+
+ waitForTracesToComplete();
+
+ List list = datastore.fetch(KEY1, KEY2);
+ assertEquals(list.get(0), entity1);
+ assertEquals(list.get(1), entity2);
+
+ fetchAndValidateTrace(
+ customSpanContext.getTraceId(),
+ /*numExpectedSpans=*/ 3,
+ Arrays.asList(
+ Collections.singletonList(SPAN_NAME_BEGIN_TRANSACTION),
+ Collections.singletonList(SPAN_NAME_TRANSACTION_LOOKUP),
+ Collections.singletonList(SPAN_NAME_ROLLBACK)));
+ }
+
@Test
public void runInTransactionQueryTest() throws Exception {
// Set up
@@ -897,10 +1009,4 @@ public void runInTransactionQueryTest() throws Exception {
Arrays.asList(SPAN_NAME_TRANSACTION_RUN, SPAN_NAME_RUN_QUERY),
Arrays.asList(SPAN_NAME_TRANSACTION_RUN, SPAN_NAME_TRANSACTION_COMMIT)));
}
-
- @Test
- public void transactionRunQueryTest() throws Exception {}
-
- @Test
- public void readWriteTransactionTraceTest() throws Exception {}
}