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 index 78076cb9803..b08986f6c71 100644 --- 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 index 07a6302b218..2a37e7c553e 100644 --- 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 index f3a718a5a3e..ffbcde5e9df 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 @@ -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 index 8e0cb9a7326..82208e13e41 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 @@ -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,8 +376,11 @@ 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() { @@ -380,8 +388,8 @@ public void activate() { } @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 index 7dcc3e2af29..09938a648ca 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 @@ -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 index 5bad02bedae..2b753b300dd 100644 --- 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 index 96d44360d91..55f68117e23 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 @@ -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 index 5e715cc3546..dcd3977f222 100644 --- 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 index 977afc11843..28a61f71768 100644 --- 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(); }