Skip to content

Commit 1a7c76c

Browse files
amarzialiamarziali
committed
add more integrations
1 parent 974f066 commit 1a7c76c

11 files changed

Lines changed: 162 additions & 12 deletions

File tree

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

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

3+
import static datadog.trace.agent.tooling.InstrumenterModule.TargetSystem.CONTEXT_TRACKING;
34
import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.declaresMethod;
45
import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.isAnnotatedWith;
56
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
67
import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext;
8+
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext;
79
import static datadog.trace.bootstrap.instrumentation.decorator.http.HttpResourceDecorator.HTTP_RESOURCE_DECORATOR;
810
import static datadog.trace.instrumentation.azure.functions.AzureFunctionsDecorator.DECORATE;
911
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
@@ -18,6 +20,7 @@
1820
import datadog.context.ContextScope;
1921
import datadog.trace.agent.tooling.Instrumenter;
2022
import datadog.trace.agent.tooling.InstrumenterModule;
23+
import datadog.trace.agent.tooling.annotation.AppliesOn;
2124
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
2225
import net.bytebuddy.asm.Advice;
2326
import net.bytebuddy.description.type.TypeDescription;
@@ -54,20 +57,36 @@ public String[] helperClassNames() {
5457
}
5558

5659
public void methodAdvice(MethodTransformer transformer) {
57-
transformer.applyAdvice(
60+
transformer.applyAdvices(
5861
isMethod()
5962
.and(isPublic())
6063
.and(takesArgument(0, named("com.microsoft.azure.functions.HttpRequestMessage")))
6164
.and(takesArgument(1, named("com.microsoft.azure.functions.ExecutionContext"))),
65+
AzureFunctionsInstrumentation.class.getName() + "$ContextTrackingAdvice",
6266
AzureFunctionsInstrumentation.class.getName() + "$AzureFunctionsAdvice");
6367
}
6468

69+
@AppliesOn(CONTEXT_TRACKING)
70+
public static class ContextTrackingAdvice {
71+
@Advice.OnMethodEnter(suppress = Throwable.class)
72+
public static ContextScope onEnter(@Advice.Argument(0) final HttpRequestMessage<?> request) {
73+
Context parentContext = DECORATE.extract(request);
74+
return parentContext.attach();
75+
}
76+
77+
@Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
78+
public static void closeScope(@Advice.Enter final ContextScope scope) {
79+
scope.close();
80+
}
81+
}
82+
6583
public static class AzureFunctionsAdvice {
6684
@Advice.OnMethodEnter(suppress = Throwable.class)
6785
public static ContextScope methodEnter(
6886
@Advice.Argument(0) final HttpRequestMessage<?> request,
6987
@Advice.Argument(1) final ExecutionContext executionContext) {
70-
final Context parentContext = DECORATE.extract(request);
88+
final Context parentContext =
89+
getCurrentContext(); // parent context attached by ContextTrackingAdvice
7190
final Context context = DECORATE.startSpan(request, parentContext);
7291
final AgentSpan span = fromContext(context);
7392
DECORATE.afterStart(span, executionContext.getFunctionName());

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

Lines changed: 29 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
package datadog.trace.instrumentation.grizzly;
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.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE;
68
import static datadog.trace.instrumentation.grizzly.GrizzlyDecorator.DECORATE;
@@ -10,6 +12,7 @@
1012
import datadog.context.Context;
1113
import datadog.context.ContextScope;
1214
import datadog.trace.agent.tooling.Instrumenter;
15+
import datadog.trace.agent.tooling.annotation.AppliesOn;
1316
import datadog.trace.api.CorrelationIdentifier;
1417
import datadog.trace.api.gateway.Flow;
1518
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -27,14 +30,37 @@ public String instrumentedType() {
2730

2831
@Override
2932
public void methodAdvice(MethodTransformer transformer) {
30-
transformer.applyAdvice(
33+
transformer.applyAdvices(
3134
isMethod()
3235
.and(named("doHandle"))
3336
.and(takesArgument(0, named("org.glassfish.grizzly.http.server.Request")))
3437
.and(takesArgument(1, named("org.glassfish.grizzly.http.server.Response"))),
38+
GrizzlyHttpHandlerInstrumentation.class.getName() + "$ContextTrackingAdvice",
3539
GrizzlyHttpHandlerInstrumentation.class.getName() + "$HandleAdvice");
3640
}
3741

42+
@AppliesOn(CONTEXT_TRACKING)
43+
public static class ContextTrackingAdvice {
44+
45+
@Advice.OnMethodEnter(suppress = Throwable.class)
46+
public static void onEnter(
47+
@Advice.Local("parentScope") ContextScope parentScope,
48+
@Advice.Argument(0) final Request request) {
49+
if (request.getAttribute(DD_CONTEXT_ATTRIBUTE) != null) {
50+
return; // re-entry: HandleAdvice will return false (no-op)
51+
}
52+
Context parentContext = DECORATE.extract(request);
53+
parentScope = parentContext.attach();
54+
}
55+
56+
@Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
57+
public static void closeScope(@Advice.Local("parentScope") ContextScope parentScope) {
58+
if (parentScope != null) {
59+
parentScope.close();
60+
}
61+
}
62+
}
63+
3864
public static class HandleAdvice {
3965

4066
@Advice.OnMethodEnter(suppress = Throwable.class, skipOn = Advice.OnNonDefaultValue.class)
@@ -46,7 +72,8 @@ public static class HandleAdvice {
4672
return false;
4773
}
4874

49-
final Context parentContext = DECORATE.extract(request);
75+
final Context parentContext =
76+
getCurrentContext(); // parent context attached by ContextTrackingAdvice
5077
final Context context = DECORATE.startSpan(request, parentContext);
5178
final AgentSpan span = spanFromContext(context);
5279
DECORATE.afterStart(span);

‎dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-10.0/src/main/java/datadog/trace/instrumentation/jetty10/JettyServerInstrumentation.java‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -89,7 +89,10 @@ public void typeAdvice(TypeTransformer transformer) {
8989

9090
@Override
9191
public void methodAdvice(MethodTransformer transformer) {
92-
transformer.applyAdvice(takesNoArguments().and(named("handle")), packageName + ".HandleAdvice");
92+
transformer.applyAdvices(
93+
takesNoArguments().and(named("handle")),
94+
packageName + ".HandleAdvice$ContextTrackingAdvice",
95+
packageName + ".HandleAdvice");
9396
transformer.applyAdvice(named("recycle").and(takesNoArguments()), packageName + ".ResetAdvice");
9497

9598
if (InstrumenterConfig.get().getAppSecActivation() != ProductActivation.FULLY_DISABLED) {

‎dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-10.0/src/main/java11/datadog/trace/instrumentation/jetty10/HandleAdvice.java‎

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

3+
import static datadog.trace.agent.tooling.InstrumenterModule.TargetSystem.CONTEXT_TRACKING;
4+
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getRootContext;
35
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext;
46
import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE;
7+
import static datadog.trace.instrumentation.jetty10.JettyDecorator.DD_PARENT_CONTEXT_ATTRIBUTE;
58
import static datadog.trace.instrumentation.jetty10.JettyDecorator.DECORATE;
69

710
import datadog.context.Context;
811
import datadog.context.ContextScope;
12+
import datadog.trace.agent.tooling.annotation.AppliesOn;
913
import datadog.trace.api.CorrelationIdentifier;
1014
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
1115
import net.bytebuddy.asm.Advice;
@@ -14,6 +18,30 @@
1418

1519
public class HandleAdvice {
1620

21+
@AppliesOn(CONTEXT_TRACKING)
22+
public static class ContextTrackingAdvice {
23+
24+
@Advice.OnMethodEnter(suppress = Throwable.class)
25+
public static void onEnter(
26+
@Advice.This final HttpChannel channel,
27+
@Advice.Local("parentScope") ContextScope parentScope) {
28+
Request req = channel.getRequest();
29+
if (req.getAttribute(DD_CONTEXT_ATTRIBUTE) instanceof Context) {
30+
return; // re-entry: HandleAdvice will attach existing context
31+
}
32+
Context parentContext = DECORATE.extract(req);
33+
req.setAttribute(DD_PARENT_CONTEXT_ATTRIBUTE, parentContext);
34+
parentScope = parentContext.attach();
35+
}
36+
37+
@Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
38+
public static void closeScope(@Advice.Local("parentScope") ContextScope parentScope) {
39+
if (parentScope != null) {
40+
parentScope.close();
41+
}
42+
}
43+
}
44+
1745
@Advice.OnMethodEnter(suppress = Throwable.class)
1846
public static ContextScope onEnter(
1947
@Advice.This final HttpChannel channel, @Advice.Local("agentSpan") AgentSpan span) {
@@ -24,7 +52,9 @@ public static ContextScope onEnter(
2452
return ((Context) existingContext).attach();
2553
}
2654

27-
final Context parentContext = DECORATE.extract(req);
55+
final Object parentContextObj = req.getAttribute(DD_PARENT_CONTEXT_ATTRIBUTE);
56+
final Context parentContext =
57+
(parentContextObj instanceof Context) ? (Context) parentContextObj : getRootContext();
2858
final Context context = DECORATE.startSpan(req, parentContext);
2959
span = spanFromContext(context);
3060
DECORATE.afterStart(span);

‎dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-10.0/src/main/java11/datadog/trace/instrumentation/jetty10/JettyDecorator.java‎

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

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

‎dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-11.0/src/main/java/datadog/trace/instrumentation/jetty11/JettyServerInstrumentation.java‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@ public String[] helperClassNames() {
5252
packageName + ".JettyDecorator",
5353
packageName + ".RequestURIDataAdapter",
5454
packageName + ".JettyServerAdvice",
55+
packageName + ".JettyServerAdvice$ContextTrackingAdvice",
5556
packageName + ".JettyServerAdvice$HandleAdvice",
5657
packageName + ".JettyServerAdvice$ResetAdvice",
5758
"datadog.trace.instrumentation.jetty.JettyBlockResponseFunction",
@@ -66,8 +67,10 @@ public void typeAdvice(TypeTransformer transformer) {
6667

6768
@Override
6869
public void methodAdvice(MethodTransformer transformer) {
69-
transformer.applyAdvice(
70-
takesNoArguments().and(named("handle")), packageName + ".JettyServerAdvice$HandleAdvice");
70+
transformer.applyAdvices(
71+
takesNoArguments().and(named("handle")),
72+
packageName + ".JettyServerAdvice$ContextTrackingAdvice",
73+
packageName + ".JettyServerAdvice$HandleAdvice");
7174
transformer.applyAdvice(
7275
named("recycle").and(takesNoArguments()), packageName + ".JettyServerAdvice$ResetAdvice");
7376
}

‎dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-11.0/src/main/java11/datadog/trace/instrumentation/jetty11/JettyDecorator.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ public class JettyDecorator extends HttpServerDecorator<Request, Request, Respon
2222

2323
public static final String DD_CONTEXT_PATH_ATTRIBUTE = "datadog.context.path";
2424
public static final String DD_SERVLET_PATH_ATTRIBUTE = "datadog.servlet.path";
25+
public static final String DD_PARENT_CONTEXT_ATTRIBUTE = "datadog.parent-context";
2526

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

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

Lines changed: 32 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,47 @@
11
package datadog.trace.instrumentation.jetty11;
22

3+
import static datadog.trace.agent.tooling.InstrumenterModule.TargetSystem.CONTEXT_TRACKING;
34
import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext;
5+
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getRootContext;
46
import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE;
7+
import static datadog.trace.instrumentation.jetty11.JettyDecorator.DD_PARENT_CONTEXT_ATTRIBUTE;
58
import static datadog.trace.instrumentation.jetty11.JettyDecorator.DECORATE;
69

710
import datadog.context.Context;
811
import datadog.context.ContextScope;
12+
import datadog.trace.agent.tooling.annotation.AppliesOn;
913
import datadog.trace.api.CorrelationIdentifier;
1014
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
1115
import net.bytebuddy.asm.Advice;
1216
import org.eclipse.jetty.server.HttpChannel;
1317
import org.eclipse.jetty.server.Request;
1418

1519
public class JettyServerAdvice {
20+
21+
@AppliesOn(CONTEXT_TRACKING)
22+
public static class ContextTrackingAdvice {
23+
24+
@Advice.OnMethodEnter(suppress = Throwable.class)
25+
public static void onEnter(
26+
@Advice.This final HttpChannel channel,
27+
@Advice.Local("parentScope") ContextScope parentScope) {
28+
Request req = channel.getRequest();
29+
if (req.getAttribute(DD_CONTEXT_ATTRIBUTE) instanceof Context) {
30+
return; // re-entry: HandleAdvice will attach existing context
31+
}
32+
Context parentContext = DECORATE.extract(req);
33+
req.setAttribute(DD_PARENT_CONTEXT_ATTRIBUTE, parentContext);
34+
parentScope = parentContext.attach();
35+
}
36+
37+
@Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
38+
public static void closeScope(@Advice.Local("parentScope") ContextScope parentScope) {
39+
if (parentScope != null) {
40+
parentScope.close();
41+
}
42+
}
43+
}
44+
1645
public static class HandleAdvice {
1746

1847
@Advice.OnMethodEnter(suppress = Throwable.class)
@@ -25,7 +54,9 @@ public static ContextScope onEnter(
2554
return ((Context) existingContext).attach();
2655
}
2756

28-
final Context parentContext = DECORATE.extract(req);
57+
final Object parentContextObj = req.getAttribute(DD_PARENT_CONTEXT_ATTRIBUTE);
58+
final Context parentContext =
59+
(parentContextObj instanceof Context) ? (Context) parentContextObj : getRootContext();
2960
final Context context = DECORATE.startSpan(req, parentContext);
3061
final ContextScope scope = context.attach();
3162
span = fromContext(context);

‎dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-12.0/src/main/java/datadog/trace/instrumentation/jetty12/JettyServerInstrumentation.java‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ public String[] helperClassNames() {
3737
packageName + ".JettyDecorator",
3838
packageName + ".RequestURIDataAdapter",
3939
packageName + ".JettyServerAdvice",
40+
packageName + ".JettyServerAdvice$ContextTrackingAdvice",
4041
packageName + ".JettyServerAdvice$HandleAdvice",
4142
packageName + ".JettyServerAdvice$ResetAdvice",
4243
packageName + ".JettyRunnableWrapper"
@@ -45,8 +46,10 @@ public String[] helperClassNames() {
4546

4647
@Override
4748
public void methodAdvice(MethodTransformer transformer) {
48-
transformer.applyAdvice(
49-
named("onRequest").and(takesArguments(1)), packageName + ".JettyServerAdvice$HandleAdvice");
49+
transformer.applyAdvices(
50+
named("onRequest").and(takesArguments(1)),
51+
packageName + ".JettyServerAdvice$ContextTrackingAdvice",
52+
packageName + ".JettyServerAdvice$HandleAdvice");
5053
transformer.applyAdvice(
5154
named("recycle").and(takesNoArguments()), packageName + ".JettyServerAdvice$ResetAdvice");
5255
}

‎dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-12.0/src/main/java17/datadog/trace/instrumentation/jetty12/JettyDecorator.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ public class JettyDecorator extends HttpServerDecorator<Request, Request, Respon
1919

2020
public static final String DD_CONTEXT_PATH_ATTRIBUTE = "datadog.context.path";
2121
public static final String DD_SERVLET_PATH_ATTRIBUTE = "datadog.servlet.path";
22+
public static final String DD_PARENT_CONTEXT_ATTRIBUTE = "datadog.parent-context";
2223

2324
private static final Class<?> JAVAX_SERVLET_EXCEPTION_CLS =
2425
findClassIfExists("javax.servlet.ServletException");

0 commit comments

Comments
 (0)