diff --git a/components/context/src/main/java/datadog/context/ContextBinder.java b/components/context/src/main/java/datadog/context/ContextBinder.java index 385cca1f6c5..8a268028c52 100644 --- a/components/context/src/main/java/datadog/context/ContextBinder.java +++ b/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). + *

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(); + } } diff --git a/components/context/src/main/java/datadog/context/ContextManager.java b/components/context/src/main/java/datadog/context/ContextManager.java index 9c4a58cfa19..af0a2b9289a 100644 --- a/components/context/src/main/java/datadog/context/ContextManager.java +++ b/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). + *

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(); + } } diff --git a/components/context/src/main/java/datadog/context/ContextProviders.java b/components/context/src/main/java/datadog/context/ContextProviders.java index 6895446947a..c4421b0a2ab 100644 --- a/components/context/src/main/java/datadog/context/ContextProviders.java +++ b/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() { diff --git a/components/context/src/main/java/datadog/context/TestContextBinder.java b/components/context/src/main/java/datadog/context/TestContextBinder.java new file mode 100644 index 00000000000..e1fbbcf0a7f --- /dev/null +++ b/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; + } + } +} diff --git a/components/context/src/main/java/datadog/context/TestContextManager.java b/components/context/src/main/java/datadog/context/TestContextManager.java new file mode 100644 index 00000000000..9c60f4dc2e0 --- /dev/null +++ b/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; + } + } +} diff --git a/components/context/src/main/java/datadog/context/ThreadLocalContextManager.java b/components/context/src/main/java/datadog/context/ThreadLocalContextManager.java index 27c17445d14..0d652868e52 100644 --- a/components/context/src/main/java/datadog/context/ThreadLocalContextManager.java +++ b/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 CURRENT_HOLDER = ThreadLocal.withInitial(() -> new Context[] {EmptyContext.INSTANCE}); diff --git a/components/context/src/main/java/datadog/context/WeakMapContextBinder.java b/components/context/src/main/java/datadog/context/WeakMapContextBinder.java index eea3728fc8a..9b5fd24299e 100644 --- a/components/context/src/main/java/datadog/context/WeakMapContextBinder.java +++ b/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 TRACKED = synchronizedMap(new WeakHashMap<>()); @Override diff --git a/components/context/src/test/java/datadog/context/ContextProviderForkedTest.java b/components/context/src/test/java/datadog/context/ContextProvidersForkedTest.java similarity index 64% rename from components/context/src/test/java/datadog/context/ContextProviderForkedTest.java rename to components/context/src/test/java/datadog/context/ContextProvidersForkedTest.java index db5d2791057..915186554a6 100644 --- a/components/context/src/test/java/datadog/context/ContextProviderForkedTest.java +++ b/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(); } } diff --git a/components/context/src/test/java/datadog/context/ContextProvidersTest.java b/components/context/src/test/java/datadog/context/ContextProvidersTest.java new file mode 100644 index 00000000000..31cff5d3254 --- /dev/null +++ b/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()); + } +} diff --git a/dd-java-agent/instrumentation/opentelemetry/opentelemetry-1.4/src/test/groovy/opentelemetry14/context/ContextForkedTest.groovy b/dd-java-agent/instrumentation/opentelemetry/opentelemetry-1.4/src/test/groovy/opentelemetry14/context/ContextTest.groovy similarity index 99% rename from dd-java-agent/instrumentation/opentelemetry/opentelemetry-1.4/src/test/groovy/opentelemetry14/context/ContextForkedTest.groovy rename to dd-java-agent/instrumentation/opentelemetry/opentelemetry-1.4/src/test/groovy/opentelemetry14/context/ContextTest.groovy index 4437f7f3f72..46309f077c4 100644 --- a/dd-java-agent/instrumentation/opentelemetry/opentelemetry-1.4/src/test/groovy/opentelemetry14/context/ContextForkedTest.groovy +++ b/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") diff --git a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java index ffbcde5e9df..bc2fc9d14bc 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java +++ b/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 localRootSpanTags; private Map defaultSpanTags; private Map serviceNameMappings; diff --git a/utils/test-utils/src/main/groovy/datadog/trace/test/util/DDSpecification.groovy b/utils/test-utils/src/main/groovy/datadog/trace/test/util/DDSpecification.groovy index 734f347e36d..43693470451 100644 --- a/utils/test-utils/src/main/groovy/datadog/trace/test/util/DDSpecification.groovy +++ b/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