Skip to content

Commit abff8f0

Browse files
mccullsmcculls
committed
Migrate kotlin coroutines instrumentation to new Context API (#8528)
1 parent b829739 commit abff8f0

File tree

5 files changed

+69
-75
lines changed

5 files changed

+69
-75
lines changed

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

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
package datadog.trace.instrumentation.kotlin.coroutines;
22

33
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4-
import static datadog.trace.instrumentation.kotlin.coroutines.CoroutineContextHelper.getScopeStateContext;
5-
import static datadog.trace.instrumentation.kotlin.coroutines.CoroutineContextHelper.initializeScopeStateContextIfActive;
4+
import static datadog.trace.instrumentation.kotlin.coroutines.CoroutineContextHelper.getDatadogContext;
5+
import static datadog.trace.instrumentation.kotlin.coroutines.CoroutineContextHelper.initializeDatadogContextIfActive;
66
import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
77
import static net.bytebuddy.matcher.ElementMatchers.isDeclaredBy;
88
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
@@ -11,7 +11,7 @@
1111
import com.google.auto.service.AutoService;
1212
import datadog.trace.agent.tooling.InstrumenterModule;
1313
import datadog.trace.bootstrap.InstrumentationContext;
14-
import datadog.trace.instrumentation.kotlin.coroutines.ScopeStateCoroutineContext.ScopeStateCoroutineContextItem;
14+
import datadog.trace.instrumentation.kotlin.coroutines.DatadogCoroutineContext.DatadogCoroutineContextItem;
1515
import kotlin.coroutines.CoroutineContext;
1616
import kotlinx.coroutines.AbstractCoroutine;
1717
import kotlinx.coroutines.Job;
@@ -33,21 +33,21 @@ public void methodAdvice(MethodTransformer transformer) {
3333
}
3434

3535
/**
36-
* Guarantees every coroutine created has an instance of ScopeStateCoroutineContext
36+
* Guarantees every coroutine created has an instance of DatadogCoroutineContext
3737
*
38-
* @see ScopeStateCoroutineContext
38+
* @see DatadogCoroutineContext
3939
* @see AbstractCoroutine#AbstractCoroutine(CoroutineContext, boolean)
4040
*/
4141
public static class AbstractCoroutineConstructorAdvice {
4242
@Advice.OnMethodEnter
4343
public static void constructorInvocation(
4444
@Advice.Argument(value = 0, readOnly = false) CoroutineContext parentContext) {
45-
final ScopeStateCoroutineContext scopeStackContext = getScopeStateContext(parentContext);
46-
if (scopeStackContext == null) {
45+
final DatadogCoroutineContext datadogContext = getDatadogContext(parentContext);
46+
if (datadogContext == null) {
4747
parentContext =
4848
parentContext.plus(
49-
new ScopeStateCoroutineContext(
50-
InstrumentationContext.get(Job.class, ScopeStateCoroutineContextItem.class)));
49+
new DatadogCoroutineContext(
50+
InstrumentationContext.get(Job.class, DatadogCoroutineContextItem.class)));
5151
}
5252
}
5353

@@ -57,7 +57,7 @@ public static void constructorInvocationOnMethodExit(
5757
@Advice.Argument(value = 1) final boolean active) {
5858
// if this is not a lazy coroutine, inherit parent span from
5959
// the coroutine constructor call site
60-
initializeScopeStateContextIfActive(coroutine, active);
60+
initializeDatadogContextIfActive(coroutine, active);
6161
}
6262
}
6363
}

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

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
package datadog.trace.instrumentation.kotlin.coroutines;
22

33
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4-
import static datadog.trace.instrumentation.kotlin.coroutines.CoroutineContextHelper.getScopeStateContext;
5-
import static datadog.trace.instrumentation.kotlin.coroutines.CoroutineContextHelper.initializeScopeStateContextIfActive;
4+
import static datadog.trace.instrumentation.kotlin.coroutines.CoroutineContextHelper.getDatadogContext;
5+
import static datadog.trace.instrumentation.kotlin.coroutines.CoroutineContextHelper.initializeDatadogContextIfActive;
66
import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
77
import static net.bytebuddy.matcher.ElementMatchers.isDeclaredBy;
88
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
@@ -11,7 +11,7 @@
1111
import com.google.auto.service.AutoService;
1212
import datadog.trace.agent.tooling.InstrumenterModule;
1313
import datadog.trace.bootstrap.InstrumentationContext;
14-
import datadog.trace.instrumentation.kotlin.coroutines.ScopeStateCoroutineContext.ScopeStateCoroutineContextItem;
14+
import datadog.trace.instrumentation.kotlin.coroutines.DatadogCoroutineContext.DatadogCoroutineContextItem;
1515
import kotlin.coroutines.CoroutineContext;
1616
import kotlinx.coroutines.AbstractCoroutine;
1717
import kotlinx.coroutines.Job;
@@ -33,21 +33,21 @@ public void methodAdvice(MethodTransformer transformer) {
3333
}
3434

