aws-sdk: Use compile time reference+@NoMuzzle instead of reflection.

instrumentation/aws-sdk/aws-sdk-2.2/javaagent-unit-tests/README.md

@@ -0,0 +1,2 @@
+Contains some tests that need to be split to a separate module to avoid references to to the
+autoconfigure library from the main library test module.

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

@@ -0,0 +1,10 @@
+plugins {
+  id("otel.java-conventions")
+}
+
+dependencies {
+  testImplementation(project(":muzzle"))
+  testImplementation(project(":javaagent-extension-api"))
+  testImplementation(project(":instrumentation:aws-sdk:aws-sdk-2.2:javaagent"))
+  testCompileOnly("com.google.auto.service:auto-service-annotations") // Avoid compiler warnings
+}

instrumentation/aws-sdk/aws-sdk-2.2/javaagent-unit-tests/src/test/java/io/opentelemetry/instrumentation/awssdk/v2_2/SqsNoMuzzleReferenceTest.java

@@ -0,0 +1,51 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.awssdk.v2_2;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
+import io.opentelemetry.javaagent.instrumentation.awssdk.v2_2.AwsSdkInstrumentationModule;
+import io.opentelemetry.javaagent.instrumentation.awssdk.v2_2.SqsInstrumentationModule;
+import io.opentelemetry.javaagent.tooling.muzzle.InstrumentationModuleMuzzle;
+import io.opentelemetry.javaagent.tooling.muzzle.references.ClassRef;
+import java.util.Collection;
+import java.util.List;
+import java.util.Locale;
+import java.util.stream.Collectors;
+import org.junit.Test;
+
+/**
+ * This unit test should be a muzzle directive. It verifies that certain references are *not*
+ * present (or sources annotated with {@code @NoMuzzle}).
+ */
+public class SqsNoMuzzleReferenceTest {
+  @Test
+  public void shouldNotReferenceSqsFromMainModule() {
+    List<ClassRef> badRefs = getSqsRefs(new AwsSdkInstrumentationModule());
+    assertThat(badRefs)
+        .as("Unexpected references, use printMuzzleReferences Gradle task to find sources")
+        .isEmpty();
+  }
+
+  @Test
+  public void shouldReferenceSqsFromSqsModule() {
+    List<ClassRef> expectedRefs = getSqsRefs(new SqsInstrumentationModule());
+    assertThat(expectedRefs)
+        .as("References from Sqs module should not be empty, maybe bad use of @NoMuzzle")
+        .isNotEmpty();
+  }
+
+  private static List<ClassRef> getSqsRefs(InstrumentationModule module) {
+    Collection<ClassRef> refs = InstrumentationModuleMuzzle.getMuzzleReferences(module).values();
+    return refs.stream()
+        .filter(
+            ref ->
+                ref.getClassName().startsWith("software.amazon.")
+                    && ref.getClassName().toLowerCase(Locale.ROOT).contains("sqs"))
+        .collect(Collectors.toList());
+  }
+}

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

