caveman/hunter/src/scheduler.rs

93 lines
2.6 KiB
Rust

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, host: String) {
let now = Instant::now();
if self.instances.get(&host).is_none() {
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 {:.0?}", 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))
}
}
}