Skip to content

Commit 90ca83f

Browse files
mccullsmcculls
authored
Support re-registration of different ContextManagers/Binders during testing (#8793)
1 parent 8bfbf8e commit 90ca83f

12 files changed

Lines changed: 201 additions & 14 deletions

File tree

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

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,9 +29,21 @@ public interface ContextBinder {
2929
/**
3030
* Requests use of a custom {@link ContextBinder}.
3131
*
32-
* @param binder the binder to use (will replace any other binder in use).
32+
* <p>Once the registered binder is used it cannot be replaced and this method will have no
33+
* effect. To test different binders, make sure {@link #allowTesting()} is called early on.
34+
*
35+
* @param binder the binder to use.
3336
*/
3437
static void register(ContextBinder binder) {
3538
ContextProviders.customBinder = binder;
3639
}
40+
41+
/**
42+
* Allow re-registration of custom {@link ContextBinder}s for testing.
43+
*
44+
* @return {@code true} if re-registration is allowed; otherwise {@code false}
45+
*/
46+
static boolean allowTesting() {
47+
return TestContextBinder.register();
48+
}
3749
}

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

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,9 +28,21 @@ public interface ContextManager {
2828
/**
2929
* Requests use of a custom {@link ContextManager}.
3030
*
31-
* @param manager the manager to use (will replace any other manager in use).
31+
* <p>Once the registered manager is used it cannot be replaced and this method will have no
32+
* effect. To test different managers, make sure {@link #allowTesting()} is called early on.
33+
*
34+
* @param manager the manager to use.
3235
*/
3336
static void register(ContextManager manager) {
3437
ContextProviders.customManager = manager;
3538
}
39+
40+
/**
41+
* Allow re-registration of custom {@link ContextManager}s for testing.
42+
*
43+
* @return {@code true} if re-registration is allowed; otherwise {@code false}
44+
*/
45+
static boolean allowTesting() {
46+
return TestContextManager.register();
47+
}
3648
}

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,14 +10,14 @@ private static final class ProvidedManager {
1010
static final ContextManager INSTANCE =
1111
null != ContextProviders.customManager
1212
? ContextProviders.customManager
13-
: new ThreadLocalContextManager();
13+
: ThreadLocalContextManager.INSTANCE;
1414
}
1515

1616
private static final class ProvidedBinder {
1717
static final ContextBinder INSTANCE =
1818
null != ContextProviders.customBinder
1919
? ContextProviders.customBinder
20-
: new WeakMapContextBinder();
20+
: WeakMapContextBinder.INSTANCE;
2121
}
2222

2323
static ContextManager manager() {

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

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
package datadog.context;
2+
3+
/** Test class that always delegates to the latest registered {@link ContextBinder}. */
4+
final class TestContextBinder implements ContextBinder {
5+
private static final ContextBinder TEST_INSTANCE = new TestContextBinder();
6+
7+
private TestContextBinder() {}
8+
9+
static boolean register() {
10+
// attempt to register before binder choice is locked, then check if we succeeded
11+
ContextProviders.customBinder = TEST_INSTANCE;
12+
return ContextProviders.binder() == TEST_INSTANCE;
13+
}
14+
15+
@Override
16+
public Context from(Object carrier) {
17+
return delegate().from(carrier);
18+
}
19+
20+
@Override
21+
public void attachTo(Object carrier, Context context) {
22+
delegate().attachTo(carrier, context);
23+
}
24+
25+
@Override
26+
public Context detachFrom(Object carrier) {
27+
return delegate().detachFrom(carrier);
28+
}
29+
30+
private static ContextBinder delegate() {
31+
ContextBinder delegate = ContextProviders.customBinder;
32+
if (delegate == TEST_INSTANCE) {
33+
// fall back to default context binder
34+
return WeakMapContextBinder.INSTANCE;
35+
} else {
36+
return delegate;
37+
}
38+
}
39+
}

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

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
package datadog.context;
2+
3+
/** Test class that always delegates to the latest registered {@link ContextManager}. */
4+
final class TestContextManager implements ContextManager {
5+
private static final ContextManager TEST_INSTANCE = new TestContextManager();
6+
7+
private TestContextManager() {}
8+
9+
static boolean register() {
10+
// attempt to register before manager choice is locked, then check if we succeeded
11+
ContextProviders.customManager = TEST_INSTANCE;
12+
return ContextProviders.manager() == TEST_INSTANCE;
13+
}
14+
15+
@Override
16+
public Context current() {
17+
return delegate().current();
18+
}
19+
20+
@Override
21+
public ContextScope attach(Context context) {
22+
return delegate().attach(context);
23+
}
24+
25+
@Override
26+
public Context swap(Context context) {
27+
return delegate().swap(context);
28+
}
29+
30+
private static ContextManager delegate() {
31+
ContextManager delegate = ContextProviders.customManager;
32+
if (delegate == TEST_INSTANCE) {
33+
// fall back to default context manager
34+
return ThreadLocalContextManager.INSTANCE;
35+
} else {
36+
return delegate;
37+
}
38+
}
39+
}

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@
22

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

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@
1111
/** {@link ContextBinder} that uses a global weak map of carriers to contexts. */
1212
@ParametersAreNonnullByDefault
1313
final class WeakMapContextBinder implements ContextBinder {
14+
static final ContextBinder INSTANCE = new WeakMapContextBinder();
15+
1416
private static final Map<Object, Context> TRACKED = synchronizedMap(new WeakHashMap<>());
1517

1618
@Override

components/context/src/test/java/datadog/context/ContextProviderForkedTest.java renamed to components/context/src/test/java/datadog/context/ContextProvidersForkedTest.java

Lines changed: 35 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -3,14 +3,27 @@
33
import static datadog.context.Context.root;
44
import static datadog.context.ContextTest.STRING_KEY;
55
import static org.junit.jupiter.api.Assertions.assertEquals;
6+
import static org.junit.jupiter.api.Assertions.assertNotEquals;
7+
import static org.junit.jupiter.api.Assertions.assertTrue;
68

79
import javax.annotation.Nonnull;
810
import org.junit.jupiter.api.Test;
911

10-
class ContextProviderForkedTest {
12+
class ContextProvidersForkedTest {
1113
@Test
1214
void testCustomBinder() {
13-
// register a NOOP context binder
15+
assertTrue(ContextBinder.allowTesting());
16+
17+
Context context = root().with(STRING_KEY, "value");
18+
Object carrier = new Object();
19+
20+
// should delegate to the default binder
21+
context.attachTo(carrier);
22+
assertNotEquals(root(), Context.from(carrier));
23+
assertEquals(context, Context.detachFrom(carrier));
24+
assertEquals(root(), Context.from(carrier));
25+
26+
// now register a NOOP context binder
1427
ContextBinder.register(
1528
new ContextBinder() {
1629
@Override
@@ -29,17 +42,28 @@ public Context detachFrom(@Nonnull Object carrier) {
2942
}
3043
});
3144

32-
Context context = root().with(STRING_KEY, "value");
33-
3445
// NOOP binder, context will always be root
35-
Object carrier = new Object();
3646
context.attachTo(carrier);
3747
assertEquals(root(), Context.from(carrier));
48+
assertEquals(root(), Context.detachFrom(carrier));
3849
}
3950

4051
@Test
4152
void testCustomManager() {
42-
// register a NOOP context manager
53+
assertTrue(ContextManager.allowTesting());
54+
55+
Context context = root().with(STRING_KEY, "value");
56+
57+
// should delegate to the default manager
58+
try (ContextScope ignored = context.attach()) {
59+
assertNotEquals(root(), Context.current());
60+
}
61+
62+
Context swapped = context.swap();
63+
assertNotEquals(root(), Context.current());
64+
swapped.swap();
65+
66+
// now register a NOOP context manager
4367
ContextManager.register(
4468
new ContextManager() {
4569
@Override
@@ -68,11 +92,14 @@ public Context swap(Context context) {
6892
}
6993
});
7094

71-
Context context = root().with(STRING_KEY, "value");
72-
7395
// NOOP manager, context will always be root
7496
try (ContextScope ignored = context.attach()) {
7597
assertEquals(root(), Context.current());
7698
}
99+
100+
// NOOP manager, context will always be root
101+
swapped = context.swap();
102+
assertEquals(root(), Context.current());
103+
swapped.swap();
77104
}
78105
}

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

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
package datadog.context;
2+
3+
import static datadog.context.Context.root;
4+
import static datadog.context.ContextTest.STRING_KEY;
5+
import static org.junit.jupiter.api.Assertions.assertFalse;
6+
import static org.junit.jupiter.api.Assertions.assertNotEquals;
7+
8+
import org.junit.jupiter.api.Test;
9+
10+
class ContextProvidersTest {
11+
@Test
12+
void testCannotChangeBinderAfterUse() {
13+
Context context = root().with(STRING_KEY, "value");
14+
Object carrier = new Object();
15+
16+
context.attachTo(carrier);
17+
Context.detachFrom(carrier);
18+
19+
// cannot change binder at this late stage
20+
assertFalse(ContextBinder.allowTesting());
21+
}
22+
23+
@Test
24+
void testCannotChangeManagerAfterUse() {
25+
Context context = root().with(STRING_KEY, "value");
26+
27+
try (ContextScope ignored = context.attach()) {
28+
assertNotEquals(root(), Context.current());
29+
}
30+
31+
// cannot change manager at this late stage
32+
assertFalse(ContextManager.allowTesting());
33+
}
34+
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ import static datadog.opentelemetry.shim.context.OtelContext.OTEL_CONTEXT_ROOT_S
1414
import static datadog.opentelemetry.shim.context.OtelContext.OTEL_CONTEXT_SPAN_KEY
1515
import static datadog.opentelemetry.shim.trace.OtelConventions.SPAN_KIND_INTERNAL
1616

17-
class ContextForkedTest extends AgentTestRunner {
17+
class ContextTest extends AgentTestRunner {
1818
@Subject
1919
def tracer = GlobalOpenTelemetry.get().tracerProvider.get("context-instrumentation")
2020

0 commit comments

Comments
 (0)