diff --git a/src/spider/scheduler/FifoPolicy.cpp b/src/spider/scheduler/FifoPolicy.cpp index f58a5d467..a50c07a6e 100644 --- a/src/spider/scheduler/FifoPolicy.cpp +++ b/src/spider/scheduler/FifoPolicy.cpp @@ -17,11 +17,13 @@ namespace spider::scheduler { FifoPolicy::FifoPolicy( + boost::uuids::uuid const scheduler_id, std::shared_ptr const& metadata_store, std::shared_ptr const& data_store, std::shared_ptr const& conn ) - : m_metadata_store{metadata_store}, + : m_scheduler_id{scheduler_id}, + m_metadata_store{metadata_store}, m_data_store{data_store}, m_conn{conn} {} @@ -63,7 +65,7 @@ auto FifoPolicy::pop_next_task(std::string const& worker_addr) } auto FifoPolicy::fetch_tasks() -> void { - m_metadata_store->get_ready_tasks(*m_conn, &m_tasks); + m_metadata_store->get_ready_tasks(*m_conn, m_scheduler_id, &m_tasks); m_metadata_store->get_task_timeout(*m_conn, &m_tasks); // Sort tasks based on job creation time in descending order. diff --git a/src/spider/scheduler/FifoPolicy.hpp b/src/spider/scheduler/FifoPolicy.hpp index 924c863d6..adf91842e 100644 --- a/src/spider/scheduler/FifoPolicy.hpp +++ b/src/spider/scheduler/FifoPolicy.hpp @@ -18,6 +18,7 @@ namespace spider::scheduler { class FifoPolicy final : public SchedulerPolicy { public: FifoPolicy( + boost::uuids::uuid scheduler_id, std::shared_ptr const& metadata_store, std::shared_ptr const& data_store, std::shared_ptr const& conn @@ -31,6 +32,8 @@ class FifoPolicy final : public SchedulerPolicy { auto pop_next_task(std::string const& worker_addr) -> std::optional; + boost::uuids::uuid m_scheduler_id; + std::shared_ptr m_metadata_store; std::shared_ptr m_data_store; std::shared_ptr m_conn; diff --git a/src/spider/scheduler/scheduler.cpp b/src/spider/scheduler/scheduler.cpp index 91775de94..4dd39ec44 100644 --- a/src/spider/scheduler/scheduler.cpp +++ b/src/spider/scheduler/scheduler.cpp @@ -206,13 +206,6 @@ auto main(int argc, char** argv) -> int { boost::uuids::random_generator gen; boost::uuids::uuid const scheduler_id = gen(); - // Start scheduler server - spider::core::StopToken stop_token; - std::shared_ptr const policy - = std::make_shared(metadata_store, data_store, conn); - spider::scheduler::SchedulerServer - server{port, policy, metadata_store, data_store, conn, stop_token}; - // Register scheduler with storage spider::core::Scheduler const scheduler{scheduler_id, scheduler_addr, port}; err = metadata_store->add_scheduler(*conn, scheduler); @@ -221,6 +214,18 @@ auto main(int argc, char** argv) -> int { return cStorageErr; } + // Start scheduler server + spider::core::StopToken stop_token; + std::shared_ptr const policy + = std::make_shared( + scheduler_id, + metadata_store, + data_store, + conn + ); + spider::scheduler::SchedulerServer + server{port, policy, metadata_store, data_store, conn, stop_token}; + try { // Start a thread that periodically updates the scheduler's heartbeat std::thread heartbeat_thread{ diff --git a/src/spider/storage/MetadataStorage.hpp b/src/spider/storage/MetadataStorage.hpp index 6fdec146c..91f733b70 100644 --- a/src/spider/storage/MetadataStorage.hpp +++ b/src/spider/storage/MetadataStorage.hpp @@ -83,8 +83,11 @@ class MetadataStorage { get_task_job_id(StorageConnection& conn, boost::uuids::uuid id, boost::uuids::uuid* job_id) -> StorageErr = 0; - virtual auto get_ready_tasks(StorageConnection& conn, std::vector* tasks) - -> StorageErr + virtual auto get_ready_tasks( + StorageConnection& conn, + boost::uuids::uuid scheduler_id, + std::vector* tasks + ) -> StorageErr = 0; virtual auto set_task_state(StorageConnection& conn, boost::uuids::uuid id, TaskState state) -> StorageErr diff --git a/src/spider/storage/mysql/MySqlStorage.cpp b/src/spider/storage/mysql/MySqlStorage.cpp index ea96173f2..1f472bb92 100644 --- a/src/spider/storage/mysql/MySqlStorage.cpp +++ b/src/spider/storage/mysql/MySqlStorage.cpp @@ -1226,11 +1226,24 @@ auto MySqlMetadataStorage::get_task_job_id( return StorageErr{}; } +constexpr int cLeaseExpireTime = 1000 * 10; // 10 ms + auto MySqlMetadataStorage::get_ready_tasks( StorageConnection& conn, + boost::uuids::uuid scheduler_id, std::vector* tasks ) -> StorageErr { try { + // Remove timeout scheduler leases + std::unique_ptr lease_timeout_statement( + static_cast(conn)->prepareStatement( + "DELETE FROM `scheduler_leases` WHERE TIMESTAMPDIFF(MICROSECOND, " + "`lease_time`, CURRENT_TIMESTAMP()) > ?" + ) + ); + lease_timeout_statement->setInt(1, cLeaseExpireTime); + lease_timeout_statement->executeUpdate(); + // Get all ready tasks from job that has not failed or cancelled std::unique_ptr task_statement( static_cast(conn)->createStatement() @@ -1238,7 +1251,7 @@ auto MySqlMetadataStorage::get_ready_tasks( std::unique_ptr const res{task_statement->executeQuery( "SELECT `id`, `func_name`, `job_id` FROM `tasks` WHERE `state` = 'ready' " "AND `job_id` NOT IN (SELECT `job_id` FROM `tasks` WHERE `state` = 'fail' OR " - "`state` = 'cancel')" + "`state` = 'cancel') AND `id` NOT IN (SELECT `task_id` FROM `scheduler_leases`)" )}; if (res->rowsCount() == 0) { @@ -1270,6 +1283,7 @@ auto MySqlMetadataStorage::get_ready_tasks( "(SELECT `job_id` FROM `tasks` WHERE `state` = 'fail' OR `state` = 'cancel'))" )}; + // Get job metadata while (job_res->next()) { boost::uuids::uuid const job_id = read_id(job_res->getBinaryStream("id")); boost::uuids::uuid const client_id = read_id(job_res->getBinaryStream("client_id")); @@ -1285,6 +1299,10 @@ auto MySqlMetadataStorage::get_ready_tasks( ) }; } + // Job id will not be in job_id_to_task_ids if the job's tasks are leased + if (job_id_to_task_ids.find(job_id) == job_id_to_task_ids.end()) { + continue; + } for (boost::uuids::uuid const& task_id : job_id_to_task_ids[job_id]) { new_tasks[task_id].set_client_id(client_id); new_tasks[task_id].set_job_creation_time(optional_creation_time.value()); @@ -1301,7 +1319,8 @@ auto MySqlMetadataStorage::get_ready_tasks( "`task_inputs`.`data_id` = `data`.`id` JOIN `data_locality` ON `data`.`id` " "= `data_locality`.`id` WHERE `task_inputs`.`task_id` IN (SELECT `id` " "FROM `tasks` WHERE `state` = 'ready' AND `job_id` NOT IN (SELECT `job_id` " - "FROM `tasks` WHERE `state` = 'fail' OR `state` = 'cancel'))" + "FROM `tasks` WHERE `state` = 'fail' OR `state` = 'cancel')) AND " + "`task_inputs`.`task_id` NOT IN (SELECT `task_id` FROM `scheduler_leases`)" )}; while (locality_res->next()) { @@ -1315,6 +1334,21 @@ auto MySqlMetadataStorage::get_ready_tasks( } } + // Add scheduler lease + std::unique_ptr lease_statement( + static_cast(conn)->prepareStatement( + "INSERT INTO `scheduler_leases` (`scheduler_id`, `task_id`) VALUES (?, ?)" + ) + ); + sql::bytes scheduler_id_bytes = uuid_get_bytes(scheduler_id); + for (auto const& [task_id, task] : new_tasks) { + sql::bytes task_id_bytes = uuid_get_bytes(task_id); + lease_statement->setBytes(1, &scheduler_id_bytes); + lease_statement->setBytes(2, &task_id_bytes); + lease_statement->addBatch(); + } + lease_statement->executeBatch(); + // Add all tasks to the output absl::flat_hash_set task_ids; for (ScheduleTaskMetadata const& task : *tasks) { @@ -1458,6 +1492,14 @@ MySqlMetadataStorage::create_task_instance(StorageConnection& conn, TaskInstance instance_statement->setBytes(1, &instance_id_bytes); instance_statement->setBytes(2, &id_bytes); instance_statement->executeUpdate(); + // Remove task from scheduler leases + std::unique_ptr const lease_statement( + static_cast(conn)->prepareStatement( + "DELETE FROM `scheduler_leases` WHERE `task_id` = ?" + ) + ); + lease_statement->setBytes(1, &id_bytes); + lease_statement->executeUpdate(); } catch (sql::SQLException& e) { static_cast(conn)->rollback(); return StorageErr{StorageErrType::OtherErr, e.what()}; diff --git a/src/spider/storage/mysql/MySqlStorage.hpp b/src/spider/storage/mysql/MySqlStorage.hpp index 8b8e23161..2dcc30f49 100644 --- a/src/spider/storage/mysql/MySqlStorage.hpp +++ b/src/spider/storage/mysql/MySqlStorage.hpp @@ -79,8 +79,11 @@ class MySqlMetadataStorage : public MetadataStorage { -> StorageErr override; auto get_task_job_id(StorageConnection& conn, boost::uuids::uuid id, boost::uuids::uuid* job_id) -> StorageErr override; - auto get_ready_tasks(StorageConnection& conn, std::vector* tasks) - -> StorageErr override; + auto get_ready_tasks( + StorageConnection& conn, + boost::uuids::uuid scheduler_id, + std::vector* tasks + ) -> StorageErr override; auto set_task_state(StorageConnection& conn, boost::uuids::uuid id, TaskState state) -> StorageErr override; auto set_task_running(StorageConnection& conn, boost::uuids::uuid id) -> StorageErr override; diff --git a/src/spider/storage/mysql/mysql_stmt.hpp b/src/spider/storage/mysql/mysql_stmt.hpp index 24a7d5be1..402b481b5 100644 --- a/src/spider/storage/mysql/mysql_stmt.hpp +++ b/src/spider/storage/mysql/mysql_stmt.hpp @@ -105,6 +105,16 @@ std::string const cCreateTaskInstanceTable = R"(CREATE TABLE IF NOT EXISTS `task PRIMARY KEY (`id`) ))"; +std::string const cCreateSchedulerLeaseTable = R"(CREATE TABLE IF NOT EXISTS `scheduler_leases` ( + `scheduler_id` BINARY(16) NOT NULL, + `task_id` BINARY(16) NOT NULL, + `lease_time` TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + CONSTRAINT `lease_scheduler_id` FOREIGN KEY (`scheduler_id`) REFERENCES `schedulers` (`id`) ON UPDATE NO ACTION ON DELETE CASCADE, + CONSTRAINT `lease_task_id` FOREIGN KEY (`task_id`) REFERENCES `tasks` (`id`) ON UPDATE NO ACTION ON DELETE CASCADE, + INDEX (`scheduler_id`), + PRIMARY KEY (`scheduler_id`, `task_id`) +))"; + std::string const cCreateDataTable = R"(CREATE TABLE IF NOT EXISTS `data` ( `id` BINARY(16) NOT NULL, `value` VARBINARY(999) NOT NULL, @@ -153,7 +163,7 @@ std::string const cCreateTaskKVDataTable = R"(CREATE TABLE IF NOT EXISTS `task_k CONSTRAINT `kv_data_task_id` FOREIGN KEY (`task_id`) REFERENCES `tasks` (`id`) ON UPDATE NO ACTION ON DELETE CASCADE ))"; -std::array const cCreateStorage = { +std::array const cCreateStorage = { cCreateDriverTable, // drivers table must be created before data_ref_driver cCreateSchedulerTable, cCreateJobTable, // jobs table must be created before task @@ -170,6 +180,8 @@ std::array const cCreateStorage = { cCreateTaskInputTable, cCreateTaskDependencyTable, cCreateTaskInstanceTable, + cCreateSchedulerLeaseTable // scheduler_lease table must be created after scheduler and + // task }; std::string const cInsertJob = R"(INSERT INTO `jobs` (`id`, `client_id`) VALUES (?, ?))"; diff --git a/tests/scheduler/test-SchedulerPolicy.cpp b/tests/scheduler/test-SchedulerPolicy.cpp index db4dfb8bd..a6fd555a5 100644 --- a/tests/scheduler/test-SchedulerPolicy.cpp +++ b/tests/scheduler/test-SchedulerPolicy.cpp @@ -44,6 +44,13 @@ TEMPLATE_LIST_TEST_CASE( = std::move(std::get>(conn_result)); boost::uuids::random_generator gen; + + // Add scheduler + boost::uuids::uuid const scheduler_id = gen(); + REQUIRE(metadata_store + ->add_scheduler(*conn, spider::core::Scheduler{scheduler_id, "127.0.0.1", 8080}) + .success()); + boost::uuids::uuid const client_id = gen(); // Submit tasks spider::core::Task const task_1{"task_1"}; @@ -62,7 +69,7 @@ TEMPLATE_LIST_TEST_CASE( boost::uuids::uuid const job_id_2 = gen(); REQUIRE(metadata_store->add_job(*conn, job_id_2, client_id, graph_2).success()); - spider::scheduler::FifoPolicy policy{metadata_store, data_store, conn}; + spider::scheduler::FifoPolicy policy{scheduler_id, metadata_store, data_store, conn}; // Schedule the earlier task std::optional optional_task_id = policy.schedule_next(gen(), ""); @@ -107,6 +114,13 @@ TEMPLATE_LIST_TEST_CASE( = std::move(std::get>(conn_result)); boost::uuids::random_generator gen; + + // Add scheduler + boost::uuids::uuid const scheduler_id = gen(); + REQUIRE(metadata_store + ->add_scheduler(*conn, spider::core::Scheduler{scheduler_id, "127.0.0.1", 8080}) + .success()); + boost::uuids::uuid const job_id = gen(); boost::uuids::uuid const client_id = gen(); // Submit task with hard locality @@ -123,7 +137,7 @@ TEMPLATE_LIST_TEST_CASE( graph.add_output_task(task.get_id()); REQUIRE(metadata_store->add_job(*conn, job_id, client_id, graph).success()); - spider::scheduler::FifoPolicy policy{metadata_store, data_store, conn}; + spider::scheduler::FifoPolicy policy{scheduler_id, metadata_store, data_store, conn}; // Schedule with wrong address REQUIRE_FALSE(policy.schedule_next(gen(), "").has_value()); // Schedule with correct address @@ -156,8 +170,15 @@ TEMPLATE_LIST_TEST_CASE( std::shared_ptr const conn = std::move(std::get>(conn_result)); - // Add task boost::uuids::random_generator gen; + + // Add scheduler + boost::uuids::uuid const scheduler_id = gen(); + REQUIRE(metadata_store + ->add_scheduler(*conn, spider::core::Scheduler{scheduler_id, "127.0.0.1", 8080}) + .success()); + + // Add task boost::uuids::uuid const job_id = gen(); boost::uuids::uuid const client_id = gen(); spider::core::Task task{"task"}; @@ -173,7 +194,8 @@ TEMPLATE_LIST_TEST_CASE( graph.add_output_task(task.get_id()); REQUIRE(metadata_store->add_job(*conn, job_id, client_id, graph).success()); - spider::scheduler::FifoPolicy policy{metadata_store, data_store, conn}; + spider::scheduler::FifoPolicy policy{scheduler_id, metadata_store, data_store, conn}; + // Schedule with wrong address std::optional const optional_task_id = policy.schedule_next(gen(), ""); REQUIRE(optional_task_id.has_value()); diff --git a/tests/scheduler/test-SchedulerServer.cpp b/tests/scheduler/test-SchedulerServer.cpp index 3e3be219b..a0a561170 100644 --- a/tests/scheduler/test-SchedulerServer.cpp +++ b/tests/scheduler/test-SchedulerServer.cpp @@ -50,8 +50,20 @@ TEMPLATE_LIST_TEST_CASE( std::shared_ptr const conn = std::move(std::get>(conn_result)); + // Add scheduler + boost::uuids::random_generator gen; + boost::uuids::uuid const scheduler_id = gen(); + REQUIRE(metadata_store + ->add_scheduler(*conn, spider::core::Scheduler{scheduler_id, "127.0.0.1", 8080}) + .success()); + std::shared_ptr const policy - = std::make_shared(metadata_store, data_store, conn); + = std::make_shared( + scheduler_id, + metadata_store, + data_store, + conn + ); constexpr unsigned short cPort = 6021; spider::core::StopToken stop_token; @@ -79,7 +91,6 @@ TEMPLATE_LIST_TEST_CASE( graph.add_dependency(parent_task.get_id(), child_task.get_id()); graph.add_input_task(parent_task.get_id()); graph.add_output_task(child_task.get_id()); - boost::uuids::random_generator gen; boost::uuids::uuid const job_id = gen(); REQUIRE(metadata_store->add_job(*conn, job_id, gen(), graph).success()); diff --git a/tests/storage/test-MetadataStorage.cpp b/tests/storage/test-MetadataStorage.cpp index c88d2eeb3..8bf7c6808 100644 --- a/tests/storage/test-MetadataStorage.cpp +++ b/tests/storage/test-MetadataStorage.cpp @@ -538,6 +538,61 @@ TEMPLATE_LIST_TEST_CASE("Job reset", "[storage]", spider::test::StorageFactoryTy // Clean up REQUIRE(storage->remove_job(*conn, job_id).success()); } + +TEMPLATE_LIST_TEST_CASE( + "Scheduler lease timeout", + "[storage]", + spider::test::StorageFactoryTypeList +) { + std::unique_ptr storage_factory + = spider::test::create_storage_factory(); + std::unique_ptr storage + = storage_factory->provide_metadata_storage(); + + std::variant, spider::core::StorageErr> + conn_result = storage_factory->provide_storage_connection(); + REQUIRE(std::holds_alternative>(conn_result)); + auto conn = std::move(std::get>(conn_result)); + + boost::uuids::random_generator gen; + + // Register scheduler + boost::uuids::uuid const scheduler_id = gen(); + constexpr int cPort = 3306; + REQUIRE(storage->add_scheduler(*conn, spider::core::Scheduler{scheduler_id, "127.0.0.1", cPort}) + .success()); + + // Add simple job + boost::uuids::uuid const job_id = gen(); + spider::core::Task const task{"simple"}; + spider::core::TaskGraph graph; + graph.add_task(task); + graph.add_input_task(task.get_id()); + graph.add_output_task(task.get_id()); + REQUIRE(storage->add_job(*conn, job_id, gen(), graph).success()); + + // Get ready tasks should schedule the task + std::vector tasks; + REQUIRE(storage->get_ready_tasks(*conn, scheduler_id, &tasks).success()); + REQUIRE(1 == tasks.size()); + REQUIRE(tasks[0].get_id() == task.get_id()); + + // Get again should not schedule the task + tasks.clear(); + REQUIRE(storage->get_ready_tasks(*conn, scheduler_id, &tasks).success()); + REQUIRE(tasks.empty()); + + // Wait for lease timeout + std::this_thread::sleep_for(std::chrono::seconds(2)); + // Get ready tasks should schedule the task again + tasks.clear(); + REQUIRE(storage->get_ready_tasks(*conn, scheduler_id, &tasks).success()); + REQUIRE(1 == tasks.size()); + REQUIRE(tasks[0].get_id() == task.get_id()); + + // Clean up + REQUIRE(storage->remove_job(*conn, job_id).success()); +} } // namespace // NOLINTEND(cert-err58-cpp,cppcoreguidelines-avoid-do-while,readability-function-cognitive-complexity,cppcoreguidelines-avoid-non-const-global-variables,cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays) diff --git a/tools/scripts/storage/init_db.sql b/tools/scripts/storage/init_db.sql index 4d9100bf0..60a8f14d1 100644 --- a/tools/scripts/storage/init_db.sql +++ b/tools/scripts/storage/init_db.sql @@ -105,6 +105,16 @@ CREATE TABLE IF NOT EXISTS `task_instances` CONSTRAINT `instance_task_id` FOREIGN KEY (`task_id`) REFERENCES `tasks` (`id`) ON UPDATE NO ACTION ON DELETE CASCADE, PRIMARY KEY (`id`) ); +CREATE TABLE IF NOT EXISTS `scheduler_leases` +( + `scheduler_id` BINARY(16) NOT NULL, + `task_id` BINARY(16) NOT NULL, + `lease_time` TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + CONSTRAINT `lease_scheduler_id` FOREIGN KEY (`scheduler_id`) REFERENCES `schedulers` (`id`) ON UPDATE NO ACTION ON DELETE CASCADE, + CONSTRAINT `lease_task_id` FOREIGN KEY (`task_id`) REFERENCES `tasks` (`id`) ON UPDATE NO ACTION ON DELETE CASCADE, + INDEX (`scheduler_id`), + PRIMARY KEY (`scheduler_id`, `task_id`) +); CREATE TABLE IF NOT EXISTS `data_locality` ( `id` BINARY(16) NOT NULL,