Atomic asyncronous stop/start swarm

Author: ianton-ru

src/Interpreters/ClusterDiscovery.cpp

@@ -108,6 +108,13 @@ class ClusterDiscovery::Flags
         cv.notify_one();
     }
 
+    void wakeup()
+    {
+        std::unique_lock<std::mutex> lk(mu);
+        any_need_update = true;
+        cv.notify_one();
+    }
+
 private:
     std::condition_variable cv;
     std::mutex mu;
@@ -527,20 +534,14 @@ void ClusterDiscovery::initialUpdate()
 
 void ClusterDiscovery::registerAll()
 {
-    for (auto & [_, info] : clusters_info)
-    {
-        auto zk = context->getDefaultOrAuxiliaryZooKeeper(info.zk_name);
-        registerInZk(zk, info);
-    }
+    register_change_flag = RegisterChangeFlag::RCF_REGISTER_ALL;
+    clusters_to_update->wakeup();
 }
 
 void ClusterDiscovery::unregisterAll()
 {
-    for (auto & [_, info] : clusters_info)
-    {
-        auto zk = context->getDefaultOrAuxiliaryZooKeeper(info.zk_name);
-        unregisterFromZk(zk, info);
-    }
+    register_change_flag = RegisterChangeFlag::RCF_UNREGISTER_ALL;
+    clusters_to_update->wakeup();
 }
 
 void ClusterDiscovery::findDynamicClusters(
@@ -766,6 +767,27 @@ bool ClusterDiscovery::runMainThread(std::function<void()> up_to_date_callback)
         {
             up_to_date_callback();
         }
+
+        RegisterChangeFlag flag = register_change_flag.exchange(RegisterChangeFlag::RCF_NONE);
+
+        if (flag == RegisterChangeFlag::RCF_REGISTER_ALL)
+        {
+            LOG_DEBUG(log, "Register in all dynamic clusters");
+            for (auto & [_, info] : clusters_info)
+            {
+                auto zk = context->getDefaultOrAuxiliaryZooKeeper(info.zk_name);
+                registerInZk(zk, info);
+            }
+        }
+        else if (flag == RegisterChangeFlag::RCF_UNREGISTER_ALL)
+        {
+            LOG_DEBUG(log, "Unregister in all dynamic clusters");
+            for (auto & [_, info] : clusters_info)
+            {
+                auto zk = context->getDefaultOrAuxiliaryZooKeeper(info.zk_name);
+                unregisterFromZk(zk, info);
+            }
+        }
     }
     LOG_DEBUG(log, "Worker thread stopped");
     return finished;

src/Interpreters/ClusterDiscovery.h

@@ -211,6 +211,15 @@ class ClusterDiscovery
     std::shared_ptr<std::vector<std::shared_ptr<MulticlusterDiscovery>>> multicluster_discovery_paths;
 
     MultiVersion<Macros>::Version macros;
+
+    enum RegisterChangeFlag
+    {
+        RCF_NONE,
+        RCF_REGISTER_ALL,
+        RCF_UNREGISTER_ALL,
+    };
+
+    std::atomic<RegisterChangeFlag> register_change_flag = RegisterChangeFlag::RCF_NONE;
 };
 
 }

src/Interpreters/Context.cpp

