Take defensive copy of parent scope stack when closing nested coroutines (#8749)

mcculls · web-flow · commit e8a21567c515 · 2025-04-29T21:32:43.000+01:00
* Take defensive copy of parent scope stack when closing nested coroutines
* Only copy when we detect the parent scope stack is on a different thread
diff --git a/dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/ScopeStateCoroutineContext.java b/dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/ScopeStateCoroutineContext.java
@@ -39,8 +39,7 @@ public void restoreThreadContext(
 
   @Override
   public ScopeState updateThreadContext(@NotNull final CoroutineContext coroutineContext) {
-    final ScopeState oldScopeState = AgentTracer.get().newScopeState();
-    oldScopeState.fetchFromActive();
+    final ScopeState oldScopeState = AgentTracer.get().oldScopeState();
 
     final Job coroutine = CoroutineContextHelper.getJob(coroutineContext);
     final ScopeStateCoroutineContextItem contextItem = contextItemPerCoroutine.get(coroutine);
@@ -53,22 +52,20 @@ public ScopeState updateThreadContext(@NotNull final CoroutineContext coroutineC
 
   /** If there's a context item for the coroutine then try to close it */
   public void maybeCloseScopeAndCancelContinuation(final Job coroutine, final Job parent) {
-    final ScopeStateCoroutineContextItem contextItem = contextItemPerCoroutine.get(coroutine);
+    final ScopeStateCoroutineContextItem contextItem = contextItemPerCoroutine.remove(coroutine);
     if (contextItem != null) {
       ScopeState currentThreadScopeState = null;
       if (parent != null) {
         final ScopeStateCoroutineContextItem parentItem = contextItemPerCoroutine.get(parent);
         if (parentItem != null) {
-          currentThreadScopeState = parentItem.getScopeState();
+          currentThreadScopeState = parentItem.maybeCopyScopeState();
         }
       }
       if (currentThreadScopeState == null) {
-        currentThreadScopeState = AgentTracer.get().newScopeState();
-        currentThreadScopeState.fetchFromActive();
+        currentThreadScopeState = AgentTracer.get().oldScopeState();
       }
 
       contextItem.maybeCloseScopeAndCancelContinuation();
-      contextItemPerCoroutine.remove(coroutine);
 
       currentThreadScopeState.activate();
     }
@@ -111,16 +108,23 @@ public static class ScopeStateCoroutineContextItem {
     @Nullable private AgentScope.Continuation continuation;
     @Nullable private AgentScope continuationScope;
     private boolean isInitialized = false;
+    private volatile Thread activatedOn;
 
     public ScopeStateCoroutineContextItem() {
       coroutineScopeState = AgentTracer.get().newScopeState();
     }
 
-    public ScopeState getScopeState() {
-      return coroutineScopeState;
+    public ScopeState maybeCopyScopeState() {
+      // take defensive copy of scope stack if on different thread
+      if (activatedOn != null && activatedOn != Thread.currentThread()) {
+        return coroutineScopeState.copy();
+      } else {
+        return coroutineScopeState;
+      }
     }
 
     public void activate() {
+      activatedOn = Thread.currentThread();
       coroutineScopeState.activate();
 
       if (continuation != null && continuationScope == null) {
@@ -144,6 +148,7 @@ public void maybeInitialize() {
      * scope and cancels the continuation.
      */
     public void maybeCloseScopeAndCancelContinuation() {
+      // only temporary activation, will be replaced by another activate in caller
       coroutineScopeState.activate();
 
       if (continuationScope != null) {
diff --git a/dd-java-agent/instrumentation/zio/zio-2.0/src/main/java/datadog/trace/instrumentation/zio/v2_0/FiberContext.java b/dd-java-agent/instrumentation/zio/zio-2.0/src/main/java/datadog/trace/instrumentation/zio/v2_0/FiberContext.java
@@ -52,8 +52,7 @@ public void onSuspend() {
   }
 
   public void onResume() {
-    this.oldState = AgentTracer.get().newScopeState();
-    this.oldState.fetchFromActive();
+    this.oldState = AgentTracer.get().oldScopeState();
 
     this.state.activate();
 
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
@@ -294,6 +294,11 @@ public EndpointTracker onRootSpanStarted(AgentSpan root) {
     return null;
   }
 
+  @Override
+  public ScopeState oldScopeState() {
+    return scopeManager.oldScopeState();
+  }
+
   @Override
   public ScopeState newScopeState() {
     return scopeManager.newScopeState();
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
@@ -349,9 +349,14 @@ ScopeStack scopeStack() {
     return this.tlsScopeStack.get();
   }
 
+  @Override
+  public ScopeState oldScopeState() {
+    return new ContinuableScopeState(tlsScopeStack.get());
+  }
+
   @Override
   public ScopeState newScopeState() {
-    return new ContinuableScopeState();
+    return new ContinuableScopeState(tlsScopeStack.initialValue());
   }
 
   @Override
@@ -371,17 +376,20 @@ public Context swap(Context context) {
   }
 
   private class ContinuableScopeState implements ScopeState {
+    private final ScopeStack localScopeStack;
 
-    private ScopeStack localScopeStack = tlsScopeStack.initialValue();
+    ContinuableScopeState(ScopeStack scopeStack) {
+      this.localScopeStack = scopeStack;
+    }
 
     @Override
     public void activate() {
       tlsScopeStack.set(localScopeStack);
     }
 
     @Override
-    public void fetchFromActive() {
-      localScopeStack = tlsScopeStack.get();
+    public ScopeState copy() {
+      return new ContinuableScopeState(localScopeStack.copy());
     }
   }
 
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
@@ -24,6 +24,14 @@ final class ScopeStack {
     this.profilingContextIntegration = profilingContextIntegration;
   }
 
+  ScopeStack copy() {
+    ScopeStack copy = new ScopeStack(profilingContextIntegration);
+    copy.stack.addAll(stack);
+    copy.top = top;
+    copy.overdueRootScope = overdueRootScope;
+    return copy;
+  }
+
   ContinuableScope active() {
     // avoid attaching further spans to the root scope when it's been marked as overdue
     return top != overdueRootScope ? top : null;
diff --git a/dd-trace-core/src/test/groovy/datadog/trace/core/scopemanager/ScopeManagerTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/core/scopemanager/ScopeManagerTest.groovy
@@ -72,8 +72,7 @@ class ScopeManagerTest extends DDCoreSpecification {
 
   def "scope state should be able to fetch and activate state when there is no active span"() {
     when:
-    def initialScopeState = scopeManager.newScopeState()
-    initialScopeState.fetchFromActive()
+    def initialScopeState = scopeManager.oldScopeState()
 
     then:
     scopeManager.active() == null
@@ -125,8 +124,7 @@ class ScopeManagerTest extends DDCoreSpecification {
     when:
     def span = tracer.buildSpan("test", "test").start()
     def scope = tracer.activateSpan(span)
-    def initialScopeState = scopeManager.newScopeState()
-    initialScopeState.fetchFromActive()
+    def initialScopeState = scopeManager.oldScopeState()
 
     then:
     scope.span() == span
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
@@ -638,6 +638,11 @@ public AgentSpanContext notifyExtensionStart(Object event) {
     @Override
     public void notifyExtensionEnd(AgentSpan span, Object result, boolean isError) {}
 
+    @Override
+    public ScopeState oldScopeState() {
+      return null;
+    }
+
     @Override
     public ScopeState newScopeState() {
       return null;
diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/ScopeState.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/ScopeState.java
@@ -3,5 +3,5 @@
 public interface ScopeState {
   void activate();
 
-  void fetchFromActive();
+  ScopeState copy();
 }
diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/ScopeStateAware.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/ScopeStateAware.java
@@ -1,5 +1,9 @@
 package datadog.trace.bootstrap.instrumentation.api;
 
 public interface ScopeStateAware {
+  /** @return The old connected scope stack */
+  ScopeState oldScopeState();
+
+  /** @return A new disconnected scope stack */
   ScopeState newScopeState();
 }

Original file line number	Diff line number	Diff line change
`@@ -52,8 +52,7 @@ public void onSuspend() {`
`52`	`52`	`}`
`53`	`53`
`54`	`54`	`public void onResume() {`
`55`		`- this.oldState = AgentTracer.get().newScopeState();`
`56`		`- this.oldState.fetchFromActive();`
	`55`	`+ this.oldState = AgentTracer.get().oldScopeState();`
`57`	`56`
`58`	`57`	`this.state.activate();`
`59`	`58`
Original file line number	Diff line number	Diff line change
`@@ -349,9 +349,14 @@ ScopeStack scopeStack() {`
`349`	`349`	`return this.tlsScopeStack.get();`
`350`	`350`	`}`
`351`	`351`
	`352`	`+ @Override`
	`353`	`+ public ScopeState oldScopeState() {`
	`354`	`+ return new ContinuableScopeState(tlsScopeStack.get());`
	`355`	`+ }`
	`356`	`+`
`352`	`357`	`@Override`
`353`	`358`	`public ScopeState newScopeState() {`
`354`		`- return new ContinuableScopeState();`
	`359`	`+ return new ContinuableScopeState(tlsScopeStack.initialValue());`
`355`	`360`	`}`
`356`	`361`
`357`	`362`	`@Override`
`@@ -371,17 +376,20 @@ public Context swap(Context context) {`
`371`	`376`	`}`
`372`	`377`
`373`	`378`	`private class ContinuableScopeState implements ScopeState {`
	`379`	`+ private final ScopeStack localScopeStack;`
`374`	`380`
`375`		`- private ScopeStack localScopeStack = tlsScopeStack.initialValue();`
	`381`	`+ ContinuableScopeState(ScopeStack scopeStack) {`
	`382`	`+ this.localScopeStack = scopeStack;`
	`383`	`+ }`
`376`	`384`
`377`	`385`	`@Override`
`378`	`386`	`public void activate() {`
`379`	`387`	`tlsScopeStack.set(localScopeStack);`
`380`	`388`	`}`
`381`	`389`
`382`	`390`	`@Override`
`383`		`- public void fetchFromActive() {`
`384`		`- localScopeStack = tlsScopeStack.get();`
	`391`	`+ public ScopeState copy() {`
	`392`	`+ return new ContinuableScopeState(localScopeStack.copy());`
`385`	`393`	`}`
`386`	`394`	`}`
`387`	`395`
Original file line number	Diff line number	Diff line change
`@@ -3,5 +3,5 @@`
`3`	`3`	`public interface ScopeState {`
`4`	`4`	`void activate();`
`5`	`5`
`6`		`- void fetchFromActive();`
	`6`	`+ ScopeState copy();`
`7`	`7`	`}`
Original file line number	Diff line number	Diff line change
`@@ -1,5 +1,9 @@`
`1`	`1`	`package datadog.trace.bootstrap.instrumentation.api;`
`2`	`2`
`3`	`3`	`public interface ScopeStateAware {`
	`4`	`+ /** @return The old connected scope stack */`
	`5`	`+ ScopeState oldScopeState();`
	`6`	`+`
	`7`	`+ /** @return A new disconnected scope stack */`
`4`	`8`	`ScopeState newScopeState();`
`5`	`9`	`}`