Support graceful shutdown in WN (#10430)

ref #10266, close #10425

Support graceful shutdown in WN

Signed-off-by: gengliqi <gengliqiii@gmail.com>

Author: gengliqi

dbms/src/Flash/Mpp/MPPTaskManager.cpp

@@ -22,7 +22,9 @@
 #include <Flash/Mpp/MPPTaskManager.h>
 #include <Interpreters/Context.h>
 #include <Interpreters/ProcessList.h>
+#include <Interpreters/SharedContexts/Disagg.h>
 #include <Interpreters/executeQuery.h>
+#include <Storages/DeltaMerge/Remote/WNDisaggSnapshotManager.h>
 #include <fmt/core.h>
 
 #include <magic_enum.hpp>
@@ -86,39 +88,58 @@ MPPGatherTaskSetPtr MPPQuery::addMPPGatherTaskSet(const MPPGatherId & gather_id)
 
 void MPPTaskMonitor::waitAllMPPTasksFinish(const std::unique_ptr<Context> & context)
 {
-    // The maximum seconds TiFlash will wait for all current MPP tasks to finish before shutting down
+    // The maximum seconds TiFlash will wait for all current MPP tasks or disagg snapshots to finish before shutting down
     static constexpr const char * GRACEFUL_WAIT_SHUTDOWN_TIMEOUT = "flash.graceful_wait_shutdown_timeout";
     // The default value of flash.graceful_wait_shutdown_timeout
     static constexpr UInt64 DEFAULT_GRACEFUL_WAIT_SHUTDOWN_TIMEOUT = 600;
     auto graceful_wait_shutdown_timeout
         = context->getUsersConfig()->getUInt64(GRACEFUL_WAIT_SHUTDOWN_TIMEOUT, DEFAULT_GRACEFUL_WAIT_SHUTDOWN_TIMEOUT);
-    LOG_INFO(log, "Start to wait all MPP tasks to finish, timeout={}s", graceful_wait_shutdown_timeout);
+
+    bool is_disagg_storage_mode = context->getSharedContextDisagg()->isDisaggregatedStorageMode();
+    const String & wait_target = is_disagg_storage_mode ? "disagg snapshots" : "MPP tasks";
+    LOG_INFO(log, "Start to wait all {} to finish, timeout={}s", wait_target, graceful_wait_shutdown_timeout);
+
     UInt64 graceful_wait_shutdown_timeout_ms = graceful_wait_shutdown_timeout * 1000;
     Stopwatch watch;
     // The first sleep before checking to reduce the chance of missing MPP tasks that are still in the process of being dispatched
-    std::this_thread::sleep_for(std::chrono::seconds(1));
+    std::this_thread::sleep_for(std::chrono::seconds(2));
     bool all_tasks_finished = false;
     while (true)
     {
         auto elapsed_ms = watch.elapsedMilliseconds();
-        if (!all_tasks_finished)
+        if (is_disagg_storage_mode)
         {
-            std::unique_lock lock(mu);
-            if (monitored_tasks.empty())
-                all_tasks_finished = true;
+            // For write node under disagg arch, should wait for all disagg establish task rpc being finished
+            // and all snapshot being released.
+            // Otherwise compute nodes may meet error when calling `FetchDisaggPages` on write nodes.
+            if (GET_METRIC(tiflash_coprocessor_handling_request_count, type_disagg_establish_task).Value() == 0
+                && context->getSharedContextDisagg()->wn_snapshot_manager->getActiveSnapshotCount() == 0)
+            {
+                LOG_INFO(log, "All disagg snapshots have finished after {}ms", elapsed_ms);
+                break;
+            }
         }
-        if (all_tasks_finished)
+        else
         {
-            // Also needs to check if all MPP gRPC connections are finished
-            if (GET_METRIC(tiflash_coprocessor_handling_request_count, type_mpp_establish_conn).Value() == 0)
+            if (!all_tasks_finished)
             {
-                LOG_INFO(log, "All MPP tasks have finished after {}ms", elapsed_ms);
-                break;
+                std::unique_lock lock(mu);
+                if (monitored_tasks.empty())
+                    all_tasks_finished = true;
+            }
+            if (all_tasks_finished)
+            {
+                // Also needs to check if all MPP gRPC connections are finished
+                if (GET_METRIC(tiflash_coprocessor_handling_request_count, type_mpp_establish_conn).Value() == 0)
+                {
+                    LOG_INFO(log, "All MPP tasks have finished after {}ms", elapsed_ms);
+                    break;
+                }
             }
         }
         if (elapsed_ms >= graceful_wait_shutdown_timeout_ms)
         {
-            LOG_WARNING(log, "Timed out waiting for all MPP tasks to finish after {}ms", elapsed_ms);
+            LOG_WARNING(log, "Timed out waiting for all {} to finish after {}ms", wait_target, elapsed_ms);
             break;
         }
         std::this_thread::sleep_for(std::chrono::milliseconds(200));

dbms/src/Storages/DeltaMerge/Remote/WNDisaggSnapshotManager.cpp

@@ -81,6 +81,11 @@ bool WNDisaggSnapshotManager::unregisterSnapshot(const DisaggTaskId & task_id)
     });
 }
 
+size_t WNDisaggSnapshotManager::getActiveSnapshotCount() const
+{
+    return snapshots.withShared([&](auto & snapshots) { return snapshots.size(); });
+}
+
 void WNDisaggSnapshotManager::clearExpiredSnapshots()
 {
     Timepoint now = Clock::now();

dbms/src/Storages/DeltaMerge/Remote/WNDisaggSnapshotManager.h

@@ -102,6 +102,8 @@ class WNDisaggSnapshotManager
 
     bool unregisterSnapshotIfEmpty(const DisaggTaskId & task_id);
 
+    size_t getActiveSnapshotCount() const;
+
     DISALLOW_COPY_AND_MOVE(WNDisaggSnapshotManager);
 
 private: