Skip to content

Commit

Permalink
feat(context): Add non default propagator registration
Browse files Browse the repository at this point in the history
  • Loading branch information
PerfectSlayer committed Jan 30, 2025
1 parent d3cfbcb commit 41e2d51
Show file tree
Hide file tree
Showing 4 changed files with 72 additions and 13 deletions.
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() {
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
@@ -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;

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

0 comments on commit 41e2d51

Please sign in to comment.