diff --git a/components/context/src/main/java/datadog/context/SingletonContext.java b/components/context/src/main/java/datadog/context/SingletonContext.java index 7a8a4e98b6f..605dbdecd28 100644 --- 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 index f4b765fb4fb..76a6abda835 100644 --- 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 { /** @@ -11,5 +9,5 @@ public interface CarrierSetter { * @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 index 96ac069692a..3b6a4584f07 100644 --- a/components/context/src/main/java/datadog/context/propagation/Propagators.java +++ b/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 PROPAGATORS = + private static final Map 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,8 +97,22 @@ 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. */ @@ -98,4 +120,26 @@ 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; + } + } } diff --git a/components/context/src/test/java/datadog/context/propagation/PropagatorsTest.java b/components/context/src/test/java/datadog/context/propagation/PropagatorsTest.java index 21d369c8f73..a8611014ede 100644 --- a/components/context/src/test/java/datadog/context/propagation/PropagatorsTest.java +++ b/components/context/src/test/java/datadog/context/propagation/PropagatorsTest.java @@ -6,6 +6,7 @@ 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; @@ -13,7 +14,7 @@ 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>, CarrierVisitor> { @Override - public void set(@Nullable Map carrier, String key, String value) { - if (carrier != null && key != null) { + public void set(Map 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 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"); + } + } }