Skip to content

Commit c7b4468

Browse files
monosoulmonosoul
committed
Add a naive Kotlin coroutines instrumentation
Signed-off-by: monosoul <[email protected]>
1 parent df6d722 commit c7b4468

5 files changed

Lines changed: 224 additions & 0 deletions

File tree

‎dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/core/NoOpContextElement.java‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
package datadog.trace.core;
2+
3+
import kotlin.coroutines.CoroutineContext;
4+
import kotlin.jvm.functions.Function2;
5+
import kotlinx.coroutines.ThreadContextElement;
6+
import org.jetbrains.annotations.NotNull;
7+
import org.jetbrains.annotations.Nullable;
8+
9+
public class NoOpContextElement implements ThreadContextElement<Object> {
10+
static final Key<NoOpContextElement> KEY = new Key<NoOpContextElement>() {};
11+
12+
public NoOpContextElement() {}
13+
14+
@Override
15+
public void restoreThreadContext(@NotNull CoroutineContext coroutineContext, Object oldState) {}
16+
17+
@Override
18+
public Object updateThreadContext(@NotNull CoroutineContext coroutineContext) {
19+
return null;
20+
}
21+
22+
@Nullable
23+
@Override
24+
public <E extends Element> E get(@NotNull Key<E> key) {
25+
return CoroutineContext.Element.DefaultImpls.get(this, key);
26+
}
27+
28+
@NotNull
29+
@Override
30+
public CoroutineContext minusKey(@NotNull Key<?> key) {
31+
return CoroutineContext.Element.DefaultImpls.minusKey(this, key);
32+
}
33+
34+
@NotNull
35+
@Override
36+
public CoroutineContext plus(@NotNull CoroutineContext coroutineContext) {
37+
return CoroutineContext.DefaultImpls.plus(this, coroutineContext);
38+
}
39+
40+
@Override
41+
public <R> R fold(
42+
R initial, @NotNull Function2<? super R, ? super Element, ? extends R> operation) {
43+
return CoroutineContext.Element.DefaultImpls.fold(this, initial, operation);
44+
}
45+
46+
@NotNull
47+
@Override
48+
public Key<?> getKey() {
49+
return KEY;
50+
}
51+
}

‎dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/core/ScopeStackCoroutineContextHelper.java‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
package datadog.trace.core;
2+
3+
import datadog.trace.bootstrap.instrumentation.api.AgentScopeManager;
4+
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
5+
import datadog.trace.core.scopemanager.ContinuableScopeManager;
6+
import datadog.trace.core.scopemanager.ScopeStackCoroutineContext;
7+
import kotlin.coroutines.CoroutineContext;
8+
import org.slf4j.Logger;
9+
import org.slf4j.LoggerFactory;
10+
11+
public class ScopeStackCoroutineContextHelper {
12+
13+
private static final NoOpContextElement NO_OP_CONTEXT_ELEMENT = new NoOpContextElement();
14+
private static final Logger logger =
15+
LoggerFactory.getLogger(ScopeStackCoroutineContextHelper.class);
16+
17+
public static CoroutineContext addScopeStackContext(final CoroutineContext other) {
18+
final AgentTracer.TracerAPI agentTracer = AgentTracer.get();
19+
final AgentScopeManager agentScopeManager =
20+
agentTracer instanceof CoreTracer ? ((CoreTracer) agentTracer).scopeManager : null;
21+
22+
if (agentScopeManager instanceof ContinuableScopeManager) {
23+
return other.plus(
24+
new ScopeStackCoroutineContext((ContinuableScopeManager) agentScopeManager));
25+
}
26+
27+
logger.warn(
28+
"Unexpected Tracer or Scope Manager implementation. Tracer[expected={}, got={}], ScopeManager[expected={}, got={}]",
29+
CoreTracer.class.getName(),
30+
agentTracer.getClass().getName(),
31+
ContinuableScopeManager.class.getName(),
32+
agentScopeManager != null ? agentScopeManager.getClass() : "null");
33+
34+
return NO_OP_CONTEXT_ELEMENT;
35+
}
36+
}

