Add helper for message accumulator

acogoluegnes · acogoluegnes · commit 8c5460798564 · 2024-11-06T11:48:47.000+01:00
diff --git a/src/main/java/com/rabbitmq/stream/impl/DynamicBatchMessageAccumulator.java b/src/main/java/com/rabbitmq/stream/impl/DynamicBatchMessageAccumulator.java
@@ -29,17 +29,10 @@
 
 final class DynamicBatchMessageAccumulator implements MessageAccumulator {
 
-  private static final Function<Message, String> NULL_FILTER_VALUE_EXTRACTOR = m -> null;
-
   private final DynamicBatch<Object> dynamicBatch;
   private final ObservationCollector<Object> observationCollector;
-  private final ToLongFunction<Message> publishSequenceFunction;
-  private final String stream;
   private final StreamProducer producer;
-  private final Codec codec;
-  private final int maxFrameSize;
-  private final Clock clock;
-  private final Function<Message, String> filterValueExtractor;
+  private final ProducerUtils.MessageAccumulatorHelper helper;
 
   @SuppressWarnings("unchecked")
   DynamicBatchMessageAccumulator(
@@ -55,15 +48,17 @@ final class DynamicBatchMessageAccumulator implements MessageAccumulator {
       ByteBufAllocator byteBufAllocator,
       ObservationCollector<?> observationCollector,
       StreamProducer producer) {
+    this.helper =
+        new ProducerUtils.MessageAccumulatorHelper(
+            codec,
+            maxFrameSize,
+            publishSequenceFunction,
+            filterValueExtractor,
+            clock,
+            stream,
+            observationCollector);
     this.producer = producer;
-    this.stream = stream;
-    this.publishSequenceFunction = publishSequenceFunction;
     this.observationCollector = (ObservationCollector<Object>) observationCollector;
-    this.codec = codec;
-    this.clock = clock;
-    this.maxFrameSize = maxFrameSize;
-    this.filterValueExtractor =
-        filterValueExtractor == null ? NULL_FILTER_VALUE_EXTRACTOR : filterValueExtractor;
     if (subEntrySize <= 1) {
       this.dynamicBatch =
           new DynamicBatch<>(
@@ -93,11 +88,8 @@ final class DynamicBatchMessageAccumulator implements MessageAccumulator {
                   List<Object> subBatches = new ArrayList<>();
                   int count = 0;
                   ProducerUtils.Batch batch =
-                      new ProducerUtils.Batch(
-                          Client.EncodedMessageBatch.create(
-                              byteBufAllocator, compressionCode, compressionCodec, subEntrySize),
-                          new ProducerUtils.CompositeConfirmationCallback(
-                              new ArrayList<>(subEntrySize)));
+                      this.helper.batch(
+                          byteBufAllocator, compressionCode, compressionCodec, subEntrySize);
                   AccumulatedEntity lastMessageInBatch = null;
                   for (Object msg : items) {
                     AccumulatedEntity message = (AccumulatedEntity) msg;
@@ -115,14 +107,8 @@ final class DynamicBatchMessageAccumulator implements MessageAccumulator {
                       subBatches.add(batch);
                       lastMessageInBatch = null;
                       batch =
-                          new ProducerUtils.Batch(
-                              Client.EncodedMessageBatch.create(
-                                  byteBufAllocator,
-                                  compressionCode,
-                                  compressionCodec,
-                                  subEntrySize),
-                              new ProducerUtils.CompositeConfirmationCallback(
-                                  new ArrayList<>(subEntrySize)));
+                          this.helper.batch(
+                              byteBufAllocator, compressionCode, compressionCodec, subEntrySize);
                       count = 0;
                     }
                   }
@@ -145,18 +131,7 @@ final class DynamicBatchMessageAccumulator implements MessageAccumulator {
 
   @Override
   public void add(Message message, ConfirmationHandler confirmationHandler) {
-    Object observationContext = this.observationCollector.prePublish(this.stream, message);
-    Codec.EncodedMessage encodedMessage = this.codec.encode(message);
-    Client.checkMessageFitsInFrame(this.maxFrameSize, encodedMessage);
-    long publishingId = this.publishSequenceFunction.applyAsLong(message);
-    this.dynamicBatch.add(
-        new ProducerUtils.SimpleAccumulatedEntity(
-            this.clock.time(),
-            publishingId,
-            this.filterValueExtractor.apply(message),
-            this.codec.encode(message),
-            new ProducerUtils.SimpleConfirmationCallback(message, confirmationHandler),
-            observationContext));
+    this.dynamicBatch.add(helper.entity(message, confirmationHandler));
   }
 
   @Override
diff --git a/src/main/java/com/rabbitmq/stream/impl/ProducerUtils.java b/src/main/java/com/rabbitmq/stream/impl/ProducerUtils.java
@@ -17,6 +17,7 @@
 import com.rabbitmq.stream.*;
 import com.rabbitmq.stream.compression.CompressionCodec;
 import io.netty.buffer.ByteBufAllocator;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.function.Function;
 import java.util.function.ToLongFunction;
@@ -262,4 +263,60 @@ public Object observationContext() {
           "batch entity does not contain only one observation context");
     }
   }
+
+  static final class MessageAccumulatorHelper {
+
+    private static final Function<Message, String> NULL_FILTER_VALUE_EXTRACTOR = m -> null;
+
+    private final ObservationCollector<Object> observationCollector;
+    private final ToLongFunction<Message> publishSequenceFunction;
+    private final String stream;
+    private final Codec codec;
+    private final int maxFrameSize;
+    private final Clock clock;
+    private final Function<Message, String> filterValueExtractor;
+
+    @SuppressWarnings("unchecked")
+    MessageAccumulatorHelper(
+        Codec codec,
+        int maxFrameSize,
+        ToLongFunction<Message> publishSequenceFunction,
+        Function<Message, String> filterValueExtractor,
+        Clock clock,
+        String stream,
+        ObservationCollector<?> observationCollector) {
+      this.publishSequenceFunction = publishSequenceFunction;
+      this.codec = codec;
+      this.clock = clock;
+      this.maxFrameSize = maxFrameSize;
+      this.filterValueExtractor =
+          filterValueExtractor == null ? NULL_FILTER_VALUE_EXTRACTOR : filterValueExtractor;
+      this.observationCollector = (ObservationCollector<Object>) observationCollector;
+      this.stream = stream;
+    }
+
+    AccumulatedEntity entity(Message message, ConfirmationHandler confirmationHandler) {
+      Object observationContext = this.observationCollector.prePublish(this.stream, message);
+      Codec.EncodedMessage encodedMessage = this.codec.encode(message);
+      Client.checkMessageFitsInFrame(this.maxFrameSize, encodedMessage);
+      long publishingId = this.publishSequenceFunction.applyAsLong(message);
+      return new ProducerUtils.SimpleAccumulatedEntity(
+          this.clock.time(),
+          publishingId,
+          this.filterValueExtractor.apply(message),
+          this.codec.encode(message),
+          new ProducerUtils.SimpleConfirmationCallback(message, confirmationHandler),
+          observationContext);
+    }
+
+    Batch batch(
+        ByteBufAllocator bba,
+        byte compressionCode,
+        CompressionCodec compressionCodec,
+        int subEntrySize) {
+      return new ProducerUtils.Batch(
+          Client.EncodedMessageBatch.create(bba, compressionCode, compressionCodec, subEntrySize),
+          new ProducerUtils.CompositeConfirmationCallback(new ArrayList<>(subEntrySize)));
+    }
+  }
 }
diff --git a/src/main/java/com/rabbitmq/stream/impl/SimpleMessageAccumulator.java b/src/main/java/com/rabbitmq/stream/impl/SimpleMessageAccumulator.java
@@ -26,18 +26,11 @@
 
 class SimpleMessageAccumulator implements MessageAccumulator {
 
-  private static final Function<Message, String> NULL_FILTER_VALUE_EXTRACTOR = m -> null;
-
   protected final BlockingQueue<AccumulatedEntity> messages;
-  protected final Clock clock;
   private final int capacity;
-  protected final Codec codec;
-  private final int maxFrameSize;
-  private final ToLongFunction<Message> publishSequenceFunction;
-  private final Function<Message, String> filterValueExtractor;
-  final String stream;
   final ObservationCollector<Object> observationCollector;
   private final StreamProducer producer;
+  final ProducerUtils.MessageAccumulatorHelper helper;
 
   @SuppressWarnings("unchecked")
   SimpleMessageAccumulator(
@@ -50,36 +43,25 @@ class SimpleMessageAccumulator implements MessageAccumulator {
       String stream,
       ObservationCollector<?> observationCollector,
       StreamProducer producer) {
+    this.helper =
+        new ProducerUtils.MessageAccumulatorHelper(
+            codec,
+            maxFrameSize,
+            publishSequenceFunction,
+            filterValueExtractor,
+            clock,
+            stream,
+            observationCollector);
     this.capacity = capacity;
-    this.messages = new LinkedBlockingQueue<>(capacity);
-    this.codec = codec;
-    this.maxFrameSize = maxFrameSize;
-    this.publishSequenceFunction = publishSequenceFunction;
-    this.filterValueExtractor =
-        filterValueExtractor == null ? NULL_FILTER_VALUE_EXTRACTOR : filterValueExtractor;
-    this.clock = clock;
-    this.stream = stream;
+    this.messages = new LinkedBlockingQueue<>(this.capacity);
     this.observationCollector = (ObservationCollector<Object>) observationCollector;
     this.producer = producer;
   }
 
   public void add(Message message, ConfirmationHandler confirmationHandler) {
-    Object observationContext = this.observationCollector.prePublish(this.stream, message);
-    Codec.EncodedMessage encodedMessage = this.codec.encode(message);
-    Client.checkMessageFitsInFrame(this.maxFrameSize, encodedMessage);
-    long publishingId = this.publishSequenceFunction.applyAsLong(message);
+    AccumulatedEntity entity = this.helper.entity(message, confirmationHandler);
     try {
-      boolean offered =
-          messages.offer(
-              new ProducerUtils.SimpleAccumulatedEntity(
-                  clock.time(),
-                  publishingId,
-                  this.filterValueExtractor.apply(message),
-                  encodedMessage,
-                  new ProducerUtils.SimpleConfirmationCallback(message, confirmationHandler),
-                  observationContext),
-              60,
-              TimeUnit.SECONDS);
+      boolean offered = messages.offer(entity, 60, TimeUnit.SECONDS);
       if (!offered) {
         throw new StreamException("Could not accumulate outbound message");
       }
diff --git a/src/main/java/com/rabbitmq/stream/impl/StreamProducer.java b/src/main/java/com/rabbitmq/stream/impl/StreamProducer.java
@@ -172,9 +172,10 @@ public int fragmentLength(Object entity) {
     if (compression != null) {
       compressionCodec = environment.compressionCodecFactory().get(compression);
     }
+    boolean dynamicBatch = true;
     this.accumulator =
         ProducerUtils.createMessageAccumulator(
-            true,
+            dynamicBatch,
             subEntrySize,
             batchSize,
             compressionCodec,
@@ -188,7 +189,11 @@ public int fragmentLength(Object entity) {
             environment.observationCollector(),
             this);
 
-    if (!batchPublishingDelay.isNegative() && !batchPublishingDelay.isZero()) {
+    boolean backgroundBatchPublishingTaskRequired =
+        !dynamicBatch && batchPublishingDelay.toMillis() > 0;
+    LOGGER.debug(
+        "Background batch publishing task required? {}", backgroundBatchPublishingTaskRequired);
+    if (backgroundBatchPublishingTaskRequired) {
       AtomicReference<Runnable> taskReference = new AtomicReference<>();
       Runnable task =
           () -> {
diff --git a/src/main/java/com/rabbitmq/stream/impl/SubEntryMessageAccumulator.java b/src/main/java/com/rabbitmq/stream/impl/SubEntryMessageAccumulator.java
@@ -20,17 +20,15 @@
 import com.rabbitmq.stream.ObservationCollector;
 import com.rabbitmq.stream.compression.Compression;
 import com.rabbitmq.stream.compression.CompressionCodec;
-import com.rabbitmq.stream.impl.Client.EncodedMessageBatch;
 import io.netty.buffer.ByteBufAllocator;
-import java.util.ArrayList;
 import java.util.function.ToLongFunction;
 
 final class SubEntryMessageAccumulator extends SimpleMessageAccumulator {
 
   private final int subEntrySize;
   private final CompressionCodec compressionCodec;
   private final ByteBufAllocator byteBufAllocator;
-  private final byte compression;
+  private final byte compressionCode;
 
   public SubEntryMessageAccumulator(
       int subEntrySize,
@@ -56,15 +54,14 @@ public SubEntryMessageAccumulator(
         producer);
     this.subEntrySize = subEntrySize;
     this.compressionCodec = compressionCodec;
-    this.compression = compressionCodec == null ? Compression.NONE.code() : compressionCodec.code();
+    this.compressionCode =
+        compressionCodec == null ? Compression.NONE.code() : compressionCodec.code();
     this.byteBufAllocator = byteBufAllocator;
   }
 
   private ProducerUtils.Batch createBatch() {
-    return new ProducerUtils.Batch(
-        EncodedMessageBatch.create(
-            byteBufAllocator, compression, compressionCodec, this.subEntrySize),
-        new ProducerUtils.CompositeConfirmationCallback(new ArrayList<>(this.subEntrySize)));
+    return this.helper.batch(
+        this.byteBufAllocator, this.compressionCode, this.compressionCodec, this.subEntrySize);
   }
 
   @Override
@@ -73,7 +70,7 @@ protected ProducerUtils.AccumulatedEntity get() {
       return null;
     }
     int count = 0;
-    ProducerUtils.Batch batch = createBatch();
+    ProducerUtils.Batch batch = this.createBatch();
     ProducerUtils.AccumulatedEntity lastMessageInBatch = null;
     while (count != this.subEntrySize) {
       ProducerUtils.AccumulatedEntity message = messages.poll();
diff --git a/src/test/java/com/rabbitmq/stream/impl/StreamProducerTest.java b/src/test/java/com/rabbitmq/stream/impl/StreamProducerTest.java
diff --git a/src/test/resources/logback-test.xml b/src/test/resources/logback-test.xml
