append events to the same span.

googleapis · May 1, 2024 · 4d3064d · 4d3064d
1 parent 735b3e1
commit 4d3064d
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 27 deletions.
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
@@ -1729,9 +1729,8 @@ private void internalStream(
       request.setReadTime(readTime.toProto());
     }
 
-    traceUtil
-        .currentSpan()
-        .addEvent(
+    TraceUtil.Span span = traceUtil.currentSpan();
+    span.addEvent(
             TraceUtil.SPAN_NAME_RUN_QUERY,
             new ImmutableMap.Builder<String, Object>()
                 .put("isTransactional", transactionId != null)
@@ -1757,21 +1756,15 @@ public void onStart(StreamController streamController) {}
           public void onResponse(RunQueryResponse response) {
             if (!firstResponse) {
               firstResponse = true;
-              traceUtil.currentSpan().addEvent(TraceUtil.SPAN_NAME_RUN_QUERY + ": First Response");
+              span.addEvent(TraceUtil.SPAN_NAME_RUN_QUERY + ": First Response");
             }
 
             runQueryResponseObserver.onNext(response);
 
             if (response.hasDocument()) {
               numDocuments++;
               if (numDocuments % NUM_RESPONSES_PER_TRACE_EVENT == 0) {
-                traceUtil
-                    .currentSpan()
-                    .addEvent(
-                        TraceUtil.SPAN_NAME_RUN_QUERY
-                            + ": Received "
-                            + numDocuments
-                            + " documents");
+                span.addEvent(TraceUtil.SPAN_NAME_RUN_QUERY + ": Received " + numDocuments + " documents");
               }
               Document document = response.getDocument();
               QueryDocumentSnapshot documentSnapshot =
@@ -1781,9 +1774,7 @@ public void onResponse(RunQueryResponse response) {
             }
 
             if (response.getDone()) {
-              traceUtil
-                  .currentSpan()
-                  .addEvent(TraceUtil.SPAN_NAME_RUN_QUERY + ": Received RunQueryResponse.Done");
+              span.addEvent(TraceUtil.SPAN_NAME_RUN_QUERY + ": Received RunQueryResponse.Done");
               onComplete();
             }
           }
@@ -1792,9 +1783,7 @@ public void onResponse(RunQueryResponse response) {
           public void onError(Throwable throwable) {
             QueryDocumentSnapshot cursor = lastReceivedDocument.get();
             if (shouldRetry(cursor, throwable)) {
-              traceUtil
-                  .currentSpan()
-                  .addEvent(
+                  span.addEvent(
                       TraceUtil.SPAN_NAME_RUN_QUERY + ": Retryable Error",
                       Collections.singletonMap("error.message", throwable.getMessage()));
 
@@ -1808,9 +1797,7 @@ public void onError(Throwable throwable) {
                       /* explainOptions= */ explainOptions,
                       /* isRetryRequestWithCursor= */ true);
             } else {
-              traceUtil
-                  .currentSpan()
-                  .addEvent(
+              span.addEvent(
                       TraceUtil.SPAN_NAME_RUN_QUERY + ": Error",
                       Collections.singletonMap("error.message", throwable.getMessage()));
               runQueryResponseObserver.onError(throwable);
@@ -1821,9 +1808,7 @@ public void onError(Throwable throwable) {
           public void onComplete() {
             if (hasCompleted) return;
             hasCompleted = true;
-            traceUtil
-                .currentSpan()
-                .addEvent(
+                span.addEvent(
                     TraceUtil.SPAN_NAME_RUN_QUERY + ": Completed",
                     Collections.singletonMap("numDocuments", numDocuments));
             runQueryResponseObserver.onCompleted();

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
@@ -63,16 +63,13 @@
 import org.junit.After;
 import org.junit.Before;
 import org.junit.BeforeClass;
-import org.junit.FixMethodOrder;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.TestName;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
-import org.junit.runners.MethodSorters;
 
 @RunWith(JUnit4.class)
-@FixMethodOrder(MethodSorters.NAME_ASCENDING)
 public class ITTracingTest {
   private static final Logger logger =
       Logger.getLogger(com.google.cloud.firestore.it.ITTracingTest.class.getName());
@@ -598,7 +595,7 @@ public void queryGet() throws Exception {
                 .build()));
     assertTrue(
         hasEvent(
-            getGrpcSpanByName(RUN_QUERY_RPC_NAME),
+            span,
             "RunQuery: Completed",
             Attributes.builder().put("numDocuments", 0).build()));
   }