‎dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/core/scopemanager/ScopeStackCoroutineContext.java‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
package datadog.trace.core.scopemanager;
2+
3+
import static datadog.trace.bootstrap.instrumentation.api.ScopeSource.INSTRUMENTATION;
4+
5+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
6+
import datadog.trace.core.scopemanager.ContinuableScopeManager.ScopeStack;
7+
import kotlin.coroutines.CoroutineContext;
8+
import kotlin.jvm.functions.Function2;
9+
import kotlinx.coroutines.ThreadContextElement;
10+
import org.jetbrains.annotations.NotNull;
11+
import org.jetbrains.annotations.Nullable;
12+
13+
public class ScopeStackCoroutineContext implements ThreadContextElement<ScopeStack> {
14+
15+
static final Key<ScopeStackCoroutineContext> KEY = new Key<ScopeStackCoroutineContext>() {};
16+
private final ContinuableScopeManager scopeManager;
17+
private final ScopeStack scopeStack;
18+
@Nullable private final AgentSpan span;
19+
20+
public ScopeStackCoroutineContext(ContinuableScopeManager scopeManager) {
21+
this.scopeManager = scopeManager;
22+
this.span = scopeManager.activeSpan();
23+
/*
24+
* initial scope stack for the context element should be empty to prevent the spans created within the coroutine
25+
* from having a wrong parent span
26+
*/
27+
this.scopeStack = scopeManager.tlsScopeStack.initialValue();
28+
}
29+
30+
@Override
31+
public void restoreThreadContext(
32+
@NotNull CoroutineContext coroutineContext, ScopeStack oldState) {
33+
scopeManager.tlsScopeStack.set(oldState);
34+
}
35+
36+
@Override
37+
public ScopeStack updateThreadContext(@NotNull CoroutineContext coroutineContext) {
38+
final ScopeStack oldScopeStack = scopeManager.tlsScopeStack.get();
39+
scopeManager.tlsScopeStack.set(scopeStack);
40+
41+
if (scopeStack.depth() == 0 && span != null) {
42+
/*
43+
* This is necessary for the spans created within the coroutine to properly inherit the spans hierarchy.
44+
* It's not necessary to close the scope created here as it will be destroyed along with the context element and
45+
* the scope stack.
46+
*/
47+
scopeManager.activate(span, INSTRUMENTATION);
48+
}
49+
50+
return oldScopeStack;
51+
}
52+
53+
@Nullable
54+
@Override
55+
public <E extends Element> E get(@NotNull Key<E> key) {
56+
return CoroutineContext.Element.DefaultImpls.get(this, key);
57+
}
58+
59+
@NotNull
60+
@Override
61+
public CoroutineContext minusKey(@NotNull Key<?> key) {
62+
return CoroutineContext.Element.DefaultImpls.minusKey(this, key);
63+
}
64+
65+
@NotNull
66+
@Override
67+
public CoroutineContext plus(@NotNull CoroutineContext coroutineContext) {
68+
return CoroutineContext.DefaultImpls.plus(this, coroutineContext);
69+
}
70+
71+
@Override
72+
public <R> R fold(
73+
R initial, @NotNull Function2<? super R, ? super Element, ? extends R> operation) {
74+
return CoroutineContext.Element.DefaultImpls.fold(this, initial, operation);
75+
}
76+
77+
@NotNull
78+
@Override
79+
public Key<?> getKey() {
80+
return KEY;
81+
}
82+
}

‎dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/CoroutineContextAdvice.java‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
package datadog.trace.instrumentation.kotlin.coroutines;
2+
3+
import datadog.trace.core.ScopeStackCoroutineContextHelper;
4+
import kotlin.coroutines.CoroutineContext;
5+
import net.bytebuddy.asm.Advice;
6+
7+
public class CoroutineContextAdvice {
8+
@Advice.OnMethodEnter
9+
public static void enter(
10+
@Advice.Argument(value = 1, readOnly = false) CoroutineContext coroutineContext) {
11+
if (coroutineContext != null) {
12+
coroutineContext = ScopeStackCoroutineContextHelper.addScopeStackContext(coroutineContext);
13+
}
14+
}
15+
}

‎dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/KotlinCoroutinesInstrumentation.java‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
package datadog.trace.instrumentation.kotlin.coroutines;
2+
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4+
import static datadog.trace.util.Strings.getPackageName;
5+
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
6+
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
7+
8+
import com.google.auto.service.AutoService;
9+
import datadog.trace.agent.tooling.Instrumenter;
10+
import datadog.trace.core.CoreTracer;
11+
12+
@AutoService(Instrumenter.class)
13+
public class KotlinCoroutinesInstrumentation extends Instrumenter.Tracing
14+
implements Instrumenter.ForSingleType {
15+
16+
public KotlinCoroutinesInstrumentation() {
17+
super("kotlin-coroutines");
18+
}
19+
20+
@Override
21+
public String[] helperClassNames() {
22+
return new String[] {
23+
getPackageName(CoreTracer.class.getName()) + ".ScopeStackCoroutineContextHelper",
24+
};
25+
}
26+
27+
@Override
28+
public String instrumentedType() {
29+
return "kotlinx.coroutines.CoroutineContextKt";
30+
}
31+
32+
@Override
33+
public void adviceTransformations(AdviceTransformation transformation) {
34+
transformation.applyAdvice(
35+
isMethod()
36+
.and(named("newCoroutineContext"))
37+
.and(takesArgument(1, named("kotlin.coroutines.CoroutineContext"))),
38+
packageName + ".CoroutineContextAdvice");
39+
}
40+
}

0 commit comments

Comments
 (0)