Skip to content

Commit 3f7c076

Browse files
committed
feat(dsm): Migrate DSM injection call to propagator API 1/3
1 parent 778cadd commit 3f7c076

File tree

35 files changed

+192
-215
lines changed

35 files changed

+192
-215
lines changed

dd-java-agent/instrumentation/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/AkkaHttpSingleRequestInstrumentation.java

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,8 @@
22

33
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
44
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
5-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
65
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
6+
import static datadog.trace.bootstrap.instrumentation.decorator.HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS;
77
import static datadog.trace.instrumentation.akkahttp.AkkaHttpClientDecorator.AKKA_CLIENT_REQUEST;
88
import static datadog.trace.instrumentation.akkahttp.AkkaHttpClientDecorator.DECORATE;
99
import static datadog.trace.instrumentation.akkahttp.AkkaHttpClientHelpers.AkkaHttpHeaders;
@@ -17,9 +17,9 @@
1717
import datadog.context.propagation.Propagators;
1818
import datadog.trace.agent.tooling.Instrumenter;
1919
import datadog.trace.agent.tooling.InstrumenterModule;
20+
import datadog.trace.api.datastreams.DataStreamsContext;
2021
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
2122
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
22-
import datadog.trace.bootstrap.instrumentation.decorator.HttpClientDecorator;
2323
import net.bytebuddy.asm.Advice;
2424
import scala.concurrent.Future;
2525

@@ -79,10 +79,8 @@ public static AgentScope methodEnter(
7979
DECORATE.onRequest(span, request);
8080

8181
if (request != null) {
82-
Propagators.defaultPropagator().inject(span, request, headers);
83-
propagate()
84-
.injectPathwayContext(
85-
span, request, headers, HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS);
82+
DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS);
83+
Propagators.defaultPropagator().inject(span.with(dsmContext), request, headers);
8684
// Request is immutable, so we have to assign new value once we update headers
8785
request = headers.getRequest();
8886
}

dd-java-agent/instrumentation/akka-http/akka-http-10.6/src/main/java11/datadog/trace/instrumentation/akkahttp106/SingleRequestAdvice.java

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,16 @@
11
package datadog.trace.instrumentation.akkahttp106;
22

33
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
4-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
54
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
5+
import static datadog.trace.bootstrap.instrumentation.decorator.HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS;
66

77
import akka.http.scaladsl.HttpExt;
88
import akka.http.scaladsl.model.HttpRequest;
99
import akka.http.scaladsl.model.HttpResponse;
1010
import datadog.context.propagation.Propagators;
11+
import datadog.trace.api.datastreams.DataStreamsContext;
1112
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
1213
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
13-
import datadog.trace.bootstrap.instrumentation.decorator.HttpClientDecorator;
1414
import net.bytebuddy.asm.Advice;
1515
import scala.concurrent.Future;
1616

@@ -29,10 +29,8 @@ public static AgentScope methodEnter(
2929
AkkaHttpClientDecorator.DECORATE.onRequest(span, request);
3030

3131
if (request != null) {
32-
Propagators.defaultPropagator().inject(span, request, headers);
33-
propagate()
34-
.injectPathwayContext(
35-
span, request, headers, HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS);
32+
DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS);
33+
Propagators.defaultPropagator().inject(span.with(dsmContext), request, headers);
3634
// Request is immutable, so we have to assign new value once we update headers
3735
request = headers.getRequest();
3836
}

dd-java-agent/instrumentation/apache-httpasyncclient-4/src/main/java/datadog/trace/instrumentation/apachehttpasyncclient/DelegatingRequestProducer.java

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,12 @@
11
package datadog.trace.instrumentation.apachehttpasyncclient;
22

3-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
3+
import static datadog.trace.bootstrap.instrumentation.decorator.HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS;
44
import static datadog.trace.instrumentation.apachehttpasyncclient.ApacheHttpAsyncClientDecorator.DECORATE;
55
import static datadog.trace.instrumentation.apachehttpasyncclient.HttpHeadersInjectAdapter.SETTER;
66

