migrate httpclient5

Author: amarziali

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,6 +13,7 @@
 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;
@@ -61,15 +63,16 @@ public ElementMatcher<TypeDescription> hierarchyMatcher() {
 
   @Override
   public void methodAdvice(MethodTransformer transformer) {
-    transformer.applyAdvice(
+    transformer.applyAdvices(
         isMethod()
             .and(named("execute"))
             .and(takesArguments(4))
             .and(takesArgument(0, named("org.apache.http.nio.protocol.HttpAsyncRequestProducer")))
             .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() + "$ClientAdvice");
+        ApacheHttpAsyncClientInstrumentation.class.getName() + "$ClientAdvice",
+        ApacheHttpAsyncClientInstrumentation.class.getName() + "$ClientContextPropagationAdvice");
   }
 
   public static class ClientAdvice {
@@ -104,4 +107,14 @@ public static void methodExit(
       }
     }
   }
+
+  @AppliesOn(CONTEXT_TRACKING)
+  public static class ClientContextPropagationAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void onEnter(@Advice.Argument(0) final HttpAsyncRequestProducer requestProducer) {
+      if (requestProducer instanceof DelegatingRequestProducer) {
+        ((DelegatingRequestProducer) requestProducer).setInjectContext(true);
+      }
+    }
+  }
 }

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

@@ -17,12 +17,17 @@
 public class DelegatingRequestProducer implements HttpAsyncRequestProducer {
   final AgentSpan span;
   final HttpAsyncRequestProducer delegate;
+  boolean injectContext = false;
 
   public DelegatingRequestProducer(final AgentSpan span, final HttpAsyncRequestProducer delegate) {
     this.span = span;
     this.delegate = delegate;
   }
 
+  public void setInjectContext(boolean injectContext) {
+    this.injectContext = injectContext;
+  }
+
   @Override
   public HttpHost getTarget() {
     return delegate.getTarget();
@@ -32,7 +37,9 @@ 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) {
+      DECORATE.injectContext(current().with(span), request, SETTER);
+    }
     return request;
   }
 

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,7 +87,8 @@ 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() + "$ClientAdvice");
+        this.getClass().getName() + "$ClientAdvice",
+        this.getClass().getName() + "$ClientContextPropagationAdvice");
   }
 
   public static class ClientAdvice {
@@ -133,4 +136,14 @@ public static void methodExit(
       }
     }
   }
+
+  @AppliesOn(CONTEXT_TRACKING)
+  public static class ClientContextPropagationAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void onEnter(@Advice.Argument(0) final AsyncRequestProducer requestProducer) {
+      if (requestProducer instanceof DelegatingRequestProducer) {
+        ((DelegatingRequestProducer) requestProducer).setInjectContext(true);
+      }
+    }
+  }
 }

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);
   }
 }

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

@@ -11,12 +11,17 @@
 public class DelegatingRequestProducer implements AsyncRequestProducer {
   final AgentSpan span;
   final AsyncRequestProducer delegate;
+  boolean injectContext = false;
 
   public DelegatingRequestProducer(final AgentSpan span, final AsyncRequestProducer delegate) {
     this.span = span;
     this.delegate = delegate;
   }
 
+  public void setInjectContext(boolean injectContext) {
+    this.injectContext = injectContext;
+  }
+
   @Override
   public void failed(final Exception ex) {
     delegate.failed(ex);
@@ -25,7 +30,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);
   }