When continuing a span merge it with any non-span context (#8658)

Author: mcculls

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

@@ -20,6 +20,7 @@
 import datadog.trace.api.scopemanager.ScopeListener;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentTraceCollector;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.bootstrap.instrumentation.api.ProfilerContext;
 import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration;
@@ -104,18 +105,21 @@ public AgentScope.Continuation captureActiveSpan() {
     if (null != activeScope && activeScope.isAsyncPropagating()) {
       AgentSpan span = activeScope.span();
       if (span != null) {
-        return captureSpan(span, activeScope.source());
+        return captureSpan(activeScope.context, activeScope.source(), span);
       }
     }
     return AgentTracer.noopContinuation();
   }
 
   public AgentScope.Continuation captureSpan(final AgentSpan span) {
-    return captureSpan(span, INSTRUMENTATION);
+    ContinuableScope top = scopeStack().top;
+    Context context = top != null ? top.context.with(span) : span;
+    return captureSpan(context, INSTRUMENTATION, span);
   }
 
-  private AgentScope.Continuation captureSpan(final AgentSpan span, byte source) {
-    return new ScopeContinuation(this, span, source).register();
+  private AgentScope.Continuation captureSpan(Context context, byte source, AgentSpan span) {
+    AgentTraceCollector traceCollector = span.context().getTraceCollector();
+    return new ScopeContinuation(this, context, source, traceCollector).register();
   }
 
   private AgentScope activate(

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

@@ -2,6 +2,7 @@
 
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopScope;
 
+import datadog.context.Context;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentTraceCollector;
@@ -24,7 +25,7 @@ final class ScopeContinuation implements AgentScope.Continuation {
   private static final int HELD = (Integer.MAX_VALUE >> 1) + 1;
 
   final ContinuableScopeManager scopeManager;
-  final AgentSpan spanUnderScope;
+  final Context context;
   final byte source;
   final AgentTraceCollector traceCollector;
 
@@ -51,13 +52,13 @@ final class ScopeContinuation implements AgentScope.Continuation {
 
   ScopeContinuation(
       final ContinuableScopeManager scopeManager,
-      final AgentSpan spanUnderScope,
-      final byte source) {
+      final Context context,
+      final byte source,
+      final AgentTraceCollector traceCollector) {
     this.scopeManager = scopeManager;
-    this.spanUnderScope = spanUnderScope;
+    this.context = context;
     this.source = source;
-
-    this.traceCollector = spanUnderScope.context().getTraceCollector();
+    this.traceCollector = traceCollector;
   }
 
   ScopeContinuation register() {
@@ -77,7 +78,7 @@ public AgentScope.Continuation hold() {
   public AgentScope activate() {
     if (COUNT.incrementAndGet(this) > 0) {
       // speculative update succeeded, continuation can be activated
-      return scopeManager.continueSpan(this, spanUnderScope, source);
+      return scopeManager.continueSpan(this, context, source);
     } else {
       // continuation cancelled or too many activations; rollback count
       COUNT.decrementAndGet(this);
@@ -118,15 +119,11 @@ void cancelFromContinuedScopeClose() {
 
   @Override
   public AgentSpan span() {
-    return spanUnderScope;
+    return AgentSpan.fromContext(context);
   }
 
   @Override
   public String toString() {
-    return getClass().getSimpleName()
-        + "@"
-        + Integer.toHexString(hashCode())
-        + "->"
-        + spanUnderScope;
+    return getClass().getSimpleName() + "@" + Integer.toHexString(hashCode()) + "->" + context;
   }
 }

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

@@ -1075,6 +1075,88 @@ class ScopeManagerTest extends DDCoreSpecification {
     scopeManager.activeSpan() == null
   }
 
+  def "capturing and continuing 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.captureSpan(span).activate()
+
+    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
+  }
+
+  def "capturing and continuing the active 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).withCloseable {
+      tracer.captureActiveSpan().activate()
+    }
+
+    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()
   }