Fix Kotlin coroutines losing scopes when resumed on another thread

[mcculls]

What Does This Do

Preserves Datadog scopes when Kotlin coroutines suspend on one worker and resume on another with the legacy context manager, keeping nested scopes and span parenting intact across dispatcher changes and structured coroutine boundaries.

Motivation

Kotlin can resume a coroutine on worker B before worker A finishes restoreThreadContext. Saving the coroutine's context during restoration can therefore publish it too late for resumption. The legacy manager also restores only a shallow snapshot when a saved context crosses threads, losing deeper scopes.

Each coroutine with its own Datadog context element now keeps a persistent SharedScopeStack. Workers swap that same live stack in while executing the coroutine and restore their own stack afterward. This uses the existing ThreadContextElement lifecycle without instrumenting individual suspension primitives.

Additional Notes

Context and restoration

• Add a legacy-only AgentTracer.sharedContextStack(initial) operation. Independent child elements get their own stack seeded from the inherited stack's current context; scoped coroutines that borrow their caller's element continue sharing its stack. Clear the temporary seed after initialization so children do not retain the parent's stack.
• Keep the shared handle stable across suspensions, overlapping restoration callbacks, and replayed undispatched restoration tokens. Restoration returns the worker to its original stack without rebinding the coroutine, including after an unbalanced context swap.
• Preserve context refresh and swap-based saving for the non-legacy manager, including scoped coroutines and flow collection. Legacy flow collection retains the caller's live handle. Elements initialized before tracer registration still save context when shared stacks are unavailable.
• Retain existing continuation ownership and local-root tracking so traces stay open until coroutine completion.
• Separate ordinary saved snapshots, shared handles, and reusable worker-restoration tokens. Shared-stack handoffs avoid taking a snapshot for every suspension; ordinary saved contexts keep their existing behavior. Clarify that the existing context-swap benchmark measures ordinary contexts.
• Update virtual-thread preloading for the ScopeContext → SavedScopeContext rename, with a forked regression checking that virtual-thread handoffs do not load the coroutine shared-stack classes.

Overlapping callbacks

Coroutine execution passes between workers in sequence, but the previous worker's restoration callbacks can overlap the next worker's execution. Other context elements may open or close Datadog scopes during those callbacks.

Shared-stack mutations use the stack monitor; frequent active-scope reads use a volatile top. Shared scopes synchronize reference-count changes on the scope itself, including closes from a thread with another stack installed. Closed or cleared scopes cannot be retained again, and overdue iteration scopes are retired before removal.

Profiler hooks, scope listeners, and span finishing run outside the stack lock so callbacks can reenter tracing or wait for another worker. Activation notifications reread the top after profiler hooks, and cleanup preserves scopes opened by close listeners. SpotBugs suppressions document the publication ordering and concurrency assumptions.

Legacy listeners receive no context argument, so another worker can still change the top during a notification. Scopes opened by other context elements during restoration can briefly remain visible until closed. Cleanup of scopes closed from another stack is best effort and retried on later operations or swaps. The shared-stack implementation and tracer API hook are specific to the legacy manager.

Regression coverage

• Core tests cover cross-thread visibility and closing, concurrent scope reuse, independent snapshots, nested swaps, live-context seeding, overdue-scope retirement, and reusable worker tokens.
• Callback tests cover reentrancy, waiting for another worker, profiler hooks changing the top, trace interceptors, and scopes opened during close notifications.
• Coroutine tests cover cancellable and safe continuations, yield, delay, select, nested scopes, dispatcher changes with delayed restoration, child isolation, suspension after flowOn collection, and trace retention until completion.
• Structured child-wait tests cover coroutineScope, supervisorScope, withContext, and withTimeout, including scopes closed inside the block, exceptions, and cancellation.
• Additional tests cover overlapping scope-touching context elements and replayed undispatched restoration tokens. The replay case is excluded for coroutines 1.6.1, which lacks the required replay behavior. Test worker shutdown uses bounded waits.
• Context-element tests cover initialization before tracer registration; a forked virtual-thread test covers class loading.