3535
/**
36-
* Guarantees every coroutine created has an instance of ScopeStateCoroutineContext
36+
* Guarantees every coroutine created has an instance of DatadogCoroutineContext
3737
*
38-
* @see ScopeStateCoroutineContext
38+
* @see DatadogCoroutineContext
3939
* @see AbstractCoroutine#AbstractCoroutine(CoroutineContext, boolean, boolean)
4040
*/
4141
public static class AbstractCoroutineConstructorAdvice {
4242
@Advice.OnMethodEnter
4343
public static void constructorInvocation(
4444
@Advice.Argument(value = 0, readOnly = false) CoroutineContext parentContext) {
45-
final ScopeStateCoroutineContext scopeStackContext = getScopeStateContext(parentContext);
46-
if (scopeStackContext == null) {
45+
final DatadogCoroutineContext datadogContext = getDatadogContext(parentContext);
46+
if (datadogContext == null) {
4747
parentContext =
4848
parentContext.plus(
49-
new ScopeStateCoroutineContext(
50-
InstrumentationContext.get(Job.class, ScopeStateCoroutineContextItem.class)));
49+
new DatadogCoroutineContext(
50+
InstrumentationContext.get(Job.class, DatadogCoroutineContextItem.class)));
5151
}
5252
}
5353

@@ -57,7 +57,7 @@ public static void constructorInvocationOnMethodExit(
5757
@Advice.Argument(value = 2) final boolean active) {
5858
// if this is not a lazy coroutine, inherit parent span from
5959
// the coroutine constructor call site
60-
initializeScopeStateContextIfActive(coroutine, active);
60+
initializeDatadogContextIfActive(coroutine, active);
6161
}
6262
}
6363
}

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

Lines changed: 14 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,8 @@
22

33
import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.extendsClass;
44
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
5-
import static datadog.trace.instrumentation.kotlin.coroutines.CoroutineContextHelper.closeScopeStateContext;
6-
import static datadog.trace.instrumentation.kotlin.coroutines.CoroutineContextHelper.initializeScopeStateContext;
5+
import static datadog.trace.instrumentation.kotlin.coroutines.CoroutineContextHelper.closeDatadogContext;
6+
import static datadog.trace.instrumentation.kotlin.coroutines.CoroutineContextHelper.initializeDatadogContext;
77
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
88
import static net.bytebuddy.matcher.ElementMatchers.isOverriddenFrom;
99
import static net.bytebuddy.matcher.ElementMatchers.returns;
@@ -42,9 +42,9 @@ protected final boolean defaultEnabled() {
4242
@Override
4343
public String[] helperClassNames() {
4444
return new String[] {
45-
packageName + ".ScopeStateCoroutineContext",
46-
packageName + ".ScopeStateCoroutineContext$ContextElementKey",
47-
packageName + ".ScopeStateCoroutineContext$ScopeStateCoroutineContextItem",
45+
packageName + ".DatadogCoroutineContext",
46+
packageName + ".DatadogCoroutineContext$ContextElementKey",
47+
packageName + ".DatadogCoroutineContext$DatadogCoroutineContextItem",
4848
packageName + ".CoroutineContextHelper",
4949
};
5050
}
@@ -84,36 +84,35 @@ public ElementMatcher<TypeDescription> structureMatcher() {
8484
public Map<String, String> contextStore() {
8585
return Collections.singletonMap(
8686
"kotlinx.coroutines.Job",
87-
packageName + ".ScopeStateCoroutineContext$ScopeStateCoroutineContextItem");
87+
packageName + ".DatadogCoroutineContext$DatadogCoroutineContextItem");
8888
}
8989

9090
/**
91-
* If/when coroutine is started lazily, initializes ScopeStateCoroutineContext element on
92-
* coroutine start
91+
* If/when coroutine is started lazily, initializes DatadogCoroutineContext element on coroutine
92+
* start
9393
*
94-
* @see ScopeStateCoroutineContext
94+
* @see DatadogCoroutineContext
9595
* @see AbstractCoroutine#onStart()
9696
*/
9797
public static class AbstractCoroutineOnStartAdvice {
9898
@Advice.OnMethodEnter
9999
public static void onStartInvocation(@Advice.This final AbstractCoroutine<?> coroutine) {
100100
// try to inherit parent span from the coroutine start call site
101-
initializeScopeStateContext(coroutine);
101+
initializeDatadogContext(coroutine);
102102
}
103103
}
104104

105105
/**
106-
* Guarantees a ScopeStateCoroutineContext element is always closed when coroutine transitions
107-
* into a terminal state.
106+
* Guarantees a DatadogCoroutineContext element is always closed when coroutine transitions into a
107+
* terminal state.
108108
*
109-
* @see ScopeStateCoroutineContext
109+
* @see DatadogCoroutineContext
110110
* @see AbstractCoroutine#onCompletionInternal(Object)
111111
*/
112112
public static class JobSupportAfterCompletionInternalAdvice {
113113
@Advice.OnMethodEnter
114114
public static void onCompletionInternal(@Advice.This final AbstractCoroutine<?> coroutine) {
115-
// close the scope if needed
116-
closeScopeStateContext(coroutine);
115+
closeDatadogContext(coroutine);
117116
}
118117
}
119118
}

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

