From 94b73e5803636b163dc84297cc00dff03682b634 Mon Sep 17 00:00:00 2001
From: Andrea Marziali <andrea.marziali@datadoghq.com>
Date: Wed, 27 May 2026 11:45:00 +0200
Subject: [PATCH] Add lazy-transformer to KeyClassifier for deferred header
 decoding

---
 .../instrumentation/api/AgentPropagation.java | 19 +++++
 .../api/KeyClassifierTest.java                | 78 +++++++++++++++++++
 2 files changed, 97 insertions(+)
 create mode 100644 internal-api/src/test/java/datadog/trace/bootstrap/instrumentation/api/KeyClassifierTest.java

diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentPropagation.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentPropagation.java
index 586b2f08612..094dc7e9fbd 100644
--- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentPropagation.java
+++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentPropagation.java
@@ -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 {
@@ -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> {
diff --git a/internal-api/src/test/java/datadog/trace/bootstrap/instrumentation/api/KeyClassifierTest.java b/internal-api/src/test/java/datadog/trace/bootstrap/instrumentation/api/KeyClassifierTest.java
new file mode 100644
index 00000000000..35e085a7892
--- /dev/null
+++ b/internal-api/src/test/java/datadog/trace/bootstrap/instrumentation/api/KeyClassifierTest.java
@@ -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);
+  }
+}