Skip to content

Commit

Permalink
feat: add opencensus tracing/stats support for Datastore RPC operatio…
Browse files Browse the repository at this point in the history 
…ns (#130)
  • Loading branch information
@suraj-qlogic
suraj-qlogic authored Jun 8, 2020
1 parent 425f6db commit 5aee52f
Show file tree
Hide file tree
Showing 5 changed files with 139 additions and 16 deletions.
4 changes: 4 additions & 0 deletions google-cloud-datastore/pom.xml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,10 @@
<groupId>com.google.auth</groupId>
<artifactId>google-auth-library-oauth2-http</artifactId>
</dependency>
<dependency>
<groupId>io.opencensus</groupId>
<artifactId>opencensus-api</artifactId>
</dependency>

<!-- Test dependencies -->
<dependency>
Expand Down
64 changes: 51 additions & 13 deletions google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -53,6 +56,7 @@ final class DatastoreImpl extends BaseService<DatastoreOptions> 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);
Expand Down Expand Up @@ -132,30 +136,36 @@ public T call() throws DatastoreException {

@Override
public <T> T runInTransaction(final TransactionCallable<T> 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<T>(self, callable, null),
new ReadWriteTransactionCallable<T>(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> T runInTransaction(
final TransactionCallable<T> 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<T>(self, callable, transactionOptions),
new ReadWriteTransactionCallable<T>(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);
}
}

Expand All @@ -175,7 +185,8 @@ <T> QueryResults<T> 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<com.google.datastore.v1.RunQueryResponse>() {
@Override
Expand All @@ -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);
}
}

Expand Down Expand Up @@ -229,7 +243,8 @@ public List<Key> 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<com.google.datastore.v1.AllocateIdsResponse>() {
@Override
Expand All @@ -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);
}
}

Expand Down Expand Up @@ -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<com.google.datastore.v1.LookupResponse>() {
@Override
Expand All @@ -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);
}
}

Expand All @@ -425,7 +447,8 @@ public List<Key> 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<com.google.datastore.v1.ReserveIdsResponse>() {
@Override
Expand All @@ -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);
}
}

Expand Down Expand Up @@ -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<com.google.datastore.v1.CommitResponse>() {
@Override
Expand All @@ -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);
}
}

Expand All @@ -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<com.google.datastore.v1.BeginTransactionResponse>() {
@Override
Expand All @@ -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);
}
}

Expand All @@ -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<Void>() {
@Override
Expand All @@ -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);
}
}
}
74 changes: 74 additions & 0 deletions google-cloud-datastore/src/main/java/com/google/cloud/datastore/TraceUtil.java
View file
Original file line number Diff line number Diff line change
@@ -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.
*
* <p>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() {}
}
12 changes: 10 additions & 2 deletions google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/HttpDatastoreRpc.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
}
};
Expand Down
1 change: 0 additions & 1 deletion pom.xml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -191,7 +191,6 @@
<artifactId>google-oauth-client</artifactId>
<version>1.30.6</version>
</dependency>

<!-- Test dependencies -->
<dependency>
<groupId>junit</groupId>
Expand Down

0 comments on commit 5aee52f

Please sign in to comment.