From 90829f3f633edf4ade8690b89fa757d391b2c208 Mon Sep 17 00:00:00 2001
From: Allen Wang <allenwang@linkedin.com>
Date: Fri, 22 Nov 2024 18:04:16 -0800
Subject: [PATCH] WIP

---
 .../executor/ExecutionTaskPlanner.java        |  8 ++--
 .../cruisecontrol/executor/Executor.java      |  6 ++-
 .../async/runnable/RemoveBrokersRunnable.java |  7 ++-
 .../executor/ExecutionTaskPlannerTest.java    | 45 +++++++++++++++++++
 4 files changed, 60 insertions(+), 6 deletions(-)

diff --git a/cruise-control/src/main/java/com/linkedin/kafka/cruisecontrol/executor/ExecutionTaskPlanner.java b/cruise-control/src/main/java/com/linkedin/kafka/cruisecontrol/executor/ExecutionTaskPlanner.java
index 32db67f0d..8b6f7137a 100644
--- a/cruise-control/src/main/java/com/linkedin/kafka/cruisecontrol/executor/ExecutionTaskPlanner.java
+++ b/cruise-control/src/main/java/com/linkedin/kafka/cruisecontrol/executor/ExecutionTaskPlanner.java
@@ -388,7 +388,7 @@ public List<ExecutionTask> getInterBrokerReplicaMovementTasks(Map<Integer, Integ
           continue;
         }
         // Check the available balancing proposals of this broker to see if we can find one ready to execute.
-        SortedSet<ExecutionTask> proposalsForBroker = _interPartMoveTasksByBrokerId.get(brokerId);
+        SortedSet<ExecutionTask> proposalsForBroker = new TreeSet<>(_interPartMoveTasksByBrokerId.get(brokerId));
         LOG.trace("Execution task for broker {} are {}", brokerId, proposalsForBroker);
         for (ExecutionTask task : proposalsForBroker) {
           // Break if max cap reached
@@ -434,8 +434,10 @@ public List<ExecutionTask> getInterBrokerReplicaMovementTasks(Map<Integer, Integ
             newTaskAdded = true;
             numInProgressPartitions++;
             LOG.debug("Found ready task {} for broker {}. Broker concurrency state: {}", task, brokerId, readyBrokers);
-            // We can stop the check for proposals for this broker because we have found a proposal.
-            break;
+            if (_preferRoundRobin) {
+              // We can stop the check for proposals for this broker because we have found a proposal.
+              break;
+            }
           }
         }
       }
diff --git a/cruise-control/src/main/java/com/linkedin/kafka/cruisecontrol/executor/Executor.java b/cruise-control/src/main/java/com/linkedin/kafka/cruisecontrol/executor/Executor.java
index f8100cde8..56cbc85bd 100644
--- a/cruise-control/src/main/java/com/linkedin/kafka/cruisecontrol/executor/Executor.java
+++ b/cruise-control/src/main/java/com/linkedin/kafka/cruisecontrol/executor/Executor.java
@@ -71,6 +71,7 @@
 import static com.linkedin.kafka.cruisecontrol.executor.ExecutorAdminUtils.*;
 import static com.linkedin.kafka.cruisecontrol.monitor.MonitorUtils.UNIT_INTERVAL_TO_PERCENTAGE;
 import static com.linkedin.kafka.cruisecontrol.monitor.sampling.MetricSampler.SamplingMode.*;
+import static com.linkedin.kafka.cruisecontrol.servlet.UserTaskManager.TaskState.COMPLETED;
 import static org.apache.kafka.clients.admin.DescribeReplicaLogDirsResult.ReplicaLogDirInfo;
 import static com.linkedin.cruisecontrol.common.utils.Utils.validateNotNull;
 
