Merge pull request #770 from DataDog/mar-kolya/aws-sdk-v0-memory-leak

AWS v0: Close span non AWS SDK errors

Author: tylerbenson

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

@@ -0,0 +1,110 @@
+package datadog.trace.instrumentation.aws.v0;
+
+import static datadog.trace.instrumentation.aws.v0.AwsSdkClientDecorator.DECORATE;
+import static java.util.Collections.singletonMap;
+import static net.bytebuddy.matcher.ElementMatchers.isAbstract;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.not;
+
+import com.amazonaws.AmazonClientException;
+import com.amazonaws.Request;
+import com.amazonaws.handlers.RequestHandler2;
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import io.opentracing.Scope;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+/**
+ * This is additional 'helper' to catch cases when HTTP request throws exception different from
+ * {@link AmazonClientException} (for example an error thrown by another handler). In these cases
+ * {@link RequestHandler2#afterError} is not called.
+ */
+@AutoService(Instrumenter.class)
+public class AWSHttpClientInstrumentation extends Instrumenter.Default {
+
+  public AWSHttpClientInstrumentation() {
+    super("aws-sdk");
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("com.amazonaws.http.AmazonHttpClient");
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      "datadog.trace.agent.decorator.BaseDecorator",
+      "datadog.trace.agent.decorator.ClientDecorator",
+      "datadog.trace.agent.decorator.HttpClientDecorator",
+      packageName + ".AwsSdkClientDecorator",
+      packageName + ".TracingRequestHandler",
+    };
+  }
+
+  @Override
+  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    return singletonMap(
+        isMethod().and(not(isAbstract())).and(named("doExecute")),
+        HttpClientAdvice.class.getName());
+  }
+
+  public static class HttpClientAdvice {
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void methodExit(
+        @Advice.Argument(value = 0, optional = true) final Request<?> request,
+        @Advice.Thrown final Throwable throwable) {
+      if (throwable != null) {
+        final Scope scope = request.getHandlerContext(TracingRequestHandler.SCOPE_CONTEXT_KEY);
+        if (scope != null) {
+          request.addHandlerContext(TracingRequestHandler.SCOPE_CONTEXT_KEY, null);
+          DECORATE.onError(scope.span(), throwable);
+          DECORATE.beforeFinish(scope.span());
+          scope.close();
+        }
+      }
+    }
+  }
+
+  /**
+   * Due to a change in the AmazonHttpClient class, this instrumentation is needed to support newer
+   * versions. The above class should cover older versions.
+   */
+  @AutoService(Instrumenter.class)
+  public static final class RequestExecutorInstrumentation extends AWSHttpClientInstrumentation {
+
+    @Override
+    public ElementMatcher<TypeDescription> typeMatcher() {
+      return named("com.amazonaws.http.AmazonHttpClient$RequestExecutor");
+    }
+
+    @Override
+    public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+      return singletonMap(
+          isMethod().and(not(isAbstract())).and(named("doExecute")),
+          RequestExecutorAdvice.class.getName());
+    }
+
+    public static class RequestExecutorAdvice {
+      @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+      public static void methodExit(
+          @Advice.FieldValue("request") final Request<?> request,
+          @Advice.Thrown final Throwable throwable) {
+        if (throwable != null) {
+          final Scope scope = request.getHandlerContext(TracingRequestHandler.SCOPE_CONTEXT_KEY);
+          if (scope != null) {
+            request.addHandlerContext(TracingRequestHandler.SCOPE_CONTEXT_KEY, null);
+            DECORATE.onError(scope.span(), throwable);
+            DECORATE.beforeFinish(scope.span());
+            scope.close();
+          }
+        }
+      }
+    }
+  }
+}

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