77
import datadog.context.propagation.Propagators;
8+
import datadog.trace.api.datastreams.DataStreamsContext;
89
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
9-
import datadog.trace.bootstrap.instrumentation.decorator.HttpClientDecorator;
1010
import java.io.IOException;
1111
import org.apache.http.HttpException;
1212
import org.apache.http.HttpHost;
@@ -35,9 +35,8 @@ public HttpRequest generateRequest() throws IOException, HttpException {
3535
final HttpRequest request = delegate.generateRequest();
3636
DECORATE.onRequest(span, new HostAndRequestAsHttpUriRequest(delegate.getTarget(), request));
3737

38-
Propagators.defaultPropagator().inject(span, request, SETTER);
39-
propagate()
40-
.injectPathwayContext(span, request, SETTER, HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS);
38+
DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS);
39+
Propagators.defaultPropagator().inject(span.with(dsmContext), request, SETTER);
4140

4241
return request;
4342
}

dd-java-agent/instrumentation/apache-httpclient-4/src/main/java/datadog/trace/instrumentation/apachehttpclient/HelperMethods.java

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,17 @@
11
package datadog.trace.instrumentation.apachehttpclient;
22

33
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
4-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
54
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
5+
import static datadog.trace.bootstrap.instrumentation.decorator.HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS;
66
import static datadog.trace.instrumentation.apachehttpclient.ApacheHttpClientDecorator.DECORATE;
77
import static datadog.trace.instrumentation.apachehttpclient.ApacheHttpClientDecorator.HTTP_REQUEST;
88
import static datadog.trace.instrumentation.apachehttpclient.HttpHeadersInjectAdapter.SETTER;
99

1010
import datadog.context.propagation.Propagators;
11+
import datadog.trace.api.datastreams.DataStreamsContext;
1112
import datadog.trace.bootstrap.CallDepthThreadLocalMap;
1213
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
1314
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
14-
import datadog.trace.bootstrap.instrumentation.decorator.HttpClientDecorator;
1515
import org.apache.http.HttpHost;
1616
import org.apache.http.HttpRequest;
1717
import org.apache.http.HttpResponse;
@@ -46,10 +46,8 @@ private static AgentScope activateHttpSpan(final HttpUriRequest request) {
4646

4747
// AWS calls are often signed, so we can't add headers without breaking the signature.
4848
if (!awsClientCall) {
49-
Propagators.defaultPropagator().inject(span, request, SETTER);
50-
propagate()
51-
.injectPathwayContext(
52-
span, request, SETTER, HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS);
49+
DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS);
50+
Propagators.defaultPropagator().inject(span.with(dsmContext), request, SETTER);
5351
}
5452

5553
return scope;

dd-java-agent/instrumentation/apache-httpclient-5/src/main/java/datadog/trace/instrumentation/apachehttpclient5/DelegatingRequestChannel.java

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,12 @@
11
package datadog.trace.instrumentation.apachehttpclient5;
22

3-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
3+
import static datadog.trace.bootstrap.instrumentation.decorator.HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS;
44
import static datadog.trace.instrumentation.apachehttpclient5.ApacheHttpClientDecorator.DECORATE;
55
import static datadog.trace.instrumentation.apachehttpclient5.HttpHeadersInjectAdapter.SETTER;
66

77
import datadog.context.propagation.Propagators;
8+
import datadog.trace.api.datastreams.DataStreamsContext;
89
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
9-
import datadog.trace.bootstrap.instrumentation.decorator.HttpClientDecorator;
1010
import java.io.IOException;
1111
import org.apache.hc.core5.http.EntityDetails;
1212
import org.apache.hc.core5.http.HttpException;
@@ -28,9 +28,8 @@ public void sendRequest(HttpRequest request, EntityDetails entityDetails, HttpCo
2828
throws HttpException, IOException {
2929
DECORATE.onRequest(span, request);
3030

31-
Propagators.defaultPropagator().inject(span, request, SETTER);
32-
propagate()
33-
.injectPathwayContext(span, request, SETTER, HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS);
31+
DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS);
32+
Propagators.defaultPropagator().inject(span.with(dsmContext), request, SETTER);
3433
delegate.sendRequest(request, entityDetails, context);
3534
}
3635
}

