8000 Migrate kotlin coroutines instrumentation to new Context API by mcculls · Pull Request #8528 · DataDog/dd-trace-java · GitHub
[go: up one dir, main page]
Skip to content

Migrate kotlin coroutines instrumentation to new Context API #8528

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Mar 10, 2025
Merged

Migrate kotlin coroutines instrumentation to new Context API #8528

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 10 additions & 10 deletions ...va/datadog/trace/instrumentation/kotlin/coroutines/KotlinCoroutines13Instrumentation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
package datadog.trace.instrumentation.kotlin.coroutines;

import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
import static datadog.trace.instrumentation.kotlin.coroutines.CoroutineContextHelper.getScopeStateContext;
import static datadog.trace.instrumentation.kotlin.coroutines.CoroutineContextHelper.initializeScopeStateContextIfActive;
import static datadog.trace.instrumentation.kotlin.coroutines.CoroutineContextHelper.getDatadogContext;
import static datadog.trace.instrumentation.kotlin.coroutines.CoroutineContextHelper.initializeDatadogContextIfActive;
import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
import static net.bytebuddy.matcher.ElementMatchers.isDeclaredBy;
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
Expand All @@ -11,7 +11,7 @@
import com.google.auto.service.AutoService;
import datadog.trace.agent.tooling.InstrumenterModule;
import datadog.trace.bootstrap.InstrumentationContext;
import datadog.trace.instrumentation.kotlin.coroutines.ScopeStateCoroutineContext.ScopeStateCoroutineContextItem;
import datadog.trace.instrumentation.kotlin.coroutines.DatadogCoroutineContext.DatadogCoroutineContextItem;
import kotlin.coroutines.CoroutineContext;
import kotlinx.coroutines.AbstractCoroutine;
import kotlinx.coroutines.Job;
Expand All @@ -33,21 +33,21 @@ public void methodAdvice(MethodTransformer transformer) {
}

/**
* Guarantees every coroutine created has an instance of ScopeStateCoroutineContext
* Guarantees every coroutine created has an instance of DatadogCoroutineContext
*
* @see ScopeStateCoroutineContext
* @see DatadogCoroutineContext
* @see AbstractCoroutine#AbstractCoroutine(CoroutineContext, boolean)
*/
public static class AbstractCoroutineConstructorAdvice {
@Advice.OnMethodEnter
public static void constructorInvocation(
@Advice.Argument(value = 0, readOnly = false) CoroutineContext parentContext) {
final ScopeStateCoroutineContext scopeStackContext = getScopeStateContext(parentContext);
if (scopeStackContext == null) {
final DatadogCoroutineContext datadogContext = getDatadogContext(parentContext);
if (datadogContext == null) {
parentContext =
parentContext.plus(
new ScopeStateCoroutineContext(
InstrumentationContext.get(Job.class, ScopeStateCoroutineContextItem.class)));
new DatadogCoroutineContext(
InstrumentationContext.get(Job.class, DatadogCoroutineContextItem.class)));
}
}

Expand All @@ -57,7 +57,7 @@ public static void constructorInvocationOnMethodExit(
@Advice.Argument(value = 1) final boolean active) {
// if this is not a lazy coroutine, inherit parent span from
// the coroutine constructor call site
initializeScopeStateContextIfActive(coroutine, active);
initializeDatadogContextIfActive(coroutine, active);
}
}
}
20 changes: 10 additions & 10 deletions ...va/datadog/trace/instrumentation/kotlin/coroutines/KotlinCoroutines15Instrumentation.java
View file
Open in desktop
8000
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
package datadog.trace.instrumentation.kotlin.coroutines;

