Skip to content
This repository was archived by the owner on May 8, 2026. It is now read-only.

Commit ee18489

Browse files
ehsannasehsannas
committed
feat: tracing for aggregate queries, bulkwriter, partition queries, and listDocuments.
1 parent 77107e2 commit ee18489

6 files changed

Lines changed: 506 additions & 72 deletions

File tree

‎google-cloud-firestore/pom.xml‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -182,12 +182,36 @@
182182
<scope>test</scope>
183183
</dependency>
184184
<!-- OpenTelemetry -->
185+
<dependency>
186+
<groupId>io.opentelemetry</groupId>
187+
<artifactId>opentelemetry-sdk-testing</artifactId>
188+
<version>1.32.0</version>
189+
<scope>test</scope>
190+
</dependency>
191+
<dependency>
192+
<groupId>io.opentelemetry</groupId>
193+
<artifactId>opentelemetry-semconv</artifactId>
194+
<version>1.29.0-alpha</version>
195+
<scope>test</scope>
196+
</dependency>
197+
<dependency>
198+
<groupId>io.opentelemetry</groupId>
199+
<artifactId>opentelemetry-sdk-trace</artifactId>
200+
<version>1.29.0</version>
201+
<scope>test</scope>
202+
</dependency>
185203
<dependency>
186204
<groupId>io.opentelemetry</groupId>
187205
<artifactId>opentelemetry-sdk</artifactId>
188206
<version>1.29.0</version>
189207
<scope>test</scope>
190208
</dependency>
209+
<dependency>
210+
<groupId>io.opentelemetry</groupId>
211+
<artifactId>opentelemetry-sdk-common</artifactId>
212+
<version>1.29.0</version>
213+
<scope>test</scope>
214+
</dependency>
191215
<!-- END OpenTelemetry -->
192216
</dependencies>
193217

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

Lines changed: 62 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@
1616

1717
package com.google.cloud.firestore;
1818

19+
import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_RUN_AGGREGATION_QUERY;
20+
1921
import com.google.api.core.ApiFuture;
2022
import com.google.api.core.InternalExtensionOnly;
2123
import com.google.api.core.SettableApiFuture;
@@ -24,7 +26,10 @@
2426
import com.google.api.gax.rpc.StatusCode;
2527
import com.google.api.gax.rpc.StreamController;
2628
import com.google.cloud.Timestamp;
29+
import com.google.cloud.firestore.telemetry.TraceUtil;
30+
import com.google.cloud.firestore.telemetry.TraceUtil.Scope;
2731
import com.google.cloud.firestore.v1.FirestoreSettings;
32+
import com.google.common.collect.ImmutableMap;
2833
import com.google.firestore.v1.RunAggregationQueryRequest;
2934
import com.google.firestore.v1.RunAggregationQueryResponse;
3035
import com.google.firestore.v1.RunQueryRequest;
@@ -34,6 +39,7 @@
3439
import com.google.firestore.v1.Value;
3540
import com.google.protobuf.ByteString;
3641
import java.util.ArrayList;
42+
import java.util.Collections;
3743
import java.util.HashMap;
3844
import java.util.HashSet;
3945
import java.util.List;
@@ -59,6 +65,11 @@ public class AggregateQuery {
5965
this.aliasMap = new HashMap<>();
6066
}
6167

