Rewrite using ManagedScope abstraction for scope stack management

Signed-off-by: monosoul <[email protected]>
DataDog · Dec 9, 2022 · b06ac78 · b06ac78
1 parent c7b4468
commit b06ac78
Show file tree

Hide file tree

Showing 12 changed files with 99 additions and 132 deletions.
diff --git a/.../kotlin-coroutines/src/main/java/datadog/trace/core/ScopeStackCoroutineContextHelper.java b/.../kotlin-coroutines/src/main/java/datadog/trace/core/ScopeStackCoroutineContextHelper.java
diff --git a/...-coroutines/src/main/java/datadog/trace/core/scopemanager/ScopeStackCoroutineContext.java b/...-coroutines/src/main/java/datadog/trace/core/scopemanager/ScopeStackCoroutineContext.java
diff --git a/...src/main/java/datadog/trace/instrumentation/kotlin/coroutines/CoroutineContextAdvice.java b/...src/main/java/datadog/trace/instrumentation/kotlin/coroutines/CoroutineContextAdvice.java
@@ -1,6 +1,5 @@
 package datadog.trace.instrumentation.kotlin.coroutines;
 
-import datadog.trace.core.ScopeStackCoroutineContextHelper;
 import kotlin.coroutines.CoroutineContext;
 import net.bytebuddy.asm.Advice;
 
@@ -9,7 +8,7 @@ public class CoroutineContextAdvice {
   public static void enter(
       @Advice.Argument(value = 1, readOnly = false) CoroutineContext coroutineContext) {
     if (coroutineContext != null) {
-      coroutineContext = ScopeStackCoroutineContextHelper.addScopeStackContext(coroutineContext);
+      coroutineContext = coroutineContext.plus(new ManagedScopeCoroutineContext());
     }
   }
 }
diff --git a/...java/datadog/trace/instrumentation/kotlin/coroutines/KotlinCoroutinesInstrumentation.java b/...java/datadog/trace/instrumentation/kotlin/coroutines/KotlinCoroutinesInstrumentation.java
@@ -1,13 +1,11 @@
 package datadog.trace.instrumentation.kotlin.coroutines;
 
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
-import static datadog.trace.util.Strings.getPackageName;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
-import datadog.trace.core.CoreTracer;
 
 @AutoService(Instrumenter.class)
 public class KotlinCoroutinesInstrumentation extends Instrumenter.Tracing
