Skip to content

Commit 5380221

Browse files
PerfectSlayerPerfectSlayer
authored
Migrate instrumentations to Context API (#9378)
* feat(dsm): Improve code style * feat(aws): Refactor AWS instrumentation around Context
1 parent 8daf352 commit 5380221

File tree

36 files changed

+230
-186
lines changed

36 files changed

+230
-186
lines changed

‎dd-java-agent/instrumentation/armeria/armeria-grpc-0.84/src/main/java/datadog/trace/instrumentation/armeria/grpc/client/GrpcClientDecorator.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package datadog.trace.instrumentation.armeria.grpc.client;
22

33
import static datadog.context.propagation.Propagators.defaultPropagator;
4+
import static datadog.trace.api.datastreams.DataStreamsTags.Direction.OUTBOUND;
45
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
56
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.traceConfig;
67

@@ -32,8 +33,7 @@ public class GrpcClientDecorator extends ClientDecorator {
3233
public static final CharSequence GRPC_MESSAGE = UTF8BytesString.create("grpc.message");
3334

3435
private static DataStreamsContext createDsmContext() {
35-
return DataStreamsContext.fromTags(
36-
DataStreamsTags.create("grpc", DataStreamsTags.Direction.Outbound));
36+
return DataStreamsContext.fromTags(DataStreamsTags.create("grpc", OUTBOUND));
3737
}
3838

3939
public static final GrpcClientDecorator DECORATE = new GrpcClientDecorator();

‎dd-java-agent/instrumentation/armeria/armeria-grpc-0.84/src/main/java/datadog/trace/instrumentation/armeria/grpc/server/GrpcServerDecorator.java

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
package datadog.trace.instrumentation.armeria.grpc.server;
22

3+
import static datadog.trace.api.datastreams.DataStreamsTags.Direction.INBOUND;
4+
35
import datadog.trace.api.Config;
46
import datadog.trace.api.cache.DDCache;
57
import datadog.trace.api.cache.DDCaches;
@@ -30,7 +32,7 @@ public class GrpcServerDecorator extends ServerDecorator {
3032
public static final CharSequence GRPC_MESSAGE = UTF8BytesString.create("grpc.message");
3133

3234
private static DataStreamsTags createServerPathwaySortedTags() {
33-
return DataStreamsTags.create("grpc", DataStreamsTags.Direction.Inbound);
35+
return DataStreamsTags.create("grpc", INBOUND);
3436
}
3537

3638
public static final DataStreamsTags SERVER_PATHWAY_EDGE_TAGS = createServerPathwaySortedTags();

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

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

33
import static datadog.context.propagation.Propagators.defaultPropagator;
4+
import static datadog.trace.api.datastreams.DataStreamsTags.Direction.OUTBOUND;
45
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.traceConfig;
56
import static datadog.trace.instrumentation.aws.v2.eventbridge.TextMapInjectAdapter.SETTER;
67

@@ -85,8 +86,7 @@ private String getTraceContextToInject(
8586
// Inject context
8687
datadog.context.Context context = span;
8788
if (traceConfig().isDataStreamsEnabled()) {
88-
DataStreamsTags tags =
89-
DataStreamsTags.createWithBus(DataStreamsTags.Direction.Outbound, eventBusName);
89+
DataStreamsTags tags = DataStreamsTags.createWithBus(OUTBOUND, eventBusName);
9090
DataStreamsContext dsmContext = DataStreamsContext.fromTags(tags);
9191
context = context.with(dsmContext);
9292
}

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

Lines changed: 12 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3,13 +3,15 @@
33
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
44
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
55
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.closeActive;
6+
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext;
7+
import static datadog.trace.instrumentation.aws.v0.OnErrorDecorator.CONTEXT_CONTEXT_KEY;
68
import static datadog.trace.instrumentation.aws.v0.OnErrorDecorator.DECORATE;
7-
import static datadog.trace.instrumentation.aws.v0.OnErrorDecorator.SPAN_CONTEXT_KEY;
89
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
910

1011
import com.amazonaws.AmazonClientException;
1112
import com.amazonaws.Request;
1213
import com.amazonaws.handlers.RequestHandler2;
14+
import datadog.context.Context;
1315
import datadog.trace.agent.tooling.Instrumenter;
1416
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
1517
import net.bytebuddy.asm.Advice;
@@ -54,12 +56,15 @@ public static void methodExit(
5456
}
5557

5658
if (throwable != null && request != null) {
57-
final AgentSpan span = request.getHandlerContext(SPAN_CONTEXT_KEY);
58-
if (span != null) {
59-
request.addHandlerContext(SPAN_CONTEXT_KEY, null);
60-
DECORATE.onError(span, throwable);
61-
DECORATE.beforeFinish(span);
62-
span.finish();
59+
final Context context = request.getHandlerContext(CONTEXT_CONTEXT_KEY);
60+
if (context != null) {
61+
request.addHandlerContext(CONTEXT_CONTEXT_KEY, null);
62+
final AgentSpan span = spanFromContext(context);
63+
if (span != null) {
64+
DECORATE.onError(span, throwable);
65+
DECORATE.beforeFinish(span);
66+
span.finish();
67+
}
6368
}
6469
}
6570
}

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

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

33
import static datadog.trace.api.datastreams.DataStreamsContext.create;
4+
import static datadog.trace.api.datastreams.DataStreamsTags.Direction.INBOUND;
5+
import static datadog.trace.api.datastreams.DataStreamsTags.Direction.OUTBOUND;
46
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.traceConfig;
57
import static datadog.trace.bootstrap.instrumentation.api.ResourceNamePriorities.RPC_COMMAND_NAME;
68

@@ -264,8 +266,7 @@ && traceConfig().isDataStreamsEnabled()) {
264266
&& ("GetObjectMetadataRequest".equalsIgnoreCase(awsOperation)
265267
|| "GetObjectRequest".equalsIgnoreCase(awsOperation))) {
266268
DataStreamsTags tags =
267-
DataStreamsTags.createWithDataset(
268-
"s3", DataStreamsTags.Direction.Inbound, bucket, key, bucket);
269+
DataStreamsTags.createWithDataset("s3", INBOUND, bucket, key, bucket);
269270
AgentTracer.get()
270271
.getDataStreamsMonitoring()
271272
.setCheckpoint(span, create(tags, 0, responseSize));
@@ -279,8 +280,7 @@ && traceConfig().isDataStreamsEnabled()) {
279280
payloadSize = (long) requestSize;
280281
}
281282
DataStreamsTags tags =
282-
DataStreamsTags.createWithDataset(
283-
"s3", DataStreamsTags.Direction.Outbound, bucket, key, bucket);
283+
DataStreamsTags.createWithDataset("s3", OUTBOUND, bucket, key, bucket);
284284
AgentTracer.get()
285285
.getDataStreamsMonitoring()
286286
.setCheckpoint(span, create(tags, 0, payloadSize));

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

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -37,9 +37,7 @@ public String[] helperClassNames() {
3737
public Map<String, String> contextStore() {
3838
Map<String, String> map = new java.util.HashMap<>();
3939
map.put(namespace + ".services.sqs.model.ReceiveMessageResult", "java.lang.String");
40-
map.put(
41-
namespace + ".AmazonWebServiceRequest",
42-
"datadog.trace.bootstrap.instrumentation.api.AgentSpan");
40+
map.put(namespace + ".AmazonWebServiceRequest", "datadog.context.Context");
4341
return map;
4442
}
4543

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -46,8 +46,7 @@ public static void addHandler(@Advice.Return final List<RequestHandler2> handler
4646
InstrumentationContext.get(
4747
"com.amazonaws.services.sqs.model.ReceiveMessageResult", "java.lang.String"),
4848
InstrumentationContext.get(
49-
"com.amazonaws.AmazonWebServiceRequest",
50-
"datadog.trace.bootstrap.instrumentation.api.AgentSpan")));
49+
"com.amazonaws.AmazonWebServiceRequest", "datadog.context.Context")));
5150
}
5251
}
5352
}

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

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,15 @@
11
package datadog.trace.instrumentation.aws.v0;
22

