Skip to content

feat: Add worker client that submit tasks and get next task from schedulers #35

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sitaowang1998 merged 8 commits into from
Dec 12, 2024
Merged

feat: Add worker client that submit tasks and get next task from schedulers #35

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
932dd82
Add basic worker client interface
sitaowang1998 Dec 11, 2024
28e3960
Add get active scheduler in metadata storage
sitaowang1998 Dec 12, 2024
b340477
Fix clang tidy
sitaowang1998 Dec 12, 2024
9ffb1d4
Add update task output, update dependent task input & state to task f…
sitaowang1998 Dec 12, 2024
0cff600
Add fetch task from scheduler in task finish
sitaowang1998 Dec 12, 2024
29ca2e2
Split get next task into separate function
sitaowang1998 Dec 12, 2024
b691e37
Fix clang-tidy
sitaowang1998 Dec 12, 2024
97d46a2
Improve code based on code rabbit suggestion
sitaowang1998 Dec 12, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions src/spider/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ set(SPIDER_CORE_SOURCES
set(SPIDER_CORE_HEADERS
core/Error.hpp
core/Data.hpp
core/Driver.hpp
core/KeyValueData.hpp
core/Task.hpp
core/TaskGraph.hpp
Expand Down Expand Up @@ -48,6 +49,8 @@ set(SPIDER_WORKER_SOURCES
worker/TaskExecutorMessage.hpp
worker/message_pipe.cpp
worker/message_pipe.hpp
worker/WorkerClient.hpp
worker/WorkerClient.cpp
CACHE INTERNAL
"spider worker source files"
)
Expand Down
45 changes: 45 additions & 0 deletions src/spider/core/Driver.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,45 @@
#ifndef SPIDER_CORE_DRIVER_HPP
#define SPIDER_CORE_DRIVER_HPP

#include <string>
#include <utility>

#include <boost/uuid/uuid.hpp>

namespace spider::core {

class Driver {
public:
Driver(boost::uuids::uuid const id, std::string addr) : m_id{id}, m_addr{std::move(addr)} {}

@coderabbitai coderabbitai Bot Dec 12, 2024

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

Add address validation in constructor.

Consider validating that the address isn't empty to ensure valid driver instances.

-    Driver(boost::uuids::uuid const id, std::string addr) : m_id{id}, m_addr{std::move(addr)} {}
+    Driver(boost::uuids::uuid const id, std::string addr) {
+        if (addr.empty()) {
+            throw std::invalid_argument("Driver address cannot be empty");
+        }
+        m_id = id;
+        m_addr = std::move(addr);
+    }
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
Driver(boost::uuids::uuid const id, std::string addr) : m_id{id}, m_addr{std::move(addr)} {}
Driver(boost::uuids::uuid const id, std::string addr) {
if (addr.empty()) {
throw std::invalid_argument("Driver address cannot be empty");
}
m_id = id;
m_addr = std::move(addr);
}


[[nodiscard]] auto get_id() const -> boost::uuids::uuid const& { return m_id; }

[[nodiscard]] auto get_addr() const -> std::string const& { return m_addr; }

private:
boost::uuids::uuid m_id;
std::string m_addr;
};

class Scheduler {
public:
Scheduler(boost::uuids::uuid const id, std::string addr, int port)
: m_id{id},
m_addr{std::move(addr)},
m_port{port} {}

Comment on lines +24 to +30

@coderabbitai coderabbitai Bot Dec 12, 2024

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

Add documentation and parameter validation for Scheduler class.

The Scheduler class needs documentation and parameter validation.

+/**
+ * Represents a scheduler node in the Spider system.
+ * Manages task distribution and scheduling operations.
+ */
 class Scheduler {
 public:
+    /**
+     * Constructs a Scheduler instance.
+     * @param id Unique identifier for the scheduler
+     * @param addr Network address (ownership transferred)
+     * @param port Network port number (must be between 1 and 65535)
+     * @throws std::invalid_argument if port is invalid or address is empty
+     */
-    Scheduler(boost::uuids::uuid const id, std::string addr, int port)
-            : m_id{id},
-              m_addr{std::move(addr)},
-              m_port{port} {}
+    Scheduler(boost::uuids::uuid const id, std::string&& addr, int port) {
+        if (addr.empty()) {
+            throw std::invalid_argument("Scheduler address cannot be empty");
+        }
+        if (port <= 0 || port > 65535) {
+            throw std::invalid_argument("Invalid port number");
+        }
+        m_id = id;
+        m_addr = std::move(addr);
+        m_port = port;
+    }
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
class Scheduler {
public:
Scheduler(boost::uuids::uuid const id, std::string addr, int port)
: m_id{id},
m_addr{std::move(addr)},
m_port{port} {}
/**
* Represents a scheduler node in the Spider system.
* Manages task distribution and scheduling operations.
*/
class Scheduler {
public:
/**
* Constructs a Scheduler instance.
* @param id Unique identifier for the scheduler
* @param addr Network address (ownership transferred)
* @param port Network port number (must be between 1 and 65535)
* @throws std::invalid_argument if port is invalid or address is empty
*/
Scheduler(boost::uuids::uuid const id, std::string&& addr, int port) {
if (addr.empty()) {
throw std::invalid_argument("Scheduler address cannot be empty");
}
if (port <= 0 || port > 65535) {
throw std::invalid_argument("Invalid port number");
}
m_id = id;
m_addr = std::move(addr);
m_port = port;
}

[[nodiscard]] auto get_id() const -> boost::uuids::uuid const& { return m_id; }

[[nodiscard]] auto get_addr() const -> std::string const& { return m_addr; }

[[nodiscard]] auto get_port() const -> int { return m_port; }

private:
boost::uuids::uuid m_id;
std::string m_addr;
int m_port;
};

} // namespace spider::core

#endif // SPIDER_CORE_DRIVER_HPP
10 changes: 6 additions & 4 deletions src/spider/storage/MetadataStorage.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@

#include <boost/uuid/uuid.hpp>

#include "../core/Driver.hpp"
#include "../core/Error.hpp"
#include "../core/JobMetadata.hpp"
#include "../core/Task.hpp"
Expand All @@ -25,10 +26,10 @@ class MetadataStorage {
virtual void close() = 0;
virtual auto initialize() -> StorageErr = 0;

virtual auto add_driver(boost::uuids::uuid id, std::string const& addr) -> StorageErr = 0;
virtual auto add_driver(boost::uuids::uuid id, std::string const& addr, int port) -> StorageErr
= 0;
virtual auto add_driver(Driver const& driver) -> StorageErr = 0;
virtual auto add_scheduler(Scheduler const& scheduler) -> StorageErr = 0;
virtual auto get_driver(boost::uuids::uuid id, std::string* addr) -> StorageErr = 0;
virtual auto get_active_scheduler(std::vector<Scheduler>* schedulers) -> StorageErr = 0;

virtual auto
add_job(boost::uuids::uuid job_id, boost::uuids::uuid client_id, TaskGraph const& task_graph
Expand All @@ -47,7 +48,8 @@ class MetadataStorage {
virtual auto get_ready_tasks(std::vector<Task>* tasks) -> StorageErr = 0;
virtual auto set_task_state(boost::uuids::uuid id, TaskState state) -> StorageErr = 0;
virtual auto add_task_instance(TaskInstance const& instance) -> StorageErr = 0;
virtual auto task_finish(TaskInstance const& instance) -> StorageErr = 0;
virtual auto task_finish(TaskInstance const& instance, std::vector<TaskOutput> const& outputs)
-> StorageErr = 0;
virtual auto get_task_timeout(std::vector<TaskInstance>* tasks) -> StorageErr = 0;
virtual auto get_child_tasks(boost::uuids::uuid id, std::vector<Task>* children) -> StorageErr
= 0;
Expand Down
138 changes: 115 additions & 23 deletions src/spider/storage/MysqlStorage.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,8 +26,10 @@
#include <mariadb/conncpp/Properties.hpp>
#include <mariadb/conncpp/ResultSet.hpp>
#include <mariadb/conncpp/Statement.hpp>
#include <mariadb/conncpp/Types.hpp>

#include "../core/Data.hpp"
#include "../core/Driver.hpp"
#include "../core/Error.hpp"
#include "../core/JobMetadata.hpp"
#include "../core/KeyValueData.hpp"
Expand Down Expand Up @@ -274,15 +276,25 @@ auto MySqlMetadataStorage::initialize() -> StorageErr {
return StorageErr{};
}

auto MySqlMetadataStorage::add_driver(boost::uuids::uuid id, std::string const& addr)
-> StorageErr {
namespace {
// NOLINTBEGIN
auto read_id(std::istream* stream) -> boost::uuids::uuid {
std::uint8_t id_bytes[16];
stream->read((char*)id_bytes, 16);
return {id_bytes};
}

// NOLINTEND
} // namespace

auto MySqlMetadataStorage::add_driver(Driver const& driver) -> StorageErr {
try {
std::unique_ptr<sql::PreparedStatement> statement(
m_conn->prepareStatement("INSERT INTO `drivers` (`id`, `address`) VALUES (?, ?)")
);
sql::bytes id_bytes = uuid_get_bytes(id);
sql::bytes id_bytes = uuid_get_bytes(driver.get_id());
statement->setBytes(1, &id_bytes);
statement->setString(2, addr);
statement->setString(2, driver.get_addr());
statement->executeUpdate();
} catch (sql::SQLException& e) {
m_conn->rollback();
Expand All @@ -295,21 +307,20 @@ auto MySqlMetadataStorage::add_driver(boost::uuids::uuid id, std::string const&
return StorageErr{};
}

auto MySqlMetadataStorage::add_driver(boost::uuids::uuid id, std::string const& addr, int port)
-> StorageErr {
auto MySqlMetadataStorage::add_scheduler(Scheduler const& scheduler) -> StorageErr {
try {
std::unique_ptr<sql::PreparedStatement> driver_statement(
m_conn->prepareStatement("INSERT INTO `drivers` (`id`, `address`) VALUES (?, ?)")
);
sql::bytes id_bytes = uuid_get_bytes(id);
sql::bytes id_bytes = uuid_get_bytes(scheduler.get_id());
driver_statement->setBytes(1, &id_bytes);
driver_statement->setString(2, addr);
driver_statement->setString(2, scheduler.get_addr());
driver_statement->executeUpdate();
std::unique_ptr<sql::PreparedStatement> scheduler_statement(m_conn->prepareStatement(
"INSERT INTO `schedulers` (`id`, `port`, `state`) VALUES (?, ?, 'normal')"
));
scheduler_statement->setBytes(1, &id_bytes);
scheduler_statement->setInt(2, port);
scheduler_statement->setInt(2, scheduler.get_port());
scheduler_statement->executeUpdate();
} catch (sql::SQLException& e) {
m_conn->rollback();
Expand Down Expand Up @@ -347,6 +358,27 @@ auto MySqlMetadataStorage::get_driver(boost::uuids::uuid id, std::string* addr)
return StorageErr{};
}

auto MySqlMetadataStorage::get_active_scheduler(std::vector<Scheduler>* schedulers) -> StorageErr {
try {
std::unique_ptr<sql::Statement> statement(m_conn->createStatement());
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'"
));
while (res->next()) {
boost::uuids::uuid const id = read_id(res->getBinaryStream(1));
std::string const addr = res->getString(2).c_str();
int const port = res->getInt(3);
schedulers->emplace_back(id, addr, port);
}
} catch (sql::SQLException& e) {
m_conn->rollback();
return StorageErr{StorageErrType::OtherErr, e.what()};
}
m_conn->commit();
return StorageErr{};
}

void MySqlMetadataStorage::add_task(sql::bytes job_id, Task const& task) {
// Add task
std::unique_ptr<sql::PreparedStatement> task_statement(
Expand Down Expand Up @@ -509,17 +541,6 @@ auto MySqlMetadataStorage::add_job(
return StorageErr{};
}

namespace {
// NOLINTBEGIN
auto read_id(std::istream* stream) -> boost::uuids::uuid {
std::uint8_t id_bytes[16];
stream->read((char*)id_bytes, 16);
return {id_bytes};
}

// NOLINTEND
} // namespace

namespace {

auto fetch_task(std::unique_ptr<sql::ResultSet> const& res) -> Task {
Expand Down Expand Up @@ -960,16 +981,87 @@ auto MySqlMetadataStorage::add_task_instance(TaskInstance const& instance) -> St
return StorageErr{};
}

auto MySqlMetadataStorage::task_finish(TaskInstance const& instance) -> StorageErr {
auto MySqlMetadataStorage::task_finish(
TaskInstance const& instance,
std::vector<TaskOutput> const& outputs
) -> StorageErr {
try {
// Try to submit task instance
std::unique_ptr<sql::PreparedStatement> const statement(m_conn->prepareStatement(
"UPDATE `tasks` SET `instance_id` = ? WHERE `id` = ? AND `instance_id` is NULL"
"UPDATE `tasks` SET `instance_id` = ?, `state` = 'success' WHERE `id` = ? AND "
"`instance_id` is NULL AND `state` = 'running'"
));
sql::bytes id_bytes = uuid_get_bytes(instance.id);
sql::bytes task_id_bytes = uuid_get_bytes(instance.task_id);
statement->setBytes(1, &id_bytes);
statement->setBytes(2, &task_id_bytes);
statement->executeUpdate();
int32_t const update_count = statement->executeUpdate();
if (update_count == 0) {
m_conn->commit();
return StorageErr{};
}

// Update task outputs
std::unique_ptr<sql::PreparedStatement> output_statement(m_conn->prepareStatement(
"UPDATE `task_outputs` SET `value` = ?, `data_id` = ? WHERE `task_id` = ? AND "
"`position` = ?"
));
for (size_t i = 0; i < outputs.size(); ++i) {
TaskOutput const& output = outputs[i];
std::optional<std::string> const& value = output.get_value();
if (value.has_value()) {
output_statement->setString(1, value.value());
} else {
output_statement->setNull(1, sql::DataType::VARCHAR);
}
std::optional<boost::uuids::uuid> const& data_id = output.get_data_id();
if (data_id.has_value()) {
sql::bytes data_id_bytes = uuid_get_bytes(data_id.value());
output_statement->setBytes(2, &data_id_bytes);
} else {
output_statement->setNull(2, sql::DataType::BINARY);
}
output_statement->setBytes(3, &task_id_bytes);
output_statement->setUInt(4, i);
output_statement->executeUpdate();
}

// Update task inputs
std::unique_ptr<sql::PreparedStatement> input_statement(m_conn->prepareStatement(
"UPDATE `task_inputs` SET `value` = ?, `data_id` = ? WHERE `output_task_id` = ? "
"AND `output_task_position` = ?"
));
for (size_t i = 0; i < outputs.size(); ++i) {
TaskOutput const& output = outputs[i];
std::optional<std::string> const& value = output.get_value();
if (value.has_value()) {
input_statement->setString(1, value.value());
} else {
input_statement->setNull(1, sql::DataType::VARCHAR);
}
std::optional<boost::uuids::uuid> const& data_id = output.get_data_id();
if (data_id.has_value()) {
sql::bytes data_id_bytes = uuid_get_bytes(data_id.value());
input_statement->setBytes(2, &data_id_bytes);
} else {
input_statement->setNull(2, sql::DataType::BINARY);
}
input_statement->setBytes(3, &task_id_bytes);
input_statement->setUInt(4, i);
input_statement->executeUpdate();
}

// Set task states to ready if all inputs are available
std::unique_ptr<sql::PreparedStatement> ready_statement(m_conn->prepareStatement(
"UPDATE `tasks` SET `state` = 'ready' WHERE `id` IN (SELECT `task_id` FROM "
"`task_inputs` WHERE `output_task_id` = ?) AND `state` = 'pending' AND NOT EXISTS "
"(SELECT `task_id` FROM `task_inputs` WHERE `task_id` IN (SELECT `task_id` FROM "
"`task_inputs` WHERE `output_task_id` = ?) AND `value` IS NULL AND `data_id` IS "
"NULL)"
));
ready_statement->setBytes(1, &task_id_bytes);
ready_statement->setBytes(2, &task_id_bytes);
ready_statement->executeUpdate();
} catch (sql::SQLException& e) {
m_conn->rollback();
if (e.getErrorCode() == ErDupKey || e.getErrorCode() == ErDupEntry) {
Expand Down
10 changes: 6 additions & 4 deletions src/spider/storage/MysqlStorage.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
#include <mariadb/conncpp/ResultSet.hpp>

#include "../core/Data.hpp"
#include "../core/Driver.hpp"
#include "../core/Error.hpp"
#include "../core/JobMetadata.hpp"
#include "../core/KeyValueData.hpp"
Expand All @@ -31,10 +32,10 @@ class MySqlMetadataStorage : public MetadataStorage {
auto connect(std::string const& url) -> StorageErr override;
void close() override;
auto initialize() -> StorageErr override;
auto add_driver(boost::uuids::uuid id, std::string const& addr) -> StorageErr override;
auto
add_driver(boost::uuids::uuid id, std::string const& addr, int port) -> StorageErr override;
auto add_driver(Driver const& driver) -> StorageErr override;
auto add_scheduler(Scheduler const& scheduler) -> StorageErr override;
auto get_driver(boost::uuids::uuid id, std::string* addr) -> StorageErr override;
auto get_active_scheduler(std::vector<Scheduler>* schedulers) -> StorageErr override;
auto
add_job(boost::uuids::uuid job_id, boost::uuids::uuid client_id, TaskGraph const& task_graph
) -> StorageErr override;
Expand All @@ -51,7 +52,8 @@ class MySqlMetadataStorage : public MetadataStorage {
auto get_ready_tasks(std::vector<Task>* tasks) -> StorageErr override;
auto set_task_state(boost::uuids::uuid id, TaskState state) -> StorageErr override;
auto add_task_instance(TaskInstance const& instance) -> StorageErr override;
auto task_finish(TaskInstance const& instance) -> StorageErr override;
auto task_finish(TaskInstance const& instance, std::vector<TaskOutput> const& outputs)
-> StorageErr override;
auto get_task_timeout(std::vector<TaskInstance>* tasks) -> StorageErr override;
auto get_child_tasks(boost::uuids::uuid id, std::vector<Task>* children) -> StorageErr override;
auto get_parent_tasks(boost::uuids::uuid id, std::vector<Task>* tasks) -> StorageErr override;
Expand Down
Loading