grafos_cache/

elastic.rs

//! Elastic sharded hash map backed by fabric memory.
//!
//! [`ElasticShardSet`] distributes key-value pairs across multiple shards,
//! each backed by its own [`grafos_std::mem::MemLease`] +
//! [`FabricHashMap`]. Shards can be
//! added or removed at runtime, and [`rehash`](ElasticShardSet::rehash)
//! redistributes entries across the current shard set.
//!
//! Shard selection uses `hash(key) % shard_count` (FNV-1a after postcard
//! serialization, matching [`FabricHashMap`]'s internal hash function).

extern crate alloc;
use alloc::vec::Vec;

use grafos_collections::map::FabricHashMap;
use grafos_std::error::{FabricError, Result};

use serde::{de::DeserializeOwned, Serialize};

/// A set of shards that can grow and shrink at runtime.
///
/// Each shard is a [`FabricHashMap`] backed by its own
/// [`grafos_std::mem::MemLease`].
/// Keys are routed to shards via `hash(key) % shard_count`.
///
/// # Example
///
/// ```rust
/// use grafos_cache::elastic::ElasticShardSet;
///
/// # grafos_std::host::reset_mock();
/// # grafos_std::host::mock_set_fbmu_arena_size(65536);
/// let mut shards: ElasticShardSet<u32, u32> =
///     ElasticShardSet::new(2, 64, 8, 8)?;
///
/// shards.put(&1, &100)?;
/// shards.put(&2, &200)?;
/// assert_eq!(shards.get(&1)?, Some(100));
/// assert_eq!(shards.shard_count(), 2);
///
/// // Grow by one shard and redistribute
/// shards.add_shard()?;
/// shards.rehash()?;
/// assert_eq!(shards.shard_count(), 3);
/// assert_eq!(shards.get(&1)?, Some(100));
/// # Ok::<(), grafos_std::FabricError>(())
/// ```
pub struct ElasticShardSet<K, V> {
    shards: Vec<FabricHashMap<K, V>>,
    buckets_per_shard: usize,
    key_stride: usize,
    val_stride: usize,
}

