Spring scheduling: ensure spans have no parent (#7583)

* Spring scheduling: ensure spans have no parent

* simplify

* add more tests

Author: amarziali

dd-java-agent/instrumentation/spring-scheduling-3.1/src/main/java/datadog/trace/instrumentation/springscheduling/SpringSchedulingRunnableWrapper.java

@@ -7,10 +7,13 @@
 import static datadog.trace.instrumentation.springscheduling.SpringSchedulingDecorator.DECORATE;
 import static datadog.trace.instrumentation.springscheduling.SpringSchedulingDecorator.SCHEDULED_CALL;
 
+import datadog.trace.api.Config;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 
 public class SpringSchedulingRunnableWrapper implements Runnable {
+  private static final boolean LEGACY_TRACING =
+      Config.get().isLegacyTracingEnabled(false, "spring-scheduling");
   private final Runnable runnable;
 
   private SpringSchedulingRunnableWrapper(final Runnable runnable) {
@@ -19,7 +22,8 @@ private SpringSchedulingRunnableWrapper(final Runnable runnable) {
 
   @Override
   public void run() {
-    final AgentSpan span = startSpan(SCHEDULED_CALL);
+    final AgentSpan span =
+        LEGACY_TRACING ? startSpan(SCHEDULED_CALL) : startSpan(SCHEDULED_CALL, null);
     DECORATE.afterStart(span);
 
     try (final AgentScope scope = activateSpan(span)) {

dd-java-agent/instrumentation/spring-scheduling-3.1/src/test/groovy/SpringSchedulingTest.groovy

@@ -1,3 +1,5 @@
+import static datadog.trace.agent.test.utils.TraceUtils.basicSpan
+
 import datadog.trace.agent.test.AgentTestRunner
 import datadog.trace.bootstrap.instrumentation.api.Tags
 import org.springframework.boot.actuate.scheduling.ScheduledTasksEndpoint
@@ -7,21 +9,42 @@ import java.util.concurrent.TimeUnit
 
 class SpringSchedulingTest extends AgentTestRunner {
 
+  def legacyTracing() {
+    false
+  }
+
+  @Override
+  protected void configurePreAgent() {
+    super.configurePreAgent()
+    if (legacyTracing()) {
+      injectSysConfig("spring-scheduling.legacy.tracing.enabled", "true")
+    }
+  }
+
   def "schedule trigger test according to cron expression"() {
     setup:
-    def context = new AnnotationConfigApplicationContext(TriggerTaskConfig)
+    def context = new AnnotationConfigApplicationContext(TriggerTaskConfig, SchedulingConfig)
     def task = context.getBean(TriggerTask)
 
     task.blockUntilExecute()
 
     expect:
     assert task != null
-    assertTraces(1) {
-      trace(1) {
+    def hasParent = legacyTracing()
+    assertTraces(hasParent ? 1 : 2) {
+      if (!hasParent) {
+        trace(1) {
+          basicSpan(it, "parent")
+        }
+      }
+      trace(hasParent ? 2 : 1) {
+        if (hasParent) {
+          basicSpan(it, "parent")
+        }
         span {
           resourceName "TriggerTask.run"
           operationName "scheduled.call"
-          parent()
+          hasParent ? childOfPrevious() : parent()
           errored false
           tags {
             "$Tags.COMPONENT" "spring-scheduling"
@@ -42,19 +65,29 @@ class SpringSchedulingTest extends AgentTestRunner {
 
   def "schedule interval test"() {
     setup:
-    def context = new AnnotationConfigApplicationContext(IntervalTaskConfig)
+    def context = new AnnotationConfigApplicationContext(IntervalTaskConfig, SchedulingConfig)
     def task = context.getBean(IntervalTask)
 
     task.blockUntilExecute()
 
     expect:
     assert task != null
-    assertTraces(1) {
-      trace(1) {
+    def hasParent = legacyTracing()
+
+    assertTraces(hasParent ? 1 : 2) {
+      if (!hasParent) {
+        trace(1) {
+          basicSpan(it, "parent")
+        }
+      }
+      trace(hasParent ? 2 : 1) {
+        if (hasParent) {
+          basicSpan(it, "parent")
+        }
         span {
           resourceName "IntervalTask.run"
           operationName "scheduled.call"
-          parent()
+          hasParent ? childOfPrevious() : parent()
           errored false
           tags {
             "$Tags.COMPONENT" "spring-scheduling"
@@ -69,18 +102,27 @@ class SpringSchedulingTest extends AgentTestRunner {
 
   def "schedule lambda test"() {
     setup:
-    def context = new AnnotationConfigApplicationContext(LambdaTaskConfig)
+    def context = new AnnotationConfigApplicationContext(LambdaTaskConfig, SchedulingConfig)
     def configurer = context.getBean(LambdaTaskConfigurer)
 
     configurer.singleUseLatch.await(2000, TimeUnit.MILLISECONDS)
 
     expect:
-    assertTraces(1) {
-      trace(1) {
+    def hasParent = legacyTracing()
+    assertTraces(hasParent ? 1 : 2) {
+      if (!hasParent) {
+        trace(1) {
+          basicSpan(it, "parent")
+        }
+      }
+      trace(hasParent ? 2 : 1) {
+        if (hasParent) {
+          basicSpan(it, "parent")
+        }
         span {
           resourceNameContains("LambdaTaskConfigurer\$\$Lambda")
           operationName "scheduled.call"
-          parent()
+          hasParent ? childOfPrevious() : parent()
           errored false
           tags {
             "$Tags.COMPONENT" "spring-scheduling"
@@ -94,3 +136,11 @@ class SpringSchedulingTest extends AgentTestRunner {
     context.close()
   }
 }
+
+class SpringSchedulingLegacyTracingForkedTest extends SpringSchedulingTest {
+  @Override
+  def legacyTracing() {
+    true
+  }
+}
+

dd-java-agent/instrumentation/spring-scheduling-3.1/src/test/java/SchedulingConfig.java

@@ -0,0 +1,95 @@
+import static datadog.trace.agent.test.utils.TraceUtils.runnableUnderTrace;
+
+import java.time.Clock;
+import java.time.Duration;
+import java.time.Instant;
+import java.util.Date;
+import java.util.concurrent.ScheduledFuture;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.scheduling.TaskScheduler;
+import org.springframework.scheduling.Trigger;
+import org.springframework.scheduling.annotation.SchedulingConfigurer;
+import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
+import org.springframework.scheduling.config.ScheduledTaskRegistrar;
+
+@Configuration
+public class SchedulingConfig implements SchedulingConfigurer {
+  static class TracingTaskScheduler implements TaskScheduler {
+    final ThreadPoolTaskScheduler scheduler;
+
+    TracingTaskScheduler() {
+      scheduler = new ThreadPoolTaskScheduler();
+      scheduler.initialize();
+    }
+
+    @Override
+    public ScheduledFuture<?> scheduleWithFixedDelay(Runnable task, Duration delay) {
+      return scheduler.scheduleWithFixedDelay(() -> runnableUnderTrace("parent", task), delay);
+    }
+
+    @Override
+    public Clock getClock() {
+      return Clock.systemUTC();
+    }
+
+    @Override
+    public ScheduledFuture<?> schedule(Runnable task, Instant startTime) {
+      return scheduler.schedule(() -> runnableUnderTrace("parent", task), startTime);
+    }
+
+    @Override
+    public ScheduledFuture<?> scheduleAtFixedRate(
+        Runnable task, Instant startTime, Duration period) {
+      return scheduler.scheduleAtFixedRate(
+          () -> runnableUnderTrace("parent", task), startTime, period);
+    }
+
+    @Override
+    public ScheduledFuture<?> scheduleAtFixedRate(Runnable task, Duration period) {
+      return scheduler.scheduleAtFixedRate(() -> runnableUnderTrace("parent", task), period);
+    }
+
+    @Override
+    public ScheduledFuture<?> scheduleWithFixedDelay(
+        Runnable task, Instant startTime, Duration delay) {
+      return scheduler.scheduleWithFixedDelay(
+          () -> runnableUnderTrace("parent", task), startTime, delay);
+    }
+
+    @Override
+    public ScheduledFuture<?> schedule(Runnable runnable, Trigger trigger) {
+      return scheduler.schedule(() -> runnableUnderTrace("parent", runnable), trigger);
+    }
+
+    @Override
+    public ScheduledFuture<?> schedule(Runnable runnable, Date date) {
+      return scheduler.schedule(() -> runnableUnderTrace("parent", runnable), date);
+    }
+
+    @Override
+    public ScheduledFuture<?> scheduleAtFixedRate(Runnable runnable, Date date, long l) {
+      return scheduler.scheduleAtFixedRate(() -> runnableUnderTrace("parent", runnable), date, l);
+    }
+
+    @Override
+    public ScheduledFuture<?> scheduleAtFixedRate(Runnable runnable, long l) {
+      return scheduler.scheduleAtFixedRate(() -> runnableUnderTrace("parent", runnable), l);
+    }
+
+    @Override
+    public ScheduledFuture<?> scheduleWithFixedDelay(Runnable runnable, Date date, long l) {
+      return scheduler.scheduleWithFixedDelay(
+          () -> runnableUnderTrace("parent", runnable), date, l);
+    }
+
+    @Override
+    public ScheduledFuture<?> scheduleWithFixedDelay(Runnable runnable, long l) {
+      return scheduler.scheduleWithFixedDelay(() -> runnableUnderTrace("parent", runnable), l);
+    }
+  }
+
+  @Override
+  public void configureTasks(ScheduledTaskRegistrar scheduledTaskRegistrar) {
+    scheduledTaskRegistrar.setTaskScheduler(new TracingTaskScheduler());
+  }
+}