Add support for distributed tracing

With Micrometer Observation and Tracing.

Author: acogoluegnes

pom.xml

@@ -45,11 +45,12 @@
         <slf4j.version>1.7.36</slf4j.version>
         <logback.version>1.2.13</logback.version>
         <protonj2.version>1.0.0-M21-SNAPSHOT</protonj2.version>
+        <micrometer.version>1.13.0</micrometer.version>
         <junit.jupiter.version>5.10.2</junit.jupiter.version>
         <assertj.version>3.26.0</assertj.version>
         <mockito.version>5.12.0</mockito.version>
         <amqp-client.version>5.20.0</amqp-client.version>
-        <micrometer.version>1.13.0</micrometer.version>
+        <micrometer-tracing-test.version>1.2.5</micrometer-tracing-test.version>
         <maven.compiler.plugin.version>3.13.0</maven.compiler.plugin.version>
         <maven-surefire-plugin.version>3.2.5</maven-surefire-plugin.version>
         <maven-resources-plugin.version>3.3.1</maven-resources-plugin.version>
@@ -166,6 +167,14 @@
             <scope>test</scope>
         </dependency>
 
+        <dependency>
+            <groupId>io.micrometer</groupId>
+            <artifactId>micrometer-tracing-integration-test</artifactId>
+            <version>${micrometer-tracing-test.version}</version>
+            <scope>test</scope>
+            <optional>true</optional>
+        </dependency>
+
         <dependency>
             <groupId>com.github.spotbugs</groupId>
             <artifactId>spotbugs-annotations</artifactId>

src/main/java/com/rabbitmq/model/Message.java

@@ -19,6 +19,7 @@
 
 import java.math.BigDecimal;
 import java.util.UUID;
