Merge pull request #317 from DataDog/tyler/tagged-headers

Add setting for collecting headers as tags during extraction

dd-trace-ot/src/main/java/datadog/opentracing/DDTracer.java

@@ -88,7 +88,8 @@ public class DDTracer implements io.opentracing.Tracer {
         Writer.Builder.forConfig(config),
         Sampler.Builder.forConfig(config),
         DDTraceConfig.parseMap(config.getProperty(DDTraceConfig.SPAN_TAGS)),
-        DDTraceConfig.parseMap(config.getProperty(DDTraceConfig.SERVICE_MAPPING)));
+        DDTraceConfig.parseMap(config.getProperty(DDTraceConfig.SERVICE_MAPPING)),
+        DDTraceConfig.parseMap(config.getProperty(DDTraceConfig.HEADER_TAGS)));
     log.debug("Using config: {}", config);
   }
 
@@ -98,6 +99,7 @@ public class DDTracer implements io.opentracing.Tracer {
         writer,
         sampler,
         Collections.<String, String>emptyMap(),
+        Collections.<String, String>emptyMap(),
         Collections.<String, String>emptyMap());
   }
 
@@ -106,7 +108,8 @@ public class DDTracer implements io.opentracing.Tracer {
       final Writer writer,
       final Sampler sampler,
       final Map<String, String> defaultSpanTags,
-      final Map<String, String> serviceNameMappings) {
+      final Map<String, String> serviceNameMappings,
+      final Map<String, String> taggedHeaders) {
     this.serviceName = serviceName;
     this.writer = writer;
     this.writer.start();
@@ -123,8 +126,8 @@ public class DDTracer implements io.opentracing.Tracer {
             });
 
     registry = new CodecRegistry();
