When swapping context stacks, restore full stack on original thread.

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;