caveman/hunter/src/scheduler.rs

use std::collections::{HashMap, BTreeMap};
use std::time::Duration;
use tokio::time::Instant;

pub struct Instance {
    last_fetch: Option<Instant>,
    no_updates: u32,
    error: bool,
}

/// Scheduler
pub struct Scheduler {
    instances: HashMap<String, Instance>,
    queue: BTreeMap<Instant, String>,
}

impl Scheduler {
    pub fn new() -> Self {
        Scheduler {
            instances: HashMap::new(),
            queue: BTreeMap::new(),
        }
    }

    pub fn size(&self) -> usize {
        self.instances.len()
    }

    pub fn queue_len(&self) -> usize {
        self.queue.len()
    }

    pub async fn introduce(&mut self, redis_man: &mut redis::aio::ConnectionManager, host: String) {
        let now = Instant::now();

        if self.instances.get(&host).is_none() {
            // save for later
            crate::redis_store::save_host(redis_man, &host).await;

            self.instances.insert(host.clone(), Instance {
                last_fetch: None,
                no_updates: 0,
                error: false,
            });
            self.queue.insert(now, host);
        }
    }

    pub fn enqueue(&mut self, host: String, fetched_anything: bool, next_interval: Duration) {
        let now = Instant::now();

        let instance = self.instances.get_mut(&host).unwrap();
        instance.last_fetch = Some(now);
        instance.error = false;
        if fetched_anything {
            instance.no_updates = 0;
        } else {
            instance.no_updates += 1;
        }

        let mut next = now + (1 + instance.no_updates) * next_interval;
        let mut d = 1;
        // avoid timestamp collision in self.queue
        while self.queue.get(&next).is_some() {
            d *= 2;
            next += Duration::from_micros(d);
        }
        self.queue.insert(next, host);
    }

    pub fn dequeue(&mut self) -> Result<String, Duration> {
        let now = Instant::now();

        if let Some(time) = self.queue.keys().next().cloned() {
            if time <= now {
                self.queue.remove(&time)
                    .map(|host| Ok(host))
                    .unwrap_or(Err(Duration::from_secs(1)))
                    .map(|host| {
                        if let Some(last_fetch) = self.instances.get(&host).and_then(|i| i.last_fetch) {
                            log::debug!("Fetch {} - last before {:?}", host, now - last_fetch);
                        } else {
                            log::debug!("Fetch {} - NEW", host);
                        }
                        host
                    })
            } else {
                Err(time - now)
            }
        } else {
            log::warn!("empty queue");
            Err(Duration::from_secs(60))
        }
    }
}