Fold setAsyncPropagation into preceding activateSpan (#8032)

Author: mcculls

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

@@ -98,8 +98,7 @@ public TestImpl(
 
     span = spanBuilder.start();
 
-    final AgentScope scope = activateSpan(span);
-    scope.setAsyncPropagation(true);
+    activateSpan(span, true);
 
     span.setSpanType(InternalSpanTypes.TEST);
     span.setTag(Tags.SPAN_KIND, Tags.SPAN_KIND_TEST);

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

@@ -121,8 +121,7 @@ public TestSuiteImpl(
     testDecorator.afterStart(span);
 
     if (!parallelized) {
-      final AgentScope scope = activateSpan(span);
-      scope.setAsyncPropagation(true);
+      activateSpan(span, true);
     }
 
     metricCollector.add(CiVisibilityCountMetric.EVENT_CREATED, 1, instrumentation, EventType.SUITE);

dd-java-agent/instrumentation/akka-concurrent/src/test/java/LinearTask.java

@@ -32,9 +32,7 @@ protected Integer compute() {
     } else {
       int next = parent + 1;
       AgentSpan span = startSpan(Integer.toString(next));
-      try (AgentScope scope = activateSpan(span)) {
-        // shouldn't be necessary with a decent API
-        scope.setAsyncPropagation(true);
+      try (AgentScope scope = activateSpan(span, true)) {
         LinearTask child = new LinearTask(next, depth);
         return child.fork().join();
       } finally {

dd-java-agent/instrumentation/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogWrapperHelper.java

@@ -15,9 +15,7 @@ public static AgentScope createSpan(final HttpRequest request) {
     DECORATE.afterStart(span);
     DECORATE.onRequest(span, request, request, extractedContext);
 
-    final AgentScope scope = activateSpan(span);
-    scope.setAsyncPropagation(true);
-    return scope;
+    return activateSpan(span, true);
   }
 
   public static void finishSpan(final AgentSpan span, final HttpResponse response) {

dd-java-agent/instrumentation/finatra-2.9/src/main/java/datadog/trace/instrumentation/finatra/FinatraInstrumentation.java

@@ -76,9 +76,7 @@ public static AgentScope nameSpan(
       DECORATE.afterStart(span);
       span.setResourceName(DECORATE.className(clazz));
 
-      final AgentScope scope = activateSpan(span);
-      scope.setAsyncPropagation(true);
-      return scope;
+      return activateSpan(span, true);
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)

dd-java-agent/instrumentation/grizzly-2/src/main/java/datadog/trace/instrumentation/grizzly/GrizzlyHttpHandlerInstrumentation.java

@@ -78,8 +78,7 @@ public static class HandleAdvice {
       DECORATE.afterStart(span);
       DECORATE.onRequest(span, request, request, parentContext);
 
-      scope = activateSpan(span);
-      scope.setAsyncPropagation(true);
+      scope = activateSpan(span, true);
 
       request.setAttribute(DD_SPAN_ATTRIBUTE, span);
       request.setAttribute(CorrelationIdentifier.getTraceIdKey(), GlobalTracer.get().getTraceId());

dd-java-agent/instrumentation/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/FilterAdvice.java

@@ -7,27 +7,22 @@
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import net.bytebuddy.asm.Advice;
-import org.glassfish.grizzly.filterchain.BaseFilter;
 import org.glassfish.grizzly.filterchain.FilterChainContext;
 
 public class FilterAdvice {
 
   @Advice.OnMethodEnter(suppress = Throwable.class)
-  public static AgentScope onEnter(
-      @Advice.This BaseFilter it, @Advice.Argument(0) final FilterChainContext ctx) {
-    if (ctx.getAttributes().getAttribute(DD_SPAN_ATTRIBUTE) == null || activeSpan() != null) {
+  public static AgentScope onEnter(@Advice.Argument(0) final FilterChainContext ctx) {
+    Object span = ctx.getAttributes().getAttribute(DD_SPAN_ATTRIBUTE);
+    if (span == null || activeSpan() != null) {
       return null;
     }
-    AgentScope scope =
-        activateSpan((AgentSpan) ctx.getAttributes().getAttribute(DD_SPAN_ATTRIBUTE));
-    scope.setAsyncPropagation(true);
-    return scope;
+    return activateSpan((AgentSpan) span, true);
   }
 
   @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
-  public static void onExit(@Advice.This BaseFilter it, @Advice.Enter final AgentScope scope) {
+  public static void onExit(@Advice.Enter final AgentScope scope) {
     if (scope != null) {
-      scope.setAsyncPropagation(false);
       scope.close();
     }
   }

dd-java-agent/instrumentation/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/GrizzlyDecorator.java

@@ -115,8 +115,7 @@ public static NextAction onHttpCodecFilterExit(
     HttpResponsePacket httpResponse = httpRequest.getResponse();
     AgentSpan.Context.Extracted context = DECORATE.extract(httpRequest);
     AgentSpan span = DECORATE.startSpan(httpRequest, context);
-    AgentScope scope = activateSpan(span);
-    scope.setAsyncPropagation(true);
+    AgentScope scope = activateSpan(span, true);
     DECORATE.afterStart(span);
     ctx.getAttributes().setAttribute(DD_SPAN_ATTRIBUTE, span);
     ctx.getAttributes().setAttribute(DD_RESPONSE_ATTRIBUTE, httpResponse);

dd-java-agent/instrumentation/ignite-2.0/src/main/java/datadog/trace/instrumentation/ignite/v2/cache/IgniteCacheAsyncInstrumentation.java

@@ -91,11 +91,9 @@ public static AgentScope onEnter(
           span, InstrumentationContext.get(IgniteCache.class, Ignite.class).get(that));
       DECORATE.onOperation(span, that.getName(), methodName);
 
-      final AgentScope scope = activateSpan(span);
       // Enable async propagation, so the newly spawned task will be associated back with this
       // original trace.
-      scope.setAsyncPropagation(true);
-      return scope;
+      return activateSpan(span, true);
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
@@ -149,11 +147,9 @@ public static AgentScope onEnter(
           span, InstrumentationContext.get(IgniteCache.class, Ignite.class).get(that));
       DECORATE.onOperation(span, that.getName(), methodName, key);
 
-      final AgentScope scope = activateSpan(span);
       // Enable async propagation, so the newly spawned task will be associated back with this
       // original trace.
-      scope.setAsyncPropagation(true);
-      return scope;
+      return activateSpan(span, true);
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)

dd-java-agent/instrumentation/jakarta-rs-annotations-3/src/main/java/datadog/trace/instrumentation/jakarta3/DefaultRequestContextInstrumentation.java

@@ -47,8 +47,7 @@ public static AgentScope createGenericSpan(@Advice.This final ContainerRequestCo
           // can only be aborted inside the filter method
         }
 
-        final AgentScope scope = activateSpan(span);
-        scope.setAsyncPropagation(true);
+        final AgentScope scope = activateSpan(span, true);
 
         DECORATE.afterStart(span);
         DECORATE.onJakartaRsSpan(span, parent, filterClass, method);