storage/

rocksdb.rs

//! RocksDB-backed vector storage for the persistent hot tier.
//!
//! Implements `VectorStorage` on top of RocksDB, providing the same semantics
//! as `InMemoryStorage` but with durability across restarts. Vectors survive
//! server crashes without needing a warm/cold tier promotion round-trip.

use async_trait::async_trait;
use common::{DakeraError, NamespaceId, Result, Vector, VectorId};
use rocksdb::{BlockBasedOptions, Cache, Options, DB};
use std::sync::Arc;

use crate::traits::VectorStorage;

/// Configuration for the RocksDB hot tier.
#[derive(Debug, Clone)]
pub struct RocksDbStorageConfig {
    /// Path to the RocksDB database directory.
    pub path: String,
    /// RocksDB block cache size in MB (default: 64).
    pub block_cache_mb: usize,
    /// Enable LZ4 compression (default: true).
    pub compression: bool,
    /// Write buffer size in bytes (default: 64 MB).
    pub write_buffer_size: usize,
    /// Max number of concurrent write buffers (default: 3).
    pub max_write_buffer_number: i32,
}

impl Default for RocksDbStorageConfig {
    fn default() -> Self {
        Self {
            path: "/data/hot".to_string(),
            block_cache_mb: 64,
            compression: true,
            write_buffer_size: 64 * 1024 * 1024,
            max_write_buffer_number: 3,
        }
    }
}

/// Persistent hot-tier storage backed by RocksDB.
///
/// Key layout: `ns:{namespace}:v:{vector_id}` → bincode-serialized `Vector`
/// Namespace existence tracked via sentinel keys: `ns:{namespace}:__meta__` → `"1"`
pub struct RocksDbStorage {
    db: Arc<DB>,
}

impl RocksDbStorage {
    /// Open (or create) a RocksDB-backed storage at the configured path.
    pub fn new(config: RocksDbStorageConfig) -> Result<Self> {
        Self::open_at(&config.path, &config)
    }

    /// Open a RocksDB instance at the given path with the supplied config.
    pub fn open_at(path: &str, config: &RocksDbStorageConfig) -> 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);
        }

        let cache = Cache::new_lru_cache(config.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);

        opts.set_level_compaction_dynamic_level_bytes(true);
        opts.set_max_background_jobs(4);

        // Use prefix extractor for efficient namespace scans
        opts.set_prefix_extractor(rocksdb::SliceTransform::create_fixed_prefix(
            // "ns:" prefix minimum — we scan by "ns:{namespace}:" which varies in
            // length, but setting a fixed prefix of 3 enables the bloom filter on
            // the constant "ns:" tag, which is good enough for our workload.
            3,
        ));

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

        tracing::info!(path, "RocksDB hot tier initialized");

        Ok(Self { db: Arc::new(db) })
    }

    // ── key helpers ────────────────────────────────────────────────────

    fn vector_key(ns: &NamespaceId, id: &VectorId) -> Vec<u8> {
        format!("ns:{ns}:v:{id}").into_bytes()
    }

    fn vector_prefix(ns: &NamespaceId) -> Vec<u8> {
        format!("ns:{ns}:v:").into_bytes()
    }

    fn meta_key(ns: &NamespaceId) -> Vec<u8> {
        format!("ns:{ns}:__meta__").into_bytes()
    }

    fn ns_prefix() -> Vec<u8> {
        b"ns:".to_vec()
    }

    // ── serde helpers ──────────────────────────────────────────────────

    fn serialize_vector(v: &Vector) -> Result<Vec<u8>> {
        serde_json::to_vec(v).map_err(|e| DakeraError::Storage(format!("vector serialize: {e}")))
    }

    fn deserialize_vector(bytes: &[u8]) -> Result<Vector> {
        serde_json::from_slice(bytes)
            .map_err(|e| DakeraError::Storage(format!("vector deserialize: {e}")))
    }
}

fn now_secs() -> u64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs()
}

