storage/

disk_cache.rs

//! L2 Disk Cache using RocksDB
//!
//! Provides persistent local caching between L1 RAM cache (moka) and L3 object storage (S3).
//! Optimized for NVMe SSDs with configurable cache size limits.

use std::path::Path;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;

use rocksdb::{BlockBasedOptions, Cache, IteratorMode, Options, WriteBatch, DB};
use serde::{Deserialize, Serialize};
use tracing::{debug, warn};

use common::{DakeraError, Result, Vector, VectorId};

/// Configuration for the disk cache
#[derive(Debug, Clone)]
pub struct DiskCacheConfig {
    /// Path to the RocksDB database
    pub path: String,
    /// Maximum cache size in bytes (0 = unlimited)
    pub max_size_bytes: u64,
    /// Enable compression (LZ4)
    pub compression: bool,
    /// Write buffer size in bytes
    pub write_buffer_size: usize,
    /// Max number of write buffers
    pub max_write_buffer_number: i32,
}

impl Default for DiskCacheConfig {
    fn default() -> Self {
        Self {
            path: "./cache".to_string(),
            max_size_bytes: 10 * 1024 * 1024 * 1024, // 10GB default
            compression: true,
            write_buffer_size: 64 * 1024 * 1024, // 64MB
            max_write_buffer_number: 3,
        }
    }
}

/// Cache entry with metadata for eviction
#[derive(Debug, Serialize, Deserialize)]
struct CacheEntry {
    vector: Vector,
    access_count: u64,
    created_at: u64,
}

/// L2 Disk Cache backed by RocksDB
pub struct DiskCache {
    db: Arc<DB>,
    config: DiskCacheConfig,
    /// Writes accumulated since the last size-eviction sweep. The `max_size_bytes`
    /// cap is enforced periodically (not on every write) to amortise the O(N) scan
    /// (DAK-7428).
    writes_since_eviction: AtomicU64,
}

impl DiskCache {
    /// Create a new disk cache
    pub fn new(config: DiskCacheConfig) -> Result<Self> {
        let mut opts = Options::default();
        opts.create_if_missing(true);
        opts.set_write_buffer_size(config.write_buffer_size);
        opts.set_max_write_buffer_number(config.max_write_buffer_number);

        if config.compression {
            opts.set_compression_type(rocksdb::DBCompressionType::Lz4);
        }

        // Explicit block cache — prevents RocksDB from using unbounded default (DAK-5716).
        let block_cache_mb: usize = std::env::var("DAKERA_ROCKSDB_BLOCK_CACHE_MB")
            .ok()
            .and_then(|v| v.parse().ok())
            .unwrap_or(64);
        let cache = Cache::new_lru_cache(block_cache_mb * 1024 * 1024);
        let mut block_opts = BlockBasedOptions::default();
        block_opts.set_block_cache(&cache);
        block_opts.set_optimize_filters_for_memory(true);
        opts.set_block_based_table_factory(&block_opts);

        // Optimize for SSD
        opts.set_level_compaction_dynamic_level_bytes(true);
        opts.set_max_background_jobs(4);

        let db = DB::open(&opts, &config.path)
            .map_err(|e| DakeraError::Storage(format!("Failed to open RocksDB: {}", e)))?;

        debug!(path = %config.path, "Disk cache initialized");

        Ok(Self {
            db: Arc::new(db),
            config,
            writes_since_eviction: AtomicU64::new(0),
        })
    }

    /// Number of writes between size-eviction sweeps. Bounds the amortised cost
    /// of the O(N) cap-enforcement scan.
    const EVICTION_CHECK_INTERVAL: u64 = 1024;

