use std::cmp::Ordering;
use std::collections::BTreeMap;
use std::fmt;
use std::fmt::Display;
use rand::RngExt as _;
use rand::SeedableRng;
use rand::distr::uniform::SampleUniform;
use rand_pcg::Pcg64Mcg;
use crate::config::SchedHeuristic;
use crate::detlog;
use crate::types::DetTid;
pub type Priority = u64;
const EAGER_IO_REPOLL_PRIORITY: Priority = Priority::MIN;
pub const FIRST_PRIORITY: Priority = EAGER_IO_REPOLL_PRIORITY + 1;
pub const LAST_PRIORITY: Priority = 10000;
pub const REPLAY_FOREGROUND_PRIORITY: Priority = FIRST_PRIORITY;
pub const REPLAY_DEFERRED_PRIORITY: Priority = LAST_PRIORITY - 1;
pub const DEFAULT_PRIORITY: Priority = 1000;
pub fn is_ordinary_priority(prio: Priority) -> bool {
(FIRST_PRIORITY..=LAST_PRIORITY).contains(&prio)
}
pub fn entropy_to_priority(entropy: u64) -> Priority {
let range = LAST_PRIORITY - FIRST_PRIORITY + 1;
let offset = entropy % range;
FIRST_PRIORITY + offset
}
type RoundRobinTurn = i64;
#[derive(Debug, Copy, Clone)]
pub struct PrioritizedOrder {
priority: Priority,
turn: RoundRobinTurn,
}
impl Ord for PrioritizedOrder {
fn cmp(&self, other: &Self) -> Ordering {
self.priority
.cmp(&other.priority)
.then(self.turn.cmp(&other.turn))
}
}
impl PartialOrd for PrioritizedOrder {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}
impl PartialEq for PrioritizedOrder {
fn eq(&self, other: &Self) -> bool {
self.cmp(other).is_eq()
}
}
impl Eq for PrioritizedOrder {}
impl fmt::Display for PrioritizedOrder {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "(p: {}, t: {})", self.priority, self.turn)
}
}
const POLLING_UPGRADE_INTERVAL: u64 = 200;
#[derive(Debug, Copy, Clone)]
struct QueueValue {
tid: DetTid,
poll_upgrade: Option<Priority>,
}
#[derive(Debug, Clone)]
pub(super) struct SuspendedRunQueueEntry {
key: PrioritizedOrder,
value: QueueValue,
persistent_priority: Priority,
}
#[derive(Debug, Clone)]
pub struct RunQueue {
queue: BTreeMap<PrioritizedOrder, QueueValue>,
last_back_turn: RoundRobinTurn,
last_front_turn: RoundRobinTurn,
tentative_selection: Option<DetTid>,
tentative_selection_is_exact: bool,
yielded_skip: Option<DetTid>,
sched_strategy: SchedHeuristic,
prng: Pcg64Mcg,
sticky_random_param: f64,
sticky_random_selection: Option<DetTid>,
}
impl fmt::Display for RunQueue {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
writeln!(
f,
"Run queue, size={}, last_back_turn={}, last_front_turn={}:",
self.queue.len(),
self.last_back_turn,
self.last_front_turn,
)?;
for x in self.queue.iter() {
writeln!(f, " {:.500?}", x)?;
}
Ok(())
}
}
impl RunQueue {
pub fn new(ss: SchedHeuristic, seed: u64, srp: f64) -> Self {
detlog!("SCHEDRAND: seeding scheduler runqueue with seed {}", seed);
Self {
queue: BTreeMap::new(),
last_back_turn: 0,
last_front_turn: 0,
sched_strategy: ss,
tentative_selection: None,
tentative_selection_is_exact: false,
yielded_skip: None,
prng: Pcg64Mcg::seed_from_u64(seed),
sticky_random_param: srp,
sticky_random_selection: None,
}
}
fn push_safety_check(&self, tid: DetTid) {
if cfg!(debug_assertions) {
for qv in self.queue.values() {
if qv.tid == tid {
panic!(
"Invariant violation! Tried to add {} to runqueue, but it's already present:\n {:?}",
tid, self
);
}
}
}
}
pub fn first_priority(&self) -> Option<Priority> {
let (k, _) = self.queue.first_key_value()?;
Some(k.priority)
}
pub fn has_runnable_besides(&self, exclude: DetTid) -> bool {
self.queue
.iter()
.any(|(k, v)| v.tid != exclude && k.priority < LAST_PRIORITY)
}
pub fn tentative_pop_in_progress(&self) -> bool {
self.tentative_selection.is_some()
}
pub fn push_back(&mut self, tid: DetTid, priority: Priority) -> PrioritizedOrder {
assert!(self.tentative_selection.is_none());
self.push_safety_check(tid);
if !is_ordinary_priority(priority) {
panic!("This is not an acceptable priority value: {}", priority);
}
self.push_back_inner(tid, priority, None)
}
pub fn push_yielded(&mut self, tid: DetTid, priority: Priority) -> PrioritizedOrder {
assert!(self.yielded_skip.is_none());
self.yielded_skip = Some(tid);
self.push_back(tid, priority)
}
pub fn push_poller(
&mut self,
tid: DetTid,
normal_priority: Priority,
poll_attempt: u32,
) -> PrioritizedOrder {
assert!(self.tentative_selection.is_none());
self.push_safety_check(tid);
let priority = 1u64
.checked_shl(poll_attempt)
.and_then(|f| f.checked_mul(normal_priority))
.unwrap_or(Priority::MAX)
.min(LAST_PRIORITY);
self.push_back_inner(tid, priority, Some(normal_priority))
}
fn push_back_inner(
&mut self,
tid: DetTid,
priority: Priority,
poll_upgrade: Option<Priority>,
) -> PrioritizedOrder {
self.last_back_turn += 1;
let turn = self.last_back_turn;
let prio = PrioritizedOrder { priority, turn };
self.push_inner(tid, prio, poll_upgrade)
}
pub fn push_front(&mut self, tid: DetTid, priority: Priority) -> PrioritizedOrder {
assert!(self.tentative_selection.is_none());
self.push_safety_check(tid);
assert!(is_ordinary_priority(priority));
self.push_front_inner(tid, priority, None)
}
pub fn push_eager_io_repoll(&mut self, tid: DetTid) -> PrioritizedOrder {
assert!(self.tentative_selection.is_none());
self.push_safety_check(tid);
let priority = EAGER_IO_REPOLL_PRIORITY;
self.push_front_inner(tid, priority, None)
}
fn push_front_inner(
&mut self,
tid: DetTid,
priority: Priority,
poll_upgrade: Option<Priority>,
) -> PrioritizedOrder {
self.last_front_turn -= 1;
let turn = self.last_front_turn;
let prio = PrioritizedOrder { priority, turn };
self.push_inner(tid, prio, poll_upgrade)
}
fn push_inner(
&mut self,
tid: DetTid,
prio: PrioritizedOrder,
poll_upgrade: Option<Priority>,
) -> PrioritizedOrder {
let qval = QueueValue { tid, poll_upgrade };
let old = self.queue.insert(prio, qval);
assert!(old.is_none()); self.check_poll_upgrade();
prio
}
pub fn is_empty(&self) -> bool {
self.queue.is_empty()
}
pub fn len(&self) -> usize {
self.queue.len()
}
pub fn tids(&self) -> impl Iterator<Item = &DetTid> {
self.queue.values().map(|v| &v.tid)
}
pub fn contains_tid(&self, tid: DetTid) -> bool {
self.tids().any(|t| t == &tid)
}
pub(super) fn suspend(
&mut self,
tid: DetTid,
persistent_priority: Priority,
) -> Option<SuspendedRunQueueEntry> {
assert!(self.tentative_selection.is_none());
let key = *self.queue.iter().find(|(_, v)| v.tid == tid)?.0;
let value = self.queue.remove(&key).expect("located queue entry");
Some(SuspendedRunQueueEntry {
key,
value,
persistent_priority,
})
}
pub(super) fn restore(
&mut self,
mut entry: SuspendedRunQueueEntry,
current_priority: Priority,
) {
assert!(self.tentative_selection.is_none());
assert!(!self.contains_tid(entry.value.tid));
if entry.persistent_priority != current_priority {
entry.key.priority = current_priority;
if entry.value.poll_upgrade.is_some() {
entry.value.poll_upgrade = Some(current_priority);
}
}
assert!(self.queue.insert(entry.key, entry.value).is_none());
}
pub fn remove_tid(&mut self, tid: DetTid) -> bool {
assert!(self.tentative_selection.is_none());
let mut kept_all = true;
self.queue.retain(|_k, v| {
let ret = v.tid != tid;
kept_all = kept_all && ret;
ret
});
if self.yielded_skip == Some(tid) {
self.yielded_skip = None;
}
if self.sticky_random_selection == Some(tid) {
self.sticky_random_selection = None;
}
!kept_all
}
fn random_range<T>(&mut self, start: T, end: T) -> T
where
T: SampleUniform + Display + PartialOrd + Copy,
{
let r = self.prng.random_range(start..end);
detlog!("SCHEDRAND: [{},{}) => {}", start, end, r);
r
}
pub fn tentative_pop_next(&mut self) -> Option<DetTid> {
assert!(!self.tentative_selection_is_exact);
let skip = self
.yielded_skip
.filter(|tid| self.queue.len() > 1 && self.contains_tid(*tid));
self.tentative_selection = match self.sched_strategy {
SchedHeuristic::None | SchedHeuristic::ConnectBind => self
.queue
.values()
.find(|value| Some(value.tid) != skip)
.map(|value| value.tid),
SchedHeuristic::Random => {
if self.queue.is_empty() {
return None;
}
if self.tentative_selection.is_none() {
let eligible = self.queue.len() - usize::from(skip.is_some());
let random_idx = self.random_range(0, eligible);
self.tentative_selection = self
.queue
.values()
.filter(|value| Some(value.tid) != skip)
.nth(random_idx)
.map(|value| value.tid);
};
self.tentative_selection
}
SchedHeuristic::StickyRandom => {
if self.queue.is_empty() {
return None;
}
if self.sticky_random_selection == skip {
self.sticky_random_selection = None;
}
if self.sticky_random_selection.is_none()
|| !self.contains_tid(self.sticky_random_selection.unwrap())
{
let eligible = self.queue.len() - usize::from(skip.is_some());
let random_idx = self.random_range(0, eligible);
self.sticky_random_selection = self
.queue
.values()
.filter(|value| Some(value.tid) != skip)
.nth(random_idx)
.map(|value| value.tid);
}
self.sticky_random_selection
}
};
self.tentative_selection
}
pub fn tentative_pop_tid(&mut self, tid: DetTid) -> Option<DetTid> {
assert!(self.tentative_selection.is_none());
if self.contains_tid(tid) {
self.tentative_selection = Some(tid);
self.tentative_selection_is_exact = true;
}
self.tentative_selection
}
pub fn commit_tentative_pop(&mut self) -> DetTid {
let tentative_selection = self
.tentative_selection
.take()
.expect("tentative_pop to already returned a `Some`");
let exact = std::mem::take(&mut self.tentative_selection_is_exact);
let ret = if exact {
let key = *self
.queue
.iter()
.find(|(_key, value)| value.tid == tentative_selection)
.map(|(key, _value)| key)
.unwrap();
self.queue.remove(&key).map(|value| value.tid)
} else {
match self.sched_strategy {
SchedHeuristic::None | SchedHeuristic::ConnectBind | SchedHeuristic::Random => {
let key = *self
.queue
.iter()
.find(|(_k, v)| v.tid == tentative_selection)
.map(|(k, _v)| k)
.unwrap();
self.queue.remove(&key).map(|v| v.tid)
}
SchedHeuristic::StickyRandom => {
let tid = self.sticky_random_selection.unwrap();
if self.random_range(0f64, 1f64) <= 1.0 - self.sticky_random_param {
self.sticky_random_selection = None;
}
let key = *self
.queue
.iter()
.find(|(_k, v)| v.tid == tid)
.map(|(k, _v)| k)
.unwrap();
self.queue.remove(&key).map(|v| v.tid)
}
}
}
.expect("to always return a DetTid");
debug_assert!(ret == tentative_selection);
ret
}
pub fn commit_tentative_pop_completed_turn(&mut self) -> DetTid {
let tid = self.commit_tentative_pop();
self.consume_yield_exclusion();
tid
}
pub fn consume_yield_exclusion(&mut self) {
self.yielded_skip = None;
}
pub fn undo_tentative_pop(&mut self) {
assert!(self.tentative_selection.is_some());
self.tentative_selection = None;
self.tentative_selection_is_exact = false;
}
fn turn_counter(&self) -> u64 {
debug_assert!(self.last_back_turn >= 0);
debug_assert!(self.last_front_turn <= 0);
self.last_back_turn as u64 + self.last_front_turn.unsigned_abs()
}
fn check_poll_upgrade(&mut self) {
if self.turn_counter().is_multiple_of(POLLING_UPGRADE_INTERVAL) {
self.do_poll_upgrade()
}
}
#[cold]
fn do_poll_upgrade(&mut self) {
let upgrades = self
.queue
.iter()
.filter_map(|(k, v)| v.poll_upgrade.map(|upgd| (*k, upgd)))
.collect::<Vec<(PrioritizedOrder, Priority)>>();
for (key, upgrade_prio) in upgrades {
let mut new_key = key;
new_key.priority = upgrade_prio;
let mut qval = self.queue.remove(&key).unwrap();
qval.poll_upgrade = None; let old = self.queue.insert(new_key, qval);
assert!(old.is_none()); }
}
}
impl Default for RunQueue {
fn default() -> Self {
Self::new(SchedHeuristic::None, 0, 0.0)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn transport_suspend_restore_preserves_all_queue_state() {
for strategy in [
SchedHeuristic::None,
SchedHeuristic::Random,
SchedHeuristic::StickyRandom,
] {
let tid = DetTid::from_raw(1);
let peer = DetTid::from_raw(2);
let mut queue = RunQueue::new(strategy, 391, 0.5);
queue.push_yielded(tid, DEFAULT_PRIORITY);
queue.push_back(peer, DEFAULT_PRIORITY);
let key = *queue
.queue
.iter()
.find(|(_, value)| value.tid == tid)
.unwrap()
.0;
queue.queue.get_mut(&key).unwrap().poll_upgrade = Some(DEFAULT_PRIORITY);
queue.sticky_random_selection = Some(tid);
let before = format!("{queue:?}");
let saved = queue.suspend(tid, DEFAULT_PRIORITY).unwrap();
assert_eq!(queue.yielded_skip, Some(tid));
assert_eq!(queue.sticky_random_selection, Some(tid));
queue.restore(saved, DEFAULT_PRIORITY);
assert_eq!(format!("{queue:?}"), before, "{strategy:?}");
}
}
#[test]
fn observation_restore_keeps_real_selection_progress_and_new_priority() {
for strategy in [
SchedHeuristic::None,
SchedHeuristic::Random,
SchedHeuristic::StickyRandom,
] {
let tid = DetTid::from_raw(1);
let peer = DetTid::from_raw(2);
let mut queue = RunQueue::new(strategy, 391, 0.5);
queue.push_yielded(tid, DEFAULT_PRIORITY);
queue.push_back(peer, DEFAULT_PRIORITY);
let key = *queue
.queue
.iter()
.find(|(_, value)| value.tid == tid)
.unwrap()
.0;
queue.queue.get_mut(&key).unwrap().poll_upgrade = Some(DEFAULT_PRIORITY);
let saved = queue.suspend(tid, DEFAULT_PRIORITY).unwrap();
assert_eq!(queue.tentative_pop_next(), Some(peer));
assert_eq!(queue.commit_tentative_pop_completed_turn(), peer);
assert_eq!(queue.yielded_skip, None);
queue.push_back(peer, DEFAULT_PRIORITY);
let random_after_turn = format!("{:?}", queue.prng);
let sticky_after_turn = queue.sticky_random_selection;
let turns_after_turn = (queue.last_back_turn, queue.last_front_turn);
queue.restore(saved, DEFAULT_PRIORITY + 3);
let (restored, value) = queue
.queue
.iter()
.find(|(_, value)| value.tid == tid)
.unwrap();
assert_eq!(restored.turn, key.turn);
assert_eq!(restored.priority, DEFAULT_PRIORITY + 3);
assert_eq!(value.poll_upgrade, Some(DEFAULT_PRIORITY + 3));
assert_eq!(queue.yielded_skip, None);
assert_eq!(queue.sticky_random_selection, sticky_after_turn);
assert_eq!(
(queue.last_back_turn, queue.last_front_turn),
turns_after_turn
);
assert_eq!(format!("{:?}", queue.prng), random_after_turn);
}
}
#[test]
fn yielded_thread_cedes_exactly_one_turn_under_every_heuristic() {
for strategy in [
SchedHeuristic::None,
SchedHeuristic::ConnectBind,
SchedHeuristic::Random,
SchedHeuristic::StickyRandom,
] {
let yielding = DetTid::from_raw(1);
let peer = DetTid::from_raw(2);
let mut queue = RunQueue::new(strategy, 0, 1.0);
queue.push_back(yielding, DEFAULT_PRIORITY - 1);
assert_eq!(queue.tentative_pop_next(), Some(yielding));
assert_eq!(queue.commit_tentative_pop(), yielding);
queue.push_yielded(yielding, DEFAULT_PRIORITY - 1);
queue.push_back(peer, LAST_PRIORITY);
assert_eq!(queue.tentative_pop_next(), Some(peer), "{strategy:?}");
assert_eq!(
queue.commit_tentative_pop_completed_turn(),
peer,
"{strategy:?}"
);
assert_eq!(queue.yielded_skip, None, "{strategy:?}");
let restored_priority = queue
.queue
.iter()
.find(|(_key, value)| value.tid == yielding)
.map(|(key, _value)| key.priority);
assert_eq!(
restored_priority,
Some(DEFAULT_PRIORITY - 1),
"{strategy:?}"
);
if matches!(strategy, SchedHeuristic::None | SchedHeuristic::ConnectBind) {
queue.push_back(peer, LAST_PRIORITY);
assert_eq!(queue.tentative_pop_next(), Some(yielding), "{strategy:?}");
}
}
}
#[test]
fn exact_selection_bypasses_priority_and_heuristic() {
for strategy in [
SchedHeuristic::None,
SchedHeuristic::ConnectBind,
SchedHeuristic::Random,
SchedHeuristic::StickyRandom,
] {
let higher_priority = DetTid::from_raw(1);
let selected = DetTid::from_raw(2);
let mut queue = RunQueue::new(strategy, 0, 1.0);
queue.push_back(higher_priority, FIRST_PRIORITY);
queue.push_back(selected, LAST_PRIORITY);
assert_eq!(queue.tentative_pop_tid(selected), Some(selected));
assert_eq!(queue.commit_tentative_pop(), selected);
assert!(queue.contains_tid(higher_priority));
assert!(!queue.contains_tid(selected));
}
}
#[test]
fn scheduler_only_commit_does_not_consume_yield_exclusion() {
let yielding = DetTid::from_raw(1);
let peer = DetTid::from_raw(2);
let mut queue = RunQueue::default();
queue.push_back(yielding, DEFAULT_PRIORITY);
assert_eq!(queue.tentative_pop_next(), Some(yielding));
assert_eq!(queue.commit_tentative_pop(), yielding);
queue.push_yielded(yielding, DEFAULT_PRIORITY);
queue.push_back(peer, DEFAULT_PRIORITY);
assert_eq!(queue.tentative_pop_next(), Some(peer));
assert_eq!(queue.commit_tentative_pop(), peer);
queue.push_back(peer, DEFAULT_PRIORITY);
assert_eq!(queue.yielded_skip, Some(yielding));
assert_eq!(queue.tentative_pop_next(), Some(peer));
assert_eq!(queue.commit_tentative_pop_completed_turn(), peer);
assert_eq!(queue.yielded_skip, None);
}
#[test]
fn tentative_pop_in_progress_tracks_the_transaction() {
let a = DetTid::from_raw(1);
let b = DetTid::from_raw(2);
let mut queue = RunQueue::default();
assert!(!queue.tentative_pop_in_progress());
queue.push_back(a, DEFAULT_PRIORITY);
queue.push_back(b, DEFAULT_PRIORITY);
assert!(!queue.tentative_pop_in_progress());
assert_eq!(queue.tentative_pop_next(), Some(a));
assert!(queue.tentative_pop_in_progress());
assert_eq!(queue.commit_tentative_pop(), a);
assert!(!queue.tentative_pop_in_progress());
assert_eq!(queue.tentative_pop_tid(b), Some(b));
assert!(queue.tentative_pop_in_progress());
queue.undo_tentative_pop();
assert!(!queue.tentative_pop_in_progress());
}
#[test]
fn removal_clears_per_incarnation_selection_state_before_tid_reuse() {
let tid = DetTid::from_raw(7);
let mut queue = RunQueue::new(SchedHeuristic::StickyRandom, 0x5107, 1.0);
queue.push_back(tid, DEFAULT_PRIORITY);
assert_eq!(queue.tentative_pop_next(), Some(tid));
queue.undo_tentative_pop();
assert_eq!(queue.sticky_random_selection, Some(tid));
queue.yielded_skip = Some(tid);
assert!(queue.remove_tid(tid));
assert_eq!(queue.sticky_random_selection, None);
assert_eq!(queue.yielded_skip, None);
queue.push_back(tid, DEFAULT_PRIORITY);
assert_eq!(queue.sticky_random_selection, None);
assert_eq!(queue.yielded_skip, None);
}
}