Repository: skiptests/astraea
Pull request #43 (2021-10-07 15:49:21)

ReplicaCollie - used to migrate replicas between brokers
這個工具可以方便地將某些replicas從指定的brokers中移到其他節點，相比kafka官方工具需要去寫json檔案，這隻工具可以快速實現“隔離”某些brokers的需求

chia7712: @wycccccc @qoo332001 這也是我這邊一個業務需求。本來應該是offline balancer的一部分，不過我先把業務急需的功能抽出來做成一隻tool

Patch:
diff --git a/README.md b/README.md
index 3e347a50..7f06bd23 100644
--- a/README.md
+++ b/README.md
@@ -16,6 +16,7 @@ This project offers many kafka tools to simplify the life for kafka users.
 2. [Kafka benchmark](#latency-benchmark): run producers/consumers to test the performance and consistency for kafka cluster
 3. [Kafka offset explorer](#offset-explorer): check the start/end offsets of kafka topics
 4. [Kafka official tool](#kafka-official-tool): run any one specific kafka official tool. All you have to prepare is the docker env.
+4. [Replica Collie](#replica-collie): move replicas from brokers to others. You can use this tool to obstruct specific brokers from hosting specific topics. 
 
 [Release page](https://github.com/skiptests/astraea/releases) offers the uber jar including all tools.
 ```shell
@@ -73,12 +74,12 @@ This tool is used to test following latencies.
 
 Run the benchmark from source code
 ```shell
-./gradlew run --args="Latency --bootstrap.servers 192.168.50.224:18878"
+./gradlew run --args="latency --bootstrap.servers 192.168.50.224:18878"
 ```
 
 Run the benchmark from release
 ```shell
-java -jar app-0.0.1-SNAPSHOT-all.jar Latency --bootstrap.servers 192.168.50.224:18878
+java -jar app-0.0.1-SNAPSHOT-all.jar latency --bootstrap.servers 192.168.50.224:18878
 ```
 
 ### Latency Benchmark Configurations
@@ -125,4 +126,22 @@ This project offers a way to run kafka official tool by container. For example:
 
 ```shell
 ./docker/start_kafka_tool.sh help
+```
+
+---
+
+## Replica Collie
+
+This tool offers an effective way to migrate all replicas from specific brokers to others.
+
+### Move all replicas from broker_0 and broker_1 to other brokers
+
+```shell
+./gradlew run --args="replica --bootstrap.servers 192.168.50.178:19993 --from 0,1"
+```
+
+### Move all replicas of topic "abc" from broker_0 to broker_1
+
+```shell
+./gradlew run --args="replica --bootstrap.servers 192.168.50.178:19993 --from 0 --to 1 --topics abc"
 ```
\ No newline at end of file
diff --git a/app/build.gradle b/app/build.gradle
index 87413e81..8ef701fe 100644
--- a/app/build.gradle
+++ b/app/build.gradle
@@ -21,6 +21,7 @@ def versions = [
     kafka: project.properties['kafka.version'] ?: "2.8.0",
     junit: project.properties['junit.version'] ?: "5.7.2",
     jcommander: project.properties['jcommander.version'] ?: "1.81",
+    slf4j: project.properties['jcommander.version'] ?: "1.7.32",
 ]
 
 dependencies {
@@ -28,6 +29,8 @@ dependencies {
 
     implementation "org.apache.kafka:kafka-clients:${versions["kafka"]}"
     implementation "com.beust:jcommander:${versions["jcommander"]}"
+    // we don't use slf4j actually, and it is used by kafka so we swallow the log.
+    implementation "org.slf4j:slf4j-nop:${versions["slf4j"]}"
 }
 
 application {
diff --git a/app/src/main/java/org/astraea/App.java b/app/src/main/java/org/astraea/App.java
index 46b121f6..c962aea4 100644
--- a/app/src/main/java/org/astraea/App.java
+++ b/app/src/main/java/org/astraea/App.java
@@ -4,40 +4,37 @@ import com.beust.jcommander.ParameterException;
 import java.lang.reflect.InvocationTargetException;
 import java.util.Arrays;
 import java.util.List;
-import java.util.function.Supplier;
+import java.util.Map;
 import java.util.stream.Collectors;
 import org.astraea.metrics.kafka.KafkaMetricClientApp;
 import org.astraea.offset.OffsetExplorer;
 import org.astraea.performance.latency.End2EndLatency;
+import org.astraea.topic.ReplicaCollie;
 
 public class App {
-  private static final List<Class<?>> MAIN_CLASSES =
-      Arrays.asList(End2EndLatency.class, OffsetExplorer.class, KafkaMetricClientApp.class);
+  private static final Map<String, Class<?>> MAIN_CLASSES =
+      Map.of(
+          "latency", End2EndLatency.class,
+          "offset", OffsetExplorer.class,
+          "metrics", KafkaMetricClientApp.class,
+          "replica", ReplicaCollie.class);
 
   private static String toString(List<Class<?>> mains) {
     return mains.stream().map(Class::getSimpleName).collect(Collectors.joining(","));
   }
 
-  private static String usage(List<Class<?>> mains) {
-    return "Usage: [" + toString(mains) + "] [args ...]";
-  }
+  static void execute(Map<String, Class<?>> mains, List<String> args) throws Throwable {
+
+    var usage = "Usage: " + mains.keySet() + " [args ...]";
 
-  static void execute(List<Class<?>> mains, List<String> args) throws Throwable {
     if (args.size() < 1) {
-      throw new IllegalArgumentException(usage(mains));
+      throw new IllegalArgumentException(usage);
     }
 
     var className = args.get(0);
 
-    var targetClass =
-        mains.stream()
-            .filter(clz -> clz.getName().contains(className))
-            .findFirst()
-            .orElseThrow(
-                (Supplier<Throwable>)
-                    () ->
-                        new IllegalArgumentException(
-                            "className: " + className + " is not matched to " + toString(mains)));
+    var targetClass = mains.get(className);
+    if (targetClass == null) throw new IllegalArgumentException(usage);
 
     var method = targetClass.getDeclaredMethod("main", String[].class);
     try {
diff --git a/app/src/main/java/org/astraea/argument/ArgumentUtil.java b/app/src/main/java/org/astraea/argument/ArgumentUtil.java
index 55fb57b0..947db752 100644
--- a/app/src/main/java/org/astraea/argument/ArgumentUtil.java
+++ b/app/src/main/java/org/astraea/argument/ArgumentUtil.java
@@ -6,7 +6,10 @@ import com.beust.jcommander.JCommander;
 import com.beust.jcommander.ParameterException;
 import com.beust.jcommander.UnixStyleUsageFormatter;
 import java.time.Duration;
+import java.util.Collection;
 import java.util.Set;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 /*
  * A tool used to parse command line arguments.
@@ -71,10 +74,30 @@ public class ArgumentUtil {
     }
   }
 
-  public static class SetConverter implements IStringConverter<Set<String>> {
+  public static class StringSetConverter implements IStringConverter<Set<String>> {
     @Override
     public Set<String> convert(String value) {
-      return Set.of(value.split(","));
+      return requireNonEmpty(Set.of(value.split(",")));
     }
   }
+
+  public static class IntegerSetConverter implements IStringConverter<Set<Integer>> {
+    @Override
+    public Set<Integer> convert(String value) {
+      return requireNonEmpty(
+          Stream.of(value.split(",")).map(Integer::valueOf).collect(Collectors.toSet()));
+    }
+  }
+
+  public static class BooleanConverter implements IStringConverter<Boolean> {
+    @Override
+    public Boolean convert(String value) {
+      return Boolean.valueOf(value);
+    }
+  }
+
+  private static <C extends Collection<T>, T> C requireNonEmpty(C collection) {
+    if (collection.isEmpty()) throw new ParameterException("array type can't be empty");
+    return collection;
+  }
 }
diff --git a/app/src/main/java/org/astraea/metrics/kafka/KafkaMetricClientAppArgument.java b/app/src/main/java/org/astraea/metrics/kafka/KafkaMetricClientAppArgument.java
index 491cd90f..198c1e13 100644
--- a/app/src/main/java/org/astraea/metrics/kafka/KafkaMetricClientAppArgument.java
+++ b/app/src/main/java/org/astraea/metrics/kafka/KafkaMetricClientAppArgument.java
@@ -17,7 +17,7 @@ public class KafkaMetricClientAppArgument {
   @Parameter(
       names = {"--metrics"},
       variableArity = true,
-      converter = ArgumentUtil.SetConverter.class,
+      converter = ArgumentUtil.StringSetConverter.class,
       description =
           "The metric names you want. If no metric name specified in argument, all metrics will be selected.")
   Set<String> metrics = BrokerTopicMetrics.AllMetricNames;
diff --git a/app/src/main/java/org/astraea/offset/OffsetExplorerArgument.java b/app/src/main/java/org/astraea/offset/OffsetExplorerArgument.java
index f3bca996..2a1eb9d9 100644
--- a/app/src/main/java/org/astraea/offset/OffsetExplorerArgument.java
+++ b/app/src/main/java/org/astraea/offset/OffsetExplorerArgument.java
@@ -17,6 +17,6 @@ public class OffsetExplorerArgument {
       names = {"--topics"},
       description = "String: topic names, comma separate",
       validateWith = ArgumentUtil.NotEmptyString.class,
-      converter = ArgumentUtil.SetConverter.class)
+      converter = ArgumentUtil.StringSetConverter.class)
   Set<String> topics = Collections.emptySet();
 }
diff --git a/app/src/main/java/org/astraea/topic/ReplicaCollie.java b/app/src/main/java/org/astraea/topic/ReplicaCollie.java
new file mode 100644
index 00000000..e70e25ac
--- /dev/null
+++ b/app/src/main/java/org/astraea/topic/ReplicaCollie.java
@@ -0,0 +1,121 @@
+package org.astraea.topic;
+
+import com.beust.jcommander.Parameter;
+import java.io.IOException;
+import java.util.*;
+import java.util.stream.Collectors;
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.common.TopicPartition;
+import org.astraea.argument.ArgumentUtil;
+
+public class ReplicaCollie {
+
+  static Map<TopicPartition, Map.Entry<Set<Integer>, Set<Integer>>> execute(
+      TopicAdmin admin, Argument args) {
+    var topics = args.topics.isEmpty() ? admin.topics() : args.topics;
+    var allBrokers = admin.brokerIds();
+    if (!args.toBrokers.isEmpty() && !allBrokers.containsAll(args.toBrokers))
+      throw new IllegalArgumentException(
+          "those brokers: "
+              + args.toBrokers.stream()
+                  .filter(i -> !allBrokers.contains(i))
+                  .map(String::valueOf)
+                  .collect(Collectors.joining(","))
+              + " are nonexistent");
+
+    var targetBrokers = args.toBrokers.isEmpty() ? allBrokers : args.toBrokers;
+    var result = new HashMap<TopicPartition, Map.Entry<Set<Integer>, Set<Integer>>>();
+    admin
+        .replicas(topics)
+        .forEach(
+            (tp, replicas) -> {
+              var currentBrokers = replicas.stream().map(r -> r.broker).collect(Collectors.toSet());
+              var keptBrokers =
+                  currentBrokers.stream()
+                      .filter(i -> !args.fromBrokers.contains(i))
+                      .collect(Collectors.toSet());
+              var numberOfMigratedReplicas = currentBrokers.size() - keptBrokers.size();
+              if (numberOfMigratedReplicas > 0) {
+                var availableBrokers =
+                    targetBrokers.stream()
+                        .filter(i -> !keptBrokers.contains(i) && !args.fromBrokers.contains(i))
+                        .collect(Collectors.toSet());
+                if (availableBrokers.size() < numberOfMigratedReplicas)
+                  throw new IllegalArgumentException(
+                      "No enough available brokers! Available: "
+                          + targetBrokers
+                          + " current: "
+                          + currentBrokers
+                          + " removed: "
+                          + args.fromBrokers);
+                var finalBrokers = new HashSet<>(keptBrokers);
+                finalBrokers.addAll(
+                    new ArrayList<>(availableBrokers).subList(0, numberOfMigratedReplicas));
+                result.put(tp, Map.entry(currentBrokers, finalBrokers));
+              }
+            });
+    result.forEach(
+        (tp, assignments) -> {
+          if (!args.verify) admin.reassign(tp.topic(), tp.partition(), assignments.getValue());
+        });
+    return result;
+  }
+
+  public static void main(String[] args) throws IOException {
+    var arguments = ArgumentUtil.parseArgument(new Argument(), args);
+    try (var admin =
+        TopicAdmin.of(Map.of(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, arguments.brokers))) {
+      execute(admin, arguments)
+          .forEach(
+              (tp, assignments) ->
+                  System.out.println(
+                      "topic: "
+                          + tp.topic()
+                          + ", partition: "
+                          + tp.partition()
+                          + " before: "
+                          + assignments.getKey()
+                          + " after: "
+                          + assignments.getValue()));
+    }
+  }
+
+  static class Argument {
+    @Parameter(
+        names = {"--bootstrap.servers"},
+        description = "String: server to connect to",
+        validateWith = ArgumentUtil.NotEmptyString.class,
+        required = true)
+    String brokers;
+
+    @Parameter(
+        names = {"--topics"},
+        description = "Those topics' partitions will get reassigned",
+        validateWith = ArgumentUtil.NotEmptyString.class,
+        converter = ArgumentUtil.StringSetConverter.class)
+    Set<String> topics = Collections.emptySet();
+
+    @Parameter(
+        names = {"--from"},
+        description = "Those brokers won't hold any replicas of topics (defined by --topics)",
+        validateWith = ArgumentUtil.NotEmptyString.class,
+        converter = ArgumentUtil.IntegerSetConverter.class,
+        required = true)
+    Set<Integer> fromBrokers;
+
+    @Parameter(
+        names = {"--to"},
+        description = "The replicas of topics (defined by --topic) will be moved to those brokers",
+        validateWith = ArgumentUtil.NotEmptyString.class,
+        converter = ArgumentUtil.IntegerSetConverter.class)
+    Set<Integer> toBrokers = Collections.emptySet();
+
+    @Parameter(
+        names = {"--verify"},
+        description =
+            "True if you just want to see the new assignment instead of executing the plan",
+        validateWith = ArgumentUtil.NotEmptyString.class,
+        converter = ArgumentUtil.BooleanConverter.class)
+    boolean verify = false;
+  }
+}
diff --git a/app/src/main/java/org/astraea/topic/TopicAdmin.java b/app/src/main/java/org/astraea/topic/TopicAdmin.java
index 5b274883..9be69332 100644
--- a/app/src/main/java/org/astraea/topic/TopicAdmin.java
+++ b/app/src/main/java/org/astraea/topic/TopicAdmin.java
@@ -3,10 +3,7 @@ package org.astraea.topic;
 import java.io.Closeable;
 import java.util.*;
 import java.util.stream.Collectors;
-import org.apache.kafka.clients.admin.Admin;
-import org.apache.kafka.clients.admin.ConsumerGroupListing;
-import org.apache.kafka.clients.admin.OffsetSpec;
-import org.apache.kafka.clients.admin.ReplicaInfo;
+import org.apache.kafka.clients.admin.*;
 import org.apache.kafka.common.Node;
 import org.apache.kafka.common.TopicPartition;
 import org.astraea.Utils;
@@ -22,6 +19,28 @@ public interface TopicAdmin extends Closeable {
         admin.close();
       }
 
+      @Override
+      public Set<Integer> brokerIds() {
+        return Utils.handleException(
+            () ->
+                admin.describeCluster().nodes().get().stream()
+                    .map(Node::id)
+                    .collect(Collectors.toSet()));
+      }
+
+      @Override
+      public void reassign(String topicName, int partition, Set<Integer> brokers) {
+        Utils.handleException(
+            () ->
+                admin
+                    .alterPartitionReassignments(
+                        Map.of(
+                            new TopicPartition(topicName, partition),
+                            Optional.of(new NewPartitionReassignment(new ArrayList<>(brokers)))))
+                    .all()
+                    .get());
+      }
+
       @Override
       public Map<TopicPartition, List<Group>> groups(Set<String> topics) {
         return Utils.handleException(
@@ -114,17 +133,10 @@ public interface TopicAdmin extends Closeable {
 
       @Override
       public Map<TopicPartition, List<Replica>> replicas(Set<String> topics) {
-        var brokers =
-            Utils.handleException(
-                () ->
-                    admin.describeCluster().nodes().get().stream()
-                        .map(Node::id)
-                        .collect(Collectors.toSet()));
-
         var lags =
             Utils.handleException(
                 () ->
-                    admin.describeLogDirs(brokers).allDescriptions().get().entrySet().stream()
+                    admin.describeLogDirs(brokerIds()).allDescriptions().get().entrySet().stream()
                         .collect(
                             Collectors.toMap(
                                 Map.Entry::getKey,
@@ -200,6 +212,16 @@ public interface TopicAdmin extends Closeable {
    */
   Map<TopicPartition, List<Replica>> replicas(Set<String> topics);
 
+  /** @return all brokers' ids */
+  Set<Integer> brokerIds();
+
+  /**
+   * @param topicName topic name
+   * @param partition partition
+   * @param brokers to hold all the
+   */
+  void reassign(String topicName, int partition, Set<Integer> brokers);
+
   class Group {
     public final String id;
     public final long offset;


Test patch:
diff --git a/app/src/test/java/org/astraea/AppTest.java b/app/src/test/java/org/astraea/AppTest.java
index 18df3aab..95ad5977 100644
--- a/app/src/test/java/org/astraea/AppTest.java
+++ b/app/src/test/java/org/astraea/AppTest.java
@@ -3,7 +3,7 @@ package org.astraea;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 
 import java.util.Arrays;
-import java.util.Collections;
+import java.util.Map;
 import org.junit.jupiter.api.Test;
 
 class AppTest {
@@ -15,7 +15,7 @@ class AppTest {
 
   @Test
   void testExecute() throws Throwable {
-    App.execute(Collections.singletonList(AppTest.class), Arrays.asList("AppTest", "1", "2"));
+    App.execute(Map.of("abc", AppTest.class), Arrays.asList("abc", "1", "2"));
     assertEquals(2, ARGS.length);
     assertEquals("1", ARGS[0]);
     assertEquals("2", ARGS[1]);
diff --git a/app/src/test/java/org/astraea/argument/ArgumentUtilTest.java b/app/src/test/java/org/astraea/argument/ArgumentUtilTest.java
index 94bf4809..6b53c4ff 100644
--- a/app/src/test/java/org/astraea/argument/ArgumentUtilTest.java
+++ b/app/src/test/java/org/astraea/argument/ArgumentUtilTest.java
@@ -32,7 +32,7 @@ public class ArgumentUtilTest {
 
     @Parameter(
         names = {"--setConverter"},
-        converter = ArgumentUtil.SetConverter.class,
+        converter = ArgumentUtil.StringSetConverter.class,
         variableArity = true)
     public Set<String> setConverter;
   }
diff --git a/app/src/test/java/org/astraea/offset/OffsetExplorerTest.java b/app/src/test/java/org/astraea/offset/OffsetExplorerTest.java
index 4d7d0680..ef9e8703 100644
--- a/app/src/test/java/org/astraea/offset/OffsetExplorerTest.java
+++ b/app/src/test/java/org/astraea/offset/OffsetExplorerTest.java
@@ -28,7 +28,7 @@ public class OffsetExplorerTest {
         new TopicAdmin() {
           @Override
           public Set<String> topics() {
-            return null;
+            throw new UnsupportedOperationException();
           }
 
           @Override
@@ -46,6 +46,16 @@ public class OffsetExplorerTest {
             return Map.of(topicPartition, List.of(new Replica(brokerId, lag, leader, inSync)));
           }
 
+          @Override
+          public Set<Integer> brokerIds() {
+            throw new UnsupportedOperationException();
+          }
+
+          @Override
+          public void reassign(String topicName, int partition, Set<Integer> brokers) {
+            throw new UnsupportedOperationException();
+          }
+
           @Override
           public void close() {}
         }) {
diff --git a/app/src/test/java/org/astraea/topic/ReplicaCollieTest.java b/app/src/test/java/org/astraea/topic/ReplicaCollieTest.java
new file mode 100644
index 00000000..6da333b0
--- /dev/null
+++ b/app/src/test/java/org/astraea/topic/ReplicaCollieTest.java
@@ -0,0 +1,90 @@
+package org.astraea.topic;
+
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import org.apache.kafka.common.TopicPartition;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+public class ReplicaCollieTest {
+
+  private final String topicName = "ReplicaCollieTest";
+  private final int partition = 0;
+  private final int badBroker = 10;
+  private final int goodBroker = 100;
+  private final int earliestOffset = 199;
+  private final int latestOffset = 299;
+  private final Set<Integer> reassignment = new HashSet<>();
+  private final TopicAdmin admin =
+      new TopicAdmin() {
+        @Override
+        public Set<String> topics() {
+          return Set.of(topicName);
+        }
+
+        @Override
+        public Map<TopicPartition, Offset> offset(Set<String> topics) {
+          throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public Map<TopicPartition, List<Group>> groups(Set<String> topics) {
+          throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public Map<TopicPartition, List<Replica>> replicas(Set<String> topics) {
+          return Map.of(
+              new TopicPartition(topicName, partition),
+              List.of(new Replica(badBroker, 0, true, true)));
+        }
+
+        @Override
+        public Set<Integer> brokerIds() {
+          return Set.of(goodBroker, badBroker);
+        }
+
+        @Override
+        public void reassign(String topicName, int partition, Set<Integer> brokers) {
+          reassignment.addAll(brokers);
+        }
+
+        @Override
+        public void close() {
+          throw new UnsupportedOperationException();
+        }
+      };
+
+  @Test
+  void testVerify() {
+    test(true);
+  }
+
+  @Test
+  void testExecute() {
+    test(false);
+  }
+
+  private void test(boolean verify) {
+    var argument = new ReplicaCollie.Argument();
+    argument.fromBrokers = Set.of(badBroker);
+    argument.toBrokers = Set.of();
+    argument.brokers = "just unit test";
+    argument.verify = verify;
+    var result = ReplicaCollie.execute(admin, argument);
+    Assertions.assertEquals(1, result.size());
+    var assignment = result.get(new TopicPartition(topicName, partition));
+    Assertions.assertEquals(1, assignment.getKey().size());
+    Assertions.assertEquals(badBroker, assignment.getKey().iterator().next());
+    Assertions.assertEquals(1, assignment.getValue().size());
+    Assertions.assertEquals(goodBroker, assignment.getValue().iterator().next());
+    if (verify) {
+      Assertions.assertEquals(0, reassignment.size());
+    } else {
+      Assertions.assertEquals(1, reassignment.size());
+      Assertions.assertEquals(goodBroker, reassignment.iterator().next());
+    }
+  }
+}