fix(gax): manage OpenTelemetry scope lifecycle internally in OpenTelemetryTracingTracer

Author: jinseopkim0

sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/tracing/OpenTelemetryTracingTracer.java

@@ -53,6 +53,7 @@ class OpenTelemetryTracingTracer implements ApiTracer {
   private final String attemptSpanName;
   private final ApiTracerContext apiTracerContext;
   private @Nullable Span attemptSpan;
+  private io.opentelemetry.context.@Nullable Scope attemptScope;
 
   @Override
   public void injectTraceContext(java.util.Map<String, String> carrier) {
@@ -71,18 +72,6 @@ public void injectTraceContext(java.util.Map<String, String> carrier) {
     }
   }
 
-  @Override
-  public Scope inScope() {
-    if (attemptSpan == null) {
-      return () -> {};
-    }
-    // Suppressing to make MustBeClosedChecker happy. The scope will be closed indirectly
-    // via the returned ApiTracer.Scope, and the checker cannot recognize it.
-    @SuppressWarnings("MustBeClosedChecker")
-    io.opentelemetry.context.Scope otelScope = attemptSpan.makeCurrent();
-    return otelScope::close;
-  }
-
   /**
    * Creates a new instance of {@code OpenTelemetryTracingTracer}.
    *
@@ -136,6 +125,24 @@ private void buildAttributes() {
 
   @Override
   public void attemptStarted(Object request, int attemptNumber) {
+    Span oldSpan = null;
+    io.opentelemetry.context.Scope oldScope = null;
+    if (attemptSpan != null) {
+      oldSpan = attemptSpan;
+      attemptSpan = null;
+      oldScope = attemptScope;
+      attemptScope = null;
+    }
+    if (oldScope != null) {
+      try {
+        oldScope.close();
+      } catch (Exception ignored) {
+      }
+    }
+    if (oldSpan != null) {
+      oldSpan.end();
+    }
+
     Map<String, Object> currentAttemptAttributes = new HashMap<>(this.attemptAttributes);
 
     if (attemptNumber > 0) {
@@ -158,6 +165,10 @@ public void attemptStarted(Object request, int attemptNumber) {
     spanBuilder.setAllAttributes(ObservabilityUtils.toOtelAttributes(currentAttemptAttributes));
 
     this.attemptSpan = spanBuilder.startSpan();
+    // Suppressing to make MustBeClosedChecker happy. The scope will be closed in endAttempt().
+    @SuppressWarnings("MustBeClosedChecker")
+    io.opentelemetry.context.Scope localScope = this.attemptSpan.makeCurrent();
+    this.attemptScope = localScope;
   }
 
   @Override
@@ -250,10 +261,34 @@ private void endAttempt() {
       return;
     }
 
+    if (attemptScope != null) {
+      try {
+        attemptScope.close();
+      } catch (Exception ignored) {
+      } finally {
+        attemptScope = null;
+      }
+    }
+
     attemptSpan.end();
     attemptSpan = null;
   }
 
+  @Override
+  public void operationSucceeded() {
+    recordErrorAndEndAttempt(null);
+  }
+
+  @Override
+  public void operationCancelled() {
+    recordErrorAndEndAttempt(new CancellationException());
+  }
+
+  @Override
+  public void operationFailed(Throwable error) {
+    recordErrorAndEndAttempt(error);
+  }
+
   @Override
   public void requestUrlResolved(String url) {
     if (attemptSpan == null) {

sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/tracing/OpenTelemetryTracingTracerTest.java

@@ -687,21 +687,56 @@ void testInjectTraceContext_addsHeaders() {
   }
 
   @Test
-  void testInScope_whenAttemptSpanActive_makesSpanCurrentAndClosesScope() {
+  void testAttemptStarted_makesSpanCurrent() {
     openTelemetryTracingTracer.attemptStarted(new Object(), 1);
-    ApiTracer.Scope tracerScope = openTelemetryTracingTracer.inScope();
     verify(span).makeCurrent();
+  }
+
+  @Test
+  void testAttemptEnded_closesScope_succeeded() {
+    openTelemetryTracingTracer.attemptStarted(new Object(), 1);
+    openTelemetryTracingTracer.attemptSucceeded();
+    verify(scope).close();
+    verify(span).end();
+  }
 
-    tracerScope.close();
+  @Test
+  void testAttemptEnded_closesScope_failed() {
+    openTelemetryTracingTracer.attemptStarted(new Object(), 1);
+    openTelemetryTracingTracer.attemptFailedRetriesExhausted(new RuntimeException());
     verify(scope).close();
+    verify(span).end();
   }
 
   @Test
-  void testInScope_whenAttemptSpanNull_returnsNoopScope() {
-    ApiTracer.Scope tracerScope = openTelemetryTracingTracer.inScope();
-    verify(span, never()).makeCurrent();
+  void testAttemptEnded_closesScope_cancelled() {
+    openTelemetryTracingTracer.attemptStarted(new Object(), 1);
+    openTelemetryTracingTracer.attemptCancelled();
+    verify(scope).close();
+    verify(span).end();
+  }
 
-    tracerScope.close();
-    verify(scope, never()).close();
+  @Test
+  void testOperationSucceeded_closesScope() {
+    openTelemetryTracingTracer.attemptStarted(new Object(), 1);
+    openTelemetryTracingTracer.operationSucceeded();
+    verify(scope).close();
+    verify(span).end();
+  }
+
+  @Test
+  void testOperationCancelled_closesScope() {
+    openTelemetryTracingTracer.attemptStarted(new Object(), 1);
+    openTelemetryTracingTracer.operationCancelled();
+    verify(scope).close();
+    verify(span).end();
+  }
+
+  @Test
+  void testOperationFailed_closesScope() {
+    openTelemetryTracingTracer.attemptStarted(new Object(), 1);
+    openTelemetryTracingTracer.operationFailed(new RuntimeException());
+    verify(scope).close();
+    verify(span).end();
   }
 }