aws-sdk-2.2: SNS.Publish support with experimental messaging propagator flag (#8830)

instrumentation/aws-sdk/README.md

@@ -5,7 +5,7 @@ For more information, see the respective public setters in the `AwsSdkTelemetryB
 * [SDK v1](./aws-sdk-1.11/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/AwsSdkTelemetryBuilder.java)
 * [SDK v2](./aws-sdk-2.2/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/AwsSdkTelemetryBuilder.java)
 
-| System property | Type | Default | Description                                                                                                                                    |
-|---|---|---|------------------------------------------------------------------------------------------------------------------------------------------------|
-| `otel.instrumentation.aws-sdk.experimental-span-attributes` | Boolean | `false` | Enable the capture of experimental span attributes.                                                                                            |
-| `otel.instrumentation.aws-sdk.experimental-use-propagator-for-messaging` | Boolean | `false` | Enable propagation via message attributes using configured propagator (in addition to X-Ray). At the moment, Supports only SQS and the v2 SDK. |
+| System property | Type | Default | Description                                                                                                                           |
+|---|---|---|---------------------------------------------------------------------------------------------------------------------------------------|
+| `otel.instrumentation.aws-sdk.experimental-span-attributes` | Boolean | `false` | Enable the capture of experimental span attributes.                                                                                   |
+| `otel.instrumentation.aws-sdk.experimental-use-propagator-for-messaging` | Boolean | `false` | v2 only, inject into SNS/SQS attributes with configured propagator: See [v2 README](aws-sdk-2.2/library/README.md#trace-propagation). |

instrumentation/aws-sdk/aws-sdk-2.2/javaagent/build.gradle.kts

@@ -11,6 +11,7 @@ muzzle {
     // client, which is not target of instrumentation anyways.
     extraDependency("software.amazon.awssdk:protocol-core")
     excludeInstrumentationName("aws-sdk-2.2-sqs")
+    excludeInstrumentationName("aws-sdk-2.2-sns")
 
     // several software.amazon.awssdk artifacts are missing for this version
     skip("2.17.200")
@@ -40,6 +41,22 @@ muzzle {
     // client, which is not target of instrumentation anyways.
     extraDependency("software.amazon.awssdk:protocol-core")
 
+    excludeInstrumentationName("aws-sdk-2.2-sns")
+
+    // several software.amazon.awssdk artifacts are missing for this version
+    skip("2.17.200")
+  }
+
+  pass {
+    group.set("software.amazon.awssdk")
+    module.set("sns")
+    versions.set("[2.2.0,)")
+    // Used by all SDK services, the only case it isn't is an SDK extension such as a custom HTTP
+    // client, which is not target of instrumentation anyways.
+    extraDependency("software.amazon.awssdk:protocol-core")
+
+    excludeInstrumentationName("aws-sdk-2.2-sqs")
+
     // several software.amazon.awssdk artifacts are missing for this version
     skip("2.17.200")
   }
@@ -63,6 +80,7 @@ dependencies {
   testLibrary("software.amazon.awssdk:rds:2.2.0")
   testLibrary("software.amazon.awssdk:s3:2.2.0")
   testLibrary("software.amazon.awssdk:sqs:2.2.0")
+  testLibrary("software.amazon.awssdk:sns:2.2.0")
 }
 
 tasks {

instrumentation/aws-sdk/aws-sdk-2.2/javaagent/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/SnsAdviceBridge.java

@@ -0,0 +1,15 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.awssdk.v2_2;
+
+public final class SnsAdviceBridge {
+  private SnsAdviceBridge() {}
+
+  public static void referenceForMuzzleOnly() {
+    throw new UnsupportedOperationException(
+        SnsImpl.class.getName() + " referencing for muzzle, should never be actually called");
+  }
+}

instrumentation/aws-sdk/aws-sdk-2.2/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/awssdk/v2_2/SnsInstrumentationModule.java

@@ -0,0 +1,38 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.awssdk.v2_2;
+
+import static net.bytebuddy.matcher.ElementMatchers.none;
+
+import com.google.auto.service.AutoService;
+import io.opentelemetry.instrumentation.awssdk.v2_2.SnsAdviceBridge;
+import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import net.bytebuddy.asm.Advice;
+
+@AutoService(InstrumentationModule.class)
+public class SnsInstrumentationModule extends AbstractAwsSdkInstrumentationModule {
+
+  public SnsInstrumentationModule() {
+    super("aws-sdk-2.2-sns");
+  }
+
+  @Override
+  public void doTransform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        none(), SnsInstrumentationModule.class.getName() + "$RegisterAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class RegisterAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void onExit() {
+      // (indirectly) using SnsImpl class here to make sure it is available from SnsAccess
+      // (injected into app classloader) and checked by Muzzle
+      SnsAdviceBridge.referenceForMuzzleOnly();
+    }
+  }
+}

instrumentation/aws-sdk/aws-sdk-2.2/library-autoconfigure/build.gradle.kts

@@ -18,6 +18,7 @@ dependencies {
   testLibrary("software.amazon.awssdk:rds:2.2.0")
   testLibrary("software.amazon.awssdk:s3:2.2.0")
   testLibrary("software.amazon.awssdk:sqs:2.2.0")
+  testLibrary("software.amazon.awssdk:sns:2.2.0")
 }
 
 tasks {

instrumentation/aws-sdk/aws-sdk-2.2/library/README.md

@@ -18,9 +18,23 @@ DynamoDbClient client = DynamoDbClient.builder()
 
 ## Trace propagation
 
-The AWS SDK instrumentation currently only supports injecting the trace header into the request
+The AWS SDK instrumentation always injects the trace header into the request
 using the [AWS Trace Header](https://docs.aws.amazon.com/xray/latest/devguide/xray-concepts.html#xray-concepts-tracingheader) format.
 This format is the only format recognized by AWS managed services, and populating will allow
-propagating the trace through them. If this does not fulfill your use case, perhaps because you are
+propagating the trace through them.
+
+Additionally, you can enable an experimental option to use the configured propagator to inject into
+message attributes (see [parent README](../../README.md)). This currently supports the following AWS APIs:
+
+* SQS.SendMessage
+* SQS.SendMessageBatch
+* SNS.Publish
+  (SNS.PublishBatch is not supported at the moment because it is not available in the minimum SDK
+  version targeted by the instrumentation)
+
+Note that injection will only happen if, after injection, a maximum of 10 attributes is used to not
+run over API limitations set by AWS.
+
+If this does not fulfill your use case, perhaps because you are
 using the same SDK with a different non-AWS managed service, let us know so we can provide
 configuration for this behavior.

instrumentation/aws-sdk/aws-sdk-2.2/library/build.gradle.kts

@@ -7,6 +7,7 @@ dependencies {
 
   library("software.amazon.awssdk:aws-core:2.2.0")
   library("software.amazon.awssdk:sqs:2.2.0")
+  library("software.amazon.awssdk:sns:2.2.0")
   library("software.amazon.awssdk:aws-json-protocol:2.2.0")
   compileOnly(project(":muzzle")) // For @NoMuzzle
 

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

@@ -0,0 +1,23 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.awssdk.v2_2;
+
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.propagation.TextMapPropagator;
+import io.opentelemetry.javaagent.tooling.muzzle.NoMuzzle;
+import software.amazon.awssdk.core.SdkRequest;
+
+final class SnsAccess {
+  private SnsAccess() {}
+
+  private static final boolean enabled = PluginImplUtil.isImplPresent("SnsImpl");
+
+  @NoMuzzle
+  public static SdkRequest modifyRequest(
+      SdkRequest request, Context otelContext, TextMapPropagator messagingPropagator) {
+    return enabled ? SnsImpl.modifyRequest(request, otelContext, messagingPropagator) : null;
+  }
+}

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

@@ -0,0 +1,75 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.awssdk.v2_2;
+
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.propagation.TextMapPropagator;
+import java.util.HashMap;
+import java.util.Map;
+import software.amazon.awssdk.core.SdkRequest;
+import software.amazon.awssdk.services.sns.SnsClient;
+import software.amazon.awssdk.services.sns.model.MessageAttributeValue;
+import software.amazon.awssdk.services.sns.model.PublishRequest;
+
+// this class is only used from SnsAccess from method with @NoMuzzle annotation
+class SnsImpl {
+  static {
+    // Force loading of SnsClient; this ensures that an exception is thrown at this point when the
+    // SNS library is not present, which will cause SnsAccess to have enabled=false in library mode.
+    @SuppressWarnings("unused")
+    String ensureLoadedDummy = SnsClient.class.getName();
+  }
+
+  private SnsImpl() {}
+
+  static SdkRequest modifyRequest(
+      SdkRequest request, Context otelContext, TextMapPropagator messagingPropagator) {
+    if (messagingPropagator == null) {
+      return null;
+    } else if (request instanceof PublishRequest) {
+      return injectIntoPublishRequest((PublishRequest) request, otelContext, messagingPropagator);
+    } else {
+      // NB: We do not support PublishBatchRequest which was only introduced in 2.17.84.
+      // To add support, some targeted use of @NoMuzzle + checks that the needed class
+      // is available should work. See
+      // https://github.com/open-telemetry/opentelemetry-java-instrumentation/pull/8830#discussion_r1247570985
+      return null;
+    }
+  }
+
+  private static SdkRequest injectIntoPublishRequest(
+      PublishRequest request, Context otelContext, TextMapPropagator messagingPropagator) {
+    // Note: Code is 1:1 copy & paste from SQS, but due to different types (packages) cannot be
+    // reused.
+    Map<String, MessageAttributeValue> messageAttributes =
+        new HashMap<>(request.messageAttributes());
+    if (!injectIntoMessageAttributes(messageAttributes, otelContext, messagingPropagator)) {
+      return request;
+    }
+    return request.toBuilder().messageAttributes(messageAttributes).build();
+  }
+
+  private static boolean injectIntoMessageAttributes(
+      Map<String, MessageAttributeValue> messageAttributes,
+      io.opentelemetry.context.Context otelContext,
+      TextMapPropagator messagingPropagator) {
+    // Note: Code is 1:1 copy & paste from SQS, but due to different types (packages) cannot be
+    // reused.
+    messagingPropagator.inject(
+        otelContext,
+        messageAttributes,
+        (carrier, k, v) -> {
+          carrier.put(k, MessageAttributeValue.builder().stringValue(v).dataType("String").build());
+        });
+
+    // Return whether the injection resulted in an attribute count that is still supported.
+    // See https://docs.aws.amazon.com/sns/latest/dg/sns-message-attributes.html
+    // While non-raw delivery would support an arbitrary number, that is something configured in
+    // the subscription, and adding more attributes might result in odd behavior (e.g. we might
+    // push out other attributes)
+    return messageAttributes.size() <= 10;
+  }
+}

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

@@ -121,13 +121,17 @@ final class TracingExecutionInterceptor implements ExecutionInterceptor {
       throw throwable;
     }
 
-    SdkRequest sqsModifiedRequest =
+    SdkRequest modifiedRequest =
         SqsAccess.modifyRequest(request, otelContext, useXrayPropagator, messagingPropagator);
-    if (sqsModifiedRequest != null) {
-      return sqsModifiedRequest;
+    if (modifiedRequest != null) {
+      return modifiedRequest;
+    }
+    modifiedRequest = SnsAccess.modifyRequest(request, otelContext, messagingPropagator);
+    if (modifiedRequest != null) {
+      return modifiedRequest;
     }
 
-    // Insert other special handling here, following the same pattern as SQS.
+    // Insert other special handling here, following the same pattern as SQS and SNS.
 
     return request;
   }

instrumentation/aws-sdk/aws-sdk-2.2/testing/build.gradle.kts

@@ -17,6 +17,7 @@ dependencies {
   compileOnly("software.amazon.awssdk:rds:2.2.0")
   compileOnly("software.amazon.awssdk:s3:2.2.0")
   compileOnly("software.amazon.awssdk:sqs:2.2.0")
+  compileOnly("software.amazon.awssdk:sns:2.2.0")
 
   // needed for SQS - using emq directly as localstack references emq v0.15.7 ie WITHOUT AWS trace header propagation
   implementation("org.elasticmq:elasticmq-rest-sqs_2.12:1.0.0")

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

@@ -28,6 +28,7 @@ import software.amazon.awssdk.services.s3.S3AsyncClient
 import software.amazon.awssdk.services.s3.S3Client
 import software.amazon.awssdk.services.s3.model.CreateBucketRequest
 import software.amazon.awssdk.services.s3.model.GetObjectRequest
+import software.amazon.awssdk.services.sns.SnsAsyncClient
 import software.amazon.awssdk.services.sqs.SqsAsyncClient
 import software.amazon.awssdk.services.sqs.SqsClient
 import software.amazon.awssdk.services.sqs.model.CreateQueueRequest
@@ -200,6 +201,17 @@ abstract class AbstractAws2ClientTest extends AbstractAws2ClientCoreTest {
     request.request().headers().get("X-Amzn-Trace-Id") != null
     request.request().headers().get("traceparent") == null
 
+    if (service == "Sns" && operation == "Publish") {
+      def content = request.request().content().toStringUtf8()
+      def containsId = content.contains("${traces[0][0].traceId}-${traces[0][0].spanId}")
+      def containsTp = content.contains("=traceparent")
+      if (isSqsAttributeInjectionEnabled()) {
+        assert containsId && containsTp
+      } else {
+        assert !containsId && !containsTp
+      }
+    }
+
     where:
     service | operation           | method | path                          | requestId                              | builder                  | call                                                                                                                             | body
     "S3"    | "CreateBucket"      | "PUT"  | path("somebucket")            | "UNKNOWN"                              | S3AsyncClient.builder()  | { c -> c.createBucket(CreateBucketRequest.builder().bucket("somebucket").build()) }                                              | ""
@@ -234,6 +246,16 @@ abstract class AbstractAws2ClientTest extends AbstractAws2ClientCoreTest {
           <ResponseMetadata><RequestId>0ac9cda2-bbf4-11d3-f92b-31fa5e8dbc99</RequestId></ResponseMetadata>
         </DeleteOptionGroupResponse>
         """
+    "Sns" | "Publish" | "POST" | "" | "f187a3c1-376f-11df-8963-01868b7c937a" | SnsAsyncClient.builder() | { SnsAsyncClient c -> c.publish(r -> r.message("hello")) } | """
+      <PublishResponse xmlns="https://sns.amazonaws.com/doc/2010-03-31/">
+          <PublishResult>
+              <MessageId>94f20ce6-13c5-43a0-9a9e-ca52d816e90b</MessageId>
+          </PublishResult>
+          <ResponseMetadata>
+              <RequestId>f187a3c1-376f-11df-8963-01868b7c937a</RequestId>
+          </ResponseMetadata>
+      </PublishResponse> 
+      """
   }
 
   // TODO(anuraaga): Without AOP instrumentation of the HTTP client, we cannot model retries as