aquatic-crawler/src/ban.rs

use chrono::{DateTime, Local};
use librqbit::dht::Id20;
use std::{collections::HashMap, time::Duration};

pub struct Ban {
    index: HashMap<Id20, DateTime<Local>>,
    timeout: u64,
}

impl Ban {
    pub fn init(timeout: u64, capacity: usize) -> Self {
        Self {
            index: HashMap::with_capacity(capacity),
            timeout,
        }
    }

    pub fn has(&self, key: &Id20) -> bool {
        self.index.contains_key(key)
    }

    pub fn total(&self) -> usize {
        self.index.len()
    }

    pub fn update(&mut self, time: DateTime<Local>) {
        self.index.retain(|i, &mut expires| {
            if time > expires {
                log::debug!(
                    "remove ban for `{}` by the timeout expiration ({expires})",
                    i.as_string()
                );
                false
            } else {
                true
            }
        })
    }

    pub fn add(&mut self, key: Id20) -> DateTime<Local> {
        let t = Local::now() + Duration::from_secs(self.index.len() as u64 * self.timeout);
        assert!(self.index.insert(key, t).is_none());
        t
    }
}