Skip to content

Commit 27e144e

Browse files
monosoulmonosoul
committed
Make sure maybeCloseScopeAndCancelContinuation() is always invoked when coroutine completes
Signed-off-by: monosoul <[email protected]>
1 parent 48922a7 commit 27e144e

4 files changed

Lines changed: 54 additions & 114 deletions

File tree

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

Lines changed: 50 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,9 +21,33 @@
2121
import net.bytebuddy.matcher.ElementMatcher;
2222

2323
@AutoService(Instrumenter.class)
24-
public class AbstractCoroutineInstrumentation extends AbstractKotlinCoroutinesInstrumentation
24+
public class AbstractCoroutineInstrumentation extends Instrumenter.Tracing
2525
implements Instrumenter.ForTypeHierarchy {
2626

27+
private static final String ABSTRACT_COROUTINE_CLASS_NAME =
28+
"kotlinx.coroutines.AbstractCoroutine";
29+
30+
private static final String JOB_SUPPORT_CLASS_NAME = "kotlinx.coroutines.JobSupport";
31+
private static final String COROUTINE_CONTEXT_CLASS_NAME = "kotlin.coroutines.CoroutineContext";
32+
33+
public AbstractCoroutineInstrumentation() {
34+
super("kotlin_coroutine.experimental");
35+
}
36+
37+
@Override
38+
protected final boolean defaultEnabled() {
39+
return false;
40+
}
41+
42+
@Override
43+
public String[] helperClassNames() {
44+
return new String[] {
45+
packageName + ".ScopeStateCoroutineContext",
46+
packageName + ".ScopeStateCoroutineContext$ContextElementKey",
47+
packageName + ".CoroutineContextHelper",
48+
};
49+
}
50+
2751
@Override
2852
public void adviceTransformations(AdviceTransformation transformation) {
2953
transformation.applyAdvice(
@@ -41,6 +65,15 @@ public void adviceTransformations(AdviceTransformation transformation) {
4165
.and(takesNoArguments())
4266
.and(returns(void.class)),
4367
AbstractCoroutineInstrumentation.class.getName() + "$AbstractCoroutineOnStartAdvice");
68+
69+
transformation.applyAdvice(
70+
isMethod()
71+
.and(isOverriddenFrom(named(JOB_SUPPORT_CLASS_NAME)))
72+
.and(named("onCompletionInternal"))
73+
.and(takesArguments(1))
74+
.and(returns(void.class)),
75+
AbstractCoroutineInstrumentation.class.getName()
76+
+ "$JobSupportAfterCompletionInternalAdvice");
4477
}
4578

4679
@Override
@@ -56,10 +89,11 @@ public ElementMatcher<TypeDescription> hierarchyMatcher() {
5689
public static class AbstractCoroutineConstructorAdvice {
5790
@Advice.OnMethodEnter
5891
public static void constructorInvocation(
59-
@Advice.Argument(value = 0) final CoroutineContext parentContext,
92+
@Advice.Argument(value = 0, readOnly = false) CoroutineContext parentContext,
6093
@Advice.Argument(value = 1) final boolean active) {
61-
final ScopeStateCoroutineContext scopeStackContext = getScopeStateContext(parentContext);
62-
if (scopeStackContext != null && active) {
94+
final ScopeStateCoroutineContext scopeStackContext = new ScopeStateCoroutineContext();
95+
parentContext = parentContext.plus(scopeStackContext);
96+
if (active) {
6397
// if this is not a lazy coroutine, inherit parent span from the coroutine constructor call
6498
// site
6599
scopeStackContext.maybeInitialize();
@@ -78,4 +112,16 @@ public static void onStartInvocation(@Advice.This final AbstractCoroutine<?> cor
78112
}
79113
}
80114
}
115+
116+
public static class JobSupportAfterCompletionInternalAdvice {
117+
@Advice.OnMethodEnter
118+
public static void onCompletionInternal(@Advice.This final AbstractCoroutine<?> coroutine) {
119+
final ScopeStateCoroutineContext scopeStackContext =
120+
getScopeStateContext(coroutine.getContext());
121+
if (scopeStackContext != null) {
122+
// close the scope if needed
123+
scopeStackContext.maybeCloseScopeAndCancelContinuation();
124+
}
125+
}
126+
}
81127
}

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

Lines changed: 0 additions & 36 deletions
This file was deleted.

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

Lines changed: 0 additions & 44 deletions
This file was deleted.

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

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

3-
import static datadog.trace.instrumentation.kotlin.coroutines.AbstractKotlinCoroutinesInstrumentation.prefixedPropertyName;
4-
import static datadog.trace.instrumentation.kotlin.coroutines.CoroutineContextHelper.getJob;
5-
import static java.lang.Boolean.getBoolean;
6-
73
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
84
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
95
import datadog.trace.bootstrap.instrumentation.api.ScopeState;
10-
import kotlin.Unit;
116
import kotlin.coroutines.CoroutineContext;
12-
import kotlin.jvm.functions.Function1;
137
import kotlin.jvm.functions.Function2;
148
import kotlinx.coroutines.ThreadContextElement;
159
import org.jetbrains.annotations.NotNull;
1610
import org.jetbrains.annotations.Nullable;
1711

18-
public class ScopeStateCoroutineContext
19-
implements ThreadContextElement<ScopeState>, Function1<Throwable, Unit> {
12+
public class ScopeStateCoroutineContext implements ThreadContextElement<ScopeState> {
2013

2114
public static final Key<ScopeStateCoroutineContext> KEY = new ContextElementKey();
2215
private final ScopeState coroutineScopeState;
@@ -31,15 +24,8 @@ public ScopeStateCoroutineContext() {
3124
public void maybeInitialize() {
3225
if (!isInitialized) {
3326
final AgentScope activeScope = AgentTracer.get().activeScope();
34-
final boolean autoPropagationEnabled =
35-
getBoolean(prefixedPropertyName("auto_propagation.enabled"));
36-
if (activeScope != null) {
37-
if (autoPropagationEnabled) {
38-
activeScope.setAsyncPropagation(true);
39-
}
40-
if (activeScope.isAsyncPropagating()) {
41-
continuation = activeScope.captureConcurrent();
42-
}
27+
if (activeScope != null && activeScope.isAsyncPropagating()) {
28+
continuation = activeScope.captureConcurrent();
4329
}
4430
isInitialized = true;
4531
}
@@ -60,7 +46,6 @@ public ScopeState updateThreadContext(@NotNull final CoroutineContext coroutineC
6046

6147
if (continuation != null && continuationScope == null) {
6248
continuationScope = continuation.activate();
63-
registerOnCompletionCallback(coroutineContext);
6449
}
6550

6651
return oldScopeState;
@@ -96,11 +81,7 @@ public Key<?> getKey() {
9681
return KEY;
9782
}
9883

99-
public void registerOnCompletionCallback(final CoroutineContext coroutineContext) {
100-
getJob(coroutineContext).invokeOnCompletion(this);
101-
}
102-
103-
private void maybeCloseScopeAndCancelContinuation() {
84+
public void maybeCloseScopeAndCancelContinuation() {
10485
final ScopeState currentThreadScopeState = AgentTracer.get().newScopeState();
10586
currentThreadScopeState.fetchFromActive();
10687

@@ -116,12 +97,5 @@ private void maybeCloseScopeAndCancelContinuation() {
11697
currentThreadScopeState.activate();
11798
}
11899

119-
@Override
120-
public Unit invoke(Throwable throwable) {
121-
maybeCloseScopeAndCancelContinuation();
122-
123-
return Unit.INSTANCE;
124-
}
125-
126100
static class ContextElementKey implements Key<ScopeStateCoroutineContext> {}
127101
}

0 commit comments

Comments
 (0)