diff --git a/google-cloud-datastore/pom.xml b/google-cloud-datastore/pom.xml index 6de85518f..703612e6f 100644 --- a/google-cloud-datastore/pom.xml +++ b/google-cloud-datastore/pom.xml @@ -78,6 +78,10 @@ com.google.auth google-auth-library-oauth2-http + + io.opencensus + opencensus-api + diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java index 25e434227..f4bd6b0be 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java @@ -34,6 +34,9 @@ import com.google.datastore.v1.ReserveIdsRequest; import com.google.datastore.v1.TransactionOptions; import com.google.protobuf.ByteString; +import io.opencensus.common.Scope; +import io.opencensus.trace.Span; +import io.opencensus.trace.Status; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; @@ -53,6 +56,7 @@ final class DatastoreImpl extends BaseService implements Datas TransactionExceptionHandler.build(); private static final ExceptionHandler TRANSACTION_OPERATION_EXCEPTION_HANDLER = TransactionOperationExceptionHandler.build(); + private final TraceUtil traceUtil = TraceUtil.getInstance();; DatastoreImpl(DatastoreOptions options) { super(options); @@ -132,30 +136,36 @@ public T call() throws DatastoreException { @Override public T runInTransaction(final TransactionCallable callable) { - final DatastoreImpl self = this; - try { + Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_TRANSACTION); + try (Scope scope = traceUtil.getTracer().withSpan(span)) { return RetryHelper.runWithRetries( - new ReadWriteTransactionCallable(self, callable, null), + new ReadWriteTransactionCallable(this, callable, null), retrySettings, TRANSACTION_EXCEPTION_HANDLER, getOptions().getClock()); } catch (RetryHelperException e) { + span.setStatus(Status.UNKNOWN.withDescription(e.getMessage())); throw DatastoreException.translateAndThrow(e); + } finally { + span.end(TraceUtil.END_SPAN_OPTIONS); } } @Override public T runInTransaction( final TransactionCallable callable, TransactionOptions transactionOptions) { - final DatastoreImpl self = this; - try { + Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_TRANSACTION); + try (Scope scope = traceUtil.getTracer().withSpan(span)) { return RetryHelper.runWithRetries( - new ReadWriteTransactionCallable(self, callable, transactionOptions), + new ReadWriteTransactionCallable(this, callable, transactionOptions), retrySettings, TRANSACTION_EXCEPTION_HANDLER, getOptions().getClock()); } catch (RetryHelperException e) { + span.setStatus(Status.UNKNOWN.withDescription(e.getMessage())); throw DatastoreException.translateAndThrow(e); + } finally { + span.end(TraceUtil.END_SPAN_OPTIONS); } } @@ -175,7 +185,8 @@ QueryResults run(com.google.datastore.v1.ReadOptions readOptionsPb, Query com.google.datastore.v1.RunQueryResponse runQuery( final com.google.datastore.v1.RunQueryRequest requestPb) { - try { + Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_RUNQUERY); + try (Scope scope = traceUtil.getTracer().withSpan(span)) { return RetryHelper.runWithRetries( new Callable() { @Override @@ -189,7 +200,10 @@ public com.google.datastore.v1.RunQueryResponse call() throws DatastoreException : TRANSACTION_OPERATION_EXCEPTION_HANDLER, getOptions().getClock()); } catch (RetryHelperException e) { + span.setStatus(Status.UNKNOWN.withDescription(e.getMessage())); throw DatastoreException.translateAndThrow(e); + } finally { + span.end(TraceUtil.END_SPAN_OPTIONS); } } @@ -229,7 +243,8 @@ public List allocateId(IncompleteKey... keys) { private com.google.datastore.v1.AllocateIdsResponse allocateIds( final com.google.datastore.v1.AllocateIdsRequest requestPb) { - try { + Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_ALLOCATEIDS); + try (Scope scope = traceUtil.getTracer().withSpan(span)) { return RetryHelper.runWithRetries( new Callable() { @Override @@ -241,7 +256,10 @@ public com.google.datastore.v1.AllocateIdsResponse call() throws DatastoreExcept EXCEPTION_HANDLER, getOptions().getClock()); } catch (RetryHelperException e) { + span.setStatus(Status.UNKNOWN.withDescription(e.getMessage())); throw DatastoreException.translateAndThrow(e); + } finally { + span.end(TraceUtil.END_SPAN_OPTIONS); } } @@ -389,7 +407,8 @@ protected Entity computeNext() { com.google.datastore.v1.LookupResponse lookup( final com.google.datastore.v1.LookupRequest requestPb) { - try { + Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_LOOKUP); + try (Scope scope = traceUtil.getTracer().withSpan(span)) { return RetryHelper.runWithRetries( new Callable() { @Override @@ -403,7 +422,10 @@ public com.google.datastore.v1.LookupResponse call() throws DatastoreException { : TRANSACTION_OPERATION_EXCEPTION_HANDLER, getOptions().getClock()); } catch (RetryHelperException e) { + span.setStatus(Status.UNKNOWN.withDescription(e.getMessage())); throw DatastoreException.translateAndThrow(e); + } finally { + span.end(TraceUtil.END_SPAN_OPTIONS); } } @@ -425,7 +447,8 @@ public List reserveIds(Key... keys) { com.google.datastore.v1.ReserveIdsResponse reserveIds( final com.google.datastore.v1.ReserveIdsRequest requestPb) { - try { + Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_RESERVEIDS); + try (Scope scope = traceUtil.getTracer().withSpan(span)) { return RetryHelper.runWithRetries( new Callable() { @Override @@ -437,7 +460,10 @@ public com.google.datastore.v1.ReserveIdsResponse call() throws DatastoreExcepti EXCEPTION_HANDLER, getOptions().getClock()); } catch (RetryHelperException e) { + span.setStatus(Status.UNKNOWN.withDescription(e.getMessage())); throw DatastoreException.translateAndThrow(e); + } finally { + span.end(TraceUtil.END_SPAN_OPTIONS); } } @@ -529,7 +555,8 @@ private com.google.datastore.v1.CommitResponse commitMutation( com.google.datastore.v1.CommitResponse commit( final com.google.datastore.v1.CommitRequest requestPb) { - try { + Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_COMMIT); + try (Scope scope = traceUtil.getTracer().withSpan(span)) { return RetryHelper.runWithRetries( new Callable() { @Override @@ -543,7 +570,10 @@ public com.google.datastore.v1.CommitResponse call() throws DatastoreException { : TRANSACTION_OPERATION_EXCEPTION_HANDLER, getOptions().getClock()); } catch (RetryHelperException e) { + span.setStatus(Status.UNKNOWN.withDescription(e.getMessage())); throw DatastoreException.translateAndThrow(e); + } finally { + span.end(TraceUtil.END_SPAN_OPTIONS); } } @@ -554,7 +584,8 @@ ByteString requestTransactionId( com.google.datastore.v1.BeginTransactionResponse beginTransaction( final com.google.datastore.v1.BeginTransactionRequest requestPb) { - try { + Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_BEGINTRANSACTION); + try (Scope scope = traceUtil.getTracer().withSpan(span)) { return RetryHelper.runWithRetries( new Callable() { @Override @@ -567,7 +598,10 @@ public com.google.datastore.v1.BeginTransactionResponse call() EXCEPTION_HANDLER, getOptions().getClock()); } catch (RetryHelperException e) { + span.setStatus(Status.UNKNOWN.withDescription(e.getMessage())); throw DatastoreException.translateAndThrow(e); + } finally { + span.end(TraceUtil.END_SPAN_OPTIONS); } } @@ -579,7 +613,8 @@ void rollbackTransaction(ByteString transaction) { } void rollback(final com.google.datastore.v1.RollbackRequest requestPb) { - try { + Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_ROLLBACK); + try (Scope scope = traceUtil.getTracer().withSpan(span)) { RetryHelper.runWithRetries( new Callable() { @Override @@ -592,7 +627,10 @@ public Void call() throws DatastoreException { EXCEPTION_HANDLER, getOptions().getClock()); } catch (RetryHelperException e) { + span.setStatus(Status.UNKNOWN.withDescription(e.getMessage())); throw DatastoreException.translateAndThrow(e); + } finally { + span.end(TraceUtil.END_SPAN_OPTIONS); } } } 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 new file mode 100644 index 000000000..1f28b2e80 --- /dev/null +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TraceUtil.java @@ -0,0 +1,74 @@ +/* + * 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_ALLOCATEIDS = "CloudDatastoreOperation.allocateIds"; + 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 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 7d3434108..4f13b4600 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,6 +21,8 @@ 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; import com.google.datastore.v1.AllocateIdsResponse; @@ -75,12 +77,18 @@ public HttpDatastoreRpc(DatastoreOptions options) { private HttpRequestInitializer getHttpRequestInitializer( final DatastoreOptions options, HttpTransportOptions httpTransportOptions) { - final HttpRequestInitializer delegate = httpTransportOptions.getHttpRequestInitializer(options); + // Open Census initialization + CensusHttpModule censusHttpModule = + new CensusHttpModule(TraceUtil.getInstance().getTracer(), true); + final HttpRequestInitializer censusHttpModuleHttpRequestInitializer = + censusHttpModule.getHttpRequestInitializer( + httpTransportOptions.getHttpRequestInitializer(options)); + final String applicationName = options.getApplicationName(); return new HttpRequestInitializer() { @Override public void initialize(HttpRequest httpRequest) throws IOException { - delegate.initialize(httpRequest); + censusHttpModuleHttpRequestInitializer.initialize(httpRequest); httpRequest.getHeaders().setUserAgent(applicationName); } }; diff --git a/pom.xml b/pom.xml index e35aab729..5a493fe29 100644 --- a/pom.xml +++ b/pom.xml @@ -191,7 +191,6 @@ google-oauth-client 1.30.6 - junit