hayson_engine/table/

cache.rs

use std::collections::HashMap;
use std::sync::{Arc, Mutex};

use serde_json::Value;

use crate::schema::HaysonRecord;

/// Approximate document size at or above which a record is cached in the large-document tier
/// instead of the main one — four kilobytes, which no record of scalars and short strings reaches
/// and every sealed `cypher` envelope or paragraph-length `fulltext` field does.
pub const DEFAULT_LARGE_BYTES: usize = 4 * 1024;

/// How many parsed documents a table keeps in memory in front of its log — the read cache
/// [`crate::table::Table::get`] answers a repeated read from instead of re-reading and re-parsing
/// the line the registry points at.
///
/// `Disabled` is the default everywhere, including [`crate::database::OpenOptions`]: caching a
/// document is a memory bet a caller has to place on purpose.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum CachePolicy {
    #[default]
    Disabled,
    Enabled {
        /// Most documents held below the `large_bytes` line, evicting least-recently-read first.
        entries: usize,
        /// Most documents held at or above the `large_bytes` line — its own, smaller budget, so a
        /// handful of big records cannot evict every small one.
        large_entries: usize,
        /// The size, as [`approx_bytes`] measures it, that separates the two tiers.
        large_bytes: usize,
    },
}

impl CachePolicy {
    /// `entries` small documents, an eighth as many large ones (never fewer than one) and
    /// `DEFAULT_LARGE_BYTES` between the tiers — the shape for a caller who has a rough memory
    /// budget rather than an opinion about tiers. `0` is [`CachePolicy::Disabled`].
    pub fn with_entries(entries: usize) -> Self {
        if entries == 0 {
            return CachePolicy::Disabled;
        }
        CachePolicy::Enabled {
            entries,
            large_entries: (entries / 8).max(1),
            large_bytes: DEFAULT_LARGE_BYTES,
        }
    }
}

/// Roughly how much a parsed document holds, without serializing it: string and field-name lengths
/// counted for real, every scalar counted as eight bytes. It only has to tell a long `fulltext`
/// field or a base64 `cypher` envelope from a handful of numbers, which is all the tier split asks
/// of it.
fn approx_bytes(record: &HaysonRecord) -> usize {
    record
        .iter()
        .map(|(field, value)| field.len() + value_bytes(value))
        .sum()
}

fn value_bytes(value: &Value) -> usize {
    match value {
        Value::Null | Value::Bool(_) | Value::Number(_) => 8,
        Value::String(text) => text.len(),
        Value::Array(items) => items.iter().map(value_bytes).sum::<usize>() + 8,
        Value::Object(fields) => {
            fields
                .iter()
                .map(|(field, value)| field.len() + value_bytes(value))
                .sum::<usize>()
                + 8
        }
    }
}

struct Node {
    key: Arc<str>,
    record: HaysonRecord,
    prev: Option<usize>,
    next: Option<usize>,
}

/// A fixed-capacity LRU over a slab of nodes: the map finds a key, the nodes are threaded
/// most-recently-used first, and eviction hands its slot back to `free` for the next insert.
///
/// Every operation is O(1) on purpose — a cache sized in the tens of thousands must not turn a
/// read miss into a scan of everything cached, which is the shape a "least recently used wins"
/// timestamp sweep would have.
struct Lru {
    capacity: usize,
    map: HashMap<Arc<str>, usize>,
    nodes: Vec<Node>,
    free: Vec<usize>,
    head: Option<usize>,
    tail: Option<usize>,
}

impl Lru {
    fn new(capacity: usize) -> Self {
        Lru {
            capacity,
            map: HashMap::new(),
            nodes: Vec::new(),
            free: Vec::new(),
            head: None,
            tail: None,
        }
    }

    fn get(&mut self, key: &str) -> Option<&HaysonRecord> {
        let at = *self.map.get(key)?;
        self.unlink(at);
        self.push_front(at);
        Some(&self.nodes[at].record)
    }

    fn insert(&mut self, key: &str, record: HaysonRecord) {
        if self.capacity == 0 {
            return;
        }
        if let Some(&at) = self.map.get(key) {
            self.nodes[at].record = record;
            self.unlink(at);
            self.push_front(at);
            return;
        }
        if self.map.len() >= self.capacity {
            self.evict_tail();
        }
        let key: Arc<str> = Arc::from(key);
        let node = Node {
            key: Arc::clone(&key),
            record,
            prev: None,
            next: None,
        };
        let at = match self.free.pop() {
            Some(at) => {
                self.nodes[at] = node;
                at
            }
            None => {
                self.nodes.push(node);
                self.nodes.len() - 1
            }
        };
        self.map.insert(key, at);
        self.push_front(at);
    }

    fn remove(&mut self, key: &str) {
        let Some(at) = self.map.remove(key) else { return };
        self.unlink(at);
        self.release(at);
    }

    fn clear(&mut self) {
        self.map.clear();
        self.nodes.clear();
        self.free.clear();
        self.head = None;
        self.tail = None;
    }

    #[cfg(test)]
    fn len(&self) -> usize {
        self.map.len()
    }

