Repository: skiptests/astraea
Pull request #116 (2021-11-21 08:54:09)

Add TopicConfig interface
Convert the `Map` to our interface so as to add more useful functions in the future.


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..236f8125 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
@@ -160,7 +152,7 @@ public class Builder {
     }
 
     @Override
-    public Map<String, Map<String, String>> topics() {
+    public Map<String, TopicConfig> topics() {
       var topics =
           Utils.handleException(
               () -> admin.listTopics(new ListTopicsOptions().listInternal(true)).names().get());
@@ -179,8 +171,9 @@ public class Builder {
               Collectors.toMap(
                   e -> e.getKey().name(),
                   e ->
-                      e.getValue().entries().stream()
-                          .collect(Collectors.toMap(ConfigEntry::name, ConfigEntry::value))));
+                      new TopicConfigImpl(
+                          e.getValue().entries().stream()
+                              .collect(Collectors.toMap(ConfigEntry::name, ConfigEntry::value)))));
     }
 
     @Override
@@ -294,10 +287,38 @@ public class Builder {
     }
   }
 
+  private static class TopicConfigImpl implements TopicConfig {
+    private final Map<String, String> configs;
+
+    TopicConfigImpl(Map<String, String> configs) {
+      this.configs = Collections.unmodifiableMap(configs);
+    }
+
+    @Override
+    public Optional<String> value(String key) {
+      return Optional.ofNullable(configs.get(key));
+    }
+
+    @Override
+    public Set<String> keys() {
+      return configs.keySet();
+    }
+
+    @Override
+    public Collection<String> values() {
+      return configs.values();
+    }
+
+    @Override
+    public Iterator<Map.Entry<String, String>> iterator() {
+      return configs.entrySet().iterator();
+    }
+  }
+
   private static class CreatorImpl implements Creator {
     private final Admin admin;
     private final Function<String, Map<TopicPartition, List<Replica>>> replicasGetter;
-    private final Function<String, Map<String, String>> configsGetter;
+    private final Function<String, TopicConfig> configsGetter;
     private String topic;
     private int numberOfPartitions = 1;
     private short numberOfReplicas = 1;
@@ -306,7 +327,7 @@ public class Builder {
     CreatorImpl(
         Admin admin,
         Function<String, Map<TopicPartition, List<Replica>>> replicasGetter,
-        Function<String, Map<String, String>> configsGetter) {
+        Function<String, TopicConfig> configsGetter) {
       this.admin = admin;
       this.replicasGetter = replicasGetter;
       this.configsGetter = configsGetter;
@@ -363,13 +384,13 @@ public class Builder {
         var actualConfigs = configsGetter.apply(topic);
         this.configs.forEach(
             (key, value) -> {
-              if (!actualConfigs.get(key).equals(value))
+              if (actualConfigs.value(key).filter(actual -> actual.equals(value)).isEmpty())
                 throw new IllegalArgumentException(
                     topic
                         + " is existent but its config: <"
                         + key
                         + ", "
-                        + actualConfigs.get(key)
+                        + actualConfigs.value(key)
                         + "> is not equal to expected: "
                         + key
                         + ", "
@@ -391,4 +412,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..e6680509 100644
--- a/app/src/main/java/org/astraea/topic/TopicAdmin.java
+++ b/app/src/main/java/org/astraea/topic/TopicAdmin.java
@@ -24,7 +24,7 @@ public interface TopicAdmin extends Closeable {
   }
 
   /** @return the topic name and its configurations. */
-  Map<String, Map<String, String>> topics();
+  Map<String, TopicConfig> topics();
 
   /** @return a topic creator to set all topic configs and then run the procedure. */
   Creator creator();
@@ -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();
 }
diff --git a/app/src/main/java/org/astraea/topic/TopicConfig.java b/app/src/main/java/org/astraea/topic/TopicConfig.java
new file mode 100644
index 00000000..1322cc9d
--- /dev/null
+++ b/app/src/main/java/org/astraea/topic/TopicConfig.java
@@ -0,0 +1,15 @@
+package org.astraea.topic;
+
+import java.util.Collection;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+
+public interface TopicConfig extends Iterable<Map.Entry<String, String>> {
+
+  Optional<String> value(String key);
+
+  Set<String> keys();
+
+  Collection<String> values();
+}
diff --git a/app/src/main/java/org/astraea/topic/cost/GetPartitionInf.java b/app/src/main/java/org/astraea/topic/cost/GetPartitionInf.java
index f2f79511..3a597872 100644
--- a/app/src/main/java/org/astraea/topic/cost/GetPartitionInf.java
+++ b/app/src/main/java/org/astraea/topic/cost/GetPartitionInf.java
@@ -1,6 +1,7 @@
 package org.astraea.topic.cost;
 
 import java.util.*;
+import java.util.stream.Collectors;
 import org.apache.kafka.common.TopicPartition;
 import org.astraea.topic.TopicAdmin;
 
@@ -20,10 +21,17 @@ public class GetPartitionInf {
   }
 
   static Map<String, Integer> getRetentionMillis(TopicAdmin client) {
-    Map<String, Integer> retentionTime = new HashMap<>();
-    for (var topic : client.topics().keySet()) {
-      retentionTime.put(topic, Integer.parseInt(client.topics().get(topic).get("retention.ms")));
-    }
-    return retentionTime;
+    return client.topics().entrySet().stream()
+        .collect(
+            Collectors.toMap(
+                Map.Entry::getKey,
+                entry ->
+                    client
+                        .topics()
+                        .get(entry.getKey())
+                        .value("retention.ms")
+                        .map(Integer::parseInt)
+                        .orElseThrow(
+                            () -> new NoSuchElementException("retention.ms does not exist"))));
   }
 }


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..ccad4dcb 100644
--- a/app/src/test/java/org/astraea/topic/TopicAdminTest.java
+++ b/app/src/test/java/org/astraea/topic/TopicAdminTest.java
@@ -10,6 +10,7 @@ import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
+import java.util.stream.StreamSupport;
 import org.apache.kafka.common.config.TopicConfig;
 import org.astraea.Utils;
 import org.astraea.producer.Producer;
@@ -33,9 +34,16 @@ public class TopicAdminTest extends RequireBrokerCluster {
           () ->
               topicAdmin
                   .topics()
-                  .getOrDefault(topicName, Map.of())
-                  .get(TopicConfig.COMPRESSION_TYPE_CONFIG)
-                  .equals("lz4"));
+                  .get(topicName)
+                  .value(TopicConfig.COMPRESSION_TYPE_CONFIG)
+                  .filter(value -> value.equals("lz4"))
+                  .isPresent());
+
+      var config = topicAdmin.topics().get(topicName);
+      Assertions.assertEquals(
+          config.keys().size(), (int) StreamSupport.stream(config.spliterator(), false).count());
+      config.keys().forEach(key -> Assertions.assertTrue(config.value(key).isPresent()));
+      Assertions.assertTrue(config.values().contains("lz4"));
     }
   }
 
@@ -136,14 +144,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 +163,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";