Move affinity algorithm in dedicated class

It is isolated enough to be outside of the AMQP connection class and
tested independently.

Author: acogoluegnes

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

@@ -30,10 +30,8 @@
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BiConsumer;
-import java.util.function.Function;
 import java.util.function.Predicate;
 import java.util.function.Supplier;
-import java.util.stream.Collectors;
 import org.apache.qpid.protonj2.client.ConnectionOptions;
 import org.apache.qpid.protonj2.client.DisconnectionEvent;
 import org.apache.qpid.protonj2.client.Session;
@@ -123,19 +121,17 @@ final class AmqpConnection extends ResourceBase implements Connection {
             : null;
     this.management = createManagement();
     NativeConnectionWrapper ncw =
-        enforceAffinity(
+        ConnectionUtils.enforceAffinity(
             addrs -> {
               NativeConnectionWrapper wrapper =
                   connect(this.connectionSettings, builder.name(), disconnectHandler, addrs);
-              this.nativeConnection = wrapper.connection;
+              this.nativeConnection = wrapper.connection();
               return wrapper;
             },
             this.management,
             this.affinity,
             this.environment.affinityCache());
-    this.connectionAddress = ncw.address;
-    this.connectionNodename = ncw.nodename;
-    this.nativeConnection = ncw.connection;
+    this.sync(ncw);
     this.state(OPEN);
     this.environment.metricsCollector().openConnection();
   }
@@ -235,96 +231,10 @@ private NativeConnectionWrapper connect(
     }
   }
 
