Instrument spring-messaging to reduce chance of disconnected traces

dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie

@@ -299,6 +299,7 @@
 2 org.springframework.jndi.*
 2 org.springframework.lang.*
 2 org.springframework.messaging.*
+0 org.springframework.messaging.handler.invocation.InvocableHandlerMethod
 2 org.springframework.objenesis.*
 2 org.springframework.orm.*
 2 org.springframework.remoting.*

dd-java-agent/instrumentation/spring-messaging-4/build.gradle

@@ -0,0 +1,36 @@
+
+muzzle {
+  pass {
+    group = 'org.springframework'
+    module = 'spring-messaging'
+    versions = "[4.0.0.RELEASE,)"
+    assertInverse = true
+  }
+}
+
+ext {
+  minJavaVersionForTests = JavaVersion.VERSION_17
+}
+
+apply from: "$rootDir/gradle/java.gradle"
+
+addTestSuiteForDir('latestDepTest', 'test')
+
+[compileTestGroovy, compileLatestDepTestGroovy].each {
+  it.javaLauncher = getJavaLauncherFor(17)
+}
+
+dependencies {
+  compileOnly group: 'org.springframework', name: 'spring-messaging', version: '4.0.0.RELEASE'
+
+  // capture SQS send and receive spans, propagate trace details in messages
+  testImplementation project(':dd-java-agent:instrumentation:apache-httpclient-4')
+  testImplementation project(':dd-java-agent:instrumentation:aws-java-sdk-2.2')
+  testImplementation project(':dd-java-agent:instrumentation:aws-java-sqs-2.0')
+
+  testImplementation group: 'jakarta.annotation', name: 'jakarta.annotation-api', version: '2.0.0'
+  testImplementation group: 'io.awspring.cloud', name: 'spring-cloud-aws-sqs', version: '3.0.1'
+  testImplementation group: 'org.elasticmq', name: 'elasticmq-rest-sqs_2.13', version: '1.2.3'
+
+  latestDepTestImplementation group: 'org.springframework', name: 'spring-messaging', version: '+'
+}

dd-java-agent/instrumentation/spring-messaging-4/src/main/java/datadog/trace/instrumentation/springmessaging/SpringMessageDecorator.java

@@ -0,0 +1,43 @@
+package datadog.trace.instrumentation.springmessaging;
+
+import static datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes.MESSAGE_CONSUMER;
+import static datadog.trace.bootstrap.instrumentation.api.Tags.SPAN_KIND_CONSUMER;
+
+import datadog.trace.api.naming.SpanNaming;
+import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
+import datadog.trace.bootstrap.instrumentation.decorator.MessagingClientDecorator;
+
+public final class SpringMessageDecorator extends MessagingClientDecorator {
+  public static final SpringMessageDecorator DECORATE = new SpringMessageDecorator();
+
+  public static final CharSequence SPRING_INBOUND =
+      UTF8BytesString.create(
+          SpanNaming.instance().namingSchema().messaging().inboundOperation("spring"));
+
+  public static final CharSequence COMPONENT_NAME = UTF8BytesString.create("spring-messaging");
+
+  @Override
+  protected CharSequence spanType() {
+    return MESSAGE_CONSUMER;
+  }
+
+  @Override
+  protected String[] instrumentationNames() {
+    return new String[] {"spring-messaging"};
+  }
+
+  @Override
+  protected CharSequence component() {
+    return COMPONENT_NAME;
+  }
+
+  @Override
+  protected String service() {
+    return null;
+  }
+
+  @Override
+  protected String spanKind() {
+    return SPAN_KIND_CONSUMER;
+  }
+}

dd-java-agent/instrumentation/spring-messaging-4/src/main/java/datadog/trace/instrumentation/springmessaging/SpringMessageExtractAdapter.java

