Fix throttler quota removal for in-progress tasks (#1616)

linkedin · Jul 16, 2021 · 025ca33 · 025ca33
1 parent 577eedd
commit 025ca33
Showing 1 changed file with 6 additions and 2 deletions.
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
@@ -1282,8 +1282,12 @@ private void interBrokerMoveReplicas() {
                  finishedDataMovementInMB, totalDataToMoveInMB,
                  totalDataToMoveInMB == 0 ? 100 : String.format("%.2f", finishedDataMovementInMB * UNIT_INTERVAL_TO_PERCENTAGE
                                                                         / totalDataToMoveInMB));
-        throttleHelper.clearThrottles(completedTasks,
-                tasksToExecute.stream().filter(t -> t.state() == ExecutionTaskState.IN_PROGRESS).collect(Collectors.toList()));
+        List<ExecutionTask> inProgressTasks = tasksToExecute.stream()
+            .filter(t -> t.state() == ExecutionTaskState.IN_PROGRESS)
+            .collect(Collectors.toList());
+        inProgressTasks.addAll(inExecutionTasks());
+
+        throttleHelper.clearThrottles(completedTasks, inProgressTasks);
       }
       // After the partition movement finishes, wait for the controller to clean the reassignment zkPath. This also
       // ensures a clean stop when the execution is stopped in the middle.