Add support for observability

WIP

References #379

Author: acogoluegnes

pom.xml

@@ -74,6 +74,7 @@
     <gson.version>2.10.1</gson.version>
     <vavr.version>0.10.4</vavr.version>
     <paho.version>1.2.5</paho.version>
+    <micrometer-tracing-test.version>1.1.3</micrometer-tracing-test.version>
     <maven.compiler.plugin.version>3.11.0</maven.compiler.plugin.version>
     <maven-surefire-plugin.version>3.1.2</maven-surefire-plugin.version>
     <spring-boot-maven-plugin.version>2.7.6</spring-boot-maven-plugin.version>
@@ -281,7 +282,7 @@
       <groupId>org.eclipse.paho</groupId>
       <artifactId>org.eclipse.paho.client.mqttv3</artifactId>
       <version>${paho.version}</version>
-    <scope>test</scope>
+      <scope>test</scope>
     </dependency>
 
     <dependency>
@@ -312,6 +313,13 @@
       <scope>test</scope>
     </dependency>
 
+    <dependency>
+      <groupId>io.micrometer</groupId>
+      <artifactId>micrometer-tracing-integration-test</artifactId>
+      <version>${micrometer-tracing-test.version}</version>
+      <scope>test</scope>
+    </dependency>
+
     <dependency>
       <groupId>org.openjdk.jmh</groupId>
       <artifactId>jmh-core</artifactId>
@@ -554,7 +562,7 @@
               <style>GOOGLE</style>
             </googleJavaFormat>
           </java>
-<!--          <ratchetFrom>origin/main</ratchetFrom>-->
+          <ratchetFrom>origin/main</ratchetFrom>
           <licenseHeader> <!-- specify either content or file, but not both -->
             <content>// Copyright (c) $YEAR VMware, Inc. or its affiliates.  All rights reserved.
               //

src/main/java/com/rabbitmq/stream/Codec.java

@@ -1,4 +1,4 @@
-// Copyright (c) 2020-2021 VMware, Inc. or its affiliates.  All rights reserved.
+// Copyright (c) 2020-2023 VMware, Inc. or its affiliates.  All rights reserved.
 //
 // This software, the RabbitMQ Stream Java client library, is dual-licensed under the
 // Mozilla Public License 2.0 ("MPL"), and the Apache License version 2 ("ASL").
@@ -28,6 +28,12 @@ public interface Codec {
 
   MessageBuilder messageBuilder();
 
+  MessageBuilder messageBuilder(String stream);
+
+  Codec messageBuilderListener(MessageBuilderListener listener);
+
+  Object listenerContext(Message message);
+
   class EncodedMessage {
 
     private final int size;
@@ -46,4 +52,10 @@ public int getSize() {
       return size;
     }
   }
+
+  @FunctionalInterface
+  interface MessageBuilderListener {
+
+    Object accept(String stream, MessageBuilder builder);
+  }
 }

src/main/java/com/rabbitmq/stream/EnvironmentBuilder.java

