From edaa5395be0353fb261d954429c624623bc4e346 Mon Sep 17 00:00:00 2001 From: Ajit Thakor <49403056+athakor@users.noreply.github.com> Date: Thu, 10 Sep 2020 01:24:27 +0530 Subject: [PATCH] feat: add opencensus tracing support (#360) --- .../cloud/firestore/CollectionGroup.java | 10 ++- .../cloud/firestore/CollectionReference.java | 11 +++- .../cloud/firestore/DocumentReference.java | 11 +++- .../google/cloud/firestore/FirestoreImpl.java | 19 ++++-- .../com/google/cloud/firestore/Query.java | 2 +- .../com/google/cloud/firestore/TraceUtil.java | 62 ++++++++++++++++++- .../google/cloud/firestore/Transaction.java | 5 +- .../google/cloud/firestore/UpdateBuilder.java | 8 ++- .../com/google/cloud/firestore/Watch.java | 2 + 9 files changed, 114 insertions(+), 16 deletions(-) diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/CollectionGroup.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/CollectionGroup.java index c2677f087..23dd7bbe3 100644 --- a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/CollectionGroup.java +++ b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/CollectionGroup.java @@ -22,6 +22,9 @@ import com.google.cloud.firestore.v1.FirestoreClient; import com.google.firestore.v1.Cursor; import com.google.firestore.v1.PartitionQueryRequest; +import io.opencensus.common.Scope; +import io.opencensus.trace.Span; +import io.opencensus.trace.Status; import javax.annotation.Nullable; /** @@ -62,13 +65,18 @@ public void getPartitions( request.setPartitionCount(desiredPartitionCount - 1); final FirestoreClient.PartitionQueryPagedResponse response; - try { + final TraceUtil traceUtil = TraceUtil.getInstance(); + Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_PARTITIONQUERY); + try (Scope scope = traceUtil.getTracer().withSpan(span)) { response = ApiExceptions.callAndTranslateApiException( rpcContext.sendRequest( request.build(), rpcContext.getClient().partitionQueryPagedCallable())); } catch (ApiException exception) { + span.setStatus(Status.UNKNOWN.withDescription(exception.getMessage())); throw FirestoreException.apiException(exception); + } finally { + span.end(TraceUtil.END_SPAN_OPTIONS); } @Nullable Object[] lastCursor = null; diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/CollectionReference.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/CollectionReference.java index c2367bed1..f0ea6a7ef 100644 --- a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/CollectionReference.java +++ b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/CollectionReference.java @@ -30,6 +30,9 @@ import com.google.firestore.v1.Document; import com.google.firestore.v1.DocumentMask; import com.google.firestore.v1.ListDocumentsRequest; +import io.opencensus.common.Scope; +import io.opencensus.trace.Span; +import io.opencensus.trace.Status; import java.util.Iterator; import java.util.Map; import javax.annotation.Nonnull; @@ -135,8 +138,9 @@ public Iterable listDocuments() { request.setShowMissing(true); final ListDocumentsPagedResponse response; - - try { + final TraceUtil traceUtil = TraceUtil.getInstance(); + Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_LISTDOCUMENTS); + try (Scope scope = traceUtil.getTracer().withSpan(span)) { FirestoreRpc client = rpcContext.getClient(); UnaryCallable callable = client.listDocumentsPagedCallable(); @@ -144,7 +148,10 @@ public Iterable listDocuments() { ApiFuture future = rpcContext.sendRequest(build, callable); response = ApiExceptions.callAndTranslateApiException(future); } catch (ApiException exception) { + span.setStatus(Status.UNKNOWN.withDescription(exception.getMessage())); throw FirestoreException.apiException(exception); + } finally { + span.end(TraceUtil.END_SPAN_OPTIONS); } return new Iterable() { diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/DocumentReference.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/DocumentReference.java index 27c10be57..51c126aaf 100644 --- a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/DocumentReference.java +++ b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/DocumentReference.java @@ -25,6 +25,9 @@ import com.google.cloud.firestore.v1.FirestoreClient.ListCollectionIdsPagedResponse; import com.google.common.util.concurrent.MoreExecutors; import com.google.firestore.v1.ListCollectionIdsRequest; +import io.opencensus.common.Scope; +import io.opencensus.trace.Span; +import io.opencensus.trace.Status; import java.util.Iterator; import java.util.List; import java.util.Map; @@ -381,14 +384,18 @@ public Iterable listCollections() { ListCollectionIdsRequest.Builder request = ListCollectionIdsRequest.newBuilder(); request.setParent(path.toString()); final ListCollectionIdsPagedResponse response; - - try { + final TraceUtil traceUtil = TraceUtil.getInstance(); + Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_LISTCOLLECTIONIDS); + try (Scope scope = traceUtil.getTracer().withSpan(span)) { response = ApiExceptions.callAndTranslateApiException( rpcContext.sendRequest( request.build(), rpcContext.getClient().listCollectionIdsPagedCallable())); } catch (ApiException exception) { + span.setStatus(Status.UNKNOWN.withDescription(exception.getMessage())); throw FirestoreException.apiException(exception); + } finally { + span.end(TraceUtil.END_SPAN_OPTIONS); } return new Iterable() { diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/FirestoreImpl.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/FirestoreImpl.java index f96960952..57afa9938 100644 --- a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/FirestoreImpl.java +++ b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/FirestoreImpl.java @@ -168,9 +168,14 @@ public void onNext(BatchGetDocumentsResponse response) { numResponses++; if (numResponses == 1) { - tracer.getCurrentSpan().addAnnotation("Firestore.BatchGet: First response"); + tracer + .getCurrentSpan() + .addAnnotation(TraceUtil.SPAN_NAME_BATCHGETDOCUMENTS + ": First response"); } else if (numResponses % 100 == 0) { - tracer.getCurrentSpan().addAnnotation("Firestore.BatchGet: Received 100 responses"); + tracer + .getCurrentSpan() + .addAnnotation( + TraceUtil.SPAN_NAME_BATCHGETDOCUMENTS + ": Received 100 responses"); } switch (response.getResultCase()) { @@ -199,13 +204,17 @@ public void onNext(BatchGetDocumentsResponse response) { @Override public void onError(Throwable throwable) { - tracer.getCurrentSpan().addAnnotation("Firestore.BatchGet: Error"); + tracer + .getCurrentSpan() + .addAnnotation(TraceUtil.SPAN_NAME_BATCHGETDOCUMENTS + ": Error"); apiStreamObserver.onError(throwable); } @Override public void onCompleted() { - tracer.getCurrentSpan().addAnnotation("Firestore.BatchGet: Complete"); + tracer + .getCurrentSpan() + .addAnnotation(TraceUtil.SPAN_NAME_BATCHGETDOCUMENTS + ": Complete"); apiStreamObserver.onCompleted(); } }; @@ -228,7 +237,7 @@ public void onCompleted() { tracer .getCurrentSpan() .addAnnotation( - "Firestore.BatchGet: Start", + TraceUtil.SPAN_NAME_BATCHGETDOCUMENTS + ": Start", ImmutableMap.of( "numDocuments", AttributeValue.longAttributeValue(documentReferences.length))); diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Query.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Query.java index 035074461..5e205d580 100644 --- a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Query.java +++ b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Query.java @@ -1417,7 +1417,7 @@ private void internalStream( Tracing.getTracer() .getCurrentSpan() .addAnnotation( - "Firestore.Query: Start", + TraceUtil.SPAN_NAME_RUNQUERY + ": Start", ImmutableMap.of( "transactional", AttributeValue.booleanAttributeValue(transactionId != null))); diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/TraceUtil.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/TraceUtil.java index 1f7c55af4..66c278bc7 100644 --- a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/TraceUtil.java +++ b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/TraceUtil.java @@ -17,11 +17,71 @@ package com.google.cloud.firestore; import com.google.api.gax.rpc.ApiException; +import com.google.cloud.firestore.spi.v1.GrpcFirestoreRpc; import io.opencensus.contrib.grpc.util.StatusConverter; +import io.opencensus.trace.EndSpanOptions; +import io.opencensus.trace.Span; import io.opencensus.trace.Status; +import io.opencensus.trace.Tracer; +import io.opencensus.trace.Tracing; -/** Census tracing utilities. */ +/** + * Helper class for tracing utility. It is used for instrumenting {@link GrpcFirestoreRpc} with + * OpenCensus APIs. + * + *

TraceUtil instances are created by the {@link TraceUtil#getInstance()} method. + */ final class TraceUtil { + + private final Tracer tracer = Tracing.getTracer(); + private static final TraceUtil traceUtil = new TraceUtil(); + static final String SPAN_NAME_GETDOCUMENT = "CloudFirestoreOperation.GetDocument"; + static final String SPAN_NAME_CREATEDOCUMENT = "CloudFirestoreOperation.CreateDocument"; + static final String SPAN_NAME_UPDATEDOCUMENT = "CloudFirestoreOperation.UpdateDocument"; + static final String SPAN_NAME_DELETEDOCUMENT = "CloudFirestoreOperation.DeleteDocument"; + static final String SPAN_NAME_LISTCOLLECTIONIDS = "CloudFirestoreOperation.ListCollectionIds"; + static final String SPAN_NAME_LISTDOCUMENTS = "CloudFirestoreOperation.ListDocuments"; + static final String SPAN_NAME_BEGINTRANSACTION = "CloudFirestoreOperation.BeginTransaction"; + static final String SPAN_NAME_COMMIT = "CloudFirestoreOperation.Commit"; + static final String SPAN_NAME_ROLLBACK = "CloudFirestoreOperation.Rollback"; + static final String SPAN_NAME_RUNQUERY = "CloudFirestoreOperation.RunQuery"; + static final String SPAN_NAME_PARTITIONQUERY = "CloudFirestoreOperation.partitionQuery"; + static final String SPAN_NAME_LISTEN = "CloudFirestoreOperation.Listen"; + static final String SPAN_NAME_BATCHGETDOCUMENTS = "CloudFirestoreOperation.BatchGetDocuments"; + static final String SPAN_NAME_BATCHWRITE = "CloudFirestoreOperation.BatchWrite"; + static final String SPAN_NAME_WRITE = "CloudFirestoreOperation.Write"; + + 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() {} public static Status statusFromApiException(ApiException exception) { diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Transaction.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Transaction.java index 6a9fd9714..e8308b57d 100644 --- a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Transaction.java +++ b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Transaction.java @@ -28,6 +28,7 @@ import com.google.firestore.v1.TransactionOptions.ReadOnly; import com.google.protobuf.ByteString; import com.google.protobuf.Empty; +import io.opencensus.trace.Tracing; import java.util.List; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -83,6 +84,7 @@ Transaction wrapResult(ApiFuture result) { /** Starts a transaction and obtains the transaction id. */ ApiFuture begin() { + Tracing.getTracer().getCurrentSpan().addAnnotation(TraceUtil.SPAN_NAME_BEGINTRANSACTION); BeginTransactionRequest.Builder beginTransaction = BeginTransactionRequest.newBuilder(); beginTransaction.setDatabase(firestore.getDatabaseName()); @@ -123,6 +125,7 @@ ApiFuture> commit() { /** Rolls a transaction back and releases all read locks. */ ApiFuture rollback() { + Tracing.getTracer().getCurrentSpan().addAnnotation(TraceUtil.SPAN_NAME_ROLLBACK); RollbackRequest.Builder reqBuilder = RollbackRequest.newBuilder(); reqBuilder.setTransaction(transactionId); reqBuilder.setDatabase(firestore.getDatabaseName()); @@ -150,7 +153,7 @@ public Void apply(Empty beginTransactionResponse) { @Nonnull public ApiFuture get(@Nonnull DocumentReference documentRef) { Preconditions.checkState(isEmpty(), READ_BEFORE_WRITE_ERROR_MSG); - + Tracing.getTracer().getCurrentSpan().addAnnotation(TraceUtil.SPAN_NAME_GETDOCUMENT); return ApiFutures.transform( firestore.getAll(new DocumentReference[] {documentRef}, /*fieldMask=*/ null, transactionId), new ApiFunction, DocumentSnapshot>() { diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/UpdateBuilder.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/UpdateBuilder.java index 7439fe081..056ce6840 100644 --- a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/UpdateBuilder.java +++ b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/UpdateBuilder.java @@ -166,6 +166,7 @@ public T create( private T performCreate( @Nonnull DocumentReference documentReference, @Nonnull Map fields) { verifyNotCommitted(); + Tracing.getTracer().getCurrentSpan().addAnnotation(TraceUtil.SPAN_NAME_CREATEDOCUMENT); DocumentSnapshot documentSnapshot = DocumentSnapshot.fromObject( firestore, documentReference, fields, UserDataConverter.NO_DELETES); @@ -548,7 +549,7 @@ private T performUpdate( @Nonnull Precondition precondition) { verifyNotCommitted(); Preconditions.checkArgument(!fields.isEmpty(), "Data for update() cannot be empty."); - + Tracing.getTracer().getCurrentSpan().addAnnotation(TraceUtil.SPAN_NAME_UPDATEDOCUMENT); Map deconstructedMap = expandObject(fields); DocumentSnapshot documentSnapshot = DocumentSnapshot.fromObject( @@ -611,6 +612,7 @@ public T delete(@Nonnull DocumentReference documentReference) { private T performDelete( @Nonnull DocumentReference documentReference, @Nonnull Precondition precondition) { verifyNotCommitted(); + Tracing.getTracer().getCurrentSpan().addAnnotation(TraceUtil.SPAN_NAME_DELETEDOCUMENT); Write.Builder write = Write.newBuilder().setDelete(documentReference.getName()); if (!precondition.isEmpty()) { @@ -626,7 +628,7 @@ ApiFuture> commit(@Nullable ByteString transactionId) { Tracing.getTracer() .getCurrentSpan() .addAnnotation( - "CloudFirestore.Commit", + TraceUtil.SPAN_NAME_COMMIT, ImmutableMap.of("numDocuments", AttributeValue.longAttributeValue(writes.size()))); final CommitRequest.Builder request = CommitRequest.newBuilder(); @@ -674,7 +676,7 @@ ApiFuture> bulkCommit() { Tracing.getTracer() .getCurrentSpan() .addAnnotation( - "CloudFirestore.BatchWrite", + TraceUtil.SPAN_NAME_BATCHWRITE, ImmutableMap.of("numDocuments", AttributeValue.longAttributeValue(writes.size()))); final BatchWriteRequest.Builder request = BatchWriteRequest.newBuilder(); diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Watch.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Watch.java index fc07cc585..d1e0f4948 100644 --- a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Watch.java +++ b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Watch.java @@ -36,6 +36,7 @@ import io.grpc.Status.Code; import io.grpc.StatusException; import io.grpc.StatusRuntimeException; +import io.opencensus.trace.Tracing; import java.util.ArrayList; import java.util.Collections; import java.util.Comparator; @@ -395,6 +396,7 @@ public void run() { current = false; nextAttempt = backoff.createNextAttempt(nextAttempt); + Tracing.getTracer().getCurrentSpan().addAnnotation(TraceUtil.SPAN_NAME_LISTEN); stream = firestore.streamRequest(Watch.this, firestore.getClient().listenCallable()); ListenRequest.Builder request = ListenRequest.newBuilder();