grafos_kv/

lib.rs

//! Key-value store over leased fabric resources.
//!
//! [`FabricKvStore`] provides a tiered key-value store where the hot tier
//! lives in leased DRAM (via [`FabricHashMap`]) and the cold tier (behind
//! the `persistence` feature) spills to leased block storage.
//!
//! Each key carries a TTL. On access, the TTL is refreshed. The
//! [`tick()`](FabricKvStore::tick) method evicts expired entries. When all
//! keys in the hot tier expire, the backing shard lease can be left to
//! expire naturally, freeing the remote memory.
//!
//! # Quick start
//!
//! ```rust
//! use grafos_kv::{KvBuilder, FabricKvStore};
//!
//! # grafos_std::host::reset_mock();
//! # grafos_std::host::mock_set_fbmu_arena_size(65536);
//! let mut kv: FabricKvStore = KvBuilder::new()
//!     .hot_buckets(64)
//!     .default_ttl_secs(300)
//!     .build()?;
//!
//! kv.put(b"hello", b"world")?;
//! assert_eq!(kv.get(b"hello")?, Some(b"world".to_vec()));
//! # Ok::<(), grafos_std::FabricError>(())
//! ```
//!
//! # Feature flags
//!
//! | Feature | Default | Effect |
//! |---------|---------|--------|
//! | `std` | Yes | Enables `std` in grafos-std |
//! | `persistence` | No | Enables cold tier (block storage spill/promote) |

#![cfg_attr(not(feature = "std"), no_std)]

extern crate alloc;

use alloc::vec::Vec;

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

use serde::{Deserialize, Serialize};

#[cfg(feature = "persistence")]
use grafos_std::block::{BlockBuilder, BlockLease, BLOCK_SIZE};

/// Maximum serialized size for keys in the hot-tier FabricHashMap.
const KEY_STRIDE: usize = 128;
/// Maximum serialized size for values (KvEntry) in the hot-tier FabricHashMap.
const VALUE_STRIDE: usize = 512;

/// Internal entry stored in the hot tier alongside the raw value bytes.
#[derive(Serialize, Deserialize, Clone, Debug, PartialEq)]
struct KvEntry {
    value: Vec<u8>,
    created_at: u64,
    ttl_secs: u32,
    access_count: u64,
}

/// Builder for [`FabricKvStore`].
///
/// # Example
///
/// ```rust
/// use grafos_kv::KvBuilder;
///
/// # grafos_std::host::reset_mock();
/// # grafos_std::host::mock_set_fbmu_arena_size(131072);
/// let mut kv = KvBuilder::new()
///     .hot_buckets(128)
///     .default_ttl_secs(60)
///     .build()?;
/// # Ok::<(), grafos_std::FabricError>(())
/// ```
pub struct KvBuilder {
    hot_buckets: usize,
    default_ttl_secs: u32,
    key_stride: usize,
    value_stride: usize,
    #[cfg(feature = "persistence")]
    cold_num_blocks: Option<u64>,
}

impl KvBuilder {
    pub fn new() -> Self {
        KvBuilder {
            hot_buckets: 64,
            default_ttl_secs: 300,
            key_stride: KEY_STRIDE,
            value_stride: VALUE_STRIDE,
            #[cfg(feature = "persistence")]
            cold_num_blocks: None,
        }
    }

    /// Set the number of buckets for the hot-tier hash map.
    pub fn hot_buckets(mut self, n: usize) -> Self {
        self.hot_buckets = n;
        self
    }

    /// Set the default TTL in seconds for keys without an explicit TTL.
    pub fn default_ttl_secs(mut self, secs: u32) -> Self {
        self.default_ttl_secs = secs;
        self
    }

    /// Set the key stride (max serialized key size in bytes).
    pub fn key_stride(mut self, stride: usize) -> Self {
        self.key_stride = stride;
        self
    }

    /// Set the value stride (max serialized KvEntry size in bytes).
    pub fn value_stride(mut self, stride: usize) -> Self {
        self.value_stride = stride;
        self
    }

    /// Enable the cold tier with the given number of blocks.
    #[cfg(feature = "persistence")]
    pub fn cold_num_blocks(mut self, n: u64) -> Self {
        self.cold_num_blocks = Some(n);
        self
    }

    /// Build the [`FabricKvStore`].
    pub fn build(self) -> Result<FabricKvStore> {
        let hot =
            FabricHashMap::with_capacity(self.hot_buckets, self.key_stride, self.value_stride)?;

        #[cfg(feature = "persistence")]
        let cold = if let Some(num_blocks) = self.cold_num_blocks {
            Some(ColdTier::new(num_blocks)?)
        } else {
            None
        };

        Ok(FabricKvStore {
            hot,
            default_ttl_secs: self.default_ttl_secs,
            #[cfg(feature = "persistence")]
            cold,
        })
    }

    /// Recover a [`FabricKvStore`] from existing hot and cold leases.
    ///
    /// The hot tier is reconstructed by reading the `FabricHashMap` header
    /// from the memory lease. The cold tier is reconstructed by scanning
    /// the append-only log to find the write cursor.
    ///
    /// This enables full crash recovery when both leases survived.
    #[cfg(feature = "persistence")]
    pub fn from_leases(
        hot_lease: grafos_std::mem::MemLease,
        cold_lease: BlockLease,
    ) -> Result<FabricKvStore> {
        let hot = FabricHashMap::from_lease(hot_lease)?;
        let cold = Some(ColdTier::from_lease(cold_lease)?);
        Ok(FabricKvStore {
            hot,
            default_ttl_secs: 300,
            cold,
        })
    }

