Skip to content

Commit 46fd85c

Browse files
PerfectSlayerPerfectSlayer
authored
Migrate context extraction calls to context-first APIs (#8368)
* feat(core): Remove core propagation injection API * feat(context): Improve composite propagator order * chore(api): Simplify span link test * feat(core): Move DSM extractor to its own propagator
1 parent 429031c commit 46fd85c

File tree

14 files changed

+139
-351
lines changed

14 files changed

+139
-351
lines changed

‎components/context/src/main/java/datadog/context/propagation/CompositePropagator.java

+2-2
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,8 @@ class CompositePropagator implements Propagator {
1414

1515
@Override
1616
public <C> void inject(Context context, C carrier, CarrierSetter<C> setter) {
17-
for (Propagator propagator : this.propagators) {
18-
propagator.inject(context, carrier, setter);
17+
for (int i = this.propagators.length - 1; i >= 0; i--) {
18+
this.propagators[i].inject(context, carrier, setter);
1919
}
2020
}
2121

‎components/context/src/main/java/datadog/context/propagation/Propagators.java

+2-1
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,8 @@ public static Propagator noop() {
7878
* Creates a composite propagator.
7979
*
8080
* @param propagators the elements that composes the returned propagator.
81-
* @return the composite propagator that will apply the propagators in their given order.
81+
* @return the composite propagator that will apply the propagators in their given order for
82+
* context extraction, and reverse given order for context injection.
8283
*/
8384
public static Propagator composite(Propagator... propagators) {
8485
if (propagators.length == 0) {

‎dd-trace-core/build.gradle

+4
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,10 @@ excludedClassesCoverage += [
4848
'datadog.trace.core.TracingConfigPoller.Updater',
4949
// covered with dd-trace-core/src/test/groovy/datadog/trace/core/datastreams/CheckpointerTest.groovy
5050
'datadog.trace.core.datastreams.DefaultDataStreamsMonitoring',
51+
// TODO CorePropagation will be removed during context refactoring
52+
'datadog.trace.core.propagation.CorePropagation',
53+
// TODO DSM propagator will be tested once fully migrated
54+
'datadog.trace.core.datastreams.DataStreamPropagator'
5155
]
5256

5357
addTestSuite('traceAgentTest')

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

+12-15
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.DSM_CONCERN;
89
import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.STANDALONE_ASM_CONCERN;
910
import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.TRACING_CONCERN;
1011
import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.XRAY_TRACING_CONCERN;
@@ -32,7 +33,6 @@
3233
import datadog.trace.api.InstrumenterConfig;
3334
import datadog.trace.api.StatsDClient;
3435
import datadog.trace.api.TraceConfig;
35-
import datadog.trace.api.TracePropagationStyle;
3636
import datadog.trace.api.config.GeneralConfig;
3737
import datadog.trace.api.experimental.DataStreamsCheckpointer;
3838
import datadog.trace.api.flare.TracerFlare;
@@ -79,7 +79,6 @@
7979
import datadog.trace.common.writer.WriterFactory;
8080
import datadog.trace.common.writer.ddintake.DDIntakeTraceInterceptor;
8181
import datadog.trace.context.TraceScope;
82-
import datadog.trace.core.datastreams.DataStreamContextInjector;
8382
import datadog.trace.core.datastreams.DataStreamsMonitoring;
8483
import datadog.trace.core.datastreams.DefaultDataStreamsMonitoring;
8584
import datadog.trace.core.flare.TracerFlarePoller;
@@ -715,27 +714,25 @@ private CoreTracer(
715714

716715
sharedCommunicationObjects.whenReady(this.dataStreamsMonitoring::start);
717716

718-
// Create default extractor from config if not provided and decorate it with DSM extractor
719-
HttpCodec.Extractor builtExtractor =
720-
extractor == null ? HttpCodec.createExtractor(config, this::captureTraceConfig) : extractor;
721-
builtExtractor = this.dataStreamsMonitoring.extractor(builtExtractor);
722-
// Create all HTTP injectors plus the DSM one
723-
Map<TracePropagationStyle, HttpCodec.Injector> injectors =
724-
HttpCodec.allInjectorsFor(config, invertMap(baggageMapping));
725-
DataStreamContextInjector dataStreamContextInjector = this.dataStreamsMonitoring.injector();
726-
// Store all propagators to propagation
727-
this.propagation =
728-
new CorePropagation(builtExtractor, injector, injectors, dataStreamContextInjector);
717+
// Store all propagators to propagation -- only DSM injection left
718+
this.propagation = new CorePropagation(this.dataStreamsMonitoring.injector());
729719

720+
// Register context propagators
721+
HttpCodec.Extractor tracingExtractor =
722+
extractor == null ? HttpCodec.createExtractor(config, this::captureTraceConfig) : extractor;
723+
TracingPropagator tracingPropagator = new TracingPropagator(injector, tracingExtractor);
730724
// Check if standalone AppSec is enabled:
731725
// If enabled, use the standalone AppSec propagator by default that will limit tracing concern
732726
// injection and delegate to the tracing propagator if needed,
733727
// If disabled, the most common case, use the usual tracing propagator by default.
734728
boolean standaloneAppSec = config.isAppSecStandaloneEnabled();
729+
boolean dsm = config.isDataStreamsEnabled();
735730
Propagators.register(STANDALONE_ASM_CONCERN, new StandaloneAsmPropagator(), standaloneAppSec);
736-
Propagators.register(
737-
TRACING_CONCERN, new TracingPropagator(injector, extractor), !standaloneAppSec);
731+
Propagators.register(TRACING_CONCERN, tracingPropagator, !standaloneAppSec);
738732
Propagators.register(XRAY_TRACING_CONCERN, new XRayPropagator(config), false);
733+
if (dsm) {
734+
Propagators.register(DSM_CONCERN, this.dataStreamsMonitoring.propagator());
735+
}
739736

740737
this.tagInterceptor =
741738
null == tagInterceptor ? new TagInterceptor(new RuleFlags(config)) : tagInterceptor;

‎dd-trace-core/src/main/java/datadog/trace/core/datastreams/DataStreamPropagator.java

+83
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
package datadog.trace.core.datastreams;
2+
3+
import datadog.context.Context;
4+
import datadog.context.propagation.CarrierSetter;
5+
import datadog.context.propagation.CarrierVisitor;
6+
import datadog.context.propagation.Propagator;
7+
import datadog.trace.api.TraceConfig;
8+
import datadog.trace.api.time.TimeSource;
9+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
10+
import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
11+
import datadog.trace.bootstrap.instrumentation.api.PathwayContext;
12+
import datadog.trace.bootstrap.instrumentation.api.TagContext;
13+
import java.util.function.Supplier;
14+
import javax.annotation.Nullable;
15+
import javax.annotation.ParametersAreNonnullByDefault;
16+
17+
// TODO Javadoc
18+
@ParametersAreNonnullByDefault
19+
public class DataStreamPropagator implements Propagator {
20+
private final Supplier<TraceConfig> traceConfigSupplier;
21+
private final TimeSource timeSource;
22+
private final long hashOfKnownTags;
23+
private final String serviceNameOverride;
24+
25+
public DataStreamPropagator(
26+
Supplier<TraceConfig> traceConfigSupplier,
27+
TimeSource timeSource,
28+
long hashOfKnownTags,
29+
String serviceNameOverride) {
30+
this.traceConfigSupplier = traceConfigSupplier;
31+
this.timeSource = timeSource;
32+
this.hashOfKnownTags = hashOfKnownTags;
33+
this.serviceNameOverride = serviceNameOverride;
34+
}
35+
36+
@Override
37+
public <C> void inject(Context context, C carrier, CarrierSetter<C> setter) {
38+
// TODO Still in CorePropagation, not migrated yet
39+
}
40+
41+
@Override
42+
public <C> Context extract(Context context, C carrier, CarrierVisitor<C> visitor) {
43+
// TODO Pathway context needs to be stored into its own context element
44+
// Get span context to store pathway context into
45+
TagContext spanContext = getSpanContextOrNull(context);
46+
PathwayContext pathwayContext;
47+
// Ensure if DSM is enabled and look for pathway context
48+
if (isDsmEnabled(spanContext)
49+
&& (pathwayContext = extractDsmPathwayContext(carrier, visitor)) != null) {
50+
// Store pathway context into span context
51+
if (spanContext == null) {
52+
spanContext = new TagContext();
53+
AgentSpan span = AgentSpan.fromSpanContext(spanContext);
54+
context = Context.root().with(span);
55+
}
56+
spanContext.withPathwayContext(pathwayContext);
57+
}
58+
return context;
59+
}
60+
61+
private TagContext getSpanContextOrNull(Context context) {
62+
AgentSpan extractedSpan = AgentSpan.fromContext(context);
63+
AgentSpanContext extractedSpanContext;
64+
if (extractedSpan != null
65+
&& (extractedSpanContext = extractedSpan.context()) instanceof TagContext) {
66+
return (TagContext) extractedSpanContext;
67+
}
68+
return null;
69+
}
70+
71+
private boolean isDsmEnabled(@Nullable TagContext tagContext) {
72+
TraceConfig traceConfig = tagContext == null ? null : tagContext.getTraceConfig();
73+
if (traceConfig == null) {
74+
traceConfig = this.traceConfigSupplier.get();
75+
}
76+
return traceConfig.isDataStreamsEnabled();
77+
}
78+
79+
private <C> PathwayContext extractDsmPathwayContext(C carrier, CarrierVisitor<C> visitor) {
80+
return DefaultPathwayContext.extract(
81+
carrier, visitor, this.timeSource, this.hashOfKnownTags, this.serviceNameOverride);
82+
}
83+
}

‎dd-trace-core/src/main/java/datadog/trace/core/datastreams/DataStreamsMonitoring.java

+8
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package datadog.trace.core.datastreams;
22

3+
import datadog.context.propagation.Propagator;
34
import datadog.trace.api.experimental.DataStreamsContextCarrier;
45
import datadog.trace.bootstrap.instrumentation.api.AgentDataStreamsMonitoring;
56
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -10,6 +11,13 @@
1011
public interface DataStreamsMonitoring extends AgentDataStreamsMonitoring, AutoCloseable {
1112
void start();
1213

14+
/**
15+
* Gets the propagator for DSM concern.
16+
*
17+
* @return The propagator for DSM concern.
18+
*/
19+
Propagator propagator();
20+
1321
/**
1422
* Get a context extractor that support {@link PathwayContext} extraction.
1523
*

‎dd-trace-core/src/main/java/datadog/trace/core/datastreams/DefaultDataStreamsMonitoring.java

+7
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414

1515
import datadog.communication.ddagent.DDAgentFeaturesDiscovery;
1616
import datadog.communication.ddagent.SharedCommunicationObjects;
17+
import datadog.context.propagation.Propagator;
1718
import datadog.trace.api.Config;
1819
import datadog.trace.api.TraceConfig;
1920
import datadog.trace.api.WellKnownTags;
@@ -200,6 +201,12 @@ public PathwayContext newPathwayContext() {
200201
}
201202
}
202203

204+
@Override
205+
public Propagator propagator() {
206+
return new DataStreamPropagator(
207+
this.traceConfigSupplier, this.timeSource, this.hashOfKnownTags, getThreadServiceName());
208+
}
209+
203210
@Override
204211
public HttpCodec.Extractor extractor(HttpCodec.Extractor delegate) {
205212
return new DataStreamContextExtractor(

‎dd-trace-core/src/main/java/datadog/trace/core/datastreams/DefaultPathwayContext.java

+7-8
Original file line numberDiff line numberDiff line change
@@ -7,10 +7,10 @@
77
import com.datadoghq.sketch.ddsketch.encoding.ByteArrayInput;
88
import com.datadoghq.sketch.ddsketch.encoding.GrowingByteArrayOutput;
99
import com.datadoghq.sketch.ddsketch.encoding.VarEncodingHelper;
10+
import datadog.context.propagation.CarrierVisitor;
1011
import datadog.trace.api.Config;
1112
import datadog.trace.api.WellKnownTags;
1213
import datadog.trace.api.time.TimeSource;
13-
import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
1414
import datadog.trace.bootstrap.instrumentation.api.PathwayContext;
1515
import datadog.trace.bootstrap.instrumentation.api.StatsPoint;
1616
import datadog.trace.util.FNV64Hash;
@@ -26,6 +26,7 @@
2626
import java.util.concurrent.TimeUnit;
2727
import java.util.concurrent.locks.Lock;
2828
import java.util.concurrent.locks.ReentrantLock;
29+
import java.util.function.BiConsumer;
2930
import java.util.function.Consumer;
3031
import org.slf4j.Logger;
3132
import org.slf4j.LoggerFactory;
@@ -267,7 +268,7 @@ public String toString() {
267268
}
268269
}
269270

270-
private static class PathwayContextExtractor implements AgentPropagation.KeyClassifier {
271+
private static class PathwayContextExtractor implements BiConsumer<String, String> {
271272
private final TimeSource timeSource;
272273
private final long hashOfKnownTags;
273274
private final String serviceNameOverride;
@@ -281,27 +282,25 @@ private static class PathwayContextExtractor implements AgentPropagation.KeyClas
281282
}
282283

283284
@Override
284-
public boolean accept(String key, String value) {
285+
public void accept(String key, String value) {
285286
if (PROPAGATION_KEY_BASE64.equalsIgnoreCase(key)) {
286287
try {
287288
extractedContext = decode(timeSource, hashOfKnownTags, serviceNameOverride, value);
288-
} catch (IOException e) {
289-
return false;
289+
} catch (IOException ignored) {
290290
}
291291
}
292-
return true;
293292
}
294293
}
295294

296295
static <C> DefaultPathwayContext extract(
297296
C carrier,
298-
AgentPropagation.ContextVisitor<C> getter,
297+
CarrierVisitor<C> getter,
299298
TimeSource timeSource,
300299
long hashOfKnownTags,
301300
String serviceNameOverride) {
302301
PathwayContextExtractor pathwayContextExtractor =
303302
new PathwayContextExtractor(timeSource, hashOfKnownTags, serviceNameOverride);
304-
getter.forEachKey(carrier, pathwayContextExtractor);
303+
getter.forEachKeyValue(carrier, pathwayContextExtractor);
305304
if (pathwayContextExtractor.extractedContext == null) {
306305
log.debug("No context extracted");
307306
} else {

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

+1-64
Original file line numberDiff line numberDiff line change
@@ -1,81 +1,23 @@
11
package datadog.trace.core.propagation;
22

3-
import datadog.trace.api.Config;
4-
import datadog.trace.api.TracePropagationStyle;
53
import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
64
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
7-
import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
8-
import datadog.trace.core.DDSpanContext;
95
import datadog.trace.core.datastreams.DataStreamContextInjector;
106
import java.util.LinkedHashMap;
11-
import java.util.Map;
127

138
public class CorePropagation implements AgentPropagation {
14-
private final HttpCodec.Injector injector;
15-
private final Map<TracePropagationStyle, HttpCodec.Injector> injectors;
169
private final DataStreamContextInjector dataStreamContextInjector;
17-
private final HttpCodec.Extractor extractor;
1810

1911
/**
2012
* Constructor
2113
*
22-
* @param extractor The context extractor.
23-
* @param defaultInjector The default injector when no {@link TracePropagationStyle} given.
24-
* @param injectors All the other injectors available for context injection.
2514
* @param dataStreamContextInjector The DSM context injector, as a specific object until generic
2615
* context injection is available.
2716
*/
28-
public CorePropagation(
29-
HttpCodec.Extractor extractor,
30-
HttpCodec.Injector defaultInjector,
31-
Map<TracePropagationStyle, HttpCodec.Injector> injectors,
32-
DataStreamContextInjector dataStreamContextInjector) {
33-
this.extractor = extractor;
34-
this.injector = defaultInjector;
35-
this.injectors = injectors;
17+
public CorePropagation(DataStreamContextInjector dataStreamContextInjector) {
3618
this.dataStreamContextInjector = dataStreamContextInjector;
3719
}
3820

39-
@Override
40-
public <C> void inject(final AgentSpan span, final C carrier, final Setter<C> setter) {
41-
inject(span.context(), carrier, setter, null);
42-
}
43-
44-
@Override
45-
public <C> void inject(AgentSpanContext context, C carrier, Setter<C> setter) {
46-
inject(context, carrier, setter, null);
47-
}
48-
49-
@Override
50-
public <C> void inject(AgentSpan span, C carrier, Setter<C> setter, TracePropagationStyle style) {
51-
inject(span.context(), carrier, setter, style);
52-
}
53-
54-
private <C> void inject(
55-
AgentSpanContext context, C carrier, Setter<C> setter, TracePropagationStyle style) {
56-
if (!(context instanceof DDSpanContext)) {
57-
return;
58-
}
59-
60-
final DDSpanContext ddSpanContext = (DDSpanContext) context;
61-
ddSpanContext.getTraceCollector().setSamplingPriorityIfNecessary();
62-
63-
/**
64-
* If the experimental appsec standalone feature is enabled and appsec propagation is disabled
65-
* (no ASM events), stop propagation
66-
*/
67-
if (Config.get().isAppSecStandaloneEnabled()
68-
&& !ddSpanContext.getPropagationTags().isAppsecPropagationEnabled()) {
69-
return;
70-
}
71-
72-
if (null == style) {
73-
injector.inject(ddSpanContext, carrier, setter);
74-
} else {
75-
injectors.get(style).inject(ddSpanContext, carrier, setter);
76-
}
77-
}
78-
7921
@Override
8022
public <C> void injectPathwayContext(
8123
AgentSpan span, C carrier, Setter<C> setter, LinkedHashMap<String, String> sortedTags) {
@@ -100,9 +42,4 @@ public <C> void injectPathwayContextWithoutSendingStats(
10042
this.dataStreamContextInjector.injectPathwayContextWithoutSendingStats(
10143
span, carrier, setter, sortedTags);
10244
}
103-
104-
@Override
105-
public <C> AgentSpanContext.Extracted extract(final C carrier, final ContextVisitor<C> getter) {
106-
return extractor.extract(carrier, getter);
107-
}
10845
}

‎dd-trace-core/src/test/groovy/datadog/trace/core/CoreTracerTest.groovy

-5
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,6 @@ import datadog.trace.common.sampling.Sampler
1919
import datadog.trace.common.writer.DDAgentWriter
2020
import datadog.trace.common.writer.ListWriter
2121
import datadog.trace.common.writer.LoggingWriter
22-
import datadog.trace.core.datastreams.DataStreamContextExtractor
23-
import datadog.trace.core.propagation.HttpCodec
2422
import datadog.trace.core.tagprocessor.TagsPostProcessorFactory
2523
import datadog.trace.core.test.DDCoreSpecification
2624
import okhttp3.HttpUrl
@@ -55,9 +53,6 @@ class CoreTracerTest extends DDCoreSpecification {
5553
tracer.writer instanceof DDAgentWriter
5654
tracer.statsDClient != null && tracer.statsDClient != StatsDClient.NO_OP
5755

58-
tracer.propagate().injector instanceof HttpCodec.CompoundInjector
59-
tracer.propagate().extractor instanceof DataStreamContextExtractor
60-
6156
cleanup:
6257
tracer.close()
6358
}

0 commit comments

Comments
 (0)