Skip to content

Commit 6930b27

Browse files
amarzialiamarziali
committed
fix bad ai stuff
1 parent 460265b commit 6930b27

File tree

22 files changed

+102
-152
lines changed

22 files changed

+102
-152
lines changed

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

Lines changed: 4 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@
33
import static datadog.trace.agent.tooling.InstrumenterModule.TargetSystem.CONTEXT_TRACKING;
44
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
55
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
6-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
76
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
87
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext;
98
import static datadog.trace.instrumentation.akkahttp.AkkaHttpClientDecorator.AKKA_CLIENT_REQUEST;
@@ -66,8 +65,7 @@ public void methodAdvice(MethodTransformer transformer) {
6665

6766
public static class SingleRequestAdvice {
6867
@Advice.OnMethodEnter(suppress = Throwable.class)
69-
public static AgentScope methodEnter(
70-
@Advice.Argument(value = 0, readOnly = false) HttpRequest request) {
68+
public static AgentScope methodEnter(@Advice.Argument(value = 0) final HttpRequest request) {
7169
/*
7270
Versions 10.0 and 10.1 have slightly different structure that is hard to distinguish so here
7371
we cast 'wider net' and avoid instrumenting twice.
@@ -113,15 +111,11 @@ public static class SingleRequestContextPropagationAdvice {
113111
@Advice.OnMethodEnter(suppress = Throwable.class)
114112
public static void methodEnter(
115113
@Advice.Argument(value = 0, readOnly = false) HttpRequest request) {
116-
final AkkaHttpHeaders headers = new AkkaHttpHeaders(request);
117-
if (headers.hadSpan()) {
114+
if (request == null) {
118115
return;
119116
}
120-
AgentSpan span = activeSpan();
121-
if (span == null) {
122-
return;
123-
}
124-
DECORATE.injectContext(getCurrentContext().with(span), request, headers);
117+
final AkkaHttpHeaders headers = new AkkaHttpHeaders(request);
118+
DECORATE.injectContext(getCurrentContext(), request, headers);
125119
request = headers.getRequest();
126120
}
127121
}

‎dd-java-agent/instrumentation/akka/akka-http/akka-http-10.6/src/main/java11/datadog/trace/instrumentation/akkahttp106/SingleRequestAdvice.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
public class SingleRequestAdvice {
1717
@Advice.OnMethodEnter(suppress = Throwable.class)
1818
public static AgentScope methodEnter(
19-
@Advice.Argument(value = 0, readOnly = false) HttpRequest request) {
19+
@Advice.Argument(value = 0) final HttpRequest request) {
2020
final AkkaHttpClientHelpers.AkkaHttpHeaders headers =
2121
new AkkaHttpClientHelpers.AkkaHttpHeaders(request);
2222
if (headers.hadSpan()) {

‎dd-java-agent/instrumentation/akka/akka-http/akka-http-10.6/src/main/java11/datadog/trace/instrumentation/akkahttp106/SingleRequestContextPropagationAdvice.java‎

Lines changed: 5 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -15,16 +15,13 @@ public class SingleRequestContextPropagationAdvice {
1515
@Advice.OnMethodEnter(suppress = Throwable.class)
1616
public static void methodEnter(
1717
@Advice.Argument(value = 0, readOnly = false) HttpRequest request) {
18-
final AkkaHttpClientHelpers.AkkaHttpHeaders headers =
19-
new AkkaHttpClientHelpers.AkkaHttpHeaders(request);
20-
if (headers.hadSpan()) {
21-
return;
22-
}
23-
AgentSpan span = activeSpan();
24-
if (span == null) {
18+
if (request == null) {
2519
return;
2620
}
27-
DECORATE.injectContext(getCurrentContext().with(span), request, headers);
21+
final AkkaHttpClientHelpers.AkkaHttpHeaders headers =
22+
new AkkaHttpClientHelpers.AkkaHttpHeaders(request);
23+
24+
DECORATE.injectContext(getCurrentContext(), request, headers);
2825
request = headers.getRequest();
2926
}
3027
}

‎dd-java-agent/instrumentation/apache-httpclient/apache-httpasyncclient-4.0/src/main/java/datadog/trace/instrumentation/apachehttpasyncclient/ApacheHttpAsyncClientInstrumentation.java‎

Lines changed: 19 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -71,23 +71,37 @@ public void methodAdvice(MethodTransformer transformer) {
7171
.and(takesArgument(1, named("org.apache.http.nio.protocol.HttpAsyncResponseConsumer")))
7272
.and(takesArgument(2, named("org.apache.http.protocol.HttpContext")))
7373
.and(takesArgument(3, named("org.apache.http.concurrent.FutureCallback"))),
74-
ApacheHttpAsyncClientInstrumentation.class.getName() + "$ClientAdvice",
75-
ApacheHttpAsyncClientInstrumentation.class.getName() + "$ClientContextPropagationAdvice");
74+
ApacheHttpAsyncClientInstrumentation.class.getName() + "$ClientContextPropagationAdvice",
75+
ApacheHttpAsyncClientInstrumentation.class.getName() + "$ClientAdvice");
7676
}
7777

78-
public static class ClientAdvice {
78+
@AppliesOn(CONTEXT_TRACKING)
79+
public static class ClientContextPropagationAdvice {
80+
@Advice.OnMethodEnter(suppress = Throwable.class)
81+
public static void methodEnter(
82+
@Advice.Argument(value = 0, readOnly = false) HttpAsyncRequestProducer requestProducer) {
83+
final DelegatingRequestProducer delegatingRequestProducer =
84+
new DelegatingRequestProducer(requestProducer);
85+
delegatingRequestProducer.setInjectContext(true);
86+
requestProducer = delegatingRequestProducer;
87+
}
88+
}
7989

90+
public static class ClientAdvice {
8091
@Advice.OnMethodEnter(suppress = Throwable.class)
8192
public static AgentSpan methodEnter(
8293
@Advice.Argument(value = 0, readOnly = false) HttpAsyncRequestProducer requestProducer,
8394
@Advice.Argument(2) HttpContext context,
8495
@Advice.Argument(value = 3, readOnly = false) FutureCallback<?> futureCallback) {
8596

97+
if (!(requestProducer instanceof DelegatingRequestProducer)) {
98+
requestProducer = new DelegatingRequestProducer(requestProducer);
99+
}
100+
86101
final AgentScope.Continuation parentContinuation = captureActiveSpan();
87102
final AgentSpan clientSpan = startSpan(HTTP_REQUEST);
88103
DECORATE.afterStart(clientSpan);
89-
90-
requestProducer = new DelegatingRequestProducer(clientSpan, requestProducer);
104+
((DelegatingRequestProducer) requestProducer).setSpan(clientSpan);
91105
futureCallback =
92106
new TraceContinuedFutureCallback<>(
93107
parentContinuation, clientSpan, context, futureCallback);
@@ -107,14 +121,4 @@ public static void methodExit(
107121
}
108122
}
109123
}
110-
111-
@AppliesOn(CONTEXT_TRACKING)
112-
public static class ClientContextPropagationAdvice {
113-
@Advice.OnMethodEnter(suppress = Throwable.class)
114-
public static void onEnter(@Advice.Argument(0) final HttpAsyncRequestProducer requestProducer) {
115-
if (requestProducer instanceof DelegatingRequestProducer) {
116-
((DelegatingRequestProducer) requestProducer).setInjectContext(true);
117-
}
118-
}
119-
}
120124
}

‎dd-java-agent/instrumentation/apache-httpclient/apache-httpasyncclient-4.0/src/main/java/datadog/trace/instrumentation/apachehttpasyncclient/DelegatingRequestProducer.java‎

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import static datadog.trace.instrumentation.apachehttpasyncclient.ApacheHttpAsyncClientDecorator.DECORATE;
55
import static datadog.trace.instrumentation.apachehttpasyncclient.HttpHeadersInjectAdapter.SETTER;
66

7+
import datadog.context.Context;
78
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
89
import java.io.IOException;
910
import org.apache.http.HttpException;
@@ -15,19 +16,22 @@
1516
import org.apache.http.protocol.HttpContext;
1617

1718
public class DelegatingRequestProducer implements HttpAsyncRequestProducer {
18-
final AgentSpan span;
19+
AgentSpan span;
1920
final HttpAsyncRequestProducer delegate;
2021
boolean injectContext = false;
2122

22-
public DelegatingRequestProducer(final AgentSpan span, final HttpAsyncRequestProducer delegate) {
23-
this.span = span;
23+
public DelegatingRequestProducer(final HttpAsyncRequestProducer delegate) {
2424
this.delegate = delegate;
2525
}
2626

2727
public void setInjectContext(boolean injectContext) {
2828
this.injectContext = injectContext;
2929
}
3030

31+
public void setSpan(final AgentSpan span) {
32+
this.span = span;
33+
}
34+
3135
@Override
3236
public HttpHost getTarget() {
3337
return delegate.getTarget();
@@ -38,7 +42,11 @@ public HttpRequest generateRequest() throws IOException, HttpException {
3842
final HttpRequest request = delegate.generateRequest();
3943
DECORATE.onRequest(span, new HostAndRequestAsHttpUriRequest(delegate.getTarget(), request));
4044
if (injectContext) {
41-
DECORATE.injectContext(current().with(span), request, SETTER);
45+
Context receiver = current();
46+
if (span != null) {
47+
receiver = receiver.with(span);
48+
}
49+
DECORATE.injectContext(receiver, request, SETTER);
4250
}
4351
return request;
4452
}

‎dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-4.0/src/main/java/datadog/trace/instrumentation/apachehttpclient/HelperMethods.java‎

Lines changed: 2 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@
22

33
import static datadog.context.Context.current;
44
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
5-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
65
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
76
import static datadog.trace.instrumentation.apachehttpclient.ApacheHttpClientDecorator.DECORATE;
87
import static datadog.trace.instrumentation.apachehttpclient.ApacheHttpClientDecorator.HTTP_REQUEST;
@@ -49,11 +48,7 @@ public static void doInjectContext(final HttpUriRequest request) {
4948
if (request.containsHeader("amz-sdk-invocation-id")) {
5049
return;
5150
}
52-
final AgentSpan span = activeSpan();
53-
if (span == null) {
54-
return;
55-
}
56-
DECORATE.injectContext(current().with(span), request, SETTER);
51+
DECORATE.injectContext(current(), request, SETTER);
5752
}
5853

5954
public static void doInjectContext(final HttpHost host, final HttpRequest request) {
@@ -63,14 +58,7 @@ public static void doInjectContext(final HttpHost host, final HttpRequest reques
6358
} else {
6459
uriRequest = new HostAndRequestAsHttpUriRequest(host, request);
6560
}
66-
if (uriRequest.containsHeader("amz-sdk-invocation-id")) {
67-
return;
68-
}
69-
final AgentSpan span = activeSpan();
70-
if (span == null) {
71-
return;
72-
}
73-
DECORATE.injectContext(current().with(span), uriRequest, SETTER);
61+
doInjectContext(uriRequest);
7462
}
7563

7664
public static void doMethodExit(

‎dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/ApacheHttpAsyncClientInstrumentation.java‎

Lines changed: 19 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -87,8 +87,20 @@ public void methodAdvice(MethodTransformer transformer) {
8787
.and(takesArgument(2, named("org.apache.hc.core5.http.nio.HandlerFactory")))
8888
.and(takesArgument(3, named("org.apache.hc.core5.http.protocol.HttpContext")))
8989
.and(takesArgument(4, named("org.apache.hc.core5.concurrent.FutureCallback"))),
90-
this.getClass().getName() + "$ClientAdvice",
91-
this.getClass().getName() + "$ClientContextPropagationAdvice");
90+
this.getClass().getName() + "$ClientContextPropagationAdvice",
91+
this.getClass().getName() + "$ClientAdvice");
92+
}
93+
94+
@AppliesOn(CONTEXT_TRACKING)
95+
public static class ClientContextPropagationAdvice {
96+
@Advice.OnMethodEnter(suppress = Throwable.class)
97+
public static void onEnter(
98+
@Advice.Argument(value = 0, readOnly = false) AsyncRequestProducer requestProducer) {
99+
final DelegatingRequestProducer delegatingRequestProducer =
100+
new DelegatingRequestProducer(requestProducer);
101+
delegatingRequestProducer.setInjectContext(true);
102+
requestProducer = delegatingRequestProducer;
103+
}
92104
}
93105

94106
public static class ClientAdvice {
@@ -108,7 +120,11 @@ public static AgentScope methodEnter(
108120
context = new BasicHttpContext();
109121
}
110122

111-
requestProducer = new DelegatingRequestProducer(clientSpan, requestProducer);
123+
if (!(requestProducer instanceof DelegatingRequestProducer)) {
124+
requestProducer = new DelegatingRequestProducer(requestProducer);
125+
}
126+
127+
((DelegatingRequestProducer) requestProducer).setSpan(clientSpan);
112128
futureCallback =
113129
new TraceContinuedFutureCallback<>(
114130
parentContinuation, clientSpan, context, futureCallback);
@@ -136,14 +152,4 @@ public static void methodExit(
136152
}
137153
}
138154
}
139-
140-
@AppliesOn(CONTEXT_TRACKING)
141-
public static class ClientContextPropagationAdvice {
142-
@Advice.OnMethodEnter(suppress = Throwable.class)
143-
public static void onEnter(@Advice.Argument(0) final AsyncRequestProducer requestProducer) {
144-
if (requestProducer instanceof DelegatingRequestProducer) {
145-
((DelegatingRequestProducer) requestProducer).setInjectContext(true);
146-
}
147-
}
148-
}
149155
}

‎dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/ApacheHttpClientInstrumentation.java‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,7 @@ public void methodAdvice(MethodTransformer transformer) {
7878
.and(takesArguments(1))
7979
.and(takesArgument(0, named("org.apache.hc.core5.http.ClassicHttpRequest"))),
8080
ApacheHttpClientInstrumentation.class.getName() + "$RequestAdvice",
81-
ApacheHttpClientInstrumentation.class.getName() + "$RequestContextPropagationAdvice");
81+
ApacheHttpClientInstrumentation.class.getName() + "$ContextPropagationAdviceArg0");
8282

8383
transformer.applyAdvices(
8484
isMethod()
@@ -87,7 +87,7 @@ public void methodAdvice(MethodTransformer transformer) {
8787
.and(takesArgument(0, named("org.apache.hc.core5.http.ClassicHttpRequest")))
8888
.and(takesArgument(1, named("org.apache.hc.core5.http.protocol.HttpContext"))),
8989
ApacheHttpClientInstrumentation.class.getName() + "$RequestAdvice",
90-
ApacheHttpClientInstrumentation.class.getName() + "$RequestContextPropagationAdvice");
90+
ApacheHttpClientInstrumentation.class.getName() + "$ContextPropagationAdviceArg0");
9191

9292
transformer.applyAdvices(
9393
isMethod()
@@ -97,7 +97,7 @@ public void methodAdvice(MethodTransformer transformer) {
9797
.and(takesArgument(1, named("org.apache.hc.core5.http.ClassicHttpRequest")))
9898
.and(takesArgument(2, named("org.apache.hc.core5.http.protocol.HttpContext"))),
9999
ApacheHttpClientInstrumentation.class.getName() + "$HostRequestAdvice",
100-
ApacheHttpClientInstrumentation.class.getName() + "$HostRequestContextPropagationAdvice");
100+
ApacheHttpClientInstrumentation.class.getName() + "$ContextPropagationAdviceArg1");
101101

102102
transformer.applyAdvices(
103103
isMethod()
@@ -106,7 +106,7 @@ public void methodAdvice(MethodTransformer transformer) {
106106
.and(takesArgument(0, named("org.apache.hc.core5.http.HttpHost")))
107107
.and(takesArgument(1, named("org.apache.hc.core5.http.ClassicHttpRequest"))),
108108
ApacheHttpClientInstrumentation.class.getName() + "$HostRequestAdvice",
109-
ApacheHttpClientInstrumentation.class.getName() + "$HostRequestContextPropagationAdvice");
109+
ApacheHttpClientInstrumentation.class.getName() + "$ContextPropagationAdviceArg1");
110110

111111
transformer.applyAdvices(
112112
isMethod()
@@ -117,7 +117,7 @@ public void methodAdvice(MethodTransformer transformer) {
117117
.and(takesArgument(2, named("org.apache.hc.core5.http.protocol.HttpContext")))
118118
.and(takesArgument(3, named("org.apache.hc.core5.http.io.HttpClientResponseHandler"))),
119119
ApacheHttpClientInstrumentation.class.getName() + "$ResponseHandlerAdvice",
120-
ApacheHttpClientInstrumentation.class.getName() + "$HostRequestContextPropagationAdvice");
120+
ApacheHttpClientInstrumentation.class.getName() + "$ContextPropagationAdviceArg1");
121121
}
122122

123123
public static class RequestAdvice {
@@ -203,15 +203,15 @@ public static void methodExit(
203203
}
204204

205205
@AppliesOn(CONTEXT_TRACKING)
206-
public static class RequestContextPropagationAdvice {
206+
public static class ContextPropagationAdviceArg0 {
207207
@Advice.OnMethodEnter(suppress = Throwable.class)
208208
public static void methodEnter(@Advice.Argument(0) final ClassicHttpRequest request) {
209209
HelperMethods.doInjectContext(request);
210210
}
211211
}
212212

213213
@AppliesOn(CONTEXT_TRACKING)
214-
public static class HostRequestContextPropagationAdvice {
214+
public static class ContextPropagationAdviceArg1 {
215215
@Advice.OnMethodEnter(suppress = Throwable.class)
216216
public static void methodEnter(@Advice.Argument(1) final ClassicHttpRequest request) {
217217
HelperMethods.doInjectContext(request);

‎dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/DelegatingRequestProducer.java‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,15 +9,18 @@
99
import org.apache.hc.core5.http.protocol.HttpContext;
1010

1111
public class DelegatingRequestProducer implements AsyncRequestProducer {
12-
final AgentSpan span;
12+
AgentSpan span;
1313
final AsyncRequestProducer delegate;
1414
boolean injectContext = false;
1515

16-
public DelegatingRequestProducer(final AgentSpan span, final AsyncRequestProducer delegate) {
17-
this.span = span;
16+
public DelegatingRequestProducer(final AsyncRequestProducer delegate) {
1817
this.delegate = delegate;
1918
}
2019

20+
public void setSpan(AgentSpan span) {
21+
this.span = span;
22+
}
23+
2124
public void setInjectContext(boolean injectContext) {
2225
this.injectContext = injectContext;
2326
}

‎dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/HelperMethods.java‎

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@
22

33
import static datadog.context.Context.current;
44
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
5-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
65
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
76
import static datadog.trace.instrumentation.apachehttpclient5.ApacheHttpClientDecorator.DECORATE;
87
import static datadog.trace.instrumentation.apachehttpclient5.ApacheHttpClientDecorator.HTTP_REQUEST;
@@ -49,11 +48,7 @@ public static void doInjectContext(final HttpRequest request) {
4948
if (request.containsHeader("amz-sdk-invocation-id")) {
5049
return;
5150
}
52-
final AgentSpan span = activeSpan();
53-
if (span == null) {
54-
return;
55-
}
56-
DECORATE.injectContext(current().with(span), request, SETTER);
51+
DECORATE.injectContext(current(), request, SETTER);
5752
}
5853

5954
public static void doMethodExit(

0 commit comments

Comments
 (0)