@@ -20,7 +18,8 @@ public KotlinCoroutinesInstrumentation() {
   @Override
   public String[] helperClassNames() {
     return new String[] {
-      getPackageName(CoreTracer.class.getName()) + ".ScopeStackCoroutineContextHelper",
+      packageName + ".ManagedScopeCoroutineContext",
+      packageName + ".ManagedScopeCoroutineContext$ContextElementKey",
     };
   }
 

diff --git a/...atadog/trace/core/NoOpContextElement.java → ...outines/ManagedScopeCoroutineContext.java b/...atadog/trace/core/NoOpContextElement.java → ...outines/ManagedScopeCoroutineContext.java
@@ -1,22 +1,32 @@
-package datadog.trace.core;
+package datadog.trace.instrumentation.kotlin.coroutines;
 
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import datadog.trace.bootstrap.instrumentation.api.ManagedScope;
 import kotlin.coroutines.CoroutineContext;
 import kotlin.jvm.functions.Function2;
 import kotlinx.coroutines.ThreadContextElement;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
-public class NoOpContextElement implements ThreadContextElement<Object> {
-  static final Key<NoOpContextElement> KEY = new Key<NoOpContextElement>() {};
+public class ManagedScopeCoroutineContext implements ThreadContextElement<ManagedScope> {
 
-  public NoOpContextElement() {}
+  private static final Key<ManagedScopeCoroutineContext> KEY = new ContextElementKey();
+  private final ManagedScope managedScope = AgentTracer.get().delegateManagedScope();
 
   @Override
-  public void restoreThreadContext(@NotNull CoroutineContext coroutineContext, Object oldState) {}
+  public void restoreThreadContext(
+      @NotNull CoroutineContext coroutineContext, ManagedScope oldState) {
+    oldState.activate();
+  }
 
   @Override
-  public Object updateThreadContext(@NotNull CoroutineContext coroutineContext) {
-    return null;
+  public ManagedScope updateThreadContext(@NotNull CoroutineContext coroutineContext) {
+    final ManagedScope oldManagedScope = AgentTracer.get().delegateManagedScope();
+    oldManagedScope.fetch();
+
+    managedScope.activate();
+
+    return oldManagedScope;
   }
 
   @Nullable
@@ -48,4 +58,7 @@ public <R> R fold(
   public Key<?> getKey() {
     return KEY;
   }
+
+  static class ContextElementKey implements Key<ManagedScopeCoroutineContext> {
+  }
 }
diff --git a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java
@@ -40,6 +40,7 @@
 import datadog.trace.bootstrap.instrumentation.api.AgentScopeManager;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import datadog.trace.bootstrap.instrumentation.api.ManagedScope;
 import datadog.trace.bootstrap.instrumentation.api.PathwayContext;
 import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration;
 import datadog.trace.bootstrap.instrumentation.api.ScopeSource;
@@ -199,6 +200,11 @@ public void onRootSpanStarted(AgentSpan root) {
     endpointCheckpointer.onRootSpanStarted(root);
   }
 
+  @Override
+  public ManagedScope delegateManagedScope() {
+    return scopeManager.delegateManagedScope();
+  }
+
   public static class CoreTracerBuilder {
 
     private Config config;

diff --git a/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuableScopeManager.java b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuableScopeManager.java
@@ -2,6 +2,7 @@
 
 import static datadog.trace.api.ConfigDefaults.DEFAULT_ASYNC_PROPAGATING;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.NoopAgentSpan;
+import static datadog.trace.bootstrap.instrumentation.api.ScopeSource.INSTRUMENTATION;
 import static java.util.concurrent.TimeUnit.NANOSECONDS;
 import static java.util.concurrent.TimeUnit.SECONDS;
 
@@ -15,6 +16,7 @@
 import datadog.trace.bootstrap.instrumentation.api.AgentTrace;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.bootstrap.instrumentation.api.AttachableWrapper;
+import datadog.trace.bootstrap.instrumentation.api.ManagedScope;
 import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration;
 import datadog.trace.bootstrap.instrumentation.api.ScopeSource;
 import datadog.trace.util.AgentTaskScheduler;
@@ -253,6 +255,30 @@ ScopeStack scopeStack() {
     return this.tlsScopeStack.get();
   }
 
+  @Override
+  public ManagedScope delegateManagedScope() {
+    return new ContinuableManagedScope();
+  }
+
+  private class ContinuableManagedScope implements ManagedScope {
+
+    private ScopeStack localScopeStack = tlsScopeStack.initialValue();
+    private AgentSpan span = activeSpan();
+
+    @Override
+    public void activate() {
+      tlsScopeStack.set(localScopeStack);
+      if (localScopeStack.depth() == 0 && span != null) {
+        ContinuableScopeManager.this.activate(span, INSTRUMENTATION);
+      }
+    }
+
+    @Override
+    public void fetch() {
+      localScopeStack = tlsScopeStack.get();
+    }
+  }
+
   private static class ContinuableScope implements AgentScope, AttachableWrapper {
     private final ContinuableScopeManager scopeManager;
 

diff --git a/dd-trace-ot/src/main/java/datadog/opentracing/CustomScopeManagerWrapper.java b/dd-trace-ot/src/main/java/datadog/opentracing/CustomScopeManagerWrapper.java
@@ -7,6 +7,7 @@
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentScopeManager;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.ManagedScope;
 import datadog.trace.bootstrap.instrumentation.api.ScopeSource;
 import datadog.trace.util.AgentTaskScheduler;
 import io.opentracing.Scope;
@@ -147,6 +148,26 @@ public AgentScope activateNext(final AgentSpan agentSpan) {
     return converter.toAgentScope(span, scope);
   }
 
+  @Override
+  public ManagedScope delegateManagedScope() {
+    return new CustomManagedScope();
+  }
+
+  private class CustomManagedScope implements ManagedScope {
+
+    private AgentSpan span = activeSpan();
+
+    @Override
+    public void activate() {
+      CustomScopeManagerWrapper.this.activate(span, ScopeSource.INSTRUMENTATION);
+    }
+
+    @Override
+    public void fetch() {
+      span = activeSpan();
+    }
+  }
+
   private void scheduleIterationSpanCleanup(final AgentSpan span) {
     if (iterationSpans == null) {
       synchronized (this) {

diff --git a/...rnal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentScopeManager.java b/...rnal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentScopeManager.java
@@ -3,7 +3,7 @@
 /**
  * Allows custom scope managers. See OTScopeManager, CustomScopeManager, and ContextualScopeManager
  */
-public interface AgentScopeManager {
+public interface AgentScopeManager extends ManagedScopeAware {
 
   AgentScope activate(AgentSpan span, ScopeSource source);
 

diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java
@@ -116,7 +116,11 @@ public static TracerAPI get() {
   private AgentTracer() {}
 
   public interface TracerAPI
-      extends datadog.trace.api.Tracer, InternalTracer, AgentPropagation, EndpointCheckpointer {
+      extends datadog.trace.api.Tracer,
+          InternalTracer,
+          AgentPropagation,
+          EndpointCheckpointer,
+          ManagedScopeAware {
     AgentSpan startSpan(CharSequence spanName);
 
     AgentSpan startSpan(CharSequence spanName, long startTimeMicros);
@@ -376,6 +380,11 @@ public AgentSpan.Context notifyExtensionStart(Object event) {
 
     @Override
     public void notifyExtensionEnd(AgentSpan span, Object result, boolean isError) {}
+
+    @Override
+    public ManagedScope delegateManagedScope() {
+      return null;
+    }
   }
 
   public static final class NoopAgentSpan implements AgentSpan {

diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/ManagedScope.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/ManagedScope.java
@@ -0,0 +1,7 @@
+package datadog.trace.bootstrap.instrumentation.api;
+
+public interface ManagedScope {
+  void activate();
+
+  void fetch();
+}
diff --git a/...rnal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/ManagedScopeAware.java b/...rnal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/ManagedScopeAware.java
@@ -0,0 +1,5 @@
+package datadog.trace.bootstrap.instrumentation.api;
+
+public interface ManagedScopeAware {
+  ManagedScope delegateManagedScope();
+}