impl<K, V> ElasticShardSet<K, V>
where
    K: Serialize + DeserializeOwned + PartialEq + Clone,
    V: Serialize + DeserializeOwned + Clone,
{
    /// Create a new shard set with `initial_shards` shards.
    ///
    /// Each shard is a [`FabricHashMap`] allocated via
    /// [`grafos_std::mem::MemBuilder`] with
    /// `buckets_per_shard` buckets. `key_stride` and `val_stride` control the
    /// maximum serialized size of keys and values respectively.
    ///
    /// # Errors
    ///
    /// Returns [`FabricError::CapacityExceeded`] if the host cannot provide
    /// enough memory for the requested shards.
    pub fn new(
        initial_shards: usize,
        buckets_per_shard: usize,
        key_stride: usize,
        val_stride: usize,
    ) -> Result<Self> {
        if initial_shards == 0 {
            return Err(FabricError::CapacityExceeded);
        }
        let mut shards = Vec::with_capacity(initial_shards);
        for _ in 0..initial_shards {
            let map = FabricHashMap::with_capacity(buckets_per_shard, key_stride, val_stride)?;
            shards.push(map);
        }
        Ok(ElasticShardSet {
            shards,
            buckets_per_shard,
            key_stride,
            val_stride,
        })
    }

    /// Look up a value by key.
    ///
    /// Hashes the key to select the correct shard, then performs a lookup
    /// within that shard.
    pub fn get(&self, key: &K) -> Result<Option<V>> {
        let idx = self.shard_index(key)?;
        self.shards[idx].get(key)
    }

    /// Insert or update a key-value pair.
    ///
    /// Hashes the key to select the correct shard, then inserts into that
    /// shard.
    pub fn put(&mut self, key: &K, value: &V) -> Result<()> {
        let idx = self.shard_index(key)?;
        self.shards[idx].insert(key, value)?;
        Ok(())
    }

    /// Remove a key-value pair. Returns `true` if the key existed.
    pub fn remove(&mut self, key: &K) -> Result<bool> {
        let idx = self.shard_index(key)?;
        Ok(self.shards[idx].remove(key)?.is_some())
    }

    /// Add one shard to the set.
    ///
    /// The new shard is empty. Call [`rehash`](Self::rehash) afterwards to
    /// redistribute existing entries across all shards.
    pub fn add_shard(&mut self) -> Result<()> {
        let map =
            FabricHashMap::with_capacity(self.buckets_per_shard, self.key_stride, self.val_stride)?;
        self.shards.push(map);
        Ok(())
    }

    /// Remove the last shard from the set.
    ///
    /// All entries across every shard are drained and redistributed among
    /// the remaining shards (since `hash % shard_count` changes for all
    /// keys when the shard count changes). The removed shard's lease is
    /// released on drop.
    ///
    /// # Errors
    ///
    /// Returns [`FabricError::CapacityExceeded`] if the set has only one
    /// shard (cannot shrink below 1).
    pub fn remove_shard(&mut self) -> Result<()> {
        if self.shards.len() <= 1 {
            return Err(FabricError::CapacityExceeded);
        }
        // Drain all entries from every shard (routing changes for all keys)
        let mut all_entries: Vec<(K, V)> = Vec::new();
        for shard in &self.shards {
            for entry in shard.iter() {
                all_entries.push(entry?);
            }
        }
        // Drop the last shard
        self.shards.pop();
        // Replace remaining shards with fresh empty ones
        let new_count = self.shards.len();
        self.shards.clear();
        for _ in 0..new_count {
            let map = FabricHashMap::with_capacity(
                self.buckets_per_shard,
                self.key_stride,
                self.val_stride,
            )?;
            self.shards.push(map);
        }
        // Re-insert all entries with the new shard count
        for (k, v) in &all_entries {
            let idx = self.shard_index(k)?;
            self.shards[idx].insert(k, v)?;
        }
        Ok(())
    }

    /// Redistribute all entries across the current shard set.
    ///
    /// Drains every shard, then re-inserts all entries using the current
    /// shard count. This is useful after [`add_shard`](Self::add_shard) to
    /// balance load.
    pub fn rehash(&mut self) -> Result<()> {
        // Collect all entries from all shards
        let mut all_entries: Vec<(K, V)> = Vec::new();
        for shard in &self.shards {
            for entry in shard.iter() {
                all_entries.push(entry?);
            }
        }

        // Replace all shards with fresh empty ones
        let shard_count = self.shards.len();
        self.shards.clear();
        for _ in 0..shard_count {
            let map = FabricHashMap::with_capacity(
                self.buckets_per_shard,
                self.key_stride,
                self.val_stride,
            )?;
            self.shards.push(map);
        }

        // Re-insert all entries
        for (k, v) in &all_entries {
            let idx = self.shard_index(k)?;
            self.shards[idx].insert(k, v)?;
        }
        Ok(())
    }

    /// Number of shards in the set.
    pub fn shard_count(&self) -> usize {
        self.shards.len()
    }

    /// Total number of entries across all shards.
    pub fn len(&self) -> Result<usize> {
        let mut total = 0;
        for shard in &self.shards {
            total += shard.len();
        }
        Ok(total)
    }

    /// Returns `true` if no entries exist across all shards.
    pub fn is_empty(&self) -> Result<bool> {
        Ok(self.len()? == 0)
    }

    /// Compute the shard index for a given key using FNV-1a.
    fn shard_index(&self, key: &K) -> Result<usize> {
        let key_bytes = postcard::to_allocvec(key).map_err(|_| FabricError::IoError(-1))?;
        let hash = fnv1a(&key_bytes);
        Ok((hash as usize) % self.shards.len())
    }
}

/// FNV-1a hash (matches FabricHashMap's internal hash).
fn fnv1a(data: &[u8]) -> u64 {
    let mut hash: u64 = 0xcbf29ce484222325;
    for &b in data {
        hash ^= b as u64;
        hash = hash.wrapping_mul(0x100000001b3);
    }
    hash
}

#[cfg(test)]
mod tests {
    use super::*;
    use grafos_std::host;

    fn setup() {
        host::reset_mock();
        host::mock_set_fbmu_arena_size(65536);
    }

