Add cluster tests against restarted stream

A stream restarts when a node it has a member goes down,
which happens on every cluster upgrade.

Author: acogoluegnes

src/test/java/com/rabbitmq/client/amqp/impl/Assertions.java

@@ -230,6 +230,17 @@ ConnectionAssert hasNodename(String nodename) {
       return this;
     }
 
+    ConnectionAssert isOnLeader(Management.QueueInfo info) {
+      Assert.notNull(info, "Queue info cannot be null");
+      String actualLeader = info.leader();
+      if (!actualLeader.equals(actual.connectionNodename())) {
+        fail(
+            "Connection is expected to be on leader node '%s' but is on '%s'",
+            actualLeader, actual.connectionNodename());
+      }
+      return this;
+    }
+
     ConnectionAssert isOnFollower(Management.QueueInfo info) {
       Assert.notNull(info, "Queue info cannot be null");
       List<String> followers =

src/test/java/com/rabbitmq/client/amqp/impl/Cli.java

@@ -62,6 +62,15 @@ static String rabbitmqQueuesCommand() {
     return rabbitmqctl.substring(0, lastIndex) + "rabbitmq-queues";
   }
 
+  static String rabbitmqStreamsCommand() {
+    String rabbitmqctl = rabbitmqctlCommand();
+    int lastIndex = rabbitmqctl.lastIndexOf("rabbitmqctl");
+    if (lastIndex == -1) {
+      throw new IllegalArgumentException("Not a valid rabbitqmctl command: " + rabbitmqctl);
+    }
+    return rabbitmqctl.substring(0, lastIndex) + "rabbitmq-streams";
+  }
+
   static ProcessState rabbitmqctl(String command) {
     return executeCommand(rabbitmqctlCommand() + " " + command);
   }
@@ -70,6 +79,10 @@ static ProcessState rabbitmqQueues(String command) {
     return executeCommand(rabbitmqQueuesCommand() + " " + command);
   }
 
+  static ProcessState rabbitmqStreams(String command) {
+    return executeCommand(rabbitmqStreamsCommand() + " " + command);
+  }
+
   static ProcessState rabbitmqctlIgnoreError(String command) {
     return executeCommand(rabbitmqctlCommand() + " " + command, true);
   }
@@ -198,6 +211,18 @@ static void deleteQuorumQueueMember(String queue, String node) {
     rabbitmqQueues(" delete_member " + queue + " " + node);
   }
 
+  static void addStreamMember(String stream, String node) {
+    rabbitmqStreams(" add_replica " + stream + " " + node);
+  }
+
+  static void deleteStreamMember(String stream, String node) {
+    rabbitmqStreams(" delete_replica " + stream + " " + node);
+  }
+
+  static void restartStream(String stream) {
+    rabbitmqStreams(" restart_stream " + stream);
+  }
+
   static List<ConnectionInfo> listConnections() {
     String output = rabbitmqctl("list_connections -q pid peer_port client_properties").output();
     // output (header line presence depends on broker version):

src/test/java/com/rabbitmq/client/amqp/impl/ClusterTest.java

@@ -21,6 +21,7 @@
 import static com.rabbitmq.client.amqp.ConnectionSettings.Affinity.Operation.PUBLISH;
 import static com.rabbitmq.client.amqp.impl.Assertions.assertThat;
 import static com.rabbitmq.client.amqp.impl.TestUtils.sync;
+import static com.rabbitmq.client.amqp.impl.TestUtils.waitAtMost;
 import static java.time.Duration.ofMillis;
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -211,6 +212,90 @@ void consumeFromMovingQq() {
     }
   }
 
+  @Test
+  void publishToRestartedStream() {
+    try {
+      management.queue(q).type(Management.QueueType.STREAM).declare();
+
+      AmqpConnection publishConnection = connection(b -> b.affinity().queue(q).operation(PUBLISH));
+      assertThat(publishConnection).isOnLeader(queueInfo());
+
+      Publisher publisher = publishConnection.publisherBuilder().queue(q).build();
+      Sync publishSync = sync();
+      publisher.publish(publisher.message().messageId(1L), ctx -> publishSync.down());
+      assertThat(publishSync).completes();
+
+      restartStream();
+
+      publishSync.reset();
+      publisher.publish(publisher.message().messageId(2L), ctx -> publishSync.down());
+      assertThat(publishSync).completes();
+
+      int messageCount = 2;
+      waitAtMost(() -> queueInfo().messageCount() == messageCount);
+      Sync consumeSync = sync(messageCount);
+      Set<Long> messageIds = ConcurrentHashMap.newKeySet(messageCount);
+      connection.consumerBuilder().queue(q).stream()
+          .offset(ConsumerBuilder.StreamOffsetSpecification.FIRST)
+          .builder()
+          .messageHandler(
+              (ctx, msg) -> {
+                messageIds.add(msg.messageIdAsLong());
+                consumeSync.down();
+                ctx.accept();
+              })
+          .build();
+      assertThat(consumeSync).completes();
+      assertThat(messageIds).containsExactlyInAnyOrder(1L, 2L);
+    } finally {
+      management.queueDeletion().delete(q);
+    }
+  }
+
+  @Test
+  void consumeFromRestartedStream() {
+    try {
+      management.queue(q).type(Management.QueueType.STREAM).declare();
+
+      AmqpConnection consumeConnection = connection(b -> b.affinity().queue(q).operation(CONSUME));
+      assertThat(consumeConnection).isOnFollower(queueInfo());
+
+      Set<Long> messageIds = ConcurrentHashMap.newKeySet();
+      Sync consumeSync = sync();
+      consumeConnection.consumerBuilder().queue(q).stream()
+          .offset(ConsumerBuilder.StreamOffsetSpecification.FIRST)
+          .builder()
+          .messageHandler(
+              (ctx, msg) -> {
+                messageIds.add(msg.messageIdAsLong());
+                consumeSync.down();
+                ctx.accept();
+              })
+          .build();
+
+      Publisher publisher = connection.publisherBuilder().queue(q).build();
+      Sync publishSync = sync();
+      publisher.publish(publisher.message().messageId(1L), ctx -> publishSync.down());
+      assertThat(publishSync).completes();
+      publishSync.reset();
+
+      assertThat(consumeSync).completes();
+      assertThat(messageIds).containsExactlyInAnyOrder(1L);
+      consumeSync.reset();
+
+      restartStream();
+
+      publisher.publish(publisher.message().messageId(2L), ctx -> publishSync.down());
+      assertThat(publishSync).completes();
+      publishSync.reset();
+
+      assertThat(consumeSync).completes();
+      assertThat(messageIds).containsExactlyInAnyOrder(1L, 2L);
+    } finally {
+      management.queueDeletion().delete(q);
+    }
+  }
+
   String moveQqLeader() {
     String initialLeader = deleteQqLeader();
     addQqMember(initialLeader);
@@ -220,23 +305,47 @@ String moveQqLeader() {
   }
 
   String deleteQqLeader() {
+    return deleteLeader(this::deleteQqMember);
+  }
+
+  String deleteStreamLeader() {
+    return deleteLeader(leader -> Cli.deleteStreamMember(q, leader));
+  }
+
+  String deleteLeader(Consumer<String> deleteMemberOperation) {
     Management.QueueInfo info = queueInfo();
     String initialLeader = info.leader();
     int initialReplicaCount = info.replicas().size();
-    deleteQqMember(initialLeader);
+    deleteMemberOperation.accept(initialLeader);
     TestUtils.waitAtMost(() -> !queueInfo().leader().equals(initialLeader));
     assertThat(queueInfo().replicas()).hasSize(initialReplicaCount - 1);
     return initialLeader;
   }
 
+  void restartStream() {
+    Cli.restartStream(this.q);
+  }
+
   void deleteQqMember(String member) {
     Cli.deleteQuorumQueueMember(q, member);
   }
 
+  void deleteStreamMember(String member) {
+    Cli.deleteStreamMember(q, member);
+  }
+
   void addQqMember(String newMember) {
+    addMember(() -> Cli.addQuorumQueueMember(q, newMember));
+  }
+
+  void addStreamMember(String newMember) {
+    addMember(() -> Cli.addStreamMember(q, newMember));
+  }
+
+  void addMember(Runnable addMemberOperation) {
     Management.QueueInfo info = queueInfo();
     int initialReplicaCount = info.replicas().size();
-    Cli.addQuorumQueueMember(q, newMember);
+    addMemberOperation.run();
     TestUtils.waitAtMost(() -> queueInfo().replicas().size() == initialReplicaCount + 1);
   }