Repository: skiptests/astraea
Pull request #115 (2021-11-21 08:33:46)

Add migrator interface
考量我們要完整支援partition assignments，因此將該操作的設定獨立成一個介面(`migrator`)以便可擴充支援所有行為


Patch:
diff --git a/app/src/main/java/org/astraea/topic/Builder.java b/app/src/main/java/org/astraea/topic/Builder.java
index ef3e7a20..9e6f6293 100644
--- a/app/src/main/java/org/astraea/topic/Builder.java
+++ b/app/src/main/java/org/astraea/topic/Builder.java
@@ -51,16 +51,8 @@ public class Builder {
     }
 
     @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());
+    public Migrator migrator() {
+      return new MigratorImpl(admin, this::partitions);
     }
 
     @Override
@@ -391,4 +383,44 @@ public class Builder {
                   .get());
     }
   }
+
+  private static class MigratorImpl implements Migrator {
+    private final Admin admin;
+    private final Function<Set<String>, Set<TopicPartition>> partitionGetter;
+    private final Set<TopicPartition> partitions = new HashSet<>();
+
+    MigratorImpl(Admin admin, Function<Set<String>, Set<TopicPartition>> partitionGetter) {
+      this.admin = admin;
+      this.partitionGetter = partitionGetter;
+    }
+
+    @Override
+    public Migrator topic(String topic) {
+      partitions.addAll(partitionGetter.apply(Set.of(topic)));
+      return this;
+    }
+
+    @Override
+    public Migrator partition(String topic, int partition) {
+      partitions.add(new TopicPartition(topic, partition));
+      return this;
+    }
+
+    @Override
+    public void moveTo(Set<Integer> brokers) {
+      Utils.handleException(
+          () ->
+              admin
+                  .alterPartitionReassignments(
+                      partitions.stream()
+                          .collect(
+                              Collectors.toMap(
+                                  Function.identity(),
+                                  ignore ->
+                                      Optional.of(
+                                          new NewPartitionReassignment(new ArrayList<>(brokers))))))
+                  .all()
+                  .get());
+    }
+  }
 }
diff --git a/app/src/main/java/org/astraea/topic/Migrator.java b/app/src/main/java/org/astraea/topic/Migrator.java
new file mode 100644
index 00000000..e8463759
--- /dev/null
+++ b/app/src/main/java/org/astraea/topic/Migrator.java
@@ -0,0 +1,31 @@
+package org.astraea.topic;
+
+import java.util.Set;
+
+/** used to migrate partitions to another broker or broker folder. */
+public interface Migrator {
+  /**
+   * move all partitions (leader replica and follower replicas) of topic
+   *
+   * @param topic topic name
+   * @return this migrator
+   */
+  Migrator topic(String topic);
+
+  /**
+   * move one partition (leader replica and follower replicas) of topic
+   *
+   * @param topic topic name
+   * @param partition partition id
+   * @return this migrator
+   */
+  Migrator partition(String topic, int partition);
+
+  /**
+   * move partitions to specify brokers. Noted that the number of brokers must be equal to the
+   * number of replicas
+   *
+   * @param brokers to host partitions
+   */
+  void moveTo(Set<Integer> brokers);
+}
diff --git a/app/src/main/java/org/astraea/topic/ReplicaCollie.java b/app/src/main/java/org/astraea/topic/ReplicaCollie.java
index 3ec87d6d..6abca971 100644
--- a/app/src/main/java/org/astraea/topic/ReplicaCollie.java
+++ b/app/src/main/java/org/astraea/topic/ReplicaCollie.java
@@ -59,7 +59,8 @@ public class ReplicaCollie {
             });
     result.forEach(
         (tp, assignments) -> {
-          if (!args.verify) admin.reassign(tp.topic(), tp.partition(), assignments.getValue());
+          if (!args.verify)
+            admin.migrator().partition(tp.topic(), tp.partition()).moveTo(assignments.getValue());
         });
     return result;
   }
diff --git a/app/src/main/java/org/astraea/topic/TopicAdmin.java b/app/src/main/java/org/astraea/topic/TopicAdmin.java
index 8a8615dc..afa600af 100644
--- a/app/src/main/java/org/astraea/topic/TopicAdmin.java
+++ b/app/src/main/java/org/astraea/topic/TopicAdmin.java
@@ -50,12 +50,6 @@ public interface TopicAdmin extends Closeable {
   /** @return all brokers' ids */
   Set<Integer> brokerIds();
 
-  /**
-   * Assign the topic partition to specific brokers.
-   *
-   * @param topicName topic name
-   * @param partition partition
-   * @param brokers to hold all the
-   */
-  void reassign(String topicName, int partition, Set<Integer> brokers);
+  /** @return a partition migrator used to move partitions to another broker or folder. */
+  Migrator migrator();
 }


Test patch:
diff --git a/app/src/test/java/org/astraea/topic/TopicAdminTest.java b/app/src/test/java/org/astraea/topic/TopicAdminTest.java
index 4f5e6d7c..cc3480e2 100644
--- a/app/src/test/java/org/astraea/topic/TopicAdminTest.java
+++ b/app/src/test/java/org/astraea/topic/TopicAdminTest.java
@@ -136,14 +136,14 @@ public class TopicAdminTest extends RequireBrokerCluster {
 
   @Test
   @DisabledOnOs(WINDOWS)
-  void testReassign() throws IOException, InterruptedException {
-    var topicName = "testReassign";
+  void testMigrateSinglePartition() throws IOException, InterruptedException {
+    var topicName = "testMigrateSinglePartition";
     try (var topicAdmin = TopicAdmin.of(bootstrapServers())) {
       topicAdmin.creator().topic(topicName).numberOfPartitions(1).create();
       // wait for syncing topic creation
       TimeUnit.SECONDS.sleep(5);
       var broker = topicAdmin.brokerIds().iterator().next();
-      topicAdmin.reassign(topicName, 0, Set.of(broker));
+      topicAdmin.migrator().partition(topicName, 0).moveTo(Set.of(broker));
       Utils.waitFor(
           () -> {
             var replicas = topicAdmin.replicas(Set.of(topicName));
@@ -155,6 +155,27 @@ public class TopicAdminTest extends RequireBrokerCluster {
     }
   }
 
+  @Test
+  @DisabledOnOs(WINDOWS)
+  void testMigrateAllPartitions() throws IOException, InterruptedException {
+    var topicName = "testMigrateAllPartitions";
+    try (var topicAdmin = TopicAdmin.of(bootstrapServers())) {
+      topicAdmin.creator().topic(topicName).numberOfPartitions(3).create();
+      // wait for syncing topic creation
+      TimeUnit.SECONDS.sleep(5);
+      var broker = topicAdmin.brokerIds().iterator().next();
+      topicAdmin.migrator().topic(topicName).moveTo(Set.of(broker));
+      Utils.waitFor(
+          () -> {
+            var replicas = topicAdmin.replicas(Set.of(topicName));
+            if (replicas.size() != 3) return false;
+            if (!replicas.values().stream().allMatch(rs -> rs.size() == 1)) return false;
+            return replicas.values().stream()
+                .allMatch(rs -> rs.stream().allMatch(r -> r.broker() == broker));
+          });
+    }
+  }
+
   @Test
   void testReplicaSize() throws IOException, ExecutionException, InterruptedException {
     var topicName = "testReplicaSize";