When swapping context stacks, restore full stack on original thread. (#9491)

For other threads restore shallow copy, using parent and active scopes captured when the original context stack was swapped out.

Author: mcculls

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

@@ -367,7 +367,7 @@ public Context swap(Context context) {
     ContinuableScope newScope;
     if (context instanceof ScopeContext) {
       // restore previously swapped context stack
-      newStack = ((ScopeContext) context).restore();
+      newStack = ((ScopeContext) context).restore(profilingContextIntegration);
       newScope = newStack.top;
     } else if (context != Context.root()) {
       // start a new stack and record the new context as active

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

@@ -2,12 +2,15 @@
 
 import datadog.context.Context;
 import datadog.context.ContextKey;
+import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration;
 import javax.annotation.Nullable;
 
 /** Wraps a {@link ScopeStack} as a {@link Context} so it can be swapped back later. */
 final class ScopeContext implements Context {
   private final Thread originalThread = Thread.currentThread();
   private final ScopeStack scopeStack;
+  private final ContinuableScope parent;
+  private final ContinuableScope active;
   private final Context context;
 
   ScopeContext(ScopeStack scopeStack) {
@@ -16,12 +19,18 @@ final class ScopeContext implements Context {
 
   private ScopeContext(ScopeStack scopeStack, Context context) {
     this.scopeStack = scopeStack;
+    this.parent = scopeStack.parent();
+    this.active = scopeStack.active();
     this.context = context;
   }
 
-  ScopeStack restore() {
-    // take defensive copy of original scope stack when restoring on different thread
-    return originalThread == Thread.currentThread() ? scopeStack : scopeStack.copy();
+  ScopeStack restore(ProfilingContextIntegration profilingContextIntegration) {
+    // restore full stack on original thread, for other threads restore shallow copy
+    if (Thread.currentThread() == originalThread) {
+      return scopeStack;
+    } else {
+      return new ScopeStack(profilingContextIntegration, parent, active);
+    }
   }
 
   @Nullable

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

@@ -24,19 +24,27 @@ 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;
+  /** Restore a shallow stack for async propagation purposes. */
+  ScopeStack(
+      ProfilingContextIntegration profilingContextIntegration,
+      ContinuableScope parent,
+      ContinuableScope active) {
+    this(profilingContextIntegration);
+    if (parent != null) {
+      stack.push(parent);
+    }
+    top = active;
   }
 
   ContinuableScope active() {
     // avoid attaching further spans to the root scope when it's been marked as overdue
     return top != overdueRootScope ? top : null;
   }
 
+  ContinuableScope parent() {
+    return stack.peek();
+  }
+
   /** Removes and closes all scopes up to the nearest live scope */
   void cleanup() {
     ContinuableScope curScope = top;