Provide ability to capture continuations using the Context API

components/context/src/main/java/datadog/context/Context.java

@@ -101,6 +101,15 @@ static Context detachFrom(Object carrier) {
     return binder().detachFrom(carrier);
   }
 
+  /**
+   * Captures this (attached) context so it can be resumed in another execution unit.
+   *
+   * @return continuation capturing this context.
+   */
+  default ContextContinuation capture() {
+    return manager().capture(this);
+  }
+
   /**
    * Gets the value stored in this context under the given key.
    *
@@ -159,4 +168,13 @@ default Context with(@Nullable ImplicitContextKeyed value) {
     }
     return value.storeInto(this);
   }
+
+  /**
+   * Wraps context as a scope without attaching it to the current execution unit.
+   *
+   * @return a scope that has no effect on execution units.
+   */
+  default ContextScope asScope() {
+    return new NoopContextScope(this);
+  }
 }

components/context/src/main/java/datadog/context/ContextContinuation.java

@@ -0,0 +1,79 @@
+package datadog.context;
+
+/**
+ * Captures a context attached to one execution unit so it can be resumed in another.
+ *
+ * <p>To propagate context to a single background task:
+ *
+ * <pre>{@code
+ * ContextContinuation continuation = Context.current().capture();
+ * executor.execute(() -> {
+ *   try (ContextScope scope = continuation.resume()) {
+ *     // ... Context.current() here returns the captured context
+ *   }
+ *   // context implicitly released from continuation when resumed scope closes
+ * });
+ * }</pre>
+ *
+ * <p>If a continuation is never resumed (e.g. a task is cancelled before it runs), you must release
+ * it explicitly to avoid resource leaks:
+ *
+ * <pre>{@code
+ * ContextContinuation continuation = Context.current().capture();
+ * Future<?> future = executor.submit(() -> {
+ *   try (ContextScope scope = continuation.resume()) {
+ *     // ...
+ *   }
+ * });
+ * // ...
+ * if (future.cancel(false)) {
+ *   continuation.release(); // task will never resume, so release manually
+ * }
+ * }</pre>
+ *
+ * <p>When the same context is resumed concurrently across multiple threads, call {@link #hold()}
+ * immediately after capture to prevent the first {@link #resume()} from releasing the context:
+ *
+ * <pre>{@code
+ * ContextContinuation continuation = Context.current().capture().hold();
+ * for (int i = 0; i < N; i++) {
+ *   executor.execute(() -> {
+ *     try (ContextScope scope = continuation.resume()) {
+ *       // ...
+ *     }
+ *   });
+ * }
+ * // ...
+ * continuation.release(); // remember to release the hold once all tasks are resumed/done
+ * }</pre>
+ */
+public interface ContextContinuation {
+
+  /**
+   * Optional builder method to stop {@link #resume()} from implicitly releasing the captured
+   * context. This is useful when multiple threads may concurrently resume the context. You must
+   * then explicitly {@link #release() release} the context once all threads are resumed/done.
+   *
+   * @return this continuation, but with implicit release-after-resume turned off.
+   */
+  ContextContinuation hold();
+
+  /**
+   * Returns the context captured by this continuation.
+   *
+   * @return the captured context.
+   */
+  Context context();
+
+  /**
+   * Resumes the context captured by this continuation by attaching it to the current execution
+   * unit. Implicitly {@link #release() releases} the captured context at the end of the resumed
+   * scope, unless {@link #hold()} was called when creating the continuation.
+   *
+   * @return a scope to be closed when the resumed context is invalid.
+   */
+  ContextScope resume();
+
+  /** Explicitly releases the context captured by this continuation. */
+  void release();
+}

components/context/src/main/java/datadog/context/ContextListener.java

@@ -0,0 +1,33 @@
+package datadog.context;
+
+/** Listener of context events. */
+public interface ContextListener {
+
+  /**
+   * Notifies that the given context has been attached to the current execution unit.
+   *
+   * @param context the attached context.
+   */
+  default void onAttach(Context context) {}
+
+  /**
+   * Notifies that the given context has been detached from the current execution unit.
+   *
+   * @param context the detached context.
+   */
+  default void onDetach(Context context) {}
+
+  /**
+   * Notifies that the given context has been captured by a continuation.
+   *
+   * @param context the captured context.
+   */
+  default void onCapture(Context context) {}
+
+  /**
+   * Notifies that the given context has been released from a continuation.
+   *
+   * @param context the released context.
+   */
+  default void onRelease(Context context) {}
+}

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

@@ -1,5 +1,7 @@
 package datadog.context;
 
