use std::{
cmp::Ordering,
collections::BTreeSet,
task::{Poll, Waker},
};
use slab::Slab;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Priority {
track: u8,
subscribe: u64,
group: u64,
}
impl Priority {
pub fn new(track: u8, subscribe: u64, group: u64) -> Self {
Self {
track,
subscribe,
group,
}
}
}
impl Ord for Priority {
fn cmp(&self, other: &Self) -> Ordering {
other
.track
.cmp(&self.track)
.then(self.subscribe.cmp(&other.subscribe))
.then(other.group.cmp(&self.group))
}
}
impl PartialOrd for Priority {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}
#[derive(Debug, Clone)]
struct PriorityItem {
id: usize,
priority: Priority,
}
impl PartialEq for PriorityItem {
fn eq(&self, other: &Self) -> bool {
self.priority == other.priority && self.id == other.id
}
}
impl Eq for PriorityItem {}
impl PartialOrd for PriorityItem {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}
impl Ord for PriorityItem {
fn cmp(&self, other: &Self) -> Ordering {
self.priority.cmp(&other.priority).then(self.id.cmp(&other.id))
}
}
#[derive(Clone)]
pub struct PriorityQueue {
state: kio::Lock<PriorityState>,
}
impl Default for PriorityQueue {
fn default() -> Self {
Self {
state: kio::Lock::new(PriorityState::default()),
}
}
}
impl PriorityQueue {
pub fn insert(&self, priority: Priority) -> PriorityHandle {
self.lock().insert(priority, self.clone())
}
fn lock(&self) -> Guard<'_> {
Guard {
lock: &self.state,
state: Some(self.state.lock()),
}
}
}
struct Guard<'a> {
lock: &'a kio::Lock<PriorityState>,
state: Option<kio::LockGuard<'a, PriorityState>>,
}
impl Drop for Guard<'_> {
fn drop(&mut self) {
let mut state = self.state.take().expect("guard already dropped");
if state.pending.is_empty() {
return;
}
let mut pending = std::mem::take(&mut state.pending);
drop(state);
for waker in pending.drain(..) {
waker.wake();
}
self.lock.lock().pending = pending;
}
}
impl std::ops::Deref for Guard<'_> {
type Target = PriorityState;
fn deref(&self) -> &Self::Target {
self.state.as_ref().expect("guard already dropped")
}
}
impl std::ops::DerefMut for Guard<'_> {
fn deref_mut(&mut self) -> &mut Self::Target {
self.state.as_mut().expect("guard already dropped")
}
}
const MAX_VEC_SIZE: usize = 255;
enum Location {
Vec(usize), Overflow, }
struct PriorityEntry {
location: Location,
priority: Priority,
rank: u8,
waker: Option<Waker>,
}
#[derive(Default)]
struct PriorityState {
vec: Vec<PriorityItem>,
overflow: BTreeSet<PriorityItem>,
entries: Slab<PriorityEntry>,
pending: Vec<Waker>,
}
impl PriorityState {
pub fn insert(&mut self, priority: Priority, myself: PriorityQueue) -> PriorityHandle {
let id = self.entries.insert(PriorityEntry {
location: Location::Overflow,
priority,
rank: u8::MAX,
waker: None,
});
self.place(PriorityItem { id, priority });
PriorityHandle {
id,
track: priority.track,
seen: u8::MAX,
queue: myself,
}
}
fn update_indices_from(&mut self, start: usize) {
for (idx, item) in self.vec.iter().enumerate().skip(start) {
Self::update_location(&mut self.entries, &mut self.pending, item.id, Location::Vec(idx));
}
}
fn update_location(entries: &mut Slab<PriorityEntry>, pending: &mut Vec<Waker>, id: usize, location: Location) {
let entry = entries.get_mut(id).expect("item not in entries");
entry.location = location;
let rank = Self::rank_of(&entry.location);
if entry.rank != rank {
entry.rank = rank;
pending.extend(entry.waker.take());
}
}
fn rank_of(location: &Location) -> u8 {
match location {
Location::Vec(idx) => (*idx).try_into().unwrap_or(u8::MAX),
Location::Overflow => u8::MAX,
}
}
fn rank(&self, id: usize) -> u8 {
Self::rank_of(&self.entries.get(id).expect("item not in entries").location)
}
fn place(&mut self, item: PriorityItem) {
let id = item.id;
self.entries[id].priority = item.priority;
if self.vec.len() < MAX_VEC_SIZE {
if let Some(top) = self.overflow.first()
&& *top < item
{
let promoted = self.overflow.pop_first().unwrap();
assert!(self.overflow.insert(item));
Self::update_location(&mut self.entries, &mut self.pending, id, Location::Overflow);
let insert_pos = self.vec.binary_search(&promoted).unwrap_or_else(|pos| pos);
let promoted_id = promoted.id;
self.vec.insert(insert_pos, promoted);
Self::update_location(
&mut self.entries,
&mut self.pending,
promoted_id,
Location::Vec(insert_pos),
);
self.update_indices_from(insert_pos + 1);
return;
}
let insert_pos = self.vec.binary_search(&item).unwrap_or_else(|pos| pos);
self.vec.insert(insert_pos, item);
Self::update_location(&mut self.entries, &mut self.pending, id, Location::Vec(insert_pos));
self.update_indices_from(insert_pos + 1);
return;
}
let lowest_in_vec = self.vec.last().unwrap();
if item > *lowest_in_vec {
assert!(self.overflow.insert(item));
Self::update_location(&mut self.entries, &mut self.pending, id, Location::Overflow);
return;
}
let removed = self.vec.pop().unwrap();
let removed_id = removed.id;
assert!(self.overflow.insert(removed));
Self::update_location(&mut self.entries, &mut self.pending, removed_id, Location::Overflow);
let insert_pos = self.vec.binary_search(&item).unwrap_or_else(|pos| pos);
self.vec.insert(insert_pos, item);
Self::update_location(&mut self.entries, &mut self.pending, id, Location::Vec(insert_pos));
self.update_indices_from(insert_pos + 1);
}
fn extract(&mut self, id: usize) -> PriorityItem {
let location = &self.entries.get(id).expect("item not in entries").location;
match location {
Location::Vec(idx) => {
let idx = *idx;
let item = self.vec.remove(idx);
self.update_indices_from(idx);
item
}
Location::Overflow => {
let priority = self.entries[id].priority;
self.overflow
.take(&PriorityItem { id, priority })
.expect("item not found in overflow set")
}
}
}
fn set_track(&mut self, id: usize, track: u8) {
let mut item = self.extract(id);
item.priority.track = track;
self.place(item);
}
fn remove(&mut self, id: usize) {
let was_in_vec = matches!(
self.entries.get(id).map(|entry| &entry.location),
Some(Location::Vec(_))
);
self.extract(id);
self.entries.remove(id);
if was_in_vec && let Some(overflow_item) = self.overflow.pop_first() {
let overflow_id = overflow_item.id;
self.vec.push(overflow_item);
let tail = self.vec.len() - 1;
Self::update_location(&mut self.entries, &mut self.pending, overflow_id, Location::Vec(tail));
}
}
}
pub struct PriorityHandle {
id: usize,
track: u8,
seen: u8,
queue: PriorityQueue,
}
impl Drop for PriorityHandle {
fn drop(&mut self) {
self.queue.lock().remove(self.id);
}
}
impl PriorityHandle {
pub fn current(&mut self) -> u8 {
let rank = self.queue.state.lock().rank(self.id);
self.seen = rank;
rank
}
pub fn send_order_of(rank: u8) -> u8 {
u8::MAX - rank
}
pub fn send_order(&mut self) -> u8 {
Self::send_order_of(self.current())
}
pub fn poll_next(&mut self, waiter: &kio::Waiter) -> Poll<u8> {
let mut state = self.queue.state.lock();
let entry = state.entries.get_mut(self.id).expect("item not in entries");
if entry.rank == self.seen {
let waker = waiter.waker();
if !entry.waker.as_ref().is_some_and(|parked| parked.will_wake(waker)) {
entry.waker = Some(waker.clone());
}
return Poll::Pending;
}
let rank = entry.rank;
drop(state);
self.seen = rank;
Poll::Ready(rank)
}
#[cfg(test)]
pub async fn next(&mut self) -> u8 {
kio::wait(|waiter| self.poll_next(waiter)).await
}
pub fn set_track(&mut self, new_track: u8) -> u8 {
if self.track == new_track {
return self.current();
}
self.track = new_track;
let rank = {
let mut state = self.queue.lock();
state.set_track(self.id, new_track);
state.rank(self.id)
};
self.seen = rank;
rank
}
}
#[cfg(test)]
mod tests {
use std::sync::{
Arc,
atomic::{AtomicUsize, Ordering as AtomicOrdering},
};
use std::task::Wake;
use super::*;
fn live(track: u8, group: u64) -> Priority {
Priority::new(track, 0, group)
}
struct WakeCount(AtomicUsize);
impl Wake for WakeCount {
fn wake(self: Arc<Self>) {
self.0.fetch_add(1, AtomicOrdering::Relaxed);
}
}
struct Parked {
handles: Vec<PriorityHandle>,
wakes: Vec<Arc<WakeCount>>,
waiters: Vec<kio::Waiter>,
}
impl Parked {
fn new(handles: Vec<PriorityHandle>) -> Self {
let wakes: Vec<_> = handles
.iter()
.map(|_| Arc::new(WakeCount(AtomicUsize::new(0))))
.collect();
let waiters = wakes
.iter()
.map(|wake| kio::Waiter::new(std::task::Waker::from(wake.clone())))
.collect();
let mut parked = Self {
handles,
wakes,
waiters,
};
parked.park();
parked
}
fn park(&mut self) {
for (handle, waiter) in self.handles.iter_mut().zip(&self.waiters) {
handle.current();
assert!(handle.poll_next(waiter).is_pending());
}
}
fn counts(&self) -> Vec<usize> {
self.wakes.iter().map(|w| w.0.load(AtomicOrdering::Relaxed)).collect()
}
}
#[test]
fn a_reorder_wakes_only_the_ranks_that_moved() {
let queue = PriorityQueue::default();
let parked = Parked::new((0..64).map(|group| queue.insert(live(100, group))).collect());
let tail = queue.insert(live(50, 0));
assert_eq!(parked.counts(), vec![0; 64], "a tail insert moved no other rank");
let front = queue.insert(live(u8::MAX, u64::MAX));
assert_eq!(parked.counts(), vec![1; 64], "a front insert moved every other rank");
drop((tail, front));
}
#[test]
fn an_overflow_removal_wakes_nobody() {
let queue = PriorityQueue::default();
let fillers = (0..MAX_VEC_SIZE as u64).map(|group| queue.insert(live(200, group)));
let overflow = (0..2).map(|group| queue.insert(live(100, group)));
let mut parked = Parked::new(fillers.chain(overflow).collect());
let doomed = parked.handles.pop().expect("overflow handle");
assert_eq!(doomed.seen, u8::MAX, "the dropped entry was in overflow");
drop(doomed);
let survivors = &parked.counts()[..parked.handles.len()];
assert!(
survivors.iter().all(|&count| count == 0),
"an overflow removal woke {survivors:?}"
);
}
#[test]
fn test_single_item() {
let queue = PriorityQueue::default();
let mut handle = queue.insert(live(100, 5));
assert_eq!(handle.current(), 0); }
#[test]
fn test_send_order_inverts_rank() {
let queue = PriorityQueue::default();
let mut top = queue.insert(live(200, 0));
let mut low = queue.insert(live(100, 0));
assert_eq!(top.current(), 0);
assert_eq!(top.send_order(), 255, "most urgent rank gets the highest send order");
assert_eq!(low.current(), 1);
assert_eq!(low.send_order(), 254);
}
#[test]
fn test_track_priority_ordering() {
let queue = PriorityQueue::default();
let mut low = queue.insert(live(50, 0));
let mut high = queue.insert(live(255, 0));
let mut mid = queue.insert(live(100, 0));
assert_eq!(high.current(), 0); assert_eq!(mid.current(), 1); assert_eq!(low.current(), 2); }
#[test]
fn test_group_priority_on_same_track() {
let queue = PriorityQueue::default();
let mut group10 = queue.insert(live(100, 10));
let mut group5 = queue.insert(live(100, 5));
let mut group1 = queue.insert(live(100, 1));
assert_eq!(group10.current(), 0);
assert_eq!(group5.current(), 1);
assert_eq!(group1.current(), 2);
}
#[test]
fn test_ord_is_total() {
let mixed = [
Priority::new(100, 7, 1),
Priority::new(100, 7, 2),
Priority::new(100, 8, 1),
Priority::new(200, 7, 1),
];
for a in mixed {
for b in mixed {
assert_eq!(a.cmp(&b), b.cmp(&a).reverse(), "{a:?} vs {b:?}");
assert_eq!(a.cmp(&b) == Ordering::Equal, a == b, "{a:?} vs {b:?}");
for c in mixed {
if a.cmp(&b) != Ordering::Greater && b.cmp(&c) != Ordering::Greater {
assert_ne!(a.cmp(&c), Ordering::Greater, "{a:?} <= {b:?} <= {c:?}");
}
}
}
}
}
#[test]
fn test_track_priority_overrides_group() {
let queue = PriorityQueue::default();
let mut low_track_high_group = queue.insert(live(50, 1000));
let mut high_track_low_group = queue.insert(live(255, 1));
assert_eq!(high_track_low_group.current(), 0);
assert_eq!(low_track_high_group.current(), 1);
}
#[test]
fn test_removal_on_drop() {
let queue = PriorityQueue::default();
let mut first = queue.insert(live(255, 0));
let mut second = queue.insert(live(100, 0));
let mut third = queue.insert(live(50, 0));
assert_eq!(first.current(), 0);
assert_eq!(second.current(), 1);
assert_eq!(third.current(), 2);
drop(second);
assert_eq!(first.current(), 0);
assert_eq!(third.current(), 1);
}
#[test]
fn test_removal_of_highest_priority() {
let queue = PriorityQueue::default();
let mut first = queue.insert(live(255, 0));
let mut second = queue.insert(live(100, 0));
assert_eq!(first.current(), 0);
assert_eq!(second.current(), 1);
drop(first);
assert_eq!(second.current(), 0);
}
#[test]
fn test_removal_of_lowest_priority() {
let queue = PriorityQueue::default();
let mut first = queue.insert(live(255, 0));
let mut second = queue.insert(live(100, 0));
assert_eq!(first.current(), 0);
assert_eq!(second.current(), 1);
drop(second);
assert_eq!(first.current(), 0);
}
#[test]
fn test_many_items_with_same_priority() {
let queue = PriorityQueue::default();
let mut handles: Vec<_> = (0..10).rev().map(|i| queue.insert(live(100, i))).collect();
assert_eq!(handles[0].current(), 0);
for handle in handles.iter_mut() {
assert!(handle.current() < 10);
}
}
#[test]
fn test_max_priority_value_overflow() {
let queue = PriorityQueue::default();
let mut handles: Vec<_> = (0..300).rev().map(|i| queue.insert(live(100, i))).collect();
assert_eq!(handles[0].current(), 0);
let mut low_priority_count = 0;
for handle in handles.iter_mut() {
if handle.current() == u8::MAX {
low_priority_count += 1;
}
}
assert!(low_priority_count > 0, "Should have some items beyond u8::MAX index");
assert_eq!(low_priority_count, 45, "Exactly 45 items should overflow (300-255)");
}
#[test]
fn test_complex_ordering() {
let queue = PriorityQueue::default();
let mut high_track_high_group = queue.insert(live(255, 10));
let mut high_track_low_group = queue.insert(live(255, 1));
let mut mid_track_high_group = queue.insert(live(100, 5));
let mut mid_track_low_group = queue.insert(live(100, 1));
let mut low_track_high_group = queue.insert(live(50, 100));
assert_eq!(high_track_high_group.current(), 0); assert_eq!(high_track_low_group.current(), 1); assert_eq!(mid_track_high_group.current(), 2); assert_eq!(mid_track_low_group.current(), 3); assert_eq!(low_track_high_group.current(), 4); }
#[tokio::test]
async fn test_watch_notification_on_overflow_promotion() {
let queue = PriorityQueue::default();
let mut fillers: Vec<_> = (0..255).rev().map(|i| queue.insert(live(100, i + 100))).collect();
let mut overflow_item = queue.insert(live(100, 50));
assert_eq!(overflow_item.current(), u8::MAX);
let task = tokio::spawn(async move { overflow_item.next().await });
tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
fillers.remove(0);
let result = task.await.unwrap();
assert!(result < u8::MAX, "Should be promoted from overflow");
}
#[test]
fn test_interleaved_insertions_and_removals() {
let queue = PriorityQueue::default();
let mut h1 = queue.insert(live(200, 0));
let h2 = queue.insert(live(150, 0));
let mut h3 = queue.insert(live(100, 0));
assert_eq!(h1.current(), 0);
drop(h2);
assert_eq!(h1.current(), 0);
assert!(h3.current() < 2);
let mut h4 = queue.insert(live(250, 0));
assert_eq!(h4.current(), 0);
assert_eq!(h1.current(), 1);
drop(h4);
assert_eq!(h1.current(), 0);
}
#[test]
fn test_same_track_and_group() {
let queue = PriorityQueue::default();
let mut h1 = queue.insert(live(100, 5));
let mut h2 = queue.insert(live(100, 5));
let mut h3 = queue.insert(live(100, 5));
let indices = [h1.current(), h2.current(), h3.current()];
assert_eq!(indices.len(), 3);
assert!(indices.contains(&0));
assert!(indices.contains(&1));
assert!(indices.contains(&2));
}
#[test]
fn test_removal_updates_siblings() {
let queue = PriorityQueue::default();
let mut root = queue.insert(live(255, 0));
let left = queue.insert(live(100, 0));
let mut right = queue.insert(live(100, 0));
assert_eq!(root.current(), 0);
drop(left);
assert_eq!(root.current(), 0);
assert_eq!(right.current(), 1);
}
#[test]
fn test_heap_property_maintained() {
let queue = PriorityQueue::default();
let mut handles = vec![
queue.insert(live(100, 5)),
queue.insert(live(200, 3)),
queue.insert(live(50, 10)),
queue.insert(live(200, 8)),
queue.insert(live(100, 1)),
];
assert_eq!(handles[3].current(), 0);
drop(handles.remove(3));
assert_eq!(handles[1].current(), 0);
}
#[tokio::test]
async fn test_notification_on_demotion_to_overflow() {
let queue = PriorityQueue::default();
let _fillers: Vec<_> = (0..254).map(|i| queue.insert(live(100, i + 100))).collect();
let mut at_edge = queue.insert(live(100, 50));
assert_eq!(at_edge.current(), 254);
let task = tokio::spawn(async move { at_edge.next().await });
tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
let _high = queue.insert(live(255, 1000));
let new_priority = task.await.unwrap();
assert_eq!(new_priority, u8::MAX, "Should be demoted to overflow");
}
#[test]
fn test_empty_after_all_removed() {
let queue = PriorityQueue::default();
let h1 = queue.insert(live(100, 0));
let h2 = queue.insert(live(200, 0));
let h3 = queue.insert(live(50, 0));
drop(h1);
drop(h2);
drop(h3);
let mut h4 = queue.insert(live(100, 0));
assert_eq!(h4.current(), 0);
}
#[test]
fn test_set_track_reorders() {
let queue = PriorityQueue::default();
let mut s1_g1 = queue.insert(live(255, 1));
let mut s1_g2 = queue.insert(live(255, 2));
let mut s2_g1 = queue.insert(live(55, 1));
let mut s2_g2 = queue.insert(live(55, 2));
assert_eq!(s1_g2.current(), 0); assert_eq!(s1_g1.current(), 1);
assert_eq!(s2_g2.current(), 2); assert_eq!(s2_g1.current(), 3);
s1_g1.set_track(55);
s1_g2.set_track(55);
s2_g1.set_track(255);
s2_g2.set_track(255);
assert_eq!(s2_g2.current(), 0); assert_eq!(s2_g1.current(), 1);
assert_eq!(s1_g2.current(), 2); assert_eq!(s1_g1.current(), 3);
}
#[tokio::test]
async fn test_set_track_notifies_other_handles() {
let queue = PriorityQueue::default();
let mut h_high = queue.insert(live(255, 1));
let mut h_low = queue.insert(live(50, 1));
assert_eq!(h_low.current(), 1);
let task = tokio::spawn(async move { h_low.next().await });
tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
h_high.set_track(10);
let new_priority = task.await.unwrap();
assert_eq!(new_priority, 0, "h_low should be promoted to the top");
}
#[test]
fn test_set_track_self() {
let queue = PriorityQueue::default();
let mut h_high = queue.insert(live(255, 1));
let mut h_mid = queue.insert(live(100, 1));
let mut h_low = queue.insert(live(50, 1));
assert_eq!(h_high.current(), 0);
assert_eq!(h_mid.current(), 1);
assert_eq!(h_low.current(), 2);
h_high.set_track(10);
assert_eq!(h_mid.current(), 0);
assert_eq!(h_low.current(), 1);
assert_eq!(h_high.current(), 2);
}
#[test]
fn test_set_track_swaps_demoted_vec_item_with_overflow() {
let queue = PriorityQueue::default();
let mut fillers: Vec<_> = (1..=255u64).map(|g| queue.insert(live(100, g))).collect();
let mut top = queue.insert(live(200, 0));
assert_eq!(top.current(), 0);
assert_eq!(fillers[0].current(), u8::MAX, "f1 was kicked into overflow");
top.set_track(0);
assert!(fillers[0].current() < u8::MAX, "f1 should be promoted back into vec");
assert_eq!(top.current(), u8::MAX, "demoted top should land in overflow");
}
#[test]
fn test_set_track_lowered_within_vec_no_overflow_disruption() {
let queue = PriorityQueue::default();
let mut a = queue.insert(live(200, 0));
let mut b = queue.insert(live(100, 0));
let mut c = queue.insert(live(50, 0));
assert_eq!(a.current(), 0);
assert_eq!(b.current(), 1);
assert_eq!(c.current(), 2);
a.set_track(75);
assert_eq!(b.current(), 0);
assert_eq!(a.current(), 1);
assert_eq!(c.current(), 2);
}
#[test]
fn test_remove_promotes_highest_priority_overflow_item() {
let queue = PriorityQueue::default();
let fillers: Vec<_> = (100..355u64).map(|g| queue.insert(live(200, g))).collect();
let mut low = queue.insert(live(100, 1));
let mut mid = queue.insert(live(100, 2));
let mut high = queue.insert(live(100, 3));
assert_eq!(low.current(), u8::MAX);
assert_eq!(mid.current(), u8::MAX);
assert_eq!(high.current(), u8::MAX);
drop(fillers);
assert_eq!(
high.current(),
0,
"highest-priority overflow item should land at index 0"
);
assert_eq!(mid.current(), 1);
assert_eq!(low.current(), 2);
}
#[tokio::test]
async fn test_set_track_notifies_swapped_overflow_item() {
tokio::time::pause();
let queue = PriorityQueue::default();
let mut fillers: Vec<_> = (1..=255u64).map(|g| queue.insert(live(100, g))).collect();
let mut top = queue.insert(live(200, 0));
assert_eq!(top.current(), 0);
let mut f1 = fillers.remove(0);
assert_eq!(f1.current(), u8::MAX);
let task = tokio::spawn(async move { f1.next().await });
tokio::task::yield_now().await;
top.set_track(0);
let promoted = task.await.unwrap();
assert!(promoted < u8::MAX, "f1 should be notified of promotion");
}
}