use std::collections::VecDeque;
use std::task::{RawWaker, RawWakerVTable, Waker};
thread_local! {
pub(crate) static READY_QUEUE: std::cell::RefCell<VecDeque<u32>> =
const { std::cell::RefCell::new(VecDeque::new()) };
}
pub(crate) const STANDALONE_BIT: u32 = 1 << 31;
pub(crate) fn conn_waker(conn_index: u32) -> Waker {
debug_assert!(
conn_index & STANDALONE_BIT == 0,
"conn_index has standalone bit set"
);
let data = conn_index as usize as *const ();
unsafe { Waker::from_raw(RawWaker::new(data, &VTABLE)) }
}
pub(crate) fn standalone_waker(task_idx: u32) -> Waker {
debug_assert!(
task_idx & STANDALONE_BIT == 0,
"task_idx already has standalone bit"
);
let data = (task_idx | STANDALONE_BIT) as usize as *const ();
unsafe { Waker::from_raw(RawWaker::new(data, &VTABLE)) }
}
const VTABLE: RawWakerVTable = RawWakerVTable::new(clone_fn, wake_fn, wake_by_ref_fn, drop_fn);
unsafe fn clone_fn(data: *const ()) -> RawWaker {
RawWaker::new(data, &VTABLE)
}
unsafe fn wake_fn(data: *const ()) {
unsafe { wake_by_ref_fn(data) };
}
unsafe fn wake_by_ref_fn(data: *const ()) {
let conn_index = data as usize as u32;
READY_QUEUE.with(|q| {
q.borrow_mut().push_back(conn_index);
});
}
unsafe fn drop_fn(_data: *const ()) {
}
pub(crate) fn drain_ready_queue(buf: &mut VecDeque<u32>) {
READY_QUEUE.with(|q| {
buf.append(&mut q.borrow_mut());
});
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn waker_pushes_to_ready_queue() {
READY_QUEUE.with(|q| q.borrow_mut().clear());
let waker = conn_waker(42);
waker.wake_by_ref();
waker.wake_by_ref();
let mut buf = VecDeque::new();
drain_ready_queue(&mut buf);
assert_eq!(buf.len(), 2);
assert_eq!(buf[0], 42);
assert_eq!(buf[1], 42);
}
#[test]
fn waker_clone_works() {
READY_QUEUE.with(|q| q.borrow_mut().clear());
let waker = conn_waker(7);
let cloned = waker.clone();
waker.wake_by_ref();
cloned.wake();
let mut buf = VecDeque::new();
drain_ready_queue(&mut buf);
assert_eq!(buf.len(), 2);
assert_eq!(buf[0], 7);
assert_eq!(buf[1], 7);
}
#[test]
fn drain_empty_queue() {
READY_QUEUE.with(|q| q.borrow_mut().clear());
let mut buf = VecDeque::new();
drain_ready_queue(&mut buf);
assert!(buf.is_empty());
}
}