Test against dynamic-batch message accumulator

Author: acogoluegnes

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

@@ -65,7 +65,24 @@ final class DynamicBatchMessageAccumulator implements MessageAccumulator {
     this.filterValueExtractor =
         filterValueExtractor == null ? NULL_FILTER_VALUE_EXTRACTOR : filterValueExtractor;
     if (subEntrySize <= 1) {
-      this.dynamicBatch = new DynamicBatch<>(this::publish, batchSize);
+      this.dynamicBatch =
+          new DynamicBatch<>(
+              items -> {
+                // TODO add a "replay" flag to DynamicBatch to avoid checking the producer status
+                // the status check helps to avoid collecting the observation another time
+                if (producer.canSend()) {
+                  items.forEach(
+                      i -> {
+                        AccumulatedEntity entity = (AccumulatedEntity) i;
+                        this.observationCollector.published(
+                            entity.observationContext(), entity.confirmationCallback().message());
+                      });
+                  return this.publish(items);
+                } else {
+                  return false;
+                }
+              },
+              batchSize);
     } else {
       byte compressionCode =
           compressionCodec == null ? Compression.NONE.code() : compressionCodec.code();

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

@@ -174,7 +174,7 @@ public int fragmentLength(Object entity) {
     }
     this.accumulator =
         ProducerUtils.createMessageAccumulator(
-            false,
+            true,
             subEntrySize,
             batchSize,
             compressionCodec,
@@ -306,8 +306,10 @@ private long computeFirstValueOfPublishingSequence() {
     }
   }
 
+  // visible for testing
   void confirm(long publishingId) {
     AccumulatedEntity accumulatedEntity = this.unconfirmedMessages.remove(publishingId);
+
     if (accumulatedEntity != null) {
       int confirmedCount =
           accumulatedEntity.confirmationCallback().handle(true, Constants.RESPONSE_CODE_OK);
@@ -317,6 +319,11 @@ void confirm(long publishingId) {
     }
   }
 
+  // for testing
+  int unconfirmedCount() {
+    return this.unconfirmedMessages.size();
+  }
+
   void error(long publishingId, short errorCode) {
     AccumulatedEntity accumulatedEntity = unconfirmedMessages.remove(publishingId);
     if (accumulatedEntity != null) {

src/test/java/com/rabbitmq/stream/impl/DynamicBatchTest.java

@@ -18,7 +18,6 @@
 import com.codahale.metrics.Histogram;
 import com.codahale.metrics.MetricRegistry;
 import com.google.common.util.concurrent.RateLimiter;
-import java.time.Duration;
 import java.util.Random;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.IntStream;
@@ -62,7 +61,7 @@ void test() {
             });
     Assertions.assertThat(sync).completes();
     long end = System.nanoTime();
-    System.out.println("Done in " + Duration.ofNanos(end - start));
-    reporter.report();
+    //    System.out.println("Done in " + Duration.ofNanos(end - start));
+    //    reporter.report();
   }
 }

src/test/java/com/rabbitmq/stream/impl/StreamProducerUnitTest.java

@@ -15,6 +15,7 @@
 package com.rabbitmq.stream.impl;
 
 import static com.rabbitmq.stream.impl.TestUtils.waitAtMost;
+import static java.util.stream.IntStream.range;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.mockito.ArgumentMatchers.*;
@@ -43,8 +44,8 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.ToLongFunction;
-import java.util.stream.IntStream;
 import org.assertj.core.api.ThrowableAssert.ThrowingCallable;
+import org.assertj.core.data.Offset;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.params.ParameterizedTest;
@@ -184,26 +185,33 @@ void confirmTimeoutTaskShouldFailMessagesAfterTimeout(
             null,
             env);
 
-    IntStream.range(0, messageCount)
+    range(0, messageCount)
         .forEach(
             i ->
                 producer.send(
                     producer.messageBuilder().addData("".getBytes()).build(), confirmationHandler));
 
-    IntStream.range(0, confirmedPart).forEach(publishingId -> producer.confirm(publishingId));
-    assertThat(confirmedCount.get()).isEqualTo(expectedConfirmed);
+    waitAtMost(() -> producer.unconfirmedCount() >= messageCount / subEntrySize);
+    range(0, confirmedPart).forEach(producer::confirm);
+    if (subEntrySize == 1) {
+      assertThat(confirmedCount.get()).isEqualTo(expectedConfirmed);
+    } else {
+      assertThat(confirmedCount.get()).isCloseTo(confirmedCount.get(), Offset.offset(subEntrySize));
+    }
     assertThat(erroredCount.get()).isZero();
+    int confirmedPreviously = confirmedCount.get();
 
     executorService.scheduleAtFixedRate(() -> clock.refresh(), 100, 100, TimeUnit.MILLISECONDS);
 
     Thread.sleep(waitTime.toMillis());
-    assertThat(confirmedCount.get()).isEqualTo(expectedConfirmed);
+    assertThat(confirmedCount.get()).isEqualTo(confirmedPreviously);
     if (confirmTimeout.isZero()) {
       assertThat(erroredCount.get()).isZero();
       assertThat(responseCodes).isEmpty();
     } else {
       waitAtMost(
-          waitTime.multipliedBy(2), () -> erroredCount.get() == (messageCount - expectedConfirmed));
+          waitTime.multipliedBy(2),
+          () -> erroredCount.get() == (messageCount - confirmedPreviously));
       assertThat(responseCodes).hasSize(1).contains(Constants.CODE_PUBLISH_CONFIRM_TIMEOUT);
     }
   }