Introduce a shared context component, independent of tracing. (#8117)

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

Author: mcculls

components/context/build.gradle.kts

@@ -0,0 +1,13 @@
+plugins {
+  id("me.champeau.jmh")
+}
+
+apply(from = "$rootDir/gradle/java.gradle")
+
+jmh {
+  version = "1.28"
+}
+
+val excludedClassesInstructionCoverage by extra {
+  listOf("datadog.context.ContextProviders") // covered by forked test
+}

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

@@ -0,0 +1,99 @@
+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.
+ *
+ * <p>Each element of the context is accessible by its {@link ContextKey}. Keys represents product
+ * or functional areas and should be created sparingly. Elements in the context may themselves be
+ * mutable.
+ */
+public interface Context {
+
+  /**
+   * Returns the root context.
+   *
+   * <p>This is the initial local context that all contexts extend.
+   */
+  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);
+  }
+
+  /**
+   * 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 from the same elements, except the key is now mapped to the given value.
+   *
+   * @return New context with the key-value mapping.
+   */
+  <T> Context with(ContextKey<T> key, T value);
+
+  /**
+   * Creates a new context from the same elements, except the implicit key is mapped to this value.
+   *
+   * @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,27 @@
+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;
+  }
+}

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

@@ -0,0 +1,38 @@
+package datadog.context;
+
+import java.util.concurrent.atomic.AtomicInteger;
+
+/**
+ * {@link Context} key that maps to a value of type {@link T}.
+ *
+ * <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;
+  }
+
+  // we want identity equality, so no need to override equals()
+
+  @Override
+  public String toString() {
+    return name;
+  }
+}

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

@@ -0,0 +1,38 @@
+package datadog.context;
+
+/** Manages context across execution units. */
+public interface ContextManager {
+
+  /**
+   * Returns the root context.
+   *
+   * <p>This is the initial local context that all contexts extend.
+   */
+  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);
+
+  /** Requests use of a custom {@link ContextManager}. */
+  static void register(ContextManager manager) {
+    ContextProviders.customManager = manager;
+  }
+}

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

@@ -0,0 +1,30 @@
+package datadog.context;
+
+/** Provides {@link ContextManager} and {@link ContextBinder} implementations. */
+final class ContextProviders {
+
+  static volatile ContextManager customManager;
+  static volatile ContextBinder customBinder;
+
+  private static final class ProvidedManager {
+    static final ContextManager INSTANCE =
+        null != ContextProviders.customManager
+            ? ContextProviders.customManager
+            : new ThreadLocalContextManager();
+  }
+
+  private static final class ProvidedBinder {
+    static final ContextBinder INSTANCE =
+        null != ContextProviders.customBinder
+            ? ContextProviders.customBinder
+            : new WeakMapContextBinder();
+  }
+
+  static ContextManager manager() {
+    return ProvidedManager.INSTANCE; // may be overridden by instrumentation
+  }
+
+  static ContextBinder binder() {
+    return ProvidedBinder.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 SingletonContext(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;
+  }
+}

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

@@ -0,0 +1,51 @@
+package datadog.context;
+
+import static java.lang.Math.max;
+
+import java.util.Objects;
+
+/** {@link Context} containing a single value. */
+final class SingletonContext implements Context {
+  private final int index;
+  private final Object value;
+
+  SingletonContext(int index, Object value) {
+    this.index = index;
+    this.value = value;
+  }
+
+  @Override
+  @SuppressWarnings("unchecked")
+  public <V> V get(ContextKey<V> key) {
+    return index == key.index ? (V) value : null;
+  }
+
+  @Override
+  public <V> Context with(ContextKey<V> secondKey, V secondValue) {
+    int secondIndex = secondKey.index;
+    if (index == secondIndex) {
+      return new SingletonContext(index, secondValue);
+    } else {
+      Object[] store = new Object[max(index, secondIndex) + 1];
+      store[index] = value;
+      store[secondIndex] = secondValue;
+      return new IndexedContext(store);
+    }
+  }
+
+  @Override
+  public boolean equals(Object o) {
+    if (this == o) return true;
+    if (o == null || getClass() != o.getClass()) return false;
+    SingletonContext that = (SingletonContext) o;
+    return index == that.index && Objects.equals(value, that.value);
+  }
+
+  @Override
+  public int hashCode() {
+    int result = 31;
+    result = 31 * result + index;
+    result = 31 * result + Objects.hashCode(value);
+    return result;
+  }
+}