diff --git a/.agents/skills/apm-integrations/references/context-tracking.md b/.agents/skills/apm-integrations/references/context-tracking.md index 744acc757b5..43918ba2a66 100644 --- a/.agents/skills/apm-integrations/references/context-tracking.md +++ b/.agents/skills/apm-integrations/references/context-tracking.md @@ -31,6 +31,12 @@ For each boundary-crossing type in the library, the instrumentation needs to ide - **The tracing wrapper** — a new class implementing the callback interface that reattaches the captured context before delegating (e.g. `TracingObserver` implements `Observer`). - **The wrapper's methods to guard** — the callback methods that must reattach context before running user code (e.g. `onNext`, `onError`, `onComplete` for RxJava observers). +## Capture controls and context switches + +Follow the [capture controls and context switches](../../../../docs/how_instrumentations_work.md#capture-controls-and-context-switches) +contract when choosing how to cross a boundary. For coroutine/fiber switching, capture the original +context before checkpointing it when trace retention is needed; use the checkpoint for swaps. + ## Tests for context-tracking instrumentations Context-tracking tests assert that **a span created by user code inside the wrapped callback becomes a child of the parent span active when the boundary was created**: @@ -45,6 +51,10 @@ runUnderTrace("parent") { Do NOT assert span kinds, operation names, span tags, or span types on the *target* methods — there are no spans on those methods. +Also verify that disabled capture prevents implicit propagation and that temporary suppression +restores the previous flag. For coroutine/fiber switching, check flag preservation across suspension +and resumption and restoration of the displaced execution unit's flag. + ## Choosing between method overloads When a boundary type exposes multiple overloads of the subscribe / invoke method — e.g. `Flowable.subscribe(Subscriber)` vs `Flowable.subscribe(FlowableSubscriber)`, or `Mono.subscribe(Subscriber)` vs `Mono.subscribe(CoreSubscriber)` — hook the **most specific framework-internal interface**, not the public wrapper. diff --git a/.agents/skills/fix-continuation-leakage/SKILL.md b/.agents/skills/fix-continuation-leakage/SKILL.md index 1d0b6d621c6..2e71be2a371 100644 --- a/.agents/skills/fix-continuation-leakage/SKILL.md +++ b/.agents/skills/fix-continuation-leakage/SKILL.md @@ -60,8 +60,10 @@ set -o pipefail preserving capture for the task actually submitted. - For intentionally delayed work, wait for its documented terminal event rather than suppressing propagation. - - For a context swap, verify both restoration and resource cleanup. Restore or close the - returned ownership object in `finally`; do not ignore every swap. + - For a context swap, restore the returned checkpoint by swapping it back in `finally`, or at + the matching suspension/exit hook. Release separately owned continuations when their owner + ends; checkpoints do not keep a trace open. See + [capture controls and context switches](../../../docs/how_instrumentations_work.md#capture-controls-and-context-switches). 5. Prefer a test-lifecycle fix when production behavior is correct. Otherwise fix ownership where it breaks, with one owner and `try/finally` cleanup across every exit. 6. Rerun the failing test, then its module. Validate the leaked record and root-trace publication diff --git a/components/context/src/main/java/datadog/context/Context.java b/components/context/src/main/java/datadog/context/Context.java index 34a07524de8..3b0bdc81a2c 100644 --- a/components/context/src/main/java/datadog/context/Context.java +++ b/components/context/src/main/java/datadog/context/Context.java @@ -35,7 +35,7 @@ * * @see ContextKey */ -public interface Context { +public interface Context extends ContextHandle { /** * Returns the root context. * @@ -64,12 +64,16 @@ default ContextScope attach() { } /** - * Swaps this context with the one attached to current execution unit. + * Saves this context and the current execution unit's flags for a future {@link #swap}. * - * @return the previously attached context; {@link #root()} if there was none. + *

