Delay Queue Task Scheduler

use std::cmp::Ordering;
use std::collections::BinaryHeap;
use std::time::{Duration, Instant};

#[derive(Debug)]
struct ScheduledTask {
    run_at: Instant,
    sequence: u64,
    name: String,
}

impl PartialEq for ScheduledTask {
    fn eq(&self, other: &Self) -> bool {
        self.run_at == other.run_at && self.sequence == other.sequence
    }
}

impl Eq for ScheduledTask {}

impl Ord for ScheduledTask {
    fn cmp(&self, other: &Self) -> Ordering {
        other
            .run_at
            .cmp(&self.run_at)
            .then_with(|| other.sequence.cmp(&self.sequence))
    }
}

impl PartialOrd for ScheduledTask {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

struct DelayQueue {
    heap: BinaryHeap<ScheduledTask>,
    sequence: u64,
}

impl DelayQueue {
    fn new() -> Self {
        Self {
            heap: BinaryHeap::new(),
            sequence: 0,
        }
    }

    fn schedule(&mut self, name: &str, delay: Duration) {
        self.sequence += 1;
        self.heap.push(ScheduledTask {
            run_at: Instant::now() + delay,
            sequence: self.sequence,
            name: name.to_string(),
        });
    }

    fn pop_ready(&mut self, now: Instant) -> Vec<String> {
        let mut ready = Vec::new();
        loop {
            let due = match self.heap.peek() {
                Some(task) => task.run_at <= now,
                None => false,
            };
            if !due {
                break;
            }
            if let Some(task) = self.heap.pop() {
                ready.push(task.name);
            }
        }
        ready
    }

    fn next_deadline(&self) -> Option<Instant> {
        self.heap.peek().map(|task| task.run_at)
    }
}

fn main() {
    let start = Instant::now();
    let mut queue = DelayQueue::new();
    queue.schedule("warm-cache", Duration::from_millis(50));
    queue.schedule("send-email", Duration::from_millis(150));
    queue.schedule("refresh-token", Duration::from_millis(100));

    println!("ready at start: {:?}", queue.pop_ready(start));
    println!("has deadline: {}", queue.next_deadline().is_some());
    println!(
        "ready at 120ms: {:?}",
        queue.pop_ready(start + Duration::from_millis(120))
    );
    println!(
        "ready at 200ms: {:?}",
        queue.pop_ready(start + Duration::from_millis(200))
    );
}