import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
import static datadog.trace.instrumentation.kotlin.coroutines.CoroutineContextHelper.getScopeStateContext;
import static datadog.trace.instrumentation.kotlin.coroutines.CoroutineContextHelper.initializeScopeStateContextIfActive;
import static datadog.trace.instrumentation.kotlin.coroutines.CoroutineContextHelper.getDatadogContext;
import static datadog.trace.instrumentation.kotlin.coroutines.CoroutineContextHelper.initializeDatadogContextIfActive;
import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
import static net.bytebuddy.matcher.ElementMatchers.isDeclaredBy;
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
Expand All @@ -11,7 +11,7 @@
import com.google.auto.service.AutoService;
import datadog.trace.agent.tooling.InstrumenterModule;
import datadog.trace.bootstrap.InstrumentationContext;
import datadog.trace.instrumentation.kotlin.coroutines.ScopeStateCoroutineContext.ScopeStateCoroutineContextItem;
import datadog.trace.instrumentation.kotlin.coroutines.DatadogCoroutineContext.DatadogCoroutineContextItem;
import kotlin.coroutines.CoroutineContext;
import kotlinx.coroutines.AbstractCoroutine;
import kotlinx.coroutines.Job;
Expand All @@ -33,21 +33,21 @@ public void methodAdvice(MethodTransformer transformer) {
}

/**
* Guarantees every coroutine created has an instance of ScopeStateCoroutineContext
* Guarantees every coroutine created has an instance of DatadogCoroutineContext
*
* @see ScopeStateCoroutineContext
* @see DatadogCoroutineContext
* @see AbstractCoroutine#AbstractCoroutine(CoroutineContext, boolean, boolean)
*/
public static class AbstractCoroutineConstructorAdvice {
@Advice.OnMethodEnter
public static void constructorInvocation(
@Advice.Argument(value = 0, readOnly = false) CoroutineContext parentContext) {
final ScopeStateCoroutineContext scopeStackContext = getScopeStateContext(parentContext);
if (scopeStackContext == null) {
final DatadogCoroutineContext datadogContext = getDatadogContext(parentContext);
if (datadogContext == null) {
parentContext =
parentContext.plus(
new ScopeStateCoroutineContext(
InstrumentationContext.get(Job.class, ScopeStateCoroutineContextItem.class)));
new DatadogCoroutineContext(
InstrumentationContext.get(Job.class, DatadogCoroutineContextItem.class)));
}
}

Expand All @@ -57,7 +57,7 @@ public static void constructorInvocationOnMethodExit(
@Advice.Argument(value = 2) final boolean active) {
// if this is not a lazy coroutine, inherit parent span from
// the coroutine constructor call site
initializeScopeStateContextIfActive(coroutine, active);
initializeDatadogContextIfActive(coroutine, active);
}
}
}
29 changes: 14 additions & 15 deletions ...va/datadog/trace/instrumentation/kotlin/coroutines/AbstractCoroutinesInstrumentation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,8 @@