This does not change the current context or execution-unit flags, or keep a trace open. + * Capture the original context first if a continuation is needed. Handles returned by swap are + * already checkpointed. + * + * @return a handle to this context and any saved flags for a future swap. */ - default Context swap() { - return manager().swap(this); + default ContextHandle checkpoint() { + return manager().checkpoint(this); } /** @@ -104,21 +108,38 @@ static Context detachFrom(Object carrier) { /** * Captures this (attached) context so it can be resumed in another execution unit. * - * @return continuation capturing this context. + * @return continuation capturing this context; no-op continuation if capture is disabled. */ default ContextContinuation capture() { return manager().capture(this); } /** - * Gets the value stored in this context under the given key. + * Returns whether the current execution unit can capture contexts. Capture is enabled by default. * - * @param the type of the value. - * @param key the key used to store the value. - * @return the value stored under the key; {@code null} if there is none. + *

Explicit attach and swap operations still work when capture is disabled. + * + * @return {@code true} if capture is enabled, {@code false} otherwise. + */ + static boolean isCaptureEnabled() { + return manager().isCaptureEnabled(); + } + + /** + * Enables or disables context capture for the current execution unit. + * + *

Attaching a different context preserves the capture flag; closing its scope restores the + * flag from before the attach. Reattaching the current context does not restore the flag on + * close. When suppressing capture outside an attached scope, save the previous flag and restore + * it in {@code finally}. + * + *

See {@link #swap()} for capture flag behavior across swaps. + * + * @param enabled {@code true} to allow capture, {@code false} to disable it. */ - @Nullable - T get(ContextKey key); + static void setCaptureEnabled(boolean enabled) { + manager().setCaptureEnabled(enabled); + } /** * Creates a copy of this context with the given key-value set. diff --git a/components/context/src/main/java/datadog/context/ContextHandle.java b/components/context/src/main/java/datadog/context/ContextHandle.java new file mode 100644 index 00000000000..dbbc7b093a6 --- /dev/null +++ b/components/context/src/main/java/datadog/context/ContextHandle.java @@ -0,0 +1,36 @@ +package datadog.context; + +import static datadog.context.ContextProviders.manager; + +import javax.annotation.Nullable; + +/** + * Read-only access to a context that can be swapped into the current execution unit. + * + *

A {@link Context} is its own handle. Other handles, like those returned by {@link #swap()} or + * {@link Context#checkpoint()}, save execution-unit flags for restoration and may wrap the context. + */ +public interface ContextHandle { + /** + * Swaps this context with the one attached to the current execution unit. + * + *

Swapping to a different context resets the current execution unit's flags unless this handle + * is a checkpoint, in which case the saved flags are restored. Swapping to the current context + * leaves the current execution unit's flags unchanged. + * + * @return a handle to the previously attached context and any saved flags. + */ + default ContextHandle swap() { + return manager().swap(this); + } + + /** + * Gets the value stored in this context under the given key. + * + * @param the type of the value. + * @param key the key used to store the value. + * @return the value stored under the key; {@code null} if there is none. + */ + @Nullable + T get(ContextKey key); +} diff --git a/components/context/src/main/java/datadog/context/ContextManager.java b/components/context/src/main/java/datadog/context/ContextManager.java index 6828352d1d4..6bd2ff94f9d 100644 --- a/components/context/src/main/java/datadog/context/ContextManager.java +++ b/components/context/src/main/java/datadog/context/ContextManager.java @@ -20,20 +20,59 @@ public interface ContextManager { ContextScope attach(Context context); /** - * Swaps the given context with the one attached to current execution unit. + * Swaps the given context with the one attached to the current execution unit. * - * @param context the context to swap. - * @return the previously attached context; {@link Context#root()} if there was none. + *

Swapping to a different context resets the current execution unit's flags unless the handle + * is a checkpoint, in which case the saved flags are restored. Swapping to the current context + * leaves the current execution unit's flags unchanged. + * + * @param handle a context or a handle returned by this manager. + * @return a handle to the previously attached context and any saved flags. */ - Context swap(Context context); + ContextHandle swap(ContextHandle handle); + + /** + * Saves the given context and the current execution unit's flags for a future {@link #swap}. + * + *

This does not change the current context or execution-unit flags, or keep a trace open. + * Capture the original context first if a continuation is needed. Handles returned by swap are + * already checkpointed. + * + * @param context the context to checkpoint. + * @return a handle to the given context and any saved flags for a future swap. + */ + ContextHandle checkpoint(Context context); /** * Captures the given (attached) context so it can be resumed in another execution unit. * - * @return continuation capturing the context. + * @return continuation capturing the context; no-op continuation if capture is disabled. */ ContextContinuation capture(Context context); + /** + * Returns whether the current execution unit can capture contexts. Capture is enabled by default. + * + *

Explicit attach and swap operations still work when capture is disabled. + * + * @return {@code true} if capture is enabled, {@code false} otherwise. + */ + boolean isCaptureEnabled(); + + /** + * Enables or disables context capture for the current execution unit. + * + *

Attaching a different context preserves the capture flag; closing its scope restores the + * flag from before the attach. Reattaching the current context does not restore the flag on + * close. When suppressing capture outside an attached scope, save the previous flag and restore + * it in {@code finally}. + * + *

See {@link #swap(ContextHandle)} for capture flag behavior across swaps. + * + * @param enabled {@code true} to allow capture, {@code false} to disable it. + */ + void setCaptureEnabled(boolean enabled); + /** * Registers the given listener to receive context events. * diff --git a/components/context/src/main/java/datadog/context/FlaggedContext.java b/components/context/src/main/java/datadog/context/FlaggedContext.java new file mode 100644 index 00000000000..2df11b7cd36 --- /dev/null +++ b/components/context/src/main/java/datadog/context/FlaggedContext.java @@ -0,0 +1,20 @@ +package datadog.context; + +import javax.annotation.Nullable; + +/** Saves non-zero execution-unit flags for a later swap to a different underlying context. */ +final class FlaggedContext implements ContextHandle { + final Context context; + final byte flags; + + FlaggedContext(Context context, byte flags) { + this.context = context; + this.flags = flags; + } + + @Override + @Nullable + public T get(ContextKey key) { + return context.get(key); + } +} diff --git a/components/context/src/main/java/datadog/context/TestContextManager.java b/components/context/src/main/java/datadog/context/TestContextManager.java index a0b1302653c..eec0ca00c54 100644 --- a/components/context/src/main/java/datadog/context/TestContextManager.java +++ b/components/context/src/main/java/datadog/context/TestContextManager.java @@ -23,15 +23,30 @@ public ContextScope attach(Context context) { } @Override - public Context swap(Context context) { + public ContextHandle swap(ContextHandle context) { return delegate().swap(context); } + @Override + public ContextHandle checkpoint(Context context) { + return delegate().checkpoint(context); + } + @Override public ContextContinuation capture(Context context) { return delegate().capture(context); } + @Override + public boolean isCaptureEnabled() { + return delegate().isCaptureEnabled(); + } + + @Override + public void setCaptureEnabled(boolean enabled) { + delegate().setCaptureEnabled(enabled); + } + @Override public void addListener(ContextListener listener) { delegate().addListener(listener); diff --git a/components/context/src/main/java/datadog/context/ThreadLocalContextManager.java b/components/context/src/main/java/datadog/context/ThreadLocalContextManager.java index c7580fc8e0f..18ce30f8414 100644 --- a/components/context/src/main/java/datadog/context/ThreadLocalContextManager.java +++ b/components/context/src/main/java/datadog/context/ThreadLocalContextManager.java @@ -14,6 +14,8 @@ final class ThreadLocalContextManager implements ContextManager { private static final NoopContextContinuation ROOT_CONTINUATION = new NoopContextContinuation(Context.root()); + private static final byte SUPPRESS_CAPTURE = 1; + private final Object listenersWriteLock = new Object(); volatile ContextListener[] listeners = {}; @@ -41,29 +43,67 @@ ContextScope doAttach(Context context, @Nullable ContextContinuationImpl continu } holder.current = context; + byte flags = holder.flags; notifyUpdate(listeners, beforeAttach, context); return continuation == null - ? new ContextScopeImpl(context, beforeAttach) - : new ResumedScopeImpl(context, beforeAttach, continuation); + ? new ContextScopeImpl(context, beforeAttach, flags) + : new ResumedScopeImpl(context, beforeAttach, flags, continuation); } @Override - public Context swap(Context context) { - ContextHolder holder = CONTEXT_HOLDER.get(); + public ContextHandle swap(ContextHandle handle) { + byte seedFlags = 0; + Context context; + if (handle instanceof Context) { + context = (Context) handle; + } else { + seedFlags = ((FlaggedContext) handle).flags; + context = ((FlaggedContext) handle).context; + } + ContextHolder holder = CONTEXT_HOLDER.get(); Context beforeSwap = holder.current; - if (context == beforeSwap) { - return beforeSwap; + byte flags = holder.flags; + + // Same-context skips resetting or restoring flags. + if (context != beforeSwap) { + holder.flags = seedFlags; + holder.current = context; + notifyUpdate(listeners, beforeSwap, context); } - holder.current = context; - notifyUpdate(listeners, beforeSwap, context); - return beforeSwap; + return checkpoint(beforeSwap, flags); + } + + @Override + public ContextHandle checkpoint(Context context) { + return checkpoint(context, CONTEXT_HOLDER.get().flags); + } + + private static ContextHandle checkpoint(Context context, byte flags) { + return flags == 0 ? context : new FlaggedContext(context, flags); } @Override public ContextContinuation capture(Context context) { - return context == Context.root() ? ROOT_CONTINUATION : new ContextContinuationImpl(context); + return context == Context.root() || !isCaptureEnabled() + ? ROOT_CONTINUATION + : new ContextContinuationImpl(context); + } + + @Override + public boolean isCaptureEnabled() { + return (CONTEXT_HOLDER.get().flags & SUPPRESS_CAPTURE) == 0; + } + + @Override + public void setCaptureEnabled(boolean enabled) { + ContextHolder holder = CONTEXT_HOLDER.get(); + if (enabled) { + holder.flags &= ~SUPPRESS_CAPTURE; + } else { + holder.flags |= SUPPRESS_CAPTURE; + } } @Override @@ -119,11 +159,14 @@ static void notifyRelease(ContextListener[] listeners, Context context) { private static class ContextScopeImpl implements ContextScope { private final Context context; private final Context beforeAttach; + private final byte flags; + private boolean closed; - ContextScopeImpl(Context context, Context beforeAttach) { + ContextScopeImpl(Context context, Context beforeAttach, byte flags) { this.context = context; this.beforeAttach = beforeAttach; + this.flags = flags; } @Override @@ -138,6 +181,7 @@ public void close() { // check for out-of-order close to avoid corrupting the current state if (context == holder.current) { holder.current = beforeAttach; + holder.flags = flags; notifyUpdate(INSTANCE.listeners, context, beforeAttach); closed = true; } @@ -149,8 +193,11 @@ private static final class ResumedScopeImpl extends ContextScopeImpl { @Nullable private ContextContinuationImpl continuation; ResumedScopeImpl( - Context context, Context beforeAttach, @Nullable ContextContinuationImpl continuation) { - super(context, beforeAttach); + Context context, + Context beforeAttach, + byte flags, + @Nullable ContextContinuationImpl continuation) { + super(context, beforeAttach, flags); this.continuation = continuation; } @@ -259,5 +306,6 @@ public void close() {} private static final class ContextHolder { Context current = Context.root(); + byte flags; // execution-unit flags } } diff --git a/components/context/src/test/java/datadog/context/ContextCaptureTest.java b/components/context/src/test/java/datadog/context/ContextCaptureTest.java new file mode 100644 index 00000000000..f06e55d8f8e --- /dev/null +++ b/components/context/src/test/java/datadog/context/ContextCaptureTest.java @@ -0,0 +1,333 @@ +package datadog.context; + +import static datadog.context.Context.current; +import static datadog.context.Context.isCaptureEnabled; +import static datadog.context.Context.root; +import static datadog.context.Context.setCaptureEnabled; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.locks.LockSupport; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.condition.EnabledForJreRange; +import org.junit.jupiter.api.condition.JRE; +import org.tabletest.junit.TableTest; + +class ContextCaptureTest extends ContextTestBase { + @AfterEach + void resetCapture() { + setCaptureEnabled(true); + } + + @Test + void enabledByDefault() { + assertTrue(isCaptureEnabled()); + } + + @TableTest({ + "Scenario | enabled | captureCurrent", + "capture active enabled | true | true ", + "capture other enabled | true | false ", + "capture active suppressed | false | true ", + "capture other suppressed | false | false " + }) + void captureHonorsExecutionUnitFlag(boolean enabled, boolean captureCurrent) { + Context active = root().with(TEST_KEY, "active"); + Context captured = captureCurrent ? active : root().with(TEST_KEY, "other"); + TrackingListener listener = trackingListener(); + ContextManager.register(listener); + ContextContinuation noop = root().capture(); + try (ContextScope ignored = active.attach()) { + listener.assertNewEvents("update:{root}->active"); + setCaptureEnabled(enabled); + ContextContinuation continuation = captured.capture(); + if (enabled) { + assertSame(captured, continuation.context()); + listener.assertNewEvents("capture:" + (captureCurrent ? "active" : "other")); + continuation.release(); + listener.assertNewEvents("release:" + (captureCurrent ? "active" : "other")); + } else { + assertSame(noop, continuation); + assertSame(noop, continuation.hold()); + try (ContextScope scope = continuation.resume()) { + assertSame(active, current()); + } + continuation.release(); + listener.assertNoNewEvents(); + } + } + assertTrue(isCaptureEnabled()); + } + + @TableTest({ + "Scenario | enabled | checkpointCurrent", + "current enabled | true | true ", + "other enabled | true | false ", + "current suppressed | false | true ", + "other suppressed | false | false " + }) + void checkpointSavesFlagsWithoutChangingCurrentState(boolean enabled, boolean checkpointCurrent) { + Context active = root().with(TEST_KEY, "active"); + Context saved = checkpointCurrent ? active : root().with(TEST_KEY, "other"); + TrackingListener listener = trackingListener(); + ContextManager.register(listener); + try (ContextScope ignored = active.attach()) { + listener.assertNewEvents("update:{root}->active"); + setCaptureEnabled(enabled); + ContextHandle checkpoint = saved.checkpoint(); + assertSame(active, current()); + assertEquals(enabled, isCaptureEnabled()); + assertSavedFlags(saved, enabled, checkpoint); + listener.assertNoNewEvents(); + + root().swap(); + setCaptureEnabled(!enabled); + checkpoint.swap(); + assertSame(saved, current()); + assertEquals(enabled, isCaptureEnabled()); + } finally { + root().swap(); + } + } + + @TableTest({ + "Scenario | before | during", + "enabled throughout | true | true ", + "restore enabled | true | false ", + "restore suppressed | false | true ", + "suppressed throughout | false | false " + }) + void differentContextResetsCaptureAndRestoresSavedFlags(boolean before, boolean during) { + Context first = root().with(TEST_KEY, "first"); + Context second = root().with(TEST_KEY, "second"); + try (ContextScope ignored = first.attach()) { + setCaptureEnabled(before); + ContextHandle saved = second.swap(); + assertSame(second, current()); + assertTrue(isCaptureEnabled()); + assertSavedFlags(first, before, saved); + setCaptureEnabled(during); + ContextHandle displaced = saved.swap(); + assertSame(first, current()); + assertEquals(before, isCaptureEnabled()); + assertSavedFlags(second, during, displaced); + } + } + + @TableTest({ + "Scenario | before | during", + "enabled throughout | true | true ", + "retain suppression | true | false ", + "retain enabled | false | true ", + "suppressed throughout | false | false " + }) + void sameContextCheckpointsWithoutResettingOrRestoringFlags(boolean before, boolean during) { + Context context = root().with(TEST_KEY, "checkpoint"); + TrackingListener listener = trackingListener(); + ContextManager.register(listener); + try (ContextScope ignored = context.attach()) { + listener.assertNewEvents("update:{root}->checkpoint"); + setCaptureEnabled(before); + ContextHandle checkpoint = current().swap(); + assertSame(context, current()); + assertEquals(before, isCaptureEnabled()); + assertSavedFlags(context, before, checkpoint); + setCaptureEnabled(during); + ContextHandle displaced = checkpoint.swap(); + assertSame(context, current()); + assertEquals(during, isCaptureEnabled()); + assertSavedFlags(context, during, displaced); + listener.assertNoNewEvents(); + // The same checkpoint restores its saved flags once the underlying context changes. + root().swap(); + checkpoint.swap(); + assertSame(context, current()); + assertEquals(before, isCaptureEnabled()); + listener.assertNewEvents("update:checkpoint->{root}", "update:{root}->checkpoint"); + } + } + + @Test + void suppressedRootCanBeCheckpointedAndRestored() { + setCaptureEnabled(false); + ContextHandle checkpoint = root().checkpoint(); + assertSame(root(), current()); + assertFalse(isCaptureEnabled()); + Context context = root().with(TEST_KEY, "active"); + context.swap(); + try { + assertTrue(isCaptureEnabled()); + checkpoint.swap(); + assertSame(root(), current()); + assertFalse(isCaptureEnabled()); + } finally { + root().swap(); + } + } + + @Test + void savedFlagsFollowContextAcrossThreads() { + Context context = root().with(TEST_KEY, "migrating"); + try (ContextScope ignored = context.attach()) { + setCaptureEnabled(false); + ContextHandle checkpoint = current().checkpoint(); + CompletableFuture.runAsync( + () -> { + assertSame(root(), current()); + assertTrue(isCaptureEnabled()); + ContextHandle previous = checkpoint.swap(); + try { + assertSame(context, current()); + assertFalse(isCaptureEnabled()); + } finally { + previous.swap(); + } + assertSame(root(), current()); + assertTrue(isCaptureEnabled()); + }) + .join(); + assertSame(context, current()); + assertFalse(isCaptureEnabled()); + } + } + + private static void assertSavedFlags(Context context, boolean enabled, ContextHandle saved) { + assertEquals(context.get(TEST_KEY), saved.get(TEST_KEY)); + assertNull(saved.get(ContextKey.named("absent"))); + if (enabled) { + assertSame(context, saved); + } else { + FlaggedContext flagged = assertInstanceOf(FlaggedContext.class, saved); + assertSame(context, flagged.context); + assertEquals(1, flagged.flags); + } + } + + @Test + void nestedScopesRestoreFlagsAfterSwap() { + Context outer = root().with(TEST_KEY, "outer"); + Context inner = root().with(TEST_KEY, "inner"); + try (ContextScope outerScope = outer.attach()) { + setCaptureEnabled(false); + try (ContextScope innerScope = inner.attach()) { + assertFalse(isCaptureEnabled()); + setCaptureEnabled(true); + ContextHandle previous = root().swap(); + assertTrue(isCaptureEnabled()); + previous.swap(); + } + assertSame(outer, current()); + assertFalse(isCaptureEnabled()); + } + assertTrue(isCaptureEnabled()); + } + + @Test + void sameContextAttachDoesNotRestoreFlag() { + Context context = root().with(TEST_KEY, "value"); + try (ContextScope outer = context.attach()) { + try (ContextScope noop = context.attach()) { + setCaptureEnabled(false); + } + assertFalse(isCaptureEnabled()); + } + assertTrue(isCaptureEnabled()); + } + + @Test + void existingContinuationCanResumeWhileSuppressed() { + Context context = root().with(TEST_KEY, "captured"); + ContextContinuation continuation = context.capture(); + setCaptureEnabled(false); + try (ContextScope ignored = continuation.resume()) { + assertSame(context, current()); + assertFalse(isCaptureEnabled()); + } + assertSame(root(), current()); + assertFalse(isCaptureEnabled()); + } + + @TableTest({ + "Scenario | enabled", + "restore enabled | true ", + "restore suppressed | false " + }) + void resumedScopeRestoresCaptureFlagAfterChange(boolean enabled) { + Context context = root().with(TEST_KEY, "captured"); + ContextContinuation continuation = context.capture(); + setCaptureEnabled(enabled); + try (ContextScope ignored = continuation.resume()) { + assertSame(context, current()); + assertEquals(enabled, isCaptureEnabled()); + setCaptureEnabled(!enabled); + } + assertSame(root(), current()); + assertEquals(enabled, isCaptureEnabled()); + } + + @Test + void outOfOrderClosePreservesCaptureFlagUntilScopesUnwind() { + Context outer = root().with(TEST_KEY, "outer"); + Context inner = root().with(TEST_KEY, "inner"); + try (ContextScope outerScope = outer.attach()) { + try (ContextScope innerScope = inner.attach()) { + setCaptureEnabled(false); + outerScope.close(); + assertSame(inner, current()); + assertFalse(isCaptureEnabled()); + } + assertSame(outer, current()); + assertTrue(isCaptureEnabled()); + } + assertSame(root(), current()); + assertTrue(isCaptureEnabled()); + } + + @Test + @EnabledForJreRange(min = JRE.JAVA_21) + void virtualThreadRetainsSuppressionAcrossParkAndSwap() throws Exception { + Context context = root().with(TEST_KEY, "virtual"); + CompletableFuture result = new CompletableFuture<>(); + try (ContextScope parentScope = context.attach()) { + setCaptureEnabled(false); + Runnable task = + () -> { + try { + assertTrue(isCaptureEnabled()); + try (ContextScope scope = context.attach()) { + setCaptureEnabled(false); + for (int i = 0; i < 5; i++) { + ContextHandle previous = root().swap(); + LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(1)); + previous.swap(); + assertSame(context, current()); + assertFalse(isCaptureEnabled()); + } + try (ContextScope inner = context.with(TEST_KEY, "nested").attach()) { + setCaptureEnabled(true); + } + assertFalse(isCaptureEnabled()); + } + assertTrue(isCaptureEnabled()); + result.complete(null); + } catch (Throwable failure) { + result.completeExceptionally(failure); + } + }; + Thread thread = + (Thread) Thread.class.getMethod("startVirtualThread", Runnable.class).invoke(null, task); + result.get(10, TimeUnit.SECONDS); + thread.join(TimeUnit.SECONDS.toMillis(10)); + assertFalse(thread.isAlive()); + assertFalse(isCaptureEnabled()); + } + assertTrue(isCaptureEnabled()); + } +} diff --git a/components/context/src/test/java/datadog/context/ContextListenerEventTest.java b/components/context/src/test/java/datadog/context/ContextListenerEventTest.java index db71cc8438f..29e13e62add 100644 --- a/components/context/src/test/java/datadog/context/ContextListenerEventTest.java +++ b/components/context/src/test/java/datadog/context/ContextListenerEventTest.java @@ -88,7 +88,7 @@ void testSwapNotifiesListeners() { TrackingListener listener = trackingListener(); ContextManager.register(listener); Context context = root().with(TEST_KEY, "value"); - Context previous = context.swap(); + ContextHandle previous = context.swap(); assertSame(root(), previous); listener.assertNewEvents("update:{root}->value"); previous = root().swap(); diff --git a/components/context/src/test/java/datadog/context/ContextProvidersForkedTest.java b/components/context/src/test/java/datadog/context/ContextProvidersForkedTest.java index 73c55a28c6f..f06887ee2ea 100644 --- a/components/context/src/test/java/datadog/context/ContextProvidersForkedTest.java +++ b/components/context/src/test/java/datadog/context/ContextProvidersForkedTest.java @@ -3,6 +3,7 @@ import static datadog.context.Context.root; import static datadog.context.ContextTest.STRING_KEY; import static datadog.context.ContextTestBase.trackingListener; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotEquals; import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -59,6 +60,16 @@ void testCustomManager() { Context context = root().with(STRING_KEY, "value"); assertNotEquals(root(), context); + // should delegate flag changes and suppression to the default manager + ContextManager manager = ContextProviders.manager(); + assertTrue(manager.isCaptureEnabled()); + manager.setCaptureEnabled(false); + assertFalse(manager.isCaptureEnabled()); + assertNotEquals(context, context.checkpoint()); + assertSame(root().capture(), context.capture()); + manager.setCaptureEnabled(true); + assertTrue(manager.isCaptureEnabled()); + // should delegate to the default manager try (ContextScope scope = context.attach()) { assertSame(context, scope.context()); @@ -68,7 +79,7 @@ void testCustomManager() { cont.release(); } - Context swapped = context.swap(); + ContextHandle swapped = context.swap(); assertSame(root(), swapped); assertSame(context, Context.current()); assertSame(context, swapped.swap()); @@ -88,15 +99,28 @@ public ContextScope attach(@Nonnull Context context) { } @Override - public Context swap(@Nonnull Context context) { + public ContextHandle swap(@Nonnull ContextHandle context) { return root(); } + @Override + public ContextHandle checkpoint(@Nonnull Context context) { + return context; + } + @Override public ContextContinuation capture(@Nonnull Context context) { return new NoopContextContinuation(root()); } + @Override + public boolean isCaptureEnabled() { + return true; + } + + @Override + public void setCaptureEnabled(boolean enabled) {} + @Override public void addListener(@Nonnull ContextListener listener) {} }); @@ -114,6 +138,11 @@ public void addListener(@Nonnull ContextListener listener) {} } // NOOP manager, context will always be root + assertTrue(manager.isCaptureEnabled()); + manager.setCaptureEnabled(false); // custom manager ignores the capture flag + assertTrue(manager.isCaptureEnabled()); + assertSame(context, context.checkpoint()); + swapped = context.swap(); assertSame(root(), swapped); assertSame(root(), Context.current()); diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/api/Java8BytecodeBridge.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/api/Java8BytecodeBridge.java index d749dd4c372..3a8cf302ca2 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/api/Java8BytecodeBridge.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/api/Java8BytecodeBridge.java @@ -1,6 +1,7 @@ package datadog.trace.bootstrap.instrumentation.api; import datadog.context.Context; +import datadog.context.ContextHandle; import datadog.context.ContextScope; /** @@ -25,6 +26,20 @@ public static Context currentContext() { return Context.current(); } + /** + * @see Context#isCaptureEnabled() + */ + public static boolean isContextCaptureEnabled() { + return Context.isCaptureEnabled(); + } + + /** + * @see Context#setCaptureEnabled(boolean) + */ + public static void setContextCaptureEnabled(boolean enabled) { + Context.setCaptureEnabled(enabled); + } + /** * @see AgentSpan#current() */ @@ -33,9 +48,9 @@ public static AgentSpan currentSpan() { } /** - * @see AgentSpan#fromContext(Context) + * @see AgentSpan#fromContext(ContextHandle) */ - public static AgentSpan spanFromContext(Context context) { + public static AgentSpan spanFromContext(ContextHandle context) { return AgentSpan.fromContext(context); } diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/VirtualThreadState.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/VirtualThreadState.java index bec26bd9fcb..f4a530fb017 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/VirtualThreadState.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/VirtualThreadState.java @@ -4,6 +4,7 @@ import datadog.context.Context; import datadog.context.ContextContinuation; +import datadog.context.ContextHandle; import datadog.trace.api.Config; import datadog.trace.api.GenericClassValue; import datadog.trace.api.InstrumenterConfig; @@ -51,13 +52,13 @@ public static boolean shouldPropagateContext(Runnable task) { || Config.get().isCwsEnabled(); /** The virtual thread's saved context (scope stack snapshot). */ - private Context context; + private ContextHandle context; /** Prevents the enclosing context scope from completing before the virtual thread finishes. */ private final ContextContinuation continuation; /** The carrier thread's saved context, set between mount and unmount. */ - private Context previousContext; + private ContextHandle previousContext; public VirtualThreadState(Context context, ContextContinuation continuation) { this.context = context; diff --git a/dd-java-agent/instrumentation-testing/src/test/java/datadog/trace/agent/test/scopediag/ScopeDiagnosticsIntegrationTest.java b/dd-java-agent/instrumentation-testing/src/test/java/datadog/trace/agent/test/scopediag/ScopeDiagnosticsIntegrationTest.java index 23358c3d542..38acabb796c 100644 --- a/dd-java-agent/instrumentation-testing/src/test/java/datadog/trace/agent/test/scopediag/ScopeDiagnosticsIntegrationTest.java +++ b/dd-java-agent/instrumentation-testing/src/test/java/datadog/trace/agent/test/scopediag/ScopeDiagnosticsIntegrationTest.java @@ -7,8 +7,8 @@ import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertTrue; -import datadog.context.Context; import datadog.context.ContextContinuation; +import datadog.context.ContextHandle; import datadog.context.ContextScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.common.writer.ListWriter; @@ -318,8 +318,8 @@ void swappedContextDoesNotCreateCloseOwnedScope() { ScopeDiagnostics.startRecording(); AgentSpan span = tracer.startSpan("test", "op"); - Context previous = tracer.swap(span); - previous.swap(); + ContextHandle previous = tracer.swap(span); + tracer.swap(previous); span.finish(); ScopeDiagnosticsReport report = ScopeDiagnostics.report(); diff --git a/dd-java-agent/instrumentation/akka/akka-actor-2.5/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaActorCellInstrumentation.java b/dd-java-agent/instrumentation/akka/akka-actor-2.5/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaActorCellInstrumentation.java index 4343ae0bd47..d2859649d5f 100644 --- a/dd-java-agent/instrumentation/akka/akka-actor-2.5/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaActorCellInstrumentation.java +++ b/dd-java-agent/instrumentation/akka/akka-actor-2.5/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaActorCellInstrumentation.java @@ -9,7 +9,7 @@ import akka.dispatch.Envelope; import com.google.auto.service.AutoService; -import datadog.context.Context; +import datadog.context.ContextHandle; import datadog.context.ContextScope; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; @@ -57,7 +57,7 @@ public void methodAdvice(MethodTransformer transformer) { */ public static class InvokeAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static Context enter( + public static ContextHandle enter( @Advice.Argument(value = 0) Envelope envelope, @Advice.Local("taskScope") ContextScope taskScope) { @@ -78,7 +78,7 @@ public static Context enter( @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void exit( @Advice.Local("taskScope") ContextScope taskScope, - @Advice.Enter Context checkpointContext) { + @Advice.Enter ContextHandle checkpointContext) { if (checkpointContext == null) { // Clean up any leaking scopes from akka-streams/akka-http etc. diff --git a/dd-java-agent/instrumentation/akka/akka-actor-2.5/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaMailboxInstrumentation.java b/dd-java-agent/instrumentation/akka/akka-actor-2.5/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaMailboxInstrumentation.java index 96efa140128..f9d77d8d87d 100644 --- a/dd-java-agent/instrumentation/akka/akka-actor-2.5/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaMailboxInstrumentation.java +++ b/dd-java-agent/instrumentation/akka/akka-actor-2.5/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaMailboxInstrumentation.java @@ -7,7 +7,7 @@ import static net.bytebuddy.matcher.ElementMatchers.isMethod; import com.google.auto.service.AutoService; -import datadog.context.Context; +import datadog.context.ContextHandle; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.api.InstrumenterConfig; @@ -45,7 +45,7 @@ public void methodAdvice(MethodTransformer transformer) { */ public static final class SuppressMailboxRunAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static Context enter() { + public static ContextHandle enter() { if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) { // remember the currently active scope so we can roll back to this point checkpointActiveForRollback(); @@ -56,7 +56,7 @@ public static Context enter() { } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) - public static void exit(@Advice.Enter final Context checkpointContext) { + public static void exit(@Advice.Enter final ContextHandle checkpointContext) { if (checkpointContext == null) { // Clean up any leaking scopes from akka-streams/akka-http etc. rollbackActiveToCheckpoint(); diff --git a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogWrapperHelper.java b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogWrapperHelper.java index cf0a5d68258..717972ce4c9 100644 --- a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogWrapperHelper.java +++ b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogWrapperHelper.java @@ -1,11 +1,14 @@ package datadog.trace.instrumentation.akkahttp; +import static datadog.context.Context.isCaptureEnabled; +import static datadog.context.Context.setCaptureEnabled; import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; import static datadog.trace.instrumentation.akkahttp.AkkaHttpServerDecorator.DECORATE; import akka.http.scaladsl.model.HttpRequest; import akka.http.scaladsl.model.HttpResponse; import datadog.context.Context; +import datadog.context.ContextHandle; import datadog.context.ContextScope; import datadog.trace.api.InstrumenterConfig; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -38,14 +41,17 @@ private static Context startSpan(final HttpRequest request) { static final class SwappedContextScope implements ContextScope { private final Context context; - private final Context previousContext; + private final ContextHandle previousContext; private final Thread ownerThread; private boolean closed; SwappedContextScope(final Context context) { this.context = context; this.ownerThread = Thread.currentThread(); + boolean allowCapture = isCaptureEnabled(); this.previousContext = context.swap(); + // Request handling inherits the enclosing flow's capture flag. + setCaptureEnabled(allowCapture); } @Override diff --git a/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/main/java/datadog/trace/instrumentation/java/lang/jdk21/PollerInstrumentation.java b/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/main/java/datadog/trace/instrumentation/java/lang/jdk21/PollerInstrumentation.java index 2b4d878a8a6..d40cb724544 100644 --- a/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/main/java/datadog/trace/instrumentation/java/lang/jdk21/PollerInstrumentation.java +++ b/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/main/java/datadog/trace/instrumentation/java/lang/jdk21/PollerInstrumentation.java @@ -6,7 +6,7 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArguments; import com.google.auto.service.AutoService; -import datadog.context.Context; +import datadog.context.ContextHandle; import datadog.environment.JavaVirtualMachine; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; @@ -49,13 +49,13 @@ public void methodAdvice(MethodTransformer transformer) { public static final class StartAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static Context enter() { + public static ContextHandle enter() { // Swap clears the raw context even when the caller has reached the scope-depth limit. return rootContext().swap(); } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) - public static void exit(@Advice.Enter Context previous) { + public static void exit(@Advice.Enter ContextHandle previous) { if (previous != null) { previous.swap(); } diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/DatadogThreadContextElement.java b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/DatadogThreadContextElement.java index 30517323846..40abc0651f5 100644 --- a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/DatadogThreadContextElement.java +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/DatadogThreadContextElement.java @@ -2,6 +2,7 @@ import datadog.context.Context; import datadog.context.ContextContinuation; +import datadog.context.ContextHandle; import datadog.trace.api.GenericClassValue; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import java.util.concurrent.atomic.AtomicReferenceFieldUpdater; @@ -17,7 +18,7 @@ /** Manages the Datadog context for coroutines, switching contexts as coroutines switch threads. */ public final class DatadogThreadContextElement - implements ThreadContextElement, + implements ThreadContextElement, Function2 { private static final AtomicReferenceFieldUpdater CONTINUATION = @@ -85,7 +86,7 @@ public static CoroutineContext forFlowCollection( DatadogThreadContextElement caller = downstream.get(DATADOG_KEY); if (caller != null) { // Downstream reinstalls the caller's element without constructing a coroutine. - caller.context = Context.current(); + caller.context = Context.current().checkpoint(); } DatadogThreadContextElement element = upstream.get(DATADOG_KEY); if (element == null || element == caller) { @@ -113,7 +114,7 @@ public CoroutineContext invoke(CoroutineContext rebuilt, CoroutineContext.Elemen } @Nullable private AgentSpan parentLocalRootSpan; - private volatile Context context; + private volatile ContextHandle context; private volatile ContextContinuation continuation; /** The owning Job, null before initialization, or COMPLETED after releasing the owner. */ @@ -158,7 +159,7 @@ public static void captureDatadogContext(@Nonnull AbstractCoroutine coroutine } else if (datadog.owner != coroutine && datadog.isScopedBy(coroutine)) { // scoped coroutines (coroutineScope, withTimeout, or withContext before 1.6.1) share their // caller's element, so refresh it with the active context - datadog.context = Context.current(); + datadog.context = Context.current().checkpoint(); } } } @@ -182,7 +183,7 @@ private synchronized void captureInitialContext(boolean owned) { } } } - context = initial; + context = initial.checkpoint(); parentLocalRootSpan = null; } } @@ -209,8 +210,8 @@ public static void cancelDatadogContext(@Nonnull AbstractCoroutine coroutine) } @Override - public Context updateThreadContext(@Nonnull CoroutineContext coroutineContext) { - Context current = context; + public ContextHandle updateThreadContext(@Nonnull CoroutineContext coroutineContext) { + ContextHandle current = context; if (current == null) { if (owner == null) { Job job = coroutineContext.get(Job.Key); @@ -230,7 +231,7 @@ public Context updateThreadContext(@Nonnull CoroutineContext coroutineContext) { @Override public void restoreThreadContext( - @Nonnull CoroutineContext coroutineContext, Context originalContext) { + @Nonnull CoroutineContext coroutineContext, ContextHandle originalContext) { context = originalContext.swap(); } diff --git a/dd-java-agent/instrumentation/pekko/pekko-concurrent-1.0/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoActorCellInstrumentation.java b/dd-java-agent/instrumentation/pekko/pekko-concurrent-1.0/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoActorCellInstrumentation.java index 2778831de22..1b37a2c7748 100644 --- a/dd-java-agent/instrumentation/pekko/pekko-concurrent-1.0/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoActorCellInstrumentation.java +++ b/dd-java-agent/instrumentation/pekko/pekko-concurrent-1.0/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoActorCellInstrumentation.java @@ -8,7 +8,7 @@ import static net.bytebuddy.matcher.ElementMatchers.isMethod; import com.google.auto.service.AutoService; -import datadog.context.Context; +import datadog.context.ContextHandle; import datadog.context.ContextScope; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; @@ -57,7 +57,7 @@ public void methodAdvice(MethodTransformer transformer) { */ public static class InvokeAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static Context enter( + public static ContextHandle enter( @Advice.Argument(value = 0) Envelope envelope, @Advice.Local("taskScope") ContextScope taskScope) { @@ -78,7 +78,7 @@ public static Context enter( @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void exit( @Advice.Local("taskScope") ContextScope taskScope, - @Advice.Enter Context checkpointContext) { + @Advice.Enter ContextHandle checkpointContext) { if (checkpointContext == null) { // Clean up any leaking scopes from pekko-streams/pekko-http etc. diff --git a/dd-java-agent/instrumentation/pekko/pekko-concurrent-1.0/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoMailboxInstrumentation.java b/dd-java-agent/instrumentation/pekko/pekko-concurrent-1.0/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoMailboxInstrumentation.java index ab33ee024c7..17be8e7476f 100644 --- a/dd-java-agent/instrumentation/pekko/pekko-concurrent-1.0/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoMailboxInstrumentation.java +++ b/dd-java-agent/instrumentation/pekko/pekko-concurrent-1.0/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoMailboxInstrumentation.java @@ -7,7 +7,7 @@ import static net.bytebuddy.matcher.ElementMatchers.isMethod; import com.google.auto.service.AutoService; -import datadog.context.Context; +import datadog.context.ContextHandle; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.api.InstrumenterConfig; @@ -45,7 +45,7 @@ public void methodAdvice(MethodTransformer transformer) { */ public static final class SuppressMailboxRunAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static Context enter() { + public static ContextHandle enter() { if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) { // remember the currently active scope so we can roll back to this point checkpointActiveForRollback(); @@ -56,7 +56,7 @@ public static Context enter() { } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) - public static void exit(@Advice.Enter final Context checkpointContext) { + public static void exit(@Advice.Enter final ContextHandle checkpointContext) { if (checkpointContext == null) { // Clean up any leaking scopes from pekko-streams/pekko-http etc. rollbackActiveToCheckpoint(); diff --git a/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogWrapperHelper.java b/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogWrapperHelper.java index 7766d5018e6..99d7cee15ca 100644 --- a/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogWrapperHelper.java +++ b/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogWrapperHelper.java @@ -1,9 +1,12 @@ package datadog.trace.instrumentation.pekkohttp; +import static datadog.context.Context.isCaptureEnabled; +import static datadog.context.Context.setCaptureEnabled; import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; import static datadog.trace.instrumentation.pekkohttp.PekkoHttpServerDecorator.DECORATE; import datadog.context.Context; +import datadog.context.ContextHandle; import datadog.context.ContextScope; import datadog.trace.api.InstrumenterConfig; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -38,14 +41,17 @@ private static Context startSpan(final HttpRequest request) { static final class SwappedContextScope implements ContextScope { private final Context context; - private final Context previousContext; + private final ContextHandle previousContext; private final Thread ownerThread; private boolean closed; SwappedContextScope(final Context context) { this.context = context; this.ownerThread = Thread.currentThread(); + boolean allowCapture = isCaptureEnabled(); this.previousContext = context.swap(); + // Request handling inherits the enclosing flow's capture flag. + setCaptureEnabled(allowCapture); } @Override diff --git a/dd-java-agent/instrumentation/zio/zio-2.0/src/main/java/datadog/trace/instrumentation/zio/v2_0/FiberContext.java b/dd-java-agent/instrumentation/zio/zio-2.0/src/main/java/datadog/trace/instrumentation/zio/v2_0/FiberContext.java index d423dca0d4a..8d97474b3dc 100644 --- a/dd-java-agent/instrumentation/zio/zio-2.0/src/main/java/datadog/trace/instrumentation/zio/v2_0/FiberContext.java +++ b/dd-java-agent/instrumentation/zio/zio-2.0/src/main/java/datadog/trace/instrumentation/zio/v2_0/FiberContext.java @@ -2,28 +2,29 @@ import datadog.context.Context; import datadog.context.ContextContinuation; +import datadog.context.ContextHandle; public class FiberContext { - private Context context; + private ContextHandle fiberContext; private final ContextContinuation continuation; - private Context originalContext; + private ContextHandle previousContext; public FiberContext() { - // record context to use for this coroutine - this.context = Context.current(); + Context initial = Context.current(); // stop enclosing trace from finishing early - this.continuation = this.context.capture(); + this.continuation = initial.capture(); + this.fiberContext = initial.checkpoint(); } public void onResume() { - originalContext = context.swap(); + previousContext = fiberContext.swap(); } public void onSuspend() { - if (originalContext != null) { - context = originalContext.swap(); - originalContext = null; + if (previousContext != null) { + fiberContext = previousContext.swap(); + previousContext = null; } } diff --git a/dd-trace-core/src/jmh/java/datadog/context/ContextManagerBenchmark.java b/dd-trace-core/src/jmh/java/datadog/context/ContextManagerBenchmark.java index 82b15ce62b0..44e930b4716 100644 --- a/dd-trace-core/src/jmh/java/datadog/context/ContextManagerBenchmark.java +++ b/dd-trace-core/src/jmh/java/datadog/context/ContextManagerBenchmark.java @@ -9,6 +9,7 @@ import java.util.concurrent.Executors; import java.util.concurrent.Semaphore; import java.util.concurrent.ThreadFactory; +import javax.annotation.Nonnull; import org.openjdk.jmh.annotations.Benchmark; import org.openjdk.jmh.annotations.BenchmarkMode; import org.openjdk.jmh.annotations.Fork; @@ -116,22 +117,35 @@ public Context current() { } @Override - public ContextScope attach(Context ctx) { + public ContextScope attach(@Nonnull Context ctx) { return csm.attach(ctx); } @Override - public Context swap(Context ctx) { + public ContextHandle swap(@Nonnull ContextHandle ctx) { return csm.swap(ctx); } @Override - public ContextContinuation capture(Context ctx) { + public ContextHandle checkpoint(@Nonnull Context ctx) { + return ctx; + } + + @Override + public ContextContinuation capture(@Nonnull Context ctx) { return csm.capture(ctx); } @Override - public void addListener(ContextListener l) {} + public boolean isCaptureEnabled() { + return true; + } + + @Override + public void setCaptureEnabled(boolean enabled) {} + + @Override + public void addListener(@Nonnull ContextListener l) {} }; } return ThreadLocalContextManager.INSTANCE; @@ -232,8 +246,8 @@ public void nestedAttachAndClose(ThreadState thread) { /** * Swap in a new context then swap back the previous one. {@link - * ContinuableScopeManager#swap(Context)} replaces the entire scope stack, making this a heavier - * operation than in {@link ThreadLocalContextManager}. + * ContinuableScopeManager#swap(ContextHandle)} replaces the entire scope stack, making this a + * heavier operation than in {@link ThreadLocalContextManager}. * *

Note: GCProfiler will show allocation asymmetry here by design. {@link * ContinuableScopeManager} swap allocates a {@code ScopeStack}, a {@code ContinuableScope}, and a @@ -243,7 +257,7 @@ public void nestedAttachAndClose(ThreadState thread) { @Benchmark public void swapContexts(ThreadState thread) { Context ctx = thread.nextContext(contexts); - Context previous = manager.swap(ctx); + ContextHandle previous = manager.swap(ctx); manager.swap(previous); } diff --git a/dd-trace-core/src/jmh/java/datadog/trace/core/VirtualThreadContextBenchmark.java b/dd-trace-core/src/jmh/java/datadog/trace/core/VirtualThreadContextBenchmark.java index 6c4098c3169..f9829f9d334 100644 --- a/dd-trace-core/src/jmh/java/datadog/trace/core/VirtualThreadContextBenchmark.java +++ b/dd-trace-core/src/jmh/java/datadog/trace/core/VirtualThreadContextBenchmark.java @@ -3,6 +3,7 @@ import static java.util.concurrent.TimeUnit.MICROSECONDS; import datadog.context.Context; +import datadog.context.ContextHandle; import datadog.trace.api.EndpointTracker; import datadog.trace.api.Stateful; import datadog.trace.api.profiling.ProfilingContextAttribute; @@ -30,7 +31,7 @@ /** * Java allocation and control-flow cost of virtual-thread context propagation: the current full - * {@link ContinuableScopeManager#swap(Context)} on every mount/unmount versus the proposed + * {@link ContinuableScopeManager#swap(ContextHandle)} on every mount/unmount versus the proposed * seed-once path (nothing when profiling is off, a profiler rebind/unbind when it is on). The * current path retains the context returned by each reverse swap, matching repeated mounts after * the first rather than reconstructing the virtual-thread context on every cycle. ddprof's native @@ -72,8 +73,8 @@ public void setup() { public static class ThreadState { AgentSpan span; Context spanContext; - Context plainContext; - Context profiledContext; + ContextHandle plainContext; + ContextHandle profiledContext; @Setup(Level.Trial) public void setup(VirtualThreadContextBenchmark bench) { @@ -93,13 +94,13 @@ public void tearDown(VirtualThreadContextBenchmark bench) { @Benchmark public void currentCycle_profilingOff(ThreadState t) { - Context previous = plainManager.swap(t.plainContext); + ContextHandle previous = plainManager.swap(t.plainContext); t.plainContext = plainManager.swap(previous); } @Benchmark public long currentCycle_profilingOn_javaStub(ThreadState t) { - Context previous = profiledManager.swap(t.profiledContext); + ContextHandle previous = profiledManager.swap(t.profiledContext); t.profiledContext = profiledManager.swap(previous); return stubProfiling.currentValue(); } 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 20d4e83dd86..2609f5359bf 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 @@ -25,6 +25,7 @@ import datadog.communication.ddagent.SharedCommunicationObjects; import datadog.context.Context; import datadog.context.ContextContinuation; +import datadog.context.ContextHandle; import datadog.context.ContextScope; import datadog.context.propagation.Propagators; import datadog.environment.ThreadSupport; @@ -2733,7 +2734,7 @@ public ContextScope attach(Context context) { } @Override - public Context swap(Context context) { + public ContextHandle swap(ContextHandle context) { return scopeManager.swap(context); } diff --git a/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuableScopeManager.java b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuableScopeManager.java index e0ac7ea14fa..685694d3904 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuableScopeManager.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuableScopeManager.java @@ -13,6 +13,7 @@ import datadog.context.Context; import datadog.context.ContextContinuation; +import datadog.context.ContextHandle; import datadog.context.ContextScope; import datadog.logging.RatelimitedLogger; import datadog.trace.api.Config; @@ -217,7 +218,7 @@ ContinuableScope continueSpan( public boolean isAsyncPropagationEnabled() { ContinuableScope activeScope = scopeStack().active(); - return activeScope != null && activeScope.isAsyncPropagating(); + return activeScope == null || activeScope.isAsyncPropagating(); } public void setAsyncPropagationEnabled(boolean asyncPropagationEnabled) { @@ -360,17 +361,18 @@ public ContextScope attach(@Nonnull Context context) { return activate(context); } - public Context swap(@Nonnull Context context) { + public ContextHandle swap(@Nonnull ContextHandle handle) { ScopeStack oldStack = tlsScopeStack.get(); ContinuableScope oldScope = oldStack.top; ScopeStack newStack; ContinuableScope newScope; - if (context instanceof ScopeContext) { + if (handle instanceof ScopeContext) { // restore previously swapped context stack - newStack = ((ScopeContext) context).restore(profilingContextIntegration); + newStack = ((ScopeContext) handle).restore(profilingContextIntegration); newScope = newStack.top; - } else if (context != Context.root()) { + } else if (handle != Context.root()) { + Context context = (Context) handle; // start a new stack and record the new context as active newStack = new ScopeStack(profilingContextIntegration); newScope = new ContinuableScope(this, context, CONTEXT, true, createScopeState(context)); diff --git a/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ScopeContext.java b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ScopeContext.java index 8919fae96f4..e00ec5ab7b9 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ScopeContext.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ScopeContext.java @@ -1,12 +1,14 @@ package datadog.trace.core.scopemanager; import datadog.context.Context; +import datadog.context.ContextHandle; import datadog.context.ContextKey; import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration; +import javax.annotation.Nonnull; import javax.annotation.Nullable; -/** Wraps a {@link ScopeStack} as a {@link Context} so it can be swapped back later. */ -final class ScopeContext implements Context { +/** Wraps a {@link ScopeStack} as a {@link ContextHandle} so it can be swapped back later. */ +final class ScopeContext implements ContextHandle { private final Thread originalThread = Thread.currentThread(); private final ScopeStack scopeStack; private final ContinuableScope parent; @@ -35,12 +37,7 @@ ScopeStack restore(ProfilingContextIntegration profilingContextIntegration) { @Nullable @Override - public T get(ContextKey key) { + public T get(@Nonnull ContextKey key) { return context.get(key); } - - @Override - public Context with(ContextKey key, @Nullable T value) { - return context.with(key, value); - } } diff --git a/dd-trace-core/src/test/java/datadog/trace/core/scopemanager/ScopeManagerForkedTest.java b/dd-trace-core/src/test/java/datadog/trace/core/scopemanager/ScopeManagerForkedTest.java index a6e833a215e..2089a88255c 100644 --- a/dd-trace-core/src/test/java/datadog/trace/core/scopemanager/ScopeManagerForkedTest.java +++ b/dd-trace-core/src/test/java/datadog/trace/core/scopemanager/ScopeManagerForkedTest.java @@ -25,6 +25,7 @@ import datadog.context.Context; import datadog.context.ContextContinuation; +import datadog.context.ContextHandle; import datadog.context.ContextKey; import datadog.context.ContextScope; import datadog.trace.api.DDTraceId; @@ -1059,7 +1060,7 @@ void contextsCanBeSwappedOutAndBack() { Context context1 = Context.root().with(testKey, "first-value"); Context context2 = context1.with(testKey, "second-value"); - Context swappedOut = scopeManager.swap(Context.root()); + ContextHandle swappedOut = scopeManager.swap(Context.root()); assertNull(scopeManager.active()); assertEquals(Context.root(), scopeManager.currentContext()); @@ -1086,7 +1087,7 @@ void contextsCanBeSwappedOutAndBack() { assertEquals(context2, scopeManager.currentContext()); assertEquals("first-value", swappedOut.get(testKey)); - Context context3 = swappedOut.with(testKey, "third-value"); + Context context3 = context1.with(testKey, "third-value"); scopeManager.swap(context3); assertNotNull(scopeManager.active()); diff --git a/docs/how_instrumentations_work.md b/docs/how_instrumentations_work.md index 92c963c3e53..686df15e354 100644 --- a/docs/how_instrumentations_work.md +++ b/docs/how_instrumentations_work.md @@ -883,6 +883,29 @@ in [`HttpClientRequestTracingHandler`](https://github.com/DataDog/dd-trace-java/ how the Continuation is used to obtain the `parentScope` which is finally [closed](https://github.com/DataDog/dd-trace-java/blob/3fe1b2d6010e50f61518fa25af3bdeb03ae7712b/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientRequestTracingHandler.java#L111). +### Capture controls and context switches + +`Context.isCaptureEnabled()` and `Context.setCaptureEnabled(boolean)` control capture for the +current execution unit; capture is enabled by default. When disabled, `Context.capture()` returns +a no-op continuation. Explicit `attach()` and `swap()` operations still work. + +Attaching a different context preserves the capture flag; closing its scope restores the flag +from before the attach. Reattaching the current context does not restore the flag on close. +When suppressing capture outside an attached scope, save the previous flag and restore it in `finally`. + +`Context.checkpoint()` returns a `ContextHandle` saving the context and current execution-unit flags +for a future `swap()`. It leaves the active context and flags unchanged and does not keep a trace +open; capture the original context first if a continuation is needed. + +Swapping to a different context resets execution-unit flags unless the context is checkpointed, +in which case its saved flags are restored. A swap to the current context leaves flags unchanged. +Every swap returns a `ContextHandle` for restoration by swapping back. Handles expose +`get(ContextKey)` for reading values, including the span used by messaging instrumentations. +`Context` extends `ContextHandle`, so ordinary contexts can also be swapped directly. + +See the [Context API](../components/context/src/main/java/datadog/context/Context.java) for the +full contract. + ### Continuation effects An unresolved continuation keeps `PendingTrace`'s reference count positive and prevents its normal diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpan.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpan.java index f29f17d6629..98b0ac0a38a 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpan.java +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpan.java @@ -4,6 +4,7 @@ import datadog.context.Context; import datadog.context.ContextContinuation; +import datadog.context.ContextHandle; import datadog.context.ContextKey; import datadog.context.ContextScope; import datadog.context.ImplicitContextKeyed; @@ -39,12 +40,12 @@ static AgentSpan current() { } /** - * Extracts the span from the given {@link Context}. + * Extracts the span from the given {@link ContextHandle}. * * @param context the context to extract the span from. * @return the span if existing, {@code null} otherwise. */ - static AgentSpan fromContext(Context context) { + static AgentSpan fromContext(ContextHandle context) { return context == null ? null : context.get(SPAN_KEY); } diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java index 71767ea7fae..3b7ad08e225 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java @@ -2,6 +2,7 @@ import datadog.context.Context; import datadog.context.ContextContinuation; +import datadog.context.ContextHandle; import datadog.context.ContextListener; import datadog.context.ContextManager; import datadog.context.ContextScope; @@ -410,7 +411,7 @@ default AgentSpan startSpan(@Nonnull SpanPrototype prototype, CharSequence opera ContextScope attach(Context context); @Deprecated - Context swap(Context context); + ContextHandle swap(ContextHandle context); @Deprecated ContextContinuation capture(Context context); @@ -683,7 +684,7 @@ public ContextScope attach(Context context) { } @Override - public Context swap(Context context) { + public ContextHandle swap(ContextHandle context) { return Context.root(); } @@ -797,6 +798,7 @@ public static void installLegacyContextManager() { } /** Shim mapping new Context API to legacy scope manager. */ + @SuppressWarnings("deprecation") static final class LegacyContextManager implements ContextManager { static final LegacyContextManager INSTANCE = new LegacyContextManager(); @@ -811,15 +813,30 @@ public ContextScope attach(@Nonnull Context context) { } @Override - public Context swap(@Nonnull Context context) { + public ContextHandle swap(@Nonnull ContextHandle context) { return AgentTracer.get().swap(context); } + @Override + public ContextHandle checkpoint(@Nonnull Context context) { + return context; // execution-unit flags are maintained by the scope stack + } + @Override public ContextContinuation capture(@Nonnull Context context) { return AgentTracer.get().capture(context); } + @Override + public boolean isCaptureEnabled() { + return AgentTracer.get().isAsyncPropagationEnabled(); + } + + @Override + public void setCaptureEnabled(boolean enabled) { + AgentTracer.get().setAsyncPropagationEnabled(enabled); + } + @Override public void addListener(@Nonnull ContextListener listener) { // this method is never used in legacy mode...