diff --git a/dd-java-agent/instrumentation/kotlin-coroutines/coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/KotlinCoroutines13Instrumentation.java b/dd-java-agent/instrumentation/kotlin-coroutines/coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/KotlinCoroutines13Instrumentation.java index 6e07f8fe686..660fdb47a2b 100644 --- a/dd-java-agent/instrumentation/kotlin-coroutines/coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/KotlinCoroutines13Instrumentation.java +++ b/dd-java-agent/instrumentation/kotlin-coroutines/coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/KotlinCoroutines13Instrumentation.java @@ -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; @@ -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; @@ -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))); } } @@ -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); } } } diff --git a/dd-java-agent/instrumentation/kotlin-coroutines/coroutines-1.5/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/KotlinCoroutines15Instrumentation.java b/dd-java-agent/instrumentation/kotlin-coroutines/coroutines-1.5/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/KotlinCoroutines15Instrumentation.java index f8f825103cf..b3ad36b6cc7 100644 --- a/dd-java-agent/instrumentation/kotlin-coroutines/coroutines-1.5/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/KotlinCoroutines15Instrumentation.java +++ b/dd-java-agent/instrumentation/kotlin-coroutines/coroutines-1.5/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/KotlinCoroutines15Instrumentation.java @@ -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; @@ -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; @@ -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))); } } @@ -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); } } } diff --git a/dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/AbstractCoroutinesInstrumentation.java b/dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/AbstractCoroutinesInstrumentation.java index 897e85079f7..c6159d0b8f9 100644 --- a/dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/AbstractCoroutinesInstrumentation.java +++ b/dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/AbstractCoroutinesInstrumentation.java @@ -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; @@ -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", }; } @@ -84,36 +84,35 @@ public ElementMatcher structureMatcher() { public Map 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); } } } diff --git a/dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/CoroutineContextHelper.java b/dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/CoroutineContextHelper.java index c8eb09d50f9..15ce695b320 100644 --- a/dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/CoroutineContextHelper.java +++ b/dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/CoroutineContextHelper.java @@ -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); } } } diff --git a/dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/ScopeStateCoroutineContext.java b/dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/DatadogCoroutineContext.java similarity index 64% rename from dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/ScopeStateCoroutineContext.java rename to dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/DatadogCoroutineContext.java index fe49748ea98..50e5ef15fb2 100644 --- a/dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/ScopeStateCoroutineContext.java +++ b/dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/DatadogCoroutineContext.java @@ -2,10 +2,9 @@ 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; @@ -13,55 +12,53 @@ import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; -public class ScopeStateCoroutineContext implements ThreadContextElement { +public class DatadogCoroutineContext implements ThreadContextElement { - public static final Key KEY = new ContextElementKey(); + public static final Key KEY = new ContextElementKey(); - private final ContextStore contextItemPerCoroutine; + private final ContextStore contextItemPerCoroutine; - public ScopeStateCoroutineContext( - final ContextStore contextItemPerCoroutine) { + public DatadogCoroutineContext( + final ContextStore 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(); } } @@ -95,20 +92,20 @@ public Key getKey() { return KEY; } - static class ContextElementKey implements Key {} + static class ContextElementKey implements Key {} - 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(); @@ -131,7 +128,7 @@ public void maybeInitialize() { * scope and cancels the continuation. */ public void maybeCloseScopeAndCancelContinuation() { - coroutineScopeState.activate(); + datadogContext.swap(); if (continuationScope != null) { continuationScope.close();