fix: Batch MySQL queries in getting ready tasks - #83
Conversation
WalkthroughThe pull request introduces the Changes
Sequence Diagram(s)sequenceDiagram
participant Worker
participant FifoPolicy
participant TaskMetadata as ScheduleTaskMetadata
Worker->>FifoPolicy: Request next task
FifoPolicy->>TaskMetadata: Retrieve hard_localities and other metadata
alt Task has no locality restrictions or meets requirements
FifoPolicy-->>Worker: Assign task
else
FifoPolicy-->>Worker: Skip task
end
sequenceDiagram
participant Caller
participant Storage as MySqlMetadataStorage
participant Database
Caller->>Storage: get_ready_tasks(conn, tasks)
Storage->>Database: Execute SQL query (select job_id, function_name, etc.)
Database-->>Storage: Return task records
Storage->>Caller: Return populated ScheduleTaskMetadata list
Possibly related PRs
📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (1)
⏰ Context from checks skipped due to timeout of 90000ms (3)
🔇 Additional comments (10)
✨ Finishing Touches
🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🔭 Outside diff range comments (1)
src/spider/storage/mysql/MySqlStorage.cpp (1)
1629-1653:⚠️ Potential issuePotential SQL syntax error in the
SELECT FROMquery.
In the snippet below, the statement is incomplete and may not compile under typical SQL engines:"SELECT FROM `task_instances` as `t1` JOIN `tasks` ON`t1`.`task_id` = `tasks`.`id` WHERE `t1.task_id` = ? ..."You likely need to specify the selected columns. For example:
- "SELECT FROM `task_instances` as `t1` JOIN `tasks` ON`t1`.`task_id` = `tasks`.`id` WHERE `t1.task_id` = ? ..." + "SELECT t1.`task_id` FROM `task_instances` AS `t1` JOIN `tasks` ON t1.`task_id` = `tasks`.`id` WHERE t1.`task_id` = ? AND ..."This fix ensures the query is valid and returns the correct columns.
🧹 Nitpick comments (1)
src/spider/core/Task.hpp (1)
146-159: Consider adding validation in setter and adder methods.While the implementation is functional, adding validation for inputs in methods like
add_hard_localityandadd_soft_localitywould improve robustness.auto add_hard_locality(std::string const& locality) -> void { + if (!locality.empty()) { m_hard_localities.push_back(locality); + } } auto add_soft_locality(std::string const& locality) -> void { + if (!locality.empty()) { m_soft_localities.push_back(locality); + } }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
src/spider/core/Task.hpp(2 hunks)src/spider/scheduler/FifoPolicy.cpp(1 hunks)src/spider/scheduler/FifoPolicy.hpp(1 hunks)src/spider/storage/MetadataStorage.hpp(2 hunks)src/spider/storage/mysql/MySqlStorage.cpp(4 hunks)src/spider/storage/mysql/MySqlStorage.hpp(2 hunks)
🧰 Additional context used
🧬 Code Definitions (4)
src/spider/storage/MetadataStorage.hpp (1)
src/spider/storage/DataStorage.hpp (15)
conn(23-23)conn(25-29)conn(30-34)conn(35-35)conn(37-37)conn(38-38)conn(39-43)conn(44-48)conn(49-53)conn(54-58)conn(59-59)conn(61-61)conn(63-63)conn(65-70)conn(71-76)
src/spider/storage/mysql/MySqlStorage.hpp (1)
src/spider/storage/MetadataStorage.hpp (16)
conn(27-27)conn(29-29)conn(30-30)conn(32-33)conn(35-40)conn(41-47)conn(48-49)conn(50-51)conn(52-53)conn(54-58)conn(59-63)conn(64-68)conn(69-69)conn(70-70)conn(71-72)conn(73-73)
src/spider/storage/mysql/MySqlStorage.cpp (2)
src/spider/storage/MetadataStorage.hpp (16)
conn(27-27)conn(29-29)conn(30-30)conn(32-33)conn(35-40)conn(41-47)conn(48-49)conn(50-51)conn(52-53)conn(54-58)conn(59-63)conn(64-68)conn(69-69)conn(70-70)conn(71-72)conn(73-73)src/spider/core/Task.hpp (2)
client_id(146-146)client_id(146-146)
src/spider/core/Task.hpp (1)
src/spider/core/Data.hpp (6)
nodiscard(20-20)nodiscard(22-22)nodiscard(24-26)nodiscard(28-28)locality(30-30)locality(30-30)
⏰ Context from checks skipped due to timeout of 90000ms (3)
- GitHub Check: non-storage-unit-tests (ubuntu-24.04)
- GitHub Check: non-storage-unit-tests (ubuntu-22.04)
- GitHub Check: lint
🔇 Additional comments (15)
src/spider/core/Task.hpp (4)
4-4: Include of<chrono>added for time point handling.The
<chrono>header is correctly added to support thestd::chrono::system_clock::time_pointused in the newScheduleTaskMetadataclass.
113-125: NewScheduleTaskMetadataclass looks well-structured.The class design effectively encapsulates all necessary metadata for scheduled tasks with appropriate constructor parameters and member initialization.
126-144: Well-implemented accessor methods with appropriate const qualifiers.The getter methods properly use
[[nodiscard]]attribute, return by value or const reference as appropriate, and maintain const correctness.
161-169: Member variables layout is clean and well-organized.The private section properly encapsulates all necessary metadata fields with clear naming conventions.
src/spider/scheduler/FifoPolicy.hpp (1)
37-37: Updated task collection type to use the new metadata structure.The change from
std::vector<core::Task>tostd::vector<core::ScheduleTaskMetadata>aligns with the PR objective of optimizing task metadata handling.src/spider/scheduler/FifoPolicy.cpp (3)
40-49: Simplified locality check with direct access to hard_localities.The lambda function efficiently determines if a task can be scheduled on a worker by checking if the worker's address is in the task's hard localities list or if no hard localities are specified.
63-68: Improved task sorting based on job creation time.The sort comparison now directly uses the job creation time from the metadata structure, simplifying the code and eliminating the need for additional job metadata queries.
58-61:Details
❓ Verification inconclusive
Verify if two separate function calls are intentional.
Both
get_ready_tasksandget_task_timeoutare adding tasks to the samem_tasksvector. Ensure this is intentional and doesn't result in duplicate tasks.
🏁 Script executed:
#!/bin/bash # Check if both methods can add the same task to the vector # Look for implementation of get_ready_tasks and get_task_timeout rg -A 10 "get_ready_tasks.*ScheduleTaskMetadata" --type cpp rg -A 10 "get_task_timeout.*ScheduleTaskMetadata" --type cpp # Look for how tasks are identified in these methods rg "task.*id.*unique" --type cppLength of output: 4305
Action Required: Confirm Duplicate Handling in Task Fetching
At
src/spider/scheduler/FifoPolicy.cpp(lines 58-61), there are two distinct calls—get_ready_tasksandget_task_timeout—both appending to them_tasksvector. The grep results confirm their declarations in the storage interfaces but did not reveal any explicit duplicate filtering logic in the searched patterns. This leaves some uncertainty regarding whether the potential overlap of tasks is intentional or if additional de-duplication safeguards should be implemented.
- Review the implementations of both functions (in
MetadataStorage.hppandMySqlStorage.hpp) to confirm that duplicate tasks are either prevented or intentionally allowed.- Manually verify that tasks identified as ready and those flagged for timeout do not overlap, or that any overlap is handled appropriately elsewhere in the code.
src/spider/storage/MetadataStorage.hpp (2)
80-81: Method signature updated to useScheduleTaskMetadata.The
get_ready_tasksmethod now accepts a vector ofScheduleTaskMetadatainstead ofTask, aligning with the objective of batching queries and reducing storage accesses.
100-101: Method signature simplified to useScheduleTaskMetadata.The
get_task_timeoutmethod now returns tasks asScheduleTaskMetadataobjects instead of tuples ofTaskInstanceandTask, which streamlines the API and reduces the need for multiple queries.src/spider/storage/mysql/MySqlStorage.hpp (2)
83-84: Adoption ofScheduleTaskMetadataaligns well with the batching approach.
Switching the parameter type tostd::vector<ScheduleTaskMetadata>*is consistent with the new design goals. It eliminates the overhead of fetching additional information per task during scheduling.
99-100: Consistent signature update forget_task_timeout.
Ensuring that all downstream usages now rely onScheduleTaskMetadatawill help standardize metadata handling in timeouts.src/spider/storage/mysql/MySqlStorage.cpp (3)
17-17: Include statement forabsl::flat_hash_map.
Including<absl/container/flat_hash_map.h>reflects the shift towards more efficient hashing structures. Please confirm that your build environment reliably provides Abseil libraries and that you have no linking issues.
1234-1340: Refinedget_ready_tasksmethod enhances performance with fewer queries and richer metadata.
- Bulk retrieval of job metadata via batched prepared statements (
job_statement) reduces round-trips to the database.- Collecting data locality pointers in a single pass is efficient; however, consider verifying correct handling of tasks that have no associated data. This ensures you don’t inadvertently miss tasks or introduce errors when data is absent.
Overall, this method effectively consolidates relevant scheduling metadata.
1654-1776: Timed-out tasks now return enriched scheduling metadata.
- The batching of job lookups (lines 1705–1716) is a solid approach to reduce multiple queries, but ensure that large batches are handled efficiently by the driver.
- Data locality gathering (lines 1739–1764) matches the logic from
get_ready_tasks, maintaining consistency for tasks that time out.- Confirm that you properly handle corner cases where tasks may not hold locality data.
Overall, this method neatly parallels
get_ready_tasksfor timeouts, improving performance and uniformity.
There was a problem hiding this comment.
Actionable comments posted: 1
🔭 Outside diff range comments (1)
src/spider/storage/mysql/MySqlStorage.cpp (1)
1234-1341:⚠️ Potential issueFix missing insertion for the first task ID.
When adding tojob_id_to_task_idsat lines 1261-1266, the vector is only populated in theelsebranch, leaving out the first task.Apply this diff to ensure you always store the task ID:
-if (job_id_to_task_ids.find(job_id) == job_id_to_task_ids.end()) { - job_id_to_task_ids[job_id] = std::vector<boost::uuids::uuid>{}; -} else { - job_id_to_task_ids[job_id].emplace_back(task_id); -} +if (job_id_to_task_ids.find(job_id) == job_id_to_task_ids.end()) { + job_id_to_task_ids[job_id] = std::vector<boost::uuids::uuid>{}; +} +job_id_to_task_ids[job_id].emplace_back(task_id);
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
src/spider/storage/mysql/MySqlStorage.cpp(3 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (3)
- GitHub Check: non-storage-unit-tests (ubuntu-24.04)
- GitHub Check: non-storage-unit-tests (ubuntu-22.04)
- GitHub Check: lint
🔇 Additional comments (2)
src/spider/storage/mysql/MySqlStorage.cpp (2)
17-17: Header inclusion looks good.
No issues here.
1629-1630: Signature change looks consistent.
Switching from tuples toScheduleTaskMetadataaligns well with the new structure.
| std::unique_ptr<sql::ResultSet> const task_timeout_res(statement->executeQuery( | ||
| "SELECT `t1`.`task_id` FROM `task_instances` as `t1` JOIN `tasks` ON " | ||
| "`t1`.`task_id` = `tasks`.`id` WHERE `tasks`.`timeout` > 0.0001 AND " | ||
| "TIMESTAMPDIFF(MICROSECOND, `t1`.`start_time`, CURRENT_TIMESTAMP()) > " | ||
| "`tasks`.`timeout` * 1000" | ||
| )); | ||
| if (task_timeout_res->rowsCount() == 0) { | ||
| static_cast<MySqlConnection&>(conn)->commit(); | ||
| return StorageErr{}; | ||
| } | ||
|
|
||
| std::unique_ptr<sql::PreparedStatement> not_timeout_statement( | ||
| static_cast<MySqlConnection&>(conn)->prepareStatement( | ||
| "SELECT FROM `task_instances` as `t1` JOIN `tasks` ON`t1`.`task_id` = " | ||
| "`tasks`.`id` WHERE `t1.task_id` = ? AND TIMESTAMPDIFF(MICROSECOND, " | ||
| "`t1`.`start_time`, CURRENT_TIMESTAMP()) < `tasks`.`timeout` * 1000" | ||
| "SELECT `t1`.`task_id` FROM `task_instances` as `t1` JOIN `tasks` ON " | ||
| "`t1`.`task_id` = `tasks`.`id` WHERE `t1.task_id` = ? AND " | ||
| "TIMESTAMPDIFF(MICROSECOND, `t1`.`start_time`, CURRENT_TIMESTAMP()) < " | ||
| "`tasks`.`timeout` * 1000" | ||
| ) | ||
| ); | ||
|
|
||
| absl::flat_hash_set<boost::uuids::uuid> task_ids; | ||
| while (task_timeout_res->next()) { | ||
| boost::uuids::uuid const task_id | ||
| = read_id(task_timeout_res->getBinaryStream("task_id")); | ||
| task_ids.insert(task_id); | ||
| sql::bytes task_id_bytes = uuid_get_bytes(task_id); | ||
| not_timeout_statement->setBytes(1, &task_id_bytes); | ||
| not_timeout_statement->addBatch(); | ||
| } | ||
| not_timeout_statement->execute(); | ||
| std::unique_ptr<sql::ResultSet> const not_timeout_res(not_timeout_statement->getResultSet() | ||
| ); | ||
| while (not_timeout_res->next()) { | ||
| boost::uuids::uuid const task_id = read_id(not_timeout_res->getBinaryStream("task_id")); | ||
| task_ids.erase(task_id); | ||
| } | ||
|
|
||
| if (task_ids.empty()) { | ||
| static_cast<MySqlConnection&>(conn)->commit(); | ||
| return StorageErr{}; | ||
| } | ||
|
|
||
| // Get task metadata | ||
| std::unique_ptr<sql::PreparedStatement> task_statement( | ||
| static_cast<MySqlConnection&>(conn)->prepareStatement( | ||
| "SELECT `id`, `func_name`, `state`, `timeout` FROM `tasks` WHERE `id` = ?" | ||
| "SELECT `id`, `func_name`, `job_id` FROM `tasks` WHERE `id` = ?" | ||
| ) | ||
| ); | ||
| while (res->next()) { | ||
| boost::uuids::uuid const task_instance_id = read_id(res->getBinaryStream("id")); | ||
| boost::uuids::uuid const task_id = read_id(res->getBinaryStream("task_id")); | ||
| for (boost::uuids::uuid const& task_id : task_ids) { | ||
| sql::bytes task_id_bytes = uuid_get_bytes(task_id); | ||
| // Check all task instance have timed out | ||
| not_timeout_statement->setBytes(1, &task_id_bytes); | ||
| std::unique_ptr<sql::ResultSet> not_timeout_res(not_timeout_statement->executeQuery()); | ||
| if (not_timeout_res->rowsCount() > 0) { | ||
| continue; | ||
| task_statement->setBytes(1, &task_id_bytes); | ||
| task_statement->addBatch(); | ||
| } | ||
| task_statement->execute(); | ||
| std::unique_ptr<sql::ResultSet> const task_res(task_statement->getResultSet()); | ||
|
|
||
| absl::flat_hash_map<boost::uuids::uuid, ScheduleTaskMetadata> new_tasks; | ||
| absl::flat_hash_map<boost::uuids::uuid, std::vector<boost::uuids::uuid>> job_id_to_task_ids; | ||
| while (task_res->next()) { | ||
| boost::uuids::uuid const task_id = read_id(task_res->getBinaryStream("id")); | ||
| boost::uuids::uuid const job_id = read_id(task_res->getBinaryStream("job_id")); | ||
| std::string const function_name = get_sql_string(task_res->getString("func_name")); | ||
| new_tasks.emplace(task_id, ScheduleTaskMetadata{task_id, function_name, job_id}); | ||
| if (job_id_to_task_ids.find(job_id) == job_id_to_task_ids.end()) { | ||
| job_id_to_task_ids[job_id] = std::vector<boost::uuids::uuid>{}; | ||
| } else { | ||
| job_id_to_task_ids[job_id].emplace_back(task_id); | ||
| } | ||
| } | ||
|
|
||
| // Fetch task | ||
| task_statement->setBytes(1, &task_id_bytes); | ||
| std::unique_ptr<sql::ResultSet> task_res(task_statement->executeQuery()); | ||
| if (task_res->next()) { | ||
| Task const task = fetch_full_task(static_cast<MySqlConnection&>(conn), task_res); | ||
| tasks->emplace_back(TaskInstance{task_instance_id, task_id}, task); | ||
| // Get all job metadata | ||
| std::unique_ptr<sql::PreparedStatement> job_statement( | ||
| static_cast<MySqlConnection&>(conn)->prepareStatement( | ||
| "SELECT `id`, `client_id`, `creation_time` FROM `jobs` WHERE `id` = ?" | ||
| ) | ||
| ); | ||
| for (auto const& iter : job_id_to_task_ids) { | ||
| sql::bytes job_id_bytes = uuid_get_bytes(iter.first); | ||
| job_statement->setBytes(1, &job_id_bytes); | ||
| job_statement->addBatch(); | ||
| } | ||
| job_statement->execute(); | ||
| std::unique_ptr<sql::ResultSet> const job_res(job_statement->getResultSet()); | ||
| 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")); | ||
| std::optional<std::chrono::system_clock::time_point> const optional_creation_time | ||
| = parse_timestamp(get_sql_string(job_res->getString("creation_time"))); | ||
| if (false == optional_creation_time.has_value()) { | ||
| static_cast<MySqlConnection&>(conn)->rollback(); | ||
| return StorageErr{ | ||
| StorageErrType::OtherErr, | ||
| fmt::format( | ||
| "Cannot parse timestamp {}", | ||
| get_sql_string(job_res->getString("creation_time")) | ||
| ) | ||
| }; | ||
| } | ||
| 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()); | ||
| } | ||
| } | ||
|
|
||
| // Get all data localities | ||
| std::unique_ptr<sql::PreparedStatement> locality_statement( | ||
| static_cast<MySqlConnection&>(conn)->prepareStatement( | ||
| "SELECT `task_inputs`.`task_id`, `data`.`hard_locality`, " | ||
| "`data_locality`.`address` FROM `task_inputs` JOIN `data` ON " | ||
| "`task_inputs`.`data_id` = `data`.`id` JOIN `data_locality` ON `data`.`id` " | ||
| "= `data_locality`.`id` WHERE `task_inputs`.`task_id` = ? AND " | ||
| "`task_inputs`.`task_id` IS NOT NULL" | ||
| ) | ||
| ); | ||
| for (auto const& iter : new_tasks) { | ||
| sql::bytes task_id_bytes = uuid_get_bytes(iter.first); | ||
| locality_statement->setBytes(1, &task_id_bytes); | ||
| locality_statement->addBatch(); | ||
| } | ||
| locality_statement->execute(); | ||
| std::unique_ptr<sql::ResultSet> const locality_res(locality_statement->getResultSet()); | ||
| while (locality_res->next()) { | ||
| boost::uuids::uuid const task_id = read_id(locality_res->getBinaryStream("task_id")); | ||
| bool const hard_locality = locality_res->getBoolean("hard_locality"); | ||
| std::string const address = get_sql_string(locality_res->getString("address")); | ||
| if (hard_locality) { | ||
| new_tasks[task_id].add_hard_locality(address); | ||
| } else { | ||
| new_tasks[task_id].add_soft_locality(address); | ||
| } | ||
| } | ||
|
|
||
| // Add all tasks to the output | ||
| for (auto const& iter : new_tasks) { | ||
| tasks->emplace_back(iter.second); | ||
| } | ||
| } catch (sql::SQLException& e) { |
There was a problem hiding this comment.
Resolve the identical bug in job_id_to_task_ids.
Similar to get_ready_tasks, the first task for each job is left out in lines 1698-1703.
Proposed fix:
-if (job_id_to_task_ids.find(job_id) == job_id_to_task_ids.end()) {
- job_id_to_task_ids[job_id] = std::vector<boost::uuids::uuid>{};
-} else {
- job_id_to_task_ids[job_id].emplace_back(task_id);
-}
+if (job_id_to_task_ids.find(job_id) == job_id_to_task_ids.end()) {
+ job_id_to_task_ids[job_id] = std::vector<boost::uuids::uuid>{};
+}
+job_id_to_task_ids[job_id].emplace_back(task_id);📝 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.
| std::unique_ptr<sql::ResultSet> const task_timeout_res(statement->executeQuery( | |
| "SELECT `t1`.`task_id` FROM `task_instances` as `t1` JOIN `tasks` ON " | |
| "`t1`.`task_id` = `tasks`.`id` WHERE `tasks`.`timeout` > 0.0001 AND " | |
| "TIMESTAMPDIFF(MICROSECOND, `t1`.`start_time`, CURRENT_TIMESTAMP()) > " | |
| "`tasks`.`timeout` * 1000" | |
| )); | |
| if (task_timeout_res->rowsCount() == 0) { | |
| static_cast<MySqlConnection&>(conn)->commit(); | |
| return StorageErr{}; | |
| } | |
| std::unique_ptr<sql::PreparedStatement> not_timeout_statement( | |
| static_cast<MySqlConnection&>(conn)->prepareStatement( | |
| "SELECT FROM `task_instances` as `t1` JOIN `tasks` ON`t1`.`task_id` = " | |
| "`tasks`.`id` WHERE `t1.task_id` = ? AND TIMESTAMPDIFF(MICROSECOND, " | |
| "`t1`.`start_time`, CURRENT_TIMESTAMP()) < `tasks`.`timeout` * 1000" | |
| "SELECT `t1`.`task_id` FROM `task_instances` as `t1` JOIN `tasks` ON " | |
| "`t1`.`task_id` = `tasks`.`id` WHERE `t1.task_id` = ? AND " | |
| "TIMESTAMPDIFF(MICROSECOND, `t1`.`start_time`, CURRENT_TIMESTAMP()) < " | |
| "`tasks`.`timeout` * 1000" | |
| ) | |
| ); | |
| absl::flat_hash_set<boost::uuids::uuid> task_ids; | |
| while (task_timeout_res->next()) { | |
| boost::uuids::uuid const task_id | |
| = read_id(task_timeout_res->getBinaryStream("task_id")); | |
| task_ids.insert(task_id); | |
| sql::bytes task_id_bytes = uuid_get_bytes(task_id); | |
| not_timeout_statement->setBytes(1, &task_id_bytes); | |
| not_timeout_statement->addBatch(); | |
| } | |
| not_timeout_statement->execute(); | |
| std::unique_ptr<sql::ResultSet> const not_timeout_res(not_timeout_statement->getResultSet() | |
| ); | |
| while (not_timeout_res->next()) { | |
| boost::uuids::uuid const task_id = read_id(not_timeout_res->getBinaryStream("task_id")); | |
| task_ids.erase(task_id); | |
| } | |
| if (task_ids.empty()) { | |
| static_cast<MySqlConnection&>(conn)->commit(); | |
| return StorageErr{}; | |
| } | |
| // Get task metadata | |
| std::unique_ptr<sql::PreparedStatement> task_statement( | |
| static_cast<MySqlConnection&>(conn)->prepareStatement( | |
| "SELECT `id`, `func_name`, `state`, `timeout` FROM `tasks` WHERE `id` = ?" | |
| "SELECT `id`, `func_name`, `job_id` FROM `tasks` WHERE `id` = ?" | |
| ) | |
| ); | |
| while (res->next()) { | |
| boost::uuids::uuid const task_instance_id = read_id(res->getBinaryStream("id")); | |
| boost::uuids::uuid const task_id = read_id(res->getBinaryStream("task_id")); | |
| for (boost::uuids::uuid const& task_id : task_ids) { | |
| sql::bytes task_id_bytes = uuid_get_bytes(task_id); | |
| // Check all task instance have timed out | |
| not_timeout_statement->setBytes(1, &task_id_bytes); | |
| std::unique_ptr<sql::ResultSet> not_timeout_res(not_timeout_statement->executeQuery()); | |
| if (not_timeout_res->rowsCount() > 0) { | |
| continue; | |
| task_statement->setBytes(1, &task_id_bytes); | |
| task_statement->addBatch(); | |
| } | |
| task_statement->execute(); | |
| std::unique_ptr<sql::ResultSet> const task_res(task_statement->getResultSet()); | |
| absl::flat_hash_map<boost::uuids::uuid, ScheduleTaskMetadata> new_tasks; | |
| absl::flat_hash_map<boost::uuids::uuid, std::vector<boost::uuids::uuid>> job_id_to_task_ids; | |
| while (task_res->next()) { | |
| boost::uuids::uuid const task_id = read_id(task_res->getBinaryStream("id")); | |
| boost::uuids::uuid const job_id = read_id(task_res->getBinaryStream("job_id")); | |
| std::string const function_name = get_sql_string(task_res->getString("func_name")); | |
| new_tasks.emplace(task_id, ScheduleTaskMetadata{task_id, function_name, job_id}); | |
| if (job_id_to_task_ids.find(job_id) == job_id_to_task_ids.end()) { | |
| job_id_to_task_ids[job_id] = std::vector<boost::uuids::uuid>{}; | |
| } else { | |
| job_id_to_task_ids[job_id].emplace_back(task_id); | |
| } | |
| } | |
| // Fetch task | |
| task_statement->setBytes(1, &task_id_bytes); | |
| std::unique_ptr<sql::ResultSet> task_res(task_statement->executeQuery()); | |
| if (task_res->next()) { | |
| Task const task = fetch_full_task(static_cast<MySqlConnection&>(conn), task_res); | |
| tasks->emplace_back(TaskInstance{task_instance_id, task_id}, task); | |
| // Get all job metadata | |
| std::unique_ptr<sql::PreparedStatement> job_statement( | |
| static_cast<MySqlConnection&>(conn)->prepareStatement( | |
| "SELECT `id`, `client_id`, `creation_time` FROM `jobs` WHERE `id` = ?" | |
| ) | |
| ); | |
| for (auto const& iter : job_id_to_task_ids) { | |
| sql::bytes job_id_bytes = uuid_get_bytes(iter.first); | |
| job_statement->setBytes(1, &job_id_bytes); | |
| job_statement->addBatch(); | |
| } | |
| job_statement->execute(); | |
| std::unique_ptr<sql::ResultSet> const job_res(job_statement->getResultSet()); | |
| 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")); | |
| std::optional<std::chrono::system_clock::time_point> const optional_creation_time | |
| = parse_timestamp(get_sql_string(job_res->getString("creation_time"))); | |
| if (false == optional_creation_time.has_value()) { | |
| static_cast<MySqlConnection&>(conn)->rollback(); | |
| return StorageErr{ | |
| StorageErrType::OtherErr, | |
| fmt::format( | |
| "Cannot parse timestamp {}", | |
| get_sql_string(job_res->getString("creation_time")) | |
| ) | |
| }; | |
| } | |
| 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()); | |
| } | |
| } | |
| // Get all data localities | |
| std::unique_ptr<sql::PreparedStatement> locality_statement( | |
| static_cast<MySqlConnection&>(conn)->prepareStatement( | |
| "SELECT `task_inputs`.`task_id`, `data`.`hard_locality`, " | |
| "`data_locality`.`address` FROM `task_inputs` JOIN `data` ON " | |
| "`task_inputs`.`data_id` = `data`.`id` JOIN `data_locality` ON `data`.`id` " | |
| "= `data_locality`.`id` WHERE `task_inputs`.`task_id` = ? AND " | |
| "`task_inputs`.`task_id` IS NOT NULL" | |
| ) | |
| ); | |
| for (auto const& iter : new_tasks) { | |
| sql::bytes task_id_bytes = uuid_get_bytes(iter.first); | |
| locality_statement->setBytes(1, &task_id_bytes); | |
| locality_statement->addBatch(); | |
| } | |
| locality_statement->execute(); | |
| std::unique_ptr<sql::ResultSet> const locality_res(locality_statement->getResultSet()); | |
| while (locality_res->next()) { | |
| boost::uuids::uuid const task_id = read_id(locality_res->getBinaryStream("task_id")); | |
| bool const hard_locality = locality_res->getBoolean("hard_locality"); | |
| std::string const address = get_sql_string(locality_res->getString("address")); | |
| if (hard_locality) { | |
| new_tasks[task_id].add_hard_locality(address); | |
| } else { | |
| new_tasks[task_id].add_soft_locality(address); | |
| } | |
| } | |
| // Add all tasks to the output | |
| for (auto const& iter : new_tasks) { | |
| tasks->emplace_back(iter.second); | |
| } | |
| } catch (sql::SQLException& e) { | |
| while (task_res->next()) { | |
| boost::uuids::uuid const task_id = read_id(task_res->getBinaryStream("id")); | |
| boost::uuids::uuid const job_id = read_id(task_res->getBinaryStream("job_id")); | |
| std::string const function_name = get_sql_string(task_res->getString("func_name")); | |
| new_tasks.emplace(task_id, ScheduleTaskMetadata{task_id, function_name, job_id}); | |
| - if (job_id_to_task_ids.find(job_id) == job_id_to_task_ids.end()) { | |
| - job_id_to_task_ids[job_id] = std::vector<boost::uuids::uuid>{}; | |
| - } else { | |
| - job_id_to_task_ids[job_id].emplace_back(task_id); | |
| - } | |
| + if (job_id_to_task_ids.find(job_id) == job_id_to_task_ids.end()) { | |
| + job_id_to_task_ids[job_id] = std::vector<boost::uuids::uuid>{}; | |
| + } | |
| + job_id_to_task_ids[job_id].emplace_back(task_id); | |
| } |
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (2)
src/spider/storage/mysql/MySqlStorage.cpp (2)
1261-1265:⚠️ Potential issueFix task ID insertion logic for the first task in a job
There's a bug in the logic for populating job_id_to_task_ids. The first task for each job is added to the vector, but subsequent tasks for the same job may not be correctly added due to the conditional structure.
Apply this fix:
-if (job_id_to_task_ids.find(job_id) == job_id_to_task_ids.end()) { - job_id_to_task_ids[job_id] = std::vector<boost::uuids::uuid>{task_id}; -} else { - job_id_to_task_ids[job_id].emplace_back(task_id); -} +if (job_id_to_task_ids.find(job_id) == job_id_to_task_ids.end()) { + job_id_to_task_ids[job_id] = std::vector<boost::uuids::uuid>{}; +} +job_id_to_task_ids[job_id].emplace_back(task_id);
1698-1702:⚠️ Potential issueFix the same task ID insertion bug in get_task_timeout
The same bug identified earlier is also present in the get_task_timeout method.
Apply this fix:
-if (job_id_to_task_ids.find(job_id) == job_id_to_task_ids.end()) { - job_id_to_task_ids[job_id] = std::vector<boost::uuids::uuid>{task_id}; -} else { - job_id_to_task_ids[job_id].emplace_back(task_id); -} +if (job_id_to_task_ids.find(job_id) == job_id_to_task_ids.end()) { + job_id_to_task_ids[job_id] = std::vector<boost::uuids::uuid>{}; +} +job_id_to_task_ids[job_id].emplace_back(task_id);
🧹 Nitpick comments (1)
src/spider/storage/mysql/MySqlStorage.cpp (1)
1234-1770: Consider extracting common task metadata fetching logicBoth get_ready_tasks and get_task_timeout methods share significant code duplication in how they fetch and populate task metadata. Consider extracting this common logic into a helper method to improve maintainability.
Example refactoring approach:
// Helper function to populate ScheduleTaskMetadata objects void fetch_and_populate_task_metadata( MySqlConnection& conn, absl::flat_hash_set<boost::uuids::uuid>& task_ids, std::vector<ScheduleTaskMetadata>* tasks) { // Get task metadata absl::flat_hash_map<boost::uuids::uuid, ScheduleTaskMetadata> new_tasks; absl::flat_hash_map<boost::uuids::uuid, std::vector<boost::uuids::uuid>> job_id_to_task_ids; // Batch fetch task basic info // Batch fetch job metadata // Batch fetch locality info // Populate output vector for (auto const& iter : new_tasks) { tasks->emplace_back(iter.second); } }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
src/spider/storage/mysql/MySqlStorage.cpp(3 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (3)
- GitHub Check: lint
- GitHub Check: non-storage-unit-tests (ubuntu-24.04)
- GitHub Check: non-storage-unit-tests (ubuntu-22.04)
🔇 Additional comments (7)
src/spider/storage/mysql/MySqlStorage.cpp (7)
17-17: Good inclusion of flat_hash_map for improved performanceThe addition of the Abseil flat_hash_map library is a great choice for this optimization task, as it provides better performance characteristics than std::unordered_map for this use case.
1234-1237: Nice API improvement changing return type to ScheduleTaskMetadataThis function signature change aligns well with the PR objectives of encapsulating all necessary task metadata for the scheduler, reducing the need for additional storage accesses.
1244-1247: Optimized SQL query selecting only required fieldsGood job on modifying the SQL query to only select the fields needed for constructing the ScheduleTaskMetadata objects (id, func_name, job_id) rather than fetching unnecessary fields.
1254-1255: Effective use of hash maps for batchingUsing hash maps to organize the data by task ID and job ID is a good approach for preparing the batch queries to reduce the number of database accesses.
1268-1300: Excellent batch processing for job metadataGood implementation of batched queries for retrieving job metadata. This reduces the number of database round-trips compared to fetching metadata individually for each task.
1302-1328: Efficient batch processing for data localityThe implementation efficiently retrieves data locality information for multiple tasks in a single batch, which aligns well with the PR objective of reducing the number of queries.
1627-1630: Consistent API improvement for get_task_timeoutGood consistency in applying the same ScheduleTaskMetadata pattern to both get_ready_tasks and get_task_timeout methods.
1 participant
Description
MySQL storage currently executes 2 queries for each task in
get_ready_tasks. Scheduler needs 1 more query for each task to get job id, and 1 more query for each job to get job metadata. If tasks have data, then 1 more query is executed for each data to get locality. Same storage accesses are needed forget_task_timeout.This pr introduces
ScheduleTaskMetadatato encapsulate all metadata needed by scheduler.get_ready_tasksandget_task_timeoutnow returns vector ofScheduleTaskMetadataso that scheduler does not need extra storage access.MySQL storage also batches queries in
get_ready_tasksandget_task_timeto reduce number of queries and runtime.Checklist
breaking change.
Validation performed
Summary by CodeRabbit