Fix aws sdk instrumentation

CHANGELOG.asciidoc

@@ -34,6 +34,7 @@ Use subheadings with the "=====" level for adding notes for unreleased changes:
 [float]
 ===== Features
 * Added protection against invalid timestamps provided by manual instrumentation - {pull}3363[#3363]
+* Added support for AWS SDK 2.21 - {pull}3373[#3373]
 
 [float]
 ===== Bug fixes

apm-agent-plugins/apm-aws-sdk/apm-aws-sdk-2-plugin/pom.xml

@@ -13,7 +13,7 @@
 
     <properties>
         <apm-agent-parent.base.dir>${project.basedir}/../../..</apm-agent-parent.base.dir>
-        <version.aws.sdk>2.20.157</version.aws.sdk>
+        <version.aws.sdk>2.21.1</version.aws.sdk>
         <version.aws.jms>2.0.0</version.aws.jms>
         <!-- AWS SDK v2 is compiled for Java 8 -->
         <maven.compiler.target>8</maven.compiler.target>

apm-agent-plugins/apm-aws-sdk/apm-aws-sdk-2-plugin/src/main/java/co/elastic/apm/agent/awssdk/v2/BaseAsyncClientHandlerInstrumentation.java

@@ -22,6 +22,8 @@
 import co.elastic.apm.agent.awssdk.v2.helper.S3Helper;
 import co.elastic.apm.agent.awssdk.v2.helper.SQSHelper;
 import co.elastic.apm.agent.sdk.ElasticApmInstrumentation;
+import co.elastic.apm.agent.sdk.logging.Logger;
+import co.elastic.apm.agent.sdk.logging.LoggerFactory;
 import co.elastic.apm.agent.tracer.GlobalTracer;
 import co.elastic.apm.agent.tracer.Span;
 import co.elastic.apm.agent.tracer.Tracer;
@@ -35,6 +37,7 @@
 import software.amazon.awssdk.core.client.config.SdkClientOption;
 import software.amazon.awssdk.core.client.handler.ClientExecutionParams;
 import software.amazon.awssdk.core.http.ExecutionContext;
+import software.amazon.awssdk.core.internal.handler.BaseAsyncClientHandler;
 import software.amazon.awssdk.core.internal.http.TransformingAsyncResponseHandler;
 
 import java.net.URI;
@@ -70,12 +73,21 @@ public Collection<String> getInstrumentationGroupNames() {
     @SuppressWarnings({"unchecked", "rawtypes"})
     public static class AdviceClass {
 
+        private static final Logger logger = LoggerFactory.getLogger(AdviceClass.class);
+
         @Advice.OnMethodEnter(suppress = Throwable.class, inline = false)
         @Advice.AssignReturned.ToArguments(@Advice.AssignReturned.ToArguments.ToArgument(value = 2))
         public static TransformingAsyncResponseHandler<?> enterDoExecute(@Advice.Argument(value = 0) ClientExecutionParams clientExecutionParams,
                                                                          @Advice.Argument(value = 1) ExecutionContext executionContext,
                                                                          @Advice.Argument(value = 2) TransformingAsyncResponseHandler<?> responseHandler,
-                                                                         @Advice.FieldValue("clientConfiguration") SdkClientConfiguration clientConfiguration) {
+                                                                         @Advice.This BaseAsyncClientHandler handler) {
+
+            SdkClientConfiguration clientConfiguration = ClientHandlerConfigInstrumentation.AdviceClass.getConfig(handler);
+            if(clientConfiguration == null) {
+                logger.warn("Not tracing AWS request due to being unable to resolve the configuration");
+                return responseHandler;
+            }
+
             String awsService = executionContext.executionAttributes().getAttribute(AwsSignerExecutionAttribute.SERVICE_NAME);
             SdkRequest sdkRequest = clientExecutionParams.getInput();
             URI uri = clientConfiguration.option(SdkClientOption.ENDPOINT);

apm-agent-plugins/apm-aws-sdk/apm-aws-sdk-2-plugin/src/main/java/co/elastic/apm/agent/awssdk/v2/BaseSyncClientHandlerInstrumentation.java

@@ -24,6 +24,8 @@
 import co.elastic.apm.agent.awssdk.v2.helper.sqs.wrapper.MessageListWrapper;
 import co.elastic.apm.agent.common.JvmRuntimeInfo;
 import co.elastic.apm.agent.sdk.ElasticApmInstrumentation;
+import co.elastic.apm.agent.sdk.logging.Logger;
+import co.elastic.apm.agent.sdk.logging.LoggerFactory;
 import co.elastic.apm.agent.tracer.GlobalTracer;
 import co.elastic.apm.agent.tracer.Outcome;
 import co.elastic.apm.agent.tracer.Span;
@@ -39,6 +41,8 @@
 import software.amazon.awssdk.core.client.config.SdkClientOption;
 import software.amazon.awssdk.core.client.handler.ClientExecutionParams;
 import software.amazon.awssdk.core.http.ExecutionContext;
+import software.amazon.awssdk.core.internal.handler.BaseAsyncClientHandler;
+import software.amazon.awssdk.core.internal.handler.BaseSyncClientHandler;
 
 import javax.annotation.Nullable;
 import java.net.URI;
@@ -81,11 +85,20 @@ public Collection<String> getInstrumentationGroupNames() {
     @SuppressWarnings("rawtypes")
     public static class AdviceClass {
 
+        private static final Logger logger = LoggerFactory.getLogger(AdviceClass.class);
+
         @Nullable
         @Advice.OnMethodEnter(suppress = Throwable.class, inline = false)
         public static Object enterDoExecute(@Advice.Argument(value = 0) ClientExecutionParams clientExecutionParams,
                                             @Advice.Argument(value = 1) ExecutionContext executionContext,
-                                            @Advice.FieldValue("clientConfiguration") SdkClientConfiguration clientConfiguration) {
+                                            @Advice.This BaseSyncClientHandler handler) {
+
+            SdkClientConfiguration clientConfiguration = ClientHandlerConfigInstrumentation.AdviceClass.getConfig(handler);
+            if(clientConfiguration == null) {
+                logger.warn("Not tracing AWS request due to being unable to resolve the configuration");
+                return null;
+            }
+
             String awsService = executionContext.executionAttributes().getAttribute(AwsSignerExecutionAttribute.SERVICE_NAME);
             SdkRequest sdkRequest = clientExecutionParams.getInput();
             URI uri = clientConfiguration.option(SdkClientOption.ENDPOINT);

apm-agent-plugins/apm-aws-sdk/apm-aws-sdk-2-plugin/src/main/java/co/elastic/apm/agent/awssdk/v2/ClientHandlerConfigInstrumentation.java

@@ -0,0 +1,95 @@
+/*
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package co.elastic.apm.agent.awssdk.v2;
+
+import co.elastic.apm.agent.awssdk.v2.helper.DynamoDbHelper;
+import co.elastic.apm.agent.awssdk.v2.helper.S3Helper;
+import co.elastic.apm.agent.awssdk.v2.helper.SQSHelper;
+import co.elastic.apm.agent.awssdk.v2.helper.sqs.wrapper.MessageListWrapper;
+import co.elastic.apm.agent.common.JvmRuntimeInfo;
+import co.elastic.apm.agent.sdk.ElasticApmInstrumentation;
+import co.elastic.apm.agent.sdk.weakconcurrent.WeakConcurrent;
+import co.elastic.apm.agent.sdk.weakconcurrent.WeakMap;
+import co.elastic.apm.agent.tracer.GlobalTracer;
+import co.elastic.apm.agent.tracer.Outcome;
+import co.elastic.apm.agent.tracer.Span;
+import co.elastic.apm.agent.tracer.Tracer;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import software.amazon.awssdk.auth.signer.AwsSignerExecutionAttribute;
+import software.amazon.awssdk.core.SdkRequest;
+import software.amazon.awssdk.core.SdkResponse;
+import software.amazon.awssdk.core.client.config.SdkClientConfiguration;
+import software.amazon.awssdk.core.client.config.SdkClientOption;
+import software.amazon.awssdk.core.client.handler.ClientExecutionParams;
+import software.amazon.awssdk.core.http.ExecutionContext;
+
+import javax.annotation.Nullable;
+import java.net.URI;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+
+import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
+import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+public class ClientHandlerConfigInstrumentation extends ElasticApmInstrumentation {
+
+
+    @Override
+    public ElementMatcher<? super TypeDescription> getTypeMatcher() {
+        return named("software.amazon.awssdk.core.internal.handler.BaseSyncClientHandler")
+            .or(named("software.amazon.awssdk.core.internal.handler.BaseAsyncClientHandler"));
+    }
+
+    @Override
+    public ElementMatcher<? super MethodDescription> getMethodMatcher() {
+        return isConstructor()
+            .and(takesArgument(0, named("software.amazon.awssdk.core.client.config.SdkClientConfiguration")));
+    }
+
+    @Override
+    public Collection<String> getInstrumentationGroupNames() {
+        return Collections.singleton("aws-sdk");
+    }
+
+
+    @SuppressWarnings("rawtypes")
+    public static class AdviceClass {
+
+        private static final WeakMap<Object, SdkClientConfiguration> configMap = WeakConcurrent.buildMap();
+
+        @Nullable
+        public static SdkClientConfiguration getConfig(Object handler) {
+            return configMap.get(handler);
+        }
+
+
+        @Advice.OnMethodExit(suppress = Throwable.class, inline = false)
+        public static void exit(@Advice.This Object handler, @Advice.Argument(value = 0) SdkClientConfiguration config) {
+            configMap.put(handler, config);
+        }
+
+    }
+}

apm-agent-plugins/apm-aws-sdk/apm-aws-sdk-2-plugin/src/main/resources/META-INF/services/co.elastic.apm.agent.sdk.ElasticApmInstrumentation

@@ -1,4 +1,5 @@
 co.elastic.apm.agent.awssdk.v2.BaseSyncClientHandlerInstrumentation
 co.elastic.apm.agent.awssdk.v2.BaseAsyncClientHandlerInstrumentation
+co.elastic.apm.agent.awssdk.v2.ClientHandlerConfigInstrumentation
 co.elastic.apm.agent.awssdk.v2.GetMessagesInstrumentation
-co.elastic.apm.agent.awssdk.v2.AmazonSQSMessagingClientWrapperInstrumentation
+co.elastic.apm.agent.awssdk.v2.AmazonSQSMessagingClientWrapperInstrumentation

docs/supported-technologies.asciidoc

@@ -294,12 +294,12 @@ This provides support for `com.datastax.cassandra:cassandra-driver-core` and
 |AWS DynamoDB
 |1.x, 2.x
 |The agent creates spans for interactions with the AWS DynamoDb service through the AWS Java SDK.
-|1.31.0
+|1.31.0, 2.21+ since 1.44.0
 
 |AWS S3
 |1.x, 2.x
 |The agent creates spans for interactions with the AWS S3 service through the AWS Java SDK.
-|1.31.0
+|1.31.0, 2.21+ since 1.44.0
 
 |===
 
@@ -459,7 +459,7 @@ same trace.
 |AWS SQS
 |1.x, 2.x
 |The agent captures SQS Message sends and polling as well as SQS message sends and consumption through JMS.
-|1.34.0
+|1.34.0, 2.21+ since 1.44.0
 
 |===
 

pom.xml

@@ -140,7 +140,7 @@
         <!-- used both for plugin & annotations dependency -->
         <version.animal-sniffer>1.17</version.animal-sniffer>
 
-        <version.testcontainers>1.16.3</version.testcontainers>
+        <version.testcontainers>1.19.1</version.testcontainers>
 
         <!-- latest version compiled for java 11 -->
         <version.jsonunit>2.38.0</version.jsonunit>