AWS v0: Close span non AWS SDK errors

Looks like AWS SDK doesn't call interceptor when non-sdk exception
occurs. This leads to leaking open spans. Fix that by instrumenting
http client.

Note: currently has no tests.

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

@@ -0,0 +1,76 @@
+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.declaresField;
+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}. In these cases {@link RequestHandler2#afterError} is not called.
+ *
+ * <p>FIXME: come up with tests for this - maybe some test that mimics timeout?
+ */
+@AutoService(Instrumenter.class)
+public final class AWSHttpClientInstrumentation extends Instrumenter.Default {
+
+  public AWSHttpClientInstrumentation() {
+    super("aws-sdk");
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("com.amazonaws.http.AmazonHttpClient.RequestExecutor")
+        .and(declaresField(named("request")));
+  }
+
+  @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.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,8 +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 =
+  static final HandlerContextKey<Scope> SCOPE_CONTEXT_KEY = new HandlerContextKey<>("DatadogScope");
-      new HandlerContextKey<>("DatadogScope");
 
   @Override
   public AmazonWebServiceRequest beforeMarshalling(final AmazonWebServiceRequest request) {
@@ -35,16 +34,22 @@ public class TracingRequestHandler extends RequestHandler2 {
   @Override
   public void afterResponse(final Request<?> request, final Response<?> response) {
     final Scope scope = request.getHandlerContext(SCOPE_CONTEXT_KEY);
+    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);
+    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,14 +1,24 @@
+import com.amazonaws.AmazonClientException
+import com.amazonaws.ClientConfiguration
 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.agent.test.utils.PortUtils
 import datadog.trace.api.DDSpanTypes
 import io.opentracing.tag.Tags
+import org.apache.http.conn.HttpHostConnectException
 import spock.lang.AutoCleanup
 import spock.lang.Shared
 
@@ -17,6 +27,13 @@ import java.util.concurrent.atomic.AtomicReference
 import static datadog.trace.agent.test.server.http.TestHttpServer.httpServer
 
 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 +80,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 +154,61 @@ 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:${PortUtils.UNUSABLE_PORT}"
+            "$Tags.HTTP_METHOD.key" "$method"
+            "$Tags.SPAN_KIND.key" Tags.SPAN_KIND_CLIENT
+            "aws.service" { it.contains(service) }
+            "aws.endpoint" "http://localhost:${PortUtils.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:${PortUtils.UNUSABLE_PORT}/$url"
+            "$Tags.PEER_HOSTNAME.key" "localhost"
+            "$Tags.PEER_PORT.key" PortUtils.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:${PortUtils.UNUSABLE_PORT}")
+  }
 }

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

@@ -1,19 +1,29 @@
 import com.amazonaws.AmazonWebServiceClient
+import com.amazonaws.ClientConfiguration
 import com.amazonaws.SDKGlobalConfiguration
+import com.amazonaws.SdkClientException
+import com.amazonaws.auth.AWSCredentialsProviderChain
 import com.amazonaws.auth.AWSStaticCredentialsProvider
 import com.amazonaws.auth.AnonymousAWSCredentials
 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.client.builder.AwsClientBuilder
 import com.amazonaws.handlers.RequestHandler2
 import com.amazonaws.regions.Regions
+import com.amazonaws.retry.PredefinedRetryPolicies
 import com.amazonaws.services.ec2.AmazonEC2ClientBuilder
 import com.amazonaws.services.rds.AmazonRDSClientBuilder
 import com.amazonaws.services.rds.model.DeleteOptionGroupRequest
 import com.amazonaws.services.s3.AmazonS3Client
 import com.amazonaws.services.s3.AmazonS3ClientBuilder
 import datadog.trace.agent.test.AgentTestRunner
+import datadog.trace.agent.test.utils.PortUtils
 import datadog.trace.api.DDSpanTypes
 import io.opentracing.tag.Tags
+import org.apache.http.conn.HttpHostConnectException
 import spock.lang.AutoCleanup
 import spock.lang.Shared
 
@@ -22,6 +32,13 @@ import java.util.concurrent.atomic.AtomicReference
 import static datadog.trace.agent.test.server.http.TestHttpServer.httpServer
 
 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")
@@ -164,4 +181,61 @@ class AWSClientTest extends AgentTestRunner {
         </DeleteOptionGroupResponse>
       """       | AmazonRDSClientBuilder.standard().withEndpointConfiguration(endpoint).withCredentials(credentialsProvider).build()
   }
+
+  def "send #operation request to closed port"() {
+    setup:
+    responseBody.set(body)
+
+    when:
+    call.call(client)
+
+    then:
+    thrown SdkClientException
+
+    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:${PortUtils.UNUSABLE_PORT}"
+            "$Tags.HTTP_METHOD.key" "$method"
+            "$Tags.SPAN_KIND.key" Tags.SPAN_KIND_CLIENT
+            "aws.service" { it.contains(service) }
+            "aws.endpoint" "http://localhost:${PortUtils.UNUSABLE_PORT}"
+            "aws.operation" "${operation}Request"
+            "aws.agent" "java-aws-sdk"
+            errorTags SdkClientException, ~/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:${PortUtils.UNUSABLE_PORT}/$url"
+            "$Tags.PEER_HOSTNAME.key" "localhost"
+            "$Tags.PEER_PORT.key" PortUtils.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:${PortUtils.UNUSABLE_PORT}")
+  }
 }