Repository: skiptests/astraea
Pull request #453 (2022-07-03 17:23:51)

Integrate StrictCostDispatcher with SMRR
這個PR包含兩個工作：
1. 建立RR介面 and SMRR實作
2. 將SMRR整合到`StrictCostDispatcher`


Patch:
diff --git a/app/src/main/java/org/astraea/app/partitioner/RoundRobin.java b/app/src/main/java/org/astraea/app/partitioner/RoundRobin.java
new file mode 100644
index 00000000..d920587c
--- /dev/null
+++ b/app/src/main/java/org/astraea/app/partitioner/RoundRobin.java
@@ -0,0 +1,87 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.astraea.app.partitioner;
+
+import java.util.Collections;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+
+/**
+ * a way to pick best target.
+ *
+ * @param <E> target type
+ */
+public interface RoundRobin<E> {
+
+  static <E> RoundRobin<E> smoothWeighted(Map<E, Double> weights) {
+    return new SmoothWeightedRoundRobin<>(weights);
+  }
+
+  /**
+   * calculate the next object to use
+   *
+   * @param availableTargets available targets
+   * @return obj to send, or null if this algorithm fails to get answer for current condition.
+   */
+  Optional<E> next(Set<E> availableTargets);
+
+  class SmoothWeightedRoundRobin<E> implements RoundRobin<E> {
+
+    private final Map<E, Double> effectiveWeights;
+    private volatile Map<E, Double> currentWeights;
+
+    private SmoothWeightedRoundRobin(Map<E, Double> weights) {
+      this.effectiveWeights = Collections.unmodifiableMap(weights);
+      this.currentWeights =
+          weights.keySet().stream()
+              .collect(Collectors.toUnmodifiableMap(Function.identity(), ignored -> 0D));
+    }
+
+    @Override
+    public Optional<E> next(Set<E> availableTargets) {
+      // no data no answer
+      if (effectiveWeights.isEmpty() || availableTargets.isEmpty()) return Optional.empty();
+
+      // 1) calculate the sum of all effective weights
+      var sum = effectiveWeights.values().stream().mapToDouble(d -> d).sum();
+      // 2) add effective weight to each current weight
+      var nextWeights =
+          currentWeights.entrySet().stream()
+              .collect(
+                  Collectors.toMap(
+                      Map.Entry::getKey,
+                      e -> effectiveWeights.getOrDefault(e.getKey(), 0D) + e.getValue()));
+      // 3) get the E which has max value
+      var maxObj =
+          nextWeights.entrySet().stream()
+              .filter(e -> availableTargets.contains(e.getKey()))
+              .max(Map.Entry.comparingByValue())
+              .map(Map.Entry::getKey);
+      // 4) subtract weight from max (effective weight)
+      maxObj.ifPresent(
+          o -> {
+            nextWeights.put(o, nextWeights.get(o) - sum);
+            // 5) update current weights by next weights
+            currentWeights = Collections.unmodifiableMap(nextWeights);
+          });
+      return maxObj;
+    }
+  }
+}
diff --git a/app/src/main/java/org/astraea/app/partitioner/StrictCostDispatcher.java b/app/src/main/java/org/astraea/app/partitioner/StrictCostDispatcher.java
index 50f1c009..bdb4b6f5 100644
--- a/app/src/main/java/org/astraea/app/partitioner/StrictCostDispatcher.java
+++ b/app/src/main/java/org/astraea/app/partitioner/StrictCostDispatcher.java
@@ -18,17 +18,20 @@ package org.astraea.app.partitioner;
 
 import java.time.Duration;
 import java.util.Collection;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.NoSuchElementException;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.TreeMap;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 import org.astraea.app.common.Utils;
 import org.astraea.app.cost.ClusterInfo;
 import org.astraea.app.cost.CostFunction;
 import org.astraea.app.cost.HasBrokerCost;
+import org.astraea.app.cost.NodeInfo;
 import org.astraea.app.cost.ReplicaInfo;
 import org.astraea.app.metrics.collector.BeanCollector;
 import org.astraea.app.metrics.collector.Fetcher;