    /// Recover a [`FabricKvStore`] from a surviving cold-tier block lease.
    ///
    /// A fresh hot tier is allocated (via `FabricHashMap::with_capacity`).
    /// Cold-tier entries are promoted to the hot tier on first access.
    ///
    /// This is the common crash-recovery path: the tasklet's hot-tier
    /// memory lease expired with the tasklet, but the cold-tier block
    /// lease survived.
    #[cfg(feature = "persistence")]
    pub fn from_cold(cold_lease: BlockLease, hot_buckets: usize) -> Result<FabricKvStore> {
        let hot = FabricHashMap::with_capacity(hot_buckets, KEY_STRIDE, VALUE_STRIDE)?;
        let cold = Some(ColdTier::from_lease(cold_lease)?);
        Ok(FabricKvStore {
            hot,
            default_ttl_secs: 300,
            cold,
        })
    }
}

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

/// Cold tier: append-only log of key-value pairs in block storage.
///
/// Layout per record:
/// ```text
/// key_len: u32 (4) | key: [u8; key_len] | value_len: u32 (4) | value: [u8; value_len] | flags: u8 (1)
/// ```
///
/// Flags: 0x00 = live, 0x01 = tombstoned.
#[cfg(feature = "persistence")]
struct ColdTier {
    lease: BlockLease,
    /// Next byte offset within the block device for appending.
    write_cursor: u64,
}

#[cfg(feature = "persistence")]
const COLD_FLAG_LIVE: u8 = 0x00;
#[cfg(feature = "persistence")]
const COLD_FLAG_TOMBSTONE: u8 = 0x01;

#[cfg(feature = "persistence")]
impl ColdTier {
    fn new(num_blocks: u64) -> Result<Self> {
        let lease = BlockBuilder::new().min_blocks(num_blocks).acquire()?;
        Ok(ColdTier {
            lease,
            write_cursor: 0,
        })
    }

    /// Recover a cold tier from an existing block lease.
    ///
    /// Scans the append-only log to find the write cursor position. This
    /// enables crash recovery: a replacement tasklet can attach to a
    /// surviving block lease and reconstruct the cold tier index.
    fn from_lease(lease: BlockLease) -> Result<Self> {
        let capacity = lease.block().num_blocks() * BLOCK_SIZE as u64;
        let mut cursor: u64 = 0;

        // Scan forward to find the end of written records.
        // Each record: key_len(4) + key + val_len(4) + val + flags(1).
        // A key_len of 0 with no valid record structure indicates end-of-log.
        loop {
            if cursor + 4 > capacity {
                break;
            }
            // Read key_len (may span a block boundary).
            let mut key_len_bytes = [0u8; 4];
            let mut bytes_read = 0;
            let mut off = cursor;
            while bytes_read < 4 {
                let bi = off / BLOCK_SIZE as u64;
                let bo = (off % BLOCK_SIZE as u64) as usize;
                let b = lease.block().read_block(bi)?;
                let avail = BLOCK_SIZE - bo;
                let take = (4 - bytes_read).min(avail);
                key_len_bytes[bytes_read..bytes_read + take].copy_from_slice(&b[bo..bo + take]);
                bytes_read += take;
                off += take as u64;
            }
            let key_len = u32::from_le_bytes(key_len_bytes) as u64;

            // A key_len of 0 means we've hit unwritten space (end of log).
            if key_len == 0 {
                break;
            }
            // Sanity: key_len shouldn't be absurdly large
            if key_len > capacity {
                break;
            }
            cursor += 4 + key_len;

            // Read val_len
            if cursor + 4 > capacity {
                break;
            }
            let mut val_len_bytes = [0u8; 4];
            bytes_read = 0;
            off = cursor;
            while bytes_read < 4 {
                let bi = off / BLOCK_SIZE as u64;
                let bo = (off % BLOCK_SIZE as u64) as usize;
                let b = lease.block().read_block(bi)?;
                let avail = BLOCK_SIZE - bo;
                let take = (4 - bytes_read).min(avail);
                val_len_bytes[bytes_read..bytes_read + take].copy_from_slice(&b[bo..bo + take]);
                bytes_read += take;
                off += take as u64;
            }
            let val_len = u32::from_le_bytes(val_len_bytes) as u64;
            if val_len > capacity {
                break;
            }
            cursor += 4 + val_len;

            // Skip flags byte
            if cursor + 1 > capacity {
                break;
            }
            cursor += 1;
        }

        Ok(ColdTier {
            lease,
            write_cursor: cursor,
        })
    }

    fn capacity_bytes(&self) -> u64 {
        self.lease.block().num_blocks() * BLOCK_SIZE as u64
    }

    fn append_live(&mut self, key: &[u8], value: &[u8]) -> Result<()> {
        self.append_with_flag(key, value, COLD_FLAG_LIVE)
    }

    fn append_tombstone(&mut self, key: &[u8]) -> Result<()> {
        self.append_with_flag(key, &[], COLD_FLAG_TOMBSTONE)
    }

