[COST] Fix NetworkCost raises an exception when some brokers have no replica

common/src/main/java/org/astraea/common/cost/NetworkCost.java

@@ -31,13 +31,15 @@
 import org.astraea.common.admin.BrokerTopic;
 import org.astraea.common.admin.ClusterBean;
 import org.astraea.common.admin.ClusterInfo;
+import org.astraea.common.admin.NodeInfo;
 import org.astraea.common.admin.Replica;
 import org.astraea.common.admin.TopicPartition;
 import org.astraea.common.metrics.HasBeanObject;
 import org.astraea.common.metrics.broker.HasRate;
 import org.astraea.common.metrics.broker.LogMetrics;
 import org.astraea.common.metrics.broker.ServerMetrics;
 import org.astraea.common.metrics.collector.MetricSensor;
+import org.astraea.common.metrics.platform.HostMetrics;
 
 /**
  * This cost function calculate the load balance score in terms of network ingress or network
@@ -139,6 +141,10 @@ public ClusterCost clusterCost(ClusterInfo clusterInfo, ClusterBean clusterBean)
                           }
                         },
                         Collectors.summingDouble(x -> x))));
+    // add the brokers having no replicas into map
+    clusterInfo.nodes().stream()
+        .filter(node -> !brokerRate.containsKey(node))
+        .forEach(node -> brokerRate.put(node, 0.0));
 
     var summary = brokerRate.values().stream().mapToDouble(x -> x).summaryStatistics();
     if (summary.getMax() < 0)
@@ -151,13 +157,8 @@ public ClusterCost clusterCost(ClusterInfo clusterInfo, ClusterBean clusterBean)
       return ClusterCost.of(
           0, () -> "network load zero"); // edge case to avoid divided by zero error
     double score = (summary.getMax() - summary.getMin()) / (summary.getMax());
-    return ClusterCost.of(
-        score,
-        () ->
-            brokerRate.values().stream()
-                .map(x -> DataRate.Byte.of(x.longValue()).perSecond())
-                .map(DataRate::toString)
-                .collect(Collectors.joining(", ", "{", "}")));
+
+    return new NetworkClusterCost(score, brokerRate);
   }
 
   @Override
@@ -168,6 +169,7 @@ public Optional<MetricSensor> metricSensor() {
     return Optional.of(
         (client, clusterBean) ->
             Stream.of(
+                    List.of(HostMetrics.jvmMemory(client)),
                     ServerMetrics.Topic.BYTES_IN_PER_SEC.fetch(client),
                     ServerMetrics.Topic.BYTES_OUT_PER_SEC.fetch(client),
                     LogMetrics.Log.SIZE.fetch(client))
@@ -258,4 +260,26 @@ public String toString() {
       return alias();
     }
   }
+
+  static class NetworkClusterCost implements ClusterCost {
+    final double score;
+    final Map<NodeInfo, Double> brokerRate;
+
+    NetworkClusterCost(double score, Map<NodeInfo, Double> brokerRate) {
+      this.score = score;
+      this.brokerRate = brokerRate;
+    }
+
+    public double value() {
+      return score;
+    }
+
+    @Override
+    public String toString() {
+      return brokerRate.values().stream()
+          .map(x -> DataRate.Byte.of(x.longValue()).perSecond())
+          .map(DataRate::toString)
+          .collect(Collectors.joining(", ", "{", "}"));
+    }
+  }
 }

common/src/test/java/org/astraea/common/cost/NetworkCostTest.java

@@ -33,6 +33,7 @@
 import java.util.stream.Stream;
 import org.astraea.common.Configuration;
 import org.astraea.common.DataRate;
+import org.astraea.common.Utils;
 import org.astraea.common.admin.BrokerTopic;
 import org.astraea.common.admin.ClusterBean;
 import org.astraea.common.admin.ClusterInfo;
@@ -46,6 +47,9 @@
 import org.astraea.common.metrics.MetricSeriesBuilder;
 import org.astraea.common.metrics.broker.LogMetrics;
 import org.astraea.common.metrics.broker.ServerMetrics;
+import org.astraea.common.metrics.collector.MetricCollector;
+import org.astraea.it.Service;
+import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.DisplayName;
@@ -57,6 +61,13 @@
 
 class NetworkCostTest {
 
+  private static final Service SERVICE = Service.builder().numberOfBrokers(3).build();
+
+  @AfterAll
+  static void close() {
+    SERVICE.close();
+  }
+
   private static NetworkIngressCost ingressCost() {
     return new NetworkIngressCost() {
       @Override
@@ -421,6 +432,47 @@ void testSingleStepImprovement(int seed) {
     System.out.println("False test: " + counting.get(false));
   }
 
+  @Test
+  void testZeroReplicaBroker() {
+    var testcase = new LargeTestCase(1, 1, 0);
+    var beans = testcase.clusterBean();
+    var cluster = testcase.clusterInfo();
+    var scaledCluster =
+        ClusterInfoBuilder.builder(cluster)
+            .addNode(Set.of(4321))
+            .addFolders(Map.of(4321, Set.of("/folder")))
+            .build();
+    var node = scaledCluster.node(4321);
+
+    var costI =
+        (NetworkCost.NetworkClusterCost) new NetworkIngressCost().clusterCost(scaledCluster, beans);
+    Assertions.assertEquals(2, costI.brokerRate.size());
+    Assertions.assertEquals(0.0, costI.brokerRate.get(node));
+    var costE =
+        (NetworkCost.NetworkClusterCost) new NetworkEgressCost().clusterCost(scaledCluster, beans);
+    Assertions.assertEquals(2, costE.brokerRate.size());
+    Assertions.assertEquals(0.0, costE.brokerRate.get(node));
+  }
+
+  @Test
+  void testNoMetricCheck() {
+    try (var collector = MetricCollector.builder().interval(Duration.ofMillis(100)).build()) {
+      var ingressCost = new NetworkIngressCost();
+
+      // setup sampling
+      SERVICE.dataFolders().keySet().forEach(collector::registerLocalJmx);
+      ingressCost.metricSensor().ifPresent(collector::addMetricSensor);
+
+      // sample metrics for a while.
+      Utils.sleep(Duration.ofMillis(500));
+
+      var emptyCluster = ClusterInfoBuilder.builder().addNode(Set.of(1, 2, 3)).build();
+      Assertions.assertDoesNotThrow(
+          () -> ingressCost.clusterCost(emptyCluster, collector.clusterBean()),
+          "Should not raise an exception");
+    }
+  }
+
   interface TestCase {
 
     ClusterInfo clusterInfo();