Skip to content
Open
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 @@ -10,6 +10,7 @@
import datadog.context.propagation.Concern;
import datadog.context.propagation.Propagators;
import java.util.function.BiConsumer;
import java.util.function.Function;
import javax.annotation.ParametersAreNonnullByDefault;

public final class AgentPropagation {
Expand Down Expand Up @@ -38,6 +39,24 @@ public static <C> AgentSpanContext.Extracted extractContextAndGetSpanContext(

public interface KeyClassifier {
boolean accept(String key, String value);

/**
* Variant of {@link #accept(String, String)} for carriers that store header values in a raw
* form (e.g. {@code byte[]}) and want to defer string conversion until after the key is known
* to be relevant.
*
* <p>The default implementation applies {@code transformer} eagerly and delegates to {@link
* #accept(String, String)}, so existing classifiers work without any changes.
*
* @param key the header name
* @param value the raw header value, in whatever form the carrier provides
* @param transformer converts {@code value} to a string; called at most once by the default
* implementation
* @return {@code false} to stop iteration, {@code true} to continue
*/
default <T> boolean accept(String key, T value, Function<T, String> transformer) {
return accept(key, transformer.apply(value));
}
}

public interface ContextVisitor<C> extends CarrierVisitor<C> {
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,78 @@
package datadog.trace.bootstrap.instrumentation.api;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertTrue;

import java.nio.charset.StandardCharsets;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;
import org.junit.jupiter.api.Test;

class KeyClassifierTest {

static class RecordingClassifier implements AgentPropagation.KeyClassifier {
String lastKey;
String lastValue;
boolean returnValue;

RecordingClassifier(boolean returnValue) {
this.returnValue = returnValue;
}

@Override
public boolean accept(String key, String value) {
lastKey = key;
lastValue = value;
return returnValue;
}
}

@Test
void defaultTransformerMethodAppliesTransformerAndDelegates() {
RecordingClassifier classifier = new RecordingClassifier(true);

boolean result =
classifier.accept(
"my-key",
"raw".getBytes(StandardCharsets.UTF_8),
bytes -> new String(bytes, StandardCharsets.UTF_8));

assertEquals("my-key", classifier.lastKey);
assertEquals("raw", classifier.lastValue);
assertTrue(result);
}

@Test
void transformerIsCalledExactlyOnce() {
AtomicInteger callCount = new AtomicInteger(0);
AtomicReference<String> transformed = new AtomicReference<>();

AgentPropagation.KeyClassifier classifier =
(key, value) -> {
transformed.set(value);
return true;
};

classifier.accept(
"key",
"input",
v -> {
callCount.incrementAndGet();
return v.toUpperCase();
});

assertEquals(1, callCount.get());
assertEquals("INPUT", transformed.get());
}

@Test
void existingAcceptStringStringContractUnchanged() {
RecordingClassifier classifier = new RecordingClassifier(true);

boolean result = classifier.accept("trace-id", "abc123");

assertEquals("trace-id", classifier.lastKey);
assertEquals("abc123", classifier.lastValue);
assertTrue(result);
}
}
Loading