    fn append_with_flag(&mut self, key: &[u8], value: &[u8], flag: u8) -> Result<()> {
        let record_len = 4 + key.len() + 4 + value.len() + 1;
        if self.write_cursor + record_len as u64 > self.capacity_bytes() {
            return Err(FabricError::CapacityExceeded);
        }

        let mut record = Vec::with_capacity(record_len);
        record.extend_from_slice(&(key.len() as u32).to_le_bytes());
        record.extend_from_slice(key);
        record.extend_from_slice(&(value.len() as u32).to_le_bytes());
        record.extend_from_slice(value);
        record.push(flag);

        // Write record across blocks
        let mut offset = self.write_cursor;
        let mut remaining = &record[..];
        while !remaining.is_empty() {
            let block_idx = offset / BLOCK_SIZE as u64;
            let block_offset = (offset % BLOCK_SIZE as u64) as usize;
            let space = BLOCK_SIZE - block_offset;
            let chunk_len = remaining.len().min(space);

            // Read-modify-write the block
            let mut block = self.lease.block().read_block(block_idx)?;
            block[block_offset..block_offset + chunk_len].copy_from_slice(&remaining[..chunk_len]);
            self.lease.block().write_block(block_idx, &block)?;

            remaining = &remaining[chunk_len..];
            offset += chunk_len as u64;
        }

        self.write_cursor += record_len as u64;
        Ok(())
    }

    fn scan(&self, target_key: &[u8]) -> Result<Option<Vec<u8>>> {
        let mut offset: u64 = 0;
        let mut result: Option<Vec<u8>> = None;

        while offset < self.write_cursor {
            // Read key_len
            let key_len_bytes = self.read_bytes(offset, 4)?;
            let key_len = u32::from_le_bytes([
                key_len_bytes[0],
                key_len_bytes[1],
                key_len_bytes[2],
                key_len_bytes[3],
            ]) as usize;
            offset += 4;

            let key = self.read_bytes(offset, key_len)?;
            offset += key_len as u64;

            let val_len_bytes = self.read_bytes(offset, 4)?;
            let val_len = u32::from_le_bytes([
                val_len_bytes[0],
                val_len_bytes[1],
                val_len_bytes[2],
                val_len_bytes[3],
            ]) as usize;
            offset += 4;

            let value = self.read_bytes(offset, val_len)?;
            offset += val_len as u64;

            let flags_bytes = self.read_bytes(offset, 1)?;
            let flags = flags_bytes[0];
            offset += 1;

            if key == target_key {
                if flags == COLD_FLAG_LIVE {
                    result = Some(value);
                } else {
                    result = None;
                }
            }
        }

        Ok(result)
    }

    fn read_bytes(&self, offset: u64, len: usize) -> Result<Vec<u8>> {
        let mut buf = vec![0u8; len];
        let mut buf_offset = 0;
        let mut disk_offset = offset;

        while buf_offset < len {
            let block_idx = disk_offset / BLOCK_SIZE as u64;
            let block_off = (disk_offset % BLOCK_SIZE as u64) as usize;
            let space = BLOCK_SIZE - block_off;
            let chunk_len = (len - buf_offset).min(space);

            let block = self.lease.block().read_block(block_idx)?;
            buf[buf_offset..buf_offset + chunk_len]
                .copy_from_slice(&block[block_off..block_off + chunk_len]);

            buf_offset += chunk_len;
            disk_offset += chunk_len as u64;
        }

        Ok(buf)
    }

    /// Compact the cold tier by removing tombstoned entries.
    /// Rewrites all live entries from the beginning.
    fn compact(&mut self) -> Result<()> {
        // Collect all live entries
        let mut live_entries: Vec<(Vec<u8>, Vec<u8>)> = Vec::new();
        let mut offset: u64 = 0;

        // Track the latest state per key (last-writer-wins)
        while offset < self.write_cursor {
            let key_len_bytes = self.read_bytes(offset, 4)?;
            let key_len = u32::from_le_bytes([
                key_len_bytes[0],
                key_len_bytes[1],
                key_len_bytes[2],
                key_len_bytes[3],
            ]) as usize;
            offset += 4;

            let key = self.read_bytes(offset, key_len)?;
            offset += key_len as u64;

            let val_len_bytes = self.read_bytes(offset, 4)?;
            let val_len = u32::from_le_bytes([
                val_len_bytes[0],
                val_len_bytes[1],
                val_len_bytes[2],
                val_len_bytes[3],
            ]) as usize;
            offset += 4;

            let value = self.read_bytes(offset, val_len)?;
            offset += val_len as u64;

            let flags_bytes = self.read_bytes(offset, 1)?;
            let flags = flags_bytes[0];
            offset += 1;

            if flags == COLD_FLAG_LIVE {
                // Remove any previous entry for this key
                live_entries.retain(|(k, _)| k != &key);
                live_entries.push((key, value));
            } else {
                live_entries.retain(|(k, _)| k != &key);
            }
        }

        // Zero the write cursor and rewrite
        self.write_cursor = 0;
        for (key, value) in &live_entries {
            self.append_live(key, value)?;
        }

        Ok(())
    }
}

/// A key-value store backed by leased fabric resources.
///
/// The hot tier stores entries in a [`FabricHashMap`] over leased DRAM.
/// Each entry tracks creation time, TTL, and access count. The
/// [`tick()`](FabricKvStore::tick) method evicts expired entries.
///
/// With the `persistence` feature, a cold tier backs overflow to block
/// storage. On a hot-tier miss, the cold tier is scanned and matching
/// entries are promoted back.
pub struct FabricKvStore {
    hot: FabricHashMap<Vec<u8>, KvEntry>,
    default_ttl_secs: u32,
    #[cfg(feature = "persistence")]
    cold: Option<ColdTier>,
}

