Skip to content

Commit 6ce5559

Browse files
PerfectSlayerPerfectSlayer
authored
Refactor AWS instrumentation around Context (#9384)
1 parent 8799a82 commit 6ce5559

File tree

8 files changed

+95
-76
lines changed

8 files changed

+95
-76
lines changed

‎dd-java-agent/instrumentation/aws-java/aws-java-dynamodb-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/dynamodb/DynamoDbInterceptor.java‎

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

3+
import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext;
4+
5+
import datadog.context.Context;
36
import datadog.trace.api.Config;
47
import datadog.trace.bootstrap.InstanceStore;
58
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
69
import java.util.Map;
710
import org.slf4j.Logger;
811
import org.slf4j.LoggerFactory;
912
import software.amazon.awssdk.core.SdkRequest;
10-
import software.amazon.awssdk.core.interceptor.Context;
13+
import software.amazon.awssdk.core.interceptor.Context.AfterExecution;
1114
import software.amazon.awssdk.core.interceptor.ExecutionAttribute;
1215
import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
1316
import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
@@ -18,20 +21,20 @@
1821
public class DynamoDbInterceptor implements ExecutionInterceptor {
1922
private static final Logger log = LoggerFactory.getLogger(DynamoDbInterceptor.class);
2023

21-
public static final ExecutionAttribute<AgentSpan> SPAN_ATTRIBUTE =
24+
public static final ExecutionAttribute<Context> CONTEXT_ATTRIBUTE =
2225
InstanceStore.of(ExecutionAttribute.class)
23-
.putIfAbsent("DatadogSpan", () -> new ExecutionAttribute<>("DatadogSpan"));
26+
.putIfAbsent("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext"));
2427

2528
private static final boolean CAN_ADD_SPAN_POINTERS = Config.get().isAddSpanPointers("aws");
2629

2730
@Override
28-
public void afterExecution(
29-
Context.AfterExecution context, ExecutionAttributes executionAttributes) {
31+
public void afterExecution(AfterExecution context, ExecutionAttributes executionAttributes) {
3032
if (!CAN_ADD_SPAN_POINTERS) {
3133
return;
3234
}
3335

34-
AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE);
36+
Context ddContext = executionAttributes.getAttribute(CONTEXT_ATTRIBUTE);
37+
AgentSpan span = fromContext(ddContext);
3538
if (span == null) {
3639
log.debug("Unable to find DynamoDb request span. Not creating span pointer.");
3740
return;

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

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -5,17 +5,17 @@
55
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.traceConfig;
66
import static datadog.trace.instrumentation.aws.v2.eventbridge.TextMapInjectAdapter.SETTER;
77

8+
import datadog.context.Context;
89
import datadog.trace.api.datastreams.DataStreamsContext;
910
import datadog.trace.api.datastreams.DataStreamsTags;
1011
import datadog.trace.api.datastreams.PathwayContext;
1112
import datadog.trace.bootstrap.InstanceStore;
12-
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
1313
import java.util.ArrayList;
1414
import java.util.List;
1515
import org.slf4j.Logger;
1616
import org.slf4j.LoggerFactory;
1717
import software.amazon.awssdk.core.SdkRequest;
18-
import software.amazon.awssdk.core.interceptor.Context;
18+
import software.amazon.awssdk.core.interceptor.Context.ModifyRequest;
1919
import software.amazon.awssdk.core.interceptor.ExecutionAttribute;
2020
import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
2121
import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
@@ -25,16 +25,15 @@
2525
public class EventBridgeInterceptor implements ExecutionInterceptor {
2626
private static final Logger log = LoggerFactory.getLogger(EventBridgeInterceptor.class);
2727

28-
public static final ExecutionAttribute<AgentSpan> SPAN_ATTRIBUTE =
28+
public static final ExecutionAttribute<Context> CONTEXT_ATTRIBUTE =
2929
InstanceStore.of(ExecutionAttribute.class)
30-
.putIfAbsent("DatadogSpan", () -> new ExecutionAttribute<>("DatadogSpan"));
30+
.putIfAbsent("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext"));
3131

3232
private static final String START_TIME_KEY = "x-datadog-start-time";
3333
private static final String RESOURCE_NAME_KEY = "x-datadog-resource-name";
3434

3535
@Override
36-
public SdkRequest modifyRequest(
37-
Context.ModifyRequest context, ExecutionAttributes executionAttributes) {
36+
public SdkRequest modifyRequest(ModifyRequest context, ExecutionAttributes executionAttributes) {
3837
if (!(context.request() instanceof PutEventsRequest)) {
3938
return context.request();
4039
}
@@ -79,12 +78,11 @@ public SdkRequest modifyRequest(
7978

8079
private String getTraceContextToInject(
8180
ExecutionAttributes executionAttributes, String eventBusName, long startTime) {
82-
final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE);
81+
Context context = executionAttributes.getAttribute(CONTEXT_ATTRIBUTE);
8382
StringBuilder jsonBuilder = new StringBuilder();
8483
jsonBuilder.append('{');
8584

8685
// Inject context
87-
datadog.context.Context context = span;
8886
if (traceConfig().isDataStreamsEnabled()) {
8987
DataStreamsTags tags = DataStreamsTags.createWithBus(OUTBOUND, eventBusName);
9088
DataStreamsContext dsmContext = DataStreamsContext.fromTags(tags);

‎dd-java-agent/instrumentation/aws-java/aws-java-s3-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/s3/S3Interceptor.java‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,13 @@
22

33
import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.S3_ETAG;
44

5+
import datadog.context.Context;
56
import datadog.trace.api.Config;
67
import datadog.trace.bootstrap.InstanceStore;
78
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
89
import org.slf4j.Logger;
910
import org.slf4j.LoggerFactory;
10-
import software.amazon.awssdk.core.interceptor.Context;
11+
import software.amazon.awssdk.core.interceptor.Context.AfterExecution;
1112
import software.amazon.awssdk.core.interceptor.ExecutionAttribute;
1213
import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
1314
import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
@@ -18,20 +19,20 @@
1819
public class S3Interceptor implements ExecutionInterceptor {
1920
private static final Logger log = LoggerFactory.getLogger(S3Interceptor.class);
2021

21-
public static final ExecutionAttribute<AgentSpan> SPAN_ATTRIBUTE =
22+
public static final ExecutionAttribute<Context> CONTEXT_ATTRIBUTE =
2223
InstanceStore.of(ExecutionAttribute.class)
23-
.putIfAbsent("DatadogSpan", () -> new ExecutionAttribute<>("DatadogSpan"));
24+
.putIfAbsent("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext"));
2425

2526
private static final boolean CAN_ADD_SPAN_POINTERS = Config.get().isAddSpanPointers("aws");
2627

2728
@Override
28-
public void afterExecution(
29-
Context.AfterExecution context, ExecutionAttributes executionAttributes) {
29+
public void afterExecution(AfterExecution context, ExecutionAttributes executionAttributes) {
3030
if (!CAN_ADD_SPAN_POINTERS) {
3131
return;
3232
}
3333

34-
AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE);
34+
Context ddContext = executionAttributes.getAttribute(CONTEXT_ATTRIBUTE);
35+
AgentSpan span = AgentSpan.fromContext(ddContext);
3536
if (span == null) {
3637
log.debug("Unable to find S3 request span. Not creating span pointer.");
3738
return;

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

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,10 @@
55
import static datadog.trace.api.datastreams.DataStreamsTags.Direction.OUTBOUND;
66
import static datadog.trace.api.datastreams.DataStreamsTags.create;
77
import static datadog.trace.api.datastreams.DataStreamsTags.createWithDataset;
8+
import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext;
89
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.traceConfig;
910

11+
import datadog.context.Context;
1012
import datadog.context.propagation.CarrierSetter;
1113
import datadog.trace.api.Config;
1214
import datadog.trace.api.ConfigDefaults;
@@ -113,11 +115,12 @@ public CharSequence spanName(final ExecutionAttributes attributes) {
113115
"aws", attributes.getAttribute(SdkExecutionAttribute.SERVICE_NAME), s));
114116
}
115117

116-
public AgentSpan onSdkRequest(
117-
final AgentSpan span,
118+
public Context onSdkRequest(
119+
final Context context,
118120
final SdkRequest request,
119121
final SdkHttpRequest httpRequest,
120122
final ExecutionAttributes attributes) {
123+
final AgentSpan span = fromContext(context);
121124
final String awsServiceName = attributes.getAttribute(SdkExecutionAttribute.SERVICE_NAME);
122125
final String awsOperationName = attributes.getAttribute(SdkExecutionAttribute.OPERATION_NAME);
123126
onOperation(span, awsServiceName, awsOperationName);
@@ -224,7 +227,7 @@ public AgentSpan onSdkRequest(
224227
span.setTag(DDTags.PEER_SERVICE_SOURCE, "peer.service");
225228
}
226229

227-
return span;
230+
return context;
228231
}
229232

230233
private static AgentSpan onOperation(
@@ -315,12 +318,13 @@ private static void setTableName(AgentSpan span, String name) {
315318
setPeerService(span, InstrumentationTags.AWS_TABLE_NAME, name);
316319
}
317320

318-
public AgentSpan onSdkResponse(
319-
final AgentSpan span,
321+
public Context onSdkResponse(
322+
final Context context,
320323
final SdkResponse response,
321324
final SdkHttpResponse httpResponse,
322325
final ExecutionAttributes attributes) {
323326

327+
final AgentSpan span = fromContext(context);
324328
Config config = Config.get();
325329
String serviceName = attributes.getAttribute(SdkExecutionAttribute.SERVICE_NAME);
326330
if (config.isCloudResponsePayloadTaggingEnabled()

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

Lines changed: 38 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,15 @@
11
package datadog.trace.instrumentation.aws.v2;
22

33
import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.XRAY_TRACING_CONCERN;
4+
import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext;
45
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
56
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpanWithoutScope;
67
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.blackholeSpan;
78
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
89
import static datadog.trace.instrumentation.aws.v2.AwsSdkClientDecorator.AWS_LEGACY_TRACING;
910
import static datadog.trace.instrumentation.aws.v2.AwsSdkClientDecorator.DECORATE;
1011

12+
import datadog.context.Context;
1113
import datadog.context.propagation.Propagators;
1214
import datadog.trace.api.Config;
1315
import datadog.trace.bootstrap.ContextStore;
@@ -19,7 +21,12 @@
1921
import org.slf4j.LoggerFactory;
2022
import software.amazon.awssdk.core.SdkRequest;
2123
import software.amazon.awssdk.core.SdkResponse;
22-
import software.amazon.awssdk.core.interceptor.Context;
24+
import software.amazon.awssdk.core.interceptor.Context.AfterExecution;
25+
import software.amazon.awssdk.core.interceptor.Context.AfterMarshalling;
26+
import software.amazon.awssdk.core.interceptor.Context.BeforeExecution;
27+
import software.amazon.awssdk.core.interceptor.Context.BeforeTransmission;
28+
import software.amazon.awssdk.core.interceptor.Context.FailedExecution;
29+
import software.amazon.awssdk.core.interceptor.Context.ModifyHttpRequest;
2330
import software.amazon.awssdk.core.interceptor.ExecutionAttribute;
2431
import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
2532
import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
@@ -28,9 +35,9 @@
2835
/** AWS request execution interceptor */
2936
public class TracingExecutionInterceptor implements ExecutionInterceptor {
3037

31-
public static final ExecutionAttribute<AgentSpan> SPAN_ATTRIBUTE =
38+
public static final ExecutionAttribute<Context> CONTEXT_ATTRIBUTE =
3239
InstanceStore.of(ExecutionAttribute.class)
33-
.putIfAbsent("DatadogSpan", () -> new ExecutionAttribute<>("DatadogSpan"));
40+
.putIfAbsent("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext"));
3441

3542
private static final Logger log = LoggerFactory.getLogger(TracingExecutionInterceptor.class);
3643

@@ -42,37 +49,40 @@ public TracingExecutionInterceptor(ContextStore<Object, String> responseQueueSto
4249

4350
@Override
4451
public void beforeExecution(
45-
final Context.BeforeExecution context, final ExecutionAttributes executionAttributes) {
52+
final BeforeExecution context, final ExecutionAttributes executionAttributes) {
4653
if (!AWS_LEGACY_TRACING && isPollingRequest(context.request())) {
4754
return; // SQS messages spans are created by aws-java-sqs-2.0
4855
}
4956

50-
final AgentSpan span = startSpan(DECORATE.spanName(executionAttributes));
57+
final AgentSpan span = startSpan("aws-sdk", DECORATE.spanName(executionAttributes));
58+
// TODO If DSM is enabled, add DSM context here too
5159
DECORATE.afterStart(span);
52-
executionAttributes.putAttribute(SPAN_ATTRIBUTE, span);
60+
executionAttributes.putAttribute(CONTEXT_ATTRIBUTE, span);
5361
}
5462

5563
@Override
5664
public void afterMarshalling(
57-
final Context.AfterMarshalling context, final ExecutionAttributes executionAttributes) {
58-
final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE);
59-
if (span != null) {
65+
final AfterMarshalling context, final ExecutionAttributes executionAttributes) {
66+
final Context ddContext = executionAttributes.getAttribute(CONTEXT_ATTRIBUTE);
67+
final AgentSpan span = fromContext(ddContext);
68+
if (context != null && span != null) {
6069
try (AgentScope ignored = activateSpan(span)) {
6170
DECORATE.onRequest(span, context.httpRequest());
62-
DECORATE.onSdkRequest(span, context.request(), context.httpRequest(), executionAttributes);
71+
DECORATE.onSdkRequest(
72+
ddContext, context.request(), context.httpRequest(), executionAttributes);
6373
}
6474
}
6575
}
6676

6777
@Override
6878
public SdkHttpRequest modifyHttpRequest(
69-
Context.ModifyHttpRequest context, ExecutionAttributes executionAttributes) {
79+
ModifyHttpRequest context, ExecutionAttributes executionAttributes) {
7080
if (Config.get().isAwsPropagationEnabled()) {
7181
try {
72-
final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE);
73-
if (span != null) {
82+
final Context ddContext = executionAttributes.getAttribute(CONTEXT_ATTRIBUTE);
83+
if (ddContext != null) {
7484
SdkHttpRequest.Builder requestBuilder = context.httpRequest().toBuilder();
75-
Propagators.forConcern(XRAY_TRACING_CONCERN).inject(span, requestBuilder, DECORATE);
85+
Propagators.forConcern(XRAY_TRACING_CONCERN).inject(ddContext, requestBuilder, DECORATE);
7686
return requestBuilder.build();
7787
}
7888
} catch (Throwable e) {
@@ -84,12 +94,13 @@ public SdkHttpRequest modifyHttpRequest(
8494

8595
@Override
8696
public void beforeTransmission(
87-
final Context.BeforeTransmission context, final ExecutionAttributes executionAttributes) {
97+
final BeforeTransmission context, final ExecutionAttributes executionAttributes) {
8898
final AgentSpan span;
8999
if (!AWS_LEGACY_TRACING) {
90100
span = blackholeSpan();
91101
} else {
92-
span = executionAttributes.getAttribute(SPAN_ATTRIBUTE);
102+
final Context ddContext = executionAttributes.getAttribute(CONTEXT_ATTRIBUTE);
103+
span = fromContext(ddContext);
93104
}
94105
if (span != null) {
95106
// This scope will be closed by AwsHttpClientInstrumentation since ExecutionInterceptor API
@@ -100,10 +111,11 @@ public void beforeTransmission(
100111

101112
@Override
102113
public void afterExecution(
103-
final Context.AfterExecution context, final ExecutionAttributes executionAttributes) {
104-
final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE);
114+
final AfterExecution context, final ExecutionAttributes executionAttributes) {
115+
final Context ddContext = executionAttributes.getAttribute(CONTEXT_ATTRIBUTE);
116+
final AgentSpan span = fromContext(ddContext);
105117
if (span != null) {
106-
executionAttributes.putAttribute(SPAN_ATTRIBUTE, null);
118+
executionAttributes.putAttribute(CONTEXT_ATTRIBUTE, null);
107119
// Call onResponse on both types of responses:
108120
DECORATE.onSdkResponse(span, context.response(), context.httpResponse(), executionAttributes);
109121
DECORATE.onResponse(span, context.httpResponse());
@@ -121,14 +133,16 @@ public void afterExecution(
121133

122134
@Override
123135
public void onExecutionFailure(
124-
final Context.FailedExecution context, final ExecutionAttributes executionAttributes) {
125-
final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE);
126-
if (span != null) {
127-
executionAttributes.putAttribute(SPAN_ATTRIBUTE, null);
136+
final FailedExecution context, final ExecutionAttributes executionAttributes) {
137+
final Context ddContext = executionAttributes.getAttribute(CONTEXT_ATTRIBUTE);
138+
final AgentSpan span = fromContext(ddContext);
139+
if (ddContext != null && span != null) {
140+
executionAttributes.putAttribute(CONTEXT_ATTRIBUTE, null);
128141
Optional<SdkResponse> responseOpt = context.response();
129142
if (responseOpt.isPresent()) {
130143
SdkResponse response = responseOpt.get();
131-
DECORATE.onSdkResponse(span, response, response.sdkHttpResponse(), executionAttributes);
144+
DECORATE.onSdkResponse(
145+
ddContext, response, response.sdkHttpResponse(), executionAttributes);
132146
DECORATE.onResponse(span, response.sdkHttpResponse());
133147
if (span.isError()) {
134148
DECORATE.onError(span, context.exception());

‎dd-java-agent/instrumentation/aws-java/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnInterceptor.java‎

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,12 @@
11
package datadog.trace.instrumentation.aws.v2.sfn;
22

3+
import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext;
4+
5+
import datadog.context.Context;
36
import datadog.trace.bootstrap.InstanceStore;
47
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
58
import software.amazon.awssdk.core.SdkRequest;
6-
import software.amazon.awssdk.core.interceptor.Context;
9+
import software.amazon.awssdk.core.interceptor.Context.ModifyRequest;
710
import software.amazon.awssdk.core.interceptor.ExecutionAttribute;
811
import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
912
import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
@@ -12,15 +15,14 @@
1215

1316
public class SfnInterceptor implements ExecutionInterceptor {
1417

15-
public static final ExecutionAttribute<AgentSpan> SPAN_ATTRIBUTE =
18+
public static final ExecutionAttribute<Context> CONTEXT_ATTRIBUTE =
1619
InstanceStore.of(ExecutionAttribute.class)
17-
.putIfAbsent("DatadogSpan", () -> new ExecutionAttribute<>("DatadogSpan"));
20+
.putIfAbsent("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext"));
1821

1922
public SfnInterceptor() {}
2023

2124
@Override
22-
public SdkRequest modifyRequest(
23-
Context.ModifyRequest context, ExecutionAttributes executionAttributes) {
25+
public SdkRequest modifyRequest(ModifyRequest context, ExecutionAttributes executionAttributes) {
2426
try {
2527
return modifyRequestImpl(context, executionAttributes);
2628
} catch (Exception e) {
@@ -29,8 +31,9 @@ public SdkRequest modifyRequest(
2931
}
3032

3133
public SdkRequest modifyRequestImpl(
32-
Context.ModifyRequest context, ExecutionAttributes executionAttributes) {
33-
final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE);
34+
ModifyRequest context, ExecutionAttributes executionAttributes) {
35+
final Context ddContext = executionAttributes.getAttribute(CONTEXT_ATTRIBUTE);
36+
final AgentSpan span = fromContext(ddContext);
3437
// StartExecutionRequest
3538
if (context.request() instanceof StartExecutionRequest) {
3639
StartExecutionRequest request = (StartExecutionRequest) context.request();

0 commit comments

Comments
 (0)