Mute tracing for agent http calls

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpClientDecorator.java

@@ -38,6 +38,8 @@ public abstract class HttpClientDecorator<REQUEST, RESPONSE> extends UriBasedCli
 
   private static final Logger log = LoggerFactory.getLogger(HttpClientDecorator.class);
 
+  private static final String DATADOG_META_LANG = "Datadog-Meta-Lang";
+
   private static final BitSet CLIENT_ERROR_STATUSES = Config.get().getHttpClientErrorStatuses();
 
   private static final UTF8BytesString DEFAULT_RESOURCE_NAME = UTF8BytesString.create("/");
@@ -50,6 +52,14 @@ public abstract class HttpClientDecorator<REQUEST, RESPONSE> extends UriBasedCli
 
   protected abstract URI url(REQUEST request) throws URISyntaxException;
 
+  /**
+   * Returns {@code true} if the request was made by the Datadog agent itself. Such requests must
+   * not be traced to avoid self-tracing loops.
+   */
+  public boolean isAgentRequest(final REQUEST request) {
+    return getRequestHeader(request, DATADOG_META_LANG) != null;
+  }
+
   protected abstract int status(RESPONSE response);
 
   protected abstract String getRequestHeader(REQUEST request, String headerName);

dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/main/java11/datadog/trace/instrumentation/httpclient/HeadersAdvice.java

@@ -15,11 +15,14 @@
 public class HeadersAdvice {
   @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
   public static void methodExit(@Advice.Return(readOnly = false) HttpHeaders headers) {
-    // Note: adding duplicate keys will throw an IllegalArgumentException so we need to dedupe
-    // case insensitively
-    final Map<String, List<String>> headerMap = new TreeMap<>(CASE_INSENSITIVE_ORDER);
-    headerMap.putAll(headers.map());
-    DECORATE.injectContext(getCurrentContext(), headerMap, SETTER);
-    headers = HttpHeaders.of(headerMap, KEEP);
+    // Check if we should be injecting context into the headers first
+    if (DECORATE.isContextInjectionAllowed()) {
+      // Note: adding duplicate keys will throw an IllegalArgumentException so we need to dedupe
+      // case insensitively
+      final Map<String, List<String>> headerMap = new TreeMap<>(CASE_INSENSITIVE_ORDER);
+      headerMap.putAll(headers.map());
+      DECORATE.injectContext(getCurrentContext(), headerMap, SETTER);
+      headers = HttpHeaders.of(headerMap, KEEP);
+    }
   }
 }

dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/main/java11/datadog/trace/instrumentation/httpclient/JavaNetClientDecorator.java

@@ -3,22 +3,23 @@
 import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
 import datadog.trace.bootstrap.instrumentation.decorator.HttpClientDecorator;
 import java.net.URI;
