caveman/src/world.rs

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

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

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

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

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

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

    pub fn introduce(&mut self, host: String) {
        let now = Instant::now();

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

    pub fn enqueue(&mut self, host: String, next_interval: Duration, latest_timestamp: Option<DateTime<FixedOffset>>) {
        let now = Instant::now();

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

        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)))
            } else {
                Err(time - now)
            }
        } else {
            println!("empty queue");
            Err(Duration::from_secs(60))
        }
    }
}