aws-sdk-1.1: Copy SQS plugin/NoMuzzle approach from 2.2 (#8866)

Co-authored-by: Mateusz Rzeszutek <[email protected]>

Author: Oberon00

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

@@ -15,6 +15,22 @@ muzzle {
     module.set("aws-java-sdk-core")
     versions.set("[1.10.33,)")
     assertInverse.set(true)
+
+    excludeInstrumentationName("aws-sdk-1.11-sqs")
+  }
+
+  fail {
+    group.set("com.amazonaws")
+    module.set("aws-java-sdk-core")
+    versions.set("[1.10.33,)")
+
+    excludeInstrumentationName("aws-sdk-1.11-core")
+  }
+
+  pass {
+    group.set("com.amazonaws")
+    module.set("aws-java-sdk-sqs")
+    versions.set("[1.10.33,)")
   }
 }
 

instrumentation/aws-sdk/aws-sdk-1.11/javaagent/src/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/SqsAdviceBridge.java

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

instrumentation/aws-sdk/aws-sdk-1.11/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/awssdk/v1_11/AbstractAwsSdkInstrumentationModule.java

@@ -0,0 +1,60 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.awssdk.v1_11;
+
+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;
+
+// TODO: Copy & paste with only trivial adaptions from v2
+abstract class AbstractAwsSdkInstrumentationModule extends InstrumentationModule {
+
+  protected AbstractAwsSdkInstrumentationModule(String additionalInstrumentationName) {
+    super("aws-sdk", "aws-sdk-1.11", additionalInstrumentationName);
+  }
+
+  @Override
+  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("com.amazonaws.AmazonWebServiceClient");
+  }
+
+  @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("com.amazonaws.AmazonWebServiceClient");
+    }
+
+    @Override
+    public void transform(TypeTransformer transformer) {
+      doTransform(transformer);
+    }
+  }
+}

instrumentation/aws-sdk/aws-sdk-1.11/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/awssdk/v1_11/AwsSdkInstrumentationModule.java

@@ -15,7 +15,7 @@
 @AutoService(InstrumentationModule.class)
 public class AwsSdkInstrumentationModule extends InstrumentationModule {
   public AwsSdkInstrumentationModule() {
-    super("aws-sdk", "aws-sdk-1.11");
+    super("aws-sdk", "aws-sdk-1.11", "aws-sdk-1.11-core");
   }
 
   @Override

instrumentation/aws-sdk/aws-sdk-1.11/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/awssdk/v1_11/SqsInstrumentationModule.java

@@ -0,0 +1,38 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.awssdk.v1_11;
+
+import static net.bytebuddy.matcher.ElementMatchers.none;
+
+import com.google.auto.service.AutoService;
+import io.opentelemetry.instrumentation.awssdk.v1_11.SqsAdviceBridge;
+import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import net.bytebuddy.asm.Advice;
+
+@AutoService(InstrumentationModule.class)
+public class SqsInstrumentationModule extends AbstractAwsSdkInstrumentationModule {
+
+  public SqsInstrumentationModule() {
+    super("aws-sdk-1.11-sqs");
+  }
+
+  @Override
+  public void doTransform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        none(), SqsInstrumentationModule.class.getName() + "$RegisterAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class RegisterAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    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.referenceForMuzzleOnly();
+    }
+  }
+}

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

