Skip to content

Commit 9f546b9

Browse files
PerfectSlayerPerfectSlayer
authored
Migrate instrumentations to Context API (#9358)
* feat(bootstrap): Migrate Spray instrumentation to Context API * feat(bootstrap): Migrate Tomcat instrumentation to Context API * feat(bootstrap): Migrate Liberty instrumentations to Context API * feat(bootstrap): Migrate instrumentations to Context API * fix(bootstrap): Fix unexpected call to IG callback * feat(play): Simplify testing * fix(bootstrap): Fix import
1 parent a11444c commit 9f546b9

File tree

33 files changed

+174
-192
lines changed

33 files changed

+174
-192
lines changed

‎dd-java-agent/agent-bootstrap/src/jmh/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecoratorBenchmark.java‎

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

3+
import static datadog.context.Context.root;
34
import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext;
5+
import static java.util.Collections.emptyMap;
46
import static java.util.concurrent.TimeUnit.MICROSECONDS;
57
import static java.util.concurrent.TimeUnit.SECONDS;
68

79
import datadog.context.Context;
810
import datadog.trace.api.GlobalTracer;
911
import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
1012
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
11-
import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
1213
import datadog.trace.bootstrap.instrumentation.api.ContextVisitors;
1314
import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter;
1415
import datadog.trace.bootstrap.instrumentation.api.URIDefaultDataAdapter;
@@ -17,7 +18,6 @@
1718
import datadog.trace.core.CoreTracer;
1819
import datadog.trace.core.DDSpan;
1920
import java.net.URI;
20-
import java.util.Collections;
2121
import java.util.List;
2222
import java.util.Map;
2323
import org.openjdk.jmh.annotations.Benchmark;
@@ -59,13 +59,13 @@ public void setUp() {
5959
.build();
6060
GlobalTracer.forceRegister(tracer);
6161
decorator = new BenchmarkHttpServerDecorator();
62-
Context context = decorator.startSpan(Collections.emptyMap(), Context.root());
62+
Context context = decorator.startSpan(emptyMap(), root());
6363
span = fromContext(context);
6464
}
6565

6666
@Benchmark
6767
public AgentSpan onRequest() {
68-
return decorator.onRequest(span, null, request, (AgentSpanContext.Extracted) null);
68+
return decorator.onRequest(span, null, request, root());
6969
}
7070

7171
public static class Request {

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

Lines changed: 36 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import static datadog.trace.api.cache.RadixTreeCache.UNSET_STATUS;
55
import static datadog.trace.api.datastreams.DataStreamsContext.forHttpServer;
66
import static datadog.trace.api.gateway.Events.EVENTS;
7+
import static datadog.trace.bootstrap.ActiveSubsystems.APPSEC_ACTIVE;
78
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.traceConfig;
89
import static datadog.trace.bootstrap.instrumentation.decorator.http.HttpResourceDecorator.HTTP_RESOURCE_DECORATOR;
910

@@ -17,11 +18,11 @@
1718
import datadog.trace.api.gateway.BlockResponseFunction;
1819
import datadog.trace.api.gateway.CallbackProvider;
1920
import datadog.trace.api.gateway.Flow;
21+
import datadog.trace.api.gateway.Flow.Action.RequestBlockingAction;
2022
import datadog.trace.api.gateway.IGSpanInfo;
2123
import datadog.trace.api.gateway.RequestContext;
2224
import datadog.trace.api.gateway.RequestContextSlot;
2325
import datadog.trace.api.naming.SpanNaming;
24-
import datadog.trace.bootstrap.ActiveSubsystems;
2526
import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
2627
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
2728
import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
@@ -146,12 +147,12 @@ public Context startSpan(REQUEST_CARRIER carrier, Context context) {
146147
instrumentationNames != null && instrumentationNames.length > 0
147148
? instrumentationNames[0]
148149
: DEFAULT_INSTRUMENTATION_NAME;
149-
AgentSpanContext.Extracted extracted = callIGCallbackStart(context);
150+
AgentSpanContext.Extracted extracted = callIGCallbackStart(getExtractedSpanContext(context));
150151
AgentSpan span =
151152
tracer().startSpan(instrumentationName, spanName(), extracted).setMeasured(true);
152153
Flow<Void> flow = callIGCallbackRequestHeaders(span, carrier);
153-
if (flow.getAction() instanceof Flow.Action.RequestBlockingAction) {
154-
span.setRequestBlockingAction((Flow.Action.RequestBlockingAction) flow.getAction());
154+
if (flow.getAction() instanceof RequestBlockingAction) {
155+
span.setRequestBlockingAction((RequestBlockingAction) flow.getAction());
155156
}
156157
AgentPropagation.ContextVisitor<REQUEST_CARRIER> getter = getter();
157158
if (null != carrier && null != getter) {
@@ -165,25 +166,9 @@ public AgentSpan onRequest(
165166
final CONNECTION connection,
166167
final REQUEST request,
167168
final Context context) {
168-
return onRequest(span, connection, request, getExtractedSpanContext(context));
169-
}
170-
171-
public AgentSpanContext.Extracted getExtractedSpanContext(Context context) {
172-
AgentSpan extractedSpan = AgentSpan.fromContext(context);
173-
return extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context();
174-
}
175-
176-
public AgentSpan onRequest(
177-
final AgentSpan span,
178-
final CONNECTION connection,
179-
final REQUEST request,
180-
@Nullable final AgentSpanContext.Extracted context) {
181169
Config config = Config.get();
182-
boolean clientIpResolverEnabled =
183-
config.isClientIpEnabled()
184-
|| traceClientIpResolverEnabled && ActiveSubsystems.APPSEC_ACTIVE;
185170

186-
if (ActiveSubsystems.APPSEC_ACTIVE) {
171+
if (APPSEC_ACTIVE) {
187172
RequestContext requestContext = span.getRequestContext();
188173
if (requestContext != null) {
189174
BlockResponseFunction brf = createBlockResponseFunction(request, connection);
@@ -193,35 +178,38 @@ public AgentSpan onRequest(
193178
}
194179
Flow<Void> flow = callIGCallbackRequestSessionId(span, request);
195180
Flow.Action action = flow.getAction();
196-
if (action instanceof Flow.Action.RequestBlockingAction) {
197-
span.setRequestBlockingAction((Flow.Action.RequestBlockingAction) flow.getAction());
181+
if (action instanceof RequestBlockingAction) {
182+
span.setRequestBlockingAction((RequestBlockingAction) flow.getAction());
198183
}
199184
}
200185

201-
if (context != null) {
186+
AgentSpanContext.Extracted extracted = getExtractedSpanContext(context);
187+
boolean clientIpResolverEnabled =
188+
config.isClientIpEnabled() || traceClientIpResolverEnabled && APPSEC_ACTIVE;
189+
if (extracted != null) {
202190
if (clientIpResolverEnabled) {
203-
String forwarded = context.getForwarded();
191+
String forwarded = extracted.getForwarded();
204192
if (forwarded != null) {
205193
span.setTag(Tags.HTTP_FORWARDED, forwarded);
206194
}
207-
String forwardedProto = context.getXForwardedProto();
195+
String forwardedProto = extracted.getXForwardedProto();
208196
if (forwardedProto != null) {
209197
span.setTag(Tags.HTTP_FORWARDED_PROTO, forwardedProto);
210198
}
211-
String forwardedHost = context.getXForwardedHost();
199+
String forwardedHost = extracted.getXForwardedHost();
212200
if (forwardedHost != null) {
213201
span.setTag(Tags.HTTP_FORWARDED_HOST, forwardedHost);
214202
}
215-
String forwardedIp = context.getXForwardedFor();
203+
String forwardedIp = extracted.getXForwardedFor();
216204
if (forwardedIp != null) {
217205
span.setTag(Tags.HTTP_FORWARDED_IP, forwardedIp);
218206
}
219-
String forwardedPort = context.getXForwardedPort();
207+
String forwardedPort = extracted.getXForwardedPort();
220208
if (forwardedPort != null) {
221209
span.setTag(Tags.HTTP_FORWARDED_PORT, forwardedPort);
222210
}
223211
}
224-
String userAgent = context.getUserAgent();
212+
String userAgent = extracted.getUserAgent();
225213
if (userAgent != null) {
226214
span.setTag(Tags.HTTP_USER_AGENT, userAgent);
227215
}
@@ -245,8 +233,8 @@ public AgentSpan onRequest(
245233
} else if (supportsRaw) {
246234
span.setTag(Tags.HTTP_URL, URIUtils.lazyInvalidUrl(url.raw()));
247235
}
248-
if (context != null && context.getXForwardedHost() != null) {
249-
span.setTag(Tags.HTTP_HOSTNAME, context.getXForwardedHost());
236+
if (extracted != null && extracted.getXForwardedHost() != null) {
237+
span.setTag(Tags.HTTP_HOSTNAME, extracted.getXForwardedHost());
250238
} else if (url.host() != null) {
251239
span.setTag(Tags.HTTP_HOSTNAME, url.host());
252240
}
@@ -258,8 +246,8 @@ public AgentSpan onRequest(
258246
span.setTag(DDTags.HTTP_FRAGMENT, url.fragment());
259247
}
260248
Flow<Void> flow = callIGCallbackURI(span, url, method);
261-
if (flow.getAction() instanceof Flow.Action.RequestBlockingAction) {
262-
span.setRequestBlockingAction((Flow.Action.RequestBlockingAction) flow.getAction());
249+
if (flow.getAction() instanceof RequestBlockingAction) {
250+
span.setRequestBlockingAction((RequestBlockingAction) flow.getAction());
263251
}
264252
if (valid && SHOULD_SET_URL_RESOURCE_NAME) {
265253
HTTP_RESOURCE_DECORATOR.withServerPath(span, method, path, encoded);
@@ -280,8 +268,8 @@ public AgentSpan onRequest(
280268
}
281269

282270
String inferredAddressStr = null;
283-
if (clientIpResolverEnabled && context != null) {
284-
InetAddress inferredAddress = ClientIpAddressResolver.resolve(context, span);
271+
if (clientIpResolverEnabled && extracted != null) {
272+
InetAddress inferredAddress = ClientIpAddressResolver.resolve(extracted, span);
285273
// the peer address should be used if:
286274
// 1. the headers yield nothing, regardless of whether it is public or not
287275
// 2. it is public and the headers yield a private address
@@ -300,9 +288,9 @@ public AgentSpan onRequest(
300288
span.setTag(Tags.HTTP_CLIENT_IP, inferredAddressStr);
301289
}
302290
} else if (clientIpResolverEnabled && span.getLocalRootSpan() != span) {
303-
// in this case context == null
304-
// If there is no context we can't do anything but use the peer addr.
305-
// Additionally, context == null arises on subspans for which the resolution
291+
// in this case extracted == null
292+
// If there is no extracted we can't do anything but use the peer addr.
293+
// Additionally, extracted == null arises on subspans for which the resolution
306294
// likely already happened on the top span, so we don't need to do the resolution
307295
// again. Instead, copy from the top span, should it exist
308296
AgentSpan localRootSpan = span.getLocalRootSpan();
@@ -321,13 +309,18 @@ public AgentSpan onRequest(
321309
}
322310
setPeerPort(span, peerPort);
323311
Flow<Void> flow = callIGCallbackAddressAndPort(span, peerIp, peerPort, inferredAddressStr);
324-
if (flow.getAction() instanceof Flow.Action.RequestBlockingAction) {
325-
span.setRequestBlockingAction((Flow.Action.RequestBlockingAction) flow.getAction());
312+
if (flow.getAction() instanceof RequestBlockingAction) {
313+
span.setRequestBlockingAction((RequestBlockingAction) flow.getAction());
326314
}
327315

328316
return span;
329317
}
330318

319+
protected static AgentSpanContext.Extracted getExtractedSpanContext(Context context) {
320+
AgentSpan extractedSpan = AgentSpan.fromContext(context);
321+
return extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context();
322+
}
323+
331324
protected BlockResponseFunction createBlockResponseFunction(
332325
REQUEST request, CONNECTION connection) {
333326
return null;
@@ -388,8 +381,8 @@ public AgentSpan onResponse(final AgentSpan span, final RESPONSE response) {
388381
return span;
389382
}
390383

391-
private AgentSpanContext.Extracted callIGCallbackStart(final Context context) {
392-
AgentSpanContext.Extracted extracted = getExtractedSpanContext(context);
384+
private AgentSpanContext.Extracted callIGCallbackStart(
385+
@Nullable final AgentSpanContext.Extracted extracted) {
393386
AgentTracer.TracerAPI tracer = tracer();
394387
Supplier<Flow<Object>> startedCbAppSec =
395388
tracer.getCallbackProvider(RequestContextSlot.APPSEC).getCallback(EVENTS.requestStarted());

0 commit comments

Comments
 (0)