    /// Enforce `max_size_bytes` by evicting the oldest entries (by `created_at`)
    /// until the cache is back under ~90% of the cap. No-op when the cap is 0
    /// (unlimited). Returns the number of entries evicted (DAK-7428).
    ///
    /// Approximate: entry size is the serialized key+value length, and the
    /// estimate is recomputed from a full scan, so a single sweep is O(N). It is
    /// therefore triggered only periodically via `note_writes`.
    pub fn enforce_size_limit(&self) -> Result<usize> {
        let max = self.config.max_size_bytes;
        if max == 0 {
            return Ok(0);
        }

        let mut entries: Vec<(Box<[u8]>, u64, u64)> = Vec::new();
        let mut total: u64 = 0;
        for item in self.db.iterator(IteratorMode::Start) {
            let (key, value) =
                item.map_err(|e| DakeraError::Storage(format!("disk cache scan failed: {}", e)))?;
            let bytes = (key.len() + value.len()) as u64;
            total += bytes;
            let created_at = serde_json::from_slice::<CacheEntry>(&value)
                .map(|e| e.created_at)
                .unwrap_or(0);
            entries.push((key, created_at, bytes));
        }

        if total <= max {
            return Ok(0);
        }

        // Evict down to 90% of the cap so we don't sweep again on the next write.
        let target = max - max / 10;
        entries.sort_by_key(|(_, created_at, _)| *created_at);

        let mut batch = WriteBatch::default();
        let mut evicted = 0usize;
        for (key, _, bytes) in entries {
            if total <= target {
                break;
            }
            batch.delete(key);
            total = total.saturating_sub(bytes);
            evicted += 1;
        }

        if evicted > 0 {
            self.db.write(batch).map_err(|e| {
                DakeraError::Storage(format!("disk cache eviction write failed: {}", e))
            })?;
            debug!(
                evicted,
                remaining_bytes = total,
                "disk cache evicted oldest entries to honor max_size_bytes"
            );
        }
        Ok(evicted)
    }

    /// Record `n` writes and run a size-eviction sweep once the accumulated write
    /// count crosses `EVICTION_CHECK_INTERVAL`. Best-effort: eviction errors are
    /// logged, never propagated to the caller's write path.
    fn note_writes(&self, n: u64) {
        if self.config.max_size_bytes == 0 {
            return;
        }
        let prev = self.writes_since_eviction.fetch_add(n, Ordering::Relaxed);
        if prev + n >= Self::EVICTION_CHECK_INTERVAL {
            self.writes_since_eviction.store(0, Ordering::Relaxed);
            if let Err(e) = self.enforce_size_limit() {
                warn!(error = %e, "disk cache size-limit enforcement failed");
            }
        }
    }

    /// Open existing cache or create new one
    pub fn open<P: AsRef<Path>>(path: P) -> Result<Self> {
        let config = DiskCacheConfig {
            path: path.as_ref().to_string_lossy().to_string(),
            ..Default::default()
        };
        Self::new(config)
    }

    /// Generate cache key from namespace and vector ID
    fn make_key(namespace: &str, id: &VectorId) -> Vec<u8> {
        format!("{}:{}", namespace, id).into_bytes()
    }

    /// Generate namespace prefix for scanning
    fn namespace_prefix(namespace: &str) -> Vec<u8> {
        format!("{}:", namespace).into_bytes()
    }

    /// Insert a vector into the cache
    pub fn put(&self, namespace: &str, vector: &Vector) -> Result<()> {
        let key = Self::make_key(namespace, &vector.id);
        let entry = CacheEntry {
            vector: vector.clone(),
            access_count: 1,
            created_at: std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap_or_default()
                .as_secs(),
        };

        let value = serde_json::to_vec(&entry)
            .map_err(|e| DakeraError::Storage(format!("Failed to serialize cache entry: {}", e)))?;

        self.db
            .put(&key, &value)
            .map_err(|e| DakeraError::Storage(format!("Failed to write to disk cache: {}", e)))?;

        self.note_writes(1);
        debug!(namespace = %namespace, id = %vector.id, "Cached vector to disk");
        Ok(())
    }