import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.extendsClass;
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
import static datadog.trace.instrumentation.kotlin.coroutines.CoroutineContextHelper.closeScopeStateContext;
import static datadog.trace.instrumentation.kotlin.coroutines.CoroutineContextHelper.initializeScopeStateContext;
import static datadog.trace.instrumentation.kotlin.coroutines.CoroutineContextHelper.closeDatadogContext;
import static datadog.trace.instrumentation.kotlin.coroutines.CoroutineContextHelper.initializeDatadogContext;
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
import static net.bytebuddy.matcher.ElementMatchers.isOverriddenFrom;
import static net.bytebuddy.matcher.ElementMatchers.returns;
Expand Down Expand Up @@ -42,9 +42,9 @@ protected final boolean defaultEnabled() {
@Override
public String[] helperClassNames() {
return new String[] {
packageName + ".ScopeStateCoroutineContext",
packageName + ".ScopeStateCoroutineContext$ContextElementKey",
packageName + ".ScopeStateCoroutineContext$ScopeStateCoroutineContextItem",
packageName + ".DatadogCoroutineContext",
packageName + ".DatadogCoroutineContext$ContextElementKey",
packageName + ".DatadogCoroutineContext$DatadogCoroutineContextItem",
packageName + ".CoroutineContextHelper",
};
}
Expand Down Expand Up @@ -84,36 +84,35 @@ public ElementMatcher<TypeDescription> structureMatcher() {
public Map<String, String> contextStore() {
return Collections.singletonMap(
"kotlinx.coroutines.Job",
packageName + ".ScopeStateCoroutineContext$ScopeStateCoroutineContextItem");
packageName + ".DatadogCoroutineContext$DatadogCoroutineContextItem");
}

/**
* If/when coroutine is started lazily, initializes ScopeStateCoroutineContext element on
* coroutine start
* If/when coroutine is started lazily, initializes DatadogCoroutineContext element on coroutine
* start
*
* @see ScopeStateCoroutineContext
* @see DatadogCoroutineContext
* @see AbstractCoroutine#onStart()
*/
public static class AbstractCoroutineOnStartAdvice {
@Advice.OnMethodEnter
public static void onStartInvocation(@Advice.This final AbstractCoroutine<?> coroutine) {
// try to inherit parent span from the coroutine start call site
initializeScopeStateContext(coroutine);
initializeDatadogContext(coroutine);
}
}

/**
* Guarantees a ScopeStateCoroutineContext element is always closed when coroutine transitions
* into a terminal state.
* Guarantees a DatadogCoroutineContext element is always closed when coroutine transitions into a
* terminal state.
*
* @see ScopeStateCoroutineContext
* @see DatadogCoroutineContext
* @see AbstractCoroutine#onCompletionInternal(Object)
*/
public static class JobSupportAfterCompletionInternalAdvice {
@Advice.OnMethodEnter
public static void onCompletionInternal(@Advice.This final AbstractCoroutine<?> coroutine) {
// close the scope if needed
closeScopeStateContext(coroutine);
closeDatadogContext(coroutine);
}
}
}
26 changes: 12 additions & 14 deletions ...src/main/java/datadog/trace/instrumentation/kotlin/coroutines/CoroutineContextHelper.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,30 +17,28 @@ public static Job getJob(final CoroutineContext context) {
}

