Release management if connection gets closed in affinity lookup

Author: acogoluegnes

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

@@ -468,7 +468,10 @@ private CompletableFuture<NativeConnectionWrapper> recoverNativeConnection(
                       public <T> T maybeRetry(Supplier<T> task) {
                         return RetryUtils.callAndMaybeRetry(
                             task::get,
-                            e -> true,
+                            // no need to retry if the connection is closed
+                            // the affinity task will fail and AsyncRetry will take care
+                            // of retrying later
+                            e -> RECOVERY_PREDICATE.negate().test(e),
                             Duration.ofMillis(10),
                             5,
                             "Connection affinity operation");

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

@@ -60,6 +60,7 @@ static AmqpConnection.NativeConnectionWrapper enforceAffinity(
       // no affinity asked, we create a connection and return it
       return retryStrategy.maybeRetry(() -> connectionFactory.apply(null));
     }
+    AmqpConnection.NativeConnectionWrapper connectionWrapper = null;
     try {
       AmqpConnection.NativeConnectionWrapper pickedConnection = null;
       int attemptCount = 0;
@@ -68,7 +69,7 @@ static AmqpConnection.NativeConnectionWrapper enforceAffinity(
       Management.QueueInfo info = affinityCache.queueInfo(context.queue());
       while (pickedConnection == null) {
         attemptCount++;
-        AmqpConnection.NativeConnectionWrapper connectionWrapper = null;
+        connectionWrapper = null;
         if (info == null) {
           connectionWrapper = retryStrategy.maybeRetry(() -> connectionFactory.apply(null));
           info = lookUpQueueInfo(management, context, affinityCache, retryStrategy);
@@ -147,6 +148,23 @@ static AmqpConnection.NativeConnectionWrapper enforceAffinity(
         }
       }
       return pickedConnection;
+    } catch (AmqpException.AmqpConnectionException e) {
+      management.releaseResources();
+      try {
+        if (connectionWrapper != null) {
+          connectionWrapper.connection().close();
+        }
+      } catch (Exception ex) {
+        LOGGER.debug(
+            "Error while closing native connection while enforcing affinity: {}", ex.getMessage());
+      }
+      management.markUnavailable();
+      LOGGER.warn(
+          "Cannot enforce affinity {} of '{}' because connection has been closed",
+          context,
+          connectionName,
+          e);
+      throw e;
     } catch (RuntimeException e) {
       LOGGER.warn(
           "Cannot enforce affinity {} of '{}' error when looking up queue",

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

@@ -19,6 +19,7 @@
 
 import static com.rabbitmq.client.amqp.Management.QueueType.QUORUM;
 import static com.rabbitmq.client.amqp.Management.QueueType.STREAM;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.assertj.core.api.Assertions.fail;
 import static org.mockito.ArgumentMatchers.anyList;
 import static org.mockito.Mockito.*;
@@ -201,6 +202,16 @@ void useAddressReturnedByCache() {
         .hasMapping(LEADER_NODENAME, LEADER_ADDRESS);
   }
 
+  @Test
+  void connectionClosedExceptionDuringQueueInfoLookupShouldMarkManagementAsUnavailable() {
+    when(management.queueInfo(Q)).thenThrow(new AmqpException.AmqpConnectionException("", null));
+    when(cf.apply(null)).thenReturn(follower1Connection());
+    assertThatThrownBy(() -> enforceAffinity(cf, management, affinity(), cache))
+        .isInstanceOf(AmqpException.AmqpConnectionException.class);
+    verify(management, times(1)).markUnavailable();
+    verify(nativeConnection, times(1)).close();
+  }
+
   AmqpConnection.NativeConnectionWrapper leaderConnection() {
     return new AmqpConnection.NativeConnectionWrapper(
         this.nativeConnection, LEADER_NODENAME, LEADER_ADDRESS);
@@ -216,11 +227,6 @@ AmqpConnection.NativeConnectionWrapper follower2Connection() {
         this.nativeConnection, FOLLOWER2_NODENAME, FOLLOWER2_ADDRESS);
   }
 
-  AmqpConnection.NativeConnectionWrapper connection(String nodename) {
-    return new AmqpConnection.NativeConnectionWrapper(
-        this.nativeConnection, nodename, NODES.get(nodename));
-  }
-
   static ConnectionUtils.AffinityContext affinity() {
     return new ConnectionUtils.AffinityContext(Q, ConnectionSettings.Affinity.Operation.PUBLISH);
   }

src/test/resources/logback-test.xml

@@ -21,7 +21,7 @@
     <logger name="com.rabbitmq.client.amqp.impl.AmqpManagement" level="debug" />
     <logger name="com.rabbitmq.client.amqp.impl.AmqpPublisher" level="warn" />
     <logger name="com.rabbitmq.client.amqp.impl.AmqpConsumer" level="warn" />
-    <logger name="com.rabbitmq.client.amqp.impl.RetryUtils" level="warn" />
+    <logger name="com.rabbitmq.client.amqp.impl.RetryUtils" level="debug" />
 
     <root level="warn">
         <appender-ref ref="STDOUT" />