feat(context): Add non default propagator registration (#8310)

PerfectSlayer · web-flow · commit dbbe55c183c8 · 2025-01-30T15:52:40.000Z
diff --git a/components/context/src/main/java/datadog/context/EmptyContext.java b/components/context/src/main/java/datadog/context/EmptyContext.java
@@ -24,4 +24,9 @@ public <T> Context with(ContextKey<T> key, @Nullable T value) {
     }
     return new SingletonContext(key.index, value);
   }
+
+  @Override
+  public String toString() {
+    return "EmptyContext{}";
+  }
 }
diff --git a/components/context/src/main/java/datadog/context/IndexedContext.java b/components/context/src/main/java/datadog/context/IndexedContext.java
@@ -23,14 +23,14 @@ final class IndexedContext implements Context {
   public <T> T get(ContextKey<T> key) {
     requireNonNull(key, "Context key cannot be null");
     int index = key.index;
-    return index < store.length ? (T) store[index] : null;
+    return index < this.store.length ? (T) this.store[index] : null;
   }
 
   @Override
   public <T> Context with(ContextKey<T> key, @Nullable T value) {
     requireNonNull(key, "Context key cannot be null");
     int index = key.index;
-    Object[] newStore = copyOfRange(store, 0, max(store.length, index + 1));
+    Object[] newStore = copyOfRange(this.store, 0, max(this.store.length, index + 1));
     newStore[index] = value;
     return new IndexedContext(newStore);
   }
@@ -40,13 +40,18 @@ 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);
+    return Arrays.equals(this.store, that.store);
   }
 
   @Override
   public int hashCode() {
     int result = 31;
-    result = 31 * result + Arrays.hashCode(store);
+    result = 31 * result + Arrays.hashCode(this.store);
     return result;
   }
+
+  @Override
+  public String toString() {
+    return "IndexedContext{store=" + Arrays.toString(this.store) + '}';
+  }
 }
diff --git a/components/context/src/main/java/datadog/context/SingletonContext.java b/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 + '}';
+  }
 }
diff --git a/components/context/src/main/java/datadog/context/propagation/CarrierSetter.java b/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);
 }
diff --git a/components/context/src/main/java/datadog/context/propagation/Propagators.java b/components/context/src/main/java/datadog/context/propagation/Propagators.java
@@ -6,11 +6,17 @@
 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 =
+      RegisteredPropagator.of(NoopPropagator.INSTANCE, false);
   private static volatile Propagator defaultPropagator = null;
-  private static volatile boolean defaultPropagatorSet = false;
+  private static volatile boolean rebuildDefaultPropagator = true;
 
   private Propagators() {}
 
@@ -20,14 +26,16 @@ private Propagators() {}
    * @return The default propagator.
    */
   public static Propagator defaultPropagator() {
-    if (!defaultPropagatorSet) {
+    if (rebuildDefaultPropagator) {
       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;
+      rebuildDefaultPropagator = false;
     }
     return defaultPropagator;
   }
@@ -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, RegisteredPropagator.of(propagator, usedAsDefault));
+    if (usedAsDefault) {
+      rebuildDefaultPropagator = true;
+    }
   }
 
   /** Clear all registered propagators. For testing purpose only. */
   static void reset() {
     PROPAGATORS.clear();
-    defaultPropagatorSet = false;
+    rebuildDefaultPropagator = true;
+  }
+
+  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);
+    }
+
+    Propagator propagator() {
+      return this.propagator;
+    }
+
+    boolean isUsedAsDefault() {
+      return this.usedAsDefault;
+    }
   }
 }
diff --git a/components/context/src/test/java/datadog/context/ContextTest.java b/components/context/src/test/java/datadog/context/ContextTest.java
@@ -5,6 +5,7 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
@@ -154,6 +155,16 @@ void testEqualsAndHashCode() {
     assertNotEquals(context4.hashCode(), context1.hashCode());
   }
 
+  @ParameterizedTest
+  @MethodSource("contextImplementations")
+  void testToString(Context context) {
+    String debugString = context.toString();
+    assertNotNull(debugString, "Context string representation should not be null");
+    assertTrue(
+        debugString.contains(context.getClass().getSimpleName()),
+        "Context string representation should contain implementation name");
+  }
+
   @SuppressWarnings({"SimplifiableAssertion"})
   @Test
   void testInflation() {
diff --git a/components/context/src/test/java/datadog/context/propagation/PropagatorsTest.java b/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");
+    }
+  }
 }

Original file line number	Diff line number	Diff line change
`@@ -24,4 +24,9 @@ public <T> Context with(ContextKey<T> key, @Nullable T value) {`
`24`	`24`	`}`
`25`	`25`	`return new SingletonContext(key.index, value);`
`26`	`26`	`}`
	`27`	`+`
	`28`	`+ @Override`
	`29`	`+ public String toString() {`
	`30`	`+ return "EmptyContext{}";`
	`31`	`+ }`
`27`	`32`	`}`
Original file line number	Diff line number	Diff line change
`@@ -57,4 +57,9 @@ public int hashCode() {`
`57`	`57`	`result = 31 * result + Objects.hashCode(this.value);`
`58`	`58`	`return result;`
`59`	`59`	`}`
	`60`	`+`
	`61`	`+ @Override`
	`62`	`+ public String toString() {`
	`63`	`+ return "SingletonContext{" + "index=" + this.index + ", value=" + this.value + '}';`
	`64`	`+ }`
`60`	`65`	`}`