+import static datadog.context.ContextProviders.manager;
+
 /** Manages context across execution units. */
 public interface ContextManager {
   /**
@@ -25,6 +27,29 @@ public interface ContextManager {
    */
   Context swap(Context context);
 
+  /**
+   * Captures the given (attached) context so it can be resumed in another execution unit.
+   *
+   * @return continuation capturing the context.
+   */
+  ContextContinuation capture(Context context);
+
+  /**
+   * Registers the given listener to receive context events.
+   *
+   * @param listener the listener to register
+   */
+  void addListener(ContextListener listener);
+
+  /**
+   * Registers the given listener to receive context events.
+   *
+   * @param listener the listener to register.
+   */
+  static void register(ContextListener listener) {
+    manager().addListener(listener);
+  }
+
   /**
    * Requests use of a custom {@link ContextManager}.
    *

components/context/src/main/java/datadog/context/EmptyContext.java

@@ -3,13 +3,13 @@
 import static java.util.Objects.requireNonNull;
 
 import javax.annotation.Nullable;
-import javax.annotation.ParametersAreNonnullByDefault;
 
 /** {@link Context} containing no values. */
-@ParametersAreNonnullByDefault
-final class EmptyContext implements Context {
+final class EmptyContext implements SelfScopedContext {
   static final Context INSTANCE = new EmptyContext();
 
+  private EmptyContext() {}
+
   @Override
   @Nullable
   public <T> T get(ContextKey<T> key) {

components/context/src/main/java/datadog/context/ImplicitContextKeyed.java

@@ -1,9 +1,6 @@
 package datadog.context;
 
-import javax.annotation.ParametersAreNonnullByDefault;
-
 /** {@link Context} value that has its own implicit {@link ContextKey}. */
-@ParametersAreNonnullByDefault
 public interface ImplicitContextKeyed {
   /**
    * Creates a new context with this value under its chosen key.

components/context/src/main/java/datadog/context/IndexedContext.java

@@ -6,11 +6,9 @@
 
 import java.util.Arrays;
 import javax.annotation.Nullable;
-import javax.annotation.ParametersAreNonnullByDefault;
 
 /** {@link Context} containing many values. */
-@ParametersAreNonnullByDefault
-final class IndexedContext implements Context {
+final class IndexedContext implements SelfScopedContext {
   final Object[] store;
 
   IndexedContext(Object[] store) {

components/context/src/main/java/datadog/context/NoopContextContinuation.java

@@ -0,0 +1,31 @@
+package datadog.context;
+
+/** {@link ContextContinuation} that has no effect on execution units. */
+final class NoopContextContinuation implements ContextContinuation, ContextScope {
+  private final Context context;
+
+  NoopContextContinuation(Context context) {
+    this.context = context;
+  }
+
+  @Override
+  public ContextContinuation hold() {
+    return this;
+  }
+
+  @Override
+  public Context context() {
+    return context;
+  }
+
+  @Override
+  public ContextScope resume() {
+    return this; // acts as no-op scope, avoiding allocation
+  }
+
+  @Override
+  public void release() {}
+
+  @Override
+  public void close() {}
+}

components/context/src/main/java/datadog/context/NoopContextScope.java

@@ -0,0 +1,18 @@
+package datadog.context;
+
+/** {@link ContextScope} that has no effect on execution units. */
+final class NoopContextScope implements ContextScope {
+  private final Context context;
+
+  NoopContextScope(Context context) {
+    this.context = context;
+  }
+
+  @Override
+  public Context context() {
+    return context;
+  }
+
+  @Override
+  public void close() {}
+}

components/context/src/main/java/datadog/context/SelfScopedContext.java

@@ -0,0 +1,17 @@
+package datadog.context;
+
+/** Context that acts as its own unattached scope. */
+interface SelfScopedContext extends Context, ContextScope {
+  @Override
+  default ContextScope asScope() {
+    return this; // acts as no-op scope, avoiding allocation
+  }
+
+  @Override
+  default Context context() {
+    return this;
+  }
+
+  @Override
+  default void close() {}
+}

components/context/src/main/java/datadog/context/SingletonContext.java

@@ -5,11 +5,9 @@
 
 import java.util.Objects;
 import javax.annotation.Nullable;
-import javax.annotation.ParametersAreNonnullByDefault;
 
 /** {@link Context} containing a single value. */
-@ParametersAreNonnullByDefault
-final class SingletonContext implements Context {
+final class SingletonContext implements SelfScopedContext {
   final int index;
   final Object value;
 

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

@@ -27,6 +27,20 @@ public Context swap(Context context) {
     return delegate().swap(context);
   }
 
+  @Override
+  public ContextContinuation capture(Context context) {
+    return delegate().capture(context);
+  }
+
+  @Override
+  public void addListener(ContextListener listener) {
+    delegate().addListener(listener);
+  }
+
+  static void clearListeners() {
+    ThreadLocalContextManager.INSTANCE.clearListeners();
+  }
+
   private static ContextManager delegate() {
     ContextManager delegate = ContextProviders.customManager;
     if (delegate == TEST_INSTANCE) {