Add context capture controls and preserve flags across swaps

.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.

.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

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.
+   * <p>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 <T> 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.
+   * <p>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.
+   *
+   * <p>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}.
+   *
+   * <p>See {@link #swap()} for capture flag behavior across swaps.
+   *
+   * @param enabled {@code true} to allow capture, {@code false} to disable it.
    */
-  @Nullable
-  <T> T get(ContextKey<T> key);
+  static void setCaptureEnabled(boolean enabled) {
+    manager().setCaptureEnabled(enabled);
+  }
 
   /**
    * Creates a copy of this context with the given key-value set.

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.
+ *
+ * <p>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.
+   *
+   * <p>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 <T> 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> T get(ContextKey<T> key);
+}

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.
+   * <p>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}.
+   *
+   * <p>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.
+   *
+   * <p>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.
+   *
+   * <p>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}.
+   *
+   * <p>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.
    *

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> T get(ContextKey<T> key) {
+    return context.get(key);
+  }
+}

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);

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
   }
 }