Preserve getQualifier from spring scheduling runnables

Author: amarziali

dd-java-agent/instrumentation/spring-scheduling-3.1/build.gradle

@@ -1,7 +1,9 @@
 ext {
+  latestDepForkedTestMinJavaVersionForTests = JavaVersion.VERSION_17
+  latestDepTestMinJavaVersionForTests = JavaVersion.VERSION_17
   spring6TestMinJavaVersionForTests = JavaVersion.VERSION_17
-  spring6LatestDepTestMinJavaVersionForTests = JavaVersion.VERSION_17
 }
+
 muzzle {
   pass {
     group = 'org.springframework'
@@ -20,22 +22,30 @@ muzzle {
 
 apply from: "$rootDir/gradle/java.gradle"
 
-addTestSuiteForDir('latestDepTest', 'test')
+addTestSuite('latestDepTest')
+addTestSuiteExtendingForDir('latestDepForkedTest','latestDepTest', 'latestDepTest')
 addTestSuiteForDir('spring6Test', 'test')
-addTestSuiteExtendingForDir('spring6LatestDepTest', 'latestDepTest', 'test')
+addTestSuiteExtendingForDir('spring5LatestDepTest', 'latestDepTest', 'test')
 
-[compileSpring6TestJava, compileSpring6LatestDepTestJava].each {
+[compileSpring6TestJava, compileLatestDepTestJava].each {
   setJavaVersion(it, 17)
   sourceCompatibility = JavaVersion.VERSION_1_8
   targetCompatibility = JavaVersion.VERSION_1_8
 }
 
 
-[compileSpring6TestGroovy, compileSpring6LatestDepTestGroovy, spring6Test, spring6LatestDepTest].each {
+[
+  compileSpring6TestGroovy,
+  compileLatestDepTestGroovy,
+  compileLatestDepForkedTestGroovy,
+  spring6Test,
+  latestDepTest,
+  latestDepForkedTest
+].each {
   it.javaLauncher = getJavaLauncherFor(17)
 }
 
-[spring6Test, spring6LatestDepTest].each {
+[spring6Test, latestDepTest, latestDepForkedTest].each {
   it.jvmArgs '--add-opens', 'java.base/java.util=ALL-UNNAMED'
 }
 
@@ -55,19 +65,17 @@ dependencies {
   testImplementation group: 'org.springframework.boot', name: 'spring-boot-starter-actuator', version: '2.4.0'
 
 
-  latestDepTestImplementation group: 'org.springframework', name: 'spring-context', version: '5.+'
+  spring5LatestDepTestImplementation group: 'org.springframework', name: 'spring-context', version: '5.+'
 
-  latestDepTestImplementation group: 'net.javacrumbs.shedlock', name: 'shedlock-spring', version: '4.+'
-  latestDepTestImplementation group: 'net.javacrumbs.shedlock', name: 'shedlock-provider-jdbc-template', version: '4.+'
-  latestDepTestImplementation group: 'com.h2database', name: 'h2', version: '2.2.+' // 2.3+ requires Java 11
-  latestDepTestImplementation group: 'org.springframework.boot', name: 'spring-boot-starter-data-jpa', version: '2.+'
+  spring5LatestDepTestImplementation group: 'net.javacrumbs.shedlock', name: 'shedlock-spring', version: '4.+'
+  spring5LatestDepTestImplementation group: 'net.javacrumbs.shedlock', name: 'shedlock-provider-jdbc-template', version: '4.+'
+  spring5LatestDepTestImplementation group: 'com.h2database', name: 'h2', version: '2.2.+' // 2.3+ requires Java 11
+  spring5LatestDepTestImplementation group: 'org.springframework.boot', name: 'spring-boot-starter-data-jpa', version: '2.+'
 
   spring6TestImplementation group: 'org.springframework', name: 'spring-context', version: '6.0.0.RELEASE'
   spring6TestImplementation group: 'org.springframework.boot', name: 'spring-boot-starter-data-jpa', version: '3.0.0'
 
-  spring6LatestDepTestImplementation group: 'com.h2database', name: 'h2', version: '+'
-  spring6LatestDepTestImplementation group: 'org.springframework', name: 'spring-context', version: '6.+'
-  spring6LatestDepTestImplementation group: 'org.springframework.boot', name: 'spring-boot-starter-data-jpa', version: '3.+'
-
-
+  latestDepTestImplementation group: 'com.h2database', name: 'h2', version: '+'
+  latestDepTestImplementation group: 'org.springframework', name: 'spring-context', version: '6.+'
+  latestDepTestImplementation group: 'org.springframework.boot', name: 'spring-boot-starter-data-jpa', version: '3.+'
 }

dd-java-agent/instrumentation/spring-scheduling-3.1/src/latestDepTest/groovy/ShedlockTest.groovy

@@ -0,0 +1,27 @@
+import datadog.trace.agent.test.AgentTestRunner
+import datadog.trace.test.util.Flaky
+import org.springframework.context.annotation.AnnotationConfigApplicationContext
+import org.springframework.jdbc.core.JdbcTemplate
+
+import javax.sql.DataSource
+import java.util.concurrent.TimeUnit
+
+class ShedlockTest extends AgentTestRunner {
+
+  @Flaky("task.invocationCount() == 0")
+  def "should not disable shedlock"() {
+    setup:
+    def context = new AnnotationConfigApplicationContext(ShedlockConfig)
+    JdbcTemplate jdbcTemplate = new JdbcTemplate(context.getBean(DataSource))
+    jdbcTemplate.execute("CREATE TABLE shedlock(name VARCHAR(64) NOT NULL PRIMARY KEY, lock_until TIMESTAMP NOT NULL,\n" +
+      "    locked_at TIMESTAMP NOT NULL, locked_by VARCHAR(255) NOT NULL);")
+    def task = context.getBean(ShedLockedTask)
+
+    expect: "lock is held for more than one second"
+    !task.awaitInvocation(1000, TimeUnit.MILLISECONDS)
+    task.invocationCount() == 1
+
+    cleanup:
+    context.close()
+  }
+}

dd-java-agent/instrumentation/spring-scheduling-3.1/src/latestDepTest/groovy/SpringAsyncTest.groovy

@@ -0,0 +1,35 @@
+import datadog.trace.agent.test.AgentTestRunner
+import org.springframework.context.annotation.AnnotationConfigApplicationContext
+
+import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
+
+class SpringAsyncTest extends AgentTestRunner {
+
+  def "context propagated through @async annotation"() {
+    setup:
+    def context = new AnnotationConfigApplicationContext(AsyncTaskConfig)
+    AsyncTask asyncTask = context.getBean(AsyncTask)
+    when:
+    runUnderTrace("root") {
+      asyncTask.async().join()
+    }
+    then:
+    assertTraces(1) {
+      trace(3) {
+        span {
+          resourceName "root"
+        }
+        span {
+          resourceName "AsyncTask.async"
+          threadNameStartsWith "SimpleAsyncTaskExecutor"
+          childOf span(0)
+        }
+        span {
+          resourceName "AsyncTask.getInt"
+          threadNameStartsWith "SimpleAsyncTaskExecutor"
+          childOf span(1)
+        }
+      }
+    }
+  }
+}

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

@@ -0,0 +1,146 @@
+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
+import org.springframework.context.annotation.AnnotationConfigApplicationContext
+
+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, SchedulingConfig)
+    def task = context.getBean(TriggerTask)
+
+    task.blockUntilExecute()
+
+    expect:
+    assert task != null
+    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"
+          hasParent ? childOfPrevious() : parent()
+          errored false
+          tags {
+            "$Tags.COMPONENT" "spring-scheduling"
+            defaultTags()
+          }
+        }
+      }
+    }
+    and:
+    def scheduledTaskEndpoint = context.getBean(ScheduledTasksEndpoint)
+    assert scheduledTaskEndpoint != null
+    scheduledTaskEndpoint.scheduledTasks().getCron().each {
+      it.getRunnable().getTarget() == TriggerTask.getName()
+    }
+    cleanup:
+    context.close()
+  }
+
+  def "schedule interval test"() {
+    setup:
+    def context = new AnnotationConfigApplicationContext(IntervalTaskConfig, SchedulingConfig)
+    def task = context.getBean(IntervalTask)
+
+    task.blockUntilExecute()
+
+    expect:
+    assert task != null
+    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"
+          hasParent ? childOfPrevious() : parent()
+          errored false
+          tags {
+            "$Tags.COMPONENT" "spring-scheduling"
+            defaultTags()
+          }
+        }
+      }
+    }
+    cleanup:
+    context.close()
+  }
+
+  def "schedule lambda test"() {
+    setup:
+    def context = new AnnotationConfigApplicationContext(LambdaTaskConfig, SchedulingConfig)
+    def configurer = context.getBean(LambdaTaskConfigurer)
+
+    configurer.singleUseLatch.await(2000, TimeUnit.MILLISECONDS)
+
+    expect:
+    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"
+          hasParent ? childOfPrevious() : parent()
+          errored false
+          tags {
+            "$Tags.COMPONENT" "spring-scheduling"
+            defaultTags()
+          }
+        }
+      }
+    }
+
+    cleanup:
+    context.close()
+  }
+}
+
+class SpringSchedulingLegacyTracingForkedTest extends SpringSchedulingTest {
+  @Override
+  def legacyTracing() {
+    true
+  }
+}
+

