Reduce use of activeScope() in instrumentation

dd-java-agent/instrumentation/couchbase/couchbase-2.6/src/main/java/datadog/trace/instrumentation/couchbase/client/CouchbaseCoreInstrumentation.java

@@ -1,7 +1,7 @@
 package datadog.trace.instrumentation.couchbase.client;
 
 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 java.util.Collections.singletonMap;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
@@ -13,7 +13,6 @@
 import datadog.trace.agent.tooling.InstrumenterModule;
 import datadog.trace.bootstrap.ContextStore;
 import datadog.trace.bootstrap.InstrumentationContext;
-import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import java.util.Map;
 import net.bytebuddy.asm.Advice;
@@ -52,18 +51,15 @@ public static class CouchbaseCoreAdvice {
     @Advice.OnMethodExit(suppress = Throwable.class)
     public static void addOperationIdToSpan(@Advice.Argument(0) final CouchbaseRequest request) {
 
-      final AgentScope scope = activeScope();
-      if (scope != null) {
-        // The scope from the initial rxJava subscribe is not available to the networking layer
+      final AgentSpan span = activeSpan();
+      if (span != null) {
+        // The context from the initial rxJava subscribe is not available to the networking layer
         // To transfer the span, the span is added to the context store
 
         final ContextStore<CouchbaseRequest, AgentSpan> contextStore =
             InstrumentationContext.get(CouchbaseRequest.class, AgentSpan.class);
 
-        AgentSpan span = contextStore.get(request);
-
-        if (span == null) {
-          span = scope.span();
+        if (contextStore.get(request) == null) {
           contextStore.put(request, span);
 
           if (request.operationId() != null) {

dd-java-agent/instrumentation/cucumber/src/main/java/datadog/trace/instrumentation/cucumber/CucumberInstrumentation.java

@@ -8,6 +8,7 @@
 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 io.cucumber.core.backend.StepDefinition;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
@@ -46,14 +47,14 @@ public void methodAdvice(MethodTransformer transformer) {
 
   public static class CucumberAdvice {
     @Advice.OnMethodEnter
-    public static void onCucumberStepStart(
+    public static AgentScope onCucumberStepStart(
         @Advice.This StepDefinition step, @Advice.Argument(0) Object[] arguments) {
-      CucumberStepDecorator.DECORATE.onStepStart(step, arguments);
+      return CucumberStepDecorator.DECORATE.onStepStart(step, arguments);
     }
 
     @Advice.OnMethodExit
-    public static void onCucumberStepFinish(@Advice.This StepDefinition step) {
-      CucumberStepDecorator.DECORATE.onStepFinish(step);
+    public static void onCucumberStepFinish(@Advice.Enter AgentScope scope) {
+      CucumberStepDecorator.DECORATE.onStepFinish(scope);
     }
 
     // Cucumber 5.0.0 and above

dd-java-agent/instrumentation/cucumber/src/main/java/datadog/trace/instrumentation/cucumber/CucumberStepDecorator.java

@@ -26,9 +26,8 @@ protected CharSequence component() {
     return "cucumber";
   }
 
-  public void onStepStart(StepDefinition step, Object[] arguments) {
+  public AgentScope onStepStart(StepDefinition step, Object[] arguments) {
     AgentSpan span = AgentTracer.startSpan("cucumber", "cucumber.step");
-    AgentScope scope = AgentTracer.activateSpan(span);
     afterStart(span);
 
     span.setResourceName(step.getPattern());
@@ -38,20 +37,14 @@ public void onStepStart(StepDefinition step, Object[] arguments) {
     if (arguments != null && arguments.length > 0) {
       span.setTag("step.arguments", Arrays.toString(arguments));
     }
-  }
-
-  public void onStepFinish(StepDefinition step) {
-    AgentSpan span = AgentTracer.activeSpan();
-    if (span == null) {
-      return;
-    }
 
-    AgentScope scope = AgentTracer.activeScope();
-    if (scope != null) {
-      scope.close();
-    }
+    return AgentTracer.activateSpan(span);
+  }
 
+  public void onStepFinish(AgentScope scope) {
+    AgentSpan span = scope.span();
     beforeFinish(span);
     span.finish();
+    scope.close();
   }
 }

dd-java-agent/instrumentation/enable-wallclock-profiling/src/main/java/datadog/trace/instrumentation/wallclock/EnableWallclockProfilingInstrumentation.java

@@ -14,7 +14,7 @@
 import datadog.trace.agent.tooling.InstrumenterModule;
 import datadog.trace.api.InstrumenterConfig;
 import datadog.trace.bootstrap.config.provider.ConfigProvider;
-import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import java.util.Arrays;
 import net.bytebuddy.asm.Advice;
@@ -86,8 +86,8 @@ public static final class EnableWallclockSampling {
 
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static boolean before() {
-      AgentScope active = AgentTracer.activeScope();
-      if (active == null) {
+      AgentSpan span = AgentTracer.activeSpan();
+      if (span == null) {
         AgentTracer.get().getProfilingContext().onAttach();
         return true;
       }

dd-java-agent/instrumentation/grpc-1.5/src/main/java/datadog/trace/instrumentation/grpc/client/ClientStreamListenerImplInstrumentation.java

@@ -3,7 +3,7 @@
 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.instrumentation.grpc.client.GrpcClientDecorator.DECORATE;
 import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
@@ -66,9 +66,8 @@ public static final class Construct {
     @Advice.OnMethodExit
     public static void capture(@Advice.This ClientStreamListener listener) {
       // instrumentation of ClientCallImpl::start ensures this scope is present and valid
-      AgentScope scope = activeScope();
-      if (null != scope) {
-        AgentSpan span = scope.span();
+      AgentSpan span = activeSpan();
+      if (null != span) {
         InstrumentationContext.get(ClientStreamListener.class, AgentSpan.class).put(listener, span);
       }
     }

dd-java-agent/instrumentation/http-url-connection/src/test/groovy/HttpUrlConnectionConnectFirstTest.groovy

@@ -1,7 +1,7 @@
 import datadog.trace.agent.test.naming.TestingGenericHttpNamingConventions
 import spock.lang.Timeout
 
-import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan
 
 @Timeout(5)
 class HttpUrlConnectionConnectFirstTest extends HttpUrlConnectionTest implements TestingGenericHttpNamingConventions.ClientV0{
@@ -15,7 +15,7 @@ class HttpUrlConnectionConnectFirstTest extends HttpUrlConnectionTest implements
       connection.setRequestProperty("Connection", "close")
       connection.connectTimeout = CONNECT_TIMEOUT_MS
       connection.readTimeout = READ_TIMEOUT_MS
-      def parentSpan = activeScope()
+      def parentSpan = activeSpan()
       connection.connect() // test connect before getting stream
       def stream
       try {
@@ -24,7 +24,7 @@ class HttpUrlConnectionConnectFirstTest extends HttpUrlConnectionTest implements
         stream = connection.errorStream
         ex.printStackTrace()
       }
-      assert activeScope() == parentSpan
+      assert activeSpan() == parentSpan
       stream?.readLines()
       stream?.close()
       callback?.call()

dd-java-agent/instrumentation/http-url-connection/src/test/groovy/HttpUrlConnectionTest.groovy

@@ -10,7 +10,7 @@ import sun.net.www.protocol.https.HttpsURLConnectionImpl
 
 import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
 import static datadog.trace.api.config.TraceInstrumentationConfig.HTTP_CLIENT_HOST_SPLIT_BY_DOMAIN
-import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan
 
 @Timeout(5)
 abstract class HttpUrlConnectionTest extends HttpClientTest {
@@ -28,15 +28,15 @@ abstract class HttpUrlConnectionTest extends HttpClientTest {
       connection.useCaches = true
       connection.connectTimeout = CONNECT_TIMEOUT_MS
       connection.readTimeout = READ_TIMEOUT_MS
-      def parentSpan = activeScope()
+      def parentSpan = activeSpan()
       def stream
       try {
         stream = connection.inputStream
       } catch (Exception ex) {
         stream = connection.errorStream
         ex.printStackTrace()
       }
-      assert activeScope() == parentSpan
+      assert activeSpan() == parentSpan
       stream?.readLines()
       stream?.close()
       callback?.call()
@@ -66,7 +66,7 @@ abstract class HttpUrlConnectionTest extends HttpClientTest {
     runUnderTrace("someTrace") {
       HttpURLConnection connection = url.openConnection()
       connection.useCaches = useCaches
-      assert activeScope() != null
+      assert activeSpan() != null
       def stream = connection.inputStream
       def lines = stream.readLines()
       stream.close()
@@ -76,7 +76,7 @@ abstract class HttpUrlConnectionTest extends HttpClientTest {
       // call again to ensure the cycling is ok
       connection = url.openConnection()
       connection.useCaches = useCaches
-      assert activeScope() != null
+      assert activeSpan() != null
       assert connection.getResponseCode() == STATUS // call before input stream to test alternate behavior
       connection.inputStream
       stream = connection.inputStream // one more to ensure state is working
@@ -157,7 +157,7 @@ abstract class HttpUrlConnectionTest extends HttpClientTest {
       HttpURLConnection connection = url.openConnection()
       connection.useCaches = useCaches
       connection.addRequestProperty("is-dd-server", "false")
-      assert activeScope() != null
+      assert activeSpan() != null
       def stream = connection.inputStream
       connection.inputStream // one more to ensure state is working
       def lines = stream.readLines()
@@ -169,7 +169,7 @@ abstract class HttpUrlConnectionTest extends HttpClientTest {
       connection = url.openConnection()
       connection.useCaches = useCaches
       connection.addRequestProperty("is-dd-server", "false")
-      assert activeScope() != null
+      assert activeSpan() != null
       assert connection.getResponseCode() == STATUS // call before input stream to test alternate behavior
       stream = connection.inputStream
       lines = stream.readLines()
@@ -247,7 +247,7 @@ abstract class HttpUrlConnectionTest extends HttpClientTest {
       HttpURLConnection connection = url.openConnection()
       connection.setRequestProperty("Connection", "close")
       connection.addRequestProperty("is-dd-server", "false")
-      assert activeScope() != null
+      assert activeSpan() != null
       assert connection.getResponseCode() == STATUS
       return connection
     }

dd-java-agent/instrumentation/http-url-connection/src/test/groovy/HttpUrlConnectionUseCachesFalseTest.groovy

@@ -1,7 +1,7 @@
 import datadog.trace.agent.test.naming.TestingGenericHttpNamingConventions
 import spock.lang.Timeout
 
-import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan
 
 @Timeout(5)
 class HttpUrlConnectionUseCachesFalseTest extends HttpUrlConnectionTest implements TestingGenericHttpNamingConventions.ClientV0 {
@@ -16,15 +16,15 @@ class HttpUrlConnectionUseCachesFalseTest extends HttpUrlConnectionTest implemen
       connection.useCaches = false
       connection.connectTimeout = CONNECT_TIMEOUT_MS
       connection.readTimeout = READ_TIMEOUT_MS
-      def parentSpan = activeScope()
+      def parentSpan = activeSpan()
       def stream
       try {
         stream = connection.inputStream
       } catch (Exception ex) {
         stream = connection.errorStream
         ex.printStackTrace()
       }
-      assert activeScope() == parentSpan
+      assert activeSpan() == parentSpan
       stream?.readLines()
       stream?.close()
       callback?.call()

dd-java-agent/instrumentation/http-url-connection/src/test/groovy/UrlConnectionTest.groovy

@@ -6,7 +6,7 @@ import datadog.trace.bootstrap.instrumentation.api.Tags
 import static datadog.trace.agent.test.utils.PortUtils.UNUSABLE_PORT
 import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
 import static datadog.trace.api.config.TraceInstrumentationConfig.HTTP_CLIENT_HOST_SPLIT_BY_DOMAIN
-import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan
 
 abstract class UrlConnectionTest extends VersionedNamingTestBase {
 
@@ -17,7 +17,7 @@ abstract class UrlConnectionTest extends VersionedNamingTestBase {
       URLConnection connection = url.openConnection()
       connection.setConnectTimeout(10000)
       connection.setReadTimeout(10000)
-      assert activeScope() != null
+      assert activeSpan() != null
       connection.inputStream
     }
 

dd-java-agent/instrumentation/java-concurrent/java-completablefuture/src/test/java/Fanout.java

@@ -1,4 +1,4 @@
-import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
 
 import datadog.trace.api.Trace;
 import java.util.concurrent.CompletableFuture;
@@ -52,11 +52,11 @@ public void executeTwoLevels() {
   }
 
   private void untracedWork() {
-    assert null != activeScope();
+    assert null != activeSpan();
   }
 
   @Trace
   private void tracedWork() {
-    assert null != activeScope();
+    assert null != activeSpan();
   }
 }

dd-java-agent/instrumentation/java-concurrent/src/test/java/runnable/CheckpointTask.java

@@ -1,6 +1,6 @@
 package runnable;
 
-import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
 
 import datadog.trace.api.Trace;
 import java.util.concurrent.CountDownLatch;
@@ -29,6 +29,6 @@ public void run() {
 
   @Trace
   private void traceableChild() {
-    assert null != activeScope();
+    assert null != activeSpan();
   }
 }

dd-java-agent/instrumentation/junit-4.10/munit-junit-4/src/main/java/datadog/trace/instrumentation/junit4/MUnitTracingListener.java

@@ -5,7 +5,6 @@
 import datadog.trace.api.civisibility.execution.TestExecutionHistory;
 import datadog.trace.api.civisibility.telemetry.tag.TestFrameworkInstrumentation;
 import datadog.trace.bootstrap.ContextStore;
-import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes;
@@ -226,12 +225,11 @@ public void testIgnored(final Description description) {
   }
 
   private static boolean isSpanInProgress(UTF8BytesString type) {
-    final AgentScope scope = AgentTracer.activeScope();
-    if (scope == null) {
+    final AgentSpan span = AgentTracer.activeSpan();
+    if (span == null) {
       return false;
     }
-    AgentSpan scopeSpan = scope.span();
-    String spanType = scopeSpan.getSpanType();
+    String spanType = span.getSpanType();
     return spanType != null && spanType.contentEquals(type);
   }
 

dd-java-agent/instrumentation/junit-5.3/src/main/java/datadog/trace/instrumentation/junit5/JUnitPlatformUtils.java

@@ -5,7 +5,6 @@
 import datadog.trace.api.civisibility.config.TestIdentifier;
 import datadog.trace.api.civisibility.config.TestSourceData;
 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.AgentTracer;
 import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes;
@@ -149,11 +148,7 @@ public static boolean isAssumptionFailure(Throwable throwable) {
   }
 
   public static boolean isTestInProgress() {
-    AgentScope activeScope = AgentTracer.activeScope();
-    if (activeScope == null) {
-      return false;
-    }
-    AgentSpan span = activeScope.span();
+    AgentSpan span = AgentTracer.activeSpan();
     if (span == null) {
       return false;
     }

dd-java-agent/instrumentation/rediscala-1.8.0/src/main/java/datadog/trace/instrumentation/rediscala/OnCompleteHandler.java

@@ -1,11 +1,10 @@
 package datadog.trace.instrumentation.rediscala;
 
-import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
 import static datadog.trace.instrumentation.rediscala.RediscalaClientDecorator.DECORATE;
 
 import akka.actor.ActorRef;
 import datadog.trace.bootstrap.ContextStore;
-import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import scala.runtime.AbstractFunction1;
 import scala.util.Try;
@@ -24,9 +23,8 @@ public OnCompleteHandler(
   @Override
   public Void apply(final Try<Object> result) {
     // propagation handled by scala promise instrumentation
-    AgentScope activeScope = activeScope();
-    if (null != activeScope) {
-      AgentSpan span = activeScope.span();
+    AgentSpan span = activeSpan();
+    if (null != span) {
       try {
         if (actorRef != null) {
           DECORATE.onConnection(span, contextStore.get(actorRef));