From e124bfb75ca0b2a5b3aa8c124dd720a98635b433 Mon Sep 17 00:00:00 2001 From: ldetmer Date: Mon, 30 Mar 2026 14:07:29 -0400 Subject: [PATCH 1/4] chore(bigquery): add resend attribute + integration tests --- .../cloud/bigquery/BigQueryRetryHelper.java | 12 +- .../bigquery/telemetry/ErrorTypeUtil.java | 7 +- .../HttpTracingRequestInitializer.java | 12 +- .../bigquery/it/ITOpenTelemetryTest.java | 286 ++++++++++++++++++ .../bigquery/spi/v2/HttpBigQueryRpcTest.java | 7 + .../HttpTracingRequestInitializerTest.java | 28 ++ 6 files changed, 348 insertions(+), 4 deletions(-) create mode 100644 java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITOpenTelemetryTest.java diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryRetryHelper.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryRetryHelper.java index 9c70830465e6..36dc865cf4c3 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryRetryHelper.java +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryRetryHelper.java @@ -27,15 +27,21 @@ import com.google.cloud.RetryHelper; import io.opentelemetry.api.trace.Span; import io.opentelemetry.api.trace.Tracer; +import io.opentelemetry.context.Context; +import io.opentelemetry.context.ContextKey; import io.opentelemetry.context.Scope; import java.io.IOException; import java.util.concurrent.Callable; import java.util.concurrent.ExecutionException; +import java.util.concurrent.atomic.AtomicInteger; import java.util.logging.Level; import java.util.logging.Logger; public class BigQueryRetryHelper extends RetryHelper { + public static final ContextKey RETRY_ATTEMPT_KEY = + ContextKey.named("bq_retry_attempt"); + private static final Logger LOG = Logger.getLogger(BigQueryRetryHelper.class.getName()); public static V runWithRetries( @@ -54,7 +60,11 @@ public static V runWithRetries( .spanBuilder("com.google.cloud.bigquery.BigQueryRetryHelper.runWithRetries") .startSpan(); } - try (Scope runWithRetriesScope = runWithRetries != null ? runWithRetries.makeCurrent() : null) { + Context retryContext = Context.current().with(RETRY_ATTEMPT_KEY, new AtomicInteger(0)); + if (runWithRetries != null) { + retryContext = retryContext.with(runWithRetries); + } + try (Scope runWithRetriesScope = retryContext != null ? retryContext.makeCurrent() : null) { // Suppressing should be ok as a workaraund. Current and only ResultRetryAlgorithm // implementation does not use response at all, so ignoring its type is ok. @SuppressWarnings("unchecked") diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/telemetry/ErrorTypeUtil.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/telemetry/ErrorTypeUtil.java index 68a13ec8aa39..c769fbb19e65 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/telemetry/ErrorTypeUtil.java +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/telemetry/ErrorTypeUtil.java @@ -16,6 +16,7 @@ package com.google.cloud.bigquery.telemetry; import com.google.api.core.BetaApi; +import com.google.common.annotations.VisibleForTesting; /** * Utility class for identifying exception types for telemetry tracking. TODO: this class should get @@ -23,9 +24,11 @@ * https://github.com/googleapis/google-cloud-java/issues/12105 */ @BetaApi -class ErrorTypeUtil { +@VisibleForTesting +public class ErrorTypeUtil { - enum ErrorType { + @VisibleForTesting + public enum ErrorType { CLIENT_TIMEOUT, CLIENT_CONNECTION_ERROR, CLIENT_REQUEST_ERROR, 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 9f2dcd3e4f2e..4251613bde34 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.BigQueryRetryHelper; import com.google.common.annotations.VisibleForTesting; import io.opentelemetry.api.common.AttributeKey; import io.opentelemetry.api.trace.Span; @@ -26,6 +27,7 @@ import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator; import io.opentelemetry.context.Context; import java.io.IOException; +import java.util.concurrent.atomic.AtomicInteger; /** * HttpRequestInitializer that wraps a delegate initializer, intercepts all HTTP requests, adds @@ -50,7 +52,7 @@ public class HttpTracingRequestInitializer implements HttpRequestInitializer { public static final AttributeKey HTTP_RESPONSE_BODY_SIZE = AttributeKey.longKey("http.response.body.size"); - @VisibleForTesting static final String HTTP_RPC_SYSTEM_NAME = "http"; + @VisibleForTesting public static final String HTTP_RPC_SYSTEM_NAME = "http"; private static final java.util.Set REDACTED_QUERY_PARAMETERS = com.google.common.collect.ImmutableSet.of( @@ -84,6 +86,14 @@ public void initialize(HttpRequest request) throws IOException { addInitialHttpAttributesToSpan(span, request); + AtomicInteger attemptTracker = Context.current().get(BigQueryRetryHelper.RETRY_ATTEMPT_KEY); + if (attemptTracker != null) { + int attempt = attemptTracker.getAndIncrement(); + if (attempt > 0) { + span.setAttribute(HTTP_REQUEST_RESEND_COUNT, (long) attempt); + } + } + HttpResponseInterceptor originalInterceptor = request.getResponseInterceptor(); request.setResponseInterceptor( response -> { diff --git a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITOpenTelemetryTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITOpenTelemetryTest.java new file mode 100644 index 000000000000..034fa093047e --- /dev/null +++ b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITOpenTelemetryTest.java @@ -0,0 +1,286 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.bigquery.it; + +import static com.google.cloud.bigquery.telemetry.ErrorTypeUtil.ErrorType; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.Assertions.fail; + +import com.google.api.gax.retrying.RetrySettings; +import com.google.cloud.bigquery.BigQuery; +import com.google.cloud.bigquery.BigQueryException; +import com.google.cloud.bigquery.telemetry.BigQueryTelemetryTracer; +import com.google.cloud.bigquery.telemetry.HttpTracingRequestInitializer; +import com.google.cloud.bigquery.testing.RemoteBigQueryHelper; +import io.opentelemetry.api.OpenTelemetry; +import io.opentelemetry.api.common.AttributeKey; +import io.opentelemetry.api.trace.Tracer; +import io.opentelemetry.sdk.OpenTelemetrySdk; +import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter; +import io.opentelemetry.sdk.trace.SdkTracerProvider; +import io.opentelemetry.sdk.trace.data.SpanData; +import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; +import java.io.IOException; +import java.util.List; +import java.util.Map; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +public class ITOpenTelemetryTest { + + private static RemoteBigQueryHelper bigqueryHelper; + private InMemorySpanExporter memoryExporter; + private Tracer tracer; + + @BeforeAll + public static void setUpClass() throws IOException { + System.setProperty("com.google.cloud.bigquery.http.tracing.dev.enabled", "true"); + bigqueryHelper = RemoteBigQueryHelper.create(); + } + + @BeforeEach + public void setUp() { + memoryExporter = InMemorySpanExporter.create(); + SdkTracerProvider tracerProvider = + SdkTracerProvider.builder() + .addSpanProcessor(SimpleSpanProcessor.create(memoryExporter)) + .build(); + OpenTelemetry openTelemetry = + OpenTelemetrySdk.builder().setTracerProvider(tracerProvider).build(); + tracer = openTelemetry.getTracer("it-otel-test"); + } + + @Test + public void testListDatasetsTraced() { + BigQuery bq = + bigqueryHelper.getOptions().toBuilder() + .setEnableOpenTelemetryTracing(true) + .setOpenTelemetryTracer(tracer) + .build() + .getService(); + + bq.listDatasets(); + + List spans = memoryExporter.getFinishedSpanItems(); + assertNotNull(spans); + assertFalse(spans.isEmpty(), "Expected at least one span collected"); + + boolean foundRpcSpan = false; + for (SpanData span : spans) { + if (span.getName().equals("com.google.cloud.bigquery.BigQueryRpc.listDatasets")) { + foundRpcSpan = true; + Map, Object> attrs = span.getAttributes().asMap(); + checkGeneralAttributes(attrs); + assertEquals("GET", attrs.get(HttpTracingRequestInitializer.HTTP_REQUEST_METHOD)); + assertEquals("DatasetService", attrs.get(AttributeKey.stringKey("bq.rpc.service"))); + assertEquals("ListDatasets", attrs.get(AttributeKey.stringKey("bq.rpc.method"))); + assertEquals("bigquery.googleapis.com", attrs.get(BigQueryTelemetryTracer.SERVER_ADDRESS)); + assertEquals(200L, attrs.get(HttpTracingRequestInitializer.HTTP_RESPONSE_STATUS_CODE)); + assertEquals("bigquery.googleapis.com", attrs.get(BigQueryTelemetryTracer.URL_DOMAIN)); + assertEquals( + "https://bigquery.googleapis.com/bigquery/v2/projects/gcloud-devel/datasets?prettyPrint=false", + attrs.get(HttpTracingRequestInitializer.URL_FULL)); + assertEquals( + "//bigquery.googleapis.com/projects/gcloud-devel/datasets", + attrs.get(BigQueryTelemetryTracer.GCP_RESOURCE_DESTINATION_ID)); + assertEquals( + "projects/{+projectId}/datasets", attrs.get(BigQueryTelemetryTracer.URL_TEMPLATE)); + } + } + assertTrue(foundRpcSpan, "Expected to find BigQueryRpc.listDatasets span"); + } + + @Test + public void testGetDatasetNotFoundTraced() { + BigQuery bq = + bigqueryHelper.getOptions().toBuilder() + .setEnableOpenTelemetryTracing(true) + .setOpenTelemetryTracer(tracer) + .build() + .getService(); + + bq.getDataset("non_existent_dataset"); + + List spans = memoryExporter.getFinishedSpanItems(); + assertNotNull(spans); + assertFalse(spans.isEmpty(), "Expected at least one span collected"); + + boolean foundRpcSpan = false; + for (SpanData span : spans) { + if (span.getName().equals("com.google.cloud.bigquery.BigQueryRpc.getDataset")) { + foundRpcSpan = true; + Map, Object> attrs = span.getAttributes().asMap(); + checkGeneralAttributes(attrs); + assertEquals("GET", attrs.get(HttpTracingRequestInitializer.HTTP_REQUEST_METHOD)); + assertEquals("DatasetService", attrs.get(AttributeKey.stringKey("bq.rpc.service"))); + assertEquals("GetDataset", attrs.get(AttributeKey.stringKey("bq.rpc.method"))); + assertEquals(404L, attrs.get(HttpTracingRequestInitializer.HTTP_RESPONSE_STATUS_CODE)); + assertEquals( + "projects/{+projectId}/datasets/{+datasetId}", + attrs.get(BigQueryTelemetryTracer.URL_TEMPLATE)); + assertEquals( + "https://bigquery.googleapis.com/bigquery/v2/projects/gcloud-devel/datasets/non_existent_dataset?prettyPrint=false", + attrs.get(HttpTracingRequestInitializer.URL_FULL)); + assertEquals("bigquery.googleapis.com", attrs.get(BigQueryTelemetryTracer.SERVER_ADDRESS)); + assertEquals("bigquery.googleapis.com", attrs.get(BigQueryTelemetryTracer.URL_DOMAIN)); + assertEquals( + "//bigquery.googleapis.com/projects/gcloud-devel/datasets/non_existent_dataset", + attrs.get(BigQueryTelemetryTracer.GCP_RESOURCE_DESTINATION_ID)); + + // Error attributes + assertEquals("notFound", attrs.get(BigQueryTelemetryTracer.ERROR_TYPE)); + assertEquals( + "Not found: Dataset gcloud-devel:non_existent_dataset", + attrs.get(BigQueryTelemetryTracer.STATUS_MESSAGE)); + } + } + assertTrue(foundRpcSpan, "Expected to find BigQueryRpc.getDataset span"); + } + + @Test + public void testConnectionErrorRetriesTraced() { + // Pass invalid host to force connection error and retries + BigQuery bq = + bigqueryHelper.getOptions().toBuilder() + .setRetrySettings(RetrySettings.newBuilder().setMaxAttempts(5).build()) + .setEnableOpenTelemetryTracing(true) + .setOpenTelemetryTracer(tracer) + .setHost("https://invalid-host-name-12345.com:8080") + .build() + .getService(); + + try { + bq.listDatasets(); + fail("Expected BigQueryException due to invalid host"); + } catch (BigQueryException e) { + // Expected + } + + List spans = memoryExporter.getFinishedSpanItems(); + assertNotNull(spans); + assertFalse(spans.isEmpty(), "Expected at least one span collected"); + + int rpcSpanCount = 0; + for (SpanData span : spans) { + if (span.getName().equals("com.google.cloud.bigquery.BigQueryRpc.listDatasets")) { + rpcSpanCount++; + Map, Object> attrs = span.getAttributes().asMap(); + checkGeneralAttributes(attrs); + assertEquals( + "https://invalid-host-name-12345.com:8080/bigquery/v2/projects/gcloud-devel/datasets?prettyPrint=false", + (String) attrs.get(HttpTracingRequestInitializer.URL_FULL)); + assertEquals( + "invalid-host-name-12345.com", attrs.get(BigQueryTelemetryTracer.SERVER_ADDRESS)); + assertEquals(8080L, attrs.get(BigQueryTelemetryTracer.SERVER_PORT)); + assertEquals("invalid-host-name-12345.com", attrs.get(BigQueryTelemetryTracer.URL_DOMAIN)); + assertEquals( + "projects/{+projectId}/datasets", attrs.get(BigQueryTelemetryTracer.URL_TEMPLATE)); + assertEquals( + "//bigquery.googleapis.com/projects/gcloud-devel/datasets", + attrs.get(BigQueryTelemetryTracer.GCP_RESOURCE_DESTINATION_ID)); + checkRetryAttribute(span, rpcSpanCount); + + // Error attributes + assertEquals( + "java.net.UnknownHostException", attrs.get(BigQueryTelemetryTracer.EXCEPTION_TYPE)); + assertEquals( + ErrorType.CLIENT_CONNECTION_ERROR.toString(), + attrs.get(BigQueryTelemetryTracer.ERROR_TYPE)); + assertEquals( + "UnknownHostException: invalid-host-name-12345.com", + attrs.get(BigQueryTelemetryTracer.STATUS_MESSAGE)); + } + } + assertEquals(5, rpcSpanCount, "Expected 5 attempts total"); + } + + @Test + public void testSimultaneousCallsDoNotAffectResendCountForEachother() { + BigQuery bq = + bigqueryHelper.getOptions().toBuilder() + .setRetrySettings(RetrySettings.newBuilder().setMaxAttempts(5).build()) + .setEnableOpenTelemetryTracing(true) + .setOpenTelemetryTracer(tracer) + .setHost("https://invalid-host-name-123456.com") + .build() + .getService(); + + try { + bq.listDatasets(); + fail("Expected BigQueryException due to invalid host"); + } catch (BigQueryException e) { + // Expected + } + try { + bq.cancel("test-job-id"); + fail("Expected BigQueryException due to invalid host"); + } catch (BigQueryException e) { + // Expected + } + + List spans = memoryExporter.getFinishedSpanItems(); + assertNotNull(spans); + assertFalse(spans.isEmpty(), "Expected at least one span collected"); + + int listDataSpanCount = 0; + int cancelJobSpanCount = 0; + for (SpanData span : spans) { + if (span.getName().equals("com.google.cloud.bigquery.BigQueryRpc.listDatasets")) { + listDataSpanCount++; + checkRetryAttribute(span, listDataSpanCount); + } else if (span.getName().equals("com.google.cloud.bigquery.BigQueryRpc.cancelJob")) { + cancelJobSpanCount++; + checkRetryAttribute(span, cancelJobSpanCount); + } + } + assertEquals(5, listDataSpanCount, "Expected 5 attempts total for listDatasets call"); + assertEquals(5, cancelJobSpanCount, "Expected 5 attempts total for cancelJob call"); + } + + private static void checkRetryAttribute(SpanData span, int listDataSpanCount) { + Map, Object> attrs = span.getAttributes().asMap(); + Long resendCount = (Long) attrs.get(HttpTracingRequestInitializer.HTTP_REQUEST_RESEND_COUNT); + if (listDataSpanCount == 1) { + assertTrue(resendCount == null || resendCount == 0); + } else { + assertNotNull(resendCount, "Expected resend count for retry attempt " + listDataSpanCount); + assertEquals((long) (listDataSpanCount - 1), resendCount.longValue()); + } + } + + private void checkGeneralAttributes(Map, Object> attrs) { + assertEquals( + HttpTracingRequestInitializer.HTTP_RPC_SYSTEM_NAME, + attrs.get(BigQueryTelemetryTracer.RPC_SYSTEM_NAME)); + assertEquals( + BigQueryTelemetryTracer.BQ_GCP_CLIENT_SERVICE, + attrs.get(BigQueryTelemetryTracer.GCP_CLIENT_SERVICE)); + assertEquals( + BigQueryTelemetryTracer.BQ_GCP_CLIENT_REPO, + attrs.get(BigQueryTelemetryTracer.GCP_CLIENT_REPO)); + assertEquals( + BigQueryTelemetryTracer.BQ_GCP_CLIENT_LANGUAGE, + attrs.get(BigQueryTelemetryTracer.GCP_CLIENT_LANGUAGE)); + assertEquals( + BigQueryTelemetryTracer.BQ_GCP_CLIENT_ARTIFACT, + attrs.get(BigQueryTelemetryTracer.GCP_CLIENT_ARTIFACT)); + assertNotNull(attrs.get(BigQueryTelemetryTracer.GCP_CLIENT_VERSION)); + } +} 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 8e3caee21529..7a98b3b9dd07 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 @@ -47,6 +47,7 @@ import com.google.cloud.NoCredentials; import com.google.cloud.bigquery.BigQueryOptions; import com.google.cloud.bigquery.telemetry.BigQueryTelemetryTracer; +import com.google.cloud.bigquery.telemetry.HttpTracingRequestInitializer; import io.opentelemetry.api.common.AttributeKey; import io.opentelemetry.api.trace.Tracer; import io.opentelemetry.sdk.OpenTelemetrySdk; @@ -160,6 +161,10 @@ 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(HttpTracingRequestInitializer.HTTP_REQUEST_RESEND_COUNT)); } private void verifySpanProductionAttributes( @@ -1245,6 +1250,8 @@ public void testHttpTracingDisabledDoesNotAddAdditionalAttributes() throws Excep "url.template attribute should not be set"); assertNull(rpcSpan.getAttributes().get(BigQueryTelemetryTracer.GCP_RESOURCE_DESTINATION_ID)); assertNull(rpcSpan.getAttributes().get(BigQueryTelemetryTracer.URL_DOMAIN)); + assertNull( + rpcSpan.getAttributes().get(HttpTracingRequestInitializer.HTTP_REQUEST_RESEND_COUNT)); } @Test 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 c218b45b331c..16615e532125 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 @@ -44,9 +44,11 @@ import com.google.api.client.testing.http.MockHttpTransport; import com.google.api.client.testing.http.MockLowLevelHttpRequest; import com.google.api.client.testing.http.MockLowLevelHttpResponse; +import com.google.cloud.bigquery.BigQueryRetryHelper; import io.opentelemetry.api.common.AttributeKey; import io.opentelemetry.api.trace.Span; import io.opentelemetry.api.trace.Tracer; +import io.opentelemetry.context.Context; import io.opentelemetry.context.Scope; import io.opentelemetry.sdk.OpenTelemetrySdk; import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter; @@ -55,6 +57,7 @@ import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; import java.io.IOException; import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -315,6 +318,31 @@ public void testAddRequestBodySizeToSpan_WithEncoding() throws IOException { assertNull(span.getAttributes().get(HttpTracingRequestInitializer.HTTP_REQUEST_BODY_SIZE)); } + @Test + public void testRetryCountFromContext() throws IOException { + HttpTransport transport = createTransport(); + AtomicInteger counter = new AtomicInteger(2); + Context context = + io.opentelemetry.context.Context.current() + .with(BigQueryRetryHelper.RETRY_ATTEMPT_KEY, counter); + + try (io.opentelemetry.context.Scope scope = context.makeCurrent()) { + 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( + 2L, span.getAttributes().get(HttpTracingRequestInitializer.HTTP_REQUEST_RESEND_COUNT)); + assertEquals(3, counter.get()); + } + @Test public void testAddRequestBodySizeToSpan() throws IOException { HttpTransport transport = createTransport(); From 4399a85d246d92f2c2b3a8c95ca55637f1a1dc14 Mon Sep 17 00:00:00 2001 From: ldetmer <1771267+ldetmer@users.noreply.github.com> Date: Mon, 30 Mar 2026 14:12:15 -0400 Subject: [PATCH 2/4] remove redundant null check Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- .../java/com/google/cloud/bigquery/BigQueryRetryHelper.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryRetryHelper.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryRetryHelper.java index 36dc865cf4c3..98adb0b273e1 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryRetryHelper.java +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryRetryHelper.java @@ -64,7 +64,7 @@ public static V runWithRetries( if (runWithRetries != null) { retryContext = retryContext.with(runWithRetries); } - try (Scope runWithRetriesScope = retryContext != null ? retryContext.makeCurrent() : null) { + try (Scope runWithRetriesScope = retryContext.makeCurrent()) { // Suppressing should be ok as a workaraund. Current and only ResultRetryAlgorithm // implementation does not use response at all, so ignoring its type is ok. @SuppressWarnings("unchecked") From a7fded2ad9c91bf5ad116daf5a766c2f912e4d26 Mon Sep 17 00:00:00 2001 From: ldetmer Date: Tue, 31 Mar 2026 11:59:14 -0400 Subject: [PATCH 3/4] use ThreadLocal so that metrics can re-use this solution if span is turned off # Conflicts: # java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryRetryHelper.java --- .../cloud/bigquery/BigQueryRetryHelper.java | 13 +--- .../cloud/bigquery/BigQueryRetryTracker.java | 49 +++++++++++++ .../HttpTracingRequestInitializer.java | 4 +- .../bigquery/BigQueryRetryTrackerTest.java | 70 +++++++++++++++++++ .../HttpTracingRequestInitializerTest.java | 13 ++-- 5 files changed, 127 insertions(+), 22 deletions(-) create mode 100644 java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryRetryTracker.java create mode 100644 java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryRetryTrackerTest.java diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryRetryHelper.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryRetryHelper.java index 98adb0b273e1..7c6380d0b657 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryRetryHelper.java +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryRetryHelper.java @@ -27,21 +27,15 @@ import com.google.cloud.RetryHelper; import io.opentelemetry.api.trace.Span; import io.opentelemetry.api.trace.Tracer; -import io.opentelemetry.context.Context; -import io.opentelemetry.context.ContextKey; import io.opentelemetry.context.Scope; import java.io.IOException; import java.util.concurrent.Callable; import java.util.concurrent.ExecutionException; -import java.util.concurrent.atomic.AtomicInteger; import java.util.logging.Level; import java.util.logging.Logger; public class BigQueryRetryHelper extends RetryHelper { - public static final ContextKey RETRY_ATTEMPT_KEY = - ContextKey.named("bq_retry_attempt"); - private static final Logger LOG = Logger.getLogger(BigQueryRetryHelper.class.getName()); public static V runWithRetries( @@ -60,11 +54,8 @@ public static V runWithRetries( .spanBuilder("com.google.cloud.bigquery.BigQueryRetryHelper.runWithRetries") .startSpan(); } - Context retryContext = Context.current().with(RETRY_ATTEMPT_KEY, new AtomicInteger(0)); - if (runWithRetries != null) { - retryContext = retryContext.with(runWithRetries); - } - try (Scope runWithRetriesScope = retryContext.makeCurrent()) { + try (Scope runWithRetriesScope = runWithRetries != null ? runWithRetries.makeCurrent() : null; + BigQueryRetryTracker tracker = new BigQueryRetryTracker()) { // Suppressing should be ok as a workaraund. Current and only ResultRetryAlgorithm // implementation does not use response at all, so ignoring its type is ok. @SuppressWarnings("unchecked") diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryRetryTracker.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryRetryTracker.java new file mode 100644 index 000000000000..93edadcf3713 --- /dev/null +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryRetryTracker.java @@ -0,0 +1,49 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.bigquery; + +import java.util.concurrent.atomic.AtomicInteger; + +/** + * Tracks retry attempts for a client request using a ThreadLocal for telemetry data. + * + *

It implements {@link AutoCloseable} to ensure that the ThreadLocal is safely restored to its + * previous state (or removed) when the try-with-resources block exits, preventing memory leaks and + * state bleeding in thread-pooling environments. + */ +public class BigQueryRetryTracker implements AutoCloseable { + + private static final ThreadLocal HOLDER = new ThreadLocal<>(); + private final AtomicInteger previous; + + public BigQueryRetryTracker() { + this.previous = HOLDER.get(); + HOLDER.set(new AtomicInteger(0)); + } + + public static AtomicInteger get() { + return HOLDER.get(); + } + + @Override + public void close() { + if (previous == null) { + HOLDER.remove(); + } else { + HOLDER.set(previous); + } + } +} 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 4251613bde34..79fd5ce7e4d1 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,7 +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.BigQueryRetryHelper; +import com.google.cloud.bigquery.BigQueryRetryTracker; import com.google.common.annotations.VisibleForTesting; import io.opentelemetry.api.common.AttributeKey; import io.opentelemetry.api.trace.Span; @@ -86,7 +86,7 @@ public void initialize(HttpRequest request) throws IOException { addInitialHttpAttributesToSpan(span, request); - AtomicInteger attemptTracker = Context.current().get(BigQueryRetryHelper.RETRY_ATTEMPT_KEY); + AtomicInteger attemptTracker = BigQueryRetryTracker.get(); if (attemptTracker != null) { int attempt = attemptTracker.getAndIncrement(); if (attempt > 0) { diff --git a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryRetryTrackerTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryRetryTrackerTest.java new file mode 100644 index 000000000000..8bb5830ab578 --- /dev/null +++ b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryRetryTrackerTest.java @@ -0,0 +1,70 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.bigquery; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; + +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.jupiter.api.Test; + +public class BigQueryRetryTrackerTest { + + @Test + public void testCreateAndClose() { + assertNull(BigQueryRetryTracker.get(), "Tracker should be null before create"); + + try (BigQueryRetryTracker tracker = new BigQueryRetryTracker()) { + AtomicInteger counter = BigQueryRetryTracker.get(); + assertNotNull(counter); + assertEquals(0, counter.get()); + } + + assertNull(BigQueryRetryTracker.get(), "Tracker should be cleaned up after close"); + } + + @Test + public void testThreadIsolation() throws ExecutionException, InterruptedException { + CompletableFuture future1 = + CompletableFuture.runAsync( + () -> { + try (BigQueryRetryTracker tracker = new BigQueryRetryTracker()) { + BigQueryRetryTracker.get().set(10); + try { + Thread.sleep(100); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + assertEquals(10, BigQueryRetryTracker.get().get()); + } + }); + + CompletableFuture future2 = + CompletableFuture.runAsync( + () -> { + try (BigQueryRetryTracker tracker = new BigQueryRetryTracker()) { + BigQueryRetryTracker.get().set(20); + assertEquals(20, BigQueryRetryTracker.get().get()); + } + }); + + future1.get(); + future2.get(); + } +} 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 16615e532125..c334486c252e 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 @@ -44,11 +44,10 @@ import com.google.api.client.testing.http.MockHttpTransport; import com.google.api.client.testing.http.MockLowLevelHttpRequest; import com.google.api.client.testing.http.MockLowLevelHttpResponse; -import com.google.cloud.bigquery.BigQueryRetryHelper; +import com.google.cloud.bigquery.BigQueryRetryTracker; import io.opentelemetry.api.common.AttributeKey; import io.opentelemetry.api.trace.Span; import io.opentelemetry.api.trace.Tracer; -import io.opentelemetry.context.Context; import io.opentelemetry.context.Scope; import io.opentelemetry.sdk.OpenTelemetrySdk; import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter; @@ -57,7 +56,6 @@ import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; import java.io.IOException; import java.util.List; -import java.util.concurrent.atomic.AtomicInteger; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -321,15 +319,13 @@ public void testAddRequestBodySizeToSpan_WithEncoding() throws IOException { @Test public void testRetryCountFromContext() throws IOException { HttpTransport transport = createTransport(); - AtomicInteger counter = new AtomicInteger(2); - Context context = - io.opentelemetry.context.Context.current() - .with(BigQueryRetryHelper.RETRY_ATTEMPT_KEY, counter); - try (io.opentelemetry.context.Scope scope = context.makeCurrent()) { + try (BigQueryRetryTracker tracker = new BigQueryRetryTracker()) { + BigQueryRetryTracker.get().set(2); HttpRequest request = buildGetRequest(transport, initializer, BASE_URL); HttpResponse response = request.execute(); response.disconnect(); + assertEquals(3, BigQueryRetryTracker.get().get()); } spanScope.close(); @@ -340,7 +336,6 @@ public void testRetryCountFromContext() throws IOException { SpanData span = spans.get(0); assertEquals( 2L, span.getAttributes().get(HttpTracingRequestInitializer.HTTP_REQUEST_RESEND_COUNT)); - assertEquals(3, counter.get()); } @Test From 395263fb53104f76584628e10dd19e77ea081d7c Mon Sep 17 00:00:00 2001 From: ldetmer Date: Tue, 31 Mar 2026 13:47:40 -0400 Subject: [PATCH 4/4] Revert "use ThreadLocal so that metrics can re-use this solution if span is turned off" This reverts commit a7fded2ad9c91bf5ad116daf5a766c2f912e4d26. --- .../cloud/bigquery/BigQueryRetryHelper.java | 13 +++- .../cloud/bigquery/BigQueryRetryTracker.java | 49 ------------- .../HttpTracingRequestInitializer.java | 4 +- .../bigquery/BigQueryRetryTrackerTest.java | 70 ------------------- .../HttpTracingRequestInitializerTest.java | 13 ++-- 5 files changed, 22 insertions(+), 127 deletions(-) delete mode 100644 java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryRetryTracker.java delete mode 100644 java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryRetryTrackerTest.java diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryRetryHelper.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryRetryHelper.java index 7c6380d0b657..98adb0b273e1 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryRetryHelper.java +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryRetryHelper.java @@ -27,15 +27,21 @@ import com.google.cloud.RetryHelper; import io.opentelemetry.api.trace.Span; import io.opentelemetry.api.trace.Tracer; +import io.opentelemetry.context.Context; +import io.opentelemetry.context.ContextKey; import io.opentelemetry.context.Scope; import java.io.IOException; import java.util.concurrent.Callable; import java.util.concurrent.ExecutionException; +import java.util.concurrent.atomic.AtomicInteger; import java.util.logging.Level; import java.util.logging.Logger; public class BigQueryRetryHelper extends RetryHelper { + public static final ContextKey RETRY_ATTEMPT_KEY = + ContextKey.named("bq_retry_attempt"); + private static final Logger LOG = Logger.getLogger(BigQueryRetryHelper.class.getName()); public static V runWithRetries( @@ -54,8 +60,11 @@ public static V runWithRetries( .spanBuilder("com.google.cloud.bigquery.BigQueryRetryHelper.runWithRetries") .startSpan(); } - try (Scope runWithRetriesScope = runWithRetries != null ? runWithRetries.makeCurrent() : null; - BigQueryRetryTracker tracker = new BigQueryRetryTracker()) { + Context retryContext = Context.current().with(RETRY_ATTEMPT_KEY, new AtomicInteger(0)); + if (runWithRetries != null) { + retryContext = retryContext.with(runWithRetries); + } + try (Scope runWithRetriesScope = retryContext.makeCurrent()) { // Suppressing should be ok as a workaraund. Current and only ResultRetryAlgorithm // implementation does not use response at all, so ignoring its type is ok. @SuppressWarnings("unchecked") diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryRetryTracker.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryRetryTracker.java deleted file mode 100644 index 93edadcf3713..000000000000 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryRetryTracker.java +++ /dev/null @@ -1,49 +0,0 @@ -/* - * Copyright 2026 Google LLC - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package com.google.cloud.bigquery; - -import java.util.concurrent.atomic.AtomicInteger; - -/** - * Tracks retry attempts for a client request using a ThreadLocal for telemetry data. - * - *

It implements {@link AutoCloseable} to ensure that the ThreadLocal is safely restored to its - * previous state (or removed) when the try-with-resources block exits, preventing memory leaks and - * state bleeding in thread-pooling environments. - */ -public class BigQueryRetryTracker implements AutoCloseable { - - private static final ThreadLocal HOLDER = new ThreadLocal<>(); - private final AtomicInteger previous; - - public BigQueryRetryTracker() { - this.previous = HOLDER.get(); - HOLDER.set(new AtomicInteger(0)); - } - - public static AtomicInteger get() { - return HOLDER.get(); - } - - @Override - public void close() { - if (previous == null) { - HOLDER.remove(); - } else { - HOLDER.set(previous); - } - } -} 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 79fd5ce7e4d1..4251613bde34 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,7 +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.BigQueryRetryTracker; +import com.google.cloud.bigquery.BigQueryRetryHelper; import com.google.common.annotations.VisibleForTesting; import io.opentelemetry.api.common.AttributeKey; import io.opentelemetry.api.trace.Span; @@ -86,7 +86,7 @@ public void initialize(HttpRequest request) throws IOException { addInitialHttpAttributesToSpan(span, request); - AtomicInteger attemptTracker = BigQueryRetryTracker.get(); + AtomicInteger attemptTracker = Context.current().get(BigQueryRetryHelper.RETRY_ATTEMPT_KEY); if (attemptTracker != null) { int attempt = attemptTracker.getAndIncrement(); if (attempt > 0) { diff --git a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryRetryTrackerTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryRetryTrackerTest.java deleted file mode 100644 index 8bb5830ab578..000000000000 --- a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryRetryTrackerTest.java +++ /dev/null @@ -1,70 +0,0 @@ -/* - * Copyright 2026 Google LLC - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package com.google.cloud.bigquery; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertNotNull; -import static org.junit.jupiter.api.Assertions.assertNull; - -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.ExecutionException; -import java.util.concurrent.atomic.AtomicInteger; -import org.junit.jupiter.api.Test; - -public class BigQueryRetryTrackerTest { - - @Test - public void testCreateAndClose() { - assertNull(BigQueryRetryTracker.get(), "Tracker should be null before create"); - - try (BigQueryRetryTracker tracker = new BigQueryRetryTracker()) { - AtomicInteger counter = BigQueryRetryTracker.get(); - assertNotNull(counter); - assertEquals(0, counter.get()); - } - - assertNull(BigQueryRetryTracker.get(), "Tracker should be cleaned up after close"); - } - - @Test - public void testThreadIsolation() throws ExecutionException, InterruptedException { - CompletableFuture future1 = - CompletableFuture.runAsync( - () -> { - try (BigQueryRetryTracker tracker = new BigQueryRetryTracker()) { - BigQueryRetryTracker.get().set(10); - try { - Thread.sleep(100); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - } - assertEquals(10, BigQueryRetryTracker.get().get()); - } - }); - - CompletableFuture future2 = - CompletableFuture.runAsync( - () -> { - try (BigQueryRetryTracker tracker = new BigQueryRetryTracker()) { - BigQueryRetryTracker.get().set(20); - assertEquals(20, BigQueryRetryTracker.get().get()); - } - }); - - future1.get(); - future2.get(); - } -} 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 c334486c252e..16615e532125 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 @@ -44,10 +44,11 @@ import com.google.api.client.testing.http.MockHttpTransport; import com.google.api.client.testing.http.MockLowLevelHttpRequest; import com.google.api.client.testing.http.MockLowLevelHttpResponse; -import com.google.cloud.bigquery.BigQueryRetryTracker; +import com.google.cloud.bigquery.BigQueryRetryHelper; import io.opentelemetry.api.common.AttributeKey; import io.opentelemetry.api.trace.Span; import io.opentelemetry.api.trace.Tracer; +import io.opentelemetry.context.Context; import io.opentelemetry.context.Scope; import io.opentelemetry.sdk.OpenTelemetrySdk; import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter; @@ -56,6 +57,7 @@ import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; import java.io.IOException; import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -319,13 +321,15 @@ public void testAddRequestBodySizeToSpan_WithEncoding() throws IOException { @Test public void testRetryCountFromContext() throws IOException { HttpTransport transport = createTransport(); + AtomicInteger counter = new AtomicInteger(2); + Context context = + io.opentelemetry.context.Context.current() + .with(BigQueryRetryHelper.RETRY_ATTEMPT_KEY, counter); - try (BigQueryRetryTracker tracker = new BigQueryRetryTracker()) { - BigQueryRetryTracker.get().set(2); + try (io.opentelemetry.context.Scope scope = context.makeCurrent()) { HttpRequest request = buildGetRequest(transport, initializer, BASE_URL); HttpResponse response = request.execute(); response.disconnect(); - assertEquals(3, BigQueryRetryTracker.get().get()); } spanScope.close(); @@ -336,6 +340,7 @@ public void testRetryCountFromContext() throws IOException { SpanData span = spans.get(0); assertEquals( 2L, span.getAttributes().get(HttpTracingRequestInitializer.HTTP_REQUEST_RESEND_COUNT)); + assertEquals(3, counter.get()); } @Test