Adapt auto offset tracking for super stream consumers

The message count before storage setting is divided by the
number of partitions to make sure we store offsets
approximately as set on the consumer builder (assuming
partitions are balanced).

Author: acogoluegnes

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

@@ -39,14 +39,16 @@ class SuperStreamConsumer implements Consumer {
     this.superStream = superStream;
     List<String> partitions = environment.locatorOperation(c -> c.partitions(superStream));
 
-    // for manual tracking offset strategy only
+    // for manual offset tracking strategy only
     ConsumerState[] states = new ConsumerState[partitions.size()];
     Map<String, ConsumerState> partitionToStates = new HashMap<>(partitions.size());
     for (int i = 0; i < partitions.size(); i++) {
       ConsumerState state = new ConsumerState();
       states[i] = state;
       partitionToStates.put(partitions.get(i), state);
     }
+    // end of manual offset tracking strategy
+
     for (String partition : partitions) {
       ConsumerState state = partitionToStates.get(partition);
       MessageHandler messageHandler;
@@ -56,13 +58,21 @@ class SuperStreamConsumer implements Consumer {
       } else {
         messageHandler = builder.messageHandler();
       }
+      StreamConsumerBuilder subConsumerBuilder = builder.duplicate();
+
+      if (trackingConfiguration.enabled() && trackingConfiguration.auto()) {
+        subConsumerBuilder =
+            (StreamConsumerBuilder)
+                subConsumerBuilder
+                    .autoTrackingStrategy()
+                    .messageCountBeforeStorage(
+                        trackingConfiguration.autoMessageCountBeforeStorage() / partitions.size())
+                    .builder();
+      }
+
       Consumer consumer =
-          builder
-              .duplicate()
-              .lazyInit(true)
-              .superStream(null)
-              .messageHandler(messageHandler)
-              .stream(partition)
+          subConsumerBuilder.lazyInit(true).superStream(null).messageHandler(messageHandler).stream(
+                  partition)
               .build();
       consumers.put(partition, consumer);
       state.consumer = consumer;

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

@@ -21,7 +21,6 @@
 import com.rabbitmq.client.Channel;
 import com.rabbitmq.client.Connection;
 import com.rabbitmq.client.ConnectionFactory;
-import com.rabbitmq.stream.impl.TestUtils.BrokerVersionAtLeast;
 import java.util.Arrays;
 import java.util.List;
 import java.util.UUID;
@@ -55,7 +54,6 @@ void tearDown() throws Exception {
   }
 
   @Test
-  @BrokerVersionAtLeast("3.9.6")
   void routeShouldReturnEmptyListWhenExchangeDoesNotExist() {
     assertThat(cf.get().route("", UUID.randomUUID().toString())).isEmpty();
   }
@@ -66,7 +64,6 @@ void partitionsShouldReturnEmptyListWhenExchangeDoesNotExist() {
   }
 
   @Test
-  @BrokerVersionAtLeast("3.9.6")
   void routeShouldReturnNullWhenNoStreamForRoutingKey() throws Exception {
     declareSuperStreamTopology(connection, superStream, partitions);
 
@@ -84,7 +81,6 @@ void partitionsShouldReturnEmptyListWhenThereIsNoBinding() throws Exception {
   }
 
   @Test
-  @BrokerVersionAtLeast("3.9.6")
   void routeTopologyWithPartitionCount() throws Exception {
     declareSuperStreamTopology(connection, superStream, 3);
 
@@ -100,7 +96,6 @@ void routeTopologyWithPartitionCount() throws Exception {
   }
 
   @Test
-  @BrokerVersionAtLeast("3.9.6")
   void routeReturnsMultipleStreamsIfMultipleBindingsForSameKey() throws Exception {
     declareSuperStreamTopology(connection, superStream, 3);
     connection.createChannel().queueBind(superStream + "-1", superStream, "0");
@@ -118,7 +113,6 @@ void routeReturnsMultipleStreamsIfMultipleBindingsForSameKey() throws Exception
   }
 
   @Test
-  @BrokerVersionAtLeast("3.9.6")
   void partitionsAndRouteShouldNotReturnNonStreamQueue() throws Exception {
     declareSuperStreamTopology(connection, superStream, 3);
     Channel channel = connection.createChannel();
@@ -137,7 +131,6 @@ void partitionsAndRouteShouldNotReturnNonStreamQueue() throws Exception {
   }
 
   @Test
-  @BrokerVersionAtLeast("3.9.6")
   void partitionsReturnsCorrectOrder() throws Exception {
     String[] partitionNames = {"z", "y", "x"};
     declareSuperStreamTopology(connection, superStream, partitionNames);

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

@@ -27,7 +27,6 @@
 import com.rabbitmq.stream.EnvironmentBuilder;
 import com.rabbitmq.stream.OffsetSpecification;
 import com.rabbitmq.stream.impl.Client.ClientParameters;
-import com.rabbitmq.stream.impl.TestUtils.BrokerVersionAtLeast;
 import io.netty.channel.EventLoopGroup;
 import java.nio.charset.StandardCharsets;
 import java.util.Collections;
@@ -97,7 +96,6 @@ void tearDown() throws Exception {
   }
 
   @Test
-  @BrokerVersionAtLeast("3.9.6")
   void consumeAllMessagesFromAllPartitions() throws Exception {
     declareSuperStreamTopology(connection, superStream, partitionCount);
     Client client = cf.get();
@@ -130,7 +128,6 @@ void consumeAllMessagesFromAllPartitions() throws Exception {
   }
 
   @Test
-  @BrokerVersionAtLeast("3.9.6")
   void manualOffsetTrackingShouldStoreOnAllPartitions() throws Exception {
     declareSuperStreamTopology(connection, superStream, partitionCount);
     Client client = cf.get();
@@ -180,4 +177,56 @@ void manualOffsetTrackingShouldStoreOnAllPartitions() throws Exception {
         p -> assertThat(client.queryOffset(consumerName, p)).isGreaterThan(almostLastOffset));
     consumer.close();
   }
+
+  @Test
+  void autoOffsetTrackingShouldStoreOnAllPartitions() throws Exception {
+    declareSuperStreamTopology(connection, superStream, partitionCount);
+    Client client = cf.get();
+    List<String> partitions = client.partitions(superStream);
+    int messageCount = 10000 * partitionCount;
+    publishToPartitions(cf, partitions, messageCount);
+    ConcurrentMap<String, AtomicInteger> messagesReceived = new ConcurrentHashMap<>(partitionCount);
+    ConcurrentMap<String, Long> lastOffsets = new ConcurrentHashMap<>(partitionCount);
+    partitions.forEach(
+        p -> {
+          messagesReceived.put(p, new AtomicInteger(0));
+        });
+    CountDownLatch consumeLatch = new CountDownLatch(messageCount);
+    String consumerName = "my-app";
+    AtomicInteger totalCount = new AtomicInteger();
+    Consumer consumer =
+        environment
+            .consumerBuilder()
+            .superStream(superStream)
+            .offset(OffsetSpecification.first())
+            .name(consumerName)
+            .autoTrackingStrategy()
+            .messageCountBeforeStorage(messageCount / partitionCount / 50)
+            .builder()
+            .messageHandler(
+                (context, message) -> {
+                  String partition = new String(message.getBodyAsBinary());
+                  messagesReceived.get(partition).incrementAndGet();
+                  lastOffsets.put(partition, context.offset());
+                  totalCount.incrementAndGet();
+                  if (totalCount.get() % 50 == 0) {
+                    context.storeOffset();
+                  }
+                  consumeLatch.countDown();
+                })
+            .build();
+    latchAssert(consumeLatch).completes();
+    assertThat(messagesReceived).hasSize(partitionCount);
+    partitions.forEach(
+        p -> {
+          assertThat(messagesReceived).containsKey(p);
+          assertThat(messagesReceived.get(p).get()).isEqualTo(messageCount / partitionCount);
+        });
+    // checking stored offsets are big enough
+    // offset near the end (the message count per partition minus a few messages)
+    long almostLastOffset = messageCount / partitionCount - messageCount / (partitionCount * 10);
+    partitions.forEach(
+        p -> assertThat(client.queryOffset(consumerName, p)).isGreaterThan(almostLastOffset));
+    consumer.close();
+  }
 }

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

@@ -26,7 +26,6 @@
 import com.rabbitmq.stream.EnvironmentBuilder;
 import com.rabbitmq.stream.OffsetSpecification;
 import com.rabbitmq.stream.Producer;
-import com.rabbitmq.stream.impl.TestUtils.BrokerVersionAtLeast;
 import io.netty.channel.EventLoopGroup;
 import java.util.Map;
 import java.util.UUID;
@@ -74,7 +73,6 @@ void tearDown() throws Exception {
   }
 
   @Test
-  @BrokerVersionAtLeast("3.9.6")
   void allMessagesSentToSuperStreamWithHashRoutingShouldBeThenConsumed() throws Exception {
     int messageCount = 10_000;
     declareSuperStreamTopology(connection, superStream, partitions);
@@ -127,7 +125,6 @@ void allMessagesSentToSuperStreamWithHashRoutingShouldBeThenConsumed() throws Ex
   }
 
   @Test
-  @BrokerVersionAtLeast("3.9.6")
   void allMessagesSentToSuperStreamWithRoutingKeyRoutingShouldBeThenConsumed() throws Exception {
     int messageCount = 10_000;
     routingKeys = new String[] {"amer", "emea", "apac"};
@@ -179,7 +176,6 @@ void allMessagesSentToSuperStreamWithRoutingKeyRoutingShouldBeThenConsumed() thr
   }
 
   @Test
-  @BrokerVersionAtLeast("3.9.6")
   void getLastPublishingIdShouldReturnLowestValue() throws Exception {
     int messageCount = 10_000;
     declareSuperStreamTopology(connection, superStream, partitions);

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

@@ -14,6 +14,7 @@
 package com.rabbitmq.stream.impl;
 
 import static com.rabbitmq.stream.impl.TestUtils.declareSuperStreamTopology;
+import static com.rabbitmq.stream.impl.TestUtils.deleteSuperStreamTopology;
 import static com.rabbitmq.stream.impl.TestUtils.latchAssert;
 import static com.rabbitmq.stream.impl.TestUtils.localhost;
 import static org.assertj.core.api.Assertions.assertThat;
@@ -24,7 +25,6 @@
 import com.rabbitmq.stream.EnvironmentBuilder;
 import com.rabbitmq.stream.OffsetSpecification;
 import com.rabbitmq.stream.Producer;
-import com.rabbitmq.stream.impl.TestUtils.BrokerVersionAtLeast;
 import io.netty.channel.EventLoopGroup;
 import java.util.UUID;
 import java.util.concurrent.CountDownLatch;
@@ -61,15 +61,14 @@ void init(TestInfo info) throws Exception {
   void tearDown() throws Exception {
     environment.close();
     if (routingKeys == null) {
-      //      deleteSuperStreamTopology(connection, superStream, partitions);
+      deleteSuperStreamTopology(connection, superStream, partitions);
     } else {
-      //      deleteSuperStreamTopology(connection, superStream, routingKeys);
+      deleteSuperStreamTopology(connection, superStream, routingKeys);
     }
     connection.close();
   }
 
   @Test
-  @BrokerVersionAtLeast("3.9.6")
   void allMessagesSentWithHashRoutingShouldBeThenConsumed() throws Exception {
     int messageCount = 10_000 * partitions;
     declareSuperStreamTopology(connection, superStream, partitions);
@@ -112,7 +111,6 @@ void allMessagesSentWithHashRoutingShouldBeThenConsumed() throws Exception {
   }
 
   @Test
-  @BrokerVersionAtLeast("3.9.6")
   void allMessagesSentWithRoutingKeyRoutingShouldBeThenConsumed() throws Exception {
     int messageCount = 10_000 * partitions;
     routingKeys = new String[] {"amer", "emea", "apac"};