    /// Insert multiple vectors into the cache
    pub fn put_batch(&self, namespace: &str, vectors: &[Vector]) -> Result<usize> {
        let mut batch = rocksdb::WriteBatch::default();
        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs();

        for vector in vectors {
            let key = Self::make_key(namespace, &vector.id);
            let entry = CacheEntry {
                vector: vector.clone(),
                access_count: 1,
                created_at: now,
            };

            let value = serde_json::to_vec(&entry).map_err(|e| {
                DakeraError::Storage(format!("Failed to serialize cache entry: {}", e))
            })?;

            batch.put(&key, &value);
        }

        let count = vectors.len();
        self.db.write(batch).map_err(|e| {
            DakeraError::Storage(format!("Failed to write batch to disk cache: {}", e))
        })?;

        self.note_writes(count as u64);
        debug!(namespace = %namespace, count = count, "Batch cached vectors to disk");
        Ok(count)
    }

    /// Get a vector from the cache
    pub fn get(&self, namespace: &str, id: &VectorId) -> Result<Option<Vector>> {
        let key = Self::make_key(namespace, id);

        match self.db.get(&key) {
            Ok(Some(value)) => {
                let entry: CacheEntry = serde_json::from_slice(&value).map_err(|e| {
                    DakeraError::Storage(format!("Failed to deserialize cache entry: {}", e))
                })?;
                Ok(Some(entry.vector))
            }
            Ok(None) => Ok(None),
            Err(e) => {
                warn!(error = %e, "Failed to read from disk cache");
                Ok(None)
            }
        }
    }

    /// Get multiple vectors from the cache
    pub fn get_batch(&self, namespace: &str, ids: &[VectorId]) -> Result<Vec<Vector>> {
        let keys: Vec<Vec<u8>> = ids.iter().map(|id| Self::make_key(namespace, id)).collect();

        let results = self.db.multi_get(&keys);
        let mut vectors = Vec::with_capacity(ids.len());

        for result in results {
            if let Ok(Some(value)) = result {
                if let Ok(entry) = serde_json::from_slice::<CacheEntry>(&value) {
                    vectors.push(entry.vector);
                }
            }
        }

        Ok(vectors)
    }

    /// Get all vectors in a namespace
    pub fn get_all(&self, namespace: &str) -> Result<Vec<Vector>> {
        let prefix = Self::namespace_prefix(namespace);
        let mut vectors = Vec::new();

        let iter = self.db.prefix_iterator(&prefix);
        for item in iter {
            match item {
                Ok((key, value)) => {
                    // Check if key still starts with prefix (RocksDB iterators can overshoot)
                    if !key.starts_with(&prefix) {
                        break;
                    }

                    if let Ok(entry) = serde_json::from_slice::<CacheEntry>(&value) {
                        vectors.push(entry.vector);
                    }
                }
                Err(e) => {
                    warn!(error = %e, "Error iterating disk cache");
                    break;
                }
            }
        }

        Ok(vectors)
    }

    /// Delete a vector from the cache
    pub fn delete(&self, namespace: &str, id: &VectorId) -> Result<bool> {
        let key = Self::make_key(namespace, id);

        // Check if exists first
        let existed = self
            .db
            .get(&key)
            .map_err(|e| DakeraError::Storage(format!("Failed to check disk cache: {}", e)))?
            .is_some();

        if existed {
            self.db.delete(&key).map_err(|e| {
                DakeraError::Storage(format!("Failed to delete from disk cache: {}", e))
            })?;
        }

        Ok(existed)
    }

    /// Delete multiple vectors from the cache
    pub fn delete_batch(&self, namespace: &str, ids: &[VectorId]) -> Result<usize> {
        let mut batch = rocksdb::WriteBatch::default();
        let mut count = 0;

        for id in ids {
            let key = Self::make_key(namespace, id);
            if self.db.get(&key).ok().flatten().is_some() {
                batch.delete(&key);
                count += 1;
            }
        }

        self.db.write(batch).map_err(|e| {
            DakeraError::Storage(format!("Failed to delete batch from disk cache: {}", e))
        })?;

        Ok(count)
    }

