Skip to content

Commit f9278b6

Browse files
PerfectSlayerPerfectSlayer
authored
Improve Instrumenter API to use Context instead of Span (#9211)
1 parent ef2e9f0 commit f9278b6

File tree

40 files changed

+457
-456
lines changed

40 files changed

+457
-456
lines changed

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

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

3+
import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext;
34
import static java.util.concurrent.TimeUnit.MICROSECONDS;
45
import static java.util.concurrent.TimeUnit.SECONDS;
56

@@ -58,7 +59,8 @@ public void setUp() {
5859
.build();
5960
GlobalTracer.forceRegister(tracer);
6061
decorator = new BenchmarkHttpServerDecorator();
61-
span = decorator.startSpan(Collections.emptyMap(), (Context) null);
62+
Context context = decorator.startSpan(Collections.emptyMap(), Context.root());
63+
span = fromContext(context);
6264
}
6365

6466
@Benchmark
@@ -99,7 +101,7 @@ public BenchmarkHttpServerDecorator() {
99101

100102
@Override
101103
protected String[] instrumentationNames() {
102-
return new String[0];
104+
return new String[] {"benchmark"};
103105
}
104106

105107
@Override

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

Lines changed: 35 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,7 @@ public abstract class HttpServerDecorator<REQUEST, CONNECTION, RESPONSE, REQUEST
7272
Config.get().isRuleEnabled("URLAsResourceNameRule");
7373

7474
private static final BitSet SERVER_ERROR_STATUSES = Config.get().getHttpServerErrorStatuses();
75+
private static final String DEFAULT_INSTRUMENTATION_NAME = "http-server";
7576

7677
private final boolean traceClientIpResolverEnabled =
7778
Config.get().isTraceClientIpResolverEnabled();
@@ -118,6 +119,12 @@ protected AgentTracer.TracerAPI tracer() {
118119
return AgentTracer.get();
119120
}
120121

122+
/**
123+
* Extracts context from an upstream service.
124+
*
125+
* @param carrier The request carrier to get the context from.
126+
* @return The extracted context, {@code Context#root()} if no valid context to extract.
127+
*/
121128
public Context extract(REQUEST_CARRIER carrier) {
122129
AgentPropagation.ContextVisitor<REQUEST_CARRIER> getter = getter();
123130
if (null == carrier || null == getter) {
@@ -126,12 +133,22 @@ public Context extract(REQUEST_CARRIER carrier) {
126133
return Propagators.defaultPropagator().extract(root(), carrier, getter);
127134
}
128135

129-
public AgentSpan startSpan(
130-
String instrumentationName, REQUEST_CARRIER carrier, AgentSpanContext.Extracted context) {
136+
/**
137+
* Starts a span.
138+
*
139+
* @param carrier The request carrier.
140+
* @param context The parent context of the span to create.
141+
* @return A new context bundling the span, child of the given parent context.
142+
*/
143+
public Context startSpan(REQUEST_CARRIER carrier, Context context) {
144+
String[] instrumentationNames = instrumentationNames();
145+
String instrumentationName =
146+
instrumentationNames != null && instrumentationNames.length > 0
147+
? instrumentationNames[0]
148+
: DEFAULT_INSTRUMENTATION_NAME;
149+
AgentSpanContext.Extracted extracted = callIGCallbackStart(context);
131150
AgentSpan span =
132-
tracer()
133-
.startSpan(instrumentationName, spanName(), callIGCallbackStart(context))
134-
.setMeasured(true);
151+
tracer().startSpan(instrumentationName, spanName(), extracted).setMeasured(true);
135152
Flow<Void> flow = callIGCallbackRequestHeaders(span, carrier);
136153
if (flow.getAction() instanceof Flow.Action.RequestBlockingAction) {
137154
span.setRequestBlockingAction((Flow.Action.RequestBlockingAction) flow.getAction());
@@ -140,16 +157,7 @@ public AgentSpan startSpan(
140157
if (null != carrier && null != getter) {
141158
tracer().getDataStreamsMonitoring().setCheckpoint(span, forHttpServer());
142159
}
143-
return span;
144-
}
145-
146-
public AgentSpan startSpan(REQUEST_CARRIER carrier, Context context) {
147-
return startSpan("http-server", carrier, getExtractedSpanContext(context));
148-
}
149-
150-
public AgentSpanContext.Extracted getExtractedSpanContext(Context context) {
151-
AgentSpan extractedSpan = AgentSpan.fromContext(context);
152-
return extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context();
160+
return context.with(span);
153161
}
154162

155163
public AgentSpan onRequest(
@@ -160,6 +168,11 @@ public AgentSpan onRequest(
160168
return onRequest(span, connection, request, getExtractedSpanContext(context));
161169
}
162170

171+
public AgentSpanContext.Extracted getExtractedSpanContext(Context context) {
172+
AgentSpan extractedSpan = AgentSpan.fromContext(context);
173+
return extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context();
174+
}
175+
163176
public AgentSpan onRequest(
164177
final AgentSpan span,
165178
final CONNECTION connection,
@@ -375,34 +388,23 @@ public AgentSpan onResponse(final AgentSpan span, final RESPONSE response) {
375388
return span;
376389
}
377390

378-
// @Override
379-
// public Span onError(final Span span, final Throwable throwable) {
380-
// assert span != null;
381-
// // FIXME
382-
// final Object status = span.getTag("http.status");
383-
// if (status == null || status.equals(200)) {
384-
// // Ensure status set correctly
385-
// span.setTag("http.status", 500);
386-
// }
387-
// return super.onError(span, throwable);
388-
// }
389-
390-
private AgentSpanContext.Extracted callIGCallbackStart(AgentSpanContext.Extracted context) {
391+
private AgentSpanContext.Extracted callIGCallbackStart(final Context context) {
392+
AgentSpanContext.Extracted extracted = getExtractedSpanContext(context);
391393
AgentTracer.TracerAPI tracer = tracer();
392394
Supplier<Flow<Object>> startedCbAppSec =
393395
tracer.getCallbackProvider(RequestContextSlot.APPSEC).getCallback(EVENTS.requestStarted());
394396
Supplier<Flow<Object>> startedCbIast =
395397
tracer.getCallbackProvider(RequestContextSlot.IAST).getCallback(EVENTS.requestStarted());
396398

397399
if (startedCbAppSec == null && startedCbIast == null) {
398-
return context;
400+
return extracted;
399401
}
400402

401403
TagContext tagContext = null;
402-
if (context == null) {
404+
if (extracted == null) {
403405
tagContext = new TagContext();
404-
} else if (context instanceof TagContext) {
405-
tagContext = (TagContext) context;
406+
} else if (extracted instanceof TagContext) {
407+
tagContext = (TagContext) extracted;
406408
}
407409

408410
if (tagContext != null) {
@@ -415,7 +417,7 @@ private AgentSpanContext.Extracted callIGCallbackStart(AgentSpanContext.Extracte
415417
return tagContext;
416418
}
417419

418-
return context;
420+
return extracted;
419421
}
420422

421423
@Override

‎dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecoratorTest.groovy

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ import datadog.trace.core.datastreams.DataStreamsMonitoring
2525
import java.util.function.Function
2626
import java.util.function.Supplier
2727

28+
import static datadog.context.Context.root
2829
import static datadog.trace.api.config.TraceInstrumentationConfig.HTTP_SERVER_DECODED_RESOURCE_PRESERVE_SPACES
2930
import static datadog.trace.api.config.TraceInstrumentationConfig.HTTP_SERVER_RAW_QUERY_STRING
3031
import static datadog.trace.api.config.TraceInstrumentationConfig.HTTP_SERVER_RAW_RESOURCE
@@ -497,7 +498,7 @@ class HttpServerDecoratorTest extends ServerDecoratorTest {
497498
def decorator = newDecorator(mTracer, null)
498499

499500
when:
500-
decorator.startSpan("test", headers, null)
501+
decorator.startSpan(headers, root())
501502

502503
then:
503504
1 * mSpan.setMeasured(true) >> mSpan

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

Lines changed: 6 additions & 4 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.instrumentation.akkahttp.AkkaHttpServerDecorator.DECORATE;
45

56
import akka.http.scaladsl.model.HttpRequest;
@@ -10,12 +11,13 @@
1011

1112
public class DatadogWrapperHelper {
1213
public static ContextScope createSpan(final HttpRequest request) {
13-
final Context context = DECORATE.extract(request);
14-
final AgentSpan span = DECORATE.startSpan(request, context);
14+
final Context parentContext = DECORATE.extract(request);
15+
final Context context = DECORATE.startSpan(request, parentContext);
16+
final AgentSpan span = fromContext(context);
1517
DECORATE.afterStart(span);
16-
DECORATE.onRequest(span, request, request, context);
18+
DECORATE.onRequest(span, request, request, parentContext);
1719

18-
return context.with(span).attach();
20+
return context.attach();
1921
}
2022

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

‎dd-java-agent/instrumentation/azure-functions/src/main/java/datadog/trace/instrumentation/azure/functions/AzureFunctionsDecorator.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
package datadog.trace.instrumentation.azurefunctions;
1+
package datadog.trace.instrumentation.azure.functions;
22

33
import com.microsoft.azure.functions.HttpRequestMessage;
44
import com.microsoft.azure.functions.HttpResponseMessage;

‎dd-java-agent/instrumentation/azure-functions/src/main/java/datadog/trace/instrumentation/azure/functions/AzureFunctionsInstrumentation.java

Lines changed: 12 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,11 @@
1-
package datadog.trace.instrumentation.azurefunctions;
1+
package datadog.trace.instrumentation.azure.functions;
22

33
import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.declaresMethod;
44
import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.isAnnotatedWith;
55
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
6-
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext;
6+
import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext;
77
import static datadog.trace.bootstrap.instrumentation.decorator.http.HttpResourceDecorator.HTTP_RESOURCE_DECORATOR;
8-
import static datadog.trace.instrumentation.azurefunctions.AzureFunctionsDecorator.DECORATE;
8+
import static datadog.trace.instrumentation.azure.functions.AzureFunctionsDecorator.DECORATE;
99
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
1010
import static net.bytebuddy.matcher.ElementMatchers.isPublic;
1111
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
@@ -65,23 +65,24 @@ public void methodAdvice(MethodTransformer transformer) {
6565
public static class AzureFunctionsAdvice {
6666
@Advice.OnMethodEnter(suppress = Throwable.class)
6767
public static ContextScope methodEnter(
68-
@Advice.Argument(0) final HttpRequestMessage request,
69-
@Advice.Argument(1) final ExecutionContext context) {
70-
final Context extractedContext = DECORATE.extract(request);
71-
final AgentSpan span = DECORATE.startSpan(request, extractedContext);
72-
DECORATE.afterStart(span, context.getFunctionName());
73-
DECORATE.onRequest(span, request, request, extractedContext);
68+
@Advice.Argument(0) final HttpRequestMessage<?> request,
69+
@Advice.Argument(1) final ExecutionContext executionContext) {
70+
final Context parentContext = DECORATE.extract(request);
71+
final Context context = DECORATE.startSpan(request, parentContext);
72+
final AgentSpan span = fromContext(context);
73+
DECORATE.afterStart(span, executionContext.getFunctionName());
74+
DECORATE.onRequest(span, request, request, parentContext);
7475
HTTP_RESOURCE_DECORATOR.withRoute(
7576
span, request.getHttpMethod().name(), request.getUri().getPath());
76-
return extractedContext.with(span).attach();
77+
return context.attach();
7778
}
7879

7980
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
8081
public static void methodExit(
8182
@Advice.Enter final ContextScope scope,
8283
@Advice.Return final HttpResponseMessage response,
8384
@Advice.Thrown final Throwable throwable) {
84-
final AgentSpan span = spanFromContext(scope.context());
85+
final AgentSpan span = fromContext(scope.context());
8586
DECORATE.onError(span, throwable);
8687
DECORATE.onResponse(span, response);
8788
DECORATE.beforeFinish(span);

‎dd-java-agent/instrumentation/azure-functions/src/main/java/datadog/trace/instrumentation/azure/functions/HttpRequestMessageExtractAdapter.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
package datadog.trace.instrumentation.azurefunctions;
1+
package datadog.trace.instrumentation.azure.functions;
22

33
import com.microsoft.azure.functions.HttpRequestMessage;
44
import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;

‎dd-java-agent/instrumentation/grizzly-2/src/main/java/datadog/trace/instrumentation/grizzly/GrizzlyHttpHandlerInstrumentation.java

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@
1313
import datadog.trace.agent.tooling.Instrumenter;
1414
import datadog.trace.agent.tooling.InstrumenterModule;
1515
import datadog.trace.api.CorrelationIdentifier;
16-
import datadog.trace.api.GlobalTracer;
1716
import datadog.trace.api.gateway.Flow;
1817
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
1918
import net.bytebuddy.asm.Advice;
@@ -74,15 +73,17 @@ public static class HandleAdvice {
7473
}
7574

7675
final Context parentContext = DECORATE.extract(request);
77-
final AgentSpan span = DECORATE.startSpan(request, parentContext);
76+
final Context context = DECORATE.startSpan(request, parentContext);
77+
final AgentSpan span = spanFromContext(context);
7878
DECORATE.afterStart(span);
7979
DECORATE.onRequest(span, request, request, parentContext);
8080

81-
scope = parentContext.with(span).attach();
81+
scope = context.attach();
8282

8383
request.setAttribute(DD_SPAN_ATTRIBUTE, span);
84-
request.setAttribute(CorrelationIdentifier.getTraceIdKey(), GlobalTracer.get().getTraceId());
85-
request.setAttribute(CorrelationIdentifier.getSpanIdKey(), GlobalTracer.get().getSpanId());
84+
request.setAttribute(
85+
CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId());
86+
request.setAttribute(CorrelationIdentifier.getSpanIdKey(), CorrelationIdentifier.getSpanId());
8687

8788
Flow.Action.RequestBlockingAction rba = span.getRequestBlockingAction();
8889
if (rba != null) {

‎dd-java-agent/instrumentation/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/GrizzlyDecorator.java

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

3+
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext;
4+
35
import datadog.appsec.api.blocking.BlockingContentType;
46
import datadog.context.Context;
57
import datadog.context.ContextScope;
@@ -112,13 +114,14 @@ public static NextAction onHttpCodecFilterExit(
112114
}
113115
HttpRequestPacket httpRequest = (HttpRequestPacket) httpHeader;
114116
HttpResponsePacket httpResponse = httpRequest.getResponse();
115-
Context context = DECORATE.extract(httpRequest);
116-
AgentSpan span = DECORATE.startSpan(httpRequest, context);
117-
ContextScope scope = context.with(span).attach();
117+
Context parentContext = DECORATE.extract(httpRequest);
118+
Context context = DECORATE.startSpan(httpRequest, parentContext);
119+
ContextScope scope = context.attach();
120+
AgentSpan span = spanFromContext(context);
118121
DECORATE.afterStart(span);
119122
ctx.getAttributes().setAttribute(DD_SPAN_ATTRIBUTE, span);
120123
ctx.getAttributes().setAttribute(DD_RESPONSE_ATTRIBUTE, httpResponse);
121-
DECORATE.onRequest(span, httpRequest, httpRequest, context);
124+
DECORATE.onRequest(span, httpRequest, httpRequest, parentContext);
122125

123126
Flow.Action.RequestBlockingAction rba = span.getRequestBlockingAction();
124127
if (rba != null && thiz instanceof HttpServerFilter) {

‎dd-java-agent/instrumentation/jetty-11/src/main/java11/datadog/trace/instrumentation/jetty11/JettyServerAdvice.java

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

3+
import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext;
34
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
45
import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE;
56
import static datadog.trace.instrumentation.jetty11.JettyDecorator.DECORATE;
67

78
import datadog.context.Context;
89
import datadog.context.ContextScope;
910
import datadog.trace.api.CorrelationIdentifier;
10-
import datadog.trace.api.GlobalTracer;
1111
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
1212
import net.bytebuddy.asm.Advice;
1313
import org.eclipse.jetty.server.HttpChannel;
@@ -26,16 +26,17 @@ public static ContextScope onEnter(
2626
return activateSpan((AgentSpan) existingSpan);
2727
}
2828

29-
final Context context = DECORATE.extract(req);
30-
span = DECORATE.startSpan(req, context);
31-
final ContextScope scope = context.with(span).attach();
29+
final Context parentContext = DECORATE.extract(req);
30+
final Context context = DECORATE.startSpan(req, parentContext);
31+
final ContextScope scope = context.attach();
32+
span = fromContext(context);
3233
span.setMeasured(true);
3334
DECORATE.afterStart(span);
34-
DECORATE.onRequest(span, req, req, context);
35+
DECORATE.onRequest(span, req, req, parentContext);
3536

3637
req.setAttribute(DD_SPAN_ATTRIBUTE, span);
37-
req.setAttribute(CorrelationIdentifier.getTraceIdKey(), GlobalTracer.get().getTraceId());
38-
req.setAttribute(CorrelationIdentifier.getSpanIdKey(), GlobalTracer.get().getSpanId());
38+
req.setAttribute(CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId());
39+
req.setAttribute(CorrelationIdentifier.getSpanIdKey(), CorrelationIdentifier.getSpanId());
3940
return scope;
4041
}
4142

0 commit comments

Comments
 (0)