wrapResult(int writeIndex) {
* The writes in the batch are not applied atomically and can be applied out of order.
*/
ApiFuture bulkCommit() {
-
// Follows same thread safety logic as `UpdateBuilder::commit`.
committed = true;
BatchWriteRequest request = buildBatchWriteRequest();
- Tracing.getTracer()
- .getCurrentSpan()
- .addAnnotation(
- TraceUtil.SPAN_NAME_BATCHWRITE,
- ImmutableMap.of(
- "numDocuments", AttributeValue.longAttributeValue(request.getWritesCount())));
-
ApiFuture response =
processExceptions(
firestore.sendRequest(request, firestore.getClient().batchWriteCallable()));
diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/BulkWriter.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/BulkWriter.java
index 4c71a3f18..bc881979c 100644
--- a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/BulkWriter.java
+++ b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/BulkWriter.java
@@ -17,6 +17,7 @@
package com.google.cloud.firestore;
import static com.google.cloud.firestore.BulkWriterOperation.DEFAULT_BACKOFF_MAX_DELAY_MS;
+import static com.google.cloud.firestore.telemetry.TraceUtil.ATTRIBUTE_KEY_DOC_COUNT;
import com.google.api.core.ApiFunction;
import com.google.api.core.ApiFuture;
@@ -25,6 +26,9 @@
import com.google.api.core.SettableApiFuture;
import com.google.api.gax.rpc.ApiException;
import com.google.api.gax.rpc.StatusCode.Code;
+import com.google.cloud.firestore.telemetry.TraceUtil;
+import com.google.cloud.firestore.telemetry.TraceUtil.Context;
+import com.google.cloud.firestore.telemetry.TraceUtil.Scope;
import com.google.cloud.firestore.v1.FirestoreSettings;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;
@@ -219,6 +223,8 @@ enum OperationType {
@GuardedBy("lock")
private Executor errorExecutor;
+ Context traceContext;
+
/**
* Used to track when writes are enqueued. The user handler executors cannot be changed after a
* write has been enqueued.
@@ -235,6 +241,7 @@ enum OperationType {
this.successExecutor = MoreExecutors.directExecutor();
this.errorExecutor = MoreExecutors.directExecutor();
this.bulkCommitBatch = new BulkCommitBatch(firestore, bulkWriterExecutor, maxBatchSize);
+ this.traceContext = firestore.getOptions().getTraceUtil().currentContext();
if (!options.getThrottlingEnabled()) {
this.rateLimiter =
@@ -897,21 +904,32 @@ private void scheduleCurrentBatchLocked(final boolean flush) {
/** Sends the provided batch once the rate limiter does not require any delay. */
private void sendBatchLocked(final BulkCommitBatch batch, final boolean flush) {
- // Send the batch if it is does not require any delay, or schedule another attempt after the
+ // Send the batch if it does not require any delay, or schedule another attempt after the
// appropriate timeout.
boolean underRateLimit = rateLimiter.tryMakeRequest(batch.getMutationsSize());
if (underRateLimit) {
- batch
- .bulkCommit()
- .addListener(
- () -> {
- if (flush) {
- synchronized (lock) {
- scheduleCurrentBatchLocked(/* flush= */ true);
- }
+ TraceUtil.Span span =
+ firestore
+ .getOptions()
+ .getTraceUtil()
+ .startSpan(TraceUtil.SPAN_NAME_BULK_WRITER_COMMIT, traceContext)
+ .setAttribute(ATTRIBUTE_KEY_DOC_COUNT, batch.getMutationsSize());
+ try (Scope ignored = span.makeCurrent()) {
+ ApiFuture result = batch.bulkCommit();
+ result.addListener(
+ () -> {
+ if (flush) {
+ synchronized (lock) {
+ scheduleCurrentBatchLocked(/* flush= */ true);
}
- },
- bulkWriterExecutor);
+ }
+ },
+ bulkWriterExecutor);
+ span.endAtFuture(result);
+ } catch (Exception error) {
+ span.end(error);
+ throw error;
+ }
} else {
long delayMs = rateLimiter.getNextRequestDelayMs(batch.getMutationsSize());
logger.log(Level.FINE, () -> String.format("Backing off for %d seconds", delayMs / 1000));
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 8ecce62cb..3026e5183 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
@@ -21,15 +21,14 @@
import com.google.api.gax.rpc.ApiException;
import com.google.api.gax.rpc.ApiExceptions;
import com.google.api.gax.rpc.ApiStreamObserver;
+import com.google.cloud.firestore.telemetry.TraceUtil;
+import com.google.cloud.firestore.telemetry.TraceUtil.Scope;
import com.google.cloud.firestore.v1.FirestoreClient.PartitionQueryPagedResponse;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import com.google.common.util.concurrent.MoreExecutors;
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 java.util.ArrayList;
import java.util.Collections;
import java.util.List;
@@ -77,9 +76,7 @@ public void getPartitions(
PartitionQueryRequest request = buildRequest(desiredPartitionCount);
final PartitionQueryPagedResponse response;
- final TraceUtil traceUtil = TraceUtil.getInstance();
- Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_PARTITIONQUERY);
- try (Scope scope = traceUtil.getTracer().withSpan(span)) {
+ try {
response =
ApiExceptions.callAndTranslateApiException(
rpcContext.sendRequest(
@@ -94,10 +91,7 @@ public void getPartitions(
observer.onCompleted();
} catch (ApiException exception) {
- span.setStatus(Status.UNKNOWN.withDescription(exception.getMessage()));
throw FirestoreException.forApiException(exception);
- } finally {
- span.end(TraceUtil.END_SPAN_OPTIONS);
}
}
}
@@ -110,27 +104,36 @@ public ApiFuture> getPartitions(long desiredPartitionCount)
} else {
PartitionQueryRequest request = buildRequest(desiredPartitionCount);
- final TraceUtil traceUtil = TraceUtil.getInstance();
- Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_PARTITIONQUERY);
- try (Scope scope = traceUtil.getTracer().withSpan(span)) {
- return ApiFutures.transform(
- rpcContext.sendRequest(request, rpcContext.getClient().partitionQueryPagedCallable()),
- response -> {
- final ImmutableList.Builder partitions = ImmutableList.builder();
- consumePartitions(
- response,
- queryPartition -> {
- partitions.add(queryPartition);
- return null;
- });
- return partitions.build();
- },
- MoreExecutors.directExecutor());
+ TraceUtil.Span span =
+ rpcContext
+ .getFirestore()
+ .getOptions()
+ .getTraceUtil()
+ .startSpan(TraceUtil.SPAN_NAME_PARTITION_QUERY);
+ try (Scope ignored = span.makeCurrent()) {
+ ApiFuture> result =
+ ApiFutures.transform(
+ rpcContext.sendRequest(
+ request, rpcContext.getClient().partitionQueryPagedCallable()),
+ response -> {
+ final ImmutableList.Builder partitions = ImmutableList.builder();
+ consumePartitions(
+ response,
+ queryPartition -> {
+ partitions.add(queryPartition);
+ return null;
+ });
+ return partitions.build();
+ },
+ MoreExecutors.directExecutor());
+ span.endAtFuture(result);
+ return result;
} catch (ApiException exception) {
- span.setStatus(Status.UNKNOWN.withDescription(exception.getMessage()));
+ span.end(exception);
throw FirestoreException.forApiException(exception);
- } finally {
- span.end(TraceUtil.END_SPAN_OPTIONS);
+ } catch (Throwable throwable) {
+ span.end(throwable);
+ throw throwable;
}
}
}
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 a7efbea91..c736d7028 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
@@ -23,15 +23,14 @@
import com.google.api.gax.rpc.ApiExceptions;
import com.google.api.gax.rpc.UnaryCallable;
import com.google.cloud.firestore.spi.v1.FirestoreRpc;
+import com.google.cloud.firestore.telemetry.TraceUtil;
+import com.google.cloud.firestore.telemetry.TraceUtil.Scope;
import com.google.cloud.firestore.v1.FirestoreClient.ListDocumentsPagedResponse;
import com.google.common.base.Preconditions;
import com.google.common.util.concurrent.MoreExecutors;
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;
@@ -129,53 +128,59 @@ public DocumentReference document(@Nonnull String childPath) {
*/
@Nonnull
public Iterable listDocuments() {
- ListDocumentsRequest.Builder request = ListDocumentsRequest.newBuilder();
- request.setParent(options.getParentPath().toString());
- request.setCollectionId(options.getCollectionId());
- request.setMask(DocumentMask.getDefaultInstance());
- request.setShowMissing(true);
-
- final ListDocumentsPagedResponse response;
- final TraceUtil traceUtil = TraceUtil.getInstance();
- Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_LISTDOCUMENTS);
- try (Scope scope = traceUtil.getTracer().withSpan(span)) {
+ TraceUtil.Span span =
+ rpcContext
+ .getFirestore()
+ .getOptions()
+ .getTraceUtil()
+ .startSpan(TraceUtil.SPAN_NAME_COL_REF_LIST_DOCUMENTS);
+ try (Scope ignored = span.makeCurrent()) {
+ ListDocumentsRequest.Builder request = ListDocumentsRequest.newBuilder();
+ request.setParent(options.getParentPath().toString());
+ request.setCollectionId(options.getCollectionId());
+ request.setMask(DocumentMask.getDefaultInstance());
+ request.setShowMissing(true);
+ final ListDocumentsPagedResponse response;
FirestoreRpc client = rpcContext.getClient();
UnaryCallable callable =
client.listDocumentsPagedCallable();
ListDocumentsRequest build = request.build();
ApiFuture future = rpcContext.sendRequest(build, callable);
response = ApiExceptions.callAndTranslateApiException(future);
+ Iterable result =
+ new Iterable() {
+ @Override
+ @Nonnull
+ public Iterator iterator() {
+ final Iterator iterator = response.iterateAll().iterator();
+ return new Iterator() {
+ @Override
+ public boolean hasNext() {
+ return iterator.hasNext();
+ }
+
+ @Override
+ public DocumentReference next() {
+ ResourcePath path = ResourcePath.create(iterator.next().getName());
+ return document(path.getId());
+ }
+
+ @Override
+ public void remove() {
+ throw new UnsupportedOperationException("remove");
+ }
+ };
+ }
+ };
+ span.end();
+ return result;
} catch (ApiException exception) {
- span.setStatus(Status.UNKNOWN.withDescription(exception.getMessage()));
+ span.end(exception);
throw FirestoreException.forApiException(exception);
- } finally {
- span.end(TraceUtil.END_SPAN_OPTIONS);
+ } catch (Throwable throwable) {
+ span.end(throwable);
+ throw throwable;
}
-
- return new Iterable() {
- @Override
- @Nonnull
- public Iterator iterator() {
- final Iterator iterator = response.iterateAll().iterator();
- return new Iterator() {
- @Override
- public boolean hasNext() {
- return iterator.hasNext();
- }
-
- @Override
- public DocumentReference next() {
- ResourcePath path = ResourcePath.create(iterator.next().getName());
- return document(path.getId());
- }
-
- @Override
- public void remove() {
- throw new UnsupportedOperationException("remove");
- }
- };
- }
- };
}
/**
@@ -189,11 +194,24 @@ public void remove() {
*/
@Nonnull
public ApiFuture add(@Nonnull final Map fields) {
- final DocumentReference documentReference = document();
- ApiFuture createFuture = documentReference.create(fields);
-
- return ApiFutures.transform(
- createFuture, writeResult -> documentReference, MoreExecutors.directExecutor());
+ TraceUtil.Span span =
+ rpcContext
+ .getFirestore()
+ .getOptions()
+ .getTraceUtil()
+ .startSpan(TraceUtil.SPAN_NAME_COL_REF_ADD);
+ try (Scope ignored = span.makeCurrent()) {
+ final DocumentReference documentReference = document();
+ ApiFuture createFuture = documentReference.create(fields);
+ ApiFuture result =
+ ApiFutures.transform(
+ createFuture, writeResult -> documentReference, MoreExecutors.directExecutor());
+ span.endAtFuture(result);
+ return result;
+ } catch (Exception error) {
+ span.end(error);
+ throw error;
+ }
}
/**
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 7a80571a7..57254bb2b 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
@@ -21,12 +21,11 @@
import com.google.api.core.InternalExtensionOnly;
import com.google.api.gax.rpc.ApiException;
import com.google.api.gax.rpc.ApiExceptions;
+import com.google.cloud.firestore.telemetry.TraceUtil;
+import com.google.cloud.firestore.telemetry.TraceUtil.Scope;
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;
@@ -134,6 +133,12 @@ private ApiFuture extractFirst(ApiFuture> results) {
MoreExecutors.directExecutor());
}
+ /** Gets the TraceUtil object associated with this DocumentReference's Firestore instance. */
+ @Nonnull
+ private TraceUtil getTraceUtil() {
+ return getFirestore().getOptions().getTraceUtil();
+ }
+
/**
* Creates a new Document at the DocumentReference's Location. It fails the write if the document
* exists.
@@ -143,8 +148,16 @@ private ApiFuture extractFirst(ApiFuture> results) {
*/
@Nonnull
public ApiFuture create(@Nonnull Map fields) {
- WriteBatch writeBatch = rpcContext.getFirestore().batch();
- return extractFirst(writeBatch.create(this, fields).commit());
+ TraceUtil.Span span = getTraceUtil().startSpan(TraceUtil.SPAN_NAME_DOC_REF_CREATE);
+ try (Scope ignored = span.makeCurrent()) {
+ WriteBatch writeBatch = rpcContext.getFirestore().batch();
+ ApiFuture result = extractFirst(writeBatch.create(this, fields).commit());
+ span.endAtFuture(result);
+ return result;
+ } catch (Exception error) {
+ span.end(error);
+ throw error;
+ }
}
/**
@@ -156,8 +169,16 @@ public ApiFuture create(@Nonnull Map fields) {
*/
@Nonnull
public ApiFuture create(@Nonnull Object pojo) {
- WriteBatch writeBatch = rpcContext.getFirestore().batch();
- return extractFirst(writeBatch.create(this, pojo).commit());
+ TraceUtil.Span span = getTraceUtil().startSpan(TraceUtil.SPAN_NAME_DOC_REF_CREATE);
+ try (Scope ignored = span.makeCurrent()) {
+ WriteBatch writeBatch = rpcContext.getFirestore().batch();
+ ApiFuture result = extractFirst(writeBatch.create(this, pojo).commit());
+ span.endAtFuture(result);
+ return result;
+ } catch (Exception error) {
+ span.end(error);
+ throw error;
+ }
}
/**
@@ -169,8 +190,16 @@ public ApiFuture create(@Nonnull Object pojo) {
*/
@Nonnull
public ApiFuture set(@Nonnull Map fields) {
- WriteBatch writeBatch = rpcContext.getFirestore().batch();
- return extractFirst(writeBatch.set(this, fields).commit());
+ TraceUtil.Span span = getTraceUtil().startSpan(TraceUtil.SPAN_NAME_DOC_REF_SET);
+ try (Scope ignored = span.makeCurrent()) {
+ WriteBatch writeBatch = rpcContext.getFirestore().batch();
+ ApiFuture result = extractFirst(writeBatch.set(this, fields).commit());
+ span.endAtFuture(result);
+ return result;
+ } catch (Exception error) {
+ span.end(error);
+ throw error;
+ }
}
/**
@@ -185,8 +214,16 @@ public ApiFuture set(@Nonnull Map fields) {
@Nonnull
public ApiFuture set(
@Nonnull Map fields, @Nonnull SetOptions options) {
- WriteBatch writeBatch = rpcContext.getFirestore().batch();
- return extractFirst(writeBatch.set(this, fields, options).commit());
+ TraceUtil.Span span = getTraceUtil().startSpan(TraceUtil.SPAN_NAME_DOC_REF_SET);
+ try (Scope ignored = span.makeCurrent()) {
+ WriteBatch writeBatch = rpcContext.getFirestore().batch();
+ ApiFuture result = extractFirst(writeBatch.set(this, fields, options).commit());
+ span.endAtFuture(result);
+ return result;
+ } catch (Exception error) {
+ span.end(error);
+ throw error;
+ }
}
/**
@@ -198,8 +235,16 @@ public ApiFuture set(
*/
@Nonnull
public ApiFuture set(@Nonnull Object pojo) {
- WriteBatch writeBatch = rpcContext.getFirestore().batch();
- return extractFirst(writeBatch.set(this, pojo).commit());
+ TraceUtil.Span span = getTraceUtil().startSpan(TraceUtil.SPAN_NAME_DOC_REF_SET);
+ try (Scope ignored = span.makeCurrent()) {
+ WriteBatch writeBatch = rpcContext.getFirestore().batch();
+ ApiFuture result = extractFirst(writeBatch.set(this, pojo).commit());
+ span.endAtFuture(result);
+ return result;
+ } catch (Exception error) {
+ span.end(error);
+ throw error;
+ }
}
/**
@@ -213,8 +258,16 @@ public ApiFuture set(@Nonnull Object pojo) {
*/
@Nonnull
public ApiFuture set(@Nonnull Object pojo, @Nonnull SetOptions options) {
- WriteBatch writeBatch = rpcContext.getFirestore().batch();
- return extractFirst(writeBatch.set(this, pojo, options).commit());
+ TraceUtil.Span span = getTraceUtil().startSpan(TraceUtil.SPAN_NAME_DOC_REF_SET);
+ try (Scope ignored = span.makeCurrent()) {
+ WriteBatch writeBatch = rpcContext.getFirestore().batch();
+ ApiFuture result = extractFirst(writeBatch.set(this, pojo, options).commit());
+ span.endAtFuture(result);
+ return result;
+ } catch (Exception error) {
+ span.end(error);
+ throw error;
+ }
}
/**
@@ -226,8 +279,16 @@ public ApiFuture set(@Nonnull Object pojo, @Nonnull SetOptions opti
*/
@Nonnull
public ApiFuture update(@Nonnull Map fields) {
- WriteBatch writeBatch = rpcContext.getFirestore().batch();
- return extractFirst(writeBatch.update(this, fields).commit());
+ TraceUtil.Span span = getTraceUtil().startSpan(TraceUtil.SPAN_NAME_DOC_REF_UPDATE);
+ try (Scope ignored = span.makeCurrent()) {
+ WriteBatch writeBatch = rpcContext.getFirestore().batch();
+ ApiFuture result = extractFirst(writeBatch.update(this, fields).commit());
+ span.endAtFuture(result);
+ return result;
+ } catch (Exception error) {
+ span.end(error);
+ throw error;
+ }
}
/**
@@ -240,8 +301,17 @@ public ApiFuture update(@Nonnull Map fields) {
*/
@Nonnull
public ApiFuture update(@Nonnull Map fields, Precondition options) {
- WriteBatch writeBatch = rpcContext.getFirestore().batch();
- return extractFirst(writeBatch.update(this, fields, options).commit());
+ TraceUtil.Span span = getTraceUtil().startSpan(TraceUtil.SPAN_NAME_DOC_REF_UPDATE);
+ try (Scope ignored = span.makeCurrent()) {
+ WriteBatch writeBatch = rpcContext.getFirestore().batch();
+ ApiFuture result =
+ extractFirst(writeBatch.update(this, fields, options).commit());
+ span.endAtFuture(result);
+ return result;
+ } catch (Exception error) {
+ span.end(error);
+ throw error;
+ }
}
/**
@@ -256,8 +326,17 @@ public ApiFuture update(@Nonnull Map fields, Precon
@Nonnull
public ApiFuture update(
@Nonnull String field, @Nullable Object value, Object... moreFieldsAndValues) {
- WriteBatch writeBatch = rpcContext.getFirestore().batch();
- return extractFirst(writeBatch.update(this, field, value, moreFieldsAndValues).commit());
+ TraceUtil.Span span = getTraceUtil().startSpan(TraceUtil.SPAN_NAME_DOC_REF_UPDATE);
+ try (Scope ignored = span.makeCurrent()) {
+ WriteBatch writeBatch = rpcContext.getFirestore().batch();
+ ApiFuture result =
+ extractFirst(writeBatch.update(this, field, value, moreFieldsAndValues).commit());
+ span.endAtFuture(result);
+ return result;
+ } catch (Exception error) {
+ span.end(error);
+ throw error;
+ }
}
/**
@@ -272,8 +351,17 @@ public ApiFuture update(
@Nonnull
public ApiFuture update(
@Nonnull FieldPath fieldPath, @Nullable Object value, Object... moreFieldsAndValues) {
- WriteBatch writeBatch = rpcContext.getFirestore().batch();
- return extractFirst(writeBatch.update(this, fieldPath, value, moreFieldsAndValues).commit());
+ TraceUtil.Span span = getTraceUtil().startSpan(TraceUtil.SPAN_NAME_DOC_REF_UPDATE);
+ try (Scope ignored = span.makeCurrent()) {
+ WriteBatch writeBatch = rpcContext.getFirestore().batch();
+ ApiFuture result =
+ extractFirst(writeBatch.update(this, fieldPath, value, moreFieldsAndValues).commit());
+ span.endAtFuture(result);
+ return result;
+ } catch (Exception error) {
+ span.end(error);
+ throw error;
+ }
}
/**
@@ -292,9 +380,18 @@ public ApiFuture update(
@Nonnull String field,
@Nullable Object value,
Object... moreFieldsAndValues) {
- WriteBatch writeBatch = rpcContext.getFirestore().batch();
- return extractFirst(
- writeBatch.update(this, options, field, value, moreFieldsAndValues).commit());
+ TraceUtil.Span span = getTraceUtil().startSpan(TraceUtil.SPAN_NAME_DOC_REF_UPDATE);
+ try (Scope ignored = span.makeCurrent()) {
+ WriteBatch writeBatch = rpcContext.getFirestore().batch();
+ ApiFuture result =
+ extractFirst(
+ writeBatch.update(this, options, field, value, moreFieldsAndValues).commit());
+ span.endAtFuture(result);
+ return result;
+ } catch (Exception error) {
+ span.end(error);
+ throw error;
+ }
}
/**
@@ -313,9 +410,18 @@ public ApiFuture update(
@Nonnull FieldPath fieldPath,
@Nullable Object value,
Object... moreFieldsAndValues) {
- WriteBatch writeBatch = rpcContext.getFirestore().batch();
- return extractFirst(
- writeBatch.update(this, options, fieldPath, value, moreFieldsAndValues).commit());
+ TraceUtil.Span span = getTraceUtil().startSpan(TraceUtil.SPAN_NAME_DOC_REF_UPDATE);
+ try (Scope ignored = span.makeCurrent()) {
+ WriteBatch writeBatch = rpcContext.getFirestore().batch();
+ ApiFuture result =
+ extractFirst(
+ writeBatch.update(this, options, fieldPath, value, moreFieldsAndValues).commit());
+ span.endAtFuture(result);
+ return result;
+ } catch (Exception error) {
+ span.end(error);
+ throw error;
+ }
}
/**
@@ -326,8 +432,16 @@ public ApiFuture update(
*/
@Nonnull
public ApiFuture delete(@Nonnull Precondition options) {
- WriteBatch writeBatch = rpcContext.getFirestore().batch();
- return extractFirst(writeBatch.delete(this, options).commit());
+ TraceUtil.Span span = getTraceUtil().startSpan(TraceUtil.SPAN_NAME_DOC_REF_DELETE);
+ try (Scope ignored = span.makeCurrent()) {
+ WriteBatch writeBatch = rpcContext.getFirestore().batch();
+ ApiFuture result = extractFirst(writeBatch.delete(this, options).commit());
+ span.endAtFuture(result);
+ return result;
+ } catch (Exception error) {
+ span.end(error);
+ throw error;
+ }
}
/**
@@ -337,20 +451,36 @@ public ApiFuture delete(@Nonnull Precondition options) {
*/
@Nonnull
public ApiFuture delete() {
- WriteBatch writeBatch = rpcContext.getFirestore().batch();
- return extractFirst(writeBatch.delete(this).commit());
+ TraceUtil.Span span = getTraceUtil().startSpan(TraceUtil.SPAN_NAME_DOC_REF_DELETE);
+ try (Scope ignored = span.makeCurrent()) {
+ WriteBatch writeBatch = rpcContext.getFirestore().batch();
+ ApiFuture result = extractFirst(writeBatch.delete(this).commit());
+ span.endAtFuture(result);
+ return result;
+ } catch (Exception error) {
+ span.end(error);
+ throw error;
+ }
}
/**
* Reads the document referenced by this DocumentReference. If the document doesn't exist, the
- * get() will return an an empty DocumentSnapshot.
+ * get() will return an empty DocumentSnapshot.
*
* @return An ApiFuture that will be resolved with the contents of the Document at this
* DocumentReference, or a failure if the document does not exist.
*/
@Nonnull
public ApiFuture get() {
- return extractFirst(rpcContext.getFirestore().getAll(this));
+ TraceUtil.Span span = getTraceUtil().startSpan(TraceUtil.SPAN_NAME_DOC_REF_GET);
+ try (Scope ignored = span.makeCurrent()) {
+ ApiFuture result = extractFirst(rpcContext.getFirestore().getAll(this));
+ span.endAtFuture(result);
+ return result;
+ } catch (Exception error) {
+ span.end(error);
+ throw error;
+ }
}
/**
@@ -363,8 +493,16 @@ public ApiFuture get() {
*/
@Nonnull
public ApiFuture get(FieldMask fieldMask) {
- return extractFirst(
- rpcContext.getFirestore().getAll(new DocumentReference[] {this}, fieldMask));
+ TraceUtil.Span span = getTraceUtil().startSpan(TraceUtil.SPAN_NAME_DOC_REF_GET);
+ try (Scope ignored = span.makeCurrent()) {
+ ApiFuture result =
+ extractFirst(rpcContext.getFirestore().getAll(new DocumentReference[] {this}, fieldMask));
+ span.endAtFuture(result);
+ return result;
+ } catch (Exception error) {
+ span.end(error);
+ throw error;
+ }
}
/**
@@ -375,46 +513,45 @@ public ApiFuture get(FieldMask fieldMask) {
*/
@Nonnull
public Iterable listCollections() {
- ListCollectionIdsRequest.Builder request = ListCollectionIdsRequest.newBuilder();
- request.setParent(path.toString());
- final ListCollectionIdsPagedResponse response;
- final TraceUtil traceUtil = TraceUtil.getInstance();
- Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_LISTCOLLECTIONIDS);
- try (Scope scope = traceUtil.getTracer().withSpan(span)) {
+ TraceUtil.Span span = getTraceUtil().startSpan(TraceUtil.SPAN_NAME_DOC_REF_LIST_COLLECTIONS);
+ try (Scope ignored = span.makeCurrent()) {
+ ListCollectionIdsRequest.Builder request = ListCollectionIdsRequest.newBuilder();
+ request.setParent(path.toString());
+ final ListCollectionIdsPagedResponse response;
response =
ApiExceptions.callAndTranslateApiException(
rpcContext.sendRequest(
request.build(), rpcContext.getClient().listCollectionIdsPagedCallable()));
+ Iterable result =
+ new Iterable() {
+ @Override
+ @Nonnull
+ public Iterator iterator() {
+ final Iterator iterator = response.iterateAll().iterator();
+ return new Iterator() {
+ @Override
+ public boolean hasNext() {
+ return iterator.hasNext();
+ }
+
+ @Override
+ public CollectionReference next() {
+ return DocumentReference.this.collection(iterator.next());
+ }
+
+ @Override
+ public void remove() {
+ throw new UnsupportedOperationException("remove");
+ }
+ };
+ }
+ };
+ span.end();
+ return result;
} catch (ApiException exception) {
- span.setStatus(Status.UNKNOWN.withDescription(exception.getMessage()));
+ span.end(exception);
throw FirestoreException.forApiException(exception);
- } finally {
- span.end(TraceUtil.END_SPAN_OPTIONS);
}
-
- return new Iterable() {
- @Override
- @Nonnull
- public Iterator iterator() {
- final Iterator iterator = response.iterateAll().iterator();
- return new Iterator() {
- @Override
- public boolean hasNext() {
- return iterator.hasNext();
- }
-
- @Override
- public CollectionReference next() {
- return DocumentReference.this.collection(iterator.next());
- }
-
- @Override
- public void remove() {
- throw new UnsupportedOperationException("remove");
- }
- };
- }
- };
}
/**
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 82dfc5176..48c691466 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
@@ -16,6 +16,8 @@
package com.google.cloud.firestore;
+import static com.google.cloud.firestore.telemetry.TraceUtil.*;
+
import com.google.api.core.ApiClock;
import com.google.api.core.ApiFuture;
import com.google.api.core.NanoClock;
@@ -30,6 +32,7 @@
import com.google.api.gax.rpc.UnaryCallable;
import com.google.cloud.Timestamp;
import com.google.cloud.firestore.spi.v1.FirestoreRpc;
+import com.google.cloud.firestore.telemetry.TraceUtil;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableMap;
@@ -37,11 +40,9 @@
import com.google.firestore.v1.BatchGetDocumentsResponse;
import com.google.firestore.v1.DatabaseRootName;
import com.google.protobuf.ByteString;
-import io.opencensus.trace.AttributeValue;
-import io.opencensus.trace.Tracer;
-import io.opencensus.trace.Tracing;
import java.security.SecureRandom;
import java.util.ArrayList;
+import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
@@ -61,8 +62,6 @@ class FirestoreImpl implements Firestore, FirestoreRpcContext {
private static final String AUTO_ID_ALPHABET =
"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789";
- private static final Tracer tracer = Tracing.getTracer();
-
private final FirestoreRpc firestoreClient;
private final FirestoreOptions firestoreOptions;
private final ResourcePath databasePath;
@@ -90,6 +89,12 @@ class FirestoreImpl implements Firestore, FirestoreRpcContext {
ResourcePath.create(DatabaseRootName.of(options.getProjectId(), options.getDatabaseId()));
}
+ /** Gets the TraceUtil object associated with this Firestore instance. */
+ @Nonnull
+ private TraceUtil getTraceUtil() {
+ return getOptions().getTraceUtil();
+ }
+
/** Lazy-load the Firestore's default BulkWriter. */
private BulkWriter getBulkWriter() {
if (bulkWriterInstance == null) {
@@ -218,14 +223,26 @@ void getAll(
@Nullable ByteString transactionId,
@Nullable com.google.protobuf.Timestamp readTime,
final ApiStreamObserver apiStreamObserver) {
+ // To reduce the size of traces, we only register one event for every 100 responses
+ // that we receive from the server.
+ final int NUM_RESPONSES_PER_TRACE_EVENT = 100;
ResponseObserver responseObserver =
new ResponseObserver() {
- int numResponses;
+ int numResponses = 0;
boolean hasCompleted = false;
@Override
- public void onStart(StreamController streamController) {}
+ public void onStart(StreamController streamController) {
+ getTraceUtil()
+ .currentSpan()
+ .addEvent(
+ TraceUtil.SPAN_NAME_BATCH_GET_DOCUMENTS + ": Start",
+ new ImmutableMap.Builder()
+ .put(ATTRIBUTE_KEY_DOC_COUNT, documentReferences.length)
+ .put(ATTRIBUTE_KEY_IS_TRANSACTIONAL, transactionId != null)
+ .build());
+ }
@Override
public void onResponse(BatchGetDocumentsResponse response) {
@@ -234,14 +251,17 @@ public void onResponse(BatchGetDocumentsResponse response) {
numResponses++;
if (numResponses == 1) {
- tracer
- .getCurrentSpan()
- .addAnnotation(TraceUtil.SPAN_NAME_BATCHGETDOCUMENTS + ": First response");
- } else if (numResponses % 100 == 0) {
- tracer
- .getCurrentSpan()
- .addAnnotation(
- TraceUtil.SPAN_NAME_BATCHGETDOCUMENTS + ": Received 100 responses");
+ getTraceUtil()
+ .currentSpan()
+ .addEvent(TraceUtil.SPAN_NAME_BATCH_GET_DOCUMENTS + ": First response received");
+ } else if (numResponses % NUM_RESPONSES_PER_TRACE_EVENT == 0) {
+ getTraceUtil()
+ .currentSpan()
+ .addEvent(
+ TraceUtil.SPAN_NAME_BATCH_GET_DOCUMENTS
+ + ": Received "
+ + numResponses
+ + " responses");
}
switch (response.getResultCase()) {
@@ -277,9 +297,7 @@ public void onResponse(BatchGetDocumentsResponse response) {
@Override
public void onError(Throwable throwable) {
- tracer
- .getCurrentSpan()
- .addAnnotation(TraceUtil.SPAN_NAME_BATCHGETDOCUMENTS + ": Error");
+ getTraceUtil().currentSpan().end(throwable);
apiStreamObserver.onError(throwable);
}
@@ -287,9 +305,14 @@ public void onError(Throwable throwable) {
public void onComplete() {
if (hasCompleted) return;
hasCompleted = true;
- tracer
- .getCurrentSpan()
- .addAnnotation(TraceUtil.SPAN_NAME_BATCHGETDOCUMENTS + ": Complete");
+ getTraceUtil()
+ .currentSpan()
+ .addEvent(
+ TraceUtil.SPAN_NAME_BATCH_GET_DOCUMENTS
+ + ": Completed with "
+ + numResponses
+ + " responses.",
+ Collections.singletonMap(ATTRIBUTE_KEY_NUM_RESPONSES, numResponses));
apiStreamObserver.onCompleted();
}
};
@@ -313,13 +336,6 @@ public void onComplete() {
request.addDocuments(docRef.getName());
}
- tracer
- .getCurrentSpan()
- .addAnnotation(
- TraceUtil.SPAN_NAME_BATCHGETDOCUMENTS + ": Start",
- ImmutableMap.of(
- "numDocuments", AttributeValue.longAttributeValue(documentReferences.length)));
-
streamRequest(request.build(), responseObserver, firestoreClient.batchGetDocumentsCallable());
}
diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/FirestoreOpenTelemetryOptions.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/FirestoreOpenTelemetryOptions.java
new file mode 100644
index 000000000..2b4606565
--- /dev/null
+++ b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/FirestoreOpenTelemetryOptions.java
@@ -0,0 +1,102 @@
+/*
+ * Copyright 2024 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.firestore;
+
+import com.google.api.core.BetaApi;
+import io.opentelemetry.api.OpenTelemetry;
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
+/**
+ * Represents the options that are used to configure the use of OpenTelemetry for telemetry
+ * collection in the Firestore SDK.
+ */
+@BetaApi
+public class FirestoreOpenTelemetryOptions {
+ private final boolean tracingEnabled;
+ private final @Nullable OpenTelemetry openTelemetry;
+
+ FirestoreOpenTelemetryOptions(Builder builder) {
+ this.tracingEnabled = builder.tracingEnabled;
+ this.openTelemetry = builder.openTelemetry;
+ }
+
+ public boolean isTracingEnabled() {
+ return tracingEnabled;
+ }
+
+ public OpenTelemetry getOpenTelemetry() {
+ return openTelemetry;
+ }
+
+ @Nonnull
+ public FirestoreOpenTelemetryOptions.Builder toBuilder() {
+ return new FirestoreOpenTelemetryOptions.Builder(this);
+ }
+
+ @Nonnull
+ public static FirestoreOpenTelemetryOptions.Builder newBuilder() {
+ return new FirestoreOpenTelemetryOptions.Builder();
+ }
+
+ public static class Builder {
+
+ private boolean tracingEnabled;
+
+ @Nullable private OpenTelemetry openTelemetry;
+
+ private Builder() {
+ tracingEnabled = false;
+ openTelemetry = null;
+ }
+
+ private Builder(FirestoreOpenTelemetryOptions options) {
+ this.tracingEnabled = options.tracingEnabled;
+ this.openTelemetry = options.openTelemetry;
+ }
+
+ @Nonnull
+ public FirestoreOpenTelemetryOptions build() {
+ return new FirestoreOpenTelemetryOptions(this);
+ }
+
+ /**
+ * Sets whether tracing should be enabled.
+ *
+ * @param tracingEnabled Whether tracing should be enabled.
+ */
+ @Nonnull
+ public FirestoreOpenTelemetryOptions.Builder setTracingEnabled(boolean tracingEnabled) {
+ this.tracingEnabled = tracingEnabled;
+ return this;
+ }
+
+ /**
+ * Sets the {@link OpenTelemetry} to use with this Firestore instance. If telemetry collection
+ * is enabled, but an `OpenTelemetry` is not provided, the Firestore SDK will attempt to use the
+ * `GlobalOpenTelemetry`.
+ *
+ * @param openTelemetry The OpenTelemetry that should be used by this Firestore instance.
+ */
+ @Nonnull
+ public FirestoreOpenTelemetryOptions.Builder setOpenTelemetry(
+ @Nonnull OpenTelemetry openTelemetry) {
+ this.openTelemetry = openTelemetry;
+ return this;
+ }
+ }
+}
diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/FirestoreOptions.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/FirestoreOptions.java
index 47de09d39..b80b9dacd 100644
--- a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/FirestoreOptions.java
+++ b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/FirestoreOptions.java
@@ -16,6 +16,8 @@
package com.google.cloud.firestore;
+import com.google.api.core.ApiFunction;
+import com.google.api.core.BetaApi;
import com.google.api.core.InternalApi;
import com.google.api.gax.core.CredentialsProvider;
import com.google.api.gax.core.FixedCredentialsProvider;
@@ -60,6 +62,8 @@ public final class FirestoreOptions extends ServiceOptions {
@Nullable private String databaseId = null;
@Nullable private TransportChannelProvider channelProvider = null;
@Nullable private CredentialsProvider credentialsProvider = null;
@Nullable private String emulatorHost = null;
+ @Nullable private FirestoreOpenTelemetryOptions openTelemetryOptions = null;
private Builder() {}
@@ -133,6 +149,7 @@ private Builder(FirestoreOptions options) {
this.channelProvider = options.channelProvider;
this.credentialsProvider = options.credentialsProvider;
this.emulatorHost = options.emulatorHost;
+ this.openTelemetryOptions = options.openTelemetryOptions;
}
/**
@@ -201,6 +218,19 @@ public Builder setDatabaseId(@Nonnull String databaseId) {
return this;
}
+ /**
+ * Sets the {@link FirestoreOpenTelemetryOptions} to be used for this Firestore instance.
+ *
+ * @param openTelemetryOptions The `FirestoreOpenTelemetryOptions` to use.
+ */
+ @BetaApi
+ @Nonnull
+ public Builder setOpenTelemetryOptions(
+ @Nonnull FirestoreOpenTelemetryOptions openTelemetryOptions) {
+ this.openTelemetryOptions = openTelemetryOptions;
+ return this;
+ }
+
@Override
@Nonnull
public FirestoreOptions build() {
@@ -212,6 +242,10 @@ public FirestoreOptions build() {
}
}
+ if (this.openTelemetryOptions == null) {
+ this.setOpenTelemetryOptions(FirestoreOpenTelemetryOptions.newBuilder().build());
+ }
+
// Override credentials and channel provider if we are using the emulator.
if (emulatorHost == null) {
emulatorHost = System.getenv(FIRESTORE_EMULATOR_SYSTEM_VARIABLE);
@@ -278,16 +312,37 @@ public void refresh() {}
protected FirestoreOptions(Builder builder) {
super(FirestoreFactory.class, FirestoreRpcFactory.class, builder, new FirestoreDefaults());
+ // FirestoreOptions must contain non-null open-telemetry options.
+ // If the builder doesn't have any open-telemetry options, use a default (disabled) one.
+ this.openTelemetryOptions =
+ builder.openTelemetryOptions != null
+ ? builder.openTelemetryOptions
+ : FirestoreOpenTelemetryOptions.newBuilder().build();
+ this.traceUtil = com.google.cloud.firestore.telemetry.TraceUtil.getInstance(this);
+
this.databaseId =
builder.databaseId != null
? builder.databaseId
: FirestoreDefaults.INSTANCE.getDatabaseId();
- this.channelProvider =
- builder.channelProvider != null
- ? builder.channelProvider
- : GrpcTransportOptions.setUpChannelProvider(
+ if (builder.channelProvider == null) {
+ ApiFunction channelConfigurator =
+ this.traceUtil.getChannelConfigurator();
+ if (channelConfigurator == null) {
+ this.channelProvider =
+ GrpcTransportOptions.setUpChannelProvider(
FirestoreSettings.defaultGrpcTransportProviderBuilder(), this);
+ } else {
+ // Intercept the grpc channel calls to add telemetry info.
+ this.channelProvider =
+ GrpcTransportOptions.setUpChannelProvider(
+ FirestoreSettings.defaultGrpcTransportProviderBuilder()
+ .setChannelConfigurator(channelConfigurator),
+ this);
+ }
+ } else {
+ this.channelProvider = builder.channelProvider;
+ }
this.credentialsProvider =
builder.credentialsProvider != null
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 d71bc5a30..4721ba93d 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
@@ -16,6 +16,7 @@
package com.google.cloud.firestore;
+import static com.google.cloud.firestore.telemetry.TraceUtil.*;
import static com.google.common.collect.Lists.reverse;
import static com.google.firestore.v1.StructuredQuery.FieldFilter.Operator.ARRAY_CONTAINS;
import static com.google.firestore.v1.StructuredQuery.FieldFilter.Operator.ARRAY_CONTAINS_ANY;
@@ -38,6 +39,8 @@
import com.google.auto.value.AutoValue;
import com.google.cloud.Timestamp;
import com.google.cloud.firestore.Query.QueryOptions.Builder;
+import com.google.cloud.firestore.telemetry.TraceUtil;
+import com.google.cloud.firestore.telemetry.TraceUtil.Scope;
import com.google.cloud.firestore.v1.FirestoreSettings;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
@@ -58,8 +61,6 @@
import com.google.protobuf.ByteString;
import com.google.protobuf.Int32Value;
import io.grpc.Status;
-import io.opencensus.trace.AttributeValue;
-import io.opencensus.trace.Tracing;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
@@ -1526,7 +1527,8 @@ public void onCompleted() {
/* startTimeNanos= */ rpcContext.getClock().nanoTime(),
/* transactionId= */ null,
/* readTime= */ null,
- /* explainOptions= */ null);
+ /* explainOptions= */ null,
+ /* isRetryRequestWithCursor= */ false);
}
/**
@@ -1582,7 +1584,8 @@ public void onCompleted() {
/* startTimeNanos= */ rpcContext.getClock().nanoTime(),
/* transactionId= */ null,
/* readTime= */ null,
- /* explainOptions= */ options);
+ /* explainOptions= */ options,
+ /* isRetryRequestWithCursor= */ false);
return metricsFuture;
}
@@ -1706,7 +1709,13 @@ private void internalStream(
final long startTimeNanos,
@Nullable final ByteString transactionId,
@Nullable final Timestamp readTime,
- @Nullable final ExplainOptions explainOptions) {
+ @Nullable final ExplainOptions explainOptions,
+ final boolean isRetryRequestWithCursor) {
+ TraceUtil traceUtil = getFirestore().getOptions().getTraceUtil();
+ // To reduce the size of traces, we only register one event for every 100 responses
+ // that we receive from the server.
+ final int NUM_RESPONSES_PER_TRACE_EVENT = 100;
+
RunQueryRequest.Builder request = RunQueryRequest.newBuilder();
request.setStructuredQuery(buildQuery()).setParent(options.getParentPath().toString());
@@ -1721,19 +1730,21 @@ private void internalStream(
request.setReadTime(readTime.toProto());
}
- Tracing.getTracer()
- .getCurrentSpan()
- .addAnnotation(
- TraceUtil.SPAN_NAME_RUNQUERY + ": Start",
- ImmutableMap.of(
- "transactional", AttributeValue.booleanAttributeValue(transactionId != null)));
+ TraceUtil.Span currentSpan = traceUtil.currentSpan();
+ currentSpan.addEvent(
+ TraceUtil.SPAN_NAME_RUN_QUERY,
+ new ImmutableMap.Builder()
+ .put(ATTRIBUTE_KEY_IS_TRANSACTIONAL, transactionId != null)
+ .put(ATTRIBUTE_KEY_IS_RETRY_WITH_CURSOR, isRetryRequestWithCursor)
+ .build());
final AtomicReference lastReceivedDocument = new AtomicReference<>();
ResponseObserver observer =
new ResponseObserver() {
- boolean firstResponse;
- int numDocuments;
+ Timestamp readTime;
+ boolean firstResponse = false;
+ int numDocuments = 0;
// The stream's `onComplete()` could be called more than once,
// this flag makes sure only the first one is actually processed.
@@ -1746,17 +1757,16 @@ public void onStart(StreamController streamController) {}
public void onResponse(RunQueryResponse response) {
if (!firstResponse) {
firstResponse = true;
- Tracing.getTracer().getCurrentSpan().addAnnotation("Firestore.Query: First response");
+ currentSpan.addEvent(TraceUtil.SPAN_NAME_RUN_QUERY + ": First Response");
}
runQueryResponseObserver.onNext(response);
if (response.hasDocument()) {
numDocuments++;
- if (numDocuments % 100 == 0) {
- Tracing.getTracer()
- .getCurrentSpan()
- .addAnnotation("Firestore.Query: Received 100 documents");
+ if (numDocuments % NUM_RESPONSES_PER_TRACE_EVENT == 0) {
+ currentSpan.addEvent(
+ TraceUtil.SPAN_NAME_RUN_QUERY + ": Received " + numDocuments + " documents");
}
Document document = response.getDocument();
QueryDocumentSnapshot documentSnapshot =
@@ -1766,12 +1776,8 @@ public void onResponse(RunQueryResponse response) {
}
if (response.getDone()) {
- Tracing.getTracer()
- .getCurrentSpan()
- .addAnnotation(
- "Firestore.Query: Completed",
- ImmutableMap.of(
- "numDocuments", AttributeValue.longAttributeValue(numDocuments)));
+ currentSpan.addEvent(
+ TraceUtil.SPAN_NAME_RUN_QUERY + ": Received RunQueryResponse.Done");
onComplete();
}
}
@@ -1780,9 +1786,9 @@ public void onResponse(RunQueryResponse response) {
public void onError(Throwable throwable) {
QueryDocumentSnapshot cursor = lastReceivedDocument.get();
if (shouldRetry(cursor, throwable)) {
- Tracing.getTracer()
- .getCurrentSpan()
- .addAnnotation("Firestore.Query: Retryable Error");
+ currentSpan.addEvent(
+ TraceUtil.SPAN_NAME_RUN_QUERY + ": Retryable Error",
+ Collections.singletonMap("error.message", throwable.getMessage()));
Query.this
.startAfter(cursor)
@@ -1791,10 +1797,12 @@ public void onError(Throwable throwable) {
startTimeNanos,
/* transactionId= */ null,
options.getRequireConsistency() ? cursor.getReadTime() : null,
- explainOptions);
-
+ explainOptions,
+ /* isRetryRequestWithCursor= */ true);
} else {
- Tracing.getTracer().getCurrentSpan().addAnnotation("Firestore.Query: Error");
+ currentSpan.addEvent(
+ TraceUtil.SPAN_NAME_RUN_QUERY + ": Error",
+ Collections.singletonMap("error.message", throwable.getMessage()));
runQueryResponseObserver.onError(throwable);
}
}
@@ -1803,13 +1811,9 @@ public void onError(Throwable throwable) {
public void onComplete() {
if (hasCompleted) return;
hasCompleted = true;
-
- Tracing.getTracer()
- .getCurrentSpan()
- .addAnnotation(
- "Firestore.Query: Completed",
- ImmutableMap.of(
- "numDocuments", AttributeValue.longAttributeValue(numDocuments)));
+ currentSpan.addEvent(
+ TraceUtil.SPAN_NAME_RUN_QUERY + ": Completed",
+ Collections.singletonMap(ATTRIBUTE_KEY_DOC_COUNT, numDocuments));
runQueryResponseObserver.onCompleted();
}
@@ -1856,68 +1860,77 @@ public ApiFuture get() {
*/
@Nonnull
public ApiFuture> explain(ExplainOptions options) {
- final SettableApiFuture> result = SettableApiFuture.create();
-
- internalStream(
- new ApiStreamObserver() {
- @Nullable List documentSnapshots = null;
- Timestamp readTime;
- ExplainMetrics metrics;
+ TraceUtil.Span span =
+ getFirestore().getOptions().getTraceUtil().startSpan(TraceUtil.SPAN_NAME_QUERY_GET);
+
+ try (Scope ignored = span.makeCurrent()) {
+ final SettableApiFuture> result = SettableApiFuture.create();
+ internalStream(
+ new ApiStreamObserver() {
+ @Nullable List documentSnapshots = null;
+ Timestamp readTime;
+ ExplainMetrics metrics;
+
+ @Override
+ public void onNext(RunQueryResponse runQueryResponse) {
+ if (runQueryResponse.hasDocument()) {
+ if (documentSnapshots == null) {
+ documentSnapshots = new ArrayList<>();
+ }
+
+ Document document = runQueryResponse.getDocument();
+ QueryDocumentSnapshot documentSnapshot =
+ QueryDocumentSnapshot.fromDocument(
+ rpcContext, Timestamp.fromProto(runQueryResponse.getReadTime()), document);
+ documentSnapshots.add(documentSnapshot);
+ }
- @Override
- public void onNext(RunQueryResponse runQueryResponse) {
- if (runQueryResponse.hasDocument()) {
- if (documentSnapshots == null) {
- documentSnapshots = new ArrayList<>();
+ if (readTime == null) {
+ readTime = Timestamp.fromProto(runQueryResponse.getReadTime());
}
- Document document = runQueryResponse.getDocument();
- QueryDocumentSnapshot documentSnapshot =
- QueryDocumentSnapshot.fromDocument(
- rpcContext, Timestamp.fromProto(runQueryResponse.getReadTime()), document);
- documentSnapshots.add(documentSnapshot);
+ if (runQueryResponse.hasExplainMetrics()) {
+ metrics = new ExplainMetrics(runQueryResponse.getExplainMetrics());
+ if (documentSnapshots == null && metrics.getExecutionStats() != null) {
+ // This indicates that the query was executed, but no documents
+ // had matched the query. Create an empty list.
+ documentSnapshots = Collections.emptyList();
+ }
+ }
}
- if (readTime == null) {
- readTime = Timestamp.fromProto(runQueryResponse.getReadTime());
+ @Override
+ public void onError(Throwable throwable) {
+ result.setException(throwable);
}
- if (runQueryResponse.hasExplainMetrics()) {
- metrics = new ExplainMetrics(runQueryResponse.getExplainMetrics());
- if (documentSnapshots == null && metrics.getExecutionStats() != null) {
- // This indicates that the query was executed, but no documents
- // had matched the query. Create an empty list.
- documentSnapshots = Collections.emptyList();
+ @Override
+ public void onCompleted() {
+ @Nullable QuerySnapshot snapshot = null;
+ if (documentSnapshots != null) {
+ // The results for limitToLast queries need to be flipped since we reversed the
+ // ordering constraints before sending the query to the backend.
+ List resultView =
+ LimitType.Last.equals(Query.this.options.getLimitType())
+ ? reverse(documentSnapshots)
+ : documentSnapshots;
+ snapshot = QuerySnapshot.withDocuments(Query.this, readTime, resultView);
}
+ result.set(new ExplainResults<>(metrics, snapshot));
}
- }
-
- @Override
- public void onError(Throwable throwable) {
- result.setException(throwable);
- }
+ },
+ /* startTimeNanos= */ rpcContext.getClock().nanoTime(),
+ /* transactionId= */ null,
+ /* readTime= */ null,
+ /* explainOptions= */ options,
+ /* isRetryRequestWithCursor= */ false);
- @Override
- public void onCompleted() {
- @Nullable QuerySnapshot snapshot = null;
- if (documentSnapshots != null) {
- // The results for limitToLast queries need to be flipped since we reversed the
- // ordering constraints before sending the query to the backend.
- List resultView =
- LimitType.Last.equals(Query.this.options.getLimitType())
- ? reverse(documentSnapshots)
- : documentSnapshots;
- snapshot = QuerySnapshot.withDocuments(Query.this, readTime, resultView);
- }
- result.set(new ExplainResults<>(metrics, snapshot));
- }
- },
- /* startTimeNanos= */ rpcContext.getClock().nanoTime(),
- /* transactionId= */ null,
- /* readTime= */ null,
- /* explainOptions= */ options);
-
- return result;
+ span.endAtFuture(result);
+ return result;
+ } catch (Exception error) {
+ span.end(error);
+ throw error;
+ }
}
/**
@@ -1946,51 +1959,65 @@ public ListenerRegistration addSnapshotListener(
ApiFuture get(
@Nullable ByteString transactionId, @Nullable Timestamp requestReadTime) {
- final SettableApiFuture result = SettableApiFuture.create();
-
- internalStream(
- new ApiStreamObserver() {
- final List documentSnapshots = new ArrayList<>();
- Timestamp responseReadTime;
-
- @Override
- public void onNext(RunQueryResponse runQueryResponse) {
- if (runQueryResponse.hasDocument()) {
- Document document = runQueryResponse.getDocument();
- QueryDocumentSnapshot documentSnapshot =
- QueryDocumentSnapshot.fromDocument(
- rpcContext, Timestamp.fromProto(runQueryResponse.getReadTime()), document);
- documentSnapshots.add(documentSnapshot);
- }
- if (responseReadTime == null) {
- responseReadTime = Timestamp.fromProto(runQueryResponse.getReadTime());
+ TraceUtil.Span span =
+ getFirestore()
+ .getOptions()
+ .getTraceUtil()
+ .startSpan(
+ transactionId == null
+ ? TraceUtil.SPAN_NAME_QUERY_GET
+ : TraceUtil.SPAN_NAME_TRANSACTION_GET_QUERY);
+ try (Scope ignored = span.makeCurrent()) {
+ final SettableApiFuture result = SettableApiFuture.create();
+ internalStream(
+ new ApiStreamObserver() {
+ final List documentSnapshots = new ArrayList<>();
+ Timestamp responseReadTime;
+
+ @Override
+ public void onNext(RunQueryResponse runQueryResponse) {
+ if (runQueryResponse.hasDocument()) {
+ Document document = runQueryResponse.getDocument();
+ QueryDocumentSnapshot documentSnapshot =
+ QueryDocumentSnapshot.fromDocument(
+ rpcContext, Timestamp.fromProto(runQueryResponse.getReadTime()), document);
+ documentSnapshots.add(documentSnapshot);
+ }
+ if (responseReadTime == null) {
+ responseReadTime = Timestamp.fromProto(runQueryResponse.getReadTime());
+ }
}
- }
- @Override
- public void onError(Throwable throwable) {
- result.setException(throwable);
- }
-
- @Override
- public void onCompleted() {
- // The results for limitToLast queries need to be flipped since we reversed the
- // ordering constraints before sending the query to the backend.
- List resultView =
- LimitType.Last.equals(Query.this.options.getLimitType())
- ? reverse(documentSnapshots)
- : documentSnapshots;
- QuerySnapshot querySnapshot =
- QuerySnapshot.withDocuments(Query.this, responseReadTime, resultView);
- result.set(querySnapshot);
- }
- },
- /* startTimeNanos= */ rpcContext.getClock().nanoTime(),
- transactionId,
- /* readTime= */ requestReadTime,
- /* explainOptions= */ null);
+ @Override
+ public void onError(Throwable throwable) {
+ result.setException(throwable);
+ }
- return result;
+ @Override
+ public void onCompleted() {
+ // The results for limitToLast queries need to be flipped since we reversed the
+ // ordering constraints before sending the query to the backend.
+ List resultView =
+ LimitType.Last.equals(Query.this.options.getLimitType())
+ ? reverse(documentSnapshots)
+ : documentSnapshots;
+ QuerySnapshot querySnapshot =
+ QuerySnapshot.withDocuments(Query.this, responseReadTime, resultView);
+ result.set(querySnapshot);
+ }
+ },
+ /* startTimeNanos= */ rpcContext.getClock().nanoTime(),
+ transactionId,
+ /* readTime= */ requestReadTime,
+ /* explainOptions= */ null,
+ /* isRetryRequestWithCursor= */ false);
+
+ span.endAtFuture(result);
+ return result;
+ } catch (Exception error) {
+ span.end(error);
+ throw error;
+ }
}
Comparator comparator() {
diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/ReadTimeTransaction.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/ReadTimeTransaction.java
index e6f421389..0c423469a 100644
--- a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/ReadTimeTransaction.java
+++ b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/ReadTimeTransaction.java
@@ -18,10 +18,10 @@
import com.google.api.core.ApiFuture;
import com.google.api.core.ApiFutures;
+import com.google.cloud.firestore.telemetry.TraceUtil;
import com.google.common.base.Preconditions;
import com.google.common.util.concurrent.MoreExecutors;
import com.google.protobuf.Timestamp;
-import io.opencensus.trace.Tracing;
import java.util.List;
import java.util.Map;
import javax.annotation.Nonnull;
@@ -55,37 +55,74 @@ public boolean hasTransactionId() {
@Nonnull
@Override
public ApiFuture get(@Nonnull DocumentReference documentRef) {
- Tracing.getTracer().getCurrentSpan().addAnnotation(TraceUtil.SPAN_NAME_GETDOCUMENT);
- return ApiFutures.transform(
- firestore.getAll(new DocumentReference[] {documentRef}, /* fieldMask= */ null, readTime),
- snapshots -> snapshots.isEmpty() ? null : snapshots.get(0),
- MoreExecutors.directExecutor());
+ TraceUtil.Span span =
+ getTraceUtil()
+ .startSpan(TraceUtil.SPAN_NAME_TRANSACTION_GET_DOCUMENT, transactionTraceContext);
+ try (TraceUtil.Scope ignored = span.makeCurrent()) {
+ ApiFuture result =
+ ApiFutures.transform(
+ firestore.getAll(
+ new DocumentReference[] {documentRef}, /* fieldMask= */ null, readTime),
+ snapshots -> snapshots.isEmpty() ? null : snapshots.get(0),
+ MoreExecutors.directExecutor());
+ span.endAtFuture(result);
+ return result;
+ } catch (Exception error) {
+ span.end(error);
+ throw error;
+ }
}
@Nonnull
@Override
public ApiFuture> getAll(
@Nonnull DocumentReference... documentReferences) {
- return firestore.getAll(documentReferences, /* fieldMask= */ null, readTime);
+ TraceUtil.Span span =
+ getTraceUtil()
+ .startSpan(TraceUtil.SPAN_NAME_TRANSACTION_GET_DOCUMENTS, transactionTraceContext);
+ try (TraceUtil.Scope ignored = span.makeCurrent()) {
+ ApiFuture> result =
+ firestore.getAll(documentReferences, /* fieldMask= */ null, readTime);
+ span.endAtFuture(result);
+ return result;
+ } catch (Exception error) {
+ span.end(error);
+ throw error;
+ }
}
@Nonnull
@Override
public ApiFuture> getAll(
@Nonnull DocumentReference[] documentReferences, @Nullable FieldMask fieldMask) {
- return firestore.getAll(documentReferences, /* fieldMask= */ null, readTime);
+ TraceUtil.Span span =
+ getTraceUtil()
+ .startSpan(TraceUtil.SPAN_NAME_TRANSACTION_GET_DOCUMENTS, transactionTraceContext);
+ try (TraceUtil.Scope ignored = span.makeCurrent()) {
+ ApiFuture> result =
+ firestore.getAll(documentReferences, /* fieldMask= */ null, readTime);
+ span.endAtFuture(result);
+ return result;
+ } catch (Exception error) {
+ span.end(error);
+ throw error;
+ }
}
@Nonnull
@Override
public ApiFuture get(@Nonnull Query query) {
- return query.get(null, com.google.cloud.Timestamp.fromProto(readTime));
+ try (TraceUtil.Scope ignored = transactionTraceContext.makeCurrent()) {
+ return query.get(null, com.google.cloud.Timestamp.fromProto(readTime));
+ }
}
@Nonnull
@Override
public ApiFuture get(@Nonnull AggregateQuery query) {
- return query.get(null, readTime);
+ try (TraceUtil.Scope ignored = transactionTraceContext.makeCurrent()) {
+ return query.get(null, readTime);
+ }
}
@Nonnull
diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/ServerSideTransaction.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/ServerSideTransaction.java
index 92b3e56b4..5d366c965 100644
--- a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/ServerSideTransaction.java
+++ b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/ServerSideTransaction.java
@@ -19,6 +19,7 @@
import com.google.api.core.ApiFuture;
import com.google.api.core.ApiFutures;
import com.google.cloud.firestore.TransactionOptions.TransactionOptionsType;
+import com.google.cloud.firestore.telemetry.TraceUtil;
import com.google.common.base.Preconditions;
import com.google.common.util.concurrent.MoreExecutors;
import com.google.firestore.v1.BeginTransactionRequest;
@@ -27,7 +28,6 @@
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 java.util.logging.Level;
import java.util.logging.Logger;
@@ -68,7 +68,6 @@ public static ApiFuture begin(
FirestoreImpl firestore,
TransactionOptions transactionOptions,
@Nullable ServerSideTransaction previousTransaction) {
- Tracing.getTracer().getCurrentSpan().addAnnotation(TraceUtil.SPAN_NAME_BEGINTRANSACTION);
BeginTransactionRequest.Builder beginTransaction = BeginTransactionRequest.newBuilder();
beginTransaction.setDatabase(firestore.getDatabaseName());
ByteString previousTransactionId =
@@ -101,35 +100,46 @@ public static ApiFuture begin(
/** Commits a transaction. */
ApiFuture> commit() {
- return super.commit(transactionId);
+ try (TraceUtil.Scope ignored = transactionTraceContext.makeCurrent()) {
+ return super.commit(transactionId);
+ }
}
/** Rolls a transaction back and releases all read locks. */
ApiFuture rollback() {
- Tracing.getTracer().getCurrentSpan().addAnnotation(TraceUtil.SPAN_NAME_ROLLBACK);
- RollbackRequest req =
- RollbackRequest.newBuilder()
- .setTransaction(transactionId)
- .setDatabase(firestore.getDatabaseName())
- .build();
-
- ApiFuture rollbackFuture =
- firestore.sendRequest(req, firestore.getClient().rollbackCallable());
-
- ApiFuture transform =
- ApiFutures.transform(rollbackFuture, resp -> null, MoreExecutors.directExecutor());
-
- return ApiFutures.catching(
- transform,
- Throwable.class,
- (error) -> {
- LOGGER.log(
- Level.WARNING,
- "Failed best effort to rollback of transaction " + transactionId,
- error);
- return null;
- },
- MoreExecutors.directExecutor());
+ TraceUtil.Span span =
+ getTraceUtil().startSpan(TraceUtil.SPAN_NAME_TRANSACTION_ROLLBACK, transactionTraceContext);
+ try (TraceUtil.Scope ignored = span.makeCurrent()) {
+ RollbackRequest req =
+ RollbackRequest.newBuilder()
+ .setTransaction(transactionId)
+ .setDatabase(firestore.getDatabaseName())
+ .build();
+
+ ApiFuture rollbackFuture =
+ firestore.sendRequest(req, firestore.getClient().rollbackCallable());
+
+ ApiFuture transform =
+ ApiFutures.transform(rollbackFuture, resp -> null, MoreExecutors.directExecutor());
+
+ ApiFuture result =
+ ApiFutures.catching(
+ transform,
+ Throwable.class,
+ (error) -> {
+ LOGGER.log(
+ Level.WARNING,
+ "Failed best effort to rollback of transaction " + transactionId,
+ error);
+ return null;
+ },
+ MoreExecutors.directExecutor());
+ span.endAtFuture(result);
+ return result;
+ } catch (Exception error) {
+ span.end(error);
+ throw error;
+ }
}
@Override
@@ -146,16 +156,26 @@ public boolean hasTransactionId() {
@Override
@Nonnull
public ApiFuture get(@Nonnull DocumentReference documentRef) {
- Tracing.getTracer().getCurrentSpan().addAnnotation(TraceUtil.SPAN_NAME_GETDOCUMENT);
- Preconditions.checkState(isEmpty(), READ_BEFORE_WRITE_ERROR_MSG);
- return ApiFutures.transform(
- firestore.getAll(
- new DocumentReference[] {documentRef},
- /* fieldMask= */ null,
- transactionId,
- /* readTime= */ null),
- snapshots -> snapshots.isEmpty() ? null : snapshots.get(0),
- MoreExecutors.directExecutor());
+ TraceUtil.Span span =
+ getTraceUtil()
+ .startSpan(TraceUtil.SPAN_NAME_TRANSACTION_GET_DOCUMENT, transactionTraceContext);
+ try (TraceUtil.Scope ignored = span.makeCurrent()) {
+ Preconditions.checkState(isEmpty(), READ_BEFORE_WRITE_ERROR_MSG);
+ ApiFuture result =
+ ApiFutures.transform(
+ firestore.getAll(
+ new DocumentReference[] {documentRef},
+ /* fieldMask= */ null,
+ transactionId,
+ /* readTime= */ null),
+ snapshots -> snapshots.isEmpty() ? null : snapshots.get(0),
+ MoreExecutors.directExecutor());
+ span.endAtFuture(result);
+ return result;
+ } catch (Exception error) {
+ span.end(error);
+ throw error;
+ }
}
/**
@@ -169,8 +189,19 @@ public ApiFuture get(@Nonnull DocumentReference documentRef) {
public ApiFuture> getAll(
@Nonnull DocumentReference... documentReferences) {
Preconditions.checkState(isEmpty(), READ_BEFORE_WRITE_ERROR_MSG);
- return firestore.getAll(
- documentReferences, /* fieldMask= */ null, transactionId, /* readTime= */ null);
+ TraceUtil.Span span =
+ getTraceUtil()
+ .startSpan(TraceUtil.SPAN_NAME_TRANSACTION_GET_DOCUMENTS, transactionTraceContext);
+ try (TraceUtil.Scope ignored = span.makeCurrent()) {
+ ApiFuture> result =
+ firestore.getAll(
+ documentReferences, /* fieldMask= */ null, transactionId, /* readTime= */ null);
+ span.endAtFuture(result);
+ return result;
+ } catch (Exception error) {
+ span.end(error);
+ throw error;
+ }
}
/**
@@ -186,7 +217,18 @@ public ApiFuture> getAll(
public ApiFuture> getAll(
@Nonnull DocumentReference[] documentReferences, @Nullable FieldMask fieldMask) {
Preconditions.checkState(isEmpty(), READ_BEFORE_WRITE_ERROR_MSG);
- return firestore.getAll(documentReferences, fieldMask, transactionId, /* readTime= */ null);
+ TraceUtil.Span span =
+ getTraceUtil()
+ .startSpan(TraceUtil.SPAN_NAME_TRANSACTION_GET_DOCUMENTS, transactionTraceContext);
+ try (TraceUtil.Scope ignored = span.makeCurrent()) {
+ ApiFuture> result =
+ firestore.getAll(documentReferences, fieldMask, transactionId, /* readTime= */ null);
+ span.endAtFuture(result);
+ return result;
+ } catch (Exception error) {
+ span.end(error);
+ throw error;
+ }
}
/**
@@ -199,7 +241,9 @@ public ApiFuture> getAll(
@Nonnull
public ApiFuture get(@Nonnull Query query) {
Preconditions.checkState(isEmpty(), READ_BEFORE_WRITE_ERROR_MSG);
- return query.get(transactionId, /* readTime= */ null);
+ try (TraceUtil.Scope ignored = transactionTraceContext.makeCurrent()) {
+ return query.get(transactionId, /* readTime= */ null);
+ }
}
/**
@@ -212,6 +256,8 @@ public ApiFuture get(@Nonnull Query query) {
@Nonnull
public ApiFuture get(@Nonnull AggregateQuery query) {
Preconditions.checkState(isEmpty(), READ_BEFORE_WRITE_ERROR_MSG);
- return query.get(transactionId, null);
+ try (TraceUtil.Scope ignored = transactionTraceContext.makeCurrent()) {
+ return query.get(transactionId, null);
+ }
}
}
diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/ServerSideTransactionRunner.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/ServerSideTransactionRunner.java
index 660ffd3b3..db8ebff63 100644
--- a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/ServerSideTransactionRunner.java
+++ b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/ServerSideTransactionRunner.java
@@ -16,6 +16,8 @@
package com.google.cloud.firestore;
+import static com.google.cloud.firestore.telemetry.TraceUtil.*;
+
import com.google.api.core.ApiFuture;
import com.google.api.core.ApiFutureCallback;
import com.google.api.core.ApiFutures;
@@ -24,16 +26,15 @@
import com.google.api.gax.retrying.ExponentialRetryAlgorithm;
import com.google.api.gax.retrying.TimedAttemptSettings;
import com.google.api.gax.rpc.ApiException;
-import com.google.common.collect.ImmutableMap;
+import com.google.cloud.firestore.telemetry.TraceUtil;
+import com.google.cloud.firestore.telemetry.TraceUtil.Scope;
+import com.google.cloud.firestore.telemetry.TraceUtil.Span;
import com.google.common.util.concurrent.MoreExecutors;
import io.grpc.Context;
-import io.opencensus.trace.AttributeValue;
-import io.opencensus.trace.Span;
-import io.opencensus.trace.Tracer;
-import io.opencensus.trace.Tracing;
import java.util.concurrent.Executor;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
+import javax.annotation.Nonnull;
/**
* Implements backoff and retry semantics for Firestore transactions.
@@ -47,15 +48,7 @@
* customize the backoff settings, you can specify custom settings via {@link FirestoreOptions}.
*/
final class ServerSideTransactionRunner {
-
- private static final Tracer tracer = Tracing.getTracer();
- private static final io.opencensus.trace.Status TOO_MANY_RETRIES_STATUS =
- io.opencensus.trace.Status.ABORTED.withDescription("too many retries");
- private static final io.opencensus.trace.Status USER_CALLBACK_FAILED =
- io.opencensus.trace.Status.ABORTED.withDescription("user callback failed");
-
private final Transaction.AsyncFunction userCallback;
- private final Span span;
private final FirestoreImpl firestore;
private final ScheduledExecutorService firestoreExecutor;
private final Executor userCallbackExecutor;
@@ -64,6 +57,8 @@ final class ServerSideTransactionRunner {
private TimedAttemptSettings nextBackoffAttempt;
private ServerSideTransaction transaction;
private int attemptsRemaining;
+ private Span runTransactionSpan;
+ private TraceUtil.Context runTransactionContext;
/**
* @param firestore The active Firestore instance
@@ -76,7 +71,6 @@ final class ServerSideTransactionRunner {
Transaction.AsyncFunction userCallback,
TransactionOptions transactionOptions) {
this.transactionOptions = transactionOptions;
- this.span = tracer.spanBuilder("CloudFirestore.Transaction").startSpan();
this.firestore = firestore;
this.firestoreExecutor = firestore.getClient().getExecutor();
this.userCallback = userCallback;
@@ -93,25 +87,57 @@ final class ServerSideTransactionRunner {
this.nextBackoffAttempt = backoffAlgorithm.createFirstAttempt();
}
- ApiFuture run() {
- --attemptsRemaining;
-
- span.addAnnotation(
- "Start runTransaction",
- ImmutableMap.of("attemptsRemaining", AttributeValue.longAttributeValue(attemptsRemaining)));
+ @Nonnull
+ private TraceUtil getTraceUtil() {
+ return firestore.getOptions().getTraceUtil();
+ }
- return ApiFutures.catchingAsync(
- ApiFutures.transformAsync(
- maybeRollback(), this::rollbackCallback, MoreExecutors.directExecutor()),
- Throwable.class,
- this::restartTransactionCallback,
- MoreExecutors.directExecutor());
+ ApiFuture run() {
+ runTransactionSpan = getTraceUtil().startSpan(TraceUtil.SPAN_NAME_TRANSACTION_RUN);
+ runTransactionSpan.setAttribute(
+ ATTRIBUTE_KEY_TRANSACTION_TYPE, transactionOptions.getType().name());
+ runTransactionSpan.setAttribute(
+ ATTRIBUTE_KEY_ATTEMPTS_ALLOWED, transactionOptions.getNumberOfAttempts());
+ runTransactionSpan.setAttribute(ATTRIBUTE_KEY_ATTEMPTS_REMAINING, attemptsRemaining);
+ try (Scope ignored = runTransactionSpan.makeCurrent()) {
+ runTransactionContext = getTraceUtil().currentContext();
+ --attemptsRemaining;
+ ApiFuture result =
+ ApiFutures.catchingAsync(
+ ApiFutures.transformAsync(
+ maybeRollback(), this::rollbackCallback, MoreExecutors.directExecutor()),
+ Throwable.class,
+ this::restartTransactionCallback,
+ MoreExecutors.directExecutor());
+ runTransactionSpan.endAtFuture(result);
+ return result;
+ } catch (Exception error) {
+ runTransactionSpan.end(error);
+ throw error;
+ }
}
ApiFuture begin() {
- ServerSideTransaction previousTransaction = this.transaction;
- this.transaction = null;
- return ServerSideTransaction.begin(firestore, transactionOptions, previousTransaction);
+ TraceUtil.Span span =
+ getTraceUtil().startSpan(TraceUtil.SPAN_NAME_TRANSACTION_BEGIN, runTransactionContext);
+ try (Scope ignored = span.makeCurrent()) {
+ ServerSideTransaction previousTransaction = this.transaction;
+ this.transaction = null;
+ ApiFuture result =
+ ServerSideTransaction.begin(firestore, transactionOptions, previousTransaction);
+ result =
+ ApiFutures.transform(
+ result,
+ serverSideTransaction -> {
+ serverSideTransaction.setTransactionTraceContext(runTransactionContext);
+ return serverSideTransaction;
+ });
+ span.endAtFuture(result);
+ return result;
+ } catch (Exception error) {
+ span.end(error);
+ throw error;
+ }
}
private ApiFuture maybeRollback() {
@@ -192,11 +218,7 @@ private ApiFuture userFunctionCallback(T userFunctionResult) {
return ApiFutures.transform(
transaction.commit(),
// The callback that is invoked after the Commit RPC returns. It returns the user result.
- input -> {
- span.setStatus(io.opencensus.trace.Status.OK);
- span.end();
- return userFunctionResult;
- },
+ input -> userFunctionResult,
MoreExecutors.directExecutor());
}
@@ -204,24 +226,23 @@ private ApiFuture userFunctionCallback(T userFunctionResult) {
private ApiFuture restartTransactionCallback(Throwable throwable) {
if (!(throwable instanceof ApiException)) {
// This is likely a failure in the user callback.
- span.setStatus(USER_CALLBACK_FAILED);
return rollbackAndReject(throwable);
}
ApiException apiException = (ApiException) throwable;
if (isRetryableTransactionError(apiException)) {
if (attemptsRemaining > 0) {
- span.addAnnotation("retrying");
+ getTraceUtil()
+ .currentSpan()
+ .addEvent("Initiating transaction retry. Attempts remaining: " + attemptsRemaining);
return run();
} else {
- span.setStatus(TOO_MANY_RETRIES_STATUS);
final FirestoreException firestoreException =
FirestoreException.forApiException(
apiException, "Transaction was cancelled because of too many retries.");
return rollbackAndReject(firestoreException);
}
} else {
- span.setStatus(TraceUtil.statusFromApiException(apiException));
final FirestoreException firestoreException =
FirestoreException.forApiException(
apiException, "Transaction failed with non-retryable error");
@@ -262,12 +283,16 @@ private ApiFuture rollbackAndReject(final Throwable throwable) {
transaction
.rollback()
.addListener(
- () -> failedTransaction.setException(throwable), MoreExecutors.directExecutor());
+ () -> {
+ runTransactionSpan.end(throwable);
+ failedTransaction.setException(throwable);
+ },
+ MoreExecutors.directExecutor());
} else {
+ runTransactionSpan.end(throwable);
failedTransaction.setException(throwable);
}
- span.end();
return failedTransaction;
}
}
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
deleted file mode 100644
index 66c278bc7..000000000
--- a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/TraceUtil.java
+++ /dev/null
@@ -1,95 +0,0 @@
-/*
- * Copyright 2018 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.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;
-
-/**
- * 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) {
- if (exception.getStatusCode().getTransportCode() instanceof io.grpc.Status) {
- io.grpc.Status grpcStatus = (io.grpc.Status) exception.getStatusCode().getTransportCode();
- return StatusConverter.fromGrpcStatus(grpcStatus);
- }
-
- return Status.UNKNOWN.withDescription(exception.getMessage());
- }
-}
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 22fa4e065..04d83a1a1 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
@@ -18,6 +18,8 @@
import com.google.api.core.ApiFuture;
import com.google.api.core.InternalExtensionOnly;
+import com.google.cloud.firestore.telemetry.TraceUtil;
+import com.google.cloud.firestore.telemetry.TraceUtil.Context;
import java.util.List;
import java.util.logging.Logger;
import javax.annotation.Nonnull;
@@ -35,9 +37,21 @@ public abstract class Transaction extends UpdateBuilder {
private static final Logger LOGGER = Logger.getLogger(Transaction.class.getName());
private static final String READ_BEFORE_WRITE_ERROR_MSG =
"Firestore transactions require all reads to be executed before all writes";
+ protected @Nonnull Context transactionTraceContext;
protected Transaction(FirestoreImpl firestore) {
super(firestore);
+ this.transactionTraceContext = firestore.getOptions().getTraceUtil().currentContext();
+ }
+
+ @Nonnull
+ TraceUtil getTraceUtil() {
+ return firestore.getOptions().getTraceUtil();
+ }
+
+ @Nonnull
+ Context setTransactionTraceContext(Context context) {
+ return transactionTraceContext = context;
}
/**
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 27f5a497d..e93fe8310 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
@@ -16,6 +16,8 @@
package com.google.cloud.firestore;
+import static com.google.cloud.firestore.telemetry.TraceUtil.ATTRIBUTE_KEY_DOC_COUNT;
+import static com.google.cloud.firestore.telemetry.TraceUtil.ATTRIBUTE_KEY_IS_TRANSACTIONAL;
import static com.google.common.base.Predicates.not;
import static java.util.stream.Collectors.toCollection;
@@ -23,16 +25,15 @@
import com.google.api.core.ApiFutures;
import com.google.api.core.InternalExtensionOnly;
import com.google.cloud.firestore.UserDataConverter.EncodingOptions;
+import com.google.cloud.firestore.telemetry.TraceUtil;
+import com.google.cloud.firestore.telemetry.TraceUtil.Scope;
import com.google.common.base.Preconditions;
-import com.google.common.collect.ImmutableMap;
import com.google.common.util.concurrent.MoreExecutors;
import com.google.firestore.v1.CommitRequest;
import com.google.firestore.v1.CommitResponse;
import com.google.firestore.v1.Write;
import com.google.protobuf.ByteString;
import com.google.protobuf.Timestamp;
-import io.opencensus.trace.AttributeValue;
-import io.opencensus.trace.Tracing;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
@@ -149,7 +150,6 @@ public T create(
private T performCreate(
@Nonnull DocumentReference documentReference, @Nonnull Map fields) {
- Tracing.getTracer().getCurrentSpan().addAnnotation(TraceUtil.SPAN_NAME_CREATEDOCUMENT);
DocumentSnapshot documentSnapshot =
DocumentSnapshot.fromObject(
firestore, documentReference, fields, UserDataConverter.NO_DELETES);
@@ -537,7 +537,6 @@ private T performUpdate(
@Nonnull final SortedMap fields,
@Nonnull Precondition precondition) {
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(
@@ -599,7 +598,6 @@ public T delete(@Nonnull DocumentReference documentReference) {
private T performDelete(
@Nonnull DocumentReference documentReference, @Nonnull Precondition precondition) {
- Tracing.getTracer().getCurrentSpan().addAnnotation(TraceUtil.SPAN_NAME_DELETEDOCUMENT);
Write.Builder write = Write.newBuilder().setDelete(documentReference.getName());
if (!precondition.isEmpty()) {
@@ -611,44 +609,54 @@ private T performDelete(
/** Commit the current batch. */
ApiFuture> commit(@Nullable ByteString transactionId) {
-
- // Sequence is thread safe.
- //
- // 1. Set committed = true
- // 2. Build commit request
- //
- // Step 1 sets uses volatile property to ensure committed is visible to all
- // threads immediately.
- //
- // Step 2 uses `forEach(..)` that is synchronized, therefore will be blocked
- // until any writes are complete.
- //
- // Writes will verify `committed==false` within synchronized block of code
- // before appending writes. Since committed is set to true before accessing
- // writes, we are ensured that no more writes will be appended after commit
- // accesses writes.
- committed = true;
- CommitRequest request = buildCommitRequest(transactionId);
-
- Tracing.getTracer()
- .getCurrentSpan()
- .addAnnotation(
- TraceUtil.SPAN_NAME_COMMIT,
- ImmutableMap.of(
- "numDocuments", AttributeValue.longAttributeValue(request.getWritesCount())));
-
- ApiFuture response =
- firestore.sendRequest(request, firestore.getClient().commitCallable());
-
- return ApiFutures.transform(
- response,
- commitResponse -> {
- Timestamp commitTime = commitResponse.getCommitTime();
- return commitResponse.getWriteResultsList().stream()
- .map(writeResult -> WriteResult.fromProto(writeResult, commitTime))
- .collect(Collectors.toList());
- },
- MoreExecutors.directExecutor());
+ TraceUtil.Span span =
+ firestore
+ .getOptions()
+ .getTraceUtil()
+ .startSpan(
+ transactionId == null
+ ? TraceUtil.SPAN_NAME_BATCH_COMMIT
+ : TraceUtil.SPAN_NAME_TRANSACTION_COMMIT);
+ span.setAttribute(ATTRIBUTE_KEY_DOC_COUNT, writes.size());
+ span.setAttribute(ATTRIBUTE_KEY_IS_TRANSACTIONAL, transactionId != null);
+ try (Scope ignored = span.makeCurrent()) {
+ // Sequence is thread safe.
+ //
+ // 1. Set committed = true
+ // 2. Build commit request
+ //
+ // Step 1 sets uses volatile property to ensure committed is visible to all
+ // threads immediately.
+ //
+ // Step 2 uses `forEach(..)` that is synchronized, therefore will be blocked
+ // until any writes are complete.
+ //
+ // Writes will verify `committed==false` within synchronized block of code
+ // before appending writes. Since committed is set to true before accessing
+ // writes, we are ensured that no more writes will be appended after commit
+ // accesses writes.
+ committed = true;
+ CommitRequest request = buildCommitRequest(transactionId);
+
+ ApiFuture response =
+ firestore.sendRequest(request, firestore.getClient().commitCallable());
+
+ ApiFuture> returnValue =
+ ApiFutures.transform(
+ response,
+ commitResponse -> {
+ Timestamp commitTime = commitResponse.getCommitTime();
+ return commitResponse.getWriteResultsList().stream()
+ .map(writeResult -> WriteResult.fromProto(writeResult, commitTime))
+ .collect(Collectors.toList());
+ },
+ MoreExecutors.directExecutor());
+ span.endAtFuture(returnValue);
+ return returnValue;
+ } catch (Exception error) {
+ span.end(error);
+ throw error;
+ }
}
private CommitRequest buildCommitRequest(ByteString transactionId) {
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 14d84cd0e..52d18cecc 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
@@ -38,7 +38,6 @@
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.Comparator;
import java.util.HashMap;
@@ -409,7 +408,6 @@ private void initStream() {
current = false;
nextAttempt = backoff.createNextAttempt(nextAttempt);
- Tracing.getTracer().getCurrentSpan().addAnnotation(TraceUtil.SPAN_NAME_LISTEN);
stream =
new SilenceableBidiStream<>(
Watch.this,
diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/telemetry/DisabledTraceUtil.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/telemetry/DisabledTraceUtil.java
new file mode 100644
index 000000000..c2b47d536
--- /dev/null
+++ b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/telemetry/DisabledTraceUtil.java
@@ -0,0 +1,112 @@
+/*
+ * Copyright 2024 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.firestore.telemetry;
+
+import com.google.api.core.ApiFunction;
+import com.google.api.core.ApiFuture;
+import io.grpc.ManagedChannelBuilder;
+import java.util.Map;
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
+/**
+ * A fully disabled (No-op) tracing utility class that does not perform any tracing actions and has
+ * near-zero overhead.
+ */
+public class DisabledTraceUtil implements TraceUtil {
+
+ static class Span implements TraceUtil.Span {
+ @Override
+ public void end() {}
+
+ @Override
+ public void end(Throwable error) {}
+
+ @Override
+ public void endAtFuture(ApiFuture futureValue) {}
+
+ @Override
+ public TraceUtil.Span addEvent(String name) {
+ return this;
+ }
+
+ @Override
+ public TraceUtil.Span addEvent(String name, Map attributes) {
+ return this;
+ }
+
+ @Override
+ public TraceUtil.Span setAttribute(String key, int value) {
+ return this;
+ }
+
+ @Override
+ public TraceUtil.Span setAttribute(String key, String value) {
+ return this;
+ }
+
+ @Override
+ public TraceUtil.Span setAttribute(String key, boolean value) {
+ return this;
+ }
+
+ @Override
+ public Scope makeCurrent() {
+ return new Scope();
+ }
+ }
+
+ static class Context implements TraceUtil.Context {
+ @Override
+ public Scope makeCurrent() {
+ return new Scope();
+ }
+ }
+
+ static class Scope implements TraceUtil.Scope {
+ @Override
+ public void close() {}
+ }
+
+ @Nullable
+ @Override
+ public ApiFunction getChannelConfigurator() {
+ return null;
+ }
+
+ @Override
+ public Span startSpan(String spanName) {
+ return new Span();
+ }
+
+ @Override
+ public TraceUtil.Span startSpan(String spanName, TraceUtil.Context parent) {
+ return new Span();
+ }
+
+ @Nonnull
+ @Override
+ public TraceUtil.Span currentSpan() {
+ return new Span();
+ }
+
+ @Nonnull
+ @Override
+ public TraceUtil.Context currentContext() {
+ return new DisabledTraceUtil.Context();
+ }
+}
diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/telemetry/EnabledTraceUtil.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/telemetry/EnabledTraceUtil.java
new file mode 100644
index 000000000..52bb7c5e0
--- /dev/null
+++ b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/telemetry/EnabledTraceUtil.java
@@ -0,0 +1,394 @@
+/*
+ * Copyright 2024 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.firestore.telemetry;
+
+import com.google.api.core.ApiFunction;
+import com.google.api.core.ApiFuture;
+import com.google.api.core.ApiFutureCallback;
+import com.google.api.core.ApiFutures;
+import com.google.cloud.firestore.FirestoreOptions;
+import com.google.common.base.Throwables;
+import io.grpc.ManagedChannelBuilder;
+import io.opentelemetry.api.GlobalOpenTelemetry;
+import io.opentelemetry.api.OpenTelemetry;
+import io.opentelemetry.api.common.Attributes;
+import io.opentelemetry.api.common.AttributesBuilder;
+import io.opentelemetry.api.trace.SpanBuilder;
+import io.opentelemetry.api.trace.SpanKind;
+import io.opentelemetry.api.trace.StatusCode;
+import io.opentelemetry.api.trace.Tracer;
+import io.opentelemetry.instrumentation.grpc.v1_6.GrpcTelemetry;
+import java.util.Map;
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
+/**
+ * A utility class that uses OpenTelemetry for trace collection. `FirestoreOpenTelemetryOptions` in
+ * `FirestoreOptions` can be used to configure its behavior.
+ */
+public class EnabledTraceUtil implements TraceUtil {
+ private final Tracer tracer;
+ private final OpenTelemetry openTelemetry;
+ private final FirestoreOptions firestoreOptions;
+
+ EnabledTraceUtil(FirestoreOptions firestoreOptions) {
+ OpenTelemetry openTelemetry = firestoreOptions.getOpenTelemetryOptions().getOpenTelemetry();
+
+ // If tracing is enabled, but an OpenTelemetry instance is not provided, fall back
+ // to using GlobalOpenTelemetry.
+ if (openTelemetry == null) {
+ openTelemetry = GlobalOpenTelemetry.get();
+ }
+
+ this.firestoreOptions = firestoreOptions;
+ this.openTelemetry = openTelemetry;
+ Package pkg = this.getClass().getPackage();
+ if (pkg != null) {
+ // TODO(tracing): OpenTelemetry is currently missing the API for adding scope attributes in
+ // Java. We should add `gcp.client.service` as scope attributes once
+ // https://github.com/open-telemetry/opentelemetry-java/issues/4695 is resolved.
+ this.tracer = openTelemetry.getTracer(LIBRARY_NAME, pkg.getImplementationVersion());
+ } else {
+ this.tracer = openTelemetry.getTracer(LIBRARY_NAME);
+ }
+ }
+
+ public OpenTelemetry getOpenTelemetry() {
+ return openTelemetry;
+ }
+
+ // The gRPC channel configurator that intercepts gRPC calls for tracing purposes.
+ public class OpenTelemetryGrpcChannelConfigurator
+ implements ApiFunction {
+ @Override
+ public ManagedChannelBuilder apply(ManagedChannelBuilder managedChannelBuilder) {
+ GrpcTelemetry grpcTelemetry = GrpcTelemetry.create(getOpenTelemetry());
+ return managedChannelBuilder.intercept(grpcTelemetry.newClientInterceptor());
+ }
+ }
+
+ @Override
+ @Nullable
+ public ApiFunction getChannelConfigurator() {
+ return new OpenTelemetryGrpcChannelConfigurator();
+ }
+
+ // Returns a JSON String representation of the given duration. The JSON representation for a
+ // Duration is a String that
+ // ends in `s` to indicate seconds and is preceded by the number of seconds, with nanoseconds
+ // expressed as fractional
+ // seconds.
+ String durationString(org.threeten.bp.Duration duration) {
+ int nanos = duration.getNano();
+ long seconds = duration.getSeconds();
+ int numLeadingZeros = 9;
+
+ double nanosFraction = nanos;
+ while (nanosFraction >= 1) {
+ nanosFraction = nanosFraction / 10;
+ numLeadingZeros--;
+ }
+
+ // If seconds=1 and nanos=0, we don't show 1.000000000s. We want to show 1.0s.
+ if (numLeadingZeros == 9) {
+ numLeadingZeros = 0;
+ }
+
+ // Get rid of trailing zeros.
+ while (nanos > 0 && nanos % 10 == 0) {
+ nanos = nanos / 10;
+ }
+
+ StringBuilder stringBuilder = new StringBuilder().append(seconds).append(".");
+ for (int i = 0; i < numLeadingZeros; ++i) {
+ stringBuilder.append("0");
+ }
+ stringBuilder.append(nanos).append("s");
+
+ return stringBuilder.toString();
+ }
+
+ static class Span implements TraceUtil.Span {
+ private final io.opentelemetry.api.trace.Span span;
+ private final String spanName;
+
+ public Span(io.opentelemetry.api.trace.Span span, String spanName) {
+ this.span = span;
+ this.spanName = spanName;
+ }
+
+ /** Ends this span. */
+ @Override
+ public void end() {
+ span.end();
+ }
+
+ /** Ends this span in an error. */
+ @Override
+ public void end(Throwable error) {
+ span.setStatus(StatusCode.ERROR, error.getMessage());
+ span.recordException(
+ error,
+ Attributes.builder()
+ .put("exception.message", error.getMessage())
+ .put("exception.type", error.getClass().getName())
+ .put("exception.stacktrace", Throwables.getStackTraceAsString(error))
+ .build());
+ span.end();
+ }
+
+ /**
+ * If an operation ends in the future, its relevant span should end _after_ the future has been
+ * completed. This method "appends" the span completion code at the completion of the given
+ * future. In order for telemetry info to be recorded, the future returned by this method should
+ * be completed.
+ */
+ @Override
+ public void endAtFuture(ApiFuture futureValue) {
+ io.opentelemetry.context.Context asyncContext = io.opentelemetry.context.Context.current();
+ ApiFutures.addCallback(
+ futureValue,
+ new ApiFutureCallback() {
+ @Override
+ public void onFailure(Throwable t) {
+ try (io.opentelemetry.context.Scope scope = asyncContext.makeCurrent()) {
+ span.addEvent(spanName + " failed.");
+ end(t);
+ }
+ }
+
+ @Override
+ public void onSuccess(T result) {
+ try (io.opentelemetry.context.Scope scope = asyncContext.makeCurrent()) {
+ span.addEvent(spanName + " succeeded.");
+ end();
+ }
+ }
+ });
+ }
+
+ /** Adds the given event to this span. */
+ @Override
+ public TraceUtil.Span addEvent(String name) {
+ span.addEvent(name);
+ return this;
+ }
+
+ @Override
+ public TraceUtil.Span addEvent(String name, Map attributes) {
+ AttributesBuilder attributesBuilder = Attributes.builder();
+ attributes.forEach(
+ (key, value) -> {
+ if (value instanceof Integer) {
+ attributesBuilder.put(key, (int) value);
+ } else if (value instanceof Long) {
+ attributesBuilder.put(key, (long) value);
+ } else if (value instanceof Double) {
+ attributesBuilder.put(key, (double) value);
+ } else if (value instanceof Float) {
+ attributesBuilder.put(key, (float) value);
+ } else if (value instanceof Boolean) {
+ attributesBuilder.put(key, (boolean) value);
+ } else if (value instanceof String) {
+ attributesBuilder.put(key, (String) value);
+ } else {
+ // OpenTelemetry APIs do not support any other type.
+ throw new IllegalArgumentException(
+ "Unknown attribute type:" + value.getClass().getSimpleName());
+ }
+ });
+ span.addEvent(name, attributesBuilder.build());
+ return this;
+ }
+
+ @Override
+ public TraceUtil.Span setAttribute(String key, int value) {
+ span.setAttribute(ATTRIBUTE_SERVICE_PREFIX + key, value);
+ return this;
+ }
+
+ @Override
+ public TraceUtil.Span setAttribute(String key, String value) {
+ span.setAttribute(ATTRIBUTE_SERVICE_PREFIX + key, value);
+ return this;
+ }
+
+ @Override
+ public TraceUtil.Span setAttribute(String key, boolean value) {
+ span.setAttribute(ATTRIBUTE_SERVICE_PREFIX + key, value);
+ return this;
+ }
+
+ @Override
+ public Scope makeCurrent() {
+ return new Scope(span.makeCurrent());
+ }
+ }
+
+ static class Scope implements TraceUtil.Scope {
+ private final io.opentelemetry.context.Scope scope;
+
+ Scope(io.opentelemetry.context.Scope scope) {
+ this.scope = scope;
+ }
+
+ @Override
+ public void close() {
+ scope.close();
+ }
+ }
+
+ static class Context implements TraceUtil.Context {
+ private final io.opentelemetry.context.Context context;
+
+ Context(io.opentelemetry.context.Context context) {
+ this.context = context;
+ }
+
+ @Override
+ public Scope makeCurrent() {
+ return new Scope(context.makeCurrent());
+ }
+ }
+
+ /** Applies the current Firestore instance settings as attributes to the current Span */
+ private SpanBuilder addSettingsAttributesToCurrentSpan(SpanBuilder spanBuilder) {
+ // TODO(tracing): OpenTelemetry is currently missing the API for adding scope attributes in
+ // Java. We are instead adding `gcp.client.service` as span attributes here.
+ // We should remove this span attribute once
+ // https://github.com/open-telemetry/opentelemetry-java/issues/4695 is resolved.
+ spanBuilder = spanBuilder.setAttribute("gcp.client.service", "Firestore");
+
+ spanBuilder =
+ spanBuilder.setAllAttributes(
+ Attributes.builder()
+ .put(
+ ATTRIBUTE_SERVICE_PREFIX + "settings.project_id",
+ firestoreOptions.getProjectId())
+ .put(
+ ATTRIBUTE_SERVICE_PREFIX + "settings.database_id",
+ firestoreOptions.getDatabaseId())
+ .put(ATTRIBUTE_SERVICE_PREFIX + "settings.host", firestoreOptions.getHost())
+ .build());
+
+ if (firestoreOptions.getTransportChannelProvider() != null) {
+ spanBuilder =
+ spanBuilder.setAllAttributes(
+ Attributes.builder()
+ .put(
+ ATTRIBUTE_SERVICE_PREFIX + "settings.channel.transport_name",
+ firestoreOptions.getTransportChannelProvider().getTransportName())
+ .put(
+ ATTRIBUTE_SERVICE_PREFIX + "settings.channel.needs_credentials",
+ String.valueOf(
+ firestoreOptions.getTransportChannelProvider().needsCredentials()))
+ .put(
+ ATTRIBUTE_SERVICE_PREFIX + "settings.channel.needs_endpoint",
+ String.valueOf(
+ firestoreOptions.getTransportChannelProvider().needsEndpoint()))
+ .put(
+ ATTRIBUTE_SERVICE_PREFIX + "settings.channel.needs_headers",
+ String.valueOf(firestoreOptions.getTransportChannelProvider().needsHeaders()))
+ .put(
+ ATTRIBUTE_SERVICE_PREFIX + "settings.channel.should_auto_close",
+ String.valueOf(
+ firestoreOptions.getTransportChannelProvider().shouldAutoClose()))
+ .build());
+ }
+
+ if (firestoreOptions.getCredentials() != null) {
+ spanBuilder =
+ spanBuilder.setAttribute(
+ ATTRIBUTE_SERVICE_PREFIX + "settings.credentials.authentication_type",
+ firestoreOptions.getCredentials().getAuthenticationType());
+ }
+
+ if (firestoreOptions.getRetrySettings() != null) {
+ spanBuilder =
+ spanBuilder.setAllAttributes(
+ Attributes.builder()
+ .put(
+ ATTRIBUTE_SERVICE_PREFIX + "settings.retry_settings.initial_retry_delay",
+ durationString(firestoreOptions.getRetrySettings().getInitialRetryDelay()))
+ .put(
+ ATTRIBUTE_SERVICE_PREFIX + "settings.retry_settings.max_retry_delay",
+ durationString(firestoreOptions.getRetrySettings().getMaxRetryDelay()))
+ .put(
+ ATTRIBUTE_SERVICE_PREFIX + "settings.retry_settings.retry_delay_multiplier",
+ String.valueOf(firestoreOptions.getRetrySettings().getRetryDelayMultiplier()))
+ .put(
+ ATTRIBUTE_SERVICE_PREFIX + "settings.retry_settings.max_attempts",
+ String.valueOf(firestoreOptions.getRetrySettings().getMaxAttempts()))
+ .put(
+ ATTRIBUTE_SERVICE_PREFIX + "settings.retry_settings.initial_rpc_timeout",
+ durationString(firestoreOptions.getRetrySettings().getInitialRpcTimeout()))
+ .put(
+ ATTRIBUTE_SERVICE_PREFIX + "settings.retry_settings.max_rpc_timeout",
+ durationString(firestoreOptions.getRetrySettings().getMaxRpcTimeout()))
+ .put(
+ ATTRIBUTE_SERVICE_PREFIX + "settings.retry_settings.rpc_timeout_multiplier",
+ String.valueOf(firestoreOptions.getRetrySettings().getRpcTimeoutMultiplier()))
+ .put(
+ ATTRIBUTE_SERVICE_PREFIX + "settings.retry_settings.total_timeout",
+ durationString(firestoreOptions.getRetrySettings().getTotalTimeout()))
+ .build());
+ }
+
+ // Add the memory utilization of the client at the time this trace was collected.
+ long totalMemory = Runtime.getRuntime().totalMemory();
+ long freeMemory = Runtime.getRuntime().freeMemory();
+ double memoryUtilization = ((double) (totalMemory - freeMemory)) / totalMemory;
+ spanBuilder.setAttribute(
+ ATTRIBUTE_SERVICE_PREFIX + "memory_utilization",
+ String.format("%.2f", memoryUtilization * 100) + "%");
+
+ return spanBuilder;
+ }
+
+ @Override
+ public Span startSpan(String spanName) {
+ SpanBuilder spanBuilder = tracer.spanBuilder(spanName).setSpanKind(SpanKind.PRODUCER);
+ io.opentelemetry.api.trace.Span span =
+ addSettingsAttributesToCurrentSpan(spanBuilder).startSpan();
+ return new Span(span, spanName);
+ }
+
+ @Override
+ public TraceUtil.Span startSpan(String spanName, TraceUtil.Context parent) {
+ assert (parent instanceof EnabledTraceUtil.Context);
+ SpanBuilder spanBuilder =
+ tracer
+ .spanBuilder(spanName)
+ .setSpanKind(SpanKind.PRODUCER)
+ .setParent(((EnabledTraceUtil.Context) parent).context);
+ io.opentelemetry.api.trace.Span span =
+ addSettingsAttributesToCurrentSpan(spanBuilder).startSpan();
+ return new Span(span, spanName);
+ }
+
+ @Nonnull
+ @Override
+ public TraceUtil.Span currentSpan() {
+ return new Span(io.opentelemetry.api.trace.Span.current(), "");
+ }
+
+ @Nonnull
+ @Override
+ public TraceUtil.Context currentContext() {
+ return new Context(io.opentelemetry.context.Context.current());
+ }
+}
diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/telemetry/TraceUtil.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/telemetry/TraceUtil.java
new file mode 100644
index 000000000..22dfb3ed0
--- /dev/null
+++ b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/telemetry/TraceUtil.java
@@ -0,0 +1,167 @@
+/*
+ * Copyright 2024 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.firestore.telemetry;
+
+import com.google.api.core.ApiFunction;
+import com.google.api.core.ApiFuture;
+import com.google.cloud.firestore.FirestoreOptions;
+import io.grpc.ManagedChannelBuilder;
+import java.util.Map;
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
+/**
+ * A utility interface for trace collection. Classes that implement this interface may make their
+ * own design choices for how they approach trace collection. For instance, they may be no-op, or
+ * they may use a particular tracing framework such as OpenTelemetry.
+ */
+public interface TraceUtil {
+ String ATTRIBUTE_SERVICE_PREFIX = "gcp.firestore.";
+ String SPAN_NAME_DOC_REF_CREATE = "DocumentReference.Create";
+ String SPAN_NAME_DOC_REF_SET = "DocumentReference.Set";
+ String SPAN_NAME_DOC_REF_UPDATE = "DocumentReference.Update";
+ String SPAN_NAME_DOC_REF_DELETE = "DocumentReference.Delete";
+ String SPAN_NAME_DOC_REF_GET = "DocumentReference.Get";
+ String SPAN_NAME_DOC_REF_LIST_COLLECTIONS = "DocumentReference.ListCollections";
+ String SPAN_NAME_COL_REF_ADD = "CollectionReference.Add";
+ String SPAN_NAME_COL_REF_LIST_DOCUMENTS = "CollectionReference.ListDocuments";
+ String SPAN_NAME_QUERY_GET = "Query.Get";
+ String SPAN_NAME_AGGREGATION_QUERY_GET = "AggregationQuery.Get";
+ String SPAN_NAME_RUN_QUERY = "RunQuery";
+ String SPAN_NAME_RUN_AGGREGATION_QUERY = "RunAggregationQuery";
+ String SPAN_NAME_BATCH_GET_DOCUMENTS = "BatchGetDocuments";
+ String SPAN_NAME_TRANSACTION_RUN = "Transaction.Run";
+ String SPAN_NAME_TRANSACTION_BEGIN = "Transaction.Begin";
+ String SPAN_NAME_TRANSACTION_GET_QUERY = "Transaction.Get.Query";
+ String SPAN_NAME_TRANSACTION_GET_AGGREGATION_QUERY = "Transaction.Get.AggregationQuery";
+ String SPAN_NAME_TRANSACTION_GET_DOCUMENT = "Transaction.Get.Document";
+ String SPAN_NAME_TRANSACTION_GET_DOCUMENTS = "Transaction.Get.Documents";
+ String SPAN_NAME_TRANSACTION_ROLLBACK = "Transaction.Rollback";
+ String SPAN_NAME_BATCH_COMMIT = "Batch.Commit";
+ String SPAN_NAME_TRANSACTION_COMMIT = "Transaction.Commit";
+ String SPAN_NAME_PARTITION_QUERY = "PartitionQuery";
+ String SPAN_NAME_BULK_WRITER_COMMIT = "BulkWriter.Commit";
+ String ATTRIBUTE_KEY_ATTEMPT = "attempt";
+ String ATTRIBUTE_KEY_DOC_COUNT = "doc_count";
+ String ATTRIBUTE_KEY_IS_TRANSACTIONAL = "transactional";
+ String ATTRIBUTE_KEY_NUM_RESPONSES = "response_count";
+ String ATTRIBUTE_KEY_IS_RETRY_WITH_CURSOR = "retry_query_with_cursor";
+ String ATTRIBUTE_KEY_TRANSACTION_TYPE = "transaction_type";
+ String ATTRIBUTE_KEY_ATTEMPTS_ALLOWED = "attempts_allowed";
+ String ATTRIBUTE_KEY_ATTEMPTS_REMAINING = "attempts_remaining";
+
+ String ENABLE_TRACING_ENV_VAR = "FIRESTORE_ENABLE_TRACING";
+ String LIBRARY_NAME = "com.google.cloud.firestore";
+
+ /**
+ * Creates and returns an instance of the TraceUtil class.
+ *
+ * @param firestoreOptions The FirestoreOptions object that is requesting an instance of
+ * TraceUtil.
+ * @return An instance of the TraceUtil class.
+ */
+ static TraceUtil getInstance(@Nonnull FirestoreOptions firestoreOptions) {
+ boolean createEnabledInstance = firestoreOptions.getOpenTelemetryOptions().isTracingEnabled();
+
+ // The environment variable can override options to enable/disable telemetry collection.
+ String enableTracingEnvVar = System.getenv(ENABLE_TRACING_ENV_VAR);
+ if (enableTracingEnvVar != null) {
+ if (enableTracingEnvVar.equalsIgnoreCase("true")
+ || enableTracingEnvVar.equalsIgnoreCase("on")) {
+ createEnabledInstance = true;
+ }
+ if (enableTracingEnvVar.equalsIgnoreCase("false")
+ || enableTracingEnvVar.equalsIgnoreCase("off")) {
+ createEnabledInstance = false;
+ }
+ }
+
+ if (createEnabledInstance) {
+ return new EnabledTraceUtil(firestoreOptions);
+ } else {
+ return new DisabledTraceUtil();
+ }
+ }
+
+ /** Returns a channel configurator for gRPC, or {@code null} if tracing is disabled. */
+ @Nullable
+ ApiFunction getChannelConfigurator();
+
+ /** Represents a trace span. */
+ interface Span {
+ /** Adds the given event to this span. */
+ Span addEvent(String name);
+
+ /** Adds the given event with the given attributes to this span. */
+ Span addEvent(String name, Map attributes);
+
+ /** Adds the given attribute to this span. */
+ Span setAttribute(String key, int value);
+
+ /** Adds the given attribute to this span. */
+ Span setAttribute(String key, String value);
+
+ /** Adds the given attribute to this span. */
+ Span setAttribute(String key, boolean value);
+
+ /** Marks this span as the current span. */
+ Scope makeCurrent();
+
+ /** Ends this span. */
+ void end();
+
+ /** Ends this span in an error. */
+ void end(Throwable error);
+
+ /**
+ * If an operation ends in the future, its relevant span should end _after_ the future has been
+ * completed. This method "appends" the span completion code at the completion of the given
+ * future. In order for telemetry info to be recorded, the future returned by this method should
+ * be completed.
+ */
+ void endAtFuture(ApiFuture futureValue);
+ }
+
+ /** Represents a trace context. */
+ interface Context {
+ /** Makes this context the current context. */
+ Scope makeCurrent();
+ }
+
+ /** Represents a trace scope. */
+ interface Scope extends AutoCloseable {
+ /** Closes the current scope. */
+ void close();
+ }
+
+ /** Starts a new span with the given name, sets it as the current span, and returns it. */
+ Span startSpan(String spanName);
+
+ /**
+ * Starts a new span with the given name and the given context as its parent, sets it as the
+ * current span, and returns it.
+ */
+ Span startSpan(String spanName, Context parent);
+
+ /** Returns the current span. */
+ @Nonnull
+ Span currentSpan();
+
+ /** Returns the current Context. */
+ @Nonnull
+ Context currentContext();
+}
diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/OpenTelemetryOptionsTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/OpenTelemetryOptionsTest.java
new file mode 100644
index 000000000..267020761
--- /dev/null
+++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/OpenTelemetryOptionsTest.java
@@ -0,0 +1,160 @@
+/*
+ * Copyright 2024 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.firestore;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import com.google.cloud.firestore.telemetry.DisabledTraceUtil;
+import com.google.cloud.firestore.telemetry.EnabledTraceUtil;
+import io.opentelemetry.api.GlobalOpenTelemetry;
+import io.opentelemetry.api.OpenTelemetry;
+import io.opentelemetry.sdk.OpenTelemetrySdk;
+import javax.annotation.Nullable;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+public class OpenTelemetryOptionsTest {
+ @Nullable private Firestore firestore;
+
+ @Before
+ public void setUp() {
+ GlobalOpenTelemetry.resetForTest();
+ }
+
+ @After
+ public void tearDown() {
+ if (firestore != null) {
+ firestore.shutdown();
+ firestore = null;
+ }
+ }
+
+ FirestoreOptions.Builder getBaseOptions() {
+ return FirestoreOptions.newBuilder().setProjectId("test-project").setDatabaseId("(default)");
+ }
+
+ @Test
+ public void defaultOptionsDisablesTelemetryCollection() {
+ FirestoreOptions firestoreOptions = getBaseOptions().build();
+ firestore = firestoreOptions.getService();
+ assertThat(firestore.getOptions().getOpenTelemetryOptions().isTracingEnabled()).isFalse();
+ assertThat(firestore.getOptions().getOpenTelemetryOptions().getOpenTelemetry()).isNull();
+ assertThat(firestore.getOptions().getTraceUtil()).isNotNull();
+ assertThat(firestore.getOptions().getTraceUtil() instanceof DisabledTraceUtil).isTrue();
+ }
+
+ @Test
+ public void canEnableTelemetryCollectionWithoutOpenTelemetryInstance() {
+ FirestoreOptions firestoreOptions =
+ getBaseOptions()
+ .setOpenTelemetryOptions(
+ FirestoreOpenTelemetryOptions.newBuilder().setTracingEnabled(true).build())
+ .build();
+ firestore = firestoreOptions.getService();
+ assertThat(firestore.getOptions().getOpenTelemetryOptions().isTracingEnabled()).isTrue();
+ assertThat(firestore.getOptions().getOpenTelemetryOptions().getOpenTelemetry()).isNull();
+ assertThat(firestore.getOptions().getTraceUtil()).isNotNull();
+ assertThat(firestore.getOptions().getTraceUtil() instanceof EnabledTraceUtil).isTrue();
+ }
+
+ @Test
+ public void canEnableTelemetryCollectionWithOpenTelemetryInstance() {
+ OpenTelemetry openTelemetry = GlobalOpenTelemetry.get();
+ FirestoreOptions firestoreOptions =
+ getBaseOptions()
+ .setOpenTelemetryOptions(
+ FirestoreOpenTelemetryOptions.newBuilder()
+ .setTracingEnabled(true)
+ .setOpenTelemetry(openTelemetry)
+ .build())
+ .build();
+ firestore = firestoreOptions.getService();
+ assertThat(firestore.getOptions().getOpenTelemetryOptions().isTracingEnabled()).isTrue();
+ assertThat(firestore.getOptions().getOpenTelemetryOptions().getOpenTelemetry())
+ .isEqualTo(openTelemetry);
+ assertThat(firestore.getOptions().getTraceUtil()).isNotNull();
+ assertThat(firestore.getOptions().getTraceUtil() instanceof EnabledTraceUtil).isTrue();
+ }
+
+ @Test
+ public void canDisableTelemetryCollectionWhileOpenTelemetryInstanceIsNotNull() {
+ OpenTelemetry openTelemetry = GlobalOpenTelemetry.get();
+ FirestoreOptions firestoreOptions =
+ getBaseOptions()
+ .setOpenTelemetryOptions(
+ FirestoreOpenTelemetryOptions.newBuilder()
+ .setTracingEnabled(false)
+ .setOpenTelemetry(openTelemetry)
+ .build())
+ .build();
+ firestore = firestoreOptions.getService();
+ assertThat(firestore.getOptions().getOpenTelemetryOptions().isTracingEnabled()).isFalse();
+ assertThat(firestore.getOptions().getOpenTelemetryOptions().getOpenTelemetry())
+ .isEqualTo(openTelemetry);
+ assertThat(firestore.getOptions().getTraceUtil()).isNotNull();
+ assertThat(firestore.getOptions().getTraceUtil() instanceof DisabledTraceUtil).isTrue();
+ }
+
+ @Test
+ public void existenceOfGlobalOpenTelemetryDoesNotEnableTracing() {
+ // Register a global OpenTelemetry SDK.
+ OpenTelemetrySdk.builder().buildAndRegisterGlobal();
+
+ // Make sure Firestore does not use GlobalOpenTelemetry by default.
+ FirestoreOptions firestoreOptions = getBaseOptions().build();
+ firestore = firestoreOptions.getService();
+ assertThat(firestore.getOptions().getOpenTelemetryOptions().isTracingEnabled()).isFalse();
+ assertThat(firestore.getOptions().getOpenTelemetryOptions().getOpenTelemetry()).isNull();
+ assertThat(firestore.getOptions().getTraceUtil()).isNotNull();
+ assertThat(firestore.getOptions().getTraceUtil() instanceof DisabledTraceUtil).isTrue();
+ }
+
+ @Test
+ public void canPassOpenTelemetrySdkInstanceToFirestore() {
+ OpenTelemetrySdk myOpenTelemetrySdk = OpenTelemetrySdk.builder().build();
+ FirestoreOptions firestoreOptions =
+ getBaseOptions()
+ .setOpenTelemetryOptions(
+ FirestoreOpenTelemetryOptions.newBuilder()
+ .setTracingEnabled(true)
+ .setOpenTelemetry(myOpenTelemetrySdk)
+ .build())
+ .build();
+ firestore = firestoreOptions.getService();
+ EnabledTraceUtil enabledTraceUtil = (EnabledTraceUtil) firestore.getOptions().getTraceUtil();
+ assertThat(enabledTraceUtil).isNotNull();
+ assertThat(enabledTraceUtil.getOpenTelemetry()).isEqualTo(myOpenTelemetrySdk);
+ }
+
+ @Test
+ public void usesGlobalOpenTelemetryIfOpenTelemetryNotProvidedInOptions() {
+ // Register a global OpenTelemetry SDK.
+ OpenTelemetrySdk.builder().buildAndRegisterGlobal();
+
+ // Do _not_ pass it to FirestoreOptions.
+ FirestoreOptions firestoreOptions =
+ getBaseOptions()
+ .setOpenTelemetryOptions(
+ FirestoreOpenTelemetryOptions.newBuilder().setTracingEnabled(true).build())
+ .build();
+ firestore = firestoreOptions.getService();
+ EnabledTraceUtil enabledTraceUtil = (EnabledTraceUtil) firestore.getOptions().getTraceUtil();
+ assertThat(enabledTraceUtil).isNotNull();
+ assertThat(enabledTraceUtil.getOpenTelemetry()).isEqualTo(GlobalOpenTelemetry.get());
+ }
+}
diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java
new file mode 100644
index 000000000..44a8376cd
--- /dev/null
+++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java
@@ -0,0 +1,1224 @@
+/*
+ * Copyright 2024 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.firestore.it;
+
+import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_BATCH_COMMIT;
+import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_BULK_WRITER_COMMIT;
+import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_COL_REF_LIST_DOCUMENTS;
+import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_DOC_REF_CREATE;
+import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_DOC_REF_DELETE;
+import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_DOC_REF_GET;
+import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_DOC_REF_LIST_COLLECTIONS;
+import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_DOC_REF_SET;
+import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_DOC_REF_UPDATE;
+import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_PARTITION_QUERY;
+import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_QUERY_GET;
+import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_BEGIN;
+import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_COMMIT;
+import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_GET_AGGREGATION_QUERY;
+import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_GET_DOCUMENTS;
+import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_GET_QUERY;
+import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_ROLLBACK;
+import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN;
+import static io.opentelemetry.semconv.resource.attributes.ResourceAttributes.SERVICE_NAME;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+
+import com.google.api.gax.rpc.NotFoundException;
+import com.google.cloud.firestore.BulkWriter;
+import com.google.cloud.firestore.BulkWriterOptions;
+import com.google.cloud.firestore.CollectionGroup;
+import com.google.cloud.firestore.DocumentReference;
+import com.google.cloud.firestore.FieldMask;
+import com.google.cloud.firestore.FieldPath;
+import com.google.cloud.firestore.Firestore;
+import com.google.cloud.firestore.FirestoreOpenTelemetryOptions;
+import com.google.cloud.firestore.FirestoreOptions;
+import com.google.cloud.firestore.Precondition;
+import com.google.cloud.firestore.Query;
+import com.google.cloud.firestore.SetOptions;
+import com.google.cloud.firestore.WriteBatch;
+import com.google.cloud.firestore.it.ITTracingTest.Pojo;
+import com.google.cloud.opentelemetry.trace.TraceConfiguration;
+import com.google.cloud.opentelemetry.trace.TraceExporter;
+import com.google.cloud.trace.v1.TraceServiceClient;
+import com.google.common.base.Preconditions;
+import com.google.devtools.cloudtrace.v1.Trace;
+import com.google.devtools.cloudtrace.v1.TraceSpan;
+import io.opentelemetry.api.GlobalOpenTelemetry;
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanContext;
+import io.opentelemetry.api.trace.TraceFlags;
+import io.opentelemetry.api.trace.TraceState;
+import io.opentelemetry.api.trace.Tracer;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+import io.opentelemetry.sdk.OpenTelemetrySdk;
+import io.opentelemetry.sdk.common.CompletableResultCode;
+import io.opentelemetry.sdk.resources.Resource;
+import io.opentelemetry.sdk.trace.SdkTracerProvider;
+import io.opentelemetry.sdk.trace.export.BatchSpanProcessor;
+import io.opentelemetry.sdk.trace.samplers.Sampler;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Random;
+import java.util.TreeMap;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+import org.junit.After;
+import org.junit.AfterClass;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Test;
+
+// This End-to-End test verifies Client-side Tracing Functionality instrumented using the
+// OpenTelemetry API.
+// The test depends on the following external APIs/Services:
+// 1. Java OpenTelemetry SDK
+// 2. Cloud Trace Exporter
+// 3. TraceServiceClient from Cloud Trace API v1.
+//
+// Permissions required to run this test (https://cloud.google.com/trace/docs/iam#trace-roles):
+// 1. gcloud auth application-default login must be run with the test user.
+// 2. To write traces, test user must have one of roles/cloudtrace.[admin|agent|user] roles.
+// 3. To read traces, test user must have one of roles/cloudtrace.[admin|user] roles.
+//
+// Each test-case has the following workflow:
+// 1. OpenTelemetry SDK is initialized with Cloud Trace Exporter and 100% Trace Sampling
+// 2. On initialization, Firestore client is provided the OpenTelemetry SDK object from (1)
+// 3. A custom TraceID is generated and injected using a custom SpanContext
+// 4. Firestore operations are run inside a root TraceSpan created using the custom SpanContext from
+// (3).
+// 5. Traces are read-back using TraceServiceClient and verified against expected Call Stacks.
+// TODO In the future it would be great to have a single test-driver for this test and
+// ITTracingTest.
+public abstract class ITE2ETracingTest extends ITBaseTest {
+ protected abstract boolean isUsingGlobalOpenTelemetrySDK();
+
+ // Helper class to track call-stacks in a trace
+ protected static class TraceContainer {
+
+ // Maps Span ID to TraceSpan
+ private final Map idSpanMap;
+
+ // Maps Parent Span ID to a list of Child SpanIDs, useful for top-down traversal
+ private final Map> parentChildIdMap;
+
+ // Tracks the Root Span ID
+ private long rootId;
+
+ public TraceContainer(String rootSpanName, Trace trace) {
+ idSpanMap = new TreeMap<>();
+ parentChildIdMap = new TreeMap<>();
+ for (TraceSpan span : trace.getSpansList()) {
+ long spanId = span.getSpanId();
+ idSpanMap.put(spanId, span);
+ if (rootSpanName.equals(span.getName())) {
+ rootId = span.getSpanId();
+ }
+
+ // Add self as a child of the parent span
+ if (!parentChildIdMap.containsKey(span.getParentSpanId())) {
+ parentChildIdMap.put(span.getParentSpanId(), new ArrayList<>());
+ }
+ parentChildIdMap.get(span.getParentSpanId()).add(spanId);
+ }
+ }
+
+ String spanName(long spanId) {
+ return idSpanMap.get(spanId).getName();
+ }
+
+ List childSpans(long spanId) {
+ return parentChildIdMap.get(spanId);
+ }
+
+ // This method only works for matching call stacks with traces which have children of distinct
+ // type at all levels. This is good enough as the intention is to validate if the e2e path is
+ // WAI - the intention is not to validate Cloud Trace's correctness w.r.t. durability of all
+ // kinds of traces.
+ boolean containsCallStack(String... callStack) throws RuntimeException {
+ List expectedCallStack = Arrays.asList(callStack);
+ if (expectedCallStack.isEmpty()) {
+ throw new RuntimeException("Input callStack is empty");
+ }
+ return dfsContainsCallStack(rootId, expectedCallStack);
+ }
+
+ // Depth-first check for call stack in the trace
+ private boolean dfsContainsCallStack(long spanId, List expectedCallStack) {
+ logger.info(
+ "span="
+ + spanName(spanId)
+ + ", expectedCallStack[0]="
+ + (expectedCallStack.isEmpty() ? "null" : expectedCallStack.get(0)));
+ if (expectedCallStack.isEmpty()) {
+ return false;
+ }
+ if (spanName(spanId).equals(expectedCallStack.get(0))) {
+ // Recursion termination
+ if (childSpans(spanId) == null) {
+ logger.info("No more children for " + spanName(spanId));
+ return true;
+ } else {
+ // Examine the child spans
+ for (Long childSpan : childSpans(spanId)) {
+ int callStackListSize = expectedCallStack.size();
+ logger.info(
+ "childSpan="
+ + spanName(childSpan)
+ + ", expectedCallStackSize="
+ + callStackListSize);
+ if (dfsContainsCallStack(
+ childSpan,
+ expectedCallStack.subList(
+ /*fromIndexInclusive=*/ 1, /*toIndexExclusive*/ callStackListSize))) {
+ return true;
+ }
+ }
+ }
+ } else {
+ logger.info(spanName(spanId) + " didn't match " + expectedCallStack.get(0));
+ }
+ return false;
+ }
+ }
+
+ private static final Logger logger = Logger.getLogger(ITE2ETracingTest.class.getName());
+
+ private static final String SERVICE = "google.firestore.v1.Firestore/";
+
+ private static final String BATCH_GET_DOCUMENTS_RPC_NAME = "BatchGetDocuments";
+
+ private static final String BATCH_WRITE_RPC_NAME = "BatchWrite";
+
+ private static final String BEGIN_TRANSACTION_RPC_NAME = "BeginTransaction";
+
+ private static final String COMMIT_RPC_NAME = "Commit";
+
+ private static final String LIST_COLLECTIONS_RPC_NAME = "ListCollectionIds";
+
+ private static final String LIST_DOCUMENTS_RPC_NAME = "ListDocuments";
+
+ private static final String ROLLBACK_RPC_NAME = "Rollback";
+
+ private static final String RUN_AGGREGATION_QUERY_RPC_NAME = "RunAggregationQuery";
+
+ private static final String RUN_QUERY_RPC_NAME = "RunQuery";
+
+ private static final int NUM_TRACE_ID_BYTES = 32;
+
+ private static final int NUM_SPAN_ID_BYTES = 16;
+
+ private static final int GET_TRACE_RETRY_COUNT = 60;
+
+ private static final int GET_TRACE_RETRY_BACKOFF_MILLIS = 1000;
+
+ private static final int TRACE_FORCE_FLUSH_MILLIS = 3000;
+
+ private static final int TRACE_PROVIDER_SHUTDOWN_MILLIS = 1000;
+
+ // Random int generator for trace ID and span ID
+ private static Random random;
+
+ private static TraceExporter traceExporter;
+
+ // Required for reading back traces from Cloud Trace for validation
+ private static TraceServiceClient traceClient_v1;
+
+ // Custom SpanContext for each test, required for TraceID injection
+ private static SpanContext customSpanContext;
+
+ // Trace read back from Cloud Trace using traceClient_v1 for verification
+ private static Trace retrievedTrace;
+
+ private static String rootSpanName;
+ private static Tracer tracer;
+
+ // Required to set custom-root span
+ private static OpenTelemetrySdk openTelemetrySdk;
+
+ private static String projectId;
+
+ private static Firestore firestore;
+
+ @BeforeClass
+ public static void setup() throws IOException {
+ projectId = FirestoreOptions.getDefaultProjectId();
+ logger.info("projectId:" + projectId);
+
+ // TODO(jimit) Make it re-usable w/ InMemorySpanExporter
+ traceExporter =
+ TraceExporter.createWithConfiguration(
+ TraceConfiguration.builder().setProjectId(projectId).build());
+
+ traceClient_v1 = TraceServiceClient.create();
+
+ random = new Random();
+ }
+
+ @Before
+ public void before() throws Exception {
+ // Set up OTel SDK
+ Resource resource =
+ Resource.getDefault().merge(Resource.builder().put(SERVICE_NAME, "Sparky").build());
+
+ if (isUsingGlobalOpenTelemetrySDK()) {
+ GlobalOpenTelemetry.resetForTest();
+ openTelemetrySdk =
+ OpenTelemetrySdk.builder()
+ .setTracerProvider(
+ SdkTracerProvider.builder()
+ .setResource(resource)
+ .addSpanProcessor(BatchSpanProcessor.builder(traceExporter).build())
+ .setSampler(Sampler.alwaysOn())
+ .build())
+ .buildAndRegisterGlobal();
+ } else {
+ openTelemetrySdk =
+ OpenTelemetrySdk.builder()
+ .setTracerProvider(
+ SdkTracerProvider.builder()
+ .setResource(resource)
+ .addSpanProcessor(BatchSpanProcessor.builder(traceExporter).build())
+ .setSampler(Sampler.alwaysOn())
+ .build())
+ .build();
+ }
+
+ // Initialize the Firestore DB w/ the OTel SDK. Ideally we'd do this is the @BeforeAll method
+ // but because gRPC traces need to be deterministically force-flushed, firestore.shutdown()
+ // must be called in @After for each test.
+ FirestoreOptions.Builder optionsBuilder;
+ if (isUsingGlobalOpenTelemetrySDK()) {
+ optionsBuilder =
+ FirestoreOptions.newBuilder()
+ .setOpenTelemetryOptions(
+ FirestoreOpenTelemetryOptions.newBuilder().setTracingEnabled(true).build());
+ } else {
+ optionsBuilder =
+ FirestoreOptions.newBuilder()
+ .setOpenTelemetryOptions(
+ FirestoreOpenTelemetryOptions.newBuilder()
+ .setOpenTelemetry(openTelemetrySdk)
+ .setTracingEnabled(true)
+ .build());
+ }
+
+ String namedDb = System.getProperty("FIRESTORE_NAMED_DATABASE");
+ if (namedDb != null) {
+ logger.log(Level.INFO, "Integration test using named database " + namedDb);
+ optionsBuilder = optionsBuilder.setDatabaseId(namedDb);
+ } else {
+ logger.log(Level.INFO, "Integration test using default database.");
+ }
+ firestore = optionsBuilder.build().getService();
+ Preconditions.checkNotNull(
+ firestore,
+ "Error instantiating Firestore. Check that the service account credentials "
+ + "were properly set.");
+
+ // Set up the tracer for custom TraceID injection
+ rootSpanName =
+ String.format("%s%d", this.getClass().getSimpleName(), System.currentTimeMillis());
+ if (isUsingGlobalOpenTelemetrySDK()) {
+ tracer = GlobalOpenTelemetry.getTracer(rootSpanName);
+ } else {
+ tracer =
+ firestore
+ .getOptions()
+ .getOpenTelemetryOptions()
+ .getOpenTelemetry()
+ .getTracer(rootSpanName);
+ }
+
+ // Get up a new SpanContext (ergo TraceId) for each test
+ customSpanContext = getNewSpanContext();
+ assertNotNull(customSpanContext);
+ assertNull(retrievedTrace);
+ }
+
+ @After
+ public void after() throws Exception {
+ firestore.shutdown();
+ rootSpanName = null;
+ tracer = null;
+ retrievedTrace = null;
+ customSpanContext = null;
+ }
+
+ @AfterClass
+ public static void teardown() throws Exception {
+ traceClient_v1.close();
+ CompletableResultCode completableResultCode =
+ openTelemetrySdk.getSdkTracerProvider().shutdown();
+ completableResultCode.join(TRACE_PROVIDER_SHUTDOWN_MILLIS, TimeUnit.MILLISECONDS);
+ }
+
+ // Generates a random hex string of length `numBytes`
+ private String generateRandomHexString(int numBytes) {
+ StringBuffer newTraceId = new StringBuffer();
+ while (newTraceId.length() < numBytes) {
+ newTraceId.append(Integer.toHexString(random.nextInt()));
+ }
+ return newTraceId.substring(0, numBytes);
+ }
+
+ protected String generateNewTraceId() {
+ return generateRandomHexString(NUM_TRACE_ID_BYTES);
+ }
+
+ // Generates a random 16-byte hex string
+ protected String generateNewSpanId() {
+ return generateRandomHexString(NUM_SPAN_ID_BYTES);
+ }
+
+ // Generates a new SpanContext w/ random traceId,spanId
+ protected SpanContext getNewSpanContext() {
+ String traceId = generateNewTraceId();
+ String spanId = generateNewSpanId();
+ logger.info("traceId=" + traceId + ", spanId=" + spanId);
+
+ return SpanContext.create(traceId, spanId, TraceFlags.getSampled(), TraceState.getDefault());
+ }
+
+ protected Span getNewRootSpanWithContext() {
+ // Execute the DB operation in the context of the custom root span.
+ return tracer
+ .spanBuilder(rootSpanName)
+ .setParent(Context.root().with(Span.wrap(customSpanContext)))
+ .startSpan();
+ }
+
+ protected String grpcSpanName(String rpcName) {
+ return "Sent." + SERVICE + rpcName;
+ }
+
+ protected void waitForTracesToComplete() throws Exception {
+ logger.info("Flushing traces...");
+ CompletableResultCode completableResultCode =
+ openTelemetrySdk.getSdkTracerProvider().forceFlush();
+ completableResultCode.join(TRACE_FORCE_FLUSH_MILLIS, TimeUnit.MILLISECONDS);
+ }
+
+ // Validates `retrievedTrace`. Cloud Trace indexes traces w/ eventual consistency, even when
+ // indexing traceId, therefore the test must retry a few times before the complete trace is
+ // available.
+ // For Transaction traces, there may be more spans than in the trace than specified in
+ // `callStack`. So `numExpectedSpans` is the expected total number of spans (and not just the
+ // spans in `callStack`)
+ protected void fetchAndValidateTrace(
+ String traceId, int numExpectedSpans, List> callStackList)
+ throws InterruptedException {
+ // Large enough count to accommodate eventually consistent Cloud Trace backend
+ int numRetries = GET_TRACE_RETRY_COUNT;
+ // Account for rootSpanName
+ numExpectedSpans++;
+
+ // Fetch traces
+ do {
+ try {
+ retrievedTrace = traceClient_v1.getTrace(projectId, traceId);
+ assertEquals(traceId, retrievedTrace.getTraceId());
+
+ logger.info(
+ "expectedSpanCount="
+ + numExpectedSpans
+ + ", retrievedSpanCount="
+ + retrievedTrace.getSpansCount());
+ } catch (NotFoundException notFound) {
+ logger.info("Trace not found, retrying in " + GET_TRACE_RETRY_BACKOFF_MILLIS + " ms");
+ } catch (IndexOutOfBoundsException outOfBoundsException) {
+ logger.info("Call stack not found in trace. Retrying.");
+ }
+ if (retrievedTrace == null || numExpectedSpans != retrievedTrace.getSpansCount()) {
+ Thread.sleep(GET_TRACE_RETRY_BACKOFF_MILLIS);
+ }
+ } while (numRetries-- > 0
+ && (retrievedTrace == null || numExpectedSpans != retrievedTrace.getSpansCount()));
+
+ if (retrievedTrace == null || numExpectedSpans != retrievedTrace.getSpansCount()) {
+ throw new RuntimeException(
+ "Expected number of spans: "
+ + numExpectedSpans
+ + ", Actual number of spans: "
+ + (retrievedTrace != null
+ ? retrievedTrace.getSpansList().toString()
+ : "Trace NOT_FOUND"));
+ }
+
+ TraceContainer traceContainer = new TraceContainer(rootSpanName, retrievedTrace);
+
+ for (List callStack : callStackList) {
+ // Update all call stacks to be rooted at rootSpanName
+ ArrayList expectedCallStack = new ArrayList<>(callStack);
+
+ // numExpectedSpans should account for rootSpanName (not passed in callStackList)
+ expectedCallStack.add(0, rootSpanName);
+
+ // *May be* the full trace was returned
+ logger.info("Checking if TraceContainer contains the callStack");
+ String[] expectedCallList = new String[expectedCallStack.size()];
+ if (!traceContainer.containsCallStack(expectedCallStack.toArray(expectedCallList))) {
+ throw new RuntimeException(
+ "Expected spans: "
+ + expectedCallList.toString()
+ + ", Actual spans: "
+ + (retrievedTrace != null
+ ? retrievedTrace.getSpansList().toString()
+ : "Trace NOT_FOUND"));
+ }
+ logger.severe("CallStack not found in TraceContainer.");
+ }
+ }
+
+ // Validates `retrievedTrace`. Cloud Trace indexes traces w/ eventual consistency, even when
+ // indexing traceId, therefore the test must retry a few times before the complete trace is
+ // available.
+ // For Non-Transaction traces, there is a 1:1 ratio of spans in `spanNames` and in the trace.
+ protected void fetchAndValidateTrace(String traceId, String... spanNames)
+ throws InterruptedException {
+ fetchAndValidateTrace(traceId, spanNames.length, Arrays.asList(Arrays.asList(spanNames)));
+ }
+
+ @Test
+ public void traceContainerTest() throws Exception {
+ // Make sure the test has a new SpanContext (and TraceId for injection)
+ assertNotNull(customSpanContext);
+
+ // Inject new trace ID
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ firestore.collection("col").whereEqualTo("foo", "my_non_existent_value").get().get();
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+
+ Trace traceResp = null;
+ int expectedSpanCount = 3;
+
+ int numRetries = GET_TRACE_RETRY_COUNT;
+ do {
+ try {
+ traceResp = traceClient_v1.getTrace(projectId, customSpanContext.getTraceId());
+ if (traceResp.getSpansCount() == expectedSpanCount) {
+ logger.info("Success: Got " + expectedSpanCount + " spans.");
+ break;
+ }
+ } catch (NotFoundException notFoundException) {
+ Thread.sleep(GET_TRACE_RETRY_BACKOFF_MILLIS);
+ logger.info("Trace not found, retrying in " + GET_TRACE_RETRY_BACKOFF_MILLIS + " ms");
+ }
+ logger.info(
+ "Trace Found. The trace did not contain "
+ + expectedSpanCount
+ + " spans. Going to retry.");
+ numRetries--;
+ } while (numRetries > 0);
+
+ // Make sure we got as many spans as we expected.
+ assertNotNull(traceResp);
+ assertEquals(expectedSpanCount, traceResp.getSpansCount());
+
+ TraceContainer traceCont = new TraceContainer(rootSpanName, traceResp);
+
+ // Contains exact path
+ assertTrue(
+ traceCont.containsCallStack(
+ rootSpanName, SPAN_NAME_QUERY_GET, grpcSpanName(RUN_QUERY_RPC_NAME)));
+
+ // Top-level mismatch
+ assertFalse(traceCont.containsCallStack(SPAN_NAME_QUERY_GET, RUN_QUERY_RPC_NAME));
+
+ // Mid-level match
+ assertFalse(traceCont.containsCallStack(rootSpanName, SPAN_NAME_QUERY_GET));
+
+ // Leaf-level mismatch/missing
+ assertFalse(
+ traceCont.containsCallStack(
+ rootSpanName, SPAN_NAME_QUERY_GET, RUN_AGGREGATION_QUERY_RPC_NAME));
+ }
+
+ @Test
+ // Trace an Aggregation.Get request
+ public void aggregateQueryGetTraceTest() throws Exception {
+ // Make sure the test has a new SpanContext (and TraceId for injection)
+ assertNotNull(customSpanContext);
+
+ // Inject new trace ID
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ // Execute the Firestore SDK op
+ firestore.collection("col").count().get().get();
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+
+ // Read and validate traces
+ fetchAndValidateTrace(
+ customSpanContext.getTraceId(),
+ "AggregationQuery.Get",
+ grpcSpanName(RUN_AGGREGATION_QUERY_RPC_NAME));
+ }
+
+ @Test
+ public void bulkWriterCommitTraceTest() throws Exception {
+ // Make sure the test has a new SpanContext (and TraceId for injection)
+ assertNotNull(customSpanContext);
+
+ // Inject new trace ID
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ // Execute the Firestore SDK op
+ ScheduledExecutorService bulkWriterExecutor = Executors.newSingleThreadScheduledExecutor();
+ BulkWriter bulkWriter =
+ firestore.bulkWriter(BulkWriterOptions.builder().setExecutor(bulkWriterExecutor).build());
+ bulkWriter.set(
+ firestore.collection("col").document("foo"),
+ Collections.singletonMap("bulk-foo", "bulk-bar"));
+ bulkWriter.close();
+ bulkWriterExecutor.awaitTermination(100, TimeUnit.MILLISECONDS);
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+
+ // Read and validate traces
+ fetchAndValidateTrace(
+ customSpanContext.getTraceId(),
+ SPAN_NAME_BULK_WRITER_COMMIT,
+ grpcSpanName(BATCH_WRITE_RPC_NAME));
+ }
+
+ @Test
+ public void partitionQueryTraceTest() throws Exception {
+ // Make sure the test has a new SpanContext (and TraceId for injection)
+ assertNotNull(customSpanContext);
+
+ // Inject new trace ID
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ CollectionGroup collectionGroup = firestore.collectionGroup("col");
+ collectionGroup.getPartitions(3).get();
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+
+ // Read and validate traces
+ fetchAndValidateTrace(
+ customSpanContext.getTraceId(),
+ SPAN_NAME_PARTITION_QUERY,
+ grpcSpanName(SPAN_NAME_PARTITION_QUERY));
+ }
+
+ @Test
+ public void collectionListDocumentsTraceTest() throws Exception {
+ // Make sure the test has a new SpanContext (and TraceId for injection)
+ assertNotNull(customSpanContext);
+
+ // Inject new trace ID
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ firestore.collection("col").listDocuments();
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+
+ // Read and validate traces
+ fetchAndValidateTrace(
+ customSpanContext.getTraceId(),
+ SPAN_NAME_COL_REF_LIST_DOCUMENTS,
+ grpcSpanName(LIST_DOCUMENTS_RPC_NAME));
+ }
+
+ @Test
+ public void docRefCreateTraceTest() throws Exception {
+ // Make sure the test has a new SpanContext (and TraceId for injection)
+ assertNotNull(customSpanContext);
+
+ // Inject new trace ID
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ firestore.collection("col").document().create(Collections.singletonMap("foo", "bar")).get();
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+
+ // Read and validate traces
+ fetchAndValidateTrace(
+ customSpanContext.getTraceId(),
+ SPAN_NAME_DOC_REF_CREATE,
+ SPAN_NAME_BATCH_COMMIT,
+ grpcSpanName(COMMIT_RPC_NAME));
+ }
+
+ @Test
+ public void docRefCreate2TraceTest() throws Exception {
+ // Make sure the test has a new SpanContext (and TraceId for injection)
+ assertNotNull(customSpanContext);
+
+ // Inject new trace ID
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ firestore.collection("col").document().create(new Pojo(1)).get();
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+
+ // Read and validate traces
+ fetchAndValidateTrace(
+ customSpanContext.getTraceId(),
+ SPAN_NAME_DOC_REF_CREATE,
+ SPAN_NAME_BATCH_COMMIT,
+ grpcSpanName(COMMIT_RPC_NAME));
+ }
+
+ @Test
+ public void docRefSetTraceTest() throws Exception {
+ // Make sure the test has a new SpanContext (and TraceId for injection)
+ assertNotNull(customSpanContext);
+
+ // Inject new trace ID
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ firestore.collection("col").document("foo").set(Collections.singletonMap("foo", "bar")).get();
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+
+ // Read and validate traces
+ fetchAndValidateTrace(
+ customSpanContext.getTraceId(),
+ SPAN_NAME_DOC_REF_SET,
+ SPAN_NAME_BATCH_COMMIT,
+ grpcSpanName(COMMIT_RPC_NAME));
+ }
+
+ @Test
+ public void docRefSet2TraceTest() throws Exception {
+ // Make sure the test has a new SpanContext (and TraceId for injection)
+ assertNotNull(customSpanContext);
+
+ // Inject new trace ID
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ firestore
+ .collection("col")
+ .document("foo")
+ .set(Collections.singletonMap("foo", "bar"), SetOptions.merge())
+ .get();
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+
+ // Read and validate traces
+ fetchAndValidateTrace(
+ customSpanContext.getTraceId(),
+ SPAN_NAME_DOC_REF_SET,
+ SPAN_NAME_BATCH_COMMIT,
+ grpcSpanName(COMMIT_RPC_NAME));
+ }
+
+ @Test
+ public void docRefSet3TraceTest() throws Exception {
+ // Make sure the test has a new SpanContext (and TraceId for injection)
+ assertNotNull(customSpanContext);
+
+ // Inject new trace ID
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ firestore.collection("col").document("foo").set(new Pojo(1)).get();
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+
+ // Read and validate traces
+ fetchAndValidateTrace(
+ customSpanContext.getTraceId(),
+ SPAN_NAME_DOC_REF_SET,
+ SPAN_NAME_BATCH_COMMIT,
+ grpcSpanName(COMMIT_RPC_NAME));
+ }
+
+ @Test
+ public void docRefSet4TraceTest() throws Exception {
+ // Make sure the test has a new SpanContext (and TraceId for injection)
+ assertNotNull(customSpanContext);
+
+ // Inject new trace ID
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ firestore.collection("col").document("foo").set(new Pojo(1), SetOptions.merge()).get();
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+
+ // Read and validate traces
+ fetchAndValidateTrace(
+ customSpanContext.getTraceId(),
+ SPAN_NAME_DOC_REF_SET,
+ SPAN_NAME_BATCH_COMMIT,
+ grpcSpanName(COMMIT_RPC_NAME));
+ }
+
+ @Test
+ public void docRefUpdateTraceTest() throws Exception {
+ // Make sure the test has a new SpanContext (and TraceId for injection)
+ assertNotNull(customSpanContext);
+
+ // Inject new trace ID
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ firestore
+ .collection("col")
+ .document("foo")
+ .update(Collections.singletonMap("foo", "bar"))
+ .get();
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+
+ // Read and validate traces
+ fetchAndValidateTrace(
+ customSpanContext.getTraceId(),
+ SPAN_NAME_DOC_REF_UPDATE,
+ SPAN_NAME_BATCH_COMMIT,
+ grpcSpanName(COMMIT_RPC_NAME));
+ }
+
+ @Test
+ public void docRefUpdate2TraceTest() throws Exception {
+ // Make sure the test has a new SpanContext (and TraceId for injection)
+ assertNotNull(customSpanContext);
+
+ // Inject new trace ID
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ firestore
+ .collection("col")
+ .document("foo")
+ .update(Collections.singletonMap("foo", "bar"), Precondition.NONE)
+ .get();
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+
+ // Read and validate traces
+ fetchAndValidateTrace(
+ customSpanContext.getTraceId(),
+ SPAN_NAME_DOC_REF_UPDATE,
+ SPAN_NAME_BATCH_COMMIT,
+ grpcSpanName(COMMIT_RPC_NAME));
+ }
+
+ @Test
+ public void docRefUpdate3TraceTest() throws Exception {
+ // Make sure the test has a new SpanContext (and TraceId for injection)
+ assertNotNull(customSpanContext);
+
+ // Inject new trace ID
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ firestore.collection("col").document("foo").update("key", "value", "key2", "value2").get();
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+
+ // Read and validate traces
+ fetchAndValidateTrace(
+ customSpanContext.getTraceId(),
+ SPAN_NAME_DOC_REF_UPDATE,
+ SPAN_NAME_BATCH_COMMIT,
+ grpcSpanName(COMMIT_RPC_NAME));
+ }
+
+ @Test
+ public void docRefUpdate4TraceTest() throws Exception {
+ // Make sure the test has a new SpanContext (and TraceId for injection)
+ assertNotNull(customSpanContext);
+
+ // Inject new trace ID
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ firestore
+ .collection("col")
+ .document("foo")
+ .update(FieldPath.of("key"), "value", FieldPath.of("key2"), "value2")
+ .get();
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+
+ // Read and validate traces
+ fetchAndValidateTrace(
+ customSpanContext.getTraceId(),
+ SPAN_NAME_DOC_REF_UPDATE,
+ SPAN_NAME_BATCH_COMMIT,
+ grpcSpanName(COMMIT_RPC_NAME));
+ }
+
+ @Test
+ public void docRefUpdate5TraceTest() throws Exception {
+ // Make sure the test has a new SpanContext (and TraceId for injection)
+ assertNotNull(customSpanContext);
+
+ // Inject new trace ID
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ firestore
+ .collection("col")
+ .document("foo")
+ .update(Precondition.NONE, "key", "value", "key2", "value2")
+ .get();
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+
+ // Read and validate traces
+ fetchAndValidateTrace(
+ customSpanContext.getTraceId(),
+ SPAN_NAME_DOC_REF_UPDATE,
+ SPAN_NAME_BATCH_COMMIT,
+ grpcSpanName(COMMIT_RPC_NAME));
+ }
+
+ @Test
+ public void docRefUpdate6TraceTest() throws Exception {
+ // Make sure the test has a new SpanContext (and TraceId for injection)
+ assertNotNull(customSpanContext);
+
+ // Inject new trace ID
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ firestore
+ .collection("col")
+ .document("foo")
+ .update(Precondition.NONE, FieldPath.of("key"), "value", FieldPath.of("key2"), "value2")
+ .get();
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+
+ // Read and validate traces
+ fetchAndValidateTrace(
+ customSpanContext.getTraceId(),
+ SPAN_NAME_DOC_REF_UPDATE,
+ SPAN_NAME_BATCH_COMMIT,
+ grpcSpanName(COMMIT_RPC_NAME));
+ }
+
+ @Test
+ public void docRefDeleteTraceTest() throws Exception {
+ // Make sure the test has a new SpanContext (and TraceId for injection)
+ assertNotNull(customSpanContext);
+
+ // Inject new trace ID
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ firestore.collection("col").document("doc0").delete().get();
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+
+ // Read and validate traces
+ fetchAndValidateTrace(
+ customSpanContext.getTraceId(),
+ SPAN_NAME_DOC_REF_DELETE,
+ SPAN_NAME_BATCH_COMMIT,
+ grpcSpanName(COMMIT_RPC_NAME));
+ }
+
+ @Test
+ public void docRefDelete2TraceTest() throws Exception {
+ // Make sure the test has a new SpanContext (and TraceId for injection)
+ assertNotNull(customSpanContext);
+
+ // Inject new trace ID
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ firestore.collection("col").document("doc0").delete(Precondition.NONE).get();
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+
+ // Read and validate traces
+ fetchAndValidateTrace(
+ customSpanContext.getTraceId(),
+ SPAN_NAME_DOC_REF_DELETE,
+ SPAN_NAME_BATCH_COMMIT,
+ grpcSpanName(COMMIT_RPC_NAME));
+ }
+
+ @Test
+ public void docRefGetTraceTest() throws Exception {
+ // Make sure the test has a new SpanContext (and TraceId for injection)
+ assertNotNull(customSpanContext);
+
+ // Inject new trace ID
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ firestore.collection("col").document("doc0").get().get();
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+
+ // Read and validate traces
+ fetchAndValidateTrace(
+ customSpanContext.getTraceId(),
+ SPAN_NAME_DOC_REF_GET,
+ grpcSpanName(BATCH_GET_DOCUMENTS_RPC_NAME));
+ }
+
+ @Test
+ public void docRefGet2TraceTest() throws Exception {
+ // Make sure the test has a new SpanContext (and TraceId for injection)
+ assertNotNull(customSpanContext);
+
+ // Inject new trace ID
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ firestore.collection("col").document("doc0").get(FieldMask.of("foo")).get();
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+
+ // Read and validate traces
+ fetchAndValidateTrace(
+ customSpanContext.getTraceId(),
+ SPAN_NAME_DOC_REF_GET,
+ grpcSpanName(BATCH_GET_DOCUMENTS_RPC_NAME));
+ }
+
+ @Test
+ public void docListCollectionsTraceTest() throws Exception {
+ // Make sure the test has a new SpanContext (and TraceId for injection)
+ assertNotNull(customSpanContext);
+
+ // Inject new trace ID
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ firestore.collection("col").document("doc0").listCollections();
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+
+ // Read and validate traces
+ fetchAndValidateTrace(
+ customSpanContext.getTraceId(),
+ SPAN_NAME_DOC_REF_LIST_COLLECTIONS,
+ grpcSpanName(LIST_COLLECTIONS_RPC_NAME));
+ }
+
+ @Test
+ public void getAllTraceTest() throws Exception {
+ // Make sure the test has a new SpanContext (and TraceId for injection)
+ assertNotNull(customSpanContext);
+
+ // Inject new trace ID
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ DocumentReference docRef0 = firestore.collection("col").document();
+ DocumentReference docRef1 = firestore.collection("col").document();
+ DocumentReference[] docs = {docRef0, docRef1};
+ firestore.getAll(docs).get();
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+
+ fetchAndValidateTrace(
+ customSpanContext.getTraceId(), grpcSpanName(BATCH_GET_DOCUMENTS_RPC_NAME));
+ }
+
+ @Test
+ public void queryGetTraceTest() throws Exception {
+ // Make sure the test has a new SpanContext (and TraceId for injection)
+ assertNotNull(customSpanContext);
+
+ // Inject new trace ID
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ firestore.collection("col").whereEqualTo("foo", "my_non_existent_value").get().get();
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+
+ fetchAndValidateTrace(
+ customSpanContext.getTraceId(), SPAN_NAME_QUERY_GET, grpcSpanName(RUN_QUERY_RPC_NAME));
+ }
+
+ @Test
+ public void transactionTraceTest() throws Exception {
+ // Make sure the test has a new SpanContext (and TraceId for injection)
+ assertNotNull(customSpanContext);
+
+ // Inject new trace ID
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ firestore
+ .runTransaction(
+ transaction -> {
+ Query q = firestore.collection("col").whereGreaterThan("bla", "");
+ DocumentReference d = firestore.collection("col").document("foo");
+ // Document Query.
+ transaction.get(q).get();
+
+ // Aggregation Query.
+ transaction.get(q.count());
+
+ // Get multiple documents.
+ transaction.getAll(d, d).get();
+
+ // Commit 2 documents.
+ transaction.set(
+ firestore.collection("foo").document("bar"),
+ Collections.singletonMap("foo", "bar"));
+ transaction.set(
+ firestore.collection("foo").document("bar2"),
+ Collections.singletonMap("foo2", "bar2"));
+ return 0;
+ })
+ .get();
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+
+ fetchAndValidateTrace(
+ customSpanContext.getTraceId(),
+ /*numExpectedSpans=*/ 11,
+ Arrays.asList(
+ Arrays.asList(
+ SPAN_NAME_TRANSACTION_RUN,
+ SPAN_NAME_TRANSACTION_BEGIN,
+ grpcSpanName(BEGIN_TRANSACTION_RPC_NAME)),
+ Arrays.asList(
+ SPAN_NAME_TRANSACTION_RUN,
+ SPAN_NAME_TRANSACTION_BEGIN,
+ grpcSpanName(BEGIN_TRANSACTION_RPC_NAME)),
+ Arrays.asList(
+ SPAN_NAME_TRANSACTION_RUN,
+ SPAN_NAME_TRANSACTION_GET_QUERY,
+ grpcSpanName(RUN_QUERY_RPC_NAME)),
+ Arrays.asList(
+ SPAN_NAME_TRANSACTION_RUN,
+ SPAN_NAME_TRANSACTION_GET_AGGREGATION_QUERY,
+ grpcSpanName(RUN_AGGREGATION_QUERY_RPC_NAME)),
+ Arrays.asList(
+ SPAN_NAME_TRANSACTION_RUN,
+ SPAN_NAME_TRANSACTION_GET_DOCUMENTS,
+ grpcSpanName(BATCH_GET_DOCUMENTS_RPC_NAME)),
+ Arrays.asList(
+ SPAN_NAME_TRANSACTION_RUN,
+ SPAN_NAME_TRANSACTION_COMMIT,
+ grpcSpanName(COMMIT_RPC_NAME))));
+ }
+
+ @Test
+ public void transactionRollbackTraceTest() throws Exception {
+ String myErrorMessage = "My error message.";
+ // Make sure the test has a new SpanContext (and TraceId for injection)
+ assertNotNull(customSpanContext);
+
+ // Inject new trace ID
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ firestore
+ .runTransaction(
+ transaction -> {
+ if (true) {
+ throw (new Exception(myErrorMessage));
+ }
+ return 0;
+ })
+ .get();
+ } catch (Exception e) {
+ // Catch and move on.
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+
+ fetchAndValidateTrace(
+ customSpanContext.getTraceId(),
+ /*numExpectedSpans=*/ 5,
+ Arrays.asList(
+ Arrays.asList(
+ SPAN_NAME_TRANSACTION_RUN,
+ SPAN_NAME_TRANSACTION_BEGIN,
+ grpcSpanName(BEGIN_TRANSACTION_RPC_NAME)),
+ Arrays.asList(
+ SPAN_NAME_TRANSACTION_RUN,
+ SPAN_NAME_TRANSACTION_ROLLBACK,
+ grpcSpanName(ROLLBACK_RPC_NAME))));
+ }
+
+ @Test
+ public void writeBatchTraceTest() throws Exception {
+ // Make sure the test has a new SpanContext (and TraceId for injection)
+ assertNotNull(customSpanContext);
+
+ // Inject new trace ID
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ WriteBatch batch = firestore.batch();
+ DocumentReference docRef = firestore.collection("foo").document();
+ batch.create(docRef, Collections.singletonMap("foo", "bar"));
+ batch.update(docRef, Collections.singletonMap("foo", "bar"));
+ batch.delete(docRef);
+ batch.commit().get();
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+
+ fetchAndValidateTrace(
+ customSpanContext.getTraceId(), SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME));
+ }
+}
diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTestGlobalOtel.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTestGlobalOtel.java
new file mode 100644
index 000000000..85bfb5437
--- /dev/null
+++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTestGlobalOtel.java
@@ -0,0 +1,27 @@
+/*
+ * Copyright 2024 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.firestore.it;
+
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class ITE2ETracingTestGlobalOtel extends ITE2ETracingTest {
+ @Override
+ protected boolean isUsingGlobalOpenTelemetrySDK() {
+ return true;
+ }
+}
diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTestNonGlobalOtel.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTestNonGlobalOtel.java
new file mode 100644
index 000000000..21272ff2e
--- /dev/null
+++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTestNonGlobalOtel.java
@@ -0,0 +1,27 @@
+/*
+ * Copyright 2024 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.firestore.it;
+
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class ITE2ETracingTestNonGlobalOtel extends ITE2ETracingTest {
+ @Override
+ protected boolean isUsingGlobalOpenTelemetrySDK() {
+ return false;
+ }
+}
diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITTracingTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITTracingTest.java
new file mode 100644
index 000000000..46e8294e1
--- /dev/null
+++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITTracingTest.java
@@ -0,0 +1,814 @@
+/*
+ * Copyright 2024 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.firestore.it;
+
+import static com.google.cloud.firestore.telemetry.TraceUtil.*;
+import static io.opentelemetry.semconv.resource.attributes.ResourceAttributes.SERVICE_NAME;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+
+import com.google.cloud.firestore.BulkWriter;
+import com.google.cloud.firestore.BulkWriterOptions;
+import com.google.cloud.firestore.CollectionGroup;
+import com.google.cloud.firestore.DocumentReference;
+import com.google.cloud.firestore.FieldMask;
+import com.google.cloud.firestore.FieldPath;
+import com.google.cloud.firestore.Firestore;
+import com.google.cloud.firestore.FirestoreOpenTelemetryOptions;
+import com.google.cloud.firestore.FirestoreOptions;
+import com.google.cloud.firestore.Precondition;
+import com.google.cloud.firestore.Query;
+import com.google.cloud.firestore.SetOptions;
+import com.google.cloud.firestore.WriteBatch;
+import com.google.common.base.Preconditions;
+import io.opentelemetry.api.GlobalOpenTelemetry;
+import io.opentelemetry.api.common.AttributeKey;
+import io.opentelemetry.api.common.Attributes;
+import io.opentelemetry.api.trace.StatusCode;
+import io.opentelemetry.sdk.OpenTelemetrySdk;
+import io.opentelemetry.sdk.OpenTelemetrySdkBuilder;
+import io.opentelemetry.sdk.common.CompletableResultCode;
+import io.opentelemetry.sdk.resources.Resource;
+import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter;
+import io.opentelemetry.sdk.trace.SdkTracerProvider;
+import io.opentelemetry.sdk.trace.SpanProcessor;
+import io.opentelemetry.sdk.trace.data.EventData;
+import io.opentelemetry.sdk.trace.data.SpanData;
+import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor;
+import io.opentelemetry.sdk.trace.samplers.Sampler;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+import javax.annotation.Nullable;
+import org.junit.After;
+import org.junit.AfterClass;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TestName;
+
+public abstract class ITTracingTest {
+ protected abstract boolean isUsingGlobalOpenTelemetrySDK();
+
+ private static final Logger logger =
+ Logger.getLogger(com.google.cloud.firestore.it.ITTracingTest.class.getName());
+
+ private static final int TRACE_FORCE_FLUSH_MILLIS = 1000;
+ private static final int TRACE_PROVIDER_SHUTDOWN_MILLIS = 1000;
+ private static final int IN_MEMORY_SPAN_EXPORTER_DELAY_MILLIS = 50;
+ private static final String SERVICE = "google.firestore.v1.Firestore/";
+ private static final String BATCH_GET_DOCUMENTS_RPC_NAME = "BatchGetDocuments";
+ private static final String COMMIT_RPC_NAME = "Commit";
+ private static final String LIST_DOCUMENTS_RPC_NAME = "ListDocuments";
+ private static final String LIST_COLLECTIONS_RPC_NAME = "ListCollectionIds";
+ private static final String BATCH_WRITE_RPC_NAME = "BatchWrite";
+ private static final String RUN_QUERY_RPC_NAME = "RunQuery";
+ private static final String RUN_AGGREGATION_QUERY_RPC_NAME = "RunAggregationQuery";
+ private static final String BEGIN_TRANSACTION_RPC_NAME = "BeginTransaction";
+ private static final String ROLLBACK_RPC_NAME = "Rollback";
+
+ private static OpenTelemetrySdk openTelemetrySdk;
+
+ // We use an InMemorySpanExporter for testing which keeps all generated trace spans
+ // in memory so that we can check their correctness.
+ protected InMemorySpanExporter inMemorySpanExporter;
+
+ protected Firestore firestore;
+
+ Map spanNameToSpanId = new HashMap<>();
+ Map spanIdToParentSpanId = new HashMap<>();
+ Map spanNameToSpanData = new HashMap<>();
+
+ @Rule public TestName testName = new TestName();
+
+ @Before
+ public void before() {
+ inMemorySpanExporter = InMemorySpanExporter.create();
+
+ Resource resource =
+ Resource.getDefault().merge(Resource.builder().put(SERVICE_NAME, "Sparky").build());
+ SpanProcessor inMemorySpanProcessor = SimpleSpanProcessor.create(inMemorySpanExporter);
+ FirestoreOptions.Builder optionsBuilder = FirestoreOptions.newBuilder();
+ FirestoreOpenTelemetryOptions.Builder otelOptionsBuilder =
+ FirestoreOpenTelemetryOptions.newBuilder();
+ OpenTelemetrySdkBuilder openTelemetrySdkBuilder =
+ OpenTelemetrySdk.builder()
+ .setTracerProvider(
+ SdkTracerProvider.builder()
+ .setResource(resource)
+ .addSpanProcessor(inMemorySpanProcessor)
+ .setSampler(Sampler.alwaysOn())
+ .build());
+
+ if (isUsingGlobalOpenTelemetrySDK()) {
+ GlobalOpenTelemetry.resetForTest();
+ openTelemetrySdk = openTelemetrySdkBuilder.buildAndRegisterGlobal();
+ optionsBuilder.setOpenTelemetryOptions(otelOptionsBuilder.setTracingEnabled(true).build());
+ } else {
+ openTelemetrySdk = openTelemetrySdkBuilder.build();
+ optionsBuilder.setOpenTelemetryOptions(
+ otelOptionsBuilder.setTracingEnabled(true).setOpenTelemetry(openTelemetrySdk).build());
+ }
+
+ String namedDb = System.getProperty("FIRESTORE_NAMED_DATABASE");
+ if (namedDb != null) {
+ logger.log(
+ Level.INFO,
+ String.format(
+ "Integration test using named database %s for test %s",
+ namedDb, testName.getMethodName()));
+ optionsBuilder = optionsBuilder.setDatabaseId(namedDb);
+ } else {
+ logger.log(
+ Level.INFO,
+ String.format(
+ "Integration test using default database for test %s", testName.getMethodName()));
+ }
+ firestore = optionsBuilder.build().getService();
+
+ // Clean up existing maps.
+ spanNameToSpanId.clear();
+ spanIdToParentSpanId.clear();
+ spanNameToSpanData.clear();
+ }
+
+ @After
+ public void after() throws Exception {
+ Preconditions.checkNotNull(
+ firestore,
+ "Error instantiating Firestore. Check that the service account credentials were properly set.");
+ firestore.shutdown();
+ inMemorySpanExporter.reset();
+ }
+
+ @AfterClass
+ public static void teardown() {
+ CompletableResultCode completableResultCode =
+ openTelemetrySdk.getSdkTracerProvider().shutdown();
+ completableResultCode.join(TRACE_PROVIDER_SHUTDOWN_MILLIS, TimeUnit.MILLISECONDS);
+ }
+
+ void waitForTracesToComplete() throws Exception {
+ // We need to call `firestore.close()` because that will also close the
+ // gRPC channel and hence force the gRPC instrumentation library to flush
+ // its spans.
+ firestore.close();
+
+ // The same way that querying the Cloud Trace backend may not give us the
+ // full trace on the first try, querying the in-memory traces may not result
+ // in the full trace immediately. Note that performing the `flush` is not
+ // enough. This doesn't pose an issue in practice, but can make tests flaky.
+ // Therefore, we're adding a delay to make sure we avoid any flakiness.
+ inMemorySpanExporter.flush().join(IN_MEMORY_SPAN_EXPORTER_DELAY_MILLIS, TimeUnit.MILLISECONDS);
+ TimeUnit.MILLISECONDS.sleep(IN_MEMORY_SPAN_EXPORTER_DELAY_MILLIS);
+
+ CompletableResultCode completableResultCode =
+ openTelemetrySdk.getSdkTracerProvider().forceFlush();
+ completableResultCode.join(TRACE_FORCE_FLUSH_MILLIS, TimeUnit.MILLISECONDS);
+ }
+
+ // Prepares all the spans in memory for inspection.
+ List prepareSpans() throws Exception {
+ waitForTracesToComplete();
+ List spans = inMemorySpanExporter.getFinishedSpanItems();
+ buildSpanMaps(spans);
+ printSpans();
+ return spans;
+ }
+
+ void buildSpanMaps(List spans) {
+ for (SpanData spanData : spans) {
+ spanNameToSpanData.put(spanData.getName(), spanData);
+ spanNameToSpanId.put(spanData.getName(), spanData.getSpanId());
+ spanIdToParentSpanId.put(spanData.getSpanId(), spanData.getParentSpanId());
+ }
+ }
+
+ // Returns the SpanData object for the span with the given name.
+ // Returns null if no span with the given name exists.
+ @Nullable
+ SpanData getSpanByName(String spanName) {
+ return spanNameToSpanData.get(spanName);
+ }
+
+ // Returns the SpanData object for the gRPC span with the given RPC name.
+ // Returns null if no such span exists.
+ @Nullable
+ SpanData getGrpcSpanByName(String rpcName) {
+ return getSpanByName(SERVICE + rpcName);
+ }
+
+ String grpcSpanName(String rpcName) {
+ return SERVICE + rpcName;
+ }
+
+ void assertSameTrace(SpanData... spans) {
+ if (spans.length > 1) {
+ String traceId = spans[0].getTraceId();
+ for (SpanData spanData : spans) {
+ assertEquals(traceId, spanData.getTraceId());
+ }
+ }
+ }
+
+ // Helper to see the spans in standard output while developing tests
+ void printSpans() {
+ for (SpanData spanData : spanNameToSpanData.values()) {
+ logger.log(
+ Level.FINE,
+ String.format(
+ "SPAN ID:%s, ParentID:%s, KIND:%s, TRACE ID:%s, NAME:%s, ATTRIBUTES:%s, EVENTS:%s\n",
+ spanData.getSpanId(),
+ spanData.getParentSpanId(),
+ spanData.getKind(),
+ spanData.getTraceId(),
+ spanData.getName(),
+ spanData.getAttributes().toString(),
+ spanData.getEvents().toString()));
+ }
+ }
+
+ // Asserts that the span hierarchy exists for the given span names. The hierarchy starts with the
+ // root span, followed
+ // by the child span, grandchild span, and so on. It also asserts that all the given spans belong
+ // to the same trace,
+ // and that Firestore-generated spans contain the expected Firestore attributes.
+ void assertSpanHierarchy(String... spanNamesHierarchy) {
+ List spanNames = Arrays.asList(spanNamesHierarchy);
+
+ for (int i = 0; i + 1 < spanNames.size(); ++i) {
+ String parentSpanName = spanNames.get(i);
+ String childSpanName = spanNames.get(i + 1);
+ SpanData parentSpan = getSpanByName(parentSpanName);
+ SpanData childSpan = getSpanByName(childSpanName);
+ assertNotNull(parentSpan);
+ assertNotNull(childSpan);
+ assertEquals(childSpan.getParentSpanId(), parentSpan.getSpanId());
+ assertSameTrace(childSpan, parentSpan);
+ // gRPC spans do not have Firestore attributes.
+ if (!parentSpanName.startsWith(SERVICE)) {
+ assertHasExpectedAttributes(parentSpan);
+ }
+ if (!childSpanName.startsWith(SERVICE)) {
+ assertHasExpectedAttributes(childSpan);
+ }
+ }
+ }
+
+ void assertHasExpectedAttributes(SpanData spanData, String... additionalExpectedAttributes) {
+ // All Firestore-generated spans have the settings attributes.
+ List expectedAttributes =
+ Arrays.asList(
+ "gcp.firestore.memory_utilization",
+ "gcp.firestore.settings.host",
+ "gcp.firestore.settings.project_id",
+ "gcp.firestore.settings.database_id",
+ "gcp.firestore.settings.channel.needs_credentials",
+ "gcp.firestore.settings.channel.needs_endpoint",
+ "gcp.firestore.settings.channel.needs_headers",
+ "gcp.firestore.settings.channel.should_auto_close",
+ "gcp.firestore.settings.channel.transport_name",
+ "gcp.firestore.settings.retry_settings.max_rpc_timeout",
+ "gcp.firestore.settings.retry_settings.retry_delay_multiplier",
+ "gcp.firestore.settings.retry_settings.initial_retry_delay",
+ "gcp.firestore.settings.credentials.authentication_type",
+ "gcp.firestore.settings.retry_settings.max_attempts",
+ "gcp.firestore.settings.retry_settings.max_retry_delay",
+ "gcp.firestore.settings.retry_settings.rpc_timeout_multiplier",
+ "gcp.firestore.settings.retry_settings.total_timeout",
+ "gcp.firestore.settings.retry_settings.initial_rpc_timeout");
+
+ expectedAttributes.addAll(Arrays.asList(additionalExpectedAttributes));
+
+ Attributes spanAttributes = spanData.getAttributes();
+ for (String expectedAttribute : expectedAttributes) {
+ assertNotNull(spanAttributes.get(AttributeKey.stringKey(expectedAttribute)));
+ }
+ }
+
+ // Returns true if and only if the given span data contains an event with the given name and the
+ // given expected
+ // attributes.
+ boolean hasEvent(SpanData spanData, String eventName, @Nullable Attributes expectedAttributes) {
+ if (spanData == null) {
+ return false;
+ }
+
+ logger.log(
+ Level.INFO,
+ String.format(
+ "Checking if span named '%s' (ID='%s') contains an event named '%s'",
+ spanData.getName(), spanData.getSpanId(), eventName));
+
+ List events = spanData.getEvents();
+ for (EventData event : events) {
+ if (event.getName().equals(eventName)) {
+ if (expectedAttributes == null) {
+ return true;
+ }
+
+ // Make sure attributes also match.
+ Attributes eventAttributes = event.getAttributes();
+ return expectedAttributes.equals(eventAttributes);
+ }
+ }
+ return false;
+ }
+
+ // This is a POJO used for testing APIs that take a POJO.
+ public static class Pojo {
+ public int bar;
+
+ public Pojo() {
+ bar = 0;
+ }
+
+ public Pojo(int bar) {
+ this.bar = bar;
+ }
+
+ public int getBar() {
+ return bar;
+ }
+
+ public void setBar(int bar) {
+ this.bar = bar;
+ }
+ }
+
+ @Test
+ public void aggregateQueryGet() throws Exception {
+ firestore.collection("col").count().get().get();
+ waitForTracesToComplete();
+ List spans = inMemorySpanExporter.getFinishedSpanItems();
+ buildSpanMaps(spans);
+ assertEquals(2, spans.size());
+ SpanData getSpan = getSpanByName(SPAN_NAME_AGGREGATION_QUERY_GET);
+ SpanData grpcSpan = getGrpcSpanByName(SPAN_NAME_RUN_AGGREGATION_QUERY);
+ assertNotNull(getSpan);
+ assertNotNull(grpcSpan);
+ assertEquals(grpcSpan.getParentSpanId(), getSpan.getSpanId());
+ assertSameTrace(getSpan, grpcSpan);
+ assertHasExpectedAttributes(getSpan);
+ List events = getSpan.getEvents();
+ assertTrue(events.size() > 0);
+ assertTrue(events.get(0).getAttributes().size() > 0);
+ assertEquals(events.get(0).getName(), "RunAggregationQuery Stream started.");
+ assertEquals(
+ events.get(0).getAttributes().get(AttributeKey.longKey(ATTRIBUTE_KEY_ATTEMPT)).longValue(),
+ 0);
+ }
+
+ @Test
+ public void bulkWriterCommit() throws Exception {
+ ScheduledExecutorService bulkWriterExecutor = Executors.newSingleThreadScheduledExecutor();
+ BulkWriter bulkWriter =
+ firestore.bulkWriter(BulkWriterOptions.builder().setExecutor(bulkWriterExecutor).build());
+ bulkWriter.set(
+ firestore.collection("col").document("foo"),
+ Collections.singletonMap("bulk-foo", "bulk-bar"));
+ bulkWriter.close();
+ bulkWriterExecutor.awaitTermination(100, TimeUnit.MILLISECONDS);
+
+ List spans = prepareSpans();
+ assertEquals(2, spans.size());
+ assertSpanHierarchy(SPAN_NAME_BULK_WRITER_COMMIT, grpcSpanName(BATCH_WRITE_RPC_NAME));
+ }
+
+ @Test
+ public void partitionQuery() throws Exception {
+ CollectionGroup collectionGroup = firestore.collectionGroup("col");
+ collectionGroup.getPartitions(3).get();
+
+ List spans = prepareSpans();
+ assertEquals(2, spans.size());
+ assertSpanHierarchy(SPAN_NAME_PARTITION_QUERY, grpcSpanName(SPAN_NAME_PARTITION_QUERY));
+ }
+
+ @Test
+ public void collectionListDocuments() throws Exception {
+ firestore.collection("col").listDocuments();
+
+ List spans = prepareSpans();
+ assertEquals(2, spans.size());
+ assertSpanHierarchy(SPAN_NAME_COL_REF_LIST_DOCUMENTS, grpcSpanName(LIST_DOCUMENTS_RPC_NAME));
+ }
+
+ @Test
+ public void docRefCreate() throws Exception {
+ firestore.collection("col").document().create(Collections.singletonMap("foo", "bar")).get();
+
+ List spans = prepareSpans();
+ assertEquals(3, spans.size());
+ assertSpanHierarchy(
+ SPAN_NAME_DOC_REF_CREATE, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME));
+ }
+
+ @Test
+ public void docRefCreate2() throws Exception {
+ firestore.collection("col").document().create(new Pojo(1)).get();
+
+ List spans = prepareSpans();
+ assertEquals(3, spans.size());
+ assertSpanHierarchy(
+ SPAN_NAME_DOC_REF_CREATE, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME));
+ }
+
+ @Test
+ public void docRefSet() throws Exception {
+ firestore.collection("col").document("foo").set(Collections.singletonMap("foo", "bar")).get();
+
+ List spans = prepareSpans();
+ assertEquals(3, spans.size());
+ assertSpanHierarchy(
+ SPAN_NAME_DOC_REF_SET, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME));
+ }
+
+ @Test
+ public void docRefSet2() throws Exception {
+ firestore
+ .collection("col")
+ .document("foo")
+ .set(Collections.singletonMap("foo", "bar"), SetOptions.merge())
+ .get();
+
+ List spans = prepareSpans();
+ assertEquals(3, spans.size());
+ assertSpanHierarchy(
+ SPAN_NAME_DOC_REF_SET, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME));
+ }
+
+ @Test
+ public void docRefSet3() throws Exception {
+ firestore.collection("col").document("foo").set(new Pojo(1)).get();
+
+ List spans = prepareSpans();
+ assertEquals(3, spans.size());
+ assertSpanHierarchy(
+ SPAN_NAME_DOC_REF_SET, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME));
+ }
+
+ @Test
+ public void docRefSet4() throws Exception {
+ firestore.collection("col").document("foo").set(new Pojo(1), SetOptions.merge()).get();
+
+ List spans = prepareSpans();
+ assertEquals(3, spans.size());
+ assertSpanHierarchy(
+ SPAN_NAME_DOC_REF_SET, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME));
+ }
+
+ @Test
+ public void docRefUpdate() throws Exception {
+ firestore
+ .collection("col")
+ .document("foo")
+ .update(Collections.singletonMap("foo", "bar"))
+ .get();
+
+ List spans = prepareSpans();
+ assertEquals(3, spans.size());
+ assertSpanHierarchy(
+ SPAN_NAME_DOC_REF_UPDATE, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME));
+ }
+
+ @Test
+ public void docRefUpdate2() throws Exception {
+ firestore
+ .collection("col")
+ .document("foo")
+ .update(Collections.singletonMap("foo", "bar"), Precondition.NONE)
+ .get();
+
+ List spans = prepareSpans();
+ assertEquals(3, spans.size());
+ assertSpanHierarchy(
+ SPAN_NAME_DOC_REF_UPDATE, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME));
+ }
+
+ @Test
+ public void docRefUpdate3() throws Exception {
+ firestore.collection("col").document("foo").update("key", "value", "key2", "value2").get();
+
+ List spans = prepareSpans();
+ assertEquals(3, spans.size());
+ assertSpanHierarchy(
+ SPAN_NAME_DOC_REF_UPDATE, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME));
+ }
+
+ @Test
+ public void docRefUpdate4() throws Exception {
+ firestore
+ .collection("col")
+ .document("foo")
+ .update(FieldPath.of("key"), "value", FieldPath.of("key2"), "value2")
+ .get();
+
+ List spans = prepareSpans();
+ assertEquals(3, spans.size());
+ assertSpanHierarchy(
+ SPAN_NAME_DOC_REF_UPDATE, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME));
+ }
+
+ @Test
+ public void docRefUpdate5() throws Exception {
+ firestore
+ .collection("col")
+ .document("foo")
+ .update(Precondition.NONE, "key", "value", "key2", "value2")
+ .get();
+
+ List spans = prepareSpans();
+ assertEquals(3, spans.size());
+ assertSpanHierarchy(
+ SPAN_NAME_DOC_REF_UPDATE, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME));
+ }
+
+ @Test
+ public void docRefUpdate6() throws Exception {
+ firestore
+ .collection("col")
+ .document("foo")
+ .update(Precondition.NONE, FieldPath.of("key"), "value", FieldPath.of("key2"), "value2")
+ .get();
+
+ List spans = prepareSpans();
+ assertEquals(3, spans.size());
+ assertSpanHierarchy(
+ SPAN_NAME_DOC_REF_UPDATE, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME));
+ }
+
+ @Test
+ public void docRefDelete() throws Exception {
+ firestore.collection("col").document("doc0").delete().get();
+
+ List spans = prepareSpans();
+ assertEquals(3, spans.size());
+ assertSpanHierarchy(
+ SPAN_NAME_DOC_REF_DELETE, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME));
+ }
+
+ @Test
+ public void docRefDelete2() throws Exception {
+ firestore.collection("col").document("doc0").delete(Precondition.NONE).get();
+
+ List spans = prepareSpans();
+ assertEquals(3, spans.size());
+ assertSpanHierarchy(
+ SPAN_NAME_DOC_REF_DELETE, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME));
+ }
+
+ @Test
+ public void docRefGet() throws Exception {
+ firestore.collection("col").document("doc0").get().get();
+
+ List spans = prepareSpans();
+ assertEquals(2, spans.size());
+ assertSpanHierarchy(SPAN_NAME_DOC_REF_GET, grpcSpanName(BATCH_GET_DOCUMENTS_RPC_NAME));
+ }
+
+ @Test
+ public void docRefGet2() throws Exception {
+ firestore.collection("col").document("doc0").get(FieldMask.of("foo")).get();
+
+ List spans = prepareSpans();
+ assertEquals(2, spans.size());
+ assertSpanHierarchy(SPAN_NAME_DOC_REF_GET, grpcSpanName(BATCH_GET_DOCUMENTS_RPC_NAME));
+ }
+
+ @Test
+ public void docListCollections() throws Exception {
+ firestore.collection("col").document("doc0").listCollections();
+
+ List spans = prepareSpans();
+ assertEquals(2, spans.size());
+ assertSpanHierarchy(
+ SPAN_NAME_DOC_REF_LIST_COLLECTIONS, grpcSpanName(LIST_COLLECTIONS_RPC_NAME));
+ }
+
+ @Test
+ public void getAll() throws Exception {
+ DocumentReference docRef0 = firestore.collection("col").document();
+ DocumentReference docRef1 = firestore.collection("col").document();
+ DocumentReference[] docs = {docRef0, docRef1};
+ firestore.getAll(docs).get();
+ List