refactor to move attribute to HttpBigQueryRpc to easily handle reset the retry count in case of error

Author: ldetmer

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

@@ -19,7 +19,6 @@
 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;
@@ -43,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 =
@@ -111,12 +108,6 @@ 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

@@ -1172,7 +1172,7 @@ public void testResendCountOnRetry() throws Exception {
             rpcSpan
                 .getAttributes()
                 .get(
-                    com.google.cloud.bigquery.telemetry.HttpTracingRequestInitializer
+                    BigQueryTelemetryTracer
                         .HTTP_REQUEST_RESEND_COUNT));
       } finally {
         com.google.cloud.bigquery.BigQueryRetryAlgorithm.setCurrentAttempt(0);
@@ -1312,7 +1312,7 @@ public void testResendCountNotSetWhenDisabled() throws Exception {
             rpcSpan
                 .getAttributes()
                 .get(
-                    com.google.cloud.bigquery.telemetry.HttpTracingRequestInitializer
+                        BigQueryTelemetryTracer
                         .HTTP_REQUEST_RESEND_COUNT));
       } finally {
         com.google.cloud.bigquery.BigQueryRetryAlgorithm.setCurrentAttempt(0);

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

@@ -335,29 +335,6 @@ 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);
-    try {
-      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));
-    } finally {
-      com.google.cloud.bigquery.BigQueryRetryAlgorithm.setCurrentAttempt(0);
-    }
-  }
-
   private static HttpTransport createTransport() {
     return createTransport(200, null);
   }
@@ -452,6 +429,5 @@ private void closeAndVerifySpanData(
     } else {
       assertNull(span.getAttributes().get(HttpTracingRequestInitializer.HTTP_RESPONSE_BODY_SIZE));
     }
-    assertNull(span.getAttributes().get(HttpTracingRequestInitializer.HTTP_REQUEST_RESEND_COUNT));
   }
 }