@@ -16,7 +16,7 @@ public class TracingRequestHandler extends RequestHandler2 {
 
   // Note: aws1.x sdk doesn't have any truly async clients so we can store scope in request context
   // safely.
-  private static final HandlerContextKey<Scope> SCOPE_CONTEXT_KEY =
+  public static final HandlerContextKey<Scope> SCOPE_CONTEXT_KEY =
       new HandlerContextKey<>("DatadogScope");
 
   @Override
@@ -35,16 +35,22 @@ public void beforeRequest(final Request<?> request) {
   @Override
   public void afterResponse(final Request<?> request, final Response<?> response) {
     final Scope scope = request.getHandlerContext(SCOPE_CONTEXT_KEY);
-    DECORATE.onResponse(scope.span(), response);
-    DECORATE.beforeFinish(scope.span());
-    scope.close();
+    if (scope != null) {
+      request.addHandlerContext(SCOPE_CONTEXT_KEY, null);
+      DECORATE.onResponse(scope.span(), response);
+      DECORATE.beforeFinish(scope.span());
+      scope.close();
+    }
   }
 
   @Override
   public void afterError(final Request<?> request, final Response<?> response, final Exception e) {
     final Scope scope = request.getHandlerContext(SCOPE_CONTEXT_KEY);
-    DECORATE.onError(scope.span(), e);
-    DECORATE.beforeFinish(scope.span());
-    scope.close();
+    if (scope != null) {
+      request.addHandlerContext(SCOPE_CONTEXT_KEY, null);
+      DECORATE.onError(scope.span(), e);
+      DECORATE.beforeFinish(scope.span());
+      scope.close();
+    }
   }
 }

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

@@ -1,22 +1,42 @@
+import com.amazonaws.AmazonClientException
+import com.amazonaws.ClientConfiguration
+import com.amazonaws.Request
 import com.amazonaws.SDKGlobalConfiguration
+import com.amazonaws.auth.AWSCredentialsProviderChain
 import com.amazonaws.auth.BasicAWSCredentials
+import com.amazonaws.auth.EnvironmentVariableCredentialsProvider
+import com.amazonaws.auth.InstanceProfileCredentialsProvider
+import com.amazonaws.auth.SystemPropertiesCredentialsProvider
+import com.amazonaws.auth.profile.ProfileCredentialsProvider
 import com.amazonaws.handlers.RequestHandler2
+import com.amazonaws.retry.PredefinedRetryPolicies
 import com.amazonaws.services.ec2.AmazonEC2Client
 import com.amazonaws.services.rds.AmazonRDSClient
 import com.amazonaws.services.rds.model.DeleteOptionGroupRequest
 import com.amazonaws.services.s3.AmazonS3Client
 import com.amazonaws.services.s3.S3ClientOptions
 import datadog.trace.agent.test.AgentTestRunner
 import datadog.trace.api.DDSpanTypes
+import io.opentracing.Tracer
 import io.opentracing.tag.Tags
+import org.apache.http.conn.HttpHostConnectException
+import org.apache.http.impl.execchain.RequestAbortedException
 import spock.lang.AutoCleanup
 import spock.lang.Shared
 
 import java.util.concurrent.atomic.AtomicReference
 
 import static datadog.trace.agent.test.server.http.TestHttpServer.httpServer
+import static datadog.trace.agent.test.utils.PortUtils.UNUSABLE_PORT
 
 class AWSClientTest extends AgentTestRunner {
+
+  private static final CREDENTIALS_PROVIDER_CHAIN = new AWSCredentialsProviderChain(
+    new EnvironmentVariableCredentialsProvider(),
+    new SystemPropertiesCredentialsProvider(),
+    new ProfileCredentialsProvider(),
+    new InstanceProfileCredentialsProvider())
+
   def setupSpec() {
     System.setProperty(SDKGlobalConfiguration.ACCESS_KEY_SYSTEM_PROPERTY, "my-access-key")
     System.setProperty(SDKGlobalConfiguration.SECRET_KEY_SYSTEM_PROPERTY, "my-secret-key")
@@ -63,9 +83,11 @@ class AWSClientTest extends AgentTestRunner {
   def "send #operation request with mocked response"() {
     setup:
     responseBody.set(body)
+
+    when:
     def response = call.call(client)
 
-    expect:
+    then:
     response != null
 
     client.requestHandler2s != null
@@ -135,4 +157,175 @@ class AWSClientTest extends AgentTestRunner {
         </DeleteOptionGroupResponse>
       """       | new AmazonRDSClient().withEndpoint("http://localhost:$server.address.port")
   }
+
+  def "send #operation request to closed port"() {
+    setup:
+    responseBody.set(body)
+
+    when:
+    call.call(client)
+
+    then:
+    thrown AmazonClientException
+
+    assertTraces(1) {
+      trace(0, 2) {
+        span(0) {
+          serviceName "java-aws-sdk"
+          operationName "aws.http"
+          resourceName "$service.$operation"
+          spanType DDSpanTypes.HTTP_CLIENT
+          errored true
+          parent()
+          tags {
+            "$Tags.COMPONENT.key" "java-aws-sdk"
+            "$Tags.HTTP_URL.key" "http://localhost:${UNUSABLE_PORT}"
+            "$Tags.HTTP_METHOD.key" "$method"
+            "$Tags.SPAN_KIND.key" Tags.SPAN_KIND_CLIENT
+            "aws.service" { it.contains(service) }
+            "aws.endpoint" "http://localhost:${UNUSABLE_PORT}"
+            "aws.operation" "${operation}Request"
+            "aws.agent" "java-aws-sdk"
+            errorTags AmazonClientException, ~/Unable to execute HTTP request/
+            defaultTags()
+          }
+        }
+        span(1) {
+          operationName "http.request"
+          resourceName "$method /$url"
+          spanType DDSpanTypes.HTTP_CLIENT
+          errored true
+          childOf(span(0))
+          tags {
+            "$Tags.COMPONENT.key" "apache-httpclient"
+            "$Tags.HTTP_URL.key" "http://localhost:${UNUSABLE_PORT}/$url"
+            "$Tags.PEER_HOSTNAME.key" "localhost"
+            "$Tags.PEER_PORT.key" UNUSABLE_PORT
+            "$Tags.HTTP_METHOD.key" "$method"
+            "$Tags.SPAN_KIND.key" Tags.SPAN_KIND_CLIENT
+            errorTags HttpHostConnectException, ~/Connection refused/
+            defaultTags()
+          }
+        }
+      }
+    }
+
+    where:
+    service | operation   | method | url                  | call                                                    | body | client
+    "S3"    | "GetObject" | "GET"  | "someBucket/someKey" | { client -> client.getObject("someBucket", "someKey") } | ""   | new AmazonS3Client(CREDENTIALS_PROVIDER_CHAIN, new ClientConfiguration().withRetryPolicy(PredefinedRetryPolicies.getDefaultRetryPolicyWithCustomMaxRetries(0))).withEndpoint("http://localhost:${UNUSABLE_PORT}")
+  }
+
+  def "naughty request handler doesn't break the trace"() {
+    setup:
+    def client = new AmazonS3Client(CREDENTIALS_PROVIDER_CHAIN)
+    client.addRequestHandler(new RequestHandler2() {
+      void beforeRequest(Request<?> request) {
+        throw new RuntimeException("bad handler")
+      }
+    })
+
+    when:
+    client.getObject("someBucket", "someKey")
+
+    then:
+    ((Tracer) TEST_TRACER).activeSpan() == null
+    thrown RuntimeException
+
+    assertTraces(1) {
+      trace(0, 1) {
+        span(0) {
+          serviceName "java-aws-sdk"
+          operationName "aws.http"
+          resourceName "S3.GetObject"
+          spanType DDSpanTypes.HTTP_CLIENT
+          errored true
+          parent()
+          tags {
+            "$Tags.COMPONENT.key" "java-aws-sdk"
+            "$Tags.HTTP_URL.key" "https://s3.amazonaws.com"
+            "$Tags.HTTP_METHOD.key" "GET"
+            "$Tags.SPAN_KIND.key" Tags.SPAN_KIND_CLIENT
+            "aws.service" "Amazon S3"
+            "aws.endpoint" "https://s3.amazonaws.com"
+            "aws.operation" "GetObjectRequest"
+            "aws.agent" "java-aws-sdk"
+            errorTags RuntimeException, "bad handler"
+            defaultTags()
+          }
+        }
+      }
+    }
+  }
+
+  def "timeout and retry errors captured"() {
+    setup:
+    def server = httpServer {
+      handlers {
+        all {
+          Thread.sleep(500)
+          response.status(200).send()
+        }
+      }
+    }
+    AmazonS3Client client = new AmazonS3Client(new ClientConfiguration().withRequestTimeout(50 /* ms */))
+      .withEndpoint("http://localhost:$server.address.port")
+
+    when:
+    client.getObject("someBucket", "someKey")
+
+    then:
+    ((Tracer) TEST_TRACER).activeSpan() == null
+    thrown AmazonClientException
+
+    assertTraces(1) {
+      trace(0, 5) {
+        span(0) {
+          serviceName "java-aws-sdk"
+          operationName "aws.http"
+          resourceName "S3.GetObject"
+          spanType DDSpanTypes.HTTP_CLIENT
+          errored true
+          parent()
+          tags {
+            "$Tags.COMPONENT.key" "java-aws-sdk"
+            "$Tags.HTTP_URL.key" "http://localhost:$server.address.port"
+            "$Tags.HTTP_METHOD.key" "GET"
+            "$Tags.SPAN_KIND.key" Tags.SPAN_KIND_CLIENT
+            "aws.service" "Amazon S3"
+            "aws.endpoint" "http://localhost:$server.address.port"
+            "aws.operation" "GetObjectRequest"
+            "aws.agent" "java-aws-sdk"
+            errorTags AmazonClientException, ~/Unable to execute HTTP request/
+            defaultTags()
+          }
+        }
+        (1..4).each {
+          span(it) {
+            operationName "http.request"
+            resourceName "GET /someBucket/someKey"
+            spanType DDSpanTypes.HTTP_CLIENT
+            errored true
+            childOf(span(0))
+            tags {
+              "$Tags.COMPONENT.key" "apache-httpclient"
+              "$Tags.HTTP_URL.key" "http://localhost:$server.address.port/someBucket/someKey"
+              "$Tags.PEER_HOSTNAME.key" "localhost"
+              "$Tags.PEER_PORT.key" server.address.port
+              "$Tags.HTTP_METHOD.key" "GET"
+              "$Tags.SPAN_KIND.key" Tags.SPAN_KIND_CLIENT
+              try {
+                errorTags SocketException, "Socket closed"
+              } catch (AssertionError e) {
+                errorTags RequestAbortedException, "Request aborted"
+              }
+              defaultTags()
+            }
+          }
+        }
+      }
+    }
+
+    cleanup:
+    server.close()
+  }
 }