Separate context tracking from tracing on http client instrumentations (#10782)

phase 1

migrate jetty and akka http clients

wip

more clients

migrate httpclient5

spotless

fix bad ai stuff

Fix compile

spotless

Use blackhole for grpc ignored methods

final fixes

fix jetty client

add nullcheck

Co-authored-by: devflow.devflow-routing-intake <devflow.devflow-routing-intake@kubernetes.us1.ddbuild.io>

Author: amarziali

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

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

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

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

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,30 +64,47 @@ 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() + "$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);

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();
@@ -31,8 +40,16 @@ public HttpHost getTarget() {
   @Override
   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 (span != null) {
+      DECORATE.onRequest(span, new HostAndRequestAsHttpUriRequest(delegate.getTarget(), request));
+    }
+    if (injectContext) {
+      Context receiver = current();
+      if (span != null) {
+        receiver = receiver.with(span);
+      }
+      DECORATE.injectContext(receiver, request, SETTER);
+    }
     return request;
   }