Skip to content

Commit 5aee52f

Browse files
suraj-qlogicsuraj-qlogic
authored
feat: add opencensus tracing/stats support for Datastore RPC operations (#130)
1 parent 425f6db commit 5aee52f

File tree

5 files changed

+139
-16
lines changed

5 files changed

+139
-16
lines changed

‎google-cloud-datastore/pom.xml‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,10 @@
7878
<groupId>com.google.auth</groupId>
7979
<artifactId>google-auth-library-oauth2-http</artifactId>
8080
</dependency>
81+
<dependency>
82+
<groupId>io.opencensus</groupId>
83+
<artifactId>opencensus-api</artifactId>
84+
</dependency>
8185

8286
<!-- Test dependencies -->
8387
<dependency>

‎google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java‎

Lines changed: 51 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,9 @@
3434
import com.google.datastore.v1.ReserveIdsRequest;
3535
import com.google.datastore.v1.TransactionOptions;
3636
import com.google.protobuf.ByteString;
37+
import io.opencensus.common.Scope;
38+
import io.opencensus.trace.Span;
39+
import io.opencensus.trace.Status;
3740
import java.util.ArrayList;
3841
import java.util.Arrays;
3942
import java.util.Collections;
@@ -53,6 +56,7 @@ final class DatastoreImpl extends BaseService<DatastoreOptions> implements Datas
5356
TransactionExceptionHandler.build();
5457
private static final ExceptionHandler TRANSACTION_OPERATION_EXCEPTION_HANDLER =
5558
TransactionOperationExceptionHandler.build();
59+
private final TraceUtil traceUtil = TraceUtil.getInstance();;
5660

5761
DatastoreImpl(DatastoreOptions options) {
5862
super(options);
@@ -132,30 +136,36 @@ public T call() throws DatastoreException {
132136

133137
@Override
134138
public <T> T runInTransaction(final TransactionCallable<T> callable) {
135-
final DatastoreImpl self = this;
136-
try {
139+
Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_TRANSACTION);
140+
try (Scope scope = traceUtil.getTracer().withSpan(span)) {
137141
return RetryHelper.runWithRetries(
138-
new ReadWriteTransactionCallable<T>(self, callable, null),
142+
new ReadWriteTransactionCallable<T>(this, callable, null),
139143
retrySettings,
140144
TRANSACTION_EXCEPTION_HANDLER,
141145
getOptions().getClock());
142146
} catch (RetryHelperException e) {
147+
span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
143148
throw DatastoreException.translateAndThrow(e);
149+
} finally {
150+
span.end(TraceUtil.END_SPAN_OPTIONS);
144151
}
145152
}
146153

147154
@Override
148155
public <T> T runInTransaction(
149156
final TransactionCallable<T> callable, TransactionOptions transactionOptions) {
150-
final DatastoreImpl self = this;
151-
try {
157+
Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_TRANSACTION);
158+
try (Scope scope = traceUtil.getTracer().withSpan(span)) {
152159
return RetryHelper.runWithRetries(
153-
new ReadWriteTransactionCallable<T>(self, callable, transactionOptions),
160+
new ReadWriteTransactionCallable<T>(this, callable, transactionOptions),
154161
retrySettings,
155162
TRANSACTION_EXCEPTION_HANDLER,
156163
getOptions().getClock());
157164
} catch (RetryHelperException e) {
165+
span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
158166
throw DatastoreException.translateAndThrow(e);
167+
} finally {
168+
span.end(TraceUtil.END_SPAN_OPTIONS);
159169
}
160170
}
161171

@@ -175,7 +185,8 @@ <T> QueryResults<T> run(com.google.datastore.v1.ReadOptions readOptionsPb, Query
175185

176186
com.google.datastore.v1.RunQueryResponse runQuery(
177187
final com.google.datastore.v1.RunQueryRequest requestPb) {
178-
try {
188+
Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_RUNQUERY);
189+
try (Scope scope = traceUtil.getTracer().withSpan(span)) {
179190
return RetryHelper.runWithRetries(
180191
new Callable<com.google.datastore.v1.RunQueryResponse>() {
181192
@Override
@@ -189,7 +200,10 @@ public com.google.datastore.v1.RunQueryResponse call() throws DatastoreException
189200
: TRANSACTION_OPERATION_EXCEPTION_HANDLER,
190201
getOptions().getClock());
191202
} catch (RetryHelperException e) {
203+
span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
192204
throw DatastoreException.translateAndThrow(e);
205+
} finally {
206+
span.end(TraceUtil.END_SPAN_OPTIONS);
193207
}
194208
}
195209

@@ -229,7 +243,8 @@ public List<Key> allocateId(IncompleteKey... keys) {
229243

230244
private com.google.datastore.v1.AllocateIdsResponse allocateIds(
231245
final com.google.datastore.v1.AllocateIdsRequest requestPb) {
232-
try {
246+
Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_ALLOCATEIDS);
247+
try (Scope scope = traceUtil.getTracer().withSpan(span)) {
233248
return RetryHelper.runWithRetries(
234249
new Callable<com.google.datastore.v1.AllocateIdsResponse>() {
235250
@Override
@@ -241,7 +256,10 @@ public com.google.datastore.v1.AllocateIdsResponse call() throws DatastoreExcept
241256
EXCEPTION_HANDLER,
242257
getOptions().getClock());
243258
} catch (RetryHelperException e) {
259+
span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
244260
throw DatastoreException.translateAndThrow(e);
261+
} finally {
262+
span.end(TraceUtil.END_SPAN_OPTIONS);
245263
}
246264
}
247265

@@ -389,7 +407,8 @@ protected Entity computeNext() {
389407

390408
com.google.datastore.v1.LookupResponse lookup(
391409
final com.google.datastore.v1.LookupRequest requestPb) {
392-
try {
410+
Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_LOOKUP);
411+
try (Scope scope = traceUtil.getTracer().withSpan(span)) {
393412
return RetryHelper.runWithRetries(
394413
new Callable<com.google.datastore.v1.LookupResponse>() {
395414
@Override
@@ -403,7 +422,10 @@ public com.google.datastore.v1.LookupResponse call() throws DatastoreException {
403422
: TRANSACTION_OPERATION_EXCEPTION_HANDLER,
404423
getOptions().getClock());
405424
} catch (RetryHelperException e) {
425+
span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
406426
throw DatastoreException.translateAndThrow(e);
427+
} finally {
428+
span.end(TraceUtil.END_SPAN_OPTIONS);
407429
}
408430
}
409431

@@ -425,7 +447,8 @@ public List<Key> reserveIds(Key... keys) {
425447

426448
com.google.datastore.v1.ReserveIdsResponse reserveIds(
427449
final com.google.datastore.v1.ReserveIdsRequest requestPb) {
428-
try {
450+
Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_RESERVEIDS);
451+
try (Scope scope = traceUtil.getTracer().withSpan(span)) {
429452
return RetryHelper.runWithRetries(
430453
new Callable<com.google.datastore.v1.ReserveIdsResponse>() {
431454
@Override
@@ -437,7 +460,10 @@ public com.google.datastore.v1.ReserveIdsResponse call() throws DatastoreExcepti
437460
EXCEPTION_HANDLER,
438461
getOptions().getClock());
439462
} catch (RetryHelperException e) {
463+
span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
440464
throw DatastoreException.translateAndThrow(e);
465+
} finally {
466+
span.end(TraceUtil.END_SPAN_OPTIONS);
441467
}
442468
}
443469

@@ -529,7 +555,8 @@ private com.google.datastore.v1.CommitResponse commitMutation(
529555

530556
com.google.datastore.v1.CommitResponse commit(
531557
final com.google.datastore.v1.CommitRequest requestPb) {
532-
try {
558+
Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_COMMIT);
559+
try (Scope scope = traceUtil.getTracer().withSpan(span)) {
533560
return RetryHelper.runWithRetries(
534561
new Callable<com.google.datastore.v1.CommitResponse>() {
535562
@Override
@@ -543,7 +570,10 @@ public com.google.datastore.v1.CommitResponse call() throws DatastoreException {
543570
: TRANSACTION_OPERATION_EXCEPTION_HANDLER,
544571
getOptions().getClock());
545572
} catch (RetryHelperException e) {
573+
span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
546574
throw DatastoreException.translateAndThrow(e);
575+
} finally {
576+
span.end(TraceUtil.END_SPAN_OPTIONS);
547577
}
548578
}
549579

@@ -554,7 +584,8 @@ ByteString requestTransactionId(
554584

555585
com.google.datastore.v1.BeginTransactionResponse beginTransaction(
556586
final com.google.datastore.v1.BeginTransactionRequest requestPb) {
557-
try {
587+
Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_BEGINTRANSACTION);
588+
try (Scope scope = traceUtil.getTracer().withSpan(span)) {
558589
return RetryHelper.runWithRetries(
559590
new Callable<com.google.datastore.v1.BeginTransactionResponse>() {
560591
@Override
@@ -567,7 +598,10 @@ public com.google.datastore.v1.BeginTransactionResponse call()
567598
EXCEPTION_HANDLER,
568599
getOptions().getClock());
569600
} catch (RetryHelperException e) {
601+
span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
570602
throw DatastoreException.translateAndThrow(e);
603+
} finally {
604+
span.end(TraceUtil.END_SPAN_OPTIONS);
571605
}
572606
}
573607

@@ -579,7 +613,8 @@ void rollbackTransaction(ByteString transaction) {
579613
}
580614

581615
void rollback(final com.google.datastore.v1.RollbackRequest requestPb) {
582-
try {
616+
Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_ROLLBACK);
617+
try (Scope scope = traceUtil.getTracer().withSpan(span)) {
583618
RetryHelper.runWithRetries(
584619
new Callable<Void>() {
585620
@Override
@@ -592,7 +627,10 @@ public Void call() throws DatastoreException {
592627
EXCEPTION_HANDLER,
593628
getOptions().getClock());
594629
} catch (RetryHelperException e) {
630+
span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
595631
throw DatastoreException.translateAndThrow(e);
632+
} finally {
633+
span.end(TraceUtil.END_SPAN_OPTIONS);
596634
}
597635
}
598636
}

‎google-cloud-datastore/src/main/java/com/google/cloud/datastore/TraceUtil.java‎

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,74 @@
1+
/*
2+
* Copyright 2020 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.cloud.datastore;
18+
19+
import com.google.cloud.datastore.spi.v1.HttpDatastoreRpc;
20+
import io.opencensus.trace.EndSpanOptions;
21+
import io.opencensus.trace.Span;
22+
import io.opencensus.trace.Tracer;
23+
import io.opencensus.trace.Tracing;
24+
25+
/**
26+
* Helper class for tracing utility. It is used for instrumenting {@link HttpDatastoreRpc} with
27+
* OpenCensus APIs.
28+
*
29+
* <p>TraceUtil instances are created by the {@link TraceUtil#getInstance()} method.
30+
*/
31+
public class TraceUtil {
32+
private final Tracer tracer = Tracing.getTracer();
33+
private static final TraceUtil traceUtil = new TraceUtil();
34+
static final String SPAN_NAME_ALLOCATEIDS = "CloudDatastoreOperation.allocateIds";
35+
static final String SPAN_NAME_TRANSACTION = "CloudDatastoreOperation.readWriteTransaction";
36+
static final String SPAN_NAME_BEGINTRANSACTION = "CloudDatastoreOperation.beginTransaction";
37+
static final String SPAN_NAME_COMMIT = "CloudDatastoreOperation.commit";
38+
static final String SPAN_NAME_LOOKUP = "CloudDatastoreOperation.lookup";
39+
static final String SPAN_NAME_RESERVEIDS = "CloudDatastoreOperation.reserveIds";
40+
static final String SPAN_NAME_ROLLBACK = "CloudDatastoreOperation.rollback";
41+
static final String SPAN_NAME_RUNQUERY = "CloudDatastoreOperation.runQuery";
42+
static final EndSpanOptions END_SPAN_OPTIONS =
43+
EndSpanOptions.builder().setSampleToLocalSpanStore(true).build();
44+
45+
/**
46+
* Starts a new span.
47+
*
48+
* @param spanName The name of the returned Span.
49+
* @return The newly created {@link Span}.
50+
*/
51+
protected Span startSpan(String spanName) {
52+
return tracer.spanBuilder(spanName).startSpan();
53+
}
54+
55+
/**
56+
* Return the global {@link Tracer}.
57+
*
58+
* @return The global {@link Tracer}.
59+
*/
60+
public Tracer getTracer() {
61+
return tracer;
62+
}
63+
64+
/**
65+
* Return TraceUtil Object.
66+
*
67+
* @return An instance of {@link TraceUtil}
68+
*/
69+
public static TraceUtil getInstance() {
70+
return traceUtil;
71+
}
72+
73+
private TraceUtil() {}
74+
}

‎google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/HttpDatastoreRpc.java‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,8 @@
2121
import com.google.api.client.http.HttpTransport;
2222
import com.google.cloud.datastore.DatastoreException;
2323
import com.google.cloud.datastore.DatastoreOptions;
24+
import com.google.cloud.datastore.TraceUtil;
25+
import com.google.cloud.http.CensusHttpModule;
2426
import com.google.cloud.http.HttpTransportOptions;
2527
import com.google.datastore.v1.AllocateIdsRequest;
2628
import com.google.datastore.v1.AllocateIdsResponse;
@@ -75,12 +77,18 @@ public HttpDatastoreRpc(DatastoreOptions options) {
7577

7678
private HttpRequestInitializer getHttpRequestInitializer(
7779
final DatastoreOptions options, HttpTransportOptions httpTransportOptions) {
78-
final HttpRequestInitializer delegate = httpTransportOptions.getHttpRequestInitializer(options);
80+
// Open Census initialization
81+
CensusHttpModule censusHttpModule =
82+
new CensusHttpModule(TraceUtil.getInstance().getTracer(), true);
83+
final HttpRequestInitializer censusHttpModuleHttpRequestInitializer =
84+
censusHttpModule.getHttpRequestInitializer(
85+
httpTransportOptions.getHttpRequestInitializer(options));
86+
7987
final String applicationName = options.getApplicationName();
8088
return new HttpRequestInitializer() {
8189
@Override
8290
public void initialize(HttpRequest httpRequest) throws IOException {
83-
delegate.initialize(httpRequest);
91+
censusHttpModuleHttpRequestInitializer.initialize(httpRequest);
8492
httpRequest.getHeaders().setUserAgent(applicationName);
8593
}
8694
};

‎pom.xml‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -191,7 +191,6 @@
191191
<artifactId>google-oauth-client</artifactId>
192192
<version>1.30.6</version>
193193
</dependency>
194-
195194
<!-- Test dependencies -->
196195
<dependency>
197196
<groupId>junit</groupId>

0 commit comments

Comments
 (0)