Consider candidates when picking a node

And not only one node.

Author: acogoluegnes

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

@@ -1,4 +1,4 @@
-// Copyright (c) 2020-2023 Broadcom. All Rights Reserved.
+// Copyright (c) 2020-2024 Broadcom. All Rights Reserved.
 // The term "Broadcom" refers to Broadcom Inc. and/or its subsidiaries.
 //
 // This software, the RabbitMQ Stream Java client library, is dual-licensed under the
@@ -70,6 +70,7 @@ class ConsumersCoordinator {
   private final AtomicLong managerIdSequence = new AtomicLong(0);
   private final NavigableSet<ClientSubscriptionsManager> managers = new ConcurrentSkipListSet<>();
   private final AtomicLong trackerIdSequence = new AtomicLong(0);
+  private final Function<List<Broker>, Broker> brokerPicker;
 
   private final List<SubscriptionTracker> trackers = new CopyOnWriteArrayList<>();
   private final ExecutorServiceFactory executorServiceFactory =
@@ -83,16 +84,14 @@ class ConsumersCoordinator {
       int maxConsumersByConnection,
       Function<ClientConnectionType, String> connectionNamingStrategy,
       ClientFactory clientFactory,
-      boolean forceReplica) {
+      boolean forceReplica,
+      Function<List<Broker>, Broker> brokerPicker) {
     this.environment = environment;
     this.clientFactory = clientFactory;
     this.maxConsumersByConnection = maxConsumersByConnection;
     this.connectionNamingStrategy = connectionNamingStrategy;
     this.forceReplica = forceReplica;
-  }
-
-  private static String keyForClientSubscription(Client.Broker broker) {
-    return broker.getHost() + ":" + broker.getPort();
+    this.brokerPicker = brokerPicker;
   }
 
   private BackOffDelayPolicy recoveryBackOffDelayPolicy() {
@@ -138,7 +137,7 @@ Runnable subscribe(
                   flowStrategy);
 
           try {
-            addToManager(newNode, subscriptionTracker, offsetSpecification, true);
+            addToManager(newNode, candidates, subscriptionTracker, offsetSpecification, true);
           } catch (ConnectionStreamException e) {
             // these exceptions are not public
             throw new StreamException(e.getMessage());
@@ -162,6 +161,7 @@ Runnable subscribe(
 
   private void addToManager(
       Broker node,
+      List<Broker> candidates,
       SubscriptionTracker tracker,
       OffsetSpecification offsetSpecification,
       boolean isInitialSubscription) {
@@ -189,9 +189,9 @@ private void addToManager(
         }
       }
       if (pickedManager == null) {
-        String name = keyForClientSubscription(node);
+        String name = keyForNode(node);
         LOGGER.debug("Creating subscription manager on {}", name);
-        pickedManager = new ClientSubscriptionsManager(node, clientParameters);
+        pickedManager = new ClientSubscriptionsManager(node, candidates, clientParameters);
         LOGGER.debug("Created subscription manager on {}, id {}", name, pickedManager.id);
       }
       try {
@@ -571,6 +571,7 @@ private class ClientSubscriptionsManager implements Comparable<ClientSubscriptio
     private final long id;
     private final Broker node;
     private final Client client;
+    // <host>:<port> (actual or advertised)
     private final String name;
     // the 2 data structures track the subscriptions, they must remain consistent
     private final Map<String, Set<SubscriptionTracker>> streamToStreamSubscriptions =
@@ -582,12 +583,12 @@ private class ClientSubscriptionsManager implements Comparable<ClientSubscriptio
     private volatile int trackerCount;
     private final AtomicBoolean closed = new AtomicBoolean(false);
 
-    private ClientSubscriptionsManager(Broker node, Client.ClientParameters clientParameters) {
+    private ClientSubscriptionsManager(
+        Broker targetNode, List<Broker> candidates, Client.ClientParameters clientParameters) {
       this.id = managerIdSequence.getAndIncrement();
-      this.node = node;
-      this.name = keyForClientSubscription(node);
-      LOGGER.debug("creating subscription manager on {}", name);
       this.trackerCount = 0;
+      AtomicReference<String> nameReference = new AtomicReference<>();
+
       AtomicBoolean clientInitializedInManager = new AtomicBoolean(false);
       ChunkListener chunkListener =
           (client, subscriptionId, offset, messageCount, dataSize) -> {
@@ -639,7 +640,7 @@ private ClientSubscriptionsManager(Broker node, Client.ClientParameters clientPa
                   "Could not find stream subscription {} in manager {}, node {} for message listener",
                   subscriptionId,
                   this.id,
-                  this.name);
+                  nameReference.get());
             }
           };
       MessageIgnoredListener messageIgnoredListener =
@@ -663,7 +664,7 @@ private ClientSubscriptionsManager(Broker node, Client.ClientParameters clientPa
                   "Could not find stream subscription {} in manager {}, node {} for message ignored listener",
                   subscriptionId,
                   this.id,
-                  this.name);
+                  nameReference.get());
             }
           };
       ShutdownListener shutdownListener =
@@ -675,7 +676,7 @@ private ClientSubscriptionsManager(Broker node, Client.ClientParameters clientPa
             if (shutdownContext.isShutdownUnexpected()) {
               LOGGER.debug(
                   "Unexpected shutdown notification on subscription connection {}, scheduling consumers re-assignment",
-                  name);
+                  nameReference.get());
               LOGGER.debug(
                   "Subscription connection has {} consumer(s) over {} stream(s) to recover",
                   this.subscriptionTrackers.stream().filter(Objects::nonNull).count(),
@@ -718,7 +719,7 @@ private ClientSubscriptionsManager(Broker node, Client.ClientParameters clientPa
                             }
                           },
                           "Consumers re-assignment after disconnection from %s",
-                          name));
+                          nameReference.get()));
             }
           };
       MetadataListener metadataListener =
