Skip to content

Commit 3b7384c

Browse files
PerfectSlayerPerfectSlayer
committed
feat(dsm): Move context injection into decorator
1 parent d6abbbf commit 3b7384c

File tree

34 files changed

+137
-132
lines changed

34 files changed

+137
-132
lines changed

‎dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/UriBasedClientDecorator.java

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,18 @@
11
package datadog.trace.bootstrap.instrumentation.decorator;
22

3+
import static datadog.context.propagation.Propagators.defaultPropagator;
4+
5+
import datadog.context.Context;
6+
import datadog.context.propagation.CarrierSetter;
37
import datadog.trace.api.Config;
8+
import datadog.trace.api.datastreams.DataStreamsContext;
49
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
510
import datadog.trace.bootstrap.instrumentation.api.Tags;
611
import java.net.URI;
712
import javax.annotation.Nonnull;
813

914
public abstract class UriBasedClientDecorator extends ClientDecorator {
15+
private static final boolean DATA_STREAMS_ENABLED = Config.get().isDataStreamsEnabled();
1016

1117
public void onURI(@Nonnull final AgentSpan span, @Nonnull final URI uri) {
1218
String host = uri.getHost();
@@ -21,4 +27,13 @@ public void onURI(@Nonnull final AgentSpan span, @Nonnull final URI uri) {
2127
}
2228
}
2329
}
30+
31+
public <C> void injectContext(Context context, final C request, CarrierSetter<C> setter) {
32+
// Add additional default DSM context for HTTP clients if missing but DSM is enabled
33+
if (DATA_STREAMS_ENABLED && DataStreamsContext.fromContext(context) == null) {
34+
context = context.with(DataStreamsContext.forHttpClient());
35+
}
36+
// Inject context into carrier
37+
defaultPropagator().inject(context, request, setter);
38+
}
2439
}

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

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

3-
import static datadog.context.propagation.Propagators.defaultPropagator;
43
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
54
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
65
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
@@ -73,12 +72,12 @@ public static AgentScope methodEnter(
7372
return null;
7473
}
7574

76-
final AgentSpan span = startSpan(AKKA_CLIENT_REQUEST);
75+
final AgentSpan span = startSpan("akka-http", AKKA_CLIENT_REQUEST);
7776
DECORATE.afterStart(span);
7877
DECORATE.onRequest(span, request);
7978

8079
if (request != null) {
81-
defaultPropagator().inject(getCurrentContext().with(span), request, headers);
80+
DECORATE.injectContext(getCurrentContext().with(span), request, headers);
8281
// Request is immutable, so we have to assign new value once we update headers
8382
request = headers.getRequest();
8483
}

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

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

3-
import static datadog.context.propagation.Propagators.defaultPropagator;
43
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
54
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
65
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext;
6+
import static datadog.trace.instrumentation.akkahttp106.AkkaHttpClientDecorator.AKKA_CLIENT_REQUEST;
7+
import static datadog.trace.instrumentation.akkahttp106.AkkaHttpClientDecorator.DECORATE;
78

89
import akka.http.scaladsl.HttpExt;
910
import akka.http.scaladsl.model.HttpRequest;
@@ -23,12 +24,11 @@ public static AgentScope methodEnter(
2324
return null;
2425
}
2526

26-
final AgentSpan span = startSpan(AkkaHttpClientDecorator.AKKA_CLIENT_REQUEST);
27-
AkkaHttpClientDecorator.DECORATE.afterStart(span);
28-
AkkaHttpClientDecorator.DECORATE.onRequest(span, request);
29-
27+
final AgentSpan span = startSpan("akka-http", AKKA_CLIENT_REQUEST);
28+
DECORATE.afterStart(span);
29+
DECORATE.onRequest(span, request);
3030
if (request != null) {
31-
defaultPropagator().inject(getCurrentContext().with(span), request, headers);
31+
DECORATE.injectContext(getCurrentContext().with(span), request, headers);
3232
// Request is immutable, so we have to assign new value once we update headers
3333
request = headers.getRequest();
3434
}
@@ -51,8 +51,8 @@ public static void methodExit(
5151
responseFuture.onComplete(
5252
new AkkaHttpClientHelpers.OnCompleteHandler(span), thiz.system().dispatcher());
5353
} else {
54-
AkkaHttpClientDecorator.DECORATE.onError(span, throwable);
55-
AkkaHttpClientDecorator.DECORATE.beforeFinish(span);
54+
DECORATE.onError(span, throwable);
55+
DECORATE.beforeFinish(span);
5656
span.finish();
5757
}
5858
scope.close();

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

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,9 @@
11
package datadog.trace.instrumentation.apachehttpasyncclient;
22