    fn evict_tail(&mut self) {
        let Some(at) = self.tail else { return };
        let key = Arc::clone(&self.nodes[at].key);
        self.map.remove(&key);
        self.unlink(at);
        self.release(at);
    }

    /// Empties the slot without shrinking the slab, so the document it held stops taking memory
    /// before another insert reuses the slot.
    fn release(&mut self, at: usize) {
        self.nodes[at].key = Arc::from("");
        self.nodes[at].record = HaysonRecord::new();
        self.free.push(at);
    }

    // NOTE: only sound for a node currently in the list — `prev == None` is read as "this is head".
    fn unlink(&mut self, at: usize) {
        let prev = self.nodes[at].prev.take();
        let next = self.nodes[at].next.take();
        match prev {
            Some(prev) => self.nodes[prev].next = next,
            None => self.head = next,
        }
        match next {
            Some(next) => self.nodes[next].prev = prev,
            None => self.tail = prev,
        }
    }

    fn push_front(&mut self, at: usize) {
        self.nodes[at].prev = None;
        self.nodes[at].next = self.head;
        if let Some(head) = self.head {
            self.nodes[head].prev = Some(at);
        }
        self.head = Some(at);
        if self.tail.is_none() {
            self.tail = Some(at);
        }
    }
}

struct Tiers {
    small: Lru,
    large: Lru,
}

/// The bounded cache of parsed documents a table reads through, in the two tiers
/// [`CachePolicy`] describes. A key lives in exactly one of them, chosen by [`approx_bytes`] of the
/// record when it was cached.
///
/// Interior mutability, because `Table::get` takes `&self` and a cache hit still has to move the
/// key it hit to the front of the list. A `Mutex` rather than a `RefCell`: `Database` is
/// `Send + Sync` and pinned as such by a test. Nothing but this module's own code ever runs under
/// the lock, so it is a leaf — no other lock is ever taken while it is held.
pub(crate) struct RecordCache {
    policy: CachePolicy,
    tiers: Mutex<Tiers>,
}

impl RecordCache {
    pub(crate) fn new(policy: CachePolicy) -> Self {
        let (entries, large_entries) = match policy {
            CachePolicy::Disabled => (0, 0),
            CachePolicy::Enabled {
                entries, large_entries, ..
            } => (entries, large_entries),
        };
        RecordCache {
            policy,
            tiers: Mutex::new(Tiers {
                small: Lru::new(entries),
                large: Lru::new(large_entries),
            }),
        }
    }

    pub(crate) fn policy(&self) -> CachePolicy {
        self.policy
    }

    fn enabled(&self) -> bool {
        self.policy != CachePolicy::Disabled
    }

    /// The cached document under `key`, moved to the front of its tier. Cloned out: the record
    /// cannot leave the lock by reference, and cloning a parsed document is still what a read of
    /// the log would end up doing after slicing and parsing the line.
    pub(crate) fn get(&self, key: &str) -> Option<HaysonRecord> {
        if !self.enabled() {
            return None;
        }
        let mut tiers = self.locked();
        if let Some(record) = tiers.small.get(key) {
            return Some(record.clone());
        }
        tiers.large.get(key).cloned()
    }

    pub(crate) fn insert(&self, key: &str, record: &HaysonRecord) {
        let CachePolicy::Enabled { large_bytes, .. } = self.policy else {
            return;
        };
        let large = approx_bytes(record) >= large_bytes;
        let mut tiers = self.locked();
        let tier = if large { &mut tiers.large } else { &mut tiers.small };
        tier.insert(key, record.clone());
    }

    pub(crate) fn invalidate(&self, key: &str) {
        if !self.enabled() {
            return;
        }
        let mut tiers = self.locked();
        tiers.small.remove(key);
        tiers.large.remove(key);
    }

    pub(crate) fn clear(&self) {
        if !self.enabled() {
            return;
        }
        let mut tiers = self.locked();
        tiers.small.clear();
        tiers.large.clear();
    }

    #[cfg(test)]
    pub(crate) fn len(&self) -> (usize, usize) {
        let tiers = self.locked();
        (tiers.small.len(), tiers.large.len())
    }

    // NOTE: a poisoned cache is taken anyway — dropping an invalidation here would serve stale docs.
    fn locked(&self) -> std::sync::MutexGuard<'_, Tiers> {
        self.tiers.lock().unwrap_or_else(|poisoned| poisoned.into_inner())
    }
}

#[cfg(test)]
mod tests {
    use serde_json::json;

    use super::*;

    fn record(name: &str) -> HaysonRecord {
        let mut record = HaysonRecord::new();
        record.insert("name".to_string(), json!(name));
        record
    }

    fn big(bytes: usize) -> HaysonRecord {
        let mut record = HaysonRecord::new();
        record.insert("body".to_string(), json!("x".repeat(bytes)));
        record
    }

    fn cache_of(entries: usize, large_entries: usize) -> RecordCache {
        RecordCache::new(CachePolicy::Enabled {
            entries,
            large_entries,
            large_bytes: DEFAULT_LARGE_BYTES,
        })
    }