@@ -792,18 +793,23 @@ private ClientSubscriptionsManager(Broker node, Client.ClientParameters clientPa
           };
       String connectionName = connectionNamingStrategy.apply(ClientConnectionType.CONSUMER);
       ClientFactoryContext clientFactoryContext =
-          ClientFactoryContext.fromParameters(
-                  clientParameters
-                      .clientProperty("connection_name", connectionName)
-                      .chunkListener(chunkListener)
-                      .creditNotification(creditNotification)
-                      .messageListener(messageListener)
-                      .messageIgnoredListener(messageIgnoredListener)
-                      .shutdownListener(shutdownListener)
-                      .metadataListener(metadataListener)
-                      .consumerUpdateListener(consumerUpdateListener))
-              .key(name);
+          new ClientFactoryContext(
+              clientParameters
+                  .clientProperty("connection_name", connectionName)
+                  .chunkListener(chunkListener)
+                  .creditNotification(creditNotification)
+                  .messageListener(messageListener)
+                  .messageIgnoredListener(messageIgnoredListener)
+                  .shutdownListener(shutdownListener)
+                  .metadataListener(metadataListener)
+                  .consumerUpdateListener(consumerUpdateListener),
+              keyForNode(targetNode),
+              candidates);
       this.client = clientFactory.client(clientFactoryContext);
+      this.node = brokerFromClient(this.client);
+      this.name = keyForNode(this.node);
+      nameReference.set(this.name);
+      LOGGER.debug("creating subscription manager on {}", name);
       LOGGER.debug("Created consumer connection '{}'", connectionName);
       clientInitializedInManager.set(true);
     }
@@ -906,7 +912,7 @@ private void recoverSubscription(List<Broker> candidates, SubscriptionTracker tr
                 } else {
                   offsetSpecification = tracker.initialOffsetSpecification;
                 }
