Skip to content

Commit b12a14b

Browse files
PerfectSlayerPerfectSlayer
committed
feat(context): Refactor HttpServerDecorator to return Context instead of span on startSpan
1 parent ca2caa8 commit b12a14b

File tree

32 files changed

+238
-181
lines changed

32 files changed

+238
-181
lines changed

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

Lines changed: 3 additions & 1 deletion
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("benchmark", Collections.emptyMap(), Context.root());
63+
span = fromContext(context);
6264
}
6365

6466
@Benchmark

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

Lines changed: 30 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,12 @@ protected AgentTracer.TracerAPI tracer() {
118118
return AgentTracer.get();
119119
}
120120

121+
/**
122+
* Extracts context from an upstream service.
123+
*
124+
* @param carrier The request carrier to get the context from.
125+
* @return The extracted context, {@code Context#root()} if no valid context to extract.
126+
*/
121127
public Context extract(REQUEST_CARRIER carrier) {
122128
AgentPropagation.ContextVisitor<REQUEST_CARRIER> getter = getter();
123129
if (null == carrier || null == getter) {
@@ -126,11 +132,22 @@ public Context extract(REQUEST_CARRIER carrier) {
126132
return Propagators.defaultPropagator().extract(root(), carrier, getter);
127133
}
128134

129-
public AgentSpan startSpan(
130-
String instrumentationName, REQUEST_CARRIER carrier, AgentSpanContext.Extracted context) {
135+
/**
136+
* Starts a span.
137+
*
138+
* @param instrumentationName The instrumentation creating the span.
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(
144+
String instrumentationName,
145+
REQUEST_CARRIER carrier,
146+
Context context) { // TODO Rename "context" to "parent"?
147+
AgentSpanContext.Extracted spanContext = getExtractedSpanContext(context);
131148
AgentSpan span =
132149
tracer()
133-
.startSpan(instrumentationName, spanName(), callIGCallbackStart(context))
150+
.startSpan(instrumentationName, spanName(), callIGCallbackStart(spanContext))
134151
.setMeasured(true);
135152
Flow<Void> flow = callIGCallbackRequestHeaders(span, carrier);
136153
if (flow.getAction() instanceof Flow.Action.RequestBlockingAction) {
@@ -140,12 +157,17 @@ public AgentSpan startSpan(
140157
if (null != carrier && null != getter) {
141158
tracer().getDataStreamsMonitoring().setCheckpoint(span, forHttpServer());
142159
}
143-
return span;
160+
return context.with(span);
144161
}
145162

146-
public AgentSpan startSpan(REQUEST_CARRIER carrier, Context context) {
147-
return startSpan("http-server", carrier, getExtractedSpanContext(context));
148-
}
163+
// /**
164+
// * Use {@code startSpan(String instrumentationName, REQUEST_CARRIER carrier, Context context)}
165+
// instead.
166+
// */
167+
// @Deprecated
168+
// public Context startSpan(REQUEST_CARRIER carrier, Context context) {
169+
// return startSpan("http-server", carrier, context);
170+
// }
149171

150172
public AgentSpanContext.Extracted getExtractedSpanContext(Context context) {
151173
AgentSpan extractedSpan = AgentSpan.fromContext(context);
@@ -387,7 +409,7 @@ public AgentSpan onResponse(final AgentSpan span, final RESPONSE response) {
387409
// return super.onError(span, throwable);
388410
// }
389411

390-
private AgentSpanContext.Extracted callIGCallbackStart(AgentSpanContext.Extracted context) {
412+
private AgentSpanContext.Extracted callIGCallbackStart(final AgentSpanContext.Extracted context) {
391413
AgentTracer.TracerAPI tracer = tracer();
392414
Supplier<Flow<Object>> startedCbAppSec =
393415
tracer.getCallbackProvider(RequestContextSlot.APPSEC).getCallback(EVENTS.requestStarted());

‎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("test", 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("akka-http", 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: 11 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,12 @@
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.AgentSpan.fromContext;
67
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext;
78
import static datadog.trace.bootstrap.instrumentation.decorator.http.HttpResourceDecorator.HTTP_RESOURCE_DECORATOR;
8-
import static datadog.trace.instrumentation.azurefunctions.AzureFunctionsDecorator.DECORATE;
9+
import static datadog.trace.instrumentation.azure.functions.AzureFunctionsDecorator.DECORATE;
910
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
1011
import static net.bytebuddy.matcher.ElementMatchers.isPublic;
1112
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
@@ -65,15 +66,16 @@ public void methodAdvice(MethodTransformer transformer) {
6566
public static class AzureFunctionsAdvice {
6667
@Advice.OnMethodEnter(suppress = Throwable.class)
6768
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);
69+
@Advice.Argument(0) final HttpRequestMessage<?> request,
70+
@Advice.Argument(1) final ExecutionContext executionContext) {
71+
final Context parentContext = DECORATE.extract(request);
72+
final Context context = DECORATE.startSpan("azure-functions", request, parentContext);
73+
final AgentSpan span = fromContext(context);
74+
DECORATE.afterStart(span, executionContext.getFunctionName());
75+
DECORATE.onRequest(span, request, request, parentContext);
7476
HTTP_RESOURCE_DECORATOR.withRoute(
7577
span, request.getHttpMethod().name(), request.getUri().getPath());
76-
return extractedContext.with(span).attach();
78+
return context.attach();
7779
}
7880

7981
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)

‎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: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package datadog.trace.instrumentation.grizzly;
22

33
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4+
import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext;
45
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext;
56
import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE;
67
import static datadog.trace.instrumentation.grizzly.GrizzlyDecorator.DECORATE;
@@ -13,7 +14,6 @@
1314
import datadog.trace.agent.tooling.Instrumenter;
1415
import datadog.trace.agent.tooling.InstrumenterModule;
1516
import datadog.trace.api.CorrelationIdentifier;
16-
import datadog.trace.api.GlobalTracer;
1717
import datadog.trace.api.gateway.Flow;
1818
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
1919
import net.bytebuddy.asm.Advice;
@@ -74,15 +74,17 @@ public static class HandleAdvice {
7474
}
7575

7676
final Context parentContext = DECORATE.extract(request);
77-
final AgentSpan span = DECORATE.startSpan(request, parentContext);
77+
final Context context = DECORATE.startSpan("grizzly", request, parentContext);
78+
final AgentSpan span = fromContext(context);
7879
DECORATE.afterStart(span);
7980
DECORATE.onRequest(span, request, request, parentContext);
8081

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

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

8789
Flow.Action.RequestBlockingAction rba = span.getRequestBlockingAction();
8890
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.AgentSpan.fromContext;
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("grizzly", httpRequest, parentContext);
119+
ContextScope scope = context.attach();
120+
AgentSpan span = fromContext(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("jetty", 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)