monovm-whois-rust 1.0.0

//! [`MemoryCache`]: an in-process cache with a TTL and a size cap.

use std::collections::HashMap;
use std::sync::Mutex;
use std::time::{Duration, Instant};

use crate::cache::{CacheKey, ResponseCache};
use crate::transport::RawResponse;

/// Default lifetime of an entry: five minutes.
///
/// Short because registration state changes and a stale "available" is the
/// expensive kind of wrong; long enough to absorb the repeated queries a bulk
/// check makes to the same few servers.
pub const DEFAULT_TTL: Duration = Duration::from_secs(300);

/// Default cap on entries.
pub const DEFAULT_CAPACITY: usize = 1024;

#[derive(Debug, Clone)]
struct Entry {
    response: RawResponse,
    stored_at: Instant,
    /// Monotonic counter, used to evict the least recently used entry.
    touched: u64,
}

#[derive(Debug, Default)]
struct Store {
    entries: HashMap<CacheKey, Entry>,
    clock: u64,
}

/// An in-memory cache with a time-to-live and least-recently-used eviction.
///
/// ```
/// use std::time::Duration;
/// use monovm_whois::cache::{CacheKey, MemoryCache, ResponseCache};
///
/// let cache = MemoryCache::with_ttl(Duration::from_secs(60));
/// assert!(cache.get(&CacheKey::new("whois.example", "example.com")).is_none());
/// assert_eq!(cache.len(), Some(0));
/// ```
#[derive(Debug)]
pub struct MemoryCache {
    store: Mutex<Store>,
    ttl: Duration,
    capacity: usize,
}

impl MemoryCache {
    /// A cache with the default TTL and capacity.
    pub fn new() -> Self {
        MemoryCache::with_ttl(DEFAULT_TTL)
    }

    /// A cache with an explicit TTL.
    pub fn with_ttl(ttl: Duration) -> Self {
        MemoryCache {
            store: Mutex::new(Store::default()),
            ttl,
            capacity: DEFAULT_CAPACITY,
        }
    }

    /// A cache with an explicit TTL and entry cap.
    pub fn with_ttl_and_capacity(ttl: Duration, capacity: usize) -> Self {
        MemoryCache {
            store: Mutex::new(Store::default()),
            ttl,
            capacity: capacity.max(1),
        }
    }

    /// The lifetime of an entry.
    pub fn ttl(&self) -> Duration {
        self.ttl
    }

    /// The entry cap.
    pub fn capacity(&self) -> usize {
        self.capacity
    }

    /// Drop expired entries. Called automatically when the cache is full; exposed
    /// for a caller that would rather reclaim memory on a timer.
    pub fn purge_expired(&self) {
        let mut store = self.lock();
        let ttl = self.ttl;
        store
            .entries
            .retain(|_, entry| entry.stored_at.elapsed() < ttl);
    }

    fn lock(&self) -> std::sync::MutexGuard<'_, Store> {
        // A panic while holding the lock would otherwise turn a cache into a
        // permanent failure; the data behind it is disposable by definition.
        self.store
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner())
    }

    fn evict_one(store: &mut Store) {
        let victim = store
            .entries
            .iter()
            .min_by_key(|(_, entry)| entry.touched)
            .map(|(key, _)| key.clone());

        if let Some(key) = victim {
            store.entries.remove(&key);
        }
    }
}

impl Default for MemoryCache {
    fn default() -> Self {
        MemoryCache::new()
    }
}

impl ResponseCache for MemoryCache {
    fn get(&self, key: &CacheKey) -> Option<RawResponse> {
        let mut store = self.lock();
        store.clock += 1;
        let now = store.clock;

        let entry = store.entries.get_mut(key)?;
        if entry.stored_at.elapsed() >= self.ttl {
            store.entries.remove(key);
            return None;
        }

        entry.touched = now;
        Some(entry.response.clone())
    }

    fn put(&self, key: CacheKey, response: RawResponse) {
        let mut store = self.lock();
        store.clock += 1;
        let touched = store.clock;

        if store.entries.len() >= self.capacity && !store.entries.contains_key(&key) {
            // Reclaiming expired entries first usually makes room without
            // discarding anything a caller could still have used.
            let ttl = self.ttl;
            store
                .entries
                .retain(|_, entry| entry.stored_at.elapsed() < ttl);

            while store.entries.len() >= self.capacity {
                MemoryCache::evict_one(&mut store);
            }
        }

        store.entries.insert(
            key,
            Entry {
                response,
                stored_at: Instant::now(),
                touched,
            },
        );
    }

    fn clear(&self) {
        self.lock().entries.clear();
    }

