Support re-registration of different ContextManagers/Binders during testing

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

@@ -29,9 +29,21 @@ public interface ContextBinder {
   /**
    * Requests use of a custom {@link ContextBinder}.
    *
-   * @param binder the binder to use (will replace any other binder in use).
+   * <p>Once the registered binder is used it cannot be replaced and this method will have no
+   * effect. To test different binders, make sure {@link #allowTesting()} is called early on.
+   *
+   * @param binder the binder to use.
    */
   static void register(ContextBinder binder) {
     ContextProviders.customBinder = binder;
   }
+
+  /**
+   * Allow re-registration of custom {@link ContextBinder}s for testing.
+   *
+   * @return {@code true} if re-registration is allowed; otherwise {@code false}
+   */
+  static boolean allowTesting() {
+    return TestContextBinder.register();
+  }
 }

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

@@ -28,9 +28,21 @@ public interface ContextManager {
   /**
    * Requests use of a custom {@link ContextManager}.
    *
-   * @param manager the manager to use (will replace any other manager in use).
+   * <p>Once the registered manager is used it cannot be replaced and this method will have no
+   * effect. To test different managers, make sure {@link #allowTesting()} is called early on.
+   *
+   * @param manager the manager to use.
    */
   static void register(ContextManager manager) {
     ContextProviders.customManager = manager;
   }
+
+  /**
+   * Allow re-registration of custom {@link ContextManager}s for testing.
+   *
+   * @return {@code true} if re-registration is allowed; otherwise {@code false}
+   */
+  static boolean allowTesting() {
+    return TestContextManager.register();
+  }
 }

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