#[async_trait]
impl VectorStorage for RocksDbStorage {
    async fn upsert(&self, namespace: &NamespaceId, vectors: Vec<Vector>) -> Result<usize> {
        if vectors.is_empty() {
            return Ok(0);
        }

        // Dimension validation: match existing vectors' dimension or the
        // first vector in this batch if the namespace is fresh.
        let now = now_secs();
        let prefix = Self::vector_prefix(namespace);
        let expected_dim = {
            let iter = self.db.prefix_iterator(&prefix);
            let mut dim = None;
            for item in iter {
                match item {
                    Ok((k, v)) if k.starts_with(&prefix) => {
                        if let Ok(vec) = Self::deserialize_vector(&v) {
                            if !vec.is_expired_at(now) {
                                dim = Some(vec.values.len());
                                break;
                            }
                        }
                    }
                    _ => break,
                }
            }
            dim.unwrap_or_else(|| vectors[0].values.len())
        };

        for v in &vectors {
            if v.values.len() != expected_dim {
                return Err(DakeraError::DimensionMismatch {
                    expected: expected_dim,
                    actual: v.values.len(),
                });
            }
        }

        let mut batch = rocksdb::WriteBatch::default();

        // Ensure namespace sentinel exists
        batch.put(Self::meta_key(namespace), b"1");

        let count = vectors.len();
        for mut vector in vectors {
            vector.apply_ttl();
            let key = Self::vector_key(namespace, &vector.id);
            let val = Self::serialize_vector(&vector)?;
            batch.put(key, val);
        }

        self.db
            .write(batch)
            .map_err(|e| DakeraError::Storage(format!("RocksDB batch write: {e}")))?;

        tracing::debug!(namespace = %namespace, count, "RocksDB upserted vectors");
        Ok(count)
    }

    async fn get(&self, namespace: &NamespaceId, ids: &[VectorId]) -> Result<Vec<Vector>> {
        if !self.namespace_exists(namespace).await? {
            return Err(DakeraError::NamespaceNotFound(namespace.clone()));
        }

        let now = now_secs();
        let keys: Vec<Vec<u8>> = ids
            .iter()
            .map(|id| Self::vector_key(namespace, id))
            .collect();

        let results = self.db.multi_get(&keys);
        let mut out = Vec::with_capacity(ids.len());
        for result in results {
            if let Ok(Some(bytes)) = result {
                if let Ok(v) = Self::deserialize_vector(&bytes) {
                    if !v.is_expired_at(now) {
                        out.push(v);
                    }
                }
            }
        }
        Ok(out)
    }

    async fn get_all(&self, namespace: &NamespaceId) -> Result<Vec<Vector>> {
        if !self.namespace_exists(namespace).await? {
            return Err(DakeraError::NamespaceNotFound(namespace.clone()));
        }

        let prefix = Self::vector_prefix(namespace);
        let now = now_secs();
        let mut out = Vec::new();

        let iter = self.db.prefix_iterator(&prefix);
        for item in iter {
            match item {
                Ok((k, v)) if k.starts_with(&prefix) => {
                    if let Ok(vec) = Self::deserialize_vector(&v) {
                        if !vec.is_expired_at(now) {
                            out.push(vec);
                        }
                    }
                }
                _ => break,
            }
        }
        Ok(out)
    }

    async fn delete(&self, namespace: &NamespaceId, ids: &[VectorId]) -> Result<usize> {
        if !self.namespace_exists(namespace).await? {
            return Err(DakeraError::NamespaceNotFound(namespace.clone()));
        }

        let mut batch = rocksdb::WriteBatch::default();
        let mut deleted = 0usize;

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

        if deleted > 0 {
            self.db
                .write(batch)
                .map_err(|e| DakeraError::Storage(format!("RocksDB batch delete: {e}")))?;
        }

        tracing::debug!(namespace = %namespace, deleted, "RocksDB deleted vectors");
        Ok(deleted)
    }

    async fn namespace_exists(&self, namespace: &NamespaceId) -> Result<bool> {
        Ok(self
            .db
            .get(Self::meta_key(namespace))
            .ok()
            .flatten()
            .is_some())
    }

