Repository: skiptests/astraea
Pull request #273 (2022-03-03 08:22:20)

Add more tests for StrictCostDispatcher
1. more tests
2. remove `NodeId`


Patch:
diff --git a/app/src/main/java/org/astraea/automation/Automation.java b/app/src/main/java/org/astraea/automation/Automation.java
index 8a0679d3..7486c675 100644
--- a/app/src/main/java/org/astraea/automation/Automation.java
+++ b/app/src/main/java/org/astraea/automation/Automation.java
@@ -54,15 +54,14 @@ public class Automation {
       else times = Integer.parseInt(properties.getProperty("--time"));
 
       while (i < times) {
-        var str =
+        var result =
             Performance.execute(
                 org.astraea.argument.Argument.parse(
                     new Performance.Argument(), performanceArgs(properties)));
         i++;
         if (whetherDeleteTopic) {
           try (final AdminClient adminClient = KafkaAdminClient.create(config)) {
-            var topicName = str.get();
-            adminClient.deleteTopics(List.of(topicName));
+            adminClient.deleteTopics(List.of(result.topicName()));
           }
         }
         System.out.println("=============== " + i + " time Performance Complete! ===============");
diff --git a/app/src/main/java/org/astraea/partitioner/NodeId.java b/app/src/main/java/org/astraea/partitioner/NodeId.java
deleted file mode 100644
index caa0fc53..00000000
--- a/app/src/main/java/org/astraea/partitioner/NodeId.java
+++ /dev/null
@@ -1,18 +0,0 @@
-package org.astraea.partitioner;
-
-/** We use this interface to replace the integer for more readable code. */
-@FunctionalInterface
-public interface NodeId extends Comparable<NodeInfo> {
-
-  static NodeId of(int id) {
-    return () -> id;
-  }
-
-  /** @return The id for this node */
-  int id();
-
-  @Override
-  default int compareTo(NodeInfo o) {
-    return Integer.compare(id(), o.id());
-  }
-}
diff --git a/app/src/main/java/org/astraea/partitioner/NodeInfo.java b/app/src/main/java/org/astraea/partitioner/NodeInfo.java
index 9a54c41f..f0f8029a 100644
--- a/app/src/main/java/org/astraea/partitioner/NodeInfo.java
+++ b/app/src/main/java/org/astraea/partitioner/NodeInfo.java
@@ -2,7 +2,7 @@ package org.astraea.partitioner;
 
 import java.util.Objects;
 
-public interface NodeInfo extends Comparable<NodeInfo>, NodeId {
+public interface NodeInfo extends Comparable<NodeInfo> {
 
   static NodeInfo of(org.apache.kafka.common.Node node) {
     return of(node.id(), node.host(), node.port());
@@ -52,4 +52,7 @@ public interface NodeInfo extends Comparable<NodeInfo>, NodeId {
 
   /** @return The client (kafka data, jmx, etc.) port for this node */
   int port();
+
+  /** @return id of broker node. it must be unique. */
+  int id();
 }
diff --git a/app/src/main/java/org/astraea/partitioner/PartitionInfo.java b/app/src/main/java/org/astraea/partitioner/PartitionInfo.java
index 7e1b997d..197b77b5 100644
--- a/app/src/main/java/org/astraea/partitioner/PartitionInfo.java
+++ b/app/src/main/java/org/astraea/partitioner/PartitionInfo.java
@@ -3,20 +3,24 @@ package org.astraea.partitioner;
 public interface PartitionInfo {
 
   static PartitionInfo of(org.apache.kafka.common.PartitionInfo pf) {
+    return of(pf.topic(), pf.partition(), NodeInfo.of(pf.leader()));
+  }
+
+  static PartitionInfo of(String topic, int partition, NodeInfo leader) {
     return new PartitionInfo() {
       @Override
       public String topic() {
-        return pf.topic();
+        return topic;
       }
 
       @Override
       public int partition() {
-        return pf.partition();
+        return partition;
       }
 
       @Override
       public NodeInfo leader() {
-        return NodeInfo.of(pf.leader());
+        return leader;
       }
     };
   }
diff --git a/app/src/main/java/org/astraea/partitioner/StrictCostDispatcher.java b/app/src/main/java/org/astraea/partitioner/StrictCostDispatcher.java
index d5dbb3e1..d23be8d1 100644
--- a/app/src/main/java/org/astraea/partitioner/StrictCostDispatcher.java
+++ b/app/src/main/java/org/astraea/partitioner/StrictCostDispatcher.java
@@ -31,8 +31,8 @@ public class StrictCostDispatcher implements Dispatcher {
       BeanCollector.builder().interval(Duration.ofSeconds(4)).build();
   private final Collection<CostFunction> functions;
   private Optional<Integer> jmxPortDefault = Optional.empty();
-  private final Map<NodeId, Integer> jmxPorts = new TreeMap<>();
-  private final Map<NodeId, Receiver> receivers = new TreeMap<>();
+  private final Map<Integer, Integer> jmxPorts = new TreeMap<>();
+  final Map<Integer, Receiver> receivers = new TreeMap<>();
 
   public StrictCostDispatcher() {
     this(List.of(CostFunction.throughput()));
@@ -43,34 +43,22 @@ public class StrictCostDispatcher implements Dispatcher {
     this.functions = functions;
   }
 
-  private static Fetcher of(Collection<Fetcher> fs) {
-    return client ->
-        fs.stream().flatMap(f -> f.fetch(client).stream()).collect(Collectors.toUnmodifiableList());
-  }
-
   @Override
   public int partition(String topic, byte[] key, byte[] value, ClusterInfo clusterInfo) {
     var partitions = clusterInfo.availablePartitions(topic);
     // just return first partition if there is no available partitions
     if (partitions.isEmpty()) return 0;
 
+    // just return the only one available partition
+    if (partitions.size() == 1) return partitions.iterator().next().partition();
+
     // add new receivers for new brokers
     partitions.stream()
-        .filter(p -> !receivers.containsKey(NodeId.of(p.leader().id())))
+        .filter(p -> !receivers.containsKey(p.leader().id()))
         .forEach(
             p ->
                 receivers.put(
-                    p.leader(),
-                    beanCollector
-                        .register()
-                        .host(p.leader().host())
-                        .port(jmxPort(p.leader().id()))
-                        .fetcher(
-                            Fetcher.of(
-                                functions.stream()
-                                    .map(CostFunction::fetcher)
-                                    .collect(Collectors.toUnmodifiableList())))
-                        .build()));
+                    p.leader().id(), receiver(p.leader().host(), jmxPort(p.leader().id()))));
 
     // get latest beans for each node
     var beans =
@@ -83,15 +71,32 @@ public class StrictCostDispatcher implements Dispatcher {
             .map(f -> f.cost(beans, clusterInfo))
             .collect(Collectors.toUnmodifiableList());
 
-    // return the partition (node) having min score
+    return bestPartition(partitions, scores).map(e -> e.getKey().partition()).orElse(0);
+  }
+
+  // visible for testing
+  static Optional<Map.Entry<PartitionInfo, Double>> bestPartition(
+      List<PartitionInfo> partitions, List<Map<Integer, Double>> scores) {
     return partitions.stream()
         .map(
             p ->
                 Map.entry(
-                    p, scores.stream().mapToDouble(s -> s.getOrDefault(p.leader(), 0.0D)).sum()))
-        .min(Map.Entry.comparingByValue())
-        .map(e -> e.getKey().partition())
-        .orElse(0);
+                    p,
+                    scores.stream().mapToDouble(s -> s.getOrDefault(p.leader().id(), 0.0D)).sum()))
+        .min(Map.Entry.comparingByValue());
+  }
+
+  Receiver receiver(String host, int port) {
+    return beanCollector
+        .register()
+        .host(host)
+        .port(port)
+        .fetcher(
+            Fetcher.of(
+                functions.stream()
+                    .map(CostFunction::fetcher)
+                    .collect(Collectors.toUnmodifiableList())))
+        .build();
   }
 
   @Override
@@ -102,15 +107,12 @@ public class StrictCostDispatcher implements Dispatcher {
     config.entrySet().stream()
         .filter(e -> e.getKey().startsWith("broker."))
         .filter(e -> e.getKey().endsWith(JMX_PORT))
-        .forEach(
-            e ->
-                jmxPorts.put(
-                    NodeId.of(Integer.parseInt(e.getKey())), Integer.parseInt(e.getValue())));
+        .forEach(e -> jmxPorts.put(Integer.parseInt(e.getKey()), Integer.parseInt(e.getValue())));
   }
 
   // visible for testing
   int jmxPort(int id) {
-    if (jmxPorts.containsKey(NodeId.of(id))) return jmxPorts.get(NodeId.of(id));
+    if (jmxPorts.containsKey(id)) return jmxPorts.get(id);
     return jmxPortDefault.orElseThrow(
         () -> new NoSuchElementException("broker: " + id + " does not have jmx port"));
   }
@@ -118,5 +120,6 @@ public class StrictCostDispatcher implements Dispatcher {
   @Override
   public void close() {
     receivers.values().forEach(Utils::close);
+    receivers.clear();
   }
 }
diff --git a/app/src/main/java/org/astraea/partitioner/cost/CostFunction.java b/app/src/main/java/org/astraea/partitioner/cost/CostFunction.java
index c22909da..81475d73 100644
--- a/app/src/main/java/org/astraea/partitioner/cost/CostFunction.java
+++ b/app/src/main/java/org/astraea/partitioner/cost/CostFunction.java
@@ -5,7 +5,6 @@ import java.util.Map;
 import org.astraea.metrics.HasBeanObject;
 import org.astraea.metrics.collector.Fetcher;
 import org.astraea.partitioner.ClusterInfo;
-import org.astraea.partitioner.NodeId;
 
 public interface CostFunction {
 
@@ -20,7 +19,7 @@ public interface CostFunction {
    * @param clusterInfo cluster information
    * @return the score of each node. the score range is [0 - 1]
    */
-  Map<NodeId, Double> cost(Map<NodeId, List<HasBeanObject>> beans, ClusterInfo clusterInfo);
+  Map<Integer, Double> cost(Map<Integer, List<HasBeanObject>> beans, ClusterInfo clusterInfo);
 
   /** @return the metrics getters. Those getters are used to fetch mbeans. */
   Fetcher fetcher();
diff --git a/app/src/main/java/org/astraea/partitioner/cost/ThroughputCost.java b/app/src/main/java/org/astraea/partitioner/cost/ThroughputCost.java
index 46d20994..7464f92b 100644
--- a/app/src/main/java/org/astraea/partitioner/cost/ThroughputCost.java
+++ b/app/src/main/java/org/astraea/partitioner/cost/ThroughputCost.java
@@ -8,21 +8,23 @@ import org.astraea.metrics.collector.Fetcher;
 import org.astraea.metrics.kafka.BrokerTopicMetricsResult;
 import org.astraea.metrics.kafka.KafkaMetrics;
 import org.astraea.partitioner.ClusterInfo;
-import org.astraea.partitioner.NodeId;
+import org.astraea.partitioner.NodeInfo;
 
 public class ThroughputCost implements CostFunction {
 
   @Override
-  public Map<NodeId, Double> cost(Map<NodeId, List<HasBeanObject>> beans, ClusterInfo clusterInfo) {
+  public Map<Integer, Double> cost(
+      Map<Integer, List<HasBeanObject>> beans, ClusterInfo clusterInfo) {
     var score = score(beans);
 
     var max = score.values().stream().mapToDouble(v -> v).max().orElse(1);
 
     return clusterInfo.nodes().stream()
+        .map(NodeInfo::id)
         .collect(Collectors.toMap(n -> n, n -> score.getOrDefault(n, 0.0D) / max));
   }
 
-  Map<NodeId, Double> score(Map<NodeId, List<HasBeanObject>> beans) {
+  Map<Integer, Double> score(Map<Integer, List<HasBeanObject>> beans) {
     // TODO: this implementation only consider the oneMinuteRate ...
     return beans.entrySet().stream()
         .collect(
diff --git a/app/src/main/java/org/astraea/performance/Performance.java b/app/src/main/java/org/astraea/performance/Performance.java
index fbb768b1..8b521746 100644
--- a/app/src/main/java/org/astraea/performance/Performance.java
+++ b/app/src/main/java/org/astraea/performance/Performance.java
@@ -8,9 +8,7 @@ import java.util.List;
 import java.util.Map;
 import java.util.Random;
 import java.util.Set;
-import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
-import java.util.concurrent.Future;
 import java.util.function.BiConsumer;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
@@ -68,10 +66,9 @@ public class Performance {
     execute(org.astraea.argument.Argument.parse(new Argument(), args));
   }
 
-  public static Future<String> execute(final Argument param)
+  public static Result execute(final Argument param)
       throws InterruptedException, IOException, ExecutionException {
     List<Integer> partitions;
-    var future = new CompletableFuture<String>();
     try (var topicAdmin = TopicAdmin.of(param.props())) {
       topicAdmin
           .creator()
@@ -132,8 +129,7 @@ public class Performance {
             .executor((param.createCSV) ? new FileWriter(manager, tracker) : () -> State.DONE)
             .build()) {
       threadPool.waitAll();
-      future.complete(param.topic);
-      return future;
+      return new Result(param.topic);
     }
   }
 
@@ -340,4 +336,16 @@ public class Performance {
         validateWith = NonEmptyStringField.class)
     List<Integer> specifyBroker = List.of(-1);
   }
+
+  public static class Result {
+    private final String topicName;
+
+    private Result(String topicName) {
+      this.topicName = topicName;
+    }
+
+    public String topicName() {
+      return topicName;
+    }
+  }
 }


Test patch:
diff --git a/app/src/test/java/org/astraea/partitioner/StrictCostDispatcherTest.java b/app/src/test/java/org/astraea/partitioner/StrictCostDispatcherTest.java
index fd7a03eb..f194629a 100644
--- a/app/src/test/java/org/astraea/partitioner/StrictCostDispatcherTest.java
+++ b/app/src/test/java/org/astraea/partitioner/StrictCostDispatcherTest.java
@@ -1,16 +1,28 @@
 package org.astraea.partitioner;
 
+import java.util.List;
 import java.util.Map;
 import java.util.NoSuchElementException;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+import org.astraea.metrics.HasBeanObject;
+import org.astraea.metrics.collector.Fetcher;
+import org.astraea.metrics.collector.Receiver;
+import org.astraea.partitioner.cost.CostFunction;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
+import org.mockito.Mockito;
 
 public class StrictCostDispatcherTest {
 
   @Test
   void testClose() {
     var dispatcher = new StrictCostDispatcher();
+    var receiver = Mockito.mock(Receiver.class);
+    dispatcher.receivers.put(10, receiver);
     dispatcher.close();
+    Mockito.verify(receiver, Mockito.times(1)).close();
+    Assertions.assertEquals(0, dispatcher.receivers.size());
   }
 
   @Test
@@ -22,4 +34,74 @@ public class StrictCostDispatcherTest {
       Assertions.assertEquals(12345, dispatcher.jmxPort(0));
     }
   }
+
+  @Test
+  void testBestPartition() {
+    var partitions =
+        List.of(
+            PartitionInfo.of("t", 0, NodeInfo.of(10, "h0", 1000)),
+            PartitionInfo.of("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());
+  }
+
+  @Test
+  void testPartition() {
+    // n0 is busy
+    var n0 = NodeInfo.of(10, "host0", 12345);
+    var p0 = PartitionInfo.of("aa", 1, n0);
+    // n1 is free
+    var n1 = NodeInfo.of(11, "host1", 12345);
+    var p1 = PartitionInfo.of("aa", 2, n1);
+
+    var receiver = Mockito.mock(Receiver.class);
+    Mockito.when(receiver.current()).thenReturn(List.of());
+
+    var costFunction =
+        new CostFunction() {
+          @Override
+          public Map<Integer, Double> cost(
+              Map<Integer, List<HasBeanObject>> beans, ClusterInfo clusterInfo) {
+            return beans.keySet().stream()
+                .collect(
+                    Collectors.toMap(Function.identity(), id -> id.equals(n0.id()) ? 0.9D : 0.5D));
+          }
+
+          @Override
+          public Fetcher fetcher() {
+            return client -> List.of();
+          }
+        };
+
+    var dispatcher =
+        new StrictCostDispatcher(List.of(costFunction)) {
+          @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.availablePartitions("aa")).thenReturn(List.of());
+    Assertions.assertEquals(0, dispatcher.partition("aa", new byte[0], new byte[0], clusterInfo));
+
+    // there is only one available partition
+    Mockito.when(clusterInfo.availablePartitions("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.availablePartitions("aa")).thenReturn(List.of(p0, p1));
+    Assertions.assertEquals(2, dispatcher.partition("aa", new byte[0], new byte[0], clusterInfo));
+  }
 }
diff --git a/app/src/test/java/org/astraea/partitioner/cost/ThroughputCostTest.java b/app/src/test/java/org/astraea/partitioner/cost/ThroughputCostTest.java
index 53ab3b20..48cb0d9d 100644
--- a/app/src/test/java/org/astraea/partitioner/cost/ThroughputCostTest.java
+++ b/app/src/test/java/org/astraea/partitioner/cost/ThroughputCostTest.java
@@ -19,9 +19,9 @@ public class ThroughputCostTest {
     var bean = Mockito.mock(BrokerTopicMetricsResult.class);
     Mockito.when(bean.oneMinuteRate()).thenReturn(100D);
 
-    var score = throughputCost.score(Map.of(node, List.of(bean)));
+    var score = throughputCost.score(Map.of(10, List.of(bean)));
     Assertions.assertEquals(1, score.size());
-    Assertions.assertEquals(100D, score.get(node));
+    Assertions.assertEquals(100D, score.get(10));
   }
 
   @Test
@@ -35,8 +35,8 @@ public class ThroughputCostTest {
     var cluster = Mockito.mock(ClusterInfo.class);
     Mockito.when(cluster.nodes()).thenReturn(List.of(node));
 
-    var cost = throughputCost.cost(Map.of(node, List.of(bean)), cluster);
+    var cost = throughputCost.cost(Map.of(10, List.of(bean)), cluster);
     Assertions.assertEquals(1, cost.size());
-    Assertions.assertEquals(1, cost.get(node));
+    Assertions.assertEquals(1, cost.get(10));
   }
 }