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/spi/v2/HttpBigQueryRpc.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpc.java index 8eede5a0937b..575241da8524 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpc.java +++ b/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 executeWithSpan(Span span, SpanOperation 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(); } } diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/telemetry/BigQueryTelemetryTracer.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/telemetry/BigQueryTelemetryTracer.java index 82c354ef5446..e368867c052b 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/telemetry/BigQueryTelemetryTracer.java +++ b/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 SERVER_PORT = AttributeKey.longKey("server.port"); public static final AttributeKey URL_TEMPLATE = AttributeKey.stringKey("url.template"); + public static final AttributeKey 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) 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 d34409d44028..3c9302a3e3a3 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 @@ -42,8 +42,6 @@ public class HttpTracingRequestInitializer implements HttpRequestInitializer { public static final AttributeKey URL_DOMAIN = AttributeKey.stringKey("url.domain"); public static final AttributeKey HTTP_RESPONSE_STATUS_CODE = AttributeKey.longKey("http.response.status_code"); - public static final AttributeKey HTTP_REQUEST_RESEND_COUNT = - AttributeKey.longKey("http.request.resend_count"); public static final AttributeKey HTTP_REQUEST_BODY_SIZE = AttributeKey.longKey("http.request.body.size"); public static final AttributeKey HTTP_RESPONSE_BODY_SIZE = 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 acb33d2a1104..5e9c3f23ac8e 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 @@ -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 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 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