SYSTEM PRESHUTDOWN to allow graceful shutdown node

Author: ianton-ru

docs/en/sql-reference/statements/system.md

@@ -205,6 +205,12 @@ SYSTEM RELOAD USERS [ON CLUSTER cluster_name]
 
 Normally shuts down ClickHouse (like `service clickhouse-server stop` / `kill {$pid_clickhouse-server}`)
 
+## PRESHUTDOWN {#preshutdown}
+
+<CloudNotSupportedBadge/>
+
+Prepare node for graceful shutdown. Unregister in autodiscovered clusters, stop accepting distributed requests to object storages (s3Cluster, icebergCluster, etc.).
+
 ## KILL {#kill}
 
 Aborts ClickHouse process (like `kill -9 {$ pid_clickhouse-server}`)

src/Access/Common/AccessType.h

@@ -162,7 +162,7 @@ enum class AccessType : uint8_t
     \
     M(TABLE_ENGINE, "TABLE ENGINE", TABLE_ENGINE, ALL) \
     \
-    M(SYSTEM_SHUTDOWN, "SYSTEM KILL, SHUTDOWN", GLOBAL, SYSTEM) \
+    M(SYSTEM_SHUTDOWN, "SYSTEM KILL, SHUTDOWN, PRESHUTDOWN", GLOBAL, SYSTEM) \
     M(SYSTEM_DROP_DNS_CACHE, "SYSTEM DROP DNS, DROP DNS CACHE, DROP DNS", GLOBAL, SYSTEM_DROP_CACHE)  \
     M(SYSTEM_DROP_CONNECTIONS_CACHE, "SYSTEM DROP CONNECTIONS CACHE, DROP CONNECTIONS CACHE", GLOBAL, SYSTEM_DROP_CACHE)  \
     M(SYSTEM_PREWARM_MARK_CACHE, "SYSTEM PREWARM MARK, PREWARM MARK CACHE, PREWARM MARKS", GLOBAL, SYSTEM_DROP_CACHE) \

src/Interpreters/ClusterDiscovery.cpp

@@ -454,12 +454,30 @@ void ClusterDiscovery::registerInZk(zkutil::ZooKeeperPtr & zk, ClusterInfo & inf
         return;
     }
 
+    if (context->isPreShutdownCalled())
+    {
+        LOG_DEBUG(log, "PreShutdown called, skip self-registering current node {} in cluster {}", current_node_name, info.name);
+        return;
+    }
+
     LOG_DEBUG(log, "Registering current node {} in cluster {}", current_node_name, info.name);
 
     zk->createOrUpdate(node_path, info.current_node.serialize(), zkutil::CreateMode::Ephemeral);
     LOG_DEBUG(log, "Current node {} registered in cluster {}", current_node_name, info.name);
 }
 
+void ClusterDiscovery::unregisterFromZk(zkutil::ZooKeeperPtr & zk, ClusterInfo & info)
+{
+    if (info.current_node_is_observer)
+        return;
+
+    String node_path = getShardsListPath(info.zk_root) / current_node_name;
+    LOG_DEBUG(log, "Removing current node {} from cluster {}", current_node_name, info.name);
+
+    zk->remove(node_path);
+    LOG_DEBUG(log, "Current node {} removed from cluster {}", current_node_name, info.name);
+}
+
 void ClusterDiscovery::initialUpdate()
 {
     LOG_DEBUG(log, "Initializing");
@@ -505,6 +523,15 @@ void ClusterDiscovery::initialUpdate()
     is_initialized = true;
 }
 
+void ClusterDiscovery::unregisterAll()
+{
+    for (auto & [_, info] : clusters_info)
+    {
+        auto zk = context->getDefaultOrAuxiliaryZooKeeper(info.zk_name);
+        unregisterFromZk(zk, info);
+    }
+}
+
 void ClusterDiscovery::findDynamicClusters(
     std::unordered_map<String, ClusterDiscovery::ClusterInfo> & info,
     std::unordered_set<size_t> * unchanged_roots)

src/Interpreters/ClusterDiscovery.h

@@ -38,6 +38,8 @@ class ClusterDiscovery
 
     ~ClusterDiscovery();
 
