adds RocketMQTracing

Signed-off-by: Adrian Cole <adrian@tetrate.io>

brave-bom/pom.xml

@@ -226,6 +226,16 @@
         <artifactId>brave-instrumentation-okhttp3</artifactId>
         <version>${project.version}</version>
       </dependency>
+      <dependency>
+        <groupId>${project.groupId}</groupId>
+        <artifactId>brave-instrumentation-rocketmq-client</artifactId>
+        <version>${project.version}</version>
+      </dependency>
+      <dependency>
+        <groupId>${project.groupId}</groupId>
+        <artifactId>brave-instrumentation-rocketmq-clients</artifactId>
+        <version>${project.version}</version>
+      </dependency>
       <dependency>
         <groupId>${project.groupId}</groupId>
         <artifactId>brave-instrumentation-rpc</artifactId>

instrumentation/README.md

@@ -20,6 +20,7 @@ Here's a brief overview of what's packaged here:
 * [mysql8](mysql8/README.md) - Tracing MySQL v8 statement interceptor
 * [netty-codec-http](netty-codec-http/README.md) - Tracing handler for [Netty](http://netty.io/) 4.x http servers
 * [okhttp3](okhttp3/README.md) - Tracing decorators for [OkHttp](https://github.com/square/okhttp) 3.x
+* [rocketmq-client](rocketmq-client/README.md) - Tracing decorators for RocketMQ producers and consumers.
 * [servlet](servlet/README.md) - Tracing filter for Servlet 2.5+ (including Async)
 * [spring-rabbit](spring-rabbit/README.md) - Tracing MessagePostProcessor and ListenerAdvice for [Spring Rabbit](https://spring.io/guides/gs/messaging-rabbitmq/)
 * [spring-web](spring-web/README.md) - Tracing interceptor for [Spring RestTemplate](https://spring.io/guides/gs/consuming-rest/)

instrumentation/pom.xml

@@ -51,6 +51,7 @@
     <module>mysql6</module>
     <module>mysql8</module>
     <module>okhttp3</module>
+    <module>rocketmq-client</module>
     <module>rpc</module>
     <module>servlet</module>
     <module>servlet-jakarta</module>

instrumentation/rocketmq-client/README.md

@@ -0,0 +1,100 @@
+# brave-instrumentation-rocketmq-client
+
+## Tracing for RocketMQ Client
+
+This module provides instrumentation for RocketMQ based services.
+
+## example
+
+### producer
+
+The key is to register our hook to the producer
+
+```java
+package brave.rocketmq.client;
+
+import brave.Tracing;
+import brave.messaging.MessagingRequest;
+import brave.messaging.MessagingTracing;
+import brave.sampler.SamplerFunction;
+import brave.sampler.SamplerFunctions;
+import org.apache.rocketmq.client.producer.DefaultMQProducer;
+import org.apache.rocketmq.common.message.Message;
+
+public class ProducerExample {
+
+  public static void main(String[] args) throws Exception {
+    // todo Replaced with actual tracing construct
+    Tracing tracing = Tracing.newBuilder().build();
+    SamplerFunction<MessagingRequest> producerSampler = SamplerFunctions.deferDecision();
+    RocketMQTracing producerTracing = RocketMQTracing.create(
+      MessagingTracing.newBuilder(tracing).producerSampler(producerSampler).build());
+
+    String topic = "testSend";
+    Message message = new Message(topic, "JoeKerouac", "hello".getBytes());
+    DefaultMQProducer producer = new DefaultMQProducer("testSend");
+    // todo This is the key, register the hook to the producer
+    producer.getDefaultMQProducerImpl()
+      .registerSendMessageHook(new SendMessageBraveHookImpl(producerTracing));
+    // Replace with actual address
+    producer.setNamesrvAddr("127.0.0.1:9876");
+    producer.start();
+    producer.send(message);
+
+    producer.shutdown();
+  }
+}
+
+```
+
+### consumer
+
+Replace `org.apache.rocketmq.client.consumer.listener.MessageListenerConcurrently`
+with `brave.rocketmq.client.TracingMessageListenerConcurrently`
+or `org.apache.rocketmq.client.consumer.listener.MessageListenerOrderly`
+with `brave.rocketmq.client.TracingMessageListenerOrderly`;
+
+```java
+package brave.rocketmq.client;
+
+import brave.Span;
+import brave.Tracer;
+import brave.Tracing;
+import brave.messaging.MessagingRequest;
+import brave.messaging.MessagingTracing;
+import brave.sampler.SamplerFunction;
+import brave.sampler.SamplerFunctions;
+import org.apache.rocketmq.client.consumer.DefaultMQPushConsumer;
+import org.apache.rocketmq.common.message.MessageExt;
+
+import java.util.Optional;
+
+public class ProducerExample {
+
+  public static void main(String[] args) throws Exception {
+    // todo Replaced with actual tracing construct
+    Tracing tracing = Tracing.newBuilder().build();
+    SamplerFunction<MessagingRequest> producerSampler = SamplerFunctions.deferDecision();
+    RocketMQTracing producerTracing = RocketMQTracing.create(
+      MessagingTracing.newBuilder(tracing).producerSampler(producerSampler).build());
+
+    String topic = "testPushConsumer";
+    String nameserverAddr = "127.0.0.1:9876";
+
+    DefaultMQPushConsumer consumer = new DefaultMQPushConsumer("testPushConsumer");
+    consumer.setNamesrvAddr(nameserverAddr);
+    consumer.subscribe(topic, "*");
+    consumer.registerMessageListener(new TraceableMessageListenerConcurrently(0, producerTracing) {
+      @Override
+      protected void handleMessage(MessageExt messageExt) {
+        Span span =
+          Optional.ofNullable(Tracing.currentTracer()).map(Tracer::currentSpan).orElse(null);
+        // do something
+      }
+    });
+    consumer.start();
+  }
+}
+
+```
+

instrumentation/rocketmq-client/bnd.bnd

@@ -0,0 +1,6 @@
+# We use brave.internal.Nullable, but it is not used at runtime.
+Import-Package: \
+  !brave.internal*,\
+  *
+Export-Package: \
+  brave.rocketmq.client

instrumentation/rocketmq-client/pom.xml

@@ -0,0 +1,68 @@
+<?xml version="1.0"?>
+<!--
+
+    Copyright 2013-2024 The OpenZipkin Authors
+
+    Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
+    in compliance with the License. You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+    Unless required by applicable law or agreed to in writing, software distributed under the License
+    is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+    or implied. See the License for the specific language governing permissions and limitations under
+    the License.
+
+-->
+<project xmlns="http://maven.apache.org/POM/4.0.0"
+  xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+  xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+  <modelVersion>4.0.0</modelVersion>
+  <parent>
+    <groupId>io.zipkin.brave</groupId>
+    <artifactId>brave-instrumentation-parent</artifactId>
+    <version>6.0.1-SNAPSHOT</version>
+  </parent>
+
+  <artifactId>brave-instrumentation-rocketmq-client</artifactId>
+  <name>Brave Instrumentation: RocketMQ Client</name>
+
+  <properties>
+    <!-- Matches Export-Package in bnd.bnd -->
+    <module.name>brave.rocketmq.client</module.name>
+
+    <main.basedir>${project.basedir}/../..</main.basedir>
+
+    <rocketmq.version>4.7.0</rocketmq.version>
+
+    <maven-failsafe-plugin.argLine>--add-opens java.base/java.nio=ALL-UNNAMED</maven-failsafe-plugin.argLine>
+  </properties>
+
+  <dependencies>
+    <dependency>
+      <groupId>${project.groupId}</groupId>
+      <artifactId>brave-instrumentation-messaging</artifactId>
+      <version>${project.version}</version>
+    </dependency>
+    <dependency>
+      <groupId>org.apache.rocketmq</groupId>
+      <artifactId>rocketmq-client</artifactId>
+      <version>${rocketmq.version}</version>
+      <scope>provided</scope>
+    </dependency>
+
+    <dependency>
+      <groupId>${project.groupId}</groupId>
+      <artifactId>brave-tests</artifactId>
+      <scope>test</scope>
+      <version>${project.version}</version>
+    </dependency>
+
+    <dependency>
+      <groupId>org.testcontainers</groupId>
+      <artifactId>junit-jupiter</artifactId>
+      <version>${testcontainers.version}</version>
+      <scope>test</scope>
+    </dependency>
+  </dependencies>
+</project>

instrumentation/rocketmq-client/src/main/java/brave/rocketmq/client/RocketMQTracing.java

@@ -0,0 +1,109 @@
+/*
+ * Copyright 2013-2024 The OpenZipkin Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
+ * in compliance with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the License
+ * is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+ * or implied. See the License for the specific language governing permissions and limitations under
+ * the License.
+ */
+package brave.rocketmq.client;
+
+import brave.Span;
+import brave.Tracer;
+import brave.Tracing;
+import brave.messaging.MessagingRequest;
+import brave.messaging.MessagingTracing;
+import brave.propagation.Propagation;
+import brave.propagation.TraceContext.Extractor;
+import brave.propagation.TraceContext.Injector;
+import brave.propagation.TraceContextOrSamplingFlags;
+import brave.sampler.SamplerFunction;
+import java.util.Map;
+
+public class RocketMQTracing {
+  public static RocketMQTracing create(Tracing tracing) {
+    return new RocketMQTracing(MessagingTracing.create(tracing), TraceConstants.ROCKETMQ_SERVICE);
+  }
+
+  public static RocketMQTracing create(MessagingTracing messagingTracing) {
+    return new RocketMQTracing(messagingTracing, TraceConstants.ROCKETMQ_SERVICE);
+  }
+
+  public static RocketMQTracing create(MessagingTracing messagingTracing,
+    String remoteServiceName) {
+    return new RocketMQTracing(messagingTracing, remoteServiceName);
+  }
+
+  final Tracing tracing;
+  final Tracer tracer;
+  final Extractor<TracingProducerRequest> producerExtractor;
+  final Extractor<TracingConsumerRequest> consumerExtractor;
+  final Injector<TracingProducerRequest> producerInjector;
+  final Injector<TracingConsumerRequest> consumerInjector;
+  final String[] traceIdHeaders;
+  final SamplerFunction<MessagingRequest> producerSampler, consumerSampler;
+  final String remoteServiceName;
+
+  RocketMQTracing(MessagingTracing messagingTracing,
+    String remoteServiceName) { // intentionally hidden constructor
+    this.tracing = messagingTracing.tracing();
+    this.tracer = tracing.tracer();
+    Propagation<String> propagation = messagingTracing.propagation();
+    this.producerExtractor = propagation.extractor(TracingProducerRequest.GETTER);
+    this.consumerExtractor = propagation.extractor(TracingConsumerRequest.GETTER);
+    this.producerInjector = propagation.injector(TracingProducerRequest.SETTER);
+    this.consumerInjector = propagation.injector(TracingConsumerRequest.SETTER);
+    this.producerSampler = messagingTracing.producerSampler();
+    this.consumerSampler = messagingTracing.consumerSampler();
+    this.remoteServiceName = remoteServiceName;
+
+    // We clear the trace ID headers, so that a stale consumer span is not preferred over current
+    // listener. We intentionally don't clear BaggagePropagation.allKeyNames as doing so will
+    // application fields "user_id" or "country_code"
+    this.traceIdHeaders = propagation.keys().toArray(new String[0]);
+  }
+
+  <R> TraceContextOrSamplingFlags extractAndClearTraceIdHeaders(Extractor<R> extractor,
+    R request,
+    Map<String, String> properties) {
+    TraceContextOrSamplingFlags extracted = extractor.extract(request);
+    // Clear any propagation keys present in the headers
+    if (extracted.samplingFlags() == null) { // then trace IDs were extracted
+      if (properties != null) {
+        clearTraceIdHeaders(properties);
+      }
+    }
+    return extracted;
+  }
+
+  /** Creates a potentially noop remote span representing this request. */
+  Span nextMessagingSpan(SamplerFunction<MessagingRequest> sampler, MessagingRequest request,
+    TraceContextOrSamplingFlags extracted) {
+    Boolean sampled = extracted.sampled();
+    // only recreate the context if the messaging sampler made a decision
+    if (sampled == null && (sampled = sampler.trySample(request)) != null) {
+      extracted = extracted.sampled(sampled);
+    }
+    return tracer.nextSpan(extracted);
+  }
+
+  // We can't just skip clearing headers we use because we might inject B3 single, yet have stale B3
+  // multi, or vice versa.
+  void clearTraceIdHeaders(Map<String, String> headers) {
+    for (String traceIDHeader : traceIdHeaders)
+      headers.remove(traceIDHeader);
+  }
+
+  public Tracing tracing() {
+    return tracing;
+  }
+
+  public Tracer tracer() {
+    return tracer;
+  }
+}

instrumentation/rocketmq-client/src/main/java/brave/rocketmq/client/SpanUtil.java

@@ -0,0 +1,49 @@
+/*
+ * Copyright 2013-2024 The OpenZipkin Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
+ * in compliance with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the License
+ * is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+ * or implied. See the License for the specific language governing permissions and limitations under
+ * the License.
+ */
+package brave.rocketmq.client;
+
+import brave.Span;
+import brave.Tracer;
+import brave.messaging.MessagingRequest;
+import brave.propagation.CurrentTraceContext;
+import brave.propagation.TraceContext;
+import brave.propagation.TraceContextOrSamplingFlags;
+import brave.sampler.SamplerFunction;
+import java.util.Map;
+
+class SpanUtil {
+  static <T extends MessagingRequest> Span createAndStartSpan(RocketMQTracing tracing,
+    TraceContext.Extractor<T> extractor, SamplerFunction<MessagingRequest> sampler, T request,
+    Map<String, String> props) {
+    Tracer tracer = tracing.tracer;
+    CurrentTraceContext currentTraceContext = tracing.tracing.currentTraceContext();
+    TraceContext traceContext = currentTraceContext.get();
+    Span span;
+
+    if (traceContext == null) {
+      TraceContextOrSamplingFlags extracted =
+        tracing.extractAndClearTraceIdHeaders(extractor, request, props);
+      span = tracing.nextMessagingSpan(sampler, request, extracted);
+    } else {
+      span = tracer.newChild(traceContext);
+    }
+
+    span.kind(request.spanKind());
+    span.remoteServiceName(tracing.remoteServiceName);
+    span.tag(TraceConstants.ROCKETMQ_TOPIC, request.channelName());
+    long timestamp = tracing.tracing.clock(span.context()).currentTimeMicroseconds();
+    span.start(timestamp);
+    return span;
+  }
+}

instrumentation/rocketmq-client/src/main/java/brave/rocketmq/client/StringUtils.java

@@ -0,0 +1,26 @@
+/*
+ * Copyright 2013-2024 The OpenZipkin Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
+ * in compliance with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the License
+ * is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+ * or implied. See the License for the specific language governing permissions and limitations under
+ * the License.
+ */
+package brave.rocketmq.client;
+
+class StringUtils {
+
+  // TODO: we shouldn't add tags with empty values!
+  static String getOrEmpty(String obj) {
+    if (obj == null) {
+      return "";
+    } else {
+      return obj;
+    }
+  }
+}