@@ -10,14 +10,14 @@ private static final class ProvidedManager {
     static final ContextManager INSTANCE =
         null != ContextProviders.customManager
             ? ContextProviders.customManager
-            : new ThreadLocalContextManager();
+            : ThreadLocalContextManager.INSTANCE;
   }
 
   private static final class ProvidedBinder {
     static final ContextBinder INSTANCE =
         null != ContextProviders.customBinder
             ? ContextProviders.customBinder
-            : new WeakMapContextBinder();
+            : WeakMapContextBinder.INSTANCE;
   }
 
   static ContextManager manager() {

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

@@ -0,0 +1,39 @@
+package datadog.context;
+
+/** Test class that always delegates to the latest registered {@link ContextBinder}. */
+final class TestContextBinder implements ContextBinder {
+  private static final ContextBinder TEST_INSTANCE = new TestContextBinder();
+
+  private TestContextBinder() {}
+
+  static boolean register() {
+    // attempt to register before binder choice is locked, then check if we succeeded
+    ContextProviders.customBinder = TEST_INSTANCE;
+    return ContextProviders.binder() == TEST_INSTANCE;
+  }
+
+  @Override
+  public Context from(Object carrier) {
+    return delegate().from(carrier);
+  }
+
+  @Override
+  public void attachTo(Object carrier, Context context) {
+    delegate().attachTo(carrier, context);
+  }
+
+  @Override
+  public Context detachFrom(Object carrier) {
+    return delegate().detachFrom(carrier);
+  }
+
+  private static ContextBinder delegate() {
+    ContextBinder delegate = ContextProviders.customBinder;
+    if (delegate == TEST_INSTANCE) {
+      // fall back to default context binder
+      return WeakMapContextBinder.INSTANCE;
+    } else {
+      return delegate;
+    }
+  }
+}

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

@@ -0,0 +1,39 @@
+package datadog.context;
+
+/** Test class that always delegates to the latest registered {@link ContextManager}. */
+final class TestContextManager implements ContextManager {
+  private static final ContextManager TEST_INSTANCE = new TestContextManager();
+
+  private TestContextManager() {}
+
+  static boolean register() {
+    // attempt to register before manager choice is locked, then check if we succeeded
+    ContextProviders.customManager = TEST_INSTANCE;
+    return ContextProviders.manager() == TEST_INSTANCE;
+  }
+
+  @Override
+  public Context current() {
+    return delegate().current();
+  }
+
+  @Override
+  public ContextScope attach(Context context) {
+    return delegate().attach(context);
+  }
+
+  @Override
+  public Context swap(Context context) {
+    return delegate().swap(context);
+  }
+
+  private static ContextManager delegate() {
+    ContextManager delegate = ContextProviders.customManager;
+    if (delegate == TEST_INSTANCE) {
+      // fall back to default context manager
+      return ThreadLocalContextManager.INSTANCE;
+    } else {
+      return delegate;
+    }
+  }
+}

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

@@ -2,6 +2,8 @@
 
 /** {@link ContextManager} that uses a {@link ThreadLocal} to track context per thread. */
 final class ThreadLocalContextManager implements ContextManager {
+  static final ContextManager INSTANCE = new ThreadLocalContextManager();
+
   private static final ThreadLocal<Context[]> CURRENT_HOLDER =
       ThreadLocal.withInitial(() -> new Context[] {EmptyContext.INSTANCE});
 

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

@@ -11,6 +11,8 @@
 /** {@link ContextBinder} that uses a global weak map of carriers to contexts. */
 @ParametersAreNonnullByDefault
 final class WeakMapContextBinder implements ContextBinder {
+  static final ContextBinder INSTANCE = new WeakMapContextBinder();
+
   private static final Map<Object, Context> TRACKED = synchronizedMap(new WeakHashMap<>());
 
   @Override

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

@@ -3,14 +3,27 @@
 import static datadog.context.Context.root;
 import static datadog.context.ContextTest.STRING_KEY;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import javax.annotation.Nonnull;
 import org.junit.jupiter.api.Test;
 
-class ContextProviderForkedTest {
+class ContextProvidersForkedTest {
   @Test
   void testCustomBinder() {
-    // register a NOOP context binder
+    assertTrue(ContextBinder.allowTesting());
+
+    Context context = root().with(STRING_KEY, "value");
+    Object carrier = new Object();
+
+    // should delegate to the default binder
+    context.attachTo(carrier);
+    assertNotEquals(root(), Context.from(carrier));
+    assertEquals(context, Context.detachFrom(carrier));
+    assertEquals(root(), Context.from(carrier));
+
+    // now register a NOOP context binder
     ContextBinder.register(
         new ContextBinder() {
           @Override
@@ -29,17 +42,28 @@ public Context detachFrom(@Nonnull Object carrier) {
           }
         });
 
-    Context context = root().with(STRING_KEY, "value");
-
     // NOOP binder, context will always be root
-    Object carrier = new Object();
     context.attachTo(carrier);
     assertEquals(root(), Context.from(carrier));
+    assertEquals(root(), Context.detachFrom(carrier));
   }
 
   @Test
   void testCustomManager() {
-    // register a NOOP context manager
+    assertTrue(ContextManager.allowTesting());
+
+    Context context = root().with(STRING_KEY, "value");
+
+    // should delegate to the default manager
+    try (ContextScope ignored = context.attach()) {
+      assertNotEquals(root(), Context.current());
+    }
+
+    Context swapped = context.swap();
+    assertNotEquals(root(), Context.current());
+    swapped.swap();
+
+    // now register a NOOP context manager
     ContextManager.register(
         new ContextManager() {
           @Override
@@ -68,11 +92,14 @@ public Context swap(Context context) {
           }
         });
 
-    Context context = root().with(STRING_KEY, "value");
-
     // NOOP manager, context will always be root
     try (ContextScope ignored = context.attach()) {
       assertEquals(root(), Context.current());
     }
+
+    // NOOP manager, context will always be root
+    swapped = context.swap();
+    assertEquals(root(), Context.current());
+    swapped.swap();
   }
 }

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

@@ -0,0 +1,34 @@
+package datadog.context;
+
+import static datadog.context.Context.root;
+import static datadog.context.ContextTest.STRING_KEY;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotEquals;
+
+import org.junit.jupiter.api.Test;
+
+class ContextProvidersTest {
+  @Test
+  void testCannotChangeBinderAfterUse() {
+    Context context = root().with(STRING_KEY, "value");
+    Object carrier = new Object();
+
+    context.attachTo(carrier);
+    Context.detachFrom(carrier);
+
+    // cannot change binder at this late stage
+    assertFalse(ContextBinder.allowTesting());
+  }
+
+  @Test
+  void testCannotChangeManagerAfterUse() {
+    Context context = root().with(STRING_KEY, "value");
+
+    try (ContextScope ignored = context.attach()) {
+      assertNotEquals(root(), Context.current());
+    }
+
+    // cannot change manager at this late stage
+    assertFalse(ContextManager.allowTesting());
+  }
+}

dd-java-agent/instrumentation/opentelemetry/opentelemetry-1.4/src/test/groovy/opentelemetry14/context/ContextForkedTest.groovy → dd-java-agent/instrumentation/opentelemetry/opentelemetry-1.4/src/test/groovy/opentelemetry14/context/ContextTest.groovy

@@ -14,7 +14,7 @@ import static datadog.opentelemetry.shim.context.OtelContext.OTEL_CONTEXT_ROOT_S
 import static datadog.opentelemetry.shim.context.OtelContext.OTEL_CONTEXT_SPAN_KEY
 import static datadog.opentelemetry.shim.trace.OtelConventions.SPAN_KIND_INTERNAL
 
-class ContextForkedTest extends AgentTestRunner {
+class ContextTest extends AgentTestRunner {
   @Subject
   def tracer = GlobalOpenTelemetry.get().tracerProvider.get("context-instrumentation")
 

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

@@ -315,7 +315,6 @@ public static class CoreTracerBuilder {
     private SingleSpanSampler singleSpanSampler;
     private HttpCodec.Injector injector;
     private HttpCodec.Extractor extractor;
-    private ContinuableScopeManager scopeManager;
     private Map<String, ?> localRootSpanTags;
     private Map<String, ?> defaultSpanTags;
     private Map<String, String> serviceNameMappings;

utils/test-utils/src/main/groovy/datadog/trace/test/util/DDSpecification.groovy

@@ -13,6 +13,8 @@ import java.lang.reflect.Modifier
 abstract class DDSpecification extends Specification {
   private static final CHECK_TIMEOUT_MS = 3000
 
+  static final String CONTEXT_BINDER = "datadog.context.ContextBinder"
+  static final String CONTEXT_MANAGER = "datadog.context.ContextManager"
   static final String INST_CONFIG = "datadog.trace.api.InstrumenterConfig"
   static final String CONFIG = "datadog.trace.api.Config"
 
@@ -22,9 +24,12 @@ abstract class DDSpecification extends Specification {
   private static Constructor configConstructor
 
   static {
+    allowContextTesting()
     makeConfigInstanceModifiable()
   }
 
+  private static Boolean contextTestingAllowed
+
   // Keep track of config instance already made modifiable
   private static isConfigInstanceModifiable = false
   static configModificationFailed = false
@@ -42,6 +47,21 @@ abstract class DDSpecification extends Specification {
   @Shared
   private boolean ignoreThreadCleanup
 
+  static void allowContextTesting() {
+    if (contextTestingAllowed == null) {
+      try {
+        contextTestingAllowed =
+          Class.forName(CONTEXT_BINDER).allowTesting() &&
+          Class.forName(CONTEXT_MANAGER).allowTesting()
+      } catch (ClassNotFoundException e) {
+        // don't block testing if these types aren't found (project doesn't use context API)
+        contextTestingAllowed = e.message == CONTEXT_BINDER || e.message == CONTEXT_MANAGER
+      } catch (Throwable ignore) {
+        contextTestingAllowed = false
+      }
+    }
+  }
+
   static void makeConfigInstanceModifiable() {
     if (isConfigInstanceModifiable || configModificationFailed) {
       return
@@ -92,6 +112,7 @@ abstract class DDSpecification extends Specification {
     assert !configModificationFailed: "Config class modification failed.  Ensure all test classes extend DDSpecification"
     assert System.getenv().findAll { it.key.startsWith("DD_") }.isEmpty()
     assert systemPropertiesExceptAllowed().findAll { it.key.toString().startsWith("dd.") }.isEmpty()
+    assert contextTestingAllowed: "Context not ready for testing.  Ensure all test classes extend DDSpecification"
 
     if (getDDThreads().isEmpty()) {
       ignoreThreadCleanup = false