From c0c3c753a040f8a55d628c8a748998a7b8a3dc0b Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Sat, 4 Jul 2026 14:51:14 +0100 Subject: [PATCH] Replace onAttach/onDetach context event pair with onUpdate(before, after) --- .../java/datadog/context/ContextListener.java | 14 +-- .../context/ThreadLocalContextManager.java | 63 ++++-------- .../context/ContextContinuationTest.java | 96 ++++++++++--------- .../context/ContextListenerEventTest.java | 51 +++++----- .../context/ContextListenerExceptionTest.java | 14 +-- .../java/datadog/context/ContextTestBase.java | 41 +++----- 6 files changed, 114 insertions(+), 165 deletions(-) diff --git a/components/context/src/main/java/datadog/context/ContextListener.java b/components/context/src/main/java/datadog/context/ContextListener.java index 1bf06a7e974..f657fb769aa 100644 --- a/components/context/src/main/java/datadog/context/ContextListener.java +++ b/components/context/src/main/java/datadog/context/ContextListener.java @@ -4,18 +4,12 @@ public interface ContextListener { /** - * Notifies that the given context has been attached to the current execution unit. + * Notifies that the context has been updated for the current execution unit. * - * @param context the attached context. + * @param before the context before. + * @param after the context after. */ - default void onAttach(Context context) {} - - /** - * Notifies that the given context has been detached from the current execution unit. - * - * @param context the detached context. - */ - default void onDetach(Context context) {} + default void onUpdate(Context before, Context after) {} /** * Notifies that the given context has been captured by a continuation. diff --git a/components/context/src/main/java/datadog/context/ThreadLocalContextManager.java b/components/context/src/main/java/datadog/context/ThreadLocalContextManager.java index 9b94f578409..ccb3a641565 100644 --- a/components/context/src/main/java/datadog/context/ThreadLocalContextManager.java +++ b/components/context/src/main/java/datadog/context/ThreadLocalContextManager.java @@ -8,12 +8,12 @@ final class ThreadLocalContextManager implements ContextManager { static final ThreadLocalContextManager INSTANCE = new ThreadLocalContextManager(); - private static final NoopContextContinuation ROOT_CONTINUATION = - new NoopContextContinuation(Context.root()); - private static final ThreadLocal CONTEXT_HOLDER = ThreadLocal.withInitial(ContextHolder::new); + private static final NoopContextContinuation ROOT_CONTINUATION = + new NoopContextContinuation(Context.root()); + private final Object listenersWriteLock = new Object(); volatile ContextListener[] listeners = {}; @@ -40,16 +40,11 @@ ContextScope doAttach(Context context, @Nullable ContextContinuationImpl continu return context.asScope(); // convert to scope without attaching } - ContextListener[] ls = listeners; - notifyDetach(beforeAttach, ls); holder.current = context; - notifyAttach(context, ls); - - if (continuation == null) { - return new ContextScopeImpl(context, holder, beforeAttach); - } else { - return new ResumedScopeImpl(context, holder, beforeAttach, continuation); - } + notifyUpdate(listeners, beforeAttach, context); + return continuation == null + ? new ContextScopeImpl(context, holder, beforeAttach) + : new ResumedScopeImpl(context, holder, beforeAttach, continuation); } @Override @@ -61,21 +56,14 @@ public Context swap(Context context) { return beforeSwap; } - ContextListener[] ls = listeners; - notifyDetach(beforeSwap, ls); holder.current = context; - notifyAttach(context, ls); - + notifyUpdate(listeners, beforeSwap, context); return beforeSwap; } @Override public ContextContinuation capture(Context context) { - if (context == Context.root()) { - return ROOT_CONTINUATION; - } else { - return new ContextContinuationImpl(context); - } + return context == Context.root() ? ROOT_CONTINUATION : new ContextContinuationImpl(context); } @Override @@ -99,31 +87,16 @@ void clearListeners() { } } - static void notifyAttach(Context context, ContextListener[] listeners) { - if (context == Context.root()) { - return; // don't emit attach events for the default "no context" case - } - for (ContextListener l : listeners) { - try { - l.onAttach(context); - } catch (Throwable ignore) { - } - } - } - - static void notifyDetach(Context context, ContextListener[] listeners) { - if (context == Context.root()) { - return; // don't emit detach events for the default "no context" case - } + static void notifyUpdate(ContextListener[] listeners, Context before, Context after) { for (ContextListener l : listeners) { try { - l.onDetach(context); + l.onUpdate(before, after); } catch (Throwable ignore) { } } } - static void notifyCapture(Context context, ContextListener[] listeners) { + static void notifyCapture(ContextListener[] listeners, Context context) { // only called for non-empty continuations for (ContextListener l : listeners) { try { @@ -133,7 +106,7 @@ static void notifyCapture(Context context, ContextListener[] listeners) { } } - static void notifyRelease(Context context, ContextListener[] listeners) { + static void notifyRelease(ContextListener[] listeners, Context context) { // only called for non-empty continuations for (ContextListener l : listeners) { try { @@ -166,10 +139,8 @@ public final Context context() { public void close() { // check for out-of-order close to avoid corrupting the current state if (!closed && context == holder.current) { - ContextListener[] ls = INSTANCE.listeners; - notifyDetach(context, ls); holder.current = beforeAttach; - notifyAttach(beforeAttach, ls); + notifyUpdate(INSTANCE.listeners, context, beforeAttach); closed = true; } } @@ -232,7 +203,7 @@ private static final class ContextContinuationImpl implements ContextContinuatio ContextContinuationImpl(Context context) { this.context = context; - notifyCapture(context, INSTANCE.listeners); + notifyCapture(INSTANCE.listeners, context); } @Override @@ -270,7 +241,7 @@ public void release() { while (current == 0) { // no outstanding resumes and hold has been removed if (COUNT.compareAndSet(this, current, RELEASED)) { - notifyRelease(context, INSTANCE.listeners); + notifyRelease(INSTANCE.listeners, context); return; } current = count; @@ -280,7 +251,7 @@ public void release() { void releaseOnScopeClose() { if (COUNT.compareAndSet(this, 1, RELEASED)) { // fast path: only one resume of the continuation (no hold) - notifyRelease(context, INSTANCE.listeners); + notifyRelease(INSTANCE.listeners, context); } else if (COUNT.decrementAndGet(this) == 0) { // slow path: multiple resumes, all scopes now closed (no hold) release(); diff --git a/components/context/src/test/java/datadog/context/ContextContinuationTest.java b/components/context/src/test/java/datadog/context/ContextContinuationTest.java index 1b132717ac7..5f818332c89 100644 --- a/components/context/src/test/java/datadog/context/ContextContinuationTest.java +++ b/components/context/src/test/java/datadog/context/ContextContinuationTest.java @@ -20,7 +20,6 @@ @ParametersAreNonnullByDefault class ContextContinuationTest extends ContextTestBase { - private static final ContextKey CONTINUATION_KEY = ContextKey.named("continuation-key"); @Test void testCaptureRootContextIsNoop() { @@ -36,7 +35,7 @@ void testCaptureRootContextIsNoop() { @Test void testCaptureStoresContext() { - Context context = root().with(CONTINUATION_KEY, "captured"); + Context context = root().with(TEST_KEY, "captured"); try (ContextScope scope = context.attach()) { ContextContinuation continuation = context.capture(); assertEquals(context, continuation.context()); @@ -48,19 +47,19 @@ void testCaptureStoresContext() { void testCaptureFiresOnCaptureEvent() { TrackingListener listener = trackingListener(); ContextManager.register(listener); - Context context = root().with(CONTINUATION_KEY, "value"); + Context context = root().with(TEST_KEY, "value"); try (ContextScope scope = context.attach()) { ContextContinuation continuation = context.capture(); // capture while active (recommended pattern) - listener.assertNewEvents("attach", "capture"); + listener.assertNewEvents("update:{root}->value", "capture:value"); continuation.release(); } - listener.assertNewEvents("release", "detach"); + listener.assertNewEvents("release:value", "update:value->{root}"); } @Test void testResumeAttachesContextAndRestoresPreviousOnClose() { - Context context = root().with(CONTINUATION_KEY, "value"); + Context context = root().with(TEST_KEY, "value"); ContextContinuation continuation; try (ContextScope scope = context.attach()) { continuation = context.capture(); // capture while active (recommended pattern) @@ -77,56 +76,59 @@ void testResumeAttachesContextAndRestoresPreviousOnClose() { void testResumeAndScopeCloseFiresLifecycleEvents() { TrackingListener listener = trackingListener(); ContextManager.register(listener); - Context context = root().with(CONTINUATION_KEY, "value"); + Context context = root().with(TEST_KEY, "value"); ContextContinuation continuation; try (ContextScope scope = context.attach()) { continuation = context.capture(); // capture while active } - listener.assertNewEvents("attach", "capture", "detach"); + listener.assertNewEvents("update:{root}->value", "capture:value", "update:value->{root}"); try (ContextScope scope = continuation.resume()) { - listener.assertNewEvents("attach"); + listener.assertNewEvents("update:{root}->value"); } - // release fires before detach (continuation is released first inside ContextScopeImpl.close) - listener.assertNewEvents("release", "detach"); + // release fires before update (continuation is released first inside ContextScopeImpl.close) + listener.assertNewEvents("release:value", "update:value->{root}"); } @Test void testHoldPreventsAutoReleaseOnScopeClose() { TrackingListener listener = trackingListener(); ContextManager.register(listener); - Context context = root().with(CONTINUATION_KEY, "value"); + Context context = root().with(TEST_KEY, "value"); ContextContinuation continuation; try (ContextScope scope = context.attach()) { continuation = context.capture(); // capture while active continuation.hold(); + listener.assertNewEvents("update:{root}->value", "capture:value"); } + listener.assertNewEvents("update:value->{root}"); try (ContextScope scope = continuation.resume()) { assertEquals(context, current()); + listener.assertNewEvents("update:{root}->value"); } assertEquals(root(), current()); - // release should not fire while hold is active - listener.assertNewEvents("attach", "capture", "detach", "attach", "detach"); + listener.assertNewEvents( + "update:value->{root}"); // release should not fire while hold is active continuation.release(); - listener.assertNewEvents("release"); + listener.assertNewEvents("release:value"); } @Test void testExplicitReleaseWithoutResumeFiresReleaseEvent() { TrackingListener listener = trackingListener(); ContextManager.register(listener); - Context context = root().with(CONTINUATION_KEY, "value"); + Context context = root().with(TEST_KEY, "value"); ContextContinuation continuation; try (ContextScope scope = context.attach()) { continuation = context.capture(); // capture while active } - listener.assertNewEvents("attach", "capture", "detach"); + listener.assertNewEvents("update:{root}->value", "capture:value", "update:value->{root}"); continuation.release(); - listener.assertNewEvents("release"); + listener.assertNewEvents("release:value"); } @Test void testResumeAfterReleaseIsNoop() { - Context context = root().with(CONTINUATION_KEY, "value"); + Context context = root().with(TEST_KEY, "value"); ContextContinuation continuation; try (ContextScope scope = context.attach()) { continuation = context.capture(); // capture while active @@ -141,7 +143,7 @@ void testResumeAfterReleaseIsNoop() { @Test void testResumeOnDifferentThread() { - Context context = root().with(CONTINUATION_KEY, "value"); + Context context = root().with(TEST_KEY, "value"); ContextContinuation continuation; try (ContextScope scope = context.attach()) { continuation = context.capture(); // capture while active (recommended pattern) @@ -174,7 +176,7 @@ public void onRelease(Context c) { events.add("release"); } }); - Context context = root().with(CONTINUATION_KEY, "value"); + Context context = root().with(TEST_KEY, "value"); ContextContinuation continuation; try (ContextScope scope = context.attach()) { continuation = context.capture(); // capture while active @@ -225,78 +227,79 @@ public void onRelease(Context c) { void testSameContextResumeReleasesImmediately() { TrackingListener listener = trackingListener(); ContextManager.register(listener); - Context context = root().with(CONTINUATION_KEY, "value"); + Context context = root().with(TEST_KEY, "value"); try (ContextScope outer = context.attach()) { // Context is already current; resume is a noop and continuation is released immediately ContextContinuation continuation = context.capture(); try (ContextScope noop = continuation.resume()) { assertEquals(context, current()); - listener.assertNewEvents("attach", "capture", "release"); // released synchronously + listener.assertNewEvents( + "update:{root}->value", "capture:value", "release:value"); // released synchronously } assertEquals(context, current()); // outer scope still holds context } - listener.assertNewEvents("detach"); + listener.assertNewEvents("update:value->{root}"); } @Test void testOutOfOrderScopeCloseReleasesImmediately() { // Recommended pattern: attach C, capture, close original scope - Context contextC = root().with(CONTINUATION_KEY, "C"); + Context contextC = root().with(TEST_KEY, "C"); ContextContinuation continuation; try (ContextScope scope = contextC.attach()) { continuation = contextC.capture(); } - TrackingListener listener = keyedTrackingListener(CONTINUATION_KEY); + TrackingListener listener = trackingListener(); ContextManager.register(listener); - Context contextD = root().with(CONTINUATION_KEY, "D"); + Context contextD = root().with(TEST_KEY, "D"); try (ContextScope scopeR = continuation.resume()) { assertEquals(contextC, current()); - try (ContextScope scopeD = contextD.attach()) { // attaching D fires detach:C, attach:D + try (ContextScope scopeD = contextD.attach()) { // attaching D fires update:C->D assertEquals(contextD, current()); // close the resume scope out-of-order while D is still nested on top; - // release fires immediately, but detach:C does not (C is not current) + // release fires immediately, but update:C->{root} does not (C is not current) scopeR.close(); - listener.assertNewEvents("attach:C", "detach:C", "attach:D", "release:C"); + listener.assertNewEvents("update:{root}->C", "update:C->D", "release:C"); assertEquals(contextD, current()); // D is still current } // scopeD closes here: unwind D normally, restores C - listener.assertNewEvents("detach:D", "attach:C"); + listener.assertNewEvents("update:D->C"); } // try-with-resources closes scopeR again; no second release, C unwinds to root assertEquals(root(), current()); - listener.assertNewEvents("detach:C"); + listener.assertNewEvents("update:C->{root}"); } @Test void testHoldWithOutOfOrderScopeCloseFiresReleaseOnExplicitRelease() { // Regression test: hold() + out-of-order close must not corrupt the count, // which would cause release() to silently no-op and lose the release event. - Context contextC = root().with(CONTINUATION_KEY, "C"); + Context contextC = root().with(TEST_KEY, "C"); ContextContinuation continuation; try (ContextScope scope = contextC.attach()) { continuation = contextC.capture(); continuation.hold(); } - TrackingListener listener = keyedTrackingListener(CONTINUATION_KEY); + TrackingListener listener = trackingListener(); ContextManager.register(listener); - Context contextD = root().with(CONTINUATION_KEY, "D"); + Context contextD = root().with(TEST_KEY, "D"); try (ContextScope scopeR = continuation.resume()) { assertEquals(contextC, current()); - try (ContextScope scopeD = contextD.attach()) { // detach:C, attach:D + try (ContextScope scopeD = contextD.attach()) { // attaching D fires update:C->D assertEquals(contextD, current()); scopeR.close(); // out-of-order close while D is still on top; hold prevents auto-release - listener.assertNewEvents("attach:C", "detach:C", "attach:D"); + listener.assertNewEvents("update:{root}->C", "update:C->D"); assertEquals(contextD, current()); } // scopeD closes here: unwind D, restores C - } // TWR closes scopeR again (now in-order); detach:C, no release yet (hold is active) + } // TWR closes scopeR again (now in-order); update:C->{root}, no release yet (hold is active) assertEquals(root(), current()); - listener.assertNewEvents("detach:D", "attach:C", "detach:C"); + listener.assertNewEvents("update:D->C", "update:C->{root}"); continuation.release(); // explicit release must fire release:C listener.assertNewEvents("release:C"); @@ -307,7 +310,7 @@ void testMultipleHoldCallsAreIdempotent() { // Calling hold() more than once should not require more than one explicit release(). TrackingListener listener = trackingListener(); ContextManager.register(listener); - Context context = root().with(CONTINUATION_KEY, "value"); + Context context = root().with(TEST_KEY, "value"); ContextContinuation continuation; try (ContextScope scope = context.attach()) { continuation = context.capture(); @@ -316,7 +319,8 @@ void testMultipleHoldCallsAreIdempotent() { } // One explicit release() is enough — no extra releases needed for the second hold(). continuation.release(); - listener.assertNewEvents("attach", "capture", "detach", "release"); + listener.assertNewEvents( + "update:{root}->value", "capture:value", "update:value->{root}", "release:value"); continuation.release(); // still idempotent after the final release listener.assertNoNewEvents(); } @@ -326,13 +330,14 @@ void testHoldAfterReleaseIsIgnored() { // hold() on an already-released continuation must not resurrect it. TrackingListener listener = trackingListener(); ContextManager.register(listener); - Context context = root().with(CONTINUATION_KEY, "value"); + Context context = root().with(TEST_KEY, "value"); ContextContinuation continuation; try (ContextScope scope = context.attach()) { continuation = context.capture(); } continuation.release(); - listener.assertNewEvents("attach", "capture", "detach", "release"); + listener.assertNewEvents( + "update:{root}->value", "capture:value", "update:value->{root}", "release:value"); continuation.hold(); // must be silently ignored // resume() after release is already a noop, even with the spurious hold() try (ContextScope scope = continuation.resume()) { @@ -346,14 +351,15 @@ void testHoldAfterReleaseIsIgnored() { void testHoldAllowsMultipleReleaseCalls() { TrackingListener listener = trackingListener(); ContextManager.register(listener); - Context context = root().with(CONTINUATION_KEY, "value"); + Context context = root().with(TEST_KEY, "value"); ContextContinuation continuation; try (ContextScope scope = context.attach()) { continuation = context.capture(); // capture while active continuation.hold(); } continuation.release(); - listener.assertNewEvents("attach", "capture", "detach", "release"); + listener.assertNewEvents( + "update:{root}->value", "capture:value", "update:value->{root}", "release:value"); continuation.release(); // second release is a no-op listener.assertNoNewEvents(); } diff --git a/components/context/src/test/java/datadog/context/ContextListenerEventTest.java b/components/context/src/test/java/datadog/context/ContextListenerEventTest.java index 4c26d230db5..db71cc8438f 100644 --- a/components/context/src/test/java/datadog/context/ContextListenerEventTest.java +++ b/components/context/src/test/java/datadog/context/ContextListenerEventTest.java @@ -2,7 +2,6 @@ import static datadog.context.Context.current; 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.assertSame; @@ -11,57 +10,57 @@ class ContextListenerEventTest extends ContextTestBase { @Test void testListenersNotifiedOnAttachAndDetach() { - TrackingListener listener = keyedTrackingListener(STRING_KEY); + TrackingListener listener = trackingListener(); ContextManager.register(listener); - Context context = root().with(STRING_KEY, "value"); + Context context = root().with(TEST_KEY, "value"); try (ContextScope scope = context.attach()) { - listener.assertNewEvents("attach:value"); + listener.assertNewEvents("update:{root}->value"); } - listener.assertNewEvents("detach:value"); + listener.assertNewEvents("update:value->{root}"); } @Test - void testListenersNotNotifiedForRootContext() { + void testListenersNotNotifiedForSameContextAttachOrSwap() { TrackingListener listener = trackingListener(); ContextManager.register(listener); root().attach(); // current is already root, no events - listener.assertNoEvents(); // root attach should not trigger listeners + listener.assertNoEvents(); root().swap(); // current is already root, no events - listener.assertNoEvents(); // root swap should not trigger listeners - Context context = root().with(STRING_KEY, "value"); + listener.assertNoEvents(); + Context context = root().with(TEST_KEY, "value"); try (ContextScope scope = context.attach()) { - listener.assertNewEvents("attach"); // attach:non-root only + listener.assertNewEvents("update:{root}->value"); } - listener.assertNewEvents("detach"); // detach:non-root but not attach:root + listener.assertNewEvents("update:value->{root}"); } @Test void testListenersNotNotifiedOnSameContextAttach() { TrackingListener listener = trackingListener(); ContextManager.register(listener); - Context context = root().with(STRING_KEY, "same"); + Context context = root().with(TEST_KEY, "same"); try (ContextScope outer = context.attach()) { - listener.assertNewEvents("attach"); + listener.assertNewEvents("update:{root}->same"); try (ContextScope noop = context.attach()) { assertEquals(context, current()); listener.assertNoNewEvents(); // no new events on same-context attach } listener.assertNoNewEvents(); // noop close fires no events either } - listener.assertNewEvents("detach"); + listener.assertNewEvents("update:same->{root}"); } @Test void testListenersNotNotifiedOnSameContextSwap() { TrackingListener listener = trackingListener(); ContextManager.register(listener); - Context context = root().with(STRING_KEY, "same"); + Context context = root().with(TEST_KEY, "same"); context.swap(); - listener.assertNewEvents("attach"); + listener.assertNewEvents("update:{root}->same"); context.swap(); // same context again, no events listener.assertNoNewEvents(); root().swap(); - listener.assertNewEvents("detach"); + listener.assertNewEvents("update:same->{root}"); } @Test @@ -69,8 +68,8 @@ void testDuplicateListenerIgnored() { TrackingListener listener = trackingListener(); ContextManager.register(listener); ContextManager.register(listener); // should be ignored - try (ContextScope scope = root().with(STRING_KEY, "value").attach()) {} - listener.assertEvents("attach", "detach"); + try (ContextScope scope = root().with(TEST_KEY, "value").attach()) {} + listener.assertEvents("update:{root}->value", "update:value->{root}"); } @Test @@ -79,21 +78,21 @@ void testMultipleListenersAllNotified() { TrackingListener listener2 = trackingListener(); ContextManager.register(listener1); ContextManager.register(listener2); - try (ContextScope scope = root().with(STRING_KEY, "value").attach()) {} - listener1.assertEvents("attach", "detach"); - listener2.assertEvents("attach", "detach"); + try (ContextScope scope = root().with(TEST_KEY, "value").attach()) {} + listener1.assertEvents("update:{root}->value", "update:value->{root}"); + listener2.assertEvents("update:{root}->value", "update:value->{root}"); } @Test void testSwapNotifiesListeners() { - TrackingListener listener = keyedTrackingListener(STRING_KEY); + TrackingListener listener = trackingListener(); ContextManager.register(listener); - Context context = root().with(STRING_KEY, "value"); + Context context = root().with(TEST_KEY, "value"); Context previous = context.swap(); assertSame(root(), previous); - listener.assertNewEvents("attach:value"); + listener.assertNewEvents("update:{root}->value"); previous = root().swap(); assertSame(context, previous); - listener.assertNewEvents("detach:value"); + listener.assertNewEvents("update:value->{root}"); } } diff --git a/components/context/src/test/java/datadog/context/ContextListenerExceptionTest.java b/components/context/src/test/java/datadog/context/ContextListenerExceptionTest.java index 972f3733ddc..160873e7103 100644 --- a/components/context/src/test/java/datadog/context/ContextListenerExceptionTest.java +++ b/components/context/src/test/java/datadog/context/ContextListenerExceptionTest.java @@ -2,7 +2,6 @@ import static datadog.context.Context.current; import static datadog.context.Context.root; -import static datadog.context.ContextTest.STRING_KEY; import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; @@ -17,16 +16,11 @@ void testListenerExceptionSwallowed() { ContextManager.register( new ContextListener() { @Override - public void onAttach(Context c) { - throw new RuntimeException("listener failure"); - } - - @Override - public void onDetach(Context c) { + public void onUpdate(Context before, Context after) { throw new RuntimeException("listener failure"); } }); - Context context = root().with(STRING_KEY, "value"); + Context context = root().with(TEST_KEY, "value"); assertDoesNotThrow( () -> { try (ContextScope scope = context.attach()) { @@ -44,7 +38,7 @@ public void onCapture(Context c) { throw new RuntimeException("listener failure on capture"); } }); - Context context = root().with(STRING_KEY, "value"); + Context context = root().with(TEST_KEY, "value"); try (ContextScope scope = context.attach()) { assertDoesNotThrow( () -> { @@ -65,7 +59,7 @@ public void onRelease(Context c) { throw new RuntimeException("listener failure on release"); } }); - Context context = root().with(STRING_KEY, "value"); + Context context = root().with(TEST_KEY, "value"); try (ContextScope scope = context.attach()) { ContextContinuation continuation = context.capture(); assertDoesNotThrow(continuation::release); diff --git a/components/context/src/test/java/datadog/context/ContextTestBase.java b/components/context/src/test/java/datadog/context/ContextTestBase.java index 9fd2a4452cd..ac4c76f2482 100644 --- a/components/context/src/test/java/datadog/context/ContextTestBase.java +++ b/components/context/src/test/java/datadog/context/ContextTestBase.java @@ -7,13 +7,14 @@ import java.util.ArrayList; import java.util.List; -import javax.annotation.Nullable; import javax.annotation.ParametersAreNonnullByDefault; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; @ParametersAreNonnullByDefault abstract class ContextTestBase { + static final ContextKey TEST_KEY = ContextKey.named("test-key"); + @BeforeEach void verifyNoContextBefore() { assertEquals(root(), current()); @@ -26,51 +27,35 @@ void verifyNoContextAfter() { } static TrackingListener trackingListener() { - return new TrackingListener(null); - } - - static TrackingListener keyedTrackingListener(ContextKey key) { - return new TrackingListener(key); + return new TrackingListener(); } /** - * A {@link ContextListener} that records the events it receives so tests can assert on them. - * - *

With a {@link ContextKey}, each event is suffixed with that key's context value; otherwise - * only the event name is recorded (e.g. {@code "attach"}). + * A {@link ContextListener} that records events suffixed with {@link #TEST_KEY} context values. + * Uses {@code {root}} when the key is absent from the context. */ static final class TrackingListener implements ContextListener { - private final List events; - @Nullable private final ContextKey key; + private final List events = new ArrayList<>(); private int checkpoint; - private TrackingListener(@Nullable ContextKey key) { - this.events = new ArrayList<>(); - this.key = key; - } - - @Override - public void onAttach(Context context) { - record("attach", context); - } - @Override - public void onDetach(Context context) { - record("detach", context); + public void onUpdate(Context before, Context after) { + this.events.add("update:" + label(before) + "->" + label(after)); } @Override public void onCapture(Context context) { - record("capture", context); + this.events.add("capture:" + label(context)); } @Override public void onRelease(Context context) { - record("release", context); + this.events.add("release:" + label(context)); } - private void record(String event, Context context) { - this.events.add(this.key == null ? event : event + ":" + context.get(this.key)); + private static String label(Context context) { + String value = context.get(TEST_KEY); + return value != null ? value : "{root}"; } /** Asserts the full sequence of recorded events equals {@code expected}. */