storage/

segmented.rs

//! Segmented (LSM-style) storage backend built on the compaction engine.
//!
//! This adapts [`CompactionManager`] — a complete in-memory log-structured store
//! (append to an active segment, seal, tombstone-on-delete, merge garbage-heavy
//! segments) — to the [`VectorStorage`] trait so it can be selected as a real
//! backend. Deletes are tombstones reclaimed by background compaction rather than
//! in-place removals, which keeps writes append-only and bounds read amplification
//! by periodically merging segments.
//!
//! Before DAK-7428 the compaction engine existed but was never wired to a storage
//! interface — it was reachable only from unit tests. This module makes it a
//! functional backend (`DAKERA_STORAGE=segmented`) with self-contained background
//! auto-compaction, not a simulated one.

use std::sync::Arc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};

use async_trait::async_trait;
use common::{NamespaceId, Result, Vector, VectorId};

use crate::compaction::{CompactionConfig, CompactionManager};
use crate::traits::VectorStorage;

/// Current Unix time in seconds — captured once per call to avoid a syscall per
/// vector when filtering expired entries.
fn now_secs() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|d| d.as_secs())
        .unwrap_or(0)
}

/// Log-structured, compaction-backed vector storage.
pub struct SegmentedStorage {
    manager: Arc<CompactionManager>,
}

impl SegmentedStorage {
    /// Build a segmented store with the given compaction policy. When
    /// `config.auto_compact` is set and a Tokio runtime is active, a background
    /// task is spawned that periodically merges garbage-heavy/small segments so
    /// tombstones are reclaimed and read amplification stays bounded.
    pub fn new(config: CompactionConfig) -> Self {
        let manager = Arc::new(CompactionManager::new(config.clone()));
        if config.auto_compact {
            if let Ok(handle) = tokio::runtime::Handle::try_current() {
                let mgr = manager.clone();
                let interval = config.compaction_interval_secs.max(1);
                handle.spawn(async move {
                    let mut ticker = tokio::time::interval(Duration::from_secs(interval));
                    ticker.tick().await; // consume the immediate first tick
                    loop {
                        ticker.tick().await;
                        let results = mgr.compact_all();
                        if !results.is_empty() {
                            let merged: usize = results.values().map(|v| v.len()).sum();
                            tracing::debug!(
                                namespaces = results.len(),
                                compactions = merged,
                                "segmented storage auto-compaction merged segments"
                            );
                        }
                    }
                });
                tracing::info!(
                    interval_secs = interval,
                    "Segmented storage background compaction started"
                );
            }
        }
        Self { manager }
    }

    /// Build with the default compaction policy.
    pub fn with_defaults() -> Self {
        Self::new(CompactionConfig::default())
    }

    /// Access the underlying compaction manager (stats, manual compaction).
    pub fn manager(&self) -> &Arc<CompactionManager> {
        &self.manager
    }

    /// Manually run compaction across every namespace; returns the number of
    /// segment-merge operations performed.
    pub fn compact_all(&self) -> usize {
        self.manager.compact_all().values().map(|v| v.len()).sum()
    }
}

impl Default for SegmentedStorage {
    fn default() -> Self {
        Self::with_defaults()
    }
}

#[async_trait]
impl VectorStorage for SegmentedStorage {
    async fn upsert(&self, namespace: &NamespaceId, vectors: Vec<Vector>) -> Result<usize> {
        let count = vectors.len();
        for vector in vectors {
            self.manager.add(namespace, vector);
        }
        Ok(count)
    }

    async fn get(&self, namespace: &NamespaceId, ids: &[VectorId]) -> Result<Vec<Vector>> {
        let now = now_secs();
        Ok(ids
            .iter()
            .filter_map(|id| self.manager.get(namespace, id))
            .filter(|v| !v.is_expired_at(now))
            .collect())
    }

    async fn get_all(&self, namespace: &NamespaceId) -> Result<Vec<Vector>> {
        let now = now_secs();
        Ok(self
            .manager
            .get_all(namespace)
            .into_iter()
            .filter(|v| !v.is_expired_at(now))
            .collect())
    }

    async fn delete(&self, namespace: &NamespaceId, ids: &[VectorId]) -> Result<usize> {
        let mut removed = 0;
        for id in ids {
            if self.manager.delete(namespace, id) {
                removed += 1;
            }
        }
        Ok(removed)
    }

    async fn namespace_exists(&self, namespace: &NamespaceId) -> Result<bool> {
        Ok(self
            .manager
            .list_namespaces()
            .iter()
            .any(|n| n == namespace))
    }