@Nullable
public static ScopeStateCoroutineContext getScopeStateContext(final CoroutineContext context) {
return context.get(ScopeStateCoroutineContext.KEY);
public static DatadogCoroutineContext getDatadogContext(final CoroutineContext context) {
return context.get(DatadogCoroutineContext.KEY);
}

public static void initializeScopeStateContextIfActive(
public static void initializeDatadogContextIfActive(
final AbstractCoroutine<?> coroutine, final boolean active) {
if (active) {
initializeScopeStateContext(coroutine);
initializeDatadogContext(coroutine);
}
}

public static void initializeScopeStateContext(final AbstractCoroutine<?> coroutine) {
final ScopeStateCoroutineContext scopeStackContext =
getScopeStateContext(coroutine.getContext());
if (scopeStackContext != null) {
scopeStackContext.maybeInitialize(coroutine);
public static void initializeDatadogContext(final AbstractCoroutine<?> coroutine) {
final DatadogCoroutineContext datadogContext = getDatadogContext(coroutine.getContext());
if (datadogContext != null) {
datadogContext.maybeInitialize(coroutine);
}
}

public static void closeScopeStateContext(final AbstractCoroutine<?> coroutine) {
final ScopeStateCoroutineContext scopeStackContext =
getScopeStateContext(coroutine.getContext());
if (scopeStackContext != null) {
scopeStackContext.maybeCloseScopeAndCancelContinuation(coroutine);
public static void closeDatadogContext(final AbstractCoroutine<?> coroutine) {
final DatadogCoroutineContext datadogContext = getDatadogContext(coroutine.getContext());
if (datadogContext != null) {
datadogContext.maybeCloseScopeAndCancelContinuation(coroutine);
}
}
}
49 changes: 23 additions & 26 deletions ...oroutines/ScopeStateCoroutineContext.java → ...n/coroutines/DatadogCoroutineContext.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,66 +2,63 @@

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

import datadog.context.Context;
import datadog.trace.bootstrap.ContextStore;
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
import datadog.trace.bootstrap.instrumentation.api.ScopeState;
import kotlin.coroutines.CoroutineContext;
import kotlin.jvm.functions.Function2;
import kotlinx.coroutines.Job;
import kotlinx.coroutines.ThreadContextElement;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

public class ScopeStateCoroutineContext implements ThreadContextElement<ScopeState> {
public class DatadogCoroutineContext implements ThreadContextElement<Context> {

public static final Key<ScopeStateCoroutineContext> KEY = new ContextElementKey();
public static final Key<DatadogCoroutineContext> KEY = new ContextElementKey();

private final ContextStore<Job, ScopeStateCoroutineContextItem> contextItemPerCoroutine;
private final ContextStore<Job, DatadogCoroutineContextItem> contextItemPerCoroutine;

public ScopeStateCoroutineContext(
final Conte F438 xtStore<Job, ScopeStateCoroutineContextItem> contextItemPerCoroutine) {
public DatadogCoroutineContext(
final ContextStore<Job, DatadogCoroutineContextItem> contextItemPerCoroutine) {
this.contextItemPerCoroutine = contextItemPerCoroutine;
}

/** Get a context item instance for the coroutine and try to initialize it */
public void maybeInitialize(final Job coroutine) {
contextItemPerCoroutine
.putIfAbsent(coroutine, ScopeStateCoroutineContextItem::new)
.putIfAbsent(coroutine, DatadogCoroutineContextItem::new)
.maybeInitialize();
}

@Override
public void restoreThreadContext(
@NotNull final CoroutineContext coroutineContext, final ScopeState oldState) {
oldState.activate();
@NotNull final CoroutineContext coroutineContext, final Context oldDatadogContext) {
oldDatadogContext.swap();
}

@Override
public ScopeState updateThreadContext(@NotNull final CoroutineContext coroutineContext) {
final ScopeState oldScopeState = AgentTracer.get().newScopeState();
oldScopeState.fetchFromActive();
public Context updateThreadContext(@NotNull final CoroutineContext coroutineContext) {
final Context oldDatadogContext = Context.current();

final Job coroutine = CoroutineContextHelper.getJob(coroutineContext);
final ScopeStateCoroutineContextItem contextItem = contextItemPerCoroutine.get(coroutine);
final DatadogCoroutineContextItem contextItem = contextItemPerCoroutine.get(coroutine);
if (contextItem != null) {
contextItem.activate();
}

return oldScopeState;
return oldDatadogContext;
}

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

contextItem.maybeCloseScopeAndCancelContinuation();
contextItemPerCoroutine.remove(coroutine);

currentThreadScopeState.activate();
currentDatadogContext.swap();
}
}

Expand Down Expand Up @@ -95,20 +92,20 @@ public Key<?> getKey() {
return KEY;
}

static class ContextElementKey implements Key<ScopeStateCoroutineContext> {}
static class ContextElementKey implements Key<DatadogCoroutineContext> {}

public static class ScopeStateCoroutineContextItem {
private final ScopeState coroutineScopeState;
public static class DatadogCoroutineContextItem {
private final Context datadogContext;
@Nullable private AgentScope.Continuation continuation;
@Nullable private AgentScope continuationScope;
private boolean isInitialized = false;

public ScopeStateCoroutineContextItem() {
coroutineScopeState = AgentTracer.get().newScopeState();
public DatadogCoroutineContextItem() {
datadogContext = Context.root();
}

public void activate() {
coroutineScopeState.activate();
datadogContext.swap();

if (continuation != null && continuationScope == null) {
continuationScope = continuation.activate();
Expand All @@ -131,7 +128,7 @@ public void maybeInitialize() {
* scope and cancels the continuation.
*/
public void maybeCloseScopeAndCancelContinuation() {
coroutineScopeState.activate();
datadogContext.swap();

if (continuationScope != null) {
continuationScope.close();
Expand Down
0