aws-sdk-2.2: More reflection cleanup.

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

@@ -15,9 +15,23 @@ muzzle {
     // several software.amazon.awssdk artifacts are missing for this version
     skip("2.17.200")
   }
-}
 
-muzzle {
+  fail {
+    group.set("software.amazon.awssdk")
+    module.set("aws-core")
+    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")
+
+    // "fail" asserts that *all* the instrumentation modules fail to load, but the core one is
+    // actually expected to succeed, so exclude it from checks.
+    excludeInstrumentationName("aws-sdk-2.2-core")
+
+    // several software.amazon.awssdk artifacts are missing for this version
+    skip("2.17.200")
+  }
+
   pass {
     group.set("software.amazon.awssdk")
     module.set("sqs")
@@ -53,17 +67,26 @@ dependencies {
   latestDepTestLibrary("software.amazon.awssdk:sqs:+")
 }
 
-tasks.withType<Test>().configureEach {
-  systemProperty("testLatestDeps", findProperty("testLatestDeps") as Boolean)
-  // 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 {
+  val testExperimentalSqs by registering(Test::class) {
+    group = "verification"
+
+    systemProperty("otel.instrumentation.aws-sdk.experimental-use-propagator-for-messaging", "true")
+  }
+
+  check {
+    dependsOn(testExperimentalSqs)
+  }
+
+  withType<Test>().configureEach {
+    systemProperty("testLatestDeps", findProperty("testLatestDeps") as Boolean)
+    // TODO run tests both with and without experimental span attributes
+    systemProperty("otel.instrumentation.aws-sdk.experimental-span-attributes", "true")
+  }
 
-tasks.withType<com.github.jengelman.gradle.plugins.shadow.tasks.ShadowJar>().configureEach {
-  mergeServiceFiles {
-    include("software/amazon/awssdk/global/handlers/execution.interceptors")
+  withType<com.github.jengelman.gradle.plugins.shadow.tasks.ShadowJar>().configureEach {
+    mergeServiceFiles {
+      include("software/amazon/awssdk/global/handlers/execution.interceptors")
+    }
   }
 }

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

@@ -5,7 +5,16 @@
 
 package io.opentelemetry.javaagent.instrumentation.awssdk.v2_2;
 
+import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
+import static java.util.Collections.singletonList;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+
 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.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
 
 abstract class AbstractAwsSdkInstrumentationModule extends InstrumentationModule {
 
@@ -17,4 +26,34 @@ protected AbstractAwsSdkInstrumentationModule(String additionalInstrumentationNa
   public boolean isHelperClass(String className) {
     return className.startsWith("io.opentelemetry.contrib.awsxray.");
   }
+
+  @Override
+  public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
+    // We don't actually transform it but want to make sure we only apply the instrumentation when
+    // our key dependency is present.
+    return hasClassesNamed("software.amazon.awssdk.core.interceptor.ExecutionInterceptor");
+  }
+
+  @Override
+  public List<TypeInstrumentation> typeInstrumentations() {
+    return singletonList(new ResourceInjectingTypeInstrumentation());
+  }
+
+  abstract void doTransform(TypeTransformer transformer);
+
+  // A type instrumentation is needed to trigger resource injection.
+  public class ResourceInjectingTypeInstrumentation implements TypeInstrumentation {
+    @Override
+    public ElementMatcher<TypeDescription> typeMatcher() {
+      // This is essentially the entry point of the AWS SDK, all clients implement it. We can ensure
+      // our interceptor service definition is injected as early as possible if we typematch against
+      // it.
+      return named("software.amazon.awssdk.core.SdkClient");
+    }
+
+    @Override
+    public void transform(TypeTransformer transformer) {
+      doTransform(transformer);
+    }
+  }
 }

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

@@ -5,19 +5,11 @@
 
 package io.opentelemetry.javaagent.instrumentation.awssdk.v2_2;
 
-import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
-import static java.util.Collections.singletonList;
-import static net.bytebuddy.matcher.ElementMatchers.named;
-
 import com.google.auto.service.AutoService;
 import io.opentelemetry.instrumentation.awssdk.v2_2.autoconfigure.TracingExecutionInterceptor;
 import io.opentelemetry.javaagent.extension.instrumentation.HelperResourceBuilder;
 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.description.type.TypeDescription;
-import net.bytebuddy.matcher.ElementMatcher;
 
 @AutoService(InstrumentationModule.class)
 public class AwsSdkInstrumentationModule extends AbstractAwsSdkInstrumentationModule {
@@ -35,30 +27,7 @@ public void registerHelperResources(HelperResourceBuilder helperResourceBuilder)
   }
 
   @Override
-  public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
-    // We don't actually transform it but want to make sure we only apply the instrumentation when
-    // our key dependency is present.
-    return hasClassesNamed("software.amazon.awssdk.core.interceptor.ExecutionInterceptor");
-  }
-
-  @Override
-  public List<TypeInstrumentation> typeInstrumentations() {
-    return singletonList(new ResourceInjectingTypeInstrumentation());
-  }
-
-  // A type instrumentation is needed to trigger resource injection.
-  public static class ResourceInjectingTypeInstrumentation implements TypeInstrumentation {
-    @Override
-    public ElementMatcher<TypeDescription> typeMatcher() {
-      // This is essentially the entry point of the AWS SDK, all clients implement it. We can ensure
-      // our interceptor service definition is injected as early as possible if we typematch against
-      // it.
-      return named("software.amazon.awssdk.core.SdkClient");
-    }
-
-    @Override
-    public void transform(TypeTransformer transformer) {
-      // Nothing to transform, this type instrumentation is only used for injecting resources.
-    }
+  void doTransform(TypeTransformer transformer) {
+    // Nothing to transform, this type instrumentation is only used for injecting resources.
   }
 }

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

@@ -5,19 +5,13 @@
 
 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 static net.bytebuddy.matcher.ElementMatchers.none;
 
 import com.google.auto.service.AutoService;
 import io.opentelemetry.instrumentation.awssdk.v2_2.SqsAdviceBridge;
 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 AbstractAwsSdkInstrumentationModule {
@@ -27,21 +21,9 @@ public SqsInstrumentationModule() {
   }
 
   @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");
-    }
+  public void doTransform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        none(), SqsInstrumentationModule.class.getName() + "$RegisterAdvice");
   }
 
   @SuppressWarnings("unused")
@@ -50,7 +32,7 @@ public static class RegisterAdvice {
     public static void onExit() {
       // (indirectly) using SqsImpl class here to make sure it is available from SqsAccess
       // (injected into app classloader) and checked by Muzzle
-      SqsAdviceBridge.init();
+      SqsAdviceBridge.referenceForMuzzleOnly();
     }
   }
 }

