|
1 | 1 | package datadog.trace.instrumentation.kotlin.coroutines; |
2 | 2 |
|
| 3 | +import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.extendsClass; |
3 | 4 | import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; |
4 | 5 | import static net.bytebuddy.matcher.ElementMatchers.isMethod; |
| 6 | +import static net.bytebuddy.matcher.ElementMatchers.takesNoArguments; |
5 | 7 |
|
6 | 8 | import datadog.trace.agent.tooling.Instrumenter; |
7 | 9 | import kotlinx.coroutines.AbstractCoroutine; |
8 | 10 | import net.bytebuddy.asm.Advice; |
| 11 | +import net.bytebuddy.description.type.TypeDescription; |
| 12 | +import net.bytebuddy.matcher.ElementMatcher; |
9 | 13 |
|
10 | 14 | /** Captures the Datadog context when lazy coroutines start. */ |
11 | 15 | public class LazyCoroutineInstrumentation |
12 | | - implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice { |
| 16 | + implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice { |
13 | 17 |
|
14 | 18 | @Override |
15 | | - public String instrumentedType() { |
16 | | - return "kotlinx.coroutines.LazyDeferredCoroutine"; |
| 19 | + public String hierarchyMarkerType() { |
| 20 | + return "kotlinx.coroutines.AbstractCoroutine"; |
| 21 | + } |
| 22 | + |
| 23 | + @Override |
| 24 | + public ElementMatcher<TypeDescription> hierarchyMatcher() { |
| 25 | + return extendsClass(named(hierarchyMarkerType())); |
17 | 26 | } |
18 | 27 |
|
19 | 28 | @Override |
20 | 29 | public void methodAdvice(MethodTransformer transformer) { |
21 | 30 | transformer.applyAdvice( |
22 | | - isMethod().and(named("onStart")), |
| 31 | + isMethod().and(named("onStart")).and(takesNoArguments()), |
23 | 32 | LazyCoroutineInstrumentation.class.getName() + "$OnStartAdvice"); |
24 | 33 | } |
25 | 34 |
|
|
0 commit comments