Copy message if necessary for observation collector

When a message is routed to several streams in
a super stream producer.

References #379

Author: acogoluegnes

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

@@ -90,4 +90,8 @@ default Message annotate(String key, Object value) {
     this.getMessageAnnotations().put(key, value);
     return this;
   }
+
+  default Message copy() {
+    return this;
+  }
 }

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

@@ -36,4 +36,8 @@ public MessageHandler subscribe(MessageHandler handler) {
   T prePublish(String stream, Message message);
 
   MessageHandler subscribe(MessageHandler handler);
+
+  default boolean isNoop() {
+    return this == NO_OP;
+  }
 }

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

@@ -259,7 +259,7 @@ public Message decode(byte[] data) {
         createMessageAnnotations(message));
   }
 
-  protected Properties createProperties(org.apache.qpid.proton.message.Message message) {
+  protected static Properties createProperties(org.apache.qpid.proton.message.Message message) {
     if (message.getProperties() != null) {
       return new QpidProtonProperties(message.getProperties());
     } else {
@@ -511,6 +511,15 @@ public Message annotate(String key, Object value) {
       this.messageAnnotations.put(key, value);
       return this;
     }
+
+    @Override
+    public Message copy() {
+      return new QpidProtonMessage(
+          message,
+          createProperties(message),
+          createApplicationProperties(message),
+          createMessageAnnotations(message));
+    }
   }
 
   static class QpidProtonAmqpMessageWrapper implements Message {
@@ -597,6 +606,27 @@ public Message annotate(String key, Object value) {
       annotations.getValue().put(Symbol.getSymbol(key), value);
       return this;
     }
+
+    @Override
+    public Message copy() {
+      org.apache.qpid.proton.message.Message copy =
+          org.apache.qpid.proton.message.Message.Factory.create();
+      copy.setProperties(this.message.getProperties());
+      copy.setBody(this.message.getBody());
+      copy.setApplicationProperties(this.message.getApplicationProperties());
+      if (this.message.getMessageAnnotations() != null) {
+        Map<Symbol, Object> annotations = message.getMessageAnnotations().getValue();
+        Map<Symbol, Object> annotationCopy;
+        if (annotations == null) {
+          annotationCopy = null;
+        } else {
+          annotationCopy = new LinkedHashMap<>(annotations.size());
+          annotationCopy.putAll(annotations);
+        }
+        copy.setMessageAnnotations(new MessageAnnotations(annotationCopy));
+      }
+      return new QpidProtonAmqpMessageWrapper(this.hasPublishingId, this.publishingId, copy);
+    }
   }
 
   // from

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

@@ -276,7 +276,7 @@ public EncodedMessage encode(Message message) {
       outboundMessage.writeContent(output);
       return new EncodedMessage(output.getCount(), output.getBuffer());
     } catch (IOException e) {
-      throw new StreamException("Error while writing AMQP 1.0 message to output stream");
+      throw new StreamException("Error while writing AMQP 1.0 message to output stream", e);
     }
   }
 
@@ -355,7 +355,7 @@ protected Message createMessage(byte[] data) {
     try {
       amqpMessage = new AMQPMessage(data);
     } catch (Exception e) {
-      throw new StreamException("Error while decoding AMQP 1.0 message");
+      throw new StreamException("Error while decoding AMQP 1.0 message", e);
     }
 
     Object body = extractBody(amqpMessage);
@@ -664,9 +664,31 @@ public Message annotate(String key, Object value) {
         map.put(new AMQPSymbol(key), convertToSwiftMqType(value));
         annotations.setValue(map);
       } catch (IOException e) {
-        throw new StreamException("Error while annotating SwiftMQ message");
+        throw new StreamException("Error while annotating SwiftMQ message", e);
       }
       return this;
     }
+
+    @Override
+    public Message copy() {
+      AMQPMessage copy = new AMQPMessage();
+      copy.setProperties(this.message.getProperties());
+      if (this.message.getData() != null) {
+        this.message.getData().forEach(copy::addData);
+      }
+      copy.setApplicationProperties(this.message.getApplicationProperties());
+      MessageAnnotations annotations = this.message.getMessageAnnotations();
+      if (annotations != null) {
+        Map<AMQPType, AMQPType> annotationCopy = null;
+        try {
+          annotationCopy = new LinkedHashMap<>(annotations.getValue().size());
+          annotationCopy.putAll(annotations.getValue());
+          copy.setMessageAnnotations(new MessageAnnotations(annotationCopy));
+        } catch (IOException e) {
+          throw new StreamException("Error while copying SwiftMQ message annotations", e);
+        }
+      }
+      return new SwiftMqAmqpMessageWrapper(this.hasPublishingId, this.publishingId, copy);
+    }
   }
 }

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