68+
@Nonnull
69+
private TraceUtil getTraceUtil() {
70+
return query.getFirestore().getOptions().getTraceUtil();
71+
}
72+
6273
/** Returns the query whose aggregations will be calculated by this object. */
6374
@Nonnull
6475
public Query getQuery() {
@@ -77,17 +88,30 @@ public ApiFuture<AggregateQuerySnapshot> get() {
7788

7889
@Nonnull
7990
ApiFuture<AggregateQuerySnapshot> get(@Nullable final ByteString transactionId) {
80-
AggregateQueryResponseDeliverer responseDeliverer =
81-
new AggregateQueryResponseDeliverer(
82-
transactionId, /* startTimeNanos= */ query.rpcContext.getClock().nanoTime());
83-
runQuery(responseDeliverer);
84-
return responseDeliverer.getFuture();
91+
TraceUtil.Span span =
92+
getTraceUtil()
93+
.startSpan(
94+
transactionId == null
95+
? TraceUtil.SPAN_NAME_AGGREGATION_QUERY_GET
96+
: TraceUtil.SPAN_NAME_TRANSACTION_GET_AGGREGATION_QUERY);
97+
try (Scope ignored = span.makeCurrent()) {
98+
AggregateQueryResponseDeliverer responseDeliverer =
99+
new AggregateQueryResponseDeliverer(
100+
transactionId, /* startTimeNanos= */ query.rpcContext.getClock().nanoTime());
101+
runQuery(responseDeliverer, /* attempt= */ 0);
102+
ApiFuture<AggregateQuerySnapshot> result = responseDeliverer.getFuture();
103+
span.endAtFuture(result);
104+
return result;
105+
} catch (Exception error) {
106+
span.end(error);
107+
throw error;
108+
}
85109
}
86110

87-
private void runQuery(AggregateQueryResponseDeliverer responseDeliverer) {
111+
private void runQuery(AggregateQueryResponseDeliverer responseDeliverer, int attempt) {
88112
RunAggregationQueryRequest request = toProto(responseDeliverer.getTransactionId());
89113
AggregateQueryResponseObserver responseObserver =
90-
new AggregateQueryResponseObserver(responseDeliverer);
114+
new AggregateQueryResponseObserver(responseDeliverer, attempt);
91115
ServerStreamingCallable<RunAggregationQueryRequest, RunAggregationQueryResponse> callable =
92116
query.rpcContext.getClient().runAggregationQueryCallable();
93117
query.rpcContext.streamRequest(request, responseObserver, callable);
@@ -138,18 +162,36 @@ private final class AggregateQueryResponseObserver
138162

139163
private final AggregateQueryResponseDeliverer responseDeliverer;
140164
private StreamController streamController;
165+
private int attempt;
141166

142-
AggregateQueryResponseObserver(AggregateQueryResponseDeliverer responseDeliverer) {
167+
AggregateQueryResponseObserver(AggregateQueryResponseDeliverer responseDeliverer, int attempt) {
143168
this.responseDeliverer = responseDeliverer;
169+
this.attempt = attempt;
170+
}
171+
172+
Map<String, Object> getAttemptAttributes() {
173+
ImmutableMap.Builder<String, Object> builder =
174+
new ImmutableMap.Builder<String, Object>().put("isRetryAttempt", attempt > 0);
175+
if (attempt > 0) {
176+
builder.put("attemptNumber", attempt);
177+
}
178+
return builder.build();
144179
}
145180

146181
@Override
147182
public void onStart(StreamController streamController) {
183+
getTraceUtil()
184+
.currentSpan()
185+
.addEvent(SPAN_NAME_RUN_AGGREGATION_QUERY + " Stream started.", getAttemptAttributes());
148186
this.streamController = streamController;
149187
}
150188

151189
@Override
152190
public void onResponse(RunAggregationQueryResponse response) {
191+
getTraceUtil()
192+
.currentSpan()
193+
.addEvent(
194+
SPAN_NAME_RUN_AGGREGATION_QUERY + " Response Received.", getAttemptAttributes());
153195
// Close the stream to avoid it dangling, since we're not expecting any more responses.
154196
streamController.cancel();
155197

@@ -165,8 +207,19 @@ public void onResponse(RunAggregationQueryResponse response) {
165207
@Override
166208
public void onError(Throwable throwable) {
167209
if (shouldRetry(throwable)) {
168-
runQuery(responseDeliverer);
210+
getTraceUtil()
211+
.currentSpan()
212+
.addEvent(
213+
SPAN_NAME_RUN_AGGREGATION_QUERY + ": Retryable Error",
214+
Collections.singletonMap("error.message", throwable.getMessage()));
215+
216+
runQuery(responseDeliverer, attempt + 1);
169217
} else {
218+
getTraceUtil()
219+
.currentSpan()
220+
.addEvent(
221+
SPAN_NAME_RUN_AGGREGATION_QUERY + ": Error",
222+
Collections.singletonMap("error.message", throwable.getMessage()));
170223
responseDeliverer.deliverError(throwable);
171224
}
172225
}

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

Lines changed: 28 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,9 @@
2626
import com.google.api.core.SettableApiFuture;
2727
import com.google.api.gax.rpc.ApiException;
2828
import com.google.api.gax.rpc.StatusCode.Code;
29+
import com.google.cloud.firestore.telemetry.TraceUtil;
30+
import com.google.cloud.firestore.telemetry.TraceUtil.Context;
31+
import com.google.cloud.firestore.telemetry.TraceUtil.Scope;
2932
import com.google.cloud.firestore.v1.FirestoreSettings;
3033
import com.google.common.annotations.VisibleForTesting;
3134
import com.google.common.base.Preconditions;
@@ -221,6 +224,8 @@ enum OperationType {
221224
@GuardedBy("lock")
222225
private Executor errorExecutor;
223226

227+
Context traceContext;
228+
224229
/**
225230
* Used to track when writes are enqueued. The user handler executors cannot be changed after a
226231
* write has been enqueued.
@@ -237,6 +242,7 @@ enum OperationType {
237242
this.successExecutor = MoreExecutors.directExecutor();
238243
this.errorExecutor = MoreExecutors.directExecutor();
239244
this.bulkCommitBatch = new BulkCommitBatch(firestore, bulkWriterExecutor, maxBatchSize);
245+
this.traceContext = firestore.getOptions().getTraceUtil().currentContext();
240246

241247
if (!options.getThrottlingEnabled()) {
242248
this.rateLimiter =
@@ -899,21 +905,32 @@ private void scheduleCurrentBatchLocked(final boolean flush) {
899905

900906
/** Sends the provided batch once the rate limiter does not require any delay. */
901907
private void sendBatchLocked(final BulkCommitBatch batch, final boolean flush) {
902-
// Send the batch if it is does not require any delay, or schedule another attempt after the
908+
// Send the batch if it does not require any delay, or schedule another attempt after the
903909
// appropriate timeout.
904910
boolean underRateLimit = rateLimiter.tryMakeRequest(batch.getMutationsSize());
905911
if (underRateLimit) {
906-
batch
907-
.bulkCommit()
908-
.addListener(
909-
() -> {
910-
if (flush) {
911-
synchronized (lock) {
912-
scheduleCurrentBatchLocked(/* flush= */ true);
913-
}
912+
TraceUtil.Span span =
913+
firestore
914+
.getOptions()
915+
.getTraceUtil()
916+
.startSpan(TraceUtil.SPAN_NAME_BULK_WRITER_COMMIT, traceContext)
917+
.setAttribute("numDocuments", batch.getWrites().size());
918+
try (Scope ignored = span.makeCurrent()) {
919+
ApiFuture<Void> result = batch.bulkCommit();
920+
result.addListener(
921+
() -> {
922+
if (flush) {
923+
synchronized (lock) {
924+
scheduleCurrentBatchLocked(/* flush= */ true);
914925
}
915-
},
916-
bulkWriterExecutor);
926+
}
927+
},
928+
bulkWriterExecutor);
929+
span.endAtFuture(result);
930+
} catch (Exception error) {
931+
span.end(error);
932+
throw error;
933+
}
917934
} else {
918935
long delayMs = rateLimiter.getNextRequestDelayMs(batch.getMutationsSize());
919936
logger.log(Level.FINE, () -> String.format("Backing off for %d seconds", delayMs / 1000));

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

Lines changed: 27 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,8 @@
2121
import com.google.api.gax.rpc.ApiException;
2222
import com.google.api.gax.rpc.ApiExceptions;
2323
import com.google.api.gax.rpc.ApiStreamObserver;
24+
import com.google.cloud.firestore.telemetry.TraceUtil;
25+
import com.google.cloud.firestore.telemetry.TraceUtil.Scope;
2426
import com.google.cloud.firestore.v1.FirestoreClient.PartitionQueryPagedResponse;
2527
import com.google.common.base.Preconditions;
2628
import com.google.common.collect.ImmutableList;
@@ -102,21 +104,32 @@ public ApiFuture<List<QueryPartition>> getPartitions(long desiredPartitionCount)
102104
} else {
103105
PartitionQueryRequest request = buildRequest(desiredPartitionCount);
104106

105-
try {
106-
return ApiFutures.transform(
107-
rpcContext.sendRequest(request, rpcContext.getClient().partitionQueryPagedCallable()),
108-
response -> {
109-
final ImmutableList.Builder<QueryPartition> partitions = ImmutableList.builder();
110-
consumePartitions(
111-
response,
112-
queryPartition -> {
113-
partitions.add(queryPartition);
114-
return null;
115-
});
116-
return partitions.build();
117-
},
118-
MoreExecutors.directExecutor());
107+
TraceUtil.Span span =
108+
rpcContext
109+
.getFirestore()
110+
.getOptions()
111+
.getTraceUtil()
112+
.startSpan(TraceUtil.SPAN_NAME_PARTITION_QUERY);
113+
try (Scope ignored = span.makeCurrent()) {
114+
ApiFuture<List<QueryPartition>> result =
115+
ApiFutures.transform(
116+
rpcContext.sendRequest(
117+
request, rpcContext.getClient().partitionQueryPagedCallable()),
118+
response -> {
119+
final ImmutableList.Builder<QueryPartition> partitions = ImmutableList.builder();
120+
consumePartitions(
121+
response,
122+
queryPartition -> {
123+
partitions.add(queryPartition);
124+
return null;
125+
});
126+
return partitions.build();
127+
},
128+
MoreExecutors.directExecutor());
129+
span.endAtFuture(result);
130+
return result;
119131
} catch (ApiException exception) {
132+
span.end(exception);
120133
throw FirestoreException.forApiException(exception);
121134
}
122135
}

0 commit comments

Comments
 (0)