Rework shim between ContextManager and the legacy scope manager (#11938)

Rework shim between ContextManager and legacy scope manager to allow for proper swapping over to ThreadLocalContextManager.

Install legacy approach to managing contexts everywhere to begin with

(as the migration proceeds we'll make it toggleable, then change the default depending on the Java version)

Rename internal method for readability

Co-authored-by: devflow.devflow-routing-intake <devflow.devflow-routing-intake@kubernetes.us1.ddbuild.io>

Author: mcculls

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/Agent.java

@@ -337,6 +337,8 @@ public static void start(
       StaticEventLogger.end("crashtracking");
     }
 
+    AgentTracer.maybeInstallLegacyContextManager();
+
     startDatadogAgent(initTelemetry, inst);
 
     final EnumSet<Library> libraries = detectLibraries(log);

dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/nativeimage/TracerActivation.java

@@ -8,6 +8,7 @@
 import datadog.trace.agent.tooling.MeterInstaller;
 import datadog.trace.agent.tooling.ProfilerInstaller;
 import datadog.trace.agent.tooling.TracerInstaller;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -18,6 +19,7 @@ public final class TracerActivation {
 
   public static void activate() {
     try {
+      AgentTracer.maybeInstallLegacyContextManager();
       // Initialize meter
       MeterInstaller.installMeter();
       // Initialize tracer

dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/InstrumentationSpecification.groovy

@@ -10,7 +10,6 @@ import static datadog.trace.api.config.TraceInstrumentationConfig.CODE_ORIGIN_FO
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.closePrevious
 import static datadog.trace.util.AgentThreadFactory.AgentThread.TASK_SCHEDULER
 
-
 import ch.qos.logback.classic.Level
 import ch.qos.logback.classic.util.ContextInitializer
 import com.datadog.debugger.agent.ClassesToRetransformFinder
@@ -351,18 +350,20 @@ abstract class InstrumentationSpecification extends DDSpecification implements A
   void setupSpec() {
     InstrumentationErrors.resetErrors()
 
-    AgentMeter.registerIfAbsent(
-    STATS_D_CLIENT,
-    new MonitoringImpl(STATS_D_CLIENT, 10, TimeUnit.SECONDS),
-    DDSketchHistograms.FACTORY
-    )
-
     // If this fails, it's likely the result of another test loading Config before it can be
     // injected into the bootstrap classpath. If one test extends AgentTestRunner in a module, all tests must extend
     assert Config.getClassLoader() == null: "Config must load on the bootstrap classpath."
 
     configurePreAgent()
 
+    AgentTracer.maybeInstallLegacyContextManager()
+
+    AgentMeter.registerIfAbsent(
+    STATS_D_CLIENT,
+    new MonitoringImpl(STATS_D_CLIENT, 10, TimeUnit.SECONDS),
+    DDSketchHistograms.FACTORY
+    )
+
     TEST_DATA_STREAMS_WRITER = new RecordingDatastreamsPayloadWriter()
     DDAgentFeaturesDiscovery features = new MockFeaturesDiscovery(true)
 

dd-java-agent/instrumentation-testing/src/main/java/datadog/trace/agent/test/AbstractInstrumentationTest.java

@@ -23,6 +23,7 @@
 import datadog.trace.core.TraceCollector;
 import datadog.trace.junit.utils.config.WithConfig;
 import datadog.trace.junit.utils.context.AllowContextTestingExtension;
+import datadog.trace.junit.utils.context.LegacyContextTestingExtension;
 import java.lang.instrument.ClassFileTransformer;
 import java.lang.instrument.Instrumentation;
 import java.util.List;
@@ -54,7 +55,11 @@
  * </ul>
  */
 @WithConfig(key = "detailed.instrumentation.errors", value = "true")
-@ExtendWith({TestClassShadowingExtension.class, AllowContextTestingExtension.class})
+@ExtendWith({
+  TestClassShadowingExtension.class,
+  AllowContextTestingExtension.class,
+  LegacyContextTestingExtension.class
+})
 public abstract class AbstractInstrumentationTest {
   static final Instrumentation INSTRUMENTATION = ByteBuddyAgent.getInstrumentation();
 

dd-trace-core/src/jmh/java/datadog/context/ContextManagerBenchmark.java

@@ -107,9 +107,34 @@ public void setup() {
   }
 
   static ContextManager createManager(String type) {
-    return "Continuable".equals(type)
-        ? new ContinuableScopeManager(0, false)
-        : ThreadLocalContextManager.INSTANCE;
+    if ("Continuable".equals(type)) {
+      ContinuableScopeManager csm = new ContinuableScopeManager(0, false);
+      return new ContextManager() {
+        @Override
+        public Context current() {
+          return csm.currentContext();
+        }
+
+        @Override
+        public ContextScope attach(Context ctx) {
+          return csm.attach(ctx);
+        }
+
+        @Override
+        public Context swap(Context ctx) {
+          return csm.swap(ctx);
+        }
+
+        @Override
+        public ContextContinuation capture(Context ctx) {
+          return csm.capture(ctx);
+        }
+
+        @Override
+        public void addListener(ContextListener l) {}
+      };
+    }
+    return ThreadLocalContextManager.INSTANCE;
   }
 
   static Context[] createContexts() {

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

@@ -22,6 +22,9 @@
 import datadog.communication.ddagent.DDAgentFeaturesDiscovery;
 import datadog.communication.ddagent.ExternalAgentLauncher;
 import datadog.communication.ddagent.SharedCommunicationObjects;
+import datadog.context.Context;
+import datadog.context.ContextContinuation;
+import datadog.context.ContextScope;
 import datadog.context.propagation.Propagators;
 import datadog.environment.ThreadSupport;
 import datadog.logging.RatelimitedLogger;
@@ -2465,4 +2468,24 @@ static TagMap withTracerTags(
     }
     return result.freeze();
   }
+
+  @Override
+  public Context currentContext() {
+    return scopeManager.currentContext();
+  }
+
+  @Override
+  public ContextScope attach(Context context) {
+    return scopeManager.attach(context);
+  }
+
+  @Override
+  public Context swap(Context context) {
+    return scopeManager.swap(context);
+  }
+
+  @Override
+  public ContextContinuation capture(Context context) {
+    return scopeManager.capture(context);
+  }
 }

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

@@ -14,8 +14,6 @@
 
 import datadog.context.Context;
 import datadog.context.ContextContinuation;
-import datadog.context.ContextListener;
-import datadog.context.ContextManager;
 import datadog.context.ContextScope;
 import datadog.logging.RatelimitedLogger;
 import datadog.trace.api.Config;
@@ -47,7 +45,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 ContextManager {
+public final class ContinuableScopeManager {
 
   static final Logger log = LoggerFactory.getLogger(ContinuableScopeManager.class);
   static final RatelimitedLogger ratelimitedLog = new RatelimitedLogger(log, 1, MINUTES);
@@ -95,8 +93,6 @@ public ContinuableScopeManager(
     this.profilingContextIntegration = profilingContextIntegration;
     this.profilingEnabled =
         !(profilingContextIntegration instanceof ProfilingContextIntegration.NoOp);
-
-    ContextManager.register(this);
   }
 
   public AgentScope activateSpan(final AgentSpan span) {
@@ -370,18 +366,15 @@ ScopeStack scopeStack() {
     return this.tlsScopeStack.get();
   }
 
-  @Override
-  public Context current() {
+  public Context currentContext() {
     final ContinuableScope active = scopeStack().active();
     return active == null ? Context.root() : active.context;
   }
 
-  @Override
   public ContextScope attach(@NonNull Context context) {
     return activate(context);
   }
 
-  @Override
   public Context swap(@NonNull Context context) {
     ScopeStack oldStack = tlsScopeStack.get();
     ContinuableScope oldScope = oldStack.top;
@@ -412,7 +405,6 @@ public Context swap(@NonNull Context context) {
     return new ScopeContext(oldStack);
   }
 
-  @Override
   public ContextContinuation capture(@NonNull Context context) {
     // respect async propagation flag for Context.current().capture()
     ContinuableScope activeScope = scopeStack().active();
@@ -431,12 +423,6 @@ public ContextContinuation capture(@NonNull Context context) {
     return new ScopeContinuation(this, context, CONTEXT, traceCollector).register();
   }
 
-  @Override
-  public void addListener(@NonNull ContextListener unused) {
-    // this new API is not expected to be used in legacy mode...
-    log.warn("Unexpected call to ContextManager.addListener(...)");
-  }
-
   static final class ScopeStackThreadLocal extends ThreadLocal<ScopeStack> {
 
     private final ProfilingContextIntegration profilingContextIntegration;