Skip to content

Commit bda8a1d

Browse files
PerfectSlayerPerfectSlayer
committed
feat(core): Migrate X-Ray inject calls to new API
1 parent 8c4249a commit bda8a1d

File tree

8 files changed

+80
-6
lines changed

8 files changed

+80
-6
lines changed

‎dd-java-agent/instrumentation/aws-java-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/AwsSdkClientDecorator.java

+1
Original file line numberDiff line numberDiff line change
@@ -333,6 +333,7 @@ protected int status(final Response response) {
333333
return response.getHttpResponse().getStatusCode();
334334
}
335335

336+
@ParametersAreNonnullByDefault
336337
@Override
337338
public void set(Request<?> carrier, String key, String value) {
338339
carrier.addHeader(key, value);

‎dd-java-agent/instrumentation/aws-java-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/TracingRequestHandler.java

+3-3
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
package datadog.trace.instrumentation.aws.v0;
22

3+
import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.XRAY_TRACING_CONCERN;
34
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
45
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.blackholeSpan;
5-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
66
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
77
import static datadog.trace.core.datastreams.TagsProcessor.DIRECTION_IN;
88
import static datadog.trace.core.datastreams.TagsProcessor.DIRECTION_TAG;
@@ -16,8 +16,8 @@
1616
import com.amazonaws.Response;
1717
import com.amazonaws.handlers.HandlerContextKey;
1818
import com.amazonaws.handlers.RequestHandler2;
19+
import datadog.context.propagation.Propagators;
1920
import datadog.trace.api.Config;
20-
import datadog.trace.api.TracePropagationStyle;
2121
import datadog.trace.bootstrap.ContextStore;
2222
import datadog.trace.bootstrap.instrumentation.api.AgentDataStreamsMonitoring;
2323
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -69,7 +69,7 @@ public void beforeRequest(final Request<?> request) {
6969
request.addHandlerContext(SPAN_CONTEXT_KEY, span);
7070
if (Config.get().isAwsPropagationEnabled()) {
7171
try {
72-
propagate().inject(span, request, DECORATE, TracePropagationStyle.XRAY);
72+
Propagators.forConcern(XRAY_TRACING_CONCERN).inject(span, request, DECORATE);
7373
} catch (Throwable e) {
7474
log.warn("Unable to inject trace header", e);
7575
}

‎dd-java-agent/instrumentation/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/AwsSdkClientDecorator.java

+2
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
import java.util.Optional;
3838
import java.util.Set;
3939
import javax.annotation.Nonnull;
40+
import javax.annotation.ParametersAreNonnullByDefault;
4041
import software.amazon.awssdk.awscore.AwsResponse;
4142
import software.amazon.awssdk.core.SdkBytes;
4243
import software.amazon.awssdk.core.SdkField;
@@ -443,6 +444,7 @@ protected int status(final SdkHttpResponse response) {
443444
return response.statusCode();
444445
}
445446

447+
@ParametersAreNonnullByDefault
446448
@Override
447449
public void set(SdkHttpRequest.Builder carrier, String key, String value) {
448450
carrier.putHeader(key, value);

‎dd-java-agent/instrumentation/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/TracingExecutionInterceptor.java

+3-3
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,14 @@
11
package datadog.trace.instrumentation.aws.v2;
22

3+
import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.XRAY_TRACING_CONCERN;
34
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
45
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.blackholeSpan;
5-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
66
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
77
import static datadog.trace.instrumentation.aws.v2.AwsSdkClientDecorator.AWS_LEGACY_TRACING;
88
import static datadog.trace.instrumentation.aws.v2.AwsSdkClientDecorator.DECORATE;
99

10+
import datadog.context.propagation.Propagators;
1011
import datadog.trace.api.Config;
11-
import datadog.trace.api.TracePropagationStyle;
1212
import datadog.trace.bootstrap.ContextStore;
1313
import datadog.trace.bootstrap.InstanceStore;
1414
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
@@ -71,7 +71,7 @@ public SdkHttpRequest modifyHttpRequest(
7171
final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE);
7272
if (span != null) {
7373
SdkHttpRequest.Builder requestBuilder = context.httpRequest().toBuilder();
74-
propagate().inject(span, requestBuilder, DECORATE, TracePropagationStyle.XRAY);
74+
Propagators.forConcern(XRAY_TRACING_CONCERN).inject(span, requestBuilder, DECORATE);
7575
return requestBuilder.build();
7676
}
7777
} catch (Throwable e) {

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

+3
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import static datadog.trace.api.DDTags.DSM_ENABLED;
77
import static datadog.trace.api.DDTags.PROFILING_CONTEXT_ENGINE;
88
import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.TRACING_CONCERN;
9+
import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.XRAY_TRACING_CONCERN;
910
import static datadog.trace.common.metrics.MetricsAggregatorFactory.createMetricsAggregator;
1011
import static datadog.trace.util.AgentThreadFactory.AGENT_THREAD_GROUP;
1112
import static datadog.trace.util.CollectionUtils.tryMakeImmutableMap;
@@ -90,6 +91,7 @@
9091
import datadog.trace.core.propagation.HttpCodec;
9192
import datadog.trace.core.propagation.PropagationTags;
9293
import datadog.trace.core.propagation.TracingPropagator;
94+
import datadog.trace.core.propagation.XRayPropagator;
9395
import datadog.trace.core.scopemanager.ContinuableScopeManager;
9496
import datadog.trace.core.taginterceptor.RuleFlags;
9597
import datadog.trace.core.taginterceptor.TagInterceptor;
@@ -724,6 +726,7 @@ private CoreTracer(
724726
new CorePropagation(builtExtractor, injector, injectors, dataStreamContextInjector);
725727

726728
Propagators.register(TRACING_CONCERN, new TracingPropagator(injector, extractor));
729+
Propagators.register(XRAY_TRACING_CONCERN, new XRayPropagator(config), false);
727730

728731
this.tagInterceptor =
729732
null == tagInterceptor ? new TagInterceptor(new RuleFlags(config)) : tagInterceptor;

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

+47
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
package datadog.trace.core.propagation;
2+
3+
import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext;
4+
5+
import datadog.context.Context;
6+
import datadog.context.propagation.CarrierSetter;
7+
import datadog.context.propagation.CarrierVisitor;
8+
import datadog.context.propagation.Propagator;
9+
import datadog.trace.api.Config;
10+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
11+
import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
12+
import datadog.trace.core.DDSpanContext;
13+
import javax.annotation.ParametersAreNonnullByDefault;
14+
15+
// TODO Javadoc
16+
@ParametersAreNonnullByDefault
17+
public class XRayPropagator implements Propagator {
18+
private final HttpCodec.Injector injector;
19+
20+
// TODO Javadoc
21+
public XRayPropagator(Config config) {
22+
this.injector = XRayHttpCodec.newInjector(config.getBaggageMapping());
23+
}
24+
25+
@Override
26+
public <C> void inject(Context context, C carrier, CarrierSetter<C> setter) {
27+
AgentSpan span;
28+
//noinspection ConstantValue
29+
if (context == null
30+
|| carrier == null
31+
|| setter == null
32+
|| (span = fromContext(context)) == null) {
33+
return;
34+
}
35+
AgentSpanContext spanContext = span.context();
36+
if (spanContext instanceof DDSpanContext) {
37+
DDSpanContext ddSpanContext = (DDSpanContext) spanContext;
38+
ddSpanContext.getTraceCollector().setSamplingPriorityIfNecessary();
39+
this.injector.inject(ddSpanContext, carrier, setter::set);
40+
}
41+
}
42+
43+
@Override
44+
public <C> Context extract(Context context, C carrier, CarrierVisitor<C> visitor) {
45+
return context;
46+
}
47+
}

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

+20
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import datadog.trace.core.test.DDCoreSpecification
1111

1212
import static datadog.trace.api.sampling.PrioritySampling.SAMPLER_KEEP
1313
import static datadog.trace.api.sampling.PrioritySampling.USER_DROP
14+
import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.XRAY_TRACING_CONCERN
1415

1516
class TracingPropagatorTest extends DDCoreSpecification {
1617
HttpCodec.Injector injector
@@ -136,4 +137,23 @@ class TracingPropagatorTest extends DDCoreSpecification {
136137
root.finish()
137138
tracer.close()
138139
}
140+
141+
def 'test AWS X-Ray propagator'() {
142+
setup:
143+
def tracer = tracerBuilder().build()
144+
def span = tracer.buildSpan('test', 'operation').start()
145+
def propagator = Propagators.forConcerns(XRAY_TRACING_CONCERN)
146+
def setter = Mock(CarrierSetter)
147+
def carrier = new Object()
148+
149+
when:
150+
propagator.inject(span, carrier, setter)
151+
152+
then:
153+
1 * setter.set(carrier, 'X-Amzn-Trace-Id', _)
154+
155+
cleanup:
156+
span.finish()
157+
tracer.close()
158+
}
139159
}

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

+1
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010

1111
public interface AgentPropagation {
1212
Concern TRACING_CONCERN = Concern.named("tracing");
13+
Concern XRAY_TRACING_CONCERN = Concern.named("tracing-xray");
1314

1415
<C> void inject(AgentSpan span, C carrier, Setter<C> setter);
1516

0 commit comments

Comments
 (0)