@@ -0,0 +1,42 @@
+package datadog.trace.instrumentation.springmessaging;
+
+import datadog.trace.api.cache.DDCache;
+import datadog.trace.api.cache.DDCaches;
+import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
+import de.thetaphi.forbiddenapis.SuppressForbidden;
+import java.util.Locale;
+import java.util.Map;
+import java.util.function.Function;
+import org.springframework.messaging.Message;
+
+public final class SpringMessageExtractAdapter
+    implements AgentPropagation.ContextVisitor<Message<?>> {
+  private static final Function<String, String> KEY_MAPPER =
+      new Function<String, String>() {
+        @SuppressForbidden
+        @Override
+        public String apply(String key) {
+          // normalize headers from different providers; raw SQS, JMS, spring-messaging, etc.
+          if ("AWSTraceHeader".equals(key) || "Sqs_Msa_AWSTraceHeader".equals(key)) {
+            return "x-amzn-trace-id";
+          }
+          return key.replace("__dash__", "-").replace('$', '-').toLowerCase(Locale.ROOT);
+        }
+      };
+
+  private final DDCache<String, String> cache = DDCaches.newFixedSizeCache(32);
+
+  public static final SpringMessageExtractAdapter GETTER = new SpringMessageExtractAdapter();
+
+  @Override
+  public void forEachKey(Message<?> carrier, AgentPropagation.KeyClassifier classifier) {
+    for (Map.Entry<String, ?> header : carrier.getHeaders().entrySet()) {
+      if (header.getValue() instanceof String) {
+        String lowerCaseKey = cache.computeIfAbsent(header.getKey(), KEY_MAPPER);
+        if (!classifier.accept(lowerCaseKey, (String) header.getValue())) {
+          return;
+        }
+      }
+    }
+  }
+}

dd-java-agent/instrumentation/spring-messaging-4/src/main/java/datadog/trace/instrumentation/springmessaging/SpringMessageHandlerInstrumentation.java

@@ -0,0 +1,87 @@
+package datadog.trace.instrumentation.springmessaging;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
+import static datadog.trace.instrumentation.springmessaging.SpringMessageDecorator.DECORATE;
+import static datadog.trace.instrumentation.springmessaging.SpringMessageDecorator.SPRING_INBOUND;
+import static datadog.trace.instrumentation.springmessaging.SpringMessageExtractAdapter.GETTER;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import net.bytebuddy.asm.Advice;
+import org.springframework.messaging.Message;
+import org.springframework.messaging.handler.invocation.InvocableHandlerMethod;
+
+@AutoService(Instrumenter.class)
+public final class SpringMessageHandlerInstrumentation extends Instrumenter.Tracing
+    implements Instrumenter.ForSingleType {
+
+  public SpringMessageHandlerInstrumentation() {
+    super("spring-messaging", "spring-messaging-4");
+  }
+
+  @Override
+  public String instrumentedType() {
+    return "org.springframework.messaging.handler.invocation.InvocableHandlerMethod";
+  }
+
+  @Override
+  public void adviceTransformations(AdviceTransformation transformation) {
+    transformation.applyAdvice(
+        isMethod()
+            .and(
+                named("invoke")
+                    .and(takesArgument(0, named("org.springframework.messaging.Message")))),
+        SpringMessageHandlerInstrumentation.class.getName() + "$HandleMessageAdvice");
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      packageName + ".SpringMessageDecorator",
+      packageName + ".SpringMessageExtractAdapter",
+      packageName + ".SpringMessageExtractAdapter$1"
+    };
+  }
+
+  public static class HandleMessageAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static AgentScope onEnter(
+        @Advice.This InvocableHandlerMethod thiz, @Advice.Argument(0) Message<?> message) {
+      AgentSpan.Context parentContext;
+      AgentSpan parent = activeSpan();
+      if (null != parent) {
+        // prefer existing context, assume it was already extracted from this message
+        parentContext = parent.context();
+      } else {
+        // otherwise try to re-extract the message context to avoid disconnected trace
+        parentContext = propagate().extract(message, GETTER);
+      }
+      AgentSpan span = startSpan(SPRING_INBOUND, parentContext);
+      DECORATE.afterStart(span);
+      span.setResourceName(DECORATE.spanNameForMethod(thiz.getMethod()));
+      return activateSpan(span);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit(@Advice.Enter AgentScope scope, @Advice.Thrown Throwable error) {
+      if (null == scope) {
+        return;
+      }
+      AgentSpan span = scope.span();
+      if (null != error) {
+        DECORATE.onError(span, error);
+      }
+      scope.close();
+      DECORATE.beforeFinish(span);
+      span.finish();
+    }
+  }
+}