use std::collections::BinaryHeap;
use std::cmp::Ordering;
use parking_lot::Mutex;
use super::Workload;
#[derive(Debug)]
pub struct QueuedWorkload {
pub workload: Workload,
pub sequence: u64,
pub attempts: u32,
pub backoff_until: Option<chrono::DateTime<chrono::Utc>>,
}
impl QueuedWorkload {
pub fn new(workload: Workload, sequence: u64) -> Self {
Self {
workload,
sequence,
attempts: 0,
backoff_until: None,
}
}
pub fn record_failure(&mut self) {
self.attempts += 1;
let backoff_secs = (2_i64.pow(self.attempts.min(8)) as i64).min(300);
self.backoff_until = Some(chrono::Utc::now() + chrono::Duration::seconds(backoff_secs));
}
pub fn is_ready(&self) -> bool {
self.backoff_until
.map(|t| chrono::Utc::now() >= t)
.unwrap_or(true)
}
}
impl PartialEq for QueuedWorkload {
fn eq(&self, other: &Self) -> bool {
self.workload.id == other.workload.id
}
}
impl Eq for QueuedWorkload {}
impl PartialOrd for QueuedWorkload {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}
impl Ord for QueuedWorkload {
fn cmp(&self, other: &Self) -> Ordering {
match self.workload.priority.cmp(&other.workload.priority) {
Ordering::Equal => {
other.sequence.cmp(&self.sequence)
}
other => other,
}
}
}
pub struct SchedulingQueue {
queue: Mutex<BinaryHeap<QueuedWorkload>>,
sequence: Mutex<u64>,
}
impl SchedulingQueue {
pub fn new() -> Self {
Self {
queue: Mutex::new(BinaryHeap::new()),
sequence: Mutex::new(0),
}
}
pub fn enqueue(&self, workload: Workload) {
let mut seq = self.sequence.lock();
let sequence = *seq;
*seq += 1;
drop(seq);
let queued = QueuedWorkload::new(workload, sequence);
self.queue.lock().push(queued);
}
pub fn dequeue(&self) -> Option<Workload> {
let mut queue = self.queue.lock();
let mut not_ready = Vec::new();
while let Some(queued) = queue.pop() {
if queued.is_ready() {
for q in not_ready {
queue.push(q);
}
return Some(queued.workload);
} else {
not_ready.push(queued);
}
}
for q in not_ready {
queue.push(q);
}
None
}
pub fn peek(&self) -> Option<String> {
self.queue.lock().peek().map(|q| q.workload.id.clone())
}
pub fn len(&self) -> usize {
self.queue.lock().len()
}
pub fn is_empty(&self) -> bool {
self.queue.lock().is_empty()
}
pub fn remove(&self, workload_id: &str) -> bool {
let mut queue = self.queue.lock();
let items: Vec<_> = std::mem::take(&mut *queue).into_vec();
let mut found = false;
for item in items {
if item.workload.id == workload_id {
found = true;
} else {
queue.push(item);
}
}
found
}
pub fn requeue_with_backoff(&self, mut queued: QueuedWorkload) {
queued.record_failure();
self.queue.lock().push(queued);
}
}
impl Default for SchedulingQueue {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_priority_ordering() {
let queue = SchedulingQueue::new();
queue.enqueue(Workload::new("low", "low").with_priority(10));
queue.enqueue(Workload::new("high", "high").with_priority(100));
queue.enqueue(Workload::new("medium", "medium").with_priority(50));
assert_eq!(queue.dequeue().unwrap().id, "high");
assert_eq!(queue.dequeue().unwrap().id, "medium");
assert_eq!(queue.dequeue().unwrap().id, "low");
}
#[test]
fn test_fifo_within_priority() {
let queue = SchedulingQueue::new();
queue.enqueue(Workload::new("first", "first").with_priority(50));
queue.enqueue(Workload::new("second", "second").with_priority(50));
queue.enqueue(Workload::new("third", "third").with_priority(50));
assert_eq!(queue.dequeue().unwrap().id, "first");
assert_eq!(queue.dequeue().unwrap().id, "second");
assert_eq!(queue.dequeue().unwrap().id, "third");
}
#[test]
fn test_remove() {
let queue = SchedulingQueue::new();
queue.enqueue(Workload::new("w1", "w1"));
queue.enqueue(Workload::new("w2", "w2"));
queue.enqueue(Workload::new("w3", "w3"));
assert!(queue.remove("w2"));
assert_eq!(queue.len(), 2);
assert!(!queue.remove("w2")); }
}