fix: Add task leases to prevent a race scheduling the same task multiple times (fixes #105).

src/spider/scheduler/FifoPolicy.cpp

@@ -17,11 +17,13 @@
 
 namespace spider::scheduler {
 FifoPolicy::FifoPolicy(
+        boost::uuids::uuid const scheduler_id,
         std::shared_ptr<core::MetadataStorage> const& metadata_store,
         std::shared_ptr<core::DataStorage> const& data_store,
         std::shared_ptr<core::StorageConnection> 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.

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<core::MetadataStorage> const& metadata_store,
             std::shared_ptr<core::DataStorage> const& data_store,
             std::shared_ptr<core::StorageConnection> 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>;
 
+    boost::uuids::uuid m_scheduler_id;
+
     std::shared_ptr<core::MetadataStorage> m_metadata_store;
     std::shared_ptr<core::DataStorage> m_data_store;
     std::shared_ptr<core::StorageConnection> m_conn;

src/spider/scheduler/scheduler.cpp

@@ -1,5 +1,4 @@
 #include <chrono>
-#include <cstddef>
 #include <functional>
 #include <memory>
 #include <string>
@@ -111,10 +110,8 @@ auto heartbeat_loop(
 
 auto cleanup_loop(
         std::shared_ptr<spider::core::StorageFactory> const& storage_factory,
-        std::shared_ptr<spider::core::MetadataStorage> const& metadata_store,
         std::shared_ptr<spider::core::DataStorage> const& data_store,
-        spider::core::Scheduler const& scheduler,
-        spider::core::StopToken& stop_token
+        spider::core::StopToken const& stop_token
 ) -> void {
     while (!stop_token.stop_requested()) {
         std::this_thread::sleep_for(std::chrono::seconds(cCleanupInterval));
@@ -132,25 +129,7 @@ auto cleanup_loop(
                 std::get<std::unique_ptr<spider::core::StorageConnection>>(conn_result)
         );
 
-        spider::core::StorageErr err
-                = metadata_store->set_scheduler_state(*conn, scheduler.get_id(), "gc");
-        if (!err.success()) {
-            spdlog::error("Failed to set scheduler state to gc: {}", err.description);
-            continue;
-        }
         data_store->remove_dangling_data(*conn);
-        for (size_t i = 0; i < cRetryCount; ++i) {
-            err = metadata_store->set_scheduler_state(*conn, scheduler.get_id(), "normal");
-            if (!err.success()) {
-                spdlog::error("Failed to set scheduler state to normal: {}", err.description);
-                if (i >= cRetryCount - 1) {
-                    stop_token.request_stop();
-                    return;
-                }
-            } else {
-                break;
-            }
-        }
         spdlog::debug("Finished cleanup");
     }
 }
@@ -227,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<spider::scheduler::SchedulerPolicy> const policy
-            = std::make_shared<spider::scheduler::FifoPolicy>(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);
@@ -242,6 +214,18 @@ auto main(int argc, char** argv) -> int {
         return cStorageErr;
     }
 
+    // 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};
+
     try {
         // Start a thread that periodically updates the scheduler's heartbeat
         std::thread heartbeat_thread{
@@ -256,9 +240,7 @@ auto main(int argc, char** argv) -> int {
         std::thread cleanup_thread{
                 cleanup_loop,
                 std::cref(storage_factory),
-                std::cref(metadata_store),
                 std::cref(data_store),
-                std::cref(scheduler),
                 std::ref(stop_token)
         };
 

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<ScheduleTaskMetadata>* tasks)
-            -> StorageErr
+    virtual auto get_ready_tasks(
+            StorageConnection& conn,
+            boost::uuids::uuid scheduler_id,
+            std::vector<ScheduleTaskMetadata>* tasks
+    ) -> StorageErr
             = 0;
     virtual auto set_task_state(StorageConnection& conn, boost::uuids::uuid id, TaskState state)
             -> StorageErr
@@ -124,18 +127,11 @@ class MetadataStorage {
     heartbeat_timeout(StorageConnection& conn, double timeout, std::vector<boost::uuids::uuid>* ids)
             -> StorageErr
             = 0;
-    virtual auto
-    get_scheduler_state(StorageConnection& conn, boost::uuids::uuid id, std::string* state)
-            -> StorageErr
-            = 0;
+
     virtual auto
     get_scheduler_addr(StorageConnection& conn, boost::uuids::uuid id, std::string* addr, int* port)
             -> StorageErr
             = 0;
-    virtual auto
-    set_scheduler_state(StorageConnection& conn, boost::uuids::uuid id, std::string const& state)
-            -> StorageErr
-            = 0;
 };
 }  // namespace spider::core
 #endif  // SPIDER_STORAGE_METADATASTORAGE_HPP

src/spider/storage/mysql/MySqlStorage.cpp

@@ -174,8 +174,7 @@ auto MySqlMetadataStorage::add_scheduler(StorageConnection& conn, Scheduler cons
         driver_statement->executeUpdate();
         std::unique_ptr<sql::PreparedStatement> scheduler_statement(
                 static_cast<MySqlConnection&>(conn)->prepareStatement(
-                        "INSERT INTO `schedulers` (`id`, `address`, `port`, `state`) VALUES (?, ?, "
-                        "?, 'normal')"
+                        "INSERT INTO `schedulers` (`id`, `address`, `port`) VALUES (?, ?, ?)"
                 )
         );
         scheduler_statement->setBytes(1, &id_bytes);
@@ -203,7 +202,7 @@ auto MySqlMetadataStorage::get_active_scheduler(
         );
         std::unique_ptr<sql::ResultSet> res(statement->executeQuery(
                 "SELECT `schedulers`.`id`, `address`, `port` FROM `schedulers` JOIN `drivers` ON "
-                "`schedulers`.`id` = `drivers`.`id` WHERE `state` = 'normal'"
+                "`schedulers`.`id` = `drivers`.`id`"
         ));
         while (res->next()) {
             boost::uuids::uuid const id = read_id(res->getBinaryStream(1));
@@ -1227,19 +1226,32 @@ 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<ScheduleTaskMetadata>* tasks
 ) -> StorageErr {
     try {
+        // Remove timeout scheduler leases
+        std::unique_ptr<sql::PreparedStatement> lease_timeout_statement(
+                static_cast<MySqlConnection&>(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<sql::Statement> task_statement(
                 static_cast<MySqlConnection&>(conn)->createStatement()
         );
         std::unique_ptr<sql::ResultSet> 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) {
@@ -1271,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"));
@@ -1286,6 +1299,10 @@ auto MySqlMetadataStorage::get_ready_tasks(
                         )
                 };
             }
+            // Job id might 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());
@@ -1302,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()) {
@@ -1316,6 +1334,21 @@ auto MySqlMetadataStorage::get_ready_tasks(
             }
         }
 
+        // Add scheduler lease
+        std::unique_ptr<sql::PreparedStatement> lease_statement(
+                static_cast<MySqlConnection&>(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<boost::uuids::uuid> task_ids;
         for (ScheduleTaskMetadata const& task : *tasks) {
@@ -1459,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<sql::PreparedStatement> const lease_statement(
+                static_cast<MySqlConnection&>(conn)->prepareStatement(
+                        "DELETE FROM `scheduler_leases` WHERE `task_id` = ?"
+                )
+        );
+        lease_statement->setBytes(1, &id_bytes);
+        lease_statement->executeUpdate();
     } catch (sql::SQLException& e) {
         static_cast<MySqlConnection&>(conn)->rollback();
         return StorageErr{StorageErrType::OtherErr, e.what()};
@@ -1840,37 +1881,6 @@ auto MySqlMetadataStorage::heartbeat_timeout(
     return StorageErr{};
 }
 
-auto MySqlMetadataStorage::get_scheduler_state(
-        StorageConnection& conn,
-        boost::uuids::uuid id,
-        std::string* state
-) -> StorageErr {
-    try {
-        std::unique_ptr<sql::PreparedStatement> statement(
-                static_cast<MySqlConnection&>(conn)->prepareStatement(
-                        "SELECT `state` FROM `schedulers` WHERE `id` = ?"
-                )
-        );
-        sql::bytes id_bytes = uuid_get_bytes(id);
-        statement->setBytes(1, &id_bytes);
-        std::unique_ptr<sql::ResultSet> res(statement->executeQuery());
-        if (res->rowsCount() == 0) {
-            static_cast<MySqlConnection&>(conn)->rollback();
-            return StorageErr{
-                    StorageErrType::KeyNotFoundErr,
-                    fmt::format("no scheduler with id {}", boost::uuids::to_string(id))
-            };
-        }
-        res->next();
-        *state = get_sql_string(res->getString(1));
-    } catch (sql::SQLException& e) {
-        static_cast<MySqlConnection&>(conn)->rollback();
-        return StorageErr{StorageErrType::OtherErr, e.what()};
-    }
-    static_cast<MySqlConnection&>(conn)->commit();
-    return StorageErr{};
-}
-
 auto MySqlMetadataStorage::get_scheduler_addr(
         StorageConnection& conn,
         boost::uuids::uuid id,
@@ -1904,29 +1914,6 @@ auto MySqlMetadataStorage::get_scheduler_addr(
     return StorageErr{};
 }
 
-auto MySqlMetadataStorage::set_scheduler_state(
-        StorageConnection& conn,
-        boost::uuids::uuid id,
-        std::string const& state
-) -> StorageErr {
-    try {
-        std::unique_ptr<sql::PreparedStatement> statement(
-                static_cast<MySqlConnection&>(conn)->prepareStatement(
-                        "UPDATE `schedulers` SET `state` = ? WHERE `id` = ?"
-                )
-        );
-        statement->setString(1, state);
-        sql::bytes id_bytes = uuid_get_bytes(id);
-        statement->setBytes(2, &id_bytes);
-        statement->executeUpdate();
-    } catch (sql::SQLException& e) {
-        static_cast<MySqlConnection&>(conn)->rollback();
-        return StorageErr{StorageErrType::OtherErr, e.what()};
-    }
-    static_cast<MySqlConnection&>(conn)->commit();
-    return StorageErr{};
-}
-
 auto MySqlDataStorage::initialize(StorageConnection& conn) -> StorageErr {
     try {
         // Need to initialize metadata storage first so that foreign constraint is not voilated

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<ScheduleTaskMetadata>* tasks)
-            -> StorageErr override;
+    auto get_ready_tasks(
+            StorageConnection& conn,
+            boost::uuids::uuid scheduler_id,
+            std::vector<ScheduleTaskMetadata>* 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;
@@ -106,14 +109,9 @@ class MySqlMetadataStorage : public MetadataStorage {
     auto
     heartbeat_timeout(StorageConnection& conn, double timeout, std::vector<boost::uuids::uuid>* ids)
             -> StorageErr override;
-    auto get_scheduler_state(StorageConnection& conn, boost::uuids::uuid id, std::string* state)
-            -> StorageErr override;
     auto
     get_scheduler_addr(StorageConnection& conn, boost::uuids::uuid id, std::string* addr, int* port)
             -> StorageErr override;
-    auto
-    set_scheduler_state(StorageConnection& conn, boost::uuids::uuid id, std::string const& state)
-            -> StorageErr override;
 
 private:
     MySqlMetadataStorage() = default;