feat(mito): Support background task priority #4804
Conversation
- Introduced `Job` struct to encapsulate job type and task. - Updated `CompactionScheduler`, `RegionFlushTask`, and `FilePurger` to use `Job` struct. - Added `SCHEDULER_TASK_ELAPSED` metric to track task execution time.
…EDULER_TASK_ELAPSED metric in LocalScheduler
• Added rate_limit dependency to Cargo.lock and Cargo.toml. • Introduced SendToChannel error variant in error.rs. • Implemented priority-based scheduling in schedule.rs with a new priority.rs module. • Added tests for high and low priority message handling, including deadline-based scheduling.
- Introduced `try_send_high` and `try_send_low` methods in `Sender` for priority-based job sending. - Implemented `Clone` trait for `Receiver` to support cloning of receivers. - Updated `LocalScheduler` to use `priority::Sender` and `priority::unbounded` for job scheduling.
Add Priority handling to Job scheduling and update channels in priority module - Implement Priority enum with High and Low variants for job scheduling - Update Job struct to include priority field and modify constructors - Modify LocalScheduler to handle job priority when scheduling - Adjust Sender and Receiver in priority module to support priority with deadlines - Add new tests for multi-threaded receiving with priority handling - Update compaction and flush tasks to use new Priority in Job creation
…rgo.lock and common/wal/Cargo.toml
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #4804 +/- ##
==========================================
- Coverage 84.48% 84.18% -0.31%
==========================================
Files 1120 1126 +6
Lines 204412 205208 +796
==========================================
+ Hits 172707 172763 +56
- Misses 31705 32445 +740 |
|
@coderabbitai full review |
✅ Actions performedReview triggered.
|
WalkthroughThe changes primarily focus on enhancing the job scheduling system within the codebase. Key modifications include the introduction of a Changes
Poem
Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media? 🪧 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
Documentation and Community
|
✅ Actions performedFull review triggered. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 8
🧹 Outside diff range and nitpick comments (7)
src/mito2/src/engine/compaction_test.rs (1)
32-32: Summary: Minimal changes to support new job scheduling system.The modifications in this file are part of the larger refactoring to support task prioritization, as outlined in the PR objectives. The changes are minimal and focused, updating only the import statement and the scheduling mechanism in the
test_readonly_during_compactionfunction.While these changes align well with the PR goals and don't alter the test's logic, it's important to consider the following:
- Ensure that all tests in this file still pass after these changes.
- Verify that the timing and behavior of the tests, especially
test_readonly_during_compaction, remain consistent with the previous implementation.- Consider adding new tests specifically for the priority scheduling system, if not already done in other files.
- Update the documentation to reflect the new usage of
Jobin scheduling tasks, particularly noting any differences between production and test environments.These changes seem to be part of a broader architectural shift towards a more flexible job scheduling system. As you continue this refactoring:
- Ensure consistent usage of the new
Jobstruct across all relevant parts of the codebase.- Consider adding integration tests that specifically verify the priority scheduling behavior in various scenarios.
- Update any relevant documentation or comments to explain the new job scheduling system and its impact on the engine's behavior.
Also applies to: 319-321
src/mito2/src/metrics.rs (1)
53-60: LGTM! Consider adding a brief comment for clarity.The new
SCHEDULER_TASK_ELAPSEDmetric is well-implemented and consistent with the existing metrics. It will provide valuable insights into the performance of different scheduler tasks.Consider adding a brief comment above the metric declaration to explain its purpose and the significance of the
TYPE_LABEL, similar to comments for other metrics in this file. For example:/// Histogram of elapsed time for different types of scheduler tasks. pub static ref SCHEDULER_TASK_ELAPSED: HistogramVec = register_histogram_vec!( // ... (rest of the code remains the same)src/mito2/src/schedule/scheduler.rs (3)
40-47: Consider renamingr#typefield to avoid using a raw identifierUsing
r#typeas a field name requires the use of a raw identifier becausetypeis a reserved keyword in Rust. To improve readability and adhere to Rust naming conventions, consider renaming this field tojob_typeortask_type.Apply this diff to rename the field:
pub struct Job { - r#type: &'static str, + job_type: &'static str, priority: Priority, task: Pin<Box<dyn Future<Output = ()> + Send>>, }Remember to update all usages of
r#typein the codebase accordingly.
162-171: Enhance error handling with more informative messagesCurrently, errors from
try_send_highandtry_send_loware mapped toInvalidSenderSnafuwithout additional context. Consider including details about the job's priority or type in the error to aid in debugging when a job fails to be scheduled.For example, you could modify the error handling as follows:
match job.priority { Priority::High => sender .try_send_high(job) - .map_err(|_| InvalidSenderSnafu {}.build()), + .map_err(|_| InvalidSenderSnafu { + description: format!("Failed to schedule high-priority job: {}", job.job_type), + } + .build()), Priority::Low(deadline) => sender .try_send_low(job, deadline) - .map_err(|_| InvalidSenderSnafu {}.build()), + .map_err(|_| InvalidSenderSnafu { + description: format!("Failed to schedule low-priority job: {}", job.job_type), + } + .build()), }Make sure to update the
InvalidSenderSnafuerror variant to accept adescriptionfield.
Line range hint
227-300: Add test cases to verify priority scheduling logicThe current tests do not cover scenarios where both high and low priority jobs are scheduled together. To ensure that the scheduler correctly prioritizes high-priority jobs and respects deadlines for low-priority jobs, consider adding test cases that:
- Schedule a mix of high and low priority jobs.
- Verify that high-priority jobs execute before low-priority ones.
- Test the behavior of low-priority jobs with deadlines.
For example, you could add a test like this:
#[tokio::test] async fn test_priority_scheduling() { let scheduler = LocalScheduler::new(1); let execution_order = Arc::new(Mutex::new(Vec::new())); // Low priority job without deadline let order_clone = execution_order.clone(); scheduler.schedule(Job::new( "low_no_deadline", Priority::Low(None), Box::pin(async move { order_clone.lock().await.push("low_no_deadline"); }), )).unwrap(); // High priority job let order_clone = execution_order.clone(); scheduler.schedule(Job::new( "high", Priority::High, Box::pin(async move { order_clone.lock().await.push("high"); }), )).unwrap(); // Low priority job with deadline in the past let order_clone = execution_order.clone(); scheduler.schedule(Job::new( "low_with_past_deadline", Priority::Low(Some(Duration::from_secs(0))), Box::pin(async move { order_clone.lock().await.push("low_with_past_deadline"); }), )).unwrap(); scheduler.stop(true).await.unwrap(); let order = execution_order.lock().await.clone(); assert_eq!( order, vec!["high", "low_with_past_deadline", "low_no_deadline"], "Jobs did not execute in expected order" ); }This test schedules three jobs and checks that they are executed in the correct order according to their priorities and deadlines.
src/mito2/src/schedule/priority.rs (1)
114-114: Correct grammatical error in documentation commentIn the documentation comment for the
into_streammethod, there's a minor grammatical error. The word "a" should be replaced with "an" before a vowel sound.-/// Converts the receiver into a async stream. +/// Converts the receiver into an async stream.src/mito2/src/flush.rs (1)
263-269: Consider adding unit tests for priority schedulingTo verify that flush jobs are correctly prioritized over other tasks, consider adding unit tests that simulate scheduling of different job types with various priorities.
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
📒 Files selected for processing (9)
- src/mito2/src/compaction.rs (2 hunks)
- src/mito2/src/engine/compaction_test.rs (2 hunks)
- src/mito2/src/error.rs (2 hunks)
- src/mito2/src/flush.rs (2 hunks)
- src/mito2/src/metrics.rs (1 hunks)
- src/mito2/src/schedule.rs (1 hunks)
- src/mito2/src/schedule/priority.rs (1 hunks)
- src/mito2/src/schedule/scheduler.rs (8 hunks)
- src/mito2/src/sst/file_purger.rs (2 hunks)
✅ Files skipped from review due to trivial changes (1)
- src/mito2/src/schedule.rs
🧰 Additional context used
🔇 Additional comments (14)
src/mito2/src/engine/compaction_test.rs (2)
32-32: New import added forJobstruct.The addition of
use crate::schedule::scheduler::Job;suggests that theJobstruct is now being used in this test file. This change is consistent with the modification in thetest_readonly_during_compactionfunction.
319-321: Approved: Updated scheduling to useJobstruct.The modification to use
Job::new_testaligns with the new priority-based scheduling system mentioned in the PR objectives. This change doesn't alter the test logic but updates how the task is scheduled.To ensure the new job scheduling system doesn't impact the test's behavior, please verify:
- The test still passes and behaves as expected.
- The timing and execution order of tasks remain consistent with the previous implementation.
You may want to run this test multiple times to check for any potential race conditions or timing issues introduced by the new scheduling system.
src/mito2/src/metrics.rs (1)
Line range hint
1-60: The new metric aligns well with the PR objectives.The addition of the
SCHEDULER_TASK_ELAPSEDmetric is a valuable enhancement to the metrics system. It will provide insights into the performance of scheduler tasks, which is crucial for the new priority system for background tasks introduced in this PR. This metric will help monitor and optimize the handling of high-priority tasks (like flush jobs) and low-priority tasks (like compaction and file purge jobs).src/mito2/src/error.rs (2)
874-879: LGTM: New error variant added correctly.The new
SendToChannelvariant is well-implemented and consistent with the existing error handling pattern in this enum. It properly captures the error message and location information.
1012-1012: LGTM: ErrorExt trait implementation updated correctly.The
SendToChannelvariant is properly mapped toStatusCode::Internal, which is consistent with similar error types in this implementation. The placement in the match expression follows the existing pattern.src/mito2/src/sst/file_purger.rs (5)
23-23: Importing scheduler components for priority-based task schedulingThe addition of
Job,Priority, andSchedulerRefimports is appropriate to support the new priority-based scheduling mechanism for background tasks.
82-85: Scheduling file purge operations as low-priority jobsWrapping the file purge logic within a
Jobscheduled withPriority::Low(None)correctly aligns with the PR objective to process file purge tasks as low-priority background jobs.
86-94: Robust error handling during SST file deletionThe code properly handles errors when deleting the SST file by logging an error message on failure and an informational message on success, ensuring visibility into the operation's outcome.
96-98: Ensuring cache manager and write cache availabilityThe check for the presence of a cache manager and write cache before attempting cache removal operations prevents potential null pointer dereferences, enhancing the code's safety.
105-106: Verify the use ofFileType::Puffinfor inverted index cache removalEnsure that
FileType::Puffinis the correct file type for inverted index entries in the cache. If the file type for inverted indexes has changed elsewhere in the codebase, this usage may need to be updated.To confirm the correct file type, please run the following script:
✅ Verification successful
FileType::Puffin correctly used for inverted index cache removal
The usage of
FileType::Puffininsrc/mito2/src/sst/file_purger.rsfor inverted index cache removal is correct.🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Description: Verify that `FileType::Puffin` is used for inverted index cache entries. # Search for cache removals of inverted index files and check the file type. rg --type rust 'write_cache' -A 5 | rg 'inverted_index' -A 5 | rg 'FileType::Puffin'Length of output: 163
src/mito2/src/schedule/scheduler.rs (1)
32-38: Definition ofPriorityenum is appropriateThe
Priorityenum is well-defined, allowing for high-priority jobs and low-priority jobs with optional deadlines. This enhances the scheduler's capability to handle tasks based on their urgency.src/mito2/src/compaction.rs (1)
66-66: Updated imports to include Job and PriorityThe import statement now includes
JobandPriority, which are necessary for the new job scheduling functionality.src/mito2/src/flush.rs (2)
44-44: ImportingPriorityenum is appropriateThe addition of the
Priorityimport allows setting job priorities in the scheduler.
263-269: Flush jobs are now scheduled with high prioritySetting the priority of flush jobs to
Priority::Highensures they are prioritized appropriately in the job queue.
There was a problem hiding this comment.
Actionable comments posted: 6
🧹 Outside diff range and nitpick comments (4)
src/mito2/src/metrics.rs (1)
53-60: LGTM! Consider adding a brief comment for clarity.The new
SCHEDULER_TASK_ELAPSEDmetric is well-integrated and consistent with existing metrics. It follows the established naming conventions and uses appropriate types and labels.Consider adding a brief comment above the metric to explain its purpose, similar to other metrics in this file. For example:
/// Elapsed time of tasks in scheduler, categorized by task type. pub static ref SCHEDULER_TASK_ELAPSED: HistogramVec = register_histogram_vec!( // ... (rest of the code remains the same)src/mito2/src/sst/file_purger.rs (1)
82-84: Consider specifying a deadline for the purge job's prioritySince low-priority jobs can have an optional deadline to ensure they are eventually executed, consider adding a deadline to
Priority::Lowfor the file purge tasks if timely execution is desired.src/mito2/src/schedule/priority.rs (2)
40-44: Add documentation comments toSender<T>structTo improve code readability and maintainability, consider adding Rust doc comments to the
Sender<T>struct, explaining its purpose and usage.
96-99: Add documentation comments toReceiver<T>structFor better documentation, add doc comments to the
Receiver<T>struct to explain its role and how it should be used.
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
📒 Files selected for processing (9)
- src/mito2/src/compaction.rs (2 hunks)
- src/mito2/src/engine/compaction_test.rs (2 hunks)
- src/mito2/src/error.rs (2 hunks)
- src/mito2/src/flush.rs (2 hunks)
- src/mito2/src/metrics.rs (1 hunks)
- src/mito2/src/schedule.rs (1 hunks)
- src/mito2/src/schedule/priority.rs (1 hunks)
- src/mito2/src/schedule/scheduler.rs (8 hunks)
- src/mito2/src/sst/file_purger.rs (2 hunks)
🧰 Additional context used
🔇 Additional comments (12)
src/mito2/src/schedule.rs (1)
15-16: Newprioritymodule added with#[allow(dead_code)]attributeThe addition of the
prioritymodule aligns with the PR objective of introducing a priority system for background tasks. However, the#[allow(dead_code)]attribute suggests that some code within this module might currently be unused.Consider the following points:
- Ensure that the
prioritymodule is properly integrated with the existing scheduling system as development progresses.- The
#[allow(dead_code)]attribute is useful during development but should be removed once the module is fully implemented and integrated.- As the PR objectives mention high and low priority tasks, verify that this module implements the described functionality for flush, compaction, and file purge jobs.
To verify the contents and integration of the new
prioritymodule, please run the following script:This script will help us understand the implementation details of the
prioritymodule and how it's being used in the codebase.✅ Verification successful
prioritymodule is actively used and integratedThe
prioritymodule is implemented and utilized withinscheduler.rs, indicating that it serves its intended purpose. The#[allow(dead_code)]attribute may no longer be necessary and can be removed to ensure all code is actively maintained and utilized.
- Action Items:
- Consider removing the
#[allow(dead_code)]attribute from theprioritymodule.- Ensure that all functionalities within the
prioritymodule are actively used and tested.🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Description: Verify the contents of the priority module and its integration # Check the contents of the priority module echo "Contents of the priority module:" cat src/mito2/src/schedule/priority.rs # Check for references to the priority module in other files echo "References to the priority module in other files:" rg "use (crate|super)::schedule::priority" src/mito2/src/Length of output: 13781
src/mito2/src/engine/compaction_test.rs (2)
32-32: LGTM: Import statement added correctly.The addition of the
Jobimport fromcrate::schedule::scheduleris appropriate and aligns with its usage in thetest_readonly_during_compactionfunction.
319-321: LGTM: Job scheduling updated, but clarification needed.The change to use
Job::new_testis consistent with the new import and likely part of a larger refactoring. However, I have a few questions:
- Can you provide more information about the
Job::new_testmethod? Is it specifically designed for testing scenarios?- Does this change affect how the job is prioritized or executed compared to the previous implementation?
- Are there any potential impacts on the test's behavior or assertions that we should be aware of?
To ensure this change doesn't introduce any unintended side effects, could you run the following verification script?
This will help us understand if this change is part of a broader refactoring across test files and if there are any inconsistencies in how job scheduling is handled in tests.
✅ Verification successful
Verified: Job scheduling change is localized.
The usage of
Job::new_testis confined tocompaction_test.rs, indicating it's a specific change for this test case. There are no other instances ofJob::new_testor modifications toscheduleacross other test files, suggesting no broader impact on job scheduling or test behavior.🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Description: Verify the usage of Job::new_test in test files # Test: Search for other occurrences of Job::new_test in test files rg --type rust -g '*_test.rs' 'Job::new_test' # Test: Check if there are any other changes related to job scheduling in test files rg --type rust -g '*_test.rs' '\bschedule\b'Length of output: 357
src/mito2/src/error.rs (2)
874-879: LGTM: NewSendToChannelvariant added correctly.The new
SendToChannelvariant has been added to theErrorenum with appropriate fields and attributes. It follows the existing pattern and provides a clear error message for channel send failures.
1012-1012: LGTM:ErrorExttrait implementation updated correctly.The
SendToChannelvariant has been properly added to thestatus_codemethod in theErrorExttrait implementation. Associating it withStatusCode::Internalis appropriate for this type of error.src/mito2/src/flush.rs (3)
44-44: New import added for Priority enumThe addition of
Priorityto the import statement is consistent with the changes made in theinto_flush_jobmethod. This import allows the use of thePriorityenum in the file.
263-269: Approved: High priority assigned to flush jobsThe addition of
Priority::Highto theJob::newcall in theinto_flush_jobmethod is a positive change. This modification aligns with the PR objective of introducing a priority system for background tasks. By assigning high priority to flush jobs, the system ensures that these critical operations are processed before lower-priority tasks like compaction, potentially improving overall system performance and responsiveness.
Line range hint
1-1006: Summary: Successful implementation of priority for flush jobsThe changes made to this file successfully implement the priority system for flush jobs as part of the larger goal of supporting background task priority. The modifications are minimal, focused, and consistent with the PR objectives. The high priority assigned to flush jobs should help improve system performance by ensuring these critical tasks are processed before lower-priority operations.
No further changes or improvements are necessary based on the current implementation.
src/mito2/src/sst/file_purger.rs (3)
23-23: Imports updated correctlyThe new import statement correctly includes
Job,Priority, andSchedulerRef, which are necessary for the job scheduling enhancements.
82-119: Encapsulation of purge logic into a scheduled job enhances maintainabilityBy wrapping the file purge operations within a
Joband scheduling it, the code improves the organization, readability, and maintainability of asynchronous tasks.
86-94: Appropriate error handling and logging for SST deletionThe code correctly handles errors during the SST file deletion by logging them, ensuring that any issues are visible for debugging and monitoring purposes.
src/mito2/src/compaction.rs (1)
66-66: LGTM!The import statement correctly includes
JobandPriorityfrom theschedulermodule.
…ods to use map_send_error function • Implement map_send_error function to standardize error mapping in Sender • Replace inline error mapping with map_send_error calls in Sender methods • Change Receiver's into_stream to return impl Stream instead of a boxed trait object • Add missing assert_eq! in tests to verify expected behavior
…rn an unboxed stream and update usage to box the stream where necessary.
| loop { | ||
| select! { | ||
| biased; | ||
| Some(_) = maybe_timeout(&mut timer) => { |
There was a problem hiding this comment.
Maybe we can simplify the channel to always execute a pending low priority task by a fixed interval or after executing N high priority task.
We can remove the deadline argument from the prority.
There was a problem hiding this comment.
Maybe we can simplify the channel to always execute a pending low priority task by a fixed interval or after executing N high priority task.
We can remove the deadline argument from the prority.
The initial implementation was like this, but the task waiting time will be indefinite.
There was a problem hiding this comment.
but the task waiting time will be indefinite.
If we send low-pri jobs in this order: None, deadline 5m, then the job with None can still block the job with deadline 5m longer than 5 minutes.
There was a problem hiding this comment.
but the task waiting time will be indefinite.
If we send low-pri jobs in this order:
None,deadline 5m, then the job withNonecan still block the job withdeadline 5mlonger than 5 minutes.
Yes, the enqueued jobs in either high or low priority channels are still guaranteed FIFO, deadline only affects low-priority jobs competing with high-priority jobs.
…al and removed unused priority functions
…refactor priority channel send methods • Implement SCHEDULER_PENDING_JOBS metric to track pending jobs in the scheduler. • Remove bounded channel creation in priority.rs and add send_high and send_low async methods for testing. • Modify scheduler.rs to increment and decrement SCHEDULER_PENDING_JOBS metric upon job submission and completion. • Replace direct calls to send_high and send_low with try_send_high and try_send_low in tests.
Improve handling of expired low-priority items and add test for low-priority overloading
| loop { | ||
| select! { | ||
| biased; | ||
| Some(_) = maybe_timeout(&mut timer) => { |
There was a problem hiding this comment.
but the task waiting time will be indefinite.
If we send low-pri jobs in this order: None, deadline 5m, then the job with None can still block the job with deadline 5m longer than 5 minutes.
…r and remove InvalidSender error • Remove InvalidSender error variant from Error enum. • Refactor Sender methods to use a single try_send method with a Priority parameter instead of separate try_send_high and try_send_low methods. • Update tests to use the new try_send method with Priority. • Add Clone and Copy traits to Priority enum.
Introduce COMPACTION_TASK_DEADLINE constant for compaction task duration
… compactions to High • Implement manual compaction request handling with high priority scheduling. • Extend CompactionScheduler and related methods to accept a manual boolean flag to determine compaction priority. • Update unit tests to pass the manual flag as false for automated compaction scenarios. • Add Debug trait to Priority enum for improved logging.
I hereby agree to the terms of the GreptimeDB CLA.
Refer to a related PR or issue link (optional)
What's changed and what's your intention?
This PR adds priorities to background tasks like flush/compaction/file purge.
In the current local scheduler implementation, the scheduler spawns coroutines to poll background tasks from a FIFO queue, waiting for each job to finish before polling the next one. During heavy insert workloads, flush jobs can be blocked by compaction jobs unless we assign priorities to background tasks.
This PR introduces an asynchronous priority channel with two levels: High and Low. It includes an optional deadline for low-priority jobs. The channel first attempts to poll jobs from the high-priority queue, switching to the low-priority queue only if the high-priority queue is empty. If a job at the head of the low-priority queue reaches its deadline, it will be polled immediately. Although jobs in low-priority channel have deadline, all jobs enqueued still follow the FIFO convention.
Currently, flush jobs are given high priority, while compaction and file purge jobs are assigned low priority. The deadline for compaction tasks are set to 60 seconds.
Checklist
Summary by CodeRabbit
New Features
Jobstruct for enhanced job scheduling, allowing for clearer organization of tasks.Priorityenum to manage job priorities effectively.Bug Fixes
SendToChannelvariant for better context on message sending issues.Documentation
Tests