Lines changed: 12 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -17,30 +17,28 @@ public static Job getJob(final CoroutineContext context) {
1717
}
1818

1919
@Nullable
20-
public static ScopeStateCoroutineContext getScopeStateContext(final CoroutineContext context) {
21-
return context.get(ScopeStateCoroutineContext.KEY);
20+
public static DatadogCoroutineContext getDatadogContext(final CoroutineContext context) {
21+
return context.get(DatadogCoroutineContext.KEY);
2222
}
2323

24-
public static void initializeScopeStateContextIfActive(
24+
public static void initializeDatadogContextIfActive(
2525
final AbstractCoroutine<?> coroutine, final boolean active) {
2626
if (active) {
27-
initializeScopeStateContext(coroutine);
27+
initializeDatadogContext(coroutine);
2828
}
2929
}
3030

31-
public static void initializeScopeStateContext(final AbstractCoroutine<?> coroutine) {
32-
final ScopeStateCoroutineContext scopeStackContext =
33-
getScopeStateContext(coroutine.getContext());
34-
if (scopeStackContext != null) {
35-
scopeStackContext.maybeInitialize(coroutine);
31+
public static void initializeDatadogContext(final AbstractCoroutine<?> coroutine) {
32+
final DatadogCoroutineContext datadogContext = getDatadogContext(coroutine.getContext());
33+
if (datadogContext != null) {
34+
datadogContext.maybeInitialize(coroutine);
3635
}
3736
}
3837

39-
public static void closeScopeStateContext(final AbstractCoroutine<?> coroutine) {
40-
final ScopeStateCoroutineContext scopeStackContext =
41-
getScopeStateContext(coroutine.getContext());
42-
if (scopeStackContext != null) {
43-
scopeStackContext.maybeCloseScopeAndCancelContinuation(coroutine);
38+
public static void closeDatadogContext(final AbstractCoroutine<?> coroutine) {
39+
final DatadogCoroutineContext datadogContext = getDatadogContext(coroutine.getContext());
40+
if (datadogContext != null) {
41+
datadogContext.maybeCloseScopeAndCancelContinuation(coroutine);
4442
}
4543
}
4644
}

‎dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/ScopeStateCoroutineContext.java‎ renamed to ‎dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/DatadogCoroutineContext.java‎

Lines changed: 23 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -2,66 +2,63 @@
22

33
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.captureActiveSpan;
44

5+
import datadog.context.Context;
56
import datadog.trace.bootstrap.ContextStore;
67
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
7-
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
8-
import datadog.trace.bootstrap.instrumentation.api.ScopeState;
98
import kotlin.coroutines.CoroutineContext;
109
import kotlin.jvm.functions.Function2;
1110
import kotlinx.coroutines.Job;
1211
import kotlinx.coroutines.ThreadContextElement;
1312
import org.jetbrains.annotations.NotNull;
1413
import org.jetbrains.annotations.Nullable;
1514

16-
public class ScopeStateCoroutineContext implements ThreadContextElement<ScopeState> {
15+
public class DatadogCoroutineContext implements ThreadContextElement<Context> {
1716

18-
public static final Key<ScopeStateCoroutineContext> KEY = new ContextElementKey();
17+
public static final Key<DatadogCoroutineContext> KEY = new ContextElementKey();
1918

20-
private final ContextStore<Job, ScopeStateCoroutineContextItem> contextItemPerCoroutine;
19+
private final ContextStore<Job, DatadogCoroutineContextItem> contextItemPerCoroutine;
2120

22-
public ScopeStateCoroutineContext(
23-
final ContextStore<Job, ScopeStateCoroutineContextItem> contextItemPerCoroutine) {
21+
public DatadogCoroutineContext(
22+
final ContextStore<Job, DatadogCoroutineContextItem> contextItemPerCoroutine) {
2423
this.contextItemPerCoroutine = contextItemPerCoroutine;
2524
}
2625

2726
/** Get a context item instance for the coroutine and try to initialize it */
2827
public void maybeInitialize(final Job coroutine) {
2928
contextItemPerCoroutine
30-
.putIfAbsent(coroutine, ScopeStateCoroutineContextItem::new)
29+
.putIfAbsent(coroutine, DatadogCoroutineContextItem::new)
3130
.maybeInitialize();
3231
}
3332

3433
@Override
3534
public void restoreThreadContext(
36-
@NotNull final CoroutineContext coroutineContext, final ScopeState oldState) {
37-
oldState.activate();
35+
@NotNull final CoroutineContext coroutineContext, final Context oldDatadogContext) {
36+
oldDatadogContext.swap();
3837
}
3938

4039
@Override
41-
public ScopeState updateThreadContext(@NotNull final CoroutineContext coroutineContext) {
42-
final ScopeState oldScopeState = AgentTracer.get().newScopeState();
43-
oldScopeState.fetchFromActive();
40+
public Context updateThreadContext(@NotNull final CoroutineContext coroutineContext) {
41+
final Context oldDatadogContext = Context.current();
4442

4543
final Job coroutine = CoroutineContextHelper.getJob(coroutineContext);
46-
final ScopeStateCoroutineContextItem contextItem = contextItemPerCoroutine.get(coroutine);
44+
final DatadogCoroutineContextItem contextItem = contextItemPerCoroutine.get(coroutine);
4745
if (contextItem != null) {
4846
contextItem.activate();
4947
}
5048

51-
return oldScopeState;
49+
return oldDatadogContext;
5250
}
5351

5452
/** If there's a context item for the coroutine then try to close it */
5553
public void maybeCloseScopeAndCancelContinuation(final Job coroutine) {
56-
final ScopeStateCoroutineContextItem contextItem = contextItemPerCoroutine.get(coroutine);
54+
final DatadogCoroutineContextItem contextItem = contextItemPerCoroutine.get(coroutine);
5755
if (contextItem != null) {
58-
final ScopeState currentThreadScopeState = AgentTracer.get().newScopeState();
59-
currentThreadScopeState.fetchFromActive();
56+
final Context currentDatadogContext = Context.current();
6057

6158
contextItem.maybeCloseScopeAndCancelContinuation();
6259
contextItemPerCoroutine.remove(coroutine);
6360

64-
currentThreadScopeState.activate();
61+
currentDatadogContext.swap();
6562
}
6663
}
6764

@@ -95,20 +92,20 @@ public Key<?> getKey() {
9592
return KEY;
9693
}
9794

98-
static class ContextElementKey implements Key<ScopeStateCoroutineContext> {}
95+
static class ContextElementKey implements Key<DatadogCoroutineContext> {}
9996

100-
public static class ScopeStateCoroutineContextItem {
101-
private final ScopeState coroutineScopeState;
97+
public static class DatadogCoroutineContextItem {
98+
private final Context datadogContext;
10299
@Nullable private AgentScope.Continuation continuation;
103100
@Nullable private AgentScope continuationScope;
104101
private boolean isInitialized = false;
105102

106-
public ScopeStateCoroutineContextItem() {
107-
coroutineScopeState = AgentTracer.get().newScopeState();
103+
public DatadogCoroutineContextItem() {
104+
datadogContext = Context.root();
108105
}
109106

110107
public void activate() {
111-
coroutineScopeState.activate();
108+
datadogContext.swap();
112109

113110
if (continuation != null && continuationScope == null) {
114111
continuationScope = continuation.activate();
@@ -131,7 +128,7 @@ public void maybeInitialize() {
131128
* scope and cancels the continuation.
132129
*/
133130
public void maybeCloseScopeAndCancelContinuation() {
134-
coroutineScopeState.activate();
131+
datadogContext.swap();
135132

136133
if (continuationScope != null) {
137134
continuationScope.close();

0 commit comments

Comments
 (0)