Validation: GitHub PR checks pass; the build and DDCI checks are still pending at the time of this update. No additional tests were run for this description-only update.

Contributor Checklist

• Format the title according to the contribution guidelines
• Assign the type: and (comp: or inst:) labels in addition to any other useful labels
• Avoid using close, fix, or any linking keywords when referencing an issue
  Use solves instead, and assign the PR milestone to the issue
• Update the CODEOWNERS file on source file addition, migration, or deletion
• Update public documentation with any new configuration flags or behaviors
• Once approved, use merge queue to merge the PR

[dd-octo-sts]

🟢 Java Benchmark SLOs — All performance SLOs passed

Suite	Status
Startup	🟢 pass

SLO thresholds are defined here based on automatically generated metrics. A warning is raised when results are within 5% of the threshold.

PR vs. master results

Scenario	Candidate	master	Δ (95% CI of mean)
startup:insecure-bank:iast:Agent	13.99 s	13.92 s	[-0.2%; +1.3%] (no difference)
startup:insecure-bank:tracing:Agent	12.95 s	12.91 s	[-0.3%; +0.9%] (no difference)
startup:petclinic:appsec:Agent	17.08 s	16.95 s	[+0.0%; +1.6%] (maybe worse)
startup:petclinic:iast:Agent	16.82 s	16.97 s	[-1.8%; -0.0%] (maybe better)
startup:petclinic:profiling:Agent	16.64 s	16.53 s	[-0.3%; +1.7%] (no difference)
startup:petclinic:sca:Agent	17.07 s	17.02 s	[-0.8%; +1.4%] (no difference)
startup:petclinic:tracing:Agent	16.11 s	16.20 s	[-1.4%; +0.3%] (no difference)

Commit: a4269e2e · CI Pipeline · Benchmarking Platform UI

---

Load and DaCapo benchmarks can be triggered manually in the GitLab pipeline. Results will appear in the Benchmarking Platform UI after completion.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-10-14T13:21:45.125790Z	0c7bb96	Draft marked ready

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[datadog-datadog-us1-prod]

More details

The completed static review raised no reportable defect in coroutine stack ownership or trace-continuation lifetimes. Runtime handoff and concurrency behavior remain unverified.

Was this helpful? React 👍 or 👎

_{Open Bits AI session}

_{🤖 Bits Code Review · Commit 0c7bb96}

[dougqh]

I think I was able to follow what the code is doing, but the flow is hard to follow. To a large extent, I think that's inherent to the problem, but I also find that we don't have agreed upon terms to describe what we're doing. I had Claude map seed, owner, trace kept open into my own terms, so I could understand what's happening here.

Since we're doing similar changes in a number of locations, I think it might be good to add some extra conveniences into the Context system to fit these situations.
That's not to say there's anything wrong with the Context system, I think what OTel offers is largely sufficient.
I just find that we keep having to compose operations in complicated ways that make the intent hard to determine.
I think having some more expressive "macros" would reduce the amount of code and make the code easier to read.

Claude was able to reproduce an issue around withContext (and likely coroutineScope).
In my terms, there's a bug with switching context vs resolving a context.
Switching is my term for going to another previously active context -- while not altering PendingTrace accounting.
Whereas resolve / cancel do alter the PendingTrace accounting.

The scenario...
When a scoped coroutine completes, cancelDatadogContext can release the owner's continuation while the owner is still running.
And since traceKeptOpen isn't cleared, a later launch(coroutineContext) from the owner can then skip its own continuation, because it is relying on a hold that's already been released.
I'll have Claude post the reproducer.

And to be clear, adding the "macros" doesn't have to be part of this PR, I just think reviewing will remain hard until we have some common terms for these things.

[dougqh]

