Harden candidate lookup for producers

Author: acogoluegnes

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

@@ -14,6 +14,7 @@
 // [email protected].
 package com.rabbitmq.stream.impl;
 
+import static com.rabbitmq.stream.impl.Tuples.pair;
 import static com.rabbitmq.stream.impl.Utils.*;
 import static java.util.stream.Collectors.toList;
 import static java.util.stream.Collectors.toSet;
@@ -30,6 +31,7 @@
 import com.rabbitmq.stream.impl.Client.PublishErrorListener;
 import com.rabbitmq.stream.impl.Client.Response;
 import com.rabbitmq.stream.impl.Client.ShutdownListener;
+import com.rabbitmq.stream.impl.Tuples.Pair;
 import com.rabbitmq.stream.impl.Utils.ClientConnectionType;
 import com.rabbitmq.stream.impl.Utils.ClientFactory;
 import com.rabbitmq.stream.impl.Utils.ClientFactoryContext;
@@ -219,10 +221,13 @@ List<BrokerWrapper> findCandidateNodes(String stream, boolean forceLeader) {
 
     List<BrokerWrapper> candidates = new ArrayList<>();
     Client.Broker leader = streamMetadata.getLeader();
-    if (leader == null && forceLeader) {
-      throw new IllegalStateException("Not leader available for stream " + stream);
+    if (leader == null) {
+      if (forceLeader) {
+        throw new IllegalStateException("Not leader available for stream " + stream);
+      }
+    } else {
+      candidates.add(new BrokerWrapper(leader, true));
     }
-    candidates.add(new BrokerWrapper(leader, true));
 
     if (!forceLeader && streamMetadata.hasReplicas()) {
       candidates.addAll(
@@ -231,7 +236,11 @@ List<BrokerWrapper> findCandidateNodes(String stream, boolean forceLeader) {
               .collect(toList()));
     }
 
-    LOGGER.debug("Candidates to publish to {}: {}", stream, candidates);
+    if (candidates.isEmpty()) {
+      throw new IllegalStateException("No stream member available to publish for stream " + stream);
+    } else {
+      LOGGER.debug("Candidates to publish to {}: {}", stream, candidates);
+    }
 
     return List.copyOf(candidates);
   }
@@ -721,15 +730,20 @@ private ClientProducersManager(
 
     private void assignProducersToNewManagers(
         Collection<AgentTracker> trackers, String stream, BackOffDelayPolicy delayPolicy) {
-      AsyncRetry.asyncRetry(() -> findCandidateNodes(stream, forceLeader))
+      AsyncRetry.asyncRetry(
+              () -> {
+                List<BrokerWrapper> candidates = findCandidateNodes(stream, forceLeader);
+                return pair(pickBroker(candidates), candidates);
+              })
           .description("Candidate lookup to publish to " + stream)
           .scheduler(environment.scheduledExecutorService())
           .retry(ex -> !(ex instanceof StreamDoesNotExistException))
           .delayPolicy(delayPolicy)
           .build()
           .thenAccept(
-              candidates -> {
-                Broker broker = pickBroker(candidates);
+              brokerAndCandidates -> {
+                Broker broker = brokerAndCandidates.v1();
+                List<BrokerWrapper> candidates = brokerAndCandidates.v2();
                 String key = keyForNode(broker);
                 LOGGER.debug(
                     "Assigning {} producer(s) and consumer tracker(s) to {}", trackers.size(), key);
@@ -805,15 +819,19 @@ private void recoverAgent(Broker node, List<BrokerWrapper> candidates, AgentTrac
               tracker.identifiable() ? tracker.id() : "N/A",
               tracker.stream());
           // maybe not a good candidate, let's refresh and retry for this one
-          candidates =
-              Utils.callAndMaybeRetry(
-                  () -> findCandidateNodes(tracker.stream(), forceLeader),
+          Pair<Broker, List<BrokerWrapper>> brokerAndCandidates =
+              callAndMaybeRetry(
+                  () -> {
+                    List<BrokerWrapper> cs = findCandidateNodes(tracker.stream(), forceLeader);
+                    return pair(pickBroker(cs), cs);
+                  },
                   ex -> !(ex instanceof StreamDoesNotExistException),
                   environment.recoveryBackOffDelayPolicy(),
                   "Candidate lookup for %s on stream '%s'",
                   tracker.type(),
                   tracker.stream());
-          node = pickBroker(candidates);
+          node = brokerAndCandidates.v1();
+          candidates = brokerAndCandidates.v2();
         } catch (Exception e) {
           LOGGER.warn(
               "Error while re-assigning {} (stream '{}')", tracker.type(), tracker.stream(), e);

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

@@ -0,0 +1,43 @@
+// Copyright (c) 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
+// Mozilla Public License 2.0 ("MPL"), and the Apache License version 2 ("ASL").
+// For the MPL, please see LICENSE-MPL-RabbitMQ. For the ASL,
+// please see LICENSE-APACHE2.
+//
+// This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY KIND,
+// either express or implied. See the LICENSE file for specific language governing
+// rights and limitations of this software.
+//
+// If you have any questions regarding licensing, please contact us at
+// [email protected].
+package com.rabbitmq.stream.impl;
+
+final class Tuples {
+
+  private Tuples() {}
+
+  static <A, B> Pair<A, B> pair(A v1, B v2) {
+    return new Pair<>(v1, v2);
+  }
+
+  static class Pair<A, B> {
+
+    private final A v1;
+    private final B v2;
+
+    private Pair(A v1, B v2) {
+      this.v1 = v1;
+      this.v2 = v2;
+    }
+
+    A v1() {
+      return this.v1;
+    }
+
+    B v2() {
+      return this.v2;
+    }
+  }
+}

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

@@ -57,7 +57,7 @@ public class LoadBalancerClusterTest {
   String stream;
   EventLoopGroup eventLoopGroup;
   Client locator;
-  Address loadBalancerAddress = new Address("localhost", LB_PORT);
+  static final Address LOAD_BALANCER_ADDRESS = new Address("localhost", LB_PORT);
 
   @BeforeEach
   void init() {
@@ -66,7 +66,7 @@ void init() {
     when(environment.locator()).thenReturn(locator);
     when(environment.clientParametersCopy())
         .thenReturn(new Client.ClientParameters().eventLoopGroup(eventLoopGroup).port(LB_PORT));
-    when(environment.addressResolver()).thenReturn(address -> loadBalancerAddress);
+    when(environment.addressResolver()).thenReturn(address -> LOAD_BALANCER_ADDRESS);
     when(environment.locatorOperation(any())).thenCallRealMethod();
   }
 
@@ -175,7 +175,7 @@ void producersConsumersShouldSpreadAccordingToDataLocalitySettings(boolean force
             .port(LB_PORT)
             .forceReplicaForConsumers(forceLocality)
             .forceReplicaForConsumers(forceLocality)
-            .addressResolver(addr -> loadBalancerAddress)
+            .addressResolver(addr -> LOAD_BALANCER_ADDRESS)
             .maxProducersByConnection(maxPerConnection)
             .maxConsumersByConnection(maxPerConnection)
             .forceLeaderForProducers(forceLocality)

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

@@ -750,12 +750,27 @@ void findCandidateNodesShouldReturnLeaderAndReplicasWhenForceLeaderIsFalse() {
   }
 
   @Test
-  void findCandidateNodesShouldThrowIfThereIsNoLeader() {
+  void findCandidateNodesShouldThrowIfThereIsNoLeaderAndForceLeaderIsTrue() {
     when(locator.metadata("stream")).thenReturn(metadata(null, replicas()));
     assertThatThrownBy(() -> coordinator.findCandidateNodes("stream", true))
         .isInstanceOf(IllegalStateException.class);
   }
 
+  @Test
+  void findCandidateNodesShouldThrowIfNoMembersAndForceLeaderIsFalse() {
+    when(locator.metadata("stream")).thenReturn(metadata(null, List.of()));
+    assertThatThrownBy(() -> coordinator.findCandidateNodes("stream", false))
+        .isInstanceOf(IllegalStateException.class);
+  }
+
+  @Test
+  void findCandidateNodesShouldReturnOnlyReplicasIfNoLeaderAndForceLeaderIsFalse() {
+    when(locator.metadata("stream")).thenReturn(metadata(null, replicas()));
+    assertThat(coordinator.findCandidateNodes("stream", false))
+        .hasSize(2)
+        .containsAll(replicaWrappers());
+  }
+
   private static ScheduledExecutorService createScheduledExecutorService() {
     return new ScheduledExecutorServiceWrapper(Executors.newSingleThreadScheduledExecutor());
   }