    #[test]
    fn a_disabled_cache_answers_nothing_and_keeps_nothing() {
        let cache = RecordCache::new(CachePolicy::Disabled);
        cache.insert("u1", &record("Ania"));
        assert_eq!(cache.get("u1"), None);
        assert_eq!(cache.len(), (0, 0));
    }

    #[test]
    fn with_entries_derives_a_smaller_large_tier_and_reads_zero_as_disabled() {
        assert_eq!(CachePolicy::with_entries(0), CachePolicy::Disabled);
        assert_eq!(
            CachePolicy::with_entries(64),
            CachePolicy::Enabled {
                entries: 64,
                large_entries: 8,
                large_bytes: DEFAULT_LARGE_BYTES,
            }
        );
        assert_eq!(
            CachePolicy::with_entries(1),
            CachePolicy::Enabled {
                entries: 1,
                large_entries: 1,
                large_bytes: DEFAULT_LARGE_BYTES,
            }
        );
    }

    #[test]
    fn what_was_inserted_reads_back_until_it_is_invalidated() {
        let cache = cache_of(4, 2);
        cache.insert("u1", &record("Ania"));
        assert_eq!(cache.get("u1"), Some(record("Ania")));
        cache.invalidate("u1");
        assert_eq!(cache.get("u1"), None);
    }

    #[test]
    fn a_second_insert_on_one_key_replaces_the_document_without_taking_a_second_slot() {
        let cache = cache_of(4, 2);
        cache.insert("u1", &record("Ania"));
        cache.insert("u1", &record("Zosia"));
        assert_eq!(cache.get("u1"), Some(record("Zosia")));
        assert_eq!(cache.len(), (1, 0));
    }

    #[test]
    fn a_full_tier_evicts_the_least_recently_read_key_and_keeps_the_one_just_read() {
        let cache = cache_of(2, 1);
        cache.insert("u1", &record("Ania"));
        cache.insert("u2", &record("Piotr"));
        // NOTE: this read is what makes u2, not u1, the eviction candidate.
        assert_eq!(cache.get("u1"), Some(record("Ania")));
        cache.insert("u3", &record("Zosia"));

        assert_eq!(cache.get("u2"), None);
        assert_eq!(cache.get("u1"), Some(record("Ania")));
        assert_eq!(cache.get("u3"), Some(record("Zosia")));
        assert_eq!(cache.len(), (2, 0));
    }

    #[test]
    fn a_large_document_lands_in_its_own_tier_and_evicts_only_large_ones() {
        let cache = cache_of(2, 1);
        cache.insert("u1", &record("Ania"));
        cache.insert("d1", &big(DEFAULT_LARGE_BYTES));
        cache.insert("d2", &big(DEFAULT_LARGE_BYTES));

        assert_eq!(cache.len(), (1, 1));
        assert_eq!(cache.get("d1"), None);
        assert_eq!(cache.get("d2"), Some(big(DEFAULT_LARGE_BYTES)));
        assert_eq!(cache.get("u1"), Some(record("Ania")));
    }

    #[test]
    fn a_document_one_byte_under_the_line_is_still_a_small_one() {
        let cache = cache_of(2, 1);
        cache.insert("d1", &big(DEFAULT_LARGE_BYTES - "body".len() - 1));
        assert_eq!(cache.len(), (1, 0));
    }

    #[test]
    fn clear_empties_both_tiers() {
        let cache = cache_of(2, 1);
        cache.insert("u1", &record("Ania"));
        cache.insert("d1", &big(DEFAULT_LARGE_BYTES));
        cache.clear();
        assert_eq!(cache.len(), (0, 0));
        assert_eq!(cache.get("u1"), None);
        assert_eq!(cache.get("d1"), None);
    }

    #[test]
    fn invalidating_a_key_nobody_cached_changes_nothing() {
        let cache = cache_of(2, 1);
        cache.insert("u1", &record("Ania"));
        cache.invalidate("u404");
        assert_eq!(cache.len(), (1, 0));
        assert_eq!(cache.get("u1"), Some(record("Ania")));
    }

    #[test]
    fn a_tier_keeps_its_capacity_across_eviction_and_reinsertion() {
        let cache = cache_of(2, 1);
        for round in 0..10 {
            cache.insert(&format!("u{round}"), &record("Ania"));
            assert!(cache.len().0 <= 2);
        }
        assert_eq!(cache.len(), (2, 0));
        assert_eq!(cache.get("u9"), Some(record("Ania")));
        assert_eq!(cache.get("u8"), Some(record("Ania")));
        assert_eq!(cache.get("u7"), None);
    }

    #[test]
    fn approx_bytes_counts_strings_for_real_and_nested_values_too() {
        let mut record = HaysonRecord::new();
        record.insert("id".to_string(), json!("u1"));
        record.insert("tags".to_string(), json!(["aaa", "bbb"]));
        record.insert("meta".to_string(), json!({ "note": "hello" }));
        assert_eq!(approx_bytes(&record), 2 + 2 + 4 + (3 + 3 + 8) + 4 + (4 + 5 + 8));
    }
}