@@ -15,14 +15,7 @@
 
 import static com.rabbitmq.stream.impl.Utils.namedFunction;
 
-import com.rabbitmq.stream.Codec;
-import com.rabbitmq.stream.ConfirmationHandler;
-import com.rabbitmq.stream.ConfirmationStatus;
-import com.rabbitmq.stream.Constants;
-import com.rabbitmq.stream.Message;
-import com.rabbitmq.stream.MessageBuilder;
-import com.rabbitmq.stream.Producer;
-import com.rabbitmq.stream.RoutingStrategy;
+import com.rabbitmq.stream.*;
 import com.rabbitmq.stream.RoutingStrategy.Metadata;
 import java.util.List;
 import java.util.Map;
@@ -47,6 +40,13 @@ class SuperStreamProducer implements Producer {
   private final Metadata superStreamMetadata;
   private final AtomicBoolean closed = new AtomicBoolean(false);
 
+  /**
+   * passthrough, except when observation collector is not no-op and a message must be sent to
+   * several streams. In this it creates a copy of the message with distinct message annotations for
+   * each message, so that the collector can populate these messages annotations without collision
+   */
+  private final MessageInterceptor messageInterceptor;
+
   SuperStreamProducer(
       StreamProducerBuilder producerBuilder,
       String name,
@@ -62,6 +62,10 @@ class SuperStreamProducer implements Producer {
     this.producerBuilder = producerBuilder.duplicate();
     this.producerBuilder.stream(null);
     this.producerBuilder.resetRouting();
+    this.messageInterceptor =
+        environment.observationCollector().isNoop()
+            ? (i, msg) -> msg
+            : (i, msg) -> i == 0 ? msg : msg.copy();
   }
 
   @Override
@@ -111,17 +115,12 @@ public void send(Message message, ConfirmationHandler confirmationHandler) {
       if (streams.isEmpty()) {
         confirmationHandler.handle(
             new ConfirmationStatus(message, false, Constants.CODE_NO_ROUTE_FOUND));
+      } else if (streams.size() == 1) {
+        producer(streams.get(0)).send(message, confirmationHandler);
       } else {
-        for (String stream : streams) {
-          Producer producer =
-              producers.computeIfAbsent(
-                  stream,
-                  stream1 -> {
-                    Producer p =
-                        producerBuilder.duplicate().superStream(null).stream(stream1).build();
-                    return p;
-                  });
-          producer.send(message, confirmationHandler);
+        for (int i = 0; i < streams.size(); i++) {
+          Producer producer = producer(streams.get(i));
+          producer(streams.get(i)).send(messageInterceptor.apply(i, message), confirmationHandler);
         }
       }
     } else {
@@ -130,6 +129,15 @@ public void send(Message message, ConfirmationHandler confirmationHandler) {
     }
   }
 
+  private Producer producer(String stream) {
+    return producers.computeIfAbsent(
+        stream,
+        stream1 -> {
+          Producer p = producerBuilder.duplicate().superStream(null).stream(stream1).build();
+          return p;
+        });
+  }
+
   private boolean canSend() {
     return !this.closed.get();
   }
@@ -191,4 +199,10 @@ public List<String> route(String routingKey) {
                       routingKey1)));
     }
   }
