Expose QueryScheduler from ServerInstance & Custom interface (apache#…

…14630)

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
    *

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;
+  }
 }