only update if shouldRetry is true, update to use clearCurrentAttempt helper method

Author: ldetmer

java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryRetryAlgorithm.java

@@ -57,6 +57,11 @@ public static void setCurrentAttempt(int attempt) {
     currentAttempt.set(attempt);
   }
 
+  @InternalApi("internal to java-bigquery")
+  public static void clearCurrentAttempt() {
+    currentAttempt.remove();
+  }
+
   public BigQueryRetryAlgorithm(
       ResultRetryAlgorithm<ResponseT> resultAlgorithm,
       TimedRetryAlgorithm timedAlgorithm,
@@ -91,7 +96,9 @@ public boolean shouldRetry(
             && shouldRetryBasedOnTiming(context, nextAttemptSettings);
 
     // Store retry attempt count in thread-local storage for tracing
-    setCurrentAttempt(attemptCount);
+    if (shouldRetry) {
+      setCurrentAttempt(attemptCount);
+    }
 
     if (LOG.isLoggable(Level.FINEST)) {
       LOG.log(

java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpc.java

@@ -2189,7 +2189,7 @@ 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);
+      BigQueryRetryAlgorithm.clearCurrentAttempt();
       span.end();
     }
   }

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

@@ -60,6 +60,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Nested;
 import org.junit.jupiter.api.Test;
@@ -227,6 +228,14 @@ public void setUp() {
       rpc = createRpc(true);
     }
 
+    @AfterEach
+    public void tearDown() {
+      assertEquals(
+          0,
+          com.google.cloud.bigquery.BigQueryRetryAlgorithm.getCurrentAttempt(),
+          "ThreadLocal should be 0 at the end of the request!");
+    }
+
     @Test
     public void testSpanEndOnError() {
       assertThrows(