pageman 0.1.0

Disk-based page manager/store
Documentation
use eyre::{bail, ensure, Context, OptionExt, Result};
use lru::LruCache;
use std::{hash::Hash, num::NonZeroUsize};

enum Value<V> {
    Detached,
    Pinned(V),
    Regular(V),
}

impl<V> Value<V> {
    fn get(&self) -> Result<&V> {
        Ok(match self {
            Value::Detached => bail!("Value is detached"),
            Value::Pinned(v) => v,
            Value::Regular(v) => v,
        })
    }

    fn get_mut(&mut self) -> Result<&mut V> {
        Ok(match self {
            Value::Detached => bail!("Value is detached"),
            Value::Pinned(v) => v,
            Value::Regular(v) => v,
        })
    }

    fn out(self, pin_count: &mut usize) -> Result<V> {
        Ok(match self {
            Value::Detached => bail!("Value is detached"),
            Value::Pinned(v) => {
                *pin_count -= 1;
                v
            }
            Value::Regular(v) => v,
        })
    }
}

pub struct Cache<K: Clone + PartialEq + Eq + Hash, V> {
    entries: LruCache<K, Value<V>>,
    detach_count: usize,
    pin_count: usize,
}

impl<K: Clone + PartialEq + Eq + Hash, V> Cache<K, V> {
    pub fn attach(&mut self, key: &K, val: V) -> Result<()> {
        let (key, d_val) = self
            .entries
            .pop_entry(key)
            .ok_or_eyre("cannot attach: requested entry is not cached")?;
        ensure!(
            matches!(d_val, Value::Detached),
            "only detached entries can be attached"
        );
        self.entries.push(key, Value::Regular(val));
        self.detach_count -= 1;
        Ok(())
    }

    // NOTE: This will also detach pinned entries
    pub fn detach(&mut self, key: &K) -> Result<V> {
        let (key, val) = self
            .entries
            .pop_entry(key)
            .ok_or_eyre("cannot detach: requested entry is not cached")?;
        self.entries.push(key, Value::Detached);
        self.detach_count += 1;
        val.out(&mut self.pin_count)
    }

    // NOTE: This will also evict pinned entries
    pub fn empty(&mut self) -> Result<Vec<(K, V)>> {
        let mut entries = Vec::new();
        while let Some((key, val)) = self.entries.pop_lru() {
            entries.push((key, val.out(&mut self.pin_count)?))
        }
        Ok(entries)
    }

    pub fn get(&mut self, key: &K) -> Result<&V> {
        self.entries
            .get(key)
            .ok_or_eyre("cannot get: requested entry is not cached")?
            .get()
    }

    pub fn get_mut(&mut self, key: &K) -> Result<&mut V> {
        self.entries
            .get_mut(key)
            .ok_or_eyre("cannot get_mut: requested entry is not cached")?
            .get_mut()
    }

    pub fn is_cached(&self, key: &K) -> bool {
        self.entries.contains(key)
    }

    pub fn is_pinned(&self, key: &K) -> bool {
        matches!(self.entries.peek(key), Some(Value::Pinned(_)))
    }

    #[allow(dead_code)]
    pub fn len(&self) -> usize {
        self.entries.len()
    }

    pub fn margin(&self) -> usize {
        self.entries.cap().get() - self.detach_count - self.pin_count
    }

    pub fn new(cap: NonZeroUsize) -> Self {
        Self {
            entries: LruCache::new(cap),
            detach_count: 0,
            pin_count: 0,
        }
    }

    pub fn pin(&mut self, key: &K) -> Result<()> {
        let (key, val) = self
            .entries
            .pop_entry(key)
            .ok_or_eyre("cannot pin: requested entry is not cached")?;
        self.entries
            .push(key, Value::Pinned(val.out(&mut self.pin_count)?));
        self.pin_count += 1;
        Ok(())
    }

    pub fn pop(&mut self, key: &K) -> Result<V> {
        self.entries
            .pop(key)
            .ok_or_eyre("cannot pop: requested entry is not cached")?
            .out(&mut self.pin_count)
    }

    #[allow(dead_code)]
    pub fn pop_lru(&mut self) -> Option<(K, V)> {
        self.refresh().expect("failed to refresh the cache");
        self.entries
            .pop_lru()
            .map(|(k, v)| (k, v.out(&mut 0).unwrap()))
    }

    pub fn push(&mut self, key: K, val: V) -> Result<Option<(K, V)>> {
        self.refresh()
            .wrap_err("failed to refresh the cache when pushing new entry")?;
        Ok(
            if let Some((k, v)) = self.entries.push(key, Value::Regular(val)) {
                Some((k, v.out(&mut 0)?))
            } else {
                None
            },
        )
    }

    fn refresh(&mut self) -> Result<()> {
        if self.entries.len() < self.entries.cap().get() {
            return Ok(());
        }
        ensure!(
            self.margin() > 0,
            "cache overflow -- capacity: {}, detached: {}, pinned: {}",
            self.entries.cap().get(),
            self.detach_count,
            self.pin_count
        );
        loop {
            let key = match self.entries.peek_lru() {
                Some((key, val)) if !matches!(val, Value::Regular(_)) => key.clone(),
                _ => break,
            };
            self.entries.promote(&key);
        }
        Ok(())
    }

    pub fn unpin(&mut self, key: &K) -> Result<()> {
        let (key, val) = self
            .entries
            .pop_entry(key)
            .ok_or_eyre("cannot unpin: requested entry is not cached")?;
        self.entries
            .push(key, Value::Regular(val.out(&mut self.pin_count)?));
        Ok(())
    }
}