feat(tt): tag server spans with _dd.tt.extraction_sources from RC patterns

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java

@@ -26,12 +26,14 @@
 import datadog.trace.api.gateway.RequestContext;
 import datadog.trace.api.gateway.RequestContextSlot;
 import datadog.trace.api.naming.SpanNaming;
+import datadog.trace.api.tt.TransactionTrackingCandidateSources;
 import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.bootstrap.instrumentation.api.ClientIpAddressData;
 import datadog.trace.bootstrap.instrumentation.api.ErrorPriorities;
+import datadog.trace.bootstrap.instrumentation.api.InstrumentationTags;
 import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes;
 import datadog.trace.bootstrap.instrumentation.api.ResourceNamePriorities;
 import datadog.trace.bootstrap.instrumentation.api.TagContext;
@@ -44,8 +46,10 @@
 import java.util.BitSet;
 import java.util.Locale;
 import java.util.Map;
+import java.util.TreeSet;
 import java.util.concurrent.ExecutionException;
 import java.util.function.BiFunction;
+import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Supplier;
 import javax.annotation.Nonnull;
@@ -119,6 +123,17 @@ protected String getRequestHeader(REQUEST request, String key) {
     return null;
   }
 
+  /**
+   * Iterates the names of every inbound HTTP request header, invoking {@code consumer} once per
+   * name. Default no-op implementation: subclasses with cheap access to the underlying request's
+   * header enumeration should override this so the Transaction Tracking candidate-sources tag works
+   * for that stack. Used only when {@link TransactionTrackingCandidateSources#isEmpty()} returns
+   * false.
+   */
+  protected void forEachRequestHeaderName(REQUEST request, Consumer<String> consumer) {
+    // no-op: stacks without cheap header enumeration silently produce no tag.
+  }
+
   protected String requestedSessionId(REQUEST request) {
     return null;
   }
@@ -344,6 +359,15 @@ public AgentSpan onRequest(
       } catch (final Exception e) {
         log.debug("Error tagging url", e);
       }
+      // Transaction Tracking: tag span with matching header / query-param names when the
+      // remote-config snapshot is non-empty. Fast path is a single volatile read + isEmpty().
+      if (!TransactionTrackingCandidateSources.isEmpty()) {
+        try {
+          tagCandidateSources(span, request);
+        } catch (Exception e) {
+          log.debug("Error tagging tt candidate sources", e);
+        }
+      }
     }
 
     String peerIp = null;
@@ -420,6 +444,111 @@ public AgentSpan onRequest(
     return span;
   }
 
