[BugFix] Fix Warehouse Idle Check bug when switching warehouse (backport #56552) (#56605)

Signed-off-by: gengjun-git <[email protected]>
Co-authored-by: gengjun-git <[email protected]>

Author: mergify[bot]

fe/fe-core/src/main/java/com/starrocks/qe/StmtExecutor.java

@@ -495,8 +495,10 @@ public void execute() throws Exception {
             httpResultSender = new HttpResultSender((HttpConnectContext) context);
         }
 
-        if (shouldMarkIdleCheck(parsedStmt)) {
-            WarehouseIdleChecker.increaseRunningSQL(context.getCurrentWarehouseId());
+        final boolean shouldMarkIdleCheck = shouldMarkIdleCheck(parsedStmt);
+        final long originWarehouseId = context.getCurrentWarehouseId();
+        if (shouldMarkIdleCheck) {
+            WarehouseIdleChecker.increaseRunningSQL(originWarehouseId);
         }
         try {
             context.getState().setIsQuery(parsedStmt instanceof QueryStatement);
@@ -807,8 +809,8 @@ public void execute() throws Exception {
             // restore session variable in connect context
             context.setSessionVariable(sessionVariableBackup);
 
-            if (shouldMarkIdleCheck(parsedStmt)) {
-                WarehouseIdleChecker.decreaseRunningSQL(context.getCurrentWarehouseId());
+            if (shouldMarkIdleCheck) {
+                WarehouseIdleChecker.decreaseRunningSQL(originWarehouseId);
             }
 
             recordExecStatsIntoContext();

fe/fe-core/src/main/java/com/starrocks/server/WarehouseManager.java

@@ -285,6 +285,17 @@ public Optional<Long> selectWorkerGroupByWarehouseId(long warehouseId) {
         return workerGroupId;
     }
 
+    public long getWarehouseResumeTime(long warehouseId) {
+        try (LockCloseable ignored = new LockCloseable(rwLock.readLock())) {
+            Warehouse warehouse = idToWh.get(warehouseId);
+            if (warehouse == null) {
+                return -1;
+            } else {
+                return warehouse.getResumeTime();
+            }
+        }
+    }
+
     private Optional<Long> selectWorkerGroupInternal(long warehouseId) {
         Warehouse warehouse = getWarehouse(warehouseId);
         List<Long> ids = warehouse.getWorkerGroupIds();

fe/fe-core/src/main/java/com/starrocks/warehouse/DefaultWarehouse.java

@@ -72,4 +72,9 @@ public List<List<String>> getWarehouseNodesInfo() {
     public ProcResult fetchResult() {
         return new BaseProcResult();
     }
+
+    @Override
+    public long getResumeTime() {
+        return -1L;
+    }
 }

fe/fe-core/src/main/java/com/starrocks/warehouse/Warehouse.java

@@ -56,6 +56,8 @@ public void write(DataOutput out) throws IOException {
         Text.writeString(out, json);
     }
 
+    public abstract long getResumeTime();
+
     public abstract Long getAnyWorkerGroupId();
 
     public abstract List<Long> getWorkerGroupIds();

fe/fe-core/src/main/java/com/starrocks/warehouse/WarehouseIdleChecker.java

@@ -76,7 +76,16 @@ protected void runAfterCatalogReady() {
             if (runningJobCnt == 0
                     && lastFinishedJobTime <
                         System.currentTimeMillis() - Config.warehouse_idle_check_interval_seconds * 2000) {
-                warehouseIdleTime.putIfAbsent(wId, System.currentTimeMillis());
+                long resumeTime = GlobalStateMgr.getCurrentState().getWarehouseMgr().getWarehouseResumeTime(wId);
+                warehouseIdleTime.compute(wId, (k, v) -> {
+                    // If this is the first time to become idle, set idleTime to now.
+                    // If resumed during an idle period, change idleTime to resumeTime.
+                    if (v == null) {
+                        return System.currentTimeMillis();
+                    } else {
+                        return v < resumeTime ? resumeTime : v;
+                    }
+                });
                 LOG.info("warehouse: {} is idle, idle start time: {}",
                         wId, TimeUtils.longToTimeString(warehouseIdleTime.get(wId)));
             } else {