Add support for Spring Scheduling

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

@@ -0,0 +1,28 @@
+muzzle {
+  pass {
+    group = 'org.springframework'
+    module = 'spring-context'
+    versions = "[3.1.0.RELEASE,]"
+    assertInverse = true
+  }
+}
+
+apply from: "${rootDir}/gradle/java.gradle"
+
+apply plugin: 'org.unbroken-dome.test-sets'
+
+testSets {
+  latestDepTest {
+    dirName = 'test'
+  }
+}
+
+dependencies {
+  // 3.2.3 is the first version with which the tests will run. Lower versions require other
+  // classes and packages to be imported. Versions 3.1.0+ work with the instrumentation.
+  compileOnly group: 'org.springframework', name: 'spring-context', version: '3.1.0.RELEASE'
+  testCompile group: 'org.springframework', name: 'spring-context', version: '3.2.3.RELEASE'
+
+  // this is the latest version that supports Java 7
+  latestDepTestCompile group: 'org.springframework', name: 'spring-context', version: '4.+'
+}

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

@@ -0,0 +1,45 @@
+package datadog.trace.instrumentation.springscheduling;
+
+import datadog.trace.agent.decorator.BaseDecorator;
+import datadog.trace.api.DDTags;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import lombok.extern.slf4j.Slf4j;
+import org.springframework.scheduling.support.ScheduledMethodRunnable;
+
+@Slf4j
+public class SpringSchedulingDecorator extends BaseDecorator {
+  public static final SpringSchedulingDecorator DECORATE = new SpringSchedulingDecorator();
+
+  private SpringSchedulingDecorator() {}
+
+  @Override
+  protected String[] instrumentationNames() {
+    return new String[] {"spring-scheduling"};
+  }
+
+  @Override
+  protected String spanType() {
+    return null;
+  }
+
+  @Override
+  protected String component() {
+    return "spring-scheduling";
+  }
+
+  public AgentSpan onRun(final AgentSpan span, final Runnable runnable) {
+    if (runnable != null) {
+      String resourceName = "";
+      if (runnable instanceof ScheduledMethodRunnable) {
+        final ScheduledMethodRunnable scheduledMethodRunnable = (ScheduledMethodRunnable) runnable;
+        resourceName = spanNameForMethod(scheduledMethodRunnable.getMethod());
+      } else {
+        final String className = spanNameForClass(runnable.getClass());
+        final String methodName = "run";
+        resourceName = className + "." + methodName;
+      }
+      span.setTag(DDTags.RESOURCE_NAME, resourceName);
+    }
+    return span;
+  }
+}

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

@@ -0,0 +1,96 @@
+package datadog.trace.instrumentation.springscheduling;
+
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
+import static datadog.trace.instrumentation.springscheduling.SpringSchedulingDecorator.DECORATE;
+import static java.util.Collections.singletonMap;
+import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
+import static net.bytebuddy.matcher.ElementMatchers.isInterface;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.not;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@AutoService(Instrumenter.class)
+public final class SpringSchedulingInstrumentation extends Instrumenter.Default {
+
+  public SpringSchedulingInstrumentation() {
+    super("spring-scheduling");
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return not(isInterface()).and(named("org.springframework.scheduling.config.Task"));
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      "datadog.trace.agent.decorator.BaseDecorator",
+      packageName + ".SpringSchedulingDecorator",
+      getClass().getName() + "$RunnableWrapper",
+    };
+  }
+
+  @Override
+  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    return singletonMap(
+        isConstructor().and(takesArgument(0, Runnable.class)),
+        SpringSchedulingInstrumentation.class.getName() + "$SpringSchedulingAdvice");
+  }
+
+  public static class SpringSchedulingAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void onConstruction(
+        @Advice.Argument(value = 0, readOnly = false) Runnable runnable) {
+      runnable = RunnableWrapper.wrapIfNeeded(runnable);
+    }
+  }
+
+  public static class RunnableWrapper implements Runnable {
+    private final Runnable runnable;
+
+    private RunnableWrapper(final Runnable runnable) {
+      this.runnable = runnable;
+    }
+
+    @Override
+    public void run() {
+      final AgentSpan span = startSpan("scheduled.call");
+      DECORATE.afterStart(span);
+
+      try (final AgentScope scope = activateSpan(span, false)) {
+        DECORATE.onRun(span, runnable);
+        scope.setAsyncPropagation(true);
+
+        try {
+          runnable.run();
+        } catch (final Throwable throwable) {
+          DECORATE.onError(span, throwable);
+          throw throwable;
+        }
+      } finally {
+        DECORATE.beforeFinish(span);
+        span.finish();
+      }
+    }
+
+    public static Runnable wrapIfNeeded(final Runnable task) {
+      // We wrap only lambdas' anonymous classes and if given object has not already been wrapped.
+      // Anonymous classes have '/' in class name which is not allowed in 'normal' classes.
+      if (task instanceof RunnableWrapper) {
+        return task;
+      }
+      return new RunnableWrapper(task);
+    }
+  }
+}

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

@@ -0,0 +1,57 @@
+import datadog.trace.agent.test.AgentTestRunner
+import datadog.trace.bootstrap.instrumentation.api.Tags
+import org.springframework.context.annotation.AnnotationConfigApplicationContext
+
+class SpringSchedulingTest extends AgentTestRunner {
+
+  def "schedule trigger test according to cron expression"() {
+    setup:
+    def context = new AnnotationConfigApplicationContext(TriggerTaskConfig)
+    def task = context.getBean(TriggerTask)
+
+    task.blockUntilExecute()
+
+    expect:
+    assert task != null
+    assertTraces(1) {
+      trace(0, 1) {
+        span(0) {
+          resourceName "TriggerTask.run"
+          operationName "scheduled.call"
+          parent()
+          errored false
+          tags {
+            "$Tags.COMPONENT" "spring-scheduling"
+            defaultTags()
+          }
+        }
+      }
+    }
+  }
+
+  def "schedule interval test"() {
+    setup:
+    def context = new AnnotationConfigApplicationContext(IntervalTaskConfig)
+    def task = context.getBean(IntervalTask)
+
+    task.blockUntilExecute()
+
+    expect:
+    assert task != null
+    assertTraces(1) {
+      trace(0, 1) {
+        span(0) {
+          resourceName "IntervalTask.run"
+          operationName "scheduled.call"
+          parent()
+          errored false
+          tags {
+            "$Tags.COMPONENT" "spring-scheduling"
+            defaultTags()
+          }
+        }
+      }
+    }
+
+  }
+}

dd-java-agent/instrumentation/spring-scheduling-3.1/src/test/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)
+  @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/test/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/test/java/TriggerTask.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 TriggerTask implements Runnable {
+
+  private final CountDownLatch latch = new CountDownLatch(1);
+
+  @Scheduled(cron = "0/5 * * * * *")
+  @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/test/java/TriggerTaskConfig.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 TriggerTaskConfig {
+  @Bean
+  public TriggerTask triggerTasks() {
+    return new TriggerTask();
+  }
+}

settings.gradle

@@ -130,6 +130,7 @@ include ':dd-java-agent:instrumentation:servlet:request-3'
 include ':dd-java-agent:instrumentation:slf4j-mdc'
 include ':dd-java-agent:instrumentation:sparkjava-2.3'
 include ':dd-java-agent:instrumentation:spring-data-1.8'
+include ':dd-java-agent:instrumentation:spring-scheduling-3.1'
 include ':dd-java-agent:instrumentation:spring-webmvc-3.1'
 include ':dd-java-agent:instrumentation:spring-webflux-5'
 include ':dd-java-agent:instrumentation:spymemcached-2.12'