Bug update aws sync async client handler instrumentation

CHANGELOG.asciidoc

@@ -38,6 +38,7 @@ Use subheadings with the "=====" level for adding notes for unreleased changes:
 [float]
 ===== Bug fixes
 * Fixed too many spans being created for `HTTPUrlConnection` requests with method `HEAD` - {pull}3353[#3353]
+* Fixed aws sdk v2 instrumentation due to internal changes - {pull}3374[#3374]
 
 [[release-notes-1.x]]
 === Java Agent version 1.x

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.0</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

@@ -35,8 +35,14 @@
 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.BaseClientHandler;
 import software.amazon.awssdk.core.internal.http.TransformingAsyncResponseHandler;
+import software.amazon.awssdk.utils.internal.ReflectionUtils;
 
+import javax.annotation.Nonnull;
+import java.lang.reflect.Field;
+import java.lang.reflect.Modifier;
 import java.net.URI;
 import java.util.Collection;
 import java.util.Collections;
@@ -75,9 +81,13 @@ public static class AdviceClass {
         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 thiz) throws Throwable {
             String awsService = executionContext.executionAttributes().getAttribute(AwsSignerExecutionAttribute.SERVICE_NAME);
             SdkRequest sdkRequest = clientExecutionParams.getInput();
+            SdkClientConfiguration clientConfiguration = SQSHelper.getInstance().findClientConfiguration(clientExecutionParams, thiz);
+            if (null == clientConfiguration) {
+                return responseHandler;
+            }
             URI uri = clientConfiguration.option(SdkClientOption.ENDPOINT);
             Span<?> span = null;
             boolean isSqs = false;

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

@@ -39,6 +39,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.BaseSyncClientHandler;
 
 import javax.annotation.Nullable;
 import java.net.URI;
@@ -85,9 +86,13 @@ public static class AdviceClass {
         @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 thiz) throws Throwable {
             String awsService = executionContext.executionAttributes().getAttribute(AwsSignerExecutionAttribute.SERVICE_NAME);
             SdkRequest sdkRequest = clientExecutionParams.getInput();
+            SdkClientConfiguration clientConfiguration = SQSHelper.getInstance().findClientConfiguration(clientExecutionParams, thiz);
+            if (null == clientConfiguration) {
+                return null;
+            }
             URI uri = clientConfiguration.option(SdkClientOption.ENDPOINT);
             Span<?> span = null;
             boolean isSqs = false;

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

@@ -18,16 +18,21 @@
  */
 package co.elastic.apm.agent.awssdk.v2.helper;
 
+import co.elastic.apm.agent.awssdk.common.AbstractMessageIteratorWrapper;
 import co.elastic.apm.agent.awssdk.common.AbstractSQSInstrumentationHelper;
 import co.elastic.apm.agent.tracer.ElasticContext;
 import co.elastic.apm.agent.tracer.GlobalTracer;
 import co.elastic.apm.agent.tracer.Span;
 import co.elastic.apm.agent.tracer.Tracer;
 import co.elastic.apm.agent.tracer.dispatch.TextHeaderSetter;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 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.handler.ClientExecutionParams;
 import software.amazon.awssdk.core.http.ExecutionContext;
+import software.amazon.awssdk.core.internal.handler.BaseClientHandler;
 import software.amazon.awssdk.services.sqs.model.Message;
 import software.amazon.awssdk.services.sqs.model.MessageAttributeValue;
 import software.amazon.awssdk.services.sqs.model.ReceiveMessageRequest;
@@ -36,7 +41,10 @@
 import software.amazon.awssdk.services.sqs.model.SendMessageBatchRequestEntry;
 import software.amazon.awssdk.services.sqs.model.SendMessageRequest;
 
+import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
+import java.lang.reflect.Field;
+import java.lang.reflect.Modifier;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
@@ -49,6 +57,8 @@ public class SQSHelper extends AbstractSQSInstrumentationHelper<SdkRequest, Exec
 
     private static final SQSHelper INSTANCE = new SQSHelper(GlobalTracer.get());
 
+    public static final Logger logger = LoggerFactory.getLogger(SQSHelper.class);
+
     public static SQSHelper getInstance() {
         return INSTANCE;
     }
@@ -182,4 +192,30 @@ public void handleReceivedMessages(Span<?> span, @Nullable SdkRequest sdkRequest
             SQSHelper.getInstance().handleReceivedMessages(span, ((ReceiveMessageRequest) sdkRequest).queueUrl(), ((ReceiveMessageResponse) sdkResponse).messages());
         }
     }
+
+    @Nullable
+    public SdkClientConfiguration findClientConfiguration(ClientExecutionParams clientExecutionParams, @Nonnull BaseClientHandler thiz) throws Throwable {
+        SdkClientConfiguration sdkClientConfiguration = clientExecutionParams.requestConfiguration();
+        if (sdkClientConfiguration != null) {
+            return sdkClientConfiguration;
+        }
+        Class<?> targetClass = BaseClientHandler.class;
+        Class<?> currentClass = thiz.getClass();
+        while (null != currentClass) {
+            if (targetClass.equals(currentClass)) {
+                Field clientConfigurationField = currentClass.getDeclaredField("clientConfiguration");
+                int fieldModifier = clientConfigurationField.getModifiers();
+                if (Modifier.isPrivate(fieldModifier)) {
+                    clientConfigurationField.setAccessible(true);
+                }
+                sdkClientConfiguration = (SdkClientConfiguration) clientConfigurationField.get(thiz);
+                break;
+            }
+            currentClass = currentClass.getSuperclass();
+        }
+        if (sdkClientConfiguration == null) {
+            logger.warn("Cannot identify SdkClientConfiguration");
+        }
+        return sdkClientConfiguration;
+    }
 }