refactor replica migrate size moveCost

app/src/main/java/org/astraea/app/web/BalancerHandler.java

@@ -259,18 +259,18 @@ private static List<MigrationCost> migrationCosts(Balancer.Plan solution) {
                         Collectors.toMap(
                             e -> String.valueOf(e.getKey()), e -> (double) e.getValue()))),
             new MigrationCost(
-                MOVED_SIZE,
+                MOVED_IN_SIZE,
                 ClusterInfo.changedRecordSize(
-                        solution.initialClusterInfo(), solution.proposal(), ignored -> true)
+                        solution.initialClusterInfo(), solution.proposal(), ignored -> true, false)
                     .entrySet()
                     .stream()
                     .collect(
                         Collectors.toMap(
                             e -> String.valueOf(e.getKey()), e -> (double) e.getValue().bytes()))),
             new MigrationCost(
-                MOVED_LEADER_SIZE,
+                MOVED_OUT_SIZE,
                 ClusterInfo.changedRecordSize(
-                        solution.initialClusterInfo(), solution.proposal(), Replica::isLeader)
+                        solution.initialClusterInfo(), solution.proposal(), ignored -> true, true)
                     .entrySet()
                     .stream()
                     .collect(
@@ -471,8 +471,8 @@ static Change from(Collection<Replica> before, Collection<Replica> after) {
   // visible for testing
   static final String CHANGED_REPLICAS = "changed replicas";
   static final String CHANGED_LEADERS = "changed leaders";
-  static final String MOVED_SIZE = "moved size (bytes)";
-  static final String MOVED_LEADER_SIZE = "moved leader size (bytes)";
+  static final String MOVED_IN_SIZE = "moved in size (bytes)";
+  static final String MOVED_OUT_SIZE = "moved out size (bytes)";
 
   static class MigrationCost {
     final String name;

app/src/test/java/org/astraea/app/web/BalancerHandlerTest.java

@@ -147,7 +147,7 @@ void testReport() {
           .forEach(p -> Assertions.assertEquals(Optional.empty(), p.size));
       var sizeMigration =
           report.migrationCosts.stream()
-              .filter(x -> x.name.equals(BalancerHandler.MOVED_SIZE))
+              .filter(x -> x.name.equals(BalancerHandler.MOVED_IN_SIZE))
               .findFirst()
               .get();
       Assertions.assertNotEquals(0, sizeMigration.brokerCosts.size());
@@ -175,7 +175,7 @@ void testTopics() {
           "Only allowed topics been altered");
       var sizeMigration =
           report.migrationCosts.stream()
-              .filter(x -> x.name.equals(BalancerHandler.MOVED_SIZE))
+              .filter(x -> x.name.equals(BalancerHandler.MOVED_IN_SIZE))
               .findFirst()
               .get();
       Assertions.assertNotEquals(0, sizeMigration.brokerCosts.size());
@@ -361,7 +361,8 @@ void testMoveCost(String leaderLimit, String sizeLimit) {
       report.migrationCosts.forEach(
           migrationCost -> {
             switch (migrationCost.name) {
-              case BalancerHandler.MOVED_SIZE:
+              case BalancerHandler.MOVED_IN_SIZE:
+              case BalancerHandler.MOVED_OUT_SIZE:
                 Assertions.assertTrue(
                     migrationCost.brokerCosts.values().stream()
                             .map(Math::abs)

common/src/main/java/org/astraea/common/admin/ClusterInfo.java

@@ -39,23 +39,59 @@ static ClusterInfo empty() {
 
   // ---------------------[helpers]---------------------//
 
+  /**
+   * @param before the ClusterInfo before migrated replicas
+   * @param after the ClusterInfo after migrated replicas
+   * @param predicate used with filter to filter some replicas
+   * @param migrateOut if data log need fetch from replica leader, set this true
+   * @return the data size to migrated by all brokers
+   */
   static Map<Integer, DataSize> changedRecordSize(
-      ClusterInfo before, ClusterInfo after, Predicate<Replica> predicate) {
-    return Stream.concat(before.nodes().stream(), after.nodes().stream())
+      ClusterInfo before, ClusterInfo after, Predicate<Replica> predicate, boolean migrateOut) {
+    final ClusterInfo sourceClusterInfo;
+    final ClusterInfo destClusterInfo;
+    if (migrateOut) {
+      sourceClusterInfo = after;
+      destClusterInfo = before;
+    } else {
+      sourceClusterInfo = before;
+      destClusterInfo = after;
+    }
+    var changePartitions =
+        ClusterInfo.findNonFulfilledAllocation(sourceClusterInfo, destClusterInfo);
+    var cost =
+        changePartitions.stream()
+            .flatMap(
+                p ->
+                    destClusterInfo.replicas(p).stream()
+                        .filter(predicate)
+                        .filter(r -> !sourceClusterInfo.replicas(p).contains(r)))
+            .map(
+                r -> {
+                  if (migrateOut)
+                    return destClusterInfo.replicaLeader(r.topicPartition()).orElse(r);
+                  return r;
+                })
+            .collect(
+                Collectors.groupingBy(
+                    r -> r.nodeInfo().id(),
+                    Collectors.mapping(
+                        Function.identity(), Collectors.summingLong(Replica::size))));
+    return Stream.concat(destClusterInfo.nodes().stream(), sourceClusterInfo.nodes().stream())
         .map(NodeInfo::id)
         .distinct()
         .parallel()
         .collect(
-            Collectors.toUnmodifiableMap(
-                Function.identity(),
-                id ->
-                    DataSize.Byte.of(
-                        after.replicaStream(id).filter(predicate).mapToLong(Replica::size).sum()
-                            - before
-                                .replicaStream(id)
-                                .filter(predicate)
-                                .mapToLong(Replica::size)
-                                .sum())));
+            Collectors.toMap(Function.identity(), n -> DataSize.Byte.of(cost.getOrDefault(n, 0L))));
+  }
+
+  static Long totalChangedRecordSize(
+      ClusterInfo before, ClusterInfo after, Predicate<Replica> predicate) {
+    var addingSize = changedRecordSize(before, after, predicate, false);
+    var removedSize = changedRecordSize(before, after, predicate, true);
+    return Math.max(
+        addingSize.values().stream().mapToLong(DataSize::bytes).sum(),
+        removedSize.values().stream().mapToLong(DataSize::bytes).sum());
   }
 
   static Map<Integer, Integer> changedReplicaNumber(

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

@@ -51,20 +51,14 @@ public BrokerCost brokerCost(ClusterInfo clusterInfo, ClusterBean clusterBean) {
 
   @Override
   public MoveCost moveCost(ClusterInfo before, ClusterInfo after, ClusterBean clusterBean) {
-    var moveCost = ClusterInfo.changedRecordSize(before, after, ignored -> true);
+    var moveCost = ClusterInfo.totalChangedRecordSize(before, after, ignored -> true);
     var maxMigratedSize =
         config
             .string(MAX_MIGRATE_SIZE_KEY)
             .map(DataSize::of)
             .map(DataSize::bytes)
             .orElse(Long.MAX_VALUE);
-    var overflow =
-        maxMigratedSize
-            < moveCost.values().stream()
-                .map(DataSize::bytes)
-                .map(Math::abs)
-                .mapToLong(s -> s)
-                .sum();
+    var overflow = maxMigratedSize < moveCost;
     return () -> overflow;
   }
 

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

@@ -58,16 +58,10 @@ public Optional<MetricSensor> metricSensor() {
 
   @Override
   public MoveCost moveCost(ClusterInfo before, ClusterInfo after, ClusterBean clusterBean) {
-    var moveCost = ClusterInfo.changedRecordSize(before, after, Replica::isLeader);
+    var moveCost = ClusterInfo.totalChangedRecordSize(before, after, Replica::isLeader);
     var maxMigratedLeaderSize =
         config.string(COST_LIMIT_KEY).map(DataSize::of).map(DataSize::bytes).orElse(Long.MAX_VALUE);
-    var overflow =
-        maxMigratedLeaderSize
-            < moveCost.values().stream()
-                .map(DataSize::bytes)
-                .map(Math::abs)
-                .mapToLong(s -> s)
-                .sum();
+    var overflow = maxMigratedLeaderSize < moveCost;
     return () -> overflow;
   }