Reproducer for the scoped-coroutine issue above. With an empty withContext {} before a structured launch, the trace is written in two pieces ([root, seed], then [child] separately, same trace id). The same test without the withContext always passes.

Over four runs: Unconfined failed 4/4, IO 3/4, Default 1/4; the fixed and single-thread pools passed. So it's timing-dependent but reliable on Unconfined.

CoreKotlinCoroutineTests.kt:

/**
 * A child launched by its parent after a scoped coroutine in the parent has completed. The
 * parent is started under a root span that finishes before the child runs.
 */
open fun structuredChildAfterScopedCoroutine(withScoped: Boolean): Int {
  runTest {
    childSpan("root").activateAndUse {
      detached = launch(Job() + jobName("detached")) {
        childSpan("seed").activateAndUse {
          if (withScoped) {
            withContext(jobName("scoped")) {}
          }
          launch(coroutineContext + jobName("child")) {
            gate.await()
            childSpan("child").finish()
          }
        }
      }
    }
  }
  return 3
}

AbstractKotlinCoroutineInstrumentationTest.groovy:

def "child launched after a scoped coroutine keeps the trace open withScoped=#withScoped #dispatcherName"() {
  setup:
  CoreKotlinCoroutineTests kotlinTest = getCoreKotlinCoroutineTestsInstance(dispatcher)
  int expectedNumberOfSpans = kotlinTest.structuredChildAfterScopedCoroutine(withScoped)
  kotlinTest.finishDetached()
  TEST_WRITER.waitForTraces(1)
  List<DDSpan> trace = TEST_WRITER.get(0)

  expect:
  TEST_WRITER.size() == 1
  trace.size() == expectedNumberOfSpans
  findSpan(trace, "child").spanContext().getParentId() == findSpan(trace, "seed").spanContext().getSpanId()
  findSpan(trace, "seed").spanContext().getParentId() == findSpan(trace, "root").spanContext().getSpanId()

  where:
  [dispatcherName, dispatcher, withScoped] << [dispatchersToTest, [true, false]].combinations().collect { it[0] + [it[1]] }
}

Written by Claude, reviewed by @dougqh

[mcculls]

@dougqh actually the majority of changes here have nothing to do with the Context component - it's all about making the legacy scope manager (and scope stacks) work with Kotlin suspension points. Much of the changes here will go away when we retire the legacy scope manager and switch over to the new context manager.

[mcculls]

@dougqh @amarziali one option could be to declare that we cannot support Kotlin coroutines/suspensions properly with the legacy scope manager and instead pivot to replacing that as soon as we can with the new context manager.

[mcculls]

@dougqh @amarziali regarding the scoped-coroutine issue - I confirmed the issue and fixed it in f3f8dc8. Scoped coroutines share their caller’s Datadog element, so their completion must not release the owning coroutine’s continuation. Release now checks ownership, and owner completion also clears traceKeptOpen.

Note: this issue predates this PR: running the regression tests with the previous element implementation also reproduces premature trace splitting. The PR’s continuation optimization relied on that existing lifetime accounting being correct.

I also added coverage for a child launched after withContext, coroutineScope, or withTimeout completes, plus a control without a scoped coroutine. Both full suites pass on kotlinx 1.3 and 1.8.

[amarziali]

looks OK for my pov

[mcculls]

I split off the general fixes (the ones that don't rely on AgentTracer / ContinuableScopeManager changes) into #12914

[dougqh]

Confirmed the scoped-coroutine trace-splitting issue from my earlier reproducer is fixed: KotlinCoroutineInstrumentationTest passed 5/5 runs locally (127 tests each), including all child launched after a scoped coroutine cases on Unconfined, which used to fail reliably.

A few further findings inline. The preload one matters most to me since it can affect non-Kotlin virtual-thread users.

These findings came from a Claude Code review that I've read through; they haven't all been reproduced.

[dougqh]

Thanks for fixing that issue - Claude has several other findings which I had it post inline