instrumentation/aws-sdk/aws-sdk-2.2/library/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/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/SqsAccess.java

@@ -8,12 +8,18 @@
 import io.opentelemetry.context.propagation.TextMapPropagator;
 import io.opentelemetry.javaagent.tooling.muzzle.NoMuzzle;
 import software.amazon.awssdk.core.SdkRequest;
+import software.amazon.awssdk.core.SdkResponse;
 import software.amazon.awssdk.core.interceptor.Context;
 import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
+import software.amazon.awssdk.services.sqs.model.ReceiveMessageRequest;
+import software.amazon.awssdk.services.sqs.model.ReceiveMessageResponse;
+import software.amazon.awssdk.services.sqs.model.SendMessageRequest;
 
 // helper class for calling methods that use sqs types in SqsImpl
 // if SqsImpl is not present these methods are no op
 final class SqsAccess {
+  private SqsAccess() {}
+
   private static final boolean enabled = isSqsImplPresent();
 
   private static boolean isSqsImplPresent() {
@@ -31,26 +37,42 @@ private static boolean isSqsImplPresent() {
   }
 
   @NoMuzzle
-  static SdkRequest injectIntoSqsSendMessageRequest(
+  static boolean isSendMessageRequest(SdkRequest request) {
+    return enabled && request instanceof SendMessageRequest;
+  }
+
+  @NoMuzzle
+  static SdkRequest injectIntoSendMessageRequest(
       TextMapPropagator messagingPropagator,
       SdkRequest rawRequest,
       io.opentelemetry.context.Context otelContext) {
-    if (!enabled) {
-      return rawRequest;
-    }
-    return SqsImpl.injectIntoSqsSendMessageRequest(messagingPropagator, rawRequest, otelContext);
+    assert enabled; // enabled checked already in instance check.
+    return SqsImpl.injectIntoSendMessageRequest(messagingPropagator, rawRequest, otelContext);
+  }
+
+  @NoMuzzle
+  static boolean isReceiveMessageRequest(SdkRequest request) {
+    return enabled && request instanceof ReceiveMessageRequest;
+  }
+
+  @NoMuzzle
+  public static SdkRequest modifyReceiveMessageRequest(
+      SdkRequest request, boolean useXrayPropagator, TextMapPropagator messagingPropagator) {
+    assert enabled; // enabled checked already in instance check.
+    return SqsImpl.modifyReceiveMessageRequest(request, useXrayPropagator, messagingPropagator);
+  }
+
+  @NoMuzzle
+  static boolean isReceiveMessageResponse(SdkResponse response) {
+    return enabled && response instanceof ReceiveMessageResponse;
   }
 
   @NoMuzzle
   static void afterReceiveMessageExecution(
       TracingExecutionInterceptor config,
       Context.AfterExecution context,
       ExecutionAttributes executionAttributes) {
-    if (!enabled) {
-      return;
-    }
-    SqsImpl.afterConsumerResponse(config, executionAttributes, context);
+    assert enabled; // enabled checked already in instance check.
+    SqsImpl.afterReceiveMessageExecution(config, executionAttributes, context);
   }
-
-  private SqsAccess() {}
 }

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

@@ -8,8 +8,8 @@
 public final class SqsAdviceBridge {
   private SqsAdviceBridge() {}
 
-  public static void init() {
-    // called from advice
-    SqsImpl.init(); // Reference the actual, package-private, implementation class for Muzzle
+  public static void referenceForMuzzleOnly() {
+    throw new UnsupportedOperationException(
+        SqsImpl.class.getName() + " referencing for muzzle, should never be actually called");
   }
 }