Deprecate AgentScope.Continuation and migrate call sites to ContextContinuation

components/context/src/main/java/datadog/context/ContextBinder.java

@@ -46,4 +46,13 @@ static void register(ContextBinder binder) {
   static boolean allowTesting() {
     return TestContextBinder.register();
   }
+
+  /**
+   * Clears any custom {@link ContextBinder} registered with {@link #register(ContextBinder)}.
+   *
+   * @see #allowTesting()
+   */
+  static void resetForTesting() {
+    ContextProviders.customBinder = null;
+  }
 }

components/context/src/main/java/datadog/context/ContextManager.java

@@ -70,4 +70,13 @@ static void register(ContextManager manager) {
   static boolean allowTesting() {
     return TestContextManager.register();
   }
+
+  /**
+   * Clears any custom {@link ContextManager} registered with {@link #register(ContextManager)}.
+   *
+   * @see #allowTesting()
+   */
+  static void resetForTesting() {
+    ContextProviders.customManager = null;
+  }
 }

components/context/src/main/java/datadog/context/TestContextBinder.java

@@ -29,7 +29,7 @@ public Context detachFrom(Object carrier) {
 
   private static ContextBinder delegate() {
     ContextBinder delegate = ContextProviders.customBinder;
-    if (delegate == TEST_INSTANCE) {
+    if (delegate == TEST_INSTANCE || delegate == null) {
       // fall back to default context binder
       return WeakMapContextBinder.INSTANCE;
     } else {

components/context/src/main/java/datadog/context/TestContextManager.java

@@ -43,7 +43,7 @@ static void clearListeners() {
 
   private static ContextManager delegate() {
     ContextManager delegate = ContextProviders.customManager;
-    if (delegate == TEST_INSTANCE) {
+    if (delegate == TEST_INSTANCE || delegate == null) {
       // fall back to default context manager
       return ThreadLocalContextManager.INSTANCE;
     } else {

components/context/src/test/java/datadog/context/ContextProvidersForkedTest.java

@@ -123,4 +123,90 @@ public void addListener(@Nonnull ContextListener listener) {}
     // NOOP manager, no events emitted
     listener.assertNoEvents();
   }
+
+  @Test
+  void testResetBinder() {
+    assertTrue(ContextBinder.allowTesting());
+
+    Context context = root().with(STRING_KEY, "value");
+    Object carrier = new Object();
+
+    // register a NOOP context binder
+    ContextBinder.register(
+        new ContextBinder() {
+          @Override
+          public Context from(@Nonnull Object carrier) {
+            return root();
+          }
+
+          @Override
+          public void attachTo(@Nonnull Object carrier, @Nonnull Context context) {
+            // no-op
+          }
+
+          @Override
+          public Context detachFrom(@Nonnull Object carrier) {
+            return root();
+          }
+        });
+
+    // NOOP binder, context will always be root
+    context.attachTo(carrier);
+    assertSame(root(), Context.from(carrier));
+
+    // reset back to the default binder
+    ContextBinder.resetForTesting();
+
+    context.attachTo(carrier);
+    assertSame(context, Context.from(carrier));
+    assertSame(context, Context.detachFrom(carrier));
+    assertSame(root(), Context.from(carrier));
+  }
+
+  @Test
+  void testResetManager() {
+    assertTrue(ContextManager.allowTesting());
+
+    Context context = root().with(STRING_KEY, "value");
+
+    // register a NOOP context manager
+    ContextManager.register(
+        new ContextManager() {
+          @Override
+          public Context current() {
+            return root();
+          }
+
+          @Override
+          public ContextScope attach(@Nonnull Context context) {
+            return new NoopContextScope(root());
+          }
+
+          @Override
+          public Context swap(@Nonnull Context context) {
+            return root();
+          }
+
+          @Override
+          public ContextContinuation capture(@Nonnull Context context) {
+            return new NoopContextContinuation(root());
+          }
+
+          @Override
+          public void addListener(@Nonnull ContextListener listener) {}
+        });
+
+    // NOOP manager, context will always be root
+    try (ContextScope scope = context.attach()) {
+      assertSame(root(), Context.current());
+    }
+
+    // reset back to the default manager
+    ContextManager.resetForTesting();
+
+    try (ContextScope scope = context.attach()) {
+      assertSame(context, scope.context());
+      assertSame(context, Context.current());
+    }
+  }
 }

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

@@ -26,6 +26,7 @@
 import datadog.instrument.utils.ClassLoaderValue;
 import datadog.metrics.api.statsd.StatsDClientManager;
 import datadog.trace.api.Config;
+import datadog.trace.api.InstrumenterConfig;
 import datadog.trace.api.Platform;
 import datadog.trace.api.WithGlobalTracer;
 import datadog.trace.api.appsec.AppSecEventTracker;
@@ -337,6 +338,10 @@ public static void start(
       StaticEventLogger.end("crashtracking");
     }
 
+    if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) {
+      AgentTracer.installLegacyContextManager();
+    }
+
     startDatadogAgent(initTelemetry, inst);
 
     final EnumSet<Library> libraries = detectLibraries(log);

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/AdviceUtils.java

@@ -3,8 +3,9 @@
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.isAsyncPropagationEnabled;
 
+import datadog.context.ContextContinuation;
+import datadog.context.ContextScope;
 import datadog.trace.bootstrap.ContextStore;
-import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 
 /** Helper utils for Runnable/Callable instrumentation */
@@ -18,16 +19,16 @@ public class AdviceUtils {
    * @param <T> task's type
    * @return scope if scope was started, or null
    */
-  public static <T> AgentScope startTaskScope(
+  public static <T> ContextScope startTaskScope(
       final ContextStore<T, State> contextStore, final T task) {
     return startTaskScope(contextStore.get(task));
   }
 
-  public static AgentScope startTaskScope(State state) {
+  public static ContextScope startTaskScope(State state) {
     if (state != null) {
-      final AgentScope.Continuation continuation = state.getAndResetContinuation();
+      final ContextContinuation continuation = state.getAndResetContinuation();
       if (continuation != null) {
-        final AgentScope scope = continuation.activate();
+        final ContextScope scope = continuation.resume();
         // important - stop timing after the scope has been activated so the time in the queue can
         // be attributed to the correct context without duplicating the propagated information
         state.stopTiming();
@@ -37,7 +38,7 @@ public static AgentScope startTaskScope(State state) {
     return null;
   }
 
-  public static void endTaskScope(final AgentScope scope) {
+  public static void endTaskScope(final ContextScope scope) {
     if (null != scope) {
       scope.close();
     }

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/ComparableRunnable.java

@@ -1,11 +1,11 @@
 package datadog.trace.bootstrap.instrumentation.java.concurrent;
 
-import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.context.ContextContinuation;
 
 public final class ComparableRunnable<T extends Runnable & Comparable<T>> extends Wrapper<T>
     implements Comparable<ComparableRunnable<T>> {
 
-  public ComparableRunnable(T delegate, AgentScope.Continuation continuation) {
+  public ComparableRunnable(T delegate, ContextContinuation continuation) {
     super(delegate, continuation);
   }
 

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/ConcurrentState.java

@@ -4,8 +4,9 @@
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.isAsyncPropagationEnabled;
 import static datadog.trace.bootstrap.instrumentation.java.concurrent.ContinuationClaim.CLAIMED;
 
+import datadog.context.ContextContinuation;
+import datadog.context.ContextScope;
 import datadog.trace.bootstrap.ContextStore;
-import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;
 import org.slf4j.Logger;
@@ -22,12 +23,12 @@ public final class ConcurrentState {
 
   public static ContextStore.Factory<ConcurrentState> FACTORY = ConcurrentState::new;
 
-  private volatile AgentScope.Continuation continuation = null;
+  private volatile ContextContinuation continuation = null;
 
-  private static final AtomicReferenceFieldUpdater<ConcurrentState, AgentScope.Continuation>
+  private static final AtomicReferenceFieldUpdater<ConcurrentState, ContextContinuation>
       CONTINUATION =
           AtomicReferenceFieldUpdater.newUpdater(
-              ConcurrentState.class, AgentScope.Continuation.class, "continuation");
+              ConcurrentState.class, ContextContinuation.class, "continuation");
 
   private ConcurrentState() {}
 
@@ -44,7 +45,7 @@ public static <K> ConcurrentState captureContinuation(
     return state;
   }
 
-  public static <K> AgentScope activateAndContinueContinuation(
+  public static <K> ContextScope activateAndContinueContinuation(
       ContextStore<K, ConcurrentState> contextStore, K key) {
     final ConcurrentState state = contextStore.get(key);
     if (state == null) {
@@ -54,7 +55,10 @@ public static <K> AgentScope activateAndContinueContinuation(
   }
 
   public static <K> void closeScope(
-      ContextStore<K, ConcurrentState> contextStore, K key, AgentScope scope, Throwable throwable) {
+      ContextStore<K, ConcurrentState> contextStore,
+      K key,
+      ContextScope scope,
+      Throwable throwable) {
     final ConcurrentState state = contextStore.get(key);
     if (scope != null) {
       scope.close();
@@ -88,27 +92,27 @@ private boolean captureAndSetContinuation(final AgentSpan span) {
     return false;
   }
 
-  private AgentScope activateAndContinueContinuation() {
-    final AgentScope.Continuation continuation = CONTINUATION.get(this);
+  private ContextScope activateAndContinueContinuation() {
+    final ContextContinuation continuation = CONTINUATION.get(this);
     if (continuation != null && continuation != CLAIMED) {
-      return continuation.activate();
+      return continuation.resume();
     }
     return null;
   }
 
   private void cancelContinuation() {
-    final AgentScope.Continuation continuation = CONTINUATION.get(this);
+    final ContextContinuation continuation = CONTINUATION.get(this);
     if (continuation != null && continuation != CLAIMED) {
-      continuation.cancel();
+      continuation.release();
     }
   }
 
   private void cancelAndClearContinuation() {
-    final AgentScope.Continuation continuation = CONTINUATION.get(this);
+    final ContextContinuation continuation = CONTINUATION.get(this);
     if (continuation != null && continuation != CLAIMED) {
       // We should never be able to reuse this state
       CONTINUATION.compareAndSet(this, continuation, CLAIMED);
-      continuation.cancel();
+      continuation.release();
     }
   }
 }

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/ContinuationClaim.java

@@ -1,29 +1,30 @@
 package datadog.trace.bootstrap.instrumentation.java.concurrent;
 
-import datadog.trace.bootstrap.instrumentation.api.AgentScope;
-import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.context.Context;
+import datadog.context.ContextContinuation;
+import datadog.context.ContextScope;
 
-final class ContinuationClaim implements AgentScope.Continuation {
+final class ContinuationClaim implements ContextContinuation {
 
   public static final ContinuationClaim CLAIMED = new ContinuationClaim();
 
   @Override
-  public AgentScope.Continuation hold() {
+  public ContextContinuation hold() {
     throw new IllegalStateException();
   }
 
   @Override
-  public AgentScope activate() {
+  public ContextScope resume() {
     throw new IllegalStateException();
   }
 
   @Override
-  public AgentSpan span() {
+  public Context context() {
     throw new IllegalStateException();
   }
 
   @Override
-  public void cancel() {
+  public void release() {
     throw new IllegalStateException();
   }
 }