From 35848996abae32637327ab90e930eabbe2aeb624 Mon Sep 17 00:00:00 2001
From: praveenc7 <praveenkchaganlal@gmail.com>
Date: Fri, 13 Dec 2024 00:02:25 +0530
Subject: [PATCH] Expose QueryScheduler from ServerInstance & Custom interface
 (#14630)

---
 .../helix/core/PinotHelixResourceManager.java | 28 ++++++++++++++++++-
 .../pinot/server/starter/ServerInstance.java  |  4 +++
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManager.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManager.java
index 42bd8c4ac4d1..77762c2ee8fb 100644
--- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManager.java
+++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManager.java
@@ -4032,7 +4032,8 @@ public void endReplaceSegments(String tableNameWithType, String segmentLineageEn
         if (!waitForSegmentsBecomeOnline(tableNameWithType, segmentsTo)) {
           return false;
         }
-
+        // Could be used to perform operation before segments replacement
+        preSegmentReplaceUpdateRouting(tableNameWithType, segmentsTo, lineageEntry.getSegmentsFrom());
         // Update lineage entry
         LineageEntry lineageEntryToUpdate =
             new LineageEntry(lineageEntry.getSegmentsFrom(), segmentsTo, LineageEntryState.COMPLETED,
@@ -4068,6 +4069,31 @@ public void endReplaceSegments(String tableNameWithType, String segmentLineageEn
         tableNameWithType, segmentLineageEntryId);
   }
 
+  /**
+   * This method can be overridden to perform custom operations before updating the routing table
+   * to switch routing from the old segments (`segmentsFrom`) to the new segments (`segmentsTo`).
+   * One example usage of this method could be triggering a pageCache warmup operation on the server
+   * for the specified table and segments. For refresh tables, this ensures that the new segments
+   * are warmed up and ready for query availability before the routing table is updated.
+   *
+   * Example:
+   * To warm up specific segments of the "salesData_OFFLINE" table:
+   *   - tableNameWithType: "salesData_OFFLINE"
+   *   - segmentsTo: ["newSegment1", "newSegment2", "newSegment3"]
+   *   - segmentsFrom: ["oldSegment1", "oldSegment2", "oldSegment3"]
+   *
+   * @param tableNameWithType The name and type of the table for which operations need to be performed
+   *                          before switching routing from `segmentsFrom` to `segmentsTo`.
+   * @param segmentsTo A list of new segments that need to be prepared (e.g., warmed up) before
+   *                   they are made available for querying.
+   * @param segmentsFrom A list of old segments that are currently routed for queries, in case some operation is needed
+   *                     on it.
+   */
+  protected void preSegmentReplaceUpdateRouting(String tableNameWithType, List<String> segmentsTo,
+      List<String> segmentsFrom) {
+    // No-op by default
+  }
+
   /**
    * List the segment lineage
    *
diff --git a/pinot-server/src/main/java/org/apache/pinot/server/starter/ServerInstance.java b/pinot-server/src/main/java/org/apache/pinot/server/starter/ServerInstance.java
index fde3f893fd60..be59b01b9b61 100644
--- a/pinot-server/src/main/java/org/apache/pinot/server/starter/ServerInstance.java
+++ b/pinot-server/src/main/java/org/apache/pinot/server/starter/ServerInstance.java
@@ -291,4 +291,8 @@ public InstanceRequestHandler getInstanceRequestHandler() {
   public HelixManager getHelixManager() {
     return _helixManager;
   }
+
+  public QueryScheduler getQueryScheduler() {
+    return _queryScheduler;
+  }
 }