KAFKA-1792: change behavior of --generate to produce assignment config with fair replica distribution and minimal number of reassignments

build.gradle

@@ -2473,9 +2473,10 @@ project(':tools') {
     testImplementation libs.apachedsProtocolLdap
     testImplementation libs.apachedsLdifPartition
     testImplementation testLog4j2Libs
+    testImplementation libs.jqwik
+    testImplementation libs.hamcrest
 
     testRuntimeOnly runtimeTestLibs
-    testRuntimeOnly libs.hamcrest
   }
 
   javadoc {
@@ -2491,6 +2492,12 @@ project(':tools') {
     duplicatesStrategy 'exclude'
   }
 
+  test {
+    useJUnitPlatform {
+      includeEngines 'jqwik', 'junit-jupiter'
+    }
+  }
+
   jar {
     dependsOn 'copyDependantLibs'
   }

tools/src/main/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommand.java

@@ -16,6 +16,10 @@
  */
 package org.apache.kafka.tools.reassign;
 
+import com.fasterxml.jackson.core.JsonParseException;
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.JsonMappingException;
+import joptsimple.OptionSpec;
 import org.apache.kafka.admin.AdminUtils;
 import org.apache.kafka.admin.BrokerMetadata;
 import org.apache.kafka.clients.admin.Admin;
@@ -47,10 +51,6 @@
 import org.apache.kafka.tools.TerseException;
 import org.apache.kafka.tools.ToolsUtils;
 
-import com.fasterxml.jackson.core.JsonParseException;
-import com.fasterxml.jackson.core.JsonProcessingException;
-import com.fasterxml.jackson.databind.JsonMappingException;
-
 import java.io.IOException;
 import java.util.AbstractMap.SimpleImmutableEntry;
 import java.util.ArrayList;
@@ -72,8 +72,6 @@
 import java.util.concurrent.ExecutionException;
 import java.util.stream.Collectors;
 
-import joptsimple.OptionSpec;
-
 @SuppressWarnings("ClassDataAbstractionCoupling")
 public class ReassignPartitionsCommand {
     private static final String ANY_LOG_DIR = "any";
@@ -164,7 +162,8 @@ private static void handleAction(Admin adminClient, ReassignPartitionsCommandOpt
             generateAssignment(adminClient,
                 Utils.readFileAsString(opts.options.valueOf(opts.topicsToMoveJsonFileOpt)),
                 opts.options.valueOf(opts.brokerListOpt),
-                !opts.options.has(opts.disableRackAware));
+                !opts.options.has(opts.disableRackAware),
+                opts.options.has(opts.sticky));
         } else if (opts.options.has(opts.executeOpt)) {
             executeAssignment(adminClient,
                 opts.options.has(opts.additionalOpt),
@@ -549,7 +548,8 @@ private static void clearTopicLevelThrottles(Admin adminClient, Set<String> topi
     public static Entry<Map<TopicPartition, List<Integer>>, Map<TopicPartition, List<Integer>>> generateAssignment(Admin adminClient,
                                                                                                              String reassignmentJson,
                                                                                                              String brokerListString,
-                                                                                                             Boolean enableRackAwareness
+                                                                                                             Boolean enableRackAwareness,
+                                                                                                             Boolean sticky
     ) throws ExecutionException, InterruptedException, JsonProcessingException {
         Entry<List<Integer>, List<String>> t0 = parseGenerateAssignmentArgs(reassignmentJson, brokerListString);
 
@@ -558,7 +558,13 @@ public static Entry<Map<TopicPartition, List<Integer>>, Map<TopicPartition, List
 
         Map<TopicPartition, List<Integer>> currentAssignments = getReplicaAssignmentForTopics(adminClient, topicsToReassign);
         List<BrokerMetadata> brokerMetadatas = getBrokerMetadata(adminClient, brokersToReassign, enableRackAwareness);
-        Map<TopicPartition, List<Integer>> proposedAssignments = calculateAssignment(currentAssignments, brokerMetadatas);
+        Map<TopicPartition, List<Integer>> proposedAssignments;
+        if (sticky) {
+            final StickyPartitionReassignor assignor = new StickyPartitionReassignor(currentAssignments, brokerMetadatas);
+            proposedAssignments = assignor.reassign();
+        } else {
+            proposedAssignments = calculateAssignment(currentAssignments, brokerMetadatas);
+        }
         System.out.printf("Current partition replica assignment%n%s%n%n",
             formatAsReassignmentJson(currentAssignments, Collections.emptyMap()));
         System.out.printf("Proposed partition reassignment configuration%n%s%n",
@@ -574,7 +580,7 @@ public static Entry<Map<TopicPartition, List<Integer>>, Map<TopicPartition, List
      *
      * @return                   A map from partitions to the proposed assignments for each.
      */
-    private static Map<TopicPartition, List<Integer>> calculateAssignment(Map<TopicPartition, List<Integer>> currentAssignment,
+    public static Map<TopicPartition, List<Integer>> calculateAssignment(Map<TopicPartition, List<Integer>> currentAssignment,
                                                                           List<BrokerMetadata> brokerMetadatas) {
         Map<String, List<Entry<TopicPartition, List<Integer>>>> groupedByTopic = new HashMap<>();
         for (Entry<TopicPartition, List<Integer>> e : currentAssignment.entrySet())
@@ -1409,7 +1415,7 @@ static ReassignPartitionsCommandOptions validateAndParseArgs(String[] args) {
         }
 
         OptionSpec<?> action = allActions.get(0);
-        
+
         if (opts.options.has(opts.bootstrapServerOpt) && opts.options.has(opts.bootstrapControllerOpt))
             CommandLineUtils.printUsageAndExit(opts.parser, "Please don't specify both --bootstrap-server and --bootstrap-controller");
         else if (!opts.options.has(opts.bootstrapServerOpt) && !opts.options.has(opts.bootstrapControllerOpt))
@@ -1448,7 +1454,8 @@ else if (!opts.options.has(opts.bootstrapServerOpt) && !opts.options.has(opts.bo
             opts.bootstrapServerOpt,
             opts.brokerListOpt,
             opts.commandConfigOpt,
-            opts.disableRackAware
+            opts.disableRackAware,
+            opts.sticky
         ));
         permittedArgs.put(opts.executeOpt, Arrays.asList(
             opts.additionalOpt,

tools/src/main/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommandOptions.java

@@ -37,6 +37,7 @@ public class ReassignPartitionsCommandOptions extends CommandDefaultOptions {
     final OptionSpec<String> brokerListOpt;
     final OptionSpec<String> bootstrapControllerOpt;
     final OptionSpec<?> disableRackAware;
+    final OptionSpec<?> sticky;
     final OptionSpec<Long> interBrokerThrottleOpt;
     final OptionSpec<Long> replicaAlterLogDirsThrottleOpt;
     final OptionSpec<Long> timeoutOpt;
@@ -92,6 +93,7 @@ public ReassignPartitionsCommandOptions(String[] args) {
             .ofType(String.class);
 
         disableRackAware = parser.accepts("disable-rack-aware", "Disable rack aware replica assignment");
+        sticky = parser.accepts("sticky", "Use a sticky reassignment strategy to reduce the number of replica moves needed to balance out brokers");
         interBrokerThrottleOpt = parser.accepts("throttle", "The movement of partitions between brokers will be throttled to this value (bytes/sec). " +
                 "This option can be included with --execute when a reassignment is started, and it can be altered by resubmitting the current reassignment " +
                 "along with the --additional flag. The throttle rate should be at least 1 KB/s.")