fix: Batch MySQL queries in getting ready tasks

src/spider/core/Task.hpp

@@ -1,6 +1,7 @@
 #ifndef SPIDER_CORE_TASK_HPP
 #define SPIDER_CORE_TASK_HPP
 
+#include <chrono>
 #include <cstddef>
 #include <cstdint>
 #include <optional>
@@ -109,6 +110,64 @@ enum class TaskState : std::uint8_t {
     Canceled,
 };
 
+class ScheduleTaskMetadata {
+public:
+    ScheduleTaskMetadata(
+            boost::uuids::uuid id,
+            std::string function_name,
+            boost::uuids::uuid job_id
+    )
+            : m_id(id),
+              m_function_name(std::move(function_name)),
+              m_job_id(job_id) {}
+
+    ScheduleTaskMetadata() = default;
+
+    [[nodiscard]] auto get_id() const -> boost::uuids::uuid { return m_id; }
+
+    [[nodiscard]] auto get_function_name() const -> std::string const& { return m_function_name; }
+
+    [[nodiscard]] auto get_job_id() const -> boost::uuids::uuid { return m_job_id; }
+
+    [[nodiscard]] auto get_client_id() const -> boost::uuids::uuid { return m_client_id; }
+
+    [[nodiscard]] auto get_job_creation_time() const -> std::chrono::system_clock::time_point {
+        return m_job_creation_time;
+    }
+
+    [[nodiscard]] auto get_hard_localities() const -> std::vector<std::string> const& {
+        return m_hard_localities;
+    }
+
+    [[nodiscard]] auto get_soft_localities() const -> std::vector<std::string> const& {
+        return m_soft_localities;
+    }
+
+    auto set_client_id(boost::uuids::uuid const client_id) -> void { m_client_id = client_id; }
+
+    auto set_job_creation_time(std::chrono::system_clock::time_point const job_creation_time
+    ) -> void {
+        m_job_creation_time = job_creation_time;
+    }
+
+    auto add_hard_locality(std::string const& locality) -> void {
+        m_hard_localities.push_back(locality);
+    }
+
+    auto add_soft_locality(std::string const& locality) -> void {
+        m_soft_localities.push_back(locality);
+    }
+
+private:
+    boost::uuids::uuid m_id;
+    std::string m_function_name;
+    boost::uuids::uuid m_job_id;
+    boost::uuids::uuid m_client_id;
+    std::chrono::system_clock::time_point m_job_creation_time;
+    std::vector<std::string> m_hard_localities;
+    std::vector<std::string> m_soft_localities;
+};
+
 class Task {
 public:
     explicit Task(std::string function_name) : m_function_name(std::move(function_name)) {

src/spider/scheduler/FifoPolicy.cpp

@@ -1,65 +1,21 @@
 #include "FifoPolicy.hpp"
 
 #include <algorithm>
-#include <chrono>
 #include <iterator>
 #include <memory>
 #include <optional>
-#include <stdexcept>
 #include <string>
-#include <tuple>
 #include <vector>
 
-#include <absl/container/flat_hash_map.h>
 #include <boost/uuid/uuid.hpp>
-#include <boost/uuid/uuid_io.hpp>
-#include <fmt/format.h>
 
-#include "../core/Data.hpp"
-#include "../core/JobMetadata.hpp"
 #include "../core/Task.hpp"
 #include "../storage/DataStorage.hpp"
 #include "../storage/MetadataStorage.hpp"
 #include "../storage/StorageConnection.hpp"
 
 namespace spider::scheduler {
 
-auto FifoPolicy::task_locality_satisfied(spider::core::Task const& task, std::string const& addr)
-        -> bool {
-    for (auto const& input : task.get_inputs()) {
-        if (input.get_value().has_value()) {
-            continue;
-        }
-        std::optional<boost::uuids::uuid> optional_data_id = input.get_data_id();
-        if (!optional_data_id.has_value()) {
-            continue;
-        }
-        boost::uuids::uuid const data_id = optional_data_id.value();
-        core::Data data;
-        if (m_data_cache.contains(data_id)) {
-            data = m_data_cache[data_id];
-        } else {
-            if (false == m_data_store->get_data(*m_conn, data_id, &data).success()) {
-                throw std::runtime_error(
-                        fmt::format("Data with id {} not exists.", to_string((data_id)))
-                );
-            }
-            m_data_cache.emplace(data_id, data);
-        }
-        if (false == data.is_hard_locality()) {
-            continue;
-        }
-        std::vector<std::string> const& locality = data.get_locality();
-        if (locality.empty()) {
-            continue;
-        }
-        if (std::ranges::find(locality, addr) == locality.end()) {
-            return false;
-        }
-    }
-    return true;
-}
-
 FifoPolicy::FifoPolicy(
         std::shared_ptr<core::MetadataStorage> const& metadata_store,
         std::shared_ptr<core::DataStorage> const& data_store,
@@ -81,56 +37,35 @@ auto FifoPolicy::schedule_next(
     }
     auto const reverse_begin = std::reverse_iterator(m_tasks.end());
     auto const reverse_end = std::reverse_iterator(m_tasks.begin());
-    auto const it = std::find_if(reverse_begin, reverse_end, [&](core::Task const& task) {
-        return task_locality_satisfied(task, worker_addr);
-    });
+    auto const it
+            = std::find_if(reverse_begin, reverse_end, [&](core::ScheduleTaskMetadata const& task) {
+                  std::vector<std::string> const& hard_localities = task.get_hard_localities();
+                  if (hard_localities.empty()) {
+                      return true;
+                  }
+                  // If the worker address is in the hard localities, then the task can be
+                  // scheduled.
+                  return std::ranges::find(hard_localities, worker_addr) != hard_localities.end();
+              });
     if (it == reverse_end) {
         return std::nullopt;
     }
     boost::uuids::uuid const task_id = it->get_id();
-    for (core::TaskInput const& input : it->get_inputs()) {
-        std::optional<boost::uuids::uuid> const data_id = input.get_data_id();
-        if (data_id.has_value()) {
-            m_data_cache.erase(data_id.value());
-        }
-    }
     m_tasks.erase(std::next(it).base());
     return task_id;
 }
 
 auto FifoPolicy::fetch_tasks() -> void {
-    m_data_cache.clear();
     m_metadata_store->get_ready_tasks(*m_conn, &m_tasks);
-    std::vector<std::tuple<core::TaskInstance, core::Task>> instances;
-    m_metadata_store->get_task_timeout(*m_conn, &instances);
-    for (auto const& [instance, task] : instances) {
-        m_tasks.emplace_back(task);
-    }
+    m_metadata_store->get_task_timeout(*m_conn, &m_tasks);
 
     // Sort tasks based on job creation time in descending order.
-    // NOLINTNEXTLINE(misc-include-cleaner)
-    absl::flat_hash_map<boost::uuids::uuid, core::JobMetadata, std::hash<boost::uuids::uuid>>
-            job_metadata_map;
-    auto get_task_job_creation_time
-            = [&](boost::uuids::uuid const task_id) -> std::chrono::system_clock::time_point {
-        boost::uuids::uuid job_id;
-        if (false == m_metadata_store->get_task_job_id(*m_conn, task_id, &job_id).success()) {
-            throw std::runtime_error(fmt::format("Task with id {} not exists.", to_string(task_id))
-            );
-        }
-        if (job_metadata_map.contains(job_id)) {
-            return job_metadata_map[job_id].get_creation_time();
-        }
-        core::JobMetadata job_metadata;
-        if (false == m_metadata_store->get_job_metadata(*m_conn, job_id, &job_metadata).success()) {
-            throw std::runtime_error(fmt::format("Job with id {} not exists.", to_string(job_id)));
-        }
-        job_metadata_map[job_id] = job_metadata;
-        return job_metadata.get_creation_time();
-    };
-    std::ranges::sort(m_tasks, [&](core::Task const& a, core::Task const& b) {
-        return get_task_job_creation_time(a.get_id()) > get_task_job_creation_time(b.get_id());
-    });
+    std::ranges::sort(
+            m_tasks,
+            [&](core::ScheduleTaskMetadata const& a, core::ScheduleTaskMetadata const& b) {
+                return a.get_job_creation_time() > b.get_job_creation_time();
+            }
+    );
 }
 
 }  // namespace spider::scheduler

src/spider/scheduler/FifoPolicy.hpp

@@ -6,7 +6,6 @@
 #include <string>
 #include <vector>
 
-#include <absl/container/flat_hash_map.h>
 #include <boost/uuid/uuid.hpp>
 
 #include "../core/Task.hpp"
@@ -30,15 +29,12 @@ class FifoPolicy final : public SchedulerPolicy {
 
 private:
     auto fetch_tasks() -> void;
-    auto task_locality_satisfied(core::Task const& task, std::string const& addr) -> bool;
 
     std::shared_ptr<core::MetadataStorage> m_metadata_store;
     std::shared_ptr<core::DataStorage> m_data_store;
     std::shared_ptr<core::StorageConnection> m_conn;
 
-    std::vector<core::Task> m_tasks;
-    // NOLINTNEXTLINE(misc-include-cleaner)
-    absl::flat_hash_map<boost::uuids::uuid, core::Data, std::hash<boost::uuids::uuid>> m_data_cache;
+    std::vector<core::ScheduleTaskMetadata> m_tasks;
 };
 
 }  // namespace spider::scheduler

src/spider/storage/MetadataStorage.hpp

@@ -2,7 +2,6 @@
 #define SPIDER_STORAGE_METADATASTORAGE_HPP
 
 #include <string>
-#include <tuple>
 #include <vector>
 
 #include <boost/uuid/uuid.hpp>
@@ -78,8 +77,8 @@ class MetadataStorage {
             boost::uuids::uuid id,
             boost::uuids::uuid* job_id
     ) -> StorageErr = 0;
-    virtual auto get_ready_tasks(StorageConnection& conn, std::vector<Task>* tasks) -> StorageErr
-                                                                                       = 0;
+    virtual auto get_ready_tasks(StorageConnection& conn, std::vector<ScheduleTaskMetadata>* tasks)
+            -> StorageErr = 0;
     virtual auto set_task_state(StorageConnection& conn, boost::uuids::uuid id, TaskState state)
             -> StorageErr = 0;
     virtual auto set_task_running(StorageConnection& conn, boost::uuids::uuid id) -> StorageErr = 0;
@@ -98,10 +97,8 @@ class MetadataStorage {
             TaskInstance const& instance,
             std::string const& error
     ) -> StorageErr = 0;
-    virtual auto get_task_timeout(
-            StorageConnection& conn,
-            std::vector<std::tuple<TaskInstance, Task>>* tasks
-    ) -> StorageErr = 0;
+    virtual auto get_task_timeout(StorageConnection& conn, std::vector<ScheduleTaskMetadata>* tasks)
+            -> StorageErr = 0;
     virtual auto get_child_tasks(
             StorageConnection& conn,
             boost::uuids::uuid id,