Update Providing number of assignments by which the number of partiti…

…ons was increased (#263)

The partition count you are increasing the partitions by should equal the actual size of the new partition assignments.

This PR achieves this.

For instance:
old partitions count = 5
new partitions count = 11

increase in partitions = 6
the size of the new partition assignments ought to be be 6 also.

Note: Random assignment is adequate here because the periodic runnable for maybeReassignPartitions runs the reassignment operation on an interval.




Signed-off-by: Andrew Choi <[email protected]>

src/main/java/com/linkedin/kmf/services/MultiClusterTopicManagementService.java

@@ -23,6 +23,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Properties;
+import java.util.Random;
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
@@ -338,14 +339,13 @@ static List<List<Integer>> newPartitionAssignments(int minPartitionNum, int part
       int partitionDifference = minPartitionNum - partitionNum;
 
       // leader assignments -
-      for (BrokerMetadata brokerMetadata : brokers) {
+      while (newPartitionAssignments.size() != partitionDifference) {
         List replicas = new ArrayList<>();
         // leader replica/broker -
-        replicas.add(brokerMetadata.id());
+        int brokerMetadata = randomBroker(brokers).id();
+        replicas.add(brokerMetadata);
+
         newPartitionAssignments.add(replicas);
-        if (newPartitionAssignments.size() == partitionDifference) {
-          break;
-        }
       }
 
       // follower assignments -
@@ -364,6 +364,20 @@ static List<List<Integer>> newPartitionAssignments(int minPartitionNum, int part
       return newPartitionAssignments;
     }
 
+    private static BrokerMetadata randomBroker(Set<BrokerMetadata> brokers) {
+      int brokerSetSize = brokers.size();
+      // In practicality, the Random object should be rather more shared than this.
+      int random = new Random().nextInt(brokerSetSize);
+      int index = 0;
+      for (BrokerMetadata brokerMetadata : brokers) {
+        if (index == random)
+          return brokerMetadata;
+        index++;
+      }
+
+      throw new IllegalStateException("Couldn't find random broker.");
+    }
+
     /**
      * Exposed package-private access for testing. Get the total number of partitions for a Kafka topic.
      * @return total number of topic partitions

src/test/java/com/linkedin/kmf/services/MultiClusterTopicManagementServiceTest.java

@@ -60,9 +60,11 @@ private void startTest() {
     nodeSet.add(new Node(3, "host-3", 2134));
     nodeSet.add(new Node(4, "host-4", 2135));
     nodeSet.add(new Node(5, "host-5", 2136));
-    nodeSet.add(new Node(6, "host-5", 2136));
-    nodeSet.add(new Node(7, "host-5", 2136));
-    nodeSet.add(new Node(8, "host-5", 2136));
+    nodeSet.add(new Node(6, "host-5", 2137));
+    nodeSet.add(new Node(7, "host-5", 2138));
+    nodeSet.add(new Node(8, "host-5", 2139));
+    nodeSet.add(new Node(9, "host-5", 2140));
+    nodeSet.add(new Node(10, "host-5", 2141));
 
     _topicManagementHelper = Mockito.mock(MultiClusterTopicManagementService.TopicManagementHelper.class);
     _topicManagementHelper._topic = SERVICE_TEST_TOPIC;
@@ -82,17 +84,18 @@ protected void maybeAddPartitionsTest() {
     for (Node broker : nodeSet) {
       brokerMetadataSet.add(new BrokerMetadata(broker.id(), Option.apply(broker.rack())));
     }
+
+    int minPartitionNum = 14;
+    int partitionNum = 5;
+    int rf = 4;
+
     List<List<Integer>> newPartitionAssignments =
-        MultiClusterTopicManagementService.TopicManagementHelper.newPartitionAssignments(11, 5, brokerMetadataSet, 4);
+        MultiClusterTopicManagementService.TopicManagementHelper.newPartitionAssignments(minPartitionNum, partitionNum, brokerMetadataSet, rf);
     Assert.assertNotNull(newPartitionAssignments);
 
     System.out.println(newPartitionAssignments);
-    Assert.assertEquals(newPartitionAssignments.get(0).get(0).intValue(), 1);
-    Assert.assertEquals(newPartitionAssignments.get(1).get(0).intValue(), 2);
-    Assert.assertEquals(newPartitionAssignments.get(2).get(0).intValue(), 3);
-    Assert.assertEquals(newPartitionAssignments.get(3).get(0).intValue(), 4);
-    Assert.assertEquals(newPartitionAssignments.get(4).get(0).intValue(), 5);
-    Assert.assertEquals(newPartitionAssignments.get(5).get(0).intValue(), 6);
+    Assert.assertEquals(newPartitionAssignments.size(), minPartitionNum - partitionNum);
+    Assert.assertEquals(newPartitionAssignments.get(0).size(), rf);
   }
 
   @Test