WIP

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/apm/context/Context.java

@@ -0,0 +1,95 @@
+package datadog.apm.context;
+
+import javax.annotation.Nullable;
+
+/** Immutable context scoped to an execution unit. */
+public interface Context {
+
+  /** Returns the root context. */
+  static Context root() {
+    return ContextProvider.manager().root();
+  }
+
+  /**
+   * Returns the context attached to the current execution unit.
+   *
+   * @return Attached context; {@link #root()} if there is none
+   */
+  static Context current() {
+    return ContextProvider.manager().current();
+  }
+
+  /**
+   * Attaches this context to the current execution unit.
+   *
+   * @return Scope to be closed when the context is invalid.
+   */
+  default ContextScope attach() {
+    return ContextProvider.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 ContextProvider.manager().swap(this);
+  }
+
+  /**
+   * Detaches and returns the context attached to the current execution unit.
+   *
+   * @return Previously attached context; {@link #root()} if there was none
+   */
+  static Context detach() {
+    return ContextProvider.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 ContextProvider.binder().from(carrier);
+  }
+
+  /** Attaches this context to the given carrier object. */
+  default void attachTo(Object carrier) {
+    ContextProvider.binder().attachTo(carrier, this);
+  }
+
+  /**
+   * Detaches and returns the context attached to the given carrier object.
+   *
+   * @return Previously attached context; {@link #root()} if there was none
+   */
+  static Context detachFrom(Object carrier) {
+    return ContextProvider.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/apm/context/ContextBinder.java

@@ -0,0 +1,22 @@
+package datadog.apm.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 and returns the context attached to the given carrier object.
+   *
+   * @return Previously attached context; {@link Context#root()} if there was none
+   */
+  Context detachFrom(Object carrier);
+}

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

@@ -0,0 +1,49 @@
+package datadog.apm.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 INDEX_GENERATOR = new AtomicInteger(0);
+
+  // internal context keys
+
+  private final String name;
+  final int index;
+
+  private ContextKey(String name) {
+    this.name = name;
+    this.index = INDEX_GENERATOR.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 + '@' + index;
+  }
+}

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

@@ -0,0 +1,36 @@
+package datadog.apm.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 and returns the context attached to the current execution unit.
+   *
+   * @return Previously attached context; {@link #root()} if there was none
+   */
+  Context detach();
+}

components/context/src/main/java/datadog/apm/context/ContextProvider.java

@@ -0,0 +1,15 @@
+package datadog.apm.context;
+
+final class ContextProvider {
+  private static final ContextManager MANAGER = new DefaultContextManager();
+
+  private static final ContextBinder BINDER = new DefaultContextBinder();
+
+  static ContextManager manager() {
+    return MANAGER;
+  }
+
+  static ContextBinder binder() {
+    return BINDER;
+  }
+}

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

@@ -0,0 +1,12 @@
+package datadog.apm.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/apm/context/DefaultContextBinder.java

@@ -0,0 +1,27 @@
+package datadog.apm.context;
+
+import java.util.Collections;
+import java.util.Map;
+import java.util.WeakHashMap;
+
+final class DefaultContextBinder implements ContextBinder {
+  private static final Map<Object, Context> tracked =
+      Collections.synchronizedMap(new WeakHashMap<>());
+
+  @Override
+  public Context from(Object carrier) {
+    Context bound = tracked.get(carrier);
+    return null != bound ? bound : Context.root();
+  }
+
+  @Override
+  public void attachTo(Object carrier, Context context) {
+    tracked.put(carrier, context);
+  }
+
+  @Override
+  public Context detachFrom(Object carrier) {
+    Context previous = tracked.remove(carrier);
+    return null != previous ? previous : Context.root();
+  }
+}

components/context/src/main/java/datadog/apm/context/DefaultContextManager.java

@@ -0,0 +1,64 @@
+package datadog.apm.context;
+
+final class DefaultContextManager implements ContextManager {
+  private static final ThreadLocal<Context[]> current =
+      ThreadLocal.withInitial(() -> new Context[1]);
+
+  @Override
+  public Context root() {
+    return RootContext.INSTANCE;
+  }
+
+  @Override
+  public Context current() {
+    return current.get()[0];
+  }
+
+  @Override
+  public ContextScope attach(Context context) {
+
+    Context[] holder = current.get();
+    Context previous = holder[0];
+    holder[0] = context;
+
+    return new ContextScope() {
+      @Override
+      public Context context() {
+        return context;
+      }
+
+      @Override
+      public void close() {
+        holder[0] = previous;
+      }
+    };
+  }
+
+  @Override
+  public Context swap(Context context) {
+    Context[] holder = current.get();
+    Context previous = holder[0];
+    holder[0] = context;
+    return previous;
+  }
+
+  @Override
+  public Context detach() {
+    return swap(root());
+  }
+
+  @Override
+  public Context from(Object carrier) {
+    return root();
+  }
+
+  @Override
+  public boolean attachTo(Object carrier, Context context) {
+    return false;
+  }
+
+  @Override
+  public Context detachFrom(Object carrier) {
+    return root();
+  }
+}

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

@@ -0,0 +1,12 @@
+package datadog.apm.context;
+
+/** A {@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);
+}