When activating a span merge it with any non-span context (#8649)

Author: mcculls

dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuableScopeManager.java

@@ -119,14 +119,14 @@ private AgentScope.Continuation captureSpan(final AgentSpan span, byte source) {
   }
 
   private AgentScope activate(
-      final Context context,
+      final AgentSpan span,
       final byte source,
       final boolean overrideAsyncPropagation,
       final boolean isAsyncPropagating) {
     ScopeStack scopeStack = scopeStack();
 
     final ContinuableScope top = scopeStack.top;
-    if (top != null && top.context.equals(context)) {
+    if (top != null && span.equals(top.span())) {
       top.incrementReferences();
       return top;
     }
@@ -139,14 +139,16 @@ private AgentScope activate(
       return noopScope();
     }
 
-    assert context != null;
+    assert span != null;
 
     // Inherit the async propagation from the active scope unless the value is overridden
     boolean asyncPropagation =
         overrideAsyncPropagation
             ? isAsyncPropagating
             : top != null ? top.isAsyncPropagating() : DEFAULT_ASYNC_PROPAGATING;
 
+    Context context = top != null ? top.context.with(span) : span;
+
     final ContinuableScope scope =
         new ContinuableScope(this, context, source, asyncPropagation, createScopeState(context));
     scopeStack.push(scope);
@@ -155,6 +157,36 @@ private AgentScope activate(
     return scope;
   }
 
+  private AgentScope activate(final Context context) {
+    ScopeStack scopeStack = scopeStack();
+
+    final ContinuableScope top = scopeStack.top;
+    if (top != null && top.context.equals(context)) {
+      top.incrementReferences();
+      return top;
+    }
+
+    // DQH - This check could go before the check above, since depth limit checking is fast
+    final int currentDepth = scopeStack.depth();
+    if (depthLimit <= currentDepth) {
+      healthMetrics.onScopeStackOverflow();
+      log.debug("Scope depth limit exceeded ({}).  Returning NoopScope.", currentDepth);
+      return noopScope();
+    }
+
+    assert context != null;
+
+    // Inherit the async propagation from the active scope
+    boolean asyncPropagation = top != null ? top.isAsyncPropagating() : DEFAULT_ASYNC_PROPAGATING;
+
+    final ContinuableScope scope =
+        new ContinuableScope(this, context, CONTEXT, asyncPropagation, createScopeState(context));
+    scopeStack.push(scope);
+    healthMetrics.onActivateScope();
+
+    return scope;
+  }
+
   /**
    * Activates a scope for the given {@link ScopeContinuation}.
    *
@@ -326,7 +358,7 @@ public Context current() {
 
   @Override
   public ContextScope attach(Context context) {
-    return activate(context, CONTEXT, false, true);
+    return activate(context);
   }
 
   @Override

dd-trace-core/src/test/groovy/datadog/trace/core/scopemanager/ScopeManagerTest.groovy

@@ -1,5 +1,7 @@
 package datadog.trace.core.scopemanager
 
+import datadog.context.Context
+import datadog.context.ContextKey
 import datadog.trace.agent.test.utils.ThreadUtils
 import datadog.trace.api.DDTraceId
 import datadog.trace.api.Stateful
@@ -1033,6 +1035,46 @@ class ScopeManagerTest extends DDCoreSpecification {
     executor.shutdown()
   }
 
+  def "activating a span merges it with existing context"() {
+    when:
+    def span = tracer.buildSpan("test", "test").start()
+    def testKey = ContextKey.named("test")
+    def context = Context.root().with(testKey, "test-value")
+    def contextScope = scopeManager.attach(context)
+
+    then:
+    scopeManager.active() == contextScope
+    scopeManager.current() == context
+    scopeManager.activeSpan() == null
+    scopeManager.current().get(testKey) == "test-value"
+
+    when:
+    def scope = tracer.activateSpan(span)
+
+    then:
+    scopeManager.active() == scope
+    scopeManager.current() != context
+    scopeManager.activeSpan() == span
+    scopeManager.current().get(testKey) == "test-value"
+
+    when:
+    scope.close()
+
+    then:
+    scopeManager.active() == contextScope
+    scopeManager.current() == context
+    scopeManager.activeSpan() == null
+    scopeManager.current().get(testKey) == "test-value"
+
+    when:
+    contextScope.close()
+
+    then:
+    scopeManager.active() == null
+    scopeManager.current() == Context.root()
+    scopeManager.activeSpan() == null
+  }
+
   boolean spanFinished(AgentSpan span) {
     return ((DDSpan) span)?.isFinished()
   }