feat: Add 'isTransactional' attribute. (#1657)

googleapis · May 2, 2024 · 8981c00 · 8981c00
1 parent b24e679
commit 8981c00
Show file tree

Hide file tree

Showing 7 changed files with 43 additions and 5 deletions.
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
@@ -33,6 +33,7 @@
 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;
 import com.google.firestore.v1.BatchGetDocumentsRequest;
 import com.google.firestore.v1.BatchGetDocumentsResponse;
 import com.google.firestore.v1.DatabaseRootName;
@@ -235,7 +236,10 @@ public void onStart(StreamController streamController) {
                 .currentSpan()
                 .addEvent(
                     TraceUtil.SPAN_NAME_BATCH_GET_DOCUMENTS + ": Start",
-                    Collections.singletonMap("numDocuments", documentReferences.length));
+                    new ImmutableMap.Builder<String, Object>()
+                        .put("numDocuments", documentReferences.length)
+                        .put("isTransactional", transactionId != null)
+                        .build());
           }
 
           @Override

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
@@ -620,6 +620,8 @@ ApiFuture<List<WriteResult>> commit(@Nullable ByteString transactionId) {
                     ? TraceUtil.SPAN_NAME_BATCH_COMMIT
                     : TraceUtil.SPAN_NAME_TRANSACTION_COMMIT);
     span.setAttribute("numDocuments", writes.size());
+    span.setAttribute("isTransactional", transactionId != null);
+
     try (Scope ignored = span.makeCurrent()) {
       ApiFuture<CommitResponse> response =
           firestore.sendRequest(request.build(), firestore.getClient().commitCallable());

diff --git a/...cloud-firestore/src/main/java/com/google/cloud/firestore/telemetry/DisabledTraceUtil.java b/...cloud-firestore/src/main/java/com/google/cloud/firestore/telemetry/DisabledTraceUtil.java
@@ -55,6 +55,11 @@ 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();

diff --git a/...-cloud-firestore/src/main/java/com/google/cloud/firestore/telemetry/EnabledTraceUtil.java b/...-cloud-firestore/src/main/java/com/google/cloud/firestore/telemetry/EnabledTraceUtil.java
@@ -180,6 +180,12 @@ public TraceUtil.Span setAttribute(String key, String 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());

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
@@ -102,6 +102,9 @@ interface Span {
     /** 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();
 

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
@@ -189,7 +189,7 @@ private boolean dfsContainsCallStack(long spanId, List<String> expectedCallStack
       if (spanName(spanId).equals(expectedCallStack.get(0))) {
         // Recursion termination
         if (childSpans(spanId) == null) {
-          logger.info("No more chilren for " + spanName(spanId));
+          logger.info("No more children for " + spanName(spanId));
           return true;
         } else {
           // Examine the child spans
@@ -241,7 +241,7 @@ private boolean dfsContainsCallStack(long spanId, List<String> expectedCallStack
 
   private static final int NUM_SPAN_ID_BYTES = 16;
 
-  private static final int GET_TRACE_RETRY_COUNT = 15;
+  private static final int GET_TRACE_RETRY_COUNT = 60;
 
   private static final int GET_TRACE_RETRY_BACKOFF_MILLIS = 1000;
 
@@ -538,15 +538,24 @@ public void traceContainerTest() throws Exception {
       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

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
@@ -580,10 +580,13 @@ public void queryGet() throws Exception {
             "RunQuery",
             Attributes.builder()
                 .put("isRetryRequestWithCursor", false)
-                .put("transactional", false)
+                .put("isTransactional", false)
                 .build()));
     assertTrue(
-        hasEvent(span, "RunQuery: Completed", Attributes.builder().put("numDocuments", 0).build()));
+        hasEvent(
+            getGrpcSpanByName(RUN_QUERY_RPC_NAME),
+            "RunQuery: Completed",
+            Attributes.builder().put("numDocuments", 0).build()));
   }
 
   @Test
@@ -733,6 +736,12 @@ public void writeBatch() throws Exception {
     List<SpanData> spans = prepareSpans();
     assertEquals(2, spans.size());
     assertSpanHierarchy(SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME));
+    assertEquals(
+        false,
+        getSpanByName(SPAN_NAME_BATCH_COMMIT)
+            .getAttributes()
+            .get(AttributeKey.booleanKey("gcp.firestore.isTransactional"))
+            .booleanValue());
     assertEquals(
         3L,
         getSpanByName(SPAN_NAME_BATCH_COMMIT)