Skip to content

Commit edaa539

Browse files
athakorathakor
authored
feat: add opencensus tracing support (googleapis#360)
1 parent acae383 commit edaa539

9 files changed

Lines changed: 114 additions & 16 deletions

File tree

‎google-cloud-firestore/src/main/java/com/google/cloud/firestore/CollectionGroup.java‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,9 @@
2222
import com.google.cloud.firestore.v1.FirestoreClient;
2323
import com.google.firestore.v1.Cursor;
2424
import com.google.firestore.v1.PartitionQueryRequest;
25+
import io.opencensus.common.Scope;
26+
import io.opencensus.trace.Span;
27+
import io.opencensus.trace.Status;
2528
import javax.annotation.Nullable;
2629

2730
/**
@@ -62,13 +65,18 @@ public void getPartitions(
6265
request.setPartitionCount(desiredPartitionCount - 1);
6366

6467
final FirestoreClient.PartitionQueryPagedResponse response;
65-
try {
68+
final TraceUtil traceUtil = TraceUtil.getInstance();
69+
Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_PARTITIONQUERY);
70+
try (Scope scope = traceUtil.getTracer().withSpan(span)) {
6671
response =
6772
ApiExceptions.callAndTranslateApiException(
6873
rpcContext.sendRequest(
6974
request.build(), rpcContext.getClient().partitionQueryPagedCallable()));
7075
} catch (ApiException exception) {
76+
span.setStatus(Status.UNKNOWN.withDescription(exception.getMessage()));
7177
throw FirestoreException.apiException(exception);
78+
} finally {
79+
span.end(TraceUtil.END_SPAN_OPTIONS);
7280
}
7381

7482
@Nullable Object[] lastCursor = null;

‎google-cloud-firestore/src/main/java/com/google/cloud/firestore/CollectionReference.java‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,9 @@
3030
import com.google.firestore.v1.Document;
3131
import com.google.firestore.v1.DocumentMask;
3232
import com.google.firestore.v1.ListDocumentsRequest;
33+
import io.opencensus.common.Scope;
34+
import io.opencensus.trace.Span;
35+
import io.opencensus.trace.Status;
3336
import java.util.Iterator;
3437
import java.util.Map;
3538
import javax.annotation.Nonnull;
@@ -135,16 +138,20 @@ public Iterable<DocumentReference> listDocuments() {
135138
request.setShowMissing(true);
136139

137140
final ListDocumentsPagedResponse response;
138-
139-
try {
141+
final TraceUtil traceUtil = TraceUtil.getInstance();
142+
Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_LISTDOCUMENTS);
143+
try (Scope scope = traceUtil.getTracer().withSpan(span)) {
140144
FirestoreRpc client = rpcContext.getClient();
141145
UnaryCallable<ListDocumentsRequest, ListDocumentsPagedResponse> callable =
142146
client.listDocumentsPagedCallable();
143147
ListDocumentsRequest build = request.build();
144148
ApiFuture<ListDocumentsPagedResponse> future = rpcContext.sendRequest(build, callable);
145149
response = ApiExceptions.callAndTranslateApiException(future);
146150
} catch (ApiException exception) {
151+
span.setStatus(Status.UNKNOWN.withDescription(exception.getMessage()));
147152
throw FirestoreException.apiException(exception);
153+
} finally {
154+
span.end(TraceUtil.END_SPAN_OPTIONS);
148155
}
149156

150157
return new Iterable<DocumentReference>() {

‎google-cloud-firestore/src/main/java/com/google/cloud/firestore/DocumentReference.java‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,9 @@
2525
import com.google.cloud.firestore.v1.FirestoreClient.ListCollectionIdsPagedResponse;
2626
import com.google.common.util.concurrent.MoreExecutors;
2727
import com.google.firestore.v1.ListCollectionIdsRequest;
28+
import io.opencensus.common.Scope;
29+
import io.opencensus.trace.Span;
30+
import io.opencensus.trace.Status;
2831
import java.util.Iterator;
2932
import java.util.List;
3033
import java.util.Map;
@@ -381,14 +384,18 @@ public Iterable<CollectionReference> listCollections() {
381384
ListCollectionIdsRequest.Builder request = ListCollectionIdsRequest.newBuilder();
382385
request.setParent(path.toString());
383386
final ListCollectionIdsPagedResponse response;
384-
385-
try {
387+
final TraceUtil traceUtil = TraceUtil.getInstance();
388+
Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_LISTCOLLECTIONIDS);
389+
try (Scope scope = traceUtil.getTracer().withSpan(span)) {
386390
response =
387391
ApiExceptions.callAndTranslateApiException(
388392
rpcContext.sendRequest(
389393
request.build(), rpcContext.getClient().listCollectionIdsPagedCallable()));
390394
} catch (ApiException exception) {
395+
span.setStatus(Status.UNKNOWN.withDescription(exception.getMessage()));
391396
throw FirestoreException.apiException(exception);
397+
} finally {
398+
span.end(TraceUtil.END_SPAN_OPTIONS);
392399
}
393400

394401
return new Iterable<CollectionReference>() {

‎google-cloud-firestore/src/main/java/com/google/cloud/firestore/FirestoreImpl.java‎

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -168,9 +168,14 @@ public void onNext(BatchGetDocumentsResponse response) {
168168

169169
numResponses++;
170170
if (numResponses == 1) {
171-
tracer.getCurrentSpan().addAnnotation("Firestore.BatchGet: First response");
171+
tracer
172+
.getCurrentSpan()
173+
.addAnnotation(TraceUtil.SPAN_NAME_BATCHGETDOCUMENTS + ": First response");
172174
} else if (numResponses % 100 == 0) {
173-
tracer.getCurrentSpan().addAnnotation("Firestore.BatchGet: Received 100 responses");
175+
tracer
176+
.getCurrentSpan()
177+
.addAnnotation(
178+
TraceUtil.SPAN_NAME_BATCHGETDOCUMENTS + ": Received 100 responses");
174179
}
175180

176181
switch (response.getResultCase()) {
@@ -199,13 +204,17 @@ public void onNext(BatchGetDocumentsResponse response) {
199204

200205
@Override
201206
public void onError(Throwable throwable) {
202-
tracer.getCurrentSpan().addAnnotation("Firestore.BatchGet: Error");
207+
tracer
208+
.getCurrentSpan()
209+
.addAnnotation(TraceUtil.SPAN_NAME_BATCHGETDOCUMENTS + ": Error");
203210
apiStreamObserver.onError(throwable);
204211
}
205212

206213
@Override
207214
public void onCompleted() {
208-
tracer.getCurrentSpan().addAnnotation("Firestore.BatchGet: Complete");
215+
tracer
216+
.getCurrentSpan()
217+
.addAnnotation(TraceUtil.SPAN_NAME_BATCHGETDOCUMENTS + ": Complete");
209218
apiStreamObserver.onCompleted();
210219
}
211220
};
@@ -228,7 +237,7 @@ public void onCompleted() {
228237
tracer
229238
.getCurrentSpan()
230239
.addAnnotation(
231-
"Firestore.BatchGet: Start",
240+
TraceUtil.SPAN_NAME_BATCHGETDOCUMENTS + ": Start",
232241
ImmutableMap.of(
233242
"numDocuments", AttributeValue.longAttributeValue(documentReferences.length)));
234243

‎google-cloud-firestore/src/main/java/com/google/cloud/firestore/Query.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1417,7 +1417,7 @@ private void internalStream(
14171417
Tracing.getTracer()
14181418
.getCurrentSpan()
14191419
.addAnnotation(
1420-
"Firestore.Query: Start",
1420+
TraceUtil.SPAN_NAME_RUNQUERY + ": Start",
14211421
ImmutableMap.of(
14221422
"transactional", AttributeValue.booleanAttributeValue(transactionId != null)));
14231423

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

Lines changed: 61 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,11 +17,71 @@
1717
package com.google.cloud.firestore;
1818

1919
import com.google.api.gax.rpc.ApiException;
20+
import com.google.cloud.firestore.spi.v1.GrpcFirestoreRpc;
2021
import io.opencensus.contrib.grpc.util.StatusConverter;
22+
import io.opencensus.trace.EndSpanOptions;
23+
import io.opencensus.trace.Span;
2124
import io.opencensus.trace.Status;
25+
import io.opencensus.trace.Tracer;
26+
import io.opencensus.trace.Tracing;
2227

23-
/** Census tracing utilities. */
28+
/**
29+
* Helper class for tracing utility. It is used for instrumenting {@link GrpcFirestoreRpc} with
30+
* OpenCensus APIs.
31+
*
32+
* <p>TraceUtil instances are created by the {@link TraceUtil#getInstance()} method.
33+
*/
2434
final class TraceUtil {
35+
36+
private final Tracer tracer = Tracing.getTracer();
37+
private static final TraceUtil traceUtil = new TraceUtil();
38+
static final String SPAN_NAME_GETDOCUMENT = "CloudFirestoreOperation.GetDocument";
39+
static final String SPAN_NAME_CREATEDOCUMENT = "CloudFirestoreOperation.CreateDocument";
40+
static final String SPAN_NAME_UPDATEDOCUMENT = "CloudFirestoreOperation.UpdateDocument";
41+
static final String SPAN_NAME_DELETEDOCUMENT = "CloudFirestoreOperation.DeleteDocument";
42+
static final String SPAN_NAME_LISTCOLLECTIONIDS = "CloudFirestoreOperation.ListCollectionIds";
43+
static final String SPAN_NAME_LISTDOCUMENTS = "CloudFirestoreOperation.ListDocuments";
44+
static final String SPAN_NAME_BEGINTRANSACTION = "CloudFirestoreOperation.BeginTransaction";
45+
static final String SPAN_NAME_COMMIT = "CloudFirestoreOperation.Commit";
46+
static final String SPAN_NAME_ROLLBACK = "CloudFirestoreOperation.Rollback";
47+
static final String SPAN_NAME_RUNQUERY = "CloudFirestoreOperation.RunQuery";
48+
static final String SPAN_NAME_PARTITIONQUERY = "CloudFirestoreOperation.partitionQuery";
49+
static final String SPAN_NAME_LISTEN = "CloudFirestoreOperation.Listen";
50+
static final String SPAN_NAME_BATCHGETDOCUMENTS = "CloudFirestoreOperation.BatchGetDocuments";
51+
static final String SPAN_NAME_BATCHWRITE = "CloudFirestoreOperation.BatchWrite";
52+
static final String SPAN_NAME_WRITE = "CloudFirestoreOperation.Write";
53+
54+
static final EndSpanOptions END_SPAN_OPTIONS =
55+
EndSpanOptions.builder().setSampleToLocalSpanStore(true).build();
56+
57+
/**
58+
* Starts a new span.
59+
*
60+
* @param spanName The name of the returned Span.
61+
* @return The newly created {@link Span}.
62+
*/
63+
protected Span startSpan(String spanName) {
64+
return tracer.spanBuilder(spanName).startSpan();
65+
}
66+
67+
/**
68+
* Return the global {@link Tracer}.
69+
*
70+
* @return The global {@link Tracer}.
71+
*/
72+
public Tracer getTracer() {
73+
return tracer;
74+
}
75+
76+
/**
77+
* Return TraceUtil Object.
78+
*
79+
* @return An instance of {@link TraceUtil}
80+
*/
81+
public static TraceUtil getInstance() {
82+
return traceUtil;
83+
}
84+
2585
private TraceUtil() {}
2686

2787
public static Status statusFromApiException(ApiException exception) {

‎google-cloud-firestore/src/main/java/com/google/cloud/firestore/Transaction.java‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
import com.google.firestore.v1.TransactionOptions.ReadOnly;
2929
import com.google.protobuf.ByteString;
3030
import com.google.protobuf.Empty;
31+
import io.opencensus.trace.Tracing;
3132
import java.util.List;
3233
import javax.annotation.Nonnull;
3334
import javax.annotation.Nullable;
@@ -83,6 +84,7 @@ Transaction wrapResult(ApiFuture<WriteResult> result) {
8384

8485
/** Starts a transaction and obtains the transaction id. */
8586
ApiFuture<Void> begin() {
87+
Tracing.getTracer().getCurrentSpan().addAnnotation(TraceUtil.SPAN_NAME_BEGINTRANSACTION);
8688
BeginTransactionRequest.Builder beginTransaction = BeginTransactionRequest.newBuilder();
8789
beginTransaction.setDatabase(firestore.getDatabaseName());
8890

@@ -123,6 +125,7 @@ ApiFuture<List<WriteResult>> commit() {
123125

124126
/** Rolls a transaction back and releases all read locks. */
125127
ApiFuture<Void> rollback() {
128+
Tracing.getTracer().getCurrentSpan().addAnnotation(TraceUtil.SPAN_NAME_ROLLBACK);
126129
RollbackRequest.Builder reqBuilder = RollbackRequest.newBuilder();
127130
reqBuilder.setTransaction(transactionId);
128131
reqBuilder.setDatabase(firestore.getDatabaseName());
@@ -150,7 +153,7 @@ public Void apply(Empty beginTransactionResponse) {
150153
@Nonnull
151154
public ApiFuture<DocumentSnapshot> get(@Nonnull DocumentReference documentRef) {
152155
Preconditions.checkState(isEmpty(), READ_BEFORE_WRITE_ERROR_MSG);
153-
156+
Tracing.getTracer().getCurrentSpan().addAnnotation(TraceUtil.SPAN_NAME_GETDOCUMENT);
154157
return ApiFutures.transform(
155158
firestore.getAll(new DocumentReference[] {documentRef}, /*fieldMask=*/ null, transactionId),
156159
new ApiFunction<List<DocumentSnapshot>, DocumentSnapshot>() {

‎google-cloud-firestore/src/main/java/com/google/cloud/firestore/UpdateBuilder.java‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -166,6 +166,7 @@ public T create(
166166
private T performCreate(
167167
@Nonnull DocumentReference documentReference, @Nonnull Map<String, Object> fields) {
168168
verifyNotCommitted();
169+
Tracing.getTracer().getCurrentSpan().addAnnotation(TraceUtil.SPAN_NAME_CREATEDOCUMENT);
169170
DocumentSnapshot documentSnapshot =
170171
DocumentSnapshot.fromObject(
171172
firestore, documentReference, fields, UserDataConverter.NO_DELETES);
@@ -548,7 +549,7 @@ private T performUpdate(
548549
@Nonnull Precondition precondition) {
549550
verifyNotCommitted();
550551
Preconditions.checkArgument(!fields.isEmpty(), "Data for update() cannot be empty.");
551-
552+
Tracing.getTracer().getCurrentSpan().addAnnotation(TraceUtil.SPAN_NAME_UPDATEDOCUMENT);
552553
Map<String, Object> deconstructedMap = expandObject(fields);
553554
DocumentSnapshot documentSnapshot =
554555
DocumentSnapshot.fromObject(
@@ -611,6 +612,7 @@ public T delete(@Nonnull DocumentReference documentReference) {
611612
private T performDelete(
612613
@Nonnull DocumentReference documentReference, @Nonnull Precondition precondition) {
613614
verifyNotCommitted();
615+
Tracing.getTracer().getCurrentSpan().addAnnotation(TraceUtil.SPAN_NAME_DELETEDOCUMENT);
614616
Write.Builder write = Write.newBuilder().setDelete(documentReference.getName());
615617

616618
if (!precondition.isEmpty()) {
@@ -626,7 +628,7 @@ ApiFuture<List<WriteResult>> commit(@Nullable ByteString transactionId) {
626628
Tracing.getTracer()
627629
.getCurrentSpan()
628630
.addAnnotation(
629-
"CloudFirestore.Commit",
631+
TraceUtil.SPAN_NAME_COMMIT,
630632
ImmutableMap.of("numDocuments", AttributeValue.longAttributeValue(writes.size())));
631633

632634
final CommitRequest.Builder request = CommitRequest.newBuilder();
@@ -674,7 +676,7 @@ ApiFuture<List<BatchWriteResult>> bulkCommit() {
674676
Tracing.getTracer()
675677
.getCurrentSpan()
676678
.addAnnotation(
677-
"CloudFirestore.BatchWrite",
679+
TraceUtil.SPAN_NAME_BATCHWRITE,
678680
ImmutableMap.of("numDocuments", AttributeValue.longAttributeValue(writes.size())));
679681

680682
final BatchWriteRequest.Builder request = BatchWriteRequest.newBuilder();

‎google-cloud-firestore/src/main/java/com/google/cloud/firestore/Watch.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
import io.grpc.Status.Code;
3737
import io.grpc.StatusException;
3838
import io.grpc.StatusRuntimeException;
39+
import io.opencensus.trace.Tracing;
3940
import java.util.ArrayList;
4041
import java.util.Collections;
4142
import java.util.Comparator;
@@ -395,6 +396,7 @@ public void run() {
395396
current = false;
396397
nextAttempt = backoff.createNextAttempt(nextAttempt);
397398

399+
Tracing.getTracer().getCurrentSpan().addAnnotation(TraceUtil.SPAN_NAME_LISTEN);
398400
stream = firestore.streamRequest(Watch.this, firestore.getClient().listenCallable());
399401

400402
ListenRequest.Builder request = ListenRequest.newBuilder();

0 commit comments

Comments
 (0)