+  /**
+   * Adds the {@code _dd.tt.candidate_sources} tag based on the currently active {@link
+   * TransactionTrackingCandidateSources} snapshot. Caller must have already verified that the
+   * snapshot is non-empty.
+   *
+   * <p>The tag value is a CSV with deterministic ordering: {@code header:} entries (sorted), then
+   * {@code qs:} entries (sorted). Names are lowercased and de-duplicated within each bucket. The
+   * tag is only set if at least one match is found.
+   */
+  private void tagCandidateSources(AgentSpan span, REQUEST request) {
+    if (request == null) {
+      return;
+    }
+    TreeSet<String> headerHits = null;
+    TreeSet<String> qsHits = null;
+
+    // 1. Header names.
+    HeaderNameCollector<REQUEST> collector = new HeaderNameCollector<>();
+    forEachRequestHeaderName(request, collector);
+    if (collector.matches != null) {
+      headerHits = collector.matches;
+    }
+
+    // 2. Query-string parameter names. Re-resolve the URL adapter so this code path does not
+    // depend on whether the URL block above succeeded.
+    try {
+      URIDataAdapter url = url(request);
+      String rawQuery = url == null ? null : url.rawQuery();
+      if (rawQuery != null && !rawQuery.isEmpty()) {
+        qsHits = collectQueryParameterMatches(rawQuery);
+      }
+    } catch (Exception e) {
+      log.debug("Error resolving URL for tt candidate sources", e);
+    }
+
+    if (headerHits == null && qsHits == null) {
+      return;
+    }
+    StringBuilder sb = new StringBuilder();
+    if (headerHits != null) {
+      for (String name : headerHits) {
+        if (sb.length() > 0) {
+          sb.append(',');
+        }
+        sb.append("header:").append(name);
+      }
+    }
+    if (qsHits != null) {
+      for (String name : qsHits) {
+        if (sb.length() > 0) {
+          sb.append(',');
+        }
+        sb.append("qs:").append(name);
+      }
+    }
+    if (sb.length() > 0) {
+      span.setTag(InstrumentationTags.TT_CANDIDATE_SOURCES, sb.toString());
+    }
+  }
+
+  private static TreeSet<String> collectQueryParameterMatches(String rawQuery) {
+    TreeSet<String> hits = null;
+    int len = rawQuery.length();
+    int start = 0;
+    while (start <= len) {
+      int amp = rawQuery.indexOf('&', start);
+      int end = amp < 0 ? len : amp;
+      if (end > start) {
+        int eq = rawQuery.indexOf('=', start);
+        int nameEnd = (eq < 0 || eq > end) ? end : eq;
+        if (nameEnd > start) {
+          String name = rawQuery.substring(start, nameEnd);
+          if (TransactionTrackingCandidateSources.matchesAny(name)) {
+            if (hits == null) {
+              hits = new TreeSet<>();
+            }
+            hits.add(name.toLowerCase(Locale.ROOT));
+          }
+        }
+      }
+      if (amp < 0) {
+        break;
+      }
+      start = amp + 1;
+    }
+    return hits;
+  }
+
+  private static final class HeaderNameCollector<R> implements Consumer<String> {
+    TreeSet<String> matches;
+
+    @Override
+    public void accept(String name) {
+      if (name == null) {
+        return;
+      }
+      if (TransactionTrackingCandidateSources.matchesAny(name)) {
+        if (matches == null) {
+          matches = new TreeSet<>();
+        }
+        matches.add(name.toLowerCase(Locale.ROOT));
+      }
+    }
+  }
+
   protected static AgentSpanContext.Extracted getExtractedSpanContext(Context parentContext) {
     AgentSpan extractedSpan = AgentSpan.fromContext(parentContext);
     if (extractedSpan != null) {

dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecoratorTtCandidateSourcesTest.groovy

@@ -0,0 +1,181 @@
+package datadog.trace.bootstrap.instrumentation.decorator
+
+import datadog.trace.api.tt.TransactionTrackingCandidateSources
+import datadog.trace.bootstrap.instrumentation.api.AgentPropagation
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer.TracerAPI
+import datadog.trace.bootstrap.instrumentation.api.ContextVisitors
+import datadog.trace.bootstrap.instrumentation.api.InstrumentationTags
+import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter
+import datadog.trace.bootstrap.instrumentation.api.URIDefaultDataAdapter
+import datadog.trace.config.inversion.ConfigHelper
+import datadog.trace.test.util.DDSpecification
+
+import java.util.function.Consumer
+
+class HttpServerDecoratorTtCandidateSourcesTest extends DDSpecification {
+
+  def setupSpec() {
+    ConfigHelper.get().setConfigInversionStrict(ConfigHelper.StrictnessPolicy.TEST)
+  }
+
+  def span = Mock(AgentSpan)
+  Map<String, Object> setTags = [:]
+
+  void setup() {
+    TransactionTrackingCandidateSources.resetForTest()
+    span.setTag(_, _) >> { String k, Object v -> setTags[k] = v; null }
+    span.getTag(_) >> { String k -> setTags[k] }
+  }
+
+  void cleanup() {
+    TransactionTrackingCandidateSources.resetForTest()
+  }
+
+  def "no tag when pattern list is empty regardless of headers / qs"() {
+    setup:
+    def decorator = newDecorator(["X-Trace-Id", "tenant"], URI.create("http://h/p?tenant=42&debug=1"))
+
+    when:
+    decorator.onRequest(span, null, [marker: "anything"], datadog.context.Context.root())
+
+    then:
+    // Fast path: no allocation, no tag.
+    setTags[InstrumentationTags.TT_CANDIDATE_SOURCES] == null
+  }
+
+  def "tags matching headers and qs with deterministic order and lowercasing"() {
+    setup:
+    TransactionTrackingCandidateSources.update(["x-trace-*", "tenant", "*-id"])
+    def decorator = newDecorator(
+      ["X-Trace-Id", "X-Trace-Source", "Authorization", "USER-ID"],
+      URI.create("http://h/p?tenant=42&debug=1&request-id=abc"))
+
+    when:
+    decorator.onRequest(span, null, [marker: "anything"], datadog.context.Context.root())
+
+    then:
+    def csv = setTags[InstrumentationTags.TT_CANDIDATE_SOURCES]
+    csv != null
+    // headers first (sorted), then qs (sorted), all lowercased + deduped per bucket
+    csv == "header:user-id,header:x-trace-id,header:x-trace-source,qs:request-id,qs:tenant"
+  }
+
+  def "headers only (no query string)"() {
+    setup:
+    TransactionTrackingCandidateSources.update(["x-foo"])
+    def decorator = newDecorator(["X-FOO", "X-Bar"], URI.create("http://h/p"))
+
+    when:
+    decorator.onRequest(span, null, [:], datadog.context.Context.root())
+
+    then:
+    setTags[InstrumentationTags.TT_CANDIDATE_SOURCES] == "header:x-foo"
+  }
+
+  def "qs only (no header overrides)"() {
+    setup:
+    TransactionTrackingCandidateSources.update(["tenant*"])
+    def decorator = newDecorator([], URI.create("http://h/p?tenantId=7&other=x"))
+
+    when:
+    decorator.onRequest(span, null, [:], datadog.context.Context.root())
+
+    then:
+    setTags[InstrumentationTags.TT_CANDIDATE_SOURCES] == "qs:tenantid"
+  }
+
+  def "no match means no tag even with non-empty patterns"() {
+    setup:
+    TransactionTrackingCandidateSources.update(["nope-*"])
+    def decorator = newDecorator(["X-Foo"], URI.create("http://h/p?a=1"))
+
+    when:
+    decorator.onRequest(span, null, [:], datadog.context.Context.root())
+
+    then:
+    setTags[InstrumentationTags.TT_CANDIDATE_SOURCES] == null
+  }
+
+  def "duplicates within a bucket collapse to one entry"() {
+    setup:
+    TransactionTrackingCandidateSources.update(["x-trace-*"])
+    def decorator = newDecorator(["X-Trace-Id", "x-trace-id", "X-TRACE-ID"], URI.create("http://h/p"))
+
+    when:
+    decorator.onRequest(span, null, [:], datadog.context.Context.root())
+
+    then:
+    setTags[InstrumentationTags.TT_CANDIDATE_SOURCES] == "header:x-trace-id"
+  }
+
+  def newDecorator(List<String> headerNames, URI uri) {
+    return new HttpServerDecorator<Map, Map, Map, Map<String, String>>() {
+        @Override
+        protected TracerAPI tracer() {
+          return AgentTracer.NOOP_TRACER
+        }
+
+        @Override
+        protected String[] instrumentationNames() {
+          ["test1", "test2"]
+        }
+
+        @Override
+        protected CharSequence component() {
+          "test-component"
+        }
+
+        @Override
+        protected AgentPropagation.ContextVisitor<Map<String, String>> getter() {
+          return ContextVisitors.stringValuesMap()
+        }
+
+        @Override
+        protected AgentPropagation.ContextVisitor<Map> responseGetter() {
+          null
+        }
+
+        @Override
+        CharSequence spanName() {
+          "http-tt-span"
+        }
+
+        @Override
+        protected String method(Map m) {
+          "GET"
+        }
+
+        @Override
+        protected URIDataAdapter url(Map m) {
+          new URIDefaultDataAdapter(uri)
+        }
+
+        @Override
+        protected String peerHostIP(Map m) {
+          null
+        }
+
+        @Override
+        protected int peerPort(Map m) {
+          0
+        }
+
+        @Override
+        protected int status(Map m) {
+          0
+        }
+
+        @Override
+        protected String getRequestHeader(Map m, String key) {
+          null
+        }
+
+        @Override
+        protected void forEachRequestHeaderName(Map m, Consumer<String> consumer) {
+          headerNames.each { consumer.accept(it) }
+        }
+      }
+  }
+}

dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/InstrumentationSpecification.groovy

@@ -260,6 +260,11 @@ abstract class InstrumentationSpecification extends DDSpecification implements A
     List<DataStreamsTransactionExtractor> getDataStreamsTransactionExtractors() {
       return null
     }
+
+    @Override
+    List<String> getTransactionTrackingCandidateSourcePatterns() {
+      return Collections.emptyList()
+    }
   }
 
   @SuppressFBWarnings(value = "AT_STALE_THREAD_WRITE_OF_PRIMITIVE", justification = "The variable is accessed only by the test thread in setup and cleanup.")