@@ -10,15 +10,33 @@ muzzle {
     // 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")
+  }
+}
+
+muzzle {
+  pass {
+    group.set("software.amazon.awssdk")
+    module.set("sqs")
+    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")
+
     // several software.amazon.awssdk artifacts are missing for this version
     skip("2.17.200")
   }
 }
 
 dependencies {
   implementation(project(":instrumentation:aws-sdk:aws-sdk-2.2:library-autoconfigure"))
+  implementation(project(":instrumentation:aws-sdk:aws-sdk-2.2:library"))
 
   library("software.amazon.awssdk:aws-core:2.2.0")
+  library("software.amazon.awssdk:sqs:2.2.0")
 
   testImplementation(project(":instrumentation:aws-sdk:aws-sdk-2.2:testing"))
   // Make sure these don't add HTTP headers
@@ -37,8 +55,11 @@ dependencies {
 
 tasks.withType<Test>().configureEach {
   systemProperty("testLatestDeps", findProperty("testLatestDeps") as Boolean)
-  // TODO run tests both with and without experimental span attributes
-  jvmArgs("-Dotel.instrumentation.aws-sdk.experimental-span-attributes=true")
+  // TODO run tests both with and without experimental span attributes, with & without extra propagation
+  systemProperties(mapOf(
+    "otel.instrumentation.aws-sdk.experimental-span-attributes" to "true",
+    "otel.instrumentation.aws-sdk.experimental-use-propagator-for-messaging" to "true",
+  ))
 }
 
 tasks.withType<com.github.jengelman.gradle.plugins.shadow.tasks.ShadowJar>().configureEach {

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

@@ -0,0 +1,55 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.awssdk.v2_2;
+
+import static java.util.Collections.singletonList;
+import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+
+import com.google.auto.service.AutoService;
+import io.opentelemetry.instrumentation.awssdk.v2_2.SqsImpl;
+import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import java.util.List;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@AutoService(InstrumentationModule.class)
+public class SqsInstrumentationModule extends InstrumentationModule {
+
+  public SqsInstrumentationModule() {
+    super("aws-sdk", "aws-sdk-2.2", "aws-sdk-2.2-sqs");
+  }
+
+  @Override
+  public List<TypeInstrumentation> typeInstrumentations() {
+    return singletonList(new DefaultSqsClientTypeInstrumentation());
+  }
+
+  public static class DefaultSqsClientTypeInstrumentation implements TypeInstrumentation {
+    @Override
+    public ElementMatcher<TypeDescription> typeMatcher() {
+      return named("software.amazon.awssdk.services.sqs.DefaultSqsClient");
+    }
+
+    @Override
+    public void transform(TypeTransformer transformer) {
+      transformer.applyAdviceToMethod(
+          isConstructor(), SqsInstrumentationModule.class.getName() + "$RegisterAdvice");
+    }
+  }
+
+  @SuppressWarnings("unused")
+  public static class RegisterAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void onExit() {
+      // using SqsImpl class here to make sure it is available from SqsAccess
+      SqsImpl.init();
+    }
+  }
+}

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

@@ -26,5 +26,6 @@ tasks {
   test {
     systemProperty("testLatestDeps", findProperty("testLatestDeps") as Boolean)
     systemProperty("otel.instrumentation.aws-sdk.experimental-span-attributes", true)
+    systemProperty("otel.instrumentation.aws-sdk.experimental-use-propagator-for-messaging", true)
   }
 }

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

@@ -6,7 +6,9 @@ dependencies {
   implementation("io.opentelemetry.contrib:opentelemetry-aws-xray-propagator")
 
   library("software.amazon.awssdk:aws-core:2.2.0")
+  library("software.amazon.awssdk:sqs:2.2.0")
   library("software.amazon.awssdk:aws-json-protocol:2.2.0")
+  compileOnly(project(":muzzle")) // For @NoMuzzle
 
   testImplementation(project(":instrumentation:aws-sdk:aws-sdk-2.2:testing"))
 

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

@@ -0,0 +1,43 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.awssdk.v2_2;
+
+import io.opentelemetry.context.propagation.TextMapPropagator;
+import io.opentelemetry.javaagent.tooling.muzzle.NoMuzzle;
+import software.amazon.awssdk.core.SdkRequest;
+
+// helper class for calling methods that use sqs types in SqsImpl
+// if SqsImpl is not present these methods are no op
+public final class SqsAccess {
+  private static final boolean enabled = isSqsImplPresent();
+
+  private static boolean isSqsImplPresent() {
+    try {
+      // for library instrumentation SqsImpl is always available
+      // for javaagent instrumentation SqsImpl is available only when SqsInstrumentationModule was
+      // successfully applied (muzzle passed)
+      // using package name here because library instrumentation classes are relocated when embedded
+      // in the agent
+      Class.forName(SqsAccess.class.getPackage().getName() + ".SqsImpl");
+      return true;
+    } catch (ClassNotFoundException e) {
+      return false;
+    }
+  }
+
+  @NoMuzzle
+  public static SdkRequest injectIntoSqsSendMessageRequest(
+      TextMapPropagator messagingPropagator,
+      SdkRequest rawRequest,
+      io.opentelemetry.context.Context otelContext) {
+    if (!enabled) {
+      return rawRequest;
+    }
+    return SqsImpl.injectIntoSqsSendMessageRequest(messagingPropagator, rawRequest, otelContext);
+  }
+
+  private SqsAccess() {}
+}

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

@@ -0,0 +1,48 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.awssdk.v2_2;
+
+import io.opentelemetry.context.propagation.TextMapPropagator;
+import java.util.HashMap;
+import java.util.Map;
+import javax.annotation.Nullable;
+import software.amazon.awssdk.core.SdkRequest;
+import software.amazon.awssdk.services.sqs.model.MessageAttributeValue;
+import software.amazon.awssdk.services.sqs.model.SendMessageRequest;
+
+// this class is only used from SqsAccess from method with @NoMuzzle annotation
+public class SqsImpl {
+  private SqsImpl() {}
+
+  public static void init() {
+    // called from advice
+  }
+
+  public static SdkRequest injectIntoSqsSendMessageRequest(
+      TextMapPropagator messagingPropagator,
+      SdkRequest rawRequest,
+      io.opentelemetry.context.Context otelContext) {
+    SendMessageRequest request = (SendMessageRequest) rawRequest;
+    Map<String, MessageAttributeValue> messageAttributes =
+        new HashMap<>(request.messageAttributes());
+
+    // Need to use a full method to allow @NoMuzzle annotation (@NoMuzzle is not transitively
+    // applied to called methods)
+    messagingPropagator.inject(otelContext, messageAttributes, SqsImpl::injectSqsAttribute);
+
+    if (messageAttributes.size() > 10) { // Too many attributes, we don't want to break the call.
+      return request;
+    }
+    return request.toBuilder().messageAttributes(messageAttributes).build();
+  }
+
+  private static void injectSqsAttribute(
+      @Nullable Map<String, MessageAttributeValue> carrier, String k, String v) {
+    if (carrier != null) {
+      carrier.put(k, MessageAttributeValue.builder().stringValue(v).dataType("String").build());
+    }
+  }
+}

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

@@ -9,9 +9,10 @@
 import io.opentelemetry.context.propagation.TextMapGetter;
 import io.opentelemetry.context.propagation.TextMapPropagator;
 import io.opentelemetry.contrib.awsxray.propagator.AwsXrayPropagator;
+import io.opentelemetry.javaagent.tooling.muzzle.NoMuzzle;
 import java.util.Collections;
 import java.util.Map;
-import software.amazon.awssdk.core.SdkPojo;
+import software.amazon.awssdk.services.sqs.model.MessageAttributeValue;
 
 final class SqsParentContext {
 
@@ -29,31 +30,32 @@ public String get(Map<String, String> map, String s) {
     }
   }
 
-  enum MessageAttributeValueMapGetter implements TextMapGetter<Map<String, SdkPojo>> {
+  enum MessageAttributeValueMapGetter implements TextMapGetter<Map<String, MessageAttributeValue>> {
     INSTANCE;
 
     @Override
-    public Iterable<String> keys(Map<String, SdkPojo> map) {
+    public Iterable<String> keys(Map<String, MessageAttributeValue> map) {
       return map.keySet();
     }
 
     @Override
-    public String get(Map<String, SdkPojo> map, String s) {
+    @NoMuzzle
+    public String get(Map<String, MessageAttributeValue> map, String s) {
       if (map == null) {
         return null;
       }
-      SdkPojo value = map.get(s);
+      MessageAttributeValue value = map.get(s);
       if (value == null) {
         return null;
       }
-      return SqsMessageAttributeValueAccess.getStringValue(value);
+      return value.stringValue();
     }
   }
 
   static final String AWS_TRACE_SYSTEM_ATTRIBUTE = "AWSTraceHeader";
 
   static Context ofMessageAttributes(
-      Map<String, SdkPojo> messageAttributes, TextMapPropagator propagator) {
+      Map<String, MessageAttributeValue> messageAttributes, TextMapPropagator propagator) {
     return propagator.extract(
         Context.root(), messageAttributes, MessageAttributeValueMapGetter.INSTANCE);
   }