From 78b998b1f90958dbfeedbbe272ba9e6482b6805f Mon Sep 17 00:00:00 2001
From: blakeli <blakeli@google.com>
Date: Sat, 4 Apr 2026 00:52:34 -0400
Subject: [PATCH] feat(gax): link span to current context and add unit tests

---
 .../google/api/gax/tracing/SpanTracer.java    |  5 +++
 .../api/gax/tracing/SpanTracerTest.java       | 34 +++++++++++++++++--
 2 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/tracing/SpanTracer.java b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/tracing/SpanTracer.java
index 01fc07531442..82ff10864eba 100644
--- a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/tracing/SpanTracer.java
+++ b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/tracing/SpanTracer.java
@@ -37,6 +37,8 @@
 import io.opentelemetry.api.trace.SpanBuilder;
 import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.api.trace.Tracer;
+import io.opentelemetry.context.Scope;
+
 import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.CancellationException;
@@ -53,6 +55,7 @@ public class SpanTracer implements ApiTracer {
   private final String attemptSpanName;
   private final ApiTracerContext apiTracerContext;
   private Span attemptSpan;
+  private io.opentelemetry.context.Scope scope;
 
   @Override
   public void injectTraceContext(java.util.Map<String, String> carrier) {
@@ -144,6 +147,7 @@ public void attemptStarted(Object request, int attemptNumber) {
     spanBuilder.setAllAttributes(ObservabilityUtils.toOtelAttributes(currentAttemptAttributes));
 
     this.attemptSpan = spanBuilder.startSpan();
+    this.scope = attemptSpan.makeCurrent();
   }
 
   @Override
@@ -249,6 +253,7 @@ private void endAttempt() {
       return;
     }
 
+    scope.close();
     attemptSpan.end();
     attemptSpan = null;
   }
diff --git a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/tracing/SpanTracerTest.java b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/tracing/SpanTracerTest.java
index a14d6a6777e0..4d38489e4027 100644
--- a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/tracing/SpanTracerTest.java
+++ b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/tracing/SpanTracerTest.java
@@ -33,9 +33,7 @@
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
-import static org.mockito.Mockito.never;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
+import static org.mockito.Mockito.*;
 
 import com.google.api.gax.rpc.ApiException;
 import com.google.api.gax.rpc.ErrorDetails;
@@ -48,6 +46,7 @@
 import io.opentelemetry.api.trace.SpanBuilder;
 import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.api.trace.Tracer;
+import io.opentelemetry.context.Scope;
 import java.net.ConnectException;
 import java.net.SocketTimeoutException;
 import org.junit.jupiter.api.BeforeEach;
@@ -62,6 +61,7 @@ class SpanTracerTest {
   @Mock private Tracer tracer;
   @Mock private SpanBuilder spanBuilder;
   @Mock private Span span;
+  @Mock private Scope scope;
   private SpanTracer spanTracer;
   private static final String ATTEMPT_SPAN_NAME = "Service/Method/attempt";
 
@@ -71,6 +71,7 @@ void setUp() {
     when(spanBuilder.setSpanKind(any(SpanKind.class))).thenReturn(spanBuilder);
     when(spanBuilder.setAllAttributes(any(Attributes.class))).thenReturn(spanBuilder);
     when(spanBuilder.startSpan()).thenReturn(span);
+    lenient().when(span.makeCurrent()).thenReturn(scope);
     spanTracer = new SpanTracer(tracer, ApiTracerContext.empty(), ATTEMPT_SPAN_NAME);
   }
 
@@ -625,4 +626,31 @@ void testInjectTraceContext_addsHeaders() {
     assertThat(carrier.get("traceparent")).contains("00000000000000000000000000000001");
     assertThat(carrier.get("traceparent")).contains("0000000000000002");
   }
+
+  @Test
+  void testAttemptStarted_makesSpanCurrent() {
+    spanTracer.attemptStarted(new Object(), 1);
+    verify(span).makeCurrent();
+  }
+
+  @Test
+  void testAttemptEnded_closesScope_succeeded() {
+    spanTracer.attemptStarted(new Object(), 1);
+    spanTracer.attemptSucceeded();
+    verify(scope).close();
+  }
+
+  @Test
+  void testAttemptEnded_closesScope_failed() {
+    spanTracer.attemptStarted(new Object(), 1);
+    spanTracer.attemptFailedRetriesExhausted(new RuntimeException());
+    verify(scope).close();
+  }
+
+  @Test
+  void testAttemptEnded_closesScope_cancelled() {
+    spanTracer.attemptStarted(new Object(), 1);
+    spanTracer.attemptCancelled();
+    verify(scope).close();
+  }
 }