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/java/java-lang/java-lang-21.0/src/test/java/testdog/trace/instrumentation/java/lang/jdk21/VirtualThreadScopeClassLoadingForkedTest.java b/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/test/java/testdog/trace/instrumentation/java/lang/jdk21/VirtualThreadScopeClassLoadingForkedTest.java new file mode 100644 index 00000000000..56bc4c5772a --- /dev/null +++ b/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/test/java/testdog/trace/instrumentation/java/lang/jdk21/VirtualThreadScopeClassLoadingForkedTest.java @@ -0,0 +1,53 @@ +package testdog.trace.instrumentation.java.lang.jdk21; + +import static datadog.trace.agent.test.assertions.SpanMatcher.span; +import static datadog.trace.agent.test.assertions.TraceMatcher.trace; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; +import static java.util.concurrent.TimeUnit.MILLISECONDS; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import datadog.context.ContextScope; +import datadog.trace.agent.test.AbstractInstrumentationTest; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import java.time.Duration; +import java.util.concurrent.locks.LockSupport; +import net.bytebuddy.agent.ByteBuddyAgent; +import org.junit.jupiter.api.Test; + +/** Isolates class-loading assertions from tests that may deliberately load coroutine helpers. */ +class VirtualThreadScopeClassLoadingForkedTest extends AbstractInstrumentationTest { + @Test + void virtualThreadHandoffsDoNotLoadSharedScopeStacks() throws InterruptedException { + AgentSpan parent = startSpan("test", "parent"); + AgentSpan[] active = new AgentSpan[2]; + try (ContextScope ignored = activateSpan(parent)) { + Thread child = + Thread.startVirtualThread( + () -> { + active[0] = activeSpan(); + LockSupport.parkNanos(MILLISECONDS.toNanos(1)); + active[1] = activeSpan(); + }); + assertTrue(child.join(Duration.ofSeconds(10))); + assertSame(parent, active[0]); + assertSame(parent, active[1]); + } finally { + parent.finish(); + } + assertTraces(trace(span().root().operationName("parent"))); + + for (Class loaded : ByteBuddyAgent.getInstrumentation().getAllLoadedClasses()) { + if (loaded != null) { + String name = loaded.getName(); + assertFalse( + "datadog.trace.core.scopemanager.SharedScopeStack".equals(name) + || "datadog.trace.core.scopemanager.SharedScopeContext".equals(name), + name); + } + } + } +} 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 e0e53c2db59..4cef874e9fc 100644 --- a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/build.gradle +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/build.gradle @@ -41,6 +41,13 @@ addTestSuiteForDir('coroutines164Test', 'latestDepTest') addTestSuiteForDir('coroutines18Test', 'latestDepTest') addTestSuite('latestDepTest') +tasks.named('coroutines16Test') { + filter { + // kotlinx only replays undispatched restore tokens from 1.6.2 (kotlinx.coroutines#2930) + excludeTestsMatching 'StandardSuspensionHandoffTest.undispatched coroutine can replay its restoration token' + } +} + dependencies { api project(':dd-java-agent:instrumentation:java:java-concurrent:java-concurrent-1.8') compileOnly libs.kotlin 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..a348f202500 --- /dev/null +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/latestDepTest/groovy/StandardSuspensionHandoffTest.groovy @@ -0,0 +1,43 @@ +import datadog.trace.instrumentation.kotlin.coroutines.AbstractStandardSuspensionHandoffTest +import datadog.trace.instrumentation.kotlin.coroutines.StandardSuspensionHandoffTests + +class StandardSuspensionHandoffTest extends AbstractStandardSuspensionHandoffTest { + // withContext uses the context merge hook from coroutines 1.6.1 onward. + def "context merge preserves an explicitly swapped span"() { + when: + new StandardSuspensionHandoffTests().runContextSwapAcrossMerge() + TEST_WRITER.waitForTraces(1) + def trace = TEST_WRITER.get(0) + def spans = trace.collectEntries { [(it.operationName.toString()): it] } + + then: + trace.size() == 3 + spans.swapped.parentId == spans.parent.spanId + spans.child.parentId == spans.swapped.spanId + } + + 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/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 99899787e84..bb9d34f305d 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 @@ -3,7 +3,10 @@ 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.AgentTracer; +import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; import java.util.concurrent.atomic.AtomicReferenceFieldUpdater; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -29,6 +32,9 @@ public final class DatadogThreadContextElement private static final Object COMPLETED = new Object(); + 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( @@ -51,15 +57,23 @@ public final class DatadogThreadContextElement public static CoroutineContext addDatadogElement(CoroutineContext coroutineContext) { DatadogThreadContextElement existing = coroutineContext.get(DATADOG_KEY); if (existing == null) { - return coroutineContext.plus(new DatadogThreadContextElement(null)); + return coroutineContext.plus(new DatadogThreadContextElement(null, null)); } return replaceDatadogElement( - coroutineContext, existing.localRootFor(coroutineContext.get(Job.Key))); + coroutineContext, + LEGACY_CONTEXT_MANAGER ? Context.current() : null, + existing.localRootFor(coroutineContext.get(Job.Key))); } private static CoroutineContext replaceDatadogElement( - CoroutineContext coroutineContext, @Nullable AgentSpan parentLocalRootSpan) { - DatadogThreadContextElement fresh = new DatadogThreadContextElement(parentLocalRootSpan); + CoroutineContext coroutineContext, + @Nullable Context legacySeed, + @Nullable AgentSpan parentLocalRootSpan) { + // Only legacy mode inherits a live stack. Non-legacy mode captures Context.current() + // later on during child initialization because the parent's saved context may be stale. + DatadogThreadContextElement fresh = + new DatadogThreadContextElement( + LEGACY_CONTEXT_MANAGER ? legacySeed : null, parentLocalRootSpan); if (!Boolean.TRUE.equals(coroutineContext.fold(null, HAS_LATER_HOOK))) { return coroutineContext.plus(fresh); } @@ -76,22 +90,26 @@ public static CoroutineContext forContextMerge( DatadogThreadContextElement caller = original.get(DATADOG_KEY); // Scoped coroutines can borrow the caller's hold even when the requested context replaces Job. // Flow merges may install this element without constructing a coroutine to own it. + // An explicit context swap can make the active context differ from the caller's shared stack. return replaceDatadogElement( - merged, caller == null ? null : caller.localRootFor(original.get(Job.Key))); + merged, + LEGACY_CONTEXT_MANAGER ? Context.current() : null, + caller == null ? null : caller.localRootFor(original.get(Job.Key))); } /** Separates upstream flow state while retaining the caller's active context for downstream. */ public static CoroutineContext forFlowCollection( CoroutineContext upstream, CoroutineContext downstream) { DatadogThreadContextElement caller = downstream.get(DATADOG_KEY); - if (caller != null) { + // Read context before legacySharedStack to acquire the initialized stack state. + if (caller != null && caller.context != null && !caller.legacySharedStack) { // Downstream reinstalls the caller's element without constructing a coroutine. caller.context = Context.current(); } DatadogThreadContextElement element = upstream.get(DATADOG_KEY); if (element == null || element == caller) { // Before 1.7.2, flowOn merges with plus and shares this mutable element with its caller. - return replaceDatadogElement(upstream, null); + return replaceDatadogElement(upstream, null, null); } // Newer releases already received a fresh element from the context merge advice. return upstream; @@ -113,7 +131,17 @@ public CoroutineContext invoke(CoroutineContext rebuilt, CoroutineContext.Elemen return rebuilt.plus(element.getKey() == DATADOG_KEY ? this : element); } + // 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 volatile Context legacySeed; @Nullable private AgentSpan parentLocalRootSpan; + + @SuppressFBWarnings( + value = "IS2_INCONSISTENT_SYNC", + justification = + "Written once before publishing volatile context; updateThreadContext reads context before restoration") + private boolean legacySharedStack; + private volatile Context context; private volatile ContextContinuation continuation; @@ -123,7 +151,9 @@ public CoroutineContext invoke(CoroutineContext rebuilt, CoroutineContext.Elemen /** The local root of the trace kept open by this coroutine or its parent. */ @Nullable private volatile AgentSpan localRootSpan; - private DatadogThreadContextElement(@Nullable AgentSpan parentLocalRootSpan) { + private DatadogThreadContextElement( + @Nullable Context legacySeed, @Nullable AgentSpan parentLocalRootSpan) { + this.legacySeed = legacySeed; this.parentLocalRootSpan = parentLocalRootSpan; } @@ -147,6 +177,15 @@ public Key getKey() { return DATADOG_KEY; } + /** Lazy children inherit the starting thread's context rather than their creation-time seed. */ + public static void captureLazyDatadogContext(@Nonnull AbstractCoroutine coroutine) { + DatadogThreadContextElement datadog = coroutine.getContext().get(DATADOG_KEY); + if (datadog != null) { + datadog.legacySeed = null; + } + captureDatadogContext(coroutine); + } + public static void captureDatadogContext(@Nonnull AbstractCoroutine coroutine) { DatadogThreadContextElement datadog = coroutine.getContext().get(DATADOG_KEY); if (datadog != null) { @@ -156,7 +195,9 @@ public static void captureDatadogContext(@Nonnull AbstractCoroutine coroutine } if (datadog.context == null) { datadog.captureInitialContext(true); - } else if (datadog.owner != coroutine && datadog.isScopedBy(coroutine)) { + } else if (!datadog.legacySharedStack + && datadog.owner != coroutine + && datadog.isScopedBy(coroutine)) { // scoped coroutines (coroutineScope, withTimeout, or withContext before 1.6.1) share their // caller's element, so refresh it with the active context datadog.context = Context.current(); @@ -166,7 +207,17 @@ public static void captureDatadogContext(@Nonnull AbstractCoroutine coroutine private synchronized void captureInitialContext(boolean owned) { if (context == null) { - Context initial = Context.current(); + Context initial = legacySeed; + if (initial == null) { + initial = Context.current(); + } + if (LEGACY_CONTEXT_MANAGER) { + // workers share this coroutine's scope stack, rather than copying it between them + Context shared = AgentTracer.get().sharedContextStack(initial); + // the noop tracer returns initial unchanged, without creating a shared stack + legacySharedStack = shared != initial; + initial = shared; + } // Initialize before the first context installation; later swaps update the stored context. AgentSpan rootSpan = localRootOf(initial); if (rootSpan != null) { @@ -184,6 +235,7 @@ private synchronized void captureInitialContext(boolean owned) { } } context = initial; + legacySeed = null; parentLocalRootSpan = null; } } @@ -232,7 +284,13 @@ public Context updateThreadContext(@Nonnull CoroutineContext coroutineContext) { @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 (!legacySharedStack) { + context = coroutineContextAfter; + } } /** Tracing state does not change the logical coroutine context used by flow equality checks. */ diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/LazyCoroutineInstrumentation.java b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/LazyCoroutineInstrumentation.java index a629a3fdcb4..db42c63fc5d 100644 --- a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/LazyCoroutineInstrumentation.java +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/LazyCoroutineInstrumentation.java @@ -35,7 +35,7 @@ public void methodAdvice(MethodTransformer transformer) { public static class OnStartAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static void onStart(@Advice.This AbstractCoroutine coroutine) { - DatadogThreadContextElement.captureDatadogContext(coroutine); + DatadogThreadContextElement.captureLazyDatadogContext(coroutine); } } } 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/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..009680c405d --- /dev/null +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/groovy/datadog/trace/instrumentation/kotlin/coroutines/AbstractStandardSuspensionHandoffTest.groovy @@ -0,0 +1,150 @@ +package datadog.trace.instrumentation.kotlin.coroutines + +import datadog.trace.agent.test.InstrumentationSpecification + +abstract class AbstractStandardSuspensionHandoffTest extends InstrumentationSpecification { + def "#builder with an enclosing context captures active span with swap #swapped"() { + when: + new StandardSuspensionHandoffTests().runChildWithEnclosingContext(builder, swapped) + TEST_WRITER.waitForTraces(1) + def trace = TEST_WRITER.get(0) + def spans = trace.collectEntries { [(it.operationName.toString()): it] } + + then: + trace.size() == 3 + spans.active.parentId == spans.parent.spanId + spans.child.parentId == spans.active.spanId + + where: + builder | swapped + 'async' | false + 'async' | true + 'launch' | false + 'launch' | true + } + + def "lazy #builder inherits starting context with automatic start #automatic and active span #tracedStarter"() { + when: + new StandardSuspensionHandoffTests().runLazyChildWithInheritedContext(builder, automatic, tracedStarter) + + then: + noExceptionThrown() + + where: + builder | automatic | tracedStarter + 'async' | false | false + 'async' | false | true + 'async' | true | false + 'async' | true | true + 'launch' | false | false + 'launch' | false | true + 'launch' | true | false + 'launch' | true | true + } + + 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 "manual scope survives suspension after flowOn collection"() { + when: + new StandardSuspensionHandoffTests().runFlowOnCollectionThenSuspend() + + 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 index 6247ba507a0..05bbed4db97 100644 --- 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 @@ -3,11 +3,15 @@ import static datadog.trace.instrumentation.kotlin.coroutines.DatadogThreadContextElement.addDatadogElement; import static datadog.trace.instrumentation.kotlin.coroutines.DatadogThreadContextElement.cancelDatadogContext; import static datadog.trace.instrumentation.kotlin.coroutines.DatadogThreadContextElement.captureDatadogContext; +import static datadog.trace.instrumentation.kotlin.coroutines.DatadogThreadContextElement.forFlowCollection; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNotSame; import static org.junit.jupiter.api.Assertions.assertSame; +import datadog.context.Context; +import datadog.context.ContextKey; +import datadog.trace.bootstrap.instrumentation.api.AgentTracer; import java.lang.reflect.Field; import java.util.ArrayList; import java.util.List; @@ -25,6 +29,47 @@ /** Context replacement checks shared by the oldest and latest dependency suites. */ public abstract class AbstractDatadogThreadContextElementTest { + @TableTest({ + "Scenario | Flow collection", + "Suspension | false ", + "Flow collection | true " + }) + void contextCapturedBeforeTracerRegistrationPreservesChanges(boolean flowCollection) { + DatadogThreadContextElement element = + (DatadogThreadContextElement) addDatadogElement(EmptyCoroutineContext.INSTANCE); + AgentTracer.TracerAPI tracer = AgentTracer.get(); + Context original; + try { + AgentTracer.forceRegister(AgentTracer.NOOP_TRACER); + original = element.updateThreadContext(element); + } finally { + AgentTracer.forceRegister(tracer); + } + + ContextKey key = ContextKey.named("startup-coroutine-context"); + try { + Context.current().with(key, "opened before suspension").swap(); + if (flowCollection) { + forFlowCollection(element, element); + Context beforeCollection = element.updateThreadContext(element); + try { + assertEquals("opened before suspension", Context.current().get(key)); + } finally { + element.restoreThreadContext(element, beforeCollection); + } + } + } finally { + element.restoreThreadContext(element, original); + } + + original = element.updateThreadContext(element); + try { + assertEquals("opened before suspension", Context.current().get(key)); + } finally { + element.restoreThreadContext(element, original); + } + } + @Test void completionReleasesOwnerWithoutAllowingAnotherCoroutineToClaimElement() throws Exception { DatadogThreadContextElement element = 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..d34ec270479 --- /dev/null +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/kotlin/datadog/trace/instrumentation/kotlin/coroutines/StandardSuspensionHandoffTests.kt @@ -0,0 +1,697 @@ +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.Deferred +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.ExperimentalCoroutinesApi +import kotlinx.coroutines.Job +import kotlinx.coroutines.ThreadContextElement +import kotlinx.coroutines.asCoroutineDispatcher +import kotlinx.coroutines.async +import kotlinx.coroutines.coroutineScope +import kotlinx.coroutines.delay +import kotlinx.coroutines.flow.collect +import kotlinx.coroutines.flow.flow +import kotlinx.coroutines.flow.flowOn +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 runChildWithEnclosingContext(builder: String, swapped: Boolean) = runBlocking { + val enclosingContext = coroutineContext + val parent = get().buildSpan("kotlin_coroutine", "parent").start() + val parentScope = get().activateManualSpan(parent) + try { + withContext(CoroutineName("nested")) { + val active = get().buildSpan("kotlin_coroutine", "active").start() + val previous = if (swapped) active.swap() else null + val scope = if (!swapped) get().activateManualSpan(active) else null + try { + val block: suspend CoroutineScope.() -> Unit = { + check(activeSpan() === active) { "Child inherited the enclosing context's older span" } + delay(1) + check(activeSpan() === active) { "Child lost its active parent after suspension" } + get().buildSpan("kotlin_coroutine", "child").start().finish() + } + val child = try { + when (builder) { + "async" -> async(enclosingContext, block = block) + "launch" -> launch(enclosingContext, block = block) + else -> error("Unknown child builder: $builder") + } + } finally { + scope?.close() + previous?.swap() + } + child.join() + } finally { + active.finish() + } + } + } finally { + parentScope.close() + parent.finish() + } + } + + fun runLazyChildWithInheritedContext(builder: String, automatic: Boolean, tracedStarter: Boolean) = runBlocking { + val creator = get().buildSpan("kotlin_coroutine", "creator").start() + val creatorScope = get().activateManualSpan(creator) + val starter = if (tracedStarter) get().buildSpan("kotlin_coroutine", "starter").ignoreActiveSpan().start() else null + val worker = Executors.newSingleThreadExecutor() + try { + val block: suspend CoroutineScope.() -> Unit = { + check(activeSpan() === starter) { "Lazy child did not inherit the starting thread's context" } + delay(1) + check(activeSpan() === starter) { "Lazy child lost its starting context after suspension" } + } + val child = when (builder) { + "async" -> async(coroutineContext, CoroutineStart.LAZY, block) + "launch" -> launch(coroutineContext, CoroutineStart.LAZY, block) + else -> error("Unknown lazy builder: $builder") + } + val started = worker.submit { + // Isolate the starter from executor instrumentation's capture of the creator's context. + val original = Context.root().swap() + val scope = starter?.let { get().activateManualSpan(it) } + try { + if (automatic) { + runBlocking { + if (child is Deferred<*>) child.await() else child.join() + } + } else { + child.start() + } + } finally { + scope?.close() + original.swap() + } + } + // Let the creator's dispatcher run its child while the starter waits for it. + withContext(Dispatchers.IO) { started.get(10, SECONDS) } + child.join() + } finally { + worker.shutdownNow() + check(worker.awaitTermination(5, SECONDS)) { "Lazy starter did not terminate" } + creatorScope.close() + creator.finish() + starter?.finish() + } + } + + fun runContextSwapAcrossMerge() { + runBlocking { + val parent = get().buildSpan("kotlin_coroutine", "parent").start() + val parentScope = get().activateManualSpan(parent) + try { + val swapped = get().buildSpan("kotlin_coroutine", "swapped").start() + val previous = swapped.swap() + try { + withContext(CoroutineName("nested")) { + check(activeSpan() === swapped) { "Context merge discarded the explicitly swapped span" } + get().buildSpan("kotlin_coroutine", "child").start().finish() + } + check(activeSpan() === swapped) { "Context merge did not restore the swapped span" } + } finally { + previous.swap() + swapped.finish() + } + check(activeSpan() === parent) { "Balanced swap did not restore the parent" } + } finally { + parentScope.close() + parent.finish() + } + } + } + + 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 runFlowOnCollectionThenSuspend() { + val workers = Executors.newFixedThreadPool(2) + val dispatcher = workers.asCoroutineDispatcher() + try { + runBlocking(dispatcher) { + // Keeping the dispatcher makes flowOn collect upstream undispatched on this coroutine. + flow { emit(1) }.flowOn(CoroutineName("producer")).collect { } + val manual = get().buildSpan("kotlin_coroutine", "manual").start() + val scope = get().activateManualSpan(manual) + try { + repeat(10) { delay(1) } + check(activeSpan() === manual) { "Suspension after flowOn lost the manual scope" } + } finally { + scope.close() + manual.finish() + } + } + } finally { + dispatcher.close() + check(workers.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..d60b1f458be 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,30 @@ 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 (oldScope != newScope && newScope != null) { - newScope.beforeActivated(); - newScope.afterActivated(); + newStack.pruneClosedTop(); + if (oldScope != newStack.top()) { + newStack.activateCurrentTop(); } - return new ScopeContext(oldStack); + return newStack.replacing(oldStack); } public ContextContinuation capture(@Nonnull Context context) { @@ -410,6 +405,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 ScopeStackContext) { + initial = ((ScopeStackContext) 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 66% 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..4a51bb108bb 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 @@ -1,29 +1,28 @@ package datadog.trace.core.scopemanager; import datadog.context.Context; -import datadog.context.ContextKey; import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration; -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) { @@ -33,14 +32,8 @@ ScopeStack restore(ProfilingContextIntegration profilingContextIntegration) { } } - @Nullable - @Override - public T get(ContextKey key) { - return context.get(key); - } - @Override - public Context with(ContextKey key, @Nullable T value) { - return context.with(key, value); + Context currentContext() { + return context; } } 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..c1860fb3949 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,42 @@ int depth() { return top != null ? 1 + stack.size() : 0; } + /** Returns a context that swaps this stack back in. */ + ScopeStackContext swappedOut() { + return new SavedScopeContext(this); + } + + /** Removes closed scopes after this stack is installed on the current thread. */ + void pruneClosedTop() {} + + /** Notifies profiler hooks and scope listeners about the current top. */ + void activateCurrentTop() { + ContinuableScope current = top; + if (current != null) { + current.beforeActivated(); + current.afterActivated(); + } + } + + /** Returns the context to restore after this stack replaces another stack. */ + ScopeStackContext replacing(ScopeStack oldStack) { + return oldStack.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 +179,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 +188,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..9e3ffd4c910 --- /dev/null +++ b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ScopeStackContext.java @@ -0,0 +1,64 @@ +package datadog.trace.core.scopemanager; + +import datadog.context.Context; +import datadog.context.ContextContinuation; +import datadog.context.ContextKey; +import datadog.context.ContextScope; +import datadog.context.ImplicitContextKeyed; +import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration; +import javax.annotation.Nullable; + +/** + * A {@link Context} that swaps a {@link ScopeStack} back in. + * + *

Value access, updates and propagation use the underlying context, keeping live stack tokens + * out of scopes, continuations and carriers. + */ +abstract class ScopeStackContext implements Context { + + /** Returns the stack to install on the current thread. */ + abstract ScopeStack restore(ProfilingContextIntegration profilingContextIntegration); + + /** The context that lookups and updates on this token are forwarded to. */ + abstract Context currentContext(); + + @Nullable + @Override + public final T get(ContextKey key) { + return currentContext().get(key); + } + + @Override + public final Context with(ContextKey key, @Nullable T value) { + return currentContext().with(key, value); + } + + @Override + public final Context with(@Nullable ImplicitContextKeyed value) { + return currentContext().with(value); + } + + @Override + public final ContextScope attach() { + return currentContext().attach(); + } + + @Override + public final void attachTo(Object carrier) { + currentContext().attachTo(carrier); + } + + @Override + public final ContextContinuation capture() { + return currentContext().capture(); + } + + @Override + public final ContextScope asScope() { + return currentContext().asScope(); + } + + 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..7c225ce9884 --- /dev/null +++ b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/SharedScopeContext.java @@ -0,0 +1,24 @@ +package datadog.trace.core.scopemanager; + +import datadog.context.Context; +import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration; + +/** 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. */ + @Override + Context currentContext() { + return contextOf(sharedStack.top()); + } + + @Override + ScopeStack restore(ProfilingContextIntegration profilingContextIntegration) { + return sharedStack; + } +} 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..910b326677f --- /dev/null +++ b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/SharedScopeStack.java @@ -0,0 +1,208 @@ +package datadog.trace.core.scopemanager; + +import static datadog.trace.core.scopemanager.ContinuableScope.ITERATION; + +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration; +import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; +import java.util.ArrayList; +import java.util.List; + +/** + * 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. + * + *

Scope listeners, profiler hooks and span finishing run outside the stack lock to avoid + * deadlocks with application code. Notifications reread the live top, but legacy listeners receive + * no context argument, so another worker can still change the active scope during notification. + */ +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; + } + + /** Removes closed scopes after swapping in; the swap activates the surviving top. */ + @Override + void pruneClosedTop() { + ContinuableScope current = liveTop; + // only take the lock when there's something to remove + if (current != null && !current.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 oldStack) { + // Coroutines always swap the replaced stack back on this thread, so need no snapshot. + return oldStack.swappedOutForSameThread(); + } + + @Override + void cleanup() { + cleanup(true); + } + + // only reactivates the new top when asked, as swap() activates it after swapping in + private void cleanup(boolean activateTop) { + ContinuableScope closedScope; + int closedCount = 0; + synchronized (this) { + ContinuableScope curScope = top; + closedScope = curScope; + while (curScope != null && !curScope.alive()) { + closedCount++; + curScope = stack.poll(); + } + if (curScope != null && curScope == overdueRootScope) { + // we know this scope is the last on the stack and is overdue + // prevent a worker that already read the top from retaining it after removal + curScope.clearReferences(); + closedCount++; + overdueRootScope = null; + curScope = null; + } + if (closedCount > 0) { + publishTop(curScope); + } + } + // Close callbacks carry no scope identity and use the same manager, so only a count is needed. + for (int i = 0; i < closedCount; i++) { + closedScope.onProperClose(); + } + if (activateTop && closedCount > 0) { + activateCurrentTop(); + } + // A close listener may have opened another scope while notifications were running. + if (liveTop == null) { + onBecomeEmpty(); + } + } + + @Override + void push(ContinuableScope scope) { + scope.markShared(); + boolean wasEmpty; + synchronized (this) { + wasEmpty = top == null; + if (!wasEmpty) { + stack.push(top); + } + publishTop(scope); + } + if (wasEmpty) { + onBecomeNonEmpty(); + } + activateCurrentTop(); + } + + @Override + void activateCurrentTop() { + ContinuableScope current = liveTop; + if (current != null) { + current.beforeActivated(); + // Profiler hooks can block or reenter; use the latest top for scope listeners. + current = liveTop; + if (current != null) { + current.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 + boolean checkOverdueScopes(ContinuableScope expectedScope) { + List overdueSpans; + boolean found = false; + synchronized (this) { + if (top == null || top.source() != ITERATION) { + return false; + } + // Preserve ScopeStack's iteration checks, but defer finishing until after releasing the lock. + overdueSpans = new ArrayList<>(); + top.clearReferences(); + AgentSpan span = top.span(); + if (span != null) { + overdueSpans.add(span); + } + for (ContinuableScope scope : stack) { + if (scope.source() != ITERATION) { + found = expectedScope.equals(scope); + break; + } + scope.clearReferences(); + span = scope.span(); + if (span != null) { + overdueSpans.add(span); + } + } + } + for (AgentSpan span : overdueSpans) { + span.finishWithEndToEnd(); + } + return found; + } + + @Override + int depth() { + // only used to enforce the depth limit, so a slightly stale value is fine + return liveDepth; + } +} 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..d70d1fbe877 --- /dev/null +++ b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ThreadScopeContext.java @@ -0,0 +1,33 @@ +package datadog.trace.core.scopemanager; + +import datadog.context.Context; +import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration; + +/** + * 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); + } + } + + @Override + Context currentContext() { + return contextOf(scopeStack.top()); + } +} diff --git a/dd-trace-core/src/test/java/datadog/trace/core/scopemanager/ScopeManagerForkedTest.java b/dd-trace-core/src/test/java/datadog/trace/core/scopemanager/ScopeManagerForkedTest.java index 4a1eda54ac2..442b1ace94a 100644 --- a/dd-trace-core/src/test/java/datadog/trace/core/scopemanager/ScopeManagerForkedTest.java +++ b/dd-trace-core/src/test/java/datadog/trace/core/scopemanager/ScopeManagerForkedTest.java @@ -27,6 +27,7 @@ import datadog.context.ContextContinuation; import datadog.context.ContextKey; import datadog.context.ContextScope; +import datadog.context.ImplicitContextKeyed; import datadog.trace.api.DDTraceId; import datadog.trace.api.Stateful; import datadog.trace.api.interceptor.MutableSpan; @@ -112,6 +113,137 @@ void cleanup() { tracer.close(); } + @TableTest({ + "Scenario | Token | Capture", + "Attach shared handle | 'shared' | false ", + "Resume shared handle | 'shared' | true ", + "Attach worker token | 'worker' | false ", + "Resume worker token | 'worker' | true ", + "Attach saved snapshot | 'saved' | false ", + "Resume saved snapshot | 'saved' | true " + }) + void stackTokensCanBeAttachedAndCaptured(String tokenKind, boolean capture) { + ContextKey key = ContextKey.named("stack-token"); + Context base = Context.root().with(key, "base"); + Context original = base.swap(); + Context sharedBase = base.with(key, "coroutine"); + Context token; + boolean sharedToken = "shared".equals(tokenKind); + if ("saved".equals(tokenKind)) { + token = Context.root().swap(); + } else { + Context worker = scopeManager.swap(scopeManager.sharedContextStack(sharedBase)); + Context shared = worker.swap(); + token = sharedToken ? shared : worker; + } + Context expected = sharedToken ? sharedBase : base; + token.swap(); + ContextContinuation continuation = capture ? token.capture() : null; + try { + ContextScope activated = capture ? continuation.resume() : token.attach(); + try { + assertSame(expected, Context.current()); + assertEquals(sharedToken ? "coroutine" : "base", Context.current().get(key)); + } finally { + activated.close(); + } + assertSame(expected, Context.current()); + } finally { + if (continuation != null) { + continuation.release(); + } + original.swap(); + } + } + + @Test + void capturingLiveHandleKeepsContextAtCaptureTime() { + ContextKey key = ContextKey.named("live-token-capture"); + Context base = Context.root().with(key, "base"); + Context original = scopeManager.sharedContextStack(base).swap(); + Context shared = Context.root().swap(); + shared.swap(); + ContextContinuation captured = shared.capture(); + Context later = base.with(key, "later"); + try (ContextScope changed = later.attach()) { + try (ContextScope resumed = captured.resume()) { + assertSame(base, Context.current()); + } + assertSame(later, Context.current()); + } finally { + captured.release(); + original.swap(); + } + } + + @ParameterizedTest + @ValueSource(strings = {"shared", "worker", "saved"}) + void tokenDerivedContextsKeepTheirValuesAfterStackChanges(String tokenKind) { + ContextKey key = ContextKey.named("derived-token-context"); + Context base = Context.root().with(key, "base"); + Context original = base.swap(); + Context token; + if ("saved".equals(tokenKind)) { + token = Context.root().swap(); + } else { + Context worker = scopeManager.sharedContextStack(base).swap(); + Context shared = worker.swap(); + token = "shared".equals(tokenKind) ? shared : worker; + } + token.swap(); + Object carrier = new Object(); + try { + token.attachTo(carrier); + ContextScope passive = token.asScope(); + Context unchanged = token.with((ImplicitContextKeyed) null); + AtomicReference supplied = new AtomicReference<>(); + Context stored = + token.with( + context -> { + supplied.set(context); + return context; + }); + Context later = base.with(key, "later"); + try (ContextScope changed = later.attach()) { + assertSame(base, Context.from(carrier)); + assertSame(base, passive.context()); + assertSame(base, unchanged); + assertSame(base, supplied.get()); + assertSame(base, stored); + passive.close(); + assertSame(later, Context.current()); + } + assertSame(base, Context.detachFrom(carrier)); + } finally { + Context.detachFrom(carrier); + original.swap(); + } + } + + @Test + void savedTokenStillAttachesAndCapturesItsSnapshotAfterStackChanges() { + ContextKey key = ContextKey.named("saved-token-snapshot"); + Context snapshot = Context.root().with(key, "snapshot"); + Context original = snapshot.swap(); + Context saved = Context.root().swap(); + saved.swap(); + Context later = snapshot.with(key, "later"); + try (ContextScope changed = later.attach()) { + try (ContextScope attached = saved.attach()) { + assertSame(snapshot, Context.current()); + } + ContextContinuation captured = saved.capture(); + try (ContextScope resumed = captured.resume()) { + assertSame(snapshot, Context.current()); + } finally { + captured.release(); + } + assertSame(later, Context.current()); + } finally { + original.swap(); + } + } + @Test void nonDdspanActivationResultsInAContinuableScope() { ContextScope scope = scopeManager.activateSpan(noopSpan()); 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..7f6ffc679c1 --- /dev/null +++ b/dd-trace-core/src/test/java/datadog/trace/core/scopemanager/SharedScopeStackTest.java @@ -0,0 +1,1151 @@ +package datadog.trace.core.scopemanager; + +import static java.util.Arrays.asList; +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 static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; + +import datadog.context.Context; +import datadog.context.ContextKey; +import datadog.context.ContextScope; +import datadog.trace.api.DDTraceId; +import datadog.trace.api.Stateful; +import datadog.trace.api.interceptor.MutableSpan; +import datadog.trace.api.interceptor.TraceInterceptor; +import datadog.trace.api.scopemanager.ExtendedScopeListener; +import datadog.trace.api.scopemanager.ScopeListener; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration; +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.Collection; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.CyclicBarrier; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; +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()); + } + + @Test + void overdueRootCannotBeRetainedAfterRemoval() throws Exception { + SharedScopeStack stack = new SharedScopeStack(ProfilingContextIntegration.NoOp.INSTANCE); + CountDownLatch topRead = new CountDownLatch(1); + CountDownLatch rootRemoved = new CountDownLatch(1); + AtomicBoolean pause = new AtomicBoolean(true); + Context context = + new Context() { + @Override + public T get(ContextKey key) { + return base.get(key); + } + + @Override + public Context with(ContextKey key, T value) { + return base.with(key, value); + } + + @Override + public boolean equals(Object other) { + // Activation has read the top but has not yet tried to retain it. + if (pause.compareAndSet(true, false)) { + topRead.countDown(); + try { + assertTrue(rootRemoved.await(10, SECONDS)); + } catch (InterruptedException e) { + throw new AssertionError(e); + } + } + return this == other; + } + + @Override + public int hashCode() { + return System.identityHashCode(this); + } + }; + ContinuableScope root = + new ContinuableScope( + scopeManager, context, ContinuableScope.ITERATION, true, Stateful.DEFAULT); + stack.seed(root); + + Thread worker = + new Thread( + () -> { + Context original = scopeManager.swap(stack.handle); + try (ContextScope attached = scopeManager.attach(context)) { + assertNotSame(root, attached); + assertSame(context, scopeManager.currentContext()); + } finally { + scopeManager.swap(original); + } + }, + "activating-worker"); + Thread cleaner = + new Thread( + () -> { + try { + assertTrue(topRead.await(10, SECONDS)); + stack.overdueRootScope = root; + stack.cleanup(); + } catch (InterruptedException e) { + throw new AssertionError(e); + } finally { + rootRemoved.countDown(); + } + }, + "overdue-root-cleaner"); + + runThreads(worker, cleaner); + assertFalse(root.alive()); + } + + @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); + }); + } + + @TableTest({ + "Scenario | Same thread token | Seed present", + "Saved context | false | true ", + "Thread context | true | true ", + "Empty saved token | false | false ", + "Empty thread token | true | false " + }) + void seedingFromRestoreTokenUsesItsContext(boolean sameThreadToken, boolean seedPresent) { + ContextScope parent = seedPresent ? scopeManager.attach(base) : null; + ScopeStack parentStack = scopeManager.scopeStack(); + Context token = + scopeManager.swap( + sameThreadToken ? scopeManager.sharedContextStack(Context.root()) : Context.root()); + try { + Context child = scopeManager.sharedContextStack(token); + Context original = scopeManager.swap(child); + try { + assertSame(seedPresent ? base : Context.root(), scopeManager.currentContext()); + assertNotSame(parentStack, scopeManager.scopeStack()); + if (!seedPresent) { + assertNull(scopeManager.active()); + } + try (ContextScope scope = scopeManager.attach(base.with(KEY, "child"))) { + assertEquals("child", scopeManager.currentContext().get(KEY)); + assertSame(parent, parentStack.top()); + } + } finally { + scopeManager.swap(original); + } + } finally { + scopeManager.swap(token); + if (parent != null) { + parent.close(); + } + } + } + + @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 | Operation", + "New scope | push ", + "Scope close | cleanup ", + "Deferred close | prune " + }) + void listenerCanWaitForAnotherWorkerUsingTheSameStack(String operation) throws Exception { + AgentSpan parent = tracer.buildSpan("test", "parent").start(); + AgentSpan child = tracer.buildSpan("test", "child").start(); + Context shared = scopeManager.sharedContextStack(parent); + CountDownLatch callbackStarted = new CountDownLatch(1); + CountDownLatch peerFinished = new CountDownLatch(1); + AtomicBoolean armed = new AtomicBoolean(); + AtomicBoolean peerFinishedDuringCallback = new AtomicBoolean(); + Runnable waitForPeer = + () -> { + if (armed.compareAndSet(true, false)) { + callbackStarted.countDown(); + try { + peerFinishedDuringCallback.set(peerFinished.await(5, SECONDS)); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + } + }; + scopeManager.addScopeListener( + new ScopeListener() { + @Override + public void afterScopeActivated() { + if ("push".equals(operation)) { + waitForPeer.run(); + } + } + + @Override + public void afterScopeClosed() { + if (!"push".equals(operation)) { + waitForPeer.run(); + } + } + }); + Thread notifyingWorker = + new Thread( + () -> { + Context original = scopeManager.swap(shared); + try { + if ("push".equals(operation)) { + armed.set(true); + scopeManager.activateSpan(child).close(); + } else { + ContextScope scope = scopeManager.activateSpan(child); + if ("prune".equals(operation)) { + scopeManager.swap(original); + scope.close(); + armed.set(true); + original = scopeManager.swap(shared); + } else { + armed.set(true); + scope.close(); + } + } + assertSame(parent, scopeManager.activeSpan()); + } finally { + scopeManager.swap(original); + } + }, + "notifying-worker"); + Thread peer = + new Thread( + () -> { + try { + assertTrue(callbackStarted.await(10, SECONDS)); + Context original = scopeManager.swap(shared); + try (ContextScope ignored = scopeManager.attach(base.with(KEY, "peer"))) { + assertEquals("peer", scopeManager.currentContext().get(KEY)); + } finally { + scopeManager.swap(original); + } + } catch (InterruptedException e) { + throw new AssertionError(e); + } finally { + peerFinished.countDown(); + } + }, + "peer-worker"); + try { + runThreads(notifyingWorker, peer); + assertTrue(peerFinishedDuringCallback.get(), "The listener must not block peer mutations"); + } finally { + child.finish(); + parent.finish(); + } + } + + @Test + void scopeListenersAndProfilerHooksRunOutsideTheStackLock() { + ProfilingContextIntegration profiling = mock(ProfilingContextIntegration.class); + SharedScopeStack shared = new SharedScopeStack(profiling); + List callbacks = new ArrayList<>(); + List callbacksUnderLock = new ArrayList<>(); + Consumer record = + callback -> { + callbacks.add(callback); + if (Thread.holdsLock(shared)) { + callbacksUnderLock.add(callback); + } + }; + doAnswer( + invocation -> { + record.accept("attach"); + return null; + }) + .when(profiling) + .onAttach(); + doAnswer( + invocation -> { + record.accept("detach"); + return null; + }) + .when(profiling) + .onDetach(); + Stateful state = + new Stateful() { + @Override + public void activate(Object context) { + record.accept("profiler activate"); + } + + @Override + public void close() {} + }; + scopeManager.addScopeListener( + new ScopeListener() { + @Override + public void afterScopeActivated() { + record.accept("scope activate"); + } + + @Override + public void afterScopeClosed() { + record.accept("scope close"); + } + }); + AgentSpan parent = tracer.buildSpan("test", "parent").start(); + AgentSpan child = tracer.buildSpan("test", "child").start(); + Context original = scopeManager.swap(shared.handle); + try { + ContinuableScope parentScope = + new ContinuableScope(scopeManager, parent, ContinuableScope.CONTEXT, true, state); + ContinuableScope childScope = + new ContinuableScope(scopeManager, child, ContinuableScope.CONTEXT, true, state); + shared.push(parentScope); + shared.push(childScope); + childScope.release(); + shared.cleanup(); + assertSame(parent, scopeManager.activeSpan()); + parentScope.release(); + shared.pruneClosedTop(); + assertNull(scopeManager.activeSpan()); + assertEquals(emptyList(), callbacksUnderLock); + assertEquals( + asList( + "attach", + "profiler activate", + "scope activate", + "profiler activate", + "scope activate", + "scope close", + "profiler activate", + "scope activate", + "scope close", + "detach"), + callbacks); + } finally { + scopeManager.swap(original); + child.finish(); + parent.finish(); + } + } + + @TableTest({ + "Scenario | Parent present", + "Restore parent | true ", + "Empty stack | false " + }) + void swapListenersUseTheTopAfterTheProfilerHook(boolean parentPresent) throws Exception { + AgentSpan parent = tracer.buildSpan("test", "parent").start(); + AgentSpan child = tracer.buildSpan("test", "child").start(); + SharedScopeStack shared = new SharedScopeStack(ProfilingContextIntegration.NoOp.INSTANCE); + CountDownLatch profilerStarted = new CountDownLatch(1); + CountDownLatch peerFinished = new CountDownLatch(1); + AtomicBoolean armed = new AtomicBoolean(); + AtomicBoolean peerFinishedDuringHook = new AtomicBoolean(); + Stateful state = + new Stateful() { + @Override + public void activate(Object context) { + if (armed.compareAndSet(true, false)) { + profilerStarted.countDown(); + try { + peerFinishedDuringHook.set(peerFinished.await(10, SECONDS)); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + } + } + + @Override + public void close() {} + }; + if (parentPresent) { + shared.seed( + new ContinuableScope( + scopeManager, parent, ContinuableScope.CONTEXT, true, Stateful.DEFAULT)); + } + ContinuableScope childScope = + new ContinuableScope(scopeManager, child, ContinuableScope.CONTEXT, true, state); + shared.push(childScope); + + List activatedSpans = new ArrayList<>(); + Thread notifyingWorker = + new Thread( + () -> { + Context original = scopeManager.swap(shared.handle); + try { + assertTrue(peerFinishedDuringHook.get()); + assertSame(parentPresent ? parent : null, scopeManager.activeSpan()); + assertEquals( + parentPresent ? singletonList(parent.getSpanId()) : emptyList(), + activatedSpans); + } finally { + scopeManager.swap(original); + } + }, + "notifying-worker"); + scopeManager.addScopeListener( + new ExtendedScopeListener() { + @Override + public void afterScopeActivated() {} + + @Override + public void afterScopeActivated(DDTraceId traceId, long spanId) { + if (Thread.currentThread() == notifyingWorker) { + activatedSpans.add(spanId); + } + } + + @Override + public void afterScopeClosed() {} + }); + Thread peer = + new Thread( + () -> { + try { + assertTrue(profilerStarted.await(10, SECONDS)); + Context original = scopeManager.swap(shared.handle); + try { + childScope.close(); + } finally { + scopeManager.swap(original); + } + } catch (InterruptedException e) { + throw new AssertionError(e); + } finally { + peerFinished.countDown(); + } + }, + "peer-worker"); + armed.set(true); + try { + runThreads(notifyingWorker, peer); + } finally { + child.finish(); + parent.finish(); + } + } + + @Test + void cleanupKeepsScopeOpenedByCloseListener() { + AgentSpan parent = tracer.buildSpan("test", "parent").start(); + AgentSpan child = tracer.buildSpan("test", "child").start(); + AgentSpan replacement = tracer.buildSpan("test", "replacement").start(); + Context shared = scopeManager.sharedContextStack(parent); + AtomicReference replacementScope = new AtomicReference<>(); + AtomicBoolean openReplacement = new AtomicBoolean(true); + scopeManager.addScopeListener( + new ScopeListener() { + @Override + public void afterScopeActivated() {} + + @Override + public void afterScopeClosed() { + if (openReplacement.compareAndSet(true, false)) { + replacementScope.set(scopeManager.activateSpan(replacement)); + } + } + }); + Context original = scopeManager.swap(shared); + try { + scopeManager.activateSpan(child).close(); + assertSame(replacement, scopeManager.activeSpan()); + replacementScope.get().close(); + assertSame(parent, scopeManager.activeSpan()); + } finally { + scopeManager.swap(original); + replacement.finish(); + child.finish(); + parent.finish(); + } + } + + @Test + void cleanupNotifiesOnceForEachRemovedScope() { + AgentSpan parent = tracer.buildSpan("test", "parent").start(); + AgentSpan child = tracer.buildSpan("test", "child").start(); + AgentSpan grandchild = tracer.buildSpan("test", "grandchild").start(); + Context shared = scopeManager.sharedContextStack(parent); + AtomicInteger closes = new AtomicInteger(); + scopeManager.addScopeListener( + new ScopeListener() { + @Override + public void afterScopeActivated() {} + + @Override + public void afterScopeClosed() { + closes.incrementAndGet(); + } + }); + Context original = scopeManager.swap(shared); + try { + ContextScope childScope = scopeManager.activateSpan(child); + ContextScope grandchildScope = scopeManager.activateSpan(grandchild); + childScope.close(); + assertEquals(0, closes.get()); + grandchildScope.close(); + assertEquals(2, closes.get()); + assertSame(parent, scopeManager.activeSpan()); + } finally { + scopeManager.swap(original); + grandchild.finish(); + child.finish(); + parent.finish(); + } + } + + @Test + void finishingOverdueScopesInvokesTraceInterceptorsOutsideTheStackLock() { + SharedScopeStack shared = new SharedScopeStack(ProfilingContextIntegration.NoOp.INSTANCE); + AtomicInteger interceptedTraces = new AtomicInteger(); + AtomicBoolean interceptedUnderLock = new AtomicBoolean(); + tracer.addTraceInterceptor( + new TraceInterceptor() { + @Override + public Collection onTraceComplete( + Collection trace) { + interceptedTraces.incrementAndGet(); + if (Thread.holdsLock(shared)) { + interceptedUnderLock.set(true); + } + return trace; + } + + @Override + public int priority() { + return 0; + } + }); + AgentSpan first = tracer.buildSpan("test", "first-iteration").start(); + AgentSpan second = tracer.buildSpan("test", "second-iteration").start(); + ContinuableScope parent = + new ContinuableScope(scopeManager, base, ContinuableScope.CONTEXT, true, Stateful.DEFAULT); + ContinuableScope firstIteration = + new ContinuableScope( + scopeManager, first, ContinuableScope.ITERATION, true, Stateful.DEFAULT); + ContinuableScope secondIteration = + new ContinuableScope( + scopeManager, second, ContinuableScope.ITERATION, true, Stateful.DEFAULT); + shared.push(parent); + shared.push(firstIteration); + shared.push(secondIteration); + try { + assertTrue(shared.checkOverdueScopes(parent)); + assertFalse(firstIteration.alive()); + assertFalse(secondIteration.alive()); + assertEquals(2, interceptedTraces.get()); + assertFalse(interceptedUnderLock.get()); + shared.cleanup(); + assertSame(parent, shared.top()); + } finally { + second.finish(); + first.finish(); + } + } + + @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..3a267edb556 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,14 @@ 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. + * Returns a new context when supported, or {@code initial} unchanged when unsupported. + */ + @Deprecated + Context sharedContextStack(Context initial); } public interface SpanBuilder { @@ -691,6 +699,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 {