Refactor observation publishing

Add method to annotate a created message.

Author: acogoluegnes

pom.xml

@@ -562,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

@@ -28,12 +28,6 @@ public interface Codec {
 
   MessageBuilder messageBuilder();
 
-  MessageBuilder messageBuilder(String stream);
-
-  Codec messageBuilderListener(MessageBuilderListener listener);
-
-  Object listenerContext(Message message);
-
   class EncodedMessage {
 
     private final int size;
@@ -52,10 +46,4 @@ public int getSize() {
       return size;
     }
   }
-
-  @FunctionalInterface
-  interface MessageBuilderListener {
-
-    Object accept(String stream, MessageBuilder builder);
-  }
 }

src/main/java/com/rabbitmq/stream/Message.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").
@@ -85,4 +85,9 @@ public interface Message {
    * @return the message annotations
    */
   Map<String, Object> getMessageAnnotations();
+
+  default Message annotate(String key, Object value) {
+    this.getMessageAnnotations().put(key, value);
+    return this;
+  }
 }

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

@@ -18,22 +18,22 @@ public interface ObservationCollector {
   ObservationCollector NO_OP =
       new ObservationCollector() {
         @Override
-        public Codec register(Codec codec) {
-          return codec;
+        public Object prePublish(String stream, Message message) {
+          return null;
         }
 
         @Override
-        public void publish(Codec codec, Message message) {}
+        public void published(Object context, Message message) {}
 
         @Override
         public MessageHandler subscribe(MessageHandler handler) {
           return handler;
         }
       };
 
-  Codec register(Codec codec);
+  void published(Object context, Message message);
 
-  void publish(Codec codec, Message message);
+  Object prePublish(String stream, Message message);
 
   MessageHandler subscribe(MessageHandler handler);
 }

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

@@ -32,19 +32,11 @@
 import org.apache.qpid.proton.codec.WritableBuffer;
 
 public class QpidProtonCodec implements Codec {
-  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) {
@@ -61,7 +53,7 @@ private static Map<String, Object> createMessageAnnotations(
       return createMapFromAmqpMap(
           MESSAGE_ANNOTATIONS_SYMBOL_KEY_EXTRACTOR, message.getMessageAnnotations().getValue());
     } else {
-      return null;
+      return new LinkedHashMap<>();
     }
   }
 
@@ -113,22 +105,6 @@ 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;
@@ -326,12 +302,7 @@ protected Object convertToQpidType(Object value) {
 
   @Override
   public MessageBuilder messageBuilder() {
-    return this.messageBuilder(null);
-  }
-
-  @Override
-  public MessageBuilder messageBuilder(String stream) {
-    return new QpidProtonMessageBuilder(stream, this.messageBuilderListener);
+    return new QpidProtonMessageBuilder();
   }
 
   private static final class QpidProtonProperties implements Properties {
@@ -534,27 +505,30 @@ public Map<String, Object> getApplicationProperties() {
     public Map<String, Object> getMessageAnnotations() {
       return messageAnnotations;
     }
+
+    @Override
+    public Message annotate(String key, Object value) {
+      this.messageAnnotations.put(key, value);
+      return this;
+    }
   }
 
   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,
-        Object listenerContext) {
+        org.apache.qpid.proton.message.Message message) {
       this.hasPublishingId = hasPublishingId;
       this.publishingId = publishingId;
       this.message = message;
-      this.listenerContext = listenerContext;
     }
 
     @Override
@@ -612,6 +586,17 @@ public Map<String, Object> getMessageAnnotations() {
         return null;
       }
     }
+
+    @Override
+    public Message annotate(String key, Object value) {
+      MessageAnnotations annotations = this.message.getMessageAnnotations();
+      if (annotations == null) {
+        annotations = new MessageAnnotations(new LinkedHashMap<>());
+        this.message.setMessageAnnotations(annotations);
+      }
+      annotations.getValue().put(Symbol.getSymbol(key), value);
+      return this;
+    }
   }
 
   // from

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

@@ -13,7 +13,6 @@
 // [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;
@@ -34,8 +33,6 @@
 
 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);
@@ -45,15 +42,9 @@ 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);
       }
@@ -66,7 +57,7 @@ public Message build() {
             new MessageAnnotations(messageAnnotationsBuilder.messageAnnotations));
       }
       return new QpidProtonCodec.QpidProtonAmqpMessageWrapper(
-          hasPublishingId, publishingId, message, context);
+          hasPublishingId, publishingId, message);
     } else {
       throw new IllegalStateException("A message builder can build only one message");
     }

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

@@ -37,21 +37,6 @@ public MessageBuilder messageBuilder() {
     return new SimpleMessageBuilder();
   }
 
-  @Override
-  public MessageBuilder messageBuilder(String stream) {
-    return new SimpleMessageBuilder();
-  }
-
-  @Override
-  public Codec messageBuilderListener(MessageBuilderListener listener) {
-    return new SimpleCodec();
-  }
-
-  @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

@@ -122,21 +122,6 @@ 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;
@@ -295,7 +280,7 @@ public EncodedMessage encode(Message message) {
     }
   }
 
-  protected AMQPType convertToSwiftMqType(Object value) {
+  protected static AMQPType convertToSwiftMqType(Object value) {
     if (value instanceof Boolean) {
       return ((Boolean) value).booleanValue() ? AMQPBoolean.TRUE : AMQPBoolean.FALSE;
     } else if (value instanceof Byte) {
@@ -663,5 +648,25 @@ public Map<String, Object> getMessageAnnotations() {
         return null;
       }
     }
+
+    @Override
+    public Message annotate(String key, Object value) {
+      MessageAnnotations annotations = this.message.getMessageAnnotations();
+      Map<AMQPType, AMQPType> map;
+      try {
+        if (annotations == null) {
+          map = new LinkedHashMap<>();
+          annotations = new MessageAnnotations(map);
+          this.message.setMessageAnnotations(annotations);
+        } else {
+          map = annotations.getValue();
+        }
+        map.put(new AMQPSymbol(key), convertToSwiftMqType(value));
+        annotations.setValue(map);
+      } catch (IOException e) {
+        throw new StreamException("Error while annotating SwiftMQ message");
+      }
+      return this;
+    }
   }
 }

src/main/java/com/rabbitmq/stream/impl/MessageAccumulator.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").
@@ -30,12 +30,14 @@ interface AccumulatedEntity {
 
     long time();
 
-    long publishindId();
+    long publishingId();
 
     String filterValue();
 
     Object encodedEntity();
 
     StreamProducer.ConfirmationCallback confirmationCallback();
+
+    Object observationContext();
   }
 }