feat(context): Add non default propagator registration

Author: PerfectSlayer

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

@@ -57,4 +57,9 @@ public int hashCode() {
     result = 31 * result + Objects.hashCode(this.value);
     return result;
   }
+
+  @Override
+  public String toString() {
+    return "SingletonContext{" + "index=" + this.index + ", value=" + this.value + '}';
+  }
 }

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

@@ -1,7 +1,5 @@
 package datadog.context.propagation;
 
-import javax.annotation.Nullable;
-
 @FunctionalInterface
 public interface CarrierSetter<C> {
   /**
@@ -11,5 +9,5 @@ public interface CarrierSetter<C> {
    * @param key the key to set.
    * @param value the value to set.
    */
-  void set(@Nullable C carrier, String key, String value);
+  void set(C carrier, String key, String value);
 }

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

@@ -1,14 +1,20 @@
 package datadog.context.propagation;
 
+import static datadog.context.propagation.Propagators.RegisteredPropagator.of;
 import static java.util.Collections.synchronizedMap;
 import static java.util.Comparator.comparingInt;
 
 import java.util.IdentityHashMap;
 import java.util.Map;
 
+/**
+ * This class is the entrypoint of the context propagation API allowing to retrieve the {@link
+ * Propagator} to use.
+ */
 public final class Propagators {
-  private static final Map<Concern, Propagator> PROPAGATORS =
+  private static final Map<Concern, RegisteredPropagator> PROPAGATORS =
       synchronizedMap(new IdentityHashMap<>());
+  private static final RegisteredPropagator NOOP = of(NoopPropagator.INSTANCE, false);
   private static volatile Propagator defaultPropagator = null;
   private static volatile boolean defaultPropagatorSet = false;
 
@@ -23,8 +29,10 @@ public static Propagator defaultPropagator() {
     if (!defaultPropagatorSet) {
       Propagator[] propagatorsByPriority =
           PROPAGATORS.entrySet().stream()
+              .filter(entry -> entry.getValue().isUsedAsDefault())
               .sorted(comparingInt(entry -> entry.getKey().priority()))
               .map(Map.Entry::getValue)
+              .map(RegisteredPropagator::propagator)
               .toArray(Propagator[]::new);
       defaultPropagator = composite(propagatorsByPriority);
       defaultPropagatorSet = true;
@@ -39,7 +47,7 @@ public static Propagator defaultPropagator() {
    * @return the related propagator if registered, a {@link #noop()} propagator otherwise.
    */
   public static Propagator forConcern(Concern concern) {
-    return PROPAGATORS.getOrDefault(concern, NoopPropagator.INSTANCE);
+    return PROPAGATORS.getOrDefault(concern, NOOP).propagator();
   }
 
   /**
@@ -89,13 +97,49 @@ public static Propagator composite(Propagator... propagators) {
    * @param propagator The propagator to register.
    */
   public static void register(Concern concern, Propagator propagator) {
-    PROPAGATORS.put(concern, propagator);
-    defaultPropagatorSet = false;
+    register(concern, propagator, true);
+  }
+
+  /**
+   * Registers a propagator for concern.
+   *
+   * @param concern The concern to register a propagator for.
+   * @param propagator The propagator to register.
+   * @param usedAsDefault Whether the propagator should be used as default propagator.
+   * @see Propagators#defaultPropagator()
+   */
+  public static void register(Concern concern, Propagator propagator, boolean usedAsDefault) {
+    PROPAGATORS.put(concern, of(propagator, usedAsDefault));
+    if (usedAsDefault) {
+      defaultPropagatorSet = false;
+    }
   }
 
   /** Clear all registered propagators. For testing purpose only. */
   static void reset() {
     PROPAGATORS.clear();
     defaultPropagatorSet = false;
   }
+
+  static class RegisteredPropagator {
+    private final Propagator propagator;
+    private final boolean usedAsDefault;
+
+    private RegisteredPropagator(Propagator propagator, boolean usedAsDefault) {
+      this.propagator = propagator;
+      this.usedAsDefault = usedAsDefault;
+    }
+
+    static RegisteredPropagator of(Propagator propagator, boolean useAsDefault) {
+      return new RegisteredPropagator(propagator, useAsDefault);
+    }
+
+    protected Propagator propagator() {
+      return this.propagator;
+    }
+
+    boolean isUsedAsDefault() {
+      return this.usedAsDefault;
+    }
+  }
 }

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

@@ -6,14 +6,15 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import datadog.context.Context;
 import datadog.context.ContextKey;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.function.BiConsumer;
-import javax.annotation.Nullable;
+import javax.annotation.ParametersAreNonnullByDefault;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -46,11 +47,12 @@ class PropagatorsTest {
           .with(DEBUGGER_KEY, "debug")
           .with(PROFILING_KEY, "profile");
 
+  @ParametersAreNonnullByDefault
   static class MapCarrierAccessor
       implements CarrierSetter<Map<String, String>>, CarrierVisitor<Map<String, String>> {
     @Override
-    public void set(@Nullable Map<String, String> carrier, String key, String value) {
-      if (carrier != null && key != null) {
+    public void set(Map<String, String> carrier, String key, String value) {
+      if (carrier != null && key != null && value != null) {
         carrier.put(key, value);
       }
     }
@@ -136,12 +138,12 @@ void testDefaultPropagator() {
     Propagator single = Propagators.defaultPropagator();
     assertInjectExtractContext(CONTEXT, single, TRACING_KEY);
 
-    Propagators.register(IAST, IAST_PROPAGATOR);
+    Propagators.register(IAST, IAST_PROPAGATOR, false);
     Propagators.register(DEBUGGER, DEBUGGER_PROPAGATOR);
     Propagators.register(PROFILING, PROFILING_PROPAGATOR);
     Propagator composite = Propagators.defaultPropagator();
-    assertInjectExtractContext(
-        CONTEXT, composite, TRACING_KEY, IAST_KEY, DEBUGGER_KEY, PROFILING_KEY);
+    assertDoNotInjectExtractContext(CONTEXT, composite, IAST_KEY);
+    assertInjectExtractContext(CONTEXT, composite, TRACING_KEY, DEBUGGER_KEY, PROFILING_KEY);
     assertFalse(
         DEBUGGER_PROPAGATOR.keyFound,
         "Debugger propagator should have run before tracing propagator");
@@ -219,4 +221,14 @@ void assertInjectExtractContext(Context context, Propagator propagator, ContextK
           context.get(key), extracted.get(key), "Key " + key + " not injected nor extracted");
     }
   }
+
+  private void assertDoNotInjectExtractContext(
+      Context context, Propagator propagator, ContextKey<?>... keys) {
+    Map<String, String> carrier = new HashMap<>();
+    propagator.inject(context, carrier, ACCESSOR);
+    Context extracted = propagator.extract(root(), carrier, ACCESSOR);
+    for (ContextKey<?> key : keys) {
+      assertNull(extracted.get(key), "Key " + key + " was injected");
+    }
+  }
 }