dd-java-agent/instrumentation/spring-scheduling-3.1/src/latestDepTest/java/AsyncTask.java

@@ -0,0 +1,17 @@
+import datadog.trace.api.Trace;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ThreadLocalRandom;
+import org.springframework.scheduling.annotation.Async;
+
+public class AsyncTask {
+
+  @Async
+  public CompletableFuture<Integer> async() {
+    return CompletableFuture.completedFuture(getInt());
+  }
+
+  @Trace
+  public int getInt() {
+    return ThreadLocalRandom.current().nextInt();
+  }
+}

dd-java-agent/instrumentation/spring-scheduling-3.1/src/latestDepTest/java/AsyncTaskConfig.java

@@ -0,0 +1,15 @@
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.scheduling.annotation.EnableAsync;
+import org.springframework.scheduling.annotation.EnableScheduling;
+
+@Configuration
+@EnableScheduling
+@EnableAsync
+public class AsyncTaskConfig {
+
+  @Bean
+  AsyncTask asyncTask() {
+    return new AsyncTask();
+  }
+}

dd-java-agent/instrumentation/spring-scheduling-3.1/src/latestDepTest/java/IntervalTask.java

@@ -0,0 +1,20 @@
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import org.springframework.scheduling.annotation.Scheduled;
+import org.springframework.stereotype.Component;
+
+@Component
+public class IntervalTask implements Runnable {
+
+  private final CountDownLatch latch = new CountDownLatch(1);
+
+  @Scheduled(fixedRate = 5000, scheduler = "tracingTaskScheduler")
+  @Override
+  public void run() {
+    latch.countDown();
+  }
+
+  public void blockUntilExecute() throws InterruptedException {
+    latch.await(5, TimeUnit.SECONDS);
+  }
+}