@@ -50,16 +53,24 @@ import org.astraea.app.metrics.collector.Receiver;
  */
 public class StrictCostDispatcher implements Dispatcher {
   public static final String JMX_PORT = "jmx.port";
+  private static final Duration ROUND_ROBIN_LEASE = Duration.ofSeconds(30);
 
   private final BeanCollector beanCollector =
       BeanCollector.builder().interval(Duration.ofSeconds(4)).build();
 
   /* The cost-functions we consider and the weight of them. It is visible for test.*/
   Map<CostFunction, Double> functions;
-  private Optional<Integer> jmxPortDefault = Optional.empty();
-  private final Map<Integer, Integer> jmxPorts = new TreeMap<>();
+
+  Function<Integer, Integer> jmxPortGetter =
+      (id) -> {
+        throw new NoSuchElementException("broker: " + id + " does not have jmx port");
+      };
+
   final Map<Integer, Receiver> receivers = new TreeMap<>();
 
+  volatile RoundRobin<Integer> roundRobin;
+  volatile long timeToUpdateRoundRobin = -1;
+
   public StrictCostDispatcher() {
     this(List.of(CostFunction.throughput()));
   }
@@ -72,6 +83,11 @@ public class StrictCostDispatcher implements Dispatcher {
             .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
   }
 
+  // visible for testing
+  Receiver receiver(String host, int port, Fetcher fetcher) {
+    return beanCollector.register().host(host).port(port).fetcher(fetcher).build();
+  }
+
   @Override
   public int partition(String topic, byte[] key, byte[] value, ClusterInfo clusterInfo) {
     var partitionLeaders = clusterInfo.availableReplicaLeaders(topic);
@@ -82,84 +98,81 @@ public class StrictCostDispatcher implements Dispatcher {
     if (partitionLeaders.size() == 1) return partitionLeaders.iterator().next().partition();
 
     // add new receivers for new brokers
-    partitionLeaders.stream()
-        .filter(p -> !receivers.containsKey(p.nodeInfo().id()))
-        .forEach(
-            p ->
-                receivers.put(
-                    p.nodeInfo().id(), receiver(p.nodeInfo().host(), jmxPort(p.nodeInfo().id()))));
+    receivers.putAll(
+        Fetcher.of(functions.keySet())
+            .map(
+                fetcher ->
+                    partitionLeaders.stream()
+                        .map(ReplicaInfo::nodeInfo)
+                        .filter(nodeInfo -> !receivers.containsKey(nodeInfo.id()))
+                        .distinct()
+                        .collect(
+                            Collectors.toMap(
+                                NodeInfo::id,
+                                nodeInfo ->
+                                    receiver(
+                                        nodeInfo.host(),
+                                        jmxPortGetter.apply(nodeInfo.id()),
+                                        fetcher))))
+            .orElse(Map.of()));
 
     // get latest beans for each node
     var beans =
         receivers.entrySet().stream()
             .collect(Collectors.toMap(Map.Entry::getKey, e -> e.getValue().current()));
 
-    // get scores from all cost functions
-    var scores = computeScore(functions, ClusterInfo.of(clusterInfo, beans));
-
-    return bestPartition(partitionLeaders, scores).map(e -> e.getKey().partition()).orElse(0);
+    if (roundRobin == null || System.currentTimeMillis() >= timeToUpdateRoundRobin) {
+      roundRobin = newRoundRobin(functions, ClusterInfo.of(clusterInfo, beans));
+      timeToUpdateRoundRobin = System.currentTimeMillis() + ROUND_ROBIN_LEASE.toMillis();
+    }
+    return roundRobin
+        .next(partitionLeaders.stream().map(r -> r.nodeInfo().id()).collect(Collectors.toSet()))
+        .orElse(0);
   }
 
   /**
-   * Pass clusterInfo into all cost-functions. The result of each cost-function will multiply on
-   * their corresponding weight.
+   * create new Round-Robin based.
    *
-   * @param functions the cost-function objects and their corresponding weight
-   * @return cost-function result multiplied on their corresponding weight
+   * @param costFunctions used to calculate weights
+   * @param clusterInfo used to calculate costs
+   * @return SmoothWeightedRoundRobin
    */
-  static List<Map<Integer, Double>> computeScore(
-      Map<CostFunction, Double> functions, ClusterInfo clusterInfo) {
-    return functions.entrySet().stream()
-        .filter(e -> e.getKey() instanceof HasBrokerCost)
-        .map(e -> Map.entry((HasBrokerCost) e.getKey(), e.getValue()))
-        .map(
-            functionWeight ->
-                functionWeight
-                    .getKey()
-                    // Execute all cost functions
-                    .brokerCost(clusterInfo)
-                    .value()
-                    .entrySet()
-                    .stream()
-                    // Weight on cost functions result
-                    .map(
-                        IdScore ->
-                            Map.entry(
-                                IdScore.getKey(), IdScore.getValue() * functionWeight.getValue()))
-                    .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)))
-        .collect(Collectors.toUnmodifiableList());
-  }
-
-  // visible for testing
-  static Optional<Map.Entry<ReplicaInfo, Double>> bestPartition(
-      List<ReplicaInfo> partitions, List<Map<Integer, Double>> scores) {
-    return partitions.stream()
-        .map(
-            p ->
-                Map.entry(
-                    p,
-                    scores.stream()
-                        .mapToDouble(s -> s.getOrDefault(p.nodeInfo().id(), 0.0D))
-                        .sum()))
-        .min(Map.Entry.comparingByValue());
-  }
-
-  Receiver receiver(String host, int port) {
-    return beanCollector
-        .register()
-        .host(host)
-        .port(port)
-        // TODO: return optional receiver
-        .fetcher(Fetcher.of(functions.keySet()).get())
-        .build();
+  static RoundRobin<Integer> newRoundRobin(
+      Map<CostFunction, Double> costFunctions, ClusterInfo clusterInfo) {
+    return RoundRobin.smoothWeighted(
+        costFunctions.entrySet().stream()
+            .filter(e -> e.getKey() instanceof HasBrokerCost)
+            .map(e -> Map.entry((HasBrokerCost) e.getKey(), e.getValue()))
+            .flatMap(
+                functionWeight ->
+                    functionWeight
+                        .getKey()
+                        // Execute all cost functions
+                        .brokerCost(clusterInfo)
+                        .value()
+                        .entrySet()
+                        .stream()
+                        // Weight on cost functions result
+                        .map(
+                            IdScore ->
+                                Map.entry(
+                                    IdScore.getKey(),
+                                    IdScore.getValue() * functionWeight.getValue())))
+            .collect(
+                Collectors.toMap(
+                    Map.Entry::getKey, Map.Entry::getValue, Double::sum, HashMap::new)));
   }
 
   @Override
   public void configure(Configuration config) {
-    jmxPortDefault = config.integer(JMX_PORT);
-
-    // seeks for custom jmx ports.
-    jmxPorts.putAll(PartitionerUtils.parseIdJMXPort(config));
+    var jmxPortDefault = config.integer(JMX_PORT);
+    var customJmxPort = PartitionerUtils.parseIdJMXPort(config);
+    jmxPortGetter =
+        id ->
+            Optional.ofNullable(customJmxPort.get(id))
+                .or(() -> jmxPortDefault)
+                .orElseThrow(
+                    () -> new NoSuchElementException("broker: " + id + " does not have jmx port"));
 
     functions = parseCostFunctionWeight(config);
   }
@@ -204,13 +217,6 @@ public class StrictCostDispatcher implements Dispatcher {
         .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
   }
 
-  // visible for testing
-  int jmxPort(int id) {
-    if (jmxPorts.containsKey(id)) return jmxPorts.get(id);
-    return jmxPortDefault.orElseThrow(
-        () -> new NoSuchElementException("broker: " + id + " does not have jmx port"));
-  }
-
   @Override
   public void close() {
     receivers.values().forEach(r -> Utils.swallowException(r::close));


Test patch:
diff --git a/app/src/test/java/org/astraea/app/partitioner/RoundRobinTest.java b/app/src/test/java/org/astraea/app/partitioner/RoundRobinTest.java
new file mode 100644
index 00000000..b0eddde5
--- /dev/null
+++ b/app/src/test/java/org/astraea/app/partitioner/RoundRobinTest.java
@@ -0,0 +1,50 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.astraea.app.partitioner;
+
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+public class RoundRobinTest {
+
+  @Test
+  void testSmoothWeighted() {
+    var weight = Map.of(0, 10D, 1, 5D);
+
+    var rr = RoundRobin.smoothWeighted(weight);
+
+    // nothing to pick
+    Assertions.assertEquals(Optional.empty(), rr.next(Set.of()));
+    Assertions.assertEquals(Optional.empty(), rr.next(Set.of(3)));
+
+    // chosen one
+    Assertions.assertEquals(Optional.of(0), rr.next(Set.of(0)));
+    Assertions.assertEquals(Optional.of(1), rr.next(Set.of(1)));
+
+    // reset
+    rr = RoundRobin.smoothWeighted(weight);
+    Assertions.assertEquals(Optional.of(0), rr.next(Set.of(0, 1)));
+    Assertions.assertEquals(Optional.of(1), rr.next(Set.of(0, 1)));
+    Assertions.assertEquals(Optional.of(0), rr.next(Set.of(0, 1)));
+    Assertions.assertEquals(Optional.of(0), rr.next(Set.of(0, 1)));
+    Assertions.assertEquals(Optional.of(1), rr.next(Set.of(0, 1)));
+    Assertions.assertEquals(Optional.of(0), rr.next(Set.of(0, 1)));
+  }
+}
diff --git a/app/src/test/java/org/astraea/app/partitioner/StrictCostDispatcherTest.java b/app/src/test/java/org/astraea/app/partitioner/StrictCostDispatcherTest.java
index 2ef04315..4d8d5dfc 100644
--- a/app/src/test/java/org/astraea/app/partitioner/StrictCostDispatcherTest.java
+++ b/app/src/test/java/org/astraea/app/partitioner/StrictCostDispatcherTest.java
@@ -19,17 +19,16 @@ package org.astraea.app.partitioner;
 import java.util.List;
 import java.util.Map;
 import java.util.NoSuchElementException;
-import java.util.Set;
-import java.util.function.Function;
-import java.util.stream.Collectors;
+import java.util.Optional;
+import java.util.concurrent.atomic.AtomicInteger;
 import org.astraea.app.admin.ClusterBean;
-import org.astraea.app.cost.BrokerCost;
 import org.astraea.app.cost.BrokerInputCost;
 import org.astraea.app.cost.ClusterInfo;
-import org.astraea.app.cost.HasBrokerCost;
+import org.astraea.app.cost.CostFunction;
 import org.astraea.app.cost.NodeInfo;
 import org.astraea.app.cost.ReplicaInfo;
 import org.astraea.app.cost.ThroughputCost;
+import org.astraea.app.metrics.collector.Fetcher;
 import org.astraea.app.metrics.collector.Receiver;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
@@ -37,11 +36,6 @@ import org.mockito.Mockito;
 
 public class StrictCostDispatcherTest {
 
-  private static ReplicaInfo createFakeReplicaLeaderInfo(
-      String topic, int partition, NodeInfo replicaNode) {
-    return ReplicaInfo.of(topic, partition, replicaNode, true, true, false);
-  }
-
   @Test
   void testClose() {
     var dispatcher = new StrictCostDispatcher();
@@ -53,16 +47,19 @@ public class StrictCostDispatcherTest {
   }
 
   @Test
-  void testConfigure() {
+  void testJmxPort() {
     try (var dispatcher = new StrictCostDispatcher()) {
       dispatcher.configure(Configuration.of(Map.of()));
-      Assertions.assertThrows(NoSuchElementException.class, () -> dispatcher.jmxPort(0));
+      Assertions.assertThrows(
+          NoSuchElementException.class, () -> dispatcher.jmxPortGetter.apply(0));
       dispatcher.configure(Configuration.of(Map.of(StrictCostDispatcher.JMX_PORT, "12345")));
-      Assertions.assertEquals(12345, dispatcher.jmxPort(0));
+      Assertions.assertEquals(12345, dispatcher.jmxPortGetter.apply(0));
     }
+  }
 
+  @Test
+  void testNegativeWeight() {
     try (var dispatcher = new StrictCostDispatcher()) {
-      // Test for negative weight.
       Assertions.assertThrows(
           IllegalArgumentException.class,
           () ->
@@ -77,88 +74,35 @@ public class StrictCostDispatcherTest {
   }
 
   @Test
-  void testBestPartition() {
-    var partitions =
-        List.of(
-            createFakeReplicaLeaderInfo("t", 0, NodeInfo.of(10, "h0", 1000)),
-            createFakeReplicaLeaderInfo("t", 1, NodeInfo.of(11, "h1", 1000)));
-
-    var score = List.of(Map.of(10, 0.8D), Map.of(11, 0.7D));
-
-    var result = StrictCostDispatcher.bestPartition(partitions, score).get();
-
-    Assertions.assertEquals(0.7D, result.getValue());
-    Assertions.assertEquals(1, result.getKey().partition());
+  void testConfigureCostFunctions() {
+    try (var dispatcher = new StrictCostDispatcher()) {
+      dispatcher.configure(
+          Configuration.of(
+              Map.of(ThroughputCost.class.getName(), "0.1", BrokerInputCost.class.getName(), "2")));
+      Assertions.assertEquals(2, dispatcher.functions.size());
+    }
   }
 
   @Test
-  void testPartition() {
-    // n0 is busy
-    var n0 = NodeInfo.of(10, "host0", 12345);
-    var p0 = createFakeReplicaLeaderInfo("aa", 1, n0);
-    // n1 is free
-    var n1 = NodeInfo.of(11, "host1", 12345);
-    var p1 = createFakeReplicaLeaderInfo("aa", 2, n1);
-
-    var receiver = Mockito.mock(Receiver.class);
-    Mockito.when(receiver.current()).thenReturn(List.of());
-
-    var costFunction1 =
-        new HasBrokerCost() {
-          @Override
-          public BrokerCost brokerCost(ClusterInfo clusterInfo) {
-            var brokerCost =
-                clusterInfo.clusterBean().all().keySet().stream()
-                    .collect(
-                        Collectors.toMap(
-                            Function.identity(), id -> id.equals(n0.id()) ? 0.9D : 0.5D));
-            return () -> brokerCost;
-          }
-        };
-    var costFunction2 =
-        new HasBrokerCost() {
-          @Override
-          public BrokerCost brokerCost(ClusterInfo clusterInfo) {
-            var brokerCost =
-                clusterInfo.clusterBean().all().keySet().stream()
-                    .collect(
-                        Collectors.toMap(
-                            Function.identity(), id -> id.equals(n0.id()) ? 0.6D : 0.8D));
-            return () -> brokerCost;
-          }
-        };
-
-    try (var dispatcher =
-        new StrictCostDispatcher(List.of(costFunction1, costFunction2)) {
-          @Override
-          Receiver receiver(String host, int port) {
-            return receiver;
-          }
-
-          @Override
-          int jmxPort(int id) {
-            return 0;
-          }
-        }) {
-      var clusterInfo = Mockito.mock(ClusterInfo.class);
-
-      // there is no available partition
-      Mockito.when(clusterInfo.availableReplicaLeaders("aa")).thenReturn(List.of());
-      Mockito.when(clusterInfo.topics()).thenReturn(Set.of("aa"));
-      Mockito.when(clusterInfo.clusterBean()).thenReturn(ClusterBean.of(Map.of()));
-      Assertions.assertEquals(0, dispatcher.partition("aa", new byte[0], new byte[0], clusterInfo));
-
-      // there is only one available partition
-      Mockito.when(clusterInfo.availableReplicaLeaders("aa")).thenReturn(List.of(p0));
-      Assertions.assertEquals(1, dispatcher.partition("aa", new byte[0], new byte[0], clusterInfo));
-
-      // there is no beans, so it just returns the partition.
-      Mockito.when(clusterInfo.availableReplicaLeaders("aa")).thenReturn(List.of(p0, p1));
-      Assertions.assertEquals(2, dispatcher.partition("aa", new byte[0], new byte[0], clusterInfo));
+  void testNoAvailableBrokers() {
+    var clusterInfo = Mockito.mock(ClusterInfo.class);
+    Mockito.when(clusterInfo.availableReplicaLeaders(Mockito.anyString())).thenReturn(List.of());
+    try (var dispatcher = new StrictCostDispatcher(List.of())) {
+      Assertions.assertEquals(
+          0, dispatcher.partition("topic", new byte[0], new byte[0], clusterInfo));
+    }
+  }
 
-      // cost functions with weight
-      dispatcher.functions = Map.of(costFunction1, 0.1, costFunction2, 0.9);
-      Assertions.assertEquals(1, dispatcher.partition("aa", new byte[0], new byte[0], clusterInfo));
+  @Test
+  void testSingleBroker() {
+    var nodeInfo = NodeInfo.of(10, "host", 11111);
+    var replicaInfo = ReplicaInfo.of("topic", 10, nodeInfo, true, true, true);
+    var clusterInfo = Mockito.mock(ClusterInfo.class);
+    Mockito.when(clusterInfo.availableReplicaLeaders(Mockito.anyString()))
+        .thenReturn(List.of(replicaInfo));
+    try (var dispatcher = new StrictCostDispatcher(List.of())) {
+      Assertions.assertEquals(
+          10, dispatcher.partition("topic", new byte[0], new byte[0], clusterInfo));
     }
   }
 
@@ -199,4 +143,64 @@ public class StrictCostDispatcherTest {
         IllegalArgumentException.class,
         () -> StrictCostDispatcher.parseCostFunctionWeight(config2));
   }
+
+  @Test
+  void testCostFunctionWithoutFetcher() {
+    var costFunction = new CostFunction() {};
+    var replicaInfo0 = ReplicaInfo.of("topic", 0, NodeInfo.of(10, "host", 11111), true, true, true);
+    var replicaInfo1 =
+        ReplicaInfo.of("topic", 1, NodeInfo.of(12, "host2", 11111), true, true, true);
+    var clusterInfo = Mockito.mock(ClusterInfo.class);
+    Mockito.when(clusterInfo.availableReplicaLeaders(Mockito.anyString()))
+        .thenReturn(List.of(replicaInfo0, replicaInfo1));
+    Mockito.when(clusterInfo.clusterBean()).thenReturn(ClusterBean.of(Map.of()));
+    try (var dispatcher = new StrictCostDispatcher(List.of(costFunction))) {
+      dispatcher.partition("topic", new byte[0], new byte[0], clusterInfo);
+      Assertions.assertNotNull(dispatcher.roundRobin);
+      Assertions.assertEquals(0, dispatcher.receivers.size());
+    }
+  }
+
+  @Test
+  void testReceivers() {
+    var costFunction =
+        new CostFunction() {
+          @Override
+          public Optional<Fetcher> fetcher() {
+            return Optional.of(Mockito.mock(Fetcher.class));
+          }
+        };
+    var jmxPort = 12345;
+    var count = new AtomicInteger();
+    var dispatcher =
+        new StrictCostDispatcher(List.of(costFunction)) {
+          @Override
+          Receiver receiver(String host, int port, Fetcher fetcher) {
+            Assertions.assertEquals(jmxPort, port);
+            count.incrementAndGet();
+            return Mockito.mock(Receiver.class);
+          }
+        };
+    dispatcher.jmxPortGetter = ignored -> jmxPort;
+    var replicaInfo0 = ReplicaInfo.of("topic", 0, NodeInfo.of(10, "host", 11111), true, true, true);
+    var replicaInfo1 = ReplicaInfo.of("topic", 1, NodeInfo.of(10, "host", 11111), true, true, true);
+    var replicaInfo2 =
+        ReplicaInfo.of("topic", 1, NodeInfo.of(11, "host2", 11111), true, true, true);
+    var clusterInfo = Mockito.mock(ClusterInfo.class);
+    Mockito.when(clusterInfo.availableReplicaLeaders(Mockito.anyString()))
+        .thenReturn(List.of(replicaInfo0, replicaInfo1, replicaInfo2));
+    Mockito.when(clusterInfo.clusterBean()).thenReturn(ClusterBean.of(Map.of()));
+    // there is no receivers by default
+    Assertions.assertEquals(0, dispatcher.receivers.size());
+
+    // generate two receivers since there are two brokers (hosting three replicas)
+    dispatcher.partition("topic", new byte[0], new byte[0], clusterInfo);
+    Assertions.assertEquals(2, dispatcher.receivers.size());
+    Assertions.assertEquals(2, count.get());
+
+    // all brokers have receivers already so no new receiver is born
+    dispatcher.partition("topic", new byte[0], new byte[0], clusterInfo);
+    Assertions.assertEquals(2, dispatcher.receivers.size());
+    Assertions.assertEquals(2, count.get());
+  }
 }