33
import com.amazonaws.handlers.HandlerContextKey;
4-
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
4+
import datadog.context.Context;
55
import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
66
import datadog.trace.bootstrap.instrumentation.decorator.BaseDecorator;
77

88
public class OnErrorDecorator extends BaseDecorator {
99

10-
public static final HandlerContextKey<AgentSpan> SPAN_CONTEXT_KEY =
11-
new HandlerContextKey<>("DatadogSpan"); // same as TracingRequestHandler.SPAN_CONTEXT_KEY
10+
public static final HandlerContextKey<Context> CONTEXT_CONTEXT_KEY =
11+
new HandlerContextKey<>(
12+
"DatadogContext"); // same as TracingRequestHandler.CONTEXT_CONTEXT_KEY
1213

1314
public static final OnErrorDecorator DECORATE = new OnErrorDecorator();
1415

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

Lines changed: 12 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,13 @@
33
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
44
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
55
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.closeActive;
6+
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext;
7+
import static datadog.trace.instrumentation.aws.v0.OnErrorDecorator.CONTEXT_CONTEXT_KEY;
68
import static datadog.trace.instrumentation.aws.v0.OnErrorDecorator.DECORATE;
7-
import static datadog.trace.instrumentation.aws.v0.OnErrorDecorator.SPAN_CONTEXT_KEY;
89
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
910

1011
import com.amazonaws.Request;
12+
import datadog.context.Context;
1113
import datadog.trace.agent.tooling.Instrumenter;
1214
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
1315
import net.bytebuddy.asm.Advice;
@@ -51,12 +53,15 @@ public static void methodExit(
5153
}
5254

5355
if (throwable != null) {
54-
final AgentSpan span = request.getHandlerContext(SPAN_CONTEXT_KEY);
55-
if (span != null) {
56-
request.addHandlerContext(SPAN_CONTEXT_KEY, null);
57-
DECORATE.onError(span, throwable);
58-
DECORATE.beforeFinish(span);
59-
span.finish();
56+
final Context context = request.getHandlerContext(CONTEXT_CONTEXT_KEY);
57+
if (context != null) {
58+
request.addHandlerContext(CONTEXT_CONTEXT_KEY, null);
59+
final AgentSpan span = spanFromContext(context);
60+
if (span != null) {
61+
DECORATE.onError(span, throwable);
62+
DECORATE.beforeFinish(span);
63+
span.finish();
64+
}
6065
}
6166
}
6267
}

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

Lines changed: 65 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,14 @@
11
package datadog.trace.instrumentation.aws.v0;
22

33
import static datadog.trace.api.datastreams.DataStreamsContext.create;
4+
import static datadog.trace.api.datastreams.DataStreamsTags.Direction.INBOUND;
5+
import static datadog.trace.api.datastreams.DataStreamsTags.create;
46
import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.XRAY_TRACING_CONCERN;
57
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpanWithoutScope;
68
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.blackholeSpan;
79
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
810
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.traceConfig;
11+
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext;
912
import static datadog.trace.instrumentation.aws.v0.AwsSdkClientDecorator.AWS_LEGACY_TRACING;
1013
import static datadog.trace.instrumentation.aws.v0.AwsSdkClientDecorator.DECORATE;
1114

@@ -14,9 +17,13 @@
1417
import com.amazonaws.Response;
1518
import com.amazonaws.handlers.HandlerContextKey;
1619
import com.amazonaws.handlers.RequestHandler2;
20+
import datadog.context.Context;
1721
import datadog.context.propagation.Propagators;
1822
import datadog.trace.api.Config;
19-
import datadog.trace.api.datastreams.*;
23+
import datadog.trace.api.datastreams.AgentDataStreamsMonitoring;
24+
import datadog.trace.api.datastreams.DataStreamsContext;
25+
import datadog.trace.api.datastreams.DataStreamsTags;
26+
import datadog.trace.api.datastreams.PathwayContext;
2027
import datadog.trace.bootstrap.ContextStore;
2128
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
2229
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
@@ -28,41 +35,42 @@
2835
/** Tracing Request Handler */
2936
public class TracingRequestHandler extends RequestHandler2 {
3037

31-
public static final HandlerContextKey<AgentSpan> SPAN_CONTEXT_KEY =
32-
new HandlerContextKey<>("DatadogSpan"); // same as OnErrorDecorator.SPAN_CONTEXT_KEY
38+
public static final HandlerContextKey<Context> CONTEXT_CONTEXT_KEY =
39+
new HandlerContextKey<>("DatadogContext"); // same as OnErrorDecorator.CONTEXT_CONTEXT_KEY
3340

3441
private static final Logger log = LoggerFactory.getLogger(TracingRequestHandler.class);
3542

3643
private final ContextStore<Object, String> responseQueueStore;
37-
private final ContextStore<AmazonWebServiceRequest, AgentSpan> requestSpanStore;
44+
private final ContextStore<AmazonWebServiceRequest, Context> requestContextStore;
3845

3946
public TracingRequestHandler(
4047
ContextStore<Object, String> responseQueueStore,
41-
ContextStore<AmazonWebServiceRequest, AgentSpan> requestSpanStore) {
48+
ContextStore<AmazonWebServiceRequest, Context> requestContextStore) {
4249
this.responseQueueStore = responseQueueStore;
43-
this.requestSpanStore = requestSpanStore;
50+
this.requestContextStore = requestContextStore;
4451
}
4552

4653
@Override
4754
public void beforeRequest(final Request<?> request) {
48-
AgentSpan span;
4955
if (!AWS_LEGACY_TRACING && isPollingRequest(request.getOriginalRequest())) {
5056
// SQS messages spans are created by aws-java-sqs-1.0 - replace client scope with no-op,
5157
// so we can tell when receive call is complete without affecting the rest of the trace
5258
activateSpanWithoutScope(blackholeSpan());
5359
} else {
54-
span = requestSpanStore.remove(request.getOriginalRequest());
60+
Context context = requestContextStore.remove(request.getOriginalRequest());
61+
AgentSpan span = spanFromContext(context);
5562
if (span != null) {
5663
// we'll land here for SQS send requests when DSM is enabled. In that case, we create the
5764
// span in SqsInterceptor to inject DSM tags.
5865
span.setOperationName(AwsNameCache.spanName(request));
5966
} else {
6067
// this is the most common code path
61-
span = startSpan(AwsNameCache.spanName(request));
68+
span = startSpan("aws-sdk", AwsNameCache.spanName(request));
69+
context = span; // TODO If DSM is enabled, add DSM context here too
6270
}
6371
DECORATE.afterStart(span);
6472
DECORATE.onRequest(span, request);
65-
request.addHandlerContext(SPAN_CONTEXT_KEY, span);
73+
request.addHandlerContext(CONTEXT_CONTEXT_KEY, context);
6674
if (Config.get().isAwsPropagationEnabled()) {
6775
try {
6876
Propagators.forConcern(XRAY_TRACING_CONCERN).inject(span, request, DECORATE);
@@ -81,13 +89,18 @@ public void beforeRequest(final Request<?> request) {
8189

8290
@Override
8391
public void afterResponse(final Request<?> request, final Response<?> response) {
84-
final AgentSpan span = request.getHandlerContext(SPAN_CONTEXT_KEY);
85-
if (span != null) {
86-
request.addHandlerContext(SPAN_CONTEXT_KEY, null);
87-
DECORATE.onResponse(span, response);
88-
DECORATE.onServiceResponse(span, request.getServiceName(), response);
89-
DECORATE.beforeFinish(span);
90-
span.finish();
92+
final Context context = request.getHandlerContext(CONTEXT_CONTEXT_KEY);
93+
log.warn("context {}", context);
94+
AgentSpan span = null;
95+
if (context != null) {
96+
request.addHandlerContext(CONTEXT_CONTEXT_KEY, null);
97+
span = spanFromContext(context);
98+
if (span != null) {
99+
DECORATE.onResponse(span, response);
100+
DECORATE.onServiceResponse(span, request.getServiceName(), response);
101+
DECORATE.beforeFinish(span);
102+
span.finish();
103+
}
91104
}
92105
AmazonWebServiceRequest originalRequest = request.getOriginalRequest();
93106
GetterAccess requestAccess = GetterAccess.of(originalRequest);
@@ -105,22 +118,25 @@ && traceConfig().isDataStreamsEnabled()
105118
&& "AmazonKinesis".equals(request.getServiceName())
106119
&& "GetRecords".equals(requestAccess.getOperationNameFromType())) {
107120
String streamArn = requestAccess.getStreamARN(originalRequest);
108-
if (null != streamArn) {
109-
List<?> records =
110-
GetterAccess.of(response.getAwsResponse()).getRecords(response.getAwsResponse());
111-
if (null != records) {
112-
DataStreamsTags tags =
113-
DataStreamsTags.create("kinesis", DataStreamsTags.Direction.Inbound, streamArn);
114-
for (Object record : records) {
115-
Date arrivalTime = GetterAccess.of(record).getApproximateArrivalTimestamp(record);
116-
AgentDataStreamsMonitoring dataStreamsMonitoring =
117-
AgentTracer.get().getDataStreamsMonitoring();
118-
PathwayContext pathwayContext = dataStreamsMonitoring.newPathwayContext();
119-
DataStreamsContext context = create(tags, arrivalTime.getTime(), 0);
120-
pathwayContext.setCheckpoint(context, dataStreamsMonitoring::add);
121-
if (!span.context().getPathwayContext().isStarted()) {
122-
span.context().mergePathwayContext(pathwayContext);
123-
}
121+
dsmCheckpoint(span, streamArn, response);
122+
}
123+
}
124+
125+
private void dsmCheckpoint(AgentSpan span, String streamArn, Response<?> response) {
126+
if (null != streamArn) {
127+
List<?> records =
128+
GetterAccess.of(response.getAwsResponse()).getRecords(response.getAwsResponse());
129+
if (null != records) {
130+
DataStreamsTags tags = create("kinesis", INBOUND, streamArn);
131+
for (Object record : records) {
132+
Date arrivalTime = GetterAccess.of(record).getApproximateArrivalTimestamp(record);
133+
AgentDataStreamsMonitoring dataStreamsMonitoring =
134+
AgentTracer.get().getDataStreamsMonitoring();
135+
PathwayContext pathwayContext = dataStreamsMonitoring.newPathwayContext();
136+
DataStreamsContext dataStreamsContext = create(tags, arrivalTime.getTime(), 0);
137+
pathwayContext.setCheckpoint(dataStreamsContext, dataStreamsMonitoring::add);
138+
if (!span.context().getPathwayContext().isStarted()) {
139+
span.context().mergePathwayContext(pathwayContext);
124140
}
125141
}
126142
}
@@ -129,24 +145,27 @@ && traceConfig().isDataStreamsEnabled()
129145

130146
@Override
131147
public void afterError(final Request<?> request, final Response<?> response, final Exception e) {
132-
AgentSpan span = request.getHandlerContext(SPAN_CONTEXT_KEY);
133-
if (span == null) {
134-
// also try getting the span from the context store, if the error happened early
135-
span = requestSpanStore.remove(request.getOriginalRequest());
148+
Context context = request.getHandlerContext(CONTEXT_CONTEXT_KEY);
149+
if (context == null) {
150+
// also try getting the context from the context store, if the error happened early
151+
context = requestContextStore.remove(request.getOriginalRequest());
136152
}
137153

138-
if (span != null) {
139-
request.addHandlerContext(SPAN_CONTEXT_KEY, null);
140-
if (response != null) {
141-
DECORATE.onResponse(span, response);
142-
if (span.isError()) {
154+
if (context != null) {
155+
request.addHandlerContext(CONTEXT_CONTEXT_KEY, null);
156+
final AgentSpan span = spanFromContext(context);
157+
if (span != null) {
158+
if (response != null) {
159+
DECORATE.onResponse(span, response);
160+
if (span.isError()) {
161+
DECORATE.onError(span, e);
162+
}
163+
} else {
143164
DECORATE.onError(span, e);
144165
}
145-
} else {
146-
DECORATE.onError(span, e);
166+
DECORATE.beforeFinish(span);
167+
span.finish();
147168
}
148-
DECORATE.beforeFinish(span);
149-
span.finish();
150169
}
151170
}
152171

0 commit comments

Comments
 (0)