Introduce a shared context component, independent of tracing.

Co-authored-by: Bruce Bujon <[email protected]>

Author: mcculls

components/context/build.gradle.kts

@@ -0,0 +1,9 @@
+plugins {
+  id("me.champeau.jmh")
+}
+
+apply(from = "$rootDir/gradle/java.gradle")
+
+jmh {
+  version = "1.28"
+}

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

@@ -0,0 +1,100 @@
+package datadog.context;
+
+import static datadog.context.ContextProviders.binder;
+import static datadog.context.ContextProviders.manager;
+
+import javax.annotation.Nullable;
+
+/** Immutable context scoped to an execution unit or carrier object. */
+public interface Context {
+
+  /** Returns the root context. */
+  static Context root() {
+    return manager().root();
+  }
+
+  /**
+   * Returns the context attached to the current execution unit.
+   *
+   * @return Attached context; {@link #root()} if there is none
+   */
+  static Context current() {
+    return manager().current();
+  }
+
+  /**
+   * Attaches this context to the current execution unit.
+   *
+   * @return Scope to be closed when the context is invalid.
+   */
+  default ContextScope attach() {
+    return manager().attach(this);
+  }
+
+  /**
+   * Swaps this context with the one attached to current execution unit.
+   *
+   * @return Previously attached context; {@link #root()} if there was none
+   */
+  default Context swap() {
+    return manager().swap(this);
+  }
+
+  /**
+   * Detaches the context attached to the current execution unit, leaving it context-less.
+   *
+   * <p>WARNING: prefer {@link ContextScope#close()} to properly restore the surrounding context.
+   *
+   * @return Previously attached context; {@link #root()} if there was none
+   */
+  static Context detach() {
+    return manager().detach();
+  }
+
+  /**
+   * Returns the context attached to the given carrier object.
+   *
+   * @return Attached context; {@link #root()} if there is none
+   */
+  static Context from(Object carrier) {
+    return binder().from(carrier);
+  }
+
+  /** Attaches this context to the given carrier object. */
+  default void attachTo(Object carrier) {
+    binder().attachTo(carrier, this);
+  }
+
+  /**
+   * Detaches the context attached to the given carrier object, leaving it context-less.
+   *
+   * @return Previously attached context; {@link #root()} if there was none
+   */
+  static Context detachFrom(Object carrier) {
+    return binder().detachFrom(carrier);
+  }
+
+  /**
+   * Gets the value stored in this context under the given key.
+   *
+   * @return Value stored under the key; {@code null} if there is no value.
+   */
+  @Nullable
+  <T> T get(ContextKey<T> key);
+
+  /**
+   * Creates a new context with the given key-value mapping.
+   *
+   * @return New context with the key-value mapping.
+   */
+  <T> Context with(ContextKey<T> key, T value);
+
+  /**
+   * Creates a new context with a value that has its own implicit key.
+   *
+   * @return New context with the implicitly keyed value.
+   */
+  default Context with(ImplicitContextKeyed value) {
+    return value.storeInto(this);
+  }
+}

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

@@ -0,0 +1,34 @@
+package datadog.context;
+
+/** Binds context to carrier objects. */
+public interface ContextBinder {
+
+  /**
+   * Returns the context attached to the given carrier object.
+   *
+   * @return Attached context; {@link Context#root()} if there is none
+   */
+  Context from(Object carrier);
+
+  /** Attaches the given context to the given carrier object. */
+  void attachTo(Object carrier, Context context);
+
+  /**
+   * Detaches the context attached to the given carrier object, leaving it context-less.
+   *
+   * @return Previously attached context; {@link Context#root()} if there was none
+   */
+  Context detachFrom(Object carrier);
+
+  /** Requests use of a custom {@link ContextBinder}. */
+  static void register(ContextBinder binder) {
+    ContextProviders.customBinder = binder;
+  }
+
+  final class Provided {
+    static final ContextBinder INSTANCE =
+        null != ContextProviders.customBinder
+            ? ContextProviders.customBinder
+            : new WeakMapContextBinder();
+  }
+}

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

@@ -0,0 +1,47 @@
+package datadog.context;
+
+import java.util.concurrent.atomic.AtomicInteger;
+
+/**
+ * Key for indexing values of type {@link T} stored in a {@link Context}.
+ *
+ * <p>Keys are compared by identity rather than by name. Each stored context type should either
+ * share its key for re-use or implement {@link ImplicitContextKeyed} to keep its key private.
+ */
+public final class ContextKey<T> {
+  private static final AtomicInteger NEXT_INDEX = new AtomicInteger(0);
+
+  private final String name;
+  final int index;
+
+  private ContextKey(String name) {
+    this.name = name;
+    this.index = NEXT_INDEX.getAndIncrement();
+  }
+
+  /** Creates a new key with the given name. */
+  public static <T> ContextKey<T> named(String name) {
+    return new ContextKey<>(name);
+  }
+
+  @Override
+  public int hashCode() {
+    return index;
+  }
+
+  @Override
+  public boolean equals(Object o) {
+    if (this == o) {
+      return true;
+    } else if (o == null || getClass() != o.getClass()) {
+      return false;
+    } else {
+      return index == ((ContextKey<?>) o).index;
+    }
+  }
+
+  @Override
+  public String toString() {
+    return name;
+  }
+}

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