    #[test]
    fn create_put_get_remove() {
        setup();
        let mut shards: ElasticShardSet<u32, u32> = ElasticShardSet::new(2, 64, 8, 8).expect("new");

        assert!(shards.is_empty().unwrap());
        shards.put(&1, &100).expect("put 1");
        shards.put(&2, &200).expect("put 2");
        shards.put(&3, &300).expect("put 3");

        assert_eq!(shards.len().unwrap(), 3);
        assert_eq!(shards.get(&1).unwrap(), Some(100));
        assert_eq!(shards.get(&2).unwrap(), Some(200));
        assert_eq!(shards.get(&3).unwrap(), Some(300));
        assert_eq!(shards.get(&99).unwrap(), None);

        assert!(shards.remove(&2).unwrap());
        assert!(!shards.remove(&99).unwrap());
        assert_eq!(shards.len().unwrap(), 2);
        assert_eq!(shards.get(&2).unwrap(), None);
    }

    #[test]
    fn correct_shard_routing() {
        setup();
        let mut shards: ElasticShardSet<u32, u32> = ElasticShardSet::new(4, 64, 8, 8).expect("new");

        // Insert many keys and verify they all round-trip
        for i in 0..20u32 {
            shards.put(&i, &(i * 10)).expect("put");
        }
        assert_eq!(shards.len().unwrap(), 20);

        for i in 0..20u32 {
            assert_eq!(shards.get(&i).unwrap(), Some(i * 10), "key {i}");
        }
    }

    #[test]
    fn add_shard_entries_still_accessible() {
        setup();
        let mut shards: ElasticShardSet<u32, u32> = ElasticShardSet::new(2, 64, 8, 8).expect("new");

        shards.put(&10, &1000).expect("put");
        shards.put(&20, &2000).expect("put");
        shards.put(&30, &3000).expect("put");

        shards.add_shard().expect("add_shard");
        assert_eq!(shards.shard_count(), 3);

        // Before rehash, entries are still in their old shards
        // After rehash, they should be redistributed but still accessible
        shards.rehash().expect("rehash");
        assert_eq!(shards.len().unwrap(), 3);
        assert_eq!(shards.get(&10).unwrap(), Some(1000));
        assert_eq!(shards.get(&20).unwrap(), Some(2000));
        assert_eq!(shards.get(&30).unwrap(), Some(3000));
    }

    #[test]
    fn remove_shard_entries_migrated() {
        setup();
        let mut shards: ElasticShardSet<u32, u32> = ElasticShardSet::new(3, 64, 8, 8).expect("new");

        for i in 0..10u32 {
            shards.put(&i, &(i * 100)).expect("put");
        }
        assert_eq!(shards.shard_count(), 3);
        assert_eq!(shards.len().unwrap(), 10);

        shards.remove_shard().expect("remove_shard");
        assert_eq!(shards.shard_count(), 2);
        assert_eq!(shards.len().unwrap(), 10);

        for i in 0..10u32 {
            assert_eq!(shards.get(&i).unwrap(), Some(i * 100), "key {i}");
        }
    }

    #[test]
    fn remove_shard_below_one_fails() {
        setup();
        let mut shards: ElasticShardSet<u32, u32> = ElasticShardSet::new(1, 64, 8, 8).expect("new");

        assert_eq!(
            shards.remove_shard().unwrap_err(),
            FabricError::CapacityExceeded
        );
    }

    #[test]
    fn rehash_redistributes() {
        setup();
        let mut shards: ElasticShardSet<u32, u32> = ElasticShardSet::new(2, 64, 8, 8).expect("new");

        for i in 0..8u32 {
            shards.put(&i, &(i + 1)).expect("put");
        }

        // Add two more shards and rehash
        shards.add_shard().expect("add");
        shards.add_shard().expect("add");
        assert_eq!(shards.shard_count(), 4);

        shards.rehash().expect("rehash");
        assert_eq!(shards.len().unwrap(), 8);

        // All entries still accessible
        for i in 0..8u32 {
            assert_eq!(shards.get(&i).unwrap(), Some(i + 1), "key {i}");
        }
    }

    #[test]
    fn zero_initial_shards_fails() {
        setup();
        let result: Result<ElasticShardSet<u32, u32>> = ElasticShardSet::new(0, 64, 8, 8);
        match result {
            Err(FabricError::CapacityExceeded) => {}
            Err(e) => panic!("expected CapacityExceeded, got {e:?}"),
            Ok(_) => panic!("expected error, got Ok"),
        }
    }
}