From 6f3b8dd704584f52623a949ea4a6a049a4e752eb Mon Sep 17 00:00:00 2001 From: Bruce Bujon Date: Tue, 28 Jan 2025 17:57:15 +0100 Subject: [PATCH] feat(context): Add non default propagator registration --- .../java/datadog/context/EmptyContext.java | 5 ++ .../java/datadog/context/IndexedContext.java | 13 ++-- .../datadog/context/SingletonContext.java | 5 ++ .../context/propagation/CarrierSetter.java | 4 +- .../context/propagation/Propagators.java | 60 ++++++++++++++++--- .../java/datadog/context/ContextTest.java | 11 ++++ .../context/propagation/PropagatorsTest.java | 24 ++++++-- 7 files changed, 101 insertions(+), 21 deletions(-) diff --git a/components/context/src/main/java/datadog/context/EmptyContext.java b/components/context/src/main/java/datadog/context/EmptyContext.java index fc810c757eb..20023482647 100644 --- 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 Context with(ContextKey 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 index cadc481d707..c3c7684465d 100644 --- 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 get(ContextKey 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 Context with(ContextKey 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 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..d29867251b5 100644 --- 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 PROPAGATORS = + private static final Map 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 index 52cc7da0371..8bf69b645cb 100644 --- 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 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"); + } + } }