Add retry for affinity during recovery

Author: acogoluegnes

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

@@ -136,7 +136,8 @@ final class AmqpConnection extends ResourceBase implements Connection {
             this.management,
             this.affinity,
             this.environment.affinityCache(),
-            this.affinityStrategy);
+            this.affinityStrategy,
+            ConnectionUtils.NO_RETRY_STRATEGY);
     this.sync(ncw);
     this.state(OPEN);
     this.environment.metricsCollector().openConnection();
@@ -404,7 +405,17 @@ private NativeConnectionWrapper recoverNativeConnection(
                 this.management,
                 this.affinity,
                 this.environment.affinityCache(),
-                this.affinityStrategy);
+                this.affinityStrategy,
+                new ConnectionUtils.RetryStrategy() {
+                  @Override
+                  public <T> T maybeRetry(Supplier<T> task) {
+                    return RetryUtils.callAndMaybeRetry(
+                        task::get,
+                        e -> true,
+                        recoveryConfiguration.backOffDelayPolicy(),
+                        "Connection affinity operation");
+                  }
+                });
         return result;
       } catch (Exception ex) {
         LOGGER.info("Error while trying to recover connection", ex);

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

@@ -27,12 +27,15 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import java.util.function.Function;
+import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 final class ConnectionUtils {
 
+  static final RetryStrategy NO_RETRY_STRATEGY = Supplier::get;
+
   static final ConnectionSettings.AffinityStrategy
       LEADER_FOR_PUBLISHING_FOLLOWERS_FOR_CONSUMING_STRATEGY =
           new LeaderForPublishingFollowersForConsumingStrategy();
@@ -50,11 +53,11 @@ static AmqpConnection.NativeConnectionWrapper enforceAffinity(
       AmqpManagement management,
       AffinityContext context,
       AffinityCache affinityCache,
-      ConnectionSettings.AffinityStrategy strategy) {
-    // TODO add retry for sensitive operations in affinity mechanism
+      ConnectionSettings.AffinityStrategy strategy,
+      RetryStrategy retryStrategy) {
     if (context == null) {
       // no affinity asked, we create a connection and return it
-      return connectionFactory.apply(null);
+      return retryStrategy.maybeRetry(() -> connectionFactory.apply(null));
     }
     try {
       AmqpConnection.NativeConnectionWrapper pickedConnection = null;
@@ -66,8 +69,8 @@ static AmqpConnection.NativeConnectionWrapper enforceAffinity(
         attemptCount++;
         AmqpConnection.NativeConnectionWrapper connectionWrapper = null;
         if (info == null) {
-          connectionWrapper = connectionFactory.apply(null);
-          info = lookUpQueueInfo(management, context, affinityCache);
+          connectionWrapper = retryStrategy.maybeRetry(() -> connectionFactory.apply(null));
+          info = lookUpQueueInfo(management, context, affinityCache, retryStrategy);
           queueInfoRefreshed = true;
         }
         if (info == null) {
@@ -89,7 +92,7 @@ static AmqpConnection.NativeConnectionWrapper enforceAffinity(
                   .map(affinityCache::nodenameToAddress)
                   .filter(Objects::nonNull)
                   .collect(Collectors.toList());
-          connectionWrapper = connectionFactory.apply(addressHints);
+          connectionWrapper = retryStrategy.maybeRetry(() -> connectionFactory.apply(addressHints));
         }
         LOGGER.debug("Nodes matching affinity {}: {}.", context, nodesWithAffinity);
         LOGGER.debug("Currently connected to node {}.", connectionWrapper.nodename());
@@ -98,7 +101,7 @@ static AmqpConnection.NativeConnectionWrapper enforceAffinity(
           if (!queueInfoRefreshed) {
             LOGGER.debug(
                 "Found affinity, but refreshing queue information to check affinity is still valid.");
-            info = lookUpQueueInfo(management, context, affinityCache);
+            info = lookUpQueueInfo(management, context, affinityCache, retryStrategy);
             if (info == null) {
               LOGGER.debug("Could not look up info for queue '{}'", context.queue());
               pickedConnection = connectionWrapper;
@@ -129,7 +132,7 @@ static AmqpConnection.NativeConnectionWrapper enforceAffinity(
           LOGGER.debug(
               "Affinity {} not found with node {}.", context, connectionWrapper.nodename());
           if (!queueInfoRefreshed) {
-            info = lookUpQueueInfo(management, context, affinityCache);
+            info = lookUpQueueInfo(management, context, affinityCache, retryStrategy);
             if (info != null) {
               nodesWithAffinity = strategy.nodesWithAffinity(context, info);
               queueInfoRefreshed = true;
@@ -147,18 +150,24 @@ static AmqpConnection.NativeConnectionWrapper enforceAffinity(
   }
 
   private static Management.QueueInfo lookUpQueueInfo(
-      AmqpManagement management, AffinityContext affinity, AffinityCache cache) {
-    Management.QueueInfo info = null;
-    management.init();
-    try {
-      info = management.queueInfo(affinity.queue());
-      cache.queueInfo(info);
-    } catch (AmqpException.AmqpEntityDoesNotExistException e) {
-      LOGGER.debug("Queue '{}' does not exist.", affinity.queue());
-      cache.clearQueueInfoEntry(affinity.queue());
-      // we just return null, caller will have to return the last connection
-    }
-    return info;
+      AmqpManagement management,
+      AffinityContext affinity,
+      AffinityCache cache,
+      RetryStrategy retryStrategy) {
+    return retryStrategy.maybeRetry(
+        () -> {
+          Management.QueueInfo info = null;
+          management.init();
+          try {
+            info = management.queueInfo(affinity.queue());
+            cache.queueInfo(info);
+          } catch (AmqpException.AmqpEntityDoesNotExistException e) {
+            LOGGER.debug("Queue '{}' does not exist.", affinity.queue());
+            cache.clearQueueInfoEntry(affinity.queue());
+            // we just return null, caller will have to return the last connection
+          }
+          return info;
+        });
   }
 
   // TODO clean affinity cache (LRU or size-based)
@@ -298,4 +307,10 @@ public List<String> nodesWithAffinity(
       return nodesWithAffinity;
     }
   }
+
+  @FunctionalInterface
+  interface RetryStrategy {
+
+    <T> T maybeRetry(Supplier<T> task);
+  }
 }

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

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