use std::collections::{HashMap, HashSet, VecDeque};
use url::Url;
mod clean;
mod learn;
#[cfg(test)]
mod tests;
pub use clean::*;
pub use learn::*;
#[derive(Default)]
pub struct Tiered {
normal: HashMap<String, VecDeque<(Url, usize, u64)>>,
ring: VecDeque<String>,
normal_len: usize,
low: VecDeque<(Url, usize, u64)>,
sorted_at: u64,
}
fn host_of(u: &Url) -> String {
u.host_str().unwrap_or_default().to_string()
}
impl Tiered {
pub fn len(&self) -> usize {
self.normal_len + self.low.len()
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn push(&mut self, learner: &Learner, url: Url, depth: usize) {
let key = shape_key(&url);
if learner.demoted_key(key) {
self.low.push_back((url, depth, key));
return;
}
let host = host_of(&url);
let q = self.normal.entry(host.clone()).or_default();
if q.is_empty() {
self.ring.push_back(host);
}
q.push_back((url, depth, key));
self.normal_len += 1;
}
pub fn pop(
&mut self,
learner: &Learner,
allow_low: bool,
busy: impl Fn(&str) -> bool,
) -> Option<(Url, usize, bool)> {
for _ in 0..self.ring.len() {
let host = self.ring.pop_front()?;
if busy(&host) {
self.ring.push_back(host);
continue;
}
let q = self.normal.get_mut(&host).expect("ring hosts have queues");
let mut found = None;
while let Some(item) = q.pop_front() {
self.normal_len -= 1;
if learner.demoted_key(item.2) {
self.low.push_back(item);
} else {
found = Some((item.0, item.1, false));
break;
}
}
if q.is_empty() {
self.normal.remove(&host);
} else {
self.ring.push_back(host);
}
if found.is_some() {
return found;
}
}
if !allow_low || self.low.front().is_some_and(|i| busy(&host_of(&i.0))) {
return None;
}
self.low.pop_front().map(|(u, d, _)| (u, d, true))
}
pub fn evict_low(&mut self, learner: &Learner) -> Option<Url> {
if self.sorted_at != learner.generation {
self.sorted_at = learner.generation;
for q in self.normal.values_mut() {
let before = q.len();
let mut keep = VecDeque::with_capacity(before);
for item in q.drain(..) {
if learner.demoted_key(item.2) {
self.low.push_back(item);
} else {
keep.push_back(item);
}
}
self.normal_len -= before - keep.len();
*q = keep;
}
self.normal.retain(|_, q| !q.is_empty());
let normal = &self.normal;
self.ring.retain(|h| normal.contains_key(h));
}
self.low.pop_back().map(|(u, _, _)| u)
}
pub fn iter(&self) -> impl Iterator<Item = (&Url, usize)> {
self.ring
.iter()
.flat_map(|h| self.normal[h].iter())
.chain(&self.low)
.map(|(u, d, _)| (u, *d))
}
}