    async fn ensure_namespace(&self, namespace: &NamespaceId) -> Result<()> {
        self.db
            .put(Self::meta_key(namespace), b"1")
            .map_err(|e| DakeraError::Storage(format!("RocksDB put meta: {e}")))?;
        Ok(())
    }

    async fn count(&self, namespace: &NamespaceId) -> Result<usize> {
        let prefix = Self::vector_prefix(namespace);
        let now = now_secs();
        let mut n = 0usize;
        let iter = self.db.prefix_iterator(&prefix);
        for item in iter {
            match item {
                Ok((k, v)) if k.starts_with(&prefix) => {
                    if let Ok(vec) = Self::deserialize_vector(&v) {
                        if !vec.is_expired_at(now) {
                            n += 1;
                        }
                    }
                }
                _ => break,
            }
        }
        Ok(n)
    }

    async fn dimension(&self, namespace: &NamespaceId) -> Result<Option<usize>> {
        let prefix = Self::vector_prefix(namespace);
        let now = now_secs();
        let iter = self.db.prefix_iterator(&prefix);
        for item in iter {
            match item {
                Ok((k, v)) if k.starts_with(&prefix) => {
                    if let Ok(vec) = Self::deserialize_vector(&v) {
                        if !vec.is_expired_at(now) {
                            return Ok(Some(vec.values.len()));
                        }
                    }
                }
                _ => break,
            }
        }
        Ok(None)
    }

    async fn list_namespaces(&self) -> Result<Vec<NamespaceId>> {
        let ns_prefix = Self::ns_prefix();
        let mut namespaces = Vec::new();

        let iter = self.db.prefix_iterator(&ns_prefix);
        for item in iter {
            match item {
                Ok((k, _)) if k.starts_with(&ns_prefix) => {
                    if let Ok(key_str) = std::str::from_utf8(&k) {
                        if key_str.ends_with(":__meta__") {
                            // Extract namespace: "ns:{name}:__meta__"
                            if let Some(name) = key_str
                                .strip_prefix("ns:")
                                .and_then(|s| s.strip_suffix(":__meta__"))
                            {
                                namespaces.push(name.to_string());
                            }
                        }
                    }
                }
                Ok((k, _)) if !k.starts_with(&ns_prefix) => break,
                Err(_) => break,
                _ => {}
            }
        }
        // Deduplicate — the iterator may encounter meta keys out of order
        // depending on prefix extractor settings. In practice the "ns:" fixed
        // prefix keeps them ordered, but `sort + dedup` is cheap insurance.
        namespaces.sort_unstable();
        namespaces.dedup();
        Ok(namespaces)
    }

    async fn delete_namespace(&self, namespace: &NamespaceId) -> Result<bool> {
        let existed = self.namespace_exists(namespace).await?;
        if !existed {
            return Ok(false);
        }

        let mut batch = rocksdb::WriteBatch::default();

        // Delete all vectors
        let prefix = Self::vector_prefix(namespace);
        let iter = self.db.prefix_iterator(&prefix);
        for item in iter {
            match item {
                Ok((k, _)) if k.starts_with(&prefix) => batch.delete(&k),
                _ => break,
            }
        }

        // Delete the meta sentinel
        batch.delete(Self::meta_key(namespace));

        self.db
            .write(batch)
            .map_err(|e| DakeraError::Storage(format!("RocksDB delete namespace: {e}")))?;

        tracing::debug!(namespace = %namespace, "RocksDB deleted namespace");
        Ok(true)
    }

    async fn cleanup_expired(&self, namespace: &NamespaceId) -> Result<usize> {
        let prefix = Self::vector_prefix(namespace);
        let now = now_secs();
        let mut batch = rocksdb::WriteBatch::default();
        let mut removed = 0usize;

        let iter = self.db.prefix_iterator(&prefix);
        for item in iter {
            match item {
                Ok((k, v)) if k.starts_with(&prefix) => {
                    if let Ok(vec) = Self::deserialize_vector(&v) {
                        if vec.is_expired_at(now) {
                            batch.delete(&k);
                            removed += 1;
                        }
                    }
                }
                _ => break,
            }
        }

        if removed > 0 {
            self.db
                .write(batch)
                .map_err(|e| DakeraError::Storage(format!("RocksDB cleanup expired: {e}")))?;
            tracing::debug!(namespace = %namespace, removed, "RocksDB cleaned up expired vectors");
        }
        Ok(removed)
    }

