feat: Catch and handle SIGTERM:

docs/src/user-docs/guides-quick-start.md

@@ -222,6 +222,21 @@ NOTE:
 If you used a different set of arguments to set up the storage backend, ensure you update the
 storage backend URL in the command.
 
+# Exiting the cluster
+
+To stop the cluster, send `SIGTERM` to the scheduler and all workers.
+
+The scheduler finishes the current tasks (e.g., scheduling tasks to workers, garbage collection,
+failure recovery, etc.), then exits with `SIGTERM`.
+
+When a worker receives `SIGTERM`, if it has no task executor, it exits immediately with `SIGTERM`.
+
+If the worker has a task executor, it sends a `SIGTERM` to the task executor and waits for it to
+exit.  
+Normally, the task executor exits immediately, and the worker sets the task as failed. If the task
+executor has a signal handler installed and catches `SIGTERM`, it completes the execution of the
+task, and the worker handles the task output as usual. Then the worker exits with `SIGTERM`.
+
 # Next steps
 
 In future guides, we'll explain how to write more complex tasks, as well as how to leverage Spider's

src/spider/CMakeLists.txt

@@ -56,6 +56,8 @@ target_link_libraries(
 target_link_libraries(spider_core PRIVATE fmt::fmt)
 
 set(SPIDER_WORKER_SOURCES
+    worker/ChildPid.hpp
+    worker/ChildPid.cpp
     worker/DllLoader.hpp
     worker/DllLoader.cpp
     worker/Process.hpp
@@ -67,7 +69,8 @@ set(SPIDER_WORKER_SOURCES
     worker/message_pipe.hpp
     worker/WorkerClient.hpp
     worker/WorkerClient.cpp
-    utils/StopToken.hpp
+    utils/StopFlag.hpp
+    utils/StopFlag.cpp
     CACHE INTERNAL
     "spider worker source files"
 )
@@ -126,7 +129,8 @@ set(SPIDER_SCHEDULER_SOURCES
     scheduler/SchedulerMessage.hpp
     scheduler/SchedulerServer.cpp
     scheduler/SchedulerServer.hpp
-    utils/StopToken.hpp
+    utils/StopFlag.hpp
+    utils/StopFlag.cpp
     CACHE INTERNAL
     "spider scheduler source files"
 )

src/spider/scheduler/SchedulerServer.cpp

@@ -19,7 +19,7 @@
 #include "../storage/DataStorage.hpp"
 #include "../storage/MetadataStorage.hpp"
 #include "../storage/StorageConnection.hpp"
-#include "../utils/StopToken.hpp"
+#include "../utils/StopFlag.hpp"
 #include "SchedulerMessage.hpp"
 #include "SchedulerPolicy.hpp"
 
@@ -29,15 +29,13 @@ SchedulerServer::SchedulerServer(
         std::shared_ptr<SchedulerPolicy> policy,
         std::shared_ptr<core::MetadataStorage> metadata_store,
         std::shared_ptr<core::DataStorage> data_store,
-        std::shared_ptr<core::StorageConnection> conn,
-        core::StopToken& stop_token
+        std::shared_ptr<core::StorageConnection> conn
 )
         : m_port{port},
           m_policy{std::move(policy)},
           m_metadata_store{std::move(metadata_store)},
           m_data_store{std::move(data_store)},
-          m_conn{std::move(conn)},
-          m_stop_token{stop_token} {
+          m_conn{std::move(conn)} {
     boost::asio::co_spawn(m_context, receive_message(), boost::asio::detached);
     std::lock_guard const lock{m_mutex};
     m_thread = std::make_unique<std::thread>([&] { m_context.run(); });
@@ -96,7 +94,7 @@ auto SchedulerServer::receive_message() -> boost::asio::awaitable<void> {
         co_return;
     } catch (boost::system::system_error& e) {
         spdlog::error("Fail to accept connection: {}", e.what());
-        m_stop_token.request_stop();
+        spider::core::StopFlag::request_stop();
         co_return;
     }
 }

src/spider/scheduler/SchedulerServer.hpp

@@ -9,7 +9,6 @@
 #include "../storage/DataStorage.hpp"
 #include "../storage/MetadataStorage.hpp"
 #include "../storage/StorageConnection.hpp"
-#include "../utils/StopToken.hpp"
 #include "SchedulerPolicy.hpp"
 
 namespace spider::scheduler {
@@ -27,8 +26,7 @@ class SchedulerServer {
             std::shared_ptr<SchedulerPolicy> policy,
             std::shared_ptr<core::MetadataStorage> metadata_store,
             std::shared_ptr<core::DataStorage> data_store,
-            std::shared_ptr<core::StorageConnection> conn,
-            core::StopToken& stop_token
+            std::shared_ptr<core::StorageConnection> conn
     );
 
     auto pause() -> void;
@@ -51,8 +49,6 @@ class SchedulerServer {
 
     std::mutex m_mutex;
     std::unique_ptr<std::thread> m_thread;
-
-    core::StopToken& m_stop_token;
 };
 }  // namespace spider::scheduler
 

src/spider/scheduler/scheduler.cpp

@@ -1,4 +1,6 @@
+#include <cerrno>
 #include <chrono>
+#include <csignal>
 #include <functional>
 #include <memory>
 #include <string>
@@ -26,20 +28,31 @@
 #include "../storage/mysql/MySqlStorageFactory.hpp"
 #include "../storage/StorageConnection.hpp"
 #include "../storage/StorageFactory.hpp"
-#include "../utils/StopToken.hpp"
+#include "../utils/StopFlag.hpp"
 #include "FifoPolicy.hpp"
 #include "SchedulerPolicy.hpp"
 #include "SchedulerServer.hpp"
 
 constexpr int cCmdArgParseErr = 1;
-constexpr int cStorageConnectionErr = 2;
-constexpr int cSchedulerAddrErr = 3;
-constexpr int cStorageErr = 4;
+constexpr int cSignalHandleErr = 2;
+constexpr int cStorageConnectionErr = 3;
+constexpr int cSchedulerAddrErr = 4;
+constexpr int cStorageErr = 5;
 
 constexpr int cCleanupInterval = 1000;
 constexpr int cRetryCount = 5;
 
 namespace {
+/*
+ * Signal handler for SIGTERM. Sets the stop flag to request a stop.
+ * @param signal The signal number.
+ */
+auto stop_scheduler_handler(int signal) -> void {
+    if (SIGTERM == signal) {
+        spider::core::StopFlag::request_stop();
+    }
+}
+
 auto parse_args(int const argc, char** argv) -> boost::program_options::variables_map {
     boost::program_options::options_description desc;
     desc.add_options()("help", "spider scheduler");
@@ -72,11 +85,10 @@ auto parse_args(int const argc, char** argv) -> boost::program_options::variable
 auto heartbeat_loop(
         std::shared_ptr<spider::core::StorageFactory> const& storage_factory,
         std::shared_ptr<spider::core::MetadataStorage> const& metadata_store,
-        spider::core::Scheduler const& scheduler,
-        spider::core::StopToken& stop_token
+        spider::core::Scheduler const& scheduler
 ) -> void {
     int fail_count = 0;
-    while (!stop_token.stop_requested()) {
+    while (!spider::core::StopFlag::is_stop_requested()) {
         std::this_thread::sleep_for(std::chrono::seconds(1));
         spdlog::debug("Updating heartbeat");
         std::variant<std::unique_ptr<spider::core::StorageConnection>, spider::core::StorageErr>
@@ -102,18 +114,17 @@ auto heartbeat_loop(
             fail_count = 0;
         }
         if (fail_count >= cRetryCount - 1) {
-            stop_token.request_stop();
+            spider::core::StopFlag::request_stop();
             break;
         }
     }
 }
 
 auto cleanup_loop(
         std::shared_ptr<spider::core::StorageFactory> const& storage_factory,
-        std::shared_ptr<spider::core::DataStorage> const& data_store,
-        spider::core::StopToken const& stop_token
+        std::shared_ptr<spider::core::DataStorage> const& data_store
 ) -> void {
-    while (!stop_token.stop_requested()) {
+    while (!spider::core::StopFlag::is_stop_requested()) {
         std::this_thread::sleep_for(std::chrono::seconds(cCleanupInterval));
         spdlog::debug("Starting cleanup");
         std::variant<std::unique_ptr<spider::core::StorageConnection>, spider::core::StorageErr>
@@ -133,6 +144,8 @@ auto cleanup_loop(
         spdlog::debug("Finished cleanup");
     }
 }
+
+constexpr int cSignalExitBase = 128;
 }  // namespace
 
 // NOLINTNEXTLINE(bugprone-exception-escape)
@@ -171,6 +184,18 @@ auto main(int argc, char** argv) -> int {
         return cCmdArgParseErr;
     }
 
+    // Ignore SIGTERM
+    // NOLINTBEGIN(misc-include-cleaner)
+    struct sigaction sig_action{};
+    sig_action.sa_handler = stop_scheduler_handler;
+    sigemptyset(&sig_action.sa_mask);
+    sig_action.sa_flags |= SA_RESTART;
+    if (0 != sigaction(SIGTERM, &sig_action, nullptr)) {
+        spdlog::error("Fail to install signal handler for SIGTERM: errno {}", errno);
+        return cSignalHandleErr;
+    }
+    // NOLINTEND(misc-include-cleaner)
+
     // Create storages
     std::shared_ptr<spider::core::StorageFactory> const storage_factory
             = std::make_unique<spider::core::MySqlStorageFactory>(storage_url);
@@ -215,34 +240,26 @@ auto main(int argc, char** argv) -> int {
     }
 
     // Start scheduler server
-    spider::core::StopToken stop_token;
     std::shared_ptr<spider::scheduler::SchedulerPolicy> const policy
             = std::make_shared<spider::scheduler::FifoPolicy>(
                     scheduler_id,
                     metadata_store,
                     data_store,
                     conn
             );
-    spider::scheduler::SchedulerServer
-            server{port, policy, metadata_store, data_store, conn, stop_token};
+    spider::scheduler::SchedulerServer server{port, policy, metadata_store, data_store, conn};
 
     try {
         // Start a thread that periodically updates the scheduler's heartbeat
         std::thread heartbeat_thread{
                 heartbeat_loop,
                 std::cref(storage_factory),
                 std::cref(metadata_store),
-                std::ref(scheduler),
-                std::ref(stop_token),
+                std::ref(scheduler)
         };
 
         // Start a thread that periodically starts cleanup
-        std::thread cleanup_thread{
-                cleanup_loop,
-                std::cref(storage_factory),
-                std::cref(data_store),
-                std::ref(stop_token)
-        };
+        std::thread cleanup_thread{cleanup_loop, std::cref(storage_factory), std::cref(data_store)};
 
         heartbeat_thread.join();
         cleanup_thread.join();
@@ -251,5 +268,10 @@ auto main(int argc, char** argv) -> int {
         spdlog::error("Failed to join thread: {}", e.what());
     }
 
+    // If SIGTERM was caught and StopFlag is requested, set the exit value corresponding to SIGTERM.
+    if (spider::core::StopFlag::is_stop_requested()) {
+        return cSignalExitBase + SIGTERM;
+    }
+
     return 0;
 }

src/spider/utils/StopFlag.cpp

@@ -0,0 +1,19 @@
+#include "StopFlag.hpp"
+
+#include <atomic>
+
+namespace spider::core {
+auto StopFlag::request_stop() -> void {
+    m_stop.test_and_set();
+}
+
+auto StopFlag::is_stop_requested() -> bool {
+    return m_stop.test();
+}
+
+auto StopFlag::reset() -> void {
+    m_stop.clear();
+}
+
+std::atomic_flag StopFlag::m_stop = ATOMIC_FLAG_INIT;
+}  // namespace spider::core

src/spider/utils/StopFlag.hpp

@@ -0,0 +1,46 @@
+#ifndef SPIDER_UTILS_STOPTOKEN_HPP
+#define SPIDER_UTILS_STOPTOKEN_HPP
+
+#include <atomic>
+
+namespace spider::core {
+/**
+ * @brief A singleton class that provides a stop flag for threads and signal handlers.
+ *
+ * User can call request_stop() to set the stop flag, and check if the stop flag is set.
+ * This class is thread-safe and signal-safe.
+ */
+class StopFlag {
+public:
+    /*
+     * Request to token owners to stop.
+     */
+    static auto request_stop() -> void;
+
+    /*
+     * @return A boolean indicating whether the stop was requested.
+     */
+    [[nodiscard]] static auto is_stop_requested() -> bool;
+
+    /*
+     * Reset the stop token.
+     */
+    static auto reset() -> void;
+
+    // Delete constructor
+    StopFlag() = delete;
+    // Delete copy constructor and assignment operator
+    StopFlag(StopFlag const&) = delete;
+    auto operator=(StopFlag const&) -> StopFlag& = delete;
+    // Delete move constructor and assignment operator
+    StopFlag(StopFlag&&) = delete;
+    auto operator=(StopFlag&&) -> StopFlag& = delete;
+    // Default destructor
+    ~StopFlag() = default;
+
+private:
+    static std::atomic_flag m_stop;
+};
+}  // namespace spider::core
+
+#endif

src/spider/worker/ChildPid.cpp

@@ -0,0 +1,17 @@
+#include "ChildPid.hpp"
+
+#include <unistd.h>
+
+#include <csignal>
+
+namespace spider::core {
+auto ChildPid::get_pid() -> std::sig_atomic_t {
+    return m_pid;
+}
+
+auto ChildPid::set_pid(pid_t const pid) -> void {
+    m_pid = pid;
+}
+
+std::sig_atomic_t volatile ChildPid::m_pid = 0;
+}  // namespace spider::core