Migrate certain instrumentations to use closeActive() to close outstanding spans

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/domain/TestImpl.java

@@ -29,7 +29,6 @@
 import datadog.trace.api.civisibility.telemetry.tag.SkipReason;
 import datadog.trace.api.civisibility.telemetry.tag.TestFrameworkInstrumentation;
 import datadog.trace.api.gateway.RequestContextSlot;
-import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
@@ -230,23 +229,22 @@ public void setSkipReason(String skipReason) {
   public void end(@Nullable Long endTime) {
     closeOutstandingSpans();
 
-    final AgentScope scope = AgentTracer.activeScope();
-    if (scope == null) {
+    final AgentSpan activeSpan = AgentTracer.activeSpan();
+    if (activeSpan == null) {
       throw new IllegalStateException(
-          "No active scope present, it is possible that end() was called multiple times");
+          "No active span present, it is possible that end() was called multiple times");
     }
 
-    AgentSpan scopeSpan = scope.span();
-    if (scopeSpan != span) {
+    if (activeSpan != this.span) {
       throw new IllegalStateException(
-          "Active scope does not correspond to the finished test, "
+          "Active span does not correspond to the finished test, "
               + "it is possible that end() was called multiple times "
               + "or an operation that was started by the test is still in progress; "
-              + "active scope span is: "
-              + scopeSpan
+              + "active span is: "
+              + activeSpan
               + "; "
               + "expected span is: "
-              + span);
+              + this.span);
     }
 
     InstrumentationTestBridge.fireBeforeTestEnd(context);
@@ -263,7 +261,7 @@ public void end(@Nullable Long endTime) {
       }
     }
 
-    scope.close();
+    AgentTracer.closeActive();
 
     onSpanFinish.accept(span);
 
@@ -310,18 +308,17 @@ public void end(@Nullable Long endTime) {
    * spans stack until it encounters a CI Visibility span or the stack is empty.
    */
   private void closeOutstandingSpans() {
-    AgentScope scope;
-    while ((scope = AgentTracer.activeScope()) != null) {
-      AgentSpan span = scope.span();
+    AgentSpan activeSpan;
+    while ((activeSpan = AgentTracer.activeSpan()) != null) {
 
-      if (span == this.span || span.getTag(Tags.TEST_SESSION_ID) != null) {
+      if (activeSpan == this.span || activeSpan.getTag(Tags.TEST_SESSION_ID) != null) {
         // encountered this span or another CI Visibility span (test, suite, module, session)
         break;
       }
 
-      log.debug("Closing outstanding span: {}", span);
-      scope.close();
-      span.finish();
+      log.debug("Closing outstanding span: {}", activeSpan);
+      AgentTracer.closeActive();
+      activeSpan.finish();
     }
   }
 }

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/domain/TestSuiteImpl.java

@@ -11,7 +11,6 @@
 import datadog.trace.api.civisibility.telemetry.CiVisibilityMetricCollector;
 import datadog.trace.api.civisibility.telemetry.tag.EventType;
 import datadog.trace.api.civisibility.telemetry.tag.TestFrameworkInstrumentation;
-import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
@@ -193,26 +192,25 @@ public void setSkipReason(String skipReason) {
   @Override
   public void end(@Nullable Long endTime) {
     if (!parallelized) {
-      final AgentScope scope = AgentTracer.activeScope();
-      if (scope == null) {
+      final AgentSpan activeSpan = AgentTracer.activeSpan();
+      if (activeSpan == null) {
         throw new IllegalStateException(
-            "No active scope present, it is possible that end() was called multiple times");
+            "No active span present, it is possible that end() was called multiple times");
       }
 
-      AgentSpan scopeSpan = scope.span();
-      if (scopeSpan != span) {
+      if (activeSpan != this.span) {
         throw new IllegalStateException(
-            "Active scope does not correspond to the finished suite, "
+            "Active span does not correspond to the finished suite, "
                 + "it is possible that end() was called multiple times "
                 + "or an operation that was started by the suite is still in progress; "
-                + "active scope span is: "
-                + scopeSpan
+                + "active span is: "
+                + activeSpan
                 + "; "
                 + "expected span is: "
-                + span);
+                + this.span);
       }
 
-      scope.close();
+      AgentTracer.closeActive();
     }
 
     onSpanFinish.accept(span);

dd-java-agent/instrumentation/aws-java-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/AWSHttpClientInstrumentation.java

@@ -1,7 +1,8 @@
 package datadog.trace.instrumentation.aws.v0;
 
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
-import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.closeActive;
 import static datadog.trace.instrumentation.aws.v0.OnErrorDecorator.DECORATE;
 import static datadog.trace.instrumentation.aws.v0.OnErrorDecorator.SPAN_CONTEXT_KEY;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
@@ -10,7 +11,6 @@
 import com.amazonaws.Request;
 import com.amazonaws.handlers.RequestHandler2;
 import datadog.trace.agent.tooling.Instrumenter;
-import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import net.bytebuddy.asm.Advice;
 
@@ -45,12 +45,12 @@ public static void methodExit(
         @Advice.Argument(value = 0, optional = true) final Request<?> request,
         @Advice.Thrown final Throwable throwable) {
 
-      final AgentScope scope = activeScope();
+      final AgentSpan activeSpan = activeSpan();
       // check name in case TracingRequestHandler failed to activate the span
-      if (scope != null
-          && (AwsNameCache.spanName(request).equals(scope.span().getSpanName())
-              || !scope.span().isValid())) {
-        scope.close();
+      if (activeSpan != null
+          && (AwsNameCache.spanName(request).equals(activeSpan.getSpanName())
+              || !activeSpan.isValid())) {
+        closeActive();
       }
 
       if (throwable != null && request != null) {

dd-java-agent/instrumentation/aws-java-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/RequestExecutorInstrumentation.java

@@ -1,14 +1,14 @@
 package datadog.trace.instrumentation.aws.v0;
 
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
-import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.closeActive;
 import static datadog.trace.instrumentation.aws.v0.OnErrorDecorator.DECORATE;
 import static datadog.trace.instrumentation.aws.v0.OnErrorDecorator.SPAN_CONTEXT_KEY;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 
 import com.amazonaws.Request;
 import datadog.trace.agent.tooling.Instrumenter;
-import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import net.bytebuddy.asm.Advice;
 
@@ -42,12 +42,12 @@ public static void methodExit(
         @Advice.FieldValue("request") final Request<?> request,
         @Advice.Thrown final Throwable throwable) {
 
-      final AgentScope scope = activeScope();
+      final AgentSpan activeSpan = activeSpan();
       // check name in case TracingRequestHandler failed to activate the span
-      if (scope != null
-          && (AwsNameCache.spanName(request).equals(scope.span().getSpanName())
-              || !scope.span().isValid())) {
-        scope.close();
+      if (activeSpan != null
+          && (AwsNameCache.spanName(request).equals(activeSpan.getSpanName())
+              || !activeSpan.isValid())) {
+        closeActive();
       }
 
       if (throwable != null) {

dd-java-agent/instrumentation/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/AwsHttpClientInstrumentation.java

@@ -5,7 +5,8 @@
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.namedOneOf;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
-import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.closeActive;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopSpan;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
@@ -14,7 +15,9 @@
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
-import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import datadog.trace.context.TraceScope;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
@@ -69,27 +72,29 @@ public static class AwsHttpClientAdvice {
      * stored in channel attributes.
      */
     @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static AgentScope methodEnter(
+    public static TraceScope methodEnter(
         @Advice.This final Object thiz,
         @Advice.Argument(1) final RequestExecutionContext requestExecutionContext) {
-      final AgentScope scope = activeScope();
+      final AgentSpan activeSpan = activeSpan();
       // check name in case TracingExecutionInterceptor failed to activate the span
-      if (scope != null
-          && ((!scope.span().isValid())
+      if (activeSpan != null
+          && ((!activeSpan.isValid())
               || AwsSdkClientDecorator.DECORATE
                   .spanName(requestExecutionContext.executionAttributes())
-                  .equals(scope.span().getSpanName()))) {
+                  .equals(activeSpan.getSpanName()))) {
         if (thiz instanceof MakeAsyncHttpRequestStage) {
-          scope.close(); // close async legacy HTTP span to avoid Netty leak
+          // close async legacy HTTP span to avoid Netty leak...
+          closeActive(); // then drop-through and activate no-op span
         } else {
-          return scope; // keep sync legacy HTTP span alive for duration of call
+          // keep sync legacy HTTP span alive for duration of call
+          return AgentTracer::closeActive;
         }
       }
       return activateSpan(noopSpan());
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
-    public static void methodExit(@Advice.Enter final AgentScope scope) {
+    public static void methodExit(@Advice.Enter final TraceScope scope) {
       scope.close();
     }
 

dd-java-agent/instrumentation/google-http-client/src/main/java/datadog/trace/instrumentation/googlehttpclient/GoogleHttpClientInstrumentation.java

@@ -2,7 +2,7 @@
 
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
-import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
 import static datadog.trace.instrumentation.googlehttpclient.GoogleHttpClientDecorator.DECORATE;
 import static datadog.trace.instrumentation.googlehttpclient.GoogleHttpClientDecorator.HTTP_REQUEST;
@@ -60,17 +60,16 @@ public void methodAdvice(MethodTransformer transformer) {
   public static class GoogleHttpClientAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static AgentScope methodEnter(
-        @Advice.This HttpRequest request, @Advice.Local("inherited") boolean inheritedScope) {
-      AgentScope scope = activeScope();
-      // detect if scope was propagated here by java-concurrent handling
+        @Advice.This HttpRequest request, @Advice.Local("inherited") AgentSpan inheritedSpan) {
+      AgentSpan activeSpan = activeSpan();
+      // detect if span was propagated here by java-concurrent handling
       // of async requests
-      if (null != scope) {
-        AgentSpan span = scope.span();
+      if (null != activeSpan) {
         // reference equality to check this instrumentation created the span,
         // not some other HTTP client
-        if (HTTP_REQUEST == span.getOperationName()) {
-          inheritedScope = true;
-          return scope;
+        if (HTTP_REQUEST == activeSpan.getOperationName()) {
+          inheritedSpan = activeSpan;
+          return null;
         }
       }
       return activateSpan(DECORATE.prepareSpan(startSpan(HTTP_REQUEST), request));
@@ -79,18 +78,18 @@ public static AgentScope methodEnter(
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void methodExit(
         @Advice.Enter AgentScope scope,
-        @Advice.Local("inherited") boolean inheritedScope,
+        @Advice.Local("inherited") AgentSpan inheritedSpan,
         @Advice.Return final HttpResponse response,
         @Advice.Thrown final Throwable throwable) {
       try {
-        AgentSpan span = scope.span();
+        AgentSpan span = scope != null ? scope.span() : inheritedSpan;
         DECORATE.onError(span, throwable);
         DECORATE.onResponse(span, response);
 
         DECORATE.beforeFinish(span);
         span.finish();
       } finally {
-        if (!inheritedScope) {
+        if (scope != null) {
           scope.close();
         }
       }

dd-java-agent/instrumentation/karate/src/main/java/datadog/trace/instrumentation/karate/KarateTracingHook.java

@@ -221,11 +221,7 @@ public void afterStep(StepResult result, ScenarioRuntime sr) {
       return;
     }
 
-    AgentScope scope = AgentTracer.activeScope();
-    if (scope != null) {
-      scope.close();
-    }
-
+    AgentTracer.closeActive();
     span.finish();
   }
 

dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/HttpRequestParserInstrumentation.java

@@ -3,7 +3,8 @@
 import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.extendsClass;
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
-import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.closeActive;
 import static datadog.trace.instrumentation.undertow.UndertowDecorator.DECORATE;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
@@ -71,12 +72,10 @@ public static void afterRequestParse(
       // a span
       // this because undertow will just write down a http 400 raw response over the net channel.
       // Here we try to create a span to record this
-      AgentScope scope = activeScope();
-      AgentSpan span = null;
+      AgentSpan span = activeSpan();
+      AgentScope scope = null;
       try {
-        if (scope != null) {
-          span = scope.span();
-        } else {
+        if (span == null) {
           final AgentSpanContext.Extracted extractedContext = DECORATE.extract(exchange);
           span = DECORATE.startSpan(exchange, extractedContext).setMeasured(true);
           scope = activateSpan(span);
@@ -90,9 +89,11 @@ public static void afterRequestParse(
       } finally {
         if (span != null) {
           span.finish();
-        }
-        if (scope != null) {
-          scope.close();
+          if (scope != null) {
+            scope.close();
+          } else {
+            closeActive();
+          }
         }
       }
     }