    fn len(&self) -> Option<usize> {
        let store = self.lock();
        let ttl = self.ttl;
        Some(
            store
                .entries
                .values()
                .filter(|entry| entry.stored_at.elapsed() < ttl)
                .count(),
        )
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::registry::Endpoint;
    use crate::transport::ResponseKind;

    fn response(text: &str) -> RawResponse {
        RawResponse::new(
            Endpoint::whois("whois.example"),
            ResponseKind::WhoisText,
            text,
            Duration::from_millis(1),
        )
    }

    fn key(name: &str) -> CacheKey {
        CacheKey::new("whois.example", name)
    }

    #[test]
    fn stores_and_returns() {
        let cache = MemoryCache::new();
        cache.put(key("a.com"), response("record a"));

        assert_eq!(cache.get(&key("a.com")).unwrap().text(), "record a");
        assert!(cache.get(&key("b.com")).is_none());
        assert_eq!(cache.len(), Some(1));
    }

    #[test]
    fn expired_entries_are_not_returned() {
        let cache = MemoryCache::with_ttl(Duration::ZERO);
        cache.put(key("a.com"), response("record a"));

        assert!(cache.get(&key("a.com")).is_none());
        assert_eq!(cache.len(), Some(0));
    }

    #[test]
    fn a_live_entry_survives_a_long_ttl() {
        let cache = MemoryCache::with_ttl(Duration::from_secs(3600));
        cache.put(key("a.com"), response("record a"));
        assert!(cache.get(&key("a.com")).is_some());
    }

    #[test]
    fn clear_empties_the_cache() {
        let cache = MemoryCache::new();
        cache.put(key("a.com"), response("a"));
        cache.put(key("b.com"), response("b"));
        cache.clear();

        assert_eq!(cache.len(), Some(0));
        assert!(cache.get(&key("a.com")).is_none());
    }

    #[test]
    fn the_capacity_is_respected() {
        let cache = MemoryCache::with_ttl_and_capacity(Duration::from_secs(3600), 2);
        cache.put(key("a.com"), response("a"));
        cache.put(key("b.com"), response("b"));
        cache.put(key("c.com"), response("c"));

        assert_eq!(cache.len(), Some(2));
    }

    #[test]
    fn eviction_takes_the_least_recently_used() {
        let cache = MemoryCache::with_ttl_and_capacity(Duration::from_secs(3600), 2);
        cache.put(key("a.com"), response("a"));
        cache.put(key("b.com"), response("b"));

        // Touch `a`, making `b` the coldest entry.
        assert!(cache.get(&key("a.com")).is_some());
        cache.put(key("c.com"), response("c"));

        assert!(
            cache.get(&key("a.com")).is_some(),
            "recently used entry was evicted"
        );
        assert!(cache.get(&key("c.com")).is_some());
        assert!(cache.get(&key("b.com")).is_none(), "coldest entry survived");
    }

    #[test]
    fn overwriting_a_key_does_not_evict() {
        let cache = MemoryCache::with_ttl_and_capacity(Duration::from_secs(3600), 2);
        cache.put(key("a.com"), response("first"));
        cache.put(key("b.com"), response("b"));
        cache.put(key("a.com"), response("second"));

        assert_eq!(cache.len(), Some(2));
        assert_eq!(cache.get(&key("a.com")).unwrap().text(), "second");
        assert!(cache.get(&key("b.com")).is_some());
    }

    #[test]
    fn purge_expired_reclaims_only_the_dead() {
        let cache = MemoryCache::with_ttl(Duration::ZERO);
        cache.put(key("a.com"), response("a"));
        cache.purge_expired();
        assert_eq!(cache.len(), Some(0));
    }

    #[test]
    fn a_zero_capacity_is_clamped_to_one() {
        let cache = MemoryCache::with_ttl_and_capacity(Duration::from_secs(60), 0);
        assert_eq!(cache.capacity(), 1);
        cache.put(key("a.com"), response("a"));
        assert_eq!(cache.len(), Some(1));
    }

    #[test]
    fn is_usable_from_several_threads() {
        let cache = std::sync::Arc::new(MemoryCache::new());
        let handles: Vec<_> = (0..8)
            .map(|index| {
                let cache = std::sync::Arc::clone(&cache);
                std::thread::spawn(move || {
                    let name = format!("{index}.com");
                    cache.put(key(&name), response(&name));
                    cache.get(&key(&name))
                })
            })
            .collect();

        for handle in handles {
            assert!(handle.join().unwrap().is_some());
        }
        assert_eq!(cache.len(), Some(8));
    }
}