amarziali · Oct 20, 2026
@@ -4,6 +4,7 @@
 import datadog.context.ContextContinuation;
 import datadog.trace.api.GenericClassValue;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
@@ -99,9 +100,6 @@ public static void captureDatadogContext(@Nonnull AbstractCoroutine<?> coroutine
       }
       if (datadog.context == null) {
         datadog.captureInitialContext();
-      } else if (datadog.owner != coroutine && datadog.isScopedBy(coroutine)) {
-        // preserve active span changes when a scoped coroutine reinstalls this element
-        datadog.context = Context.current();
       }
     }
   }
@@ -125,7 +123,7 @@ private synchronized void captureInitialContext() {
           }
         }
       }
-      context = initial;
+      context = AgentTracer.get().forkContext(initial);
       parentLocalRootSpan = null;
     }
   }
@@ -168,7 +166,8 @@ public Context updateThreadContext(@Nonnull CoroutineContext coroutineContext) {
   @Override
   public void restoreThreadContext(
       @Nonnull CoroutineContext coroutineContext, Context originalContext) {
-    context = originalContext.swap();
+    // The next worker may already be executing against the coroutine's stack.
+    originalContext.swap();
   }
 
   @Nonnull

@@ -0,0 +1,99 @@
+package datadog.trace.instrumentation.kotlin.coroutines;
+
+import static java.util.concurrent.Executors.newSingleThreadExecutor;
+import static java.util.concurrent.TimeUnit.SECONDS;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertSame;
+
+import datadog.context.Context;
+import datadog.context.ContextScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import datadog.trace.common.writer.ListWriter;
+import datadog.trace.core.CoreTracer;
+import java.util.Properties;
+import java.util.concurrent.ExecutorService;
+import kotlin.coroutines.EmptyCoroutineContext;
+import org.junit.jupiter.api.Test;
+
+class DatadogThreadContextElementHandoffForkedTest {
+  @Test
+  void resumptionBeforeRestorationPreservesTheEntireScopeStack() throws Exception {
+    AgentTracer.installLegacyContextManager();
+    Properties properties = new Properties();
+    properties.setProperty("dd.trace.scope.strict.mode", "true");
+    CoreTracer tracer =
+        CoreTracer.builder()
+            .withProperties(properties)
+            .strictTraceWrites(true)
+            .writer(new ListWriter())
+            .build();
+    AgentTracer.forceRegister(tracer);
+    ExecutorService second = newSingleThreadExecutor();
+    ExecutorService third = newSingleThreadExecutor();
+    DatadogThreadContextElement element =
+        (DatadogThreadContextElement)
+            DatadogThreadContextElement.addDatadogElement(EmptyCoroutineContext.INSTANCE);
+    Context firstPrevious = element.updateThreadContext(element);
+    AgentSpan root = tracer.startSpan("test", "root");
+    ContextScope rootScope = tracer.activateManualSpan(root);
+    AgentSpan child = tracer.startSpan("test", "child");
+    ContextScope childScope = tracer.activateManualSpan(child);
+    AgentSpan leaf = tracer.startSpan("test", "leaf");
+    ContextScope leafScope = tracer.activateManualSpan(leaf);
+    try {
+      // Kotlin permits the next update before the previous restore has even started.
+      Context secondPrevious =
+          second
+              .submit(
+                  () -> {
+                    Context previous = element.updateThreadContext(element);
+                    assertSame(leaf, tracer.activeSpan());
+                    leafScope.close();
+                    assertSame(child, tracer.activeSpan());
+                    childScope.close();
+                    assertSame(root, tracer.activeSpan());
+                    return previous;
+                  })
+              .get(10, SECONDS);
+      element.restoreThreadContext(element, firstPrevious);
+      assertNull(tracer.activeSpan(), "The first worker must recover its previous context");
+
+      third
+          .submit(
+              () -> {
+                Context previous = element.updateThreadContext(element);
+                try {
+                  assertSame(
+                      root, tracer.activeSpan(), "A late restore must not overwrite progress");
+                  rootScope.close();
+                  assertNull(tracer.activeSpan());
+                } finally {
+                  element.restoreThreadContext(element, previous);
+                }
+              })
+          .get(10, SECONDS);
+      second
+          .submit(
+              () -> {
+                element.restoreThreadContext(element, secondPrevious);
+                assertNull(
+                    tracer.activeSpan(), "The second worker must recover its previous context");
+                Context previous = element.updateThreadContext(element);
+                try {
+                  assertNull(tracer.activeSpan(), "Late restores must not resurrect closed scopes");
+                } finally {
+                  element.restoreThreadContext(element, previous);
+                }
+              })
+          .get(10, SECONDS);
+    } finally {
+      second.shutdownNow();
+      third.shutdownNow();
+      leaf.finish();
+      child.finish();
+      root.finish();
+      tracer.close();
+    }
+  }
+}
@@ -2683,6 +2683,11 @@ public Context swap(Context context) {
     return scopeManager.swap(context);
   }
 
+  @Override
+  public Context forkContext(Context context) {
+    return scopeManager.forkContext(context);
+  }
+
   @Override
   public ContextContinuation capture(Context context) {
     return scopeManager.capture(context);

@@ -360,13 +360,32 @@ public ContextScope attach(@Nonnull Context context) {
     return activate(context);
   }
 
+  /** Creates a stack whose scope changes are visible before the next worker resumes execution. */
+  public Context forkContext(@Nonnull Context context) {
+    MovableScopeContext stack = new MovableScopeContext(profilingContextIntegration);
+    if (context != Context.root()) {
+      stack.top = new ContinuableScope(this, context, CONTEXT, true, createScopeState(context));
+    }
+    return stack;
+  }
+
   public Context swap(@Nonnull Context context) {
     ScopeStack oldStack = tlsScopeStack.get();
-    ContinuableScope oldScope = oldStack.top;
+    if (context == oldStack) {
+      // Scoped coroutines can reinstall the same stack, including during a late restore.
+      return context;
+    }
+    boolean moving = oldStack instanceof MovableScopeContext;
+    // A moving stack may already be in use by the next worker. Only detach it here.
+    ContinuableScope oldScope = moving ? null : oldStack.top;
 
     ScopeStack newStack;
     ContinuableScope newScope;
-    if (context instanceof ScopeContext) {
+    if (context instanceof MovableScopeContext) {
+      // The execution unit owns this stack; a departing worker must not copy or mutate it.
+      newStack = (MovableScopeContext) context;
+      newScope = newStack.top;
+    } else if (context instanceof ScopeContext) {
       // restore previously swapped context stack
       newStack = ((ScopeContext) context).restore(profilingContextIntegration);
       newScope = newStack.top;
@@ -382,12 +401,12 @@ public Context swap(@Nonnull Context context) {
     }
 
     tlsScopeStack.set(newStack);
-    if (oldScope != newScope && newScope != null) {
+    if (oldStack != newStack && oldScope != newScope && newScope != null) {
       newScope.beforeActivated();
       newScope.afterActivated();
     }
 
-    return new ScopeContext(oldStack);
+    return moving ? (MovableScopeContext) oldStack : new ScopeContext(oldStack);
   }
 
   public ContextContinuation capture(@Nonnull Context context) {

@@ -0,0 +1,29 @@
+package datadog.trace.core.scopemanager;
+
+import datadog.context.Context;
+import datadog.context.ContextKey;
+import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration;
+import javax.annotation.Nullable;
+
+/** A scope stack owned by a serial execution unit rather than a worker thread. */
+final class MovableScopeContext extends ScopeStack implements Context {
+  MovableScopeContext(ProfilingContextIntegration profilingContextIntegration) {
+    super(profilingContextIntegration);
+  }
+
+  private Context current() {
+    ContinuableScope active = active();
+    return active != null ? active.context : Context.root();
+  }
+
+  @Nullable
+  @Override
+  public <T> T get(ContextKey<T> key) {
+    return current().get(key);
+  }
+
+  @Override
+  public <T> Context with(ContextKey<T> key, @Nullable T value) {
+    return current().with(key, value);
+  }
+}
@@ -10,7 +10,7 @@
  * 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<ContinuableScope> stack = new ArrayDeque<>(); // previous scopes

@@ -1054,6 +1054,40 @@ void rollbackStopsAtMostRecentCheckpoint() {
     assertNull(scopeManager.activeSpan());
   }
 
+  @Test
+  void forkedContextsIsolateScopeChangesAndPreserveContextValues() {
+    ContextKey<String> key = ContextKey.named("forked-context-test");
+    Context parent = Context.root().with(key, "parent");
+    ContextScope parentScope = scopeManager.attach(parent);
+    Context first = scopeManager.forkContext(parent);
+    Context second = scopeManager.forkContext(parent);
+
+    Context previous = scopeManager.swap(first);
+    ContextScope firstScope = scopeManager.attach(parent.with(key, "first"));
+    assertEquals("first", first.get(key));
+    assertEquals("parent", second.get(key));
+    assertEquals("detached", first.with(key, "detached").get(key));
+    assertEquals("first", first.get(key));
+
+    scopeManager.swap(second);
+    assertEquals("parent", scopeManager.currentContext().get(key));
+    ContextScope secondScope = scopeManager.attach(parent.with(key, "second"));
+
+    scopeManager.swap(first);
+    assertSame(firstScope, scopeManager.active());
+    firstScope.close();
+    assertEquals("parent", first.get(key));
+    assertEquals("second", second.get(key));
+
+    scopeManager.swap(second);
+    assertSame(secondScope, scopeManager.active());
+    secondScope.close();
+    scopeManager.swap(previous);
+    assertSame(parentScope, scopeManager.active());
+    parentScope.close();
+    assertNull(scopeManager.active());
+  }
+
   @Test
   void contextsCanBeSwappedOutAndBack() {
     ContextKey<String> testKey = ContextKey.named("test");

@@ -401,6 +401,15 @@ default AgentSpan startSpan(@Nonnull SpanPrototype prototype, CharSequence opera
 
     void addShutdownListener(Runnable listener);
 
+    /**
+     * Creates an isolated scope stack for an execution unit that moves between threads. Execution
+     * must be handed off serially with a happens-before edge, but restoration of a previous worker
+     * may overlap the next activation. The returned context is a live view of that stack.
+     */
+    default Context forkContext(Context context) {
+      return context;
+    }
+
     // these methods are only used for legacy context manager migration
 
     @Deprecated