Test latest version of spring-cloud-aws

instrumentation/spring/spring-cloud-aws-3.0/javaagent/build.gradle.kts

@@ -22,7 +22,7 @@ dependencies {
 
   testImplementation("org.elasticmq:elasticmq-rest-sqs_2.13")
 
-  testImplementation("io.awspring.cloud:spring-cloud-aws-starter-sqs:3.0.0")
+  testLibrary("io.awspring.cloud:spring-cloud-aws-starter-sqs:3.0.0")
   testLibrary("org.springframework.boot:spring-boot-starter-test:3.0.0")
   testLibrary("org.springframework.boot:spring-boot-starter-web:3.0.0")
 }

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

@@ -0,0 +1,45 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.spring.aws;
+
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.namedOneOf;
+import static net.bytebuddy.matcher.ElementMatchers.returns;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import org.springframework.messaging.Message;
+
+public class MessageHeaderUtilsInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("io.awspring.cloud.sqs.MessageHeaderUtils");
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        namedOneOf("addHeaderIfAbsent", "addHeadersIfAbsent", "removeHeaderIfPresent")
+            .and(returns(named("org.springframework.messaging.Message")))
+            .and(takesArgument(0, named("org.springframework.messaging.Message"))),
+        this.getClass().getName() + "$PreserveContextAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class PreserveContextAdvice {
+
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void methodExit(
+        @Advice.Argument(0) Message<?> original, @Advice.Return Message<?> result) {
+      SpringAwsUtil.copyTracingState(original, result);
+    }
+  }
+}

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

@@ -37,6 +37,7 @@ public List<TypeInstrumentation> typeInstrumentations() {
         new AbstractMessageConvertingMessageSourceInstrumentation(),
         new MessagingMessageListenerAdapterInstrumentation(),
         new SqsTemplateInstrumentation(),
-        new AcknowledgementExecutionContextInstrumentation());
+        new AcknowledgementExecutionContextInstrumentation(),
+        new MessageHeaderUtilsInstrumentation());
   }
 }

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

@@ -53,6 +53,14 @@ public static void attachTracingState(Object originalMessage, Message<?> convert
     tracingContextField.set(convertedMessage, new TracingContext(tracingList, message));
   }
 
+  public static void copyTracingState(Message<?> original, Message<?> transformed) {
+    if (original == transformed) {
+      return;
+    }
+
+    tracingContextField.set(transformed, tracingContextField.get(original));
+  }
+
   public static MessageScope handleMessage(Message<?> message) {
     TracingContext tracingContext = tracingContextField.get(message);
     if (tracingContext == null) {