3-
import static datadog.context.propagation.Propagators.defaultPropagator;
3+
import static datadog.context.Context.current;
44
import static datadog.trace.instrumentation.apachehttpasyncclient.ApacheHttpAsyncClientDecorator.DECORATE;
55
import static datadog.trace.instrumentation.apachehttpasyncclient.HttpHeadersInjectAdapter.SETTER;
66

7-
import datadog.context.Context;
87
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
98
import java.io.IOException;
109
import org.apache.http.HttpException;
@@ -33,7 +32,7 @@ public HttpHost getTarget() {
3332
public HttpRequest generateRequest() throws IOException, HttpException {
3433
final HttpRequest request = delegate.generateRequest();
3534
DECORATE.onRequest(span, new HostAndRequestAsHttpUriRequest(delegate.getTarget(), request));
36-
defaultPropagator().inject(Context.current().with(span), request, SETTER);
35+
DECORATE.injectContext(current().with(span), request, SETTER);
3736
return request;
3837
}
3938

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

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,12 @@
11
package datadog.trace.instrumentation.apachehttpclient;
22

3-
import static datadog.context.propagation.Propagators.defaultPropagator;
3+
import static datadog.context.Context.current;
44
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
55
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
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

10-
import datadog.context.Context;
1110
import datadog.trace.bootstrap.CallDepthThreadLocalMap;
1211
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
1312
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -45,7 +44,7 @@ private static AgentScope activateHttpSpan(final HttpUriRequest request) {
4544

4645
// AWS calls are often signed, so we can't add headers without breaking the signature.
4746
if (!awsClientCall) {
48-
defaultPropagator().inject(Context.current().with(span), request, SETTER);
47+
DECORATE.injectContext(current().with(span), request, SETTER);
4948
}
5049

5150
return scope;

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

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

3-
import static datadog.context.propagation.Propagators.defaultPropagator;
3+
import static datadog.context.Context.current;
44
import static datadog.trace.instrumentation.apachehttpclient5.ApacheHttpClientDecorator.DECORATE;
55
import static datadog.trace.instrumentation.apachehttpclient5.HttpHeadersInjectAdapter.SETTER;
66

7-
import datadog.context.Context;
8-
import datadog.trace.api.datastreams.DataStreamsContext;
97
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
108
import java.io.IOException;
119
import org.apache.hc.core5.http.EntityDetails;
@@ -27,7 +25,7 @@ public DelegatingRequestChannel(RequestChannel requestChannel, AgentSpan span) {
2725
public void sendRequest(HttpRequest request, EntityDetails entityDetails, HttpContext context)
2826
throws HttpException, IOException {
2927
DECORATE.onRequest(span, request);
30-
defaultPropagator().inject(Context.current().with(span), request, SETTER);
28+
DECORATE.injectContext(current().with(span), request, SETTER);
3129
delegate.sendRequest(request, entityDetails, context);
3230
}
3331
}

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

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

3-
import static datadog.context.propagation.Propagators.defaultPropagator;
3+
import static datadog.context.Context.current;
44
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
55
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
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

10-
import datadog.context.Context;
1110
import datadog.trace.bootstrap.CallDepthThreadLocalMap;
1211
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
1312
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -45,7 +44,7 @@ private static AgentScope activateHttpSpan(final HttpRequest request) {
4544
final boolean awsClientCall = request.containsHeader("amz-sdk-invocation-id");
4645
// AWS calls are often signed, so we can't add headers without breaking the signature.
4746
if (!awsClientCall) {
48-
defaultPropagator().inject(Context.current().with(span), request, SETTER);
47+
DECORATE.injectContext(current().with(span), request, SETTER);
4948
}
5049

5150
return scope;

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

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

3-
import static datadog.context.propagation.Propagators.defaultPropagator;
3+
import static datadog.context.Context.current;
44
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
55
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
66
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
@@ -134,7 +134,7 @@ public static <T> AgentScope before(
134134
if (null != responseListener && null != headers) {
135135
span = InstrumentationContext.get(ClientCall.class, AgentSpan.class).get(call);
136136
if (null != span) {
137-
defaultPropagator().inject(span, headers, SETTER);
137+
DECORATE.injectContext(current().with(span), headers, SETTER);
138138
return activateSpan(span);
139139
}
140140
}

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

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,18 @@
11
package datadog.trace.instrumentation.armeria.grpc.client;
22

3+
import static datadog.context.propagation.Propagators.defaultPropagator;
34
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
5+
import static datadog.trace.core.datastreams.TagsProcessor.DIRECTION_OUT;
6+
import static datadog.trace.core.datastreams.TagsProcessor.DIRECTION_TAG;
7+
import static datadog.trace.core.datastreams.TagsProcessor.TYPE_TAG;
48

9+
import datadog.context.Context;
10+
import datadog.context.propagation.CarrierSetter;
511
import datadog.trace.api.Config;
612
import datadog.trace.api.GenericClassValue;
713
import datadog.trace.api.cache.DDCache;
814
import datadog.trace.api.cache.DDCaches;
15+
import datadog.trace.api.datastreams.DataStreamsContext;
916
import datadog.trace.api.naming.SpanNaming;
1017
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
1118
import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes;
@@ -15,6 +22,7 @@
1522
import io.grpc.MethodDescriptor;
1623
import io.grpc.Status;
1724
import java.util.BitSet;
25+
import java.util.LinkedHashMap;
1826
import java.util.Set;
1927
import java.util.function.Function;
2028

@@ -25,10 +33,18 @@ public class GrpcClientDecorator extends ClientDecorator {
2533
public static final CharSequence COMPONENT_NAME = UTF8BytesString.create("armeria-grpc-client");
2634
public static final CharSequence GRPC_MESSAGE = UTF8BytesString.create("grpc.message");
2735

36+
private static DataStreamsContext createDsmContext() {
37+
LinkedHashMap<String, String> result = new LinkedHashMap<>();
38+
result.put(DIRECTION_TAG, DIRECTION_OUT);
39+
result.put(TYPE_TAG, "grpc");
40+
return DataStreamsContext.fromTags(result);
41+
}
42+
2843
public static final GrpcClientDecorator DECORATE = new GrpcClientDecorator();
2944

3045
private static final Set<String> IGNORED_METHODS = Config.get().getGrpcIgnoredOutboundMethods();
3146
private static final BitSet CLIENT_ERROR_STATUSES = Config.get().getGrpcClientErrorStatuses();
47+
private static final boolean DATA_STREAMS_ENABLED = Config.get().isDataStreamsEnabled();
3248

3349
private static final ClassValue<UTF8BytesString> MESSAGE_TYPES =
3450
GenericClassValue.of(
@@ -82,7 +98,7 @@ public <ReqT, RespT> AgentSpan startCall(MethodDescriptor<ReqT, RespT> method) {
8298
return null;
8399
}
84100
AgentSpan span =
85-
startSpan(OPERATION_NAME)
101+
startSpan("grpc", OPERATION_NAME)
86102
.setTag("request.type", requestMessageType(method))
87103
.setTag("response.type", responseMessageType(method))
88104
// method.getServiceName() may not be available on some grpc versions
@@ -94,6 +110,13 @@ public <ReqT, RespT> AgentSpan startCall(MethodDescriptor<ReqT, RespT> method) {
94110
return afterStart(span);
95111
}
96112

113+
public <C> void injectContext(Context context, final C request, CarrierSetter<C> setter) {
114+
if (DATA_STREAMS_ENABLED) {
115+
context = context.with(createDsmContext());
116+
}
117+
defaultPropagator().inject(context, request, setter);
118+
}
119+
97120
public AgentSpan onClose(final AgentSpan span, final Status status) {
98121

99122
span.setTag("status.code", status.getCode().name());

‎dd-java-agent/instrumentation/commons-httpclient-2/src/main/java/datadog/trace/instrumentation/commonshttpclient/CommonsHttpClientInstrumentation.java

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
package datadog.trace.instrumentation.commonshttpclient;
22

3-
import static datadog.context.propagation.Propagators.defaultPropagator;
43
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
54
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
65
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
@@ -68,7 +67,7 @@ public static AgentScope methodEnter(@Advice.Argument(1) final HttpMethod httpMe
6867

6968
DECORATE.afterStart(span);
7069
DECORATE.onRequest(span, httpMethod);
71-
defaultPropagator().inject(getCurrentContext().with(span), httpMethod, SETTER);
70+
DECORATE.injectContext(getCurrentContext().with(span), httpMethod, SETTER);
7271

7372
return scope;
7473
} catch (BlockingException e) {

‎dd-java-agent/instrumentation/google-http-client/src/main/java/datadog/trace/instrumentation/googlehttpclient/GoogleHttpClientDecorator.java

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,10 @@
11
package datadog.trace.instrumentation.googlehttpclient;
22

3-
import static datadog.context.propagation.Propagators.defaultPropagator;
3+
import static datadog.context.Context.current;
44
import static datadog.trace.instrumentation.googlehttpclient.HeadersInjectAdapter.SETTER;
55

66
import com.google.api.client.http.HttpRequest;
77
import com.google.api.client.http.HttpResponse;
8-
import datadog.context.Context;
98
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
109
import datadog.trace.bootstrap.instrumentation.api.URIUtils;
1110
import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
@@ -38,7 +37,7 @@ protected URI url(final HttpRequest httpRequest) throws URISyntaxException {
3837
public AgentSpan prepareSpan(AgentSpan span, HttpRequest request) {
3938
DECORATE.afterStart(span);
4039
DECORATE.onRequest(span, request);
41-
defaultPropagator().inject(Context.current().with(span), request, SETTER);
40+
DECORATE.injectContext(current().with(span), request, SETTER);
4241
return span;
4342
}
4443

‎dd-java-agent/instrumentation/grizzly-client-1.9/src/main/java/datadog/trace/instrumentation/grizzly/client/AsyncHttpClientInstrumentation.java

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,10 @@
11
package datadog.trace.instrumentation.grizzly.client;
22

3-
import static datadog.context.propagation.Propagators.defaultPropagator;
43
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
54
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
65
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
7-
import static datadog.trace.instrumentation.grizzly.client.ClientDecorator.DECORATE;
86
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext;
7+
import static datadog.trace.instrumentation.grizzly.client.ClientDecorator.DECORATE;
98
import static datadog.trace.instrumentation.grizzly.client.ClientDecorator.HTTP_REQUEST;
109
import static datadog.trace.instrumentation.grizzly.client.InjectAdapter.SETTER;
1110
import static net.bytebuddy.matcher.ElementMatchers.isPublic;
@@ -68,7 +67,7 @@ public static void onEnter(
6867
AgentSpan span = startSpan(HTTP_REQUEST);
6968
DECORATE.afterStart(span);
7069
DECORATE.onRequest(span, request);
71-
defaultPropagator().inject(getCurrentContext().with(span), request, SETTER);
70+
DECORATE.injectContext(getCurrentContext().with(span), request, SETTER);
7271
handler = new AsyncHandlerAdapter<>(span, parentSpan, handler);
7372
}
7473
}

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

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

3-
import static datadog.context.propagation.Propagators.defaultPropagator;
43
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
54
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
65
import static datadog.trace.instrumentation.grpc.client.GrpcClientDecorator.DECORATE;
@@ -91,7 +90,7 @@ public static AgentScope before(
9190
@Advice.Local("$$ddSpan") AgentSpan span) {
9291
span = InstrumentationContext.get(ClientCall.class, AgentSpan.class).get(call);
9392
if (null != span) {
94-
defaultPropagator().inject(span, headers, SETTER);
93+
DECORATE.injectContext(span, headers, SETTER);
9594
return activateSpan(span);
9695
}
9796
return null;

0 commit comments

Comments
 (0)