Skip to content

Commit 2a09977

Browse files
PerfectSlayerPerfectSlayer
authored
feat(core): Introduce tracing propagator (#8313)
1 parent 8d0bb34 commit 2a09977

File tree

8 files changed

+396
-4
lines changed

8 files changed

+396
-4
lines changed

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
import static datadog.trace.api.DDTags.DJM_ENABLED;
66
import static datadog.trace.api.DDTags.DSM_ENABLED;
77
import static datadog.trace.api.DDTags.PROFILING_CONTEXT_ENGINE;
8+
import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.TRACING_CONCERN;
89
import static datadog.trace.common.metrics.MetricsAggregatorFactory.createMetricsAggregator;
910
import static datadog.trace.util.AgentThreadFactory.AGENT_THREAD_GROUP;
1011
import static datadog.trace.util.CollectionUtils.tryMakeImmutableMap;
@@ -18,6 +19,7 @@
1819
import datadog.communication.ddagent.SharedCommunicationObjects;
1920
import datadog.communication.monitor.Monitoring;
2021
import datadog.communication.monitor.Recording;
22+
import datadog.context.propagation.Propagators;
2123
import datadog.trace.api.ClassloaderConfigurationOverrides;
2224
import datadog.trace.api.Config;
2325
import datadog.trace.api.DDSpanId;
@@ -86,6 +88,7 @@
8688
import datadog.trace.core.propagation.ExtractedContext;
8789
import datadog.trace.core.propagation.HttpCodec;
8890
import datadog.trace.core.propagation.PropagationTags;
91+
import datadog.trace.core.propagation.TracingPropagator;
8992
import datadog.trace.core.scopemanager.ContinuableScopeManager;
9093
import datadog.trace.core.taginterceptor.RuleFlags;
9194
import datadog.trace.core.taginterceptor.TagInterceptor;
@@ -719,6 +722,8 @@ private CoreTracer(
719722
this.propagation =
720723
new CorePropagation(builtExtractor, injector, injectors, dataStreamContextInjector);
721724

725+
Propagators.register(TRACING_CONCERN, new TracingPropagator(injector, extractor));
726+
722727
this.tagInterceptor =
723728
null == tagInterceptor ? new TagInterceptor(new RuleFlags(config)) : tagInterceptor;
724729

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

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
package datadog.trace.core.propagation;
2+
3+
import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext;
4+
import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromSpanContext;
5+
6+
import datadog.context.Context;
7+
import datadog.context.propagation.CarrierSetter;
8+
import datadog.context.propagation.CarrierVisitor;
9+
import datadog.context.propagation.Propagator;
10+
import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
11+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
12+
import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
13+
import datadog.trace.bootstrap.instrumentation.api.TagContext;
14+
import datadog.trace.core.DDSpanContext;
15+
import datadog.trace.core.propagation.HttpCodec.Extractor;
16+
import datadog.trace.core.propagation.HttpCodec.Injector;
17+
import javax.annotation.ParametersAreNonnullByDefault;
18+
19+
/** Propagator for tracing concern. */
20+
@ParametersAreNonnullByDefault
21+
public class TracingPropagator implements Propagator {
22+
private final Injector injector;
23+
private final Extractor extractor;
24+
25+
/**
26+
* Constructor.
27+
*
28+
* @param injector The {@link Injector} used for tracing context injection.
29+
* @param extractor The {@link Extractor} used for tracing context extraction.
30+
*/
31+
public TracingPropagator(Injector injector, Extractor extractor) {
32+
this.injector = injector;
33+
this.extractor = extractor;
34+
}
35+
36+
@Override
37+
public <C> void inject(Context context, C carrier, CarrierSetter<C> setter) {
38+
AgentSpan span;
39+
//noinspection ConstantValue
40+
if (context == null
41+
|| carrier == null
42+
|| setter == null
43+
|| (span = fromContext(context)) == null) {
44+
return;
45+
}
46+
AgentSpanContext spanContext = span.context();
47+
if (spanContext instanceof DDSpanContext) {
48+
DDSpanContext ddSpanContext = (DDSpanContext) spanContext;
49+
ddSpanContext.getTraceCollector().setSamplingPriorityIfNecessary();
50+
this.injector.inject(ddSpanContext, carrier, setter::set);
51+
}
52+
}
53+
54+
@Override
55+
public <C> Context extract(Context context, C carrier, CarrierVisitor<C> visitor) {
56+
//noinspection ConstantValue
57+
if (context == null || carrier == null || visitor == null) {
58+
return context;
59+
}
60+
TagContext spanContext = this.extractor.extract(carrier, toContextVisitor(visitor));
61+
// If the extraction fails, return the original context
62+
if (spanContext == null) {
63+
return context;
64+
}
65+
// Otherwise, append a fake span wrapper to context
66+
return context.with(fromSpanContext(spanContext));
67+
}
68+
69+
private static <C> AgentPropagation.ContextVisitor<C> toContextVisitor(
70+
CarrierVisitor<C> visitor) {
71+
if (visitor instanceof AgentPropagation.ContextVisitor) {
72+
return (AgentPropagation.ContextVisitor<C>) visitor;
73+
}
74+
return (carrier, classifier) -> visitor.forEachKeyValue(carrier, classifier::accept);
75+
}
76+
}

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

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,139 @@
1+
package datadog.trace.core.propagation
2+
3+
import datadog.context.Context
4+
import datadog.context.propagation.CarrierSetter
5+
import datadog.context.propagation.Propagators
6+
import datadog.trace.api.sampling.PrioritySampling
7+
import datadog.trace.bootstrap.instrumentation.api.AgentPropagation
8+
import datadog.trace.common.writer.LoggingWriter
9+
import datadog.trace.core.ControllableSampler
10+
import datadog.trace.core.test.DDCoreSpecification
11+
12+
import static datadog.trace.api.sampling.PrioritySampling.SAMPLER_KEEP
13+
import static datadog.trace.api.sampling.PrioritySampling.USER_DROP
14+
15+
class TracingPropagatorTest extends DDCoreSpecification {
16+
HttpCodec.Injector injector
17+
HttpCodec.Extractor extractor
18+
TracingPropagator propagator
19+
20+
def setup() {
21+
injector = Mock(HttpCodec.Injector)
22+
extractor = Mock(HttpCodec.Extractor)
23+
this.propagator = new TracingPropagator(injector, extractor)
24+
}
25+
26+
def 'test tracing propagator context injection'() {
27+
setup:
28+
def tracer = tracerBuilder().build()
29+
def span = tracer.buildSpan('test', 'operation').start()
30+
def setter = Mock(CarrierSetter)
31+
def carrier = new Object()
32+
33+
when:
34+
this.propagator.inject(span, carrier, setter)
35+
36+
then:
37+
1 * injector.inject(span.context(), carrier, _)
38+
39+
cleanup:
40+
span.finish()
41+
tracer.close()
42+
}
43+
44+
def 'test tracing propagator context extractor'() {
45+
setup:
46+
def context = Context.root()
47+
// TODO Use ContextVisitor mock as getter once extractor API is refactored
48+
def getter = Mock(AgentPropagation.ContextVisitor)
49+
def carrier = new Object()
50+
51+
when:
52+
this.propagator.extract(context, carrier, getter)
53+
54+
then:
55+
1 * extractor.extract(carrier, _)
56+
}
57+
58+
def 'span priority set when injecting'() {
59+
given:
60+
injectSysConfig('writer.type', 'LoggingWriter')
61+
def tracer = tracerBuilder().build()
62+
def setter = Mock(CarrierSetter)
63+
def carrier = new Object()
64+
65+
when:
66+
def root = tracer.buildSpan('test', 'parent').start()
67+
def child = tracer.buildSpan('test', 'child').asChildOf(root).start()
68+
Propagators.defaultPropagator().inject(child, carrier, setter)
69+
70+
then:
71+
root.getSamplingPriority() == SAMPLER_KEEP as int
72+
child.getSamplingPriority() == root.getSamplingPriority()
73+
1 * setter.set(carrier, DatadogHttpCodec.SAMPLING_PRIORITY_KEY, String.valueOf(SAMPLER_KEEP))
74+
75+
cleanup:
76+
child.finish()
77+
root.finish()
78+
tracer.close()
79+
}
80+
81+
def 'span priority only set after first injection'() {
82+
given:
83+
def sampler = new ControllableSampler()
84+
def tracer = tracerBuilder().writer(new LoggingWriter()).sampler(sampler).build()
85+
def setter = Mock(AgentPropagation.Setter)
86+
def carrier = new Object()
87+
88+
when:
89+
def root = tracer.buildSpan('test', 'parent').start()
90+
def child = tracer.buildSpan('test', 'child').asChildOf(root).start()
91+
Propagators.defaultPropagator().inject(child, carrier, setter)
92+
93+
then:
94+
root.getSamplingPriority() == SAMPLER_KEEP as int
95+
child.getSamplingPriority() == root.getSamplingPriority()
96+
1 * setter.set(carrier, DatadogHttpCodec.SAMPLING_PRIORITY_KEY, String.valueOf(SAMPLER_KEEP))
97+
98+
when:
99+
sampler.nextSamplingPriority = PrioritySampling.SAMPLER_DROP as int
100+
def child2 = tracer.buildSpan('test', 'child2').asChildOf(root).start()
101+
Propagators.defaultPropagator().inject(child2, carrier, setter)
102+
103+
then:
104+
root.getSamplingPriority() == SAMPLER_KEEP as int
105+
child.getSamplingPriority() == root.getSamplingPriority()
106+
child2.getSamplingPriority() == root.getSamplingPriority()
107+
1 * setter.set(carrier, DatadogHttpCodec.SAMPLING_PRIORITY_KEY, String.valueOf(SAMPLER_KEEP))
108+
109+
cleanup:
110+
child.finish()
111+
child2.finish()
112+
root.finish()
113+
tracer.close()
114+
}
115+
116+
def 'injection does not override set priority'() {
117+
given:
118+
def sampler = new ControllableSampler()
119+
def tracer = tracerBuilder().writer(new LoggingWriter()).sampler(sampler).build()
120+
def setter = Mock(AgentPropagation.Setter)
121+
def carrier = new Object()
122+
123+
when:
124+
def root = tracer.buildSpan('test', 'root').start()
125+
def child = tracer.buildSpan('test', 'child').asChildOf(root).start()
126+
child.setSamplingPriority(USER_DROP)
127+
Propagators.defaultPropagator().inject(child, carrier, setter)
128+
129+
then:
130+
root.getSamplingPriority() == USER_DROP as int
131+
child.getSamplingPriority() == root.getSamplingPriority()
132+
1 * setter.set(carrier, DatadogHttpCodec.SAMPLING_PRIORITY_KEY, String.valueOf(USER_DROP))
133+
134+
cleanup:
135+
child.finish()
136+
root.finish()
137+
tracer.close()
138+
}
139+
}

‎internal-api/build.gradle

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,8 @@ excludedClassesCoverage += [
6464
"datadog.trace.bootstrap.instrumentation.api.Tags",
6565
"datadog.trace.bootstrap.instrumentation.api.CommonTagValues",
6666
// Caused by empty 'default' interface method
67+
"datadog.trace.bootstrap.instrumentation.api.AgentPropagation",
68+
"datadog.trace.bootstrap.instrumentation.api.AgentPropagation.ContextVisitor",
6769
"datadog.trace.bootstrap.instrumentation.api.AgentSpan",
6870
"datadog.trace.bootstrap.instrumentation.api.AgentSpanContext",
6971
"datadog.trace.bootstrap.instrumentation.api.AgentTracer",

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

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,16 @@
11
package datadog.trace.bootstrap.instrumentation.api;
22

3+
import datadog.context.propagation.CarrierSetter;
4+
import datadog.context.propagation.CarrierVisitor;
5+
import datadog.context.propagation.Concern;
36
import datadog.trace.api.TracePropagationStyle;
47
import java.util.LinkedHashMap;
8+
import java.util.function.BiConsumer;
9+
import javax.annotation.ParametersAreNonnullByDefault;
510

611
public interface AgentPropagation {
12+
Concern TRACING_CONCERN = Concern.named("tracing");
13+
714
<C> void inject(AgentSpan span, C carrier, Setter<C> setter);
815

916
<C> void inject(AgentSpanContext context, C carrier, Setter<C> setter);
@@ -25,7 +32,7 @@ <C> void injectPathwayContext(
2532
<C> void injectPathwayContextWithoutSendingStats(
2633
AgentSpan span, C carrier, Setter<C> setter, LinkedHashMap<String, String> sortedTags);
2734

28-
interface Setter<C> {
35+
interface Setter<C> extends CarrierSetter<C> {
2936
void set(C carrier, String key, String value);
3037
}
3138

@@ -35,7 +42,18 @@ interface KeyClassifier {
3542
boolean accept(String key, String value);
3643
}
3744

38-
interface ContextVisitor<C> {
45+
interface ContextVisitor<C> extends CarrierVisitor<C> {
3946
void forEachKey(C carrier, KeyClassifier classifier);
47+
48+
@ParametersAreNonnullByDefault
49+
@Override
50+
default void forEachKeyValue(C carrier, BiConsumer<String, String> visitor) {
51+
forEachKey(
52+
carrier,
53+
(key, value) -> {
54+
visitor.accept(key, value);
55+
return true;
56+
});
57+
}
4058
}
4159
}

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

Lines changed: 32 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,12 +10,42 @@
1010
import datadog.trace.api.gateway.IGSpanInfo;
1111
import datadog.trace.api.gateway.RequestContext;
1212
import datadog.trace.api.interceptor.MutableSpan;
13+
import datadog.trace.bootstrap.instrumentation.api.AgentTracer.NoopAgentSpan;
14+
import datadog.trace.bootstrap.instrumentation.api.AgentTracer.NoopContext;
1315
import java.util.Map;
16+
import javax.annotation.Nonnull;
1417
import javax.annotation.Nullable;
1518

1619
public interface AgentSpan
1720
extends MutableSpan, ImplicitContextKeyed, Context, IGSpanInfo, WithAgentSpan {
1821

22+
/**
23+
* Extracts the span from context.
24+
*
25+
* @param context the context to extract the span from.
26+
* @return the span if existing, {@code null} otherwise.
27+
*/
28+
static AgentSpan fromContext(Context context) {
29+
return context.get(SPAN_KEY);
30+
}
31+
32+
/**
33+
* Creates a span wrapper from a span context.
34+
*
35+
* <p>Creating a such span will not create a tracing span to complete a local root trace. It gives
36+
* a span instance based on a span context for span-based API. It is usually used with an
37+
* extracted span context as parameter to represent a remove span.
38+
*
39+
* @param spanContext the span context to get a full-fledged span.
40+
* @return a span wrapped based on a span context.
41+
*/
42+
static AgentSpan fromSpanContext(AgentSpanContext spanContext) {
43+
if (spanContext == null || spanContext == NoopContext.INSTANCE) {
44+
return NoopAgentSpan.INSTANCE;
45+
}
46+
return new AgentTracer.ExtractedSpan(spanContext);
47+
}
48+
1949
DDTraceId getTraceId();
2050

2151
long getSpanId();
@@ -163,13 +193,13 @@ default Context storeInto(Context context) {
163193

164194
@Nullable
165195
@Override
166-
default <T> T get(ContextKey<T> key) {
196+
default <T> T get(@Nonnull ContextKey<T> key) {
167197
// noinspection unchecked
168198
return SPAN_KEY == key ? (T) this : Context.root().get(key);
169199
}
170200

171201
@Override
172-
default <T> Context with(ContextKey<T> key, @Nullable T value) {
202+
default <T> Context with(@Nonnull ContextKey<T> key, @Nullable T value) {
173203
return Context.root().with(SPAN_KEY, this, key, value);
174204
}
175205
}

0 commit comments

Comments
 (0)