+import java.util.function.BiConsumer;
 
 public interface Message {
 
@@ -136,6 +137,12 @@ public interface Message {
   // TODO support iteration over message application properties
 
   // TODO support message annotations
+  Object annotation(String key);
+
+  Message annotation(String key, String value);
+
+  Message forEachAnnotation(BiConsumer<String, Object> action);
+
   // TODO support message headers
 
   MessageAddressBuilder toAddress();

src/main/java/com/rabbitmq/model/ObservationCollector.java

@@ -0,0 +1,39 @@
+// Copyright (c) 2024 Broadcom. All Rights Reserved.
+// The term "Broadcom" refers to Broadcom Inc. and/or its subsidiaries.
+//
+// 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.
+//
+// If you have any questions regarding licensing, please contact us at
+// [email protected].
+package com.rabbitmq.model;
+
+import java.util.function.Function;
+
+public interface ObservationCollector {
+
+  <T> T publish(
+      String exchange,
+      String routingKey,
+      Message message,
+      ConnectionInfo connectionInfo,
+      Function<Message, T> publishCall);
+
+  Consumer.MessageHandler subscribe(String queue, Consumer.MessageHandler handler);
+
+  interface ConnectionInfo {
+
+    String peerAddress();
+
+    int peerPort();
+  }
+}

src/main/java/com/rabbitmq/model/amqp/AmqpConnection.java

@@ -23,6 +23,7 @@
 import static java.util.Collections.singletonMap;
 
 import com.rabbitmq.model.*;
+import com.rabbitmq.model.ObservationCollector;
 import com.rabbitmq.model.metrics.MetricsCollector;
 import java.time.Duration;
 import java.util.ArrayList;
@@ -479,6 +480,10 @@ MetricsCollector metricsCollector() {
     return this.environment.metricsCollector();
   }
 
+  ObservationCollector observationCollector() {
+    return this.environment.observationCollector();
+  }
+
   Publisher createPublisher(AmqpPublisherBuilder builder) {
     // TODO copy the builder properties to create the publisher
     AmqpPublisher publisher = new AmqpPublisher(builder);
@@ -544,6 +549,10 @@ private String currentConnectionLabel() {
     }
   }
 
+  Address connectionAddress() {
+    return this.connectionAddress;
+  }
+
   private void close(Throwable cause) {
     if (this.closed.compareAndSet(false, true)) {
       this.state(CLOSING, cause);

src/main/java/com/rabbitmq/model/amqp/AmqpConsumer.java

@@ -72,7 +72,11 @@ final class AmqpConsumer extends ResourceBase implements Consumer {
     super(builder.listeners());
     this.id = ID_SEQUENCE.getAndIncrement();
     this.initialCredits = builder.initialCredits();
-    this.messageHandler = builder.messageHandler();
+    this.messageHandler =
+        builder
+            .connection()
+            .observationCollector()
+            .subscribe(builder.queue(), builder.messageHandler());
     this.address = "/queue/" + builder.queue();
     this.nativeReceiver = createNativeReceiver(builder.connection().nativeSession(), this.address);
     this.initStateFromNativeReceiver(this.nativeReceiver);

src/main/java/com/rabbitmq/model/amqp/AmqpEnvironment.java

@@ -18,6 +18,7 @@
 package com.rabbitmq.model.amqp;
 
 import com.rabbitmq.model.*;
+import com.rabbitmq.model.ObservationCollector;
 import com.rabbitmq.model.metrics.MetricsCollector;
 import com.rabbitmq.model.metrics.NoOpMetricsCollector;
 import java.util.ArrayList;
@@ -49,11 +50,13 @@ class AmqpEnvironment implements Environment {
   private volatile ScheduledFuture<?> clockRefreshFuture;
   private final AtomicBoolean clockRefreshSet = new AtomicBoolean(false);
   private final MetricsCollector metricsCollector;
+  private final ObservationCollector observationCollector;
 
   AmqpEnvironment(
       ExecutorService executorService,
       DefaultConnectionSettings<?> connectionSettings,
-      MetricsCollector metricsCollector) {
+      MetricsCollector metricsCollector,
+      ObservationCollector observationCollector) {
     this.id = ID_SEQUENCE.getAndIncrement();
     connectionSettings.copyTo(this.connectionSettings);
     this.connectionSettings.consolidate();
@@ -71,6 +74,8 @@ class AmqpEnvironment implements Environment {
         Executors.newScheduledThreadPool(0, Utils.defaultThreadFactory());
     this.metricsCollector =
         metricsCollector == null ? NoOpMetricsCollector.INSTANCE : metricsCollector;
+    this.observationCollector =
+        observationCollector == null ? Utils.NO_OP_OBSERVATION_COLLECTOR : observationCollector;
   }
 
   DefaultConnectionSettings<?> connectionSettings() {
@@ -128,6 +133,10 @@ MetricsCollector metricsCollector() {
     return this.metricsCollector;
   }
 
+  ObservationCollector observationCollector() {
+    return this.observationCollector;
+  }
+
   void addConnection(AmqpConnection connection) {
     this.connections.add(connection);
   }

src/main/java/com/rabbitmq/model/amqp/AmqpEnvironmentBuilder.java

@@ -20,6 +20,7 @@
 import com.rabbitmq.model.ConnectionSettings;
 import com.rabbitmq.model.Environment;
 import com.rabbitmq.model.EnvironmentBuilder;
+import com.rabbitmq.model.ObservationCollector;
 import com.rabbitmq.model.metrics.MetricsCollector;
 import com.rabbitmq.model.metrics.NoOpMetricsCollector;
 import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
@@ -31,6 +32,7 @@ public class AmqpEnvironmentBuilder implements EnvironmentBuilder {
       new DefaultEnvironmentConnectionSettings(this);
   private ExecutorService executorService;
   private MetricsCollector metricsCollector = NoOpMetricsCollector.INSTANCE;
+  private ObservationCollector observationCollector = Utils.NO_OP_OBSERVATION_COLLECTOR;
 
   public AmqpEnvironmentBuilder() {}
 
@@ -44,14 +46,20 @@ public AmqpEnvironmentBuilder metricsCollector(MetricsCollector metricsCollector
     return this;
   }
 
+  public AmqpEnvironmentBuilder observationCollector(ObservationCollector observationCollector) {
+    this.observationCollector = observationCollector;
+    return this;
+  }
+
   @SuppressFBWarnings("EI_EXPOSE_REP")
   public EnvironmentConnectionSettings connectionSettings() {
     return this.connectionSettings;
   }
 
   @Override
   public Environment build() {
-    return new AmqpEnvironment(executorService, connectionSettings, metricsCollector);
+    return new AmqpEnvironment(
+        executorService, connectionSettings, metricsCollector, observationCollector);
   }
 
   public interface EnvironmentConnectionSettings

src/main/java/com/rabbitmq/model/amqp/AmqpMessage.java

@@ -179,13 +179,13 @@ public Message correlationId(String correlationId) {
 
   @Override
   public Message correlationId(long correlationId) {
-    callOnDelegate(m -> m.correlationId(correlationId));
+    callOnDelegate(m -> m.correlationId(UnsignedLong.valueOf(correlationId)));
     return this;
   }
 
   @Override
   public Message correlationId(byte[] correlationId) {
-    callOnDelegate(m -> m.correlationId(correlationId));
+    callOnDelegate(m -> m.correlationId(new Binary(correlationId)));
     return this;
   }
 
@@ -354,6 +354,23 @@ public boolean hasProperties() {
     return returnFromDelegate(org.apache.qpid.protonj2.client.Message::hasProperties);
   }
 
+  @Override
+  public Object annotation(String key) {
+    return returnFromDelegate(m -> m.annotation(key));
+  }
+
+  @Override
+  public Message annotation(String key, String value) {
+    callOnDelegate(m -> m.annotation(key, value));
+    return this;
+  }
+
+  @Override
+  public Message forEachAnnotation(BiConsumer<String, Object> action) {
+    callOnDelegate(m -> m.forEachAnnotation(action));
+    return this;
+  }
+
   @Override
   public Object removeProperty(String key) {
     return returnFromDelegate(m -> m.removeProperty(key));

src/main/java/com/rabbitmq/model/amqp/AmqpPublisher.java

@@ -20,11 +20,13 @@
 import static com.rabbitmq.model.Resource.State.OPEN;
 
 import com.rabbitmq.model.Message;
+import com.rabbitmq.model.ObservationCollector;
 import com.rabbitmq.model.Publisher;
 import com.rabbitmq.model.metrics.MetricsCollector;
 import java.util.concurrent.*;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Function;
 import org.apache.qpid.protonj2.client.*;
 import org.apache.qpid.protonj2.client.exceptions.ClientException;
 import org.apache.qpid.protonj2.client.exceptions.ClientIllegalStateException;
@@ -44,17 +46,38 @@ final class AmqpPublisher extends ResourceBase implements Publisher {
   private final AmqpConnection connection;
   private final AtomicBoolean closed = new AtomicBoolean(false);
   private final MetricsCollector metricsCollector;
+  private final ObservationCollector observationCollector;
+  private final Function<Message, Tracker> publishCall;
+  private final DefaultAddressBuilder.DestinationSpec destinationSpec;
+  private volatile ObservationCollector.ConnectionInfo connectionInfo;
 
   AmqpPublisher(AmqpPublisherBuilder builder) {
     super(builder.listeners());
     this.id = ID_SEQUENCE.getAndIncrement();
     this.executorService = builder.connection().executorService();
     this.address = builder.address();
+    this.destinationSpec = builder.destination();
     this.connection = builder.connection();
     this.sender = this.createSender(builder.connection().nativeSession(), this.address);
     this.metricsCollector = this.connection.metricsCollector();
+    this.observationCollector = this.connection.observationCollector();
     this.state(OPEN);
     this.metricsCollector.openPublisher();
+    this.publishCall =
+        msg -> {
+          try {
+            org.apache.qpid.protonj2.client.Message<?> nativeMessage =
+                ((AmqpMessage) msg).nativeMessage();
+            return this.sender.send(nativeMessage.durable(true));
+          } catch (ClientIllegalStateException e) {
+            // the link is closed
+            this.close(ExceptionUtils.convert(e));
+            throw ExceptionUtils.convert(e);
+          } catch (ClientException e) {
+            throw ExceptionUtils.convert(e);
+          }
+        };
+    this.connectionInfo = new Utils.ObservationConnectionInfo(this.connection.connectionAddress());
   }
 
   @Override
@@ -70,40 +93,35 @@ public Message message(byte[] body) {
   @Override
   public void publish(Message message, Callback callback) {
     checkOpen();
-    try {
-      org.apache.qpid.protonj2.client.Message<?> nativeMessage =
-          ((AmqpMessage) message).nativeMessage();
-      Tracker tracker = this.sender.send(nativeMessage.durable(true));
-      this.executorService.submit(
-          () -> {
-            Status status;
-            try {
-              tracker.settlementFuture().get();
-              status =
-                  tracker.remoteState() == DeliveryState.accepted()
-                      ? Status.ACCEPTED
-                      : Status.FAILED;
-            } catch (InterruptedException | ExecutionException e) {
-              status = Status.FAILED;
-            }
-            DefaultContext defaultContext = new DefaultContext(message, status);
-            this.metricsCollector.publishDisposition(
-                status == Status.ACCEPTED
-                    ? MetricsCollector.PublishDisposition.ACCEPTED
-                    : MetricsCollector.PublishDisposition.FAILED);
-            callback.handle(defaultContext);
-          });
-      this.metricsCollector.publish();
-    } catch (ClientIllegalStateException e) {
-      // the link is closed
-      this.close(ExceptionUtils.convert(e));
-      throw ExceptionUtils.convert(e);
-    } catch (ClientException e) {
-      throw ExceptionUtils.convert(e);
-    }
+    Tracker tracker =
+        this.observationCollector.publish(
+            destinationSpec.exchange(),
+            destinationSpec.routingKey(),
+            message,
+            this.connectionInfo,
+            publishCall);
+    this.executorService.submit(
+        () -> {
+          Status status;
+          try {
+            tracker.settlementFuture().get();
+            status =
+                tracker.remoteState() == DeliveryState.accepted() ? Status.ACCEPTED : Status.FAILED;
+          } catch (InterruptedException | ExecutionException e) {
+            status = Status.FAILED;
+          }
+          DefaultContext defaultContext = new DefaultContext(message, status);
+          this.metricsCollector.publishDisposition(
+              status == Status.ACCEPTED
+                  ? MetricsCollector.PublishDisposition.ACCEPTED
+                  : MetricsCollector.PublishDisposition.FAILED);
+          callback.handle(defaultContext);
+        });
+    this.metricsCollector.publish();
   }
 
   void recoverAfterConnectionFailure() {
+    this.connectionInfo = new Utils.ObservationConnectionInfo(this.connection.connectionAddress());
     this.sender = this.createSender(this.connection.nativeSession(false), this.address);
   }
 

src/main/java/com/rabbitmq/model/amqp/AmqpPublisherBuilder.java

@@ -80,4 +80,8 @@ AddressBuilder<?> addressBuilder() {
   String address() {
     return this.addressBuilder.address();
   }
+
+  DefaultAddressBuilder.DestinationSpec destination() {
+    return this.addressBuilder.destination();
+  }
 }