liquid_cache_storage/cache/cache_policies/

s3_fifo.rs

//! S3 FIFO cache policy implementation

use std::collections::{HashMap, HashSet, VecDeque};
use std::fmt;
use std::sync::Mutex;

use crate::cache::{cached_batch::CachedBatchType, utils::EntryID};
use crate::cache_policies::CachePolicy;

type EntryFreq = u8;

#[derive(Debug, Default)]
struct S3FifoInternalState {
    small: VecDeque<EntryID>,
    main: VecDeque<EntryID>,
    ghost: VecDeque<EntryID>,
    ghost_set: HashSet<EntryID>,

    frequency: HashMap<EntryID, EntryFreq>,

    small_queue_size: usize,
    main_queue_size: usize,
    total_size: usize,
}

impl S3FifoInternalState {
    fn cap_frequency(freq: u8) -> u8 {
        std::cmp::min(freq, 3)
    }

    fn inc_frequency(&mut self, entry_id: &EntryID) {
        if let Some(freq) = self.frequency.get_mut(entry_id) {
            *freq = Self::cap_frequency(*freq + 1);
        }
    }

    fn dec_frequency(&mut self, entry_id: &EntryID) {
        if let Some(freq) = self.frequency.get_mut(entry_id) {
            *freq = freq.saturating_sub(1);
        }
    }

    fn inc_small_queue_size(&mut self, size: usize) {
        self.small_queue_size += size;
        self.total_size += size;
    }

    fn dec_small_queue_size(&mut self, size: usize) {
        self.small_queue_size -= size;
        self.total_size -= size;
    }

    fn inc_main_queue_size(&mut self, size: usize) {
        self.main_queue_size += size;
        self.total_size += size;
    }

    fn dec_main_queue_size(&mut self, size: usize) {
        self.main_queue_size -= size;
        self.total_size -= size;
    }

    fn small_queue_fraction(&self) -> f32 {
        if self.total_size == 0 {
            0.0
        } else {
            self.small_queue_size as f32 / self.total_size as f32
        }
    }

    fn check_if_entry_exists_in_small_or_main(&self, entry_id: &EntryID) -> bool {
        self.frequency.contains_key(entry_id) && !self.ghost_set.contains(entry_id)
    }
}

/// The policy that implements object size aware S3Fifo algorithm using Deque.
#[derive(Default)]
pub struct S3FifoPolicy {
    state: Mutex<S3FifoInternalState>,
}

impl fmt::Debug for S3FifoPolicy {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("S3FifoPolicy")
            .field("state", &self.state)
            .finish()
    }
}

unsafe impl Send for S3FifoPolicy {}
unsafe impl Sync for S3FifoPolicy {}

impl S3FifoPolicy {
    /// Create a new [`S3FifoPolicy`].
    pub fn new() -> Self {
        Self {
            state: Mutex::new(S3FifoInternalState::default()),
        }
    }

    fn entry_size(&self, _entry_id: &EntryID) -> usize {
        1
    }

    fn evict_from_small(&self, state: &mut S3FifoInternalState) -> Option<EntryID> {
        let mut is_evicted = false;
        let mut victim: Option<EntryID> = None;

        while !is_evicted && !state.small.is_empty() {
            let Some(element) = state.small.pop_back() else {
                break;
            };
            let freq = state.frequency.get(&element).copied().unwrap_or(0);
            let entry_size = self.entry_size(&element);
            state.dec_small_queue_size(entry_size);

            if freq > 1 {
                state.main.push_front(element);
                state.inc_main_queue_size(entry_size);
                state.frequency.insert(element, 0);
            } else {
                // Move to ghost queue
                state.ghost.push_front(element);
                state.ghost_set.insert(element);
                state.frequency.remove(&element);
                is_evicted = true;
                victim = Some(element);
            }
        }
        victim
    }

