Skip to content

Commit aa0d527

Browse files
amarzialiamarziali
committed
more instrumentations
1 parent d3d0031 commit aa0d527

13 files changed

Lines changed: 126 additions & 23 deletions

File tree

‎dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.0/src/main/java/datadog/trace/instrumentation/jetty70/JettyDecorator.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ public class JettyDecorator extends HttpServerDecorator<Request, Request, Respon
2121
UTF8BytesString.create(DECORATE.operationName());
2222
public static final String DD_CONTEXT_PATH_ATTRIBUTE = "datadog.context.path";
2323
public static final String DD_SERVLET_PATH_ATTRIBUTE = "datadog.servlet.path";
24+
public static final String DD_PARENT_CONTEXT_ATTRIBUTE = "datadog.parent-context";
2425

2526
@Override
2627
protected String[] instrumentationNames() {

‎dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.0/src/main/java/datadog/trace/instrumentation/jetty70/JettyServerInstrumentation.java‎

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

3+
import static datadog.trace.agent.tooling.InstrumenterModule.TargetSystem.CONTEXT_TRACKING;
34
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
5+
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getRootContext;
46
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext;
57
import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE;
8+
import static datadog.trace.instrumentation.jetty70.JettyDecorator.DD_PARENT_CONTEXT_ATTRIBUTE;
69
import static datadog.trace.instrumentation.jetty70.JettyDecorator.DECORATE;
710
import static java.util.Collections.singletonMap;
811
import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
@@ -14,6 +17,7 @@
1417
import datadog.context.ContextScope;
1518
import datadog.trace.agent.tooling.Instrumenter;
1619
import datadog.trace.agent.tooling.InstrumenterModule;
20+
import datadog.trace.agent.tooling.annotation.AppliesOn;
1721
import datadog.trace.api.Config;
1822
import datadog.trace.api.CorrelationIdentifier;
1923
import datadog.trace.api.ProductActivation;
@@ -80,8 +84,9 @@ public void typeAdvice(TypeTransformer transformer) {
8084
public void methodAdvice(MethodTransformer transformer) {
8185
transformer.applyAdvice(
8286
isConstructor(), JettyServerInstrumentation.class.getName() + "$ConstructorAdvice");
83-
transformer.applyAdvice(
87+
transformer.applyAdvices(
8488
named("handleRequest").and(takesNoArguments()),
89+
JettyServerInstrumentation.class.getName() + "$ContextTrackingAdvice",
8590
JettyServerInstrumentation.class.getName() + "$HandleRequestAdvice");
8691
transformer.applyAdvice(
8792
named("reset").and(takesArgument(0, boolean.class)),
@@ -138,6 +143,30 @@ public static void link(
138143
}
139144
}
140145

146+
@AppliesOn(CONTEXT_TRACKING)
147+
public static class ContextTrackingAdvice {
148+
@Advice.OnMethodEnter(suppress = Throwable.class)
149+
public static void onEnter(
150+
@Advice.This final HttpConnection connection,
151+
@Advice.Local("parentScope") ContextScope parentScope) {
152+
Request req = connection.getRequest();
153+
Object existingContext = req.getAttribute(DD_CONTEXT_ATTRIBUTE);
154+
if (existingContext instanceof Context) {
155+
return; // re-entry: HandleRequestAdvice will attach existing context
156+
}
157+
Context parentContext = DECORATE.extract(req);
158+
req.setAttribute(DD_PARENT_CONTEXT_ATTRIBUTE, parentContext);
159+
parentScope = parentContext.attach();
160+
}
161+
162+
@Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
163+
public static void closeScope(@Advice.Local("parentScope") ContextScope parentScope) {
164+
if (parentScope != null) {
165+
parentScope.close();
166+
}
167+
}
168+
}
169+
141170
/**
142171
* The handleRequest call denotes the earliest point at which the incoming request is fully
143172
* parsed. This allows us to read the headers from the request to extract propagation info.
@@ -155,7 +184,9 @@ public static ContextScope onEnter(
155184
return ((Context) existingContext).attach();
156185
}
157186

158-
final Context parentContext = DECORATE.extract(req);
187+
final Object parentContextObj = req.getAttribute(DD_PARENT_CONTEXT_ATTRIBUTE);
188+
final Context parentContext =
189+
(parentContextObj instanceof Context) ? (Context) parentContextObj : getRootContext();
159190
final Context context = DECORATE.startSpan(req, parentContext);
160191
final ContextScope scope = context.attach();
161192
span = spanFromContext(context);

‎dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.6/src/main/java/datadog/trace/instrumentation/jetty76/JettyDecorator.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ public class JettyDecorator extends HttpServerDecorator<Request, Request, Respon
2121
UTF8BytesString.create(DECORATE.operationName());
2222
public static final String DD_CONTEXT_PATH_ATTRIBUTE = "datadog.context.path";
2323
public static final String DD_SERVLET_PATH_ATTRIBUTE = "datadog.servlet.path";
24+
public static final String DD_PARENT_CONTEXT_ATTRIBUTE = "datadog.parent-context";
2425

2526
@Override
2627
protected String[] instrumentationNames() {

‎dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.6/src/main/java/datadog/trace/instrumentation/jetty76/JettyServerInstrumentation.java‎

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

3+
import static datadog.trace.agent.tooling.InstrumenterModule.TargetSystem.CONTEXT_TRACKING;
34
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
5+
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getRootContext;
46
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext;
57
import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE;
68
import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_FIN_DISP_LIST_SPAN_ATTRIBUTE;
9+
import static datadog.trace.instrumentation.jetty76.JettyDecorator.DD_PARENT_CONTEXT_ATTRIBUTE;
710
import static datadog.trace.instrumentation.jetty76.JettyDecorator.DECORATE;
811
import static java.util.Collections.singletonMap;
912
import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
@@ -14,6 +17,7 @@
1417
import datadog.context.ContextScope;
1518
import datadog.trace.agent.tooling.Instrumenter;
1619
import datadog.trace.agent.tooling.InstrumenterModule;
20+
import datadog.trace.agent.tooling.annotation.AppliesOn;
1721
import datadog.trace.api.Config;
1822
import datadog.trace.api.CorrelationIdentifier;
1923
import datadog.trace.api.ProductActivation;
@@ -80,8 +84,9 @@ public void typeAdvice(TypeTransformer transformer) {
8084
public void methodAdvice(MethodTransformer transformer) {
8185
transformer.applyAdvice(
8286
isConstructor(), JettyServerInstrumentation.class.getName() + "$ConstructorAdvice");
83-
transformer.applyAdvice(
87+
transformer.applyAdvices(
8488
named("handleRequest").and(takesNoArguments()),
89+
JettyServerInstrumentation.class.getName() + "$ContextTrackingAdvice",
8590
JettyServerInstrumentation.class.getName() + "$HandleRequestAdvice");
8691
transformer.applyAdvice(
8792
named("reset").and(takesNoArguments()),
@@ -138,6 +143,30 @@ public static void link(
138143
}
139144
}
140145

146+
@AppliesOn(CONTEXT_TRACKING)
147+
public static class ContextTrackingAdvice {
148+
@Advice.OnMethodEnter(suppress = Throwable.class)
149+
public static void onEnter(
150+
@Advice.This final AbstractHttpConnection connection,
151+
@Advice.Local("parentScope") ContextScope parentScope) {
152+
Request req = connection.getRequest();
153+
Object existingContext = req.getAttribute(DD_CONTEXT_ATTRIBUTE);
154+
if (existingContext instanceof Context) {
155+
return; // re-entry: HandleRequestAdvice will attach existing context
156+
}
157+
Context parentContext = DECORATE.extract(req);
158+
req.setAttribute(DD_PARENT_CONTEXT_ATTRIBUTE, parentContext);
159+
parentScope = parentContext.attach();
160+
}
161+
162+
@Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
163+
public static void closeScope(@Advice.Local("parentScope") ContextScope parentScope) {
164+
if (parentScope != null) {
165+
parentScope.close();
166+
}
167+
}
168+
}
169+
141170
/**
142171
* The handleRequest call denotes the earliest point at which the incoming request is fully
143172
* parsed. This allows us to read the headers from the request to extract propagation info.
@@ -156,7 +185,9 @@ public static ContextScope onEnter(
156185
return ((Context) existingContext).attach();
157186
}
158187

159-
final Context parentContext = DECORATE.extract(req);
188+
final Object parentContextObj = req.getAttribute(DD_PARENT_CONTEXT_ATTRIBUTE);
189+
final Context parentContext =
190+
(parentContextObj instanceof Context) ? (Context) parentContextObj : getRootContext();
160191
final Context context = DECORATE.startSpan(req, parentContext);
161192
final ContextScope scope = context.attach();
162193
span = spanFromContext(context);

‎dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java/datadog/trace/instrumentation/jetty9/JettyServerInstrumentation.java‎

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

3-
import static datadog.context.Context.root;
43
import static datadog.trace.agent.tooling.InstrumenterModule.TargetSystem.CONTEXT_TRACKING;
54
import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.declaresMethod;
65
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
76
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.namedOneOf;
7+
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getRootContext;
88
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext;
99
import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE;
1010
import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType.RUNNABLE;
@@ -191,7 +191,7 @@ public static ContextScope onEnter(
191191

192192
Object parentContextObj = req.getAttribute(DD_PARENT_CONTEXT_ATTRIBUTE);
193193
Context parentContext =
194-
(parentContextObj instanceof Context) ? (Context) parentContextObj : root();
194+
(parentContextObj instanceof Context) ? (Context) parentContextObj : getRootContext();
195195
final Context context = DECORATE.startSpan(req, parentContext);
196196
final ContextScope scope = context.attach();
197197
span = spanFromContext(context);

‎dd-java-agent/instrumentation/jms/javax-jms-1.1/src/main/java/datadog/trace/instrumentation/jms/MDBMessageConsumerInstrumentation.java‎

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

3-
import static datadog.context.Context.root;
43
import static datadog.context.propagation.Propagators.defaultPropagator;
54
import static datadog.trace.agent.tooling.InstrumenterModule.TargetSystem.CONTEXT_TRACKING;
65
import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.declaresAnnotation;
@@ -9,6 +8,7 @@
98
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
109
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
1110
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
11+
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getRootContext;
1212
import static datadog.trace.instrumentation.jms.JMSDecorator.CONSUMER_DECORATE;
1313
import static datadog.trace.instrumentation.jms.JMSDecorator.JMS_CONSUME;
1414
import static datadog.trace.instrumentation.jms.JMSDecorator.logJMSException;
@@ -70,7 +70,7 @@ public static class ContextPropagationAdvice {
7070
@Advice.OnMethodEnter(suppress = Throwable.class)
7171
public static void onEnter(
7272
@Advice.Argument(0) final Message message, @Advice.Local("ctxScope") ContextScope scope) {
73-
scope = defaultPropagator().extract(root(), message, GETTER).attach();
73+
scope = defaultPropagator().extract(getRootContext(), message, GETTER).attach();
7474
}
7575

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

‎dd-java-agent/instrumentation/kafka/kafka-streams-0.11/src/main/java/datadog/trace/instrumentation/kafka_streams/KafkaStreamTaskInstrumentation.java‎

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

3-
import static datadog.context.Context.root;
43
import static datadog.context.propagation.Propagators.defaultPropagator;
54
import static datadog.trace.agent.tooling.InstrumenterModule.TargetSystem.CONTEXT_TRACKING;
65
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
@@ -11,6 +10,7 @@
1110
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
1211
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
1312
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.traceConfig;
13+
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getRootContext;
1414
import static datadog.trace.instrumentation.kafka_common.StreamingContext.STREAMING_CONTEXT;
1515
import static datadog.trace.instrumentation.kafka_common.Utils.computePayloadSizeBytes;
1616
import static datadog.trace.instrumentation.kafka_streams.KafkaStreamsDecorator.BROKER_DECORATE;
@@ -232,7 +232,7 @@ public static void onEnter(
232232
return;
233233
}
234234
if (!Config.get().isKafkaClientPropagationDisabledForTopic(record.topic())) {
235-
scope = defaultPropagator().extract(root(), record, SR_GETTER).attach();
235+
scope = defaultPropagator().extract(getRootContext(), record, SR_GETTER).attach();
236236
}
237237
}
238238

@@ -253,7 +253,7 @@ public static void onEnter(
253253
return;
254254
}
255255
if (!Config.get().isKafkaClientPropagationDisabledForTopic(record.topic())) {
256-
scope = defaultPropagator().extract(root(), record, PR_GETTER).attach();
256+
scope = defaultPropagator().extract(getRootContext(), record, PR_GETTER).attach();
257257
}
258258
}
259259

‎dd-java-agent/instrumentation/liberty/liberty-20.0/src/main/java/datadog/trace/instrumentation/liberty20/LibertyServerInstrumentation.java‎

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

3-
import static datadog.context.Context.root;
43
import static datadog.trace.agent.tooling.InstrumenterModule.TargetSystem.CONTEXT_TRACKING;
54
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
65
import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext;
6+
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getRootContext;
77
import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE;
88
import static datadog.trace.instrumentation.liberty20.HttpInboundServiceContextImplInstrumentation.REQUEST_MSG_TYPE;
99
import static datadog.trace.instrumentation.liberty20.LibertyDecorator.DD_PARENT_CONTEXT_ATTRIBUTE;
@@ -138,7 +138,7 @@ public static class HandleRequestAdvice {
138138

139139
Object parentContextObj = request.getAttribute(DD_PARENT_CONTEXT_ATTRIBUTE);
140140
final Context parentContext =
141-
(parentContextObj instanceof Context) ? (Context) parentContextObj : root();
141+
(parentContextObj instanceof Context) ? (Context) parentContextObj : getRootContext();
142142
final Context context = DECORATE.startSpan(request, parentContext);
143143
scope = context.attach();
144144
final AgentSpan span = fromContext(context);

‎dd-java-agent/instrumentation/liberty/liberty-23.0/src/main/java/datadog/trace/instrumentation/liberty23/LibertyServerInstrumentation.java‎

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

3-
import static datadog.context.Context.root;
43
import static datadog.trace.agent.tooling.InstrumenterModule.TargetSystem.CONTEXT_TRACKING;
54
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
65
import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext;
6+
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getRootContext;
77
import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE;
88
import static datadog.trace.instrumentation.liberty23.HttpInboundServiceContextImplInstrumentation.REQUEST_MSG_TYPE;
99
import static datadog.trace.instrumentation.liberty23.LibertyDecorator.DD_PARENT_CONTEXT_ATTRIBUTE;
@@ -140,7 +140,7 @@ public static class HandleRequestAdvice {
140140

141141
Object parentContextObj = request.getAttribute(DD_PARENT_CONTEXT_ATTRIBUTE);
142142
final Context parentContext =
143-
(parentContextObj instanceof Context) ? (Context) parentContextObj : root();
143+
(parentContextObj instanceof Context) ? (Context) parentContextObj : getRootContext();
144144
final Context context = DECORATE.startSpan(request, parentContext);
145145
scope = context.attach();
146146
final AgentSpan span = fromContext(context);

‎dd-java-agent/instrumentation/restlet-2.2/src/main/java/datadog/trace/instrumentation/restlet/RestletInstrumentation.java‎

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

3+
import static datadog.trace.agent.tooling.InstrumenterModule.TargetSystem.CONTEXT_TRACKING;
34
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
5+
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext;
46
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext;
57
import static datadog.trace.instrumentation.restlet.RestletDecorator.DECORATE;
68
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
@@ -12,6 +14,7 @@
1214
import datadog.context.ContextScope;
1315
import datadog.trace.agent.tooling.Instrumenter;
1416
import datadog.trace.agent.tooling.InstrumenterModule;
17+
import datadog.trace.agent.tooling.annotation.AppliesOn;
1518
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
1619
import net.bytebuddy.asm.Advice;
1720

@@ -33,10 +36,11 @@ public String[] knownMatchingTypes() {
3336

3437
@Override
3538
public void methodAdvice(MethodTransformer transformer) {
36-
transformer.applyAdvice(
39+
transformer.applyAdvices(
3740
isMethod()
3841
.and(named("handle"))
3942
.and(takesArgument(0, named("com.sun.net.httpserver.HttpExchange"))),
43+
getClass().getName() + "$ContextTrackingAdvice",
4044
getClass().getName() + "$RestletHandleAdvice");
4145
}
4246

@@ -51,10 +55,25 @@ public String[] helperClassNames() {
5155
};
5256
}
5357

58+
@AppliesOn(CONTEXT_TRACKING)
59+
public static class ContextTrackingAdvice {
60+
@Advice.OnMethodEnter(suppress = Throwable.class)
61+
public static ContextScope onEnter(@Advice.Argument(0) final HttpExchange exchange) {
62+
Context parentContext = DECORATE.extract(exchange);
63+
return parentContext.attach();
64+
}
65+
66+
@Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
67+
public static void closeScope(@Advice.Enter final ContextScope scope) {
68+
scope.close();
69+
}
70+
}
71+
5472
public static class RestletHandleAdvice {
5573
@Advice.OnMethodEnter(suppress = Throwable.class)
5674
public static ContextScope beginRequest(@Advice.Argument(0) final HttpExchange exchange) {
57-
Context parentContext = DECORATE.extract(exchange);
75+
Context parentContext =
76+
getCurrentContext(); // parent context attached by ContextTrackingAdvice
5877
Context context = DECORATE.startSpan(exchange, parentContext);
5978
AgentSpan span = spanFromContext(context);
6079
ContextScope scope = context.attach();

0 commit comments

Comments
 (0)