Add locator count setting

Author: acogoluegnes

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

@@ -18,6 +18,7 @@
 import static com.rabbitmq.stream.impl.Utils.*;
 import static java.lang.String.format;
 import static java.util.concurrent.TimeUnit.SECONDS;
+import static java.util.stream.Collectors.toList;
 
 import com.rabbitmq.stream.*;
 import com.rabbitmq.stream.MessageHandler.Context;
@@ -53,6 +54,7 @@
 import java.util.function.LongSupplier;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
+import java.util.stream.IntStream;
 import javax.net.ssl.SSLException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -81,7 +83,7 @@ class StreamEnvironment implements Environment {
   private final ByteBufAllocator byteBufAllocator;
   private final AtomicBoolean locatorsInitialized = new AtomicBoolean(false);
   private final Runnable locatorInitializationSequence;
-  private final List<Locator> locators = new CopyOnWriteArrayList<>();
+  private final List<Locator> locators;
   private final ExecutorServiceFactory executorServiceFactory;
   private final ObservationCollector<?> observationCollector;
 
@@ -105,7 +107,8 @@ class StreamEnvironment implements Environment {
       boolean forceReplicaForConsumers,
       boolean forceLeaderForProducers,
       Duration producerNodeRetryDelay,
-      Duration consumerNodeRetryDelay) {
+      Duration consumerNodeRetryDelay,
+      int expectedLocatorCount) {
     this.recoveryBackOffDelayPolicy = recoveryBackOffDelayPolicy;
     this.topologyUpdateBackOffDelayPolicy = topologyBackOffDelayPolicy;
     this.byteBufAllocator = byteBufAllocator;
@@ -147,7 +150,7 @@ class StreamEnvironment implements Environment {
                       new Address(
                           uriItem.getHost() == null ? "localhost" : uriItem.getHost(),
                           uriItem.getPort() == -1 ? defaultPort : uriItem.getPort()))
-              .collect(Collectors.toList());
+              .collect(toList());
     }
 
     AddressResolver addressResolverToUse = addressResolver;
@@ -179,7 +182,19 @@ class StreamEnvironment implements Environment {
 
     this.addressResolver = addressResolverToUse;
 
-    this.addresses.forEach(address -> this.locators.add(new Locator(address)));
+    int locatorCount = Math.max(this.addresses.size(), expectedLocatorCount);
+    LOGGER.debug("Using {} locator connection(s)", locatorCount);
+
+    List<Locator> lctrs =
+        IntStream.range(0, locatorCount)
+            .mapToObj(
+                i -> {
+                  Address addr = this.addresses.get(i % this.addresses.size());
+                  return new Locator(addr);
+                })
+            .collect(toList());
+    this.locators = List.copyOf(lctrs);
+
     this.executorServiceFactory =
         new DefaultExecutorServiceFactory(
             this.addresses.size(), 1, "rabbitmq-stream-locator-connection-");

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

@@ -70,6 +70,7 @@ public class StreamEnvironmentBuilder implements EnvironmentBuilder {
   private ObservationCollector<?> observationCollector = ObservationCollector.NO_OP;
   private Duration producerNodeRetryDelay = Duration.ofMillis(500);
   private Duration consumerNodeRetryDelay = Duration.ofMillis(1000);
+  private int locatorCount = 1;
 
   public StreamEnvironmentBuilder() {}
 
@@ -315,6 +316,11 @@ StreamEnvironmentBuilder consumerNodeRetryDelay(Duration consumerNodeRetryDelay)
     return this;
   }
 
+  StreamEnvironmentBuilder locatorCount(int locatorCount) {
+    this.locatorCount = locatorCount;
+    return this;
+  }
+
   @Override
   public Environment build() {
     if (this.compressionCodecFactory == null) {
@@ -349,7 +355,8 @@ public Environment build() {
         this.forceReplicaForConsumers,
         this.forceLeaderForProducers,
         this.producerNodeRetryDelay,
-        this.consumerNodeRetryDelay);
+        this.consumerNodeRetryDelay,
+        this.locatorCount);
   }
 
   static final class DefaultTlsConfiguration implements TlsConfiguration {

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

@@ -115,6 +115,7 @@ void clusterRestart(boolean useLoadBalancer, boolean forceLeader) throws Interru
           .addressResolver(addr -> LOAD_BALANCER_ADDRESS);
       Duration nodeRetryDelay = Duration.ofMillis(100);
       environmentBuilder.forceLeaderForProducers(forceLeader);
+      ((StreamEnvironmentBuilder) environmentBuilder).locatorCount(URIS.size());
       // to make the test faster
       ((StreamEnvironmentBuilder) environmentBuilder).producerNodeRetryDelay(nodeRetryDelay);
       ((StreamEnvironmentBuilder) environmentBuilder).consumerNodeRetryDelay(nodeRetryDelay);

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

@@ -475,8 +475,7 @@ void manualTrackingConsumerShouldRestartWhereItLeftOff() throws Exception {
 
   @Test
   @DisabledIfRabbitMqCtlNotSet
-  void consumerShouldReUseInitialOffsetSpecificationAfterDisruptionIfNoMessagesReceived()
-      throws Exception {
+  void consumerShouldReUseInitialOffsetSpecificationAfterDisruptionIfNoMessagesReceived() {
     int messageCountFirstWave = 10_000;
     Producer producer = environment.producerBuilder().stream(stream).build();
 

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

@@ -101,7 +101,8 @@ Client.ClientParameters duplicate() {
             false,
             true,
             Duration.ofMillis(100),
-            Duration.ofMillis(100));
+            Duration.ofMillis(100),
+            1);
   }
 
   @AfterEach
@@ -169,7 +170,8 @@ void shouldTryUrisOnInitializationFailure() throws Exception {
             false,
             true,
             Duration.ofMillis(100),
-            Duration.ofMillis(100));
+            Duration.ofMillis(100),
+            1);
     verify(cf, times(3)).apply(any(Client.ClientParameters.class));
   }
 
@@ -200,7 +202,8 @@ void shouldNotOpenConnectionWhenLazyInitIsEnabled(
             false,
             true,
             Duration.ofMillis(100),
-            Duration.ofMillis(100));
+            Duration.ofMillis(100),
+            1);
     verify(cf, times(expectedConnectionCreation)).apply(any(Client.ClientParameters.class));
   }