    fn evict_from_main(&self, state: &mut S3FifoInternalState) -> Option<EntryID> {
        let mut is_evicted = false;
        let mut victim: Option<EntryID> = None;

        while !is_evicted && !state.main.is_empty() {
            let Some(element) = state.main.pop_back() else {
                break;
            };

            let freq = state.frequency.get(&element).copied().unwrap_or(0);
            let entry_size = self.entry_size(&element);
            state.dec_main_queue_size(entry_size);

            if freq > 0 {
                state.main.push_front(element);
                state.dec_frequency(&element);
                state.inc_main_queue_size(entry_size);
            } else {
                state.frequency.remove(&element);
                is_evicted = true;
                victim = Some(element);
            }
        }
        victim
    }
}

impl CachePolicy for S3FifoPolicy {
    fn find_victim(&self, cnt: usize) -> Vec<EntryID> {
        let mut state = self.state.lock().unwrap();
        let mut advices = Vec::with_capacity(cnt);
        let threshold_for_small_eviction = 0.1;
        while advices.len() < cnt && state.total_size > 0 {
            let victim = if !state.small.is_empty()
                && state.small_queue_fraction() >= threshold_for_small_eviction
            {
                self.evict_from_small(&mut state)
            } else {
                self.evict_from_main(&mut state)
            };

            if let Some(v) = victim {
                advices.push(v);
            }
        }
        advices
    }

    fn notify_insert(&self, entry_id: &EntryID, _batch_type: CachedBatchType) {
        let mut state = self.state.lock().unwrap();
        let entry_size = self.entry_size(entry_id);

        if state.check_if_entry_exists_in_small_or_main(entry_id) {
            state.inc_frequency(entry_id);
        } else if state.ghost_set.contains(entry_id) {
            state.ghost_set.remove(entry_id);
            state.ghost.retain(|x| *x != *entry_id);
            state.main.push_front(*entry_id);
            state.inc_main_queue_size(entry_size);
        } else {
            state.small.push_front(*entry_id);
            state.inc_small_queue_size(entry_size);
            state.frequency.insert(*entry_id, 0);
        }
    }

    fn notify_access(&self, entry_id: &EntryID, _batch_type: CachedBatchType) {
        let mut state = self.state.lock().unwrap();
        if state.check_if_entry_exists_in_small_or_main(entry_id) {
            state.inc_frequency(entry_id);
        }
    }
}

