diff --git a/docs/src/user-docs/guides-quick-start.md b/docs/src/user-docs/guides-quick-start.md index d1d6e09ff..746da12ac 100644 --- a/docs/src/user-docs/guides-quick-start.md +++ b/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 diff --git a/src/spider/CMakeLists.txt b/src/spider/CMakeLists.txt index db0a773c0..a66f13b72 100644 --- a/src/spider/CMakeLists.txt +++ b/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" ) diff --git a/src/spider/scheduler/SchedulerServer.cpp b/src/spider/scheduler/SchedulerServer.cpp index 98913ab00..1324189e7 100644 --- a/src/spider/scheduler/SchedulerServer.cpp +++ b/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 policy, std::shared_ptr metadata_store, std::shared_ptr data_store, - std::shared_ptr conn, - core::StopToken& stop_token + std::shared_ptr 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([&] { m_context.run(); }); @@ -96,7 +94,7 @@ auto SchedulerServer::receive_message() -> boost::asio::awaitable { 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; } } diff --git a/src/spider/scheduler/SchedulerServer.hpp b/src/spider/scheduler/SchedulerServer.hpp index 3aa8cd318..933a4ac91 100644 --- a/src/spider/scheduler/SchedulerServer.hpp +++ b/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 policy, std::shared_ptr metadata_store, std::shared_ptr data_store, - std::shared_ptr conn, - core::StopToken& stop_token + std::shared_ptr conn ); auto pause() -> void; @@ -51,8 +49,6 @@ class SchedulerServer { std::mutex m_mutex; std::unique_ptr m_thread; - - core::StopToken& m_stop_token; }; } // namespace spider::scheduler diff --git a/src/spider/scheduler/scheduler.cpp b/src/spider/scheduler/scheduler.cpp index 4dd39ec44..562042870 100644 --- a/src/spider/scheduler/scheduler.cpp +++ b/src/spider/scheduler/scheduler.cpp @@ -1,4 +1,6 @@ +#include #include +#include #include #include #include @@ -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 const& storage_factory, std::shared_ptr 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, spider::core::StorageErr> @@ -102,7 +114,7 @@ auto heartbeat_loop( fail_count = 0; } if (fail_count >= cRetryCount - 1) { - stop_token.request_stop(); + spider::core::StopFlag::request_stop(); break; } } @@ -110,10 +122,9 @@ auto heartbeat_loop( auto cleanup_loop( std::shared_ptr const& storage_factory, - std::shared_ptr const& data_store, - spider::core::StopToken const& stop_token + std::shared_ptr 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, 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 const storage_factory = std::make_unique(storage_url); @@ -215,7 +240,6 @@ auto main(int argc, char** argv) -> int { } // Start scheduler server - spider::core::StopToken stop_token; std::shared_ptr const policy = std::make_shared( scheduler_id, @@ -223,8 +247,7 @@ auto main(int argc, char** argv) -> int { 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 @@ -232,17 +255,11 @@ auto main(int argc, char** argv) -> int { 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; } diff --git a/src/spider/utils/StopFlag.cpp b/src/spider/utils/StopFlag.cpp new file mode 100644 index 000000000..326e1f4d6 --- /dev/null +++ b/src/spider/utils/StopFlag.cpp @@ -0,0 +1,19 @@ +#include "StopFlag.hpp" + +#include + +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 diff --git a/src/spider/utils/StopFlag.hpp b/src/spider/utils/StopFlag.hpp new file mode 100644 index 000000000..e82a2cc62 --- /dev/null +++ b/src/spider/utils/StopFlag.hpp @@ -0,0 +1,46 @@ +#ifndef SPIDER_UTILS_STOPTOKEN_HPP +#define SPIDER_UTILS_STOPTOKEN_HPP + +#include + +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 diff --git a/src/spider/utils/StopToken.hpp b/src/spider/utils/StopToken.hpp deleted file mode 100644 index 28abb81c3..000000000 --- a/src/spider/utils/StopToken.hpp +++ /dev/null @@ -1,22 +0,0 @@ -#ifndef SPIDER_UTILS_STOPTOKEN_HPP -#define SPIDER_UTILS_STOPTOKEN_HPP - -#include - -namespace spider::core { -class StopToken { -public: - StopToken() : m_stop{false} {} - - auto request_stop() -> void { m_stop = true; } - - [[nodiscard]] auto stop_requested() const -> bool { return m_stop; } - - auto reset() -> void { m_stop = false; } - -private: - std::atomic m_stop; -}; -} // namespace spider::core - -#endif diff --git a/src/spider/worker/ChildPid.cpp b/src/spider/worker/ChildPid.cpp new file mode 100644 index 000000000..af95e1ade --- /dev/null +++ b/src/spider/worker/ChildPid.cpp @@ -0,0 +1,17 @@ +#include "ChildPid.hpp" + +#include + +#include + +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 diff --git a/src/spider/worker/ChildPid.hpp b/src/spider/worker/ChildPid.hpp new file mode 100644 index 000000000..629e07b34 --- /dev/null +++ b/src/spider/worker/ChildPid.hpp @@ -0,0 +1,44 @@ +#ifndef SPIDER_WORKER_CHILDPID_HPP +#define SPIDER_WORKER_CHILDPID_HPP + +#include + +#include + +namespace spider::core { +/** + * @brief A singleton class to manage the child process ID for signal handler. + * + * User can set the child process ID using set_pid() method, and retrieve it using get_pid() method. + * This class is signal-safe but is **not** thread-safe. + */ +class ChildPid { +public: + /* + * @return The process ID of the child process. + */ + [[nodiscard]] static auto get_pid() -> std::sig_atomic_t; + + /* + * @param pid The process ID to set. + */ + static auto set_pid(pid_t pid) -> void; + + // Delete constructor + ChildPid() = delete; + // Delete copy constructor and assignment operator + ChildPid(ChildPid const&) = delete; + auto operator=(ChildPid const&) -> ChildPid& = delete; + // Delete move constructor and assignment operator + ChildPid(ChildPid&&) = delete; + auto operator=(ChildPid&&) -> ChildPid& = delete; + + // Default destructor + ~ChildPid() = default; + +private: + static std::sig_atomic_t volatile m_pid; +}; +} // namespace spider::core + +#endif diff --git a/src/spider/worker/Process.cpp b/src/spider/worker/Process.cpp index b9b856c05..69aa9c8ff 100644 --- a/src/spider/worker/Process.cpp +++ b/src/spider/worker/Process.cpp @@ -118,4 +118,8 @@ auto Process::terminate() const -> void { throw std::runtime_error("Failed to terminate process"); } } + +auto Process::get_pid() const -> pid_t { + return m_pid; +} } // namespace spider::worker diff --git a/src/spider/worker/Process.hpp b/src/spider/worker/Process.hpp index 2e5f1a475..ca126d9d3 100644 --- a/src/spider/worker/Process.hpp +++ b/src/spider/worker/Process.hpp @@ -29,6 +29,11 @@ class Process { */ auto terminate() const -> void; + /* + * @return the process ID of the spawned process. + */ + [[nodiscard]] auto get_pid() const -> pid_t; + // Delete copy constructor and assignment operator Process(Process const&) = delete; auto operator=(Process const&) -> Process& = delete; diff --git a/src/spider/worker/TaskExecutor.cpp b/src/spider/worker/TaskExecutor.cpp index dad1d3705..b4a9f7fff 100644 --- a/src/spider/worker/TaskExecutor.cpp +++ b/src/spider/worker/TaskExecutor.cpp @@ -1,5 +1,7 @@ #include "TaskExecutor.hpp" +#include + #include #include #include @@ -15,6 +17,10 @@ #include "TaskExecutorMessage.hpp" namespace spider::worker { +auto TaskExecutor::get_pid() const -> pid_t { + return m_process->get_pid(); +} + auto TaskExecutor::completed() -> bool { std::lock_guard const lock(m_state_mutex); return TaskExecutorState::Succeed == m_state || TaskExecutorState::Error == m_state diff --git a/src/spider/worker/TaskExecutor.hpp b/src/spider/worker/TaskExecutor.hpp index f7bb960d4..e417adc70 100644 --- a/src/spider/worker/TaskExecutor.hpp +++ b/src/spider/worker/TaskExecutor.hpp @@ -150,6 +150,11 @@ class TaskExecutor { auto operator=(TaskExecutor&&) -> TaskExecutor& = delete; ~TaskExecutor() = default; + /* + * @return The process ID of the task executor. + */ + [[nodiscard]] auto get_pid() const -> pid_t; + auto completed() -> bool; auto waiting() -> bool; auto succeed() -> bool; diff --git a/src/spider/worker/task_executor.cpp b/src/spider/worker/task_executor.cpp index 064b48b1a..b5937dc12 100644 --- a/src/spider/worker/task_executor.cpp +++ b/src/spider/worker/task_executor.cpp @@ -1,5 +1,7 @@ #include +#include +#include #include #include #include @@ -63,11 +65,12 @@ auto parse_arg(int const argc, char** const& argv) -> boost::program_options::va } // namespace constexpr int cCmdArgParseErr = 1; -constexpr int cStorageErr = 2; -constexpr int cDllErr = 3; -constexpr int cFuncArgParseErr = 4; -constexpr int cResultSendErr = 5; -constexpr int cOtherErr = 6; +constexpr int cSignalHandleErr = 2; +constexpr int cStorageErr = 3; +constexpr int cDllErr = 4; +constexpr int cFuncArgParseErr = 5; +constexpr int cResultSendErr = 6; +constexpr int cOtherErr = 7; auto main(int const argc, char** argv) -> int { // Set up spdlog to write to stderr diff --git a/src/spider/worker/worker.cpp b/src/spider/worker/worker.cpp index f3d1c4ced..d494827af 100644 --- a/src/spider/worker/worker.cpp +++ b/src/spider/worker/worker.cpp @@ -1,4 +1,8 @@ +#include + +#include #include +#include #include #include #include @@ -41,19 +45,38 @@ #include "../storage/mysql/MySqlStorageFactory.hpp" #include "../storage/StorageConnection.hpp" #include "../storage/StorageFactory.hpp" -#include "../utils/StopToken.hpp" +#include "../utils/StopFlag.hpp" +#include "ChildPid.hpp" #include "TaskExecutor.hpp" #include "WorkerClient.hpp" constexpr int cCmdArgParseErr = 1; -constexpr int cWorkerAddrErr = 2; -constexpr int cStorageConnectionErr = 3; -constexpr int cStorageErr = 4; -constexpr int cTaskErr = 5; +constexpr int cSignalHandleErr = 2; +constexpr int cWorkerAddrErr = 3; +constexpr int cStorageConnectionErr = 4; +constexpr int cStorageErr = 5; +constexpr int cTaskErr = 6; constexpr int cRetryCount = 5; namespace { +/* + * Signal handler for SIGTERM. It sets the stop flag to request a stop and sends SIGTERM to the task + * executor. + * @param signal The signal number. + */ +auto stop_task_handler(int signal) -> void { + if (SIGTERM == signal) { + spider::core::StopFlag::request_stop(); + // Send SIGTERM to task executor + pid_t const pid = spider::core::ChildPid::get_pid(); + if (pid > 0) { + // NOLINTNEXTLINE(misc-include-cleaner) + kill(pid, SIGTERM); + } + } +} + 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"); @@ -103,11 +126,10 @@ auto get_environment_variable() -> absl::flat_hash_map< auto heartbeat_loop( std::shared_ptr const& storage_factory, std::shared_ptr const& metadata_store, - spider::core::Driver const& driver, - spider::core::StopToken& stop_token + spider::core::Driver const& driver ) -> 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, spider::core::StorageErr> @@ -133,7 +155,7 @@ auto heartbeat_loop( fail_count = 0; } if (fail_count >= cRetryCount - 1) { - stop_token.request_stop(); + spider::core::StopFlag::request_stop(); break; } } @@ -143,18 +165,19 @@ constexpr int cFetchTaskTimeout = 100; auto fetch_task(spider::worker::WorkerClient& client, std::optional fail_task_id) - -> std::tuple { + -> std::optional> { spdlog::debug("Fetching task"); - while (true) { + while (!spider::core::StopFlag::is_stop_requested()) { std::optional> const optional_task_ids = client.get_next_task(fail_task_id); if (optional_task_ids.has_value()) { - return optional_task_ids.value(); + return optional_task_ids; } // If the first request succeeds, later requests should not include the failed task id fail_task_id = std::nullopt; std::this_thread::sleep_for(std::chrono::milliseconds(cFetchTaskTimeout)); } + return std::nullopt; } /* @@ -183,8 +206,8 @@ auto setup_task( ); return std::nullopt; } - auto conn = std::move(std::get>(conn_result)); - + std::unique_ptr conn + = std::move(std::get>(conn_result)); // Get task details spider::core::StorageErr const err = metadata_store->get_task(*conn, instance.task_id, &task); if (!err.success()) { @@ -327,18 +350,22 @@ auto task_loop( std::vector const& libs, absl::flat_hash_map< boost::process::v2::environment::key, - boost::process::v2::environment::value> const& environment, - spider::core::StopToken const& stop_token + boost::process::v2::environment::value> const& environment ) -> void { std::optional fail_task_id = std::nullopt; - while (!stop_token.stop_requested()) { + while (!spider::core::StopFlag::is_stop_requested()) { boost::asio::io_context context; - auto const [task_id, task_instance_id] = fetch_task(client, fail_task_id); + auto const& optional_task = fetch_task(client, fail_task_id); + if (false == optional_task.has_value()) { + continue; + } + auto const [task_id, task_instance_id] = optional_task.value(); spider::core::TaskInstance const instance{task_instance_id, task_id}; spdlog::debug("Fetched task {}", boost::uuids::to_string(task_id)); // Fetch task detail from metadata storage spider::core::Task task{""}; + std::optional> optional_arg_buffers = setup_task(storage_factory, metadata_store, instance, task); if (!optional_arg_buffers.has_value()) { @@ -359,9 +386,19 @@ auto task_loop( arg_buffers }; + pid_t const pid = executor.get_pid(); + spider::core::ChildPid::set_pid(pid); + // Double check if stop token is set to avoid any missing signal + if (spider::core::StopFlag::is_stop_requested()) { + // NOLINTNEXTLINE(misc-include-cleaner) + kill(pid, SIGTERM); + } + context.run(); executor.wait(); + spider::core::ChildPid::set_pid(0); + if (handle_executor_result(storage_factory, metadata_store, instance, task, executor)) { fail_task_id = std::nullopt; } else { @@ -371,9 +408,10 @@ auto task_loop( } // NOLINTEND(clang-analyzer-unix.BlockInCriticalSection) + +constexpr int cSignalExitBase = 128; } // namespace -// NOLINTNEXTLINE(bugprone-exception-escape) auto main(int argc, char** argv) -> int { // Set up spdlog to write to stderr // NOLINTNEXTLINE(misc-include-cleaner) @@ -411,6 +449,17 @@ auto main(int argc, char** argv) -> int { return cCmdArgParseErr; } + // NOLINTBEGIN(misc-include-cleaner) + struct sigaction sig_action{}; + sig_action.sa_handler = stop_task_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 storage std::shared_ptr const storage_factory = std::make_shared(storage_url); @@ -444,8 +493,6 @@ auto main(int argc, char** argv) -> int { } } - spider::core::StopToken stop_token; - // Start client spider::worker::WorkerClient client{worker_id, worker_addr, data_store, metadata_store, storage_factory}; @@ -461,7 +508,6 @@ auto main(int argc, char** argv) -> int { std::cref(storage_factory), std::cref(metadata_store), std::ref(driver), - std::ref(stop_token) }; // Start a thread that processes tasks @@ -473,11 +519,15 @@ auto main(int argc, char** argv) -> int { std::cref(storage_url), std::cref(libs), std::cref(environment_variables), - std::cref(stop_token), }; heartbeat_thread.join(); task_thread.join(); + // 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; } diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 26f76a15d..3e9b66468 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -73,6 +73,16 @@ target_link_libraries( spdlog::spdlog ) +add_library(signal_test SHARED) +target_sources(signal_test PUBLIC worker/signal-test.hpp) +target_sources(signal_test PRIVATE worker/signal-test.cpp) +target_link_libraries( + signal_test + PRIVATE + spider_core + spider_client +) + add_custom_target(integrationTest ALL) add_custom_command( TARGET integrationTest @@ -85,4 +95,5 @@ add_dependencies( integrationTest worker_test client_test + signal_test ) diff --git a/tests/integration/client.py b/tests/integration/client.py index e3aca9f08..1b658c5ac 100644 --- a/tests/integration/client.py +++ b/tests/integration/client.py @@ -72,12 +72,12 @@ def is_head_task(task_id: uuid.UUID, dependencies: List[Tuple[uuid.UUID, uuid.UU return not any(dependency[1] == task_id for dependency in dependencies) -storage_url = "jdbc:mariadb://localhost:3306/spider_test?user=root&password=password" +g_storage_url = "jdbc:mariadb://localhost:3306/spider_test?user=root&password=password" @pytest.fixture(scope="session") def storage(): - conn = create_connection(storage_url) + conn = create_connection(g_storage_url) yield conn conn.close() diff --git a/tests/integration/test_client.py b/tests/integration/test_client.py index b6ebde74d..94b48afdf 100644 --- a/tests/integration/test_client.py +++ b/tests/integration/test_client.py @@ -6,9 +6,10 @@ import pytest from .client import ( + g_storage_url, storage, - storage_url, ) +from .utils import g_scheduler_port def start_scheduler_workers( @@ -41,13 +42,10 @@ def start_scheduler_workers( return scheduler_process, worker_process_0, worker_process_1 -scheduler_port = 6103 - - @pytest.fixture(scope="class") def scheduler_worker(storage): scheduler_process, worker_process_0, worker_process_1 = start_scheduler_workers( - storage_url=storage_url, scheduler_port=scheduler_port + storage_url=g_storage_url, scheduler_port=g_scheduler_port ) # Wait for 5 second to make sure the scheduler and worker are started time.sleep(5) @@ -64,7 +62,7 @@ def test_client(self, scheduler_worker): client_cmds = [ str(dir_path / "client_test"), "--storage_url", - storage_url, + g_storage_url, ] p = subprocess.run(client_cmds, timeout=20) assert p.returncode == 0 diff --git a/tests/integration/test_scheduler_worker.py b/tests/integration/test_scheduler_worker.py index 7f41f2c7d..9dfc359d0 100644 --- a/tests/integration/test_scheduler_worker.py +++ b/tests/integration/test_scheduler_worker.py @@ -12,18 +12,19 @@ add_driver_data, Data, Driver, + g_storage_url, get_task_outputs, get_task_state, remove_data, remove_job, storage, - storage_url, submit_job, Task, TaskGraph, TaskInput, TaskOutput, ) +from .utils import g_scheduler_port def start_scheduler_worker( @@ -55,13 +56,10 @@ def start_scheduler_worker( return scheduler_process, worker_process -scheduler_port = 6103 - - @pytest.fixture(scope="class") def scheduler_worker(storage): scheduler_process, worker_process = start_scheduler_worker( - storage_url=storage_url, scheduler_port=scheduler_port + storage_url=g_storage_url, scheduler_port=g_scheduler_port ) # Wait for 5 second to make sure the scheduler and worker are started time.sleep(5) diff --git a/tests/integration/test_signal.py b/tests/integration/test_signal.py new file mode 100644 index 000000000..a32f1c3d8 --- /dev/null +++ b/tests/integration/test_signal.py @@ -0,0 +1,178 @@ +import os +import signal +import subprocess +import time +import uuid +from pathlib import Path + +import msgpack +import pytest + +from .client import ( + g_storage_url, + get_task_outputs, + get_task_state, + remove_job, + storage, + submit_job, + Task, + TaskGraph, + TaskInput, + TaskOutput, +) +from .utils import g_scheduler_port + + +def start_scheduler_worker(storage_url: str, scheduler_port: int, lib: str): + root_dir = Path(__file__).resolve().parents[2] + bin_dir = root_dir / "src" / "spider" + popen_opts = dict(stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) + scheduler_cmds = [ + str(bin_dir / "spider_scheduler"), + "--host", + "127.0.0.1", + "--port", + str(scheduler_port), + "--storage_url", + storage_url, + ] + scheduler_process = subprocess.Popen(scheduler_cmds, **popen_opts) + worker_cmds = [ + str(bin_dir / "spider_worker"), + "--host", + "127.0.0.1", + "--storage_url", + storage_url, + "--libs", + lib, + ] + worker_process = subprocess.Popen(worker_cmds, **popen_opts) + + return scheduler_process, worker_process + + +@pytest.fixture(scope="function") +def scheduler_worker_signal(storage): + scheduler_process, worker_process = start_scheduler_worker( + storage_url=g_storage_url, scheduler_port=g_scheduler_port, lib="tests/libsignal_test.so" + ) + # Wait for 5 second to make sure the scheduler and worker are started + time.sleep(5) + yield scheduler_process, worker_process + worker_process.kill() + scheduler_process.kill() + + +class TestWorkerSignal: + + # Test that worker propagates the SIGTERM signal to the task executor. + # Submit a task that checks whether the task executor receives the SIGTERM signal. + # The task should return the SIGTERM signal number as the output. + # Later task should not be executed. + # Worker should exit with SIGTERM. + def test_task_signal(self, storage, scheduler_worker_signal): + _, worker_process = scheduler_worker_signal + + # Submit signal handler task to check for SIGTERM signal in task executor + task = Task( + id=uuid.uuid4(), + function_name="signal_handler_test", + inputs=[ + TaskInput(type="int", value=msgpack.packb(0)), + ], + outputs=[TaskOutput(type="int")], + ) + graph = TaskGraph( + id=uuid.uuid4(), + tasks={task.id: task}, + dependencies=[], + ) + client_id = uuid.uuid4() + submit_job(storage, client_id, graph) + # Sleep for 1 second to wait for the task to start + time.sleep(1) + + # Check if the task is in progress + assert get_task_state(storage, task.id) == "running" + + # Send signal to worker + os.kill(worker_process.pid, signal.SIGTERM) + + # Submit new task + new_task = Task( + id=uuid.uuid4(), + function_name="signal_handler_test", + inputs=[ + TaskInput(type="int", value=msgpack.packb(0)), + ], + outputs=[TaskOutput(type="int")], + ) + new_graph = TaskGraph( + id=uuid.uuid4(), + tasks={new_task.id: new_task}, + dependencies=[], + ) + submit_job(storage, client_id, new_graph) + + # Sleep for the signal handler task to finish + time.sleep(15) + + # Check if the task is finished + assert get_task_state(storage, task.id) == "success" + # Check if the task output is correct + results = get_task_outputs(storage, task.id) + assert results[0].value == msgpack.packb(signal.SIGTERM) + + # Check if the new task is not executed + assert get_task_state(storage, new_task.id) == "ready" + + # Check the worker process exited with SIGTERM + assert worker_process.poll() == signal.SIGTERM + 128 + + # Cleanup job + remove_job(storage, new_graph.id) + remove_job(storage, graph.id) + + # Test that worker propagates the SIGTERM signal to the task executor. + # Task executor exits immediately after receiving the signal. + # The running task should be marked as failed. + # The worker should exit with SIGTERM. + def test_task_exit(self, storage, scheduler_worker_signal): + _, worker_process = scheduler_worker_signal + + # Submit a task to sleep for 10 seconds + task = Task( + id=uuid.uuid4(), + function_name="sleep_test", + inputs=[ + TaskInput(type="int", value=msgpack.packb(10)), + ], + outputs=[TaskOutput(type="int")], + ) + graph = TaskGraph( + id=uuid.uuid4(), + tasks={task.id: task}, + dependencies=[], + ) + client_id = uuid.uuid4() + submit_job(storage, client_id, graph) + + # Wait for the task start + time.sleep(1) + + # Check if the task is running + assert get_task_state(storage, task.id) == "running" + + # Send signal to worker + os.kill(worker_process.pid, signal.SIGTERM) + + # Sleep for 3 seconds to wait for the task executor and worker to exit + time.sleep(3) + + # Check the task fails + assert get_task_state(storage, task.id) == "fail" + # Check the worker process exited with SIGTERM + assert worker_process.poll() == signal.SIGTERM + 128 + + # Cleanup job + remove_job(storage, graph.id) diff --git a/tests/integration/utils.py b/tests/integration/utils.py new file mode 100644 index 000000000..33aada438 --- /dev/null +++ b/tests/integration/utils.py @@ -0,0 +1,11 @@ +import socket + + +def _get_free_tcp_port() -> int: + """Returns a free TCP port.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +g_scheduler_port = _get_free_tcp_port() diff --git a/tests/scheduler/test-SchedulerServer.cpp b/tests/scheduler/test-SchedulerServer.cpp index a0a561170..9deb2006d 100644 --- a/tests/scheduler/test-SchedulerServer.cpp +++ b/tests/scheduler/test-SchedulerServer.cpp @@ -26,7 +26,6 @@ #include "../../src/spider/storage/MetadataStorage.hpp" #include "../../src/spider/storage/StorageConnection.hpp" #include "../../src/spider/storage/StorageFactory.hpp" -#include "../../src/spider/utils/StopToken.hpp" #include "../storage/StorageTestHelper.hpp" namespace { @@ -66,9 +65,7 @@ TEMPLATE_LIST_TEST_CASE( ); constexpr unsigned short cPort = 6021; - spider::core::StopToken stop_token; - spider::scheduler::SchedulerServer - server{cPort, policy, metadata_store, data_store, conn, stop_token}; + spider::scheduler::SchedulerServer server{cPort, policy, metadata_store, data_store, conn}; // Pause and resume server server.pause(); diff --git a/tests/worker/signal-test.cpp b/tests/worker/signal-test.cpp new file mode 100644 index 000000000..9e854da25 --- /dev/null +++ b/tests/worker/signal-test.cpp @@ -0,0 +1,54 @@ +#include "signal-test.hpp" + +#include +#include +#include +#include + +#include +#include + +auto SignalNumber::get_instance() -> SignalNumber& { + static SignalNumber instance; + return instance; +} + +auto SignalNumber::set_signal_number(int const signal_number) -> void { + m_signal_number = signal_number; +} + +auto SignalNumber::get_signal_number() const -> int { + return m_signal_number; +} + +namespace { +/* + * Signal handler function for SIGTERM. Sets the signal number in the singleton instance. + * @param signal_number The signal number to handle. + */ +auto signal_handler(int const signal_number) -> void { + SignalNumber::get_instance().set_signal_number(signal_number); +} + +constexpr int cSleepTime = 10; +} // namespace + +auto signal_handler_test(spider::TaskContext&, int const) -> int { + if (SIG_ERR == std::signal(SIGTERM, signal_handler)) { + std::cerr << "Failed to set signal handler for SIGTERM\n"; + return 1; + } + std::this_thread::sleep_for(std::chrono::seconds(cSleepTime)); + int const signal_number = SignalNumber::get_instance().get_signal_number(); + return signal_number; +} + +auto sleep_test(spider::TaskContext&, int const seconds) -> int { + std::this_thread::sleep_for(std::chrono::seconds(seconds)); + return 0; +} + +// NOLINTNEXTLINE(cert-err58-cpp) +SPIDER_REGISTER_TASK(signal_handler_test); +// NOLINTNEXTLINE(cert-err58-cpp) +SPIDER_REGISTER_TASK(sleep_test); diff --git a/tests/worker/signal-test.hpp b/tests/worker/signal-test.hpp new file mode 100644 index 000000000..e094aa275 --- /dev/null +++ b/tests/worker/signal-test.hpp @@ -0,0 +1,43 @@ +#ifndef SPIDER_TEST_SIGNAL_TEST_LIB_HPP +#define SPIDER_TEST_SIGNAL_TEST_LIB_HPP + +#include + +#include + +/* + * Singleton class to store the signal number. + */ +class SignalNumber { +public: + /* + * @return The singleton instance of SignalNumber. + */ + static auto get_instance() -> SignalNumber&; + + /* + * @return The signal number. + */ + [[nodiscard]] auto get_signal_number() const -> int; + + /* + * @param signal_number The signal number to set. + */ + auto set_signal_number(int signal_number) -> void; + +private: + std::sig_atomic_t volatile m_signal_number{0}; +}; + +/* + * Installs the signal handler on SIGTERM to watch for whether the signal handler is triggered. + * @return Signal number if the installed signal handler is triggered, 0 otherwise. + */ +auto signal_handler_test(spider::TaskContext& /*context*/, int /*x*/) -> int; + +/** + * @param seconds time to sleep + */ +auto sleep_test(spider::TaskContext& /*context*/, int seconds) -> int; + +#endif