Skip to content

Commit 3a5c47b

Browse files
PerfectSlayerPerfectSlayer
authored
Clean up DSM context injection (#8776)
* feat(dsm): Clean up context injection * feat(context): Improve Context API usage * feat(dsm): Move context injection into decorator * feat(dsm): Unify activation check using dynamic config only
1 parent e590c41 commit 3a5c47b

File tree

73 files changed

+326
-408
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

73 files changed

+326
-408
lines changed

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

Lines changed: 0 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -26,21 +26,12 @@
2626
import java.net.URI;
2727
import java.net.URISyntaxException;
2828
import java.util.BitSet;
29-
import java.util.LinkedHashMap;
3029
import java.util.Map;
3130
import java.util.function.BiFunction;
3231
import org.slf4j.Logger;
3332
import org.slf4j.LoggerFactory;
3433

3534
public abstract class HttpClientDecorator<REQUEST, RESPONSE> extends UriBasedClientDecorator {
36-
public static final LinkedHashMap<String, String> CLIENT_PATHWAY_EDGE_TAGS;
37-
38-
static {
39-
CLIENT_PATHWAY_EDGE_TAGS = new LinkedHashMap<>(2);
40-
// TODO: Refactor TagsProcessor to move it into a package that we can link the constants for.
41-
CLIENT_PATHWAY_EDGE_TAGS.put("direction", "out");
42-
CLIENT_PATHWAY_EDGE_TAGS.put("type", "http");
43-
}
4435

4536
private static final Logger log = LoggerFactory.getLogger(HttpClientDecorator.class);
4637

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

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

33
import static datadog.context.Context.root;
44
import static datadog.trace.api.cache.RadixTreeCache.UNSET_STATUS;
5-
import static datadog.trace.api.datastreams.DataStreamsContext.fromTags;
5+
import static datadog.trace.api.datastreams.DataStreamsContext.forHttpServer;
66
import static datadog.trace.api.gateway.Events.EVENTS;
77
import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.extractContextAndGetSpanContext;
88
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.traceConfig;
@@ -38,14 +38,14 @@
3838
import datadog.trace.bootstrap.instrumentation.decorator.http.ClientIpAddressResolver;
3939
import java.net.InetAddress;
4040
import java.util.BitSet;
41-
import java.util.LinkedHashMap;
4241
import java.util.Locale;
4342
import java.util.Map;
4443
import java.util.concurrent.ExecutionException;
4544
import java.util.function.BiFunction;
4645
import java.util.function.Function;
4746
import java.util.function.Supplier;
4847
import javax.annotation.Nonnull;
48+
import javax.annotation.Nullable;
4949
import org.slf4j.Logger;
5050
import org.slf4j.LoggerFactory;
5151

@@ -62,15 +62,6 @@ public abstract class HttpServerDecorator<REQUEST, CONNECTION, RESPONSE, REQUEST
6262
public static final String DD_RESPONSE_ATTRIBUTE = "datadog.response";
6363
public static final String DD_IGNORE_COMMIT_ATTRIBUTE = "datadog.commit.ignore";
6464

65-
public static final LinkedHashMap<String, String> SERVER_PATHWAY_EDGE_TAGS;
66-
67-
static {
68-
SERVER_PATHWAY_EDGE_TAGS = new LinkedHashMap<>(2);
69-
// TODO: Refactor TagsProcessor to move it into a package that we can link the constants for.
70-
SERVER_PATHWAY_EDGE_TAGS.put("direction", "in");
71-
SERVER_PATHWAY_EDGE_TAGS.put("type", "http");
72-
}
73-
7465
private static final UTF8BytesString DEFAULT_RESOURCE_NAME = UTF8BytesString.create("/");
7566
protected static final UTF8BytesString NOT_FOUND_RESOURCE_NAME = UTF8BytesString.create("404");
7667
protected static final boolean SHOULD_SET_404_RESOURCE_NAME =
@@ -165,7 +156,7 @@ public AgentSpan startSpan(
165156
}
166157
AgentPropagation.ContextVisitor<REQUEST_CARRIER> getter = getter();
167158
if (null != carrier && null != getter) {
168-
tracer().getDataStreamsMonitoring().setCheckpoint(span, fromTags(SERVER_PATHWAY_EDGE_TAGS));
159+
tracer().getDataStreamsMonitoring().setCheckpoint(span, forHttpServer());
169160
}
170161
return span;
171162
}
@@ -191,7 +182,7 @@ public AgentSpan onRequest(
191182
final AgentSpan span,
192183
final CONNECTION connection,
193184
final REQUEST request,
194-
final AgentSpanContext.Extracted context) {
185+
@Nullable final AgentSpanContext.Extracted context) {
195186
Config config = Config.get();
196187
boolean clientIpResolverEnabled =
197188
config.isClientIpEnabled()

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

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

3+
import static datadog.context.propagation.Propagators.defaultPropagator;
4+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.traceConfig;
5+
6+
import datadog.context.Context;
7+
import datadog.context.propagation.CarrierSetter;
38
import datadog.trace.api.Config;
9+
import datadog.trace.api.datastreams.DataStreamsContext;
410
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
511
import datadog.trace.bootstrap.instrumentation.api.Tags;
612
import java.net.URI;
713
import javax.annotation.Nonnull;
814

915
public abstract class UriBasedClientDecorator extends ClientDecorator {
10-
1116
public void onURI(@Nonnull final AgentSpan span, @Nonnull final URI uri) {
1217
String host = uri.getHost();
1318
int port = uri.getPort();
@@ -21,4 +26,13 @@ public void onURI(@Nonnull final AgentSpan span, @Nonnull final URI uri) {
2126
}
2227
}
2328
}
29+
30+
public <C> void injectContext(Context context, final C request, CarrierSetter<C> setter) {
31+
// Add additional default DSM context for HTTP clients if missing but DSM is enabled
32+
if (DataStreamsContext.fromContext(context) == null && traceConfig().isDataStreamsEnabled()) {
33+
context = context.with(DataStreamsContext.forHttpClient());
34+
}
35+
// Inject context into carrier
36+
defaultPropagator().inject(context, request, setter);
37+
}
2438
}

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

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,9 @@
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;
76
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext;
8-
import static datadog.trace.bootstrap.instrumentation.decorator.HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS;
97
import static datadog.trace.instrumentation.akkahttp.AkkaHttpClientDecorator.AKKA_CLIENT_REQUEST;
108
import static datadog.trace.instrumentation.akkahttp.AkkaHttpClientDecorator.DECORATE;
119
import static datadog.trace.instrumentation.akkahttp.AkkaHttpClientHelpers.AkkaHttpHeaders;
@@ -18,7 +16,6 @@
1816
import com.google.auto.service.AutoService;
1917
import datadog.trace.agent.tooling.Instrumenter;
2018
import datadog.trace.agent.tooling.InstrumenterModule;
21-
import datadog.trace.api.datastreams.DataStreamsContext;
2219
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
2320
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
2421
import net.bytebuddy.asm.Advice;
@@ -75,14 +72,12 @@ public static AgentScope methodEnter(
7572
return null;
7673
}
7774

78-
final AgentSpan span = startSpan(AKKA_CLIENT_REQUEST);
75+
final AgentSpan span = startSpan("akka-http", AKKA_CLIENT_REQUEST);
7976
DECORATE.afterStart(span);
8077
DECORATE.onRequest(span, request);
8178

8279
if (request != null) {
83-
DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS);
84-
defaultPropagator()
85-
.inject(getCurrentContext().with(span).with(dsmContext), request, headers);
80+
DECORATE.injectContext(getCurrentContext().with(span), request, headers);
8681
// Request is immutable, so we have to assign new value once we update headers
8782
request = headers.getRequest();
8883
}

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

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