-import java.net.URISyntaxException;
 import java.net.http.HttpRequest;
 import java.net.http.HttpResponse;
 
 public class JavaNetClientDecorator extends HttpClientDecorator<HttpRequest, HttpResponse<?>> {
-
-  public static final CharSequence COMPONENT = UTF8BytesString.create("java-http-client");
+  public static final String INSTRUMENTATION_NAME = "java-http-client";
+  public static final CharSequence COMPONENT = UTF8BytesString.create(INSTRUMENTATION_NAME);
 
   public static final JavaNetClientDecorator DECORATE = new JavaNetClientDecorator();
 
   public static final UTF8BytesString OPERATION_NAME =
       UTF8BytesString.create(DECORATE.operationName());
 
+  private static final ThreadLocal<Boolean> INJECT_CONTEXT = new ThreadLocal<>();
+
   @Override
   protected String[] instrumentationNames() {
-    return new String[] {"java-http-client"};
+    return new String[] {INSTRUMENTATION_NAME};
   }
 
   @Override
@@ -32,7 +33,7 @@ protected String method(HttpRequest httpRequest) {
   }
 
   @Override
-  protected URI url(HttpRequest httpRequest) throws URISyntaxException {
+  protected URI url(HttpRequest httpRequest) {
     return httpRequest.uri();
   }
 
@@ -55,4 +56,24 @@ protected String getRequestHeader(HttpRequest request, String headerName) {
   protected String getResponseHeader(HttpResponse<?> response, String headerName) {
     return response.headers().firstValue(headerName).orElse(null);
   }
+
+  /**
+   * Checks whether context injection into HTTP headers is currently allowed.
+   *
+   * @return {@code true} if context injection is allowed for the current thread, {@code false}
+   *     otherwise
+   */
+  public boolean isContextInjectionAllowed() {
+    return INJECT_CONTEXT.get() != null && INJECT_CONTEXT.get();
+  }
+
+  /** Enables context injection into HTTP headers for the current thread. */
+  public void allowContextInjection() {
+    INJECT_CONTEXT.set(true);
+  }
+
+  /** Disables context injection into HTTP headers for the current thread. */
+  public void blockContextInjection() {
+    INJECT_CONTEXT.remove();
+  }
 }

dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/main/java11/datadog/trace/instrumentation/httpclient/SendAdvice.java

@@ -3,6 +3,8 @@
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
 import static datadog.trace.instrumentation.httpclient.JavaNetClientDecorator.DECORATE;
+import static datadog.trace.instrumentation.httpclient.JavaNetClientDecorator.INSTRUMENTATION_NAME;
+import static datadog.trace.instrumentation.httpclient.JavaNetClientDecorator.OPERATION_NAME;
 
 import datadog.appsec.api.blocking.BlockingException;
 import datadog.trace.bootstrap.CallDepthThreadLocalMap;
@@ -17,16 +19,19 @@ public class SendAdvice {
   @Advice.OnMethodEnter(suppress = Throwable.class)
   public static AgentScope methodEnter(@Advice.Argument(value = 0) final HttpRequest httpRequest) {
     try {
+      if (DECORATE.isAgentRequest(httpRequest)) {
+        return null;
+      }
       // Here we avoid having the advice applied twice in case we have nested call of this
-      // intercepted
-      // method.
+      // intercepted method.
       // In this particular case, in HttpClientImpl the send method is calling sendAsync under the
-      // hood and we do not want to instrument twice.
+      // hood, and we do not want to instrument twice.
       final int callDepth = CallDepthThreadLocalMap.incrementCallDepth(HttpClient.class);
       if (callDepth > 0) {
         return null;
       }
-      final AgentSpan span = startSpan(JavaNetClientDecorator.OPERATION_NAME);
+      DECORATE.allowContextInjection();
+      final AgentSpan span = startSpan(INSTRUMENTATION_NAME, OPERATION_NAME);
       final AgentScope scope = activateSpan(span);
 
       DECORATE.afterStart(span);
@@ -36,6 +41,7 @@ public static AgentScope methodEnter(@Advice.Argument(value = 0) final HttpReque
       return scope;
     } catch (BlockingException e) {
       CallDepthThreadLocalMap.reset(HttpClient.class);
+      DECORATE.blockContextInjection();
       // re-throw blocking exceptions
       throw e;
     }
@@ -50,6 +56,7 @@ public static void methodExit(
       return;
     }
     CallDepthThreadLocalMap.reset(HttpClient.class);
+    DECORATE.blockContextInjection();
 
     AgentSpan span = scope.span();
     if (null != throwable) {

dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/main/java11/datadog/trace/instrumentation/httpclient/SendAsyncAdvice.java

@@ -2,13 +2,15 @@
 
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.captureSpan;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
 import static datadog.trace.instrumentation.httpclient.JavaNetClientDecorator.DECORATE;
+import static datadog.trace.instrumentation.httpclient.JavaNetClientDecorator.INSTRUMENTATION_NAME;
+import static datadog.trace.instrumentation.httpclient.JavaNetClientDecorator.OPERATION_NAME;
 
 import datadog.appsec.api.blocking.BlockingException;
 import datadog.trace.bootstrap.CallDepthThreadLocalMap;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
-import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import java.net.http.HttpClient;
 import java.net.http.HttpRequest;
 import java.net.http.HttpResponse;
@@ -21,16 +23,19 @@ public static AgentScope methodEnter(
       @Advice.Argument(value = 0) final HttpRequest httpRequest,
       @Advice.Argument(value = 1, readOnly = false) HttpResponse.BodyHandler<?> bodyHandler) {
     try {
+      if (DECORATE.isAgentRequest(httpRequest)) {
+        return null;
+      }
       // Here we avoid having the advice applied twice in case we have nested call of this
-      // intercepted
-      // method.
+      // intercepted method.
       // In this particular case, in HttpClientImpl the send method is calling sendAsync under the
-      // hood and we do not want to instrument twice.
+      // hood, and we do not want to instrument twice.
       final int callDepth = CallDepthThreadLocalMap.incrementCallDepth(HttpClient.class);
       if (callDepth > 0) {
         return null;
       }
-      final AgentSpan span = AgentTracer.startSpan(JavaNetClientDecorator.OPERATION_NAME);
+      DECORATE.allowContextInjection();
+      final AgentSpan span = startSpan(INSTRUMENTATION_NAME, OPERATION_NAME);
       final AgentScope scope = activateSpan(span);
       if (bodyHandler != null) {
         bodyHandler = new BodyHandlerWrapper<>(bodyHandler, captureSpan(span));
@@ -43,6 +48,7 @@ public static AgentScope methodEnter(
       return scope;
     } catch (BlockingException e) {
       CallDepthThreadLocalMap.reset(HttpClient.class);
+      DECORATE.blockContextInjection();
       // re-throw blocking exceptions
       throw e;
     }
@@ -59,6 +65,7 @@ public static void methodExit(
     }
     // clear the call depth once finished
     CallDepthThreadLocalMap.reset(HttpClient.class);
+    DECORATE.blockContextInjection();
 
     AgentSpan span = scope.span();
     if (throwable != null) {

dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/test/groovy/datadog/trace/instrumentation/httpclient/JavaHttpClientTest.groovy

@@ -44,6 +44,17 @@ abstract class JavaHttpClientTest extends HttpClientTest {
     false
   }
 
+  def "request to agent not traced"() {
+    when:
+    def status = doRequest("GET", server.address.resolve("/success"), ["Datadog-Meta-Lang": "java"])
+
+    then:
+    status == 200
+    assertTraces(1) {
+      server.distributedRequestTrace(it)
+    }
+  }
+
   def 'should not inject duplicate headers'() {
     when:
     def status = doRequest("GET", server.address.resolve("/success"),

dd-java-agent/instrumentation/okhttp/okhttp-3.0/src/main/java/datadog/trace/instrumentation/okhttp3/TracingInterceptor.java

@@ -17,7 +17,7 @@
 public class TracingInterceptor implements Interceptor {
   @Override
   public Response intercept(final Chain chain) throws IOException {
-    if (chain.request().header("Datadog-Meta-Lang") != null) {
+    if (DECORATE.isAgentRequest(chain.request())) {
       return chain.proceed(chain.request());
     }