+    void unregisterAll();
+
 private:
     struct NodeInfo
     {
@@ -125,6 +127,7 @@ class ClusterDiscovery
     void initialUpdate();
 
     void registerInZk(zkutil::ZooKeeperPtr & zk, ClusterInfo & info);
+    void unregisterFromZk(zkutil::ZooKeeperPtr & zk, ClusterInfo & info);
 
     Strings getNodeNames(zkutil::ZooKeeperPtr & zk,
                          const String & zk_root,

src/Interpreters/Context.cpp

@@ -566,6 +566,7 @@ struct ContextSharedPart : boost::noncopyable
     std::map<String, UInt16> server_ports;
 
     std::atomic<bool> shutdown_called = false;
+    std::atomic<bool> preshutdown_called = false;
 
     Stopwatch uptime_watch TSA_GUARDED_BY(mutex);
 
@@ -913,6 +914,11 @@ struct ContextSharedPart : boost::noncopyable
         total_memory_tracker.resetPageCache();
     }
 
+    void preShutdown()
+    {
+        preshutdown_called = true;
+    }
+
     bool hasTraceCollector() const
     {
         return trace_collector.has_value();
@@ -4481,7 +4487,6 @@ std::shared_ptr<Cluster> Context::getCluster(const std::string & cluster_name) c
     throw Exception(ErrorCodes::CLUSTER_DOESNT_EXIST, "Requested cluster '{}' not found", cluster_name);
 }
 
-
 std::shared_ptr<Cluster> Context::tryGetCluster(const std::string & cluster_name) const
 {
     std::shared_ptr<Cluster> res = nullptr;
@@ -4500,6 +4505,13 @@ std::shared_ptr<Cluster> Context::tryGetCluster(const std::string & cluster_name
     return res;
 }
 
+void Context::unregisterInDynamicClusters()
+{
+    std::lock_guard lock(shared->clusters_mutex);
+    if (!shared->cluster_discovery)
+        return;
+    shared->cluster_discovery->unregisterAll();
+}
 
 void Context::reloadClusterConfig() const
 {
@@ -5350,12 +5362,20 @@ void Context::stopServers(const ServerType & server_type) const
     shared->stop_servers_callback(server_type);
 }
 
-
 void Context::shutdown() TSA_NO_THREAD_SAFETY_ANALYSIS
 {
     shared->shutdown();
 }
 
+void Context::preShutdown()
+{
+    shared->preshutdown_called = true;
+}
+
+bool Context::isPreShutdownCalled() const
+{
+    return shared->preshutdown_called;
+}
 
 Context::ApplicationType Context::getApplicationType() const
 {

src/Interpreters/Context.h

@@ -1225,6 +1225,7 @@ class Context: public ContextData, public std::enable_shared_from_this<Context>
     size_t getClustersVersion() const;
 
     void startClusterDiscovery();
+    void unregisterInDynamicClusters();
 
     /// Sets custom cluster, but doesn't update configuration
     void setCluster(const String & cluster_name, const std::shared_ptr<Cluster> & cluster);
@@ -1335,6 +1336,10 @@ class Context: public ContextData, public std::enable_shared_from_this<Context>
 
     void shutdown();
 
+    /// Stop some works to allow graceful shutdown later
+    void preShutdown();
+    bool isPreShutdownCalled() const;
+
     bool isInternalQuery() const { return is_internal_query; }
     void setInternalQuery(bool internal) { is_internal_query = internal; }
 

src/Interpreters/InterpreterSystemQuery.cpp

@@ -330,6 +330,13 @@ BlockIO InterpreterSystemQuery::execute()
                 throw ErrnoException(ErrorCodes::CANNOT_KILL, "System call kill(0, SIGTERM) failed");
             break;
         }
+        case Type::PRESHUTDOWN:
+        {
+            getContext()->checkAccess(AccessType::SYSTEM_SHUTDOWN);
+            getContext()->preShutdown();
+            getContext()->unregisterInDynamicClusters();
+            break;
+        }
         case Type::KILL:
         {
             getContext()->checkAccess(AccessType::SYSTEM_SHUTDOWN);
@@ -1468,6 +1475,7 @@ AccessRightsElements InterpreterSystemQuery::getRequiredAccessForDDLOnCluster()
     switch (query.type)
     {
         case Type::SHUTDOWN:
+        case Type::PRESHUTDOWN:
         case Type::KILL:
         case Type::SUSPEND:
         {

src/Parsers/ASTSystemQuery.cpp

@@ -443,6 +443,7 @@ void ASTSystemQuery::formatImpl(WriteBuffer & ostr, const FormatSettings & setti
         }
         case Type::KILL:
         case Type::SHUTDOWN:
+        case Type::PRESHUTDOWN:
         case Type::DROP_DNS_CACHE:
         case Type::DROP_CONNECTIONS_CACHE:
         case Type::DROP_MMAP_CACHE:

src/Parsers/ASTSystemQuery.h

@@ -17,6 +17,7 @@ class ASTSystemQuery : public IAST, public ASTQueryWithOnCluster
     {
         UNKNOWN,
         SHUTDOWN,
+        PRESHUTDOWN,
         KILL,
         SUSPEND,
         DROP_DNS_CACHE,

src/QueryPipeline/RemoteQueryExecutor.cpp

@@ -984,6 +984,11 @@ void RemoteQueryExecutor::setProfileInfoCallback(ProfileInfoCallback callback)
     profile_info_callback = std::move(callback);
 }
 
+bool RemoteQueryExecutor::skipUnavailableShards() const
+{
+    return context->getSettingsRef()[Setting::skip_unavailable_shards];
+}
+
 bool RemoteQueryExecutor::needToSkipUnavailableShard() const
 {
     return context->getSettingsRef()[Setting::skip_unavailable_shards] && (0 == connections->size());