[COST] add MigrateTimeCost

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

@@ -196,7 +196,8 @@ private PlanExecutionProgress progress(String taskId) {
                         solution ->
                             new PlanReport(
                                 changes.apply(solution),
-                                MigrationCost.migrationCosts(contextCluster, solution.proposal())))
+                                MigrationCost.migrationCosts(
+                                    contextCluster, solution.proposal(), solution.clusterBean())))
                     .orElse(null);
     var phase = balancerConsole.taskPhase(taskId).orElseThrow();
     return new PlanExecutionProgress(
@@ -241,7 +242,8 @@ static class BalancerPostRequest implements Request {
             "org.astraea.common.cost.ReplicaLeaderCost",
             "org.astraea.common.cost.RecordSizeCost",
             "org.astraea.common.cost.ReplicaNumberCost",
-            "org.astraea.common.cost.ReplicaLeaderSizeCost");
+            "org.astraea.common.cost.ReplicaLeaderSizeCost",
+            "org.astraea.common.cost.PartitionMigrateTimeCost");
 
     HasClusterCost clusterCost() {
       if (clusterCosts.isEmpty())

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

@@ -32,7 +32,8 @@ public class SensorHandler implements Handler {
   private static final Set<String> DEFAULT_COSTS =
       Set.of(
           "org.astraea.common.cost.ReplicaLeaderCost",
-          "org.astraea.common.cost.NetworkIngressCost");
+          "org.astraea.common.cost.NetworkIngressCost",
+          "org.astraea.common.cost.PartitionMigrateTimeCost");
 
   SensorHandler(Sensors sensors) {
     this.sensors = sensors;

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

@@ -19,6 +19,7 @@
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
 import java.util.stream.Collectors;
 import org.apache.kafka.common.requests.DescribeLogDirsResponse;
@@ -89,6 +90,23 @@ public Map<TopicPartition, Long> orphanPartitionSizes() {
                         .map(p -> TopicPartition.of(topic.name(), p.partition())))
             .collect(Collectors.toUnmodifiableSet());
     return new Broker() {
+
+      private final int hashCode = Objects.hash(id(), host(), port());
+
+      @Override
+      public int hashCode() {
+        return hashCode;
+      }
+
+      @Override
+      public boolean equals(Object other) {
+        if (other instanceof NodeInfo) {
+          var node = (NodeInfo) other;
+          return id() == node.id() && port() == node.port() && host().equals(node.host());
+        }
+        return false;
+      }
+
       @Override
       public String host() {
         return nodeInfo.host();

common/src/main/java/org/astraea/common/balancer/Balancer.java

@@ -25,6 +25,7 @@
 import org.astraea.common.balancer.algorithms.SingleStepBalancer;
 import org.astraea.common.cost.ClusterCost;
 import org.astraea.common.cost.MoveCost;
+import org.astraea.common.metrics.ClusterBean;
 
 public interface Balancer {
 
@@ -34,6 +35,7 @@ public interface Balancer {
   Optional<Plan> offer(AlgorithmConfig config);
 
   class Plan {
+    private final ClusterBean clusterBean;
     private final ClusterInfo initialClusterInfo;
     private final ClusterCost initialClusterCost;
 
@@ -42,18 +44,24 @@ class Plan {
     private final MoveCost moveCost;
 
     public Plan(
+        ClusterBean clusterBean,
         ClusterInfo initialClusterInfo,
         ClusterCost initialClusterCost,
         ClusterInfo proposal,
         ClusterCost proposalClusterCost,
         MoveCost moveCost) {
+      this.clusterBean = clusterBean;
       this.initialClusterInfo = initialClusterInfo;
       this.initialClusterCost = initialClusterCost;
       this.proposal = proposal;
       this.proposalClusterCost = proposalClusterCost;
       this.moveCost = moveCost;
     }
 
+    public ClusterBean clusterBean() {
+      return clusterBean;
+    }
+
     public ClusterInfo initialClusterInfo() {
       return initialClusterInfo;
     }

common/src/main/java/org/astraea/common/balancer/algorithms/GreedyBalancer.java

@@ -164,6 +164,7 @@ public Optional<Plan> offer(AlgorithmConfig config) {
                 .map(
                     newAllocation ->
                         new Plan(
+                            config.clusterBean(),
                             config.clusterInfo(),
                             initialCost,
                             newAllocation,

common/src/main/java/org/astraea/common/balancer/algorithms/SingleStepBalancer.java

@@ -89,6 +89,7 @@ public Optional<Plan> offer(AlgorithmConfig config) {
         .map(
             newAllocation ->
                 new Plan(
+                    config.clusterBean(),
                     config.clusterInfo(),
                     currentCost,
                     newAllocation,

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

@@ -16,6 +16,8 @@
  */
 package org.astraea.common.cost;
 
+import static org.astraea.common.cost.PartitionMigrateTimeCost.brokerMaxRate;
+
 import java.util.List;
 import java.util.Map;
 import java.util.function.Function;
@@ -25,19 +27,21 @@
 import org.astraea.common.admin.ClusterInfo;
 import org.astraea.common.admin.NodeInfo;
 import org.astraea.common.admin.Replica;
+import org.astraea.common.metrics.ClusterBean;
 
 public class MigrationCost {
 
   public final String name;
-
   public final Map<Integer, Long> brokerCosts;
   public static final String TO_SYNC_BYTES = "record size to sync (bytes)";
   public static final String TO_FETCH_BYTES = "record size to fetch (bytes)";
   public static final String REPLICA_LEADERS_TO_ADDED = "leader number to add";
   public static final String REPLICA_LEADERS_TO_REMOVE = "leader number to remove";
   public static final String CHANGED_REPLICAS = "changed replicas";
+  public static final String PARTITION_MIGRATED_TIME = "partition migrated time";
 
-  public static List<MigrationCost> migrationCosts(ClusterInfo before, ClusterInfo after) {
+  public static List<MigrationCost> migrationCosts(
+      ClusterInfo before, ClusterInfo after, ClusterBean clusterBean) {
     var migrateInBytes = recordSizeToSync(before, after);
     var migrateOutBytes = recordSizeToFetch(before, after);
     var migrateReplicaNum = replicaNumChanged(before, after);
@@ -46,6 +50,8 @@ public static List<MigrationCost> migrationCosts(ClusterInfo before, ClusterInfo
     return List.of(
         new MigrationCost(TO_SYNC_BYTES, migrateInBytes),
         new MigrationCost(TO_FETCH_BYTES, migrateOutBytes),
+        new MigrationCost(CHANGED_REPLICAS, migrateReplicaNum),
+        new MigrationCost(PARTITION_MIGRATED_TIME, brokerMigrationTime(before, after, clusterBean)),
         new MigrationCost(REPLICA_LEADERS_TO_ADDED, migrateInLeader),
         new MigrationCost(REPLICA_LEADERS_TO_REMOVE, migrateOutLeader),
         new MigrationCost(CHANGED_REPLICAS, migrateReplicaNum));
@@ -76,6 +82,65 @@ static Map<Integer, Long> replicaLeaderToRemove(ClusterInfo before, ClusterInfo
     return migratedChanged(before, after, false, Replica::isLeader, ignore -> 1L);
   }
 
+  /**
+   * @param before the ClusterInfo before migrated replicas
+   * @param after the ClusterInfo after migrated replicas
+   * @param clusterBean cluster metrics
+   * @return estimated migrated time required by all brokers (seconds)
+   */
+  public static Map<Integer, Long> brokerMigrationTime(
+      ClusterInfo before, ClusterInfo after, ClusterBean clusterBean) {
+    var brokerInRate =
+        before.nodes().stream()
+            .collect(
+                Collectors.toMap(
+                    NodeInfo::id,
+                    nodeInfo ->
+                        brokerMaxRate(
+                            nodeInfo.id(),
+                            clusterBean,
+                            PartitionMigrateTimeCost.MaxReplicationInRateBean.class)));
+    var brokerOutRate =
+        before.nodes().stream()
+            .collect(
+                Collectors.toMap(
+                    NodeInfo::id,
+                    nodeInfo ->
+                        brokerMaxRate(
+                            nodeInfo.id(),
+                            clusterBean,
+                            PartitionMigrateTimeCost.MaxReplicationOutRateBean.class)));
+    var brokerMigrateInTime =
+        MigrationCost.recordSizeToFetch(before, after).entrySet().stream()
+            .map(
+                brokerSize ->
+                    Map.entry(
+                        brokerSize.getKey(),
+                        brokerSize.getValue()
+                            / brokerInRate.get(brokerSize.getKey()).orElse(Long.MAX_VALUE)))
+            .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+    var brokerMigrateOutTime =
+        MigrationCost.recordSizeToSync(before, after).entrySet().stream()
+            .map(
+                brokerSize ->
+                    Map.entry(
+                        brokerSize.getKey(),
+                        brokerSize.getValue()
+                            / brokerOutRate.get(brokerSize.getKey()).orElse(Long.MAX_VALUE)))
+            .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+    return Stream.concat(before.nodes().stream(), after.nodes().stream())
+        .distinct()
+        .map(
+            nodeInfo ->
+                Map.entry(
+                    nodeInfo.id(),
+                    (long)
+                        Math.max(
+                            brokerMigrateInTime.get(nodeInfo.id()),
+                            brokerMigrateOutTime.get(nodeInfo.id()))))
+        .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+  }
+
   /**
    * @param before the ClusterInfo before migrated replicas
    * @param after the ClusterInfo after migrated replicas

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

@@ -0,0 +1,123 @@
+/*
+ * 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.common.cost;
+
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Comparator;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.OptionalDouble;
+import java.util.function.Function;
+import java.util.stream.Stream;
+import org.astraea.common.Configuration;
+import org.astraea.common.admin.ClusterInfo;
+import org.astraea.common.metrics.BeanObject;
+import org.astraea.common.metrics.ClusterBean;
+import org.astraea.common.metrics.HasBeanObject;
+import org.astraea.common.metrics.Sensor;
+import org.astraea.common.metrics.broker.HasMeter;
+import org.astraea.common.metrics.broker.HasRate;
+import org.astraea.common.metrics.broker.ServerMetrics;
+import org.astraea.common.metrics.collector.MetricSensor;
+import org.astraea.common.metrics.stats.Max;
+
+/** MoveCost: more max write rate change -> higher migrate cost. */
+public class PartitionMigrateTimeCost implements HasMoveCost {
+  private static final String REPLICATION_IN_RATE = "replication_in_rate";
+  private static final String REPLICATION_OUT_RATE = "replication_out_rate";
+  public static final String MAX_MIGRATE_TIME_KEY = "max.migrated.time.limit";
+  // metrics windows size
+  private final long maxMigrateTime;
+
+  public PartitionMigrateTimeCost(Configuration config) {
+    this.maxMigrateTime =
+        config.string(MAX_MIGRATE_TIME_KEY).map(Long::parseLong).orElse(Long.MAX_VALUE);
+  }
+
+  @Override
+  public Optional<MetricSensor> metricSensor() {
+    return Optional.of(
+        (client, clusterBean) -> {
+          var oldInRate =
+              brokerMaxRate(
+                  client.identity(),
+                  clusterBean,
+                  PartitionMigrateTimeCost.MaxReplicationInRateBean.class);
+          var oldOutRate =
+              brokerMaxRate(
+                  client.identity(),
+                  clusterBean,
+                  PartitionMigrateTimeCost.MaxReplicationOutRateBean.class);
+          var newInMetrics = ServerMetrics.BrokerTopic.REPLICATION_BYTES_IN_PER_SEC.fetch(client);
+          var newOutMetrics = ServerMetrics.BrokerTopic.REPLICATION_BYTES_OUT_PER_SEC.fetch(client);
+          var current = Duration.ofMillis(System.currentTimeMillis());
+          var maxInRateSensor =
+              Sensor.builder().addStat(REPLICATION_IN_RATE, Max.<Double>of()).build();
+          var maxOutRateSensor =
+              Sensor.builder().addStat(REPLICATION_OUT_RATE, Max.<Double>of()).build();
+          maxInRateSensor.record(newInMetrics.oneMinuteRate());
+          maxOutRateSensor.record(newOutMetrics.oneMinuteRate());
+          var inRate = maxInRateSensor.measure(REPLICATION_IN_RATE);
+          var outRate = maxOutRateSensor.measure(REPLICATION_OUT_RATE);
+
+          var metrics = new ArrayList<HasBeanObject>();
+          metrics.add(
+              (MaxReplicationInRateBean)
+                  () ->
+                      new BeanObject(
+                          newInMetrics.beanObject().domainName(),
+                          newInMetrics.beanObject().properties(),
+                          Map.of(HasRate.ONE_MIN_RATE_KEY, Math.max(oldInRate.orElse(0), inRate)),
+                          current.toMillis()));
+          metrics.add(
+              (MaxReplicationOutRateBean)
+                  () ->
+                      new BeanObject(
+                          newOutMetrics.beanObject().domainName(),
+                          newOutMetrics.beanObject().properties(),
+                          Map.of(HasRate.ONE_MIN_RATE_KEY, Math.max(oldOutRate.orElse(0), outRate)),
+                          current.toMillis()));
+          return metrics;
+        });
+  }
+
+  public static OptionalDouble brokerMaxRate(
+      int identity, ClusterBean clusterBean, Class<? extends HasBeanObject> statisticMetrics) {
+    return clusterBean.all().getOrDefault(identity, List.of()).stream()
+        .filter(x -> statisticMetrics.isAssignableFrom(x.getClass()))
+        .mapToDouble(x -> ((HasMeter) x).oneMinuteRate())
+        .max();
+  }
+
+  @Override
+  public MoveCost moveCost(ClusterInfo before, ClusterInfo after, ClusterBean clusterBean) {
+    var brokerMigrateTime = MigrationCost.brokerMigrationTime(before, after, clusterBean);
+    var maxMigrateTime =
+        Stream.concat(before.nodes().stream(), after.nodes().stream())
+            .distinct()
+            .map(nodeInfo -> brokerMigrateTime.get(nodeInfo.id()))
+            .max(Comparator.comparing(Function.identity()))
+            .orElse(Long.MAX_VALUE);
+    return () -> maxMigrateTime > this.maxMigrateTime;
+  }
+
+  public interface MaxReplicationInRateBean extends HasMeter {}
+
+  public interface MaxReplicationOutRateBean extends HasMeter {}
+}

common/src/main/java/org/astraea/common/metrics/broker/HasRate.java

@@ -20,6 +20,10 @@
 import org.astraea.common.metrics.HasBeanObject;
 
 public interface HasRate extends HasBeanObject {
+  String ONE_MIN_RATE_KEY = "OneMinuteRate";
+  String FIVE_MINUTE_RATE = "FiveMinuteRate";
+  String FIFTEEN_MINUTE_RATE = "FifteenMinuteRate";
+
   default double meanRate() {
     return (double) beanObject().attributes().get("MeanRate");
   }

common/src/test/java/org/astraea/common/balancer/BalancerConsoleTest.java

@@ -402,7 +402,12 @@ public Optional<Plan> offer(AlgorithmConfig config) {
                   Duration.ofMillis(sampleTimeMs - (System.currentTimeMillis() - startMs)));
             return Optional.of(
                 new Plan(
-                    config.clusterInfo(), () -> 0, config.clusterInfo(), () -> 0, MoveCost.EMPTY));
+                    config.clusterBean(),
+                    config.clusterInfo(),
+                    () -> 0,
+                    config.clusterInfo(),
+                    () -> 0,
+                    MoveCost.EMPTY));
           }
         };