impl Drop for S3FifoPolicy {
    fn drop(&mut self) {
        let mut state = self.state.lock().unwrap();
        state.small.clear();
        state.main.clear();
        state.ghost.clear();
        state.ghost_set.clear();
        state.frequency.clear();
        state.total_size = 0;
        state.small_queue_size = 0;
        state.main_queue_size = 0;
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::cache::utils::EntryID;

    fn entry(id: usize) -> EntryID {
        id.into()
    }

    #[test]
    fn test_s3fifo_basic_insert_eviction() {
        let policy = S3FifoPolicy::new();
        let e1 = entry(1);
        let e2 = entry(2);
        let e3 = entry(3);

        policy.notify_insert(&e1, CachedBatchType::MemoryArrow);
        policy.notify_insert(&e2, CachedBatchType::MemoryArrow);
        policy.notify_insert(&e3, CachedBatchType::MemoryArrow);

        let evicted = policy.find_victim(1);
        assert_eq!(evicted.len(), 1);
    }

    #[test]
    fn test_s3fifo_frequency_increase() {
        let policy = S3FifoPolicy::new();
        let e1 = entry(1);
        policy.notify_insert(&e1, CachedBatchType::MemoryArrow);
        policy.notify_access(&e1, CachedBatchType::MemoryArrow);

        let state = policy.state.lock().unwrap();
        assert_eq!(*state.frequency.get(&e1).unwrap(), 1);
    }

    #[test]
    fn test_s3fifo_eviction_order() {
        let policy = S3FifoPolicy::new();
        let e1 = entry(1);
        let e2 = entry(2);

        policy.notify_insert(&e1, CachedBatchType::MemoryArrow);
        policy.notify_insert(&e2, CachedBatchType::MemoryArrow);

        let evicted = policy.find_victim(1);
        assert_eq!(evicted[0], e1);
    }

    #[test]
    fn test_s3fifo_ghost_promote() {
        let policy = S3FifoPolicy::new();
        let e1 = entry(1);

        policy.notify_insert(&e1, CachedBatchType::MemoryArrow);
        let evicted = policy.find_victim(1);
        assert_eq!(evicted[0], e1);

        // Re-insert evicted entry from ghost
        policy.notify_insert(&e1, CachedBatchType::MemoryArrow);
        let state = policy.state.lock().unwrap();
        assert!(state.main.contains(&e1));
        assert!(!state.ghost_set.contains(&e1));
    }

    #[test]
    fn test_s3fifo_size_aware_fraction() {
        let policy = S3FifoPolicy::new();
        let e1 = entry(20);
        let e2 = entry(30);

        policy.notify_insert(&e1, CachedBatchType::MemoryArrow);
        policy.notify_insert(&e2, CachedBatchType::MemoryArrow);

        let state = policy.state.lock().unwrap();
        assert_eq!(state.small_queue_size, 2);
        assert_eq!(state.small_queue_fraction(), 1.0);
    }

    #[test]
    fn test_insert_and_access_updates_freq() {
        let policy = S3FifoPolicy::new();
        let e1 = entry(1);

        policy.notify_insert(&e1, CachedBatchType::MemoryArrow);
        policy.notify_access(&e1, CachedBatchType::MemoryArrow);
        policy.notify_access(&e1, CachedBatchType::MemoryArrow);

        let state = policy.state.lock().unwrap();
        assert_eq!(*state.frequency.get(&e1).unwrap(), 2); // capped at 3
    }

    #[test]
    fn test_freq_cap_at_three() {
        let policy = S3FifoPolicy::new();
        let e1 = entry(1);

        policy.notify_insert(&e1, CachedBatchType::MemoryArrow);
        for _ in 0..10 {
            policy.notify_access(&e1, CachedBatchType::MemoryArrow);
        }

        let state = policy.state.lock().unwrap();
        assert_eq!(*state.frequency.get(&e1).unwrap(), 3);
    }

    #[test]
    fn test_eviction_from_s_to_ghost() {
        let policy = S3FifoPolicy::new();
        let e1 = entry(1);

        policy.notify_insert(&e1, CachedBatchType::MemoryArrow);
        let evicted = policy.find_victim(1);

        assert_eq!(evicted[0], e1);
        let state = policy.state.lock().unwrap();
        assert!(state.ghost_set.contains(&e1));
        assert!(state.ghost.contains(&e1));
    }

    #[test]
    fn test_eviction_from_main_and_reinsertion_logic() {
        let policy = S3FifoPolicy::new();
        let e1 = entry(1);
        let e2 = entry(2);
        let _e3 = entry(3);

        let mut state = policy.state.lock().unwrap();
        state.main.push_front(e1);
        state.frequency.insert(e1, 1);
        state.main_queue_size += 1;

        state.main.push_front(e2);
        state.frequency.insert(e2, 2);
        state.main_queue_size += 1;
        state.total_size += 2;
        drop(state);

        let evicted = policy.find_victim(2);
        assert_eq!(evicted.len(), 2);
    }

    #[test]
    fn test_s3fifo_reinsert_does_not_duplicate_entry() {
        let policy = S3FifoPolicy::new();
        let e1 = entry(1);
        let e2 = entry(2);

        policy.notify_insert(&e1, CachedBatchType::MemoryArrow);
        policy.notify_insert(&e2, CachedBatchType::MemoryArrow);

        policy.notify_insert(&e1, CachedBatchType::MemoryArrow);

        let state = policy.state.lock().unwrap();
        let occurrences = state.small.iter().filter(|&&id| id == e1).count()
            + state.main.iter().filter(|&&id| id == e1).count();
        assert_eq!(occurrences, 1);
        assert_eq!(*state.frequency.get(&e1).unwrap(), 1);
    }
}