@@ -0,0 +1,50 @@
+package datadog.context;
+
+/** Manages context across execution units. */
+public interface ContextManager {
+
+  /** Returns the root context. */
+  Context root();
+
+  /**
+   * Returns the context attached to the current execution unit.
+   *
+   * @return Attached context; {@link #root()} if there is none
+   */
+  Context current();
+
+  /**
+   * Attaches the given context to the current execution unit.
+   *
+   * @return Scope to be closed when the context is invalid.
+   */
+  ContextScope attach(Context context);
+
+  /**
+   * Swaps the given context with the one attached to current execution unit.
+   *
+   * @return Previously attached context; {@link #root()} if there was none
+   */
+  Context swap(Context context);
+
+  /**
+   * Detaches the context attached to the current execution unit, leaving it context-less.
+   *
+   * <p>WARNING: prefer {@link ContextScope#close()} to properly restore the surrounding context.
+   *
+   * @return Previously attached context; {@link #root()} if there was none
+   */
+  Context detach();
+
+  /** Requests use of a custom {@link ContextManager}. */
+  static void register(ContextManager manager) {
+    ContextProviders.customManager = manager;
+  }
+
+  final class Provided {
+    static final ContextManager INSTANCE =
+        null != ContextProviders.customManager
+            ? ContextProviders.customManager
+            : new ThreadLocalContextManager();
+  }
+}

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

@@ -0,0 +1,16 @@
+package datadog.context;
+
+/** Provides {@link ContextManager} and {@link ContextBinder} implementations. */
+final class ContextProviders {
+
+  static volatile ContextManager customManager;
+  static volatile ContextBinder customBinder;
+
+  static ContextManager manager() {
+    return ContextManager.Provided.INSTANCE; // may be overridden by instrumentation
+  }
+
+  static ContextBinder binder() {
+    return ContextBinder.Provided.INSTANCE; // may be overridden by instrumentation
+  }
+}

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

@@ -0,0 +1,12 @@
+package datadog.context;
+
+/** Controls the validity of context attached to an execution unit. */
+public interface ContextScope extends AutoCloseable {
+
+  /** Returns the context controlled by this scope. */
+  Context context();
+
+  /** Detaches the context from the execution unit. */
+  @Override
+  void close();
+}

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

@@ -0,0 +1,16 @@
+package datadog.context;
+
+/** {@link Context} containing no values. */
+final class EmptyContext implements Context {
+  static final Context INSTANCE = new EmptyContext();
+
+  @Override
+  public <T> T get(ContextKey<T> key) {
+    return null;
+  }
+
+  @Override
+  public <T> Context with(ContextKey<T> key, T value) {
+    return new SingleContext(key.index, value);
+  }
+}

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

@@ -0,0 +1,12 @@
+package datadog.context;
+
+/** {@link Context} value that has its own implicit {@link ContextKey}. */
+public interface ImplicitContextKeyed {
+
+  /**
+   * Creates a new context with this value under its chosen key.
+   *
+   * @return New context with the implicitly keyed value.
+   */
+  Context storeInto(Context context);
+}

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

@@ -0,0 +1,46 @@
+package datadog.context;
+
+import static java.lang.Math.max;
+import static java.util.Arrays.copyOfRange;
+
+import java.util.Arrays;
+
+/** {@link Context} containing many values. */
+final class IndexedContext implements Context {
+  private final Object[] store;
+
+  IndexedContext(Object[] store) {
+    this.store = store;
+  }
+
+  @Override
+  @SuppressWarnings("unchecked")
+  public <T> T get(ContextKey<T> key) {
+    int index = key.index;
+    return index < store.length ? (T) store[index] : null;
+  }
+
+  @Override
+  public <T> Context with(ContextKey<T> key, T value) {
+    int index = key.index;
+    Object[] newStore = copyOfRange(store, 0, max(store.length, index + 1));
+    newStore[index] = value;
+
+    return new IndexedContext(newStore);
+  }
+
+  @Override
+  public boolean equals(Object o) {
+    if (this == o) return true;
+    if (o == null || getClass() != o.getClass()) return false;
+    IndexedContext that = (IndexedContext) o;
+    return Arrays.equals(store, that.store);
+  }
+
+  @Override
+  public int hashCode() {
+    int result = 31;
+    result = 31 * result + Arrays.hashCode(store);
+    return result;
+  }
+}