    /// Clear all entries in a namespace
    pub fn clear_namespace(&self, namespace: &str) -> Result<usize> {
        let prefix = Self::namespace_prefix(namespace);
        let mut batch = rocksdb::WriteBatch::default();
        let mut count = 0;

        let iter = self.db.prefix_iterator(&prefix);
        for item in iter {
            match item {
                Ok((key, _)) => {
                    if !key.starts_with(&prefix) {
                        break;
                    }
                    batch.delete(&key);
                    count += 1;
                }
                Err(_) => break,
            }
        }

        if count > 0 {
            self.db.write(batch).map_err(|e| {
                DakeraError::Storage(format!("Failed to clear namespace from disk cache: {}", e))
            })?;
        }

        debug!(namespace = %namespace, count = count, "Cleared namespace from disk cache");
        Ok(count)
    }

    /// Get approximate size of the cache in bytes
    pub fn approximate_size(&self) -> u64 {
        self.db
            .property_int_value("rocksdb.estimate-live-data-size")
            .ok()
            .flatten()
            .unwrap_or(0)
    }

    /// Get cache statistics
    pub fn stats(&self) -> DiskCacheStats {
        DiskCacheStats {
            approximate_size_bytes: self.approximate_size(),
            approximate_num_keys: self
                .db
                .property_int_value("rocksdb.estimate-num-keys")
                .ok()
                .flatten()
                .unwrap_or(0),
        }
    }

    /// Flush all pending writes to disk
    pub fn flush(&self) -> Result<()> {
        self.db
            .flush()
            .map_err(|e| DakeraError::Storage(format!("Failed to flush disk cache: {}", e)))
    }
}

/// Statistics for the disk cache
#[derive(Debug, Clone)]
pub struct DiskCacheStats {
    pub approximate_size_bytes: u64,
    pub approximate_num_keys: u64,
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::TempDir;

    fn create_test_cache() -> (DiskCache, TempDir) {
        let temp_dir = TempDir::new().unwrap();
        let config = DiskCacheConfig {
            path: temp_dir.path().to_string_lossy().to_string(),
            ..Default::default()
        };
        let cache = DiskCache::new(config).unwrap();
        (cache, temp_dir)
    }

    fn test_vector(id: &str) -> Vector {
        Vector {
            id: id.to_string(),
            values: vec![1.0, 2.0, 3.0],
            metadata: None,
            ttl_seconds: None,
            expires_at: None,
        }
    }

    #[test]
    fn test_put_and_get() {
        let (cache, _dir) = create_test_cache();
        let namespace = "test";
        let vector = test_vector("v1");

        cache.put(namespace, &vector).unwrap();
        let result = cache.get(namespace, &"v1".to_string()).unwrap();

        assert!(result.is_some());
        let retrieved = result.unwrap();
        assert_eq!(retrieved.id, "v1");
        assert_eq!(retrieved.values, vec![1.0, 2.0, 3.0]);
    }

    #[test]
    fn test_get_nonexistent() {
        let (cache, _dir) = create_test_cache();
        let result = cache.get("test", &"nonexistent".to_string()).unwrap();
        assert!(result.is_none());
    }

    // DAK-7428: max_size_bytes must actually bound the cache — previously it was
    // stored but never enforced, so the L2 cache grew without limit.
    #[test]
    fn enforce_size_limit_evicts_when_over_cap() {
        let temp_dir = TempDir::new().unwrap();
        let config = DiskCacheConfig {
            path: temp_dir.path().to_string_lossy().to_string(),
            max_size_bytes: 2048, // tiny cap so a few hundred entries exceed it
            ..Default::default()
        };
        let cache = DiskCache::new(config).unwrap();
        let ns = "test";
        let vectors: Vec<Vector> = (0..200).map(|i| test_vector(&format!("v{i}"))).collect();
        cache.put_batch(ns, &vectors).unwrap();

        let before = cache.get_all(ns).unwrap().len();
        let evicted = cache.enforce_size_limit().unwrap();
        let after = cache.get_all(ns).unwrap().len();

        assert!(evicted > 0, "should evict when over the cap");
        assert!(after < before, "eviction should reduce entry count");
    }