dd-java-agent/instrumentation/apache-httpclient-5/src/main/java/datadog/trace/instrumentation/apachehttpclient5/HelperMethods.java

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,17 @@
11
package datadog.trace.instrumentation.apachehttpclient5;
22

33
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
4-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
54
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
5+
import static datadog.trace.bootstrap.instrumentation.decorator.HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS;
66
import static datadog.trace.instrumentation.apachehttpclient5.ApacheHttpClientDecorator.DECORATE;
77
import static datadog.trace.instrumentation.apachehttpclient5.ApacheHttpClientDecorator.HTTP_REQUEST;
88
import static datadog.trace.instrumentation.apachehttpclient5.HttpHeadersInjectAdapter.SETTER;
99

1010
import datadog.context.propagation.Propagators;
11+
import datadog.trace.api.datastreams.DataStreamsContext;
1112
import datadog.trace.bootstrap.CallDepthThreadLocalMap;
1213
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
1314
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
14-
import datadog.trace.bootstrap.instrumentation.decorator.HttpClientDecorator;
1515
import org.apache.hc.client5.http.classic.HttpClient;
1616
import org.apache.hc.core5.http.HttpHost;
1717
import org.apache.hc.core5.http.HttpRequest;
@@ -46,10 +46,8 @@ private static AgentScope activateHttpSpan(final HttpRequest request) {
4646
final boolean awsClientCall = request.containsHeader("amz-sdk-invocation-id");
4747
// AWS calls are often signed, so we can't add headers without breaking the signature.
4848
if (!awsClientCall) {
49-
Propagators.defaultPropagator().inject(span, request, SETTER);
50-
propagate()
51-
.injectPathwayContext(
52-
span, request, SETTER, HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS);
49+
DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS);
50+
Propagators.defaultPropagator().inject(span.with(dsmContext), request, SETTER);
5351
}
5452

5553
return scope;

dd-java-agent/instrumentation/armeria-grpc/src/main/java/datadog/trace/instrumentation/armeria/grpc/client/ClientCallImplInstrumentation.java

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@
33
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
44
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
55
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
6-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
76
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
87
import static datadog.trace.instrumentation.armeria.grpc.client.GrpcClientDecorator.CLIENT_PATHWAY_EDGE_TAGS;
98
import static datadog.trace.instrumentation.armeria.grpc.client.GrpcClientDecorator.DECORATE;
@@ -21,6 +20,7 @@
2120
import datadog.trace.agent.tooling.InstrumenterModule;
2221
import datadog.trace.agent.tooling.muzzle.Reference;
2322
import datadog.trace.api.InstrumenterConfig;
23+
import datadog.trace.api.datastreams.DataStreamsContext;
2424
import datadog.trace.bootstrap.InstrumentationContext;
2525
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
2626
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -121,8 +121,8 @@ public static <T> AgentScope before(
121121
if (null != responseListener && null != headers) {
122122
span = InstrumentationContext.get(ClientCall.class, AgentSpan.class).get(call);
123123
if (null != span) {
124-
Propagators.defaultPropagator().inject(span, headers, SETTER);
125-
propagate().injectPathwayContext(span, headers, SETTER, CLIENT_PATHWAY_EDGE_TAGS);
124+
DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS);
125+
Propagators.defaultPropagator().inject(span.with(dsmContext), headers, SETTER);
126126
return activateSpan(span);
127127
}
128128
}

dd-java-agent/instrumentation/aws-java-eventbridge-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/eventbridge/EventBridgeInterceptor.java

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
package datadog.trace.instrumentation.aws.v2.eventbridge;
22

3-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
43
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.traceConfig;
54
import static datadog.trace.core.datastreams.TagsProcessor.BUS_TAG;
65
import static datadog.trace.core.datastreams.TagsProcessor.DIRECTION_OUT;
@@ -9,6 +8,7 @@
98
import static datadog.trace.instrumentation.aws.v2.eventbridge.TextMapInjectAdapter.SETTER;
109