impl FabricKvStore {
    /// Store a key-value pair with the default TTL.
    pub fn put(&mut self, key: &[u8], value: &[u8]) -> Result<()> {
        self.put_with_ttl(key, value, self.default_ttl_secs)
    }

    /// Store a key-value pair with an explicit TTL in seconds.
    pub fn put_with_ttl(&mut self, key: &[u8], value: &[u8], ttl_secs: u32) -> Result<()> {
        let now = host::unix_time_secs();
        let entry = KvEntry {
            value: value.to_vec(),
            created_at: now,
            ttl_secs,
            access_count: 0,
        };
        let key_vec = key.to_vec();

        match self.hot.insert(&key_vec, &entry) {
            Ok(_) => Ok(()),
            #[cfg(feature = "persistence")]
            Err(FabricError::CapacityExceeded) => {
                // Hot tier full — spill to cold tier
                if let Some(ref mut cold) = self.cold {
                    cold.append_live(key, value)?;
                    Ok(())
                } else {
                    Err(FabricError::CapacityExceeded)
                }
            }
            Err(e) => Err(e),
        }
    }

    /// Retrieve a value by key.
    ///
    /// Checks the hot tier first. On a miss (with `persistence` enabled),
    /// scans the cold tier and promotes the entry back to the hot tier.
    ///
    /// Accessing a key refreshes its TTL (resets `created_at` to now).
    pub fn get(&mut self, key: &[u8]) -> Result<Option<Vec<u8>>> {
        let key_vec = key.to_vec();
        let now = host::unix_time_secs();

        if let Some(mut entry) = self.hot.get(&key_vec)? {
            // Check if expired
            if now >= entry.created_at.saturating_add(entry.ttl_secs as u64) {
                // Expired — remove it
                self.hot.remove(&key_vec)?;
                return Ok(None);
            }
            // Refresh TTL on access
            entry.access_count += 1;
            entry.created_at = now;
            self.hot.insert(&key_vec, &entry)?;
            return Ok(Some(entry.value));
        }

        #[cfg(feature = "persistence")]
        if let Some(ref cold) = self.cold {
            if let Some(value) = cold.scan(key)? {
                // Promote to hot tier
                let entry = KvEntry {
                    value: value.clone(),
                    created_at: now,
                    ttl_secs: self.default_ttl_secs,
                    access_count: 1,
                };
                // Best-effort promote; if hot is full, still return the value
                let _ = self.hot.insert(&key_vec, &entry);
                return Ok(Some(value));
            }
        }

        Ok(None)
    }

    /// Delete a key. Returns `true` if the key existed.
    pub fn delete(&mut self, key: &[u8]) -> Result<bool> {
        let key_vec = key.to_vec();
        let removed = self.hot.remove(&key_vec)?.is_some();

        #[cfg(feature = "persistence")]
        if let Some(ref mut cold) = self.cold {
            // Append a tombstone to cold tier
            cold.append_tombstone(key)?;
            // If something was found in cold but not hot, we still consider it deleted
            // The tombstone prevents future cold-tier scans from returning it
        }

        Ok(removed)
    }

    /// Check if a key exists (without reading the value).
    pub fn exists(&mut self, key: &[u8]) -> Result<bool> {
        // Use get() which also handles expiry checking
        Ok(self.get(key)?.is_some())
    }

    /// Return all non-expired keys in the hot tier.
    pub fn keys(&self) -> Result<Vec<Vec<u8>>> {
        let now = host::unix_time_secs();
        let mut result = Vec::new();
        for item in self.hot.iter() {
            let (key, entry) = item?;
            if now < entry.created_at.saturating_add(entry.ttl_secs as u64) {
                result.push(key);
            }
        }
        Ok(result)
    }

    /// Return the remaining TTL for a key in seconds, or `None` if the key
    /// does not exist or is expired.
    pub fn ttl(&self, key: &[u8]) -> Result<Option<u32>> {
        let key_vec = key.to_vec();
        let now = host::unix_time_secs();

        if let Some(entry) = self.hot.get(&key_vec)? {
            let expires_at = entry.created_at.saturating_add(entry.ttl_secs as u64);
            if now < expires_at {
                return Ok(Some((expires_at - now) as u32));
            }
        }
        Ok(None)
    }

    /// Scan the hot tier and evict expired entries.
    pub fn tick(&mut self) -> Result<usize> {
        let now = host::unix_time_secs();
        let mut expired_keys = Vec::new();

        for item in self.hot.iter() {
            let (key, entry) = item?;
            if now >= entry.created_at.saturating_add(entry.ttl_secs as u64) {
                expired_keys.push(key);
            }
        }

        let count = expired_keys.len();
        for key in &expired_keys {
            self.hot.remove(key)?;
        }
        Ok(count)
    }

    /// Store a serializable struct as a value.
    pub fn put_struct<T: Serialize>(&mut self, key: &[u8], value: &T) -> Result<()> {
        let bytes = postcard::to_allocvec(value).map_err(|_| FabricError::IoError(-1))?;
        self.put(key, &bytes)
    }

