Add instrumentation for rocketmq

instrumentation/rocketmq-client-4.8/javaagent/rocketmq-client-4.8-javaagent.gradle

@@ -0,0 +1,16 @@
+apply from: "$rootDir/gradle/instrumentation.gradle"
+
+muzzle {
+  pass {
+    group = "org.apache.rocketmq"
+    module = 'rocketmq-client'
+    versions = "[4.8.0,)"
+  }
+}
+dependencies {
+  library group: 'org.apache.rocketmq', name: 'rocketmq-client', version: '4.8.0'
+  implementation project(':instrumentation:rocketmq-client-4.8:library')
+  testImplementation project(':instrumentation:rocketmq-client-4.8:testing')
+
+}
+

instrumentation/rocketmq-client-4.8/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/rocketmq/RocketMqClientApiImplInstrumentation.java

@@ -0,0 +1,104 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.rocketmq;
+
+import static io.opentelemetry.instrumentation.rocketmq.RocketMqProducerTracer.tracer;
+import static io.opentelemetry.instrumentation.rocketmq.TextMapInjectAdapter.SETTER;
+import static java.util.Collections.singletonMap;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+import io.opentelemetry.instrumentation.rocketmq.SendCallbackWrapper;
+import io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge;
+import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import org.apache.rocketmq.client.producer.SendCallback;
+import org.apache.rocketmq.common.message.Message;
+import org.apache.rocketmq.common.protocol.header.SendMessageRequestHeader;
+
+public class RocketMqClientApiImplInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("org.apache.rocketmq.client.impl.MQClientAPIImpl");
+  }
+
+  @Override
+  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    return singletonMap(
+        isMethod().and(named("sendMessage")).and(takesArguments(12)),
+        RocketMqClientApiImplInstrumentation.class.getName() + "$SendMessageAdvice");
+  }
+
+  public static class SendMessageAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void onEnter(
+        @Advice.Argument(value = 0, readOnly = false) String addr,
+        @Advice.Argument(value = 2, readOnly = false) Message msg,
+        @Advice.Argument(value = 3, readOnly = false) SendMessageRequestHeader requestHeader,
+        @Advice.Argument(value = 6, readOnly = false) SendCallback sendCallback,
+        @Advice.Local("otelSpan") Span span,
+        @Advice.Local("otelScope") Scope scope) {
+
+      Context parent = Java8BytecodeBridge.currentContext();
+      span = tracer().startProducerSpan(addr, msg);
+      Context newContext = parent.with(span);
+      try {
+        Java8BytecodeBridge.getGlobalPropagators()
+            .getTextMapPropagator()
+            .inject(newContext, requestHeader, SETTER);
+      } catch (IllegalStateException e) {
+        requestHeader = new SendMessageRequestHeader();
+        requestHeader.getBornTimestamp();
+        requestHeader.getDefaultTopic();
+        requestHeader.getDefaultTopicQueueNums();
+        requestHeader.getFlag();
+        requestHeader.getProducerGroup();
+        requestHeader.getMaxReconsumeTimes();
+        requestHeader.getProperties();
+        requestHeader.getSysFlag();
+        requestHeader.getTopic();
+        requestHeader.getQueueId();
+        requestHeader.getReconsumeTimes();
+        Java8BytecodeBridge.getGlobalPropagators()
+            .getTextMapPropagator()
+            .inject(newContext, requestHeader, SETTER);
+      }
+
+      scope = newContext.makeCurrent();
+      if (sendCallback != null) {
+        sendCallback = new SendCallbackWrapper(sendCallback, span);
+      }
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void stopSpan(
+        @Advice.Thrown Throwable throwable,
+        @Advice.Argument(value = 6, readOnly = false) SendCallback sendCallback,
+        @Advice.Local("otelSpan") Span span,
+        @Advice.Local("otelScope") Scope scope) {
+      if (scope == null) {
+        return;
+      }
+      scope.close();
+      if (sendCallback == null) {
+        if (throwable == null) {
+          tracer().end(span);
+        } else {
+          tracer().endExceptionally(span, throwable);
+        }
+      }
+    }
+  }
+}

instrumentation/rocketmq-client-4.8/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/rocketmq/RocketMqConcurrentlyConsumeInstrumentation.java

@@ -0,0 +1,70 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.rocketmq;
+
+import static io.opentelemetry.instrumentation.rocketmq.RocketMqConsumerTracer.tracer;
+import static io.opentelemetry.javaagent.tooling.bytebuddy.matcher.AgentElementMatchers.implementsInterface;
+import static java.util.Collections.singletonMap;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.context.Scope;
+import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
+import java.util.List;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import org.apache.rocketmq.client.consumer.listener.ConsumeConcurrentlyStatus;
+import org.apache.rocketmq.common.message.MessageExt;
+
+public class RocketMqConcurrentlyConsumeInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return implementsInterface(
+        named("org.apache.rocketmq.client.consumer.listener.MessageListenerConcurrently"));
+  }
+
+  @Override
+  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    return singletonMap(
+        isMethod().and(named("consumeMessage")),
+        RocketMqConcurrentlyConsumeInstrumentation.class.getName() + "$ConcurrentlyConsumeAdvice");
+  }
+
+  public static class ConcurrentlyConsumeAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void onEnter(
+        @Advice.Argument(0) List<MessageExt> msgs,
+        @Advice.Local("otelSpan") Span span,
+        @Advice.Local("otelScope") Scope scope) {
+
+      span = tracer().startSpan(msgs);
+      scope = span.makeCurrent();
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void stopSpan(
+        @Advice.Return ConsumeConcurrentlyStatus status,
+        @Advice.Thrown Throwable throwable,
+        @Advice.Local("otelSpan") Span span,
+        @Advice.Local("otelScope") Scope scope) {
+      if (scope == null) {
+        return;
+      }
+      tracer().endConcurrentlySpan(span, status);
+      scope.close();
+      if (throwable == null) {
+        tracer().end(span);
+      } else {
+        tracer().endExceptionally(span, throwable);
+      }
+    }
+  }
+}

