Introduce context propagation API

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);
+}

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

@@ -0,0 +1,101 @@
+package datadog.context.propagation;
+
+import static java.util.Collections.synchronizedMap;
+import static java.util.Comparator.comparingInt;
+
+import java.util.IdentityHashMap;
+import java.util.Map;
+
+public final class Propagators {
+  private static final Map<Concern, Propagator> PROPAGATORS =
+      synchronizedMap(new IdentityHashMap<>());
+  private static volatile Propagator defaultPropagator = null;
+  private static volatile boolean defaultPropagatorSet = false;
+
+  private Propagators() {}
+
+  /**
+   * Gets the default propagator that applies all registered propagators in their priority order.
+   *
+   * @return The default propagator.
+   */
+  public static Propagator defaultPropagator() {
+    if (!defaultPropagatorSet) {
+      Propagator[] propagatorsByPriority =
+          PROPAGATORS.entrySet().stream()
+              .sorted(comparingInt(entry -> entry.getKey().priority()))
+              .map(Map.Entry::getValue)
+              .toArray(Propagator[]::new);
+      defaultPropagator = composite(propagatorsByPriority);
+      defaultPropagatorSet = true;
+    }
+    return defaultPropagator;
+  }
+
+  /**
+   * Gets the propagator for a given concern.
+   *
+   * @param concern the concern to get propagator for.
+   * @return the related propagator if registered, a {@link #noop()} propagator otherwise.
+   */
+  public static Propagator forConcern(Concern concern) {
+    return PROPAGATORS.getOrDefault(concern, NoopPropagator.INSTANCE);
+  }
+
+  /**
+   * Gets the propagator for the given concerns.
+   *
+   * @param concerns the concerns to get propagators for.
+   * @return A propagator that will apply the concern propagators if registered, in the given
+   *     concern order.
+   */
+  public static Propagator forConcerns(Concern... concerns) {
+    Propagator[] propagators = new Propagator[concerns.length];
+    for (int i = 0; i < concerns.length; i++) {
+      propagators[i] = forConcern(concerns[i]);
+    }
+    return composite(propagators);
+  }
+
+  /**
+   * Returns a noop propagator.
+   *
+   * @return a noop propagator.
+   */
+  public static Propagator noop() {
+    return NoopPropagator.INSTANCE;
+  }
+
+  /**
+   * Creates a composite propagator.
+   *
+   * @param propagators the elements that composes the returned propagator.
+   * @return the composite propagator that will apply the propagators in their given order.
+   */
+  public static Propagator composite(Propagator... propagators) {
+    if (propagators.length == 0) {
+      return NoopPropagator.INSTANCE;
+    } else if (propagators.length == 1) {
+      return propagators[0];
+    } else {
+      return new CompositePropagator(propagators);
+    }
+  }
+
+  /**
+   * Registers a propagator for concern.
+   *
+   * @param concern The concern to register a propagator for.
+   * @param propagator The propagator to register.
+   */
+  public static void register(Concern concern, Propagator propagator) {
+    PROPAGATORS.put(concern, propagator);
+    defaultPropagatorSet = false;
+  }
+
+  /** Clear all registered propagators. For testing purpose only. */
+  static void reset() {
+    PROPAGATORS.clear();
+    defaultPropagatorSet = false;
+  }
+}

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

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

components/context/src/test/java/datadog/context/propagation/ConcernTest.java

@@ -0,0 +1,38 @@
+package datadog.context.propagation;
+
+import static datadog.context.propagation.Concern.DEFAULT_PRIORITY;
+import static org.junit.jupiter.api.Assertions.*;
+
+import org.junit.jupiter.api.Test;
+
+class ConcernTest {
+  @Test
+  void testNamed() {
+    assertThrows(
+        NullPointerException.class,
+        () -> Concern.named(null),
+        "Should not create null named concern");
+    assertNotNull(Concern.named("name"));
+  }
+
+  @Test
+  void testWithPriority() {
+    assertThrows(
+        NullPointerException.class,
+        () -> Concern.withPriority(null, DEFAULT_PRIORITY),
+        "Should not create null named concern");
+    assertThrows(
+        IllegalArgumentException.class,
+        () -> Concern.withPriority("name", -1),
+        "Should not create negative priority concern");
+    assertNotNull(Concern.withPriority("high-priority", DEFAULT_PRIORITY - 10));
+    assertNotNull(Concern.withPriority("low-priority", DEFAULT_PRIORITY + 10));
+  }
+
+  @Test
+  void testName() {
+    String debugName = "name";
+    Concern concern = Concern.named(debugName);
+    assertEquals(debugName, concern.toString(), "Concern name mismatch");
+  }
+}