add test for interval task and trigger task

dd-java-agent/instrumentation/spring-data-1.8/src/main/java/datadog/trace/instrumentation/springdata/SpringRepositoryInstrumentation.java

@@ -2,21 +2,10 @@
 
 package datadog.trace.instrumentation.springdata;
 
-import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
-import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
-import static datadog.trace.instrumentation.springdata.SpringDataDecorator.DECORATOR;
-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 com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
-import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.instrumentation.api.AgentScope;
-import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.instrumentation.api.AgentSpan;
-import java.lang.reflect.Method;
-import java.util.Collections;
-import java.util.Map;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
@@ -29,6 +18,18 @@ import org.springframework.data.repository.core.RepositoryInformation;
 import org.springframework.data.repository.core.support.RepositoryFactorySupport;
 import org.springframework.data.repository.core.support.RepositoryProxyPostProcessor;
 
+import java.lang.reflect.Method;
+import java.util.Collections;
+import java.util.Map;
+
+import static datadog.trace.instrumentation.api.AgentTracer.activateSpan;
+import static datadog.trace.instrumentation.api.AgentTracer.startSpan;
+import static datadog.trace.instrumentation.springdata.SpringDataDecorator.DECORATOR;
+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;
+
 @AutoService(Instrumenter.class)
 public final class SpringRepositoryInstrumentation extends Instrumenter.Default {
 

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

@@ -3,9 +3,14 @@
 package datadog.trace.instrumentation.springscheduling;
 
 import datadog.trace.agent.decorator.BaseDecorator;
+import datadog.trace.api.DDTags;
+import datadog.trace.instrumentation.api.AgentSpan;
+import lombok.extern.slf4j.Slf4j;
+import org.springframework.scheduling.support.ScheduledMethodRunnable;
 
-public final class SpringSchedulingDecorator extends BaseDecorator {
+@Slf4j
-  public static final SpringSchedulingDecorator DECORATOR = new SpringSchedulingDecorator();
+public class SpringSchedulingDecorator extends BaseDecorator {
+  public static final SpringSchedulingDecorator DECORATE = new SpringSchedulingDecorator();
 
   private SpringSchedulingDecorator() {}
 
@@ -23,4 +28,20 @@ public final class SpringSchedulingDecorator extends BaseDecorator {
   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 = runnable.getClass().getName();
+        final String methodName = "run";
+        resourceName = className + "." + methodName;
+      }
+      span.setTag(DDTags.RESOURCE_NAME, resourceName);
+    }
+    return span;
+  }
 }

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

@@ -1,19 +1,26 @@
-// This file includes software developed at SignalFx
-
 package datadog.trace.instrumentation.springscheduling;
 
-import static java.util.Collections.singletonMap;
-import static net.bytebuddy.matcher.ElementMatchers.*;
-
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
-import datadog.trace.api.Trace;
+import datadog.trace.instrumentation.api.AgentScope;
-import java.util.Map;
+import datadog.trace.instrumentation.api.AgentSpan;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 
+import java.util.Map;
+
+import static datadog.trace.instrumentation.api.AgentTracer.activateSpan;
+import static datadog.trace.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;
+
 @AutoService(Instrumenter.class)
 public final class SpringSchedulingInstrumentation extends Instrumenter.Default {
 
@@ -39,7 +46,7 @@ public final class SpringSchedulingInstrumentation extends Instrumenter.Default
   public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
     return singletonMap(
         isConstructor().and(takesArgument(0, Runnable.class)),
-        SpringSchedulingInstrumentation.class.getName() + "$RepositoryFactorySupportAdvice");
+        SpringSchedulingInstrumentation.class.getName() + "$SpringSchedulingAdvice");
   }
 
   public static class SpringSchedulingAdvice {
@@ -57,10 +64,25 @@ public final class SpringSchedulingInstrumentation extends Instrumenter.Default
       this.runnable = runnable;
     }
 
-    @Trace
     @Override
     public void run() {
-      runnable.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) {

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

@@ -1,71 +0,0 @@
-import datadog.trace.agent.test.AgentTestRunner
-import datadog.trace.agent.test.utils.ConfigUtils
-import datadog.trace.instrumentation.api.Tags
-import org.springframework.context.annotation.AnnotationConfigApplicationContext
-
-class ReportTimeTest extends AgentTestRunner {
-
-  static {
-    ConfigUtils.updateConfig {
-      System.clearProperty("dd.trace.annotations")
-    }
-  }
-
-  def "test method executed"() {
-    setup:
-    def scheduledTask = new ScheduledTasks();
-    scheduledTask.reportCurrentTime();
-
-    expect:
-    assert scheduledTask.reportCurrentTimeExecuted == true;
-  }
-
-  def "trace fired"() {
-    setup:
-    def scheduledTask = new ScheduledTasks();
-    scheduledTask.reportCurrentTime();
-
-    expect:
-    assertTraces(1) {
-      trace(0, 1) {
-        span(0) {
-          serviceName "test"
-          resourceName "ScheduledTasks.reportCurrentTime"
-          operationName "trace.annotation"
-          parent()
-          errored false
-          tags {
-//            "$Tags.COMPONENT" "scheduled"
-            "$Tags.COMPONENT" "trace"
-            defaultTags()
-          }
-        }
-      }
-    }
-  }
-
-  def "test with context"() {
-    setup:
-    def context = new AnnotationConfigApplicationContext(ScheduledTasksConfig) // provide config as argument
-    def task = context.getBean(ScheduledTasks) // provide class we are trying to test
-    task.reportCurrentTime()
-
-    expect:
-    assertTraces(1) {
-      trace(0, 1) {
-        span(0) {
-          serviceName "test"
-          resourceName "ScheduledTasks.reportCurrentTime"
-          operationName "trace.annotation"
-          parent()
-          errored false
-          tags {
-//            "$Tags.COMPONENT" "scheduled"
-            "$Tags.COMPONENT" "trace"
-            defaultTags()
-          }
-        }
-      }
-    }
-  }
-}

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

@@ -0,0 +1,60 @@
+import datadog.trace.agent.test.AgentTestRunner
+import datadog.trace.instrumentation.api.Tags
+import org.springframework.context.annotation.AnnotationConfigApplicationContext
+
+class SpringSchedulingTest extends AgentTestRunner {
+
+  def "schedule interval test"() {
+    setup:
+    def context = new AnnotationConfigApplicationContext(IntervalTaskConfig)
+    def task = context.getBean(IntervalTask)
+
+    TEST_WRITER.clear()
+
+    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()
+          }
+        }
+      }
+    }
+  }
+
+  def "schedule trigger test according to cron expression"() {
+    setup:
+    def context = new AnnotationConfigApplicationContext(TriggerTaskConfig)
+    def task = context.getBean(TriggerTask)
+
+    TEST_WRITER.clear()
+
+    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()
+          }
+        }
+      }
+    }
+  }
+}

dd-java-agent/instrumentation/spring-scheduling/src/test/java/IntervalTask.java

@@ -0,0 +1,41 @@
+/*
+ * Copyright 2012-2015 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// import datadog.trace.api.DDTags;
+// import datadog.trace.instrumentation.api.AgentScope;
+// import datadog.trace.instrumentation.api.AgentSpan;
+
+import org.springframework.scheduling.annotation.Scheduled;
+import org.springframework.stereotype.Component;
+
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+
+@Component
+public class IntervalTask implements Runnable {
+
+  private final CountDownLatch latch = new CountDownLatch(1);
+
+  @Scheduled(initialDelay = 2, fixedRate = 5000)
+  @Override
+  public void run() {
+    latch.countDown();
+  }
+
+  public void blockUntilExecute() throws InterruptedException {
+    latch.await(5, TimeUnit.SECONDS);
+  }
+}

dd-java-agent/instrumentation/spring-scheduling/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/src/test/java/ScheduledTasks.java

@@ -1,72 +0,0 @@
-/*
- * Copyright 2012-2015 the original author or authors.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *    https://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-import datadog.trace.api.DDTags;
-import datadog.trace.instrumentation.api.AgentScope;
-import datadog.trace.instrumentation.api.AgentSpan;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-import org.springframework.scheduling.annotation.Scheduled;
-import org.springframework.stereotype.Component;
-
-import java.text.SimpleDateFormat;
-
-import static datadog.trace.instrumentation.api.AgentTracer.activateSpan;
-import static datadog.trace.instrumentation.api.AgentTracer.activeSpan;
-import static datadog.trace.instrumentation.api.AgentTracer.startSpan;
-import static datadog.trace.instrumentation.trace_annotation.TraceDecorator.DECORATE;
-
-@Component
-public class ScheduledTasks {
-
-  private static final Logger log = LoggerFactory.getLogger(ScheduledTasks.class);
-
-  private static final SimpleDateFormat dateFormat = new SimpleDateFormat("HH:mm:ss");
-
-  public static boolean reportCurrentTimeExecuted = false;
-
-  @Scheduled(fixedRate = 5000)
-  //  @Trace
-  public void reportCurrentTime() {
-    //    log.info("The time is now {}", dateFormat.format(new Date()));
-    // test that the body of method has been executed
-    // create span
-
-    //    activeSpan().setTag(DDTags.SERVICE_NAME, "test");
-    reportCurrentTimeExecuted = true;
-  }
-
-  public void runSpan() {
-    // create span
-    final AgentSpan span = startSpan("currentTime");
-    DECORATE.afterStart(span);
-
-    try (final AgentScope scope = activateSpan(span, false)) {
-      activeSpan().setTag(DDTags.SERVICE_NAME, "test");
-      DECORATE.afterStart(span);
-      scope.setAsyncPropagation(true);
-
-      try {
-        reportCurrentTime();
-      } catch (final Throwable throwable) {
-        DECORATE.onError(span, throwable);
-        DECORATE.beforeFinish(span);
-        span.finish();
-        throw throwable;
-      }
-    }
-  }
-}

dd-java-agent/instrumentation/spring-scheduling/src/test/java/ScheduledTasksConfig.java

@@ -1,10 +0,0 @@
-import org.springframework.context.annotation.Bean;
-import org.springframework.context.annotation.Configuration;
-
-@Configuration
-public class ScheduledTasksConfig {
-  @Bean
-  public ScheduledTasks scheduledTasks() {
-    return new ScheduledTasks();
-  }
-}

dd-java-agent/instrumentation/spring-scheduling/src/test/java/TriggerTask.java

@@ -0,0 +1,44 @@
+import org.springframework.scheduling.annotation.Scheduled;
+import org.springframework.stereotype.Component;
+
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+
+/*
+ * Copyright 2012-2015 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+// import datadog.trace.api.DDTags;
+// import datadog.trace.instrumentation.api.AgentScope;
+// import datadog.trace.instrumentation.api.AgentSpan;
+
+@Component
+public class TriggerTask implements Runnable {
+
+  public TriggerTask() {
+    System.out.println("in the ScheduledTasks constructor");
+  }
+
+  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/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();
+  }
+}