    /// Store a serializable struct with an explicit TTL.
    pub fn put_struct_with_ttl<T: Serialize>(
        &mut self,
        key: &[u8],
        value: &T,
        ttl_secs: u32,
    ) -> Result<()> {
        let bytes = postcard::to_allocvec(value).map_err(|_| FabricError::IoError(-1))?;
        self.put_with_ttl(key, &bytes, ttl_secs)
    }

    /// Retrieve and deserialize a struct by key.
    pub fn get_struct<T: serde::de::DeserializeOwned>(&mut self, key: &[u8]) -> Result<Option<T>> {
        match self.get(key)? {
            Some(bytes) => {
                let val: T = postcard::from_bytes(&bytes).map_err(|_| FabricError::IoError(-1))?;
                Ok(Some(val))
            }
            None => Ok(None),
        }
    }

    /// Compact the cold tier, removing tombstoned and duplicate entries.
    ///
    /// Only available with the `persistence` feature.
    #[cfg(feature = "persistence")]
    pub fn compact_cold(&mut self) -> Result<()> {
        if let Some(ref mut cold) = self.cold {
            cold.compact()
        } else {
            Ok(())
        }
    }
}

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

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

    #[test]
    fn put_get_delete_roundtrip() {
        setup();
        let mut kv = KvBuilder::new()
            .hot_buckets(64)
            .default_ttl_secs(300)
            .build()
            .expect("build");

        kv.put(b"key1", b"value1").expect("put");
        kv.put(b"key2", b"value2").expect("put");

        assert_eq!(kv.get(b"key1").expect("get"), Some(b"value1".to_vec()));
        assert_eq!(kv.get(b"key2").expect("get"), Some(b"value2".to_vec()));
        assert_eq!(kv.get(b"key3").expect("get"), None);

        assert!(kv.delete(b"key1").expect("delete"));
        assert_eq!(kv.get(b"key1").expect("get"), None);
        assert!(!kv.delete(b"nonexistent").expect("delete"));
    }

    #[test]
    fn put_struct_get_struct() {
        setup();
        let mut kv = KvBuilder::new().build().expect("build");

        #[derive(Serialize, Deserialize, Debug, PartialEq)]
        struct MyData {
            count: u64,
            name: Vec<u8>,
        }

        let data = MyData {
            count: 42,
            name: b"test".to_vec(),
        };
        kv.put_struct(b"structured", &data).expect("put_struct");

        let retrieved: MyData = kv
            .get_struct(b"structured")
            .expect("get_struct")
            .expect("some");
        assert_eq!(retrieved, data);
    }

    #[test]
    fn ttl_expiry() {
        setup();
        let mut kv = KvBuilder::new()
            .default_ttl_secs(10)
            .build()
            .expect("build");

        kv.put(b"ephemeral", b"data").expect("put");
        assert_eq!(kv.get(b"ephemeral").expect("get"), Some(b"data".to_vec()));
        assert!(kv.ttl(b"ephemeral").expect("ttl").is_some());

        // Advance time past TTL
        host::mock_advance_time_secs(11);

        assert_eq!(kv.get(b"ephemeral").expect("get"), None);
        assert_eq!(kv.ttl(b"ephemeral").expect("ttl"), None);
    }

    #[test]
    fn ttl_renewal_on_access() {
        setup();
        let mut kv = KvBuilder::new()
            .default_ttl_secs(10)
            .build()
            .expect("build");

        kv.put(b"renew-me", b"data").expect("put");

        // Advance time by 8 seconds (within TTL)
        host::mock_advance_time_secs(8);

        // Access the key — this refreshes created_at to now
        let val = kv.get(b"renew-me").expect("get");
        assert_eq!(val, Some(b"data".to_vec()));

        // Advance by another 8 seconds (16 total from start, but only 8 from last access)
        host::mock_advance_time_secs(8);

        // Should still be alive because access renewed it
        let val = kv.get(b"renew-me").expect("get");
        assert_eq!(val, Some(b"data".to_vec()));
    }

    #[test]
    fn exists_and_keys() {
        setup();
        let mut kv = KvBuilder::new()
            .default_ttl_secs(300)
            .build()
            .expect("build");

        kv.put(b"a", b"1").expect("put");
        kv.put(b"b", b"2").expect("put");
        kv.put(b"c", b"3").expect("put");

        assert!(kv.exists(b"a").expect("exists"));
        assert!(kv.exists(b"b").expect("exists"));
        assert!(!kv.exists(b"nonexistent").expect("exists"));

        let mut keys = kv.keys().expect("keys");
        keys.sort();
        assert_eq!(keys, vec![b"a".to_vec(), b"b".to_vec(), b"c".to_vec()]);
    }

    #[test]
    fn tick_evicts_expired() {
        setup();
        let mut kv = KvBuilder::new().default_ttl_secs(5).build().expect("build");

        kv.put(b"short", b"val").expect("put");
        kv.put_with_ttl(b"long", b"val", 100).expect("put");

        host::mock_advance_time_secs(6);

        let evicted = kv.tick().expect("tick");
        assert_eq!(evicted, 1);

        assert_eq!(kv.get(b"short").expect("get"), None);
        assert_eq!(kv.get(b"long").expect("get"), Some(b"val".to_vec()));
    }

    #[test]
    fn keys_excludes_expired() {
        setup();
        let mut kv = KvBuilder::new().default_ttl_secs(5).build().expect("build");

        kv.put(b"expire-soon", b"val").expect("put");
        kv.put_with_ttl(b"lives-long", b"val", 100).expect("put");

        host::mock_advance_time_secs(6);

        let keys = kv.keys().expect("keys");
        assert_eq!(keys, vec![b"lives-long".to_vec()]);
    }

    #[test]
    fn put_with_explicit_ttl() {
        setup();
        let mut kv = KvBuilder::new()
            .default_ttl_secs(300)
            .build()
            .expect("build");

        kv.put_with_ttl(b"short", b"val", 5).expect("put");
        kv.put_with_ttl(b"long", b"val", 600).expect("put");

        // After 6 seconds, short key should be expired, long should not
        host::mock_advance_time_secs(6);

        assert_eq!(kv.get(b"short").expect("get"), None);
        assert_eq!(kv.get(b"long").expect("get"), Some(b"val".to_vec()));
    }

    #[test]
    fn ttl_remaining_decreases() {
        setup();
        let mut kv = KvBuilder::new()
            .default_ttl_secs(100)
            .build()
            .expect("build");

        kv.put(b"k", b"v").expect("put");

        let ttl1 = kv.ttl(b"k").expect("ttl").expect("some");
        assert_eq!(ttl1, 100);

        host::mock_advance_time_secs(30);
        let ttl2 = kv.ttl(b"k").expect("ttl").expect("some");
        assert_eq!(ttl2, 70);

        host::mock_advance_time_secs(71);
        // Past expiry
        assert_eq!(kv.ttl(b"k").expect("ttl"), None);
    }

    #[test]
    fn tick_mixed_ttls_evicts_selectively() {
        setup();
        let mut kv = KvBuilder::new()
            .default_ttl_secs(100)
            .build()
            .expect("build");

        kv.put_with_ttl(b"a", b"1", 5).expect("put");
        kv.put_with_ttl(b"b", b"2", 20).expect("put");
        kv.put_with_ttl(b"c", b"3", 100).expect("put");

        // After 6s: "a" expired
        host::mock_advance_time_secs(6);
        let evicted = kv.tick().expect("tick");
        assert_eq!(evicted, 1);
        // "a" should be gone, "b" and "c" should remain
        assert_eq!(kv.keys().expect("keys").len(), 2);

        // After 21s total: "b" expired too
        // Note: don't use get() between ticks — it refreshes TTL on access.
        host::mock_advance_time_secs(15);
        let evicted = kv.tick().expect("tick");
        assert_eq!(evicted, 1);
        // Only "c" should remain
        let keys = kv.keys().expect("keys");
        assert_eq!(keys.len(), 1);
        assert_eq!(keys[0], b"c".to_vec());
    }

    #[test]
    fn tick_on_empty_store() {
        setup();
        let mut kv = KvBuilder::new().build().expect("build");
        let evicted = kv.tick().expect("tick");
        assert_eq!(evicted, 0);
    }

    #[test]
    fn overwrite_resets_ttl() {
        setup();
        let mut kv = KvBuilder::new()
            .default_ttl_secs(10)
            .build()
            .expect("build");

        kv.put(b"k", b"v1").expect("put");
        host::mock_advance_time_secs(8);

        // Overwrite with new value — should reset TTL
        kv.put(b"k", b"v2").expect("put");

        // After 5 more seconds (13 total), original would have expired
        // but overwrite reset TTL, so it should still be alive
        host::mock_advance_time_secs(5);
        assert_eq!(kv.get(b"k").expect("get"), Some(b"v2".to_vec()));
    }

    #[test]
    fn delete_nonexistent_returns_false() {
        setup();
        let mut kv = KvBuilder::new().build().expect("build");
        assert!(!kv.delete(b"nope").expect("delete"));
    }

    #[test]
    fn exists_returns_false_for_expired() {
        setup();
        let mut kv = KvBuilder::new().default_ttl_secs(5).build().expect("build");

        kv.put(b"k", b"v").expect("put");
        assert!(kv.exists(b"k").expect("exists"));

        host::mock_advance_time_secs(6);
        assert!(!kv.exists(b"k").expect("exists after expiry"));
    }

    #[test]
    fn builder_custom_strides() {
        setup();
        // Use custom strides
        let mut kv = KvBuilder::new()
            .hot_buckets(16)
            .key_stride(64)
            .value_stride(256)
            .default_ttl_secs(60)
            .build()
            .expect("build");

        kv.put(b"key", b"val").expect("put");
        assert_eq!(kv.get(b"key").expect("get"), Some(b"val".to_vec()));
    }

    #[test]
    fn get_on_expired_entry_removes_it_from_hot() {
        setup();
        let mut kv = KvBuilder::new().default_ttl_secs(5).build().expect("build");

        kv.put(b"k", b"v").expect("put");
        assert_eq!(kv.keys().expect("keys").len(), 1);

        host::mock_advance_time_secs(6);

        // get() on an expired entry should return None AND remove it
        assert_eq!(kv.get(b"k").expect("get"), None);

        // tick() should find nothing to evict since get() already cleaned up
        let evicted = kv.tick().expect("tick");
        assert_eq!(evicted, 0);
    }

    #[cfg(feature = "persistence")]
    mod cold_tier_tests {
        use super::*;

        fn setup_with_cold() {
            host::reset_mock();
            host::mock_set_fbmu_arena_size(65536);
            host::mock_set_fbbu_num_blocks(256);
        }

        /// Set up mock with a small hot tier that actually triggers cold spill.
        ///
        /// bucket_size = 1 + 8 + 4 + 128 + 4 + 512 = 657
        /// 4 buckets: HEADER(32) + 4 * 657 = 2660 bytes
        /// 75% load factor → max 3 entries, 4th spills to cold.
        fn setup_small_hot() {
            host::reset_mock();
            host::mock_set_fbmu_arena_size(2660);
            host::mock_set_fbbu_num_blocks(256);
        }

        #[test]
        fn cold_tier_spill_and_promote() {
            setup_small_hot();

            let mut kv = KvBuilder::new()
                .hot_buckets(4)
                .default_ttl_secs(300)
                .cold_num_blocks(64)
                .build()
                .expect("build");

            // Fill hot tier to 75% capacity (3 of 4 buckets)
            kv.put(b"a", b"1").expect("put a");
            kv.put(b"b", b"2").expect("put b");
            kv.put(b"c", b"3").expect("put c");

            // 4th insert exceeds load factor → spills to cold tier
            kv.put(b"d", b"4").expect("put d (cold spill)");
            assert!(
                kv.cold.as_ref().unwrap().write_cursor > 0,
                "d should have spilled to cold tier"
            );

            // Key "d" should be findable via cold tier scan
            let val = kv.get(b"d").expect("get d");
            assert_eq!(val, Some(b"4".to_vec()));
        }

        #[test]
        fn cold_tier_compaction() {
            setup_small_hot();
            let mut kv = KvBuilder::new()
                .hot_buckets(4)
                .default_ttl_secs(300)
                .cold_num_blocks(64)
                .build()
                .expect("build");

            // Fill hot tier (3 entries at 75% of 4 buckets)
            kv.put(b"a", b"1").expect("put");
            kv.put(b"b", b"2").expect("put");
            kv.put(b"c", b"3").expect("put");

            // Spill to cold
            kv.put(b"d", b"4").expect("put d");
            kv.put(b"e", b"5").expect("put e");

            // Delete a cold-tier key (writes tombstone)
            kv.delete(b"d").expect("delete d");

            // Compact should remove the tombstoned entry
            kv.compact_cold().expect("compact");

            // "e" should still be findable
            let val = kv.get(b"e").expect("get e");
            assert_eq!(val, Some(b"5".to_vec()));

            // "d" should be gone
            let val = kv.get(b"d").expect("get d");
            assert_eq!(val, None);
        }

        #[test]
        fn cold_tier_multiple_writes_last_writer_wins() {
            setup_small_hot();
            let mut kv = KvBuilder::new()
                .hot_buckets(4)
                .default_ttl_secs(300)
                .cold_num_blocks(64)
                .build()
                .expect("build");

            // Fill hot tier (3 entries at 75% of 4 buckets)
            kv.put(b"a", b"1").expect("put a");
            kv.put(b"b", b"2").expect("put b");
            kv.put(b"c", b"3").expect("put c");

            // Spill to cold with same key multiple times
            kv.put(b"d", b"first").expect("put d first");
            kv.put(b"d", b"second").expect("put d second");

            // Last value should win (cold tier scan returns last live entry)
            let val = kv.get(b"d").expect("get d");
            assert_eq!(val, Some(b"second".to_vec()));
        }

        #[test]
        fn cold_tier_promote_to_hot_on_read() {
            setup_small_hot();
            let mut kv = KvBuilder::new()
                .hot_buckets(4)
                .default_ttl_secs(300)
                .cold_num_blocks(64)
                .build()
                .expect("build");

            // Fill hot tier (3 entries at 75% of 4 buckets)
            kv.put(b"a", b"1").expect("put a");
            kv.put(b"b", b"2").expect("put b");
            kv.put(b"c", b"3").expect("put c");

            // Spill "d" to cold
            kv.put(b"d", b"4").expect("put d");

            // Delete a hot-tier key to free space
            kv.delete(b"a").expect("delete a");

            // Now read "d" — should promote from cold to hot
            let val = kv.get(b"d").expect("get d");
            assert_eq!(val, Some(b"4".to_vec()));

            // Reading "d" again should now hit the hot tier (promoted)
            // We can verify by checking it still works
            let val2 = kv.get(b"d").expect("get d again");
            assert_eq!(val2, Some(b"4".to_vec()));
        }

        #[test]
        fn cold_tier_tombstone_prevents_read() {
            setup_small_hot();
            let mut kv = KvBuilder::new()
                .hot_buckets(4)
                .default_ttl_secs(300)
                .cold_num_blocks(64)
                .build()
                .expect("build");

            // Fill hot tier (3 entries at 75% of 4 buckets)
            kv.put(b"a", b"1").expect("put a");
            kv.put(b"b", b"2").expect("put b");
            kv.put(b"c", b"3").expect("put c");

            // Spill to cold
            kv.put(b"d", b"4").expect("put d");

            // Delete "d" — writes tombstone to cold
            kv.delete(b"d").expect("delete d");

            // "d" should not be found (tombstone blocks it)
            let val = kv.get(b"d").expect("get d");
            assert_eq!(val, None);
        }

        #[test]
        fn compact_cold_removes_duplicates() {
            setup_small_hot();
            let mut kv = KvBuilder::new()
                .hot_buckets(4)
                .default_ttl_secs(300)
                .cold_num_blocks(64)
                .build()
                .expect("build");

            // Fill hot tier (3 entries at 75% of 4 buckets)
            kv.put(b"a", b"1").expect("put a");
            kv.put(b"b", b"2").expect("put b");
            kv.put(b"c", b"3").expect("put c");

            // Write several versions to cold
            kv.put(b"d", b"v1").expect("put d v1");
            kv.put(b"d", b"v2").expect("put d v2");
            kv.put(b"d", b"v3").expect("put d v3");
            kv.put(b"e", b"x1").expect("put e x1");
            kv.put(b"e", b"x2").expect("put e x2");

            // Compact — should deduplicate, keeping latest
            kv.compact_cold().expect("compact");

            // Values should still be correct
            let val_d = kv.get(b"d").expect("get d");
            assert_eq!(val_d, Some(b"v3".to_vec()));
            let val_e = kv.get(b"e").expect("get e");
            assert_eq!(val_e, Some(b"x2".to_vec()));
        }

        #[test]
        fn hot_full_without_cold_returns_capacity_exceeded() {
            host::reset_mock();
            // Small arena: bucket_size = 1+8+4+128+4+512 = 657.
            // For exactly 4 buckets: 32 + 4*657 = 2660 bytes.
            // 75% load factor = 3 entries max.
            host::mock_set_fbmu_arena_size(2660);
            host::mock_set_fbbu_num_blocks(256);

            // Build without cold tier
            let mut kv = KvBuilder::new()
                .hot_buckets(4)
                .default_ttl_secs(300)
                .build()
                .expect("build");

            // Fill hot tier to 75% load (3 of 4 buckets)
            kv.put(b"a", b"1").expect("put a");
            kv.put(b"b", b"2").expect("put b");
            kv.put(b"c", b"3").expect("put c");

            // 4th insert should fail — no cold tier to spill to
            let err = kv.put(b"d", b"4").unwrap_err();
            assert_eq!(err, FabricError::CapacityExceeded);
        }

        #[test]
        fn compact_cold_on_store_without_cold_is_noop() {
            setup_with_cold();
            let mut kv = KvBuilder::new()
                .hot_buckets(64)
                .default_ttl_secs(300)
                .build()
                .expect("build");

            // Should not error
            kv.compact_cold().expect("compact on no cold");
        }

        #[test]
        fn from_cold_recovers_spilled_keys() {
            host::reset_mock();
            host::mock_set_fbbu_num_blocks(256);

            // bucket_size = 1 + 8 + 4 + 128 + 4 + 512 = 657
            // For exactly 4 buckets: HEADER(32) + 4 * 657 = 2660 bytes
            // 75% load factor: max 3 entries before CapacityExceeded
            host::mock_set_fbmu_arena_size(2660);

            // Build a store, spill data to cold tier, then simulate crash.
            let mut kv = KvBuilder::new()
                .hot_buckets(4)
                .default_ttl_secs(300)
                .cold_num_blocks(64)
                .build()
                .expect("build");

            // Fill hot tier to 75% capacity (3 of 4 buckets)
            kv.put(b"a", b"1").expect("put a");
            kv.put(b"b", b"2").expect("put b");
            kv.put(b"c", b"3").expect("put c");

            // 4th insert exceeds 75% load factor → spills to cold tier
            kv.put(b"cold-key", b"cold-value").expect("put cold");
            assert!(
                kv.cold.as_ref().unwrap().write_cursor > 0,
                "cold-key should have spilled to cold tier"
            );

            // Grab lease ID, then simulate a crash (forget = no Drop = no free).
            let cold_lease_id = kv.cold.as_ref().unwrap().lease.lease_id();
            core::mem::forget(kv);

            // Bump arena size for recovery hot tier allocation (64 buckets).
            host::mock_set_fbmu_arena_size(65536);

            // Recover from cold tier only (hot tier lease "expired").
            let cold_lease =
                grafos_std::block::BlockBuilder::attach(cold_lease_id).expect("attach cold");
            let mut recovered = KvBuilder::from_cold(cold_lease, 64).expect("from_cold");

            // Cold-tier key should be accessible via promote-on-read.
            let val = recovered.get(b"cold-key").expect("get cold-key");
            assert_eq!(val, Some(b"cold-value".to_vec()));
        }

        #[test]
        fn from_leases_recovers_hot_and_cold() {
            setup_with_cold();

            let mut kv = KvBuilder::new()
                .hot_buckets(64)
                .default_ttl_secs(300)
                .cold_num_blocks(64)
                .build()
                .expect("build");

            kv.put(b"hot-key", b"hot-value").expect("put hot");

            let hot_lease_id = kv.hot.lease_id();
            let cold_lease_id = kv.cold.as_ref().unwrap().lease.lease_id();
            core::mem::forget(kv);

            // Recover from both leases.
            let hot_lease = grafos_std::mem::MemBuilder::attach(hot_lease_id).expect("attach hot");
            let cold_lease =
                grafos_std::block::BlockBuilder::attach(cold_lease_id).expect("attach cold");
            let mut recovered = KvBuilder::from_leases(hot_lease, cold_lease).expect("from_leases");

            // Hot-tier key should be directly accessible.
            let val = recovered.get(b"hot-key").expect("get hot-key");
            assert_eq!(val, Some(b"hot-value".to_vec()));
        }
    }
}