@@ -737,6 +737,7 @@ struct ContextSharedPart : boost::noncopyable
     void shutdown() TSA_NO_THREAD_SAFETY_ANALYSIS
     {
         swarm_mode_enabled = false;
+        CurrentMetrics::set(CurrentMetrics::IsSwarmModeEnabled, 0);
         bool is_shutdown_called = shutdown_called.exchange(true);
         if (is_shutdown_called)
             return;
@@ -4502,15 +4503,15 @@ std::shared_ptr<Cluster> Context::tryGetCluster(const std::string & cluster_name
     return res;
 }
 
-void Context::unregisterInDynamicClusters()
+void Context::unregisterInAutodiscoveryClusters()
 {
     std::lock_guard lock(shared->clusters_mutex);
     if (!shared->cluster_discovery)
         return;
     shared->cluster_discovery->unregisterAll();
 }
 
-void Context::registerInDynamicClusters()
+void Context::registerInAutodiscoveryClusters()
 {
     std::lock_guard lock(shared->clusters_mutex);
     if (!shared->cluster_discovery)
@@ -5372,16 +5373,24 @@ void Context::shutdown() TSA_NO_THREAD_SAFETY_ANALYSIS
     shared->shutdown();
 }
 
-void Context::stopSwarmMode()
+bool Context::stopSwarmMode()
 {
-    CurrentMetrics::set(CurrentMetrics::IsSwarmModeEnabled, 0);
-    shared->swarm_mode_enabled = false;
+    bool expected_is_enabled = true;
+    bool is_stopped_now = shared->swarm_mode_enabled.compare_exchange_strong(expected_is_enabled, false);
+    if (is_stopped_now)
+        CurrentMetrics::set(CurrentMetrics::IsSwarmModeEnabled, 0);
+    // return true if stop successful
+    return is_stopped_now;
 }
 
-void Context::startSwarmMode()
+bool Context::startSwarmMode()
 {
-    shared->swarm_mode_enabled = true;
-    CurrentMetrics::set(CurrentMetrics::IsSwarmModeEnabled, 1);
+    bool expected_is_enabled = false;
+    bool is_started_now = shared->swarm_mode_enabled.compare_exchange_strong(expected_is_enabled, true);
+    if (is_started_now)
+        CurrentMetrics::set(CurrentMetrics::IsSwarmModeEnabled, 1);
+    // return true if start successful
+    return is_started_now;
 }
 
 bool Context::isSwarmModeEnabled() const

src/Interpreters/Context.h

@@ -1225,8 +1225,8 @@ class Context: public ContextData, public std::enable_shared_from_this<Context>
     size_t getClustersVersion() const;
 
     void startClusterDiscovery();
-    void registerInDynamicClusters();
-    void unregisterInDynamicClusters();
+    void registerInAutodiscoveryClusters();
+    void unregisterInAutodiscoveryClusters();
 
     /// Sets custom cluster, but doesn't update configuration
     void setCluster(const String & cluster_name, const std::shared_ptr<Cluster> & cluster);
@@ -1337,9 +1337,13 @@ class Context: public ContextData, public std::enable_shared_from_this<Context>
 
     void shutdown();
 
-    /// Stop some works to allow graceful shutdown later
-    void stopSwarmMode();
-    void startSwarmMode();
+    /// Stop some works to allow graceful shutdown later.
+    /// Returns true if stop successful.
+    bool stopSwarmMode();
+    /// Resume some works if we change our mind.
+    /// Returns true if start successful.
+    bool startSwarmMode();
+    /// Return current swarm mode state.
     bool isSwarmModeEnabled() const;
 
     bool isInternalQuery() const { return is_internal_query; }

src/Interpreters/InterpreterSystemQuery.cpp

@@ -696,15 +696,15 @@ BlockIO InterpreterSystemQuery::execute()
         case Type::STOP_SWARM_MODE:
         {
             getContext()->checkAccess(AccessType::SYSTEM_SWARM);
-            getContext()->stopSwarmMode();
-            getContext()->unregisterInDynamicClusters();
+            if (getContext()->stopSwarmMode())
+                getContext()->unregisterInAutodiscoveryClusters();
             break;
         }
         case Type::START_SWARM_MODE:
         {
             getContext()->checkAccess(AccessType::SYSTEM_SWARM);
-            getContext()->registerInDynamicClusters();
-            getContext()->startSwarmMode();
+            if (getContext()->startSwarmMode())
+                getContext()->registerInAutodiscoveryClusters();
             break;
         }
         case Type::STOP_FETCHES:

src/QueryPipeline/RemoteQueryExecutorReadContext.cpp

@@ -74,6 +74,9 @@ void RemoteQueryExecutorReadContext::Task::run(AsyncCallback async_callback, Sus
         }
         catch (const Exception & e)
         {
+            /// If cluster node unxepectedly shutted down (kill/segfault/power off/etc.) socket just closes.
+            /// If initiator did not process any data packets before, this fact can be ignored.
+            /// Unprocessed tasks will be executed on other nodes.
             if (e.code() == ErrorCodes::ATTEMPT_TO_READ_AFTER_EOF
                 && !read_context.has_data_packets.load() && read_context.executor.skipUnavailableShards())
             {