feat: add resend attribute to span tracing

Author: ldetmer

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/telemetry/HttpTracingRequestInitializer.java

@@ -19,6 +19,7 @@
 import com.google.api.client.http.*;
 import com.google.api.core.BetaApi;
 import com.google.api.core.InternalApi;
+import com.google.cloud.bigquery.BigQueryRetryAlgorithm;
 import com.google.common.annotations.VisibleForTesting;
 import io.opentelemetry.api.common.AttributeKey;
 import io.opentelemetry.api.trace.Span;
@@ -111,6 +112,12 @@ private void addInitialHttpAttributesToSpan(Span span, HttpRequest request) {
       span.setAttribute(BigQueryTelemetryTracer.SERVER_PORT, (long) port);
     }
     span.setAttribute(URL_FULL, getSanitizedUrl(request));
+    int retryAttempt = BigQueryRetryAlgorithm.getCurrentAttempt();
+    if (retryAttempt > 0) {
+      span.setAttribute(HTTP_REQUEST_RESEND_COUNT, (long) retryAttempt);
+    }
+    // Reset attempt count to 0 to avoid carrying over state across requests on the same thread
+    BigQueryRetryAlgorithm.setCurrentAttempt(0);
   }
 
   private static void addCommonResponseAttributesToSpan(

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

@@ -1101,6 +1101,43 @@ public void testHttpTracingEnabled_JsonResponseException_SetsAttributes() throws
           "Invalid request", rpcSpan.getAttributes().get(BigQueryTelemetryTracer.STATUS_MESSAGE));
       assertNull(rpcSpan.getAttributes().get(BigQueryTelemetryTracer.EXCEPTION_TYPE));
     }
+
+    @Test
+    public void testResendCountOnRetry() throws Exception {
+      // Manually set attempt count to simulate being inside a retry loop
+      com.google.cloud.bigquery.BigQueryRetryAlgorithm.setCurrentAttempt(2);
+
+      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(
+                  com.google.cloud.bigquery.telemetry.HttpTracingRequestInitializer
+                      .HTTP_REQUEST_RESEND_COUNT));
+
+      com.google.cloud.bigquery.BigQueryRetryAlgorithm.setCurrentAttempt(0);
+    }
   }
 
   @Nested
@@ -1199,6 +1236,42 @@ 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);
+
+      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(
+                  com.google.cloud.bigquery.telemetry.HttpTracingRequestInitializer
+                      .HTTP_REQUEST_RESEND_COUNT));
+
+      com.google.cloud.bigquery.BigQueryRetryAlgorithm.setCurrentAttempt(0);
+    }
   }
 
   @Nested

java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/telemetry/HttpTracingRequestInitializerTest.java

@@ -335,6 +335,27 @@ public void testAddResponseBodySizeToSpan_NullLength() throws IOException {
     assertNull(span.getAttributes().get(HttpTracingRequestInitializer.HTTP_RESPONSE_BODY_SIZE));
   }
 
+  @Test
+  public void testResendCountIsSetFromBigQueryRetryAlgorithm() throws IOException {
+    com.google.cloud.bigquery.BigQueryRetryAlgorithm.setCurrentAttempt(3);
+    HttpTransport transport = createTransport();
+    HttpRequest request = buildGetRequest(transport, initializer, BASE_URL);
+
+    HttpResponse response = request.execute();
+    response.disconnect();
+
+    spanScope.close();
+    parentSpan.end();
+
+    List<SpanData> spans = spanExporter.getFinishedSpanItems();
+    assertEquals(1, spans.size());
+    SpanData span = spans.get(0);
+    assertEquals(
+        3L, span.getAttributes().get(HttpTracingRequestInitializer.HTTP_REQUEST_RESEND_COUNT));
+
+    com.google.cloud.bigquery.BigQueryRetryAlgorithm.setCurrentAttempt(0);
+  }
+
   private static HttpTransport createTransport() {
     return createTransport(200, null);
   }
@@ -429,5 +450,6 @@ private void closeAndVerifySpanData(
     } else {
       assertNull(span.getAttributes().get(HttpTracingRequestInitializer.HTTP_RESPONSE_BODY_SIZE));
     }
+    assertNull(span.getAttributes().get(HttpTracingRequestInitializer.HTTP_REQUEST_RESEND_COUNT));
   }
 }