instrumentation/rocketmq-client-4.8/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/rocketmq/RocketMqInstrumentationModule.java

@@ -0,0 +1,28 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.rocketmq;
+
+import static java.util.Arrays.asList;
+
+import com.google.auto.service.AutoService;
+import io.opentelemetry.javaagent.tooling.InstrumentationModule;
+import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
+import java.util.List;
+
+@AutoService(InstrumentationModule.class)
+public class RocketMqInstrumentationModule extends InstrumentationModule {
+  public RocketMqInstrumentationModule() {
+    super("rocketmq", "rockemq-client-4.3");
+  }
+
+  @Override
+  public List<TypeInstrumentation> typeInstrumentations() {
+    return asList(
+        new RocketMqClientApiImplInstrumentation(),
+        new RocketMqConcurrentlyConsumeInstrumentation(),
+        new RocketMqOrderlyConsumeInstrumentation());
+  }
+}

instrumentation/rocketmq-client-4.8/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/rocketmq/RocketMqOrderlyConsumeInstrumentation.java

@@ -0,0 +1,70 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.rocketmq;
+
+import static io.opentelemetry.instrumentation.rocketmq.RocketMqConsumerTracer.tracer;
+import static io.opentelemetry.javaagent.tooling.bytebuddy.matcher.AgentElementMatchers.implementsInterface;
+import static java.util.Collections.singletonMap;
+import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.context.Scope;
+import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
+import java.util.List;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import org.apache.rocketmq.client.consumer.listener.ConsumeOrderlyStatus;
+import org.apache.rocketmq.common.message.MessageExt;
+
+public class RocketMqOrderlyConsumeInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return implementsInterface(
+        named("org.apache.rocketmq.client.consumer.listener.MessageListenerOrderly"));
+  }
+
+  @Override
+  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    return singletonMap(
+        nameStartsWith("consumeMessage"),
+        RocketMqOrderlyConsumeInstrumentation.class.getName() + "$OrderlyConsumeAdvice");
+  }
+
+  public static class OrderlyConsumeAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void onEnter(
+        @Advice.Argument(0) List<MessageExt> msgs,
+        @Advice.Local("otelSpan") Span span,
+        @Advice.Local("otelScope") Scope scope) {
+
+      span = tracer().startSpan(msgs);
+      scope = span.makeCurrent();
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void stopSpan(
+        @Advice.Return ConsumeOrderlyStatus status,
+        @Advice.Thrown Throwable throwable,
+        @Advice.Local("otelSpan") Span span,
+        @Advice.Local("otelScope") Scope scope) {
+      if (scope == null) {
+        return;
+      }
+      tracer().endOrderlySpan(span, status);
+      scope.close();
+      if (throwable == null) {
+        tracer().end(span);
+      } else {
+        tracer().endExceptionally(span, throwable);
+      }
+    }
+  }
+}

instrumentation/rocketmq-client-4.8/javaagent/src/test/groovy/io/opentelemetry/instrumentation/rocketmq/RocketMqClientTest.groovy

@@ -0,0 +1,14 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.rocketmq
+
+import io.opentelemetery.instrumentation.rocketmq.AbstractRocketMqClientTest
+import io.opentelemetry.instrumentation.test.AgentTestTrait
+
+class RocketMqClientTest extends AbstractRocketMqClientTest implements AgentTestTrait {
+
+
+}

instrumentation/rocketmq-client-4.8/library/rocketmq-client-4.8-library.gradle

@@ -0,0 +1,8 @@
+apply from: "$rootDir/gradle/instrumentation-library.gradle"
+
+
+dependencies {
+  library group: 'org.apache.rocketmq', name: 'rocketmq-client', version: '4.8.0'
+  testImplementation project(':instrumentation:rocketmq-client-4.8:testing')
+
+}

instrumentation/rocketmq-client-4.8/library/src/main/java/io/opentelemetry/instrumentation/rocketmq/RocketMqClientConfig.java

@@ -0,0 +1,18 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.rocketmq;
+
+import io.opentelemetry.instrumentation.api.config.Config;
+
+public final class RocketMqClientConfig {
+
+  public static boolean isPropagationEnabled() {
+    return Config.get()
+        .getBooleanProperty("otel.instrumentation.rocketmq.client-propagation", true);
+  }
+
+  private RocketMqClientConfig() {}
+}