@@ -6,6 +6,8 @@ dependencies {
   implementation("io.opentelemetry.contrib:opentelemetry-aws-xray-propagator")
 
   library("com.amazonaws:aws-java-sdk-core:1.11.0")
+  library("com.amazonaws:aws-java-sdk-sqs:1.11.106")
+  compileOnly(project(":muzzle"))
 
   testImplementation(project(":instrumentation:aws-sdk:aws-sdk-1.11:testing"))
 
@@ -15,5 +17,4 @@ dependencies {
   testLibrary("com.amazonaws:aws-java-sdk-kinesis:1.11.106")
   testLibrary("com.amazonaws:aws-java-sdk-dynamodb:1.11.106")
   testLibrary("com.amazonaws:aws-java-sdk-sns:1.11.106")
-  testLibrary("com.amazonaws:aws-java-sdk-sqs:1.11.106")
 }

instrumentation/aws-sdk/aws-sdk-1.11/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/PluginImplUtil.java

@@ -0,0 +1,56 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.awssdk.v1_11;
+
+import java.util.logging.Level;
+import java.util.logging.Logger;
+
+final class PluginImplUtil { // TODO: Copy & paste from v2
+  private PluginImplUtil() {}
+
+  private static final Logger logger = Logger.getLogger(PluginImplUtil.class.getName());
+
+  /**
+   * Check if the given {@code moduleNameImpl} is present.
+   *
+   * <p>For library instrumentations, the Impls will always be available but might fail to
+   * load/link/initialize if the corresponding SDK classes are not on the class path. For javaagent,
+   * the Impl is available only when the corresponding InstrumentationModule was successfully
+   * applied (muzzle passed).
+   *
+   * <p>Note that an present-but-incompatible library can only be reliably detected by Muzzle. In
+   * library-mode, users need to ensure they are using a compatible SDK (component) versions
+   * themselves.
+   *
+   * @param implSimpleClassName The simple name of the impl class, e.g. {@code "SqsImpl"}. *
+   */
+  static boolean isImplPresent(String implSimpleClassName) {
+    // Computing the full name dynamically name here because library instrumentation classes are
+    // relocated when embedded in the agent.
+    // We use getName().replace() instead of getPackage() because the latter is not guaranteed to
+    // work in all cases (e.g., we might be loaded into a custom classloader that doesn't handle it)
+    String implFullClassName =
+        PluginImplUtil.class.getName().replace(".PluginImplUtil", "." + implSimpleClassName);
+    try {
+      Class.forName(implFullClassName);
+      return true;
+    } catch (ClassNotFoundException | LinkageError e) {
+      // ClassNotFoundException will happen when muzzle disabled us in javaagent mode; LinkageError
+      // (most likely a NoClassDefFoundError, potentially wrapped in an ExceptionInInitializerError)
+      // should be thrown when the class is loaded in library mode (where the Impl class itself can
+      // always be found) but a dependency failed to load (most likely because the corresponding SDK
+      // dependency is not on the class path).
+      logger.log(
+          Level.FINE,
+          e,
+          () ->
+              implFullClassName
+                  + " not present. "
+                  + "Most likely, corresponding SDK component is either not on classpath or incompatible.");
+      return false;
+    }
+  }
+}

instrumentation/aws-sdk/aws-sdk-1.11/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/SqsAccess.java

@@ -0,0 +1,31 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.awssdk.v1_11;
+
+import com.amazonaws.AmazonWebServiceRequest;
+import com.amazonaws.Request;
+import com.amazonaws.Response;
+import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
+import io.opentelemetry.javaagent.tooling.muzzle.NoMuzzle;
+
+final class SqsAccess {
+  private SqsAccess() {}
+
+  private static final boolean enabled = PluginImplUtil.isImplPresent("SqsImpl");
+
+  @NoMuzzle
+  static boolean afterResponse(
+      Request<?> request,
+      Response<?> response,
+      Instrumenter<Request<?>, Response<?>> consumerInstrumenter) {
+    return enabled && SqsImpl.afterResponse(request, response, consumerInstrumenter);
+  }
+
+  @NoMuzzle
+  static boolean beforeMarshalling(AmazonWebServiceRequest request) {
+    return enabled && SqsImpl.beforeMarshalling(request);
+  }
+}

instrumentation/aws-sdk/aws-sdk-1.11/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/SqsImpl.java

@@ -0,0 +1,70 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.awssdk.v1_11;
+
+import com.amazonaws.AmazonWebServiceRequest;
+import com.amazonaws.Request;
+import com.amazonaws.Response;
+import com.amazonaws.services.sqs.AmazonSQS;
+import com.amazonaws.services.sqs.model.Message;
+import com.amazonaws.services.sqs.model.ReceiveMessageRequest;
+import com.amazonaws.services.sqs.model.ReceiveMessageResult;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
+
+final class SqsImpl {
+  static {
+    // Force loading of SQS class; this ensures that an exception is thrown at this point when the
+    // SQS library is not present, which will cause SqsAccess to have enabled=false in library mode.
+    @SuppressWarnings("unused")
+    String ensureLoadedDummy = AmazonSQS.class.getName();
+  }
+
+  private SqsImpl() {}
+
+  static boolean afterResponse(
+      Request<?> request,
+      Response<?> response,
+      Instrumenter<Request<?>, Response<?>> consumerInstrumenter) {
+    if (response.getAwsResponse() instanceof ReceiveMessageResult) {
+      afterConsumerResponse(request, response, consumerInstrumenter);
+      return true;
+    }
+    return false;
+  }
+
+  /** Create and close CONSUMER span for each message consumed. */
+  private static void afterConsumerResponse(
+      Request<?> request,
+      Response<?> response,
+      Instrumenter<Request<?>, Response<?>> consumerInstrumenter) {
+    ReceiveMessageResult receiveMessageResult = (ReceiveMessageResult) response.getAwsResponse();
+    for (Message message : receiveMessageResult.getMessages()) {
+      createConsumerSpan(message, request, response, consumerInstrumenter);
+    }
+  }
+
+  private static void createConsumerSpan(
+      Message message,
+      Request<?> request,
+      Response<?> response,
+      Instrumenter<Request<?>, Response<?>> consumerInstrumenter) {
+    Context parentContext = SqsParentContext.ofSystemAttributes(message.getAttributes());
+    Context context = consumerInstrumenter.start(parentContext, request);
+    consumerInstrumenter.end(context, request, response, null);
+  }
+
+  static boolean beforeMarshalling(AmazonWebServiceRequest rawRequest) {
+    if (rawRequest instanceof ReceiveMessageRequest) {
+      ReceiveMessageRequest request = (ReceiveMessageRequest) rawRequest;
+      if (!request.getAttributeNames().contains(SqsParentContext.AWS_TRACE_SYSTEM_ATTRIBUTE)) {
+        request.withAttributeNames(SqsParentContext.AWS_TRACE_SYSTEM_ATTRIBUTE);
+      }
+      return true;
+    }
+    return false;
+  }
+}