Repository: skiptests/astraea
Pull request #430 (2022-06-26 20:03:46)

Enable to list reassignment by Admin
新增`Admin#reassignments`用來列舉目前的搬移狀況，包含兩種例子：

1. 節點之間的搬移
2. 目錄之間的搬移


Patch:
diff --git a/app/src/main/java/org/astraea/app/admin/Admin.java b/app/src/main/java/org/astraea/app/admin/Admin.java
index e101b87c..1032b094 100644
--- a/app/src/main/java/org/astraea/app/admin/Admin.java
+++ b/app/src/main/java/org/astraea/app/admin/Admin.java
@@ -232,6 +232,23 @@ public interface Admin extends Closeable {
    */
   void removeStaticMembers(String groupId, Set<String> members);
 
+  /**
+   * Get the reassignments of all topics.
+   *
+   * @return reassignment
+   */
+  default Map<TopicPartition, Reassignment> reassignments() {
+    return reassignments(topicNames());
+  }
+
+  /**
+   * Get the reassignments of topics. It returns nothing if the partitions are not migrating.
+   *
+   * @param topics to search
+   * @return reassignment
+   */
+  Map<TopicPartition, Reassignment> reassignments(Set<String> topics);
+
   @Override
   void close();
 }
diff --git a/app/src/main/java/org/astraea/app/admin/Builder.java b/app/src/main/java/org/astraea/app/admin/Builder.java
index d0f50b02..911c2972 100644
--- a/app/src/main/java/org/astraea/app/admin/Builder.java
+++ b/app/src/main/java/org/astraea/app/admin/Builder.java
@@ -577,6 +577,77 @@ public class Builder {
                   .all()
                   .get());
     }
