add deadline for grpc shutdown

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

Author: gengliqi

dbms/src/Flash/Mpp/MPPTaskManager.cpp

@@ -86,13 +86,7 @@ MPPGatherTaskSetPtr MPPQuery::addMPPGatherTaskSet(const MPPGatherId & gather_id)
 
 void MPPTaskMonitor::waitAllMPPTasksFinish(const std::unique_ptr<Context> & global_context)
 {
-    // The maximum seconds TiFlash will wait for all current MPP tasks to finish before shutting down
-    static constexpr const char * GRACEFUL_WIAT_BEFORE_SHUTDOWN = "flash.graceful_wait_before_shutdown";
-    // The default value of flash.graceful_wait_before_shutdown
-    static constexpr UInt64 DEFAULT_GRACEFUL_WAIT_BEFORE_SHUTDOWN = 600;
-    auto graceful_wait_before_shutdown = global_context->getUsersConfig()->getUInt64(
-        GRACEFUL_WIAT_BEFORE_SHUTDOWN,
-        DEFAULT_GRACEFUL_WAIT_BEFORE_SHUTDOWN);
+    auto graceful_wait_before_shutdown = global_context->getGracefulWaitBeforeShutdown();
     LOG_INFO(log, "Start to wait all MPPTasks to finish, timeout={}s", graceful_wait_before_shutdown);
     Stopwatch watch;
     // The first sleep before checking to reduce the chance of missing MPP tasks that are still in the process of being dispatched

dbms/src/Interpreters/Context.cpp

@@ -2055,6 +2055,17 @@ const std::unordered_set<uint64_t> * Context::getStoreIdBlockList() const
     return &shared->store_id_blocklist;
 }
 
+UInt64 Context::getGracefulWaitBeforeShutdown()
+{
+    // The maximum seconds TiFlash will wait for all current MPP tasks to finish before shutting down
+    static constexpr const char * GRACEFUL_WIAT_BEFORE_SHUTDOWN = "flash.graceful_wait_before_shutdown";
+    // The default value of flash.graceful_wait_before_shutdown
+    static constexpr UInt64 DEFAULT_GRACEFUL_WAIT_BEFORE_SHUTDOWN = 600;
+    return global_context->getUsersConfig()->getUInt64(
+        GRACEFUL_WIAT_BEFORE_SHUTDOWN,
+        DEFAULT_GRACEFUL_WAIT_BEFORE_SHUTDOWN);
+}
+
 // NOLINTNEXTLINE(readability-convert-member-functions-to-static)
 bool Context::initializeStoreIdBlockList(const String & comma_sep_string)
 {

dbms/src/Interpreters/Context.h

@@ -536,6 +536,8 @@ class Context
     bool initializeStoreIdBlockList(const String &);
     const std::unordered_set<uint64_t> * getStoreIdBlockList() const;
 
+    UInt64 getGracefulWaitBeforeShutdown();
+
 private:
     /** Check if the current client has access to the specified database.
       * If access is denied, throw an exception.

dbms/src/Server/FlashGrpcServerHolder.cpp

@@ -220,7 +220,8 @@ FlashGrpcServerHolder::~FlashGrpcServerHolder()
     {
         /// Shut down grpc server.
         LOG_INFO(log, "Begin to shut down flash grpc server");
-        flash_grpc_server->Shutdown();
+        auto deadline = std::chrono::system_clock::now() + std::chrono::seconds(grpc_shutdown_max_wait);
+        flash_grpc_server->Shutdown(deadline);
         *is_shutdown = true;
 
         for (auto & cq : cqs)

dbms/src/Server/FlashGrpcServerHolder.h

@@ -44,6 +44,8 @@ class FlashGrpcServerHolder
 
     std::unique_ptr<FlashService> & flashService();
 
+    void setMaxWaitDuringGRPCShutdown(UInt64 max_wait) { grpc_shutdown_max_wait = max_wait; }
+
 private:
     const LoggerPtr & log;
     std::shared_ptr<std::atomic<bool>> is_shutdown;
@@ -56,6 +58,7 @@ class FlashGrpcServerHolder
     std::vector<std::thread> cq_workers;
     std::vector<std::thread> notify_cq_workers;
     CollectProcInfoBackgroundTask background_task;
+    UInt64 grpc_shutdown_max_wait = 600;
 };
 
 } // namespace DB

dbms/src/Server/Server.cpp

@@ -1257,6 +1257,7 @@ try
         LOG_INFO(log, "Set unavailable for MPPTask");
         tmt_context.getMPPTaskManager()->setUnavailable();
         tmt_context.getMPPTaskManager()->getMPPTaskMonitor()->waitAllMPPTasksFinish(global_context);
+        flash_grpc_server_holder.setMaxWaitDuringGRPCShutdown(global_context->getGracefulWaitBeforeShutdown());
 
         {
             // Set limiters stopping and wakeup threads in waitting queue.