Skip to content

Commit 0f2d96f

Browse files
amarzialiamarziali
authored
Resolve right runnable name for spring scheduling tasks (DataDog#6249)
1 parent acd82c6 commit 0f2d96f

4 files changed

Lines changed: 88 additions & 10 deletions

File tree

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

Lines changed: 39 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,35 +1,72 @@
1-
1+
ext {
2+
spring6TestMinJavaVersionForTests = JavaVersion.VERSION_17
3+
spring6LatestDepTestMinJavaVersionForTests = JavaVersion.VERSION_17
4+
}
25
muzzle {
36
pass {
47
group = 'org.springframework'
58
module = 'spring-context'
69
versions = "[3.1.0.RELEASE,6)"
710
// assertInverse = true
811
}
12+
pass {
13+
group = 'org.springframework'
14+
module = 'spring-context'
15+
versions = "[6,)"
16+
javaVersion = "17"
17+
// assertInverse = true
18+
}
919
}
1020

1121
apply from: "$rootDir/gradle/java.gradle"
1222

1323
addTestSuiteForDir('latestDepTest', 'test')
24+
addTestSuiteForDir('spring6Test', 'test')
25+
addTestSuiteExtendingForDir('spring6LatestDepTest', 'latestDepTest', 'test')
26+
27+
[compileSpring6TestJava, compileSpring6LatestDepTestJava].each {
28+
setJavaVersion(it, 17)
29+
sourceCompatibility = JavaVersion.VERSION_1_8
30+
targetCompatibility = JavaVersion.VERSION_1_8
31+
}
32+
33+
34+
[compileSpring6TestGroovy, compileSpring6LatestDepTestGroovy, spring6Test, spring6LatestDepTest].each {
35+
it.javaLauncher = getJavaLauncherFor(17)
36+
}
37+
38+
[spring6Test, spring6LatestDepTest].each {
39+
it.jvmArgs '--add-opens', 'java.base/java.util=ALL-UNNAMED'
40+
}
1441

1542
dependencies {
1643
// choose a recent version so that we can test both lambdas (JDK8)
1744
// @Async requires proxying and older versions can't read classfile versions > 51
1845
// we muzzle older versions of spring anyway
1946
compileOnly group: 'org.springframework', name: 'spring-context', version: '5.0.0.RELEASE'
20-
testImplementation group: 'org.springframework', name: 'spring-context', version: '5.0.0.RELEASE'
2147

48+
testImplementation group: 'org.springframework', name: 'spring-context', version: '5.0.0.RELEASE'
2249
testImplementation project(':dd-java-agent:instrumentation:trace-annotation')
2350

2451
testImplementation group: 'net.javacrumbs.shedlock', name: 'shedlock-spring', version: '4.21.0'
2552
testImplementation group: 'net.javacrumbs.shedlock', name: 'shedlock-provider-jdbc-template', version: '4.21.0'
2653
testImplementation group: 'com.h2database', name: 'h2', version: '1.4.199'
2754
testImplementation group: 'org.springframework.boot', name: 'spring-boot-starter-data-jpa', version: '2.4.0'
55+
testImplementation group: 'org.springframework.boot', name: 'spring-boot-starter-actuator', version: '2.4.0'
56+
2857

2958
latestDepTestImplementation group: 'org.springframework', name: 'spring-context', version: '5.+'
3059

3160
latestDepTestImplementation group: 'net.javacrumbs.shedlock', name: 'shedlock-spring', version: '4.+'
3261
latestDepTestImplementation group: 'net.javacrumbs.shedlock', name: 'shedlock-provider-jdbc-template', version: '4.+'
3362
latestDepTestImplementation group: 'com.h2database', name: 'h2', version: '+'
3463
latestDepTestImplementation group: 'org.springframework.boot', name: 'spring-boot-starter-data-jpa', version: '2.+'
64+
65+
spring6TestImplementation group: 'org.springframework', name: 'spring-context', version: '6.0.0.RELEASE'
66+
spring6TestImplementation group: 'org.springframework.boot', name: 'spring-boot-starter-data-jpa', version: '3.0.0'
67+
68+
spring6LatestDepTestImplementation group: 'org.springframework', name: 'spring-context', version: '6.+'
69+
spring6LatestDepTestImplementation group: 'org.springframework.boot', name: 'spring-boot-starter-data-jpa', version: '3.+'
70+
71+
3572
}

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

Lines changed: 32 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,23 +1,29 @@
11
package datadog.trace.instrumentation.springscheduling;
22

3-
import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface;
4+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.nameStartsWith;
5+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
6+
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
47
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
58

69
import com.google.auto.service.AutoService;
710
import datadog.trace.agent.tooling.Instrumenter;
11+
import datadog.trace.bootstrap.CallDepthThreadLocalMap;
812
import net.bytebuddy.asm.Advice;
13+
import net.bytebuddy.description.type.TypeDescription;
14+
import net.bytebuddy.matcher.ElementMatcher;
15+
import org.springframework.scheduling.TaskScheduler;
916

1017
@AutoService(Instrumenter.class)
1118
public final class SpringSchedulingInstrumentation extends Instrumenter.Tracing
12-
implements Instrumenter.ForSingleType {
19+
implements Instrumenter.ForTypeHierarchy {
1320

1421
public SpringSchedulingInstrumentation() {
1522
super("spring-scheduling");
1623
}
1724

18-
@Override
19-
public String instrumentedType() {
20-
return "org.springframework.scheduling.config.Task";
25+
public String hierarchyMarkerType() {
26+
return "org.springframework.scheduling.TaskScheduler";
2127
}
2228

2329
@Override
@@ -27,18 +33,36 @@ public String[] helperClassNames() {
2733
};
2834
}
2935

36+
@Override
37+
public ElementMatcher<TypeDescription> hierarchyMatcher() {
38+
return implementsInterface(named(hierarchyMarkerType()));
39+
}
40+
3041
@Override
3142
public void adviceTransformations(AdviceTransformation transformation) {
3243
transformation.applyAdvice(
33-
isConstructor().and(takesArgument(0, Runnable.class)),
34-
SpringSchedulingInstrumentation.class.getName() + "$SpringSchedulingAdvice");
44+
isMethod().and(nameStartsWith("schedule")).and(takesArgument(0, Runnable.class)),
45+
getClass().getName() + "$SpringSchedulingAdvice");
3546
}
3647

3748
public static class SpringSchedulingAdvice {
3849
@Advice.OnMethodEnter(suppress = Throwable.class)
39-
public static void onConstruction(
50+
public static boolean beforeSchedule(
4051
@Advice.Argument(value = 0, readOnly = false) Runnable runnable) {
52+
if (CallDepthThreadLocalMap.incrementCallDepth(TaskScheduler.class) > 0) {
53+
return false;
54+
}
4155
runnable = SpringSchedulingRunnableWrapper.wrapIfNeeded(runnable);
56+
return true;
57+
}
58+
59+
@Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
60+
public static void afterSchedule(
61+
@Advice.Enter boolean reset,
62+
@Advice.Argument(value = 0, readOnly = false) Runnable runnable) {
63+
if (reset) {
64+
CallDepthThreadLocalMap.reset(TaskScheduler.class);
65+
}
4266
}
4367
}
4468
}

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import datadog.trace.agent.test.AgentTestRunner
22
import datadog.trace.bootstrap.instrumentation.api.Tags
3+
import org.springframework.boot.actuate.scheduling.ScheduledTasksEndpoint
34
import org.springframework.context.annotation.AnnotationConfigApplicationContext
45

56
import java.util.concurrent.TimeUnit
@@ -29,6 +30,12 @@ class SpringSchedulingTest extends AgentTestRunner {
2930
}
3031
}
3132
}
33+
and:
34+
def scheduledTaskEndpoint = context.getBean(ScheduledTasksEndpoint)
35+
assert scheduledTaskEndpoint != null
36+
scheduledTaskEndpoint.scheduledTasks().getCron().each {
37+
it.getRunnable().getTarget() == TriggerTask.getName()
38+
}
3239
cleanup:
3340
context.close()
3441
}

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,10 @@
1+
import java.util.stream.Collectors;
2+
import org.springframework.beans.factory.ObjectProvider;
3+
import org.springframework.boot.actuate.scheduling.ScheduledTasksEndpoint;
14
import org.springframework.context.annotation.Bean;
25
import org.springframework.context.annotation.Configuration;
36
import org.springframework.scheduling.annotation.EnableScheduling;
7+
import org.springframework.scheduling.config.ScheduledTaskHolder;
48

59
@Configuration
610
@EnableScheduling
@@ -9,4 +13,10 @@ public class TriggerTaskConfig {
913
public TriggerTask triggerTasks() {
1014
return new TriggerTask();
1115
}
16+
17+
@Bean
18+
public ScheduledTasksEndpoint scheduledTasksEndpoint(
19+
ObjectProvider<ScheduledTaskHolder> holders) {
20+
return new ScheduledTasksEndpoint(holders.orderedStream().collect(Collectors.toSet()));
21+
}
1222
}

0 commit comments

Comments
 (0)