3+
import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext;
4+
35
import akka.http.scaladsl.model.HttpRequest;
46
import akka.http.scaladsl.model.HttpResponse;
57
import akka.http.scaladsl.util.FastFuture$;
68
import akka.stream.Materializer;
7-
import datadog.context.Context;
89
import datadog.context.ContextScope;
910
import datadog.trace.api.gateway.Flow;
1011
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -28,8 +29,7 @@ public DatadogAsyncHandlerWrapper(
2829
@Override
2930
public Future<HttpResponse> apply(final HttpRequest request) {
3031
final ContextScope scope = DatadogWrapperHelper.createSpan(request);
31-
Context context = scope.context();
32-
final AgentSpan span = AgentSpan.fromContext(context);
32+
final AgentSpan span = fromContext(scope.context());
3333
Future<HttpResponse> futureResponse;
3434

3535
// handle blocking in the beginning of the request

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

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

3+
import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext;
34
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
45

56
import akka.http.scaladsl.model.HttpRequest;
@@ -67,7 +68,7 @@ public GraphStageLogic createLogic(final Attributes inheritedAttributes) throws
6768
public void onPush() throws Exception {
6869
final HttpRequest request = grab(requestInlet);
6970
final ContextScope scope = DatadogWrapperHelper.createSpan(request);
70-
final AgentSpan span = AgentSpan.fromContext(scope.context());
71+
final AgentSpan span = fromContext(scope.context());
7172
RequestContext requestContext = span.getRequestContext();
7273
if (requestContext != null) {
7374
HttpResponse response =
@@ -133,7 +134,7 @@ public void onPush() throws Exception {
133134
HttpResponse response = grab(responseInlet);
134135
final ContextScope scope = scopes.poll();
135136
if (scope != null) {
136-
AgentSpan span = AgentSpan.fromContext(scope.context());
137+
AgentSpan span = fromContext(scope.context());
137138
HttpResponse newResponse =
138139
BlockingResponseHelper.handleFinishForWaf(span, response);
139140
if (newResponse != response) {
@@ -159,7 +160,7 @@ public void onUpstreamFinish() throws Exception {
159160
// remaining spans
160161
ContextScope scope = scopes.poll();
161162
while (scope != null) {
162-
AgentSpan.fromContext(scope.context()).finish();
163+
fromContext(scope.context()).finish();
163164
scope = scopes.poll();
164165
}
165166
completeStage();
@@ -170,14 +171,14 @@ public void onUpstreamFailure(final Throwable ex) throws Exception {
170171
ContextScope scope = scopes.poll();
171172
if (scope != null) {
172173
// Mark the span as failed
173-
AgentSpan span = AgentSpan.fromContext(scope.context());
174+
AgentSpan span = fromContext(scope.context());
174175
DatadogWrapperHelper.finishSpan(span, ex);
175176
}
176177
// We will not receive any more responses from the user code, so clean up any
177178
// remaining spans
178179
scope = scopes.poll();
179180
while (scope != null) {
180-
AgentSpan.fromContext(scope.context()).finish();
181+
fromContext(scope.context()).finish();
181182
scope = scopes.poll();
182183
}
183184
fail(responseOutlet, ex);

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -10,12 +10,12 @@
1010

1111
public class DatadogWrapperHelper {
1212
public static ContextScope createSpan(final HttpRequest request) {
13-
final Context extractedContext = DECORATE.extractContext(request);
14-
final AgentSpan span = DECORATE.startSpan(request, extractedContext);
13+
final Context context = DECORATE.extractContext(request);
14+
final AgentSpan span = DECORATE.startSpan(request, context);
1515
DECORATE.afterStart(span);
16-
DECORATE.onRequest(span, request, request, extractedContext);
16+
DECORATE.onRequest(span, request, request, context);
1717

18-
return extractedContext.with(span).attach();
18+
return context.with(span).attach();
1919
}
2020

2121
public static void finishSpan(final AgentSpan span, final HttpResponse response) {

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

Lines changed: 8 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,14 @@
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;
7-
import static datadog.trace.bootstrap.instrumentation.decorator.HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS;
6+
import static datadog.trace.instrumentation.akkahttp106.AkkaHttpClientDecorator.AKKA_CLIENT_REQUEST;
7+
import static datadog.trace.instrumentation.akkahttp106.AkkaHttpClientDecorator.DECORATE;
88

99
import akka.http.scaladsl.HttpExt;
1010
import akka.http.scaladsl.model.HttpRequest;
1111
import akka.http.scaladsl.model.HttpResponse;
12-
import datadog.trace.api.datastreams.DataStreamsContext;
1312
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
1413
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
1514
import net.bytebuddy.asm.Advice;
@@ -25,13 +24,11 @@ public static AgentScope methodEnter(
2524
return null;
2625
}
2726

28-
final AgentSpan span = startSpan(AkkaHttpClientDecorator.AKKA_CLIENT_REQUEST);
29-
AkkaHttpClientDecorator.DECORATE.afterStart(span);
30-
AkkaHttpClientDecorator.DECORATE.onRequest(span, request);
31-
27+
final AgentSpan span = startSpan("akka-http", AKKA_CLIENT_REQUEST);
28+
DECORATE.afterStart(span);
29+
DECORATE.onRequest(span, request);
3230
if (request != null) {
33-
DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS);
34-
defaultPropagator().inject(getCurrentContext().with(span).with(dsmContext), request, headers);
31+
DECORATE.injectContext(getCurrentContext().with(span), request, headers);
3532
// Request is immutable, so we have to assign new value once we update headers
3633
request = headers.getRequest();
3734
}
@@ -54,8 +51,8 @@ public static void methodExit(
5451
responseFuture.onComplete(
5552
new AkkaHttpClientHelpers.OnCompleteHandler(span), thiz.system().dispatcher());
5653
} else {
57-
AkkaHttpClientDecorator.DECORATE.onError(span, throwable);
58-
AkkaHttpClientDecorator.DECORATE.beforeFinish(span);
54+
DECORATE.onError(span, throwable);
55+
DECORATE.beforeFinish(span);
5956
span.finish();
6057
}
6158
scope.close();

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

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

3-
import static datadog.context.propagation.Propagators.defaultPropagator;
4-
import static datadog.trace.bootstrap.instrumentation.decorator.HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS;
3+
import static datadog.context.Context.current;
54
import static datadog.trace.instrumentation.apachehttpasyncclient.ApacheHttpAsyncClientDecorator.DECORATE;
65
import static datadog.trace.instrumentation.apachehttpasyncclient.HttpHeadersInjectAdapter.SETTER;
76

8-
import datadog.context.Context;
9-
import datadog.trace.api.datastreams.DataStreamsContext;
107
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
118
import java.io.IOException;
129
import org.apache.http.HttpException;
@@ -35,10 +32,7 @@ public HttpHost getTarget() {
3532
public HttpRequest generateRequest() throws IOException, HttpException {
3633
final HttpRequest request = delegate.generateRequest();
3734
DECORATE.onRequest(span, new HostAndRequestAsHttpUriRequest(delegate.getTarget(), request));
38-
39-
DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS);
40-
defaultPropagator().inject(Context.current().with(span).with(dsmContext), request, SETTER);
41-
35+
DECORATE.injectContext(current().with(span), request, SETTER);
4236
return request;
4337
}
4438

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

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +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;
6-
import static datadog.trace.bootstrap.instrumentation.decorator.HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS;
76
import static datadog.trace.instrumentation.apachehttpclient.ApacheHttpClientDecorator.DECORATE;
87
import static datadog.trace.instrumentation.apachehttpclient.ApacheHttpClientDecorator.HTTP_REQUEST;
98
import static datadog.trace.instrumentation.apachehttpclient.HttpHeadersInjectAdapter.SETTER;
109

11-
import datadog.context.Context;
12-
import datadog.trace.api.datastreams.DataStreamsContext;
1310
import datadog.trace.bootstrap.CallDepthThreadLocalMap;
1411
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
1512
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -47,8 +44,7 @@ private static AgentScope activateHttpSpan(final HttpUriRequest request) {
4744

4845
// AWS calls are often signed, so we can't add headers without breaking the signature.
4946
if (!awsClientCall) {
50-
DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS);
51-
defaultPropagator().inject(Context.current().with(span).with(dsmContext), request, SETTER);
47+
DECORATE.injectContext(current().with(span), request, SETTER);
5248
}
5349

5450
return scope;

0 commit comments

Comments
 (0)