Implement getLastPublishingId in super stream producer

Author: acogoluegnes

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

@@ -187,7 +187,8 @@ public Producer build() {
                   this.stream, this.routingKeyExtractor, this.environment, hashFunction)
               : new RoutingKeyRoutingStrategy(
                   this.stream, this.routingKeyExtractor, this.environment);
-      producer = new SuperStreamProducer(this, stream, routingStrategy, environment);
+      producer =
+          new SuperStreamProducer(this, this.name, this.stream, routingStrategy, this.environment);
     }
     return producer;
   }

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

@@ -18,6 +18,7 @@
 import com.rabbitmq.stream.Message;
 import com.rabbitmq.stream.MessageBuilder;
 import com.rabbitmq.stream.Producer;
+import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.concurrent.ConcurrentHashMap;
@@ -33,15 +34,20 @@ class SuperStreamProducer implements Producer {
   private final String superStream;
   private final Map<String, Producer> producers = new ConcurrentHashMap<>();
   private final StreamProducerBuilder producerBuilder;
+  private final StreamEnvironment environment;
+  private final String name;
 
   SuperStreamProducer(
       StreamProducerBuilder producerBuilder,
+      String name,
       String superStream,
       RoutingStrategy routingStrategy,
       StreamEnvironment streamEnvironment) {
     this.routingStrategy = routingStrategy;
     this.codec = streamEnvironment.codec();
+    this.name = name;
     this.superStream = superStream;
+    this.environment = streamEnvironment;
     this.producerBuilder = producerBuilder.duplicate();
     this.producerBuilder.stream(null);
     this.producerBuilder.resetRouting();
@@ -54,9 +60,25 @@ public MessageBuilder messageBuilder() {
 
   @Override
   public long getLastPublishingId() {
-    // TODO get all partitions for this super stream, query the last publishing ID for each of team,
-    // return the highest (or the lowest, because duplicates will be filtered out anyway?)
-    return 0;
+    if (this.name != null && !this.name.isEmpty()) {
+      List<String> streams = this.environment.locator().partitions(superStream);
+      long publishingId = 0;
+      boolean first = true;
+      for (String partition : streams) {
+        long pubId = this.environment.locator().queryPublisherSequence(this.name, partition);
+        if (first) {
+          publishingId = pubId;
+          first = false;
+        } else {
+          if (Long.compareUnsigned(publishingId, pubId) > 0) {
+            publishingId = pubId;
+          }
+        }
+      }
+      return publishingId;
+    } else {
+      throw new IllegalStateException("The producer has no name");
+    }
   }
 
   @Override

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

@@ -51,6 +51,7 @@ public class SuperStreamProducerTest {
   int partitions = 3;
   String superStream;
   String[] routingKeys = null;
+  TestUtils.ClientFactory cf;
 
   @BeforeEach
   void init(TestInfo info) throws Exception {
@@ -173,4 +174,68 @@ void allMessagesSentToSuperStreamWithRoutingKeyRoutingShouldBeThenConsumed() thr
     assertThat(counts.values().stream().map(AtomicLong::get).reduce(0L, Long::sum))
         .isEqualTo(messageCount);
   }
+
+  @Test
+  void getLastPublishingIdShouldReturnLowestValue() throws Exception {
+    int messageCount = 10_000;
+    declareSuperStreamTopology(connection, superStream, partitions);
+    String producerName = "super-stream-application";
+    Producer producer =
+        environment.producerBuilder().name(producerName).stream(superStream)
+            .routing(message -> message.getProperties().getMessageIdAsString(), RoutingType.HASH)
+            .build();
+
+    CountDownLatch publishLatch = new CountDownLatch(messageCount);
+    IntStream.range(0, messageCount)
+        .forEach(
+            i ->
+                producer.send(
+                    producer
+                        .messageBuilder()
+                        .publishingId(i)
+                        .properties()
+                        .messageId(UUID.randomUUID().toString())
+                        .messageBuilder()
+                        .build(),
+                    confirmationStatus -> publishLatch.countDown()));
+
+    assertThat(latchAssert(publishLatch)).completes(5);
+
+    long lastPublishingId = producer.getLastPublishingId();
+    assertThat(lastPublishingId).isNotZero();
+    Client client = cf.get();
+    IntStream.range(0, partitions)
+        .mapToObj(i -> superStream + "-" + i)
+        .forEach(
+            stream -> {
+              long publishingId = client.queryPublisherSequence(producerName, stream);
+              assertThat(publishingId).isGreaterThanOrEqualTo(lastPublishingId);
+            });
+
+    Map<String, AtomicLong> counts = new ConcurrentHashMap<>();
+    AtomicLong totalCount = new AtomicLong(0);
+    IntStream.range(0, partitions)
+        .mapToObj(i -> superStream + "-" + i)
+        .forEach(
+            stream -> {
+              AtomicLong streamCount = new AtomicLong(0);
+              counts.put(stream, streamCount);
+              environment.consumerBuilder().stream(stream)
+                  .offset(OffsetSpecification.first())
+                  .messageHandler(
+                      (context, message) -> {
+                        streamCount.incrementAndGet();
+                        totalCount.incrementAndGet();
+                      })
+                  .build();
+            });
+
+    waitAtMost(10, () -> totalCount.get() == messageCount);
+
+    assertThat(counts.values().stream().map(AtomicLong::get))
+        .hasSize(partitions)
+        .doesNotContain(0L);
+    assertThat(counts.values().stream().map(AtomicLong::get).reduce(0L, Long::sum))
+        .isEqualTo(messageCount);
+  }
 }