Introduce tracing propagator

dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java

@@ -5,6 +5,7 @@
 import static datadog.trace.api.DDTags.DJM_ENABLED;
 import static datadog.trace.api.DDTags.DSM_ENABLED;
 import static datadog.trace.api.DDTags.PROFILING_CONTEXT_ENGINE;
+import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.TRACING_CONCERN;
 import static datadog.trace.common.metrics.MetricsAggregatorFactory.createMetricsAggregator;
 import static datadog.trace.util.AgentThreadFactory.AGENT_THREAD_GROUP;
 import static datadog.trace.util.CollectionUtils.tryMakeImmutableMap;
@@ -18,6 +19,7 @@
 import datadog.communication.ddagent.SharedCommunicationObjects;
 import datadog.communication.monitor.Monitoring;
 import datadog.communication.monitor.Recording;
+import datadog.context.propagation.Propagators;
 import datadog.trace.api.ClassloaderConfigurationOverrides;
 import datadog.trace.api.Config;
 import datadog.trace.api.DDSpanId;
@@ -86,6 +88,7 @@
 import datadog.trace.core.propagation.ExtractedContext;
 import datadog.trace.core.propagation.HttpCodec;
 import datadog.trace.core.propagation.PropagationTags;
+import datadog.trace.core.propagation.TracingPropagator;
 import datadog.trace.core.scopemanager.ContinuableScopeManager;
 import datadog.trace.core.taginterceptor.RuleFlags;
 import datadog.trace.core.taginterceptor.TagInterceptor;
