diff --git a/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/main/java/datadog/trace/instrumentation/java/lang/jdk21/VirtualThreadInstrumentation.java b/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/main/java/datadog/trace/instrumentation/java/lang/jdk21/VirtualThreadInstrumentation.java index e105b42e26b..81d7fa472d5 100644 --- a/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/main/java/datadog/trace/instrumentation/java/lang/jdk21/VirtualThreadInstrumentation.java +++ b/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/main/java/datadog/trace/instrumentation/java/lang/jdk21/VirtualThreadInstrumentation.java @@ -73,7 +73,7 @@ public final class VirtualThreadInstrumentation extends InstrumenterModule.Conte */ private static final String[] PRELOAD_CLASS_NAMES = { VIRTUAL_THREAD_STATE_CLASS_NAME, - "datadog.trace.core.scopemanager.ScopeContext", + "datadog.trace.core.scopemanager.SavedScopeContext", "datadog.trace.core.scopemanager.ScopeStack" }; diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/build.gradle b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/build.gradle index 86abd12eb08..0ca36dd63d6 100644 --- a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/build.gradle +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/build.gradle @@ -53,9 +53,11 @@ dependencies { testImplementation testFixtures(project(':dd-java-agent:instrumentation:kotlin-coroutines-1.3')) latestDepTestImplementation libs.kotlin - latestDepTestImplementation "org.jetbrains.kotlinx:kotlinx-coroutines-core-jvm:1.6.+" + latestDepTestImplementation "org.jetbrains.kotlinx:kotlinx-coroutines-core-jvm:1.8.+" latestDepTestImplementation testFixtures(project(':dd-java-agent:instrumentation:kotlin-coroutines-1.3')) testFixturesCompileOnly(libs.bundles.groovy) testFixturesCompileOnly(libs.bundles.spock) + testFixturesCompileOnly libs.junit.jupiter + testFixturesCompileOnly libs.tabletest } diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/gradle.lockfile b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/gradle.lockfile index 31525e3b1ca..43cea71c634 100644 --- a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/gradle.lockfile +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/gradle.lockfile @@ -99,23 +99,24 @@ org.jetbrains.kotlin:kotlin-scripting-common:2.1.20=kotlinBuildToolsApiClasspath org.jetbrains.kotlin:kotlin-scripting-compiler-embeddable:2.1.20=kotlinBuildToolsApiClasspath,kotlinCompilerPluginClasspathLatestDepTest,kotlinCompilerPluginClasspathMain,kotlinCompilerPluginClasspathTest,kotlinCompilerPluginClasspathTestFixtures org.jetbrains.kotlin:kotlin-scripting-compiler-impl-embeddable:2.1.20=kotlinBuildToolsApiClasspath,kotlinCompilerPluginClasspathLatestDepTest,kotlinCompilerPluginClasspathMain,kotlinCompilerPluginClasspathTest,kotlinCompilerPluginClasspathTestFixtures org.jetbrains.kotlin:kotlin-scripting-jvm:2.1.20=kotlinBuildToolsApiClasspath,kotlinCompilerPluginClasspathLatestDepTest,kotlinCompilerPluginClasspathMain,kotlinCompilerPluginClasspathTest,kotlinCompilerPluginClasspathTestFixtures -org.jetbrains.kotlin:kotlin-stdlib-common:1.6.21=compileClasspath,compileOnlyDependenciesMetadata,latestDepTestCompileClasspath,latestDepTestImplementationDependenciesMetadata,latestDepTestRuntimeClasspath,testCompileClasspath,testFixturesCompileClasspath,testFixturesImplementationDependenciesMetadata,testFixturesRuntimeClasspath,testImplementationDependenciesMetadata,testRuntimeClasspath -org.jetbrains.kotlin:kotlin-stdlib-jdk7:1.6.21=latestDepTestCompileClasspath,latestDepTestImplementationDependenciesMetadata,latestDepTestRuntimeClasspath -org.jetbrains.kotlin:kotlin-stdlib-jdk8:1.6.21=latestDepTestCompileClasspath,latestDepTestImplementationDependenciesMetadata,latestDepTestRuntimeClasspath -org.jetbrains.kotlin:kotlin-stdlib:1.6.21=compileClasspath,compileOnlyDependenciesMetadata,latestDepTestCompileClasspath,latestDepTestImplementationDependenciesMetadata,latestDepTestRuntimeClasspath,testCompileClasspath,testFixturesCompileClasspath,testFixturesImplementationDependenciesMetadata,testFixturesRuntimeClasspath,testImplementationDependenciesMetadata,testRuntimeClasspath +org.jetbrains.kotlin:kotlin-stdlib-common:1.6.21=compileClasspath,compileOnlyDependenciesMetadata,testCompileClasspath,testFixturesCompileClasspath,testFixturesImplementationDependenciesMetadata,testFixturesRuntimeClasspath,testImplementationDependenciesMetadata,testRuntimeClasspath +org.jetbrains.kotlin:kotlin-stdlib-common:1.9.21=latestDepTestImplementationDependenciesMetadata +org.jetbrains.kotlin:kotlin-stdlib:1.6.21=compileClasspath,compileOnlyDependenciesMetadata,testCompileClasspath,testFixturesCompileClasspath,testFixturesImplementationDependenciesMetadata,testFixturesRuntimeClasspath,testImplementationDependenciesMetadata,testRuntimeClasspath +org.jetbrains.kotlin:kotlin-stdlib:1.9.21=latestDepTestCompileClasspath,latestDepTestImplementationDependenciesMetadata,latestDepTestRuntimeClasspath org.jetbrains.kotlin:kotlin-stdlib:2.1.20=kotlinBuildToolsApiClasspath,kotlinCompilerClasspath,kotlinCompilerPluginClasspathLatestDepTest,kotlinCompilerPluginClasspathMain,kotlinCompilerPluginClasspathTest,kotlinCompilerPluginClasspathTestFixtures,kotlinKlibCommonizerClasspath -org.jetbrains.kotlinx:atomicfu:0.17.3=latestDepTestImplementationDependenciesMetadata -org.jetbrains.kotlinx:kotlinx-coroutines-bom:1.6.4=latestDepTestCompileClasspath,latestDepTestImplementationDependenciesMetadata,latestDepTestRuntimeClasspath -org.jetbrains.kotlinx:kotlinx-coroutines-core-jvm:1.6.4=latestDepTestCompileClasspath,latestDepTestImplementationDependenciesMetadata,latestDepTestRuntimeClasspath +org.jetbrains.kotlinx:atomicfu:0.23.1=latestDepTestImplementationDependenciesMetadata +org.jetbrains.kotlinx:kotlinx-coroutines-bom:1.8.1=latestDepTestCompileClasspath,latestDepTestImplementationDependenciesMetadata,latestDepTestRuntimeClasspath org.jetbrains.kotlinx:kotlinx-coroutines-core-jvm:1.8.0=kotlinBuildToolsApiClasspath,kotlinCompilerClasspath,kotlinKlibCommonizerClasspath +org.jetbrains.kotlinx:kotlinx-coroutines-core-jvm:1.8.1=latestDepTestCompileClasspath,latestDepTestImplementationDependenciesMetadata,latestDepTestRuntimeClasspath org.jetbrains.kotlinx:kotlinx-coroutines-core:1.3.0=compileClasspath,compileOnlyDependenciesMetadata,testCompileClasspath,testFixturesCompileClasspath,testFixturesImplementationDependenciesMetadata,testFixturesRuntimeClasspath,testImplementationDependenciesMetadata,testRuntimeClasspath -org.jetbrains.kotlinx:kotlinx-coroutines-core:1.6.4=latestDepTestCompileClasspath,latestDepTestImplementationDependenciesMetadata,latestDepTestRuntimeClasspath -org.jetbrains:annotations:13.0=compileClasspath,compileOnlyDependenciesMetadata,kotlinBuildToolsApiClasspath,kotlinCompilerClasspath,kotlinCompilerPluginClasspathLatestDepTest,kotlinCompilerPluginClasspathMain,kotlinCompilerPluginClasspathTest,kotlinCompilerPluginClasspathTestFixtures,kotlinKlibCommonizerClasspath,latestDepTestCompileClasspath,latestDepTestImplementationDependenciesMetadata,latestDepTestRuntimeClasspath,testCompileClasspath,testFixturesCompileClasspath,testFixturesImplementationDependenciesMetadata,testFixturesRuntimeClasspath,testImplementationDependenciesMetadata,testRuntimeClasspath +org.jetbrains.kotlinx:kotlinx-coroutines-core:1.8.1=latestDepTestCompileClasspath,latestDepTestImplementationDependenciesMetadata,latestDepTestRuntimeClasspath +org.jetbrains:annotations:13.0=compileClasspath,compileOnlyDependenciesMetadata,kotlinBuildToolsApiClasspath,kotlinCompilerClasspath,kotlinCompilerPluginClasspathLatestDepTest,kotlinCompilerPluginClasspathMain,kotlinCompilerPluginClasspathTest,kotlinCompilerPluginClasspathTestFixtures,kotlinKlibCommonizerClasspath,testCompileClasspath,testFixturesCompileClasspath,testFixturesImplementationDependenciesMetadata,testFixturesRuntimeClasspath,testImplementationDependenciesMetadata,testRuntimeClasspath +org.jetbrains:annotations:23.0.0=latestDepTestCompileClasspath,latestDepTestImplementationDependenciesMetadata,latestDepTestRuntimeClasspath org.jspecify:jspecify:1.0.0=latestDepTestCompileClasspath,latestDepTestImplementationDependenciesMetadata,latestDepTestRuntimeClasspath,testCompileClasspath,testFixturesApiDependenciesMetadata,testFixturesCompileClasspath,testFixturesImplementationDependenciesMetadata,testFixturesRuntimeClasspath,testImplementationDependenciesMetadata,testRuntimeClasspath -org.junit.jupiter:junit-jupiter-api:5.14.1=latestDepTestCompileClasspath,latestDepTestImplementationDependenciesMetadata,latestDepTestRuntimeClasspath,testCompileClasspath,testFixturesRuntimeClasspath,testImplementationDependenciesMetadata,testRuntimeClasspath +org.junit.jupiter:junit-jupiter-api:5.14.1=latestDepTestCompileClasspath,latestDepTestImplementationDependenciesMetadata,latestDepTestRuntimeClasspath,testCompileClasspath,testFixturesCompileClasspath,testFixturesRuntimeClasspath,testImplementationDependenciesMetadata,testRuntimeClasspath org.junit.jupiter:junit-jupiter-engine:5.14.1=latestDepTestRuntimeClasspath,testFixturesRuntimeClasspath,testRuntimeClasspath -org.junit.jupiter:junit-jupiter-params:5.14.1=latestDepTestCompileClasspath,latestDepTestImplementationDependenciesMetadata,latestDepTestRuntimeClasspath,testCompileClasspath,testFixturesRuntimeClasspath,testImplementationDependenciesMetadata,testRuntimeClasspath -org.junit.jupiter:junit-jupiter:5.14.1=latestDepTestCompileClasspath,latestDepTestImplementationDependenciesMetadata,latestDepTestRuntimeClasspath,testCompileClasspath,testFixturesRuntimeClasspath,testImplementationDependenciesMetadata,testRuntimeClasspath +org.junit.jupiter:junit-jupiter-params:5.14.1=latestDepTestCompileClasspath,latestDepTestImplementationDependenciesMetadata,latestDepTestRuntimeClasspath,testCompileClasspath,testFixturesCompileClasspath,testFixturesRuntimeClasspath,testImplementationDependenciesMetadata,testRuntimeClasspath +org.junit.jupiter:junit-jupiter:5.14.1=latestDepTestCompileClasspath,latestDepTestImplementationDependenciesMetadata,latestDepTestRuntimeClasspath,testCompileClasspath,testFixturesCompileClasspath,testFixturesRuntimeClasspath,testImplementationDependenciesMetadata,testRuntimeClasspath org.junit.platform:junit-platform-commons:1.14.1=latestDepTestCompileClasspath,latestDepTestImplementationDependenciesMetadata,latestDepTestRuntimeClasspath,testCompileClasspath,testFixturesCompileClasspath,testFixturesCompileOnlyDependenciesMetadata,testFixturesRuntimeClasspath,testImplementationDependenciesMetadata,testRuntimeClasspath org.junit.platform:junit-platform-engine:1.14.1=latestDepTestCompileClasspath,latestDepTestImplementationDependenciesMetadata,latestDepTestRuntimeClasspath,testCompileClasspath,testFixturesCompileClasspath,testFixturesCompileOnlyDependenciesMetadata,testFixturesRuntimeClasspath,testImplementationDependenciesMetadata,testRuntimeClasspath org.junit.platform:junit-platform-launcher:1.14.1=latestDepTestRuntimeClasspath,testFixturesRuntimeClasspath,testRuntimeClasspath @@ -144,7 +145,7 @@ org.slf4j:slf4j-simple:2.0.17=spotbugsSlf4j org.snakeyaml:snakeyaml-engine:2.9=buildTimeInstrumentationPlugin,latestDepTestRuntimeClasspath,muzzleTooling,runtimeClasspath,testFixturesRuntimeClasspath,testRuntimeClasspath org.spockframework:spock-bom:2.4-groovy-3.0=latestDepTestCompileClasspath,latestDepTestImplementationDependenciesMetadata,latestDepTestRuntimeClasspath,testCompileClasspath,testFixturesCompileClasspath,testFixturesCompileOnlyDependenciesMetadata,testImplementationDependenciesMetadata,testRuntimeClasspath org.spockframework:spock-core:2.4-groovy-3.0=latestDepTestCompileClasspath,latestDepTestImplementationDependenciesMetadata,latestDepTestRuntimeClasspath,testCompileClasspath,testFixturesCompileClasspath,testFixturesCompileOnlyDependenciesMetadata,testImplementationDependenciesMetadata,testRuntimeClasspath -org.tabletest:tabletest-junit:1.2.2=latestDepTestCompileClasspath,latestDepTestImplementationDependenciesMetadata,latestDepTestRuntimeClasspath,testCompileClasspath,testImplementationDependenciesMetadata,testRuntimeClasspath -org.tabletest:tabletest-parser:1.2.1=latestDepTestCompileClasspath,latestDepTestImplementationDependenciesMetadata,latestDepTestRuntimeClasspath,testCompileClasspath,testImplementationDependenciesMetadata,testRuntimeClasspath +org.tabletest:tabletest-junit:1.2.2=latestDepTestCompileClasspath,latestDepTestImplementationDependenciesMetadata,latestDepTestRuntimeClasspath,testCompileClasspath,testFixturesCompileClasspath,testImplementationDependenciesMetadata,testRuntimeClasspath +org.tabletest:tabletest-parser:1.2.1=latestDepTestCompileClasspath,latestDepTestImplementationDependenciesMetadata,latestDepTestRuntimeClasspath,testCompileClasspath,testFixturesCompileClasspath,testImplementationDependenciesMetadata,testRuntimeClasspath org.xmlresolver:xmlresolver:5.3.3=spotbugs empty=apiDependenciesMetadata,intransitiveDependenciesMetadata,kotlinCompilerPluginClasspath,kotlinNativeCompilerPluginClasspath,kotlinScriptDefExtensions,latestDepTestApiDependenciesMetadata,latestDepTestIntransitiveDependenciesMetadata,latestDepTestKotlinScriptDefExtensions,spotbugsPlugins,testApiDependenciesMetadata,testFixturesAnnotationProcessor,testFixturesIntransitiveDependenciesMetadata,testFixturesKotlinScriptDefExtensions,testIntransitiveDependenciesMetadata,testKotlinScriptDefExtensions diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/latestDepTest/groovy/StandardSuspensionHandoffTest.groovy b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/latestDepTest/groovy/StandardSuspensionHandoffTest.groovy new file mode 100644 index 00000000000..494714fce87 --- /dev/null +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/latestDepTest/groovy/StandardSuspensionHandoffTest.groovy @@ -0,0 +1,29 @@ +import datadog.trace.instrumentation.kotlin.coroutines.AbstractStandardSuspensionHandoffTest +import datadog.trace.instrumentation.kotlin.coroutines.StandardSuspensionHandoffTests + +class StandardSuspensionHandoffTest extends AbstractStandardSuspensionHandoffTest { + def "undispatched coroutine can replay its restoration token"() { + when: + new StandardSuspensionHandoffTests().runReplayedRestoration() + + then: + noExceptionThrown() + } + + def "scope-touching thread context elements do not corrupt the handed-off stack"() { + when: + def violations = new StandardSuspensionHandoffTests().runScopeTouchingElements(1000) + + then: + // Overlapping workers can temporarily see each other's element scopes. The coroutine must + // retain its base context, but restore callbacks can observe a different top during overlap. + def expectedDuringOverlap = [ + 'element scope leaked from previous worker', + 'before scope not on top at restore', + 'after scope not on top at restore', + 'before restore did not return to previous context', + 'after restore did not return to previous context' + ] + (violations.keySet() - expectedDuringOverlap).isEmpty() + } +} diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/latestDepTest/java/datadog/trace/instrumentation/kotlin/coroutines/DatadogThreadContextElementTest.java b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/latestDepTest/java/datadog/trace/instrumentation/kotlin/coroutines/DatadogThreadContextElementTest.java new file mode 100644 index 00000000000..e45a22d4b23 --- /dev/null +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/latestDepTest/java/datadog/trace/instrumentation/kotlin/coroutines/DatadogThreadContextElementTest.java @@ -0,0 +1,3 @@ +package datadog.trace.instrumentation.kotlin.coroutines; + +class DatadogThreadContextElementTest extends AbstractDatadogThreadContextElementTest {} diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/latestDepTest/kotlin/KotlinCoroutineTests.kt b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/latestDepTest/kotlin/KotlinCoroutineTests.kt index 77be265ed46..8bd598d68ed 100644 --- a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/latestDepTest/kotlin/KotlinCoroutineTests.kt +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/latestDepTest/kotlin/KotlinCoroutineTests.kt @@ -51,6 +51,18 @@ class KotlinCoroutineTests(dispatcher: CoroutineDispatcher) : CoreKotlinCoroutin 3 } + @Trace + override fun structuredChildrenOutliveSeedSpan(): Int = super.structuredChildrenOutliveSeedSpan() + + @Trace + override fun structuredChildrenOfScopedCoroutines(): Int = super.structuredChildrenOfScopedCoroutines() + + @Trace + override fun childWithReplacedJob(supervisor: Boolean): Int = super.childWithReplacedJob(supervisor) + + @Trace + override fun structuredChildOfCancelledCoroutine(): Int = super.structuredChildOfCancelledCoroutine() + @Trace override fun tracePreventedByCancellation(): Int = super.tracePreventedByCancellation() diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/DatadogThreadContextElement.java b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/DatadogThreadContextElement.java index 55516db89e0..5dff5afb348 100644 --- a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/DatadogThreadContextElement.java +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/DatadogThreadContextElement.java @@ -2,16 +2,26 @@ import datadog.context.Context; import datadog.context.ContextContinuation; +import datadog.trace.api.GenericClassValue; +import datadog.trace.api.InstrumenterConfig; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.AgentTraceCollector; +import datadog.trace.bootstrap.instrumentation.api.AgentTracer; import java.util.concurrent.atomic.AtomicReferenceFieldUpdater; import javax.annotation.Nonnull; import javax.annotation.Nullable; import kotlin.coroutines.CoroutineContext; +import kotlin.coroutines.EmptyCoroutineContext; import kotlin.jvm.functions.Function2; import kotlinx.coroutines.AbstractCoroutine; +import kotlinx.coroutines.Job; import kotlinx.coroutines.ThreadContextElement; +import kotlinx.coroutines.internal.ScopeCoroutine; /** Manages the Datadog context for coroutines, switching contexts as coroutines switch threads. */ -public final class DatadogThreadContextElement implements ThreadContextElement { +public final class DatadogThreadContextElement + implements ThreadContextElement, + Function2 { private static final AtomicReferenceFieldUpdater CONTINUATION = AtomicReferenceFieldUpdater.newUpdater( @@ -20,16 +30,71 @@ public final class DatadogThreadContextElement implements ThreadContextElement DATADOG_KEY = new CoroutineContext.Key() {}; + private static final boolean LEGACY_CONTEXT_MANAGER = + InstrumenterConfig.get().isLegacyContextManagerEnabled(); + + /** Coroutines built by coroutineScope, withContext, etc. that their caller waits for. */ + private static final ClassValue SCOPED_COROUTINE = + GenericClassValue.of( + type -> + ScopeCoroutine.class.isAssignableFrom(type) + // older timeout coroutines are scoped but don't extend ScopeCoroutine + || (AbstractCoroutine.class.isAssignableFrom(type) + && "kotlinx.coroutines.TimeoutCoroutine".equals(type.getName()))); + + /** Adds a fresh Datadog element, replacing any inherited element in place. */ public static CoroutineContext addDatadogElement(CoroutineContext coroutineContext) { - if (coroutineContext.get(DATADOG_KEY) != null) { - return coroutineContext; // already added + DatadogThreadContextElement existing = coroutineContext.get(DATADOG_KEY); + if (existing == null) { + return coroutineContext.plus(new DatadogThreadContextElement(null, null)); } - return coroutineContext.plus(new DatadogThreadContextElement()); + // a context passed from another coroutine starts from that coroutine's current context + DatadogThreadContextElement fresh = + new DatadogThreadContextElement( + existing.context, existing.traceKeptOpenFor(coroutineContext.get(Job.Key))); + // keep the element's position, so elements added after it can still update scopes on top + return coroutineContext.fold(EmptyCoroutineContext.INSTANCE, fresh); + } + + /** Rebuilds the context using this replacement element, avoiding a capturing fold lambda. */ + @Override + public CoroutineContext invoke(CoroutineContext rebuilt, CoroutineContext.Element element) { + return rebuilt.plus(element.getKey() == DATADOG_KEY ? this : element); } - private Context context; + // only needed until the initial context is captured; cleared so a long-lived child doesn't keep + // its parent's scope stack and trace alive + @Nullable private Context seed; + @Nullable private AgentTraceCollector traceKeptOpenByParent; + private volatile Context context; private volatile ContextContinuation continuation; + /** The coroutine this element was created for. */ + @Nullable private volatile Job owner; + + /** The trace kept open until this coroutine completes, by its own continuation or its parent. */ + @Nullable private volatile AgentTraceCollector traceKeptOpen; + + private DatadogThreadContextElement( + @Nullable Context seed, @Nullable AgentTraceCollector traceKeptOpenByParent) { + this.seed = seed; + this.traceKeptOpenByParent = traceKeptOpenByParent; + } + + /** Returns the trace this coroutine keeps open, if coroutines of this job are its children. */ + @Nullable + private AgentTraceCollector traceKeptOpenFor(@Nullable Job job) { + AgentTraceCollector trace = traceKeptOpen; + // structured children complete before their parent, which keeps its trace open until then + return trace != null && job != null && (job == owner || isScopedBy(job)) ? trace : null; + } + + /** Scoped coroutines share their caller's element and complete before the caller does. */ + private boolean isScopedBy(Job job) { + return SCOPED_COROUTINE.get(job.getClass()) + && ((AbstractCoroutine) job).getContext().get(DATADOG_KEY) == this; + } + @Nonnull @Override public Key getKey() { @@ -38,18 +103,60 @@ public Key getKey() { public static void captureDatadogContext(@Nonnull AbstractCoroutine coroutine) { DatadogThreadContextElement datadog = coroutine.getContext().get(DATADOG_KEY); - if (datadog != null && datadog.context == null) { - // record context to use for this coroutine - datadog.context = Context.current(); - // stop enclosing trace from finishing early - datadog.continuation = datadog.context.capture(); + if (datadog != null) { + if (datadog.owner == null) { + // scoped coroutines built later share this element, but aren't its owner + datadog.owner = coroutine; + } + if (datadog.context == null) { + datadog.captureInitialContext(); + } + } + } + + private synchronized void captureInitialContext() { + if (context == null) { + Context initial = seed != null ? seed : Context.current(); + if (LEGACY_CONTEXT_MANAGER) { + // workers share this coroutine's scope stack, rather than copying it between them + initial = AgentTracer.get().sharedContextStack(initial); + } + // this coroutine isn't running yet, so its context can't change under us + AgentTraceCollector trace = traceOf(initial); + if (trace != null) { + if (trace == traceKeptOpenByParent) { + // the parent keeps this trace open until after this child completes + traceKeptOpen = trace; + } else { + // stop enclosing trace from finishing early + ContextContinuation captured = initial.capture(); + continuation = captured; + if (captured.context() != Context.root()) { + // only when the capture was accepted, e.g. async propagation is enabled + traceKeptOpen = trace; + } + } + } + context = initial; + seed = null; + traceKeptOpenByParent = null; } } + @Nullable + private static AgentTraceCollector traceOf(Context context) { + AgentSpan span = AgentSpan.fromContext(context); + return span != null ? span.spanContext().getTraceCollector() : null; + } + public static void cancelDatadogContext(@Nonnull AbstractCoroutine coroutine) { DatadogThreadContextElement datadog = coroutine.getContext().get(DATADOG_KEY); - ContextContinuation continuation = - datadog == null ? null : CONTINUATION.getAndSet(datadog, null); + if (datadog == null || datadog.owner != coroutine) { + // scoped coroutines share their caller's element, but must not release its trace hold + return; + } + datadog.traceKeptOpen = null; + ContextContinuation continuation = CONTINUATION.getAndSet(datadog, null); if (continuation != null) { // release enclosing trace now the coroutine has completed continuation.release(); @@ -58,19 +165,28 @@ public static void cancelDatadogContext(@Nonnull AbstractCoroutine coroutine) @Override public Context updateThreadContext(@Nonnull CoroutineContext coroutineContext) { - if (context == null) { - // record context to use for this coroutine - context = Context.current(); - // stop enclosing trace from finishing early - continuation = context.capture(); + Context current = context; + if (current == null) { + if (owner == null) { + // an already-cancelled coroutine can skip constructor capture but still install context + owner = coroutineContext.get(Job.Key); + } + captureInitialContext(); + current = context; } - return context.swap(); + return current.swap(); } @Override public void restoreThreadContext( @Nonnull CoroutineContext coroutineContext, Context originalContext) { - context = originalContext.swap(); + Context coroutineContextAfter = originalContext.swap(); + // legacy workers share the coroutine's scope stack, so it never needs saving; writing it back + // would usually store that same stack, but after an unbalanced swap on this thread it would + // rebind the coroutine to the wrong stack + if (!LEGACY_CONTEXT_MANAGER) { + context = coroutineContextAfter; + } } @Nonnull diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/test/groovy/StandardSuspensionHandoffTest.groovy b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/test/groovy/StandardSuspensionHandoffTest.groovy new file mode 100644 index 00000000000..0f40bb4a9bd --- /dev/null +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/test/groovy/StandardSuspensionHandoffTest.groovy @@ -0,0 +1,4 @@ +import datadog.trace.instrumentation.kotlin.coroutines.AbstractStandardSuspensionHandoffTest + +class StandardSuspensionHandoffTest extends AbstractStandardSuspensionHandoffTest { +} diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/test/java/datadog/trace/instrumentation/kotlin/coroutines/DatadogThreadContextElementTest.java b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/test/java/datadog/trace/instrumentation/kotlin/coroutines/DatadogThreadContextElementTest.java new file mode 100644 index 00000000000..e45a22d4b23 --- /dev/null +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/test/java/datadog/trace/instrumentation/kotlin/coroutines/DatadogThreadContextElementTest.java @@ -0,0 +1,3 @@ +package datadog.trace.instrumentation.kotlin.coroutines; + +class DatadogThreadContextElementTest extends AbstractDatadogThreadContextElementTest {} diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/test/kotlin/KotlinCoroutineTests.kt b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/test/kotlin/KotlinCoroutineTests.kt index 8ba909ae74d..65f2352c6ed 100644 --- a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/test/kotlin/KotlinCoroutineTests.kt +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/test/kotlin/KotlinCoroutineTests.kt @@ -32,6 +32,18 @@ class KotlinCoroutineTests(dispatcher: CoroutineDispatcher) : CoreKotlinCoroutin 7 } + @Trace + override fun structuredChildrenOutliveSeedSpan(): Int = super.structuredChildrenOutliveSeedSpan() + + @Trace + override fun structuredChildrenOfScopedCoroutines(): Int = super.structuredChildrenOfScopedCoroutines() + + @Trace + override fun childWithReplacedJob(supervisor: Boolean): Int = super.childWithReplacedJob(supervisor) + + @Trace + override fun structuredChildOfCancelledCoroutine(): Int = super.structuredChildOfCancelledCoroutine() + @Trace override fun tracePreventedByCancellation(): Int = super.tracePreventedByCancellation() diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/groovy/datadog/trace/instrumentation/kotlin/coroutines/AbstractKotlinCoroutineInstrumentationTest.groovy b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/groovy/datadog/trace/instrumentation/kotlin/coroutines/AbstractKotlinCoroutineInstrumentationTest.groovy index 4298933db84..153b195d519 100644 --- a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/groovy/datadog/trace/instrumentation/kotlin/coroutines/AbstractKotlinCoroutineInstrumentationTest.groovy +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/groovy/datadog/trace/instrumentation/kotlin/coroutines/AbstractKotlinCoroutineInstrumentationTest.groovy @@ -1,5 +1,7 @@ package datadog.trace.instrumentation.kotlin.coroutines +import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace + import datadog.trace.agent.test.InstrumentationSpecification import datadog.trace.core.DDSpan import kotlinx.coroutines.CoroutineDispatcher @@ -11,6 +13,19 @@ abstract class AbstractKotlinCoroutineInstrumentationTest trace = TEST_WRITER.get(0) + + expect: + TEST_WRITER.size() == 1 + trace.size() == expectedNumberOfSpans + findSpan(trace, "child").spanContext().getParentId() == findSpan(trace, "seed").spanContext().getSpanId() + findSpan(trace, "grandchild").spanContext().getParentId() == findSpan(trace, "child").spanContext().getSpanId() + + where: + [dispatcherName, dispatcher] << dispatchersToTest + } + + def "child launched after a scoped coroutine keeps the trace open scoped=#scoped #dispatcherName"() { + setup: + CoreKotlinCoroutineTests kotlinTest = getCoreKotlinCoroutineTestsInstance(dispatcher) + int expectedNumberOfSpans = kotlinTest.structuredChildAfterScopedCoroutine(scoped) + kotlinTest.finishDetached() + TEST_WRITER.waitForTraces(1) + List trace = TEST_WRITER.get(0) + + expect: + TEST_WRITER.size() == 1 + trace.size() == expectedNumberOfSpans + findSpan(trace, "child").spanContext().getParentId() == findSpan(trace, "seed").spanContext().getSpanId() + findSpan(trace, "seed").spanContext().getParentId() == findSpan(trace, "root").spanContext().getSpanId() + + where: + [dispatcherName, dispatcher, scoped] << dispatchersToTest.collectMany { dispatcher -> + ["none", "withContext", "coroutineScope", "withTimeout"].collect { dispatcher + it } + } + } + + def "children of scoped coroutines keep the trace open #dispatcherName"() { + setup: + CoreKotlinCoroutineTests kotlinTest = getCoreKotlinCoroutineTestsInstance(dispatcher) + int expectedNumberOfSpans = kotlinTest.structuredChildrenOfScopedCoroutines() + kotlinTest.finishDetached() + TEST_WRITER.waitForTraces(1) + List trace = TEST_WRITER.get(0) + + expect: + TEST_WRITER.size() == 1 + trace.size() == expectedNumberOfSpans + findSpan(trace, "scope-child").spanContext().getParentId() == findSpan(trace, "seed").spanContext().getSpanId() + findSpan(trace, "context-child").spanContext().getParentId() == findSpan(trace, "seed").spanContext().getSpanId() + + where: + [dispatcherName, dispatcher] << dispatchersToTest + } + + def "structured child keeps the trace open when its parent started without one #dispatcherName"() { + setup: + CoreKotlinCoroutineTests kotlinTest = getCoreKotlinCoroutineTestsInstance(dispatcher) + int expectedNumberOfSpans = kotlinTest.structuredChildOfUntracedCoroutine() + kotlinTest.finishDetached() + TEST_WRITER.waitForTraces(1) + List trace = TEST_WRITER.get(0) + + expect: + TEST_WRITER.size() == 1 + trace.size() == expectedNumberOfSpans + findSpan(trace, "child").spanContext().getParentId() == findSpan(trace, "root").spanContext().getSpanId() + where: + [dispatcherName, dispatcher] << dispatchersToTest + } + + def "child with a replaced job keeps the trace open supervisor=#supervisor #dispatcherName"() { + setup: + CoreKotlinCoroutineTests kotlinTest = getCoreKotlinCoroutineTestsInstance(dispatcher) + int expectedNumberOfSpans = kotlinTest.childWithReplacedJob(supervisor) + kotlinTest.finishDetached() + TEST_WRITER.waitForTraces(1) + List trace = TEST_WRITER.get(0) + + expect: + TEST_WRITER.size() == 1 + trace.size() == expectedNumberOfSpans + findSpan(trace, "child").spanContext().getParentId() == findSpan(trace, "trace.annotation").spanContext().getSpanId() + + where: + [dispatcherName, dispatcher, supervisor] << dispatchersToTest.collectMany { [it + false, it + true] } + } + + def "cancelled coroutine keeps the trace open for its structured child #dispatcherName"() { + setup: + CoreKotlinCoroutineTests kotlinTest = getCoreKotlinCoroutineTestsInstance(dispatcher) + int expectedNumberOfSpans = kotlinTest.structuredChildOfCancelledCoroutine() + kotlinTest.cancelDetached() + TEST_WRITER.waitForTraces(1) + List trace = TEST_WRITER.get(0) + + expect: + TEST_WRITER.size() == 1 + trace.size() == expectedNumberOfSpans + findSpan(trace, "cleanup").spanContext().getParentId() == findSpan(trace, "trace.annotation").spanContext().getSpanId() + + where: + [dispatcherName, dispatcher] << dispatchersToTest + } + protected static DDSpan findSpan(List trace, String opName) { for (DDSpan span : trace) { if (span.getOperationName() == opName) { diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/groovy/datadog/trace/instrumentation/kotlin/coroutines/AbstractStandardSuspensionHandoffTest.groovy b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/groovy/datadog/trace/instrumentation/kotlin/coroutines/AbstractStandardSuspensionHandoffTest.groovy new file mode 100644 index 00000000000..c225f3b420a --- /dev/null +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/groovy/datadog/trace/instrumentation/kotlin/coroutines/AbstractStandardSuspensionHandoffTest.groovy @@ -0,0 +1,103 @@ +package datadog.trace.instrumentation.kotlin.coroutines + +import datadog.trace.agent.test.InstrumentationSpecification + +abstract class AbstractStandardSuspensionHandoffTest extends InstrumentationSpecification { + def "#kind child wait refreshes scope closed inside block #closeInside on #outcome"() { + when: + new StandardSuspensionHandoffTests().runScopedChildWait(kind, closeInside, outcome) + TEST_WRITER.waitForTraces(1) + def trace = TEST_WRITER.get(0) + def spans = trace.collectEntries { [(it.operationName.toString()): it] } + + then: + trace.size() == 3 + spans.closed.parentId == spans.parent.spanId + spans['after-child'].parentId == spans.parent.spanId + + where: + kind | closeInside | outcome + 'scope' | false | 'return' + 'scope' | true | 'return' + 'supervisor' | false | 'return' + 'supervisor' | true | 'return' + 'context' | false | 'return' + 'context' | true | 'return' + 'timeout' | false | 'return' + 'timeout' | true | 'return' + 'scope' | true | 'throw' + 'scope' | true | 'cancel' + } + + def "early dispatcher #outcome preserves scope opened afterward #openAfter with delayed restore #delayedRestore"() { + when: + new StandardSuspensionHandoffTests().runEarlyDispatcherChange(openAfter, outcome, delayedRestore) + TEST_WRITER.waitForTraces(1) + def trace = TEST_WRITER.get(0) + def spans = trace.collectEntries { [(it.operationName.toString()): it] } + + then: + trace.size() == 2 + spans.child.parentId == spans.manual.spanId + + where: + openAfter | outcome | delayedRestore + false | 'return' | false + true | 'return' | false + false | 'throw' | false + true | 'throw' | false + false | 'cancel' | false + true | 'cancel' | false + true | 'return' | true + } + + def "manual scope survives child completion after nested withContext"() { + when: + new StandardSuspensionHandoffTests().runNestedContextChange() + + then: + noExceptionThrown() + } + + def "child scope open across suspension does not leak into parent resumption"() { + when: + new StandardSuspensionHandoffTests().runChildScopeIsolation() + + then: + noExceptionThrown() + } + + def "dispatcher switch publishes manual scope before source restoration"() { + when: + new StandardSuspensionHandoffTests().runDispatcherChange() + + then: + noExceptionThrown() + } + + def "standard #kind suspension preserves manual scope with cancellation #cancelled and nesting #nested"() { + setup: + new StandardSuspensionHandoffTests().run(kind, cancelled, nested, { + assert TEST_WRITER.empty : "The finished trace must remain held until coroutine completion" + } as Runnable) + TEST_WRITER.waitForTraces(1) + def trace = TEST_WRITER.get(0) + def spans = trace.collectEntries { [(it.operationName.toString()): it] } + + expect: + TEST_WRITER.size() == 1 + trace.size() == (cancelled ? 2 : 3) + spans.manual.parentId == spans.parent.spanId + cancelled || spans.child.parentId == spans.manual.spanId + + where: + kind | cancelled | nested + 'cancellable' | false | false + 'cancellable' | true | false + 'safe' | false | false + 'yield' | false | false + 'delay' | false | false + 'select' | false | false + 'cancellable' | false | true + } +} diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/java/datadog/trace/instrumentation/kotlin/coroutines/AbstractDatadogThreadContextElementTest.java b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/java/datadog/trace/instrumentation/kotlin/coroutines/AbstractDatadogThreadContextElementTest.java new file mode 100644 index 00000000000..da1a2234bdd --- /dev/null +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/java/datadog/trace/instrumentation/kotlin/coroutines/AbstractDatadogThreadContextElementTest.java @@ -0,0 +1,166 @@ +package datadog.trace.instrumentation.kotlin.coroutines; + +import static datadog.trace.instrumentation.kotlin.coroutines.DatadogThreadContextElement.addDatadogElement; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotSame; +import static org.junit.jupiter.api.Assertions.assertSame; + +import java.util.ArrayList; +import java.util.List; +import kotlin.coroutines.AbstractCoroutineContextElement; +import kotlin.coroutines.Continuation; +import kotlin.coroutines.ContinuationInterceptor; +import kotlin.coroutines.CoroutineContext; +import kotlin.coroutines.EmptyCoroutineContext; +import kotlin.jvm.functions.Function2; +import kotlinx.coroutines.ThreadContextElement; +import org.junit.jupiter.api.Test; +import org.tabletest.junit.TableTest; + +/** Context replacement checks shared by the oldest and latest dependency suites. */ +public abstract class AbstractDatadogThreadContextElementTest { + @TableTest({ + "Scenario | Ordinary before | Ordinary after | Hooks before | Hooks after | Dispatcher | Custom view", + "Singleton | 0 | 0 | 0 | 0 | false | false ", + "Ordinary middle | 2 | 2 | 0 | 0 | true | false ", + "Later hook | 1 | 1 | 0 | 1 | true | false ", + "Hooks either side | 1 | 1 | 2 | 2 | false | false ", + "Custom later hook | 2 | 2 | 1 | 1 | true | true " + }) + void replacementPreservesContentsAndCallbackOrder( + int ordinaryBefore, + int ordinaryAfter, + int hooksBefore, + int hooksAfter, + boolean dispatcher, + boolean customView) { + DatadogThreadContextElement original = + (DatadogThreadContextElement) addDatadogElement(EmptyCoroutineContext.INSTANCE); + CoroutineContext input = EmptyCoroutineContext.INSTANCE; + for (int i = 0; i < ordinaryBefore; i++) { + input = input.plus(new Element()); + } + for (int i = 0; i < hooksBefore; i++) { + input = input.plus(new Hook()); + } + input = input.plus(original); + for (int i = 0; i < ordinaryAfter; i++) { + input = input.plus(new Element()); + } + for (int i = 0; i < hooksAfter; i++) { + input = input.plus(new Hook()); + } + if (dispatcher) { + input = input.plus(new Interceptor()); + } + if (customView) { + input = new ContextView(input); + } + assertReplacement(input, original); + } + + @Test + void contextWithoutDatadogElementRetainsItsContents() { + Element ordinary = new Element(); + Hook hook = new Hook(); + CoroutineContext input = ordinary.plus(hook); + CoroutineContext output = addDatadogElement(input); + assertSame(ordinary, lookup(output, ordinary)); + assertSame(hook, lookup(output, hook)); + assertEquals(3, keys(output).size()); + assertEquals(ordinary.getKey(), keys(output).get(0)); + assertEquals(hook.getKey(), keys(output).get(1)); + } + + private static void assertReplacement( + CoroutineContext input, DatadogThreadContextElement original) { + CoroutineContext output = addDatadogElement(input); + assertSame(original, lookup(input, original)); + assertNotSame(original, lookup(output, original)); + assertEquals(keys(input), keys(output)); + input.fold( + output, + (rebuilt, element) -> { + if (element != original) { + assertSame(element, lookup(rebuilt, element)); + } + return rebuilt; + }); + } + + @SuppressWarnings("unchecked") + private static CoroutineContext.Element lookup( + CoroutineContext context, CoroutineContext.Element element) { + return context.get((CoroutineContext.Key) element.getKey()); + } + + private static List> keys(CoroutineContext context) { + return context.fold( + new ArrayList>(), + (keys, element) -> { + keys.add(element.getKey()); + return keys; + }); + } + + private static class Element extends AbstractCoroutineContextElement { + Element() { + super(new CoroutineContext.Key() {}); + } + } + + private static final class Hook extends Element implements ThreadContextElement { + @Override + public Object updateThreadContext(CoroutineContext context) { + return null; + } + + @Override + public void restoreThreadContext(CoroutineContext context, Object state) {} + } + + private static final class Interceptor extends AbstractCoroutineContextElement + implements ContinuationInterceptor { + Interceptor() { + super(ContinuationInterceptor.Key); + } + + @SuppressWarnings("unchecked") + @Override + public Continuation interceptContinuation(Continuation continuation) { + return (Continuation) continuation; + } + + @Override + public void releaseInterceptedContinuation(Continuation continuation) {} + } + + private static final class ContextView implements CoroutineContext { + private final CoroutineContext delegate; + + ContextView(CoroutineContext delegate) { + this.delegate = delegate; + } + + @Override + public E get(CoroutineContext.Key key) { + return delegate.get(key); + } + + @Override + public R fold( + R initial, Function2 operation) { + return delegate.fold(initial, operation); + } + + @Override + public CoroutineContext minusKey(CoroutineContext.Key key) { + return delegate.minusKey(key); + } + + @Override + public CoroutineContext plus(CoroutineContext other) { + return CoroutineContext.DefaultImpls.plus(this, other); + } + } +} diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/kotlin/datadog/trace/instrumentation/kotlin/coroutines/CoreKotlinCoroutineTests.kt b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/kotlin/datadog/trace/instrumentation/kotlin/coroutines/CoreKotlinCoroutineTests.kt index 4e6e4fcda64..fbfabd5c199 100644 --- a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/kotlin/datadog/trace/instrumentation/kotlin/coroutines/CoreKotlinCoroutineTests.kt +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/kotlin/datadog/trace/instrumentation/kotlin/coroutines/CoreKotlinCoroutineTests.kt @@ -11,6 +11,9 @@ import kotlinx.coroutines.CoroutineName import kotlinx.coroutines.CoroutineScope import kotlinx.coroutines.CoroutineStart import kotlinx.coroutines.Deferred +import kotlinx.coroutines.Job +import kotlinx.coroutines.NonCancellable +import kotlinx.coroutines.SupervisorJob import kotlinx.coroutines.async import kotlinx.coroutines.awaitAll import kotlinx.coroutines.cancelAndJoin @@ -21,14 +24,36 @@ import kotlinx.coroutines.runBlocking import kotlinx.coroutines.selects.select import kotlinx.coroutines.sync.Mutex import kotlinx.coroutines.sync.withLock +import kotlinx.coroutines.withContext import kotlinx.coroutines.withTimeout import kotlinx.coroutines.yield import java.util.concurrent.CancellationException import java.util.concurrent.TimeUnit import java.util.concurrent.atomic.AtomicInteger +import kotlin.coroutines.CoroutineContext abstract class CoreKotlinCoroutineTests(private val dispatcher: CoroutineDispatcher) { + fun launchWithCancelledParent() { + val queuedDispatcher = object : CoroutineDispatcher() { + var pending: Runnable? = null + + override fun dispatch(context: CoroutineContext, block: Runnable) { + check(pending == null) { "Unexpected second dispatch" } + pending = block + } + } + val parent = Job() + parent.cancel() + + val child = CoroutineScope(parent + queuedDispatcher).launch { + error("A child of an already-cancelled parent must not execute") + } + + checkNotNull(queuedDispatcher.pending) { "Cancellation was not dispatched" }.run() + check(child.isCompleted && child.isCancelled) { "Child did not complete cancellation" } + } + @Trace open fun tracePreventedByCancellation(): Int { kotlin.runCatching { @@ -387,4 +412,155 @@ abstract class CoreKotlinCoroutineTests(private val dispatcher: CoroutineDispatc setAsyncPropagationEnabled(asyncPropagation) return runBlocking(jobName("test") + dispatcher, block = block) } + + // The following tests start a detached coroutine that outlives the method's root span. Its + // children wait on a gate that is only released by finishDetached(), after the root span has + // finished, so only coroutine continuations can keep the trace open until the children finish. + + private val gate = CompletableDeferred() + private var detached: Job? = null + + /** Releases the children of the detached coroutine and waits for it to complete. */ + fun finishDetached() { + gate.complete(Unit) + runBlocking { detached?.join() } + } + + /** Cancels the detached coroutine and waits for it to complete. */ + fun cancelDetached() { + detached?.cancel() + runBlocking { detached?.join() } + } + + /** Children launched with their parent's context, from a span the parent finishes early. */ + @Trace + open fun structuredChildrenOutliveSeedSpan(): Int { + runTest { + detached = launch(Job() + jobName("detached")) { + childSpan("seed").activateAndUse { + launch(coroutineContext + jobName("child")) { + gate.await() + childSpan("child").activateAndUse { + launch(coroutineContext + jobName("grandchild")) { + childSpan("grandchild").finish() + } + } + } + } + } + } + return 4 + } + + /** A structured child launched after a scoped coroutine has completed in its parent. */ + open fun structuredChildAfterScopedCoroutine(scoped: String): Int { + runTest { + childSpan("root").activateAndUse { + val seedFinished = CompletableDeferred() + detached = launch(Job() + jobName("detached")) { + childSpan("seed").activateAndUse { + when (scoped) { + "none" -> Unit + "withContext" -> withContext(jobName("scoped")) {} + "coroutineScope" -> coroutineScope {} + "withTimeout" -> withTimeout(30_000) {} + else -> error("Unknown scoped coroutine: $scoped") + } + launch(coroutineContext + jobName("child")) { + gate.await() + childSpan("child").finish() + } + } + seedFinished.complete(Unit) + } + // Finish both spans before opening the child's gate, regardless of dispatcher timing. + seedFinished.await() + } + } + return 3 + } + + /** Children launched with the context of scoped coroutines inside their parent. */ + @Trace + open fun structuredChildrenOfScopedCoroutines(): Int { + runTest { + detached = launch(Job() + jobName("detached")) { + childSpan("seed").activateAndUse { + coroutineScope { + launch(coroutineContext + jobName("scope-child")) { + gate.await() + childSpan("scope-child").finish() + } + } + withContext(jobName("with-context")) { + launch(coroutineContext + jobName("context-child")) { + gate.await() + childSpan("context-child").finish() + } + } + } + } + } + return 4 + } + + /** A child launched with its parent's context, when the parent started without a trace. */ + open fun structuredChildOfUntracedCoroutine(): Int { + runTest { + detached = launch(Job() + jobName("detached")) { + childSpan("root").activateAndUse { + val started = CompletableDeferred() + launch(coroutineContext + jobName("child")) { + started.complete(Unit) + gate.await() + childSpan("child").finish() + } + // the child's first dispatch can keep the trace open, so finish the root after it + started.await() + } + } + } + return 2 + } + + /** A child launched with its parent's context but a different job, so not a structured child. */ + @Trace + open fun childWithReplacedJob(supervisor: Boolean): Int { + runTest { + val job = if (supervisor) SupervisorJob() else Job() + val started = CompletableDeferred() + detached = launch(coroutineContext + job + jobName("detached")) { + started.complete(Unit) + gate.await() + childSpan("child").finish() + } + // the child's first dispatch can keep the trace open, so finish the root after it + started.await() + } + return 2 + } + + /** A structured child that only adds its span while its cancelled parent waits for it. */ + @Trace + open fun structuredChildOfCancelledCoroutine(): Int { + runTest { + val started = CompletableDeferred() + detached = launch(Job() + jobName("detached")) { + launch(coroutineContext + jobName("child")) { + try { + started.complete(Unit) + CompletableDeferred().await() + } finally { + withContext(NonCancellable) { + delay(50) + childSpan("cleanup").finish() + } + } + } + } + // a child cancelled before it starts never runs its finally block + started.await() + } + return 2 + } } diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/kotlin/datadog/trace/instrumentation/kotlin/coroutines/StandardSuspensionHandoffTests.kt b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/kotlin/datadog/trace/instrumentation/kotlin/coroutines/StandardSuspensionHandoffTests.kt new file mode 100644 index 00000000000..8a81fc7f530 --- /dev/null +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/kotlin/datadog/trace/instrumentation/kotlin/coroutines/StandardSuspensionHandoffTests.kt @@ -0,0 +1,561 @@ +package datadog.trace.instrumentation.kotlin.coroutines + +import datadog.context.Context +import datadog.context.ContextKey +import datadog.context.ContextScope +import datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan +import datadog.trace.bootstrap.instrumentation.api.AgentTracer.get +import kotlinx.coroutines.CancellationException +import kotlinx.coroutines.CoroutineDispatcher +import kotlinx.coroutines.CoroutineExceptionHandler +import kotlinx.coroutines.CoroutineName +import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.CoroutineStart +import kotlinx.coroutines.ExperimentalCoroutinesApi +import kotlinx.coroutines.Job +import kotlinx.coroutines.ThreadContextElement +import kotlinx.coroutines.asCoroutineDispatcher +import kotlinx.coroutines.coroutineScope +import kotlinx.coroutines.delay +import kotlinx.coroutines.launch +import kotlinx.coroutines.newCoroutineContext +import kotlinx.coroutines.runBlocking +import kotlinx.coroutines.selects.select +import kotlinx.coroutines.supervisorScope +import kotlinx.coroutines.suspendCancellableCoroutine +import kotlinx.coroutines.withContext +import kotlinx.coroutines.withTimeout +import kotlinx.coroutines.yield +import java.util.concurrent.ConcurrentHashMap +import java.util.concurrent.CountDownLatch +import java.util.concurrent.Executors +import java.util.concurrent.TimeUnit.SECONDS +import java.util.concurrent.atomic.AtomicInteger +import java.util.concurrent.atomic.AtomicReference +import kotlin.coroutines.AbstractCoroutineContextElement +import kotlin.coroutines.Continuation +import kotlin.coroutines.ContinuationInterceptor +import kotlin.coroutines.CoroutineContext +import kotlin.coroutines.resume +import kotlin.coroutines.suspendCoroutine + +private val TOUCH_KEY = ContextKey.named("scope-touching") +private val BASE_KEY = ContextKey.named("scope-touching-base") + +@OptIn(ExperimentalCoroutinesApi::class) +class StandardSuspensionHandoffTests { + fun runScopedChildWait(kind: String, closeInside: Boolean, outcome: String) { + val workers = Executors.newFixedThreadPool(2) + val childWorkers = Executors.newSingleThreadExecutor() + val dispatcher = workers.asCoroutineDispatcher() + val childDispatcher = childWorkers.asCoroutineDispatcher() + val restoring = CountDownLatch(1) + val released = CountDownLatch(1) + val childStarted = CountDownLatch(1) + val expected = when (outcome) { + "return" -> null + "throw" -> IllegalStateException("scoped body failed") + "cancel" -> CancellationException("scoped body cancelled") + else -> error("Unknown completion: $outcome") + } + val before = RestoreGate(restoring, released) + val after = RestoreGate(restoring, released) + val context = CoroutineScope(dispatcher).newCoroutineContext(before) + after + val parent = get().buildSpan("kotlin_coroutine", "parent").start() + val parentScope = get().activateManualSpan(parent) + try { + runBlocking(context) { + suspend fun closeCapturedScope() { + val transient = get().buildSpan("kotlin_coroutine", "closed").start() + val scope = get().activateManualSpan(transient) + try { + suspendCoroutine { it.resume(Unit) } + } finally { + scope.close() + transient.finish() + } + check(activeSpan() === parent) { "Inline completion did not restore parent" } + } + if (!closeInside) closeCapturedScope() + val block: suspend CoroutineScope.() -> Unit = { + if (closeInside) closeCapturedScope() + launch(childDispatcher) { + childStarted.countDown() + check(restoring.await(10, SECONDS)) { "Source did not enter restoration" } + } + check(childStarted.await(10, SECONDS)) { "Child did not start" } + if (expected != null) throw expected + Unit + } + try { + try { + when (kind) { + "scope" -> coroutineScope(block) + "supervisor" -> supervisorScope(block) + "context" -> withContext(CoroutineName("nested"), block) + "timeout" -> withTimeout(10000, block) + else -> error("Unknown scoped builder: $kind") + } + check(expected == null) { "Scoped failure was swallowed" } + } catch (failure: Exception) { + check(expected != null && (failure === expected || failure.cause === expected)) { + "Unexpected scoped failure: $failure" + } + } + check(restoring.count == 0L && released.count == 1L) { "Source restoration did not overlap resumption" } + check(before.updates.get() >= 2 && after.updates.get() >= 2) { "Parent did not resume" } + check(activeSpan() === parent) { "Child wait resumed ${activeSpan()?.operationName} instead of parent" } + get().buildSpan("kotlin_coroutine", "after-child").start().finish() + } finally { + released.countDown() + } + } + } finally { + released.countDown() + restoring.countDown() + parentScope.close() + parent.finish() + dispatcher.close() + childDispatcher.close() + check(workers.awaitTermination(10, SECONDS)) + check(childWorkers.awaitTermination(10, SECONDS)) + } + } + + fun runReplayedRestoration() { + val interceptor = object : AbstractCoroutineContextElement(ContinuationInterceptor), ContinuationInterceptor { + override fun interceptContinuation(continuation: Continuation): Continuation = continuation + } + var saved: Continuation? = null + var failure: Throwable? = null + var completed = false + val parent = get().buildSpan("kotlin_coroutine", "parent").start() + val parentScope = get().activateManualSpan(parent) + val handler = CoroutineExceptionHandler { _, error -> failure = error } + try { + val job = CoroutineScope(interceptor + handler).launch { + withContext(CoroutineName("nested")) { + suspendCoroutine { saved = it } + } + check(activeSpan() === parent) { "Caller continuation lost its context" } + completed = true + } + check(saved != null) { "The test did not suspend" } + // Resume without a dispatcher installing context. Kotlin must replay its saved restore token. + val original = Context.root().swap() + try { + saved!!.resume(Unit) + check(failure == null) { "Coroutine failed: $failure" } + check(completed && job.isCompleted) { "Caller continuation did not complete" } + check(activeSpan() === parent) { "Replayed restoration lost the caller thread context" } + } finally { + original.swap() + } + } finally { + parentScope.close() + parent.finish() + } + } + + fun runEarlyDispatcherChange(openAfter: Boolean, outcome: String, delayedRestore: Boolean) { + val sourceWorkers = Executors.newSingleThreadExecutor() + val destinationWorkers = Executors.newSingleThreadExecutor() + val source = sourceWorkers.asCoroutineDispatcher() + val restoring = CountDownLatch(1) + val released = CountDownLatch(1) + val destination = object : CoroutineDispatcher() { + override fun dispatch(context: CoroutineContext, block: Runnable) { + // Run the real destination task, including restoration, before source getResult can execute. + val task = destinationWorkers.submit(block) + if (delayedRestore) { + check(restoring.await(10, SECONDS)) { "Destination did not reach restoration" } + } else { + task.get(10, SECONDS) + } + } + } + val context: CoroutineContext = if (delayedRestore) { + // Gate both restoration orders: Kotlin 1.3 restores forwards, newer versions backwards. + val before = DestinationRestoreGate(destination, restoring, released) + val after = DestinationRestoreGate(destination, restoring, released) + CoroutineScope(source).newCoroutineContext(before) + after + } else { + source + } + val expected = when (outcome) { + "return" -> null + "throw" -> IllegalStateException("destination failure") + "cancel" -> CancellationException("destination cancellation") + else -> error("Unknown completion: $outcome") + } + try { + runBlocking(context) { + var span = if (openAfter) null else get().buildSpan("kotlin_coroutine", "manual").start() + var scope = span?.let { get().activateManualSpan(it) } + try { + try { + withContext(destination) { + if (!openAfter) { + check(activeSpan() === span) { "Destination lost the caller scope before getResult" } + } + if (expected != null) throw expected + } + check(expected == null) { "Destination failure was swallowed" } + } catch (failure: Exception) { + check(failure === expected) { "Unexpected destination failure: $failure" } + } + if (openAfter) { + span = get().buildSpan("kotlin_coroutine", "manual").start() + scope = get().activateManualSpan(span) + } + check(activeSpan() === span) { "Caller scope was lost on synchronous completion" } + suspendCancellableCoroutine { continuation -> + // This queue entry cannot run until source restoration returns its worker. + sourceWorkers.execute { continuation.resume(Unit) } + } + check(activeSpan() === span) { "Caller scope was lost after the next suspension" } + if (delayedRestore) { + check(restoring.count == 0L && released.count == 1L) { "Destination cleanup was not delayed" } + } + get().buildSpan("kotlin_coroutine", "child").start().finish() + } finally { + scope?.close() + span?.finish() + released.countDown() + } + } + } finally { + released.countDown() + source.close() + destinationWorkers.shutdown() + check(sourceWorkers.awaitTermination(10, SECONDS)) + check(destinationWorkers.awaitTermination(10, SECONDS)) + } + } + + fun runNestedContextChange() { + val workers = Executors.newSingleThreadExecutor() + val otherWorkers = Executors.newSingleThreadExecutor() + val dispatcher = workers.asCoroutineDispatcher() + val otherDispatcher = otherWorkers.asCoroutineDispatcher() + val childRelease = CountDownLatch(1) + try { + runBlocking(dispatcher) { + withContext(CoroutineName("nested")) { Unit } + val manual = get().buildSpan("kotlin_coroutine", "manual").start() + val scope = get().activateManualSpan(manual) + try { + coroutineScope { + launch(otherDispatcher, start = CoroutineStart.ATOMIC) { + check(childRelease.await(10, SECONDS)) + } + // The child can finish only after this coroutine returns its worker. + workers.execute { childRelease.countDown() } + } + check(activeSpan() === manual) { "Waiting for the child lost the manual scope" } + } finally { + scope.close() + manual.finish() + } + } + } finally { + childRelease.countDown() + dispatcher.close() + otherDispatcher.close() + check(workers.awaitTermination(10, SECONDS)) + check(otherWorkers.awaitTermination(10, SECONDS)) + } + } + + fun runChildScopeIsolation() { + val workers = Executors.newSingleThreadExecutor() + val childWorkers = Executors.newSingleThreadExecutor() + val dispatcher = workers.asCoroutineDispatcher() + val childDispatcher = childWorkers.asCoroutineDispatcher() + val parentSuspended = CountDownLatch(1) + val childContinuation = AtomicReference>() + val parent = get().buildSpan("kotlin_coroutine", "parent").start() + val parentScope = get().activateManualSpan(parent) + try { + runBlocking(dispatcher + SuspendGate(parentSuspended)) { + try { + suspendCancellableCoroutine { parentContinuation -> + // The child inherits the parent's Datadog element; override the gate so only the parent signals. + launch(childDispatcher + SuspendGate(CountDownLatch(0))) { + val child = get().buildSpan("kotlin_coroutine", "child").start() + val childScope = get().activateManualSpan(child) + try { + // Suspend with the child scope still open. + suspendCancellableCoroutine { continuation -> + childContinuation.set(continuation) + // Queued behind this task, so it runs only after the child has restored its worker. + childWorkers.execute { + check(parentSuspended.await(10, SECONDS)) { "Parent did not suspend" } + parentContinuation.resume(Unit) + } + } + } finally { + childScope.close() + child.finish() + } + } + } + check(activeSpan() === parent) { "Parent resumed with ${activeSpan()?.operationName} instead of parent" } + } finally { + childContinuation.get()?.resume(Unit) + } + } + } finally { + parentScope.close() + parent.finish() + dispatcher.close() + childDispatcher.close() + check(workers.awaitTermination(10, SECONDS)) + check(childWorkers.awaitTermination(10, SECONDS)) + } + } + + /** + * Runs a coroutine between two elements that attach and close Datadog scopes in their own + * update/restore. With newer kotlinx versions the element after Datadog closes its scope on the + * coroutine stack after the suspension has been published, racing with the next worker reading + * that stack. Returns counts of each kind of violation seen, empty if none. + */ + fun runScopeTouchingElements(iterations: Int): Map { + val workers = Executors.newFixedThreadPool(4) + val resumers = Executors.newSingleThreadExecutor() + val dispatcher = workers.asCoroutineDispatcher() + val violations = ConcurrentHashMap() + val record: (String) -> Unit = { kind -> violations.computeIfAbsent(kind) { AtomicInteger() }.incrementAndGet() } + val before = ScopeTouchingElement("before", record) + val after = ScopeTouchingElement("after", record) + val context = CoroutineScope(dispatcher).newCoroutineContext(before) + after + val base = Context.root().with(BASE_KEY, "base") + val baseScope = base.attach() + try { + runBlocking(context) { + repeat(iterations) { + // Suspend with no scope open in the body, so the element's restore pops the coroutine + // stack itself. Resume straight away so the next worker races that restore. + suspendCancellableCoroutine { continuation -> + resumers.execute { continuation.resume(Unit) } + } + // What Datadog installed for this activation, as seen by the element updated after it. + val installed = after.below.get() + when { + installed === base -> Unit + installed.get(TOUCH_KEY) != null -> record("element scope leaked from previous worker") + installed.get(BASE_KEY) == null -> record("coroutine base context lost") + else -> record("unexpected coroutine context") + } + } + } + } catch (failure: Throwable) { + record("coroutine failed: ${failure.javaClass.name}") + } finally { + baseScope.close() + dispatcher.close() + resumers.shutdown() + check(workers.awaitTermination(10, SECONDS)) + check(resumers.awaitTermination(10, SECONDS)) + } + return violations.mapValues { it.value.get() } + } + + fun runDispatcherChange() { + val workers = Executors.newSingleThreadExecutor() + val otherWorkers = Executors.newSingleThreadExecutor() + val dispatcher = workers.asCoroutineDispatcher() + val otherDispatcher = otherWorkers.asCoroutineDispatcher() + val restoring = CountDownLatch(1) + val released = CountDownLatch(1) + val before = RestoreGate(restoring, released) + val after = RestoreGate(restoring, released) + val context = CoroutineScope(dispatcher).newCoroutineContext(before) + after + // withContext dispatches before getResult. Delay the destination until getResult has run, + // while the source's Datadog restoration is still blocked by one of the gates. + val destinationGate = otherWorkers.submit { check(restoring.await(10, SECONDS)) } + try { + runBlocking(context) { + val manual = get().buildSpan("kotlin_coroutine", "manual").start() + val scope = get().activateManualSpan(manual) + try { + withContext(otherDispatcher) { + try { + check(restoring.count == 0L) { "The source did not suspend" } + check(released.count == 1L) { "The source restoration was not delayed" } + check(activeSpan() === manual) { "Dispatcher switch lost the manual scope" } + } finally { + released.countDown() + } + } + } finally { + scope.close() + manual.finish() + } + } + destinationGate.get(10, SECONDS) + check(before.updates.get() >= 2) + check(after.updates.get() >= 2) + } finally { + released.countDown() + restoring.countDown() + dispatcher.close() + otherDispatcher.close() + check(workers.awaitTermination(10, SECONDS)) + check(otherWorkers.awaitTermination(10, SECONDS)) + } + } + + fun run(kind: String, cancelled: Boolean, nested: Boolean, beforeCompletion: Runnable) { + val workers = Executors.newFixedThreadPool(2) + val dispatcher = workers.asCoroutineDispatcher() + val restoring = CountDownLatch(1) + val released = CountDownLatch(1) + val continuation = AtomicReference>() + val before = RestoreGate(restoring, released) + val after = RestoreGate(restoring, released) + // Place a gate on either side of the auto-installed Datadog element. Kotlin 1.3 restores + // elements forwards and 1.6 restores backwards, so either version delays Datadog restoration. + val context = CoroutineScope(dispatcher).newCoroutineContext(before) + after + val resumer = Thread { + check(restoring.await(10, SECONDS)) + continuation.get()?.let { + if (cancelled) it.context[Job]?.cancel() + it.resume(Unit) + } + } + resumer.start() + val parent = get().buildSpan("kotlin_coroutine", "parent").start() + val parentScope = get().activateManualSpan(parent) + try { + try { + runBlocking(context) { + val suspendWithManualScope: suspend () -> Unit = { + parent.finish() + val manual = get().buildSpan("kotlin_coroutine", "manual").start() + val scope = get().activateManualSpan(manual) + try { + when (kind) { + "cancellable" -> suspendCancellableCoroutine { continuation.set(it) } + "safe" -> suspendCoroutine { continuation.set(it) } + "yield" -> yield() + "delay" -> delay(20) + "select" -> select { onTimeout(20) { Unit } } + else -> error("Unknown suspension API: $kind") + } + check(activeSpan() === manual) { "Manual scope was lost across $kind" } + get().buildSpan("kotlin_coroutine", "child").start().finish() + } finally { + try { + try { + check(activeSpan() === manual) { "Cleanup lost the manual scope across $kind" } + } finally { + try { + scope.close() + } finally { + manual.finish() + } + } + resumer.join(10000) + check(!resumer.isAlive) + beforeCompletion.run() + } finally { + released.countDown() + } + } + } + if (nested) { + withContext(CoroutineName("nested")) { suspendWithManualScope() } + } else { + suspendWithManualScope() + } + } + check(!cancelled) + } catch (expected: CancellationException) { + check(cancelled) + } + check(before.updates.get() >= 2) { "The test did not suspend" } + check(after.updates.get() >= 2) { "The test did not suspend" } + } finally { + released.countDown() + restoring.countDown() + resumer.join(10000) + parentScope.close() + dispatcher.close() + check(workers.awaitTermination(10, SECONDS)) + } + } + + private class Touch(val scope: ContextScope, val below: Context, val token: Any) + + /** Attaches a Datadog scope in update and closes it in restore, checking it was on top. */ + private class ScopeTouchingElement( + private val name: String, + private val record: (String) -> Unit + ) : AbstractCoroutineContextElement(object : CoroutineContext.Key {}), + ThreadContextElement { + val below = ThreadLocal() + + override fun updateThreadContext(context: CoroutineContext): Touch { + val current = Context.current() + below.set(current) + val token = Any() + return Touch(current.with(TOUCH_KEY, token).attach(), current, token) + } + + override fun restoreThreadContext(context: CoroutineContext, oldState: Touch) { + if (Context.current().get(TOUCH_KEY) !== oldState.token) { + record("$name scope not on top at restore") + } + oldState.scope.close() + if (Context.current() !== oldState.below) { + record("$name restore did not return to previous context") + } + } + } + + private class SuspendGate( + private val suspended: CountDownLatch + ) : AbstractCoroutineContextElement(SuspendGate), + ThreadContextElement { + companion object Key : CoroutineContext.Key + + override fun updateThreadContext(context: CoroutineContext) = Unit + + override fun restoreThreadContext(context: CoroutineContext, oldState: Unit) { + suspended.countDown() + } + } + + private class DestinationRestoreGate( + private val destination: CoroutineDispatcher, + private val restoring: CountDownLatch, + private val released: CountDownLatch + ) : AbstractCoroutineContextElement(object : CoroutineContext.Key {}), + ThreadContextElement { + override fun updateThreadContext(context: CoroutineContext): Boolean = context[ContinuationInterceptor] === destination + + override fun restoreThreadContext(context: CoroutineContext, oldState: Boolean) { + if (oldState) { + restoring.countDown() + check(released.await(10, SECONDS)) { "Caller did not finish before destination restoration" } + } + } + } + + private class RestoreGate( + private val restoring: CountDownLatch, + private val released: CountDownLatch + ) : AbstractCoroutineContextElement(object : CoroutineContext.Key {}), + ThreadContextElement { + val updates = AtomicInteger() + override fun updateThreadContext(context: CoroutineContext): Boolean = updates.incrementAndGet() == 1 + override fun restoreThreadContext(context: CoroutineContext, oldState: Boolean) { + if (oldState) { + restoring.countDown() + check(released.await(10, SECONDS)) { "The new worker did not finish before the old restoration" } + } + } + } +} diff --git a/dd-trace-core/src/jmh/java/datadog/context/ContextManagerBenchmark.java b/dd-trace-core/src/jmh/java/datadog/context/ContextManagerBenchmark.java index 0128535a9d0..1764b66f32a 100644 --- a/dd-trace-core/src/jmh/java/datadog/context/ContextManagerBenchmark.java +++ b/dd-trace-core/src/jmh/java/datadog/context/ContextManagerBenchmark.java @@ -235,10 +235,10 @@ public void nestedAttachAndClose(ThreadState thread) { * ContinuableScopeManager#swap(Context)} replaces the entire scope stack, making this a heavier * operation than in {@link ThreadLocalContextManager}. * - *

Note: GCProfiler will show allocation asymmetry here by design. {@link - * ContinuableScopeManager} swap allocates a {@code ScopeStack}, a {@code ContinuableScope}, and a - * {@code ScopeContext} per invocation; {@link ThreadLocalContextManager} swap is a plain field - * write. That asymmetry is the real cost of each manager's swap operation, not scaffolding. + *

This scenario uses ordinary contexts: entering creates a new scope stack and seed scope, and + * each swap creates a saved-stack token. The unused return token may be optimized away. + * Shared-stack swaps reuse their stack and restoration tokens and are not measured here. {@link + * ThreadLocalContextManager} swap is a plain field write. */ @Benchmark public void swapContexts(ThreadState thread) { diff --git a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java index 8f52b49052c..995f5a2c8fa 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java @@ -2687,4 +2687,13 @@ public Context swap(Context context) { public ContextContinuation capture(Context context) { return scopeManager.capture(context); } + + @Override + public Context sharedContextStack(Context initial) { + if (!InstrumenterConfig.get().isLegacyContextManagerEnabled()) { + throw new IllegalStateException( + "sharedContextStack must not be called when context swap based logic is enabled"); + } + return scopeManager.sharedContextStack(initial); + } } diff --git a/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuableScope.java b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuableScope.java index 0f6feaf095c..cfacc375b47 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuableScope.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuableScope.java @@ -6,6 +6,7 @@ import datadog.trace.api.scopemanager.ExtendedScopeListener; import datadog.trace.api.scopemanager.ScopeListener; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; class ContinuableScope implements ContextScope { @@ -27,6 +28,9 @@ class ContinuableScope implements ContextScope { private short referenceCount = 1; + // set when pushed onto a coroutine's shared scope stack, which other threads can also update + private boolean shared; + private final Stateful scopeState; ContinuableScope( @@ -50,7 +54,7 @@ public final void close() { if (!scopeStack.checkTop(this) && !scopeStack.checkOverdueScopes(this)) { if (ContinuableScopeManager.log.isDebugEnabled()) { ContinuableScopeManager.log.debug( - "Tried to close {} scope when not on top. Current top: {}", this, scopeStack.top); + "Tried to close {} scope when not on top. Current top: {}", this, scopeStack.top()); } byte source = source(); @@ -60,7 +64,7 @@ public final void close() { } } - final boolean alive = decrementReferences(); + final boolean alive = release(); scopeManager.healthMetrics.onCloseScope(); if (!alive) { scopeState.close(); @@ -105,20 +109,78 @@ final void onProperClose() { } } - final void incrementReferences() { + @SuppressFBWarnings( + value = "AT_STALE_THREAD_WRITE_OF_PRIMITIVE", + justification = + "Set before publishing the scope through SharedScopeStack.liveTop; the flag never " + + "changes after publication.") + final void markShared() { + shared = true; + } + + /** Adds a reference to a live scope -- returns false if the scope was already closed */ + @SuppressFBWarnings( + value = "AT_STALE_THREAD_WRITE_OF_PRIMITIVE", + justification = + "Only non-shared scopes use the unlocked write; shared scopes retain under their " + + "own monitor in retainShared().") + final boolean retain() { + if (shared) { + return retainShared(); + } + ++referenceCount; + return true; + } + + private synchronized boolean retainShared() { + // another thread may have closed the scope since it was read from the shared stack + if (referenceCount <= 0) { + return false; + } ++referenceCount; + return true; } - /** Decrements ref count -- returns true if the scope is still alive */ - final boolean decrementReferences() { + /** Removes a reference -- returns true if the scope is still alive */ + @SuppressFBWarnings( + value = "AT_STALE_THREAD_WRITE_OF_PRIMITIVE", + justification = + "Only non-shared scopes use the unlocked write; shared scopes release under their " + + "own monitor in releaseShared(), including closes from another stack.") + final boolean release() { + return shared ? releaseShared() : --referenceCount > 0; + } + + private synchronized boolean releaseShared() { + // the scope may be closed on any thread, whichever stack that thread has swapped in return --referenceCount > 0; } + @SuppressFBWarnings( + value = "AT_STALE_THREAD_WRITE_OF_PRIMITIVE", + justification = + "Only non-shared scopes use the unlocked write; shared scopes clear under their " + + "own monitor in clearShared().") final void clearReferences() { + if (shared) { + clearShared(); + } else { + referenceCount = 0; + } + } + + private synchronized void clearShared() { + // stops another thread retaining the scope once it's been cleared referenceCount = 0; } - /** Returns true if the scope is still alive (non-zero ref count) */ + /** + * Returns true if the scope's reference count is positive. + * + *

For shared scopes, this unlocked read is a best-effort cleanup check. Once closed, a shared + * scope cannot be retained again, so a stale read can delay cleanup but cannot cause a live scope + * to be removed. + */ final boolean alive() { return referenceCount > 0; } @@ -136,6 +198,11 @@ public Context context() { return context; } + @SuppressFBWarnings( + value = "AT_STALE_THREAD_WRITE_OF_PRIMITIVE", + justification = + "Changes are serialized with scope execution, including coroutine handoff between " + + "workers. Overlapping scope close/restore callbacks do not modify this flag.") public final void setAsyncPropagation(final boolean value) { asyncPropagating = value; } @@ -145,10 +212,20 @@ public final String toString() { return super.toString() + "->" + context; } + @SuppressFBWarnings( + value = "AT_STALE_THREAD_WRITE_OF_PRIMITIVE", + justification = + "Checkpoint updates are serialized with scope execution; coroutine handoff orders " + + "them between workers. Overlapping scope close/restore does not modify checkpoints.") public void checkpoint() { checkpointCount++; } + @SuppressFBWarnings( + value = "AT_STALE_THREAD_WRITE_OF_PRIMITIVE", + justification = + "Checkpoint updates are serialized with scope execution; coroutine handoff orders " + + "them between workers. Overlapping scope close/restore does not modify checkpoints.") public boolean rollback() { if (checkpointCount > 0) { checkpointCount--; diff --git a/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuableScopeManager.java b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuableScopeManager.java index e0ac7ea14fa..7d156b1df8f 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuableScopeManager.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuableScopeManager.java @@ -119,9 +119,8 @@ private ContextScope activate( ScopeStack scopeStack = scopeStack(); - final ContinuableScope top = scopeStack.top; - if (top != null && span.equals(top.span())) { - top.incrementReferences(); + final ContinuableScope top = scopeStack.top(); + if (top != null && span.equals(top.span()) && top.retain()) { return top; } @@ -154,9 +153,8 @@ private ContextScope activate( private ContextScope activate(final Context context) { ScopeStack scopeStack = scopeStack(); - final ContinuableScope top = scopeStack.top; - if (top != null && top.context.equals(context)) { - top.incrementReferences(); + final ContinuableScope top = scopeStack.top(); + if (top != null && top.context.equals(context) && top.retain()) { return top; } @@ -194,9 +192,8 @@ ContinuableScope continueSpan( // optimization: if the top scope is already keeping the same span alive // then re-use that scope (avoids allocation) and cancel the continuation - final ContinuableScope top = scopeStack.top; - if (top != null && top.context.equals(context)) { - top.incrementReferences(); + final ContinuableScope top = scopeStack.top(); + if (top != null && top.context.equals(context) && top.retain()) { if (continuation != null) { continuation.cancelFromContinuedScopeClose(); } @@ -231,7 +228,7 @@ public void closePrevious(final boolean finishSpan) { ScopeStack scopeStack = scopeStack(); // close any immediately previous iteration scope - final ContinuableScope top = scopeStack.top; + final ContinuableScope top = scopeStack.top(); if (top != null && top.source() == ITERATION) { if (iterationKeepAlive > 0) { // skip depth check because cancelling is cheap cancelRootIterationScopeCleanup(scopeStack, top); @@ -264,7 +261,7 @@ public ContextScope activateNext(final AgentSpan span) { assert span != null; - final ContinuableScope top = scopeStack.top; + final ContinuableScope top = scopeStack.top(); boolean asyncPropagation = top != null ? top.isAsyncPropagating() : DEFAULT_ASYNC_PROPAGATING; @@ -362,32 +359,34 @@ public ContextScope attach(@Nonnull Context context) { public Context swap(@Nonnull Context context) { ScopeStack oldStack = tlsScopeStack.get(); - ContinuableScope oldScope = oldStack.top; + ContinuableScope oldScope = oldStack.top(); ScopeStack newStack; - ContinuableScope newScope; - if (context instanceof ScopeContext) { + if (context instanceof ScopeStackContext) { // restore previously swapped context stack - newStack = ((ScopeContext) context).restore(profilingContextIntegration); - newScope = newStack.top; + newStack = ((ScopeStackContext) context).restore(profilingContextIntegration); } else if (context != Context.root()) { // start a new stack and record the new context as active newStack = new ScopeStack(profilingContextIntegration); - newScope = new ContinuableScope(this, context, CONTEXT, true, createScopeState(context)); + ContinuableScope newScope = + new ContinuableScope(this, context, CONTEXT, true, createScopeState(context)); newStack.top = newScope; } else { // start a new stack with no active context newStack = new ScopeStack(profilingContextIntegration); - newScope = null; } tlsScopeStack.set(newStack); + if (newStack instanceof SharedScopeStack) { + ((SharedScopeStack) newStack).pruneClosedTop(); + } + ContinuableScope newScope = newStack.top(); if (oldScope != newScope && newScope != null) { newScope.beforeActivated(); newScope.afterActivated(); } - return new ScopeContext(oldStack); + return newStack.replacing(oldStack); } public ContextContinuation capture(@Nonnull Context context) { @@ -410,6 +409,22 @@ public ContextContinuation capture(@Nonnull Context context) { return new ScopeContinuation(this, context, CONTEXT, traceCollector).register(); } + /** + * Returns a context whose scope stack is shared by every thread it is swapped onto. If the + * initial context is already shared, the new stack starts from its current top. + */ + public Context sharedContextStack(@Nonnull Context initial) { + if (initial instanceof SharedScopeContext) { + initial = ((SharedScopeContext) initial).currentContext(); + } + SharedScopeStack sharedStack = new SharedScopeStack(profilingContextIntegration); + if (initial != Context.root()) { + sharedStack.seed( + new ContinuableScope(this, initial, CONTEXT, true, createScopeState(initial))); + } + return sharedStack.handle; + } + static final class ScopeStackThreadLocal extends ThreadLocal { private final ProfilingContextIntegration profilingContextIntegration; diff --git a/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ScopeContext.java b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/SavedScopeContext.java similarity index 83% rename from dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ScopeContext.java rename to dd-trace-core/src/main/java/datadog/trace/core/scopemanager/SavedScopeContext.java index 8919fae96f4..7288bca3985 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ScopeContext.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/SavedScopeContext.java @@ -6,24 +6,25 @@ import javax.annotation.Nullable; /** Wraps a {@link ScopeStack} as a {@link Context} so it can be swapped back later. */ -final class ScopeContext implements Context { +final class SavedScopeContext extends ScopeStackContext { private final Thread originalThread = Thread.currentThread(); private final ScopeStack scopeStack; private final ContinuableScope parent; private final ContinuableScope active; private final Context context; - ScopeContext(ScopeStack scopeStack) { - this(scopeStack, scopeStack.top != null ? scopeStack.top.context : Context.root()); + SavedScopeContext(ScopeStack scopeStack) { + this(scopeStack, contextOf(scopeStack.top())); } - private ScopeContext(ScopeStack scopeStack, Context context) { + private SavedScopeContext(ScopeStack scopeStack, Context context) { this.scopeStack = scopeStack; this.parent = scopeStack.parent(); this.active = scopeStack.active(); this.context = context; } + @Override ScopeStack restore(ProfilingContextIntegration profilingContextIntegration) { // restore full stack on original thread, for other threads restore shallow copy if (Thread.currentThread() == originalThread) { diff --git a/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ScopeStack.java b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ScopeStack.java index 84223a267d2..607f509a93b 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ScopeStack.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ScopeStack.java @@ -10,10 +10,10 @@ * The invariant is that the top of a non-empty stack is always active. Anytime a scope is closed, * cleanup() is called to ensure the invariant */ -final class ScopeStack { +class ScopeStack { private final ProfilingContextIntegration profilingContextIntegration; - private final ArrayDeque stack = new ArrayDeque<>(); // previous scopes + final ArrayDeque stack; // previous scopes ContinuableScope top; // current scope @@ -22,6 +22,13 @@ final class ScopeStack { ScopeStack(ProfilingContextIntegration profilingContextIntegration) { this.profilingContextIntegration = profilingContextIntegration; + this.stack = new ArrayDeque<>(); + } + + /** Creates an empty stack with room for the given number of previous scopes. */ + ScopeStack(ProfilingContextIntegration profilingContextIntegration, int numPrevious) { + this.profilingContextIntegration = profilingContextIntegration; + this.stack = new ArrayDeque<>(numPrevious); } /** Restore a shallow stack for async propagation purposes. */ @@ -36,6 +43,10 @@ final class ScopeStack { top = active; } + ContinuableScope top() { + return top; + } + ContinuableScope active() { // avoid attaching further spans to the root scope when it's been marked as overdue return top != overdueRootScope ? top : null; @@ -93,7 +104,7 @@ boolean checkTop(final ContinuableScope expectedScope) { * Slower check to see if overdue scopes ahead of the expected scope are all ITERATION scopes. * These represent iterations that are now out-of-scope and can be finished ready for cleanup. */ - final boolean checkOverdueScopes(final ContinuableScope expectedScope) { + boolean checkOverdueScopes(final ContinuableScope expectedScope) { // we already know 'top' isn't the expected scope, so just need to check its source if (top == null || top.source() != ITERATION) { return false; @@ -125,6 +136,30 @@ int depth() { return top != null ? 1 + stack.size() : 0; } + /** Returns a context that swaps this stack back in. */ + ScopeStackContext swappedOut() { + return new SavedScopeContext(this); + } + + /** Returns a context that swaps the replaced stack back in, after this stack replaces it. */ + ScopeStackContext replacing(ScopeStack replaced) { + return replaced.swappedOut(); + } + + private ThreadScopeContext sameThreadToken; // reused, as it takes no snapshot + + /** + * Returns a reusable context that swaps this stack back in, for callers that always swap it back + * on this thread. + */ + ScopeStackContext swappedOutForSameThread() { + ThreadScopeContext token = sameThreadToken; + if (token == null) { + sameThreadToken = token = new ThreadScopeContext(this); + } + return token; + } + // DQH - regrettably needed for pre-existing tests void clear() { stack.clear(); @@ -132,7 +167,7 @@ void clear() { } /** Notifies profiler that this thread has a context now */ - private void onBecomeNonEmpty() { + final void onBecomeNonEmpty() { try { profilingContextIntegration.onAttach(); } catch (Throwable e) { @@ -141,7 +176,7 @@ private void onBecomeNonEmpty() { } /** Notifies profiler that this thread no longer has a context */ - private void onBecomeEmpty() { + final void onBecomeEmpty() { try { profilingContextIntegration.onDetach(); } catch (Throwable e) { diff --git a/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ScopeStackContext.java b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ScopeStackContext.java new file mode 100644 index 00000000000..18e7c808567 --- /dev/null +++ b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ScopeStackContext.java @@ -0,0 +1,16 @@ +package datadog.trace.core.scopemanager; + +import datadog.context.Context; +import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration; +import javax.annotation.Nullable; + +/** A {@link Context} that swaps a {@link ScopeStack} back in. */ +abstract class ScopeStackContext implements Context { + + /** Returns the stack to install on the current thread. */ + abstract ScopeStack restore(ProfilingContextIntegration profilingContextIntegration); + + static Context contextOf(@Nullable ContinuableScope scope) { + return scope != null ? scope.context : Context.root(); + } +} diff --git a/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/SharedScopeContext.java b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/SharedScopeContext.java new file mode 100644 index 00000000000..efc882501ac --- /dev/null +++ b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/SharedScopeContext.java @@ -0,0 +1,36 @@ +package datadog.trace.core.scopemanager; + +import datadog.context.Context; +import datadog.context.ContextKey; +import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration; +import javax.annotation.Nullable; + +/** Swaps the same {@link SharedScopeStack} onto every thread, instead of a copy. */ +final class SharedScopeContext extends ScopeStackContext { + private final SharedScopeStack sharedStack; + + SharedScopeContext(SharedScopeStack sharedStack) { + this.sharedStack = sharedStack; + } + + /** The context of the shared stack's current top scope. */ + Context currentContext() { + return contextOf(sharedStack.top()); + } + + @Override + ScopeStack restore(ProfilingContextIntegration profilingContextIntegration) { + return sharedStack; + } + + @Nullable + @Override + public T get(ContextKey key) { + return currentContext().get(key); + } + + @Override + public Context with(ContextKey key, @Nullable T value) { + return currentContext().with(key, value); + } +} diff --git a/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/SharedScopeStack.java b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/SharedScopeStack.java new file mode 100644 index 00000000000..1a880f28e42 --- /dev/null +++ b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/SharedScopeStack.java @@ -0,0 +1,161 @@ +package datadog.trace.core.scopemanager; + +import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration; +import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; + +/** + * A {@link ScopeStack} owned by a coroutine and shared by whichever threads are running it. + * + *

Threads normally hand the coroutine over in sequence, but the previous thread can still be + * unwinding while the next one starts. Stack mutations take this stack's lock; reference-count + * changes take the scope's lock. The top scope is published through a volatile field so the + * frequent active-scope reads need no lock. + */ +final class SharedScopeStack extends ScopeStack { + private volatile ContinuableScope liveTop; + // written before liveTop, so it's visible to readers that read top() first + private int liveDepth; + + /** Swaps this stack in on any thread; also returned when this stack is swapped out. */ + final SharedScopeContext handle; + + SharedScopeStack(ProfilingContextIntegration profilingContextIntegration) { + super(profilingContextIntegration, 4); + this.handle = new SharedScopeContext(this); + } + + /** Seeds the stack with an initial scope, before it is shared with other threads. */ + void seed(ContinuableScope scope) { + scope.markShared(); + publishTop(scope); + } + + /** Updates the top scope, publishing it to lock-free readers before any listeners run. */ + @SuppressFBWarnings( + value = "AT_STALE_THREAD_WRITE_OF_PRIMITIVE", + justification = + "The following volatile liveTop write publishes liveDepth to readers that acquire " + + "the top. Depth is a best-effort limit check, so concurrent staleness is acceptable.") + private void publishTop(ContinuableScope scope) { + top = scope; + liveDepth = scope != null ? 1 + stack.size() : 0; + liveTop = scope; + } + + @Override + ContinuableScope top() { + return liveTop; + } + + @Override + ContinuableScope active() { + // avoid attaching further spans to the root scope when it's been marked as overdue + ContinuableScope current = liveTop; + return current != overdueRootScope ? current : null; + } + + @Override + synchronized ContinuableScope parent() { + return super.parent(); + } + + /** Removes closed scopes after swapping in; the swap activates the surviving top. */ + void pruneClosedTop() { + ContinuableScope current = liveTop; + // only take the lock when there's something to remove + if (current != null && !current.alive()) { + synchronized (this) { + if (top != null && !top.alive()) { + cleanup(false); + } + } + } + } + + @Override + ScopeStackContext swappedOut() { + // the shared handle always reads the live stack, so there's no snapshot to take + return handle; + } + + @Override + ScopeStackContext swappedOutForSameThread() { + return handle; + } + + @Override + ScopeStackContext replacing(ScopeStack replaced) { + // the coroutine always swaps the replaced stack back on this thread, so needs no snapshot + return replaced.swappedOutForSameThread(); + } + + // copies ScopeStack.cleanup(), so listeners see the new top when it's reactivated + @Override + synchronized void cleanup() { + cleanup(true); + } + + // only reactivates the new top when asked, as swap() activates it after swapping in + private void cleanup(boolean activateTop) { + ContinuableScope curScope = top; + boolean changedTop = false; + while (curScope != null && !curScope.alive()) { + // no longer alive -- trigger listener & null out + curScope.onProperClose(); + changedTop = true; + curScope = stack.poll(); + } + if (curScope != null && curScope == overdueRootScope) { + // we know this scope is the last on the stack and is overdue + curScope.onProperClose(); + overdueRootScope = null; + publishTop(null); + } else if (changedTop) { + publishTop(curScope); + if (activateTop && curScope != null) { + curScope.beforeActivated(); + curScope.afterActivated(); + } + } + if (top == null) { + onBecomeEmpty(); + } + } + + // copies ScopeStack.push(), so listeners see the new top when it's activated + @Override + synchronized void push(ContinuableScope scope) { + scope.markShared(); + scope.beforeActivated(); + if (top != null) { + stack.push(top); + } else { + onBecomeNonEmpty(); + } + publishTop(scope); + scope.afterActivated(); + } + + @Override + boolean checkTop(ContinuableScope expectedScope) { + // only a fast check; closing re-checks the stack under the lock during cleanup + return expectedScope.equals(liveTop); + } + + @Override + synchronized boolean checkOverdueScopes(ContinuableScope expectedScope) { + return super.checkOverdueScopes(expectedScope); + } + + @Override + int depth() { + // only used to enforce the depth limit, so a slightly stale value is fine + return liveDepth; + } + + @Override + synchronized void clear() { + stack.clear(); + publishTop(null); + } +} diff --git a/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ThreadScopeContext.java b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ThreadScopeContext.java new file mode 100644 index 00000000000..ba3669b7ff0 --- /dev/null +++ b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ThreadScopeContext.java @@ -0,0 +1,41 @@ +package datadog.trace.core.scopemanager; + +import datadog.context.Context; +import datadog.context.ContextKey; +import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration; +import javax.annotation.Nullable; + +/** + * Swaps a thread's own stack back in after a coroutine's shared stack ran on that thread. Unlike + * {@link SavedScopeContext} it takes no snapshot, so one instance can be reused for every + * coroutine. + */ +final class ThreadScopeContext extends ScopeStackContext { + private final Thread originalThread = Thread.currentThread(); + private final ScopeStack scopeStack; + + ThreadScopeContext(ScopeStack scopeStack) { + this.scopeStack = scopeStack; + } + + @Override + ScopeStack restore(ProfilingContextIntegration profilingContextIntegration) { + // coroutines restore on the same thread; there's no snapshot to copy for any other thread + if (Thread.currentThread() == originalThread) { + return scopeStack; + } else { + return new ScopeStack(profilingContextIntegration); + } + } + + @Nullable + @Override + public T get(ContextKey key) { + return contextOf(scopeStack.top()).get(key); + } + + @Override + public Context with(ContextKey key, @Nullable T value) { + return contextOf(scopeStack.top()).with(key, value); + } +} diff --git a/dd-trace-core/src/test/java/datadog/trace/core/scopemanager/SharedScopeStackTest.java b/dd-trace-core/src/test/java/datadog/trace/core/scopemanager/SharedScopeStackTest.java new file mode 100644 index 00000000000..28fae1a41ae --- /dev/null +++ b/dd-trace-core/src/test/java/datadog/trace/core/scopemanager/SharedScopeStackTest.java @@ -0,0 +1,635 @@ +package datadog.trace.core.scopemanager; + +import static java.util.Collections.emptyList; +import static java.util.Collections.singletonList; +import static java.util.concurrent.TimeUnit.SECONDS; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotSame; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import datadog.context.Context; +import datadog.context.ContextKey; +import datadog.context.ContextScope; +import datadog.trace.api.Stateful; +import datadog.trace.api.scopemanager.ScopeListener; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.common.writer.ListWriter; +import datadog.trace.core.CoreTracer; +import datadog.trace.core.DDCoreJavaSpecification; +import datadog.trace.core.ScopeManagerTestBridge; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CyclicBarrier; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.tabletest.junit.TableTest; + +class SharedScopeStackTest extends DDCoreJavaSpecification { + private static final ContextKey KEY = ContextKey.named("shared-scope-stack-test"); + + private final Context base = Context.root().with(KEY, "base"); + + private CoreTracer tracer; + private ContinuableScopeManager scopeManager; + + @BeforeEach + void setup() { + tracer = tracerBuilder().writer(new ListWriter()).build(); + scopeManager = ScopeManagerTestBridge.getScopeManager(tracer); + } + + @AfterEach + void cleanup() { + tracer.close(); + } + + @Test + void scopeOpenedOnOneThreadIsVisibleAndClosableOnAnother() throws Exception { + Context shared = scopeManager.sharedContextStack(base); + Context opened = base.with(KEY, "opened"); + AtomicReference scope = new AtomicReference<>(); + + runOnNewThread( + () -> { + Context original = scopeManager.swap(shared); + scope.set(scopeManager.attach(opened)); + scopeManager.swap(original); + }); + runOnNewThread( + () -> { + Context original = scopeManager.swap(shared); + assertSame(opened, scopeManager.currentContext()); + scope.get().close(); + assertSame(base, scopeManager.currentContext()); + scopeManager.swap(original); + }); + } + + @Test + void rootSeedStartsEmpty() throws Exception { + Context shared = scopeManager.sharedContextStack(Context.root()); + + runOnNewThread( + () -> { + Context original = scopeManager.swap(shared); + assertNull(scopeManager.active()); + assertEquals(0, scopeManager.scopeStack().depth()); + scopeManager.swap(original); + }); + } + + @Test + void scopeClosedElsewhereIsRemovedWhenSwappedIn() throws Exception { + Context shared = scopeManager.sharedContextStack(base); + AtomicReference scope = new AtomicReference<>(); + + runOnNewThread( + () -> { + Context original = scopeManager.swap(shared); + scope.set(scopeManager.attach(base.with(KEY, "closed elsewhere"))); + scopeManager.swap(original); + }); + // closed while the shared stack isn't swapped in, leaving it closed on top of that stack + scope.get().close(); + runOnNewThread( + () -> { + Context original = scopeManager.swap(shared); + assertSame(base, scopeManager.currentContext()); + assertEquals(1, scopeManager.scopeStack().depth()); + scopeManager.swap(original); + }); + } + + @Test + void sharedContextReadsTheLiveTopAndDerivesIndependentContexts() throws Exception { + ContextKey extra = ContextKey.named("shared-context-extra"); + Context shared = scopeManager.sharedContextStack(base); + Context before = shared.with(extra, "before"); + AtomicReference opened = new AtomicReference<>(); + + runOnNewThread( + () -> { + Context original = scopeManager.swap(shared); + opened.set(scopeManager.attach(base.with(KEY, "opened"))); + scopeManager.swap(original); + }); + + assertEquals("opened", shared.get(KEY)); + assertNull(shared.get(extra)); + assertEquals("base", before.get(KEY)); + assertEquals("before", before.get(extra)); + Context after = shared.with(extra, "after"); + assertEquals("opened", after.get(KEY)); + assertEquals("after", after.get(extra)); + + runOnNewThread( + () -> { + Context original = scopeManager.swap(shared); + opened.get().close(); + scopeManager.swap(original); + }); + + assertEquals("base", shared.get(KEY)); + assertNull(shared.get(extra)); + assertEquals("opened", after.get(KEY)); + } + + @Test + void overdueSharedRootIsRemovedBeforeAnotherWorkerRuns() throws Exception { + Context shared = scopeManager.sharedContextStack(Context.root()); + AgentSpan iteration = tracer.buildSpan("test", "iteration").start(); + AtomicInteger closes = new AtomicInteger(); + scopeManager.addScopeListener( + new ScopeListener() { + @Override + public void afterScopeActivated() {} + + @Override + public void afterScopeClosed() { + closes.incrementAndGet(); + } + }); + + runOnNewThread( + () -> { + Context original = scopeManager.swap(shared); + ScopeStack stack = scopeManager.scopeStack(); + ContinuableScope root = + new ContinuableScope( + scopeManager, iteration, ContinuableScope.ITERATION, true, Stateful.DEFAULT); + stack.push(root); + assertSame(iteration, scopeManager.activeSpan()); + // simulate the cleaner marking an expired root, without scheduling its background timer + stack.overdueRootScope = root; + iteration.finishWithEndToEnd(); + assertNull(scopeManager.active()); + stack.cleanup(); + assertNull(stack.overdueRootScope); + scopeManager.swap(original); + }); + + runOnNewThread( + () -> { + Context original = scopeManager.swap(shared); + assertNull(scopeManager.scopeStack().top()); + assertNull(scopeManager.active()); + assertEquals(0, scopeManager.scopeStack().depth()); + assertSame(Context.root(), scopeManager.currentContext()); + scopeManager.swap(original); + }); + assertEquals(1, closes.get()); + } + + @TableTest({ + "Scenario | Parent present", + "Active only | false ", + "With parent | true " + }) + void savedContextRestoresAnIndependentStackOnAnotherThread(boolean parentPresent) + throws Exception { + ContextScope parent = parentPresent ? scopeManager.attach(base) : null; + Context activeContext = base.with(KEY, "active"); + ContextScope active = scopeManager.attach(activeContext); + ScopeStack source = scopeManager.scopeStack(); + Context saved = scopeManager.swap(Context.root()); + + try { + runOnNewThread( + () -> { + Context original = scopeManager.swap(saved); + try { + ScopeStack restored = scopeManager.scopeStack(); + assertNotSame(source, restored); + assertSame(active, restored.top()); + assertSame(parent, restored.parent()); + assertEquals(parentPresent ? 2 : 1, restored.depth()); + assertSame(activeContext, scopeManager.currentContext()); + try (ContextScope child = scopeManager.attach(base.with(KEY, "worker"))) { + assertEquals("worker", scopeManager.currentContext().get(KEY)); + assertSame(active, source.top()); + assertEquals(parentPresent ? 2 : 1, source.depth()); + } + assertSame(activeContext, scopeManager.currentContext()); + } finally { + scopeManager.swap(original); + } + }); + } finally { + scopeManager.swap(saved); + active.close(); + if (parent != null) { + parent.close(); + } + } + } + + @Test + void threadsCanUseTheSharedStackAtTheSameTime() throws Exception { + Context shared = scopeManager.sharedContextStack(base); + int iterations = 10_000; + CyclicBarrier start = new CyclicBarrier(2); + Runnable attachAndClose = + () -> { + Context original = scopeManager.swap(shared); + try { + start.await(10, SECONDS); + for (int i = 0; i < iterations; i++) { + scopeManager.attach(base.with(KEY, Thread.currentThread().getName() + i)).close(); + } + } catch (Exception e) { + throw new RuntimeException(e); + } finally { + scopeManager.swap(original); + } + }; + Thread first = new Thread(attachAndClose, "first"); + Thread second = new Thread(attachAndClose, "second"); + runThreads(first, second); + + runOnNewThread( + () -> { + Context original = scopeManager.swap(shared); + assertSame(base, scopeManager.currentContext()); + assertEquals(1, scopeManager.scopeStack().depth()); + scopeManager.swap(original); + }); + } + + @Test + void threadsCanReuseTheSameScopeAtTheSameTime() throws Exception { + Context shared = scopeManager.sharedContextStack(base); + int iterations = 100_000; + CyclicBarrier start = new CyclicBarrier(2); + Runnable reattachBase = + () -> { + Context original = scopeManager.swap(shared); + try { + start.await(10, SECONDS); + for (int i = 0; i < iterations; i++) { + // reuses the seed scope, adding and removing a reference each time + scopeManager.attach(base).close(); + } + } catch (Exception e) { + throw new RuntimeException(e); + } finally { + scopeManager.swap(original); + } + }; + Thread first = new Thread(reattachBase, "first"); + Thread second = new Thread(reattachBase, "second"); + runThreads(first, second); + + assertSeedHoldsOnlyItsOriginalReference(shared); + } + + @Test + void scopeCanBeClosedFromAnotherStackWhileItsReused() throws Exception { + Context shared = scopeManager.sharedContextStack(base); + int references = 20_000; + List extraReferences = new ArrayList<>(); + runOnNewThread( + () -> { + Context original = scopeManager.swap(shared); + for (int i = 0; i < references; i++) { + extraReferences.add(scopeManager.attach(base)); + } + scopeManager.swap(original); + }); + CyclicBarrier start = new CyclicBarrier(2); + Thread reuser = + new Thread( + () -> { + Context original = scopeManager.swap(shared); + try { + start.await(10, SECONDS); + for (int i = 0; i < references; i++) { + scopeManager.attach(base).close(); + } + } catch (Exception e) { + throw new RuntimeException(e); + } finally { + scopeManager.swap(original); + } + }); + // closes the extra references on its own stack, while the other thread reuses the same scope + Thread closer = + new Thread( + () -> { + try { + start.await(10, SECONDS); + } catch (Exception e) { + throw new RuntimeException(e); + } + extraReferences.forEach(ContextScope::close); + }); + runThreads(reuser, closer); + + assertSeedHoldsOnlyItsOriginalReference(shared); + } + + @Test + void clearedIterationScopeIsNotReusedByAnotherThread() throws Exception { + int rounds = 2_000; + CyclicBarrier phase = new CyclicBarrier(2); + AtomicReference shared = new AtomicReference<>(); + AtomicReference iterationSpan = new AtomicReference<>(); + AtomicReference iterationScope = new AtomicReference<>(); + Runnable reuser = + () -> { + try { + for (int round = 0; round < rounds; round++) { + phase.await(10, SECONDS); // iteration scope is on top + Context original = scopeManager.swap(shared.get()); + for (int i = 0; i < 200; i++) { + // reuses the iteration scope until it's cleared, then pushes a new scope instead + scopeManager.activateSpan(iterationSpan.get()).close(); + } + scopeManager.swap(original); + phase.await(10, SECONDS); // round finished + } + } catch (Exception e) { + throw new RuntimeException(e); + } + }; + Runnable closer = + () -> { + try { + for (int round = 0; round < rounds; round++) { + shared.set(scopeManager.sharedContextStack(base)); + Context original = scopeManager.swap(shared.get()); + ContextScope outer = scopeManager.attach(base.with(KEY, "outer")); + iterationSpan.set(tracer.buildSpan("test", "iteration").start()); + iterationScope.set((ContinuableScope) scopeManager.activateNext(iterationSpan.get())); + phase.await(10, SECONDS); + // closing the outer scope clears the overdue iteration scope above it + outer.close(); + scopeManager.swap(original); + phase.await(10, SECONDS); + assertFalse(iterationScope.get().alive(), "cleared scope retained in round " + round); + } + } catch (Exception e) { + throw new RuntimeException(e); + } + }; + Thread first = new Thread(reuser, "reuser"); + Thread second = new Thread(closer, "closer"); + runThreads(first, second); + } + + @Test + void threadGetsItsOwnStackBackAfterNestedSwaps() { + Context outer = scopeManager.sharedContextStack(base); + Context inner = scopeManager.sharedContextStack(base.with(KEY, "inner")); + ScopeStack threadStack = scopeManager.scopeStack(); + + Context original = scopeManager.swap(outer); + ScopeStack outerStack = scopeManager.scopeStack(); + Context outerToken = scopeManager.swap(inner); + scopeManager.swap(outerToken); + assertSame(outerStack, scopeManager.scopeStack()); + scopeManager.swap(original); + + assertSame(threadStack, scopeManager.scopeStack()); + } + + @Test + void threadReusesOneTokenForItsOwnStack() throws Exception { + Context shared = scopeManager.sharedContextStack(base); + Context opened = base.with(KEY, "thread"); + AtomicReference token = new AtomicReference<>(); + + runOnNewThread( + () -> { + ContextScope scope = scopeManager.attach(opened); + Context first = scopeManager.swap(shared); + scopeManager.swap(first); + Context second = scopeManager.swap(shared); + assertSame(first, second); + scopeManager.swap(second); + assertSame(opened, scopeManager.currentContext()); + token.set(second); + scope.close(); + }); + runOnNewThread( + () -> { + // the token only restores its own thread's stack; other threads get an empty stack + scopeManager.swap(token.get()); + assertNull(scopeManager.active()); + }); + } + + @Test + void reusedThreadTokenReadsItsOwnLiveStack() { + ContextKey extra = ContextKey.named("thread-token-extra"); + Context shared = scopeManager.sharedContextStack(base.with(KEY, "coroutine")); + ContextScope parentScope = scopeManager.attach(base); + Context token = scopeManager.swap(shared); + + assertEquals("coroutine", scopeManager.currentContext().get(KEY)); + assertEquals("base", token.get(KEY)); + Context captured = token.with(extra, "captured"); + assertEquals("base", captured.get(KEY)); + assertEquals("captured", captured.get(extra)); + assertNull(token.get(extra)); + + scopeManager.swap(token); + ContextScope childScope = scopeManager.attach(base.with(KEY, "child")); + Context reused = scopeManager.swap(shared); + assertSame(token, reused); + assertEquals("child", token.get(KEY)); + assertEquals("base", captured.get(KEY)); + + scopeManager.swap(reused); + childScope.close(); + parentScope.close(); + assertNull(token.get(KEY)); + } + + @Test + void seedingFromSharedContextStartsFromItsCurrentTop() throws Exception { + Context parent = scopeManager.sharedContextStack(base); + Context opened = base.with(KEY, "opened"); + AtomicReference child = new AtomicReference<>(); + + runOnNewThread( + () -> { + Context original = scopeManager.swap(parent); + scopeManager.attach(opened); + child.set(scopeManager.sharedContextStack(parent)); + scopeManager.swap(original); + }); + runOnNewThread( + () -> { + Context original = scopeManager.swap(child.get()); + assertSame(opened, scopeManager.currentContext()); + // the child has its own stack + scopeManager.attach(base.with(KEY, "child")); + scopeManager.swap(original); + scopeManager.swap(parent); + assertSame(opened, scopeManager.currentContext()); + scopeManager.swap(original); + }); + } + + @Test + void listenersSeeTheScopeBeingActivated() throws Exception { + Context shared = scopeManager.sharedContextStack(Context.root()); + AgentSpan parent = tracer.buildSpan("test", "parent").start(); + AgentSpan child = tracer.buildSpan("test", "child").start(); + List activeSpans = new ArrayList<>(); + scopeManager.addScopeListener( + new ScopeListener() { + @Override + public void afterScopeActivated() { + activeSpans.add(scopeManager.activeSpan()); + } + + @Override + public void afterScopeClosed() {} + }); + + runOnNewThread( + () -> { + Context original = scopeManager.swap(shared); + ContextScope parentScope = scopeManager.activateSpan(parent); + scopeManager.activateSpan(child).close(); + parentScope.close(); + scopeManager.swap(original); + }); + + assertEquals(3, activeSpans.size()); + assertSame(parent, activeSpans.get(0)); + assertSame(child, activeSpans.get(1)); + assertSame(parent, activeSpans.get(2), "parent reactivated"); + } + + @TableTest({ + "Scenario | Seed parent | Close child", + "Restore parent | true | true ", + "Empty stack | false | true ", + "No closed scopes | true | false " + }) + void deferredCleanupListenersSeeTheRestoredStack(boolean seedParent, boolean closeChild) { + AgentSpan parent = tracer.buildSpan("test", "coroutine-parent").start(); + AgentSpan child = tracer.buildSpan("test", "coroutine-child").start(); + AgentSpan worker = tracer.buildSpan("test", "worker-unrelated").start(); + List activeSpans = new ArrayList<>(); + List closedStacks = new ArrayList<>(); + scopeManager.addScopeListener( + new ScopeListener() { + @Override + public void afterScopeActivated() { + activeSpans.add(scopeManager.activeSpan()); + } + + @Override + public void afterScopeClosed() { + closedStacks.add(scopeManager.scopeStack()); + } + }); + + try { + Context shared = scopeManager.sharedContextStack(seedParent ? parent : Context.root()); + ContextScope childScope; + ScopeStack sharedStack; + Context original = scopeManager.swap(shared); + try { + sharedStack = scopeManager.scopeStack(); + childScope = scopeManager.activateSpan(child); + } finally { + scopeManager.swap(original); + } + + try (ContextScope ignored = scopeManager.activateSpan(worker)) { + if (closeChild) { + // The shared stack retains the closed child until it is restored. + childScope.close(); + } + activeSpans.clear(); + closedStacks.clear(); + + Context workerContext = scopeManager.swap(shared); + try { + AgentSpan expected = closeChild ? (seedParent ? parent : null) : child; + assertSame(expected, scopeManager.activeSpan()); + assertEquals(expected == null ? emptyList() : singletonList(expected), activeSpans); + assertEquals(closeChild ? singletonList(sharedStack) : emptyList(), closedStacks); + } finally { + if (!closeChild) { + childScope.close(); + } + scopeManager.swap(workerContext); + } + } + } finally { + child.finish(); + parent.finish(); + worker.finish(); + } + } + + private void assertSeedHoldsOnlyItsOriginalReference(Context shared) throws Exception { + runOnNewThread( + () -> { + Context original = scopeManager.swap(shared); + assertSame(base, scopeManager.currentContext()); + assertEquals(1, scopeManager.scopeStack().depth()); + ContinuableScope seed = scopeManager.scopeStack().top(); + assertTrue(seed.alive()); + // releasing the original reference closes it, so there are no extra references left + assertFalse(seed.release()); + scopeManager.swap(original); + }); + } + + private static void runOnNewThread(Runnable task) throws Exception { + runThreads(new Thread(task, "worker")); + } + + /** Stops peers on failure and bounds shutdown, including when a worker cannot be interrupted. */ + private static void runThreads(Thread... threads) throws Exception { + AtomicReference failure = new AtomicReference<>(); + for (Thread thread : threads) { + // a monitor deadlock cannot be interrupted; it must not prevent the test JVM exiting + thread.setDaemon(true); + thread.setUncaughtExceptionHandler( + (failed, error) -> { + failure.compareAndSet(null, error); + for (Thread peer : threads) { + if (peer != failed) { + peer.interrupt(); + } + } + }); + } + try { + for (Thread thread : threads) { + thread.start(); + } + for (Thread thread : threads) { + thread.join(10_000); + } + } finally { + for (Thread thread : threads) { + thread.interrupt(); + } + for (Thread thread : threads) { + thread.join(10_000); + } + } + if (failure.get() instanceof Error) { + throw (Error) failure.get(); + } else if (failure.get() != null) { + throw new AssertionError(failure.get()); + } + for (Thread thread : threads) { + assertFalse(thread.isAlive(), thread.getName() + " did not terminate"); + } + } +} diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java index 71767ea7fae..ed023cdd46f 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java @@ -414,6 +414,13 @@ default AgentSpan startSpan(@Nonnull SpanPrototype prototype, CharSequence opera @Deprecated ContextContinuation capture(Context context); + + /** + * Legacy scope manager only: returns a context whose scope stack is shared by every thread it + * is swapped onto. If {@code initial} is already shared, the new stack starts from its top. + */ + @Deprecated + Context sharedContextStack(Context initial); } public interface SpanBuilder { @@ -691,6 +698,11 @@ public Context swap(Context context) { public ContextContinuation capture(Context context) { return NoopContinuation.INSTANCE; } + + @Override + public Context sharedContextStack(Context initial) { + return initial; + } } public static class NoopAgentTraceCollector implements AgentTraceCollector {