Skip to content

Commit 5dad898

Browse files
author
Daniel Chia
committed
Census tracing integration for Firestore.
1 parent 45a9978 commit 5dad898

5 files changed

Lines changed: 141 additions & 3 deletions

File tree

‎google-cloud-firestore/pom.xml‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,14 @@
6363
<groupId>io.grpc</groupId>
6464
<artifactId>grpc-auth</artifactId>
6565
</dependency>
66+
<dependency>
67+
<groupId>io.opencensus</groupId>
68+
<artifactId>opencensus-api</artifactId>
69+
</dependency>
70+
<dependency>
71+
<groupId>io.opencensus</groupId>
72+
<artifactId>opencensus-contrib-grpc-util</artifactId>
73+
</dependency>
6674
<dependency>
6775
<groupId>${project.groupId}</groupId>
6876
<artifactId>google-cloud-core</artifactId>

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

Lines changed: 62 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,17 +20,25 @@
2020
import com.google.api.core.ApiFutureCallback;
2121
import com.google.api.core.ApiFutures;
2222
import com.google.api.core.SettableApiFuture;
23+
import com.google.api.gax.rpc.ApiException;
2324
import com.google.api.gax.rpc.ApiStreamObserver;
2425
import com.google.api.gax.rpc.BidiStreamingCallable;
2526
import com.google.api.gax.rpc.ServerStreamingCallable;
2627
import com.google.api.gax.rpc.UnaryCallable;
2728
import com.google.cloud.firestore.spi.v1beta1.FirestoreRpc;
2829
import com.google.common.base.Preconditions;
30+
import com.google.common.collect.ImmutableMap;
2931
import com.google.firestore.v1beta1.BatchGetDocumentsRequest;
3032
import com.google.firestore.v1beta1.BatchGetDocumentsResponse;
3133
import com.google.firestore.v1beta1.DatabaseRootName;
3234
import com.google.protobuf.ByteString;
35+
import io.grpc.Context;
3336
import io.grpc.Status;
37+
import io.opencensus.common.Scope;
38+
import io.opencensus.trace.AttributeValue;
39+
import io.opencensus.trace.Span;
40+
import io.opencensus.trace.Tracer;
41+
import io.opencensus.trace.Tracing;
3442
import java.util.ArrayList;
3543
import java.util.HashMap;
3644
import java.util.List;
@@ -52,6 +60,10 @@ class FirestoreImpl implements Firestore {
5260
private static final String AUTO_ID_ALPHABET =
5361
"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789";
5462

63+
private static final Tracer tracer = Tracing.getTracer();
64+
private static final io.opencensus.trace.Status TOO_MANY_RETRIES_STATUS =
65+
io.opencensus.trace.Status.ABORTED.withDescription("too many retries");
66+
5567
private final FirestoreRpc firestoreClient;
5668
private final FirestoreOptions firestoreOptions;
5769
private final ResourcePath databasePath;
@@ -126,11 +138,24 @@ ApiFuture<List<DocumentSnapshot>> getAll(
126138

127139
ApiStreamObserver<BatchGetDocumentsResponse> responseObserver =
128140
new ApiStreamObserver<BatchGetDocumentsResponse>() {
141+
int numResponses;
142+
129143
@Override
130144
public void onNext(BatchGetDocumentsResponse response) {
131145
DocumentReference documentReference;
132146
DocumentSnapshot documentSnapshot;
133147

148+
numResponses++;
149+
if (numResponses == 1) {
150+
tracer
151+
.getCurrentSpan()
152+
.addAnnotation("Firestore.BatchGet: First response");
153+
} else if (numResponses % 100 == 0) {
154+
tracer
155+
.getCurrentSpan()
156+
.addAnnotation("Firestore.BatchGet: Received 100 responses");
157+
}
158+
134159
switch (response.getResultCase()) {
135160
case FOUND:
136161
documentReference =
@@ -161,11 +186,13 @@ public void onNext(BatchGetDocumentsResponse response) {
161186

162187
@Override
163188
public void onError(Throwable throwable) {
189+
tracer.getCurrentSpan().addAnnotation("Firestore.BatchGet: Error");
164190
futureList.setException(throwable);
165191
}
166192

167193
@Override
168194
public void onCompleted() {
195+
tracer.getCurrentSpan().addAnnotation("Firestore.BatchGet: Complete");
169196
List<DocumentSnapshot> documentSnapshots = new ArrayList<>();
170197

171198
for (DocumentReference documentReference : documentReferences) {
@@ -187,6 +214,13 @@ public void onCompleted() {
187214
request.addDocuments(docRef.getName());
188215
}
189216

217+
tracer
218+
.getCurrentSpan()
219+
.addAnnotation(
220+
"Firestore.BatchGet: Start",
221+
ImmutableMap.of(
222+
"numDocuments", AttributeValue.longAttributeValue(documentReferences.length)));
223+
190224
streamRequest(request.build(), responseObserver, firestoreClient.batchGetDocumentsCallable());
191225

192226
return futureList;
@@ -213,11 +247,26 @@ private <T> void runTransaction(
213247
final Transaction.Function<T> transactionCallback,
214248
final SettableApiFuture<T> resultFuture,
215249
final TransactionOptions options) {
250+
Span span = tracer.spanBuilder("CloudFirestore.Transaction").startSpan();
251+
try (Scope s = tracer.withSpan(span)) {
252+
runTransactionAttempt(transactionCallback, resultFuture, options, span);
253+
}
254+
}
255+
256+
private <T> void runTransactionAttempt(
257+
final Transaction.Function<T> transactionCallback,
258+
final SettableApiFuture<T> resultFuture,
259+
final TransactionOptions options,
260+
final Span span) {
216261
final Transaction transaction = new Transaction(this, options.getPreviousTransactionId());
217262
final Executor userCallbackExecutor =
218-
options.getExecutor() != null ? options.getExecutor() : firestoreClient.getExecutor();
263+
Context.currentContextExecutor(
264+
options.getExecutor() != null ? options.getExecutor() : firestoreClient.getExecutor());
219265

220266
final int attemptsRemaining = options.getNumberOfAttempts() - 1;
267+
span.addAnnotation(
268+
"Start runTransaction",
269+
ImmutableMap.of("attemptsRemaining", AttributeValue.longAttributeValue(attemptsRemaining)));
221270

222271
ApiFutures.addCallback(
223272
transaction.begin(),
@@ -253,6 +302,8 @@ public void onFailure(Throwable throwable) {
253302

254303
@Override
255304
public void onSuccess(List<WriteResult> writeResults) {
305+
span.setStatus(io.opencensus.trace.Status.OK);
306+
span.end();
256307
resultFuture.set(userResult);
257308
}
258309
});
@@ -279,19 +330,27 @@ public void run() {
279330

280331
private void maybeRetry() {
281332
if (attemptsRemaining > 0) {
282-
runTransaction(
333+
span.addAnnotation("retrying");
334+
runTransactionAttempt(
283335
transactionCallback,
284336
resultFuture,
285337
new TransactionOptions(
286-
attemptsRemaining, options.getExecutor(), transaction.getTransactionId()));
338+
attemptsRemaining, options.getExecutor(), transaction.getTransactionId()),
339+
span);
287340
} else {
341+
span.setStatus(TOO_MANY_RETRIES_STATUS);
342+
span.end();
288343
rejectTransaction(
289344
FirestoreException.serverRejected(
290345
Status.ABORTED, "Transaction was cancelled because of too many retries."));
291346
}
292347
}
293348

294349
private void rejectTransaction(final Throwable throwable) {
350+
if (throwable instanceof ApiException) {
351+
span.setStatus(TraceUtil.statusFromApiException((ApiException) throwable));
352+
}
353+
span.end();
295354
if (transaction.isPending()) {
296355
ApiFutures.addCallback(
297356
transaction.rollback(),

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

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
import com.google.api.gax.rpc.ApiStreamObserver;
2828
import com.google.cloud.firestore.DocumentChange.Type;
2929
import com.google.common.base.Preconditions;
30+
import com.google.common.collect.ImmutableMap;
3031
import com.google.firestore.v1beta1.Cursor;
3132
import com.google.firestore.v1beta1.Document;
3233
import com.google.firestore.v1beta1.RunQueryRequest;
@@ -39,6 +40,8 @@
3940
import com.google.firestore.v1beta1.Value;
4041
import com.google.protobuf.ByteString;
4142
import com.google.protobuf.Int32Value;
43+
import io.opencensus.trace.AttributeValue;
44+
import io.opencensus.trace.Tracing;
4245
import java.util.ArrayList;
4346
import java.util.Comparator;
4447
import java.util.Iterator;
@@ -922,13 +925,31 @@ private void stream(
922925
request.setTransaction(transactionId);
923926
}
924927

928+
Tracing.getTracer().getCurrentSpan()
929+
.addAnnotation(
930+
"Firestore.Query: Start",
931+
ImmutableMap.of(
932+
"transactional", AttributeValue.booleanAttributeValue(transactionId != null)));
933+
925934
ApiStreamObserver<RunQueryResponse> observer =
926935
new ApiStreamObserver<RunQueryResponse>() {
927936
Instant readTime;
937+
boolean firstResponse;
938+
int numDocuments;
928939

929940
@Override
930941
public void onNext(RunQueryResponse response) {
942+
if (!firstResponse) {
943+
firstResponse = true;
944+
Tracing.getTracer().getCurrentSpan()
945+
.addAnnotation("Firestore.Query: First response");
946+
}
931947
if (response.hasDocument()) {
948+
numDocuments++;
949+
if (numDocuments % 100 == 0) {
950+
Tracing.getTracer().getCurrentSpan()
951+
.addAnnotation("Firestore.Query: Received 100 documents");
952+
}
932953
Document document = response.getDocument();
933954
QueryDocumentSnapshot documentSnapshot =
934955
QueryDocumentSnapshot.fromDocument(firestore, response.getReadTime(), document);
@@ -944,11 +965,18 @@ public void onNext(RunQueryResponse response) {
944965

945966
@Override
946967
public void onError(Throwable throwable) {
968+
Tracing.getTracer().getCurrentSpan().addAnnotation("Firestore.Query: Error");
947969
documentObserver.onError(throwable);
948970
}
949971

950972
@Override
951973
public void onCompleted() {
974+
Tracing.getTracer().getCurrentSpan()
975+
.addAnnotation(
976+
"Firestore.Query: Completed",
977+
ImmutableMap.of(
978+
"numDocuments",
979+
AttributeValue.longAttributeValue(numDocuments)));
952980
documentObserver.onCompleted(readTime);
953981
}
954982
};

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

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
/*
2+
* Copyright 2017 Google Inc. All Rights Reserved.
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.firestore;
18+
19+
import com.google.api.gax.rpc.ApiException;
20+
import io.opencensus.contrib.grpc.util.StatusConverter;
21+
import io.opencensus.trace.Status;
22+
23+
/** Census tracing utilities. */
24+
class TraceUtil {
25+
private TraceUtil() {}
26+
27+
public static Status statusFromApiException(ApiException exception) {
28+
if (exception.getStatusCode().getTransportCode() instanceof io.grpc.Status) {
29+
io.grpc.Status grpcStatus = (io.grpc.Status) exception.getStatusCode().getTransportCode();
30+
return StatusConverter.fromGrpcStatus(grpcStatus);
31+
}
32+
33+
return Status.UNKNOWN.withDescription(exception.getMessage());
34+
}
35+
}

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,10 +21,15 @@
2121
import com.google.api.core.ApiFutures;
2222
import com.google.cloud.firestore.UserDataConverter.EncodingOptions;
2323
import com.google.common.base.Preconditions;
24+
import com.google.common.collect.ImmutableList;
25+
import com.google.common.collect.ImmutableMap;
26+
import com.google.common.collect.UnmodifiableIterator;
2427
import com.google.firestore.v1beta1.CommitRequest;
2528
import com.google.firestore.v1beta1.CommitResponse;
2629
import com.google.firestore.v1beta1.Write;
2730
import com.google.protobuf.ByteString;
31+
import io.opencensus.trace.AttributeValue;
32+
import io.opencensus.trace.Tracing;
2833
import java.util.ArrayList;
2934
import java.util.HashMap;
3035
import java.util.Iterator;
@@ -575,6 +580,9 @@ private T performDelete(
575580

576581
/** Commit the current batch. */
577582
ApiFuture<List<WriteResult>> commit(@Nullable ByteString transactionId) {
583+
Tracing.getTracer().getCurrentSpan().addAnnotation(
584+
"CloudFirestore.Commit",
585+
ImmutableMap.of("numDocuments", AttributeValue.longAttributeValue(mutations.size())));
578586

579587
final CommitRequest.Builder request = CommitRequest.newBuilder();
580588
request.setDatabase(firestore.getDatabaseName());

0 commit comments

Comments
 (0)