-  static NativeConnectionWrapper enforceAffinity(
-      Function<List<Address>, NativeConnectionWrapper> connectionFactory,
-      AmqpManagement management,
-      ConnectionUtils.ConnectionAffinity affinity,
-      ConnectionUtils.AffinityCache affinityCache) {
-    // TODO add retry for sensitive operations in affinity mechanism
-    if (affinity == null) {
-      return connectionFactory.apply(null);
-    }
-    try {
-      NativeConnectionWrapper pickedConnection = null;
-      int attemptCount = 0;
-      boolean queueInfoRefreshed = false;
-      List<String> nodesWithAffinity = null;
-      Management.QueueInfo info = affinityCache.queueInfo(affinity.queue());
-      while (pickedConnection == null) {
-        attemptCount++;
-        NativeConnectionWrapper connectionWrapper = null;
-        if (info == null) {
-          connectionWrapper = connectionFactory.apply(null);
-          management.init();
-          info = management.queueInfo(affinity.queue());
-          affinityCache.queueInfo(info);
-          queueInfoRefreshed = true;
-        }
-        LOGGER.debug(
-            "Looking affinity with queue '{}' (type = {}, leader = {}, replicas = {})",
-            info.name(),
-            info.type(),
-            info.leader(),
-            info.replicas());
-        if (nodesWithAffinity == null) {
-          nodesWithAffinity = ConnectionUtils.findAffinity(affinity, info);
-          LOGGER.debug("Nodes matching affinity {}: {}", affinity, nodesWithAffinity);
-        }
-        if (connectionWrapper == null) {
-          List<Address> addressHints =
-              nodesWithAffinity.stream()
-                  .map(affinityCache::nodenameToAddress)
-                  .filter(Objects::nonNull)
-                  .collect(Collectors.toList());
-          connectionWrapper = connectionFactory.apply(addressHints);
-        }
-        LOGGER.debug("Currently connected to node {}", connectionWrapper.nodename);
-        affinityCache.nodenameToAddress(connectionWrapper.nodename, connectionWrapper.address);
-        if (nodesWithAffinity.contains(connectionWrapper.nodename)) {
-          LOGGER.debug("Affinity {} found with node {}", affinity, connectionWrapper.nodename);
-          if (!queueInfoRefreshed) {
-            LOGGER.debug("Refreshing queue information.");
-            management.init();
-            info = management.queueInfo(affinity.queue());
-            affinityCache.queueInfo(info);
-            nodesWithAffinity = ConnectionUtils.findAffinity(affinity, info);
-            LOGGER.debug("Nodes matching affinity {}: {}", affinity, nodesWithAffinity);
-            queueInfoRefreshed = true;
-            if (nodesWithAffinity.contains(connectionWrapper.nodename)) {
-              pickedConnection = connectionWrapper;
-            } else {
-              management.releaseResources();
-              connectionWrapper.connection.close();
-            }
-          } else {
-            pickedConnection = connectionWrapper;
-          }
-        } else if (attemptCount == 5) {
-          LOGGER.debug(
-              "Could not find affinity {} after {} attempt(s), using last connection.",
-              affinity,
-              attemptCount);
-          pickedConnection = connectionWrapper;
-        } else {
-          LOGGER.debug("Affinity {} not found with node {}.", affinity, connectionWrapper.nodename);
-          if (!queueInfoRefreshed) {
-            LOGGER.debug("Refreshing queue information.");
-            management.init();
-            info = management.queueInfo(affinity.queue());
-            affinityCache.queueInfo(info);
-            nodesWithAffinity = ConnectionUtils.findAffinity(affinity, info);
-            LOGGER.debug("Nodes matching affinity {}: {}", affinity, nodesWithAffinity);
-            queueInfoRefreshed = true;
-          }
-          management.releaseResources();
-          connectionWrapper.connection.close();
-        }
-      }
-      return pickedConnection;
-    } catch (RuntimeException e) {
-      LOGGER.warn("Cannot enforce affinity {} of error when looking up queue", affinity, e);
-      throw e;
-    }
+  private void sync(NativeConnectionWrapper wrapper) {
+    this.connectionAddress = wrapper.address();
+    this.connectionNodename = wrapper.nodename();
+    this.nativeConnection = wrapper.connection();
   }
 
   private static void checkBrokerVersion(org.apache.qpid.protonj2.client.Connection connection)
@@ -419,9 +329,7 @@ private void recoverAfterConnectionFailure(
       NativeConnectionWrapper ncw =
           recoverNativeConnection(
               recoveryConfiguration, connectionName, disconnectedHandlerReference);
-      this.connectionAddress = ncw.address;
-      this.connectionNodename = ncw.nodename;
-      this.nativeConnection = ncw.connection;
+      this.sync(ncw);
       LOGGER.debug("Reconnected to {}", this.currentConnectionLabel());
     } catch (Exception ex) {
       if (ex instanceof InterruptedException) {
@@ -476,15 +384,15 @@ private NativeConnectionWrapper recoverNativeConnection(
 
       try {
         NativeConnectionWrapper result =
-            enforceAffinity(
+            ConnectionUtils.enforceAffinity(
                 addrs -> {
                   NativeConnectionWrapper wrapper =
                       connect(
                           this.connectionSettings,
                           connectionName,
                           disconnectedHandlerReference.get(),
                           addrs);
-                  this.nativeConnection = wrapper.connection;
+                  this.nativeConnection = wrapper.connection();
                   return wrapper;
                 },
                 this.management,
@@ -772,9 +680,13 @@ String nodename() {
       return this.nodename;
     }
 
-    public Address address() {
+    Address address() {
       return this.address;
     }
+
+    org.apache.qpid.protonj2.client.Connection connection() {
+      return this.connection;
+    }
   }
 
   @Override

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

@@ -25,6 +25,7 @@
 import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -35,6 +36,103 @@ final class ConnectionUtils {
 
   private ConnectionUtils() {}
 
+  static AmqpConnection.NativeConnectionWrapper enforceAffinity(
+      Function<List<Address>, AmqpConnection.NativeConnectionWrapper> connectionFactory,
+      AmqpManagement management,
+      ConnectionAffinity affinity,
+      AffinityCache affinityCache) {
+    // TODO add retry for sensitive operations in affinity mechanism
+    if (affinity == null) {
+      return connectionFactory.apply(null);
+    }
+    try {
+      AmqpConnection.NativeConnectionWrapper pickedConnection = null;
+      int attemptCount = 0;
+      boolean queueInfoRefreshed = false;
+      List<String> nodesWithAffinity = null;
+      Management.QueueInfo info = affinityCache.queueInfo(affinity.queue());
+      while (pickedConnection == null) {
+        attemptCount++;
+        AmqpConnection.NativeConnectionWrapper connectionWrapper = null;
+        if (info == null) {
+          connectionWrapper = connectionFactory.apply(null);
+          info = lookUpQueueInfo(management, affinity, affinityCache);
+          queueInfoRefreshed = true;
+        }
+        LOGGER.debug(
+            "Looking affinity with queue '{}' (type = {}, leader = {}, replicas = {})",
+            info.name(),
+            info.type(),
+            info.leader(),
+            info.replicas());
+        if (nodesWithAffinity == null) {
+          nodesWithAffinity = findAffinity(affinity, info);
+        }
+        if (connectionWrapper == null) {
+          List<Address> addressHints =
+              nodesWithAffinity.stream()
+                  .map(affinityCache::nodenameToAddress)
+                  .filter(Objects::nonNull)
+                  .collect(Collectors.toList());
+          connectionWrapper = connectionFactory.apply(addressHints);
+        }
+        LOGGER.debug("Nodes matching affinity {}: {}.", affinity, nodesWithAffinity);
+        LOGGER.debug("Currently connected to node {}.", connectionWrapper.nodename());
+        affinityCache.nodenameToAddress(connectionWrapper.nodename(), connectionWrapper.address());
+        if (nodesWithAffinity.contains(connectionWrapper.nodename())) {
+          if (!queueInfoRefreshed) {
+            info = lookUpQueueInfo(management, affinity, affinityCache);
+            LOGGER.debug(
+                "Found affinity, but refreshing queue information to check affinity is still valid.");
+            nodesWithAffinity = findAffinity(affinity, info);
+            queueInfoRefreshed = true;
+            if (nodesWithAffinity.contains(connectionWrapper.nodename())) {
+              pickedConnection = connectionWrapper;
+            } else {
+              LOGGER.debug("Affinity no longer valid, retrying.");
+              management.releaseResources();
+              connectionWrapper.connection().close();
+            }
+          } else {
+            pickedConnection = connectionWrapper;
+          }
+          if (pickedConnection != null) {
+            LOGGER.debug(
+                "Affinity found with node {}, returning connection", pickedConnection.nodename());
+          }
+        } else if (attemptCount == 5) {
+          LOGGER.debug(
+              "Could not find affinity {} after {} attempt(s), using last connection.",
+              affinity,
+              attemptCount);
+          pickedConnection = connectionWrapper;
+        } else {
+          LOGGER.debug(
+              "Affinity {} not found with node {}.", affinity, connectionWrapper.nodename());
+          if (!queueInfoRefreshed) {
+            info = lookUpQueueInfo(management, affinity, affinityCache);
+            nodesWithAffinity = findAffinity(affinity, info);
+            queueInfoRefreshed = true;
+          }
+          management.releaseResources();
+          connectionWrapper.connection().close();
+        }
+      }
+      return pickedConnection;
+    } catch (RuntimeException e) {
+      LOGGER.warn("Cannot enforce affinity {} of error when looking up queue", affinity, e);
+      throw e;
+    }
+  }
+
+  private static Management.QueueInfo lookUpQueueInfo(
+      AmqpManagement management, ConnectionAffinity affinity, AffinityCache cache) {
+    management.init();
+    Management.QueueInfo info = management.queueInfo(affinity.queue());
+    cache.queueInfo(info);
+    return info;
+  }
+
   static class AffinityCache {
 
     private final ConcurrentMap<String, Management.QueueInfo> queueInfoCache =

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

@@ -19,7 +19,7 @@
 
 import static com.rabbitmq.client.amqp.Management.QueueType.QUORUM;
 import static com.rabbitmq.client.amqp.Management.QueueType.STREAM;
-import static com.rabbitmq.client.amqp.impl.AmqpConnection.enforceAffinity;
+import static com.rabbitmq.client.amqp.impl.ConnectionUtils.enforceAffinity;
 import static org.assertj.core.api.Assertions.fail;
 import static org.mockito.ArgumentMatchers.anyList;
 import static org.mockito.Mockito.*;