@@ -222,6 +222,8 @@ public interface EnvironmentBuilder {
    */
   EnvironmentBuilder metricsCollector(MetricsCollector metricsCollector);
 
+  EnvironmentBuilder observationCollector(ObservationCollector observationCollector);
+
   /**
    * The maximum number of producers allocated to a single connection.
    *

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

@@ -0,0 +1,39 @@
+// Copyright (c) 2023 VMware, Inc. or its affiliates.  All rights reserved.
+//
+// This software, the RabbitMQ Stream Java client library, is dual-licensed under the
+// Mozilla Public License 2.0 ("MPL"), and the Apache License version 2 ("ASL").
+// For the MPL, please see LICENSE-MPL-RabbitMQ. For the ASL,
+// please see LICENSE-APACHE2.
+//
+// This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY KIND,
+// either express or implied. See the LICENSE file for specific language governing
+// rights and limitations of this software.
+//
+// If you have any questions regarding licensing, please contact us at
+// [email protected].
+package com.rabbitmq.stream;
+
+public interface ObservationCollector {
+
+  ObservationCollector NO_OP =
+      new ObservationCollector() {
+        @Override
+        public Codec register(Codec codec) {
+          return codec;
+        }
+
+        @Override
+        public void publish(Codec codec, Message message) {}
+
+        @Override
+        public MessageHandler subscribe(MessageHandler handler) {
+          return handler;
+        }
+      };
+
+  Codec register(Codec codec);
+
+  void publish(Codec codec, Message message);
+
+  MessageHandler subscribe(MessageHandler handler);
+}

src/main/java/com/rabbitmq/stream/codec/QpidProtonCodec.java

@@ -1,4 +1,4 @@
-// Copyright (c) 2020-2022 VMware, Inc. or its affiliates.  All rights reserved.
+// Copyright (c) 2020-2023 VMware, Inc. or its affiliates.  All rights reserved.
 //
 // This software, the RabbitMQ Stream Java client library, is dual-licensed under the
 // Mozilla Public License 2.0 ("MPL"), and the Apache License version 2 ("ASL").
@@ -32,10 +32,19 @@
 import org.apache.qpid.proton.codec.WritableBuffer;
 
 public class QpidProtonCodec implements Codec {
+  private static final MessageBuilderListener NO_OP_MESSAGE_BUILDER_LISTENER = (s, mb) -> null;
   private static final Function<String, String> MESSAGE_ANNOTATIONS_STRING_KEY_EXTRACTOR = k -> k;
   private static final Function<Symbol, String> MESSAGE_ANNOTATIONS_SYMBOL_KEY_EXTRACTOR =
       Symbol::toString;
 
+  public QpidProtonCodec() {
+    this(NO_OP_MESSAGE_BUILDER_LISTENER);
+  }
+
+  public QpidProtonCodec(MessageBuilderListener messageBuilderListener) {
+    this.messageBuilderListener = messageBuilderListener;
+  }
+
   private static Map<String, Object> createApplicationProperties(
       org.apache.qpid.proton.message.Message message) {
     if (message.getApplicationProperties() != null) {
@@ -104,6 +113,22 @@ private static Object convertApplicationProperty(Object value) {
     }
   }
 
+  private final MessageBuilderListener messageBuilderListener;
+
+  @Override
+  public Codec messageBuilderListener(MessageBuilderListener listener) {
+    return new QpidProtonCodec(listener);
+  }
+
+  @Override
+  public Object listenerContext(Message message) {
+    if (message instanceof QpidProtonAmqpMessageWrapper) {
+      return ((QpidProtonAmqpMessageWrapper) message).listenerContext;
+    } else {
+      return null;
+    }
+  }
+
   @Override
   public EncodedMessage encode(Message message) {
     org.apache.qpid.proton.message.Message qpidMessage;
@@ -301,7 +326,12 @@ protected Object convertToQpidType(Object value) {
 
   @Override
   public MessageBuilder messageBuilder() {
-    return new QpidProtonMessageBuilder();
+    return this.messageBuilder(null);
+  }
+
+  @Override
+  public MessageBuilder messageBuilder(String stream) {
+    return new QpidProtonMessageBuilder(stream, this.messageBuilderListener);
   }
 
   private static final class QpidProtonProperties implements Properties {
@@ -511,17 +541,20 @@ static class QpidProtonAmqpMessageWrapper implements Message {
     private final boolean hasPublishingId;
     private final long publishingId;
     private final org.apache.qpid.proton.message.Message message;
+    private final Object listenerContext;
     private Properties properties;
     private Map<String, Object> applicationProperties;
     private Map<String, Object> messageAnnotations;
 
     QpidProtonAmqpMessageWrapper(
         boolean hasPublishingId,
         long publishingId,
-        org.apache.qpid.proton.message.Message message) {
+        org.apache.qpid.proton.message.Message message,
+        Object listenerContext) {
       this.hasPublishingId = hasPublishingId;
       this.publishingId = publishingId;
       this.message = message;
+      this.listenerContext = listenerContext;
     }
 
     @Override

src/main/java/com/rabbitmq/stream/codec/QpidProtonMessageBuilder.java

@@ -1,4 +1,4 @@
-// Copyright (c) 2020-2021 VMware, Inc. or its affiliates.  All rights reserved.
+// Copyright (c) 2020-2023 VMware, Inc. or its affiliates.  All rights reserved.
 //
 // This software, the RabbitMQ Stream Java client library, is dual-licensed under the
 // Mozilla Public License 2.0 ("MPL"), and the Apache License version 2 ("ASL").
@@ -13,6 +13,7 @@
 // [email protected].
 package com.rabbitmq.stream.codec;
 
+import com.rabbitmq.stream.Codec;
 import com.rabbitmq.stream.Message;
 import com.rabbitmq.stream.MessageBuilder;
 import java.math.BigDecimal;
@@ -33,6 +34,8 @@
 
 class QpidProtonMessageBuilder implements MessageBuilder {
 
+  private final String stream;
+  private final Codec.MessageBuilderListener listener;
   private final org.apache.qpid.proton.message.Message message =
       org.apache.qpid.proton.message.Message.Factory.create();
   private final AtomicBoolean built = new AtomicBoolean(false);
@@ -42,9 +45,15 @@ class QpidProtonMessageBuilder implements MessageBuilder {
   private QpidProtonjApplicationPropertiesBuilder applicationPropertiesBuilder;
   private QpidProtonjMessageAnnotationsBuilder messageAnnotationsBuilder;
 
+  QpidProtonMessageBuilder(String stream, Codec.MessageBuilderListener listener) {
+    this.stream = stream;
+    this.listener = listener;
+  }
+
   @Override
   public Message build() {
     if (built.compareAndSet(false, true)) {
+      Object context = listener.accept(this.stream, this);
       if (propertiesBuilder != null) {
         message.setProperties(propertiesBuilder.properties);
       }
@@ -57,7 +66,7 @@ public Message build() {
             new MessageAnnotations(messageAnnotationsBuilder.messageAnnotations));
       }
       return new QpidProtonCodec.QpidProtonAmqpMessageWrapper(
-          hasPublishingId, publishingId, message);
+          hasPublishingId, publishingId, message, context);
     } else {
       throw new IllegalStateException("A message builder can build only one message");
     }

src/main/java/com/rabbitmq/stream/codec/SimpleCodec.java

@@ -1,4 +1,4 @@
-// Copyright (c) 2020-2021 VMware, Inc. or its affiliates.  All rights reserved.
+// Copyright (c) 2020-2023 VMware, Inc. or its affiliates.  All rights reserved.
 //
 // This software, the RabbitMQ Stream Java client library, is dual-licensed under the
 // Mozilla Public License 2.0 ("MPL"), and the Apache License version 2 ("ASL").
@@ -37,6 +37,21 @@ public MessageBuilder messageBuilder() {
     return new SimpleMessageBuilder();
   }
 
+  @Override
+  public MessageBuilder messageBuilder(String stream) {
+    return null;
+  }
+
+  @Override
+  public Codec messageBuilderListener(MessageBuilderListener listener) {
+    return null;
+  }
+
+  @Override
+  public Object listenerContext(Message message) {
+    return null;
+  }
+
   private static class SimpleMessage implements Message {
 
     private final boolean hasPublishingId;

src/main/java/com/rabbitmq/stream/codec/SwiftMqCodec.java

@@ -1,4 +1,4 @@
-// Copyright (c) 2020-2022 VMware, Inc. or its affiliates.  All rights reserved.
+// Copyright (c) 2020-2023 VMware, Inc. or its affiliates.  All rights reserved.
 //
 // This software, the RabbitMQ Stream Java client library, is dual-licensed under the
 // Mozilla Public License 2.0 ("MPL"), and the Apache License version 2 ("ASL").
@@ -122,6 +122,21 @@ public MessageBuilder messageBuilder() {
     return new SwiftMqMessageBuilder();
   }
 
+  @Override
+  public MessageBuilder messageBuilder(String stream) {
+    return null;
+  }
+
+  @Override
+  public Codec messageBuilderListener(MessageBuilderListener listener) {
+    return null;
+  }
+
+  @Override
+  public Object listenerContext(Message message) {
+    return null;
+  }
+
   @Override
   public EncodedMessage encode(Message message) {
     AMQPMessage outboundMessage;

src/main/java/com/rabbitmq/stream/impl/SimpleMessageAccumulator.java

@@ -1,4 +1,4 @@
-// Copyright (c) 2020-2021 VMware, Inc. or its affiliates.  All rights reserved.
+// Copyright (c) 2020-2023 VMware, Inc. or its affiliates.  All rights reserved.
 //
 // This software, the RabbitMQ Stream Java client library, is dual-licensed under the
 // Mozilla Public License 2.0 ("MPL"), and the Apache License version 2 ("ASL").
@@ -13,11 +13,7 @@
 // [email protected].
 package com.rabbitmq.stream.impl;
 
-import com.rabbitmq.stream.Codec;
-import com.rabbitmq.stream.ConfirmationHandler;
-import com.rabbitmq.stream.ConfirmationStatus;
-import com.rabbitmq.stream.Message;
-import com.rabbitmq.stream.StreamException;
+import com.rabbitmq.stream.*;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
@@ -31,18 +27,20 @@ class SimpleMessageAccumulator implements MessageAccumulator {
   protected final BlockingQueue<AccumulatedEntity> messages;
   protected final Clock clock;
   private final int capacity;
-  private final Codec codec;
+  protected final Codec codec;
   private final int maxFrameSize;
   private final ToLongFunction<Message> publishSequenceFunction;
   private final Function<Message, String> filterValueExtractor;
+  final ObservationCollector observationCollector;
 
   SimpleMessageAccumulator(
       int capacity,
       Codec codec,
       int maxFrameSize,
       ToLongFunction<Message> publishSequenceFunction,
       Function<Message, String> filterValueExtractor,
-      Clock clock) {
+      Clock clock,
+      ObservationCollector observationCollector) {
     this.capacity = capacity;
     this.messages = new LinkedBlockingQueue<>(capacity);
     this.codec = codec;
@@ -51,6 +49,7 @@ class SimpleMessageAccumulator implements MessageAccumulator {
     this.filterValueExtractor =
         filterValueExtractor == null ? NULL_FILTER_VALUE_EXTRACTOR : filterValueExtractor;
     this.clock = clock;
+    this.observationCollector = observationCollector;
   }
 
   public boolean add(Message message, ConfirmationHandler confirmationHandler) {
@@ -79,7 +78,11 @@ public boolean add(Message message, ConfirmationHandler confirmationHandler) {
 
   @Override
   public AccumulatedEntity get() {
-    return this.messages.poll();
+    AccumulatedEntity entity = this.messages.poll();
+    if (entity != null) {
+      this.observationCollector.publish(this.codec, entity.confirmationCallback().message());
+    }
+    return entity;
   }
 
   @Override
@@ -155,5 +158,10 @@ public int handle(boolean confirmed, short code) {
       confirmationHandler.handle(new ConfirmationStatus(message, confirmed, code));
       return 1;
     }
+
+    @Override
+    public Message message() {
+      return this.message;
+    }
   }
 }

src/main/java/com/rabbitmq/stream/impl/StreamConsumerBuilder.java

@@ -219,6 +219,8 @@ public Consumer build() {
           };
     }
 
+    handler = this.environment.observationCollector().subscribe(handler);
+
     Consumer consumer;
     if (this.stream != null) {
       consumer =