diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryRetryAlgorithm.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryRetryAlgorithm.java index 140f2c6eba3c..927f87641a9e 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryRetryAlgorithm.java +++ b/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 extends RetryAlgorithm private static final Logger LOG = Logger.getLogger(BigQueryRetryAlgorithm.class.getName()); private static final UUID RETRY_UUID = UUID.randomUUID(); + private static final ThreadLocal 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 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, diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/telemetry/HttpTracingRequestInitializer.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/telemetry/HttpTracingRequestInitializer.java index 76f664df1069..1809d14b18e9 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/telemetry/HttpTracingRequestInitializer.java +++ b/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( diff --git a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpcTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpcTest.java index bafc92b63e99..01b6141997b7 100644 --- a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpcTest.java +++ b/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 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 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 diff --git a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/telemetry/HttpTracingRequestInitializerTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/telemetry/HttpTracingRequestInitializerTest.java index 0bf0012be30e..9a06326fdb35 100644 --- a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/telemetry/HttpTracingRequestInitializerTest.java +++ b/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 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)); } }