use super::{
idle::Idle,
mailbox::{Mailbox, Message},
queue::{self, CAPACITY, Overflow, Push, STEALABLE},
runtime::{Panic, Role, Shared, Worker},
task::{Ready, Runnable},
};
use crate::utils::Panicker;
use crossbeam_utils::CachePadded;
use std::{
collections::VecDeque,
future::poll_fn,
iter, mem,
panic::{AssertUnwindSafe, catch_unwind, resume_unwind},
sync::{Arc, mpsc},
task::Poll,
thread::{self, JoinHandle},
};
cfg_if::cfg_if! {
if #[cfg(feature = "loom")] {
use loom::sync::{
Condvar, Mutex, MutexGuard,
atomic::{AtomicUsize, Ordering, fence},
};
} else {
use commonware_utils::sync::{Condvar, Mutex, MutexGuard};
use std::sync::atomic::{AtomicUsize, Ordering, fence};
}
}
pub const HALF: usize = CAPACITY / 2;
pub const MAX_WORKERS: usize = 64;
pub type LocalQueue = queue::Queue<Runnable>;
struct GlobalQueue {
open: bool,
runnables: VecDeque<Runnable>,
}
struct Closing {
closed: bool,
stopped: bool,
active: usize,
}
pub struct Pool {
idle: CachePadded<Idle>,
mailboxes: Box<[Arc<Mailbox>]>,
queues: Box<[LocalQueue]>,
global_len: CachePadded<AtomicUsize>,
global: CachePadded<Mutex<GlobalQueue>>,
closing: Mutex<Closing>,
progress: Condvar,
}
impl Pool {
pub fn new(mailboxes: Vec<Arc<Mailbox>>) -> Self {
let workers = mailboxes.len();
assert!(
(1..=MAX_WORKERS).contains(&workers),
"an io_uring pool must have 1 to {MAX_WORKERS} workers"
);
let queue = || {
if workers == 1 {
LocalQueue::exclusive()
} else {
LocalQueue::shared()
}
};
Self {
idle: CachePadded::new(Idle::new()),
mailboxes: mailboxes.into(),
queues: (0..workers).map(|_| queue()).collect(),
global_len: CachePadded::new(AtomicUsize::new(0)),
global: CachePadded::new(Mutex::new(GlobalQueue {
open: true,
runnables: VecDeque::new(),
})),
closing: Mutex::new(Closing {
closed: false,
stopped: false,
active: 0,
}),
progress: Condvar::new(),
}
}
pub fn workers(&self) -> usize {
self.mailboxes.len()
}
pub fn idle(&self) -> &Idle {
&self.idle
}
pub fn mailbox(&self, index: u32) -> &Arc<Mailbox> {
&self.mailboxes[index as usize]
}
pub fn queue(&self, index: u32) -> &LocalQueue {
&self.queues[index as usize]
}
fn global(&self) -> MutexGuard<'_, GlobalQueue> {
cfg_if::cfg_if! {
if #[cfg(feature = "loom")] {
let queue = self.global.lock().unwrap();
} else {
let queue = self.global.lock();
}
}
queue
}
fn closing(&self) -> MutexGuard<'_, Closing> {
cfg_if::cfg_if! {
if #[cfg(feature = "loom")] {
let closing = self.closing.lock().unwrap();
} else {
let closing = self.closing.lock();
}
}
closing
}
fn wait<'a>(&self, closing: MutexGuard<'a, Closing>) -> MutexGuard<'a, Closing> {
cfg_if::cfg_if! {
if #[cfg(feature = "loom")] {
self.progress.wait(closing).unwrap()
} else {
let mut closing = closing;
self.progress.wait(&mut closing);
closing
}
}
}
pub fn push(&self, runnable: Runnable) -> Result<(), Runnable> {
self.push_batch(iter::once(runnable))
.map_err(|mut refused| refused.next().unwrap())
}
fn push_batch<I: Iterator<Item = Runnable>>(&self, runnables: I) -> Result<(), I> {
{
let mut queue = self.global();
if !queue.open {
return Err(runnables);
}
queue.runnables.extend(runnables);
self.global_len
.store(queue.runnables.len(), Ordering::Release);
}
fence(Ordering::SeqCst);
self.idle.notify_global(&self.mailboxes);
Ok(())
}
pub fn has_global(&self) -> bool {
self.global_len.load(Ordering::Acquire) != 0
}
#[must_use]
pub fn take(&self, ready: &mut Ready) -> Option<Runnable> {
if !self.has_global() {
return None;
}
let mut queue = self.global();
let len = queue.runnables.len();
let share = (len / self.workers() + 1).min(len);
let admitted = ready.room().min(HALF.saturating_sub(ready.len()));
let count = share.min(admitted + 1);
let mut taken = queue.runnables.drain(..count);
let first = taken.next();
ready.push_batch(taken);
self.global_len
.store(queue.runnables.len(), Ordering::Release);
first
}
#[inline]
pub fn enqueue(&self, ready: &mut Ready, runnable: Runnable) {
match ready.push(runnable) {
Push::Pushed(len) => {
if len >= STEALABLE && self.workers() > 1 {
self.kick();
}
}
Push::Overflow(batch, runnable) => self.overflow(batch, runnable),
Push::Claimed(runnable) => {
if let Err(runnable) = self.push(runnable) {
runnable.discard();
}
}
}
}
#[inline(never)]
fn overflow(&self, batch: Overflow<'_, Runnable>, runnable: Runnable) {
#[cfg(test)]
super::runtime::tests::before_overflow(self);
if let Err(refused) = self.push_batch(batch.chain(iter::once(runnable))) {
for runnable in refused {
runnable.discard();
}
}
}
#[inline]
pub fn steal(&self, index: u32, start: usize, ready: &mut Ready) -> bool {
let workers = self.workers();
(0..workers)
.map(|offset| (start + offset) % workers)
.filter(|&victim| victim != index as usize)
.any(|victim| self.queues[victim].steal_into(ready) > 0)
}
#[inline]
pub fn kick(&self) {
fence(Ordering::SeqCst);
self.idle.notify(&self.mailboxes);
}
#[inline]
pub fn has_stealable(&self, index: u32) -> bool {
self.queues
.iter()
.enumerate()
.any(|(other, queue)| other != index as usize && queue.is_stealable())
}
pub fn enter(&self) {
self.closing().active += 1;
}
pub fn close(&self) {
let runnables = {
let mut queue = self.global();
queue.open = false;
self.global_len.store(0, Ordering::Release);
mem::take(&mut queue.runnables)
};
for runnable in runnables {
runnable.discard();
}
self.closing().closed = true;
self.progress.notify_all();
}
pub fn stop(&self) {
{
let mut closing = self.closing();
if closing.stopped {
return;
}
closing.stopped = true;
}
for mailbox in self.mailboxes.iter().skip(1) {
let _ = mailbox.send(Message::WakeRoot);
}
}
fn is_stopped(&self) -> bool {
self.closing().stopped
}
pub fn wait_closed(&self) {
let mut closing = self.closing();
while !closing.closed {
closing = self.wait(closing);
}
}
pub fn finish(&self) {
let mut closing = self.closing();
closing.active -= 1;
if closing.active == 0 {
self.progress.notify_all();
return;
}
while closing.active != 0 {
closing = self.wait(closing);
}
}
}
#[derive(Default)]
pub struct Threads {
handles: Vec<JoinHandle<()>>,
}
impl Threads {
pub fn start(&mut self, shared: &Arc<Shared>, failures: &Panicker) {
let count = shared.pool.workers();
self.handles.reserve(count - 1);
for index in 1..count {
#[cfg(test)]
let _fault = super::runtime::tests::before_pool_worker(shared, index);
let (ready, started) = mpsc::channel();
let shared = shared.clone();
let failures = failures.clone();
let handle = thread::Builder::new()
.name(format!("iouring-pool-{index}"))
.stack_size(shared.cfg.thread_stack_size())
.spawn(move || run(shared, index as u32, ready, failures))
.expect("pool worker thread must spawn");
self.handles.push(handle);
match started
.recv()
.expect("pool worker must report its startup before exiting")
{
Ok(()) => {}
Err(panic) => resume_unwind(panic),
}
}
}
pub fn join(&mut self) {
let mut first: Option<Panic> = None;
for thread in self.handles.drain(..) {
if let Err(panic) = thread.join() {
if first.is_none() {
first = Some(panic);
} else {
mem::forget(panic);
}
}
}
if let Some(panic) = first {
resume_unwind(panic);
}
}
}
impl Drop for Threads {
fn drop(&mut self) {
if let Err(panic) = catch_unwind(AssertUnwindSafe(|| self.join())) {
mem::forget(panic);
}
}
}
fn run(
shared: Arc<Shared>,
index: u32,
ready: mpsc::Sender<Result<(), Panic>>,
failures: Panicker,
) {
let pool = shared.pool.clone();
let mut ready = Some(ready);
let result = catch_unwind(AssertUnwindSafe(|| {
let (mut worker, _) = Worker::run(
shared,
Role::Pool(index),
|| {
let _ = ready.take().unwrap().send(Ok(()));
poll_fn(|_| {
if pool.is_stopped() {
Poll::Ready(())
} else {
Poll::Pending
}
})
},
None,
)?;
if let Some(panic) = worker.take_panic() {
failures.notify(panic);
let _ = pool.mailbox(0).send(Message::WakeRoot);
}
worker.cleanup();
Ok::<_, Panic>(worker.take_panic())
}));
let panic = match result {
Ok(Ok(None)) => return,
Ok(Ok(Some(panic))) | Ok(Err(panic)) | Err(panic) => panic,
};
match ready {
Some(ready) => {
if let Err(error) = ready.send(Err(panic)) {
mem::forget(error);
}
}
None => {
failures.notify(panic);
let _ = pool.mailbox(0).send(Message::WakeRoot);
}
}
}
#[cfg(test)]
pub mod tests {
use super::*;
use crate::iouring::{
idle::Searcher,
task::{
Task,
tests::{queue, refs, task_of},
},
tasks::Tasks,
};
use std::{
future::pending,
panic::{AssertUnwindSafe, catch_unwind},
sync::Weak,
thread,
time::Duration,
};
const TEST_TIMEOUT: Duration = Duration::from_secs(10);
const BLOCKED: Duration = Duration::from_millis(20);
impl Pool {
#[must_use]
pub fn pop(&self) -> Option<Runnable> {
if !self.has_global() {
return None;
}
let mut queue = self.global();
let runnable = queue.runnables.pop_front();
self.global_len
.store(queue.runnables.len(), Ordering::Release);
runnable
}
pub fn is_closed(&self) -> bool {
!self.global().open
}
pub fn mailboxes(&self) -> &[Arc<Mailbox>] {
&self.mailboxes
}
}
pub fn pool(workers: usize) -> Pool {
Pool::new(
(0..workers)
.map(|_| Arc::new(Mailbox::new().unwrap()))
.collect(),
)
}
pub fn ready(pool: &Pool, index: u32) -> Ready {
pool.queue(index).owner()
}
pub fn task(set: &Tasks) -> (Task, Runnable) {
Task::new(pending::<()>(), set, Weak::new())
}
pub fn push(pool: &Pool, set: &Tasks) -> Task {
let (task, runnable) = task(set);
assert!(pool.push(runnable).is_ok());
task
}
#[test]
fn test_take_moves_a_bounded_share_of_the_global_queue() {
for (workers, queued, taken) in [(4, 8, 3), (1, HALF + 2, HALF + 1)] {
let pool = pool(workers);
let set = Tasks::new(1);
let tasks: Vec<Task> = (0..queued).map(|_| push(&pool, &set)).collect();
let mut ready = ready(&pool, 0);
let mut moved = vec![pool.take(&mut ready).unwrap()];
while let Some(runnable) = ready.pop() {
moved.push(runnable);
}
assert_eq!(moved.len(), taken);
for (runnable, task) in moved.iter().zip(&tasks) {
assert_eq!(task_of(runnable).as_ptr(), task.as_ptr());
}
for runnable in moved {
runnable.discard();
}
pool.close();
for task in tasks {
task.clear();
}
}
}
#[test]
fn test_take_fills_the_run_queue_at_most_to_half() {
for held in [0, 100, HALF - 1, HALF, 200, CAPACITY - 1] {
let pool = pool(1);
let set = Tasks::new(1);
let mut ready = ready(&pool, 0);
let mut tasks = Vec::new();
for _ in 0..held {
let (task, runnable) = task(&set);
queue(&mut ready, runnable);
tasks.push(task);
}
tasks.extend((0..300).map(|_| push(&pool, &set)));
let first = pool.take(&mut ready).unwrap();
let moved = ready.len() - held + 1;
assert_eq!(
moved,
HALF.saturating_sub(held) + 1,
"a take into a run queue holding {held} must fill it to half and return one more"
);
first.discard();
ready.discard();
while let Some(runnable) = pool.pop() {
runnable.discard();
}
for task in tasks {
task.clear();
}
}
}
#[test]
fn test_take_leaves_the_slots_of_a_held_claim_alone() {
let pool = pool(2);
let set = Tasks::new(1);
let mut ready = ready(&pool, 0);
let mut tasks = Vec::new();
for _ in 0..CAPACITY {
let (task, runnable) = task(&set);
queue(&mut ready, runnable);
tasks.push(task);
}
let claim = pool.queue(0).hold_claim(HALF as u32);
while let Some(runnable) = ready.pop() {
runnable.discard();
}
assert_eq!(ready.room(), 0);
tasks.extend((0..200).map(|_| push(&pool, &set)));
pool.take(&mut ready).unwrap().discard();
assert!(ready.is_empty());
for runnable in pool.queue(0).finish_claim(claim) {
runnable.discard();
}
pool.take(&mut ready).unwrap().discard();
assert_eq!(ready.len() + 1, 199 / 2 + 1);
ready.discard();
while let Some(runnable) = pool.pop() {
runnable.discard();
}
for task in tasks {
task.clear();
}
}
#[test]
fn test_closed_queue_refuses_pushes_and_discards_its_runnables() {
let pool = pool(1);
let set = Tasks::new(1);
let queued = push(&pool, &set);
assert_eq!(refs(&queued), 2);
pool.close();
assert!(!pool.has_global());
assert_eq!(refs(&queued), 1);
let (late, runnable) = task(&set);
pool.push(runnable).unwrap_err().discard();
assert_eq!(refs(&late), 1);
queued.clear();
late.clear();
}
#[test]
fn test_full_run_queue_moves_its_newest_half_to_the_global_queue() {
let pool = pool(1);
let set = Tasks::new(1);
let mut ready = ready(&pool, 0);
let mut tasks = Vec::new();
for _ in 0..=CAPACITY {
let (task, runnable) = task(&set);
pool.enqueue(&mut ready, runnable);
tasks.push(task);
}
assert_eq!(ready.len(), HALF);
for task in &tasks[..HALF] {
let runnable = ready.pop().unwrap();
assert_eq!(task_of(&runnable).as_ptr(), task.as_ptr());
runnable.discard();
}
for task in &tasks[HALF..] {
let runnable = pool.pop().unwrap();
assert_eq!(task_of(&runnable).as_ptr(), task.as_ptr());
runnable.discard();
}
assert!(pool.pop().is_none());
pool.close();
for _ in 0..=CAPACITY {
let (task, runnable) = task(&set);
pool.enqueue(&mut ready, runnable);
tasks.push(task);
}
assert_eq!(ready.len(), HALF);
for task in &tasks[tasks.len() - HALF - 1..] {
assert_eq!(refs(task), 1);
}
ready.discard();
for task in tasks {
task.clear();
}
}
#[test]
fn test_full_run_queue_under_a_claim_moves_only_the_new_runnable() {
for open in [true, false] {
let pool = pool(2);
let set = Tasks::new(1);
let mut ready = ready(&pool, 0);
let mut tasks = Vec::new();
for _ in 0..CAPACITY {
let (task, runnable) = task(&set);
queue(&mut ready, runnable);
tasks.push(task);
}
let claim = pool.queue(0).hold_claim(HALF as u32);
if !open {
pool.close();
}
let (pushed, runnable) = task(&set);
pool.enqueue(&mut ready, runnable);
assert_eq!(ready.len(), CAPACITY - HALF);
if open {
let runnable = pool.pop().unwrap();
assert_eq!(task_of(&runnable).as_ptr(), pushed.as_ptr());
runnable.discard();
} else {
assert_eq!(refs(&pushed), 1);
}
assert!(pool.pop().is_none());
for runnable in pool.queue(0).finish_claim(claim) {
runnable.discard();
}
ready.discard();
tasks.push(pushed);
for task in tasks {
task.clear();
}
}
}
#[test]
fn test_search_steals_half_of_another_run_queue() {
let pool = pool(3);
let set = Tasks::new(1);
let (mut zero, mut one, mut two) = (ready(&pool, 0), ready(&pool, 1), ready(&pool, 2));
let mut tasks = Vec::new();
for _ in 0..5 {
let (task, runnable) = task(&set);
queue(&mut two, runnable);
tasks.push(task);
}
let mut searcher = Searcher::new(&pool, 0);
assert!(searcher.search(&mut zero));
for task in &tasks[..2] {
let runnable = zero.pop().unwrap();
assert_eq!(task_of(&runnable).as_ptr(), task.as_ptr());
runnable.discard();
}
assert!(zero.is_empty());
assert_eq!(two.len(), 3);
let mut other = Searcher::new(&pool, 1);
assert!(other.search(&mut one));
let second = one.pop().unwrap();
assert_eq!(task_of(&second).as_ptr(), tasks[2].as_ptr());
second.discard();
two.pop().unwrap().discard();
assert!(!other.search(&mut one));
assert_eq!(two.len(), 1);
searcher.leave_search();
other.leave_search();
two.discard();
for task in tasks {
task.clear();
}
}
#[test]
fn test_only_a_one_worker_pool_has_an_exclusive_run_queue() {
for (workers, exclusive) in [(1, true), (2, false)] {
let pool = pool(workers);
let set = Tasks::new(1);
let mut ready = ready(&pool, 0);
let mut tasks = Vec::new();
for _ in 0..2 {
let (task, runnable) = task(&set);
queue(&mut ready, runnable);
tasks.push(task);
}
let thief = LocalQueue::shared();
let mut into = thief.owner();
let stolen = catch_unwind(AssertUnwindSafe(|| pool.queue(0).steal_into(&mut into)));
assert_eq!(stolen.is_err(), exclusive);
ready.discard();
into.discard();
for task in tasks {
task.clear();
}
}
}
#[test]
fn test_stop_wakes_every_other_root_once() {
let pool = pool(3);
assert!(!pool.is_stopped());
pool.stop();
pool.stop();
assert!(pool.is_stopped());
assert!(!pool.mailbox(0).waker.pending(0));
for index in 1..3 {
let mut messages = Vec::new();
assert!(pool.mailbox(index).take(&mut messages));
assert!(matches!(messages.as_slice(), [Message::WakeRoot]));
}
}
#[test]
fn test_workers_wait_for_the_close_and_for_every_finish() {
let pool = Arc::new(pool(2));
pool.enter();
pool.enter();
let (progress, observed) = mpsc::channel();
let other = thread::spawn({
let pool = pool.clone();
move || {
pool.wait_closed();
progress.send("closed").unwrap();
pool.finish();
progress.send("finished").unwrap();
}
});
assert!(observed.recv_timeout(BLOCKED).is_err());
pool.close();
assert_eq!(observed.recv_timeout(TEST_TIMEOUT).unwrap(), "closed");
assert!(observed.recv_timeout(BLOCKED).is_err());
pool.finish();
assert_eq!(observed.recv_timeout(TEST_TIMEOUT).unwrap(), "finished");
other.join().unwrap();
}
}