Mute tracing for agent http calls (#10690)

feat(java-net): Mute tracing for agent http calls

fix(java-net): Limit context injection into headers only when making request

feat(communication): Reuse prepareRequest to deduplicate request creation and add missing marking headers

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

Author: PerfectSlayer

communication/src/main/java/datadog/communication/ddagent/DDAgentFeaturesDiscovery.java

@@ -1,8 +1,10 @@
 package datadog.communication.ddagent;
 
-import static datadog.communication.http.OkHttpUtils.DATADOG_CONTAINER_ID;
 import static datadog.communication.http.OkHttpUtils.DATADOG_CONTAINER_TAGS_HASH;
+import static datadog.communication.http.OkHttpUtils.msgpackRequestBodyOf;
+import static datadog.communication.http.OkHttpUtils.prepareRequest;
 import static datadog.communication.serialization.msgpack.MsgPackWriter.FIXARRAY;
+import static java.util.Collections.emptyMap;
 import static java.util.Collections.emptySet;
 import static java.util.Collections.singletonList;
 import static java.util.Collections.unmodifiableSet;
@@ -11,7 +13,6 @@
 import com.squareup.moshi.Moshi;
 import com.squareup.moshi.Types;
 import datadog.common.container.ContainerInfo;
-import datadog.communication.http.OkHttpUtils;
 import datadog.metrics.api.Monitoring;
 import datadog.metrics.api.Recording;
 import datadog.metrics.impl.statsd.DDAgentStatsDClientManager;
@@ -151,13 +152,9 @@ private void doDiscovery(State newState) {
     // 3. fallback if the endpoint couldn't be found or the response couldn't be parsed
     try (Recording recording = discoveryTimer.start()) {
       boolean fallback = true;
-      final Request.Builder requestBuilder =
-          new Request.Builder().url(agentBaseUrl.resolve("info").url());
-      final String containerId = ContainerInfo.get().getContainerId();
-      if (containerId != null) {
-        requestBuilder.header(DATADOG_CONTAINER_ID, containerId);
-      }
-      try (Response response = client.newCall(requestBuilder.build()).execute()) {
+      final Request request =
+          prepareRequest(agentBaseUrl.resolve("info"), emptyMap()).get().build();
+      try (Response response = client.newCall(request).execute()) {
         if (response.isSuccessful()) {
           processInfoResponseHeaders(response);
           fallback = !processInfoResponse(newState, response.body().string());
@@ -202,11 +199,8 @@ private String probeTracesEndpoint(State newState, String[] endpoints) {
       try (Response response =
           client
               .newCall(
-                  new Request.Builder()
-                      .put(
-                          OkHttpUtils.msgpackRequestBodyOf(
-                              singletonList(ByteBuffer.wrap(PROBE_MESSAGE))))
-                      .url(agentBaseUrl.resolve(candidate))
+                  prepareRequest(agentBaseUrl.resolve(candidate), emptyMap())
+                      .put(msgpackRequestBodyOf(singletonList(ByteBuffer.wrap(PROBE_MESSAGE))))
                       .build())
               .execute()) {
         if (response.code() != 404) {

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

@@ -38,6 +38,9 @@ public abstract class HttpClientDecorator<REQUEST, RESPONSE> extends UriBasedCli
 
   private static final Logger log = LoggerFactory.getLogger(HttpClientDecorator.class);
 
+  private static final String DATADOG_META_LANG_HEADER_NAME = "Datadog-Meta-Lang";
+  private static final String DD_CLIENT_LIBRARY_LANGUAGE_HEADER_NAME = "DD-Client-Library-Language";
+
   private static final BitSet CLIENT_ERROR_STATUSES = Config.get().getHttpClientErrorStatuses();
 
   private static final UTF8BytesString DEFAULT_RESOURCE_NAME = UTF8BytesString.create("/");
@@ -50,6 +53,15 @@ 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_HEADER_NAME) != null
+        || getRequestHeader(request, DD_CLIENT_LIBRARY_LANGUAGE_HEADER_NAME) != 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());
     }