Rework shim between ContextManager and the legacy scope manager

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;