Add cluster restart test with quorum queues

Author: acogoluegnes

src/main/java/com/rabbitmq/client/amqp/impl/ConnectionUtils.java

@@ -33,11 +33,13 @@
 
 final class ConnectionUtils {
 
-  static final ConnectionSettings.AffinityStrategy PREFER_LEADER_FOR_PUBLISHING_STRATEGY =
-      new PreferLeaderForPublishingAffinityStrategy();
+  static final ConnectionSettings.AffinityStrategy
+      LEADER_FOR_PUBLISHING_FOLLOWERS_FOR_CONSUMING_STRATEGY =
+          new LeaderForPublishingFollowersForConsumingStrategy();
 
-  static final ConnectionSettings.AffinityStrategy MEMBER_AFFINITY_STRATEGY =
-      new MemberAffinityStrategy();
+  static final ConnectionSettings.AffinityStrategy
+      LEADER_FOR_PUBLISHING_MEMBERS_FOR_CONSUMING_STRATEGY =
+          new LeaderForPublishingMembersForConsumingStrategy();
 
   private static final Logger LOGGER = LoggerFactory.getLogger(ConnectionUtils.class);
 
@@ -229,18 +231,28 @@ public int hashCode() {
     }
   }
 
-  static class MemberAffinityStrategy implements ConnectionSettings.AffinityStrategy {
+  static class LeaderForPublishingMembersForConsumingStrategy
+      implements ConnectionSettings.AffinityStrategy {
 
     @Override
     public List<String> nodesWithAffinity(
         ConnectionSettings.AffinityContext context, Management.QueueInfo info) {
-      return (info.replicas() == null || info.replicas().isEmpty())
-          ? Collections.emptyList()
-          : List.copyOf(info.replicas());
+      List<String> nodesWithAffinity =
+          (info.replicas() == null || info.replicas().isEmpty())
+              ? Collections.emptyList()
+              : List.copyOf(info.replicas());
+      if (context.operation() == ConnectionSettings.Affinity.Operation.PUBLISH) {
+        if (info.leader() != null && !info.leader().isBlank()) {
+          nodesWithAffinity = List.of(info.leader());
+        } else {
+          nodesWithAffinity = Collections.emptyList();
+        }
+      }
+      return nodesWithAffinity;
     }
   }
 
-  static class PreferLeaderForPublishingAffinityStrategy
+  static class LeaderForPublishingFollowersForConsumingStrategy
       implements ConnectionSettings.AffinityStrategy {
 
     @Override

src/main/java/com/rabbitmq/client/amqp/impl/DefaultConnectionSettings.java

@@ -413,7 +413,8 @@ static class DefaultAffinity<T> implements Affinity<T> {
     private String queue;
     private Operation operation;
     private boolean reuse = false;
-    private AffinityStrategy strategy = ConnectionUtils.MEMBER_AFFINITY_STRATEGY;
+    private AffinityStrategy strategy =
+        ConnectionUtils.LEADER_FOR_PUBLISHING_MEMBERS_FOR_CONSUMING_STRATEGY;
 
     DefaultAffinity(DefaultConnectionSettings<T> connectionSettings) {
       this.connectionSettings = connectionSettings;

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

@@ -443,6 +443,6 @@ private static AmqpConnection.NativeConnectionWrapper enforceAffinity(
         management,
         context,
         affinityCache,
-        ConnectionUtils.PREFER_LEADER_FOR_PUBLISHING_STRATEGY);
+        ConnectionUtils.LEADER_FOR_PUBLISHING_FOLLOWERS_FOR_CONSUMING_STRATEGY);
   }
 }

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

@@ -88,14 +88,9 @@ SyncAssert completes() {
     }
 
     SyncAssert completes(Duration timeout) {
-      try {
-        boolean completed = actual.await(timeout);
-        if (!completed) {
-          fail("Sync timed out after %d ms", timeout.toMillis());
-        }
-      } catch (InterruptedException e) {
-        Thread.interrupted();
-        throw new RuntimeException(e);
+      boolean completed = actual.await(timeout);
+      if (!completed) {
+        fail("Sync timed out after %d ms", timeout.toMillis());
       }
       return this;
     }

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

@@ -32,6 +32,7 @@
 import java.util.function.Predicate;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
+import java.util.stream.Collectors;
 
 abstract class Cli {
 
@@ -273,6 +274,26 @@ static void restartNode(String node) {
     executeCommand(dockerCommand + "rabbitmqctl status");
   }
 
+  static void rebalance() {
+    rabbitmqQueues("rebalance all");
+  }
+
+  static List<String> nodes() {
+    List<String> clusterNodes = new ArrayList<>();
+    clusterNodes.add(rabbitmqctl("eval 'node().'").output().trim());
+    List<String> nodes =
+        Arrays.stream(
+                rabbitmqctl("eval 'nodes().'")
+                    .output()
+                    .replace("[", "")
+                    .replace("]", "")
+                    .split(","))
+            .map(String::trim)
+            .collect(Collectors.toList());
+    clusterNodes.addAll(nodes);
+    return List.copyOf(clusterNodes);
+  }
+
   private static String nodeToDockerContainer(String node) {
     String containerId = DOCKER_NODES_TO_CONTAINERS.get(node);
     Assert.notNull(containerId, "No container for node " + node);

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

@@ -447,7 +447,9 @@ void connectionShouldBeOnOwningNodeWhenAffinityIsActivatedForClassicQueues(TestI
                           b ->
                               b.uri(uri)
                                   .affinity()
-                                  .strategy(ConnectionUtils.MEMBER_AFFINITY_STRATEGY)))
+                                  .strategy(
+                                      ConnectionUtils
+                                          .LEADER_FOR_PUBLISHING_MEMBERS_FOR_CONSUMING_STRATEGY)))
               .collect(toList());
       List<Management.QueueInfo> queueInfos =
           range(0, URIS.length)
@@ -541,7 +543,7 @@ AmqpConnection connection(Consumer<AmqpConnectionBuilder> operation) {
         .backOffDelayPolicy(BACK_OFF_DELAY_POLICY)
         .connectionBuilder()
         .affinity()
-        .strategy(ConnectionUtils.PREFER_LEADER_FOR_PUBLISHING_STRATEGY);
+        .strategy(ConnectionUtils.LEADER_FOR_PUBLISHING_FOLLOWERS_FOR_CONSUMING_STRATEGY);
     operation.accept(builder);
     return (AmqpConnection) builder.build();
   }