dd-java-agent/instrumentation/spring-scheduling-3.1/src/latestDepTest/java/IntervalTaskConfig.java

@@ -0,0 +1,12 @@
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.scheduling.annotation.EnableScheduling;
+
+@Configuration
+@EnableScheduling
+public class IntervalTaskConfig {
+  @Bean
+  public IntervalTask scheduledTasks() {
+    return new IntervalTask();
+  }
+}

dd-java-agent/instrumentation/spring-scheduling-3.1/src/latestDepTest/java/LambdaTaskConfig.java

@@ -0,0 +1,13 @@
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.scheduling.annotation.EnableScheduling;
+
+@Configuration
+@EnableScheduling
+public class LambdaTaskConfig {
+
+  @Bean
+  LambdaTaskConfigurer lambdaTaskConfigurer() {
+    return new LambdaTaskConfigurer();
+  }
+}

dd-java-agent/instrumentation/spring-scheduling-3.1/src/latestDepTest/java/LambdaTaskConfigurer.java

@@ -0,0 +1,17 @@
+import java.util.concurrent.CountDownLatch;
+import org.springframework.scheduling.annotation.SchedulingConfigurer;
+import org.springframework.scheduling.config.ScheduledTaskRegistrar;
+import org.springframework.stereotype.Service;
+
+@Service
+public class LambdaTaskConfigurer implements SchedulingConfigurer {
+
+  public final CountDownLatch singleUseLatch = new CountDownLatch(1);
+
+  @Override
+  public void configureTasks(ScheduledTaskRegistrar taskRegistrar) {
+    // need to manually set in this case since it won't use the scheduler in the annotation
+    taskRegistrar.setTaskScheduler(new SchedulingConfig.TracingTaskScheduler());
+    taskRegistrar.addFixedDelayTask(singleUseLatch::countDown, 500);
+  }
+}