Support re-registration of different ContextManagers/Binders during testing (#8793)

Author: mcculls

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 renamed to 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 renamed to 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")