diff --git a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/AkkaHttpSingleRequestInstrumentation.java b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/AkkaHttpSingleRequestInstrumentation.java index 29e413e9ece..cff65c5f4c7 100644 --- a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/AkkaHttpSingleRequestInstrumentation.java +++ b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/AkkaHttpSingleRequestInstrumentation.java @@ -1,5 +1,6 @@ package datadog.trace.instrumentation.akkahttp; +import static datadog.trace.agent.tooling.InstrumenterModule.TargetSystem.CONTEXT_TRACKING; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; @@ -16,6 +17,7 @@ import com.google.auto.service.AutoService; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.agent.tooling.annotation.AppliesOn; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import net.bytebuddy.asm.Advice; @@ -47,20 +49,23 @@ public String[] helperClassNames() { @Override public void methodAdvice(MethodTransformer transformer) { // This is mainly for compatibility with 10.0 - transformer.applyAdvice( + transformer.applyAdvices( named("singleRequest").and(takesArgument(0, named("akka.http.scaladsl.model.HttpRequest"))), - AkkaHttpSingleRequestInstrumentation.class.getName() + "$SingleRequestAdvice"); + AkkaHttpSingleRequestInstrumentation.class.getName() + "$SingleRequestAdvice", + AkkaHttpSingleRequestInstrumentation.class.getName() + + "$SingleRequestContextPropagationAdvice"); // This is for 10.1+ - transformer.applyAdvice( + transformer.applyAdvices( named("singleRequestImpl") .and(takesArgument(0, named("akka.http.scaladsl.model.HttpRequest"))), - AkkaHttpSingleRequestInstrumentation.class.getName() + "$SingleRequestAdvice"); + AkkaHttpSingleRequestInstrumentation.class.getName() + "$SingleRequestAdvice", + AkkaHttpSingleRequestInstrumentation.class.getName() + + "$SingleRequestContextPropagationAdvice"); } public static class SingleRequestAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static AgentScope methodEnter( - @Advice.Argument(value = 0, readOnly = false) HttpRequest request) { + public static AgentScope methodEnter(@Advice.Argument(value = 0) final HttpRequest request) { /* Versions 10.0 and 10.1 have slightly different structure that is hard to distinguish so here we cast 'wider net' and avoid instrumenting twice. @@ -75,12 +80,6 @@ public static AgentScope methodEnter( final AgentSpan span = startSpan("akka-http", AKKA_CLIENT_REQUEST); DECORATE.afterStart(span); DECORATE.onRequest(span, request); - - if (request != null) { - DECORATE.injectContext(getCurrentContext().with(span), request, headers); - // Request is immutable, so we have to assign new value once we update headers - request = headers.getRequest(); - } return activateSpan(span); } @@ -106,4 +105,18 @@ public static void methodExit( scope.close(); } } + + @AppliesOn(CONTEXT_TRACKING) + public static class SingleRequestContextPropagationAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void methodEnter( + @Advice.Argument(value = 0, readOnly = false) HttpRequest request) { + if (request == null) { + return; + } + final AkkaHttpHeaders headers = new AkkaHttpHeaders(request); + DECORATE.injectContext(getCurrentContext(), request, headers); + request = headers.getRequest(); + } + } } diff --git a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.6/src/main/java/datadog/trace/instrumentation/akkahttp106/AkkaHttpSingleRequestInstrumentation.java b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.6/src/main/java/datadog/trace/instrumentation/akkahttp106/AkkaHttpSingleRequestInstrumentation.java index 7d7596befaa..ee5b362437b 100644 --- a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.6/src/main/java/datadog/trace/instrumentation/akkahttp106/AkkaHttpSingleRequestInstrumentation.java +++ b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.6/src/main/java/datadog/trace/instrumentation/akkahttp106/AkkaHttpSingleRequestInstrumentation.java @@ -32,8 +32,9 @@ public String[] helperClassNames() { @Override public void methodAdvice(MethodTransformer transformer) { - transformer.applyAdvice( + transformer.applyAdvices( named("singleRequest").and(takesArgument(0, named("akka.http.scaladsl.model.HttpRequest"))), - packageName + ".SingleRequestAdvice"); + packageName + ".SingleRequestAdvice", + packageName + ".SingleRequestContextPropagationAdvice"); } } diff --git a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.6/src/main/java11/datadog/trace/instrumentation/akkahttp106/SingleRequestAdvice.java b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.6/src/main/java11/datadog/trace/instrumentation/akkahttp106/SingleRequestAdvice.java index e70e90208e3..9ab8e3a3cce 100644 --- a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.6/src/main/java11/datadog/trace/instrumentation/akkahttp106/SingleRequestAdvice.java +++ b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.6/src/main/java11/datadog/trace/instrumentation/akkahttp106/SingleRequestAdvice.java @@ -2,7 +2,6 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; -import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; import static datadog.trace.instrumentation.akkahttp106.AkkaHttpClientDecorator.AKKA_CLIENT_REQUEST; import static datadog.trace.instrumentation.akkahttp106.AkkaHttpClientDecorator.DECORATE; @@ -16,8 +15,7 @@ public class SingleRequestAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static AgentScope methodEnter( - @Advice.Argument(value = 0, readOnly = false) HttpRequest request) { + public static AgentScope methodEnter(@Advice.Argument(value = 0) final HttpRequest request) { final AkkaHttpClientHelpers.AkkaHttpHeaders headers = new AkkaHttpClientHelpers.AkkaHttpHeaders(request); if (headers.hadSpan()) { @@ -27,11 +25,6 @@ public static AgentScope methodEnter( final AgentSpan span = startSpan("akka-http", AKKA_CLIENT_REQUEST); DECORATE.afterStart(span); DECORATE.onRequest(span, request); - if (request != null) { - DECORATE.injectContext(getCurrentContext().with(span), request, headers); - // Request is immutable, so we have to assign new value once we update headers - request = headers.getRequest(); - } return activateSpan(span); } diff --git a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.6/src/main/java11/datadog/trace/instrumentation/akkahttp106/SingleRequestContextPropagationAdvice.java b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.6/src/main/java11/datadog/trace/instrumentation/akkahttp106/SingleRequestContextPropagationAdvice.java new file mode 100644 index 00000000000..49f9ce99e7b --- /dev/null +++ b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.6/src/main/java11/datadog/trace/instrumentation/akkahttp106/SingleRequestContextPropagationAdvice.java @@ -0,0 +1,25 @@ +package datadog.trace.instrumentation.akkahttp106; + +import static datadog.trace.agent.tooling.InstrumenterModule.TargetSystem.CONTEXT_TRACKING; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; +import static datadog.trace.instrumentation.akkahttp106.AkkaHttpClientDecorator.DECORATE; + +import akka.http.scaladsl.model.HttpRequest; +import datadog.trace.agent.tooling.annotation.AppliesOn; +import net.bytebuddy.asm.Advice; + +@AppliesOn(CONTEXT_TRACKING) +public class SingleRequestContextPropagationAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void methodEnter( + @Advice.Argument(value = 0, readOnly = false) HttpRequest request) { + if (request == null) { + return; + } + final AkkaHttpClientHelpers.AkkaHttpHeaders headers = + new AkkaHttpClientHelpers.AkkaHttpHeaders(request); + + DECORATE.injectContext(getCurrentContext(), request, headers); + request = headers.getRequest(); + } +} diff --git a/dd-java-agent/instrumentation/apache-httpclient/apache-httpasyncclient-4.0/src/main/java/datadog/trace/instrumentation/apachehttpasyncclient/ApacheHttpAsyncClientInstrumentation.java b/dd-java-agent/instrumentation/apache-httpclient/apache-httpasyncclient-4.0/src/main/java/datadog/trace/instrumentation/apachehttpasyncclient/ApacheHttpAsyncClientInstrumentation.java index 525761c110a..316cf60f6f5 100644 --- a/dd-java-agent/instrumentation/apache-httpclient/apache-httpasyncclient-4.0/src/main/java/datadog/trace/instrumentation/apachehttpasyncclient/ApacheHttpAsyncClientInstrumentation.java +++ b/dd-java-agent/instrumentation/apache-httpclient/apache-httpasyncclient-4.0/src/main/java/datadog/trace/instrumentation/apachehttpasyncclient/ApacheHttpAsyncClientInstrumentation.java @@ -1,5 +1,6 @@ package datadog.trace.instrumentation.apachehttpasyncclient; +import static datadog.trace.agent.tooling.InstrumenterModule.TargetSystem.CONTEXT_TRACKING; import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.captureActiveSpan; @@ -12,9 +13,11 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArguments; import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.agent.tooling.annotation.AppliesOn; import datadog.trace.api.InstrumenterConfig; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; import net.bytebuddy.asm.Advice; import net.bytebuddy.description.type.TypeDescription; import net.bytebuddy.matcher.ElementMatcher; @@ -61,7 +64,7 @@ public ElementMatcher hierarchyMatcher() { @Override public void methodAdvice(MethodTransformer transformer) { - transformer.applyAdvice( + transformer.applyAdvices( isMethod() .and(named("execute")) .and(takesArguments(4)) @@ -69,22 +72,39 @@ public void methodAdvice(MethodTransformer transformer) { .and(takesArgument(1, named("org.apache.http.nio.protocol.HttpAsyncResponseConsumer"))) .and(takesArgument(2, named("org.apache.http.protocol.HttpContext"))) .and(takesArgument(3, named("org.apache.http.concurrent.FutureCallback"))), + ApacheHttpAsyncClientInstrumentation.class.getName() + "$ClientContextPropagationAdvice", ApacheHttpAsyncClientInstrumentation.class.getName() + "$ClientAdvice"); } - public static class ClientAdvice { + @AppliesOn(CONTEXT_TRACKING) + @SuppressFBWarnings("UC_USELESS_OBJECT") + public static class ClientContextPropagationAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void methodEnter( + @Advice.Argument(value = 0, readOnly = false) HttpAsyncRequestProducer requestProducer) { + final DelegatingRequestProducer delegatingRequestProducer = + new DelegatingRequestProducer(requestProducer); + delegatingRequestProducer.setInjectContext(true); + requestProducer = delegatingRequestProducer; + } + } + @SuppressFBWarnings("UC_USELESS_OBJECT") + public static class ClientAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static AgentSpan methodEnter( @Advice.Argument(value = 0, readOnly = false) HttpAsyncRequestProducer requestProducer, @Advice.Argument(2) HttpContext context, @Advice.Argument(value = 3, readOnly = false) FutureCallback futureCallback) { + if (!(requestProducer instanceof DelegatingRequestProducer)) { + requestProducer = new DelegatingRequestProducer(requestProducer); + } + final AgentScope.Continuation parentContinuation = captureActiveSpan(); final AgentSpan clientSpan = startSpan(HTTP_REQUEST); DECORATE.afterStart(clientSpan); - - requestProducer = new DelegatingRequestProducer(clientSpan, requestProducer); + ((DelegatingRequestProducer) requestProducer).setSpan(clientSpan); futureCallback = new TraceContinuedFutureCallback<>( parentContinuation, clientSpan, context, futureCallback); diff --git a/dd-java-agent/instrumentation/apache-httpclient/apache-httpasyncclient-4.0/src/main/java/datadog/trace/instrumentation/apachehttpasyncclient/DelegatingRequestProducer.java b/dd-java-agent/instrumentation/apache-httpclient/apache-httpasyncclient-4.0/src/main/java/datadog/trace/instrumentation/apachehttpasyncclient/DelegatingRequestProducer.java index cc4aee9342e..21ab1157320 100644 --- a/dd-java-agent/instrumentation/apache-httpclient/apache-httpasyncclient-4.0/src/main/java/datadog/trace/instrumentation/apachehttpasyncclient/DelegatingRequestProducer.java +++ b/dd-java-agent/instrumentation/apache-httpclient/apache-httpasyncclient-4.0/src/main/java/datadog/trace/instrumentation/apachehttpasyncclient/DelegatingRequestProducer.java @@ -4,6 +4,7 @@ import static datadog.trace.instrumentation.apachehttpasyncclient.ApacheHttpAsyncClientDecorator.DECORATE; import static datadog.trace.instrumentation.apachehttpasyncclient.HttpHeadersInjectAdapter.SETTER; +import datadog.context.Context; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import java.io.IOException; import org.apache.http.HttpException; @@ -15,14 +16,22 @@ import org.apache.http.protocol.HttpContext; public class DelegatingRequestProducer implements HttpAsyncRequestProducer { - final AgentSpan span; + AgentSpan span; final HttpAsyncRequestProducer delegate; + boolean injectContext = false; - public DelegatingRequestProducer(final AgentSpan span, final HttpAsyncRequestProducer delegate) { - this.span = span; + public DelegatingRequestProducer(final HttpAsyncRequestProducer delegate) { this.delegate = delegate; } + public void setInjectContext(boolean injectContext) { + this.injectContext = injectContext; + } + + public void setSpan(final AgentSpan span) { + this.span = span; + } + @Override public HttpHost getTarget() { return delegate.getTarget(); @@ -32,7 +41,13 @@ public HttpHost getTarget() { public HttpRequest generateRequest() throws IOException, HttpException { final HttpRequest request = delegate.generateRequest(); DECORATE.onRequest(span, new HostAndRequestAsHttpUriRequest(delegate.getTarget(), request)); - DECORATE.injectContext(current().with(span), request, SETTER); + if (injectContext) { + Context receiver = current(); + if (span != null) { + receiver = receiver.with(span); + } + DECORATE.injectContext(receiver, request, SETTER); + } return request; } diff --git a/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-4.0/src/main/java/datadog/trace/instrumentation/apachehttpclient/ApacheHttpClientInstrumentation.java b/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-4.0/src/main/java/datadog/trace/instrumentation/apachehttpclient/ApacheHttpClientInstrumentation.java index 9cde4c7a2df..4f70c43a034 100644 --- a/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-4.0/src/main/java/datadog/trace/instrumentation/apachehttpclient/ApacheHttpClientInstrumentation.java +++ b/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-4.0/src/main/java/datadog/trace/instrumentation/apachehttpclient/ApacheHttpClientInstrumentation.java @@ -1,5 +1,6 @@ package datadog.trace.instrumentation.apachehttpclient; +import static datadog.trace.agent.tooling.InstrumenterModule.TargetSystem.CONTEXT_TRACKING; import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static net.bytebuddy.matcher.ElementMatchers.isMethod; @@ -10,6 +11,7 @@ import datadog.appsec.api.blocking.BlockingException; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.agent.tooling.annotation.AppliesOn; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import net.bytebuddy.asm.Advice; import net.bytebuddy.description.type.TypeDescription; @@ -82,65 +84,72 @@ public void methodAdvice(MethodTransformer transformer) { // some methods can share the same advice class. The call depth tracking ensures only 1 span is // created - transformer.applyAdvice( + transformer.applyAdvices( isMethod() .and(named("execute")) .and(takesArguments(1)) .and(takesArgument(0, named("org.apache.http.client.methods.HttpUriRequest"))), - ApacheHttpClientInstrumentation.class.getName() + "$UriRequestAdvice"); + ApacheHttpClientInstrumentation.class.getName() + "$UriRequestAdvice", + ApacheHttpClientInstrumentation.class.getName() + "$UriRequestContextPropagationAdvice"); - transformer.applyAdvice( + transformer.applyAdvices( isMethod() .and(named("execute")) .and(takesArguments(2)) .and(takesArgument(0, named("org.apache.http.client.methods.HttpUriRequest"))) .and(takesArgument(1, named("org.apache.http.protocol.HttpContext"))), - ApacheHttpClientInstrumentation.class.getName() + "$UriRequestAdvice"); + ApacheHttpClientInstrumentation.class.getName() + "$UriRequestAdvice", + ApacheHttpClientInstrumentation.class.getName() + "$UriRequestContextPropagationAdvice"); - transformer.applyAdvice( + transformer.applyAdvices( isMethod() .and(named("execute")) .and(takesArguments(2)) .and(takesArgument(0, named("org.apache.http.client.methods.HttpUriRequest"))) .and(takesArgument(1, named("org.apache.http.client.ResponseHandler"))), - ApacheHttpClientInstrumentation.class.getName() + "$UriRequestWithHandlerAdvice"); + ApacheHttpClientInstrumentation.class.getName() + "$UriRequestWithHandlerAdvice", + ApacheHttpClientInstrumentation.class.getName() + "$UriRequestContextPropagationAdvice"); - transformer.applyAdvice( + transformer.applyAdvices( isMethod() .and(named("execute")) .and(takesArguments(3)) .and(takesArgument(0, named("org.apache.http.client.methods.HttpUriRequest"))) .and(takesArgument(1, named("org.apache.http.client.ResponseHandler"))) .and(takesArgument(2, named("org.apache.http.protocol.HttpContext"))), - ApacheHttpClientInstrumentation.class.getName() + "$UriRequestWithHandlerAdvice"); + ApacheHttpClientInstrumentation.class.getName() + "$UriRequestWithHandlerAdvice", + ApacheHttpClientInstrumentation.class.getName() + "$UriRequestContextPropagationAdvice"); - transformer.applyAdvice( + transformer.applyAdvices( isMethod() .and(named("execute")) .and(takesArguments(2)) .and(takesArgument(0, named("org.apache.http.HttpHost"))) .and(takesArgument(1, named("org.apache.http.HttpRequest"))), - ApacheHttpClientInstrumentation.class.getName() + "$RequestAdvice"); + ApacheHttpClientInstrumentation.class.getName() + "$RequestAdvice", + ApacheHttpClientInstrumentation.class.getName() + "$RequestContextPropagationAdvice"); - transformer.applyAdvice( + transformer.applyAdvices( isMethod() .and(named("execute")) .and(takesArguments(3)) .and(takesArgument(0, named("org.apache.http.HttpHost"))) .and(takesArgument(1, named("org.apache.http.HttpRequest"))) .and(takesArgument(2, named("org.apache.http.protocol.HttpContext"))), - ApacheHttpClientInstrumentation.class.getName() + "$RequestAdvice"); + ApacheHttpClientInstrumentation.class.getName() + "$RequestAdvice", + ApacheHttpClientInstrumentation.class.getName() + "$RequestContextPropagationAdvice"); - transformer.applyAdvice( + transformer.applyAdvices( isMethod() .and(named("execute")) .and(takesArguments(3)) .and(takesArgument(0, named("org.apache.http.HttpHost"))) .and(takesArgument(1, named("org.apache.http.HttpRequest"))) .and(takesArgument(2, named("org.apache.http.client.ResponseHandler"))), - ApacheHttpClientInstrumentation.class.getName() + "$RequestWithHandlerAdvice"); + ApacheHttpClientInstrumentation.class.getName() + "$RequestWithHandlerAdvice", + ApacheHttpClientInstrumentation.class.getName() + "$RequestContextPropagationAdvice"); - transformer.applyAdvice( + transformer.applyAdvices( isMethod() .and(named("execute")) .and(takesArguments(4)) @@ -148,7 +157,8 @@ public void methodAdvice(MethodTransformer transformer) { .and(takesArgument(1, named("org.apache.http.HttpRequest"))) .and(takesArgument(2, named("org.apache.http.client.ResponseHandler"))) .and(takesArgument(3, named("org.apache.http.protocol.HttpContext"))), - ApacheHttpClientInstrumentation.class.getName() + "$RequestWithHandlerAdvice"); + ApacheHttpClientInstrumentation.class.getName() + "$RequestWithHandlerAdvice", + ApacheHttpClientInstrumentation.class.getName() + "$RequestContextPropagationAdvice"); } public static class UriRequestAdvice { @@ -274,4 +284,21 @@ public static void methodExit( HelperMethods.doMethodExit(scope, result, throwable); } } + + @AppliesOn(CONTEXT_TRACKING) + public static class UriRequestContextPropagationAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void methodEnter(@Advice.Argument(0) final HttpUriRequest request) { + HelperMethods.doInjectContext(request); + } + } + + @AppliesOn(CONTEXT_TRACKING) + public static class RequestContextPropagationAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void methodEnter( + @Advice.Argument(0) final HttpHost host, @Advice.Argument(1) final HttpRequest request) { + HelperMethods.doInjectContext(host, request); + } + } } diff --git a/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-4.0/src/main/java/datadog/trace/instrumentation/apachehttpclient/HelperMethods.java b/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-4.0/src/main/java/datadog/trace/instrumentation/apachehttpclient/HelperMethods.java index 389f58a50f5..758edb5f754 100644 --- a/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-4.0/src/main/java/datadog/trace/instrumentation/apachehttpclient/HelperMethods.java +++ b/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-4.0/src/main/java/datadog/trace/instrumentation/apachehttpclient/HelperMethods.java @@ -40,14 +40,25 @@ private static AgentScope activateHttpSpan(final HttpUriRequest request) { DECORATE.afterStart(span); DECORATE.onRequest(span, request); - final boolean awsClientCall = request.containsHeader("amz-sdk-invocation-id"); - // AWS calls are often signed, so we can't add headers without breaking the signature. - if (!awsClientCall) { - DECORATE.injectContext(current().with(span), request, SETTER); + return scope; + } + + public static void doInjectContext(final HttpUriRequest request) { + if (request.containsHeader("amz-sdk-invocation-id")) { + return; } + DECORATE.injectContext(current(), request, SETTER); + } - return scope; + public static void doInjectContext(final HttpHost host, final HttpRequest request) { + final HttpUriRequest uriRequest; + if (request instanceof HttpUriRequest) { + uriRequest = (HttpUriRequest) request; + } else { + uriRequest = new HostAndRequestAsHttpUriRequest(host, request); + } + doInjectContext(uriRequest); } public static void doMethodExit( diff --git a/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/ApacheHttpAsyncClientInstrumentation.java b/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/ApacheHttpAsyncClientInstrumentation.java index 6249819dee0..3bc8441f2f5 100644 --- a/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/ApacheHttpAsyncClientInstrumentation.java +++ b/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/ApacheHttpAsyncClientInstrumentation.java @@ -1,5 +1,6 @@ package datadog.trace.instrumentation.apachehttpclient5; +import static datadog.trace.agent.tooling.InstrumenterModule.TargetSystem.CONTEXT_TRACKING; import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; @@ -14,6 +15,7 @@ import com.google.auto.service.AutoService; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.agent.tooling.annotation.AppliesOn; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; @@ -76,7 +78,7 @@ public String[] helperClassNames() { @Override public void methodAdvice(MethodTransformer transformer) { - transformer.applyAdvice( + transformer.applyAdvices( isMethod() .and(named("execute")) .and(takesArguments(5)) @@ -85,9 +87,23 @@ public void methodAdvice(MethodTransformer transformer) { .and(takesArgument(2, named("org.apache.hc.core5.http.nio.HandlerFactory"))) .and(takesArgument(3, named("org.apache.hc.core5.http.protocol.HttpContext"))) .and(takesArgument(4, named("org.apache.hc.core5.concurrent.FutureCallback"))), + this.getClass().getName() + "$ClientContextPropagationAdvice", this.getClass().getName() + "$ClientAdvice"); } + @AppliesOn(CONTEXT_TRACKING) + @SuppressFBWarnings("UC_USELESS_OBJECT") + public static class ClientContextPropagationAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void onEnter( + @Advice.Argument(value = 0, readOnly = false) AsyncRequestProducer requestProducer) { + final DelegatingRequestProducer delegatingRequestProducer = + new DelegatingRequestProducer(requestProducer); + delegatingRequestProducer.setInjectContext(true); + requestProducer = delegatingRequestProducer; + } + } + public static class ClientAdvice { @SuppressFBWarnings("UC_USELESS_OBJECT") @Advice.OnMethodEnter(suppress = Throwable.class) @@ -105,7 +121,11 @@ public static AgentScope methodEnter( context = new BasicHttpContext(); } - requestProducer = new DelegatingRequestProducer(clientSpan, requestProducer); + if (!(requestProducer instanceof DelegatingRequestProducer)) { + requestProducer = new DelegatingRequestProducer(requestProducer); + } + + ((DelegatingRequestProducer) requestProducer).setSpan(clientSpan); futureCallback = new TraceContinuedFutureCallback<>( parentContinuation, clientSpan, context, futureCallback); diff --git a/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/ApacheHttpClientInstrumentation.java b/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/ApacheHttpClientInstrumentation.java index 051b5305771..d1e1d2752a7 100644 --- a/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/ApacheHttpClientInstrumentation.java +++ b/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/ApacheHttpClientInstrumentation.java @@ -1,5 +1,6 @@ package datadog.trace.instrumentation.apachehttpclient5; +import static datadog.trace.agent.tooling.InstrumenterModule.TargetSystem.CONTEXT_TRACKING; import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static net.bytebuddy.matcher.ElementMatchers.isMethod; @@ -10,6 +11,7 @@ import datadog.appsec.api.blocking.BlockingException; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.agent.tooling.annotation.AppliesOn; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import net.bytebuddy.asm.Advice; import net.bytebuddy.description.type.TypeDescription; @@ -70,39 +72,43 @@ public String[] helperClassNames() { @Override public void methodAdvice(MethodTransformer transformer) { - transformer.applyAdvice( + transformer.applyAdvices( isMethod() .and(named("execute")) .and(takesArguments(1)) .and(takesArgument(0, named("org.apache.hc.core5.http.ClassicHttpRequest"))), - ApacheHttpClientInstrumentation.class.getName() + "$RequestAdvice"); + ApacheHttpClientInstrumentation.class.getName() + "$RequestAdvice", + ApacheHttpClientInstrumentation.class.getName() + "$ContextPropagationAdviceArg0"); - transformer.applyAdvice( + transformer.applyAdvices( isMethod() .and(named("execute")) .and(takesArguments(2)) .and(takesArgument(0, named("org.apache.hc.core5.http.ClassicHttpRequest"))) .and(takesArgument(1, named("org.apache.hc.core5.http.protocol.HttpContext"))), - ApacheHttpClientInstrumentation.class.getName() + "$RequestAdvice"); + ApacheHttpClientInstrumentation.class.getName() + "$RequestAdvice", + ApacheHttpClientInstrumentation.class.getName() + "$ContextPropagationAdviceArg0"); - transformer.applyAdvice( + transformer.applyAdvices( isMethod() .and(named("execute")) .and(takesArguments(3)) .and(takesArgument(0, named("org.apache.hc.core5.http.HttpHost"))) .and(takesArgument(1, named("org.apache.hc.core5.http.ClassicHttpRequest"))) .and(takesArgument(2, named("org.apache.hc.core5.http.protocol.HttpContext"))), - ApacheHttpClientInstrumentation.class.getName() + "$HostRequestAdvice"); + ApacheHttpClientInstrumentation.class.getName() + "$HostRequestAdvice", + ApacheHttpClientInstrumentation.class.getName() + "$ContextPropagationAdviceArg1"); - transformer.applyAdvice( + transformer.applyAdvices( isMethod() .and(named("execute")) .and(takesArguments(2)) .and(takesArgument(0, named("org.apache.hc.core5.http.HttpHost"))) .and(takesArgument(1, named("org.apache.hc.core5.http.ClassicHttpRequest"))), - ApacheHttpClientInstrumentation.class.getName() + "$HostRequestAdvice"); + ApacheHttpClientInstrumentation.class.getName() + "$HostRequestAdvice", + ApacheHttpClientInstrumentation.class.getName() + "$ContextPropagationAdviceArg1"); - transformer.applyAdvice( + transformer.applyAdvices( isMethod() .and(named("execute")) .and(takesArguments(4)) @@ -110,7 +116,8 @@ public void methodAdvice(MethodTransformer transformer) { .and(takesArgument(1, named("org.apache.hc.core5.http.ClassicHttpRequest"))) .and(takesArgument(2, named("org.apache.hc.core5.http.protocol.HttpContext"))) .and(takesArgument(3, named("org.apache.hc.core5.http.io.HttpClientResponseHandler"))), - ApacheHttpClientInstrumentation.class.getName() + "$ResponseHandlerAdvice"); + ApacheHttpClientInstrumentation.class.getName() + "$ResponseHandlerAdvice", + ApacheHttpClientInstrumentation.class.getName() + "$ContextPropagationAdviceArg1"); } public static class RequestAdvice { @@ -194,4 +201,20 @@ public static void methodExit( HelperMethods.doMethodExit(scope, result, throwable); } } + + @AppliesOn(CONTEXT_TRACKING) + public static class ContextPropagationAdviceArg0 { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void methodEnter(@Advice.Argument(0) final ClassicHttpRequest request) { + HelperMethods.doInjectContext(request); + } + } + + @AppliesOn(CONTEXT_TRACKING) + public static class ContextPropagationAdviceArg1 { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void methodEnter(@Advice.Argument(1) final ClassicHttpRequest request) { + HelperMethods.doInjectContext(request); + } + } } diff --git a/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/DelegatingRequestChannel.java b/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/DelegatingRequestChannel.java index 2894796a149..367e3b45a84 100644 --- a/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/DelegatingRequestChannel.java +++ b/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/DelegatingRequestChannel.java @@ -15,17 +15,22 @@ public class DelegatingRequestChannel implements RequestChannel { private final RequestChannel delegate; private final AgentSpan span; + private final boolean injectContext; - public DelegatingRequestChannel(RequestChannel requestChannel, AgentSpan span) { + public DelegatingRequestChannel( + RequestChannel requestChannel, AgentSpan span, boolean injectContext) { this.delegate = requestChannel; this.span = span; + this.injectContext = injectContext; } @Override public void sendRequest(HttpRequest request, EntityDetails entityDetails, HttpContext context) throws HttpException, IOException { DECORATE.onRequest(span, request); - DECORATE.injectContext(current().with(span), request, SETTER); + if (injectContext) { + DECORATE.injectContext(current().with(span), request, SETTER); + } delegate.sendRequest(request, entityDetails, context); } } diff --git a/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/DelegatingRequestProducer.java b/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/DelegatingRequestProducer.java index da09ec9687b..50ecf1b729b 100644 --- a/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/DelegatingRequestProducer.java +++ b/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/DelegatingRequestProducer.java @@ -9,14 +9,22 @@ import org.apache.hc.core5.http.protocol.HttpContext; public class DelegatingRequestProducer implements AsyncRequestProducer { - final AgentSpan span; + AgentSpan span; final AsyncRequestProducer delegate; + boolean injectContext = false; - public DelegatingRequestProducer(final AgentSpan span, final AsyncRequestProducer delegate) { - this.span = span; + public DelegatingRequestProducer(final AsyncRequestProducer delegate) { this.delegate = delegate; } + public void setSpan(AgentSpan span) { + this.span = span; + } + + public void setInjectContext(boolean injectContext) { + this.injectContext = injectContext; + } + @Override public void failed(final Exception ex) { delegate.failed(ex); @@ -25,7 +33,8 @@ public void failed(final Exception ex) { @Override public void sendRequest(RequestChannel channel, HttpContext context) throws HttpException, IOException { - DelegatingRequestChannel requestChannel = new DelegatingRequestChannel(channel, span); + DelegatingRequestChannel requestChannel = + new DelegatingRequestChannel(channel, span, injectContext); delegate.sendRequest(requestChannel, context); } diff --git a/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/HelperMethods.java b/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/HelperMethods.java index e6187609433..e30233b65da 100644 --- a/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/HelperMethods.java +++ b/dd-java-agent/instrumentation/apache-httpclient/apache-httpclient-5.0/src/main/java/datadog/trace/instrumentation/apachehttpclient5/HelperMethods.java @@ -41,15 +41,16 @@ private static AgentScope activateHttpSpan(final HttpRequest request) { DECORATE.afterStart(span); DECORATE.onRequest(span, request); - final boolean awsClientCall = request.containsHeader("amz-sdk-invocation-id"); - // AWS calls are often signed, so we can't add headers without breaking the signature. - if (!awsClientCall) { - DECORATE.injectContext(current().with(span), request, SETTER); - } - return scope; } + public static void doInjectContext(final HttpRequest request) { + if (request.containsHeader("amz-sdk-invocation-id")) { + return; + } + DECORATE.injectContext(current(), request, SETTER); + } + public static void doMethodExit( final AgentScope scope, final Object result, final Throwable throwable) { if (scope == null) { diff --git a/dd-java-agent/instrumentation/armeria/armeria-grpc-0.84/src/main/java/datadog/trace/instrumentation/armeria/grpc/client/ClientCallImplInstrumentation.java b/dd-java-agent/instrumentation/armeria/armeria-grpc-0.84/src/main/java/datadog/trace/instrumentation/armeria/grpc/client/ClientCallImplInstrumentation.java index 3c503841b2c..5809e6c005c 100644 --- a/dd-java-agent/instrumentation/armeria/armeria-grpc-0.84/src/main/java/datadog/trace/instrumentation/armeria/grpc/client/ClientCallImplInstrumentation.java +++ b/dd-java-agent/instrumentation/armeria/armeria-grpc-0.84/src/main/java/datadog/trace/instrumentation/armeria/grpc/client/ClientCallImplInstrumentation.java @@ -1,6 +1,6 @@ package datadog.trace.instrumentation.armeria.grpc.client; -import static datadog.context.Context.current; +import static datadog.trace.agent.tooling.InstrumenterModule.TargetSystem.CONTEXT_TRACKING; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; @@ -15,10 +15,12 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArguments; import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.agent.tooling.annotation.AppliesOn; import datadog.trace.api.InstrumenterConfig; import datadog.trace.bootstrap.InstrumentationContext; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge; import io.grpc.ClientCall; import io.grpc.Metadata; import io.grpc.MethodDescriptor; @@ -43,7 +45,10 @@ public void methodAdvice(MethodTransformer transformer) { transformer.applyAdvice( isConstructor().and(takesArgument(2, named("io.grpc.MethodDescriptor"))), getClass().getName() + "$CaptureCallPos2"); - transformer.applyAdvice(named("start").and(isMethod()), getClass().getName() + "$Start"); + transformer.applyAdvices( + named("start").and(isMethod()), + getClass().getName() + "$Start", + getClass().getName() + "$StartContextPropagationAdvice"); transformer.applyAdvice(named("cancel").and(isMethod()), getClass().getName() + "$Cancel"); transformer.applyAdvice( named("request") @@ -95,7 +100,6 @@ public static AgentScope before( if (null != responseListener && null != headers) { span = InstrumentationContext.get(ClientCall.class, AgentSpan.class).get(call); if (null != span) { - DECORATE.injectContext(current().with(span), headers, SETTER); return activateSpan(span); } } @@ -120,6 +124,14 @@ public static void after( } } + @AppliesOn(CONTEXT_TRACKING) + public static final class StartContextPropagationAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void before(@Advice.Argument(1) Metadata headers) { + DECORATE.injectContext(Java8BytecodeBridge.getCurrentContext(), headers, SETTER); + } + } + public static final class ActivateSpan { @Advice.OnMethodEnter public static AgentScope before(@Advice.This ClientCall call) { diff --git a/dd-java-agent/instrumentation/armeria/armeria-grpc-0.84/src/main/java/datadog/trace/instrumentation/armeria/grpc/client/GrpcClientDecorator.java b/dd-java-agent/instrumentation/armeria/armeria-grpc-0.84/src/main/java/datadog/trace/instrumentation/armeria/grpc/client/GrpcClientDecorator.java index de52229988d..519b8dba427 100644 --- a/dd-java-agent/instrumentation/armeria/armeria-grpc-0.84/src/main/java/datadog/trace/instrumentation/armeria/grpc/client/GrpcClientDecorator.java +++ b/dd-java-agent/instrumentation/armeria/armeria-grpc-0.84/src/main/java/datadog/trace/instrumentation/armeria/grpc/client/GrpcClientDecorator.java @@ -15,6 +15,7 @@ import datadog.trace.api.datastreams.DataStreamsTags; import datadog.trace.api.naming.SpanNaming; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.AgentTracer; import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes; import datadog.trace.bootstrap.instrumentation.api.Tags; import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString; @@ -90,7 +91,7 @@ protected String service() { public AgentSpan startCall(MethodDescriptor method) { if (IGNORED_METHODS.contains(method.getFullMethodName())) { - return null; + return AgentTracer.blackholeSpan(); } AgentSpan span = startSpan("grpc", OPERATION_NAME) diff --git a/dd-java-agent/instrumentation/commons-httpclient-2.0/src/main/java/datadog/trace/instrumentation/commonshttpclient/CommonsHttpClientInstrumentation.java b/dd-java-agent/instrumentation/commons-httpclient-2.0/src/main/java/datadog/trace/instrumentation/commonshttpclient/CommonsHttpClientInstrumentation.java index 0ad7480901d..db4b5d4d37f 100644 --- a/dd-java-agent/instrumentation/commons-httpclient-2.0/src/main/java/datadog/trace/instrumentation/commonshttpclient/CommonsHttpClientInstrumentation.java +++ b/dd-java-agent/instrumentation/commons-httpclient-2.0/src/main/java/datadog/trace/instrumentation/commonshttpclient/CommonsHttpClientInstrumentation.java @@ -1,5 +1,6 @@ package datadog.trace.instrumentation.commonshttpclient; +import static datadog.trace.agent.tooling.InstrumenterModule.TargetSystem.CONTEXT_TRACKING; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; @@ -15,6 +16,7 @@ import datadog.appsec.api.blocking.BlockingException; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.agent.tooling.annotation.AppliesOn; import datadog.trace.bootstrap.CallDepthThreadLocalMap; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -44,12 +46,13 @@ public String[] helperClassNames() { @Override public void methodAdvice(MethodTransformer transformer) { - transformer.applyAdvice( + transformer.applyAdvices( isMethod() .and(named("executeMethod")) .and(takesArguments(3)) .and(takesArgument(1, named("org.apache.commons.httpclient.HttpMethod"))), - CommonsHttpClientInstrumentation.class.getName() + "$ExecAdvice"); + CommonsHttpClientInstrumentation.class.getName() + "$ExecAdvice", + CommonsHttpClientInstrumentation.class.getName() + "$ContextPropagationAdvice"); } public static class ExecAdvice { @@ -67,7 +70,6 @@ public static AgentScope methodEnter(@Advice.Argument(1) final HttpMethod httpMe DECORATE.afterStart(span); DECORATE.onRequest(span, httpMethod); - DECORATE.injectContext(getCurrentContext().with(span), httpMethod, SETTER); return scope; } catch (BlockingException e) { @@ -98,4 +100,12 @@ public static void methodExit( } } } + + @AppliesOn(CONTEXT_TRACKING) + public static class ContextPropagationAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void methodEnter(@Advice.Argument(1) final HttpMethod httpMethod) { + DECORATE.injectContext(getCurrentContext(), httpMethod, SETTER); + } + } } diff --git a/dd-java-agent/instrumentation/google-http-client-1.19/src/main/java/datadog/trace/instrumentation/googlehttpclient/GoogleHttpClientDecorator.java b/dd-java-agent/instrumentation/google-http-client-1.19/src/main/java/datadog/trace/instrumentation/googlehttpclient/GoogleHttpClientDecorator.java index 5847efc3f1b..6204dfda7ae 100644 --- a/dd-java-agent/instrumentation/google-http-client-1.19/src/main/java/datadog/trace/instrumentation/googlehttpclient/GoogleHttpClientDecorator.java +++ b/dd-java-agent/instrumentation/google-http-client-1.19/src/main/java/datadog/trace/instrumentation/googlehttpclient/GoogleHttpClientDecorator.java @@ -1,8 +1,5 @@ package datadog.trace.instrumentation.googlehttpclient; -import static datadog.context.Context.current; -import static datadog.trace.instrumentation.googlehttpclient.HeadersInjectAdapter.SETTER; - import com.google.api.client.http.HttpRequest; import com.google.api.client.http.HttpResponse; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -37,7 +34,6 @@ protected URI url(final HttpRequest httpRequest) throws URISyntaxException { public AgentSpan prepareSpan(AgentSpan span, HttpRequest request) { DECORATE.afterStart(span); DECORATE.onRequest(span, request); - DECORATE.injectContext(current().with(span), request, SETTER); return span; } diff --git a/dd-java-agent/instrumentation/google-http-client-1.19/src/main/java/datadog/trace/instrumentation/googlehttpclient/GoogleHttpClientInstrumentation.java b/dd-java-agent/instrumentation/google-http-client-1.19/src/main/java/datadog/trace/instrumentation/googlehttpclient/GoogleHttpClientInstrumentation.java index 870941a60d6..93edb10af51 100644 --- a/dd-java-agent/instrumentation/google-http-client-1.19/src/main/java/datadog/trace/instrumentation/googlehttpclient/GoogleHttpClientInstrumentation.java +++ b/dd-java-agent/instrumentation/google-http-client-1.19/src/main/java/datadog/trace/instrumentation/googlehttpclient/GoogleHttpClientInstrumentation.java @@ -1,11 +1,14 @@ package datadog.trace.instrumentation.googlehttpclient; +import static datadog.trace.agent.tooling.InstrumenterModule.TargetSystem.CONTEXT_TRACKING; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; import static datadog.trace.instrumentation.googlehttpclient.GoogleHttpClientDecorator.DECORATE; import static datadog.trace.instrumentation.googlehttpclient.GoogleHttpClientDecorator.HTTP_REQUEST; +import static datadog.trace.instrumentation.googlehttpclient.HeadersInjectAdapter.SETTER; import static net.bytebuddy.matcher.ElementMatchers.isMethod; import static net.bytebuddy.matcher.ElementMatchers.isPublic; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; @@ -16,6 +19,7 @@ import com.google.auto.service.AutoService; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.agent.tooling.annotation.AppliesOn; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import net.bytebuddy.asm.Advice; @@ -44,17 +48,21 @@ public String[] helperClassNames() { @Override public void methodAdvice(MethodTransformer transformer) { - transformer.applyAdvice( + transformer.applyAdvices( isMethod().and(isPublic()).and(named("execute")).and(takesArguments(0)), - GoogleHttpClientInstrumentation.class.getName() + "$GoogleHttpClientAdvice"); + GoogleHttpClientInstrumentation.class.getName() + "$GoogleHttpClientAdvice", + GoogleHttpClientInstrumentation.class.getName() + + "$GoogleHttpClientContextPropagationAdvice"); - transformer.applyAdvice( + transformer.applyAdvices( isMethod() .and(isPublic()) .and(named("executeAsync")) .and(takesArguments(1)) .and(takesArgument(0, (named("java.util.concurrent.Executor")))), - GoogleHttpClientInstrumentation.class.getName() + "$GoogleHttpClientAsyncAdvice"); + GoogleHttpClientInstrumentation.class.getName() + "$GoogleHttpClientAsyncAdvice", + GoogleHttpClientInstrumentation.class.getName() + + "$GoogleHttpClientContextPropagationAdvice"); } public static class GoogleHttpClientAdvice { @@ -118,4 +126,12 @@ public static void methodExit( } } } + + @AppliesOn(CONTEXT_TRACKING) + public static class GoogleHttpClientContextPropagationAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void methodEnter(@Advice.This HttpRequest request) { + DECORATE.injectContext(getCurrentContext(), request, SETTER); + } + } } diff --git a/dd-java-agent/instrumentation/grizzly/grizzly-client-1.9/src/main/java/datadog/trace/instrumentation/grizzly/client/AsyncHttpClientInstrumentation.java b/dd-java-agent/instrumentation/grizzly/grizzly-client-1.9/src/main/java/datadog/trace/instrumentation/grizzly/client/AsyncHttpClientInstrumentation.java index 92bdac88d09..a8b87193d1c 100644 --- a/dd-java-agent/instrumentation/grizzly/grizzly-client-1.9/src/main/java/datadog/trace/instrumentation/grizzly/client/AsyncHttpClientInstrumentation.java +++ b/dd-java-agent/instrumentation/grizzly/grizzly-client-1.9/src/main/java/datadog/trace/instrumentation/grizzly/client/AsyncHttpClientInstrumentation.java @@ -1,6 +1,8 @@ package datadog.trace.instrumentation.grizzly.client; +import static datadog.trace.agent.tooling.InstrumenterModule.TargetSystem.CONTEXT_TRACKING; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; @@ -13,6 +15,8 @@ import com.ning.http.client.AsyncHandler; import com.ning.http.client.Request; import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.agent.tooling.annotation.AppliesOn; +import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import net.bytebuddy.asm.Advice; @@ -26,26 +30,50 @@ public String instrumentedType() { @Override public void methodAdvice(MethodTransformer transformer) { - transformer.applyAdvice( + transformer.applyAdvices( named("executeRequest") .and(takesArgument(0, named("com.ning.http.client.Request"))) .and(takesArgument(1, named("com.ning.http.client.AsyncHandler"))) .and(isPublic()), - getClass().getName() + "$ExecuteRequest"); + getClass().getName() + "$ExecuteRequest", + getClass().getName() + "$ExecuteContextPropagationAdvice"); } public static class ExecuteRequest { @Advice.OnMethodEnter(suppress = Throwable.class) - public static void onEnter( + public static AgentScope onEnter( @Advice.Argument(0) final Request request, @Advice.Argument(value = 1, readOnly = false) AsyncHandler handler) { AgentSpan parentSpan = activeSpan(); AgentSpan span = startSpan(HTTP_REQUEST); DECORATE.afterStart(span); DECORATE.onRequest(span, request); - DECORATE.injectContext(getCurrentContext().with(span), request, SETTER); handler = new AsyncHandlerAdapter<>(span, parentSpan, handler); + return activateSpan(span); + } + + @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) + public static void onExit( + @Advice.Enter final AgentScope scope, @Advice.Thrown final Throwable throwable) { + if (scope == null) { + return; + } + if (throwable != null) { + AgentSpan span = scope.span(); + DECORATE.onError(span, throwable); + DECORATE.beforeFinish(span); + span.finish(); + } + scope.close(); + } + } + + @AppliesOn(CONTEXT_TRACKING) + public static class ExecuteContextPropagationAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void onEnter(@Advice.Argument(0) final Request request) { + DECORATE.injectContext(getCurrentContext(), request, SETTER); } } } diff --git a/dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/main/java/datadog/trace/instrumentation/httpclient/HttpHeadersInstrumentation.java b/dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/main/java/datadog/trace/instrumentation/httpclient/HttpHeadersInstrumentation.java index 82b802d70a8..1cd09452d12 100644 --- a/dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/main/java/datadog/trace/instrumentation/httpclient/HttpHeadersInstrumentation.java +++ b/dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/main/java/datadog/trace/instrumentation/httpclient/HttpHeadersInstrumentation.java @@ -54,6 +54,6 @@ public String[] helperClassNames() { @Override public void methodAdvice(MethodTransformer transformer) { - transformer.applyAdvice(isMethod().and(named("headers")), packageName + ".HeadersAdvice"); + transformer.applyAdvices(isMethod().and(named("headers")), packageName + ".HeadersAdvice"); } } diff --git a/dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/main/java11/datadog/trace/instrumentation/httpclient/HeadersAdvice.java b/dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/main/java11/datadog/trace/instrumentation/httpclient/HeadersAdvice.java index 493eec8bf3e..0657b3f44f5 100644 --- a/dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/main/java11/datadog/trace/instrumentation/httpclient/HeadersAdvice.java +++ b/dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/main/java11/datadog/trace/instrumentation/httpclient/HeadersAdvice.java @@ -1,17 +1,20 @@ package datadog.trace.instrumentation.httpclient; +import static datadog.trace.agent.tooling.InstrumenterModule.TargetSystem.CONTEXT_TRACKING; import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; import static datadog.trace.instrumentation.httpclient.HttpHeadersInjectAdapter.KEEP; import static datadog.trace.instrumentation.httpclient.HttpHeadersInjectAdapter.SETTER; import static datadog.trace.instrumentation.httpclient.JavaNetClientDecorator.DECORATE; import static java.lang.String.CASE_INSENSITIVE_ORDER; +import datadog.trace.agent.tooling.annotation.AppliesOn; import java.net.http.HttpHeaders; import java.util.List; import java.util.Map; import java.util.TreeMap; import net.bytebuddy.asm.Advice; +@AppliesOn(CONTEXT_TRACKING) public class HeadersAdvice { @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void methodExit(@Advice.Return(readOnly = false) HttpHeaders headers) { diff --git a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-10.0/src/main/java/datadog/trace/instrumentation/jetty_client10/JettyClientInstrumentation.java b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-10.0/src/main/java/datadog/trace/instrumentation/jetty_client10/JettyClientInstrumentation.java index 1ddf6eb654a..24a8c15cf53 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-10.0/src/main/java/datadog/trace/instrumentation/jetty_client10/JettyClientInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-10.0/src/main/java/datadog/trace/instrumentation/jetty_client10/JettyClientInstrumentation.java @@ -47,7 +47,7 @@ public Map contextStore() { @Override public void methodAdvice(MethodTransformer transformer) { - transformer.applyAdvice( + transformer.applyAdvices( isMethod() .and(named("send")) .and( @@ -57,7 +57,8 @@ public void methodAdvice(MethodTransformer transformer) { "org.eclipse.jetty.client.api.Request", "org.eclipse.jetty.client.HttpRequest"))) .and(takesArgument(1, List.class)), - packageName + ".SendAdvice"); + packageName + ".SendAdvice", + packageName + ".SendContextPropagationAdvice"); } @Override diff --git a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-10.0/src/main/java11/datadog/trace/instrumentation/jetty_client10/SendAdvice.java b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-10.0/src/main/java11/datadog/trace/instrumentation/jetty_client10/SendAdvice.java index 95159ca3033..2e92e35f91b 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-10.0/src/main/java11/datadog/trace/instrumentation/jetty_client10/SendAdvice.java +++ b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-10.0/src/main/java11/datadog/trace/instrumentation/jetty_client10/SendAdvice.java @@ -1,8 +1,6 @@ package datadog.trace.instrumentation.jetty_client10; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; -import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; -import static datadog.trace.instrumentation.jetty_client.HeadersInjectAdapter.SETTER; import static datadog.trace.instrumentation.jetty_client10.JettyClientDecorator.DECORATE; import static datadog.trace.instrumentation.jetty_client10.JettyClientDecorator.HTTP_REQUEST; @@ -24,7 +22,6 @@ public static AgentSpan methodEnter( responseListeners.add(0, new SpanFinishingCompleteListener(span)); DECORATE.afterStart(span); DECORATE.onRequest(span, request); - DECORATE.injectContext(getCurrentContext().with(span), request, SETTER); return span; } diff --git a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-10.0/src/main/java11/datadog/trace/instrumentation/jetty_client10/SendContextPropagationAdvice.java b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-10.0/src/main/java11/datadog/trace/instrumentation/jetty_client10/SendContextPropagationAdvice.java new file mode 100644 index 00000000000..8a0f4846f1e --- /dev/null +++ b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-10.0/src/main/java11/datadog/trace/instrumentation/jetty_client10/SendContextPropagationAdvice.java @@ -0,0 +1,26 @@ +package datadog.trace.instrumentation.jetty_client10; + +import static datadog.trace.agent.tooling.InstrumenterModule.TargetSystem.CONTEXT_TRACKING; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; +import static datadog.trace.instrumentation.jetty_client.HeadersInjectAdapter.SETTER; +import static datadog.trace.instrumentation.jetty_client10.JettyClientDecorator.DECORATE; + +import datadog.context.Context; +import datadog.trace.agent.tooling.annotation.AppliesOn; +import datadog.trace.bootstrap.InstrumentationContext; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import net.bytebuddy.asm.Advice; +import org.eclipse.jetty.client.api.Request; + +@AppliesOn(CONTEXT_TRACKING) +public class SendContextPropagationAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void methodEnter(@Advice.Argument(0) final Request request) { + final AgentSpan span = InstrumentationContext.get(Request.class, AgentSpan.class).get(request); + Context destination = getCurrentContext(); + if (span != null) { + destination = destination.with(span); + } + DECORATE.injectContext(destination, request, SETTER); + } +} diff --git a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java/datadog/trace/instrumentation/jetty_client12/JettyHttpClientInstrumentation.java b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java/datadog/trace/instrumentation/jetty_client12/JettyHttpClientInstrumentation.java index d8d78fb9c01..82c1f39967b 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java/datadog/trace/instrumentation/jetty_client12/JettyHttpClientInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java/datadog/trace/instrumentation/jetty_client12/JettyHttpClientInstrumentation.java @@ -50,11 +50,12 @@ public Map contextStore() { @Override public void methodAdvice(MethodTransformer transformer) { transformer.applyAdvice(isConstructor(), packageName + ".RequestCreateAdvice"); - transformer.applyAdvice( + transformer.applyAdvices( isMethod() .and(named("send")) .and(takesArgument(0, named("org.eclipse.jetty.client.Response$CompleteListener"))), - packageName + ".SendAdvice"); + packageName + ".SendAdvice", + packageName + ".SendContextPropagationAdvice"); transformer.applyAdvice( isMethod() .and(isPublic()) diff --git a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java17/datadog/trace/instrumentation/jetty_client12/SendAdvice.java b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java17/datadog/trace/instrumentation/jetty_client12/SendAdvice.java index e1682a4020e..cc08b865280 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java17/datadog/trace/instrumentation/jetty_client12/SendAdvice.java +++ b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java17/datadog/trace/instrumentation/jetty_client12/SendAdvice.java @@ -2,8 +2,6 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; -import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; -import static datadog.trace.instrumentation.jetty_client12.HeadersInjectAdapter.SETTER; import static datadog.trace.instrumentation.jetty_client12.JettyClientDecorator.DECORATE; import static datadog.trace.instrumentation.jetty_client12.JettyClientDecorator.HTTP_REQUEST; @@ -21,7 +19,6 @@ public static AgentScope methodEnter(@Advice.This final HttpRequest request) { InstrumentationContext.get(Request.class, AgentSpan.class).put(request, span); DECORATE.afterStart(span); DECORATE.onRequest(span, request); - DECORATE.injectContext(getCurrentContext().with(span), request, SETTER); return activateSpan(span); } diff --git a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java17/datadog/trace/instrumentation/jetty_client12/SendContextPropagationAdvice.java b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java17/datadog/trace/instrumentation/jetty_client12/SendContextPropagationAdvice.java new file mode 100644 index 00000000000..490763613f8 --- /dev/null +++ b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-12.0/src/main/java17/datadog/trace/instrumentation/jetty_client12/SendContextPropagationAdvice.java @@ -0,0 +1,25 @@ +package datadog.trace.instrumentation.jetty_client12; + +import static datadog.trace.agent.tooling.InstrumenterModule.TargetSystem.CONTEXT_TRACKING; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; +import static datadog.trace.instrumentation.jetty_client12.HeadersInjectAdapter.SETTER; +import static datadog.trace.instrumentation.jetty_client12.JettyClientDecorator.DECORATE; + +import datadog.trace.agent.tooling.annotation.AppliesOn; +import datadog.trace.bootstrap.InstrumentationContext; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import net.bytebuddy.asm.Advice; +import org.eclipse.jetty.client.Request; +import org.eclipse.jetty.client.transport.HttpRequest; + +@AppliesOn(CONTEXT_TRACKING) +public class SendContextPropagationAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void methodEnter(@Advice.This final HttpRequest request) { + AgentSpan span = InstrumentationContext.get(Request.class, AgentSpan.class).get(request); + if (span == null) { + return; + } + DECORATE.injectContext(getCurrentContext().with(span), request, SETTER); + } +} diff --git a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-9.1/src/main/java/datadog/trace/instrumentation/jetty_client91/JettyClientInstrumentation.java b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-9.1/src/main/java/datadog/trace/instrumentation/jetty_client91/JettyClientInstrumentation.java index 041c6d008b6..ccdfad18391 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-9.1/src/main/java/datadog/trace/instrumentation/jetty_client91/JettyClientInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty/jetty-client/jetty-client-9.1/src/main/java/datadog/trace/instrumentation/jetty_client91/JettyClientInstrumentation.java @@ -1,5 +1,6 @@ package datadog.trace.instrumentation.jetty_client91; +import static datadog.trace.agent.tooling.InstrumenterModule.TargetSystem.CONTEXT_TRACKING; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.namedOneOf; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; @@ -14,9 +15,11 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.ExcludeFilterProvider; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.agent.tooling.annotation.AppliesOn; import datadog.trace.bootstrap.InstrumentationContext; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter; @@ -61,7 +64,7 @@ public Map contextStore() { @Override public void methodAdvice(MethodTransformer transformer) { - transformer.applyAdvice( + transformer.applyAdvices( isMethod() .and(named("send")) .and( @@ -71,7 +74,8 @@ public void methodAdvice(MethodTransformer transformer) { "org.eclipse.jetty.client.api.Request", "org.eclipse.jetty.client.HttpRequest"))) .and(takesArgument(1, List.class)), - JettyClientInstrumentation.class.getName() + "$SendAdvice"); + JettyClientInstrumentation.class.getName() + "$SendAdvice", + JettyClientInstrumentation.class.getName() + "$ContextPropagationAdvice"); } @Override @@ -90,7 +94,6 @@ public static AgentSpan methodEnter( responseListeners.add(0, new SpanFinishingCompleteListener(span)); DECORATE.afterStart(span); DECORATE.onRequest(span, request); - DECORATE.injectContext(getCurrentContext().with(span), request, SETTER); return span; } @@ -104,4 +107,18 @@ public static void methodExit( } } } + + @AppliesOn(CONTEXT_TRACKING) + public static class ContextPropagationAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void methodEnter(@Advice.Argument(0) final Request request) { + final AgentSpan span = + InstrumentationContext.get(Request.class, AgentSpan.class).get(request); + Context destination = getCurrentContext(); + if (span != null) { + destination = destination.with(span); + } + DECORATE.injectContext(destination, request, SETTER); + } + } } diff --git a/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/PekkoHttpSingleRequestInstrumentation.java b/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/PekkoHttpSingleRequestInstrumentation.java index e230a604172..e286100db94 100644 --- a/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/PekkoHttpSingleRequestInstrumentation.java +++ b/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/PekkoHttpSingleRequestInstrumentation.java @@ -1,5 +1,6 @@ package datadog.trace.instrumentation.pekkohttp; +import static datadog.trace.agent.tooling.InstrumenterModule.TargetSystem.CONTEXT_TRACKING; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; @@ -13,6 +14,7 @@ import com.google.auto.service.AutoService; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.agent.tooling.annotation.AppliesOn; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import net.bytebuddy.asm.Advice; @@ -47,15 +49,19 @@ public String[] helperClassNames() { @Override public void methodAdvice(MethodTransformer transformer) { // This is mainly for compatibility with 10.0 - transformer.applyAdvice( + transformer.applyAdvices( named("singleRequest") .and(takesArgument(0, named("org.apache.pekko.http.scaladsl.model.HttpRequest"))), - PekkoHttpSingleRequestInstrumentation.class.getName() + "$SingleRequestAdvice"); + PekkoHttpSingleRequestInstrumentation.class.getName() + "$SingleRequestAdvice", + PekkoHttpSingleRequestInstrumentation.class.getName() + + "$SingleRequestContextPropagationAdvice"); // This is for 10.1+ - transformer.applyAdvice( + transformer.applyAdvices( named("singleRequestImpl") .and(takesArgument(0, named("org.apache.pekko.http.scaladsl.model.HttpRequest"))), - PekkoHttpSingleRequestInstrumentation.class.getName() + "$SingleRequestAdvice"); + PekkoHttpSingleRequestInstrumentation.class.getName() + "$SingleRequestAdvice", + PekkoHttpSingleRequestInstrumentation.class.getName() + + "$SingleRequestContextPropagationAdvice"); } public static class SingleRequestAdvice { @@ -76,12 +82,6 @@ public static AgentScope methodEnter( final AgentSpan span = startSpan("pekko-http", PEKKO_CLIENT_REQUEST); DECORATE.afterStart(span); DECORATE.onRequest(span, request); - - if (request != null) { - DECORATE.injectContext(getCurrentContext().with(span), request, headers); - // Request is immutable, so we have to assign new value once we update headers - request = headers.getRequest(); - } return activateSpan(span); } @@ -107,4 +107,15 @@ public static void methodExit( scope.close(); } } + + @AppliesOn(CONTEXT_TRACKING) + public static class SingleRequestContextPropagationAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void methodEnter( + @Advice.Argument(value = 0, readOnly = false) HttpRequest request) { + final PekkoHttpHeaders headers = new PekkoHttpHeaders(request); + DECORATE.injectContext(getCurrentContext(), request, headers); + request = headers.getRequest(); + } + } } diff --git a/dd-java-agent/instrumentation/play-ws/play-ws-1.0/src/main/java/datadog/trace/instrumentation/playws1/PlayWSClientInstrumentation.java b/dd-java-agent/instrumentation/play-ws/play-ws-1.0/src/main/java/datadog/trace/instrumentation/playws1/PlayWSClientInstrumentation.java index 82b8fa40799..154ff59edb0 100644 --- a/dd-java-agent/instrumentation/play-ws/play-ws-1.0/src/main/java/datadog/trace/instrumentation/playws1/PlayWSClientInstrumentation.java +++ b/dd-java-agent/instrumentation/play-ws/play-ws-1.0/src/main/java/datadog/trace/instrumentation/playws1/PlayWSClientInstrumentation.java @@ -1,14 +1,14 @@ package datadog.trace.instrumentation.playws1; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; -import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; -import static datadog.trace.instrumentation.playws.HeadersInjectAdapter.SETTER; import static datadog.trace.instrumentation.playws.PlayWSClientDecorator.DECORATE; import static datadog.trace.instrumentation.playws.PlayWSClientDecorator.PLAY_WS_REQUEST; import com.google.auto.service.AutoService; +import datadog.context.ContextScope; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge; import datadog.trace.instrumentation.playws.BasePlayWSClientInstrumentation; import net.bytebuddy.asm.Advice; import play.shaded.ahc.org.asynchttpclient.AsyncHandler; @@ -20,7 +20,7 @@ public class PlayWSClientInstrumentation extends BasePlayWSClientInstrumentation { public static class ClientAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static AgentSpan methodEnter( + public static ContextScope methodEnter( @Advice.Argument(0) final Request request, @Advice.Argument(value = 1, readOnly = false) AsyncHandler asyncHandler) { @@ -28,7 +28,6 @@ public static AgentSpan methodEnter( DECORATE.afterStart(span); DECORATE.onRequest(span, request); - DECORATE.injectContext(getCurrentContext().with(span), request, SETTER); if (asyncHandler instanceof StreamedAsyncHandler) { asyncHandler = new StreamedAsyncHandlerWrapper((StreamedAsyncHandler) asyncHandler, span); @@ -37,18 +36,22 @@ public static AgentSpan methodEnter( asyncHandler = new AsyncHandlerWrapper(asyncHandler, span); } - return span; + return Java8BytecodeBridge.getCurrentContext().with(span).attach(); } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void methodExit( - @Advice.Enter final AgentSpan clientSpan, @Advice.Thrown final Throwable throwable) { - + @Advice.Enter final ContextScope scope, @Advice.Thrown final Throwable throwable) { + if (scope == null) { + return; + } if (throwable != null) { - DECORATE.onError(clientSpan, throwable); - DECORATE.beforeFinish(clientSpan); - clientSpan.finish(); + final AgentSpan span = Java8BytecodeBridge.spanFromContext(scope.context()); + DECORATE.onError(span, throwable); + DECORATE.beforeFinish(span); + span.finish(); } + scope.close(); } } } diff --git a/dd-java-agent/instrumentation/play-ws/play-ws-2.0/src/main/java/datadog/trace/instrumentation/playws2/PlayWSClientInstrumentation.java b/dd-java-agent/instrumentation/play-ws/play-ws-2.0/src/main/java/datadog/trace/instrumentation/playws2/PlayWSClientInstrumentation.java index f48ba53e31b..08ea1c509f2 100644 --- a/dd-java-agent/instrumentation/play-ws/play-ws-2.0/src/main/java/datadog/trace/instrumentation/playws2/PlayWSClientInstrumentation.java +++ b/dd-java-agent/instrumentation/play-ws/play-ws-2.0/src/main/java/datadog/trace/instrumentation/playws2/PlayWSClientInstrumentation.java @@ -1,14 +1,14 @@ package datadog.trace.instrumentation.playws2; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; -import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; -import static datadog.trace.instrumentation.playws.HeadersInjectAdapter.SETTER; import static datadog.trace.instrumentation.playws.PlayWSClientDecorator.DECORATE; import static datadog.trace.instrumentation.playws.PlayWSClientDecorator.PLAY_WS_REQUEST; import com.google.auto.service.AutoService; +import datadog.context.ContextScope; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge; import datadog.trace.instrumentation.playws.BasePlayWSClientInstrumentation; import net.bytebuddy.asm.Advice; import play.shaded.ahc.org.asynchttpclient.AsyncHandler; @@ -20,7 +20,7 @@ public class PlayWSClientInstrumentation extends BasePlayWSClientInstrumentation { public static class ClientAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static AgentSpan methodEnter( + public static ContextScope methodEnter( @Advice.Argument(0) final Request request, @Advice.Argument(value = 1, readOnly = false) AsyncHandler asyncHandler) { @@ -28,7 +28,6 @@ public static AgentSpan methodEnter( DECORATE.afterStart(span); DECORATE.onRequest(span, request); - DECORATE.injectContext(getCurrentContext().with(span), request, SETTER); if (asyncHandler instanceof StreamedAsyncHandler) { asyncHandler = new StreamedAsyncHandlerWrapper((StreamedAsyncHandler) asyncHandler, span); @@ -37,18 +36,22 @@ public static AgentSpan methodEnter( asyncHandler = new AsyncHandlerWrapper(asyncHandler, span); } - return span; + return Java8BytecodeBridge.getCurrentContext().with(span).attach(); } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void methodExit( - @Advice.Enter final AgentSpan clientSpan, @Advice.Thrown final Throwable throwable) { - + @Advice.Enter final ContextScope scope, @Advice.Thrown final Throwable throwable) { + if (scope == null) { + return; + } if (throwable != null) { - DECORATE.onError(clientSpan, throwable); - DECORATE.beforeFinish(clientSpan); - clientSpan.finish(); + final AgentSpan span = Java8BytecodeBridge.spanFromContext(scope.context()); + DECORATE.onError(span, throwable); + DECORATE.beforeFinish(span); + span.finish(); } + scope.close(); } } } diff --git a/dd-java-agent/instrumentation/play-ws/play-ws-2.1/src/main/java/datadog/trace/instrumentation/playws21/PlayWSClientInstrumentation.java b/dd-java-agent/instrumentation/play-ws/play-ws-2.1/src/main/java/datadog/trace/instrumentation/playws21/PlayWSClientInstrumentation.java index d7f59046e06..a28952016ff 100644 --- a/dd-java-agent/instrumentation/play-ws/play-ws-2.1/src/main/java/datadog/trace/instrumentation/playws21/PlayWSClientInstrumentation.java +++ b/dd-java-agent/instrumentation/play-ws/play-ws-2.1/src/main/java/datadog/trace/instrumentation/playws21/PlayWSClientInstrumentation.java @@ -1,14 +1,14 @@ package datadog.trace.instrumentation.playws21; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; -import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; -import static datadog.trace.instrumentation.playws.HeadersInjectAdapter.SETTER; import static datadog.trace.instrumentation.playws.PlayWSClientDecorator.DECORATE; import static datadog.trace.instrumentation.playws.PlayWSClientDecorator.PLAY_WS_REQUEST; import com.google.auto.service.AutoService; +import datadog.context.ContextScope; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge; import datadog.trace.instrumentation.playws.BasePlayWSClientInstrumentation; import net.bytebuddy.asm.Advice; import play.shaded.ahc.org.asynchttpclient.AsyncHandler; @@ -20,7 +20,7 @@ public class PlayWSClientInstrumentation extends BasePlayWSClientInstrumentation { public static class ClientAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static AgentSpan methodEnter( + public static ContextScope methodEnter( @Advice.Argument(0) final Request request, @Advice.Argument(value = 1, readOnly = false) AsyncHandler asyncHandler) { @@ -28,7 +28,6 @@ public static AgentSpan methodEnter( DECORATE.afterStart(span); DECORATE.onRequest(span, request); - DECORATE.injectContext(getCurrentContext().with(span), request, SETTER); if (asyncHandler instanceof StreamedAsyncHandler) { asyncHandler = new StreamedAsyncHandlerWrapper((StreamedAsyncHandler) asyncHandler, span); @@ -37,18 +36,22 @@ public static AgentSpan methodEnter( asyncHandler = new AsyncHandlerWrapper(asyncHandler, span); } - return span; + return Java8BytecodeBridge.getCurrentContext().with(span).attach(); } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void methodExit( - @Advice.Enter final AgentSpan clientSpan, @Advice.Thrown final Throwable throwable) { - + @Advice.Enter final ContextScope scope, @Advice.Thrown final Throwable throwable) { + if (scope == null) { + return; + } if (throwable != null) { - DECORATE.onError(clientSpan, throwable); - DECORATE.beforeFinish(clientSpan); - clientSpan.finish(); + final AgentSpan span = Java8BytecodeBridge.spanFromContext(scope.context()); + DECORATE.onError(span, throwable); + DECORATE.beforeFinish(span); + span.finish(); } + scope.close(); } } } diff --git a/dd-java-agent/instrumentation/play-ws/play-ws-common/src/main/java/datadog/trace/instrumentation/playws/BasePlayWSClientInstrumentation.java b/dd-java-agent/instrumentation/play-ws/play-ws-common/src/main/java/datadog/trace/instrumentation/playws/BasePlayWSClientInstrumentation.java index 29770b841c3..fcf90c9088d 100644 --- a/dd-java-agent/instrumentation/play-ws/play-ws-common/src/main/java/datadog/trace/instrumentation/playws/BasePlayWSClientInstrumentation.java +++ b/dd-java-agent/instrumentation/play-ws/play-ws-common/src/main/java/datadog/trace/instrumentation/playws/BasePlayWSClientInstrumentation.java @@ -1,8 +1,12 @@ package datadog.trace.instrumentation.playws; +import static datadog.trace.agent.tooling.InstrumenterModule.TargetSystem.CONTEXT_TRACKING; import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.nameStartsWith; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; +import static datadog.trace.instrumentation.playws.HeadersInjectAdapter.SETTER; +import static datadog.trace.instrumentation.playws.PlayWSClientDecorator.DECORATE; import static net.bytebuddy.matcher.ElementMatchers.isMethod; import static net.bytebuddy.matcher.ElementMatchers.not; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; @@ -10,8 +14,11 @@ import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.agent.tooling.annotation.AppliesOn; +import net.bytebuddy.asm.Advice; import net.bytebuddy.description.type.TypeDescription; import net.bytebuddy.matcher.ElementMatcher; +import play.shaded.ahc.org.asynchttpclient.Request; public abstract class BasePlayWSClientInstrumentation extends InstrumenterModule.Tracing implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice { @@ -36,13 +43,14 @@ public ElementMatcher hierarchyMatcher() { @Override public void methodAdvice(MethodTransformer transformer) { - transformer.applyAdvice( + transformer.applyAdvices( isMethod() .and(named("execute")) .and(takesArguments(2)) .and(takesArgument(0, named("play.shaded.ahc.org.asynchttpclient.Request"))) .and(takesArgument(1, named("play.shaded.ahc.org.asynchttpclient.AsyncHandler"))), - getClass().getName() + "$ClientAdvice"); + getClass().getName() + "$ClientAdvice", + BasePlayWSClientInstrumentation.class.getName() + "$ClientContextPropagationAdvice"); } @Override @@ -54,4 +62,12 @@ public String[] helperClassNames() { packageName + ".StreamedAsyncHandlerWrapper" }; } + + @AppliesOn(CONTEXT_TRACKING) + public static class ClientContextPropagationAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void methodEnter(@Advice.Argument(0) final Request request) { + DECORATE.injectContext(getCurrentContext(), request, SETTER); + } + } } diff --git a/dd-java-agent/instrumentation/rs/jax-rs/jax-rs-client/jax-rs-client-1.1/src/main/java/datadog/trace/instrumentation/jaxrs/v1/JaxRsClientV1Instrumentation.java b/dd-java-agent/instrumentation/rs/jax-rs/jax-rs-client/jax-rs-client-1.1/src/main/java/datadog/trace/instrumentation/jaxrs/v1/JaxRsClientV1Instrumentation.java index 63f1e9a1de5..5db42c50be8 100644 --- a/dd-java-agent/instrumentation/rs/jax-rs/jax-rs-client/jax-rs-client-1.1/src/main/java/datadog/trace/instrumentation/jaxrs/v1/JaxRsClientV1Instrumentation.java +++ b/dd-java-agent/instrumentation/rs/jax-rs/jax-rs-client/jax-rs-client-1.1/src/main/java/datadog/trace/instrumentation/jaxrs/v1/JaxRsClientV1Instrumentation.java @@ -1,5 +1,6 @@ package datadog.trace.instrumentation.jaxrs.v1; +import static datadog.trace.agent.tooling.InstrumenterModule.TargetSystem.CONTEXT_TRACKING; import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.extendsClass; import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; @@ -19,6 +20,7 @@ import com.sun.jersey.api.client.ClientResponse; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.agent.tooling.annotation.AppliesOn; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import net.bytebuddy.asm.Advice; @@ -52,11 +54,12 @@ public String[] helperClassNames() { @Override public void methodAdvice(MethodTransformer transformer) { - transformer.applyAdvice( + transformer.applyAdvices( named("handle") .and(takesArgument(0, extendsClass(named("com.sun.jersey.api.client.ClientRequest")))) .and(returns(extendsClass(named("com.sun.jersey.api.client.ClientResponse")))), - JaxRsClientV1Instrumentation.class.getName() + "$HandleAdvice"); + JaxRsClientV1Instrumentation.class.getName() + "$HandleAdvice", + JaxRsClientV1Instrumentation.class.getName() + "$HandleContextPropagationAdvice"); } public static class HandleAdvice { @@ -73,7 +76,6 @@ public static AgentScope onEnter( DECORATE.afterStart(span); DECORATE.onRequest(span, request); request.getProperties().put(DD_CONTEXT_ATTRIBUTE, span); - DECORATE.injectContext(getCurrentContext().with(span), request.getHeaders(), SETTER); return activateSpan(span); } return null; @@ -95,4 +97,12 @@ public static void onExit( scope.span().finish(); } } + + @AppliesOn(CONTEXT_TRACKING) + public static class HandleContextPropagationAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void onEnter(@Advice.Argument(0) final ClientRequest request) { + DECORATE.injectContext(getCurrentContext(), request.getHeaders(), SETTER); + } + } }