diff --git a/dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/core/NoOpContextElement.java b/dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/core/NoOpContextElement.java new file mode 100644 index 00000000000..eef125d85c7 --- /dev/null +++ b/dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/core/NoOpContextElement.java @@ -0,0 +1,51 @@ +package datadog.trace.core; + +import kotlin.coroutines.CoroutineContext; +import kotlin.jvm.functions.Function2; +import kotlinx.coroutines.ThreadContextElement; +import org.jetbrains.annotations.NotNull; +import org.jetbrains.annotations.Nullable; + +public class NoOpContextElement implements ThreadContextElement { + static final Key KEY = new Key() {}; + + public NoOpContextElement() {} + + @Override + public void restoreThreadContext(@NotNull CoroutineContext coroutineContext, Object oldState) {} + + @Override + public Object updateThreadContext(@NotNull CoroutineContext coroutineContext) { + return null; + } + + @Nullable + @Override + public E get(@NotNull Key key) { + return CoroutineContext.Element.DefaultImpls.get(this, key); + } + + @NotNull + @Override + public CoroutineContext minusKey(@NotNull Key key) { + return CoroutineContext.Element.DefaultImpls.minusKey(this, key); + } + + @NotNull + @Override + public CoroutineContext plus(@NotNull CoroutineContext coroutineContext) { + return CoroutineContext.DefaultImpls.plus(this, coroutineContext); + } + + @Override + public R fold( + R initial, @NotNull Function2 operation) { + return CoroutineContext.Element.DefaultImpls.fold(this, initial, operation); + } + + @NotNull + @Override + public Key getKey() { + return KEY; + } +} diff --git a/dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/core/ScopeStackCoroutineContextHelper.java b/dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/core/ScopeStackCoroutineContextHelper.java new file mode 100644 index 00000000000..13c504905a0 --- /dev/null +++ b/dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/core/ScopeStackCoroutineContextHelper.java @@ -0,0 +1,36 @@ +package datadog.trace.core; + +import datadog.trace.bootstrap.instrumentation.api.AgentScopeManager; +import datadog.trace.bootstrap.instrumentation.api.AgentTracer; +import datadog.trace.core.scopemanager.ContinuableScopeManager; +import datadog.trace.core.scopemanager.ScopeStackCoroutineContext; +import kotlin.coroutines.CoroutineContext; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +public class ScopeStackCoroutineContextHelper { + + private static final NoOpContextElement NO_OP_CONTEXT_ELEMENT = new NoOpContextElement(); + private static final Logger logger = + LoggerFactory.getLogger(ScopeStackCoroutineContextHelper.class); + + public static CoroutineContext addScopeStackContext(final CoroutineContext other) { + final AgentTracer.TracerAPI agentTracer = AgentTracer.get(); + final AgentScopeManager agentScopeManager = + agentTracer instanceof CoreTracer ? ((CoreTracer) agentTracer).scopeManager : null; + + if (agentScopeManager instanceof ContinuableScopeManager) { + return other.plus( + new ScopeStackCoroutineContext((ContinuableScopeManager) agentScopeManager)); + } + + logger.warn( + "Unexpected Tracer or Scope Manager implementation. Tracer[expected={}, got={}], ScopeManager[expected={}, got={}]", + CoreTracer.class.getName(), + agentTracer.getClass().getName(), + ContinuableScopeManager.class.getName(), + agentScopeManager != null ? agentScopeManager.getClass() : "null"); + + return NO_OP_CONTEXT_ELEMENT; + } +} diff --git a/dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/core/scopemanager/ScopeStackCoroutineContext.java b/dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/core/scopemanager/ScopeStackCoroutineContext.java new file mode 100644 index 00000000000..aefd9118f16 --- /dev/null +++ b/dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/core/scopemanager/ScopeStackCoroutineContext.java @@ -0,0 +1,82 @@ +package datadog.trace.core.scopemanager; + +import static datadog.trace.bootstrap.instrumentation.api.ScopeSource.INSTRUMENTATION; + +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.core.scopemanager.ContinuableScopeManager.ScopeStack; +import kotlin.coroutines.CoroutineContext; +import kotlin.jvm.functions.Function2; +import kotlinx.coroutines.ThreadContextElement; +import org.jetbrains.annotations.NotNull; +import org.jetbrains.annotations.Nullable; + +public class ScopeStackCoroutineContext implements ThreadContextElement { + + static final Key KEY = new Key() {}; + private final ContinuableScopeManager scopeManager; + private final ScopeStack scopeStack; + @Nullable private final AgentSpan span; + + public ScopeStackCoroutineContext(ContinuableScopeManager scopeManager) { + this.scopeManager = scopeManager; + this.span = scopeManager.activeSpan(); + /* + * initial scope stack for the context element should be empty to prevent the spans created within the coroutine + * from having a wrong parent span + */ + this.scopeStack = scopeManager.tlsScopeStack.initialValue(); + } + + @Override + public void restoreThreadContext( + @NotNull CoroutineContext coroutineContext, ScopeStack oldState) { + scopeManager.tlsScopeStack.set(oldState); + } + + @Override + public ScopeStack updateThreadContext(@NotNull CoroutineContext coroutineContext) { + final ScopeStack oldScopeStack = scopeManager.tlsScopeStack.get(); + scopeManager.tlsScopeStack.set(scopeStack); + + if (scopeStack.depth() == 0 && span != null) { + /* + * This is necessary for the spans created within the coroutine to properly inherit the spans hierarchy. + * It's not necessary to close the scope created here as it will be destroyed along with the context element and + * the scope stack. + */ + scopeManager.activate(span, INSTRUMENTATION); + } + + return oldScopeStack; + } + + @Nullable + @Override + public E get(@NotNull Key key) { + return CoroutineContext.Element.DefaultImpls.get(this, key); + } + + @NotNull + @Override + public CoroutineContext minusKey(@NotNull Key key) { + return CoroutineContext.Element.DefaultImpls.minusKey(this, key); + } + + @NotNull + @Override + public CoroutineContext plus(@NotNull CoroutineContext coroutineContext) { + return CoroutineContext.DefaultImpls.plus(this, coroutineContext); + } + + @Override + public R fold( + R initial, @NotNull Function2 operation) { + return CoroutineContext.Element.DefaultImpls.fold(this, initial, operation); + } + + @NotNull + @Override + public Key getKey() { + return KEY; + } +} diff --git a/dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/CoroutineContextAdvice.java b/dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/CoroutineContextAdvice.java new file mode 100644 index 00000000000..8b0d6f2e6bf --- /dev/null +++ b/dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/CoroutineContextAdvice.java @@ -0,0 +1,15 @@ +package datadog.trace.instrumentation.kotlin.coroutines; + +import datadog.trace.core.ScopeStackCoroutineContextHelper; +import kotlin.coroutines.CoroutineContext; +import net.bytebuddy.asm.Advice; + +public class CoroutineContextAdvice { + @Advice.OnMethodEnter + public static void enter( + @Advice.Argument(value = 1, readOnly = false) CoroutineContext coroutineContext) { + if (coroutineContext != null) { + coroutineContext = ScopeStackCoroutineContextHelper.addScopeStackContext(coroutineContext); + } + } +} diff --git a/dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/KotlinCoroutinesInstrumentation.java b/dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/KotlinCoroutinesInstrumentation.java new file mode 100644 index 00000000000..7daea238207 --- /dev/null +++ b/dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/KotlinCoroutinesInstrumentation.java @@ -0,0 +1,40 @@ +package datadog.trace.instrumentation.kotlin.coroutines; + +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static datadog.trace.util.Strings.getPackageName; +import static net.bytebuddy.matcher.ElementMatchers.isMethod; +import static net.bytebuddy.matcher.ElementMatchers.takesArgument; + +import com.google.auto.service.AutoService; +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.core.CoreTracer; + +@AutoService(Instrumenter.class) +public class KotlinCoroutinesInstrumentation extends Instrumenter.Tracing + implements Instrumenter.ForSingleType { + + public KotlinCoroutinesInstrumentation() { + super("kotlin-coroutines"); + } + + @Override + public String[] helperClassNames() { + return new String[] { + getPackageName(CoreTracer.class.getName()) + ".ScopeStackCoroutineContextHelper", + }; + } + + @Override + public String instrumentedType() { + return "kotlinx.coroutines.CoroutineContextKt"; + } + + @Override + public void adviceTransformations(AdviceTransformation transformation) { + transformation.applyAdvice( + isMethod() + .and(named("newCoroutineContext")) + .and(takesArgument(1, named("kotlin.coroutines.CoroutineContext"))), + packageName + ".CoroutineContextAdvice"); + } +}