Introduce context propagation API (#8161)

* feat(context): Introduce context propagation API

* feat(context): Clean context API

* feat(context): Add API test coverage

* feat(context): Introduce compositor extractor carrier cache

* feat(context): Improve propagators collection

Author: PerfectSlayer

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

@@ -4,7 +4,6 @@
 import static datadog.context.ContextProviders.manager;
 
 import javax.annotation.Nullable;
-import javax.annotation.ParametersAreNonnullByDefault;
 
 /**
  * Immutable context scoped to an execution unit or carrier object.
@@ -36,7 +35,6 @@
  *
  * @see ContextKey
  */
-@ParametersAreNonnullByDefault
 public interface Context {
   /**
    * Returns the root context.

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

@@ -1,9 +1,6 @@
 package datadog.context;
 
-import javax.annotation.ParametersAreNonnullByDefault;
-
 /** Binds context to carrier objects. */
-@ParametersAreNonnullByDefault
 public interface ContextBinder {
   /**
    * Returns the context attached to the given carrier object.

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

@@ -10,7 +10,7 @@
  */
 public final class ContextKey<T> {
   private static final AtomicInteger NEXT_INDEX = new AtomicInteger(0);
-  /** The key name, for debugging purpose only . */
+  /** The key name, for debugging purpose only. */
   private final String name;
   /** The key unique context, related to {@link IndexedContext} implementation. */
   final int index;

components/context/src/main/java/datadog/context/package-info.java

@@ -0,0 +1,4 @@
+@ParametersAreNonnullByDefault
+package datadog.context;
+
+import javax.annotation.ParametersAreNonnullByDefault;

components/context/src/main/java/datadog/context/propagation/CarrierSetter.java

@@ -0,0 +1,15 @@
+package datadog.context.propagation;
+
+import javax.annotation.Nullable;
+
+@FunctionalInterface
+public interface CarrierSetter<C> {
+  /**
+   * Sets a carrier key/value pair.
+   *
+   * @param carrier the carrier to store key/value into.
+   * @param key the key to set.
+   * @param value the value to set.
+   */
+  void set(@Nullable C carrier, String key, String value);
+}

components/context/src/main/java/datadog/context/propagation/CarrierVisitor.java

@@ -0,0 +1,26 @@
+package datadog.context.propagation;
+
+import java.util.function.BiConsumer;
+
+/**
+ * This interface represents the capacity of walking through a carrier content, iterating over its
+ * key/value pairs.
+ *
+ * <p>Walking through carrier is preferred to direct access to carrier key/value pairs as some
+ * carrier implementations do not have built-in direct access and require walking over the full
+ * carrier structure to find the requested key/value pair, leading to multiple walks when multiple
+ * keys are requested, whereas the visitor is expected to walk through only once, and the
+ * propagators to keep the data they need using the visitor callback.
+ *
+ * @param <C> the type of carrier.
+ */
+@FunctionalInterface
+public interface CarrierVisitor<C> {
+  /**
+   * Iterates over the carrier content and calls the visitor callback for every key/value found.
+   *
+   * @param carrier the carrier to iterate over.
+   * @param visitor the callback to call for each carrier key/value pair found.
+   */
+  void forEachKeyValue(C carrier, BiConsumer<String, String> visitor);
+}

components/context/src/main/java/datadog/context/propagation/CompositePropagator.java

@@ -0,0 +1,55 @@
+package datadog.context.propagation;
+
+import datadog.context.Context;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.function.BiConsumer;
+
+class CompositePropagator implements Propagator {
+  private final Propagator[] propagators;
+
+  CompositePropagator(Propagator[] propagators) {
+    this.propagators = propagators;
+  }
+
+  @Override
+  public <C> void inject(Context context, C carrier, CarrierSetter<C> setter) {
+    for (Propagator propagator : this.propagators) {
+      propagator.inject(context, carrier, setter);
+    }
+  }
+
+  @Override
+  public <C> Context extract(Context context, C carrier, CarrierVisitor<C> visitor) {
+    // Extract and cache carrier key/value pairs
+    CarrierCache carrierCache = new CarrierCache();
+    visitor.forEachKeyValue(carrier, carrierCache);
+    // Run the multiple extractions on cache
+    for (Propagator propagator : this.propagators) {
+      context = propagator.extract(context, carrierCache, carrierCache);
+    }
+    return context;
+  }
+
+  static class CarrierCache implements BiConsumer<String, String>, CarrierVisitor<CarrierCache> {
+    /** Cached key/values from carrier (even indexes are keys, odd indexes are values). */
+    private final List<String> keysAndValues;
+
+    public CarrierCache() {
+      this.keysAndValues = new ArrayList<>(32);
+    }
+
+    @Override
+    public void accept(String key, String value) {
+      this.keysAndValues.add(key);
+      this.keysAndValues.add(value);
+    }
+
+    @Override
+    public void forEachKeyValue(CarrierCache carrier, BiConsumer<String, String> visitor) {
+      for (int i = 0; i < carrier.keysAndValues.size() - 1; i += 2) {
+        visitor.accept(carrier.keysAndValues.get(i), carrier.keysAndValues.get(i + 1));
+      }
+    }
+  }
+}

components/context/src/main/java/datadog/context/propagation/Concern.java

@@ -0,0 +1,57 @@
+package datadog.context.propagation;
+
+import static java.util.Objects.requireNonNull;
+
+import datadog.context.Context;
+
+/** This class defines a cross-cutting concern to be propagated from a {@link Context}. */
+public class Concern {
+  /** The concern default priority. */
+  public static final int DEFAULT_PRIORITY = 100;
+  /** The concern name, for debugging purpose only. */
+  private final String name;
+  /** The concern priority, lower value means higher priority. */
+  private final int priority;
+
+  /**
+   * Creates a concern.
+   *
+   * @param name the concern name, for debugging purpose only.
+   * @return The created concern.
+   */
+  public static Concern named(String name) {
+    return new Concern(name, DEFAULT_PRIORITY);
+  }
+
+  /**
+   * Creates a concern with a specific priority.
+   *
+   * @param name the concern name, for debugging purpose only.
+   * @param priority the concern priority (lower value means higher priority, while the default is
+   *     {@link #DEFAULT_PRIORITY}),
+   * @return The created concern.
+   */
+  public static Concern withPriority(String name, int priority) {
+    return new Concern(name, priority);
+  }
+
+  private Concern(String name, int priority) {
+    requireNonNull(name, "Concern name cannot be null");
+    if (priority < 0) {
+      throw new IllegalArgumentException("Concern priority cannot be negative");
+    }
+    this.name = name;
+    this.priority = priority;
+  }
+
+  int priority() {
+    return this.priority;
+  }
+
+  // We want identity equality, so no need to override equals().
+
+  @Override
+  public String toString() {
+    return this.name;
+  }
+}

components/context/src/main/java/datadog/context/propagation/NoopPropagator.java

@@ -0,0 +1,19 @@
+package datadog.context.propagation;
+
+import datadog.context.Context;
+
+final class NoopPropagator implements Propagator {
+  static final NoopPropagator INSTANCE = new NoopPropagator();
+
+  private NoopPropagator() {}
+
+  @Override
+  public <C> void inject(Context context, C carrier, CarrierSetter<C> setter) {
+    // noop
+  }
+
+  @Override
+  public <C> Context extract(Context context, C carrier, CarrierVisitor<C> visitor) {
+    return context;
+  }
+}

components/context/src/main/java/datadog/context/propagation/Propagator.java

@@ -0,0 +1,35 @@
+package datadog.context.propagation;
+
+import datadog.context.Context;
+
+/**
+ * This interface represents a {@link Context} propagator for a given {@link Concern}.
+ *
+ * <p>Its goal is to {@link #inject} context values into carriers, or {@link #extract} them from
+ * carriers to populate context. Carrier could be any kind of object that stores key/value pairs,
+ * like HTTP or messages headers. {@link CarrierSetter}s and {@link CarrierVisitor}s define how to
+ * store and retrieve key/value pairs from carriers.
+ */
+public interface Propagator {
+  /**
+   * Injects a context into a downstream service using the given carrier.
+   *
+   * @param context the context containing the values to be injected.
+   * @param carrier the instance that will receive the key/value pairs to propagate.
+   * @param setter the callback to set key/value pairs into the carrier.
+   * @param <C> the type of carrier instance.
+   */
+  <C> void inject(Context context, C carrier, CarrierSetter<C> setter);
+
+  /**
+   * Extracts a context from un upstream service.
+   *
+   * @param context the base context to store the extracted values on top, use {@link
+   *     Context#root()} for a default base context.
+   * @param carrier the instance to fetch the propagated key/value pairs from.
+   * @param visitor the callback to walk over the carrier and extract its key/value pais.
+   * @param <C> the type of the carrier.
+   * @return A context with the extracted values on top of the given base context.
+   */
+  <C> Context extract(Context context, C carrier, CarrierVisitor<C> visitor);
+}