append events to the same span.

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

Hide file tree

Showing 2 changed files with 38 additions and 42 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,14 +1729,13 @@ private void internalStream(
       request.setReadTime(readTime.toProto());
     }
 
-    traceUtil
-        .currentSpan()
-        .addEvent(
-            TraceUtil.SPAN_NAME_RUN_QUERY,
-            new ImmutableMap.Builder<String, Object>()
-                .put("isTransactional", transactionId != null)
-                .put("isRetryRequestWithCursor", isRetryRequestWithCursor)
-                .build());
+    TraceUtil.Span span = traceUtil.currentSpan();
+    span.addEvent(
+        TraceUtil.SPAN_NAME_RUN_QUERY,
+        new ImmutableMap.Builder<String, Object>()
+            .put("isTransactional", transactionId != null)
+            .put("isRetryRequestWithCursor", isRetryRequestWithCursor)
+            .build());
 
     final AtomicReference<QueryDocumentSnapshot> lastReceivedDocument = new AtomicReference<>();
 
@@ -1757,21 +1756,16 @@ 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 +1775,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,11 +1784,9 @@ public void onResponse(RunQueryResponse response) {
           public void onError(Throwable throwable) {
             QueryDocumentSnapshot cursor = lastReceivedDocument.get();
             if (shouldRetry(cursor, throwable)) {
-              traceUtil
-                  .currentSpan()
-                  .addEvent(
-                      TraceUtil.SPAN_NAME_RUN_QUERY + ": Retryable Error",
-                      Collections.singletonMap("error.message", throwable.getMessage()));
+              span.addEvent(
+                  TraceUtil.SPAN_NAME_RUN_QUERY + ": Retryable Error",
+                  Collections.singletonMap("error.message", throwable.getMessage()));
 
               Query.this
                   .startAfter(cursor)
@@ -1808,11 +1798,9 @@ public void onError(Throwable throwable) {
                       /* explainOptions= */ explainOptions,
                       /* isRetryRequestWithCursor= */ true);
             } else {
-              traceUtil
-                  .currentSpan()
-                  .addEvent(
-                      TraceUtil.SPAN_NAME_RUN_QUERY + ": Error",
-                      Collections.singletonMap("error.message", throwable.getMessage()));
+              span.addEvent(
+                  TraceUtil.SPAN_NAME_RUN_QUERY + ": Error",
+                  Collections.singletonMap("error.message", throwable.getMessage()));
               runQueryResponseObserver.onError(throwable);
             }
           }
@@ -1821,11 +1809,9 @@ public void onError(Throwable throwable) {
           public void onComplete() {
             if (hasCompleted) return;
             hasCompleted = true;
-            traceUtil
-                .currentSpan()
-                .addEvent(
-                    TraceUtil.SPAN_NAME_RUN_QUERY + ": Completed",
-                    Collections.singletonMap("numDocuments", numDocuments));
+            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
@@ -60,6 +60,7 @@
 import java.util.logging.Level;
 import java.util.logging.Logger;
 import javax.annotation.Nullable;
+
 import org.junit.After;
 import org.junit.Before;
 import org.junit.BeforeClass;
@@ -98,7 +99,8 @@ public class ITTracingTest {
   Map<String, String> spanIdToParentSpanId = new HashMap<>();
   Map<String, SpanData> spanNameToSpanData = new HashMap<>();
 
-  @Rule public TestName testName = new TestName();
+  @Rule
+  public TestName testName = new TestName();
 
   @BeforeClass
   public static void beforeAll() {
@@ -125,10 +127,17 @@ public void before() {
                 FirestoreOpenTelemetryOptions.newBuilder().setTracingEnabled(true).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()));
+      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()));
+      logger.log(
+          Level.INFO,
+          String.format(
+              "Integration test using default database for test %s", testName.getMethodName()));
     }
     firestore = optionsBuilder.build().getService();
 
@@ -286,7 +295,11 @@ boolean hasEvent(SpanData spanData, String eventName, @Nullable Attributes expec
       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));
+    logger.log(
+        Level.INFO,
+        String.format(
+            "Checking if span named '%s' (ID='%s') contains an event named '%s'",
+            spanData.getName(), spanData.getSpanId(), eventName));
 
     List<EventData> events = spanData.getEvents();
     for (EventData event : events) {
@@ -597,10 +610,7 @@ public void queryGet() throws Exception {
                 .put("isTransactional", false)
                 .build()));
     assertTrue(
-        hasEvent(
-            getGrpcSpanByName(RUN_QUERY_RPC_NAME),
-            "RunQuery: Completed",
-            Attributes.builder().put("numDocuments", 0).build()));
+        hasEvent(span, "RunQuery: Completed", Attributes.builder().put("numDocuments", 0).build()));
   }
 
   @Test