+
+    @Override
+    public Map<TopicPartition, Reassignment> reassignments(Set<String> topics) {
+      var assignments =
+          Utils.packException(
+              () ->
+                  admin
+                      .listPartitionReassignments(
+                          partitions(topics).stream()
+                              .map(TopicPartition::to)
+                              .collect(Collectors.toSet()))
+                      .reassignments()
+                      .get());
+
+      var dirs =
+          Utils.packException(
+              () ->
+                  admin
+                      .describeReplicaLogDirs(
+                          replicas(topics).entrySet().stream()
+                              .flatMap(
+                                  e ->
+                                      e.getValue().stream()
+                                          .map(
+                                              r ->
+                                                  new org.apache.kafka.common.TopicPartitionReplica(
+                                                      e.getKey().topic(),
+                                                      e.getKey().partition(),
+                                                      r.broker())))
+                              .collect(Collectors.toUnmodifiableList()))
+                      .all()
+                      .get());
+
+      var result =
+          new HashMap<
+              TopicPartition, Map.Entry<Set<Reassignment.Location>, Set<Reassignment.Location>>>();
+
+      dirs.forEach(
+          (replica, logDir) -> {
+            var brokerId = replica.brokerId();
+            var tp = new TopicPartition(replica.topic(), replica.partition());
+            var ls =
+                result.computeIfAbsent(tp, ignored -> Map.entry(new HashSet<>(), new HashSet<>()));
+            // the replica is moved from a folder to another folder (in the same node)
+            if (logDir.getFutureReplicaLogDir() != null)
+              ls.getValue()
+                  .add(new Reassignment.Location(brokerId, logDir.getFutureReplicaLogDir()));
+            if (logDir.getCurrentReplicaLogDir() != null) {
+              var assignment = assignments.get(TopicPartition.to(tp));
+              // the replica is moved from a node to another node
+              if (assignment != null && assignment.addingReplicas().contains(brokerId)) {
+                ls.getValue()
+                    .add(new Reassignment.Location(brokerId, logDir.getCurrentReplicaLogDir()));
+              } else {
+                ls.getKey()
+                    .add(new Reassignment.Location(brokerId, logDir.getCurrentReplicaLogDir()));
+              }
+            }
+          });
+
+      return result.entrySet().stream()
+          // empty "to" means there is no reassignment
+          .filter(e -> !e.getValue().getValue().isEmpty())
+          .collect(
+              Collectors.toMap(
+                  Map.Entry::getKey,
+                  e ->
+                      new Reassignment(
+                          Collections.unmodifiableSet(e.getValue().getKey()),
+                          Collections.unmodifiableSet(e.getValue().getValue()))));
+    }
   }
 
   private static class ConfigImpl implements Config {
diff --git a/app/src/main/java/org/astraea/app/admin/Reassignment.java b/app/src/main/java/org/astraea/app/admin/Reassignment.java
new file mode 100644
index 00000000..358dbbbc
--- /dev/null
+++ b/app/src/main/java/org/astraea/app/admin/Reassignment.java
@@ -0,0 +1,94 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.astraea.app.admin;
+
+import java.util.Collection;
+import java.util.Objects;
+
+public class Reassignment {
+  private final Collection<Location> from;
+  private final Collection<Location> to;
+
+  Reassignment(Collection<Location> from, Collection<Location> to) {
+    this.from = from;
+    this.to = to;
+  }
+
+  /** @return the partition locations before migration */
+  public Collection<Location> from() {
+    return from;
+  }
+
+  /** @return the partition locations after migration */
+  public Collection<Location> to() {
+    return to;
+  }
+
+  @Override
+  public String toString() {
+    return "Reassignment{" + "from=" + from + ", to=" + to + '}';
+  }
+
+  @Override
+  public boolean equals(Object o) {
+    if (this == o) return true;
+    if (o == null || getClass() != o.getClass()) return false;
+    Reassignment that = (Reassignment) o;
+    return Objects.equals(from, that.from) && Objects.equals(to, that.to);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(from, to);
+  }
+
+  public static class Location {
+    private final int broker;
+    private final String path;
+
+    public Location(int broker, String path) {
+      this.broker = broker;
+      this.path = Objects.requireNonNull(path);
+    }
+
+    public int broker() {
+      return broker;
+    }
+
+    public String path() {
+      return path;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      if (this == o) return true;
+      if (o == null || getClass() != o.getClass()) return false;
+      Location location = (Location) o;
+      return broker == location.broker && path.equals(location.path);
+    }
+
+    @Override
+    public int hashCode() {
+      return Objects.hash(broker, path);
+    }
+
+    @Override
+    public String toString() {
+      return "Location{" + "broker=" + broker + ", path='" + path + '\'' + '}';
+    }
+  }
+}
diff --git a/app/src/main/java/org/astraea/app/common/Utils.java b/app/src/main/java/org/astraea/app/common/Utils.java
index 63d6d55e..4180000a 100644
--- a/app/src/main/java/org/astraea/app/common/Utils.java
+++ b/app/src/main/java/org/astraea/app/common/Utils.java
@@ -134,11 +134,24 @@ public final class Utils {
    * @param done a flag indicating the result.
    */
   public static void waitFor(Supplier<Boolean> done, Duration timeout) {
+    waitForNonNull(() -> done.get() ? "good" : null, timeout);
+  }
+
+  /**
+   * loop the supplier until it returns non-null value
+   *
+   * @param supplier to loop
+   * @param timeout to break the loop
+   * @param <T> returned type
+   * @return value from supplier
+   */
+  public static <T> T waitForNonNull(Supplier<T> supplier, Duration timeout) {
     var endTime = System.currentTimeMillis() + timeout.toMillis();
     Exception lastError = null;
     while (System.currentTimeMillis() <= endTime)
       try {
-        if (done.get()) return;
+        var r = supplier.get();
+        if (r != null) return r;
         TimeUnit.SECONDS.sleep(1);
       } catch (Exception e) {
         lastError = e;


Test patch:
diff --git a/app/src/test/java/org/astraea/app/admin/AdminTest.java b/app/src/test/java/org/astraea/app/admin/AdminTest.java
index d5bd7d91..f53f0224 100644
--- a/app/src/test/java/org/astraea/app/admin/AdminTest.java
+++ b/app/src/test/java/org/astraea/app/admin/AdminTest.java
@@ -20,6 +20,7 @@ import static org.junit.jupiter.api.condition.OS.WINDOWS;
 
 import java.nio.charset.StandardCharsets;
 import java.time.Duration;
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
@@ -28,6 +29,7 @@ import java.util.Set;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ThreadLocalRandom;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
@@ -36,6 +38,8 @@ import java.util.stream.StreamSupport;
 import org.apache.kafka.common.config.TopicConfig;
 import org.apache.kafka.common.errors.GroupNotEmptyException;
 import org.astraea.app.common.Utils;
+import org.astraea.app.concurrent.State;
+import org.astraea.app.concurrent.ThreadPool;
 import org.astraea.app.consumer.Consumer;
 import org.astraea.app.consumer.Deserializer;
 import org.astraea.app.cost.NodeInfo;
@@ -44,6 +48,7 @@ import org.astraea.app.producer.Producer;
 import org.astraea.app.producer.Serializer;
 import org.astraea.app.service.RequireBrokerCluster;
 import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.RepeatedTest;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.condition.DisabledOnOs;
 import org.junit.jupiter.params.ParameterizedTest;
@@ -929,4 +934,153 @@ public class AdminTest extends RequireBrokerCluster {
       Assertions.assertFalse(admin.consumerGroupIds().contains(groupId));
     }
   }
+
+  @RepeatedTest(5) // this test is timer-based, so it needs to be looped for stability
+  void testReassignmentWhenMovingPartitionToAnotherBroker() throws InterruptedException {
+    var topicName = Utils.randomString(10);
+    try (var admin = Admin.of(bootstrapServers())) {
+      admin.creator().topic(topicName).numberOfPartitions(1).create();
+      TimeUnit.SECONDS.sleep(3);
+
+      var currentBroker =
+          admin.replicas(Set.of(topicName)).get(new TopicPartition(topicName, 0)).get(0).broker();
+      var nextBroker = brokerIds().stream().filter(i -> i != currentBroker).findAny().get();
+
+      try (var producer = Producer.of(bootstrapServers())) {
+        var done = new AtomicBoolean(false);
+        var data = new byte[1000];
+        try (var pool =
+            ThreadPool.builder()
+                .executor(
+                    () -> {
+                      producer.sender().topic(topicName).key(data).value(data).run();
+                      return done.get() ? State.DONE : State.RUNNING;
+                    })
+                .build()) {
+
+          try {
+            admin.migrator().topic(topicName).moveTo(List.of(nextBroker));
+            var reassignment =
+                admin.reassignments(Set.of(topicName)).get(new TopicPartition(topicName, 0));
+
+            // Don't verify the result if the migration is done
+            if (reassignment != null) {
+              Assertions.assertEquals(1, reassignment.from().size());
+              var from = reassignment.from().iterator().next();
+              Assertions.assertEquals(currentBroker, from.broker());
+              Assertions.assertEquals(1, reassignment.to().size());
+              var to = reassignment.to().iterator().next();
+              Assertions.assertEquals(nextBroker, to.broker());
+            }
+          } finally {
+            done.set(true);
+          }
+        }
+      }
+    }
+  }
+
+  @RepeatedTest(5) // this test is timer-based, so it needs to be looped for stability
+  void testReassignmentWhenMovingPartitionToAnotherPath() throws InterruptedException {
+    var topicName = Utils.randomString(10);
+    try (var admin = Admin.of(bootstrapServers())) {
+      admin.creator().topic(topicName).numberOfPartitions(1).create();
+      TimeUnit.SECONDS.sleep(3);
+
+      var currentReplica =
+          admin.replicas(Set.of(topicName)).get(new TopicPartition(topicName, 0)).get(0);
+      var currentBroker = currentReplica.broker();
+      var currentPath = currentReplica.path();
+      var nextPath =
+          logFolders().get(currentBroker).stream()
+              .filter(p -> !p.equals(currentPath))
+              .findFirst()
+              .get();
+
+      try (var producer = Producer.of(bootstrapServers())) {
+        var done = new AtomicBoolean(false);
+        var data = new byte[1000];
+        try (var pool =
+            ThreadPool.builder()
+                .executor(
+                    () -> {
+                      producer.sender().topic(topicName).key(data).value(data).run();
+                      return done.get() ? State.DONE : State.RUNNING;
+                    })
+                .build()) {
+
+          try {
+            admin.migrator().topic(topicName).moveTo(Map.of(currentReplica.broker(), nextPath));
+            var reassignment =
+                admin.reassignments(Set.of(topicName)).get(new TopicPartition(topicName, 0));
+            // Don't verify the result if the migration is done
+            if (reassignment != null) {
+              Assertions.assertEquals(1, reassignment.from().size());
+              var from = reassignment.from().iterator().next();
+              Assertions.assertEquals(currentBroker, from.broker());
+              Assertions.assertEquals(currentPath, from.path());
+              Assertions.assertEquals(1, reassignment.to().size());
+              var to = reassignment.to().iterator().next();
+              Assertions.assertEquals(currentBroker, to.broker());
+              Assertions.assertEquals(nextPath, to.path());
+            }
+          } finally {
+            done.set(true);
+          }
+        }
+      }
+    }
+  }
+
+  @RepeatedTest(5) // this test is timer-based, so it needs to be looped for stability
+  void testMultiReassignments() throws InterruptedException {
+    var topicName = Utils.randomString(10);
+    try (var admin = Admin.of(bootstrapServers())) {
+      admin.creator().topic(topicName).numberOfPartitions(1).numberOfReplicas((short) 3).create();
+      TimeUnit.SECONDS.sleep(3);
+
+      var brokers = new ArrayList<>(brokerIds()).subList(0, 2);
+      try (var producer = Producer.of(bootstrapServers())) {
+        var done = new AtomicBoolean(false);
+        var data = new byte[1000];
+        try (var pool =
+            ThreadPool.builder()
+                .executor(
+                    () -> {
+                      producer.sender().topic(topicName).key(data).value(data).run();
+                      return done.get() ? State.DONE : State.RUNNING;
+                    })
+                .build()) {
+
+          try {
+            admin.migrator().topic(topicName).moveTo(brokers);
+            var reassignment =
+                admin.reassignments(Set.of(topicName)).get(new TopicPartition(topicName, 0));
+            // Don't verify the result if the migration is done
+            if (reassignment != null) {
+              Assertions.assertEquals(3, reassignment.from().size());
+              Assertions.assertEquals(2, reassignment.to().size());
+            }
+          } finally {
+            done.set(true);
+          }
+        }
+      }
+    }
+  }
+
+  @Test
+  void testReassignmentWithNothing() throws InterruptedException {
+    var topicName = Utils.randomString(10);
+    try (var admin = Admin.of(bootstrapServers())) {
+      admin.creator().topic(topicName).numberOfPartitions(1).numberOfReplicas((short) 3).create();
+      TimeUnit.SECONDS.sleep(3);
+      var brokers = new ArrayList<>(brokerIds()).subList(0, 2);
+      try (var producer = Producer.of(bootstrapServers())) {
+        producer.sender().topic(topicName).value(new byte[100]).run();
+        producer.flush();
+      }
+      Assertions.assertEquals(0, admin.reassignments(Set.of(topicName)).size());
+    }
+  }
 }