-                addToManager(broker, tracker, offsetSpecification, false);
+                addToManager(broker, candidates, tracker, offsetSpecification, false);
               }
             }
           } else {

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

@@ -1,4 +1,4 @@
-// Copyright (c) 2020-2023 Broadcom. All Rights Reserved.
+// Copyright (c) 2020-2024 Broadcom. All Rights Reserved.
 // The term "Broadcom" refers to Broadcom Inc. and/or its subsidiaries.
 //
 // This software, the RabbitMQ Stream Java client library, is dual-licensed under the
@@ -32,13 +32,7 @@
 import com.rabbitmq.stream.impl.Utils.ClientConnectionType;
 import com.rabbitmq.stream.impl.Utils.ClientFactory;
 import com.rabbitmq.stream.impl.Utils.ClientFactoryContext;
-import java.util.Collection;
-import java.util.Iterator;
-import java.util.List;
-import java.util.Map;
-import java.util.NavigableSet;
-import java.util.Objects;
-import java.util.Set;
+import java.util.*;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.ConcurrentSkipListSet;
@@ -87,10 +81,6 @@ class ProducersCoordinator {
     this.connectionNamingStrategy = connectionNamingStrategy;
   }
 
-  private static String keyForNode(Client.Broker broker) {
-    return broker.getHost() + ":" + broker.getPort();
-  }
-
   Runnable registerProducer(StreamProducer producer, String reference, String stream) {
     return lock(
         this.coordinatorLock,
@@ -162,7 +152,7 @@ private void addToManager(Broker node, AgentTracker tracker) {
       }
       if (pickedManager == null) {
         String name = keyForNode(node);
-        LOGGER.debug("Creating producer manager on {}", name);
+        LOGGER.debug("Trying to create producer manager on {}", name);
         pickedManager = new ClientProducersManager(node, this.clientFactory, clientParameters);
         LOGGER.debug("Created producer manager on {}, id {}", name, pickedManager.id);
       }
@@ -578,10 +568,9 @@ private class ClientProducersManager implements Comparable<ClientProducersManage
     private final AtomicBoolean closed = new AtomicBoolean(false);
 
     private ClientProducersManager(
-        Broker node, ClientFactory cf, Client.ClientParameters clientParameters) {
+        Broker targetNode, ClientFactory cf, Client.ClientParameters clientParameters) {
       this.id = managerIdSequence.getAndIncrement();
-      this.name = keyForNode(node);
-      this.node = node;
+      AtomicReference<String> nameReference = new AtomicReference<>();
       AtomicReference<Client> ref = new AtomicReference<>();
       AtomicBoolean clientInitializedInManager = new AtomicBoolean(false);
       PublishConfirmListener publishConfirmListener =
@@ -636,7 +625,7 @@ private ClientProducersManager(
                                 });
                           },
                           "Producer recovery after disconnection from %s",
-                          name));
+                          nameReference.get()));
             }
           };
       MetadataListener metadataListener =
@@ -685,15 +674,19 @@ private ClientProducersManager(
           };
       String connectionName = connectionNamingStrategy.apply(ClientConnectionType.PRODUCER);
       ClientFactoryContext connectionFactoryContext =
-          ClientFactoryContext.fromParameters(
-                  clientParameters
-                      .publishConfirmListener(publishConfirmListener)
-                      .publishErrorListener(publishErrorListener)
-                      .shutdownListener(shutdownListener)
-                      .metadataListener(metadataListener)
-                      .clientProperty("connection_name", connectionName))
-              .key(name);
+          new ClientFactoryContext(
+              clientParameters
+                  .publishConfirmListener(publishConfirmListener)
+                  .publishErrorListener(publishErrorListener)
+                  .shutdownListener(shutdownListener)
+                  .metadataListener(metadataListener)
+                  .clientProperty("connection_name", connectionName),
+              keyForNode(targetNode),
+              Collections.emptyList());
       this.client = cf.client(connectionFactoryContext);
+      this.node = Utils.brokerFromClient(this.client);
+      this.name = keyForNode(this.node);
+      nameReference.set(this.name);
       LOGGER.debug("Created producer connection '{}'", connectionName);
       clientInitializedInManager.set(true);
       ref.set(this.client);

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

@@ -1,4 +1,4 @@
-// Copyright (c) 2020-2023 Broadcom. All Rights Reserved.
+// Copyright (c) 2020-2024 Broadcom. All Rights Reserved.
 // The term "Broadcom" refers to Broadcom Inc. and/or its subsidiaries.
 //
 // This software, the RabbitMQ Stream Java client library, is dual-licensed under the
@@ -219,7 +219,8 @@ class StreamEnvironment implements Environment {
             maxConsumersByConnection,
             connectionNamingStrategy,
             Utils.coordinatorClientFactory(this),
-            forceReplicaForConsumers);
+            forceReplicaForConsumers,
+            Utils.brokerPicker());
     this.offsetTrackingCoordinator = new OffsetTrackingCoordinator(this);
     ClientParameters clientParametersForInit = locatorParametersCopy();
     Runnable locatorInitSequence =