1110
import datadog.context.propagation.Propagators;
11+
import datadog.trace.api.datastreams.DataStreamsContext;
1212
import datadog.trace.bootstrap.InstanceStore;
1313
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
1414
import datadog.trace.bootstrap.instrumentation.api.PathwayContext;
@@ -86,12 +86,13 @@ private String getTraceContextToInject(
8686
StringBuilder jsonBuilder = new StringBuilder();
8787
jsonBuilder.append('{');
8888

89-
// Inject trace context
90-
Propagators.defaultPropagator().inject(span, jsonBuilder, SETTER);
91-
89+
// Inject context
90+
datadog.context.Context context = span;
9291
if (traceConfig().isDataStreamsEnabled()) {
93-
propagate().injectPathwayContext(span, jsonBuilder, SETTER, getTags(eventBusName));
92+
DataStreamsContext dsmContext = DataStreamsContext.fromTags(getTags(eventBusName));
93+
context = context.with(dsmContext);
9494
}
95+
Propagators.defaultPropagator().inject(context, jsonBuilder, SETTER);
9596

9697
// Add bus name and start time
9798
jsonBuilder

dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/SnsInterceptor.java

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
package datadog.trace.instrumentation.aws.v1.sns;
22

3-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
43
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.traceConfig;
54
import static datadog.trace.core.datastreams.TagsProcessor.DIRECTION_OUT;
65
import static datadog.trace.core.datastreams.TagsProcessor.DIRECTION_TAG;
@@ -14,7 +13,9 @@
1413
import com.amazonaws.services.sns.model.PublishBatchRequest;
1514
import com.amazonaws.services.sns.model.PublishBatchRequestEntry;
1615
import com.amazonaws.services.sns.model.PublishRequest;
16+
import datadog.context.Context;
1717
import datadog.context.propagation.Propagators;
18+
import datadog.trace.api.datastreams.DataStreamsContext;
1819
import datadog.trace.bootstrap.ContextStore;
1920
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
2021
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
@@ -37,10 +38,12 @@ private ByteBuffer getMessageAttributeValueToInject(
3738
final AgentSpan span = newSpan(request);
3839
StringBuilder jsonBuilder = new StringBuilder();
3940
jsonBuilder.append('{');
40-
Propagators.defaultPropagator().inject(span, jsonBuilder, SETTER);
41+
Context context = span;
4142
if (traceConfig().isDataStreamsEnabled()) {
42-
propagate().injectPathwayContext(span, jsonBuilder, SETTER, getTags(snsTopicName));
43+
DataStreamsContext dsmContext = DataStreamsContext.fromTags(getTags(snsTopicName));
44+
context = context.with(dsmContext);
4345
}
46+
Propagators.defaultPropagator().inject(context, jsonBuilder, SETTER);
4447
jsonBuilder.setLength(jsonBuilder.length() - 1); // Remove the last comma
4548
jsonBuilder.append('}');
4649
return ByteBuffer.wrap(jsonBuilder.toString().getBytes(StandardCharsets.UTF_8));

dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
package datadog.trace.instrumentation.aws.v2.sns;
22

3-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
43
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.traceConfig;
54
import static datadog.trace.core.datastreams.TagsProcessor.DIRECTION_OUT;
65
import static datadog.trace.core.datastreams.TagsProcessor.DIRECTION_TAG;
@@ -9,6 +8,7 @@
98
import static datadog.trace.instrumentation.aws.v2.sns.TextMapInjectAdapter.SETTER;
109

1110
import datadog.context.propagation.Propagators;
11+
import datadog.trace.api.datastreams.DataStreamsContext;
1212
import datadog.trace.bootstrap.InstanceStore;
1313
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
1414
import java.nio.charset.StandardCharsets;
@@ -38,10 +38,12 @@ private SdkBytes getMessageAttributeValueToInject(
3838
final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE);
3939
StringBuilder jsonBuilder = new StringBuilder();
4040
jsonBuilder.append('{');
41-
Propagators.defaultPropagator().inject(span, jsonBuilder, SETTER);
41+
datadog.context.Context context = span;
4242
if (traceConfig().isDataStreamsEnabled()) {
43-
propagate().injectPathwayContext(span, jsonBuilder, SETTER, getTags(snsTopicName));
43+
DataStreamsContext dsmContext = DataStreamsContext.fromTags(getTags(snsTopicName));
44+
context = context.with(dsmContext);
4445
}
46+
Propagators.defaultPropagator().inject(context, jsonBuilder, SETTER);
4547
jsonBuilder.setLength(jsonBuilder.length() - 1); // Remove the last comma
4648
jsonBuilder.append('}');
4749
return SdkBytes.fromString(jsonBuilder.toString(), StandardCharsets.UTF_8);

dd-java-agent/instrumentation/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/SqsInterceptor.java

Lines changed: 18 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
package datadog.trace.instrumentation.aws.v1.sqs;
22

3-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
3+
import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.DSM_CONCERN;
44
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
55
import static datadog.trace.bootstrap.instrumentation.api.PathwayContext.DATADOG_KEY;
66
import static datadog.trace.bootstrap.instrumentation.api.URIUtils.urlFileName;
@@ -16,6 +16,10 @@
1616
import com.amazonaws.services.sqs.model.SendMessageBatchRequest;
1717
import com.amazonaws.services.sqs.model.SendMessageBatchRequestEntry;
1818
import com.amazonaws.services.sqs.model.SendMessageRequest;
19+
import datadog.context.Context;
20+
import datadog.context.propagation.Propagator;
21+
import datadog.context.propagation.Propagators;
22+
import datadog.trace.api.datastreams.DataStreamsContext;
1923
import datadog.trace.bootstrap.ContextStore;
2024
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
2125
import java.util.LinkedHashMap;
@@ -36,23 +40,21 @@ public AmazonWebServiceRequest beforeMarshalling(AmazonWebServiceRequest request
3640
String queueUrl = smRequest.getQueueUrl();
3741
if (queueUrl == null) return request;
3842

39-
LinkedHashMap<String, String> sortedTags = getTags(queueUrl);
40-
41-
final AgentSpan span = newSpan(request);
43+
Propagator dsmPropagator = Propagators.forConcern(DSM_CONCERN);
44+
Context context = newContext(request, queueUrl);
4245
// note: modifying message attributes has to be done before marshalling, otherwise the changes
4346
// are not reflected in the actual request (and the MD5 check on send will fail).
44-
propagate().injectPathwayContext(span, smRequest.getMessageAttributes(), SETTER, sortedTags);
47+
dsmPropagator.inject(context, smRequest.getMessageAttributes(), SETTER);
4548
} else if (request instanceof SendMessageBatchRequest) {
4649
SendMessageBatchRequest smbRequest = (SendMessageBatchRequest) request;
4750

4851
String queueUrl = smbRequest.getQueueUrl();
4952
if (queueUrl == null) return request;
5053

51-
LinkedHashMap<String, String> sortedTags = getTags(queueUrl);
52-
53-
final AgentSpan span = newSpan(request);
54+
Propagator dsmPropagator = Propagators.forConcern(DSM_CONCERN);
55+
Context context = newContext(request, queueUrl);
5456
for (SendMessageBatchRequestEntry entry : smbRequest.getEntries()) {
55-
propagate().injectPathwayContext(span, entry.getMessageAttributes(), SETTER, sortedTags);
57+
dsmPropagator.inject(context, entry.getMessageAttributes(), SETTER);
5658
}
5759
} else if (request instanceof ReceiveMessageRequest) {
5860
ReceiveMessageRequest rmRequest = (ReceiveMessageRequest) request;
@@ -64,8 +66,14 @@ public AmazonWebServiceRequest beforeMarshalling(AmazonWebServiceRequest request
6466
return request;
6567
}
6668

69+
private Context newContext(AmazonWebServiceRequest request, String queueUrl) {
70+
AgentSpan span = newSpan(request);
71+
DataStreamsContext dsmContext = DataStreamsContext.fromTags(getTags(queueUrl));
72+
return span.with(dsmContext);
73+
}
74+
6775
private AgentSpan newSpan(AmazonWebServiceRequest request) {
68-
final AgentSpan span = startSpan("aws.sqs.send");
76+
final AgentSpan span = startSpan("sqs", "aws.sqs.send");
6977
// pass the span to TracingRequestHandler in the sdk instrumentation where it'll be enriched &
7078
// activated
7179
contextStore.put(request, span);

0 commit comments

Comments
 (0)