@@ -719,6 +722,8 @@ private CoreTracer(
     this.propagation =
         new CorePropagation(builtExtractor, injector, injectors, dataStreamContextInjector);
 
+    Propagators.register(TRACING_CONCERN, new TracingPropagator(injector, extractor));
+
     this.tagInterceptor =
         null == tagInterceptor ? new TagInterceptor(new RuleFlags(config)) : tagInterceptor;
 

dd-trace-core/src/main/java/datadog/trace/core/propagation/TracingPropagator.java

@@ -0,0 +1,76 @@
+package datadog.trace.core.propagation;
+
+import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext;
+import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromSpanContext;
+
+import datadog.context.Context;
+import datadog.context.propagation.CarrierSetter;
+import datadog.context.propagation.CarrierVisitor;
+import datadog.context.propagation.Propagator;
+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.TagContext;
+import datadog.trace.core.DDSpanContext;
+import datadog.trace.core.propagation.HttpCodec.Extractor;
+import datadog.trace.core.propagation.HttpCodec.Injector;
+import javax.annotation.ParametersAreNonnullByDefault;
+
+/** Propagator for tracing concern. */
+@ParametersAreNonnullByDefault
+public class TracingPropagator implements Propagator {
+  private final Injector injector;
+  private final Extractor extractor;
+
+  /**
+   * Constructor.
+   *
+   * @param injector The {@link Injector} used for tracing context injection.
+   * @param extractor The {@link Extractor} used for tracing context extraction.
+   */
+  public TracingPropagator(Injector injector, Extractor extractor) {
+    this.injector = injector;
+    this.extractor = extractor;
+  }
+
+  @Override
+  public <C> void inject(Context context, C carrier, CarrierSetter<C> setter) {
+    AgentSpan span;
+    //noinspection ConstantValue
+    if (context == null
+        || carrier == null
+        || setter == null
+        || (span = fromContext(context)) == null) {
+      return;
+    }
+    AgentSpanContext spanContext = span.context();
+    if (spanContext instanceof DDSpanContext) {
+      DDSpanContext ddSpanContext = (DDSpanContext) spanContext;
+      ddSpanContext.getTraceCollector().setSamplingPriorityIfNecessary();
+      this.injector.inject(ddSpanContext, carrier, setter::set);
+    }
+  }
+
+  @Override
+  public <C> Context extract(Context context, C carrier, CarrierVisitor<C> visitor) {
+    //noinspection ConstantValue
+    if (context == null || carrier == null || visitor == null) {
+      return context;
+    }
+    TagContext spanContext = this.extractor.extract(carrier, toContextVisitor(visitor));
+    // If the extraction fails, return the original context
+    if (spanContext == null) {
+      return context;
+    }
+    // Otherwise, append a fake span wrapper to context
+    return context.with(fromSpanContext(spanContext));
+  }
+
+  private static <C> AgentPropagation.ContextVisitor<C> toContextVisitor(
+      CarrierVisitor<C> visitor) {
+    if (visitor instanceof AgentPropagation.ContextVisitor) {
+      return (AgentPropagation.ContextVisitor<C>) visitor;
+    }
+    return (carrier, classifier) -> visitor.forEachKeyValue(carrier, classifier::accept);
+  }
+}

dd-trace-core/src/test/groovy/datadog/trace/core/propagation/TracingPropagatorTest.groovy

@@ -0,0 +1,139 @@
+package datadog.trace.core.propagation
+
+import datadog.context.Context
+import datadog.context.propagation.CarrierSetter
+import datadog.context.propagation.Propagators
+import datadog.trace.api.sampling.PrioritySampling
+import datadog.trace.bootstrap.instrumentation.api.AgentPropagation
+import datadog.trace.common.writer.LoggingWriter
+import datadog.trace.core.ControllableSampler
+import datadog.trace.core.test.DDCoreSpecification
+
+import static datadog.trace.api.sampling.PrioritySampling.SAMPLER_KEEP
+import static datadog.trace.api.sampling.PrioritySampling.USER_DROP
+
+class TracingPropagatorTest extends DDCoreSpecification {
+  HttpCodec.Injector injector
+  HttpCodec.Extractor extractor
+  TracingPropagator propagator
+
+  def setup() {
+    injector = Mock(HttpCodec.Injector)
+    extractor = Mock(HttpCodec.Extractor)
+    this.propagator = new TracingPropagator(injector, extractor)
+  }
+
+  def 'test tracing propagator context injection'() {
+    setup:
+    def tracer = tracerBuilder().build()
+    def span = tracer.buildSpan('test', 'operation').start()
+    def setter = Mock(CarrierSetter)
+    def carrier = new Object()
+
+    when:
+    this.propagator.inject(span, carrier, setter)
+
+    then:
+    1 * injector.inject(span.context(), carrier, _)
+
+    cleanup:
+    span.finish()
+    tracer.close()
+  }
+
+  def 'test tracing propagator context extractor'() {
+    setup:
+    def context = Context.root()
+    // TODO Use ContextVisitor mock as getter once extractor API is refactored
+    def getter = Mock(AgentPropagation.ContextVisitor)
+    def carrier = new Object()
+
+    when:
+    this.propagator.extract(context, carrier, getter)
+
+    then:
+    1 * extractor.extract(carrier, _)
+  }
+
+  def 'span priority set when injecting'() {
+    given:
+    injectSysConfig('writer.type', 'LoggingWriter')
+    def tracer = tracerBuilder().build()
+    def setter = Mock(CarrierSetter)
+    def carrier = new Object()
+
+    when:
+    def root = tracer.buildSpan('test', 'parent').start()
+    def child = tracer.buildSpan('test', 'child').asChildOf(root).start()
+    Propagators.defaultPropagator().inject(child, carrier, setter)
+
+    then:
+    root.getSamplingPriority() == SAMPLER_KEEP as int
+    child.getSamplingPriority() == root.getSamplingPriority()
+    1 * setter.set(carrier, DatadogHttpCodec.SAMPLING_PRIORITY_KEY, String.valueOf(SAMPLER_KEEP))
+
+    cleanup:
+    child.finish()
+    root.finish()
+    tracer.close()
+  }
+
+  def 'span priority only set after first injection'() {
+    given:
+    def sampler = new ControllableSampler()
+    def tracer = tracerBuilder().writer(new LoggingWriter()).sampler(sampler).build()
+    def setter = Mock(AgentPropagation.Setter)
+    def carrier = new Object()
+
+    when:
+    def root = tracer.buildSpan('test', 'parent').start()
+    def child = tracer.buildSpan('test', 'child').asChildOf(root).start()
+    Propagators.defaultPropagator().inject(child, carrier, setter)
+
+    then:
+    root.getSamplingPriority() == SAMPLER_KEEP as int
+    child.getSamplingPriority() == root.getSamplingPriority()
+    1 * setter.set(carrier, DatadogHttpCodec.SAMPLING_PRIORITY_KEY, String.valueOf(SAMPLER_KEEP))
+
+    when:
+    sampler.nextSamplingPriority = PrioritySampling.SAMPLER_DROP as int
+    def child2 = tracer.buildSpan('test', 'child2').asChildOf(root).start()
+    Propagators.defaultPropagator().inject(child2, carrier, setter)
+
+    then:
+    root.getSamplingPriority() == SAMPLER_KEEP as int
+    child.getSamplingPriority() == root.getSamplingPriority()
+    child2.getSamplingPriority() == root.getSamplingPriority()
+    1 * setter.set(carrier, DatadogHttpCodec.SAMPLING_PRIORITY_KEY, String.valueOf(SAMPLER_KEEP))
+
+    cleanup:
+    child.finish()
+    child2.finish()
+    root.finish()
+    tracer.close()
+  }
+
+  def 'injection does not override set priority'() {
+    given:
+    def sampler = new ControllableSampler()
+    def tracer = tracerBuilder().writer(new LoggingWriter()).sampler(sampler).build()
+    def setter = Mock(AgentPropagation.Setter)
+    def carrier = new Object()
+
+    when:
+    def root = tracer.buildSpan('test', 'root').start()
+    def child = tracer.buildSpan('test', 'child').asChildOf(root).start()
+    child.setSamplingPriority(USER_DROP)
+    Propagators.defaultPropagator().inject(child, carrier, setter)
+
+    then:
+    root.getSamplingPriority() == USER_DROP as int
+    child.getSamplingPriority() == root.getSamplingPriority()
+    1 * setter.set(carrier, DatadogHttpCodec.SAMPLING_PRIORITY_KEY, String.valueOf(USER_DROP))
+
+    cleanup:
+    child.finish()
+    root.finish()
+    tracer.close()
+  }
+}

internal-api/build.gradle

@@ -64,6 +64,8 @@ excludedClassesCoverage += [
   "datadog.trace.bootstrap.instrumentation.api.Tags",
   "datadog.trace.bootstrap.instrumentation.api.CommonTagValues",
   // Caused by empty 'default' interface method
+  "datadog.trace.bootstrap.instrumentation.api.AgentPropagation",
+  "datadog.trace.bootstrap.instrumentation.api.AgentPropagation.ContextVisitor",
   "datadog.trace.bootstrap.instrumentation.api.AgentSpan",
   "datadog.trace.bootstrap.instrumentation.api.AgentSpanContext",
   "datadog.trace.bootstrap.instrumentation.api.AgentTracer",

internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentPropagation.java

@@ -1,9 +1,16 @@
 package datadog.trace.bootstrap.instrumentation.api;
 
+import datadog.context.propagation.CarrierSetter;
+import datadog.context.propagation.CarrierVisitor;
+import datadog.context.propagation.Concern;
 import datadog.trace.api.TracePropagationStyle;
 import java.util.LinkedHashMap;
+import java.util.function.BiConsumer;
+import javax.annotation.ParametersAreNonnullByDefault;
 
 public interface AgentPropagation {
+  Concern TRACING_CONCERN = Concern.named("tracing");
+
   <C> void inject(AgentSpan span, C carrier, Setter<C> setter);
 
   <C> void inject(AgentSpanContext context, C carrier, Setter<C> setter);
@@ -25,7 +32,7 @@ <C> void injectPathwayContext(
   <C> void injectPathwayContextWithoutSendingStats(
       AgentSpan span, C carrier, Setter<C> setter, LinkedHashMap<String, String> sortedTags);
 
-  interface Setter<C> {
+  interface Setter<C> extends CarrierSetter<C> {
     void set(C carrier, String key, String value);
   }
 
@@ -35,7 +42,18 @@ interface KeyClassifier {
     boolean accept(String key, String value);
   }
 
-  interface ContextVisitor<C> {
+  interface ContextVisitor<C> extends CarrierVisitor<C> {
     void forEachKey(C carrier, KeyClassifier classifier);
+
+    @ParametersAreNonnullByDefault
+    @Override
+    default void forEachKeyValue(C carrier, BiConsumer<String, String> visitor) {
+      forEachKey(
+          carrier,
+          (key, value) -> {
+            visitor.accept(key, value);
+            return true;
+          });
+    }
   }
 }

internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpan.java

@@ -10,12 +10,42 @@
 import datadog.trace.api.gateway.IGSpanInfo;
 import datadog.trace.api.gateway.RequestContext;
 import datadog.trace.api.interceptor.MutableSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer.NoopAgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer.NoopContext;
 import java.util.Map;
+import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 
 public interface AgentSpan
     extends MutableSpan, ImplicitContextKeyed, Context, IGSpanInfo, WithAgentSpan {
 
+  /**
+   * Extracts the span from context.
+   *
+   * @param context the context to extract the span from.
+   * @return the span if existing, {@code null} otherwise.
+   */
+  static AgentSpan fromContext(Context context) {
+    return context.get(SPAN_KEY);
+  }
+
+  /**
+   * Creates a span wrapper from a span context.
+   *
+   * <p>Creating a such span will not create a tracing span to complete a local root trace. It gives
+   * a span instance based on a span context for span-based API. It is usually used with an
+   * extracted span context as parameter to represent a remove span.
+   *
+   * @param spanContext the span context to get a full-fledged span.
+   * @return a span wrapped based on a span context.
+   */
+  static AgentSpan fromSpanContext(AgentSpanContext spanContext) {
+    if (spanContext == null || spanContext == NoopContext.INSTANCE) {
+      return NoopAgentSpan.INSTANCE;
+    }
+    return new AgentTracer.ExtractedSpan(spanContext);
+  }
+
   DDTraceId getTraceId();
 
   long getSpanId();
@@ -163,13 +193,13 @@ default Context storeInto(Context context) {
 
   @Nullable
   @Override
-  default <T> T get(ContextKey<T> key) {
+  default <T> T get(@Nonnull ContextKey<T> key) {
     // noinspection unchecked
     return SPAN_KEY == key ? (T) this : Context.root().get(key);
   }
 
   @Override
-  default <T> Context with(ContextKey<T> key, @Nullable T value) {
+  default <T> Context with(@Nonnull ContextKey<T> key, @Nullable T value) {
     return Context.root().with(SPAN_KEY, this, key, value);
   }
 }