Simplify aws-sdk traces (#4684)

* Trace simplification: don't capture HTTP client spans if they are part of an AWS-SDK call
* Move X-Ray propagation out from individual HTTP client advice
* Move X-Ray propagation into AWS-SDK advice
* Update AWS-SDK tests to not expect http.request spans
* Run old AWS-SDK tests with legacy flag enabled
* Guard against propagation errors leaking out through the interceptor

dd-java-agent/instrumentation/apache-httpclient-4/src/main/java/datadog/trace/instrumentation/apachehttpclient/ApacheHttpClientInstrumentation.java

@@ -188,7 +188,7 @@ public static AgentScope methodEnter(
  final AgentScope scope = HelperMethods.doMethodEnter(request);
 
  // Wrap the handler so we capture the status code
- if (handler instanceof ResponseHandler) {
+ if (null != scope && handler instanceof ResponseHandler) {
  handler = new WrappingStatusSettingResponseHandler(scope.span(), (ResponseHandler) handler);
  }
  return scope;
@@ -256,7 +256,7 @@ public static AgentScope methodEnter(
  }
 
  // Wrap the handler so we capture the status code
- if (handler instanceof ResponseHandler) {
+ if (null != scope && handler instanceof ResponseHandler) {
  handler = new WrappingStatusSettingResponseHandler(scope.span(), (ResponseHandler) handler);
  }
  return scope;

dd-java-agent/instrumentation/apache-httpclient-4/src/main/java/datadog/trace/instrumentation/apachehttpclient/HelperMethods.java

@@ -8,7 +8,6 @@
 import static datadog.trace.instrumentation.apachehttpclient.HttpHeadersInjectAdapter.SETTER;
 
 import datadog.trace.api.Config;
-import datadog.trace.api.TracePropagationStyle;
 import datadog.trace.bootstrap.CallDepthThreadLocalMap;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -18,23 +17,30 @@
 import org.apache.http.client.methods.HttpUriRequest;
 
 public class HelperMethods {
+ private static final boolean AWS_LEGACY_TRACING =
+ Config.get().isLegacyTracingEnabled(false, "aws-sdk");
+
  public static AgentScope doMethodEnter(final HttpUriRequest request) {
+ boolean awsClientCall = request.containsHeader("amz-sdk-invocation-id");
+ if (!AWS_LEGACY_TRACING && awsClientCall) {
+ // avoid creating an extra HTTP client span beneath the AWS client call
+ return null;
+ }
+
  final AgentSpan span = startSpan(HTTP_REQUEST);
  final AgentScope scope = activateSpan(span);
 
  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) {
  propagate().inject(span, request, SETTER);
  propagate()
  .injectPathwayContext(
  span, request, SETTER, HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS);
- } else if (Config.get().isAwsPropagationEnabled()) {
- propagate().inject(span, request, SETTER, TracePropagationStyle.XRAY);
  }
+
  return scope;
  }
 

dd-java-agent/instrumentation/apache-httpclient-5/src/main/java/datadog/trace/instrumentation/apachehttpclient5/ApacheHttpClientInstrumentation.java

@@ -173,7 +173,7 @@ public static AgentScope methodEnter(
  HelperMethods.doMethodEnter(new HostAndRequestAsHttpUriRequest(host, request));
 
  // Wrap the handler so we capture the status code
- if (handler instanceof HttpClientResponseHandler) {
+ if (null != scope && handler instanceof HttpClientResponseHandler) {
  handler =
  new WrappingStatusSettingResponseHandler(
  scope.span(), (HttpClientResponseHandler) handler);

dd-java-agent/instrumentation/apache-httpclient-5/src/main/java/datadog/trace/instrumentation/apachehttpclient5/HelperMethods.java

@@ -7,6 +7,7 @@
 import static datadog.trace.instrumentation.apachehttpclient5.ApacheHttpClientDecorator.HTTP_REQUEST;
 import static datadog.trace.instrumentation.apachehttpclient5.HttpHeadersInjectAdapter.SETTER;
 
+import datadog.trace.api.Config;
 import datadog.trace.bootstrap.CallDepthThreadLocalMap;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -16,16 +17,30 @@
 import org.apache.hc.core5.http.HttpResponse;
 
 public class HelperMethods {
+ private static final boolean AWS_LEGACY_TRACING =
+ Config.get().isLegacyTracingEnabled(false, "aws-sdk");
+
  public static AgentScope doMethodEnter(final HttpRequest request) {
+ boolean awsClientCall = request.containsHeader("amz-sdk-invocation-id");
+ if (!AWS_LEGACY_TRACING && awsClientCall) {
+ // avoid creating an extra HTTP client span beneath the AWS client call
+ return null;
+ }
+
  final AgentSpan span = startSpan(HTTP_REQUEST);
  final AgentScope scope = activateSpan(span);
 
  DECORATE.afterStart(span);
  DECORATE.onRequest(span, request);
 
- propagate().inject(span, request, SETTER);
- propagate()
- .injectPathwayContext(span, request, SETTER, HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS);
+ // AWS calls are often signed, so we can't add headers without breaking the signature.
+ if (!awsClientCall) {
+ propagate().inject(span, request, SETTER);
+ propagate()
+ .injectPathwayContext(
+ span, request, SETTER, HttpClientDecorator.CLIENT_PATHWAY_EDGE_TAGS);
+ }
+
  return scope;
  }
 

dd-java-agent/instrumentation/aws-java-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/AwsSdkClientDecorator.java

@@ -6,14 +6,16 @@
 import com.amazonaws.Response;
 import datadog.trace.api.Functions;
 import datadog.trace.api.cache.QualifiedClassNameCache;
+import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
 import datadog.trace.bootstrap.instrumentation.decorator.HttpClientDecorator;
 import java.net.URI;
 import java.util.function.Function;
 import java.util.regex.Pattern;
 
-public class AwsSdkClientDecorator extends HttpClientDecorator<Request, Response> {
+public class AwsSdkClientDecorator extends HttpClientDecorator<Request, Response>
+ implements AgentPropagation.Setter<Request<?>> {
  public static final AwsSdkClientDecorator DECORATE = new AwsSdkClientDecorator();
 
  private static final Pattern REQUEST_PATTERN = Pattern.compile("Request", Pattern.LITERAL);
@@ -138,4 +140,9 @@ protected URI url(final Request request) {
  protected int status(final Response response) {
  return response.getHttpResponse().getStatusCode();
  }
+
+ @Override
+ public void set(Request<?> carrier, String key, String value) {
+ carrier.addHeader(key, value);
+ }
 }

dd-java-agent/instrumentation/aws-java-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/TracingRequestHandler.java

@@ -3,6 +3,7 @@
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateNext;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.closePrevious;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
 import static datadog.trace.instrumentation.aws.v0.AwsSdkClientDecorator.DECORATE;
 
@@ -11,9 +12,13 @@
 import com.amazonaws.Response;
 import com.amazonaws.handlers.HandlerContextKey;
 import com.amazonaws.handlers.RequestHandler2;
+import datadog.trace.api.Config;
+import datadog.trace.api.TracePropagationStyle;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /** Tracing Request Handler */
 public class TracingRequestHandler extends RequestHandler2 {
@@ -24,6 +29,8 @@ public class TracingRequestHandler extends RequestHandler2 {
 
  private static final CharSequence AWS_HTTP = UTF8BytesString.create("aws.http");
 
+ private static final Logger log = LoggerFactory.getLogger(TracingRequestHandler.class);
+
  @Override
  public AmazonWebServiceRequest beforeMarshalling(final AmazonWebServiceRequest request) {
  return request;
@@ -42,6 +49,13 @@ public void beforeRequest(final Request<?> request) {
  activateNext(span); // this scope will last until next poll
  }
  request.addHandlerContext(SCOPE_CONTEXT_KEY, activateSpan(span));
+ if (Config.get().isAwsPropagationEnabled()) {
+ try {
+ propagate().inject(span, request, DECORATE, TracePropagationStyle.XRAY);
+ } catch (Throwable e) {
+ log.warn("Unable to inject trace header", e);
+ }
+ }
  }
 
  @Override

dd-java-agent/instrumentation/aws-java-sdk-1.11.0/src/test/groovy/AWS1ClientTest.groovy

@@ -34,8 +34,6 @@ import datadog.trace.agent.test.AgentTestRunner
 import datadog.trace.api.DDSpanTypes
 import datadog.trace.bootstrap.instrumentation.api.Tags
 import datadog.trace.test.util.Flaky
-import org.apache.http.conn.HttpHostConnectException
-import org.apache.http.impl.execchain.RequestAbortedException
 import spock.lang.AutoCleanup
 import spock.lang.Shared
 
@@ -130,7 +128,7 @@ class AWS1ClientTest extends AgentTestRunner {
  client.requestHandler2s.findAll{ it.getClass().getSimpleName() == "TracingRequestHandler" }.size() == 1
 
  assertTraces(1) {
- trace(2) {
+ trace(1) {
  span {
  serviceName "$ddService"
  operationName "aws.http"
@@ -157,24 +155,6 @@ class AWS1ClientTest extends AgentTestRunner {
  defaultTags()
  }
  }
- span {
- operationName "http.request"
- resourceName "$method $path"
- spanType DDSpanTypes.HTTP_CLIENT
- errored false
- measured true
- childOf(span(0))
- tags {
- "$Tags.COMPONENT" "apache-httpclient"
- "$Tags.SPAN_KIND" Tags.SPAN_KIND_CLIENT
- "$Tags.PEER_HOSTNAME" "localhost"
- "$Tags.PEER_PORT" server.address.port
- "$Tags.HTTP_URL" "${server.address}${path}"
- "$Tags.HTTP_METHOD" "$method"
- "$Tags.HTTP_STATUS" 200
- defaultTags()
- }
- }
  }
  }
  server.lastRequest.headers.get("x-datadog-trace-id") == null
@@ -237,7 +217,7 @@ class AWS1ClientTest extends AgentTestRunner {
  thrown SdkClientException
 
  assertTraces(1) {
- trace(2) {
+ trace(1) {
  span {
  serviceName "java-aws-sdk"
  operationName "aws.http"
@@ -264,24 +244,6 @@ class AWS1ClientTest extends AgentTestRunner {
  defaultTags()
  }
  }
- span {
- operationName "http.request"
- resourceName "$method /$url"
- spanType DDSpanTypes.HTTP_CLIENT
- errored true
- measured true
- childOf(span(0))
- tags {
- "$Tags.COMPONENT" "apache-httpclient"
- "$Tags.SPAN_KIND" Tags.SPAN_KIND_CLIENT
- "$Tags.PEER_HOSTNAME" "localhost"
- "$Tags.PEER_PORT" UNUSABLE_PORT
- "$Tags.HTTP_URL" "http://localhost:${UNUSABLE_PORT}/$url"
- "$Tags.HTTP_METHOD" "$method"
- errorTags HttpHostConnectException, ~/Connection refused/
- defaultTags()
- }
- }
  }
  }
 
@@ -357,7 +319,7 @@ class AWS1ClientTest extends AgentTestRunner {
  thrown AmazonClientException
 
  assertTraces(1) {
- trace(5) {
+ trace(1) {
  span {
  serviceName "java-aws-sdk"
  operationName "aws.http"
@@ -386,30 +348,6 @@ class AWS1ClientTest extends AgentTestRunner {
  defaultTags()
  }
  }
- (1..4).each {
- span {
- operationName "http.request"
- resourceName "GET /someBucket/someKey"
- spanType DDSpanTypes.HTTP_CLIENT
- errored true
- measured true
- childOf(span(0))
- tags {
- "$Tags.COMPONENT" "apache-httpclient"
- "$Tags.SPAN_KIND" Tags.SPAN_KIND_CLIENT
- "$Tags.PEER_HOSTNAME" "localhost"
- "$Tags.PEER_PORT" server.address.port
- "$Tags.HTTP_URL" "$server.address/someBucket/someKey"
- "$Tags.HTTP_METHOD" "GET"
- try {
- errorTags SocketException, "Socket closed"
- } catch (AssertionError e) {
- errorTags RequestAbortedException, "Request aborted"
- }
- defaultTags()
- }
- }
- }
  }
  }