Skip to content

Add non default propagator registration #8310

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jan 30, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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{}";
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Expand All @@ -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) + '}';
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -57,4 +57,9 @@ public int hashCode() {
result = 31 * result + Objects.hashCode(this.value);
return result;
}

@Override
public String toString() {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you add toString to the empty and indexed cases for consistency? thanks

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sure, I originally kept those change in stash but I will happen them with their tests in this commit.

return "SingletonContext{" + "index=" + this.index + ", value=" + this.value + '}';
}
}
Original file line number Diff line number Diff line change
@@ -1,7 +1,5 @@
package datadog.context.propagation;

import javax.annotation.Nullable;

@FunctionalInterface
public interface CarrierSetter<C> {
/**
Expand All @@ -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);
}
Original file line number Diff line number Diff line change
Expand Up @@ -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() {}

Expand All @@ -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;
}
Expand All @@ -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();
}

/**
Expand Down Expand Up @@ -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;
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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() {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
}
}
Expand Down Expand Up @@ -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");
Expand Down Expand Up @@ -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");
}
}
}
Loading