    #[test]
    fn enforce_size_limit_noop_when_unlimited() {
        let temp_dir = TempDir::new().unwrap();
        let config = DiskCacheConfig {
            path: temp_dir.path().to_string_lossy().to_string(),
            max_size_bytes: 0, // unlimited
            ..Default::default()
        };
        let cache = DiskCache::new(config).unwrap();
        cache
            .put_batch("test", &[test_vector("a"), test_vector("b")])
            .unwrap();
        assert_eq!(cache.enforce_size_limit().unwrap(), 0);
        assert_eq!(cache.get_all("test").unwrap().len(), 2);
    }

    #[test]
    fn test_batch_operations() {
        let (cache, _dir) = create_test_cache();
        let namespace = "test";
        let vectors = vec![test_vector("v1"), test_vector("v2"), test_vector("v3")];

        let count = cache.put_batch(namespace, &vectors).unwrap();
        assert_eq!(count, 3);

        let ids: Vec<String> = vec!["v1".to_string(), "v2".to_string(), "v3".to_string()];
        let retrieved = cache.get_batch(namespace, &ids).unwrap();
        assert_eq!(retrieved.len(), 3);
    }

    #[test]
    fn test_get_all() {
        let (cache, _dir) = create_test_cache();
        let namespace = "test";
        let vectors = vec![test_vector("v1"), test_vector("v2")];

        cache.put_batch(namespace, &vectors).unwrap();
        let all = cache.get_all(namespace).unwrap();
        assert_eq!(all.len(), 2);
    }

    #[test]
    fn test_delete() {
        let (cache, _dir) = create_test_cache();
        let namespace = "test";
        let vector = test_vector("v1");

        cache.put(namespace, &vector).unwrap();
        assert!(cache.get(namespace, &"v1".to_string()).unwrap().is_some());

        let deleted = cache.delete(namespace, &"v1".to_string()).unwrap();
        assert!(deleted);

        assert!(cache.get(namespace, &"v1".to_string()).unwrap().is_none());
    }

    #[test]
    fn test_delete_batch() {
        let (cache, _dir) = create_test_cache();
        let namespace = "test";
        let vectors = vec![test_vector("v1"), test_vector("v2"), test_vector("v3")];

        cache.put_batch(namespace, &vectors).unwrap();

        let ids = vec!["v1".to_string(), "v2".to_string()];
        let deleted = cache.delete_batch(namespace, &ids).unwrap();
        assert_eq!(deleted, 2);

        assert!(cache.get(namespace, &"v1".to_string()).unwrap().is_none());
        assert!(cache.get(namespace, &"v2".to_string()).unwrap().is_none());
        assert!(cache.get(namespace, &"v3".to_string()).unwrap().is_some());
    }

    #[test]
    fn test_clear_namespace() {
        let (cache, _dir) = create_test_cache();
        let vectors = vec![test_vector("v1"), test_vector("v2")];

        cache.put_batch("ns1", &vectors).unwrap();
        cache.put_batch("ns2", &vectors).unwrap();

        let cleared = cache.clear_namespace("ns1").unwrap();
        assert_eq!(cleared, 2);

        assert!(cache.get_all("ns1").unwrap().is_empty());
        assert_eq!(cache.get_all("ns2").unwrap().len(), 2);
    }

    #[test]
    fn test_namespace_isolation() {
        let (cache, _dir) = create_test_cache();
        let vector = test_vector("v1");

        cache.put("ns1", &vector).unwrap();
        cache.put("ns2", &vector).unwrap();

        assert!(cache.get("ns1", &"v1".to_string()).unwrap().is_some());
        assert!(cache.get("ns2", &"v1".to_string()).unwrap().is_some());

        cache.delete("ns1", &"v1".to_string()).unwrap();

        assert!(cache.get("ns1", &"v1".to_string()).unwrap().is_none());
        assert!(cache.get("ns2", &"v1".to_string()).unwrap().is_some());
    }

    #[test]
    fn test_stats() {
        let (cache, _dir) = create_test_cache();
        let vectors = vec![test_vector("v1"), test_vector("v2")];

        cache.put_batch("test", &vectors).unwrap();
        cache.flush().unwrap();

        let stats = cache.stats();
        // Stats are approximate, just check they're available
        let _ = stats.approximate_num_keys;
    }
}