From f5b12881c09583973f1f53fbb7147cf9b90bc2a9 Mon Sep 17 00:00:00 2001
From: Jialiang Li <jlli@linkedin.com>
Date: Tue, 9 Jul 2024 14:19:07 -0700
Subject: [PATCH] Add nonLeaderCleanUp with tables parameters in
 PinotTaskGenerator (#13509)

---
 .../helix/core/minion/PinotTaskManager.java         | 13 +++++++++++++
 .../core/minion/generator/PinotTaskGenerator.java   |  7 +++++++
 2 files changed, 20 insertions(+)

diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/minion/PinotTaskManager.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/minion/PinotTaskManager.java
index 621e82547fca..bacaba059d2e 100644
--- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/minion/PinotTaskManager.java
+++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/minion/PinotTaskManager.java
@@ -34,6 +34,7 @@
 import java.util.concurrent.Executor;
 import javax.annotation.Nullable;
 import org.apache.commons.collections.CollectionUtils;
+import org.apache.commons.lang3.StringUtils;
 import org.apache.helix.AccessOption;
 import org.apache.helix.task.TaskState;
 import org.apache.helix.zookeeper.zkclient.IZkChildListener;
@@ -702,6 +703,18 @@ public void cleanUpTask() {
     }
   }
 
+  @Override
+  protected void nonLeaderCleanup(List<String> tableNamesWithType) {
+    LOGGER.info(
+        "Cleaning up all task generators for tables that the controller is not the leader for. Number of tables to be"
+            + " cleaned up: {}. Printing at most first 10 table names to be cleaned up: [{}].",
+        tableNamesWithType.size(),
+        StringUtils.join(tableNamesWithType.stream().limit(10).map(t -> "\"" + t + "\"").toArray(), ", "));
+    for (String taskType : _taskGeneratorRegistry.getAllTaskTypes()) {
+      _taskGeneratorRegistry.getTaskGenerator(taskType).nonLeaderCleanUp(tableNamesWithType);
+    }
+  }
+
   @Nullable
   public Scheduler getScheduler() {
     return _scheduler;
diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/minion/generator/PinotTaskGenerator.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/minion/generator/PinotTaskGenerator.java
index ad4213588608..fd2a4615edd5 100644
--- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/minion/generator/PinotTaskGenerator.java
+++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/minion/generator/PinotTaskGenerator.java
@@ -86,6 +86,13 @@ default int getMaxAttemptsPerTask() {
   default void nonLeaderCleanUp() {
   }
 
+  /**
+   * Performs necessary cleanups (e.g. remove metrics) when the controller leadership changes,
+   * given a list of tables that the current controller isn't the leader for.
+   */
+  default void nonLeaderCleanUp(List<String> tableNamesWithType) {
+  }
+
   /**
    * Gets the minionInstanceTag for the tableConfig
    */