-    registry.register(Format.Builtin.HTTP_HEADERS, new HTTPCodec());
+    registry.register(Format.Builtin.HTTP_HEADERS, new HTTPCodec(taggedHeaders));
-    registry.register(Format.Builtin.TEXT_MAP, new HTTPCodec());
+    registry.register(Format.Builtin.TEXT_MAP, new HTTPCodec(taggedHeaders));
     if (this.writer instanceof DDAgentWriter && sampler instanceof DDApi.ResponseListener) {
       final DDApi api = ((DDAgentWriter) this.writer).getApi();
       api.addResponseListener((DDApi.ResponseListener) this.sampler);
@@ -149,7 +152,8 @@ public class DDTracer implements io.opentracing.Tracer {
         writer,
         new AllSampler(),
         DDTraceConfig.parseMap(new DDTraceConfig().getProperty(DDTraceConfig.SPAN_TAGS)),
-        DDTraceConfig.parseMap(new DDTraceConfig().getProperty(DDTraceConfig.SERVICE_MAPPING)));
+        DDTraceConfig.parseMap(new DDTraceConfig().getProperty(DDTraceConfig.SERVICE_MAPPING)),
+        DDTraceConfig.parseMap(new DDTraceConfig().getProperty(DDTraceConfig.HEADER_TAGS)));
   }
 
   /**
@@ -489,6 +493,12 @@ public class DDTracer implements io.opentracing.Tracer {
         traceId = ddsc.getTraceId();
         parentSpanId = ddsc.getSpanId();
         baggage = ddsc.getBaggage();
+        if (this.tags.isEmpty() && !ddsc.getTags().isEmpty()) {
+          this.tags = new HashMap<>();
+        }
+        if (!ddsc.getTags().isEmpty()) {
+          tags.putAll(ddsc.getTags());
+        }
         parentTrace = new PendingTrace(DDTracer.this, traceId);
         samplingPriority = ddsc.getSamplingPriority();
 

dd-trace-ot/src/main/java/datadog/opentracing/propagation/ExtractedContext.java

@@ -9,17 +9,20 @@ public class ExtractedContext implements SpanContext {
   private final Long spanId;
   private final int samplingPriority;
   private final Map<String, String> baggage;
+  private final Map<String, String> tags;
   private final AtomicBoolean samplingPriorityLocked = new AtomicBoolean(false);
 
   public ExtractedContext(
       final Long traceId,
       final Long spanId,
       final int samplingPriority,
-      final Map<String, String> baggage) {
+      final Map<String, String> baggage,
+      final Map<String, String> tags) {
     this.traceId = traceId;
     this.spanId = spanId;
     this.samplingPriority = samplingPriority;
     this.baggage = baggage;
+    this.tags = tags;
   }
 
   @Override
@@ -47,6 +50,10 @@ public class ExtractedContext implements SpanContext {
     return baggage;
   }
 
+  public Map<String, String> getTags() {
+    return tags;
+  }
+
   public boolean getSamplingPriorityLocked() {
     return samplingPriorityLocked.get();
   }

dd-trace-ot/src/main/java/datadog/opentracing/propagation/HTTPCodec.java

@@ -20,6 +20,15 @@ public class HTTPCodec implements Codec<TextMap> {
   private static final String SPAN_ID_KEY = "x-datadog-parent-id";
   private static final String SAMPLING_PRIORITY_KEY = "x-datadog-sampling-priority";
 
+  private final Map<String, String> taggedHeaders;
+
+  public HTTPCodec(final Map<String, String> taggedHeaders) {
+    this.taggedHeaders = new HashMap<>();
+    for (final Map.Entry<String, String> mapping : taggedHeaders.entrySet()) {
+      this.taggedHeaders.put(mapping.getKey().trim().toLowerCase(), mapping.getValue());
+    }
+  }
+
   @Override
   public void inject(final DDSpanContext context, final TextMap carrier) {
     carrier.put(TRACE_ID_KEY, String.valueOf(context.getTraceId()));
@@ -38,6 +47,7 @@ public class HTTPCodec implements Codec<TextMap> {
   public ExtractedContext extract(final TextMap carrier) {
 
     Map<String, String> baggage = Collections.emptyMap();
+    Map<String, String> tags = Collections.emptyMap();
     Long traceId = 0L;
     Long spanId = 0L;
     int samplingPriority = PrioritySampling.UNSET;
@@ -56,10 +66,17 @@ public class HTTPCodec implements Codec<TextMap> {
       } else if (key.equalsIgnoreCase(SAMPLING_PRIORITY_KEY)) {
         samplingPriority = Integer.parseInt(entry.getValue());
       }
+
+      if (taggedHeaders.containsKey(key)) {
+        if (tags.isEmpty()) {
+          tags = new HashMap<>();
+        }
+        tags.put(taggedHeaders.get(key), decode(entry.getValue()));
+      }
     }
     ExtractedContext context = null;
     if (traceId != 0L) {
-      context = new ExtractedContext(traceId, spanId, samplingPriority, baggage);
+      context = new ExtractedContext(traceId, spanId, samplingPriority, baggage, tags);
       context.lockSamplingPriority();
 
       log.debug("{} - Parent context extracted", context.getTraceId());

dd-trace-ot/src/main/java/datadog/trace/common/DDTraceConfig.java

@@ -30,6 +30,7 @@ public class DDTraceConfig extends Properties {
   public static final String AGENT_PORT = "agent.port";
   public static final String PRIORITY_SAMPLING = "priority.sampling";
   public static final String SPAN_TAGS = "trace.span.tags";
+  public static final String HEADER_TAGS = "trace.header.tags";
 
   private final String serviceName = getPropOrEnv(PREFIX + SERVICE_NAME);
   private final String serviceMapping = getPropOrEnv(PREFIX + SERVICE_MAPPING);
@@ -38,6 +39,7 @@ public class DDTraceConfig extends Properties {
   private final String agentPort = getPropOrEnv(PREFIX + AGENT_PORT);
   private final String prioritySampling = getPropOrEnv(PREFIX + PRIORITY_SAMPLING);
   private final String spanTags = getPropOrEnv(PREFIX + SPAN_TAGS);
+  private final String headerTags = getPropOrEnv(PREFIX + HEADER_TAGS);
 
   public DDTraceConfig() {
     super();
@@ -56,6 +58,7 @@ public class DDTraceConfig extends Properties {
     setIfNotNull(AGENT_PORT, agentPort);
     setIfNotNull(PRIORITY_SAMPLING, prioritySampling);
     setIfNotNull(SPAN_TAGS, spanTags);
+    setIfNotNull(HEADER_TAGS, headerTags);
   }
 
   public DDTraceConfig(final String serviceName) {

dd-trace-ot/src/test/groovy/datadog/opentracing/DDSpanBuilderTest.groovy

@@ -1,5 +1,6 @@
 package datadog.opentracing
 
+import datadog.opentracing.propagation.ExtractedContext
 import datadog.trace.api.DDTags
 import datadog.trace.common.writer.ListWriter
 import spock.lang.Specification
@@ -222,6 +223,24 @@ class DDSpanBuilderTest extends Specification {
     spans[(int) (Math.random() * nbSamples)].context.trace.containsAll(spans)
   }
 
+  def "ExtractedContext should populate new span details"() {
+    setup:
+    final DDSpan span = tracer.buildSpan("op name")
+      .asChildOf(extractedContext).start()
+
+    expect:
+    span.traceId == extractedContext.traceId
+    span.parentId == extractedContext.spanId
+    span.samplingPriority == extractedContext.samplingPriority
+    span.context().baggageItems == extractedContext.baggage
+    span.context().@tags == extractedContext.tags
+
+    where:
+    extractedContext                                                  | _
+    new ExtractedContext(1, 2, 0, [:], [:])                           | _
+    new ExtractedContext(3, 4, 1, ["asdf": "qwer"], ["zxcv": "1234"]) | _
+  }
+
   def "global span tags populated on each span"() {
     setup:
     System.setProperty("dd.trace.span.tags", tagString)

dd-trace-ot/src/test/groovy/datadog/opentracing/propagation/HTTPCodecTest.groovy

@@ -20,6 +20,8 @@ class HTTPCodecTest extends Specification {
   @Shared
   private static final String SAMPLING_PRIORITY_KEY = "x-datadog-sampling-priority"
 
+  HTTPCodec codec = new HTTPCodec(["SOME_HEADER": "some-tag"])
+
   def "inject http headers"() {
     setup:
     def writer = new ListWriter()
@@ -47,7 +49,6 @@ class HTTPCodecTest extends Specification {
 
     final Map<String, String> carrier = new HashMap<>()
 
-    final HTTPCodec codec = new HTTPCodec()
     codec.inject(mockedContext, new TextMapInjectAdapter(carrier))
 
     expect:
@@ -70,13 +71,13 @@ class HTTPCodecTest extends Specification {
       (SPAN_ID_KEY.toUpperCase())             : "2",
       (OT_BAGGAGE_PREFIX.toUpperCase() + "k1"): "v1",
       (OT_BAGGAGE_PREFIX.toUpperCase() + "k2"): "v2",
+      SOME_HEADER                             : "my-interesting-info",
     ]
 
     if (samplingPriority != PrioritySampling.UNSET) {
       actual.put(SAMPLING_PRIORITY_KEY, String.valueOf(samplingPriority))
     }
 
-    final HTTPCodec codec = new HTTPCodec()
     final ExtractedContext context = codec.extract(new TextMapExtractAdapter(actual))
 
     expect:
@@ -84,6 +85,7 @@ class HTTPCodecTest extends Specification {
     context.getSpanId() == 2l
     context.getBaggage().get("k1") == "v1"
     context.getBaggage().get("k2") == "v2"
+    context.getTags() == ["some-tag": "my-interesting-info"]
     context.getSamplingPriority() == samplingPriority
 
     where:

dd-trace-ot/src/test/groovy/datadog/trace/DDTraceConfigTest.groovy

@@ -14,6 +14,7 @@ import spock.lang.Specification
 
 import static datadog.trace.common.DDTraceConfig.AGENT_HOST
 import static datadog.trace.common.DDTraceConfig.AGENT_PORT
+import static datadog.trace.common.DDTraceConfig.HEADER_TAGS
 import static datadog.trace.common.DDTraceConfig.PREFIX
 import static datadog.trace.common.DDTraceConfig.SERVICE_MAPPING
 import static datadog.trace.common.DDTraceConfig.SERVICE_NAME
@@ -114,16 +115,19 @@ class DDTraceConfigTest extends Specification {
     setup:
     System.setProperty(PREFIX + SERVICE_MAPPING, mapString)
     System.setProperty(PREFIX + SPAN_TAGS, mapString)
+    System.setProperty(PREFIX + HEADER_TAGS, mapString)
 
     when:
     def tracer = new DDTracer()
     ServiceNameDecorator decorator = tracer.spanContextDecorators.values().flatten().find {
       it instanceof ServiceNameDecorator
     }
+    def taggedHeaders = tracer.registry.codecs.values().first().taggedHeaders
 
     then:
     tracer.spanTags == map
     decorator.mappings == map
+    taggedHeaders == map
 
     where:
     mapString       | map