Implement Context.swap() on top of scope stacks and use it for Kotlin coroutines and Zio fibers (#8805)

Author: mcculls

dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/DatadogThreadContextElement.java

@@ -1,8 +1,8 @@
 package datadog.trace.instrumentation.kotlin.coroutines;
 
+import datadog.context.Context;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
-import datadog.trace.bootstrap.instrumentation.api.ScopeState;
 import kotlin.coroutines.CoroutineContext;
 import kotlin.jvm.functions.Function2;
 import kotlinx.coroutines.AbstractCoroutine;
@@ -11,7 +11,7 @@
 import org.jetbrains.annotations.Nullable;
 
 /** Manages the Datadog context for coroutines, switching contexts as coroutines switch threads. */
-public final class DatadogThreadContextElement implements ThreadContextElement<ScopeState> {
+public final class DatadogThreadContextElement implements ThreadContextElement<Context> {
   private static final CoroutineContext.Key<DatadogThreadContextElement> DATADOG_KEY =
       new CoroutineContext.Key<DatadogThreadContextElement>() {};
 
@@ -22,7 +22,7 @@ public static CoroutineContext addDatadogElement(CoroutineContext coroutineConte
     return coroutineContext.plus(new DatadogThreadContextElement());
   }
 
-  private ScopeState scopeState;
+  private Context context;
   private AgentScope.Continuation continuation;
 
   @NotNull
@@ -32,40 +32,38 @@ public Key<?> getKey() {
   }
 
   public static void captureDatadogContext(@NotNull AbstractCoroutine<?> coroutine) {
-    DatadogThreadContextElement datadogContext = coroutine.getContext().get(DATADOG_KEY);
-    if (datadogContext != null && datadogContext.scopeState == null) {
-      // copy scope stack to use for this coroutine
-      datadogContext.scopeState = AgentTracer.get().oldScopeState().copy();
+    DatadogThreadContextElement datadog = coroutine.getContext().get(DATADOG_KEY);
+    if (datadog != null && datadog.context == null) {
+      // record context to use for this coroutine
+      datadog.context = Context.current();
       // stop enclosing trace from finishing early
-      datadogContext.continuation = AgentTracer.captureActiveSpan();
+      datadog.continuation = AgentTracer.captureActiveSpan();
     }
   }
 
   public static void cancelDatadogContext(@NotNull AbstractCoroutine<?> coroutine) {
-    DatadogThreadContextElement datadogContext = coroutine.getContext().get(DATADOG_KEY);
-    if (datadogContext != null && datadogContext.continuation != null) {
+    DatadogThreadContextElement datadog = coroutine.getContext().get(DATADOG_KEY);
+    if (datadog != null && datadog.continuation != null) {
       // release enclosing trace now the coroutine has completed
-      datadogContext.continuation.cancel();
+      datadog.continuation.cancel();
     }
   }
 
   @Override
-  public ScopeState updateThreadContext(@NotNull CoroutineContext coroutineContext) {
-    ScopeState oldState = AgentTracer.get().oldScopeState();
-    if (scopeState == null) {
-      // copy scope stack to use for this coroutine
-      scopeState = oldState.copy();
+  public Context updateThreadContext(@NotNull CoroutineContext coroutineContext) {
+    if (context == null) {
+      // record context to use for this coroutine
+      context = Context.current();
       // stop enclosing trace from finishing early
       continuation = AgentTracer.captureActiveSpan();
     }
-    scopeState.activate(); // swap in the coroutine's scope stack
-    return oldState;
+    return context.swap();
   }
 
   @Override
   public void restoreThreadContext(
-      @NotNull CoroutineContext coroutineContext, ScopeState oldState) {
-    oldState.activate(); // swap bock the original scope stack
+      @NotNull CoroutineContext coroutineContext, Context originalContext) {
+    context = originalContext.swap();
   }
 
   @NotNull

dd-java-agent/instrumentation/zio/zio-2.0/src/main/java/datadog/trace/instrumentation/zio/v2_0/FiberContext.java

@@ -2,32 +2,30 @@
 
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.captureActiveSpan;
 
+import datadog.context.Context;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
-import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
-import datadog.trace.bootstrap.instrumentation.api.ScopeState;
 
 public class FiberContext {
-  private final ScopeState scopeState;
+  private Context context;
   private final AgentScope.Continuation continuation;
 
-  private ScopeState oldScopeState;
+  private Context originalContext;
 
   public FiberContext() {
-    // copy scope stack to use for this fiber
-    this.scopeState = AgentTracer.get().oldScopeState().copy();
+    // record context to use for this coroutine
+    this.context = Context.current();
     // stop enclosing trace from finishing early
     this.continuation = captureActiveSpan();
   }
 
   public void onResume() {
-    oldScopeState = AgentTracer.get().oldScopeState();
-    scopeState.activate(); // swap in the fiber's scope stack
+    originalContext = context.swap();
   }
 
   public void onSuspend() {
-    if (oldScopeState != null) {
-      oldScopeState.activate(); // swap bock the original scope stack
-      oldScopeState = null;
+    if (originalContext != null) {
+      context = originalContext.swap();
+      originalContext = null;
     }
   }
 

dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java

@@ -60,7 +60,6 @@
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.bootstrap.instrumentation.api.BlackHoleSpan;
 import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration;
-import datadog.trace.bootstrap.instrumentation.api.ScopeState;
 import datadog.trace.bootstrap.instrumentation.api.SpanAttributes;
 import datadog.trace.bootstrap.instrumentation.api.SpanLink;
 import datadog.trace.bootstrap.instrumentation.api.TagContext;
@@ -294,16 +293,6 @@ public EndpointTracker onRootSpanStarted(AgentSpan root) {
     return null;
   }
 
-  @Override
-  public ScopeState oldScopeState() {
-    return scopeManager.oldScopeState();
-  }
-
-  @Override
-  public ScopeState newScopeState() {
-    return scopeManager.newScopeState();
-  }
-
   public static class CoreTracerBuilder {
 
     private Config config;

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

@@ -24,8 +24,6 @@
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.bootstrap.instrumentation.api.ProfilerContext;
 import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration;
-import datadog.trace.bootstrap.instrumentation.api.ScopeState;
-import datadog.trace.bootstrap.instrumentation.api.ScopeStateAware;
 import datadog.trace.core.monitor.HealthMetrics;
 import datadog.trace.relocate.api.RatelimitedLogger;
 import datadog.trace.util.AgentTaskScheduler;
@@ -45,7 +43,7 @@
  * from being reported even if all related spans are finished. It also delegates to other
  * ScopeInterceptors to provide additional functionality.
  */
-public final class ContinuableScopeManager implements ScopeStateAware, ContextManager {
+public final class ContinuableScopeManager implements ContextManager {
 
   static final Logger log = LoggerFactory.getLogger(ContinuableScopeManager.class);
   static final RatelimitedLogger ratelimitedLog = new RatelimitedLogger(log, 1, MINUTES);
@@ -349,16 +347,6 @@ ScopeStack scopeStack() {
     return this.tlsScopeStack.get();
   }
 
-  @Override
-  public ScopeState oldScopeState() {
-    return new ContinuableScopeState(tlsScopeStack.get());
-  }
-
-  @Override
-  public ScopeState newScopeState() {
-    return new ContinuableScopeState(tlsScopeStack.initialValue());
-  }
-
   @Override
   public Context current() {
     final ContinuableScope active = scopeStack().active();
@@ -372,31 +360,33 @@ public ContextScope attach(Context context) {
 
   @Override
   public Context swap(Context context) {
-    throw new UnsupportedOperationException("Not yet implemented");
-  }
-
-  private class ContinuableScopeState implements ScopeState {
-    private final ScopeStack localScopeStack;
-
-    ContinuableScopeState(ScopeStack scopeStack) {
-      this.localScopeStack = scopeStack;
+    ScopeStack oldStack = tlsScopeStack.get();
+    ContinuableScope oldScope = oldStack.top;
+
+    ScopeStack newStack;
+    ContinuableScope newScope;
+    if (context instanceof ScopeContext) {
+      // restore previously swapped context stack
+      newStack = ((ScopeContext) context).scopeStack;
+      newScope = newStack.top;
+    } else if (context != Context.root()) {
+      // start a new stack and record the new context as active
+      newStack = new ScopeStack(profilingContextIntegration);
+      newScope = new ContinuableScope(this, context, CONTEXT, true, createScopeState(context));
+      newStack.top = newScope;
+    } else {
+      // start a new stack with no active context
+      newStack = new ScopeStack(profilingContextIntegration);
+      newScope = null;
     }
 
-    @Override
-    public void activate() {
-      ContinuableScope oldScope = tlsScopeStack.get().top;
-      tlsScopeStack.set(localScopeStack);
-      ContinuableScope newScope = localScopeStack.top;
-      if (oldScope != newScope && newScope != null) {
-        newScope.beforeActivated();
-        newScope.afterActivated();
-      }
+    tlsScopeStack.set(newStack);
+    if (oldScope != newScope && newScope != null) {
+      newScope.beforeActivated();
+      newScope.afterActivated();
     }
 
-    @Override
-    public ScopeState copy() {
-      return new ContinuableScopeState(localScopeStack.copy());
-    }
+    return new ScopeContext(oldStack);
   }
 
   static final class ScopeStackThreadLocal extends ThreadLocal<ScopeStack> {

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

@@ -0,0 +1,31 @@
+package datadog.trace.core.scopemanager;
+
+import datadog.context.Context;
+import datadog.context.ContextKey;
+import javax.annotation.Nullable;
+
+/** Wraps a {@link ScopeStack} as a {@link Context} so it can be swapped back later. */
+final class ScopeContext implements Context {
+  final ScopeStack scopeStack;
+  private final Context context;
+
+  ScopeContext(ScopeStack scopeStack) {
+    this(scopeStack, scopeStack.top != null ? scopeStack.top.context : Context.root());
+  }
+
+  private ScopeContext(ScopeStack scopeStack, Context context) {
+    this.scopeStack = scopeStack;
+    this.context = context;
+  }
+
+  @Nullable
+  @Override
+  public <T> T get(ContextKey<T> key) {
+    return context.get(key);
+  }
+
+  @Override
+  public <T> Context with(ContextKey<T> key, @Nullable T value) {
+    return context.with(key, value);
+  }
+}