@@ -1362,8 +1363,8 @@ private class ProposalExecutionRunnable implements Runnable {
     public void run() {
       LOG.info("User task {}: Starting executing balancing proposals.", _uuid);
       final long start = System.currentTimeMillis();
+      UserTaskManager.UserTaskInfo userTaskInfo = initExecution();
       try {
-        UserTaskManager.UserTaskInfo userTaskInfo = initExecution();
         execute(userTaskInfo);
       } catch (Exception e) {
         LOG.error("User task {}: ProposalExecutionRunnable got exception during run", _uuid, e);
@@ -1381,7 +1382,8 @@ public void run() {
         if (_executionException != null) {
           LOG.info("User task {}: Execution failed: {}. Exception: {}", _uuid, executionStatusString, _executionException.getMessage());
         } else {
-          LOG.info("User task {}: Execution succeeded: {}. ", _uuid, executionStatusString);
+          String status = userTaskInfo.state() == COMPLETED ? "succeeded" : userTaskInfo.state().toString();
+          LOG.info("User task {}: Execution {}: {}. ", _uuid, status, executionStatusString);
         }
         // Clear completed execution.
         clearCompletedExecution();
diff --git a/cruise-control/src/main/java/com/linkedin/kafka/cruisecontrol/servlet/handler/async/runnable/RemoveBrokersRunnable.java b/cruise-control/src/main/java/com/linkedin/kafka/cruisecontrol/servlet/handler/async/runnable/RemoveBrokersRunnable.java
index 16e05b96e..ef41c42b7 100644
--- a/cruise-control/src/main/java/com/linkedin/kafka/cruisecontrol/servlet/handler/async/runnable/RemoveBrokersRunnable.java
+++ b/cruise-control/src/main/java/com/linkedin/kafka/cruisecontrol/servlet/handler/async/runnable/RemoveBrokersRunnable.java
@@ -96,7 +96,12 @@ public RemoveBrokersRunnable(KafkaCruiseControl kafkaCruiseControl,
 
   @Override
   protected OptimizationResult getResult() throws Exception {
-    return new OptimizationResult(computeResult(), _kafkaCruiseControl.config());
+    try {
+      return new OptimizationResult(computeResult(), _kafkaCruiseControl.config());
+    } catch (Exception e) {
+      LOG.error("User task {}: failed to remove brokers due to {}", _uuid, e);
+      throw e;
+    }
   }
 
   @Override
diff --git a/cruise-control/src/test/java/com/linkedin/kafka/cruisecontrol/executor/ExecutionTaskPlannerTest.java b/cruise-control/src/test/java/com/linkedin/kafka/cruisecontrol/executor/ExecutionTaskPlannerTest.java
index 3a3a4c24d..41bca12ef 100644
--- a/cruise-control/src/test/java/com/linkedin/kafka/cruisecontrol/executor/ExecutionTaskPlannerTest.java
+++ b/cruise-control/src/test/java/com/linkedin/kafka/cruisecontrol/executor/ExecutionTaskPlannerTest.java
@@ -361,6 +361,51 @@ public void testGetInterBrokerPartitionMovementWithMinIsrTasks() {
     assertEquals("Fourth task", _rf4PartitionMovement0, partitionMovementTasks.get(3).proposal());
   }
 
+  @Test
+  public void testGetInterBrokerPartitionMovementWithMinIsrNoRoundRobinTasks() {
+    List<ExecutionProposal> proposals = new ArrayList<>();
+    proposals.add(_rf4PartitionMovement0);
+    proposals.add(_rf4PartitionMovement1);
+    proposals.add(_rf4PartitionMovement2);
+    proposals.add(_rf4PartitionMovement3);
+    // Test PrioritizeOneAboveMinIsrWithOfflineReplicasStrategy execution strategies.
+    // Create prioritizeOneAboveMinIsrMovementPlanner, chain after prioritizeMinIsr strategy
+    Properties prioritizeOneAboveMinIsrMovementProps = KafkaCruiseControlUnitTestUtils.getKafkaCruiseControlProperties();
+    prioritizeOneAboveMinIsrMovementProps.setProperty(ExecutorConfig.DEFAULT_REPLICA_MOVEMENT_STRATEGIES_CONFIG,
+        String.format("%s,%s", PrioritizeMinIsrWithOfflineReplicasStrategy.class.getName(),
+            PrioritizeOneAboveMinIsrWithOfflineReplicasStrategy.class.getName()));
+    prioritizeOneAboveMinIsrMovementProps.setProperty(ExecutorConfig.PREFER_BROKER_ROUND_ROBIN_CONFIG, "false");
+    ExecutionTaskPlanner prioritizeOneAboveMinIsrMovementPlanner
+        = new ExecutionTaskPlanner(null, new KafkaCruiseControlConfig(prioritizeOneAboveMinIsrMovementProps));
+
+    Set<PartitionInfo> partitions = new HashSet<>();
+    partitions.add(generatePartitionInfo(_rf4PartitionMovement0, false));
+    partitions.add(generatePartitionInfoWithUrpHavingOfflineReplica(_rf4PartitionMovement1, 2));
+    partitions.add(generatePartitionInfoWithUrpHavingOfflineReplica(_rf4PartitionMovement2, 3));
+    partitions.add(generatePartitionInfoWithUrpHavingOfflineReplica(_rf4PartitionMovement3, 1));
+
+    Cluster expectedCluster = new Cluster(null, _rf4ExpectedNodes, partitions, Collections.emptySet(), Collections.emptySet());
+    // Setting topic min ISR to 2
+    Map<String, MinIsrWithTime> minIsrWithTimeByTopic
+        = Collections.singletonMap(TOPIC3, new MinIsrWithTime((short) 2, 0));
+    StrategyOptions strategyOptions = new StrategyOptions.Builder(expectedCluster).minIsrWithTimeByTopic(minIsrWithTimeByTopic).build();
+
+    Map<Integer, Integer> readyBrokers = new HashMap<>();
+    readyBrokers.put(0, 5);
+    readyBrokers.put(1, 6);
+    readyBrokers.put(2, 6);
+    readyBrokers.put(3, 6);
+    readyBrokers.put(4, 5);
+    readyBrokers.put(5, 6);
+    prioritizeOneAboveMinIsrMovementPlanner.addExecutionProposals(proposals, strategyOptions, null);
+    List<ExecutionTask> partitionMovementTasks
+        = prioritizeOneAboveMinIsrMovementPlanner.getInterBrokerReplicaMovementTasks(readyBrokers, Collections.emptySet(), _defaultPartitionsMaxCap);
+    assertEquals("First task", _rf4PartitionMovement2, partitionMovementTasks.get(0).proposal());
+    assertEquals("Second task", _rf4PartitionMovement1, partitionMovementTasks.get(1).proposal());
+    assertEquals("Third task", _rf4PartitionMovement3, partitionMovementTasks.get(2).proposal());
+    assertEquals("Fourth task", _rf4PartitionMovement0, partitionMovementTasks.get(3).proposal());
+  }
+
   @Test
   public void testDynamicConfigReplicaMovementStrategy() {
     List<ExecutionProposal> proposals = new ArrayList<>();