Repository: skiptests/astraea
Pull request #329 (2022-05-12 10:46:12)

Implementation of `RebalancePlanGenerator`
這個 PR 實作產生搬移計劃提案的程式碼


Patch:
diff --git a/app/src/main/java/org/astraea/Utils.java b/app/src/main/java/org/astraea/Utils.java
index 5fa68d89..4fc5f55d 100644
--- a/app/src/main/java/org/astraea/Utils.java
+++ b/app/src/main/java/org/astraea/Utils.java
@@ -149,12 +149,15 @@ public final class Utils {
     return value;
   }
 
-  public static boolean overSecond(long lastTime, int second) {
-    return (lastTime + Duration.ofSeconds(second).toMillis()) < System.currentTimeMillis();
-  }
-
-  public static boolean overSecond(long lastTime, Duration second) {
-    return (lastTime + second.toMillis()) < System.currentTimeMillis();
+  /**
+   * Check if the time is expired.
+   *
+   * @param lastTime Check time.
+   * @param interval Interval.
+   * @return Is expired.
+   */
+  public static boolean isExpired(long lastTime, Duration interval) {
+    return (lastTime + interval.toMillis()) < System.currentTimeMillis();
   }
 
   public static void sleep(Duration duration) {
diff --git a/app/src/main/java/org/astraea/balancer/RebalancePlanProposal.java b/app/src/main/java/org/astraea/balancer/RebalancePlanProposal.java
new file mode 100644
index 00000000..a411a16a
--- /dev/null
+++ b/app/src/main/java/org/astraea/balancer/RebalancePlanProposal.java
@@ -0,0 +1,111 @@
+package org.astraea.balancer;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Objects;
+import java.util.Optional;
+import org.astraea.balancer.log.ClusterLogAllocation;
+
+public interface RebalancePlanProposal {
+
+  Optional<ClusterLogAllocation> rebalancePlan();
+
+  List<String> info();
+
+  List<String> warnings();
+
+  static Build builder() {
+    return new Build();
+  }
+
+  class Build {
+    ClusterLogAllocation allocation = null;
+    List<String> info = Collections.synchronizedList(new ArrayList<>());
+    List<String> warnings = Collections.synchronizedList(new ArrayList<>());
+
+    // guard by this
+    private boolean built = false;
+
+    public synchronized Build noRebalancePlan() {
+      ensureNotBuiltYet();
+      this.allocation = null;
+      return this;
+    }
+
+    public synchronized Build withRebalancePlan(ClusterLogAllocation clusterLogAllocation) {
+      ensureNotBuiltYet();
+      this.allocation = Objects.requireNonNull(clusterLogAllocation);
+      return this;
+    }
+
+    public synchronized Build addWarning(String warning) {
+      ensureNotBuiltYet();
+      this.warnings.add(warning);
+      return this;
+    }
+
+    public synchronized Build addInfo(String info) {
+      ensureNotBuiltYet();
+      this.info.add(info);
+      return this;
+    }
+
+    private synchronized void ensureNotBuiltYet() {
+      if (built) throw new IllegalStateException("This builder already built.");
+    }
+
+    public synchronized RebalancePlanProposal build() {
+      final var allocationRef = allocation;
+      final var infoRef = info;
+      final var warningRef = warnings;
+
+      ensureNotBuiltYet();
+
+      built = true;
+      allocation = null;
+      info = null;
+      warnings = null;
+
+      return new RebalancePlanProposal() {
+
+        @Override
+        public Optional<ClusterLogAllocation> rebalancePlan() {
+          return Optional.ofNullable(allocationRef);
+        }
+
+        @Override
+        public List<String> info() {
+          // use Collections.unmodifiableList instead of List.copyOf to avoid excessive memory
+          // footprint
+          return Collections.unmodifiableList(infoRef);
+        }
+
+        @Override
+        public List<String> warnings() {
+          // use Collections.unmodifiableList instead of List.copyOf to avoid excessive memory
+          // footprint
+          return Collections.unmodifiableList(warningRef);
+        }
+
+        @Override
+        public String toString() {
+          StringBuilder sb = new StringBuilder();
+
+          sb.append("[RebalancePlanProposal]").append(System.lineSeparator());
+
+          sb.append("  Info:").append(System.lineSeparator());
+          if (info().isEmpty()) sb.append(String.format("    no information%n"));
+          else info().forEach(infoString -> sb.append(String.format("    * %s%n", infoString)));
+          if (!warnings().isEmpty()) {
+            sb.append("  Warning:").append(System.lineSeparator());
+            warnings()
+                .forEach(warningString -> sb.append(String.format("    * %s%n", warningString)));
+          }
+
+          return sb.toString();
+        }
+      };
+    }
+  }
+}
diff --git a/app/src/main/java/org/astraea/balancer/generator/RebalancePlanGenerator.java b/app/src/main/java/org/astraea/balancer/generator/RebalancePlanGenerator.java
new file mode 100644
index 00000000..40dfa1e4
--- /dev/null
+++ b/app/src/main/java/org/astraea/balancer/generator/RebalancePlanGenerator.java
@@ -0,0 +1,37 @@
+package org.astraea.balancer.generator;
+
+import java.util.stream.Stream;
+import org.astraea.balancer.RebalancePlanProposal;
+import org.astraea.balancer.log.ClusterLogAllocation;
+import org.astraea.balancer.log.LayeredClusterLogAllocation;
+import org.astraea.cost.ClusterInfo;
+
+/** */
+public interface RebalancePlanGenerator {
+
+  /**
+   * Generate a rebalance proposal, noted that this function doesn't require proposing exactly the
+   * same plan for the same input argument. There can be some randomization that takes part in this
+   * process.
+   *
+   * @param clusterInfo the cluster state, implementation can take advantage of the data inside to
+   *     proposal the plan it feels confident to improve the cluster.
+   * @return a {@link Stream} generating rebalance plan regarding the given {@link ClusterInfo}
+   */
+  default Stream<RebalancePlanProposal> generate(ClusterInfo clusterInfo) {
+    return generate(clusterInfo, LayeredClusterLogAllocation.of(clusterInfo));
+  }
+
+  /**
+   * Generate a rebalance proposal, noted that this function doesn't require proposing exactly the
+   * same plan for the same input argument. There can be some randomization that takes part in this
+   * process.
+   *
+   * @param clusterInfo the cluster state, implementation can take advantage of the data inside to
+   *     proposal the plan it feels confident to improve the cluster.
+   * @param baseAllocation the cluster log allocation as the based of proposal generation.
+   * @return a {@link Stream} generating rebalance plan regarding the given {@link ClusterInfo}
+   */
+  Stream<RebalancePlanProposal> generate(
+      ClusterInfo clusterInfo, ClusterLogAllocation baseAllocation);
+}
diff --git a/app/src/main/java/org/astraea/balancer/generator/ShufflePlanGenerator.java b/app/src/main/java/org/astraea/balancer/generator/ShufflePlanGenerator.java
new file mode 100644
index 00000000..d1c55bb8
--- /dev/null
+++ b/app/src/main/java/org/astraea/balancer/generator/ShufflePlanGenerator.java
@@ -0,0 +1,227 @@
+package org.astraea.balancer.generator;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ThreadLocalRandom;
+import java.util.function.Consumer;
+import java.util.function.Supplier;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+import java.util.stream.Stream;
+import org.astraea.admin.TopicPartition;
+import org.astraea.balancer.RebalancePlanProposal;
+import org.astraea.balancer.log.ClusterLogAllocation;
+import org.astraea.balancer.log.LayeredClusterLogAllocation;
+import org.astraea.balancer.log.LogPlacement;
+import org.astraea.cost.ClusterInfo;
+import org.astraea.cost.NodeInfo;
+
+/**
+ * The {@link ShufflePlanGenerator} proposes a new log placement based on the current log placement,
+ * but with a few random placement changes. <br>
+ * <br>
+ * The following operations are considered as a valid shuffle action:
+ *
+ * <ol>
+ *   <li>Remove a replica from the replica set, then add another broker(must not be part of the
+ *       replica set before this action) into the replica set. Noted that this operation doesn't
+ *       specify which data directory the moving replica will eventually be on the destination
+ *       broker. It's likely that the replica will follow the default placement scheme, see the
+ *       LogManager#nextLogDirs method implementation in Apache Kafka server for more details.
+ *   <li>Change the leader/follower of a partition by a member of this replica set, the original
+ *       leader/follower becomes a follower/leader.
+ *   <li>Change the data directory of one replica from current data directory to another on the
+ *       current broker.
+ * </ol>
+ */
+public class ShufflePlanGenerator implements RebalancePlanGenerator {
+
+  private final Supplier<Integer> numberOfShuffle;
+
+  public ShufflePlanGenerator(int origin, int bound) {
+    this(() -> ThreadLocalRandom.current().nextInt(origin, bound));
+  }
+
+  public ShufflePlanGenerator(Supplier<Integer> numberOfShuffle) {
+    this.numberOfShuffle = numberOfShuffle;
+  }
+
+  private int sourceTopicPartitionSelector(List<TopicPartition> migrationCandidates) {
+    return ThreadLocalRandom.current().nextInt(0, migrationCandidates.size());
+  }
+
+  private int sourceLogPlacementSelector(List<LogPlacement> migrationCandidates) {
+    return ThreadLocalRandom.current().nextInt(0, migrationCandidates.size());
+  }
+
+  private int migrationSelector(List<Movement> movementCandidates) {
+    // There are two kinds of Movement, each operating on different kind of candidate.
+    // * The Movement#replicaSetMovement selecting target by the broker ID
+    // * The Movement#dataDirectoryMovement selecting target by the data dir on the specific broker
+    // If we mix all these movements and randomly picking one as the selected migration, doing so
+    // will cause fairness issue. Given a concrete example, there are 100 brokers, and each broker
+    // have 3 data directories. Now the leader log of a topic/partition with 3 replicas is being
+    // selected. How many possible migrations over here? The Answer is:
+    // 1. 99 from Movement#replicaSetMigration since the leader log can move to any other broker.
+    // 2. 2 from Movement#dataDirectoryMigration since the leader log can move to the other 2 dirs.
+    // So there are 101 possible migrations. But these "two ways" of migration have different
+    // probability to occur. We can expect a large number of replica set changes occur but only a
+    // few for the data directory migration. This will make the ShufflePlanGenerator hard to address
+    // specific kind of balance issue due to <strong>it rarely propose them</strong>.
+    final var migrationTypeList =
+        IntStream.range(0, movementCandidates.size())
+            .mapToObj(i -> Map.entry(movementCandidates.get(i).getClass(), i))
+            .collect(
+                Collectors.groupingBy(
+                    Map.Entry::getKey,
+                    Collectors.mapping(Map.Entry::getValue, Collectors.toUnmodifiableList())))
+            .values()
+            .stream()
+            .collect(Collectors.toUnmodifiableList());
+    final var selectedMigration =
+        migrationTypeList.get(ThreadLocalRandom.current().nextInt(0, migrationTypeList.size()));
+    return selectedMigration.get(ThreadLocalRandom.current().nextInt(0, selectedMigration.size()));
+  }
+
+  @Override
+  public Stream<RebalancePlanProposal> generate(
+      ClusterInfo clusterInfo, ClusterLogAllocation baseAllocation) {
+    return Stream.generate(
+        () -> {
+          final var rebalancePlanBuilder = RebalancePlanProposal.builder();
+          final var brokerIds =
+              clusterInfo.nodes().stream()
+                  .map(NodeInfo::id)
+                  .collect(Collectors.toUnmodifiableSet());
+
+          if (brokerIds.size() == 0)
+            return rebalancePlanBuilder
+                .addWarning("Why there is no broker?")
+                .noRebalancePlan()
+                .build();
+
+          if (brokerIds.size() == 1)
+            return rebalancePlanBuilder
+                .addWarning("Only one broker exists. There is no reason to rebalance.")
+                .noRebalancePlan()
+                .build();
+
+          if (clusterInfo.topics().size() == 0)
+            return rebalancePlanBuilder
+                .addWarning("No non-ignored topic to working on.")
+                .noRebalancePlan()
+                .build();
+
+          final var shuffleCount = numberOfShuffle.get();
+          final var newAllocation = LayeredClusterLogAllocation.of(baseAllocation);
+          final var pickingList =
+              newAllocation.topicPartitionStream().collect(Collectors.toUnmodifiableList());
+
+          rebalancePlanBuilder.addInfo(
+              "Make " + shuffleCount + (shuffleCount > 0 ? " shuffles." : " shuffle."));
+          for (int i = 0; i < shuffleCount; i++) {
+            final var sourceTopicPartitionIndex = sourceTopicPartitionSelector(pickingList);
+            final var sourceTopicPartition = pickingList.get(sourceTopicPartitionIndex);
+            final var sourceLogPlacements = newAllocation.logPlacements(sourceTopicPartition);
+            final var sourceLogPlacementIndex = sourceLogPlacementSelector(sourceLogPlacements);
+            final var sourceLogPlacement = sourceLogPlacements.get(sourceLogPlacementIndex);
+            final var sourceIsLeader = sourceLogPlacementIndex == 0;
+            final var sourceBroker = sourceLogPlacement.broker();
+
+            Consumer<Integer> replicaSetMigration =
+                (targetBroker) -> {
+                  newAllocation.migrateReplica(sourceTopicPartition, sourceBroker, targetBroker);
+                  rebalancePlanBuilder.addInfo(
+                      String.format(
+                          "Change replica set of topic %s partition %d, from %d to %d.",
+                          sourceTopicPartition.topic(),
+                          sourceTopicPartition.partition(),
+                          sourceLogPlacement.broker(),
+                          targetBroker));
+                };
+            Consumer<LogPlacement> leaderFollowerMigration =
+                (newLeaderReplicaCandidate) -> {
+                  newAllocation.letReplicaBecomeLeader(
+                      sourceTopicPartition, newLeaderReplicaCandidate.broker());
+                  rebalancePlanBuilder.addInfo(
+                      String.format(
+                          "Change the log identity of topic %s partition %d replica at broker %d, from %s to %s",
+                          sourceTopicPartition.topic(),
+                          sourceTopicPartition.partition(),
+                          sourceLogPlacement.broker(),
+                          sourceIsLeader ? "leader" : "follower",
+                          sourceIsLeader ? "follower" : "leader"));
+                };
+            Consumer<String> dataDirectoryMigration =
+                (dataDirectory) -> {
+                  newAllocation.changeDataDirectory(
+                      sourceTopicPartition, sourceBroker, dataDirectory);
+                  rebalancePlanBuilder.addInfo(
+                      String.format(
+                          "Change the data directory of topic %s partition %d replica at broker %d, from %s to %s",
+                          sourceTopicPartition.topic(),
+                          sourceTopicPartition.partition(),
+                          sourceLogPlacement.broker(),
+                          sourceLogPlacement.logDirectory().map(Object::toString).orElse("unknown"),
+                          dataDirectory));
+                };
+
+            // generate a set of valid migration broker for given placement.
+            final var validMigrationCandidates = new ArrayList<Movement>();
+            // [Valid movement 1] add all brokers and remove all broker in current replica set
+            brokerIds.stream()
+                .filter(
+                    broker -> sourceLogPlacements.stream().noneMatch(log -> log.broker() == broker))
+                .map(targetBroker -> (Runnable) () -> replicaSetMigration.accept(targetBroker))
+                .map(Movement::replicaSetMovement)
+                .forEach(validMigrationCandidates::add);
+            // [Valid movement 2] add all leader/follower change candidate
+            if (sourceLogPlacements.size() > 1) {
+              if (sourceIsLeader) {
+                // leader can migrate its identity to follower.
+                sourceLogPlacements.stream()
+                    .skip(1)
+                    .map(
+                        followerReplica ->
+                            (Runnable) () -> leaderFollowerMigration.accept(followerReplica))
+                    .map(Movement::replicaSetMovement)
+                    .forEach(validMigrationCandidates::add);
+              } else {
+                // follower can migrate its identity to leader.
+                validMigrationCandidates.add(
+                    Movement.replicaSetMovement(
+                        () -> leaderFollowerMigration.accept(sourceLogPlacement)));
+              }
+            }
+            // [Valid movement 3] change the data directory of selected replica
+            clusterInfo.dataDirectories(sourceLogPlacement.broker()).stream()
+                .filter(dir -> !dir.equals(sourceLogPlacement.logDirectory().orElse(null)))
+                .map(dir -> (Runnable) () -> dataDirectoryMigration.accept(dir))
+                .map(Movement::dataDirectoryMovement)
+                .forEach(validMigrationCandidates::add);
+
+            // pick a migration and execute
+            final var selectedMigrationIndex = migrationSelector(validMigrationCandidates);
+            validMigrationCandidates.get(selectedMigrationIndex).run();
+          }
+
+          return rebalancePlanBuilder.withRebalancePlan(newAllocation).build();
+        });
+  }
+
+  /** Action to trigger upon this migration */
+  interface Movement extends Runnable {
+    static ReplicaSetMovement replicaSetMovement(Runnable runnable) {
+      return runnable::run;
+    }
+
+    static DataDirectoryMovement dataDirectoryMovement(Runnable runnable) {
+      return runnable::run;
+    }
+  }
+
+  interface ReplicaSetMovement extends Movement {}
+
+  interface DataDirectoryMovement extends Movement {}
+}
diff --git a/app/src/main/java/org/astraea/balancer/log/ClusterLogAllocation.java b/app/src/main/java/org/astraea/balancer/log/ClusterLogAllocation.java
new file mode 100644
index 00000000..a0fe097b
--- /dev/null
+++ b/app/src/main/java/org/astraea/balancer/log/ClusterLogAllocation.java
@@ -0,0 +1,30 @@
+package org.astraea.balancer.log;
+
+import java.util.List;
+import java.util.stream.Stream;
+import org.astraea.admin.TopicPartition;
+
+/**
+ * Describe the log allocation state of a Kafka cluster. The implementation have to keep the cluster
+ * log allocation information, provide method for query the placement, and offer a set of log
+ * placement change operation.
+ */
+public interface ClusterLogAllocation {
+
+  /** let specific broker leave the replica set and let another broker join the replica set. */
+  void migrateReplica(TopicPartition topicPartition, int atBroker, int toBroker);
+
+  /** let specific follower log become the leader log of this topic/partition. */
+  void letReplicaBecomeLeader(TopicPartition topicPartition, int followerReplica);
+
+  /** change the data directory of specific log */
+  void changeDataDirectory(TopicPartition topicPartition, int atBroker, String newPath);
+
+  /** Retrieve the log placements of specific {@link TopicPartition}. */
+  List<LogPlacement> logPlacements(TopicPartition topicPartition);
+
+  /** Retrieve the stream of all topic/partition pairs in allocation. */
+  Stream<TopicPartition> topicPartitionStream();
+
+  // TODO: add a method to calculate the difference between two ClusterLogAllocation
+}
diff --git a/app/src/main/java/org/astraea/balancer/log/IllegalMigrationException.java b/app/src/main/java/org/astraea/balancer/log/IllegalMigrationException.java
new file mode 100644
index 00000000..bf3f26e8
--- /dev/null
+++ b/app/src/main/java/org/astraea/balancer/log/IllegalMigrationException.java
@@ -0,0 +1,7 @@
+package org.astraea.balancer.log;
+
+public class IllegalMigrationException extends RuntimeException {
+  public IllegalMigrationException(String s) {
+    super(s);
+  }
+}
diff --git a/app/src/main/java/org/astraea/balancer/log/LayeredClusterLogAllocation.java b/app/src/main/java/org/astraea/balancer/log/LayeredClusterLogAllocation.java
new file mode 100644
index 00000000..3588912c
--- /dev/null
+++ b/app/src/main/java/org/astraea/balancer/log/LayeredClusterLogAllocation.java
@@ -0,0 +1,263 @@
+package org.astraea.balancer.log;
+
+import java.util.Collection;
+import java.util.Comparator;
+import java.util.List;
+import java.util.Map;
+import java.util.OptionalInt;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+import java.util.stream.Stream;
+import org.astraea.admin.TopicPartition;
+import org.astraea.cost.ClusterInfo;
+import org.astraea.cost.ReplicaInfo;
+
+public class LayeredClusterLogAllocation implements ClusterLogAllocation {
+
+  // set operation is guard by this
+  private final AtomicBoolean isLocked = new AtomicBoolean(false);
+  private final LayeredClusterLogAllocation upperLayer;
+
+  // guard by this
+  private final Map<TopicPartition, List<LogPlacement>> allocation;
+
+  private LayeredClusterLogAllocation(
+      LayeredClusterLogAllocation upperLayer, Map<TopicPartition, List<LogPlacement>> allocation) {
+    allocation.keySet().stream()
+        .collect(Collectors.groupingBy(TopicPartition::topic))
+        .forEach(
+            (topic, tp) -> {
+              int maxPartitionId =
+                  tp.stream().mapToInt(TopicPartition::partition).max().orElseThrow();
+              if ((maxPartitionId + 1) != tp.size())
+                throw new IllegalArgumentException(
+                    "The partition size of " + topic + " is illegal");
+            });
+    allocation.forEach(
+        (tp, logs) -> {
+          long uniqueBrokers = logs.stream().map(LogPlacement::broker).distinct().count();
+          if (uniqueBrokers != logs.size() || logs.size() == 0)
+            throw new IllegalArgumentException(
+                "The topic "
+                    + tp.topic()
+                    + " partition "
+                    + tp.partition()
+                    + " has illegal replica set "
+                    + logs);
+        });
+    this.upperLayer = upperLayer;
+    this.allocation = allocation;
+  }
+
+  public static LayeredClusterLogAllocation of(ClusterInfo clusterInfo) {
+    final Map<TopicPartition, List<LogPlacement>> allocation =
+        clusterInfo.topics().stream()
+            .map(clusterInfo::partitions)
+            .flatMap(Collection::stream)
+            .collect(
+                Collectors.groupingBy(
+                    replica ->
+                        TopicPartition.of(replica.topic(), Integer.toString(replica.partition()))))
+            .entrySet()
+            .stream()
+            .map(
+                (entry) -> {
+                  // validate if the given log placements are valid
+                  if (entry.getValue().stream().filter(ReplicaInfo::isLeader).count() != 1)
+                    throw new IllegalArgumentException(
+                        "The " + entry.getKey() + " leader count mismatch 1.");
+
+                  final var topicPartition = entry.getKey();
+                  final var logPlacements =
+                      entry.getValue().stream()
+                          .sorted(Comparator.comparingInt(replica -> replica.isLeader() ? 0 : 1))
+                          .map(
+                              replica ->
+                                  LogPlacement.of(
+                                      replica.nodeInfo().id(), replica.dataFolder().orElse(null)))
+                          .collect(Collectors.toList());
+
+                  return Map.entry(topicPartition, logPlacements);
+                })
+            .collect(Collectors.toConcurrentMap(Map.Entry::getKey, Map.Entry::getValue));
+    return new LayeredClusterLogAllocation(null, allocation);
+  }
+
+  public static LayeredClusterLogAllocation of(
+      Map<TopicPartition, List<LogPlacement>> allocationMap) {
+    return new LayeredClusterLogAllocation(null, allocationMap);
+  }
+
+  /**
+   * Create a new {@link LayeredClusterLogAllocation}. If the given {@link ClusterLogAllocation} is
+   * an instance of {@link LayeredClusterLogAllocation}. Then the given {@link
+   * LayeredClusterLogAllocation} will be locked and become unmodifiable. The new log allocation
+   * will continue from the unmodifiable allocation. Doing so can improve the performance of new log
+   * allocation creation as we avoid the heavy data copy work.
+   */
+  public static LayeredClusterLogAllocation of(ClusterLogAllocation clusterLogAllocation) {
+    if (clusterLogAllocation instanceof LayeredClusterLogAllocation) {
+      return createNewLayer((LayeredClusterLogAllocation) clusterLogAllocation);
+    } else {
+      // no layer support, do the expensive copy work.
+      final var map = new ConcurrentHashMap<TopicPartition, List<LogPlacement>>();
+      clusterLogAllocation
+          .topicPartitionStream()
+          .forEach(
+              topicPartition ->
+                  map.put(topicPartition, clusterLogAllocation.logPlacements(topicPartition)));
+      return of(map);
+    }
+  }
+
+  /**
+   * Create a new {@link LayeredClusterLogAllocation} based on the given one. The given layered
+   * cluster log will become unmodifiable afterward.
+   */
+  private static LayeredClusterLogAllocation createNewLayer(
+      LayeredClusterLogAllocation baseLogAllocation) {
+    if (!baseLogAllocation.isLocked()) {
+      baseLogAllocation.lockLayer();
+    }
+    return new LayeredClusterLogAllocation(baseLogAllocation, new ConcurrentHashMap<>());
+  }
+
+  /** Make this {@link LayeredClusterLogAllocation} unmodifiable. */
+  private synchronized void lockLayer() {
+    isLocked.set(true);
+  }
+
+  /** Is this {@link LayeredClusterLogAllocation} unmodifiable. */
+  private boolean isLocked() {
+    return isLocked.get();
+  }
+
+  private synchronized void ensureNotLocked() {
+    if (isLocked())
+      throw new IllegalStateException("This layered cluster log allocation has been locked");
+  }
+
+  private static OptionalInt indexOfBroker(List<LogPlacement> logPlacements, int targetBroker) {
+    return IntStream.range(0, logPlacements.size())
+        .filter(index -> logPlacements.get(index).broker() == targetBroker)
+        .findFirst();
+  }
+
+  @Override
+  public synchronized void migrateReplica(
+      TopicPartition topicPartition, int broker, int destinationBroker) {
+    ensureNotLocked();
+
+    final List<LogPlacement> sourceLogPlacements = this.logPlacements(topicPartition);
+    if (sourceLogPlacements == null)
+      throw new IllegalMigrationException(
+          topicPartition.topic() + "-" + topicPartition.partition() + " no such topic/partition");
+
+    int sourceLogIndex = indexOfBroker(sourceLogPlacements, broker).orElse(-1);
+    if (sourceLogIndex == -1)
+      throw new IllegalMigrationException(
+          broker + " is not part of the replica set for " + topicPartition);
+
+    int destinationLogIndex = indexOfBroker(sourceLogPlacements, destinationBroker).orElse(-1);
+    if (destinationLogIndex != -1)
+      throw new IllegalArgumentException(
+          destinationBroker + " is already part of the replica set, no need to move");
+
+    this.allocation.put(
+        topicPartition,
+        IntStream.range(0, sourceLogPlacements.size())
+            .mapToObj(
+                index ->
+                    index == sourceLogIndex
+                        ? LogPlacement.of(destinationBroker)
+                        : sourceLogPlacements.get(index))
+            .collect(Collectors.toUnmodifiableList()));
+  }
+
+  @Override
+  public synchronized void letReplicaBecomeLeader(
+      TopicPartition topicPartition, int followerReplica) {
+    ensureNotLocked();
+
+    final List<LogPlacement> sourceLogPlacements = this.logPlacements(topicPartition);
+    if (sourceLogPlacements == null)
+      throw new IllegalMigrationException(
+          topicPartition.topic() + "-" + topicPartition.partition() + " no such topic/partition");
+
+    int leaderLogIndex = 0;
+    if (sourceLogPlacements.size() == 0)
+      throw new IllegalStateException("This partition has no log");
+
+    int followerLogIndex = indexOfBroker(sourceLogPlacements, followerReplica).orElse(-1);
+    if (followerLogIndex == -1)
+      throw new IllegalArgumentException(
+          followerReplica + " is not part of the replica set for " + topicPartition);
+
+    if (leaderLogIndex == followerLogIndex) return; // nothing to do
+
+    final var leaderLog = this.logPlacements(topicPartition).get(leaderLogIndex);
+    final var followerLog = this.logPlacements(topicPartition).get(followerLogIndex);
+
+    this.allocation.put(
+        topicPartition,
+        IntStream.range(0, sourceLogPlacements.size())
+            .mapToObj(
+                index ->
+                    index == leaderLogIndex
+                        ? followerLog
+                        : index == followerLogIndex ? leaderLog : sourceLogPlacements.get(index))
+            .collect(Collectors.toUnmodifiableList()));
+  }
+
+  @Override
+  public synchronized void changeDataDirectory(
+      TopicPartition topicPartition, int broker, String path) {
+    ensureNotLocked();
+
+    final List<LogPlacement> sourceLogPlacements = this.logPlacements(topicPartition);
+    if (sourceLogPlacements == null)
+      throw new IllegalMigrationException(
+          topicPartition.topic() + "-" + topicPartition.partition() + " no such topic/partition");
+
+    int sourceLogIndex = indexOfBroker(sourceLogPlacements, broker).orElse(-1);
+    if (sourceLogIndex == -1)
+      throw new IllegalMigrationException(
+          broker + " is not part of the replica set for " + topicPartition);
+
+    final var oldLog = this.logPlacements(topicPartition).get(sourceLogIndex);
+    final var newLog = LogPlacement.of(oldLog.broker(), path);
+    this.allocation.put(
+        topicPartition,
+        IntStream.range(0, sourceLogPlacements.size())
+            .mapToObj(index -> index == sourceLogIndex ? newLog : sourceLogPlacements.get(index))
+            .collect(Collectors.toUnmodifiableList()));
+  }
+
+  @Override
+  public List<LogPlacement> logPlacements(TopicPartition topicPartition) {
+    if (allocation.containsKey(topicPartition)) return allocation.get(topicPartition);
+    else if (upperLayer != null) return upperLayer.logPlacements(topicPartition);
+    else return null;
+  }
+
+  @Override
+  public Stream<TopicPartition> topicPartitionStream() {
+    if (upperLayer == null) return this.allocation.keySet().stream();
+    else {
+      LayeredClusterLogAllocation baseAllocationLookUp = this;
+      while (baseAllocationLookUp.upperLayer != null)
+        baseAllocationLookUp = baseAllocationLookUp.upperLayer;
+      // Ignore all the intermediate layers, just return the keys set of the base layer.
+      // THIS WILL WORK, AS LONG AS THE UPPER LAYERS DO NOT DO THE FOLLOWING.
+      // 1. delete topic (should balancer do this?)
+      // 2. create topic (should balancer do this?)
+      // 3. shrink partition size (impossible for the current Kafka)
+      // 4. expand partition size (should balancer do this?)
+      // Of course, we can consider every element in all the layers, but doing that hurt
+      // performance.
+      return baseAllocationLookUp.allocation.keySet().stream();
+    }
+  }
+}
diff --git a/app/src/main/java/org/astraea/balancer/log/LogPlacement.java b/app/src/main/java/org/astraea/balancer/log/LogPlacement.java
new file mode 100644
index 00000000..2f75a23b
--- /dev/null
+++ b/app/src/main/java/org/astraea/balancer/log/LogPlacement.java
@@ -0,0 +1,43 @@
+package org.astraea.balancer.log;
+
+import java.util.Optional;
+
+/** This class describe the placement state of one kafka log. */
+public interface LogPlacement {
+
+  int broker();
+
+  Optional<String> logDirectory();
+
+  static LogPlacement of(int broker) {
+    return of(broker, null);
+  }
+
+  static LogPlacement of(int broker, String logDirectory) {
+    return new LogPlacement() {
+      @Override
+      public int broker() {
+        return broker;
+      }
+
+      @Override
+      public Optional<String> logDirectory() {
+        return Optional.ofNullable(logDirectory);
+      }
+
+      @Override
+      public boolean equals(Object obj) {
+        if (obj instanceof LogPlacement) {
+          final var that = (LogPlacement) obj;
+          return this.broker() == that.broker() && this.logDirectory().equals(that.logDirectory());
+        }
+        return false;
+      }
+
+      @Override
+      public String toString() {
+        return "LogPlacement{broker=" + broker() + " logDir=" + logDirectory() + "}";
+      }
+    };
+  }
+}
diff --git a/app/src/main/java/org/astraea/cost/ClusterInfo.java b/app/src/main/java/org/astraea/cost/ClusterInfo.java
index b147ff26..691ea3ab 100644
--- a/app/src/main/java/org/astraea/cost/ClusterInfo.java
+++ b/app/src/main/java/org/astraea/cost/ClusterInfo.java
@@ -20,6 +20,12 @@ public interface ClusterInfo {
         return cluster.nodes().stream().map(NodeInfo::of).collect(Collectors.toUnmodifiableList());
       }
 
+      @Override
+      public Set<String> dataDirectories(int brokerId) {
+        // org.apache.kafka.common.Cluster doesn't have such information.
+        throw new UnsupportedOperationException("This information is not available");
+      }
+
       public Set<String> topics() {
         return cluster.topics();
       }
@@ -82,6 +88,11 @@ public interface ClusterInfo {
         return cluster.nodes();
       }
 
+      @Override
+      public Set<String> dataDirectories(int brokerId) {
+        return cluster.dataDirectories(brokerId);
+      }
+
       @Override
       public List<ReplicaInfo> availablePartitionLeaders(String topic) {
         return cluster.availablePartitionLeaders(topic);
@@ -147,6 +158,10 @@ public interface ClusterInfo {
   /** @return The known set of nodes */
   List<NodeInfo> nodes();
 
+  /** @return return the data directories on specific broker */
+  Set<String> dataDirectories(int brokerId);
+  // TODO: provide a ClusterInfo implementation with this info
+
   /**
    * Get the list of replica leader information of each available partition for the given topic
    *
diff --git a/app/src/main/java/org/astraea/cost/MemoryWarningCost.java b/app/src/main/java/org/astraea/cost/MemoryWarningCost.java
deleted file mode 100644
index 51521315..00000000
--- a/app/src/main/java/org/astraea/cost/MemoryWarningCost.java
+++ /dev/null
@@ -1,48 +0,0 @@
-package org.astraea.cost;
-
-import java.util.List;
-import java.util.Map;
-import java.util.stream.Collectors;
-import org.astraea.Utils;
-import org.astraea.metrics.collector.Fetcher;
-import org.astraea.metrics.java.HasJvmMemory;
-import org.astraea.metrics.kafka.KafkaMetrics;
-
-/** Warning on heap usage >= 80%. */
-public class MemoryWarningCost implements HasBrokerCost {
-  private final long constructTime = System.currentTimeMillis();
-
-  /** @return 1 if the heap usage >= 80%, 0 otherwise. */
-  @Override
-  public BrokerCost brokerCost(ClusterInfo clusterInfo) {
-    var brokerScore =
-        clusterInfo.allBeans().entrySet().stream()
-            .map(
-                e ->
-                    Map.entry(
-                        e.getKey(),
-                        e.getValue().stream()
-                            .filter(beanObject -> beanObject instanceof HasJvmMemory)
-                            .map(beanObject -> (HasJvmMemory) beanObject)
-                            .map(
-                                hasJvmMemory -> {
-                                  if (Utils.overSecond(constructTime, 10)
-                                      && (hasJvmMemory.heapMemoryUsage().getUsed() + 0.0)
-                                              / (hasJvmMemory.heapMemoryUsage().getMax() + 1)
-                                          >= 0.8) {
-                                    return 1.0;
-                                  } else {
-                                    return 0.0;
-                                  }
-                                })
-                            .findAny()
-                            .orElseThrow()))
-            .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
-    return () -> brokerScore;
-  }
-
-  @Override
-  public Fetcher fetcher() {
-    return client -> List.of(KafkaMetrics.Host.jvmMemory(client));
-  }
-}
diff --git a/app/src/main/java/org/astraea/cost/Periodic.java b/app/src/main/java/org/astraea/cost/Periodic.java
index 34f33b65..03814ff8 100644
--- a/app/src/main/java/org/astraea/cost/Periodic.java
+++ b/app/src/main/java/org/astraea/cost/Periodic.java
@@ -10,43 +10,31 @@ public abstract class Periodic<Value> {
   private Value value;
 
   /**
-   * Updates the value each second.
+   * Updates the value interval second.
    *
    * @param updater Methods that need to be updated regularly.
+   * @param interval Required interval.
    * @return an object of type Value created from the parameter value.
    */
-  protected Value tryUpdate(Supplier<Value> updater) {
-    if (Utils.overSecond(lastUpdate, 1)) {
+  protected Value tryUpdate(Supplier<Value> updater, Duration interval) {
+    if (Utils.isExpired(lastUpdate, interval)) {
       value = updater.get();
       lastUpdate = currentTime();
     }
     return value;
   }
 
-  protected long currentTime() {
-    return System.currentTimeMillis();
-  }
-
   /**
-   * Updates the value interval second.
+   * Updates the value each second.
    *
    * @param updater Methods that need to be updated regularly.
-   * @param interval Required interval.
    * @return an object of type Value created from the parameter value.
    */
-  protected Value tryUpdate(Supplier<Value> updater, int interval) {
-    if (Utils.overSecond(lastUpdate, interval)) {
-      value = updater.get();
-      lastUpdate = currentTime();
-    }
-    return value;
+  protected Value tryUpdateAfterOneSecond(Supplier<Value> updater) {
+    return tryUpdate(updater, Duration.ofSeconds(1));
   }
 
-  protected Value tryUpdate(Supplier<Value> updater, Duration interval) {
-    if (Utils.overSecond(lastUpdate, interval)) {
-      value = updater.get();
-      lastUpdate = currentTime();
-    }
-    return value;
+  protected long currentTime() {
+    return System.currentTimeMillis();
   }
 }
diff --git a/app/src/main/java/org/astraea/cost/broker/BrokerInputCost.java b/app/src/main/java/org/astraea/cost/broker/BrokerInputCost.java
index 6e3ef9ee..721a8ad2 100644
--- a/app/src/main/java/org/astraea/cost/broker/BrokerInputCost.java
+++ b/app/src/main/java/org/astraea/cost/broker/BrokerInputCost.java
@@ -10,7 +10,6 @@ import java.util.stream.IntStream;
 import org.astraea.cost.BrokerCost;
 import org.astraea.cost.ClusterInfo;
 import org.astraea.cost.HasBrokerCost;
-import org.astraea.cost.Periodic;
 import org.astraea.metrics.collector.Fetcher;
 import org.astraea.metrics.kafka.BrokerTopicMetricsResult;
 import org.astraea.metrics.kafka.KafkaMetrics;
@@ -26,44 +25,44 @@ import org.astraea.metrics.kafka.KafkaMetrics;
  *   <li>The final result is the average of the ten-second data.
  * </ol>
  */
-public class BrokerInputCost extends Periodic<Map<Integer, Double>> implements HasBrokerCost {
+public class BrokerInputCost implements HasBrokerCost {
   private final Map<Integer, BrokerMetric> brokersMetric = new HashMap<>();
 
   @Override
   public BrokerCost brokerCost(ClusterInfo clusterInfo) {
-    var brokerScore =
-        tryUpdate(
-            () -> {
-              var costMetrics =
-                  clusterInfo.allBeans().entrySet().stream()
-                      .collect(Collectors.toMap(Map.Entry::getKey, entry -> 0.0));
-              clusterInfo
-                  .allBeans()
-                  .forEach(
-                      (brokerID, value) -> {
-                        if (!brokersMetric.containsKey(brokerID)) {
-                          brokersMetric.put(brokerID, new BrokerMetric());
-                        }
-                        value.stream()
+    var costMetrics =
+        clusterInfo.allBeans().entrySet().stream()
+            .collect(
+                Collectors.toMap(
+                    Map.Entry::getKey,
+                    entry ->
+                        entry.getValue().stream()
                             .filter(
                                 hasBeanObject ->
                                     KafkaMetrics.BrokerTopic.BytesInPerSec.metricName()
                                         .equals(
                                             hasBeanObject.beanObject().getProperties().get("name")))
-                            .forEach(
-                                hasBeanObject -> {
-                                  var broker = brokersMetric.get(brokerID);
-                                  var inBean = (BrokerTopicMetricsResult) hasBeanObject;
-                                  costMetrics.put(
-                                      brokerID, (double) (inBean.count() - broker.accumulateCount));
-                                  broker.accumulateCount = inBean.count();
-                                });
-                      });
-              TScore(costMetrics).forEach((broker, v) -> brokersMetric.get(broker).updateLoad(v));
+                            .findAny()
+                            .orElseThrow()))
+            .entrySet()
+            .stream()
+            .collect(
+                Collectors.toMap(
+                    Map.Entry::getKey,
+                    entry -> {
+                      if (!brokersMetric.containsKey(entry.getKey())) {
+                        brokersMetric.put(entry.getKey(), new BrokerMetric());
+                      }
+                      var broker = brokersMetric.get(entry.getKey());
+                      var inBean = (BrokerTopicMetricsResult) entry.getValue();
+                      var count = (double) (inBean.count() - broker.accumulateCount);
+                      broker.accumulateCount = inBean.count();
+                      return count;
+                    }));
 
-              return computeLoad();
-            });
-    return () -> brokerScore;
+    TScore(costMetrics).forEach((broker, v) -> brokersMetric.get(broker).updateLoad(v));
+
+    return this::computeLoad;
   }
 
   Map<Integer, Double> computeLoad() {
diff --git a/app/src/main/java/org/astraea/cost/broker/BrokerOutputCost.java b/app/src/main/java/org/astraea/cost/broker/BrokerOutputCost.java
index a175cdd9..58bc6888 100644
--- a/app/src/main/java/org/astraea/cost/broker/BrokerOutputCost.java
+++ b/app/src/main/java/org/astraea/cost/broker/BrokerOutputCost.java
@@ -10,7 +10,6 @@ import java.util.stream.IntStream;
 import org.astraea.cost.BrokerCost;
 import org.astraea.cost.ClusterInfo;
 import org.astraea.cost.HasBrokerCost;
-import org.astraea.cost.Periodic;
 import org.astraea.metrics.collector.Fetcher;
 import org.astraea.metrics.kafka.BrokerTopicMetricsResult;
 import org.astraea.metrics.kafka.KafkaMetrics;
@@ -26,44 +25,44 @@ import org.astraea.metrics.kafka.KafkaMetrics;
  *   <li>The final result is the average of the ten-second data.
  * </ol>
  */
-public class BrokerOutputCost extends Periodic<Map<Integer, Double>> implements HasBrokerCost {
+public class BrokerOutputCost implements HasBrokerCost {
   private final Map<Integer, BrokerMetric> brokersMetric = new HashMap<>();
 
   @Override
   public BrokerCost brokerCost(ClusterInfo clusterInfo) {
-    var brokerScore =
-        tryUpdate(
-            () -> {
-              var costMetrics =
-                  clusterInfo.allBeans().entrySet().stream()
-                      .collect(Collectors.toMap(Map.Entry::getKey, entry -> 0.0));
-              clusterInfo
-                  .allBeans()
-                  .forEach(
-                      (brokerID, value) -> {
-                        if (!brokersMetric.containsKey(brokerID)) {
-                          brokersMetric.put(brokerID, new BrokerMetric());
-                        }
-                        value.stream()
+    var costMetrics =
+        clusterInfo.allBeans().entrySet().stream()
+            .collect(
+                Collectors.toMap(
+                    Map.Entry::getKey,
+                    entry ->
+                        entry.getValue().stream()
                             .filter(
                                 hasBeanObject ->
                                     KafkaMetrics.BrokerTopic.BytesOutPerSec.metricName()
                                         .equals(
                                             hasBeanObject.beanObject().getProperties().get("name")))
-                            .forEach(
-                                hasBeanObject -> {
-                                  var broker = brokersMetric.get(brokerID);
-                                  var outBean = (BrokerTopicMetricsResult) hasBeanObject;
-                                  costMetrics.put(
-                                      brokerID,
-                                      (double) (outBean.count() - broker.accumulateCount));
-                                  broker.accumulateCount = outBean.count();
-                                });
-                      });
-              TScore(costMetrics).forEach((broker, v) -> brokersMetric.get(broker).updateLoad(v));
-              return computeLoad();
-            });
-    return () -> brokerScore;
+                            .findAny()
+                            .orElseThrow()))
+            .entrySet()
+            .stream()
+            .collect(
+                Collectors.toMap(
+                    Map.Entry::getKey,
+                    entry -> {
+                      if (!brokersMetric.containsKey(entry.getKey())) {
+                        brokersMetric.put(entry.getKey(), new BrokerMetric());
+                      }
+                      var broker = brokersMetric.get(entry.getKey());
+                      var inBean = (BrokerTopicMetricsResult) entry.getValue();
+                      var count = (double) (inBean.count() - broker.accumulateCount);
+                      broker.accumulateCount = inBean.count();
+                      return count;
+                    }));
+
+    TScore(costMetrics).forEach((broker, v) -> brokersMetric.get(broker).updateLoad(v));
+
+    return this::computeLoad;
   }
 
   Map<Integer, Double> computeLoad() {
diff --git a/app/src/main/java/org/astraea/cost/broker/CpuCost.java b/app/src/main/java/org/astraea/cost/broker/CpuCost.java
index 7639f6f8..fa45fe8e 100644
--- a/app/src/main/java/org/astraea/cost/broker/CpuCost.java
+++ b/app/src/main/java/org/astraea/cost/broker/CpuCost.java
@@ -10,7 +10,6 @@ import java.util.stream.IntStream;
 import org.astraea.cost.BrokerCost;
 import org.astraea.cost.ClusterInfo;
 import org.astraea.cost.HasBrokerCost;
-import org.astraea.cost.Periodic;
 import org.astraea.metrics.collector.Fetcher;
 import org.astraea.metrics.java.OperatingSystemInfo;
 import org.astraea.metrics.kafka.KafkaMetrics;
@@ -26,36 +25,37 @@ import org.astraea.metrics.kafka.KafkaMetrics;
  *   <li>The final result is the average of the ten-second data.
  * </ol>
  */
-public class CpuCost extends Periodic<Map<Integer, Double>> implements HasBrokerCost {
+public class CpuCost implements HasBrokerCost {
   private final Map<Integer, BrokerMetric> brokersMetric = new HashMap<>();
 
   @Override
   public BrokerCost brokerCost(ClusterInfo clusterInfo) {
-    var brokerScore =
-        tryUpdate(
-            () -> {
-              var costMetrics =
-                  clusterInfo.allBeans().entrySet().stream()
-                      .collect(Collectors.toMap(Map.Entry::getKey, entry -> 0.0));
-              clusterInfo
-                  .allBeans()
-                  .forEach(
-                      (brokerID, value) -> {
-                        if (!brokersMetric.containsKey(brokerID)) {
-                          brokersMetric.put(brokerID, new BrokerMetric());
-                        }
-                        value.stream()
-                            .filter(beanObject -> beanObject instanceof OperatingSystemInfo)
-                            .forEach(
-                                hasBeanObject -> {
-                                  var cpuBean = (OperatingSystemInfo) hasBeanObject;
-                                  costMetrics.put(brokerID, cpuBean.systemCpuLoad());
-                                });
-                      });
-              TScore(costMetrics).forEach((broker, v) -> brokersMetric.get(broker).updateLoad(v));
-              return computeLoad();
-            });
-    return () -> brokerScore;
+    var costMetrics =
+        clusterInfo.allBeans().entrySet().stream()
+            .collect(
+                Collectors.toMap(
+                    Map.Entry::getKey,
+                    entry ->
+                        entry.getValue().stream()
+                            .filter(hasBeanObject -> hasBeanObject instanceof OperatingSystemInfo)
+                            .findAny()
+                            .orElseThrow()))
+            .entrySet()
+            .stream()
+            .collect(
+                Collectors.toMap(
+                    Map.Entry::getKey,
+                    entry -> {
+                      if (!brokersMetric.containsKey(entry.getKey())) {
+                        brokersMetric.put(entry.getKey(), new BrokerMetric());
+                      }
+                      var cpuBean = (OperatingSystemInfo) entry.getValue();
+                      return cpuBean.systemCpuLoad();
+                    }));
+
+    TScore(costMetrics).forEach((broker, v) -> brokersMetric.get(broker).updateLoad(v));
+
+    return this::computeLoad;
   }
 
   Map<Integer, Double> computeLoad() {
diff --git a/app/src/main/java/org/astraea/cost/broker/MemoryCost.java b/app/src/main/java/org/astraea/cost/broker/MemoryCost.java
index 8d18ae49..91b05219 100644
--- a/app/src/main/java/org/astraea/cost/broker/MemoryCost.java
+++ b/app/src/main/java/org/astraea/cost/broker/MemoryCost.java
@@ -31,34 +31,33 @@ public class MemoryCost extends Periodic<Map<Integer, Double>> implements HasBro
    */
   @Override
   public BrokerCost brokerCost(ClusterInfo clusterInfo) {
-    var brokerScore =
-        tryUpdate(
-            () -> {
-              var costMetrics =
-                  clusterInfo.allBeans().entrySet().stream()
-                      .collect(Collectors.toMap(Map.Entry::getKey, entry -> 0.0));
-              clusterInfo
-                  .allBeans()
-                  .forEach(
-                      (brokerID, value) -> {
-                        if (!brokersMetric.containsKey(brokerID)) {
-                          brokersMetric.put(brokerID, new BrokerMetric());
-                        }
-                        value.stream()
-                            .filter(beanObject -> beanObject instanceof HasJvmMemory)
-                            .forEach(
-                                hasBeanObject -> {
-                                  var jvmBean = (HasJvmMemory) hasBeanObject;
-                                  costMetrics.put(
-                                      brokerID,
-                                      ((jvmBean.heapMemoryUsage().getUsed() + 0.0)
-                                          / (jvmBean.heapMemoryUsage().getMax() + 1)));
-                                });
-                      });
-              TScore(costMetrics).forEach((broker, v) -> brokersMetric.get(broker).updateLoad(v));
-              return computeLoad();
-            });
-    return () -> brokerScore;
+    var costMetrics =
+        clusterInfo.allBeans().entrySet().stream()
+            .collect(
+                Collectors.toMap(
+                    Map.Entry::getKey,
+                    entry ->
+                        entry.getValue().stream()
+                            .filter(hasBeanObject -> hasBeanObject instanceof HasJvmMemory)
+                            .findAny()
+                            .orElseThrow()))
+            .entrySet()
+            .stream()
+            .collect(
+                Collectors.toMap(
+                    Map.Entry::getKey,
+                    entry -> {
+                      if (!brokersMetric.containsKey(entry.getKey())) {
+                        brokersMetric.put(entry.getKey(), new BrokerMetric());
+                      }
+                      var jvmBean = (HasJvmMemory) entry.getValue();
+                      return (jvmBean.heapMemoryUsage().getUsed() + 0.0)
+                          / (jvmBean.heapMemoryUsage().getMax() + 1);
+                    }));
+
+    TScore(costMetrics).forEach((broker, v) -> brokersMetric.get(broker).updateLoad(v));
+
+    return this::computeLoad;
   }
 
   Map<Integer, Double> computeLoad() {
diff --git a/app/src/main/java/org/astraea/cost/broker/NeutralIntegratedCost.java b/app/src/main/java/org/astraea/cost/broker/NeutralIntegratedCost.java
index c01f708b..10ebba84 100644
--- a/app/src/main/java/org/astraea/cost/broker/NeutralIntegratedCost.java
+++ b/app/src/main/java/org/astraea/cost/broker/NeutralIntegratedCost.java
@@ -9,7 +9,6 @@ import org.astraea.cost.BrokerCost;
 import org.astraea.cost.ClusterInfo;
 import org.astraea.cost.CostFunction;
 import org.astraea.cost.HasBrokerCost;
-import org.astraea.cost.Periodic;
 import org.astraea.metrics.collector.Fetcher;
 
 /**
@@ -19,7 +18,7 @@ import org.astraea.metrics.collector.Fetcher;
  * Because AHP is a subjective evaluation method and entropy method is an objective evaluation
  * method, so I named it NeutralIntegratedCost.
  */
-public class NeutralIntegratedCost extends Periodic<Map<Integer, Double>> implements HasBrokerCost {
+public class NeutralIntegratedCost implements HasBrokerCost {
   private final List<HasBrokerCost> metricsCost =
       List.of(new BrokerInputCost(), new BrokerOutputCost(), new CpuCost(), new MemoryCost());
   private final Map<Integer, BrokerMetrics> brokersMetric = new HashMap<>();
@@ -28,75 +27,64 @@ public class NeutralIntegratedCost extends Periodic<Map<Integer, Double>> implem
 
   @Override
   public BrokerCost brokerCost(ClusterInfo clusterInfo) {
-    var score =
-        tryUpdate(
-            () -> {
-              var costMetrics =
-                  clusterInfo.allBeans().entrySet().stream()
-                      .collect(Collectors.toMap(Map.Entry::getKey, entry -> 0.0));
-              costMetrics.forEach(
-                  (key, value) -> {
-                    if (!brokersMetric.containsKey(key)) {
-                      brokersMetric.put(key, new BrokerMetrics());
-                    }
-                  });
-
-              metricsCost.forEach(
-                  hasBrokerCost -> {
-                    if (hasBrokerCost instanceof BrokerInputCost) {
-                      hasBrokerCost
-                          .brokerCost(clusterInfo)
-                          .value()
-                          .forEach(
-                              (brokerID, value) -> brokersMetric.get(brokerID).inputScore = value);
-                    } else if (hasBrokerCost instanceof BrokerOutputCost) {
-                      hasBrokerCost
-                          .brokerCost(clusterInfo)
-                          .value()
-                          .forEach(
-                              (brokerID, value) -> brokersMetric.get(brokerID).outputScore = value);
-                    } else if (hasBrokerCost instanceof CpuCost) {
-                      hasBrokerCost
-                          .brokerCost(clusterInfo)
-                          .value()
-                          .forEach(
-                              (brokerID, value) -> brokersMetric.get(brokerID).cpuScore = value);
-                    } else if (hasBrokerCost instanceof MemoryCost) {
-                      hasBrokerCost
-                          .brokerCost(clusterInfo)
-                          .value()
-                          .forEach(
-                              (brokerID, value) -> brokersMetric.get(brokerID).memoryScore = value);
-                    }
-                  });
-
-              var integratedEmpowerment =
-                  entropyEmpowerment.empowerment(brokersMetric).entrySet().stream()
-                      .collect(
-                          Collectors.toMap(
-                              Map.Entry::getKey,
-                              entry ->
-                                  entry.getValue() * 0.5
-                                      + ahpEmpowerment.empowerment().get(entry.getKey()) * 0.5));
-
-              return brokersMetric.entrySet().stream()
-                  .collect(
-                      Collectors.toMap(
-                          Map.Entry::getKey,
-                          entry ->
-                              entry.getValue().inputScore
-                                      * integratedEmpowerment.get(
-                                          Metrics.inputThroughput.metricName)
-                                  + entry.getValue().outputScore
-                                      * integratedEmpowerment.get(
-                                          Metrics.outputThroughput.metricName)
-                                  + entry.getValue().cpuScore
-                                      * integratedEmpowerment.get(Metrics.cpu.metricName)
-                                  + entry.getValue().memoryScore
-                                      * integratedEmpowerment.get(Metrics.memory.metricName)));
-            });
-
-    return () -> score;
+    var costMetrics =
+        clusterInfo.allBeans().entrySet().stream()
+            .collect(Collectors.toMap(Map.Entry::getKey, entry -> 0.0));
+    costMetrics.forEach(
+        (key, value) -> {
+          if (!brokersMetric.containsKey(key)) {
+            brokersMetric.put(key, new BrokerMetrics());
+          }
+        });
+
+    metricsCost.forEach(
+        hasBrokerCost -> {
+          if (hasBrokerCost instanceof BrokerInputCost) {
+            hasBrokerCost
+                .brokerCost(clusterInfo)
+                .value()
+                .forEach((brokerID, value) -> brokersMetric.get(brokerID).inputScore = value);
+          } else if (hasBrokerCost instanceof BrokerOutputCost) {
+            hasBrokerCost
+                .brokerCost(clusterInfo)
+                .value()
+                .forEach((brokerID, value) -> brokersMetric.get(brokerID).outputScore = value);
+          } else if (hasBrokerCost instanceof CpuCost) {
+            hasBrokerCost
+                .brokerCost(clusterInfo)
+                .value()
+                .forEach((brokerID, value) -> brokersMetric.get(brokerID).cpuScore = value);
+          } else if (hasBrokerCost instanceof MemoryCost) {
+            hasBrokerCost
+                .brokerCost(clusterInfo)
+                .value()
+                .forEach((brokerID, value) -> brokersMetric.get(brokerID).memoryScore = value);
+          }
+        });
+
+    var integratedEmpowerment =
+        entropyEmpowerment.empowerment(brokersMetric).entrySet().stream()
+            .collect(
+                Collectors.toMap(
+                    Map.Entry::getKey,
+                    entry ->
+                        entry.getValue() * 0.5
+                            + ahpEmpowerment.empowerment().get(entry.getKey()) * 0.5));
+
+    return () ->
+        brokersMetric.entrySet().stream()
+            .collect(
+                Collectors.toMap(
+                    Map.Entry::getKey,
+                    entry ->
+                        entry.getValue().inputScore
+                                * integratedEmpowerment.get(Metrics.inputThroughput.metricName)
+                            + entry.getValue().outputScore
+                                * integratedEmpowerment.get(Metrics.outputThroughput.metricName)
+                            + entry.getValue().cpuScore
+                                * integratedEmpowerment.get(Metrics.cpu.metricName)
+                            + entry.getValue().memoryScore
+                                * integratedEmpowerment.get(Metrics.memory.metricName)));
   }
 
   @Override
diff --git a/app/src/main/java/org/astraea/partitioner/smooth/SmoothWeight.java b/app/src/main/java/org/astraea/partitioner/smooth/SmoothWeight.java
deleted file mode 100644
index e049650b..00000000
--- a/app/src/main/java/org/astraea/partitioner/smooth/SmoothWeight.java
+++ /dev/null
@@ -1,47 +0,0 @@
-package org.astraea.partitioner.smooth;
-
-import java.util.Comparator;
-import java.util.HashMap;
-import java.util.Map;
-import java.util.Objects;
-
-/**
- * Given initial key-score pair, it will output a preferred key with higher score. The score of the
- * chosen key will decrease. And the other score will increment by its original score. It may result
- * in "higher score with higher chosen rate".
- */
-public final class SmoothWeight {
-  private Map<Integer, Integer> load;
-  private int loadSum;
-  public Map<Integer, Integer> weight;
-
-  public SmoothWeight() {}
-
-  public SmoothWeight(Map<Integer, Integer> load) {
-    init(load);
-  }
-
-  public synchronized void init(Map<Integer, Integer> load) {
-    this.load = new HashMap<>(load);
-    this.weight = new HashMap<>(load);
-    this.loadSum = load.values().stream().mapToInt(d -> d).sum();
-  }
-
-  /**
-   * Get the preferred ID, and update the state.
-   *
-   * @return the preferred ID
-   */
-  public int getAndChoose() {
-    Objects.requireNonNull(this.weight);
-
-    this.weight.replaceAll((ID, value) -> value + this.load.get(ID));
-    var maxID =
-        this.weight.entrySet().stream()
-            .max(Comparator.comparingDouble(Map.Entry::getValue))
-            .orElseGet(() -> Map.entry(0, 0))
-            .getKey();
-    this.weight.computeIfPresent(maxID, (ID, value) -> value - loadSum);
-    return maxID;
-  }
-}
diff --git a/app/src/main/java/org/astraea/partitioner/smooth/SmoothWeightDispatcher.java b/app/src/main/java/org/astraea/partitioner/smooth/SmoothWeightDispatcher.java
deleted file mode 100644
index a3247359..00000000
--- a/app/src/main/java/org/astraea/partitioner/smooth/SmoothWeightDispatcher.java
+++ /dev/null
@@ -1,159 +0,0 @@
-package org.astraea.partitioner.smooth;
-
-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.Optional;
-import java.util.Random;
-import java.util.TreeMap;
-import java.util.concurrent.locks.Lock;
-import java.util.concurrent.locks.ReentrantLock;
-import java.util.stream.Collectors;
-import org.astraea.Utils;
-import org.astraea.cost.ClusterInfo;
-import org.astraea.cost.CostFunction;
-import org.astraea.cost.HasBrokerCost;
-import org.astraea.cost.LoadCost;
-import org.astraea.cost.MemoryWarningCost;
-import org.astraea.metrics.collector.BeanCollector;
-import org.astraea.metrics.collector.Fetcher;
-import org.astraea.metrics.collector.Receiver;
-import org.astraea.partitioner.Configuration;
-import org.astraea.partitioner.Dispatcher;
-
-/**
- * Based on the jmx metrics obtained from Kafka, it records the load status of the node over a
- * period of time. Predict the future status of each node through the poisson of the load status.
- * Finally, the result of poisson is used as the weight to perform smooth weighted RoundRobin.
- */
-public class SmoothWeightDispatcher implements Dispatcher {
-  public static final String JMX_PORT = "jmx.port";
-  private final BeanCollector beanCollector =
-      BeanCollector.builder().interval(Duration.ofSeconds(1)).build();
-  private Optional<Integer> jmxPortDefault = Optional.empty();
-  private final Map<Integer, Integer> jmxPorts = new TreeMap<>();
-  private final Map<Integer, Receiver> receivers = new TreeMap<>();
-
-  private final Collection<CostFunction> functions =
-      List.of(new LoadCost(), new MemoryWarningCost());
-  // Memory warning can dominate the score.
-  private final Map<Object, Double> costFraction =
-      Map.of(LoadCost.class, 1.0, MemoryWarningCost.class, -100.0);
-
-  private final SmoothWeight smoothWeightCal = new SmoothWeight();
-  private final Random random = new Random();
-  // Fetch data every n seconds
-  private long lastFetchTime = 0L;
-  private final Lock lock = new ReentrantLock();
-
-  @Override
-  public int partition(String topic, byte[] key, byte[] value, ClusterInfo clusterInfo) {
-    var partitionLeaders = clusterInfo.availablePartitionLeaders(topic);
-    // just return first partition if there is no available partitions
-    if (partitionLeaders.isEmpty()) return 0;
-
-    // just return the only one available partition
-    if (partitionLeaders.size() == 1) return partitionLeaders.iterator().next().partition();
-
-    if (Utils.overSecond(lastFetchTime, 1) && lock.tryLock()) {
-      try {
-        // 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()))));
-
-        // fetch the latest beans for each node
-        var beans =
-            receivers.entrySet().stream()
-                .collect(Collectors.toMap(Map.Entry::getKey, e -> e.getValue().current()));
-
-        var compoundScore =
-            functions.stream()
-                .filter(f -> f instanceof HasBrokerCost)
-                .map(f -> (HasBrokerCost) f)
-                .map(
-                    f -> {
-                      var map = f.brokerCost(ClusterInfo.of(clusterInfo, beans)).value();
-                      map.replaceAll((k, v) -> v * costFraction.get(f.getClass()));
-                      return map;
-                    })
-                .reduce(new HashMap<>(), SmoothWeightDispatcher::costCompound)
-                .entrySet()
-                .stream()
-                .map(e -> Map.entry(e.getKey(), (int) (double) e.getValue()))
-                .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
-
-        smoothWeightCal.init(compoundScore);
-      } finally {
-        lastFetchTime = System.currentTimeMillis();
-        lock.unlock();
-      }
-    }
-
-    // Make "smooth weight choosing" on the score.
-    var targetBroker = smoothWeightCal.getAndChoose();
-    var targetPartitions =
-        partitionLeaders.stream()
-            .filter(p -> p.nodeInfo().id() == targetBroker)
-            .collect(Collectors.toList());
-
-    return targetPartitions.get(random.nextInt(targetPartitions.size())).partition();
-  }
-
-  // Just add all cost function score together
-  static Map<Integer, Double> costCompound(
-      Map<Integer, Double> identity, Map<Integer, Double> cost) {
-    cost.forEach((id, value) -> identity.computeIfPresent(id, (k, v) -> v + value));
-    cost.forEach(identity::putIfAbsent);
-    return identity;
-  }
-
-  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
-  public void configure(Configuration config) {
-    jmxPortDefault = config.integer(JMX_PORT);
-
-    // seeks for custom jmx ports.
-    config.entrySet().stream()
-        .filter(e -> e.getKey().startsWith("broker."))
-        .filter(e -> e.getKey().endsWith(JMX_PORT))
-        .map(
-            e ->
-                Map.entry(
-                    e.getKey().replaceAll("^broker[.]", "").replaceAll("[.]" + JMX_PORT + "$", ""),
-                    e.getValue()))
-        .forEach(e -> jmxPorts.put(Integer.parseInt(e.getKey()), Integer.parseInt(e.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(Utils::close);
-    receivers.clear();
-  }
-}
diff --git a/app/src/main/java/org/astraea/partitioner/smooth/SmoothWeightRoundRobinDispatcher.java b/app/src/main/java/org/astraea/partitioner/smooth/SmoothWeightRoundRobinDispatcher.java
new file mode 100644
index 00000000..3dff8429
--- /dev/null
+++ b/app/src/main/java/org/astraea/partitioner/smooth/SmoothWeightRoundRobinDispatcher.java
@@ -0,0 +1,175 @@
+package org.astraea.partitioner.smooth;
+
+import static org.astraea.partitioner.PartitionerUtils.partitionerConfig;
+
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+import java.util.NoSuchElementException;
+import java.util.Optional;
+import java.util.TreeMap;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentLinkedDeque;
+import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
+import org.apache.kafka.common.Cluster;
+import org.astraea.Utils;
+import org.astraea.cost.ClusterInfo;
+import org.astraea.cost.NodeInfo;
+import org.astraea.cost.Periodic;
+import org.astraea.cost.ReplicaInfo;
+import org.astraea.cost.broker.NeutralIntegratedCost;
+import org.astraea.metrics.HasBeanObject;
+import org.astraea.metrics.collector.BeanCollector;
+import org.astraea.metrics.collector.Receiver;
+import org.astraea.metrics.jmx.MBeanClient;
+import org.astraea.partitioner.Configuration;
+import org.astraea.partitioner.Dispatcher;
+
+public class SmoothWeightRoundRobinDispatcher extends Periodic<Map<Integer, Double>>
+    implements Dispatcher {
+  private final ConcurrentLinkedDeque<Integer> unusedPartitions = new ConcurrentLinkedDeque<>();
+  private final ConcurrentMap<String, BrokerNextCounter> topicCounter = new ConcurrentHashMap<>();
+  private final BeanCollector beanCollector =
+      BeanCollector.builder()
+          .interval(Duration.ofSeconds(1))
+          .numberOfObjectsPerNode(1)
+          .clientCreator(MBeanClient::jndi)
+          .build();
+  private final Optional<Integer> jmxPortDefault = Optional.empty();
+  private final Map<Integer, Integer> jmxPorts = new TreeMap<>();
+  private final Map<Integer, Receiver> receivers = new TreeMap<>();
+
+  private final Map<Integer, List<Integer>> hasPartitions = new TreeMap<>();
+
+  private SmoothWeightRoundRobin smoothWeightRoundRobinCal;
+
+  private final NeutralIntegratedCost neutralIntegratedCost = new NeutralIntegratedCost();
+
+  private Map<Integer, Collection<HasBeanObject>> beans;
+  private List<ReplicaInfo> partitions;
+
+  public static final String JMX_PORT = "jmx.port";
+
+  @Override
+  public int partition(String topic, byte[] key, byte[] value, ClusterInfo clusterInfo) {
+    var targetPartition = unusedPartitions.poll();
+    tryUpdateAfterOneSecond(
+        () -> {
+          refreshPartitionMetaData(clusterInfo, topic);
+          // fetch the latest beans for each node
+          beans =
+              receivers.entrySet().stream()
+                  .collect(Collectors.toMap(Map.Entry::getKey, e -> e.getValue().current()));
+
+          var compoundScore =
+              neutralIntegratedCost.brokerCost(ClusterInfo.of(clusterInfo, beans)).value();
+
+          if (smoothWeightRoundRobinCal == null) {
+            smoothWeightRoundRobinCal = new SmoothWeightRoundRobin(compoundScore);
+          }
+          smoothWeightRoundRobinCal.init(compoundScore);
+
+          return compoundScore;
+        });
+    // 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();
+
+    if (targetPartition == null) {
+      var targetBroker = smoothWeightRoundRobinCal.getAndChoose();
+      targetPartition =
+          hasPartitions
+              .get(targetBroker)
+              .get(
+                  nextValue(topic, clusterInfo, targetBroker)
+                      % hasPartitions.get(targetBroker).size());
+    }
+
+    return targetPartition;
+  }
+
+  @Override
+  public void close() {
+    receivers.values().forEach(Utils::close);
+    receivers.clear();
+  }
+
+  @Override
+  public void configure(Map<String, ?> configs) {
+    var properties = partitionerConfig(configs);
+    var config =
+        Configuration.of(
+            properties.entrySet().stream()
+                .collect(
+                    Collectors.toMap(e -> e.getKey().toString(), e -> e.getValue().toString())));
+
+    // seeks for custom jmx ports.
+    config.entrySet().stream()
+        .filter(e -> e.getKey().startsWith("broker."))
+        .filter(e -> e.getKey().endsWith(JMX_PORT))
+        .forEach(
+            e ->
+                jmxPorts.put(
+                    Integer.parseInt(e.getKey().split("\\.")[1]), Integer.parseInt(e.getValue())));
+  }
+
+  @Override
+  public void onNewBatch(String topic, Cluster cluster, int prevPartition) {
+    unusedPartitions.add(prevPartition);
+  }
+
+  int jmxPort(int id) {
+    if (jmxPorts.containsKey(id)) return jmxPorts.get(id);
+    return jmxPortDefault.orElseThrow(
+        () -> new NoSuchElementException("broker: " + id + " does not have jmx port"));
+  }
+
+  Receiver receiver(String host, int port) {
+    return beanCollector
+        .register()
+        .host(host)
+        .port(port)
+        .fetcher(neutralIntegratedCost.fetcher())
+        .build();
+  }
+
+  private int nextValue(String topic, ClusterInfo clusterInfo, int targetBroker) {
+    return topicCounter
+        .computeIfAbsent(topic, k -> new BrokerNextCounter(clusterInfo))
+        .brokerCounter
+        .get(targetBroker)
+        .getAndIncrement();
+  }
+
+  private void refreshPartitionMetaData(ClusterInfo clusterInfo, String topic) {
+    partitions = clusterInfo.availablePartitions(topic);
+    partitions.forEach(
+        p ->
+            hasPartitions
+                .computeIfAbsent(p.nodeInfo().id(), k -> new ArrayList<>())
+                .add(p.partition()));
+
+    partitions.stream()
+        .filter(p -> !receivers.containsKey(p.nodeInfo().id()))
+        .forEach(
+            p ->
+                receivers.put(
+                    p.nodeInfo().id(), receiver(p.nodeInfo().host(), jmxPort(p.nodeInfo().id()))));
+  }
+
+  private static class BrokerNextCounter {
+    private final Map<Integer, AtomicInteger> brokerCounter;
+
+    BrokerNextCounter(ClusterInfo clusterInfo) {
+      brokerCounter =
+          clusterInfo.nodes().stream()
+              .collect(Collectors.toMap(NodeInfo::id, node -> new AtomicInteger(0)));
+    }
+  }
+}
diff --git a/docker/start_spark.sh b/docker/start_spark.sh
index 5e7da6cd..adbce3d7 100755
--- a/docker/start_spark.sh
+++ b/docker/start_spark.sh
@@ -46,7 +46,7 @@ function checkConflictContainer() {
 
 function generateDockerfileByVersion() {
   echo "# this dockerfile is generated dynamically
-FROM ubuntu:22.04 AS build
+FROM ubuntu:20.04 AS build
 
 # install tools
 RUN apt-get update && apt-get install -y wget unzip
@@ -57,7 +57,9 @@ RUN wget https://archive.apache.org/dist/spark/spark-${VERSION}/spark-${VERSION}
 RUN mkdir /opt/spark
 RUN tar -zxvf spark-${VERSION}-bin-hadoop3.2.tgz -C /opt/spark --strip-components=1
 
-FROM ubuntu:22.04
+# the python3 in ubuntu 22.04 is 3.10 by default, and it has a known issue (https://github.com/vmprof/vmprof-python/issues/240)
+# The issue obstructs us from installing 3-third python libraries, so we downgrade the ubuntu to 20.04
+FROM ubuntu:20.04
 
 # Do not ask for confirmations when running apt-get, etc.
 ENV DEBIAN_FRONTEND noninteractive
@@ -101,7 +103,9 @@ RUN mkdir /opt/spark
 RUN tar -zxvf \$(find ./ -maxdepth 1 -type f -name spark-*SNAPSHOT*.tgz) -C /opt/spark --strip-components=1
 RUN ./build/mvn install -DskipTests
 
-FROM ubuntu:22.04
+# the python3 in ubuntu 22.04 is 3.10 by default, and it has a known issue (https://github.com/vmprof/vmprof-python/issues/240)
+# The issue obstructs us from installing 3-third python libraries, so we downgrade the ubuntu to 20.04
+FROM ubuntu:20.04
 
 # Do not ask for confirmations when running apt-get, etc.
 ENV DEBIAN_FRONTEND noninteractive


Test patch:
diff --git a/app/src/test/java/org/astraea/balancer/RebalancePlanProposalTest.java b/app/src/test/java/org/astraea/balancer/RebalancePlanProposalTest.java
new file mode 100644
index 00000000..89c28eb7
--- /dev/null
+++ b/app/src/test/java/org/astraea/balancer/RebalancePlanProposalTest.java
@@ -0,0 +1,68 @@
+package org.astraea.balancer;
+
+import java.util.stream.Collectors;
+import org.astraea.balancer.log.LayeredClusterLogAllocation;
+import org.astraea.cost.ClusterInfoProvider;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+class RebalancePlanProposalTest {
+
+  @Test
+  void testBuild() {
+    final var fakeClusterInfo = ClusterInfoProvider.fakeClusterInfo(10, 10, 10, 10);
+    final var thisAllocation = LayeredClusterLogAllocation.of(fakeClusterInfo);
+    final var build =
+        RebalancePlanProposal.builder()
+            .withRebalancePlan(thisAllocation)
+            .addInfo("Info0")
+            .addInfo("Info1")
+            .addInfo("Info2")
+            .addWarning("Warning0")
+            .addWarning("Warning1")
+            .addWarning("Warning2")
+            .build();
+
+    Assertions.assertTrue(build.rebalancePlan().isPresent());
+    final var thatAllocation = build.rebalancePlan().orElseThrow();
+    final var thisTps =
+        thisAllocation.topicPartitionStream().collect(Collectors.toUnmodifiableSet());
+    final var thatTps =
+        thatAllocation.topicPartitionStream().collect(Collectors.toUnmodifiableSet());
+    Assertions.assertEquals(thisTps, thatTps);
+    thisTps.forEach(
+        tp ->
+            Assertions.assertEquals(
+                thisAllocation.logPlacements(tp), thatAllocation.logPlacements(tp)));
+    Assertions.assertEquals("Info0", build.info().get(0));
+    Assertions.assertEquals("Info1", build.info().get(1));
+    Assertions.assertEquals("Info2", build.info().get(2));
+    Assertions.assertEquals("Warning0", build.warnings().get(0));
+    Assertions.assertEquals("Warning1", build.warnings().get(1));
+    Assertions.assertEquals("Warning2", build.warnings().get(2));
+  }
+
+  @Test
+  void testNoBuildTwice() {
+    // A builder should only build once. If a builder can build multiple times then it will have to
+    // do much copy work once a new build is requested. This will harm performance.
+    final var fakeClusterInfo = ClusterInfoProvider.fakeClusterInfo(10, 10, 10, 10);
+    final var logAllocation = LayeredClusterLogAllocation.of(fakeClusterInfo);
+    final var build = RebalancePlanProposal.builder().withRebalancePlan(logAllocation);
+
+    Assertions.assertDoesNotThrow(build::build);
+    Assertions.assertThrows(IllegalStateException.class, build::build);
+  }
+
+  @Test
+  void testNoModifyAfterBuild() {
+    final var fakeClusterInfo = ClusterInfoProvider.fakeClusterInfo(10, 10, 10, 10);
+    final var logAllocation = LayeredClusterLogAllocation.of(fakeClusterInfo);
+    final var build = RebalancePlanProposal.builder().withRebalancePlan(logAllocation);
+
+    RebalancePlanProposal proposal = build.build();
+    Assertions.assertThrows(
+        IllegalStateException.class, () -> build.addInfo("modify after built."));
+    Assertions.assertFalse(proposal.info().contains("modify after built."));
+  }
+}
diff --git a/app/src/test/java/org/astraea/balancer/generator/ShufflePlanGeneratorTest.java b/app/src/test/java/org/astraea/balancer/generator/ShufflePlanGeneratorTest.java
new file mode 100644
index 00000000..e6dac26e
--- /dev/null
+++ b/app/src/test/java/org/astraea/balancer/generator/ShufflePlanGeneratorTest.java
@@ -0,0 +1,173 @@
+package org.astraea.balancer.generator;
+
+import java.util.concurrent.ForkJoinPool;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.LongAdder;
+import java.util.stream.Collectors;
+import org.astraea.balancer.log.LayeredClusterLogAllocation;
+import org.astraea.cost.ClusterInfoProvider;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Disabled;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvSource;
+import org.junit.jupiter.params.provider.ValueSource;
+
+class ShufflePlanGeneratorTest {
+
+  @Test
+  void testRun() {
+    final var shufflePlanGenerator = new ShufflePlanGenerator(5, 10);
+    final var fakeCluster = ClusterInfoProvider.fakeClusterInfo(100, 10, 10, 3);
+    final var stream = shufflePlanGenerator.generate(fakeCluster);
+    final var iterator = stream.iterator();
+
+    Assertions.assertDoesNotThrow(() -> System.out.println(iterator.next()));
+    Assertions.assertDoesNotThrow(() -> System.out.println(iterator.next()));
+    Assertions.assertDoesNotThrow(() -> System.out.println(iterator.next()));
+  }
+
+  @ParameterizedTest
+  @ValueSource(ints = {3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 301})
+  void testMovement(int shuffle) {
+    final var fakeClusterInfo = ClusterInfoProvider.fakeClusterInfo(30, 30, 20, 5);
+    final var allocation = LayeredClusterLogAllocation.of(fakeClusterInfo);
+    final var shufflePlanGenerator = new ShufflePlanGenerator(() -> shuffle);
+
+    shufflePlanGenerator
+        .generate(fakeClusterInfo)
+        .limit(100)
+        .forEach(
+            proposal -> {
+              final var that = proposal.rebalancePlan().orElseThrow();
+              final var thisTps =
+                  allocation.topicPartitionStream().collect(Collectors.toUnmodifiableSet());
+              final var thatTps =
+                  that.topicPartitionStream().collect(Collectors.toUnmodifiableSet());
+              final var thisMap =
+                  thisTps.stream()
+                      .collect(Collectors.toUnmodifiableMap(x -> x, allocation::logPlacements));
+              final var thatMap =
+                  thatTps.stream()
+                      .collect(Collectors.toUnmodifiableMap(x -> x, that::logPlacements));
+              Assertions.assertEquals(thisTps, thatTps);
+              Assertions.assertNotEquals(thisMap, thatMap);
+            });
+  }
+
+  @Test
+  void testNoNodes() {
+    final var fakeClusterInfo = ClusterInfoProvider.fakeClusterInfo(0, 0, 0, 0);
+    final var shufflePlanGenerator = new ShufflePlanGenerator(() -> 3);
+
+    final var proposal = shufflePlanGenerator.generate(fakeClusterInfo).iterator().next();
+
+    System.out.println(proposal);
+    Assertions.assertFalse(proposal.rebalancePlan().isPresent());
+    Assertions.assertTrue(proposal.warnings().size() >= 1);
+  }
+
+  @Test
+  void testOneNode() {
+    final var fakeClusterInfo = ClusterInfoProvider.fakeClusterInfo(1, 1, 1, 1);
+    final var shufflePlanGenerator = new ShufflePlanGenerator(() -> 3);
+
+    final var proposal = shufflePlanGenerator.generate(fakeClusterInfo).iterator().next();
+
+    System.out.println(proposal);
+    Assertions.assertFalse(proposal.rebalancePlan().isPresent());
+    Assertions.assertTrue(proposal.warnings().size() >= 1);
+  }
+
+  @Test
+  void testNoTopic() {
+    final var fakeClusterInfo = ClusterInfoProvider.fakeClusterInfo(3, 0, 0, 0);
+    final var shufflePlanGenerator = new ShufflePlanGenerator(() -> 3);
+
+    final var proposal = shufflePlanGenerator.generate(fakeClusterInfo).iterator().next();
+
+    System.out.println(proposal);
+    Assertions.assertFalse(proposal.rebalancePlan().isPresent());
+    Assertions.assertTrue(proposal.warnings().size() >= 1);
+  }
+
+  @ParameterizedTest(name = "[{0}] {1} nodes, {2} topics, {3} partitions, {4} replicas")
+  @CsvSource(
+      value = {
+        //      scenario, node, topic, partition, replica
+        "  small cluster,    3,    10,        30,       3",
+        " medium cluster,   30,    50,        50,       3",
+        "    big cluster,  100,   100,       100,       1",
+        "  many replicas, 1000,    30,       100,      30",
+      })
+  void performanceTest(
+      String scenario, int nodeCount, int topicCount, int partitionCount, int replicaCount) {
+    // This test is not intended for any performance guarantee.
+    // It only served the purpose of keeping track of the generator performance change in the CI
+    // log.
+    // Notice: Stream#limit() will hurt performance. the number here might not reflect the actual
+    // performance.
+    final var shufflePlanGenerator = new ShufflePlanGenerator(0, 10);
+    final var fakeClusterInfo =
+        ClusterInfoProvider.fakeClusterInfo(nodeCount, topicCount, partitionCount, replicaCount);
+    final var size = 1000;
+
+    final long s = System.nanoTime();
+    final var count = shufflePlanGenerator.generate(fakeClusterInfo).limit(size).count();
+    final long t = System.nanoTime();
+    Assertions.assertEquals(size, count);
+    System.out.printf("[%s]%n", scenario);
+    System.out.printf(
+        "%d nodes, %d topics, %d partitions, %d replicas.%n",
+        nodeCount, topicCount, partitionCount, replicaCount);
+    System.out.printf("Generate %.3f proposals per second.%n", count / (((double) (t - s) / 1e9)));
+    System.out.println();
+  }
+
+  @Test
+  void parallelStreamWorks() {
+    final var shufflePlanGenerator = new ShufflePlanGenerator(0, 10);
+    final var fakeClusterInfo = ClusterInfoProvider.fakeClusterInfo(10, 20, 10, 3);
+
+    // generator can do parallel without error.
+    Assertions.assertDoesNotThrow(
+        () -> shufflePlanGenerator.generate(fakeClusterInfo).parallel().limit(100).count());
+  }
+
+  @Test
+  @Disabled
+  void parallelPerformanceTests() throws InterruptedException {
+    final var shufflePlanGenerator = new ShufflePlanGenerator(0, 10);
+    final var fakeClusterInfo = ClusterInfoProvider.fakeClusterInfo(50, 500, 30, 2);
+    final var counter = new LongAdder();
+    final var forkJoinPool = new ForkJoinPool(ForkJoinPool.getCommonPoolParallelism());
+    final var startTime = System.nanoTime();
+
+    forkJoinPool.submit(
+        () ->
+            shufflePlanGenerator
+                .generate(fakeClusterInfo)
+                .parallel()
+                .forEach((ignore) -> counter.increment()));
+
+    // report progress
+    final var reportThread =
+        new Thread(
+            () -> {
+              while (!Thread.currentThread().isInterrupted()) {
+                try {
+                  TimeUnit.SECONDS.sleep(1);
+                  long now = System.nanoTime();
+                  System.out.printf(
+                      "%.3f proposal/s %n", counter.sum() / ((now - startTime) / 1e9));
+                } catch (InterruptedException e) {
+                  break;
+                }
+              }
+            });
+
+    reportThread.start();
+    reportThread.join();
+    forkJoinPool.shutdownNow();
+  }
+}
diff --git a/app/src/test/java/org/astraea/balancer/log/LayeredClusterLogAllocationTest.java b/app/src/test/java/org/astraea/balancer/log/LayeredClusterLogAllocationTest.java
new file mode 100644
index 00000000..d2c38609
--- /dev/null
+++ b/app/src/test/java/org/astraea/balancer/log/LayeredClusterLogAllocationTest.java
@@ -0,0 +1,166 @@
+package org.astraea.balancer.log;
+
+import java.util.Comparator;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.stream.Collectors;
+import org.astraea.admin.TopicPartition;
+import org.astraea.cost.ClusterInfoProvider;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+class LayeredClusterLogAllocationTest {
+
+  @Test
+  void creation() {
+    // empty replica set
+    var badAllocation0 = Map.of(TopicPartition.of("topic", "0"), List.<LogPlacement>of());
+    Assertions.assertThrows(
+        IllegalArgumentException.class, () -> LayeredClusterLogAllocation.of(badAllocation0));
+
+    // partial topic/partition
+    var badAllocation1 = Map.of(TopicPartition.of("topic", "999"), List.of(LogPlacement.of(1001)));
+    Assertions.assertThrows(
+        IllegalArgumentException.class, () -> LayeredClusterLogAllocation.of(badAllocation1));
+
+    // duplicate replica
+    var badAllocation2 =
+        Map.of(
+            TopicPartition.of("topic", "0"),
+            List.of(LogPlacement.of(1001), LogPlacement.of(1001), LogPlacement.of(1001)));
+    Assertions.assertThrows(
+        IllegalArgumentException.class, () -> LayeredClusterLogAllocation.of(badAllocation2));
+  }
+
+  @Test
+  void migrateReplica() {
+    final var fakeClusterInfo =
+        ClusterInfoProvider.fakeClusterInfo(3, 1, 1, 1, (i) -> Set.of("topic"));
+    final var clusterLogAllocation = LayeredClusterLogAllocation.of(fakeClusterInfo);
+    final var sourceTopicPartition = TopicPartition.of("topic", "0");
+
+    clusterLogAllocation.migrateReplica(sourceTopicPartition, 0, 1);
+
+    Assertions.assertEquals(
+        1, clusterLogAllocation.logPlacements(sourceTopicPartition).get(0).broker());
+    Assertions.assertDoesNotThrow(() -> LayeredClusterLogAllocation.of(clusterLogAllocation));
+  }
+
+  @Test
+  void letReplicaBecomeLeader() {
+    final var fakeClusterInfo =
+        ClusterInfoProvider.fakeClusterInfo(3, 1, 1, 2, (i) -> Set.of("topic"));
+    final var clusterLogAllocation = LayeredClusterLogAllocation.of(fakeClusterInfo);
+    final var sourceTopicPartition = TopicPartition.of("topic", "0");
+
+    clusterLogAllocation.letReplicaBecomeLeader(sourceTopicPartition, 1);
+
+    Assertions.assertEquals(
+        1, clusterLogAllocation.logPlacements(sourceTopicPartition).get(0).broker());
+    Assertions.assertEquals(
+        0, clusterLogAllocation.logPlacements(sourceTopicPartition).get(1).broker());
+    Assertions.assertDoesNotThrow(() -> LayeredClusterLogAllocation.of(clusterLogAllocation));
+  }
+
+  @Test
+  void changeDataDirectory() {
+    final var fakeClusterInfo =
+        ClusterInfoProvider.fakeClusterInfo(3, 1, 1, 1, (i) -> Set.of("topic"));
+    final var clusterLogAllocation = LayeredClusterLogAllocation.of(fakeClusterInfo);
+    final var sourceTopicPartition = TopicPartition.of("topic", "0");
+
+    clusterLogAllocation.changeDataDirectory(sourceTopicPartition, 0, "/path/to/somewhere");
+
+    Assertions.assertEquals(
+        "/path/to/somewhere",
+        clusterLogAllocation
+            .logPlacements(sourceTopicPartition)
+            .get(0)
+            .logDirectory()
+            .orElseThrow());
+    Assertions.assertDoesNotThrow(() -> LayeredClusterLogAllocation.of(clusterLogAllocation));
+  }
+
+  @Test
+  void logPlacements() {
+    final var allocation =
+        LayeredClusterLogAllocation.of(
+            Map.of(TopicPartition.of("topic", "0"), List.of(LogPlacement.of(0, "/nowhere"))));
+
+    Assertions.assertEquals(1, allocation.logPlacements(TopicPartition.of("topic", "0")).size());
+    Assertions.assertEquals(
+        0, allocation.logPlacements(TopicPartition.of("topic", "0")).get(0).broker());
+    Assertions.assertEquals(
+        "/nowhere",
+        allocation
+            .logPlacements(TopicPartition.of("topic", "0"))
+            .get(0)
+            .logDirectory()
+            .orElseThrow());
+    Assertions.assertNull(allocation.logPlacements(TopicPartition.of("no", "0")));
+    allocation.logPlacements(TopicPartition.of("no", "0"));
+  }
+
+  @Test
+  void topicPartitionStream() {
+    final var fakeClusterInfo = ClusterInfoProvider.fakeClusterInfo(10, 10, 10, 3);
+    final var allocation0 = LayeredClusterLogAllocation.of(fakeClusterInfo);
+    final var allocation1 = LayeredClusterLogAllocation.of(allocation0);
+    fakeClusterInfo.topics().stream()
+        .flatMap(x -> fakeClusterInfo.partitions(x).stream())
+        .forEach(
+            replica ->
+                allocation1.changeDataDirectory(
+                    TopicPartition.of(replica.topic(), Integer.toString(replica.partition())),
+                    0,
+                    "/yeah"));
+
+    final var allTopicPartitions =
+        allocation1
+            .topicPartitionStream()
+            .sorted(
+                Comparator.comparing(TopicPartition::topic)
+                    .thenComparing(TopicPartition::partition))
+            .collect(Collectors.toUnmodifiableList());
+    Assertions.assertEquals(10 * 10, allTopicPartitions.size());
+    final var expectedTopicPartitions =
+        fakeClusterInfo.topics().stream()
+            .flatMap(x -> fakeClusterInfo.partitions(x).stream())
+            .map(x -> TopicPartition.of(x.topic(), Integer.toString(x.partition())))
+            .distinct()
+            .sorted(
+                Comparator.comparing(TopicPartition::topic)
+                    .thenComparing(TopicPartition::partition))
+            .collect(Collectors.toUnmodifiableList());
+    Assertions.assertEquals(expectedTopicPartitions, allTopicPartitions);
+  }
+
+  @Test
+  void lockWorks() {
+    final var fakeClusterInfo = ClusterInfoProvider.fakeClusterInfo(10, 10, 10, 3);
+    final var allocation = LayeredClusterLogAllocation.of(fakeClusterInfo);
+    final var toModify = allocation.topicPartitionStream().findFirst().orElseThrow();
+
+    // can modify before lock
+    Assertions.assertDoesNotThrow(() -> allocation.migrateReplica(toModify, 0, 9));
+    Assertions.assertDoesNotThrow(() -> allocation.letReplicaBecomeLeader(toModify, 1));
+    Assertions.assertDoesNotThrow(() -> allocation.changeDataDirectory(toModify, 2, "/nowhere"));
+
+    final var extended = LayeredClusterLogAllocation.of(allocation);
+
+    // cannot modify after some other layer rely on it
+    Assertions.assertThrows(
+        IllegalStateException.class, () -> allocation.migrateReplica(toModify, 9, 0));
+    Assertions.assertThrows(
+        IllegalStateException.class, () -> allocation.letReplicaBecomeLeader(toModify, 0));
+    Assertions.assertThrows(
+        IllegalStateException.class,
+        () -> allocation.changeDataDirectory(toModify, 2, "/anywhere"));
+
+    // the extended one can modify
+    Assertions.assertDoesNotThrow(() -> extended.migrateReplica(toModify, 9, 0));
+    Assertions.assertDoesNotThrow(() -> extended.letReplicaBecomeLeader(toModify, 0));
+    Assertions.assertDoesNotThrow(() -> extended.changeDataDirectory(toModify, 2, "/anywhere"));
+  }
+}
diff --git a/app/src/test/java/org/astraea/cost/ClusterInfoProvider.java b/app/src/test/java/org/astraea/cost/ClusterInfoProvider.java
new file mode 100644
index 00000000..509f5540
--- /dev/null
+++ b/app/src/test/java/org/astraea/cost/ClusterInfoProvider.java
@@ -0,0 +1,124 @@
+package org.astraea.cost;
+
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+import java.util.Random;
+import java.util.Set;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+import org.astraea.admin.TopicPartition;
+import org.astraea.metrics.HasBeanObject;
+
+public class ClusterInfoProvider {
+
+  public static ClusterInfo fakeClusterInfo(
+      int nodeCount, int topicCount, int partitionCount, int replicaCount) {
+    final var random = new Random();
+    random.setSeed(0);
+
+    return fakeClusterInfo(
+        nodeCount,
+        topicCount,
+        partitionCount,
+        replicaCount,
+        (partition) ->
+            IntStream.range(0, partition)
+                .mapToObj(
+                    (ignore) -> {
+                      final var suffix =
+                          IntStream.range(0, 5)
+                              .map(i -> random.nextInt(26))
+                              .map(i -> 'a' + i)
+                              .mapToObj(i -> String.valueOf((char) i))
+                              .collect(Collectors.joining());
+                      return "fake-topic-" + suffix;
+                    })
+                .collect(Collectors.toUnmodifiableSet()));
+  }
+
+  public static ClusterInfo fakeClusterInfo(
+      int nodeCount,
+      int topicCount,
+      int partitionCount,
+      int replicaCount,
+      Function<Integer, Set<String>> topicNameGenerator) {
+    final var nodes =
+        IntStream.range(0, nodeCount)
+            .mapToObj(nodeId -> NodeInfo.of(nodeId, "host" + nodeId, 9092))
+            .collect(Collectors.toUnmodifiableList());
+    final var dataDirCount = 3;
+    final var dataDirectories =
+        IntStream.range(0, dataDirCount)
+            .mapToObj(i -> "/tmp/data-directory-" + i)
+            .collect(Collectors.toUnmodifiableSet());
+    final var dataDirectoryList = List.copyOf(dataDirectories);
+    final var topics = topicNameGenerator.apply(topicCount);
+    final var replicas =
+        topics.stream()
+            .flatMap(
+                topic ->
+                    IntStream.range(0, partitionCount)
+                        .mapToObj(p -> TopicPartition.of(topic, Integer.toString(p))))
+            .flatMap(
+                tp ->
+                    IntStream.range(0, replicaCount)
+                        .mapToObj(
+                            r ->
+                                ReplicaInfo.of(
+                                    tp.topic(),
+                                    tp.partition(),
+                                    nodes.get(r),
+                                    r == 0,
+                                    true,
+                                    false,
+                                    dataDirectoryList.get(
+                                        tp.partition() % dataDirectories.size()))))
+            .collect(Collectors.groupingBy(ReplicaInfo::topic));
+
+    return new ClusterInfo() {
+      @Override
+      public List<NodeInfo> nodes() {
+        return nodes;
+      }
+
+      @Override
+      public Set<String> dataDirectories(int brokerId) {
+        return dataDirectories;
+      }
+
+      @Override
+      public List<ReplicaInfo> availablePartitionLeaders(String topic) {
+        return partitions(topic).stream()
+            .filter(ReplicaInfo::isLeader)
+            .collect(Collectors.toUnmodifiableList());
+      }
+
+      @Override
+      public List<ReplicaInfo> availablePartitions(String topic) {
+        return partitions(topic);
+      }
+
+      @Override
+      public Set<String> topics() {
+        return topics;
+      }
+
+      @Override
+      public List<ReplicaInfo> partitions(String topic) {
+        return replicas.get(topic);
+      }
+
+      @Override
+      public Collection<HasBeanObject> beans(int brokerId) {
+        throw new UnsupportedOperationException();
+      }
+
+      @Override
+      public Map<Integer, Collection<HasBeanObject>> allBeans() {
+        throw new UnsupportedOperationException();
+      }
+    };
+  }
+}
diff --git a/app/src/test/java/org/astraea/cost/FakeClusterInfo.java b/app/src/test/java/org/astraea/cost/FakeClusterInfo.java
index fd0f1d75..bae2062e 100644
--- a/app/src/test/java/org/astraea/cost/FakeClusterInfo.java
+++ b/app/src/test/java/org/astraea/cost/FakeClusterInfo.java
@@ -12,6 +12,11 @@ public class FakeClusterInfo implements ClusterInfo {
     return List.of();
   }
 
+  @Override
+  public Set<String> dataDirectories(int brokerId) {
+    return Set.of();
+  }
+
   @Override
   public List<ReplicaInfo> availablePartitionLeaders(String topic) {
     return List.of();
diff --git a/app/src/test/java/org/astraea/cost/MemoryWarningCostTest.java b/app/src/test/java/org/astraea/cost/MemoryWarningCostTest.java
deleted file mode 100644
index 78cf5c16..00000000
--- a/app/src/test/java/org/astraea/cost/MemoryWarningCostTest.java
+++ /dev/null
@@ -1,67 +0,0 @@
-package org.astraea.cost;
-
-import java.lang.management.MemoryUsage;
-import java.util.Collection;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
-import org.astraea.metrics.HasBeanObject;
-import org.astraea.metrics.java.HasJvmMemory;
-import org.junit.jupiter.api.Assertions;
-import org.junit.jupiter.api.Test;
-import org.mockito.Mockito;
-
-public class MemoryWarningCostTest {
-  @Test
-  void testCost() throws InterruptedException {
-    var jvmMemory1 = mockResult(50L, 100L);
-    var jvmMemory2 = mockResult(81L, 100L);
-    var jvmMemory3 = mockResult(80L, 150L);
-
-    Collection<HasBeanObject> broker1 = List.of(jvmMemory1);
-    Collection<HasBeanObject> broker2 = List.of(jvmMemory2);
-    Collection<HasBeanObject> broker3 = List.of(jvmMemory3);
-    ClusterInfo clusterInfo =
-        new FakeClusterInfo() {
-          @Override
-          public Map<Integer, Collection<HasBeanObject>> allBeans() {
-            return Map.of(1, broker1, 2, broker2, 3, broker3);
-          }
-
-          @Override
-          public Set<String> topics() {
-            return Set.of("t");
-          }
-
-          @Override
-          public List<ReplicaInfo> availablePartitions(String topic) {
-            return List.of(
-                ReplicaInfo.of("t", 0, NodeInfo.of(1, "host1", 9092), true, true, false),
-                ReplicaInfo.of("t", 0, NodeInfo.of(2, "host2", 9092), false, true, false),
-                ReplicaInfo.of("t", 0, NodeInfo.of(3, "host3", 9092), false, true, false));
-          }
-        };
-
-    var memoryWarning = new MemoryWarningCost();
-    var scores = memoryWarning.brokerCost(clusterInfo).value();
-    Assertions.assertEquals(0.0, scores.get(1));
-    Assertions.assertEquals(0.0, scores.get(2));
-    Assertions.assertEquals(0.0, scores.get(3));
-
-    Thread.sleep(10000);
-
-    scores = memoryWarning.brokerCost(clusterInfo).value();
-    Assertions.assertEquals(0.0, scores.get(1));
-    Assertions.assertEquals(1.0, scores.get(2));
-    Assertions.assertEquals(0.0, scores.get(3));
-  }
-
-  private HasJvmMemory mockResult(long used, long max) {
-    var jvmMemory = Mockito.mock(HasJvmMemory.class);
-    var memoryUsage = Mockito.mock(MemoryUsage.class);
-    Mockito.when(jvmMemory.heapMemoryUsage()).thenReturn(memoryUsage);
-    Mockito.when(memoryUsage.getUsed()).thenReturn(used);
-    Mockito.when(memoryUsage.getMax()).thenReturn(max);
-    return jvmMemory;
-  }
-}
diff --git a/app/src/test/java/org/astraea/cost/PeriodicTest.java b/app/src/test/java/org/astraea/cost/PeriodicTest.java
index 603fbe28..f071e4db 100644
--- a/app/src/test/java/org/astraea/cost/PeriodicTest.java
+++ b/app/src/test/java/org/astraea/cost/PeriodicTest.java
@@ -1,5 +1,6 @@
 package org.astraea.cost;
 
+import java.time.Duration;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
@@ -10,13 +11,24 @@ public class PeriodicTest extends Periodic<Map<Integer, Double>> {
   private double brokerValue = 0.0;
 
   @Test
-  void testTryUpdate() {
-    var broker1 = tryUpdate(this::testMap);
+  void testTryUpdateAfterOneSecond() {
+    var broker1 = tryUpdateAfterOneSecond(this::testMap);
     Assertions.assertEquals(broker1.get(0), 0.0);
-    var broker2 = tryUpdate(this::testMap);
+    var broker2 = tryUpdateAfterOneSecond(this::testMap);
     Assertions.assertEquals(broker2.get(0), 0.0);
     sleep(1);
-    broker2 = tryUpdate(this::testMap);
+    broker2 = tryUpdateAfterOneSecond(this::testMap);
+    Assertions.assertEquals(broker2.get(0), 1.0);
+  }
+
+  @Test
+  void testTryUpdate() {
+    var broker1 = tryUpdate(this::testMap, Duration.ofSeconds(3));
+    Assertions.assertEquals(broker1.get(0), 0.0);
+    var broker2 = tryUpdate(this::testMap, Duration.ofSeconds(3));
+    Assertions.assertEquals(broker2.get(0), 0.0);
+    sleep(4);
+    broker2 = tryUpdate(this::testMap, Duration.ofSeconds(3));
     Assertions.assertEquals(broker2.get(0), 1.0);
   }
 
diff --git a/app/src/test/java/org/astraea/cost/broker/BrokerInputCostTest.java b/app/src/test/java/org/astraea/cost/broker/BrokerInputCostTest.java
index b1effd31..8302b4ae 100644
--- a/app/src/test/java/org/astraea/cost/broker/BrokerInputCostTest.java
+++ b/app/src/test/java/org/astraea/cost/broker/BrokerInputCostTest.java
@@ -15,7 +15,7 @@ import org.mockito.Mockito;
 
 public class BrokerInputCostTest {
   @Test
-  void testCost() throws InterruptedException {
+  void testCost() {
     ClusterInfo clusterInfo = exampleClusterInfo(10000L, 20000L, 5000L);
 
     var brokerInputCost = new BrokerInputCost();
@@ -23,8 +23,6 @@ public class BrokerInputCostTest {
     Assertions.assertEquals(0.47, scores.get(1));
     Assertions.assertEquals(0.63, scores.get(2));
     Assertions.assertEquals(0.39, scores.get(3));
-    var scores5 = brokerInputCost.brokerCost(clusterInfo).value();
-    Thread.sleep(1000);
 
     ClusterInfo clusterInfo2 = exampleClusterInfo(55555L, 25352L, 25000L);
     scores = brokerInputCost.brokerCost(clusterInfo2).value();
diff --git a/app/src/test/java/org/astraea/cost/broker/BrokerOutPutCostTest.java b/app/src/test/java/org/astraea/cost/broker/BrokerOutPutCostTest.java
index a8ca9a15..7f666090 100644
--- a/app/src/test/java/org/astraea/cost/broker/BrokerOutPutCostTest.java
+++ b/app/src/test/java/org/astraea/cost/broker/BrokerOutPutCostTest.java
@@ -24,8 +24,6 @@ public class BrokerOutPutCostTest {
     Assertions.assertEquals(0.63, scores.get(2));
     Assertions.assertEquals(0.39, scores.get(3));
 
-    Thread.sleep(1000);
-
     ClusterInfo clusterInfo2 = exampleClusterInfo(55555L, 25352L, 25000L);
     scores = brokerOutputCost.brokerCost(clusterInfo2).value();
     Assertions.assertEquals(0.55, scores.get(1));
diff --git a/app/src/test/java/org/astraea/cost/broker/MemoryCostTest.java b/app/src/test/java/org/astraea/cost/broker/MemoryCostTest.java
index 95760900..1415a87c 100644
--- a/app/src/test/java/org/astraea/cost/broker/MemoryCostTest.java
+++ b/app/src/test/java/org/astraea/cost/broker/MemoryCostTest.java
@@ -18,7 +18,7 @@ import org.mockito.Mockito;
 public class MemoryCostTest {
   @Test
   void testCost() throws InterruptedException {
-    var jvmMemory1 = mockResult(50L, 100L);
+    var jvmMemory1 = mockResult(40L, 100L);
     var jvmMemory2 = mockResult(81L, 100L);
     var jvmMemory3 = mockResult(80L, 150L);
 
@@ -48,12 +48,12 @@ public class MemoryCostTest {
 
     var memoryCost = new MemoryCost();
     var scores = memoryCost.brokerCost(clusterInfo).value();
-    Assertions.assertEquals(0.42, scores.get(1));
-    Assertions.assertEquals(0.64, scores.get(2));
-    Assertions.assertEquals(0.44, scores.get(3));
+    Assertions.assertEquals(0.39, scores.get(1));
+    Assertions.assertEquals(0.63, scores.get(2));
+    Assertions.assertEquals(0.47, scores.get(3));
 
     Thread.sleep(1000);
-    jvmMemory1 = mockResult(50L, 50L);
+    jvmMemory1 = mockResult(50L, 150L);
     jvmMemory2 = mockResult(30L, 100L);
     jvmMemory3 = mockResult(80L, 150L);
 
@@ -81,9 +81,9 @@ public class MemoryCostTest {
           }
         };
     scores = memoryCost.brokerCost(clusterInfo2).value();
-    Assertions.assertEquals(0.53, scores.get(1));
+    Assertions.assertEquals(0.42, scores.get(1));
     Assertions.assertEquals(0.52, scores.get(2));
-    Assertions.assertEquals(0.45, scores.get(3));
+    Assertions.assertEquals(0.55, scores.get(3));
   }
 
   private HasJvmMemory mockResult(long used, long max) {
diff --git a/app/src/test/java/org/astraea/cost/broker/IntegratedScoreTest.java b/app/src/test/java/org/astraea/cost/broker/NeutralIntegratedCostTest.java
similarity index 95%
rename from app/src/test/java/org/astraea/cost/broker/IntegratedScoreTest.java
rename to app/src/test/java/org/astraea/cost/broker/NeutralIntegratedCostTest.java
index 079958bb..8145ed23 100644
--- a/app/src/test/java/org/astraea/cost/broker/IntegratedScoreTest.java
+++ b/app/src/test/java/org/astraea/cost/broker/NeutralIntegratedCostTest.java
@@ -4,7 +4,8 @@ import java.util.Map;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 
-public class IntegratedScoreTest {
+// TODO Add more test.
+public class NeutralIntegratedCostTest {
   @Test
   void testEntropyEmpowerment() {
     var neutralIntegratedCost = new NeutralIntegratedCost();
diff --git a/app/src/test/java/org/astraea/partitioner/smooth/SmoothWeightRoundRobinDispatchTest.java b/app/src/test/java/org/astraea/partitioner/smooth/SmoothWeightRoundRobinDispatchTest.java
new file mode 100644
index 00000000..c0e84377
--- /dev/null
+++ b/app/src/test/java/org/astraea/partitioner/smooth/SmoothWeightRoundRobinDispatchTest.java
@@ -0,0 +1,286 @@
+package org.astraea.partitioner.smooth;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import java.io.File;
+import java.io.FileWriter;
+import java.io.IOException;
+import java.time.Duration;
+import java.util.List;
+import java.util.Map;
+import java.util.Properties;
+import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+import org.apache.kafka.clients.producer.ProducerConfig;
+import org.apache.kafka.common.serialization.ByteArraySerializer;
+import org.astraea.admin.Admin;
+import org.astraea.concurrent.Executor;
+import org.astraea.concurrent.State;
+import org.astraea.concurrent.ThreadPool;
+import org.astraea.consumer.Consumer;
+import org.astraea.consumer.Deserializer;
+import org.astraea.consumer.Header;
+import org.astraea.producer.Producer;
+import org.astraea.producer.Serializer;
+import org.astraea.service.RequireBrokerCluster;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+public class SmoothWeightRoundRobinDispatchTest extends RequireBrokerCluster {
+  private final String brokerList = bootstrapServers();
+  Admin admin = Admin.of(bootstrapServers());
+
+  private Properties initProConfig() {
+    Properties props = new Properties();
+    props.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList);
+    props.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, ByteArraySerializer.class.getName());
+    props.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, ByteArraySerializer.class.getName());
+    props.put(ProducerConfig.CLIENT_ID_CONFIG, "id1");
+    props.put(
+        ProducerConfig.PARTITIONER_CLASS_CONFIG, SmoothWeightRoundRobinDispatcher.class.getName());
+    props.put("producerID", 1);
+    var file =
+        new File(
+            SmoothWeightRoundRobinDispatchTest.class.getResource("").getPath()
+                + "PartitionerConfigTest");
+    try {
+      var fileWriter = new FileWriter(file);
+      fileWriter.write("jmx.port=" + jmxServiceURL().getPort() + "\n");
+      fileWriter.write("broker.0.jmx.port=" + jmxServiceURL().getPort() + "\n");
+      fileWriter.write("broker.1.jmx.port=" + jmxServiceURL().getPort() + "\n");
+      fileWriter.write("broker.2.jmx.port=" + jmxServiceURL().getPort());
+      fileWriter.flush();
+      fileWriter.close();
+    } catch (IOException e) {
+      e.printStackTrace();
+    }
+    props.put(
+        "partitioner.config",
+        SmoothWeightRoundRobinDispatchTest.class.getResource("").getPath()
+            + "PartitionerConfigTest");
+    return props;
+  }
+
+  @Test
+  void testPartitioner() {
+    var topicName = "address";
+    admin.creator().topic(topicName).numberOfPartitions(10).create();
+    var key = "tainan";
+    var timestamp = System.currentTimeMillis() + 10;
+    var header = Header.of("a", "b".getBytes());
+    try (var producer =
+        Producer.builder()
+            .keySerializer(Serializer.STRING)
+            .configs(
+                initProConfig().entrySet().stream()
+                    .collect(
+                        Collectors.toMap(e -> e.getKey().toString(), e -> e.getValue().toString())))
+            .build()) {
+      var i = 0;
+      while (i < 300) {
+        var metadata =
+            producer
+                .sender()
+                .topic(topicName)
+                .key(key)
+                .timestamp(timestamp)
+                .headers(List.of(header))
+                .run()
+                .toCompletableFuture()
+                .get();
+        assertEquals(topicName, metadata.topic());
+        assertEquals(timestamp, metadata.timestamp());
+        i++;
+      }
+    } catch (ExecutionException | InterruptedException e) {
+      e.printStackTrace();
+    }
+    sleep(1);
+    try (var consumer =
+        Consumer.builder()
+            .bootstrapServers(bootstrapServers())
+            .fromBeginning()
+            .topics(Set.of(topicName))
+            .keyDeserializer(Deserializer.STRING)
+            .build()) {
+      var records = consumer.poll(Duration.ofSeconds(20));
+      var recordsCount = records.size();
+      while (recordsCount < 300) {
+        recordsCount += consumer.poll(Duration.ofSeconds(20)).size();
+      }
+      assertEquals(300, recordsCount);
+      var record = records.iterator().next();
+      assertEquals(topicName, record.topic());
+      assertEquals("tainan", record.key());
+      assertEquals(1, record.headers().size());
+      var actualHeader = record.headers().iterator().next();
+      assertEquals(header.key(), actualHeader.key());
+      Assertions.assertArrayEquals(header.value(), actualHeader.value());
+    }
+  }
+
+  @Test
+  void testMultipleProducer() {
+    var topicName = "addr";
+    admin.creator().topic(topicName).numberOfPartitions(10).create();
+    var key = "tainan";
+    var timestamp = System.currentTimeMillis() + 10;
+    var header = Header.of("a", "b".getBytes());
+    try (var threadPool =
+        ThreadPool.builder()
+            .executors(
+                IntStream.range(0, 10)
+                    .mapToObj(
+                        i ->
+                            producerExecutor(
+                                Producer.builder()
+                                    .keySerializer(Serializer.STRING)
+                                    .configs(
+                                        initProConfig().entrySet().stream()
+                                            .collect(
+                                                Collectors.toMap(
+                                                    e -> e.getKey().toString(),
+                                                    e -> e.getValue().toString())))
+                                    .build(),
+                                topicName,
+                                key,
+                                header,
+                                timestamp))
+                    .collect(Collectors.toUnmodifiableList()))
+            .build()) {
+      threadPool.waitAll();
+    }
+    try (var consumer =
+        Consumer.builder()
+            .bootstrapServers(bootstrapServers())
+            .fromBeginning()
+            .topics(Set.of(topicName))
+            .keyDeserializer(Deserializer.STRING)
+            .build()) {
+      var records = consumer.poll(Duration.ofSeconds(20));
+      var recordsCount = records.size();
+      while (recordsCount < 1000) {
+        recordsCount += consumer.poll(Duration.ofSeconds(20)).size();
+      }
+      assertEquals(1000, recordsCount);
+      var record = records.iterator().next();
+      assertEquals(topicName, record.topic());
+      assertEquals("tainan", record.key());
+      assertEquals(1, record.headers().size());
+      var actualHeader = record.headers().iterator().next();
+      assertEquals(header.key(), actualHeader.key());
+      Assertions.assertArrayEquals(header.value(), actualHeader.value());
+    }
+  }
+
+  @Test
+  void testJmxConfig() {
+    var props = initProConfig();
+    var file =
+        new File(
+            SmoothWeightRoundRobinDispatchTest.class.getResource("").getPath()
+                + "PartitionerConfigTest");
+    try {
+      var fileWriter = new FileWriter(file);
+      fileWriter.write(
+          "broker.0.jmx.port="
+              + jmxServiceURL().getPort()
+              + "\n"
+              + "broker.1.jmx.port="
+              + jmxServiceURL().getPort()
+              + "\n"
+              + "broker.2.jmx.port="
+              + jmxServiceURL().getPort()
+              + "\n");
+      fileWriter.flush();
+      fileWriter.close();
+    } catch (IOException e) {
+      e.printStackTrace();
+    }
+    var topicName = "addressN";
+    admin.creator().topic(topicName).numberOfPartitions(10).create();
+    var key = "tainan";
+    var timestamp = System.currentTimeMillis() + 10;
+    var header = Header.of("a", "b".getBytes());
+
+    try (var producer =
+        Producer.builder()
+            .keySerializer(Serializer.STRING)
+            .configs(
+                props.entrySet().stream()
+                    .collect(
+                        Collectors.toMap(e -> e.getKey().toString(), e -> e.getValue().toString())))
+            .build()) {
+      var metadata =
+          producer
+              .sender()
+              .topic(topicName)
+              .key(key)
+              .timestamp(timestamp)
+              .headers(List.of(header))
+              .run()
+              .toCompletableFuture()
+              .get();
+      assertEquals(topicName, metadata.topic());
+      assertEquals(timestamp, metadata.timestamp());
+    } catch (ExecutionException | InterruptedException e) {
+      e.printStackTrace();
+    }
+  }
+
+  private Executor producerExecutor(
+      Producer<String, byte[]> producer, String topic, String key, Header header, long timeStamp) {
+    return new Executor() {
+      int i = 0;
+
+      @Override
+      public State execute() throws InterruptedException {
+        if (i > 99) return State.DONE;
+        try {
+          producer
+              .sender()
+              .topic(topic)
+              .key(key)
+              .timestamp(timeStamp)
+              .headers(List.of(header))
+              .run()
+              .toCompletableFuture()
+              .get();
+        } catch (ExecutionException e) {
+          e.printStackTrace();
+        }
+        i++;
+        return State.RUNNING;
+      }
+
+      @Override
+      public void close() {
+        producer.close();
+      }
+    };
+  }
+
+  private static void sleep(int seconds) {
+    try {
+      TimeUnit.SECONDS.sleep(seconds);
+    } catch (InterruptedException e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  @Test
+  public void testGetAndChoose() {
+    var smoothWeight = new SmoothWeightRoundRobin(Map.of(1, 5.0, 2, 3.0, 3, 1.0));
+
+    Assertions.assertEquals(1, smoothWeight.getAndChoose());
+    Assertions.assertEquals(2, smoothWeight.getAndChoose());
+    Assertions.assertEquals(3, smoothWeight.getAndChoose());
+    Assertions.assertEquals(1, smoothWeight.getAndChoose());
+    Assertions.assertEquals(2, smoothWeight.getAndChoose());
+    Assertions.assertEquals(3, smoothWeight.getAndChoose());
+    Assertions.assertEquals(1, smoothWeight.getAndChoose());
+  }
+}
diff --git a/app/src/test/java/org/astraea/partitioner/smooth/SmoothWeightTest.java b/app/src/test/java/org/astraea/partitioner/smooth/SmoothWeightTest.java
deleted file mode 100644
index 6f853448..00000000
--- a/app/src/test/java/org/astraea/partitioner/smooth/SmoothWeightTest.java
+++ /dev/null
@@ -1,20 +0,0 @@
-package org.astraea.partitioner.smooth;
-
-import java.util.Map;
-import org.junit.jupiter.api.Assertions;
-import org.junit.jupiter.api.Test;
-
-public class SmoothWeightTest {
-  @Test
-  public void testGetAndChoose() {
-    var smoothWeight = new SmoothWeight(Map.of(1, 5, 2, 3, 3, 1));
-
-    Assertions.assertEquals(1, smoothWeight.getAndChoose());
-    Assertions.assertEquals(2, smoothWeight.getAndChoose());
-    Assertions.assertEquals(1, smoothWeight.getAndChoose());
-    Assertions.assertEquals(1, smoothWeight.getAndChoose());
-    Assertions.assertEquals(2, smoothWeight.getAndChoose());
-    Assertions.assertEquals(1, smoothWeight.getAndChoose());
-    Assertions.assertEquals(3, smoothWeight.getAndChoose());
-  }
-}