From 699138751444927cc5fe4e071d02b7106252ff55 Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Wed, 4 Nov 2026 13:49:38 +0000 Subject: [PATCH] Add context capture controls and preserve flags across swaps Expose capture controls on Context and ContextManager, bridging them to legacy async propagation flags. Disabled capture returns a no-op continuation; closing an attached scope restores its previous flag. Introduce ContextHandle for saved contexts and checkpoint() to preserve execution-unit flags without changing the active context. Swapping to a different context resets flags or restores those saved in a handle, and returns a checkpoint of the displaced context. Same-context swaps leave flags unchanged. Use checkpoint handles across coroutine, fiber, actor, poller, and virtual thread context switches. Checkpoint the initial ZIO fiber context after capturing its continuation, and inherit the enclosing flow's capture flag in Akka and Pekko HTTP handlers. Cover capture suppression, checkpoints, scope restoration, context swaps, cross-thread restoration, virtual threads, and custom context providers. Update context tracking guidance and benchmarks for the expanded API. --- .../references/context-tracking.md | 10 + .../skills/fix-continuation-leakage/SKILL.md | 6 +- .../main/java/datadog/context/Context.java | 45 ++- .../java/datadog/context/ContextHandle.java | 36 ++ .../java/datadog/context/ContextManager.java | 49 ++- .../java/datadog/context/FlaggedContext.java | 20 ++ .../datadog/context/TestContextManager.java | 17 +- .../context/ThreadLocalContextManager.java | 74 +++- .../datadog/context/ContextCaptureTest.java | 333 ++++++++++++++++++ .../context/ContextListenerEventTest.java | 2 +- .../context/ContextProvidersForkedTest.java | 33 +- .../api/Java8BytecodeBridge.java | 19 +- .../java/lang/VirtualThreadState.java | 5 +- .../ScopeDiagnosticsIntegrationTest.java | 6 +- .../AkkaActorCellInstrumentation.java | 6 +- .../AkkaMailboxInstrumentation.java | 6 +- .../akkahttp/DatadogWrapperHelper.java | 8 +- .../lang/jdk21/PollerInstrumentation.java | 6 +- .../DatadogThreadContextElement.java | 17 +- .../PekkoActorCellInstrumentation.java | 6 +- .../PekkoMailboxInstrumentation.java | 6 +- .../pekkohttp/DatadogWrapperHelper.java | 8 +- .../zio/v2_0/FiberContext.java | 19 +- .../context/ContextManagerBenchmark.java | 28 +- .../core/VirtualThreadContextBenchmark.java | 11 +- .../java/datadog/trace/core/CoreTracer.java | 3 +- .../scopemanager/ContinuableScopeManager.java | 12 +- .../trace/core/scopemanager/ScopeContext.java | 13 +- .../scopemanager/ScopeManagerForkedTest.java | 5 +- docs/how_instrumentations_work.md | 23 ++ .../instrumentation/api/AgentSpan.java | 5 +- .../instrumentation/api/AgentTracer.java | 23 +- 32 files changed, 750 insertions(+), 110 deletions(-) create mode 100644 components/context/src/main/java/datadog/context/ContextHandle.java create mode 100644 components/context/src/main/java/datadog/context/FlaggedContext.java create mode 100644 components/context/src/test/java/datadog/context/ContextCaptureTest.java 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...