diff --git a/src/pool/builder.rs b/src/pool/builder.rs index 29e3be3..09752a0 100644 --- a/src/pool/builder.rs +++ b/src/pool/builder.rs @@ -230,13 +230,13 @@ impl Builder { /// Freezes the configurations and returns the task scheduler and /// a builder to for lazy spawning threads. /// - /// `queue_builder` is a closure that creates a task queue. It accepts the - /// number of local queues and returns the task injector and local queues. + /// `queue_type` selects a built-in queue, or a custom queue that implements + /// [`crate::queue::TaskQueue`]. /// /// In some cases, especially building up a large application, a task /// scheduler is required before spawning new threads. You can use this /// to separate the construction and starting. - pub fn freeze_with_queue(&self, queue_type: QueueType) -> (Remote, LazyBuilder) + pub fn freeze_with_queue(&self, queue_type: QueueType) -> (Remote, LazyBuilder) where T: TaskCell + Send, { @@ -303,13 +303,26 @@ impl Builder { self.build_with_queue_and_runner(QueueType::Priority(queue_builder), runner_builder) } + /// Spawn a custom future pool. + /// + /// It setups the pool with the given custom queue. + pub fn build_custom_future_pool( + &self, + queue: Arc>, + ) -> ThreadPool { + let fb = CloneRunnerBuilder(future::Runner::default()); + let queue_builder = queue::CustomBuilder::new(queue::CustomConfig::default(), queue); + let runner_builder = queue_builder.runner_builder(fb); + self.build_with_queue_and_runner(queue_builder.into(), runner_builder) + } + /// Spawns the thread pool immediately. /// - /// `queue_builder` is a closure that creates a task queue. It accepts the - /// number of local queues and returns the task injector and local queues. + /// `queue_type` selects a built-in queue, or a custom queue that implements + /// [`crate::queue::TaskQueue`]. pub fn build_with_queue_and_runner( &self, - queue_type: QueueType, + queue_type: QueueType, runner_builder: B, ) -> ThreadPool where diff --git a/src/pool/spawn.rs b/src/pool/spawn.rs index d648e63..2512013 100644 --- a/src/pool/spawn.rs +++ b/src/pool/spawn.rs @@ -5,13 +5,14 @@ //! tasks waiting to be handled. use crate::pool::SchedConfig; -use crate::queue::{Extras, LocalQueue, Pop, TaskCell, TaskInjector, WithExtras}; +use crate::queue::{Extras, LocalQueue, Pop, PopResult, TaskCell, TaskInjector, WithExtras}; use fail::fail_point; use parking_lot_core::{FilterOp, ParkResult, ParkToken, UnparkToken}; use std::sync::{ atomic::{AtomicUsize, Ordering}, Arc, Weak, }; +use std::time::Instant; /// An usize is used to trace the threads that are working actively. /// To save additional memory and atomic operation, the number and @@ -62,6 +63,10 @@ impl QueueCore { return; } + self.wake_one_core_worker(source); + } + + fn wake_one_core_worker(&self, source: usize) { let addr = self as *const QueueCore as usize; let mut unparked_once = false; @@ -309,47 +314,122 @@ impl Local { &self.core } - pub(crate) fn pop(&mut self) -> Option> { + pub(crate) fn pop(&mut self) -> PopResult { self.local_queue.pop() } + pub(crate) fn is_scaled_down_worker(&self) -> bool { + self.id > self.core.config.core_thread_count.load(Ordering::SeqCst) + } + + pub(crate) fn drain(&mut self) { + self.local_queue.drain(); + } + /// Pops a task from the queue. /// /// If there are no tasks at the moment, it will go to sleep until woken - /// up by other threads. - pub(crate) fn pop_or_sleep(&mut self) -> Option> { + /// up by other threads or until the next known pending-task retry time. + pub(crate) fn pop_or_sleep(&mut self, initial_retry_at: Option) -> Option> { let address = &*self.core as *const QueueCore as usize; - let mut task = None; let id = self.id; - - let res = unsafe { - parking_lot_core::park( - address, - || { - if !self.core.mark_sleep() { - return false; - } - // If this thread is above core_thread_count, go to sleep - // without popping so scaled-down threads don't keep working. - if id > self.core.config.core_thread_count.load(Ordering::SeqCst) { - return true; - } - task = self.local_queue.pop(); - task.is_none() - }, - || {}, - |_, _| {}, - ParkToken(id), - None, - ) - }; - match res { - ParkResult::Unparked(_) | ParkResult::Invalid => { + let mut timeout = initial_retry_at; + + while !self.core.is_shutdown() { + let mut task = None; + let mut next_timeout = None; + let mut marked_sleep = false; + + fail_point!("worker-pop-or-sleep-before-park"); + let park_result = unsafe { + parking_lot_core::park( + address, + || { + // Returning false from validate aborts this park and + // makes parking_lot_core return ParkResult::Invalid. + // Use it when the decision to sleep needs to be + // changed after rechecking the queue. + if !self.core.mark_sleep() { + return false; + } + marked_sleep = true; + // If this thread is above core_thread_count, go to sleep + // without popping so scaled-down threads don't keep working. + if id > self.core.config.core_thread_count.load(Ordering::SeqCst) { + return true; + } + fail_point!("worker-pop-or-sleep-before-validate-pop"); + match self.local_queue.pop() { + PopResult::Ready(t) => { + task = Some(t); + false + } + PopResult::Pending { retry_at } => match timeout { + Some(timeout) if retry_at >= timeout => true, + _ => { + // The current park call cannot change its + // timeout after validate has started. Abort + // this park so the outer loop can retry + // with a timeout that wakes no later than + // retry_at. + next_timeout = Some(retry_at); + false + } + }, + PopResult::Empty => true, + } + }, + || { + fail_point!("worker-pop-or-sleep-before-sleep"); + }, + |_, _| {}, + ParkToken(id), + timeout, + ) + }; + + // mark_sleep decreases the active worker count before the park + // decision is finalized. Whether the thread actually slept, + // timed out, was unparked, or aborted the park from validate, the + // worker is running again after park returns, so restore the count. + if marked_sleep { self.core.mark_woken(); - task } - ParkResult::TimedOut => unreachable!(), + + if self.core.is_shutdown() { + return None; + } + + // If validate found a ready task, the park was aborted before the + // thread actually slept. Return it immediately. + if task.is_some() { + return task; + } + + // If validate found pending work that needs an earlier retry time, + // retry park with the updated timeout. The current park call cannot + // change its timeout after validate has started. + if next_timeout.is_some() { + timeout = next_timeout; + continue; + } + + if matches!(park_result, ParkResult::TimedOut) && self.is_scaled_down_worker() { + // This worker carried the pending retry timeout, but it is no + // longer allowed to pop tasks after scale-down. Wake a core + // worker to re-check the queue and install its own timeout (or + // run ready work), then park this worker without a deadline. + self.core.wake_one_core_worker(id); + timeout = None; + continue; + } + + // Otherwise the thread was either unparked, timed out, or the park + // was aborted without a task, for example because shutdown made + // mark_sleep fail. Let the worker loop re-check the pool state. + return None; } + None } /// Returns whether there are preemptive tasks to run. @@ -366,7 +446,7 @@ impl Local { /// This is only for tests purpose so that a thread pool doesn't have to be /// spawned to test a Runner. pub fn build_spawn( - queue_type: impl Into, + queue_type: impl Into>, config: SchedConfig, ) -> (Remote, Vec>) where diff --git a/src/pool/worker.rs b/src/pool/worker.rs index 0af3ef8..664d68e 100644 --- a/src/pool/worker.rs +++ b/src/pool/worker.rs @@ -1,7 +1,7 @@ // Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0. use crate::pool::{Local, Runner}; -use crate::queue::{Pop, TaskCell}; +use crate::queue::{Pop, PopResult, TaskCell}; use parking_lot_core::SpinWait; pub(crate) struct WorkerThread { @@ -24,16 +24,18 @@ where fn pop(&mut self) -> Option> { // Wait some time before going to sleep, which is more expensive. let mut spin = SpinWait::new(); - loop { - if let Some(t) = self.local.pop() { - return Some(t); - } + let initial_retry_at = loop { + let retry_at = match self.local.pop() { + PopResult::Ready(task) => return Some(task), + PopResult::Pending { retry_at } => Some(retry_at), + PopResult::Empty => None, + }; if !spin.spin() { - break; + break retry_at; } - } + }; self.runner.pause(&mut self.local); - let t = self.local.pop_or_sleep(); + let t = self.local.pop_or_sleep(initial_retry_at); self.runner.resume(&mut self.local); t } @@ -50,7 +52,7 @@ where self.runner.end(&mut self.local); // Drain all futures in the queue - while self.local.pop().is_some() {} + self.local.drain(); } } @@ -59,10 +61,13 @@ mod tests { use super::*; use crate::pool::spawn::*; use crate::pool::SchedConfig; - use crate::queue::QueueType; + use crate::queue::{CustomBuilder, CustomConfig, Extras, QueueType, TaskQueue}; use crate::task::callback; - use std::sync::atomic::AtomicUsize; + use std::collections::VecDeque; + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::*; + use std::thread; + use std::thread::JoinHandle; use std::time::*; #[derive(Default, PartialEq, Debug)] @@ -116,6 +121,505 @@ mod tests { } } + struct TestTask { + id: usize, + extras: Extras, + } + + impl TestTask { + fn new(id: usize) -> TestTask { + TestTask { + id, + extras: Extras::multilevel_default(), + } + } + } + + impl TaskCell for TestTask { + fn mut_extras(&mut self) -> &mut Extras { + &mut self.extras + } + } + + struct ScriptedQueue { + scripted_results: Mutex>>, + pushed_tasks: Mutex>, + push_tx: Mutex>>, + } + + impl ScriptedQueue { + fn new(scripted_results: Vec>) -> ScriptedQueue { + ScriptedQueue { + scripted_results: Mutex::new(scripted_results.into()), + pushed_tasks: Mutex::new(VecDeque::new()), + push_tx: Mutex::new(None), + } + } + + fn with_push_signal( + scripted_results: Vec>, + push_tx: mpsc::Sender<()>, + ) -> ScriptedQueue { + ScriptedQueue { + scripted_results: Mutex::new(scripted_results.into()), + pushed_tasks: Mutex::new(VecDeque::new()), + push_tx: Mutex::new(Some(push_tx)), + } + } + } + + impl ScriptedQueue { + fn ready(id: usize) -> PopResult { + PopResult::Ready(Pop { + task_cell: TestTask::new(id), + schedule_time: Instant::now(), + from_local: false, + }) + } + } + + impl TaskQueue for ScriptedQueue { + fn push(&self, task_cell: T) { + self.pushed_tasks.lock().unwrap().push_back(task_cell); + if let Some(push_tx) = self.push_tx.lock().unwrap().as_ref() { + let _ = push_tx.send(()); + } + } + + fn pop(&self) -> PopResult { + if let Some(result) = self.scripted_results.lock().unwrap().pop_front() { + return result; + } + self.pushed_tasks + .lock() + .unwrap() + .pop_front() + .map(|task_cell| Pop { + task_cell, + schedule_time: Instant::now(), + from_local: false, + }) + .into() + } + + fn drain(&self) { + self.scripted_results.lock().unwrap().clear(); + self.pushed_tasks.lock().unwrap().clear(); + } + + fn has_ready_task(&self) -> bool { + !self.scripted_results.lock().unwrap().is_empty() + || !self.pushed_tasks.lock().unwrap().is_empty() + } + } + + fn one_thread_config() -> SchedConfig { + SchedConfig { + min_thread_count: 1, + max_thread_count: 1, + core_thread_count: AtomicUsize::new(1), + ..Default::default() + } + } + + fn two_thread_config() -> SchedConfig { + SchedConfig { + min_thread_count: 1, + max_thread_count: 2, + core_thread_count: AtomicUsize::new(2), + ..Default::default() + } + } + + fn build_scripted_local(queue: Arc>) -> Local { + let queue_builder = CustomBuilder::new(CustomConfig::default(), queue); + let (_, mut locals) = build_spawn(queue_builder, one_thread_config()); + locals.remove(0) + } + + fn assert_next_ready_task(local: &mut Local, id: usize) { + assert_eq!(local.pop().unwrap_ready().task_cell.id, id); + } + + fn callback_task( + task: impl FnOnce(&mut callback::Handle<'_>) + Send + 'static, + ) -> callback::TaskCell { + callback::TaskCell { + task: callback::Task::new_once(task), + extras: Extras::multilevel_default(), + } + } + + const WORKER_SPIN_POP_COUNT: usize = 11; + const DELAYED_TASK_DELAY: Duration = Duration::from_millis(50); + const MAX_DELAYED_TASK_LAG: Duration = Duration::from_secs(2); + const LATER_RETRY_OFFSET: Duration = Duration::from_secs(10); + + #[derive(Clone, Copy)] + enum DelayedQueueScenario { + PendingDuringSpinAndValidate, + EmptyDuringSpin, + EarlierRetryInValidate, + } + + enum DeadlineScript { + Empty, + Pending(Instant), + PendingAfter(Duration), + } + + struct DeadlineTask { + task_cell: T, + ready_at: Instant, + } + + struct DeadlineQueueState { + delays: VecDeque, + ready_ats: Vec, + scripted_results: VecDeque, + tasks: VecDeque>, + } + + struct DeadlineQueue { + state: Mutex>, + } + + struct AlwaysPendingQueue { + retry_at: Instant, + } + + impl TaskQueue for AlwaysPendingQueue { + fn push(&self, _: callback::TaskCell) {} + + fn pop(&self) -> PopResult { + PopResult::Pending { + retry_at: self.retry_at, + } + } + + fn drain(&self) {} + + fn has_ready_task(&self) -> bool { + false + } + } + + impl DeadlineQueue { + fn new(delays: Vec) -> DeadlineQueue { + DeadlineQueue { + state: Mutex::new(DeadlineQueueState { + delays: delays.into(), + ready_ats: Vec::new(), + scripted_results: VecDeque::new(), + tasks: VecDeque::new(), + }), + } + } + + fn push_scripted_result(&self, result: DeadlineScript) { + self.state + .lock() + .unwrap() + .scripted_results + .push_back(result); + } + + fn ready_at(&self, index: usize) -> Instant { + self.state.lock().unwrap().ready_ats[index] + } + } + + impl TaskQueue for DeadlineQueue { + fn push(&self, task_cell: T) { + let mut state = self.state.lock().unwrap(); + let delay = state.delays.pop_front().unwrap(); + let ready_at = Instant::now() + delay; + state.ready_ats.push(ready_at); + state.tasks.push_back(DeadlineTask { + task_cell, + ready_at, + }); + } + + fn pop(&self) -> PopResult { + let mut state = self.state.lock().unwrap(); + if let Some(result) = state.scripted_results.pop_front() { + return match result { + DeadlineScript::Empty => PopResult::Empty, + DeadlineScript::Pending(retry_at) => PopResult::Pending { retry_at }, + DeadlineScript::PendingAfter(delay) => { + let retry_at = Instant::now() + delay; + let index = state + .tasks + .iter() + .enumerate() + .min_by_key(|(_, task)| task.ready_at) + .map(|(index, _)| index); + if let Some(index) = index { + state.tasks[index].ready_at = retry_at; + state.ready_ats[index] = retry_at; + } + PopResult::Pending { retry_at } + } + }; + } + + let (index, ready_at) = match state + .tasks + .iter() + .enumerate() + .min_by_key(|(_, task)| task.ready_at) + .map(|(index, task)| (index, task.ready_at)) + { + Some(task) => task, + None => return PopResult::Empty, + }; + + if Instant::now() < ready_at { + return PopResult::Pending { retry_at: ready_at }; + } + + let task = state.tasks.remove(index).unwrap(); + PopResult::Ready(Pop { + task_cell: task.task_cell, + schedule_time: task.ready_at, + from_local: false, + }) + } + + fn drain(&self) { + let mut state = self.state.lock().unwrap(); + state.scripted_results.clear(); + state.tasks.clear(); + } + + fn has_ready_task(&self) -> bool { + self.state + .lock() + .unwrap() + .tasks + .iter() + .any(|task| Instant::now() >= task.ready_at) + } + } + + fn ready_callback_pop(tx: mpsc::Sender) -> PopResult { + PopResult::Ready(Pop { + task_cell: callback_task(move |_: &mut callback::Handle<'_>| { + tx.send(Instant::now()).unwrap(); + }), + schedule_time: Instant::now(), + from_local: false, + }) + } + + fn check_scripted_worker_runs_task( + queue: Arc>, + done_rx: mpsc::Receiver, + ) -> Instant { + let (remote, _pause_rx, metrics, handle) = build_custom_worker(queue); + let executed_value = done_rx + .recv_timeout(MAX_DELAYED_TASK_LAG + Duration::from_secs(1)) + .unwrap(); + + { + let metrics = metrics.lock().unwrap(); + assert_eq!(metrics.start, 1); + assert_eq!(metrics.handle, 1); + assert!(metrics.pause >= 1); + assert!(metrics.resume >= 1); + } + + remote.stop(); + handle.join().unwrap(); + assert_eq!(metrics.lock().unwrap().end, 1); + + executed_value + } + + fn build_custom_worker( + queue: Arc>, + ) -> ( + Remote, + mpsc::Receiver<()>, + Arc>, + JoinHandle<()>, + ) { + let queue_builder = CustomBuilder::new(CustomConfig::default(), queue); + let (remote, mut locals) = build_spawn(queue_builder, one_thread_config()); + let (pause_rx, metrics, handle) = start_custom_worker(locals.remove(0)); + + (remote, pause_rx, metrics, handle) + } + + fn start_custom_worker( + local: Local, + ) -> (mpsc::Receiver<()>, Arc>, JoinHandle<()>) { + let (pause_tx, pause_rx) = mpsc::channel(); + let metrics = Arc::new(Mutex::new(Metrics::default())); + let runner = Runner { + runner: callback::Runner::default(), + metrics: metrics.clone(), + tx: pause_tx, + }; + let worker = WorkerThread::new(local, runner); + let handle = thread::spawn(move || worker.run()); + + (pause_rx, metrics, handle) + } + + fn check_worker_runs_ready_task_inserted_while_pending() { + let _lock = lock_failpoint_tests(); + let queue = Arc::new(DeadlineQueue::new(vec![LATER_RETRY_OFFSET, Duration::ZERO])); + let _guard = fail::FailScenario::setup(); + let (entered_rx, release_tx) = + configure_blocking_failpoint("worker-pop-or-sleep-before-sleep"); + + let (unexpected_tx, unexpected_rx) = mpsc::channel(); + queue.push(callback_task(move |_: &mut callback::Handle<'_>| { + unexpected_tx.send(()).unwrap(); + })); + let pending_retry_at = queue.ready_at(0); + let (remote, _pause_rx, metrics, handle) = build_custom_worker(queue.clone()); + entered_rx.recv_timeout(Duration::from_secs(1)).unwrap(); + + let (done_tx, done_rx) = mpsc::channel(); + let ready_inserted_at = Instant::now(); + remote.spawn(move |_: &mut callback::Handle<'_>| { + done_tx.send(Instant::now()).unwrap(); + }); + release_tx.send(()).unwrap(); + + let executed_at = done_rx.recv_timeout(MAX_DELAYED_TASK_LAG).unwrap(); + assert!(executed_at >= ready_inserted_at); + assert!(executed_at.duration_since(ready_inserted_at) <= MAX_DELAYED_TASK_LAG); + assert!(executed_at < pending_retry_at); + assert!(unexpected_rx.try_recv().is_err()); + { + let metrics = metrics.lock().unwrap(); + assert_eq!(metrics.start, 1); + assert_eq!(metrics.handle, 1); + } + + remote.stop(); + handle.join().unwrap(); + { + let metrics = metrics.lock().unwrap(); + assert_eq!(metrics.handle, 1); + assert_eq!(metrics.end, 1); + } + } + + fn configure_blocking_failpoint(name: &'static str) -> (mpsc::Receiver<()>, mpsc::Sender<()>) { + let (entered_tx, entered_rx) = mpsc::channel(); + let (release_tx, release_rx) = mpsc::channel(); + let entered_tx = Arc::new(Mutex::new(Some(entered_tx))); + let release_rx = Arc::new(Mutex::new(Some(release_rx))); + let fired = Arc::new(AtomicBool::new(false)); + fail::cfg_callback(name, move || { + if fired.swap(true, Ordering::SeqCst) { + return; + } + if let Ok(mut entered_tx) = entered_tx.lock() { + if let Some(entered_tx) = entered_tx.take() { + let _ = entered_tx.send(()); + } + } + if let Ok(mut release_rx) = release_rx.lock() { + if let Some(release_rx) = release_rx.take() { + let _ = release_rx.recv_timeout(Duration::from_secs(3)); + } + } + }) + .unwrap(); + + (entered_rx, release_tx) + } + + fn configure_counting_failpoint( + name: &'static str, + blocked_count: usize, + ) -> (mpsc::Receiver, mpsc::Sender<()>) { + let (entered_tx, entered_rx) = mpsc::channel(); + let (release_tx, release_rx) = mpsc::channel(); + let release_rx = Arc::new(Mutex::new(release_rx)); + let count = Arc::new(AtomicUsize::new(0)); + + fail::cfg_callback(name, move || { + let current = count.fetch_add(1, Ordering::SeqCst) + 1; + let _ = entered_tx.send(current); + if current <= blocked_count { + let _ = release_rx + .lock() + .unwrap() + .recv_timeout(Duration::from_secs(3)); + } + }) + .unwrap(); + + (entered_rx, release_tx) + } + + fn lock_failpoint_tests() -> std::sync::MutexGuard<'static, ()> { + static FAILPOINT_TEST_LOCK: Mutex<()> = Mutex::new(()); + FAILPOINT_TEST_LOCK + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + } + + fn spawn_task_and_wait_until_pushed( + remote: Remote, + push_rx: &mpsc::Receiver<()>, + done_tx: mpsc::Sender, + value: usize, + ) -> JoinHandle<()> { + let handle = thread::spawn(move || { + remote.spawn(move |_: &mut callback::Handle<'_>| { + done_tx.send(value).unwrap(); + }); + }); + push_rx.recv_timeout(Duration::from_secs(1)).unwrap(); + handle + } + + fn check_worker_wakes_when_task_is_inserted_at(failpoint: &'static str) { + let _lock = lock_failpoint_tests(); + let _guard = fail::FailScenario::setup(); + let (entered_rx, release_tx) = configure_blocking_failpoint(failpoint); + let (push_tx, push_rx) = mpsc::channel(); + let queue = Arc::new(ScriptedQueue::with_push_signal(Vec::new(), push_tx)); + let (remote, pause_rx, metrics, handle) = build_custom_worker(queue); + + // The first pause means the worker has already observed Empty during + // spin and is about to enter pop_or_sleep. + pause_rx.recv_timeout(Duration::from_secs(1)).unwrap(); + entered_rx.recv_timeout(Duration::from_secs(1)).unwrap(); + + let (done_tx, done_rx) = mpsc::channel(); + let spawn_handle = spawn_task_and_wait_until_pushed(remote.clone(), &push_rx, done_tx, 42); + release_tx.send(()).unwrap(); + + assert_eq!(done_rx.recv_timeout(Duration::from_secs(1)).unwrap(), 42); + { + let metrics = metrics.lock().unwrap(); + assert_eq!(metrics.start, 1); + assert_eq!(metrics.handle, 1); + assert_eq!(metrics.resume, 1); + assert!(metrics.pause >= 1); + } + spawn_handle.join().unwrap(); + remote.stop(); + handle.join().unwrap(); + { + let metrics = metrics.lock().unwrap(); + assert_eq!(metrics.start, 1); + assert_eq!(metrics.handle, 1); + assert_eq!(metrics.end, 1); + assert!(metrics.pause >= 1); + assert!(metrics.resume >= 1); + } + } + #[test] fn test_hooks() { let (tx, rx) = mpsc::channel(); @@ -151,4 +655,496 @@ mod tests { expected_metrics.end = 1; assert_eq!(expected_metrics, *metrics.lock().unwrap()); } + + #[test] + fn test_pop_or_sleep_uses_pending_retry_from_validate() { + let retry_at = Instant::now() + Duration::from_millis(20); + let queue = Arc::new(ScriptedQueue::new(vec![ + PopResult::Pending { retry_at }, + PopResult::Pending { retry_at }, + ScriptedQueue::ready(1), + ])); + let mut local = build_scripted_local(queue); + + assert!(local.pop_or_sleep(None).is_none()); + assert!(Instant::now() >= retry_at); + assert_next_ready_task(&mut local, 1); + } + + #[test] + fn test_pop_or_sleep_uses_initial_retry_when_validate_empty() { + let retry_at = Instant::now() + Duration::from_millis(20); + let queue = Arc::new(ScriptedQueue::new(vec![ + PopResult::Empty, + ScriptedQueue::ready(1), + ])); + let mut local = build_scripted_local(queue); + + assert!(local.pop_or_sleep(Some(retry_at)).is_none()); + assert!(Instant::now() >= retry_at); + assert_next_ready_task(&mut local, 1); + } + + #[test] + fn test_pop_or_sleep_returns_ready_from_validate_without_sleeping() { + let retry_at = Instant::now() + LATER_RETRY_OFFSET; + let queue = Arc::new(ScriptedQueue::new(vec![ScriptedQueue::ready(1)])); + let mut local = build_scripted_local(queue); + + let pop = local.pop_or_sleep(Some(retry_at)).unwrap(); + assert_eq!(pop.task_cell.id, 1); + assert!(Instant::now() < retry_at); + } + + #[test] + fn test_pop_or_sleep_uses_min_retry_when_validate_pending() { + let earlier_retry_at = Instant::now() + Duration::from_millis(20); + let later_retry_at = Instant::now() + LATER_RETRY_OFFSET; + let queue = Arc::new(ScriptedQueue::new(vec![ + PopResult::Pending { + retry_at: earlier_retry_at, + }, + PopResult::Pending { + retry_at: earlier_retry_at, + }, + ScriptedQueue::ready(1), + ])); + let mut local = build_scripted_local(queue); + + assert!(local.pop_or_sleep(Some(later_retry_at)).is_none()); + assert!(Instant::now() >= earlier_retry_at); + assert!(Instant::now() < later_retry_at); + assert_next_ready_task(&mut local, 1); + + let earlier_retry_at = Instant::now() + Duration::from_millis(20); + let later_retry_at = Instant::now() + LATER_RETRY_OFFSET; + let queue = Arc::new(ScriptedQueue::new(vec![ + PopResult::Pending { + retry_at: later_retry_at, + }, + ScriptedQueue::ready(2), + ])); + let mut local = build_scripted_local(queue); + + assert!(local.pop_or_sleep(Some(earlier_retry_at)).is_none()); + assert!(Instant::now() >= earlier_retry_at); + assert!(Instant::now() < later_retry_at); + assert_next_ready_task(&mut local, 2); + } + + #[test] + fn test_worker_retries_immediately_when_retry_time_has_passed() { + // A custom queue may report a stale retry time. The worker should not + // block for such a Pending result; it should retry immediately and run + // the task once the queue reports it as Ready. + let retry_at = Instant::now() - Duration::from_millis(10); + let (done_tx, done_rx) = mpsc::channel(); + let mut results = Vec::new(); + for _ in 0..WORKER_SPIN_POP_COUNT { + results.push(PopResult::Pending { retry_at }); + } + results.push(PopResult::Pending { retry_at }); + results.push(ready_callback_pop(done_tx)); + let queue = Arc::new(ScriptedQueue::new(results)); + + let started_at = Instant::now(); + let executed_at = check_scripted_worker_runs_task(queue, done_rx); + assert!(executed_at.duration_since(started_at) <= MAX_DELAYED_TASK_LAG); + } + + #[test] + fn test_worker_uses_last_spin_retry_when_it_gets_shorter() { + // The spin loop should pass its last observed retry time into + // pop_or_sleep. If the last observation gets shorter, the worker should + // wake at the shorter deadline instead of an older longer one. + let earlier_retry_at = Instant::now() + Duration::from_millis(50); + let later_retry_at = Instant::now() + LATER_RETRY_OFFSET; + let (done_tx, done_rx) = mpsc::channel(); + let mut results = Vec::new(); + for _ in 1..WORKER_SPIN_POP_COUNT { + results.push(PopResult::Pending { + retry_at: later_retry_at, + }); + } + results.push(PopResult::Pending { + retry_at: earlier_retry_at, + }); + results.push(PopResult::Empty); + results.push(ready_callback_pop(done_tx)); + let queue = Arc::new(ScriptedQueue::new(results)); + + let executed_at = check_scripted_worker_runs_task(queue, done_rx); + assert!(executed_at >= earlier_retry_at); + assert!(executed_at < later_retry_at); + } + + #[test] + fn test_worker_uses_last_spin_retry_when_it_gets_longer() { + // Conversely, if the last spin observation gets longer, the worker + // should not keep a stale shorter retry time. + let earlier_retry_at = Instant::now() + Duration::from_millis(20); + let later_retry_at = Instant::now() + Duration::from_millis(100); + let (done_tx, done_rx) = mpsc::channel(); + let mut results = Vec::new(); + for _ in 1..WORKER_SPIN_POP_COUNT { + results.push(PopResult::Pending { + retry_at: earlier_retry_at, + }); + } + results.push(PopResult::Pending { + retry_at: later_retry_at, + }); + results.push(PopResult::Empty); + results.push(ready_callback_pop(done_tx)); + let queue = Arc::new(ScriptedQueue::new(results)); + + let executed_at = check_scripted_worker_runs_task(queue, done_rx); + assert!(executed_at >= later_retry_at); + assert!(executed_at.duration_since(later_retry_at) <= MAX_DELAYED_TASK_LAG); + } + + fn check_worker_stops_at_pop_or_sleep_failpoint(failpoint: &'static str, stop_in_thread: bool) { + let _lock = lock_failpoint_tests(); + let _guard = fail::FailScenario::setup(); + let (entered_rx, release_tx) = configure_blocking_failpoint(failpoint); + let queue = Arc::new(ScriptedQueue::new(Vec::new())); + let (remote, _pause_rx, metrics, handle) = build_custom_worker(queue); + + entered_rx.recv_timeout(Duration::from_secs(1)).unwrap(); + let stop_handle = if stop_in_thread { + let remote = remote.clone(); + Some(thread::spawn(move || remote.stop())) + } else { + remote.stop(); + None + }; + release_tx.send(()).unwrap(); + if let Some(stop_handle) = stop_handle { + stop_handle.join().unwrap(); + } + handle.join().unwrap(); + + let metrics = metrics.lock().unwrap(); + assert_eq!(metrics.start, 1); + assert_eq!(metrics.handle, 0); + assert_eq!(metrics.end, 1); + } + + #[cfg_attr(not(feature = "failpoints"), ignore)] + #[test] + fn test_worker_stops_when_shutdown_before_mark_sleep() { + // Shutdown before mark_sleep should make validate fail without + // decrementing the active worker count. + check_worker_stops_at_pop_or_sleep_failpoint("worker-pop-or-sleep-before-park", false); + } + + #[cfg_attr(not(feature = "failpoints"), ignore)] + #[test] + fn test_worker_stops_when_shutdown_during_validate() { + // Shutdown while validate is running happens after mark_sleep has + // succeeded. The worker should still return from park and finish. + check_worker_stops_at_pop_or_sleep_failpoint( + "worker-pop-or-sleep-before-validate-pop", + true, + ); + } + + #[cfg_attr(not(feature = "failpoints"), ignore)] + #[test] + fn test_worker_wakes_ready_task_after_pending_then_validate_empty() { + // The worker carries an initial Pending retry from spin, validate then + // sees Empty and parks with that timeout. A newly inserted ready task + // should still wake the worker immediately instead of waiting for the + // old retry deadline. + let _lock = lock_failpoint_tests(); + let _guard = fail::FailScenario::setup(); + let (entered_rx, release_tx) = + configure_blocking_failpoint("worker-pop-or-sleep-before-sleep"); + let retry_at = Instant::now() + LATER_RETRY_OFFSET; + let mut results = Vec::new(); + for _ in 0..WORKER_SPIN_POP_COUNT { + results.push(PopResult::Pending { retry_at }); + } + results.push(PopResult::Empty); + let queue = Arc::new(ScriptedQueue::new(results)); + let (remote, _pause_rx, metrics, handle) = build_custom_worker(queue); + + entered_rx.recv_timeout(Duration::from_secs(1)).unwrap(); + let (done_tx, done_rx) = mpsc::channel(); + let inserted_at = Instant::now(); + remote.spawn(move |_: &mut callback::Handle<'_>| { + done_tx.send(Instant::now()).unwrap(); + }); + release_tx.send(()).unwrap(); + + let executed_at = done_rx.recv_timeout(MAX_DELAYED_TASK_LAG).unwrap(); + assert!(executed_at >= inserted_at); + assert!(executed_at.duration_since(inserted_at) <= MAX_DELAYED_TASK_LAG); + assert!(executed_at < retry_at); + + remote.stop(); + handle.join().unwrap(); + let metrics = metrics.lock().unwrap(); + assert_eq!(metrics.start, 1); + assert_eq!(metrics.handle, 1); + assert_eq!(metrics.end, 1); + } + + #[cfg_attr(not(feature = "failpoints"), ignore)] + #[test] + fn test_worker_runs_ready_task_inserted_while_pending() { + // The worker is sleeping for a delayed task's Pending retry. A later + // ready task should wake it and run immediately, without waiting for + // the delayed task's retry time. + check_worker_runs_ready_task_inserted_while_pending(); + } + + #[cfg_attr(not(feature = "failpoints"), ignore)] + #[test] + fn test_worker_refreshes_timeout_when_earlier_pending_task_is_inserted() { + // The worker is already queued to sleep with a later Pending retry. + // Pushing another delayed task with an earlier retry should wake it so + // the worker can replace the old timeout with the earlier deadline. + let _lock = lock_failpoint_tests(); + let _guard = fail::FailScenario::setup(); + let (entered_rx, release_tx) = + configure_blocking_failpoint("worker-pop-or-sleep-before-sleep"); + let queue = Arc::new(DeadlineQueue::new(vec![ + LATER_RETRY_OFFSET, + DELAYED_TASK_DELAY, + ])); + let (unexpected_tx, unexpected_rx) = mpsc::channel(); + queue.push(callback_task(move |_: &mut callback::Handle<'_>| { + unexpected_tx.send(()).unwrap(); + })); + let later_retry_at = queue.ready_at(0); + let (remote, _pause_rx, metrics, handle) = build_custom_worker(queue.clone()); + entered_rx.recv_timeout(Duration::from_secs(1)).unwrap(); + release_tx.send(()).unwrap(); + + // Let the worker pass the before_sleep hook and block on the later + // timeout before pushing the earlier delayed task. + thread::sleep(Duration::from_millis(20)); + let (done_tx, done_rx) = mpsc::channel(); + remote.spawn(move |_: &mut callback::Handle<'_>| { + done_tx.send(Instant::now()).unwrap(); + }); + let earlier_retry_at = queue.ready_at(1); + + let executed_at = done_rx + .recv_timeout(MAX_DELAYED_TASK_LAG + Duration::from_secs(1)) + .unwrap(); + assert!(executed_at >= earlier_retry_at); + assert!(executed_at.duration_since(earlier_retry_at) <= MAX_DELAYED_TASK_LAG); + assert!(executed_at < later_retry_at); + assert!(unexpected_rx.try_recv().is_err()); + + remote.stop(); + handle.join().unwrap(); + let metrics = metrics.lock().unwrap(); + assert_eq!(metrics.start, 1); + assert_eq!(metrics.handle, 1); + assert_eq!(metrics.end, 1); + } + + #[cfg_attr(not(feature = "failpoints"), ignore)] + #[test] + fn test_scaled_down_worker_reparks_after_pending_timeout() { + // A worker that becomes above core_thread_count while carrying a + // Pending timeout should not return to the outer spin-pop path when + // the timeout fires. It should clear the timeout and park again. + let _lock = lock_failpoint_tests(); + let _guard = fail::FailScenario::setup(); + let (sleep_rx, release_tx) = + configure_counting_failpoint("worker-pop-or-sleep-before-sleep", 2); + let queue = Arc::new(AlwaysPendingQueue { + retry_at: Instant::now() + Duration::from_millis(20), + }); + let queue_builder = CustomBuilder::new(CustomConfig::default(), queue); + let (remote, mut locals) = build_spawn(queue_builder, two_thread_config()); + let (pause_rx, metrics, handle) = start_custom_worker(locals.remove(1)); + + pause_rx.recv_timeout(Duration::from_secs(1)).unwrap(); + assert_eq!(sleep_rx.recv_timeout(Duration::from_secs(1)).unwrap(), 1); + remote.scale_workers(1); + release_tx.send(()).unwrap(); + + assert_eq!(sleep_rx.recv_timeout(Duration::from_secs(1)).unwrap(), 2); + { + let metrics = metrics.lock().unwrap(); + assert_eq!(metrics.pause, 1); + assert_eq!(metrics.resume, 0); + } + release_tx.send(()).unwrap(); + thread::sleep(Duration::from_millis(20)); + + remote.stop(); + handle.join().unwrap(); + let metrics = metrics.lock().unwrap(); + assert_eq!(metrics.start, 1); + assert_eq!(metrics.handle, 0); + assert_eq!(metrics.pause, 1); + assert_eq!(metrics.resume, 1); + assert_eq!(metrics.end, 1); + } + + #[cfg_attr(not(feature = "failpoints"), ignore)] + #[test] + fn test_scaled_down_pending_timeout_wakes_core_worker() { + // Worker 1 parks on an empty queue without a timeout. A delayed task is + // then inserted directly into the custom queue so no push wakeup is + // sent. Worker 2 observes the Pending timeout, gets scaled down, and + // must wake worker 1 when the timeout fires. + let _lock = lock_failpoint_tests(); + let _guard = fail::FailScenario::setup(); + let (sleep_rx, _) = configure_counting_failpoint("worker-pop-or-sleep-before-sleep", 0); + let queue = Arc::new(DeadlineQueue::new(vec![Duration::from_millis(100)])); + let queue_builder = CustomBuilder::new(CustomConfig::default(), queue.clone()); + let (remote, mut locals) = build_spawn(queue_builder, two_thread_config()); + let local_2 = locals.remove(1); + let local_1 = locals.remove(0); + let (pause_rx_1, metrics_1, handle_1) = start_custom_worker(local_1); + + pause_rx_1.recv_timeout(Duration::from_secs(1)).unwrap(); + assert_eq!(sleep_rx.recv_timeout(Duration::from_secs(1)).unwrap(), 1); + fail::remove("worker-pop-or-sleep-before-sleep"); + thread::sleep(Duration::from_millis(20)); + + let (done_tx, done_rx) = mpsc::channel(); + queue.push(callback_task(move |_: &mut callback::Handle<'_>| { + done_tx.send(Instant::now()).unwrap(); + })); + let ready_at = queue.ready_at(0); + let (pause_rx_2, metrics_2, handle_2) = start_custom_worker(local_2); + + pause_rx_2.recv_timeout(Duration::from_secs(1)).unwrap(); + remote.scale_workers(1); + + let executed_at = done_rx + .recv_timeout(MAX_DELAYED_TASK_LAG + Duration::from_secs(1)) + .unwrap(); + assert!(executed_at >= ready_at); + assert!(executed_at.duration_since(ready_at) <= MAX_DELAYED_TASK_LAG); + + remote.stop(); + handle_1.join().unwrap(); + handle_2.join().unwrap(); + let metrics_1 = metrics_1.lock().unwrap(); + let metrics_2 = metrics_2.lock().unwrap(); + assert_eq!(metrics_1.handle, 1); + assert_eq!(metrics_1.end, 1); + assert_eq!(metrics_2.handle, 0); + assert_eq!(metrics_2.end, 1); + } + + #[cfg_attr(not(feature = "failpoints"), ignore)] + #[test] + fn test_worker_wakes_when_task_inserted_before_park() { + // The task is inserted after the worker decides to sleep, but before it + // calls parking_lot_core::park. The validate callback should pop the + // task as Ready and abort the park. + check_worker_wakes_when_task_is_inserted_at("worker-pop-or-sleep-before-park"); + } + + #[cfg_attr(not(feature = "failpoints"), ignore)] + #[test] + fn test_worker_wakes_when_task_inserted_before_validate_pop() { + // The task is inserted after mark_sleep succeeds inside validate, but + // before validate pops from the queue. The pop should see the inserted + // task immediately and avoid sleeping. + check_worker_wakes_when_task_is_inserted_at("worker-pop-or-sleep-before-validate-pop"); + } + + #[cfg_attr(not(feature = "failpoints"), ignore)] + #[test] + fn test_worker_wakes_when_task_inserted_before_sleep() { + // The task is inserted after validate has returned Empty and the worker + // has been queued for parking, but before it actually sleeps. The push + // should unpark the worker, and the next worker loop should run the + // inserted task immediately. + check_worker_wakes_when_task_is_inserted_at("worker-pop-or-sleep-before-sleep"); + } + + fn check_worker_runs_delayed_task_without_new_insert( + scenario: DelayedQueueScenario, + ) -> (Instant, Instant) { + let queue = Arc::new(DeadlineQueue::new(vec![LATER_RETRY_OFFSET])); + let (done_tx, done_rx) = mpsc::channel(); + queue.push(callback_task(move |_: &mut callback::Handle<'_>| { + done_tx.send(Instant::now()).unwrap(); + })); + match scenario { + DelayedQueueScenario::PendingDuringSpinAndValidate => { + let later_retry_at = Instant::now() + LATER_RETRY_OFFSET; + for _ in 0..WORKER_SPIN_POP_COUNT { + queue.push_scripted_result(DeadlineScript::Pending(later_retry_at)); + } + } + DelayedQueueScenario::EmptyDuringSpin => { + for _ in 0..WORKER_SPIN_POP_COUNT { + queue.push_scripted_result(DeadlineScript::Empty); + } + } + DelayedQueueScenario::EarlierRetryInValidate => { + let later_retry_at = Instant::now() + LATER_RETRY_OFFSET; + for _ in 0..WORKER_SPIN_POP_COUNT { + queue.push_scripted_result(DeadlineScript::Pending(later_retry_at)); + } + } + } + // The retry deadline is installed while the worker is in pop_or_sleep + // validation. This keeps the test from depending on how quickly the CI + // runner starts the worker thread after the task was pushed. + queue.push_scripted_result(DeadlineScript::PendingAfter(DELAYED_TASK_DELAY)); + let (remote, _pause_rx, metrics, handle) = build_custom_worker(queue.clone()); + let executed_at = done_rx + .recv_timeout(MAX_DELAYED_TASK_LAG + Duration::from_secs(1)) + .unwrap(); + let ready_at = queue.ready_at(0); + + assert!(executed_at >= ready_at); + assert!(executed_at.duration_since(ready_at) <= MAX_DELAYED_TASK_LAG); + { + let metrics = metrics.lock().unwrap(); + assert_eq!(metrics.start, 1); + assert_eq!(metrics.handle, 1); + assert!(metrics.resume >= 1); + assert!(metrics.pause >= 1); + } + + remote.stop(); + handle.join().unwrap(); + assert_eq!(metrics.lock().unwrap().end, 1); + + (executed_at, ready_at) + } + + #[test] + fn test_worker_runs_delayed_task_after_pending_then_pending() { + // The worker first observes Pending during spin, then observes Pending + // again inside pop_or_sleep. With no later push, it should wake by the + // retry timeout and run the delayed task. + check_worker_runs_delayed_task_without_new_insert( + DelayedQueueScenario::PendingDuringSpinAndValidate, + ); + } + + #[test] + fn test_worker_runs_delayed_task_after_empty_then_pending() { + // The worker first observes Empty during spin, then Pending inside + // pop_or_sleep. The validate Pending retry should still drive a timed + // park and let the delayed task run once it becomes ready. + check_worker_runs_delayed_task_without_new_insert(DelayedQueueScenario::EmptyDuringSpin); + } + + #[test] + fn test_worker_uses_shorter_retry_from_pending_validate() { + // The worker observes a later Pending retry during spin, then a shorter + // Pending retry in pop_or_sleep validate. It should use the shorter + // retry instead of sleeping until the stale later deadline. + let (executed_at, ready_at) = check_worker_runs_delayed_task_without_new_insert( + DelayedQueueScenario::EarlierRetryInValidate, + ); + assert!(executed_at < ready_at + LATER_RETRY_OFFSET); + } } diff --git a/src/queue.rs b/src/queue.rs index fefc6ec..935da54 100644 --- a/src/queue.rs +++ b/src/queue.rs @@ -11,9 +11,11 @@ pub mod multilevel; pub mod priority; +mod custom; mod extras; mod single_level; +pub use self::custom::{Builder as CustomBuilder, Config as CustomConfig, TaskQueue}; pub use self::extras::Extras; use std::time::Instant; @@ -44,24 +46,28 @@ enum InjectorInner { SingleLevel(single_level::TaskInjector), Multilevel(multilevel::TaskInjector), Priority(priority::TaskInjector), + Custom(custom::TaskInjector), } impl TaskInjector { /// Pushes a task to the queue. + #[inline] pub fn push(&self, task_cell: T) { match &self.0 { InjectorInner::SingleLevel(q) => q.push(task_cell), InjectorInner::Multilevel(q) => q.push(task_cell), InjectorInner::Priority(q) => q.push(task_cell), + InjectorInner::Custom(q) => q.push(task_cell), } } + #[inline] pub fn default_extras(&self) -> Extras { - match self.0 { + match &self.0 { InjectorInner::SingleLevel(_) => Extras::single_level(), - InjectorInner::Multilevel(_) | InjectorInner::Priority(_) => { - Extras::multilevel_default() - } + InjectorInner::Multilevel(_) + | InjectorInner::Priority(_) + | InjectorInner::Custom(_) => Extras::multilevel_default(), } } @@ -77,6 +83,7 @@ impl TaskInjector { /// /// Callers may retry on `None` if they need stronger guarantees under /// contention. + #[inline] pub fn try_evict_lowest(&self, incoming_priority: u64) -> Option { match &self.0 { InjectorInner::Priority(q) => q.try_evict(incoming_priority), @@ -98,6 +105,87 @@ pub struct Pop { pub from_local: bool, } +/// Result of attempting to pop a task from a task queue. +/// +/// A queue should return [`PopResult::Empty`] only when it has no queued +/// task at all. If the queue still owns tasks but none of them can run now, +/// return [`PopResult::Pending`] instead. A later push may make work ready +/// before `retry_at`; otherwise the worker must retry no later than `retry_at`. +pub enum PopResult { + /// A task is ready and should be scheduled immediately. + Ready(Pop), + /// The queue has at least one task, but no task can be scheduled now. + /// + /// `retry_at` specifies the earliest time when the worker should try to + /// pop from the queue again. + Pending { + /// Earliest time when the worker should retry popping from the queue. + retry_at: Instant, + }, + /// The queue has no task. + Empty, +} + +impl PopResult { + /// Returns `true` if the pop result contains a ready task. + #[inline] + pub fn is_ready(&self) -> bool { + matches!(self, PopResult::Ready(_)) + } + + /// Returns `true` if the queue has tasks but none can run yet. + #[inline] + pub fn is_pending(&self) -> bool { + matches!(self, PopResult::Pending { .. }) + } + + /// Returns `true` if the queue has no task. + #[inline] + pub fn is_empty(&self) -> bool { + matches!(self, PopResult::Empty) + } + + /// Returns the ready task, panicking if the result is not ready. + #[inline] + pub fn unwrap_ready(self) -> Pop { + match self { + PopResult::Ready(pop) => pop, + PopResult::Pending { .. } => panic!("called `PopResult::unwrap_ready()` on `Pending`"), + PopResult::Empty => panic!("called `PopResult::unwrap_ready()` on `Empty`"), + } + } + + /// Returns the retry time, panicking if the result is not pending. + #[inline] + pub fn unwrap_pending(self) -> Instant { + match self { + PopResult::Pending { retry_at } => retry_at, + PopResult::Ready(_) => panic!("called `PopResult::unwrap_pending()` on `Ready`"), + PopResult::Empty => panic!("called `PopResult::unwrap_pending()` on `Empty`"), + } + } + + /// Verifies the result is empty, panicking otherwise. + #[inline] + pub fn unwrap_empty(self) { + match self { + PopResult::Empty => {} + PopResult::Ready(_) => panic!("called `PopResult::unwrap_empty()` on `Ready`"), + PopResult::Pending { .. } => panic!("called `PopResult::unwrap_empty()` on `Pending`"), + } + } +} + +impl From>> for PopResult { + #[inline] + fn from(pop: Option>) -> PopResult { + match pop { + Some(pop) => PopResult::Ready(pop), + None => PopResult::Empty, + } + } +} + /// The local queue of a task queue. pub(crate) struct LocalQueue(LocalQueueInner); @@ -105,50 +193,71 @@ enum LocalQueueInner { SingleLevel(single_level::LocalQueue), Multilevel(multilevel::LocalQueue), Priority(priority::LocalQueue), + Custom(custom::LocalQueue), } impl LocalQueue { /// Pushes a task to the local queue. + #[inline] pub fn push(&mut self, task_cell: T) { match &mut self.0 { LocalQueueInner::SingleLevel(q) => q.push(task_cell), LocalQueueInner::Multilevel(q) => q.push(task_cell), LocalQueueInner::Priority(q) => q.push(task_cell), + LocalQueueInner::Custom(q) => q.push(task_cell), + } + } + + /// Gets a task cell from the queue. + #[inline] + pub fn pop(&mut self) -> PopResult { + match &mut self.0 { + LocalQueueInner::SingleLevel(q) => q.pop().into(), + LocalQueueInner::Multilevel(q) => q.pop().into(), + LocalQueueInner::Priority(q) => q.pop().into(), + LocalQueueInner::Custom(q) => q.pop(), } } - /// Gets a task cell from the queue. Returns `None` if there is no task cell - /// available. - pub fn pop(&mut self) -> Option> { + /// Forcefully drains all currently queued tasks. + #[inline] + pub fn drain(&mut self) { match &mut self.0 { - LocalQueueInner::SingleLevel(q) => q.pop(), - LocalQueueInner::Multilevel(q) => q.pop(), - LocalQueueInner::Priority(q) => q.pop(), + LocalQueueInner::SingleLevel(q) => while q.pop().is_some() {}, + LocalQueueInner::Multilevel(q) => while q.pop().is_some() {}, + LocalQueueInner::Priority(q) => while q.pop().is_some() {}, + LocalQueueInner::Custom(q) => q.drain(), } } + #[inline] pub fn default_extras(&self) -> Extras { - match self.0 { + match &self.0 { LocalQueueInner::SingleLevel(_) => Extras::single_level(), LocalQueueInner::Multilevel(_) => Extras::multilevel_default(), LocalQueueInner::Priority(_) => Extras::single_level(), + LocalQueueInner::Custom(_) => Extras::multilevel_default(), } } /// If there are tasks in the local queue, returns true. Otherwise, pulls /// tasks from the global queue and returns whether it succeeds. + #[inline] pub fn has_tasks_or_pull(&mut self) -> bool { match &mut self.0 { LocalQueueInner::SingleLevel(q) => q.has_tasks_or_pull(), LocalQueueInner::Multilevel(q) => q.has_tasks_or_pull(), LocalQueueInner::Priority(q) => q.has_tasks_or_pull(), + LocalQueueInner::Custom(q) => q.has_tasks_or_pull(), } } } /// Supported available queues. -pub enum QueueType { +#[derive(Default)] +pub enum QueueType { /// A single level work stealing queue. + #[default] SingleLevel, /// A multilevel feedback queue. /// @@ -156,31 +265,37 @@ pub enum QueueType { Multilevel(multilevel::Builder), /// A concurrent prioirty queue. Priority(priority::Builder), + /// A custom task queue. + Custom(CustomBuilder), } -impl Default for QueueType { - fn default() -> QueueType { - QueueType::SingleLevel +impl From for QueueType { + fn from(b: multilevel::Builder) -> QueueType { + QueueType::Multilevel(b) } } -impl From for QueueType { - fn from(b: multilevel::Builder) -> QueueType { - QueueType::Multilevel(b) +impl From for QueueType { + fn from(b: priority::Builder) -> QueueType { + QueueType::Priority(b) } } -impl From for QueueType { - fn from(b: priority::Builder) -> QueueType { - QueueType::Priority(b) +impl From> for QueueType { + fn from(b: CustomBuilder) -> QueueType { + QueueType::Custom(b) } } -pub(crate) fn build(ty: QueueType, local_num: usize) -> (TaskInjector, Vec>) { +pub(crate) fn build( + ty: QueueType, + local_num: usize, +) -> (TaskInjector, Vec>) { match ty { QueueType::SingleLevel => single_level(local_num), QueueType::Multilevel(b) => b.build(local_num), QueueType::Priority(b) => b.build(local_num), + QueueType::Custom(b) => b.build(local_num), } } diff --git a/src/queue/custom.rs b/src/queue/custom.rs new file mode 100644 index 0000000..fc33eed --- /dev/null +++ b/src/queue/custom.rs @@ -0,0 +1,709 @@ +// Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0. + +//! Custom task queue abstractions. +//! +//! This module defines the common interfaces implemented by custom task queues. + +use std::sync::Arc; + +use super::{ + multilevel::{MultiLevelMetrics, TrackedRunnerBuilder}, + PopResult, +}; + +/// Common interface implemented by custom task queues used by the thread pool. +/// +/// A task queue accepts task cells from producers and selects tasks for workers +/// to run. [`TaskQueue::pop`] describes both task availability and scheduling +/// readiness: +/// +/// - [`PopResult::Ready`] means a task is available and should be run +/// immediately. +/// - [`PopResult::Pending`] means the queue contains tasks, but none are +/// ready to run now. The worker may sleep, but it can still be woken by a +/// later push and must retry no later than `retry_at`. This is useful for +/// queues with waiting states, such as waiting for rate-limit tokens before +/// a task can run. +/// - [`PopResult::Empty`] means the queue contains no tasks. If workers +/// keep seeing an empty queue for a while, they may park until a later +/// [`TaskQueue::push`] wakes them. +/// +/// Queues must not return [`PopResult::Empty`] while they still contain +/// delayed or throttled tasks. Doing so can let workers park indefinitely and +/// leave those tasks unscheduled for a long time. Return +/// [`PopResult::Pending`] for that case instead. +pub trait TaskQueue: Send + Sync + 'static { + /// Pushes a task into the queue. + fn push(&self, task_cell: T); + + /// Pops the next task to run, or reports why no task can run now. + fn pop(&self) -> PopResult; + + /// Drains all queued tasks regardless of scheduling readiness. + /// + /// This is used by shutdown paths to drop remaining tasks. Unlike + /// [`TaskQueue::pop`], it should not leave delayed or throttled tasks in + /// the queue just because they are not ready to run. + /// + /// A custom queue is shared by all worker-local handles, so shutdown may + /// call this method multiple times or concurrently. Implementations must be + /// idempotent and thread-safe. + fn drain(&self); + + /// Returns whether the queue may have a ready task. + /// + /// This method must not remove a task from the queue. It is used as a + /// preemption hint, so returning `false` while ready tasks exist can delay + /// those tasks and hurt scheduling fairness. Returning `true` while no task + /// is ready can cause unnecessary rescheduling. + fn has_ready_task(&self) -> bool; +} + +/// The configurations of custom task queues. +#[derive(Default)] +pub struct Config { + name: Option, +} + +impl Config { + /// Sets the name of the custom task queue. Metrics are available if name is provided. + pub fn name(mut self, name: Option>) -> Self { + self.name = name.map(Into::into); + self + } +} + +/// The builder of a custom task queue. +pub struct Builder { + queue: Arc>, + metrics: MultiLevelMetrics, +} + +impl Builder { + /// Creates a custom task queue builder from a shared queue. + pub fn new(config: Config, queue: Arc>) -> Builder { + Builder { + queue, + metrics: MultiLevelMetrics::new(config.name.as_deref()), + } + } + + /// Creates a runner builder for the custom task queue with a normal runner builder. + pub fn runner_builder(&self, inner_runner_builder: B) -> TrackedRunnerBuilder { + TrackedRunnerBuilder::new(inner_runner_builder, self.metrics.clone(), true) + } +} + +impl Builder { + /// Creates the injector and local queue handles of the custom task queue. + pub(crate) fn build( + self, + local_num: usize, + ) -> (super::TaskInjector, Vec>) { + let injector = TaskInjector::new(self.queue.clone()); + let locals: Vec> = + std::iter::repeat_with(|| LocalQueue::new(self.queue.clone())) + .take(local_num) + .collect(); + + ( + super::TaskInjector(super::InjectorInner::Custom(injector)), + locals + .into_iter() + .map(|local| super::LocalQueue(super::LocalQueueInner::Custom(local))) + .collect(), + ) + } +} + +/// The injector of a custom task queue. +pub struct TaskInjector { + queue: Arc>, +} + +impl Clone for TaskInjector { + #[inline] + fn clone(&self) -> TaskInjector { + TaskInjector { + queue: self.queue.clone(), + } + } +} + +impl TaskInjector { + /// Creates a custom task queue injector from a shared queue. + #[inline] + pub fn new(queue: Arc>) -> TaskInjector { + TaskInjector { queue } + } + + /// Pushes a task into the custom queue. + #[inline] + pub fn push(&self, task_cell: T) { + self.queue.push(task_cell); + } +} + +/// The local queue handle of a custom task queue. +pub struct LocalQueue { + queue: Arc>, +} + +impl Clone for LocalQueue { + #[inline] + fn clone(&self) -> LocalQueue { + LocalQueue { + queue: self.queue.clone(), + } + } +} + +impl LocalQueue { + /// Creates a local queue handle from a shared custom queue. + #[inline] + pub fn new(queue: Arc>) -> LocalQueue { + LocalQueue { queue } + } + + /// Pushes a task into the custom queue. + #[inline] + pub fn push(&self, task_cell: T) { + self.queue.push(task_cell); + } + + /// Pops a task from the custom queue. + #[inline] + pub fn pop(&self) -> PopResult { + self.queue.pop() + } + + /// Forcefully drains all tasks from the custom queue. + /// + /// Some queue implementations have local handles backed by shared queue + /// state. This adapter forwards each local drain directly to the + /// user-provided [`TaskQueue`], so shutdown may call [`TaskQueue::drain`] + /// from multiple worker threads. Implementations must be idempotent and + /// thread-safe. + #[inline] + pub fn drain(&self) { + self.queue.drain(); + } + + /// Returns whether the custom queue may have ready work for preemption. + /// + /// This forwards to [`TaskQueue::has_ready_task`], which must be a + /// non-consuming readiness hint. + #[inline] + pub fn has_tasks_or_pull(&self) -> bool { + self.queue.has_ready_task() + } +} + +#[cfg(test)] +mod tests { + use std::{ + collections::VecDeque, + sync::{ + atomic::{AtomicBool, AtomicUsize, Ordering}, + mpsc, Arc, Barrier, Mutex, + }, + thread, + time::{Duration, Instant}, + }; + + use super::*; + use crate::{ + metrics::{ + MULTILEVEL_LEVEL_ELAPSED, TASK_EXEC_DURATION, TASK_EXEC_TIMES, TASK_POLL_DURATION, + TASK_WAIT_DURATION, + }, + pool::{build_spawn, Local, Runner, RunnerBuilder}, + queue::{Extras, Pop, TaskCell}, + }; + + struct MockTask { + id: u64, + sleep_ms: u64, + extras: Extras, + } + + impl MockTask { + fn new(id: u64) -> MockTask { + MockTask { + id, + sleep_ms: 0, + extras: Extras::multilevel_default(), + } + } + + fn with_sleep(id: u64, sleep_ms: u64) -> MockTask { + MockTask { + id, + sleep_ms, + extras: Extras::multilevel_default(), + } + } + } + + impl TaskCell for MockTask { + fn mut_extras(&mut self) -> &mut Extras { + &mut self.extras + } + } + + struct MockRunner; + + impl Runner for MockRunner { + type TaskCell = MockTask; + + fn handle(&mut self, _local: &mut Local, task_cell: MockTask) -> bool { + thread::sleep(Duration::from_millis(task_cell.sleep_ms)); + true + } + } + + struct MockRunnerBuilder; + + impl RunnerBuilder for MockRunnerBuilder { + type Runner = MockRunner; + + fn build(&mut self) -> MockRunner { + MockRunner + } + } + + struct MockQueue { + tasks: Mutex>, + scripted_results: Mutex>>, + ready_hint: AtomicBool, + drain_count: AtomicUsize, + } + + impl Default for MockQueue { + fn default() -> MockQueue { + MockQueue { + tasks: Mutex::new(VecDeque::new()), + scripted_results: Mutex::new(VecDeque::new()), + ready_hint: AtomicBool::new(false), + drain_count: AtomicUsize::new(0), + } + } + } + + impl MockQueue { + fn push_scripted_result(&self, result: PopResult) { + self.scripted_results.lock().unwrap().push_back(result); + } + + fn set_ready_hint(&self, ready: bool) { + self.ready_hint.store(ready, Ordering::SeqCst); + } + + fn len(&self) -> usize { + self.tasks.lock().unwrap().len() + } + + fn drain_count(&self) -> usize { + self.drain_count.load(Ordering::SeqCst) + } + } + + impl TaskQueue for MockQueue { + fn push(&self, task_cell: T) { + self.tasks.lock().unwrap().push_back(task_cell); + self.ready_hint.store(true, Ordering::SeqCst); + } + + fn pop(&self) -> PopResult { + if let Some(result) = self.scripted_results.lock().unwrap().pop_front() { + return result; + } + + let mut tasks = self.tasks.lock().unwrap(); + let task = tasks.pop_front(); + self.ready_hint.store(!tasks.is_empty(), Ordering::SeqCst); + task.map(|task_cell| Pop { + task_cell, + schedule_time: Instant::now(), + from_local: false, + }) + .into() + } + + fn drain(&self) { + self.drain_count.fetch_add(1, Ordering::SeqCst); + self.tasks.lock().unwrap().clear(); + self.scripted_results.lock().unwrap().clear(); + self.ready_hint.store(false, Ordering::SeqCst); + } + + fn has_ready_task(&self) -> bool { + self.ready_hint.load(Ordering::SeqCst) + } + } + + struct PendingQueue { + tasks: Mutex>, + pending_tx: Mutex>>, + drain_count: AtomicUsize, + } + + impl PendingQueue { + fn new(pending_tx: mpsc::Sender<()>) -> PendingQueue { + PendingQueue { + tasks: Mutex::new(VecDeque::new()), + pending_tx: Mutex::new(Some(pending_tx)), + drain_count: AtomicUsize::new(0), + } + } + + fn len(&self) -> usize { + self.tasks.lock().unwrap().len() + } + + fn drain_count(&self) -> usize { + self.drain_count.load(Ordering::SeqCst) + } + } + + impl TaskQueue for PendingQueue { + fn push(&self, task_cell: T) { + self.tasks.lock().unwrap().push_back(task_cell); + } + + fn pop(&self) -> PopResult { + if self.tasks.lock().unwrap().is_empty() { + return PopResult::Empty; + } + + if let Some(tx) = self.pending_tx.lock().unwrap().take() { + let _ = tx.send(()); + } + PopResult::Pending { + retry_at: Instant::now() + Duration::from_secs(60), + } + } + + fn drain(&self) { + self.drain_count.fetch_add(1, Ordering::SeqCst); + self.tasks.lock().unwrap().clear(); + } + + fn has_ready_task(&self) -> bool { + false + } + } + + #[test] + fn test_build_uses_shared_queue() { + // Custom queues do not create separate per-worker queues. This test + // verifies that the builder wires the injector and every local handle + // to the same user-provided queue. + let queue = Arc::new(MockQueue::default()); + let builder = Builder::new(Config::default(), queue); + let (injector, mut locals) = builder.build(3); + + // The injector and all local handles wrap the same custom queue, so a + // task pushed through the injector can be popped from any local handle. + injector.push(MockTask::new(1)); + assert_eq!(locals[0].pop().unwrap_ready().task_cell.id, 1); + + // A push through one local handle also goes to the shared custom queue, + // not to per-worker local storage. + locals[1].push(MockTask::new(2)); + assert_eq!(locals[2].pop().unwrap_ready().task_cell.id, 2); + assert!(locals[0].pop().is_empty()); + } + + #[test] + fn test_pop_result_forwarding() { + // Custom local queues should preserve the exact pop state reported by + // the user-provided queue. The scheduler distinguishes Ready, Pending, + // and Empty when deciding whether to run, retry later, or sleep. + let queue = Arc::new(MockQueue::default()); + let retry_at = Instant::now() + Duration::from_millis(10); + let schedule_time = Instant::now(); + queue.push_scripted_result(PopResult::Ready(Pop { + task_cell: MockTask::new(1), + schedule_time, + from_local: true, + })); + queue.push_scripted_result(PopResult::Pending { retry_at }); + queue.push_scripted_result(PopResult::Empty); + + let builder = Builder::new(Config::default(), queue); + let (_, mut locals) = builder.build(1); + + // Ready keeps the original Pop metadata. + let pop = locals[0].pop().unwrap_ready(); + assert_eq!(pop.task_cell.id, 1); + assert_eq!(pop.schedule_time, schedule_time); + assert!(pop.from_local); + // Pending and Empty must not be collapsed into each other. + assert_eq!(locals[0].pop().unwrap_pending(), retry_at); + locals[0].pop().unwrap_empty(); + } + + #[test] + fn test_drain_forwards_to_task_queue() { + // Shutdown drains through LocalQueue::drain, so custom queues must see + // the drain call instead of relying on repeated readiness-aware pops. + let queue = Arc::new(MockQueue::default()); + let builder = Builder::new(Config::default(), queue.clone()); + let (injector, mut locals) = builder.build(1); + + injector.push(MockTask::new(1)); + injector.push(MockTask::new(2)); + assert_eq!(queue.len(), 2); + + // Draining should clear all queued tasks, including tasks a custom + // queue might otherwise report as Pending. + locals[0].drain(); + assert!(queue.drain_count() >= 1); + assert_eq!(queue.len(), 0); + assert!(locals[0].pop().is_empty()); + } + + #[test] + fn test_drain_clears_tasks_that_pop_reports_pending() { + // A custom queue can own tasks that are not ready yet and report + // Pending from pop. Shutdown still needs drain to drop those tasks. + let queue = Arc::new(MockQueue::default()); + let builder = Builder::new(Config::default(), queue.clone()); + let (injector, mut locals) = builder.build(1); + let retry_at = Instant::now() + Duration::from_secs(1); + + queue.push_scripted_result(PopResult::Pending { retry_at }); + injector.push(MockTask::new(1)); + + assert!(locals[0].pop().is_pending()); + assert_eq!(queue.len(), 1); + + locals[0].drain(); + assert!(queue.drain_count() >= 1); + assert_eq!(queue.len(), 0); + assert!(locals[0].pop().is_empty()); + } + + #[test] + fn test_drain_is_idempotent_when_called_by_shared_locals() { + // Some queue implementations have local handles backed by shared queue + // state. The custom adapter exposes this contract to the user-provided + // queue by forwarding each local drain directly, so implementations must + // tolerate repeated and concurrent calls while leaving the queue fully + // cleared. + let queue = Arc::new(MockQueue::default()); + let builder = Builder::new(Config::default(), queue.clone()); + let (injector, locals) = builder.build(4); + let local_num = locals.len(); + let barrier = Arc::new(Barrier::new(local_num)); + + for i in 0..16 { + injector.push(MockTask::new(i)); + } + assert_eq!(queue.len(), 16); + + let handles: Vec<_> = locals + .into_iter() + .map(|mut local| { + let barrier = barrier.clone(); + thread::spawn(move || { + barrier.wait(); + local.drain(); + local.drain(); + }) + }) + .collect(); + + for handle in handles { + handle.join().unwrap(); + } + + assert!(queue.drain_count() >= local_num); + assert_eq!(queue.len(), 0); + } + + #[test] + fn test_has_tasks_or_pull_forwards_ready_hint() { + // Custom queues provide their own non-consuming readiness hint for + // preemption. The adapter should forward it without popping a task. + let queue = Arc::new(MockQueue::default()); + let builder = Builder::new(Config::default(), queue.clone()); + let (injector, mut locals) = builder.build(1); + + injector.push(MockTask::new(1)); + queue.set_ready_hint(true); + assert!(locals[0].has_tasks_or_pull()); + assert_eq!(queue.len(), 1); + + // A false hint should also be forwarded as-is, and checking it must + // still leave queued tasks untouched. + queue.set_ready_hint(false); + assert!(!locals[0].has_tasks_or_pull()); + assert_eq!(queue.len(), 1); + } + + #[test] + fn test_default_extras_are_multilevel_for_custom_queue() { + // Custom future pools reuse the tracked runner, so both remote spawns + // and worker-local spawns need multilevel-compatible default extras. + let queue = Arc::new(MockQueue::::default()); + let builder = Builder::new(Config::default(), queue); + let (injector, locals) = builder.build(1); + + // Remote::spawn gets defaults from the injector. + let injector_extras = injector.default_extras(); + assert!(injector_extras.running_time.is_some()); + assert_eq!(injector_extras.current_level(), 0); + + // Local::spawn gets defaults from the local queue handle. + let local_extras = locals[0].default_extras(); + assert!(local_extras.running_time.is_some()); + assert_eq!(local_extras.current_level(), 0); + } + + #[test] + fn test_custom_metrics() { + // A named custom queue reuses the tracked runner metrics. This verifies + // that the custom runner builder wires execution and wait metrics the + // same way as the built-in tracked queues. + let name = "test_custom_metrics"; + let level0_elapsed = MULTILEVEL_LEVEL_ELAPSED + .get_metric_with_label_values(&[name, "0"]) + .unwrap(); + let total_elapsed = MULTILEVEL_LEVEL_ELAPSED + .get_metric_with_label_values(&[name, "total"]) + .unwrap(); + let wait_duration = TASK_WAIT_DURATION + .get_metric_with_label_values(&[name]) + .unwrap(); + let exec_duration = TASK_EXEC_DURATION + .get_metric_with_label_values(&[name]) + .unwrap(); + let poll_duration = TASK_POLL_DURATION + .get_metric_with_label_values(&[name, "0"]) + .unwrap(); + let exec_times = TASK_EXEC_TIMES + .get_metric_with_label_values(&[name]) + .unwrap(); + let level0_elapsed_before = level0_elapsed.get(); + let total_elapsed_before = total_elapsed.get(); + let wait_count_before = wait_duration.get_sample_count(); + let exec_duration_count_before = exec_duration.get_sample_count(); + let exec_duration_sum_before = exec_duration.get_sample_sum(); + let poll_duration_count_before = poll_duration.get_sample_count(); + let poll_duration_sum_before = poll_duration.get_sample_sum(); + let exec_times_count_before = exec_times.get_sample_count(); + let exec_times_sum_before = exec_times.get_sample_sum(); + let queue = Arc::new(MockQueue::default()); + let builder = Builder::new(Config::default().name(Some(name)), queue); + let mut runner = builder.runner_builder(MockRunnerBuilder).build(); + let (remote, mut locals) = build_spawn(builder, Default::default()); + + for i in 0..4 { + remote.spawn(MockTask::with_sleep(i, 35)); + } + while let PopResult::Ready(Pop { task_cell, .. }) = locals[0].pop() { + assert!(runner.handle(&mut locals[0], task_cell)); + } + runner.flush(); + + // Explicitly flush local metrics so the assertions do not depend on + // whether the elapsed-time threshold was crossed before the last task. + assert!(level0_elapsed.get() - level0_elapsed_before > 100_000); + assert!(total_elapsed.get() - total_elapsed_before > 100_000); + assert_eq!(wait_duration.get_sample_count() - wait_count_before, 4); + assert_eq!( + exec_duration.get_sample_count() - exec_duration_count_before, + 4 + ); + assert!(exec_duration.get_sample_sum() - exec_duration_sum_before >= 0.1); + assert_eq!( + poll_duration.get_sample_count() - poll_duration_count_before, + 4 + ); + assert!(poll_duration.get_sample_sum() - poll_duration_sum_before >= 0.1); + assert_eq!(exec_times.get_sample_count() - exec_times_count_before, 4); + assert!(exec_times.get_sample_sum() - exec_times_sum_before >= 3.0); + } + + #[test] + fn test_build_custom_future_pool() { + // Smoke test the public builder path: a custom TaskQueue should be + // enough to build a future pool and execute a spawned future. + let queue = Arc::new(MockQueue::default()); + let mut builder = crate::pool::Builder::new("test-custom-future-pool"); + builder + .min_thread_count(1) + .max_thread_count(1) + .core_thread_count(1); + let pool = builder.build_custom_future_pool(queue); + let (tx, rx) = mpsc::channel(); + + pool.spawn(async move { + tx.send(()).unwrap(); + }); + + rx.recv_timeout(Duration::from_secs(1)).unwrap(); + pool.shutdown(); + } + + #[test] + fn test_custom_callback_pool_supports_worker_local_spawn() { + // This covers the worker-local spawn path. A task running on a worker + // can create a new task through Handle::spawn, which gets default + // extras from the local custom queue handle before the tracked runner + // executes it. + let queue: Arc> = + Arc::new(MockQueue::default()); + let queue_builder = Builder::new(Config::default(), queue); + let runner_builder = queue_builder.runner_builder(crate::pool::CloneRunnerBuilder( + crate::task::callback::Runner::default(), + )); + let mut pool_builder = crate::pool::Builder::new("test-custom-callback-pool"); + pool_builder + .min_thread_count(1) + .max_thread_count(1) + .core_thread_count(1); + let pool = pool_builder.build_with_queue_and_runner(queue_builder.into(), runner_builder); + let (tx, rx) = mpsc::channel(); + let child_tx = tx.clone(); + + pool.spawn(move |handle: &mut crate::task::callback::Handle<'_>| { + tx.send(1).unwrap(); + handle.spawn(move |_: &mut crate::task::callback::Handle<'_>| { + child_tx.send(2).unwrap(); + }); + }); + + assert_eq!(rx.recv_timeout(Duration::from_secs(1)).unwrap(), 1); + assert_eq!(rx.recv_timeout(Duration::from_secs(1)).unwrap(), 2); + pool.shutdown(); + } + + #[test] + fn test_shutdown_drains_pending_custom_queue() { + // This is the real worker shutdown path. The custom queue owns a task + // but reports it as Pending, so the worker must not rely on pop() to + // drain the queue when the pool shuts down. + let (pending_tx, pending_rx) = mpsc::channel(); + let queue = Arc::new(PendingQueue::new(pending_tx)); + let mut builder = crate::pool::Builder::new("test-shutdown-drains-pending-custom-queue"); + builder + .min_thread_count(1) + .max_thread_count(1) + .core_thread_count(1); + let pool = builder.build_custom_future_pool(queue.clone()); + + pool.spawn(async {}); + + // Wait until the worker has observed Pending. This makes the test check + // shutdown-from-pending rather than shutdown-before-the-worker-runs. + pending_rx.recv_timeout(Duration::from_secs(1)).unwrap(); + assert_eq!(queue.len(), 1); + + pool.shutdown(); + assert!(queue.drain_count() >= 1); + assert_eq!(queue.len(), 0); + } +} diff --git a/src/queue/multilevel.rs b/src/queue/multilevel.rs index f0f5288..1f48bd6 100644 --- a/src/queue/multilevel.rs +++ b/src/queue/multilevel.rs @@ -537,15 +537,17 @@ impl LevelManager { let total_tasks = (cur_total_tasks - last_total_tasks) as usize; // adjust the batch size after meeting enough tasks. if total_tasks > ADJUST_LEVEL_STEAL_SIZE_THRESHOLD { - let new_steal_count = if level_0_tasks == 0 { - // level 0 has no tasks, that means the current workloads are all low-priority tasks. - LEVEL_MAX_QUEUE_MAX_STEAL_SIZE - } else { - // by default level0 contains 80% of all tasks, so in the most common case, only - // pop 1 task from level max once, and increases level max batch size when the executed - // tasks are more than level0. - std::cmp::min(total_tasks / level_0_tasks, LEVEL_MAX_QUEUE_MAX_STEAL_SIZE) - }; + // When level 0 has no tasks, the current workloads are all + // low-priority tasks, so use the maximum steal size. + let new_steal_count = total_tasks.checked_div(level_0_tasks).map_or( + LEVEL_MAX_QUEUE_MAX_STEAL_SIZE, + |steal_count| { + // by default level0 contains 80% of all tasks, so in the most common case, only + // pop 1 task from level max once, and increases level max batch size when the executed + // tasks are more than level0. + std::cmp::min(steal_count, LEVEL_MAX_QUEUE_MAX_STEAL_SIZE) + }, + ); self.max_level_queue_steal_size .store(new_steal_count, SeqCst); for (i, c) in self.last_exec_tasks_per_level.iter().enumerate() { @@ -557,12 +559,22 @@ impl LevelManager { } } -pub(super) struct TaskLevelManager { +/// Tracks task running time and assigns multilevel scheduling levels. +/// +/// Custom queues can use this helper to reuse the same level calculation as the +/// built-in multilevel and priority queues before inserting a task. +pub struct TaskLevelManager { task_elapsed_map: TaskElapsedMap, level_time_threshold: [Duration; LEVEL_NUM - 1], } impl TaskLevelManager { + /// Creates a task level manager. + /// + /// `level_time_threshold` defines the accumulated running-time boundary for + /// each level. `cleanup_interval` controls automatic cleanup of old task + /// elapsed records; set it to `None` to disable automatic cleanup and call + /// [`TaskLevelManager::try_cleanup`] manually. pub fn new( level_time_threshold: [Duration; LEVEL_NUM - 1], cleanup_interval: Option, @@ -573,6 +585,11 @@ impl TaskLevelManager { } } + /// Updates the task's current level according to its accumulated running time. + /// + /// If the task has a fixed level, that level is used directly. Otherwise, + /// the manager looks up the task's accumulated running time and compares it + /// with the configured thresholds. pub fn adjust_task_level(&self, task_cell: &mut T) where T: TaskCell, @@ -597,7 +614,11 @@ impl TaskLevelManager { extras.current_level = current_level; } - pub(super) fn try_cleanup(&self) -> Option { + /// Attempts to clean up old task elapsed records. + /// + /// Returns the cleanup time if this call performed cleanup, or `None` if + /// another caller is already cleaning up. + pub fn try_cleanup(&self) -> Option { self.task_elapsed_map.try_cleanup() } @@ -920,7 +941,7 @@ pub(super) fn recent() -> Instant { mod tests { use super::*; use crate::pool::build_spawn; - use crate::queue::Extras; + use crate::queue::{Extras, PopResult}; use std::sync::atomic::AtomicU64; use std::sync::mpsc; @@ -1023,7 +1044,7 @@ mod tests { let (injector, mut locals) = builder.build(1); injector.push(MockTask::new(0, Extras::multilevel_default())); thread::sleep(SLEEP_DUR); - let schedule_time = locals[0].pop().unwrap().schedule_time; + let schedule_time = locals[0].pop().unwrap_ready().schedule_time; assert!(schedule_time.elapsed() >= SLEEP_DUR); } @@ -1120,10 +1141,10 @@ mod tests { injector.push(MockTask::new(i, Extras::multilevel_default())); } let sum: u64 = (0..100) - .map(|_| locals[2].pop().unwrap().task_cell.sleep_ms) + .map(|_| locals[2].pop().unwrap_ready().task_cell.sleep_ms) .sum(); assert_eq!(sum, (0..100).sum()); - assert!(locals.iter_mut().all(|c| c.pop().is_none())); + assert!(locals.iter_mut().all(|c| c.pop().is_empty())); } #[test] @@ -1162,7 +1183,7 @@ mod tests { .map(|mut consumer| { let sum = sum.clone(); thread::spawn(move || { - while let Some(pop) = consumer.pop() { + while let PopResult::Ready(pop) = consumer.pop() { sum.fetch_add(pop.task_cell.sleep_ms, SeqCst); } }) @@ -1183,7 +1204,7 @@ mod tests { let mut runner = runner_builder.build(); remote.spawn(MockTask::new(100, Extras::new_multilevel(1, None))); - if let Some(Pop { task_cell, .. }) = locals[0].pop() { + if let PopResult::Ready(Pop { task_cell, .. }) = locals[0].pop() { assert!(runner.handle(&mut locals[0], task_cell)); } assert!( diff --git a/src/queue/priority.rs b/src/queue/priority.rs index bcba420..43d9100 100644 --- a/src/queue/priority.rs +++ b/src/queue/priority.rs @@ -316,7 +316,7 @@ mod tests { use crate::pool::{build_spawn, Local, Runner, RunnerBuilder}; use crate::queue::{ multilevel::{now, recent}, - Extras, InjectorInner, + Extras, InjectorInner, PopResult, }; use rand::RngCore; #[derive(Debug)] @@ -365,7 +365,7 @@ mod tests { impl TaskPriorityProvider for OrderByIdProvider { fn priority_of(&self, extras: &Extras) -> u64 { - return extras.task_id(); + extras.task_id() } } @@ -663,6 +663,29 @@ mod tests { #[test] fn test_metrics() { let name = "test_priority_metrics"; + let level0_elapsed = MULTILEVEL_LEVEL_ELAPSED + .get_metric_with_label_values(&[name, "0"]) + .unwrap(); + let wait_duration = TASK_WAIT_DURATION + .get_metric_with_label_values(&[name]) + .unwrap(); + let exec_duration = TASK_EXEC_DURATION + .get_metric_with_label_values(&[name]) + .unwrap(); + let poll_duration = TASK_POLL_DURATION + .get_metric_with_label_values(&[name, "0"]) + .unwrap(); + let exec_times = TASK_EXEC_TIMES + .get_metric_with_label_values(&[name]) + .unwrap(); + let level0_elapsed_before = level0_elapsed.get(); + let wait_count_before = wait_duration.get_sample_count(); + let exec_duration_count_before = exec_duration.get_sample_count(); + let exec_duration_sum_before = exec_duration.get_sample_sum(); + let poll_duration_count_before = poll_duration.get_sample_count(); + let poll_duration_sum_before = poll_duration.get_sample_sum(); + let exec_times_count_before = exec_times.get_sample_count(); + let exec_times_sum_before = exec_times.get_sample_sum(); let builder = Builder::new( Config::default().name(Some(name)), Arc::new(OrderByIdProvider), @@ -673,66 +696,26 @@ mod tests { for i in 0..4 { remote.spawn(MockTask::new(35, i)); } - while let Some(Pop { task_cell, .. }) = locals[0].pop() { + while let PopResult::Ready(Pop { task_cell, .. }) = locals[0].pop() { assert!(runner.handle(&mut locals[0], task_cell)); } + runner.flush(); - // we spawn 4 tasks here but the metrics of the last one is not flush, so only check the first 3 here. - assert!( - MULTILEVEL_LEVEL_ELAPSED - .get_metric_with_label_values(&[name, "0"]) - .unwrap() - .get() - > 100_000 - ); - assert!( - TASK_WAIT_DURATION - .get_metric_with_label_values(&[name]) - .unwrap() - .get_sample_count() - >= 3 - ); - assert!( - TASK_EXEC_DURATION - .get_metric_with_label_values(&[name]) - .unwrap() - .get_sample_count() - >= 3 - ); - assert!( - TASK_EXEC_DURATION - .get_metric_with_label_values(&[name]) - .unwrap() - .get_sample_sum() - >= 0.1 - ); - assert!( - TASK_POLL_DURATION - .get_metric_with_label_values(&[name, "0"]) - .unwrap() - .get_sample_count() - >= 3 - ); - assert!( - TASK_POLL_DURATION - .get_metric_with_label_values(&[name, "0"]) - .unwrap() - .get_sample_sum() - >= 0.1 - ); - assert!( - TASK_EXEC_TIMES - .get_metric_with_label_values(&[name]) - .unwrap() - .get_sample_count() - >= 3 + // Explicitly flush local metrics so the assertions do not depend on + // whether the elapsed-time threshold was crossed before the last task. + assert!(level0_elapsed.get() - level0_elapsed_before > 100_000); + assert_eq!(wait_duration.get_sample_count() - wait_count_before, 4); + assert_eq!( + exec_duration.get_sample_count() - exec_duration_count_before, + 4 ); - assert!( - TASK_EXEC_TIMES - .get_metric_with_label_values(&[name]) - .unwrap() - .get_sample_sum() - >= 3.0 + assert!(exec_duration.get_sample_sum() - exec_duration_sum_before >= 0.1); + assert_eq!( + poll_duration.get_sample_count() - poll_duration_count_before, + 4 ); + assert!(poll_duration.get_sample_sum() - poll_duration_sum_before >= 0.1); + assert_eq!(exec_times.get_sample_count() - exec_times_count_before, 4); + assert!(exec_times.get_sample_sum() - exec_times_sum_before >= 3.0); } } diff --git a/src/task/callback.rs b/src/task/callback.rs index 5186d91..2cd4172 100644 --- a/src/task/callback.rs +++ b/src/task/callback.rs @@ -204,7 +204,7 @@ mod tests { ); assert_eq!(rx.recv().unwrap(), 42); assert_eq!(rx.recv().unwrap(), 42); - assert!(locals[0].pop().is_none()); + assert!(locals[0].pop().is_empty()); assert!(rx.recv().is_err()); } @@ -230,7 +230,7 @@ mod tests { ); assert_eq!(rx.recv().unwrap(), 42); assert_eq!(rx.recv().unwrap(), 42); - assert!(locals[0].pop().is_some()); + assert!(locals[0].pop().is_ready()); assert!(rx.recv().is_err()); } } diff --git a/src/task/future.rs b/src/task/future.rs index 9b9b846..1b5d884 100644 --- a/src/task/future.rs +++ b/src/task/future.rs @@ -299,7 +299,7 @@ unsafe fn clone_task(task: *const ()) -> TaskCell { thread_local! { /// Local queue reference that is set before polling and unset after polled. - static LOCAL: Cell<*mut Local> = Cell::new(std::ptr::null_mut()); + static LOCAL: Cell<*mut Local> = const { Cell::new(std::ptr::null_mut()) }; } unsafe fn wake_task(task: Cow<'_, TaskCell>, reschedule: bool) { @@ -369,8 +369,15 @@ impl Runner { } } +#[derive(Clone, Copy, PartialEq)] +enum RescheduleAction { + None, + Reschedule, + YieldToScheduler, +} + thread_local! { - static NEED_RESCHEDULE: Cell = Cell::new(false); + static RESCHEDULE_ACTION: Cell = const { Cell::new(RescheduleAction::None) }; } impl crate::pool::Runner for Runner { @@ -403,11 +410,19 @@ impl crate::pool::Runner for Runner { { Ok(_) => return false, Err(NOTIFIED) => { - let need_reschedule = NEED_RESCHEDULE.with(|r| r.replace(false)); - if (repoll_times >= self.repoll_limit || need_reschedule) + let action = RESCHEDULE_ACTION.with(|r| r.replace(RescheduleAction::None)); + if action == RescheduleAction::YieldToScheduler { + wake_task(Cow::Owned(task_cell), true); + return false; + } + if (repoll_times >= self.repoll_limit + || action == RescheduleAction::Reschedule) && scope.0.need_preempt() { - wake_task(Cow::Owned(task_cell), need_reschedule); + wake_task( + Cow::Owned(task_cell), + action == RescheduleAction::Reschedule, + ); return false; } else { repoll_times += 1; @@ -424,11 +439,30 @@ impl crate::pool::Runner for Runner { /// /// It is only guaranteed to work in yatp. pub async fn reschedule() { - Reschedule { first_poll: true }.await + Reschedule { + first_poll: true, + action: RescheduleAction::Reschedule, + } + .await +} + +/// Gives up a time slice and returns the task to the scheduler. +/// +/// Unlike [`reschedule`], this always queues the current task before polling it +/// again, even when there are no other ready tasks that need to preempt it. +/// +/// It is only guaranteed to work in yatp. +pub async fn yield_to_scheduler() { + Reschedule { + first_poll: true, + action: RescheduleAction::YieldToScheduler, + } + .await } struct Reschedule { first_poll: bool, + action: RescheduleAction, } impl Future for Reschedule { @@ -437,8 +471,8 @@ impl Future for Reschedule { fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll { if self.first_poll { self.first_poll = false; - NEED_RESCHEDULE.with(|r| { - r.set(true); + RESCHEDULE_ACTION.with(|r| { + r.set(self.action); }); cx.waker().wake_by_ref(); Poll::Pending @@ -452,11 +486,18 @@ impl Future for Reschedule { mod tests { use super::*; use crate::pool::{build_spawn, Builder, Remote, Runner as _}; - use crate::queue::QueueType; + use crate::queue::{CustomBuilder, CustomConfig, Pop, PopResult, QueueType, TaskQueue}; - use std::sync::mpsc; + use std::collections::VecDeque; + use std::sync::{ + atomic::{AtomicUsize, Ordering}, + mpsc, Arc, Mutex, + }; use std::{cell::RefCell, thread}; - use std::{rc::Rc, time::Duration}; + use std::{ + rc::Rc, + time::{Duration, Instant}, + }; struct MockLocal { runner: Rc>, @@ -466,7 +507,11 @@ mod tests { impl MockLocal { fn new(runner: Runner) -> MockLocal { - let (remote, locals) = build_spawn(QueueType::SingleLevel, Default::default()); + MockLocal::with_queue(runner, QueueType::SingleLevel) + } + + fn with_queue(runner: Runner, queue_type: QueueType) -> MockLocal { + let (remote, locals) = build_spawn(queue_type, Default::default()); MockLocal { runner: Rc::new(RefCell::new(runner)), remote, @@ -476,13 +521,49 @@ mod tests { /// Run `Runner::handle` once. fn handle_once(&mut self) { - if let Some(t) = self.locals[0].pop() { + if let PopResult::Ready(t) = self.locals[0].pop() { let runner = self.runner.clone(); runner.borrow_mut().handle(&mut self.locals[0], t.task_cell); } } } + #[derive(Default)] + struct NoReadyHintQueue { + tasks: Mutex>, + has_ready_task_calls: AtomicUsize, + } + + impl TaskQueue for NoReadyHintQueue { + fn push(&self, task_cell: TaskCell) { + self.tasks.lock().unwrap().push_back(task_cell); + } + + fn pop(&self) -> PopResult { + self.tasks + .lock() + .unwrap() + .pop_front() + .map(|task_cell| { + PopResult::Ready(Pop { + task_cell, + schedule_time: Instant::now(), + from_local: false, + }) + }) + .unwrap_or(PopResult::Empty) + } + + fn drain(&self) { + self.tasks.lock().unwrap().clear(); + } + + fn has_ready_task(&self) -> bool { + self.has_ready_task_calls.fetch_add(1, Ordering::SeqCst); + false + } + } + impl Default for MockLocal { fn default() -> Self { MockLocal::new(Default::default()) @@ -552,6 +633,7 @@ mod tests { } #[test] + #[allow(clippy::waker_clone_wake)] fn test_waker_clone() { test_wake_impl(|waker| waker.clone().wake()); } @@ -674,6 +756,28 @@ mod tests { assert_eq!(res_rx.recv().unwrap(), 3); } + #[test] + fn test_yield_to_scheduler() { + let queue = Arc::new(NoReadyHintQueue::default()); + let queue_builder = CustomBuilder::new(CustomConfig::default(), queue.clone()); + let mut local = MockLocal::with_queue(Default::default(), queue_builder.into()); + let (res_tx, res_rx) = mpsc::channel(); + + let fut = async move { + res_tx.send(1).unwrap(); + yield_to_scheduler().await; + res_tx.send(2).unwrap(); + }; + local.remote.spawn(fut); + + local.handle_once(); + assert_eq!(res_rx.recv().unwrap(), 1); + assert!(res_rx.try_recv().is_err()); + assert_eq!(queue.has_ready_task_calls.load(Ordering::SeqCst), 0); + local.handle_once(); + assert_eq!(res_rx.recv().unwrap(), 2); + } + #[cfg_attr(not(feature = "failpoints"), ignore)] #[test] fn test_no_preemptive_task() {