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);
  }
 }