    async fn ensure_namespace(&self, namespace: &NamespaceId) -> Result<()> {
        // `namespace()` is get-or-create.
        let _ = self.manager.namespace(namespace);
        Ok(())
    }

    async fn count(&self, namespace: &NamespaceId) -> Result<usize> {
        let now = now_secs();
        Ok(self
            .manager
            .get_all(namespace)
            .into_iter()
            .filter(|v| !v.is_expired_at(now))
            .count())
    }

    async fn dimension(&self, namespace: &NamespaceId) -> Result<Option<usize>> {
        let now = now_secs();
        Ok(self
            .manager
            .get_all(namespace)
            .into_iter()
            .find(|v| !v.is_expired_at(now))
            .map(|v| v.values.len()))
    }

    async fn list_namespaces(&self) -> Result<Vec<NamespaceId>> {
        Ok(self.manager.list_namespaces())
    }

    async fn delete_namespace(&self, namespace: &NamespaceId) -> Result<bool> {
        Ok(self.manager.delete_namespace(namespace))
    }

    async fn cleanup_expired(&self, namespace: &NamespaceId) -> Result<usize> {
        let now = now_secs();
        let expired: Vec<VectorId> = self
            .manager
            .get_all(namespace)
            .into_iter()
            .filter(|v| v.is_expired_at(now))
            .map(|v| v.id)
            .collect();
        let mut removed = 0;
        for id in &expired {
            if self.manager.delete(namespace, id) {
                removed += 1;
            }
        }
        Ok(removed)
    }

    async fn cleanup_all_expired(&self) -> Result<usize> {
        let mut total = 0;
        for namespace in self.manager.list_namespaces() {
            total += self.cleanup_expired(&namespace).await?;
        }
        Ok(total)
    }
}

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

    fn tv(id: &str) -> Vector {
        Vector {
            id: id.to_string(),
            values: vec![0.1, 0.2, 0.3],
            metadata: None,
            ttl_seconds: None,
            expires_at: None,
        }
    }

    // Build without background compaction so the test is deterministic and needs
    // no live runtime timer.
    fn store() -> SegmentedStorage {
        SegmentedStorage::new(CompactionConfig {
            auto_compact: false,
            ..Default::default()
        })
    }

    #[tokio::test]
    async fn upsert_get_delete_roundtrip() {
        let s = store();
        let ns = "n".to_string();
        s.ensure_namespace(&ns).await.unwrap();
        assert_eq!(s.upsert(&ns, vec![tv("a"), tv("b")]).await.unwrap(), 2);
        assert_eq!(s.count(&ns).await.unwrap(), 2);
        assert_eq!(s.dimension(&ns).await.unwrap(), Some(3));

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

        assert_eq!(s.delete(&ns, &["a".to_string()]).await.unwrap(), 1);
        assert!(s.get(&ns, &["a".to_string()]).await.unwrap().is_empty());
        assert_eq!(s.count(&ns).await.unwrap(), 1);

        let all: Vec<String> = s
            .get_all(&ns)
            .await
            .unwrap()
            .into_iter()
            .map(|v| v.id)
            .collect();
        assert_eq!(all, vec!["b".to_string()]);
    }

    #[tokio::test]
    async fn expired_vectors_are_filtered_and_reclaimed() {
        let s = store();
        let ns = "n".to_string();
        let mut expired = tv("old");
        expired.expires_at = Some(1); // far in the past
        s.upsert(&ns, vec![expired, tv("fresh")]).await.unwrap();
        // Read path hides the expired vector.
        assert_eq!(s.count(&ns).await.unwrap(), 1);
        let ids: Vec<String> = s
            .get_all(&ns)
            .await
            .unwrap()
            .into_iter()
            .map(|v| v.id)
            .collect();
        assert_eq!(ids, vec!["fresh".to_string()]);
        // Sweep physically tombstones it.
        assert_eq!(s.cleanup_expired(&ns).await.unwrap(), 1);
    }

    #[tokio::test]
    async fn namespace_lifecycle() {
        let s = store();
        let ns = "n".to_string();
        s.upsert(&ns, vec![tv("a")]).await.unwrap();
        assert!(s.namespace_exists(&ns).await.unwrap());
        assert_eq!(s.list_namespaces().await.unwrap(), vec![ns.clone()]);
        assert!(s.delete_namespace(&ns).await.unwrap());
        assert!(!s.namespace_exists(&ns).await.unwrap());
    }
}