    async fn cleanup_all_expired(&self) -> Result<usize> {
        let namespaces = self.list_namespaces().await?;
        let mut total = 0usize;
        for ns in &namespaces {
            total += self.cleanup_expired(ns).await?;
        }
        if total > 0 {
            tracing::info!(
                total,
                "RocksDB cleaned up expired vectors across all namespaces"
            );
        }
        Ok(total)
    }

    async fn get_page(
        &self,
        namespace: &NamespaceId,
        offset: usize,
        limit: usize,
    ) -> Result<Vec<Vector>> {
        let prefix = Self::vector_prefix(namespace);
        let now = now_secs();
        let mut out = Vec::with_capacity(limit);
        let mut skipped = 0usize;

        let iter = self.db.prefix_iterator(&prefix);
        for item in iter {
            match item {
                Ok((k, v)) if k.starts_with(&prefix) => {
                    if let Ok(vec) = Self::deserialize_vector(&v) {
                        if !vec.is_expired_at(now) {
                            if skipped < offset {
                                skipped += 1;
                            } else {
                                out.push(vec);
                                if out.len() >= limit {
                                    break;
                                }
                            }
                        }
                    }
                }
                _ => break,
            }
        }
        Ok(out)
    }
}

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

    fn test_config(dir: &TempDir) -> RocksDbStorageConfig {
        RocksDbStorageConfig {
            path: dir.path().to_string_lossy().to_string(),
            block_cache_mb: 8,
            compression: false,
            ..Default::default()
        }
    }

    #[tokio::test]
    async fn test_upsert_and_get() {
        let dir = TempDir::new().unwrap();
        let storage = RocksDbStorage::new(test_config(&dir)).unwrap();
        let ns = "test".to_string();

        storage.ensure_namespace(&ns).await.unwrap();
        let count = storage
            .upsert(
                &ns,
                vec![Vector {
                    id: "v1".to_string(),
                    values: vec![1.0, 2.0, 3.0],
                    metadata: None,
                    ttl_seconds: None,
                    expires_at: None,
                }],
            )
            .await
            .unwrap();
        assert_eq!(count, 1);

        let got = storage.get(&ns, &["v1".to_string()]).await.unwrap();
        assert_eq!(got.len(), 1);
        assert_eq!(got[0].id, "v1");
        assert_eq!(got[0].values, vec![1.0, 2.0, 3.0]);
    }

    #[tokio::test]
    async fn test_dimension_mismatch() {
        let dir = TempDir::new().unwrap();
        let storage = RocksDbStorage::new(test_config(&dir)).unwrap();
        let ns = "test".to_string();

        storage.ensure_namespace(&ns).await.unwrap();
        storage
            .upsert(
                &ns,
                vec![Vector {
                    id: "v1".to_string(),
                    values: vec![1.0, 2.0, 3.0],
                    metadata: None,
                    ttl_seconds: None,
                    expires_at: None,
                }],
            )
            .await
            .unwrap();

        let result = storage
            .upsert(
                &ns,
                vec![Vector {
                    id: "v2".to_string(),
                    values: vec![1.0, 2.0],
                    metadata: None,
                    ttl_seconds: None,
                    expires_at: None,
                }],
            )
            .await;
        assert!(matches!(result, Err(DakeraError::DimensionMismatch { .. })));
    }

    #[tokio::test]
    async fn test_delete() {
        let dir = TempDir::new().unwrap();
        let storage = RocksDbStorage::new(test_config(&dir)).unwrap();
        let ns = "test".to_string();

        storage.ensure_namespace(&ns).await.unwrap();
        storage
            .upsert(
                &ns,
                vec![
                    Vector {
                        id: "v1".to_string(),
                        values: vec![1.0],
                        metadata: None,
                        ttl_seconds: None,
                        expires_at: None,
                    },
                    Vector {
                        id: "v2".to_string(),
                        values: vec![2.0],
                        metadata: None,
                        ttl_seconds: None,
                        expires_at: None,
                    },
                ],
            )
            .await
            .unwrap();

        let deleted = storage.delete(&ns, &["v1".to_string()]).await.unwrap();
        assert_eq!(deleted, 1);
        assert_eq!(storage.count(&ns).await.unwrap(), 1);
    }

    #[tokio::test]
    async fn test_get_all_and_count() {
        let dir = TempDir::new().unwrap();
        let storage = RocksDbStorage::new(test_config(&dir)).unwrap();
        let ns = "test".to_string();

        storage.ensure_namespace(&ns).await.unwrap();
        storage
            .upsert(
                &ns,
                vec![
                    Vector {
                        id: "v1".to_string(),
                        values: vec![1.0, 2.0],
                        metadata: None,
                        ttl_seconds: None,
                        expires_at: None,
                    },
                    Vector {
                        id: "v2".to_string(),
                        values: vec![3.0, 4.0],
                        metadata: None,
                        ttl_seconds: None,
                        expires_at: None,
                    },
                ],
            )
            .await
            .unwrap();

        let all = storage.get_all(&ns).await.unwrap();
        assert_eq!(all.len(), 2);
        assert_eq!(storage.count(&ns).await.unwrap(), 2);
    }

    #[tokio::test]
    async fn test_namespace_lifecycle() {
        let dir = TempDir::new().unwrap();
        let storage = RocksDbStorage::new(test_config(&dir)).unwrap();
        let ns = "lifecycle".to_string();

        assert!(!storage.namespace_exists(&ns).await.unwrap());
        storage.ensure_namespace(&ns).await.unwrap();
        assert!(storage.namespace_exists(&ns).await.unwrap());

        let nss = storage.list_namespaces().await.unwrap();
        assert!(nss.contains(&ns));

        storage.delete_namespace(&ns).await.unwrap();
        assert!(!storage.namespace_exists(&ns).await.unwrap());
    }

    #[tokio::test]
    async fn test_ttl_filtering() {
        let dir = TempDir::new().unwrap();
        let storage = RocksDbStorage::new(test_config(&dir)).unwrap();
        let ns = "ttl".to_string();

        storage.ensure_namespace(&ns).await.unwrap();
        let past = now_secs().saturating_sub(100);

        storage
            .upsert(
                &ns,
                vec![
                    Vector {
                        id: "expired".to_string(),
                        values: vec![1.0, 2.0],
                        metadata: None,
                        ttl_seconds: None,
                        expires_at: Some(past),
                    },
                    Vector {
                        id: "valid".to_string(),
                        values: vec![3.0, 4.0],
                        metadata: None,
                        ttl_seconds: None,
                        expires_at: None,
                    },
                ],
            )
            .await
            .unwrap();

        let got = storage
            .get(&ns, &["expired".to_string(), "valid".to_string()])
            .await
            .unwrap();
        assert_eq!(got.len(), 1);
        assert_eq!(got[0].id, "valid");

        let all = storage.get_all(&ns).await.unwrap();
        assert_eq!(all.len(), 1);
        assert_eq!(storage.count(&ns).await.unwrap(), 1);
    }

    #[tokio::test]
    async fn test_cleanup_expired() {
        let dir = TempDir::new().unwrap();
        let storage = RocksDbStorage::new(test_config(&dir)).unwrap();
        let ns = "cleanup".to_string();

        storage.ensure_namespace(&ns).await.unwrap();
        let past = now_secs().saturating_sub(100);

        storage
            .upsert(
                &ns,
                vec![
                    Vector {
                        id: "expired1".to_string(),
                        values: vec![1.0],
                        metadata: None,
                        ttl_seconds: None,
                        expires_at: Some(past),
                    },
                    Vector {
                        id: "valid".to_string(),
                        values: vec![2.0],
                        metadata: None,
                        ttl_seconds: None,
                        expires_at: None,
                    },
                ],
            )
            .await
            .unwrap();

        let removed = storage.cleanup_expired(&ns).await.unwrap();
        assert_eq!(removed, 1);

        // Only valid vector should remain (even at raw level)
        let all = storage.get_all(&ns).await.unwrap();
        assert_eq!(all.len(), 1);
        assert_eq!(all[0].id, "valid");
    }

    #[tokio::test]
    async fn test_get_page() {
        let dir = TempDir::new().unwrap();
        let storage = RocksDbStorage::new(test_config(&dir)).unwrap();
        let ns = "paged".to_string();

        storage.ensure_namespace(&ns).await.unwrap();
        let mut vecs = Vec::new();
        for i in 0..10 {
            vecs.push(Vector {
                id: format!("v{i:02}"),
                values: vec![i as f32],
                metadata: None,
                ttl_seconds: None,
                expires_at: None,
            });
        }
        storage.upsert(&ns, vecs).await.unwrap();

        let page = storage.get_page(&ns, 3, 4).await.unwrap();
        assert_eq!(page.len(), 4);
    }

    #[tokio::test]
    async fn test_persistence_across_reopen() {
        let dir = TempDir::new().unwrap();

        // Write data with one instance
        {
            let storage = RocksDbStorage::new(test_config(&dir)).unwrap();
            let ns = "persist".to_string();
            storage.ensure_namespace(&ns).await.unwrap();
            storage
                .upsert(
                    &ns,
                    vec![Vector {
                        id: "durable".to_string(),
                        values: vec![42.0],
                        metadata: None,
                        ttl_seconds: None,
                        expires_at: None,
                    }],
                )
                .await
                .unwrap();
        }

        // Reopen and verify data survived
        {
            let storage = RocksDbStorage::new(test_config(&dir)).unwrap();
            let ns = "persist".to_string();
            assert!(storage.namespace_exists(&ns).await.unwrap());
            let got = storage.get(&ns, &["durable".to_string()]).await.unwrap();
            assert_eq!(got.len(), 1);
            assert_eq!(got[0].values, vec![42.0]);
        }
    }

    #[tokio::test]
    async fn test_multi_namespace_isolation() {
        let dir = TempDir::new().unwrap();
        let storage = RocksDbStorage::new(test_config(&dir)).unwrap();

        storage.ensure_namespace(&"ns1".to_string()).await.unwrap();
        storage.ensure_namespace(&"ns2".to_string()).await.unwrap();

        storage
            .upsert(
                &"ns1".to_string(),
                vec![Vector {
                    id: "v1".to_string(),
                    values: vec![1.0],
                    metadata: None,
                    ttl_seconds: None,
                    expires_at: None,
                }],
            )
            .await
            .unwrap();

        storage
            .upsert(
                &"ns2".to_string(),
                vec![Vector {
                    id: "v1".to_string(),
                    values: vec![2.0],
                    metadata: None,
                    ttl_seconds: None,
                    expires_at: None,
                }],
            )
            .await
            .unwrap();

        let ns1_vecs = storage.get_all(&"ns1".to_string()).await.unwrap();
        let ns2_vecs = storage.get_all(&"ns2".to_string()).await.unwrap();
        assert_eq!(ns1_vecs.len(), 1);
        assert_eq!(ns2_vecs.len(), 1);
        assert_eq!(ns1_vecs[0].values, vec![1.0]);
        assert_eq!(ns2_vecs[0].values, vec![2.0]);

        // Deleting ns1 should not affect ns2
        storage.delete_namespace(&"ns1".to_string()).await.unwrap();
        assert!(!storage.namespace_exists(&"ns1".to_string()).await.unwrap());
        assert!(storage.namespace_exists(&"ns2".to_string()).await.unwrap());
        assert_eq!(storage.get_all(&"ns2".to_string()).await.unwrap().len(), 1);
    }
}