+
+  @FunctionalInterface
+  private interface MessageInterceptor {
+
+    Message apply(int partitionIndex, Message message);
+  }
 }

src/test/java/com/rabbitmq/stream/codec/CodecsTest.java

@@ -39,6 +39,7 @@
 import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Supplier;
+import java.util.function.UnaryOperator;
 import java.util.stream.Stream;
 import org.apache.qpid.proton.amqp.messaging.AmqpValue;
 import org.assertj.core.api.InstanceOfAssertFactories;
@@ -57,8 +58,19 @@ static Iterable<CodecCouple> codecsCouples() {
     List<CodecCouple> couples = new ArrayList<>();
     for (Codec serializer : codecs) {
       for (Codec deserializer : codecs) {
-        couples.add(new CodecCouple(serializer, deserializer, () -> serializer.messageBuilder()));
-        couples.add(new CodecCouple(serializer, deserializer, () -> new WrapperMessageBuilder()));
+        couples.add(new CodecCouple(serializer, deserializer, serializer::messageBuilder));
+        couples.add(new CodecCouple(serializer, deserializer, WrapperMessageBuilder::new));
+      }
+    }
+    return couples;
+  }
+
+  static Iterable<CodecCouple> codecsCombinations() {
+    List<Codec> codecs = asList(new QpidProtonCodec(), new SwiftMqCodec());
+    List<CodecCouple> couples = new ArrayList<>();
+    for (Codec serializer : codecs) {
+      for (Codec deserializer : codecs) {
+        couples.add(new CodecCouple(serializer, deserializer, serializer::messageBuilder));
       }
     }
     return couples;
@@ -559,6 +571,58 @@ void publishingIdShouldNotBeSetOnMessageIfNotSetOnMessageBuilder(MessageBuilder
     assertThat(message.getPublishingId()).isEqualTo(0);
   }
 
+  @ParameterizedTest
+  @MethodSource("codecsCombinations")
+  void copy(CodecCouple codecCouple) {
+    Codec serializer = codecCouple.serializer;
+    Codec deserializer = codecCouple.deserializer;
+    byte[] body = "hello".getBytes(StandardCharsets.UTF_8);
+
+    Message message =
+        serializer
+            .messageBuilder()
+            .addData(body)
+            .messageAnnotations()
+            .entry("foo", "bar")
+            .messageBuilder()
+            .build();
+    Message copy = message.copy();
+
+    message.annotate("original", "original value");
+    copy.annotate("copy", "copy value");
+
+    assertThat(message.getMessageAnnotations())
+        .hasSize(2)
+        .containsEntry("foo", "bar")
+        .containsEntry("original", "original value");
+
+    assertThat(copy.getMessageAnnotations())
+        .hasSize(2)
+        .containsEntry("foo", "bar")
+        .containsEntry("copy", "copy value");
+
+    UnaryOperator<Message> encodeDecode =
+        msg -> {
+          EncodedMessage encoded = serializer.encode(msg);
+          byte[] encodedData = new byte[encoded.getSize()];
+          System.arraycopy(encoded.getData(), 0, encodedData, 0, encoded.getSize());
+          return deserializer.decode(encodedData);
+        };
+
+    message = encodeDecode.apply(message);
+
+    assertThat(message.getMessageAnnotations())
+        .hasSize(2)
+        .containsEntry("foo", "bar")
+        .containsEntry("original", "original value");
+
+    copy = encodeDecode.apply(copy);
+    assertThat(copy.getMessageAnnotations())
+        .hasSize(2)
+        .containsEntry("foo", "bar")
+        .containsEntry("copy", "copy value");
+  }
+
   MessageTestConfiguration test(
       Function<MessageBuilder, MessageBuilder> messageOperation,
       Consumer<Message> messageExpectation) {