Provide ability to capture continuations using the Context API

This is meant as a stepping stone from the old API towards a simpler foundation

Author: mcculls

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

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 {
   static final Context INSTANCE = new EmptyContext();
 
+  private EmptyContext() {}
+
   @Override
   @Nullable
   public <T> T get(ContextKey<T> key) {

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

@@ -0,0 +1,26 @@
+package datadog.context;
+
+/** {@link ContextContinuation} capturing the empty (root) context. */
+final class EmptyContextContinuation implements ContextContinuation {
+  static final ContextContinuation INSTANCE = new EmptyContextContinuation();
+
+  private EmptyContextContinuation() {}
+
+  @Override
+  public ContextContinuation hold() {
+    return this;
+  }
+
+  @Override
+  public Context context() {
+    return EmptyContext.INSTANCE;
+  }
+
+  @Override
+  public ContextScope resume() {
+    return NoopContextScope.create(context());
+  }
+
+  @Override
+  public void release() {}
+}

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,10 +6,8 @@
 
 import java.util.Arrays;
 import javax.annotation.Nullable;
-import javax.annotation.ParametersAreNonnullByDefault;
 
 /** {@link Context} containing many values. */
-@ParametersAreNonnullByDefault
 final class IndexedContext implements Context {
   final Object[] store;
 

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

@@ -0,0 +1,73 @@
+package datadog.context;
+
+import java.lang.ref.WeakReference;
+
+/** {@link ContextScope} that has no effect on execution units. */
+final class NoopContextScope extends WeakReference<Context> implements ContextScope {
+  private static final ContextScope ROOT_SCOPE = new NoopContextScope(Context.root());
+
+  private static final int CACHE_SIZE = 32; // must be power of 2
+  private static final int SLOT_MASK = CACHE_SIZE - 1;
+  private static final int MAX_HASH_ATTEMPTS = 3;
+
+  /** Bounded cache of no-op scopes to reduce (re)allocations. */
+  private static final NoopContextScope[] cache = new NoopContextScope[CACHE_SIZE];
+
+  @SuppressWarnings({"resource", "StatementWithEmptyBody"})
+  static ContextScope create(Context context) {
+    if (context == Context.root()) {
+      return ROOT_SCOPE;
+    }
+    int hash = System.identityHashCode(context);
+    int evictedSlot = -1;
+    // search by repeated hashing; stop when we find an empty slot,
+    // a matching slot, or we exhaust all attempts and re-use a slot
+    for (int i = 1, h = hash; true; i++, h = rehash(h)) {
+      int slot = SLOT_MASK & h;
+      NoopContextScope existing = cache[slot];
+      if (existing != null) {
+        // slot already used
+        Context existingContext = existing.get();
+        if (context == existingContext) {
+          return existing; // match found
+        }
+        if (i < MAX_HASH_ATTEMPTS) {
+          // still more slots to search
+          if (existingContext == null && evictedSlot < 0) {
+            // record first evicted slot for re-use later
+            evictedSlot = slot;
+          }
+          continue; // rehash and try again
+        }
+        // exhausted attempts, pick best slot to re-use
+        if (evictedSlot >= 0) {
+          slot = evictedSlot; // re-use first evicted slot
+        } else if (existingContext == null) {
+          // last hashed slot is itself evicted, re-use it
+        } else {
+          slot = SLOT_MASK & hash; // re-use first hashed slot
+        }
+      }
+      return (cache[slot] = new NoopContextScope(context));
+    }
+  }
+
+  private NoopContextScope(Context context) {
+    super(context);
+  }
+
+  @Override
+  public Context context() {
+    Context context = get();
+    // no-op scopes are used when the context is already attached so the reference
+    // value should still be there; if not then we fall back to empty (root) context
+    return context != null ? context : Context.root();
+  }
+
+  @Override
+  public void close() {}
+
+  private static int rehash(int oldHash) {
+    return Integer.reverseBytes(oldHash * 0x9e3775cd) * 0x9e3775cd;
+  }
+}

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

@@ -5,10 +5,8 @@
 
 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 int index;
   final Object value;