feat: add resend attribute to span tracing

java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryRetryAlgorithm.java

@@ -18,6 +18,7 @@
 
 import static com.google.common.base.Preconditions.checkNotNull;
 
+import com.google.api.core.InternalApi;
 import com.google.api.gax.retrying.ResultRetryAlgorithm;
 import com.google.api.gax.retrying.ResultRetryAlgorithmWithContext;
 import com.google.api.gax.retrying.RetryAlgorithm;
@@ -44,6 +45,17 @@ public class BigQueryRetryAlgorithm<ResponseT> extends RetryAlgorithm<ResponseT>
 
   private static final Logger LOG = Logger.getLogger(BigQueryRetryAlgorithm.class.getName());
   private static final UUID RETRY_UUID = UUID.randomUUID();
+  private static final ThreadLocal<Integer> currentAttempt = ThreadLocal.withInitial(() -> 0);
+
+  @InternalApi("internal to java-bigquery")
+  public static int getCurrentAttempt() {
+    return currentAttempt.get();
+  }
+
+  @InternalApi("internal to java-bigquery")
+  public static void setCurrentAttempt(int attempt) {
+    currentAttempt.set(attempt);
+  }
 
   public BigQueryRetryAlgorithm(
       ResultRetryAlgorithm<ResponseT> resultAlgorithm,
@@ -78,6 +90,9 @@ public boolean shouldRetry(
                     previousThrowable, bigQueryRetryConfig, previousResponse))
             && shouldRetryBasedOnTiming(context, nextAttemptSettings);
 
+    // Store retry attempt count in thread-local storage for tracing
+    setCurrentAttempt(attemptCount);
+
     if (LOG.isLoggable(Level.FINEST)) {
       LOG.log(
           Level.FINEST,

java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpc.java

@@ -66,6 +66,7 @@
 import com.google.cloud.Tuple;
 import com.google.cloud.bigquery.BigQueryException;
 import com.google.cloud.bigquery.BigQueryOptions;
+import com.google.cloud.bigquery.BigQueryRetryAlgorithm;
 import com.google.cloud.bigquery.telemetry.BigQueryTelemetryTracer;
 import com.google.cloud.bigquery.telemetry.HttpTracingRequestInitializer;
 import com.google.cloud.http.HttpTransportOptions;
@@ -2150,6 +2151,11 @@ private Span createRpcTracingSpan(
           .setAttribute(
               BigQueryTelemetryTracer.GCP_RESOURCE_DESTINATION_ID, gcpResourceDestinationId)
           .setAttribute(BigQueryTelemetryTracer.URL_TEMPLATE, urlTemplate);
+      int retryAttempt = BigQueryRetryAlgorithm.getCurrentAttempt();
+      if (retryAttempt > 0) {
+        builder.setAttribute(
+            BigQueryTelemetryTracer.HTTP_REQUEST_RESEND_COUNT, (long) retryAttempt);
+      }
     }
 
     if (options != null) {
@@ -2182,6 +2188,8 @@ private <T> T executeWithSpan(Span span, SpanOperation<T> operation) throws IOEx
       }
       throw e;
     } finally {
+      // Reset attempt count to 0 to avoid carrying over state across requests on the same thread
+      BigQueryRetryAlgorithm.setCurrentAttempt(0);
       span.end();
     }
   }

java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/telemetry/BigQueryTelemetryTracer.java

@@ -66,6 +66,8 @@ private BigQueryTelemetryTracer() {}
       AttributeKey.stringKey("server.address");
   public static final AttributeKey<Long> SERVER_PORT = AttributeKey.longKey("server.port");
   public static final AttributeKey<String> URL_TEMPLATE = AttributeKey.stringKey("url.template");
