Fix coordinator manager creation in load balancer mode

src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java

@@ -635,8 +635,10 @@ private ClientSubscriptionsManager(Broker node, Client.ClientParameters clientPa
           };
       ShutdownListener shutdownListener =
           shutdownContext -> {
-            this.closed.set(true);
-            managers.remove(this);
+            if (clientInitializedInManager.get()) {
+              this.closed.set(true);
+              managers.remove(this);
+            }
             if (shutdownContext.isShutdownUnexpected()) {
               LOGGER.debug(
                   "Unexpected shutdown notification on subscription connection {}, scheduling consumers re-assignment",
@@ -906,9 +908,15 @@ synchronized void add(
         OffsetSpecification offsetSpecification,
         boolean isInitialSubscription) {
       if (this.isFull()) {
+        LOGGER.debug(
+            "Cannot add subscription tracker for stream '{}', manager is full",
+            subscriptionTracker.stream);
         throw new IllegalStateException("Cannot add subscription tracker, the manager is full");
       }
       if (this.isClosed()) {
+        LOGGER.debug(
+            "Cannot add subscription tracker for stream '{}', manager is closed",
+            subscriptionTracker.stream);
         throw new IllegalStateException("Cannot add subscription tracker, the manager is closed");
       }
 

src/main/java/com/rabbitmq/stream/impl/ProducersCoordinator.java

@@ -599,7 +599,10 @@ private ClientProducersManager(
           };
       ShutdownListener shutdownListener =
           shutdownContext -> {
-            managers.remove(this);
+            if (clientInitializedInManager.get()) {
+              this.closed.set(true);
+              managers.remove(this);
+            }
             if (shutdownContext.isShutdownUnexpected()) {
               LOGGER.debug(
                   "Recovering {} producer(s) after unexpected connection termination",

src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java

@@ -67,6 +67,7 @@
 import org.mockito.Captor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
+import org.mockito.stubbing.Answer;
 
 public class ConsumersCoordinatorTest {
 
@@ -1718,6 +1719,73 @@ void subscribeUnsubscribeInDifferentThreadsShouldNotDeadlock() {
     }
   }
 
+  @Test
+  void consumerShouldBeCreatedProperlyIfManagerClientIsRetried() {
+    scheduledExecutorService = createScheduledExecutorService();
+    when(environment.scheduledExecutorService()).thenReturn(scheduledExecutorService);
+    Duration retryDelay = Duration.ofMillis(100);
+    when(environment.recoveryBackOffDelayPolicy()).thenReturn(BackOffDelayPolicy.fixed(retryDelay));
+    when(consumer.isOpen()).thenReturn(true);
+    when(locator.metadata("stream")).thenReturn(metadata(null, replica()));
+
+    when(clientFactory.client(any()))
+        .thenAnswer(
+            (Answer<Client>)
+                invocationOnMock -> {
+                  // simulates the client is not the good one (e.g. because of load balancer),
+                  // so the connection is closed (hence the call to the shutdown listener)
+                  shutdownListener.handle(
+                      new Client.ShutdownContext(
+                          Client.ShutdownContext.ShutdownReason.CLIENT_CLOSE));
+                  // and a client is returned
+                  return client;
+                });
+
+    when(client.subscribe(
+            subscriptionIdCaptor.capture(),
+            anyString(),
+            any(OffsetSpecification.class),
+            anyInt(),
+            anyMap()))
+        .thenReturn(new Client.Response(Constants.RESPONSE_CODE_OK));
+
+    AtomicInteger messageHandlerCalls = new AtomicInteger();
+    Runnable closingRunnable =
+        coordinator.subscribe(
+            consumer,
+            "stream",
+            OffsetSpecification.first(),
+            null,
+            NO_OP_SUBSCRIPTION_LISTENER,
+            NO_OP_TRACKING_CLOSING_CALLBACK,
+            (offset, message) -> messageHandlerCalls.incrementAndGet(),
+            Collections.emptyMap(),
+            flowStrategy());
+    verify(clientFactory, times(1)).client(any());
+    verify(client, times(1))
+        .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap());
+
+    assertThat(messageHandlerCalls.get()).isEqualTo(0);
+    messageListener.handle(
+        subscriptionIdCaptor.getAllValues().get(0),
+        1,
+        0,
+        0,
+        null,
+        new WrapperMessageBuilder().build());
+    assertThat(messageHandlerCalls.get()).isEqualTo(1);
+
+    when(client.unsubscribe(subscriptionIdCaptor.getValue()))
+        .thenReturn(new Client.Response(Constants.RESPONSE_CODE_OK));
+
+    closingRunnable.run();
+    verify(client, times(1)).unsubscribe(subscriptionIdCaptor.getValue());
+
+    messageListener.handle(
+        subscriptionIdCaptor.getValue(), 0, 0, 0, null, new WrapperMessageBuilder().build());
+    assertThat(messageHandlerCalls.get()).isEqualTo(1);
+  }
+
   Client.Broker leader() {
     return new Client.Broker("leader", -1);
   }

src/test/java/com/rabbitmq/stream/impl/ProducersCoordinatorTest.java

@@ -13,6 +13,7 @@
 // [email protected].
 package com.rabbitmq.stream.impl;
 
+import static com.rabbitmq.stream.impl.TestUtils.CountDownLatchConditions.completed;
 import static com.rabbitmq.stream.impl.TestUtils.answer;
 import static com.rabbitmq.stream.impl.TestUtils.metadata;
 import static org.assertj.core.api.Assertions.assertThat;
@@ -49,6 +50,7 @@
 import org.junit.jupiter.api.Test;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
+import org.mockito.stubbing.Answer;
 
 public class ProducersCoordinatorTest {
 
@@ -602,6 +604,41 @@ class TrackingConsumerInfo {
     assertThat(coordinator.clientCount()).isEqualTo(1);
   }
 
+  @Test
+  void producerShouldBeCreatedProperlyIfManagerClientIsRetried() throws Exception {
+    scheduledExecutorService = createScheduledExecutorService();
+    when(environment.scheduledExecutorService()).thenReturn(scheduledExecutorService);
+    Duration retryDelay = Duration.ofMillis(50);
+    when(environment.recoveryBackOffDelayPolicy()).thenReturn(BackOffDelayPolicy.fixed(retryDelay));
+    when(locator.metadata("stream")).thenReturn(metadata(leader(), replicas()));
+
+    when(clientFactory.client(any()))
+        .thenAnswer(
+            (Answer<Client>)
+                invocationOnMock -> {
+                  shutdownListener.handle(
+                      new Client.ShutdownContext(
+                          Client.ShutdownContext.ShutdownReason.CLIENT_CLOSE));
+
+                  return client;
+                })
+        .thenReturn(client);
+
+    when(producer.isOpen()).thenReturn(true);
+    when(trackingConsumer.isOpen()).thenReturn(true);
+
+    CountDownLatch setClientLatch = new CountDownLatch(1);
+    doAnswer(answer(() -> setClientLatch.countDown())).when(producer).setClient(client);
+
+    coordinator.registerProducer(producer, null, "stream");
+
+    verify(producer, times(1)).setClient(client);
+    assertThat(coordinator.nodesConnected()).isEqualTo(1);
+    assertThat(coordinator.clientCount()).isEqualTo(1);
+
+    assertThat(setClientLatch).is(completed());
+  }
+
   private static ScheduledExecutorService createScheduledExecutorService() {
     return new ScheduledExecutorServiceWrapper(Executors.newSingleThreadScheduledExecutor());
   }