feat(spring-cloud-aws): instrument onMessage(Collection<Message<T>>) for batch consumption

Author: Aryainguz

instrumentation/spring/spring-cloud-aws-3.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/cloud/aws/v3_0/MessagingMessageListenerAdapterInstrumentation.java

@@ -8,8 +8,10 @@
 import static net.bytebuddy.matcher.ElementMatchers.named;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 
+import io.opentelemetry.context.Scope;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import java.util.Collection;
 import javax.annotation.Nullable;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
@@ -28,7 +30,9 @@ public void transform(TypeTransformer transformer) {
     transformer.applyAdviceToMethod(
         named("onMessage").and(takesArgument(0, named("org.springframework.messaging.Message"))),
         getClass().getName() + "$OnMessageAdvice");
-    // TODO: onMessage(Collection<Message<T>> messages) not instrumented
+    transformer.applyAdviceToMethod(
+        named("onMessage").and(takesArgument(0, named("java.util.Collection"))),
+        getClass().getName() + "$OnMessagesAdvice");
   }
 
   @SuppressWarnings("unused")
@@ -48,4 +52,20 @@ public static void methodExit(
       }
     }
   }
+
+  @SuppressWarnings("unused")
+  public static class OnMessagesAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class, inline = false)
+    @Nullable
+    public static Scope methodEnter(@Advice.Argument(0) Collection<Message<?>> messages) {
+      return SpringAwsUtil.handleBatch(messages);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class, inline = false)
+    public static void methodExit(@Advice.Enter @Nullable Scope scope) {
+      if (scope != null) {
+        scope.close();
+      }
+    }
+  }
 }

instrumentation/spring/spring-cloud-aws-3.0/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/spring/cloud/aws/v3_0/AwsSqsTest.java

@@ -26,6 +26,8 @@
 import static java.util.concurrent.TimeUnit.SECONDS;
 import static org.assertj.core.api.Assertions.assertThat;
 
+import java.util.Arrays;
+
 import io.awspring.cloud.sqs.operations.SqsTemplate;
 import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.instrumentation.testing.junit.AgentInstrumentationExtension;
@@ -176,4 +178,45 @@ void sqsListener() throws Exception {
                                     + "/000000000000/test-queue"),
                             satisfies(AWS_REQUEST_ID, val -> val.isInstanceOf(String.class)))));
   }
+
+  @Test
+  void sqsBatchListener() throws Exception {
+    String messageContent1 = "hello";
+    String messageContent2 = "world";
+    CompletableFuture<java.util.List<String>> messageFuture = new CompletableFuture<>();
+    AwsSqsTestApplication.batchMessageHandler =
+        strings -> testing.runWithSpan("callback", () -> messageFuture.complete(strings));
+
+    testing.runWithSpan(
+        "parent",
+        () ->
+            sqsTemplate.sendMany(
+                "test-batch-queue", Arrays.asList(messageContent1, messageContent2)));
+
+    java.util.List<String> result = messageFuture.get(10, SECONDS);
+    assertThat(result).containsExactlyInAnyOrder(messageContent1, messageContent2);
+
+    testing.waitAndAssertTraces(
+        trace ->
+            trace.hasSpansSatisfyingExactly(
+                span -> span.hasName("parent").hasKind(SpanKind.INTERNAL).hasNoParent(),
+                span ->
+                    span.hasName("Sqs.GetQueueUrl")
+                        .hasKind(SpanKind.CLIENT)
+                        .hasParent(trace.getSpan(0)),
+                span ->
+                    span.hasName("test-batch-queue publish")
+                        .hasKind(SpanKind.PRODUCER)
+                        .hasParent(trace.getSpan(0)),
+                span ->
+                    span.hasName("test-batch-queue process")
+                        .hasKind(SpanKind.CONSUMER)
+                        .hasParent(trace.getSpan(2)),
+                span ->
+                    span.hasName("callback").hasKind(SpanKind.INTERNAL).hasParent(trace.getSpan(3)),
+                span ->
+                    span.hasName("Sqs.DeleteMessageBatch")
+                        .hasKind(SpanKind.CLIENT)
+                        .hasParent(trace.getSpan(2))));
+  }
 }

instrumentation/spring/spring-cloud-aws-3.0/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/spring/cloud/aws/v3_0/AwsSqsTestApplication.java

@@ -42,4 +42,13 @@ void receiveStringMessage(String message) {
       messageHandler.accept(message);
     }
   }
+
+  static volatile Consumer<java.util.List<String>> batchMessageHandler;
+
+  @SqsListener("test-batch-queue")
+  void receiveBatchMessages(java.util.List<String> messages) {
+    if (batchMessageHandler != null) {
+      batchMessageHandler.accept(messages);
+    }
+  }
 }