+  public static final AttributeKey<Long> HTTP_REQUEST_RESEND_COUNT =
+      AttributeKey.longKey("http.request.resend_count");
 
   public static void addCommonAttributeToSpan(Span span) {
     span.setAttribute(GCP_CLIENT_SERVICE, BQ_GCP_CLIENT_SERVICE)

java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/telemetry/HttpTracingRequestInitializer.java

@@ -42,8 +42,6 @@ public class HttpTracingRequestInitializer implements HttpRequestInitializer {
   public static final AttributeKey<String> URL_DOMAIN = AttributeKey.stringKey("url.domain");
   public static final AttributeKey<Long> HTTP_RESPONSE_STATUS_CODE =
       AttributeKey.longKey("http.response.status_code");
-  public static final AttributeKey<Long> HTTP_REQUEST_RESEND_COUNT =
-      AttributeKey.longKey("http.request.resend_count");
   public static final AttributeKey<Long> HTTP_REQUEST_BODY_SIZE =
       AttributeKey.longKey("http.request.body.size");
   public static final AttributeKey<Long> HTTP_RESPONSE_BODY_SIZE =

java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpcTest.java

@@ -160,6 +160,9 @@ private void verifySpan(
     assertEquals(
         gcpResourceDestinationId,
         rpcSpan.getAttributes().get(BigQueryTelemetryTracer.GCP_RESOURCE_DESTINATION_ID));
+
+    // this attribute should never get set in a normal success flow
+    assertNull(rpcSpan.getAttributes().get(BigQueryTelemetryTracer.HTTP_REQUEST_RESEND_COUNT));
   }
 
   private void verifySpanProductionAttributes(
@@ -1140,6 +1143,39 @@ public void testGetUriTemplateValueTelemetry() throws Exception {
           "projects/{+projectId}/datasets/{+datasetId}",
           rpcSpan.getAttributes().get(BigQueryTelemetryTracer.URL_TEMPLATE));
     }
+
+    @Test
+    public void testResendCountOnRetry() throws Exception {
+      // Manually set attempt count to simulate being inside a retry loop
+      com.google.cloud.bigquery.BigQueryRetryAlgorithm.setCurrentAttempt(2);
+      try {
+        setMockResponse(
+            "{\"kind\":\"bigquery#dataset\",\"id\":\""
+                + PROJECT_ID
+                + ":"
+                + DATASET_ID
+                + "\",\"datasetReference\":{\"projectId\":\""
+                + PROJECT_ID
+                + "\",\"datasetId\":\""
+                + DATASET_ID
+                + "\"}}");
+
+        rpc.getDatasetSkipExceptionTranslation(PROJECT_ID, DATASET_ID, new HashMap<>());
+
+        List<io.opentelemetry.sdk.trace.data.SpanData> spans = spanExporter.getFinishedSpanItems();
+        assertThat(spans).isNotEmpty();
+        io.opentelemetry.sdk.trace.data.SpanData rpcSpan =
+            spans.stream()
+                .filter(s -> s.getName().equals("com.google.cloud.bigquery.BigQueryRpc.getDataset"))
+                .findFirst()
+                .orElse(null);
+        assertNotNull(rpcSpan);
+        assertEquals(
+            2L, rpcSpan.getAttributes().get(BigQueryTelemetryTracer.HTTP_REQUEST_RESEND_COUNT));
+      } finally {
+        com.google.cloud.bigquery.BigQueryRetryAlgorithm.setCurrentAttempt(0);
+      }
+    }
   }
 
   @Nested
@@ -1243,6 +1279,38 @@ public void testHttpTracingDisabled_GoogleJsonResponseException_DoesNotSetAttrib
       assertNull(rpcSpan.getAttributes().get(BigQueryTelemetryTracer.ERROR_TYPE));
       assertNull(rpcSpan.getAttributes().get(BigQueryTelemetryTracer.STATUS_MESSAGE));
     }
+
+    @Test
+    public void testResendCountNotSetWhenDisabled() throws Exception {
+      // Manually set attempt count to simulate being inside a retry loop
+      com.google.cloud.bigquery.BigQueryRetryAlgorithm.setCurrentAttempt(2);
+      try {
+        setMockResponse(
+            "{\"kind\":\"bigquery#dataset\",\"id\":\""
+                + PROJECT_ID
+                + ":"
+                + DATASET_ID
+                + "\",\"datasetReference\":{\"projectId\":\""
+                + PROJECT_ID
+                + "\",\"datasetId\":\""
+                + DATASET_ID
+                + "\"}}");
+
+        rpc.getDatasetSkipExceptionTranslation(PROJECT_ID, DATASET_ID, new HashMap<>());
+
+        List<io.opentelemetry.sdk.trace.data.SpanData> spans = spanExporter.getFinishedSpanItems();
+        assertThat(spans).isNotEmpty();
+        io.opentelemetry.sdk.trace.data.SpanData rpcSpan =
+            spans.stream()
+                .filter(s -> s.getName().equals("com.google.cloud.bigquery.BigQueryRpc.getDataset"))
+                .findFirst()
+                .orElse(null);
+        assertNotNull(rpcSpan);
+        assertNull(rpcSpan.getAttributes().get(BigQueryTelemetryTracer.HTTP_REQUEST_RESEND_COUNT));
+      } finally {
+        com.google.cloud.bigquery.BigQueryRetryAlgorithm.setCurrentAttempt(0);
+      }
+    }
   }
 
   @Nested