Fix scope leak in aws sdk instrumentation (open-telemetry#13129)

Author: laurit

instrumentation/aws-sdk/aws-sdk-2.2/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/internal/TracingExecutionInterceptor.java

@@ -142,6 +142,11 @@ public SdkRequest modifyRequest(
     io.opentelemetry.context.Context parentOtelContext = io.opentelemetry.context.Context.current();
     SdkRequest request = context.request();
 
+    // the request has already been modified, duplicate interceptor?
+    if (executionAttributes.getAttribute(SDK_REQUEST_ATTRIBUTE) != null) {
+      return request;
+    }
+
     // Ignore presign request. These requests don't run all interceptor methods and the span created
     // here would never be ended and scope closed.
     if (executionAttributes.getAttribute(AwsSignerExecutionAttribute.PRESIGNER_EXPIRATION)
@@ -192,13 +197,6 @@ public SdkRequest modifyRequest(
     executionAttributes.putAttribute(PARENT_CONTEXT_ATTRIBUTE, parentOtelContext);
     executionAttributes.putAttribute(CONTEXT_ATTRIBUTE, otelContext);
     executionAttributes.putAttribute(REQUEST_FINISHER_ATTRIBUTE, requestFinisher);
-    if (executionAttributes
-        .getAttribute(SdkExecutionAttribute.CLIENT_TYPE)
-        .equals(ClientType.SYNC)) {
-      // We can only activate context for synchronous clients, which allows downstream
-      // instrumentation like Apache to know about the SDK span.
-      executionAttributes.putAttribute(SCOPE_ATTRIBUTE, otelContext.makeCurrent());
-    }
 
     Span span = Span.fromContext(otelContext);
 
@@ -233,6 +231,25 @@ public SdkRequest modifyRequest(
     return request;
   }
 
+  @Override
+  public void afterMarshalling(
+      Context.AfterMarshalling context, ExecutionAttributes executionAttributes) {
+    // the request has already been modified, duplicate interceptor?
+    if (executionAttributes.getAttribute(SCOPE_ATTRIBUTE) != null) {
+      return;
+    }
+
+    io.opentelemetry.context.Context otelContext = getContext(executionAttributes);
+    if (otelContext != null
+        && executionAttributes
+            .getAttribute(SdkExecutionAttribute.CLIENT_TYPE)
+            .equals(ClientType.SYNC)) {
+      // We can only activate context for synchronous clients, which allows downstream
+      // instrumentation like Apache to know about the SDK span.
+      executionAttributes.putAttribute(SCOPE_ATTRIBUTE, otelContext.makeCurrent());
+    }
+  }
+
   @Override
   public void beforeTransmission(
       Context.BeforeTransmission context, ExecutionAttributes executionAttributes) {

instrumentation/aws-sdk/aws-sdk-2.2/testing/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/AbstractAws2ClientTest.java

@@ -24,9 +24,11 @@
 import static io.opentelemetry.semconv.incubating.RpcIncubatingAttributes.RPC_SYSTEM;
 import static java.util.Arrays.asList;
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.assertj.core.api.Assertions.catchThrowable;
 
 import io.opentelemetry.api.trace.SpanKind;
+import io.opentelemetry.context.Context;
 import io.opentelemetry.sdk.testing.assertj.AttributeAssertion;
 import io.opentelemetry.sdk.trace.data.StatusData;
 import io.opentelemetry.testing.internal.armeria.common.HttpData;
@@ -56,6 +58,7 @@
 import software.amazon.awssdk.core.ResponseInputStream;
 import software.amazon.awssdk.core.async.AsyncResponseTransformer;
 import software.amazon.awssdk.core.exception.SdkClientException;
+import software.amazon.awssdk.core.exception.SdkException;
 import software.amazon.awssdk.core.retry.RetryPolicy;
 import software.amazon.awssdk.http.apache.ApacheHttpClient;
 import software.amazon.awssdk.regions.Region;
@@ -696,4 +699,24 @@ void testTimeoutAndRetryErrorsAreNotCaptured() {
                                 equalTo(stringKey("aws.agent"), "java-aws-sdk"),
                                 equalTo(stringKey("aws.bucket.name"), "somebucket"))));
   }
+
+  // regression test for
+  // https://github.com/open-telemetry/opentelemetry-java-instrumentation/issues/13124
+  // verify that scope is not leaked on exception
+  @Test
+  void testS3ListNullBucket() {
+    S3ClientBuilder builder = S3Client.builder();
+    configureSdkClient(builder);
+    S3Client client =
+        builder
+            .endpointOverride(clientUri)
+            .region(Region.AP_NORTHEAST_1)
+            .credentialsProvider(CREDENTIALS_PROVIDER)
+            .build();
+
+    assertThatThrownBy(() -> client.listObjectsV2(b -> b.bucket(null)))
+        .isInstanceOf(SdkException.class);
+
+    assertThat(Context.current()).isEqualTo(Context.root());
+  }
 }