engine/

online_update.rs

//! Online Index Updates
//!
//! Provides a unified interface for incrementally updating indexes without
//! requiring a full rebuild. Supports:
//!
//! - Incremental vector insertion
//! - Vector deletion with lazy cleanup
//! - Vector updates (delete + insert)
//! - Background maintenance operations
//! - Concurrent update handling
//!
//! # Design Principles
//!
//! 1. **Non-blocking reads**: Updates don't block ongoing searches
//! 2. **Eventual consistency**: Updates visible after brief delay
//! 3. **Lazy cleanup**: Deleted vectors marked, cleaned up in background
//! 4. **Batch efficiency**: Batching updates for better performance

use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::{Arc, RwLock};

use common::VectorId;

/// Statistics for online update operations
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct OnlineUpdateStats {
    /// Total vectors inserted since index creation
    pub total_inserts: u64,
    /// Total vectors deleted since index creation
    pub total_deletes: u64,
    /// Total vectors updated since index creation
    pub total_updates: u64,
    /// Vectors pending cleanup (deleted but not yet removed)
    pub pending_cleanup: usize,
    /// Last maintenance operation timestamp (unix millis)
    pub last_maintenance_ms: u64,
    /// Number of maintenance operations performed
    pub maintenance_count: u64,
}

/// Configuration for online updates
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OnlineUpdateConfig {
    /// Maximum pending deletes before triggering cleanup
    pub max_pending_deletes: usize,
    /// Maximum pending inserts before triggering batch commit
    pub max_pending_inserts: usize,
    /// Whether to enable background maintenance
    pub enable_background_maintenance: bool,
    /// Interval for background maintenance in milliseconds
    pub maintenance_interval_ms: u64,
}

impl Default for OnlineUpdateConfig {
    fn default() -> Self {
        Self {
            max_pending_deletes: 1000,
            max_pending_inserts: 100,
            enable_background_maintenance: true,
            maintenance_interval_ms: 60_000, // 1 minute
        }
    }
}

/// Trait for indexes that support online updates
pub trait OnlineUpdatable: Send + Sync {
    /// Insert a single vector
    fn online_insert(&self, id: VectorId, vector: Vec<f32>) -> Result<(), String>;

    /// Insert multiple vectors in a batch
    fn online_insert_batch(&self, vectors: Vec<(VectorId, Vec<f32>)>) -> Result<usize, String>;

    /// Delete a vector by ID (may be lazy)
    fn online_delete(&self, id: &VectorId) -> Result<bool, String>;

    /// Delete multiple vectors by ID
    fn online_delete_batch(&self, ids: &[VectorId]) -> Result<usize, String>;

    /// Update a vector (atomic delete + insert)
    fn online_update(&self, id: VectorId, vector: Vec<f32>) -> Result<(), String> {
        self.online_delete(&id)?;
        self.online_insert(id, vector)
    }

    /// Update multiple vectors
    fn online_update_batch(&self, vectors: Vec<(VectorId, Vec<f32>)>) -> Result<usize, String> {
        let ids: Vec<_> = vectors.iter().map(|(id, _)| id.clone()).collect();
        self.online_delete_batch(&ids)?;
        self.online_insert_batch(vectors)
    }

    /// Check if a vector exists
    fn contains(&self, id: &VectorId) -> bool;

    /// Get the current vector count
    fn len(&self) -> usize;

    /// Check if the index is empty
    fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Perform maintenance operations (cleanup deleted vectors, optimize structure)
    fn maintenance(&self) -> Result<MaintenanceResult, String>;

    /// Get online update statistics
    fn online_stats(&self) -> OnlineUpdateStats;
}

/// Result of a maintenance operation
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct MaintenanceResult {
    /// Number of vectors cleaned up
    pub vectors_cleaned: usize,
    /// Memory reclaimed in bytes
    pub memory_reclaimed_bytes: usize,
    /// Duration of maintenance in milliseconds
    pub duration_ms: u64,
    /// Whether the index was restructured
    pub restructured: bool,
}

/// Buffer for batching online updates
#[derive(Debug)]
pub struct UpdateBuffer {
    config: OnlineUpdateConfig,
    /// Pending inserts
    pending_inserts: RwLock<Vec<(VectorId, Vec<f32>)>>,
    /// Pending deletes
    pending_deletes: RwLock<HashSet<VectorId>>,
    /// Statistics
    stats: UpdateBufferStats,
}

#[derive(Debug, Default)]
struct UpdateBufferStats {
    total_inserts: AtomicU64,
    total_deletes: AtomicU64,
    total_updates: AtomicU64,
    pending_count: AtomicUsize,
}

impl UpdateBuffer {
    /// Create a new update buffer
    pub fn new(config: OnlineUpdateConfig) -> Self {
        Self {
            config,
            pending_inserts: RwLock::new(Vec::new()),
            pending_deletes: RwLock::new(HashSet::new()),
            stats: UpdateBufferStats::default(),
        }
    }

    /// Add an insert to the buffer
    pub fn buffer_insert(&self, id: VectorId, vector: Vec<f32>) {
        let mut inserts = self
            .pending_inserts
            .write()
            .expect("pending_inserts lock poisoned in buffer_insert");
        inserts.push((id, vector));
        self.stats.pending_count.fetch_add(1, Ordering::Relaxed);
    }

    /// Add a delete to the buffer
    pub fn buffer_delete(&self, id: VectorId) {
        let mut deletes = self
            .pending_deletes
            .write()
            .expect("pending_deletes lock poisoned in buffer_delete");
        deletes.insert(id);
    }

    /// Check if buffer should be flushed
    pub fn should_flush(&self) -> bool {
        let inserts = self
            .pending_inserts
            .read()
            .expect("pending_inserts lock poisoned in should_flush");
        let deletes = self
            .pending_deletes
            .read()
            .expect("pending_deletes lock poisoned in should_flush");

        inserts.len() >= self.config.max_pending_inserts
            || deletes.len() >= self.config.max_pending_deletes
    }

    /// Flush pending inserts and return them
    pub fn flush_inserts(&self) -> Vec<(VectorId, Vec<f32>)> {
        let mut inserts = self
            .pending_inserts
            .write()
            .expect("pending_inserts lock poisoned in flush_inserts");
        let flushed: Vec<_> = inserts.drain(..).collect();
        self.stats
            .total_inserts
            .fetch_add(flushed.len() as u64, Ordering::Relaxed);
        self.stats
            .pending_count
            .fetch_sub(flushed.len(), Ordering::Relaxed);
        flushed
    }

    /// Flush pending deletes and return them
    pub fn flush_deletes(&self) -> HashSet<VectorId> {
        let mut deletes = self
            .pending_deletes
            .write()
            .expect("pending_deletes lock poisoned in flush_deletes");
        let flushed: HashSet<_> = deletes.drain().collect();
        self.stats
            .total_deletes
            .fetch_add(flushed.len() as u64, Ordering::Relaxed);
        flushed
    }

    /// Get buffer statistics
    pub fn stats(&self) -> OnlineUpdateStats {
        OnlineUpdateStats {
            total_inserts: self.stats.total_inserts.load(Ordering::Relaxed),
            total_deletes: self.stats.total_deletes.load(Ordering::Relaxed),
            total_updates: self.stats.total_updates.load(Ordering::Relaxed),
            pending_cleanup: self
                .pending_deletes
                .read()
                .expect("pending_deletes lock poisoned in stats")
                .len(),
            last_maintenance_ms: 0,
            maintenance_count: 0,
        }
    }
}

/// Wrapper that adds online update capabilities to any index
pub struct OnlineIndex<T> {
    /// The underlying index
    inner: Arc<RwLock<T>>,
    /// Update buffer (reserved for batch operations)
    _buffer: UpdateBuffer,
    /// Deleted vector IDs (for lazy cleanup)
    deleted_ids: RwLock<HashSet<VectorId>>,
    /// ID to vector mapping for updates
    id_to_vector: RwLock<HashMap<VectorId, Vec<f32>>>,
    /// Statistics
    stats: Arc<RwLock<OnlineUpdateStats>>,
    /// Configuration
    config: OnlineUpdateConfig,
}

impl<T> OnlineIndex<T> {
    /// Create a new online-updatable index wrapper
    pub fn new(inner: T, config: OnlineUpdateConfig) -> Self {
        Self {
            inner: Arc::new(RwLock::new(inner)),
            _buffer: UpdateBuffer::new(config.clone()),
            deleted_ids: RwLock::new(HashSet::new()),
            id_to_vector: RwLock::new(HashMap::new()),
            stats: Arc::new(RwLock::new(OnlineUpdateStats::default())),
            config,
        }
    }

    /// Get read access to the inner index
    pub fn read(&self) -> std::sync::RwLockReadGuard<'_, T> {
        self.inner
            .read()
            .expect("inner index lock poisoned in read")
    }

    /// Get write access to the inner index
    pub fn write(&self) -> std::sync::RwLockWriteGuard<'_, T> {
        self.inner
            .write()
            .expect("inner index lock poisoned in write")
    }

    /// Check if a vector ID is deleted
    pub fn is_deleted(&self, id: &VectorId) -> bool {
        self.deleted_ids
            .read()
            .expect("deleted_ids lock poisoned in is_deleted")
            .contains(id)
    }

    /// Mark a vector as deleted
    pub fn mark_deleted(&self, id: VectorId) -> bool {
        let mut deleted = self
            .deleted_ids
            .write()
            .expect("deleted_ids lock poisoned in mark_deleted");
        let was_new = deleted.insert(id.clone());

        if was_new {
            let mut stats = self
                .stats
                .write()
                .expect("stats lock poisoned in mark_deleted");
            stats.total_deletes += 1;
            stats.pending_cleanup += 1;
        }

        was_new
    }

    /// Clear deleted marks after cleanup
    pub fn clear_deleted(&self, ids: &[VectorId]) {
        let mut deleted = self
            .deleted_ids
            .write()
            .expect("deleted_ids lock poisoned in clear_deleted");
        let mut stats = self
            .stats
            .write()
            .expect("stats lock poisoned in clear_deleted");

        for id in ids {
            if deleted.remove(id) {
                stats.pending_cleanup = stats.pending_cleanup.saturating_sub(1);
            }
        }
    }

    /// Get number of pending deletes
    pub fn pending_deletes(&self) -> usize {
        self.deleted_ids
            .read()
            .expect("deleted_ids lock poisoned in pending_deletes")
            .len()
    }

    /// Check if maintenance is needed
    pub fn needs_maintenance(&self) -> bool {
        self.pending_deletes() >= self.config.max_pending_deletes
    }

    /// Record an insert
    pub fn record_insert(&self) {
        let mut stats = self
            .stats
            .write()
            .expect("stats lock poisoned in record_insert");
        stats.total_inserts += 1;
    }

    /// Record an update
    pub fn record_update(&self) {
        let mut stats = self
            .stats
            .write()
            .expect("stats lock poisoned in record_update");
        stats.total_updates += 1;
    }

    /// Get statistics
    pub fn stats(&self) -> OnlineUpdateStats {
        self.stats
            .read()
            .expect("stats lock poisoned in stats")
            .clone()
    }

    /// Store vector for potential retrieval
    pub fn store_vector(&self, id: VectorId, vector: Vec<f32>) {
        let mut map = self
            .id_to_vector
            .write()
            .expect("id_to_vector lock poisoned in store_vector");
        map.insert(id, vector);
    }

    /// Remove stored vector
    pub fn remove_vector(&self, id: &VectorId) -> Option<Vec<f32>> {
        let mut map = self
            .id_to_vector
            .write()
            .expect("id_to_vector lock poisoned in remove_vector");
        map.remove(id)
    }

    /// Check if vector exists in storage
    pub fn has_vector(&self, id: &VectorId) -> bool {
        let map = self
            .id_to_vector
            .read()
            .expect("id_to_vector lock poisoned in has_vector");
        map.contains_key(id) && !self.is_deleted(id)
    }

    /// Get vector count (excluding deleted)
    pub fn vector_count(&self) -> usize {
        let map = self
            .id_to_vector
            .read()
            .expect("id_to_vector lock poisoned in vector_count");
        let deleted = self
            .deleted_ids
            .read()
            .expect("deleted_ids lock poisoned in vector_count");
        map.len().saturating_sub(deleted.len())
    }
}

/// Delta log for tracking changes between index snapshots
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct DeltaLog {
    /// Inserted vectors since last snapshot
    pub inserts: Vec<(VectorId, Vec<f32>)>,
    /// Deleted vector IDs since last snapshot
    pub deletes: Vec<VectorId>,
    /// Timestamp of the base snapshot
    pub base_snapshot_time: u64,
    /// Current log sequence number
    pub sequence_number: u64,
}

impl DeltaLog {
    /// Create a new delta log
    pub fn new(base_snapshot_time: u64) -> Self {
        Self {
            inserts: Vec::new(),
            deletes: Vec::new(),
            base_snapshot_time,
            sequence_number: 0,
        }
    }

    /// Record an insert
    pub fn record_insert(&mut self, id: VectorId, vector: Vec<f32>) {
        self.inserts.push((id, vector));
        self.sequence_number += 1;
    }

    /// Record a delete
    pub fn record_delete(&mut self, id: VectorId) {
        self.deletes.push(id);
        self.sequence_number += 1;
    }

    /// Check if log should be compacted
    pub fn should_compact(&self, threshold: usize) -> bool {
        self.inserts.len() + self.deletes.len() > threshold
    }

    /// Clear the log after snapshot
    pub fn clear(&mut self, new_base_time: u64) {
        self.inserts.clear();
        self.deletes.clear();
        self.base_snapshot_time = new_base_time;
        // Keep sequence number for ordering
    }

    /// Get total operations in log
    pub fn len(&self) -> usize {
        self.inserts.len() + self.deletes.len()
    }

    /// Check if log is empty
    pub fn is_empty(&self) -> bool {
        self.inserts.is_empty() && self.deletes.is_empty()
    }
}

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

    #[test]
    fn test_update_buffer() {
        let config = OnlineUpdateConfig::default();
        let buffer = UpdateBuffer::new(config);

        // Buffer inserts
        buffer.buffer_insert("v1".to_string(), vec![1.0, 2.0, 3.0]);
        buffer.buffer_insert("v2".to_string(), vec![4.0, 5.0, 6.0]);

        assert!(!buffer.should_flush()); // Not at threshold yet

        // Buffer deletes
        buffer.buffer_delete("v1".to_string());

        // Flush and check
        let inserts = buffer.flush_inserts();
        assert_eq!(inserts.len(), 2);

        let deletes = buffer.flush_deletes();
        assert_eq!(deletes.len(), 1);
        assert!(deletes.contains("v1"));
    }

    #[test]
    fn test_delta_log() {
        let mut log = DeltaLog::new(1000);

        log.record_insert("v1".to_string(), vec![1.0, 2.0]);
        log.record_insert("v2".to_string(), vec![3.0, 4.0]);
        log.record_delete("v0".to_string());

        assert_eq!(log.len(), 3);
        assert_eq!(log.inserts.len(), 2);
        assert_eq!(log.deletes.len(), 1);
        assert_eq!(log.sequence_number, 3);

        // Clear after snapshot
        log.clear(2000);
        assert!(log.is_empty());
        assert_eq!(log.base_snapshot_time, 2000);
        assert_eq!(log.sequence_number, 3); // Preserved
    }

    #[test]
    fn test_online_index_wrapper() {
        // Test with a simple Vec as inner "index"
        let inner: Vec<(String, Vec<f32>)> = Vec::new();
        let config = OnlineUpdateConfig::default();
        let online = OnlineIndex::new(inner, config);

        // Store some vectors
        online.store_vector("v1".to_string(), vec![1.0, 2.0]);
        online.store_vector("v2".to_string(), vec![3.0, 4.0]);

        assert!(online.has_vector(&"v1".to_string()));
        assert!(online.has_vector(&"v2".to_string()));
        assert_eq!(online.vector_count(), 2);

        // Mark one as deleted
        online.mark_deleted("v1".to_string());
        assert!(!online.has_vector(&"v1".to_string()));
        assert_eq!(online.vector_count(), 1);
        assert_eq!(online.pending_deletes(), 1);

        // Check stats
        let stats = online.stats();
        assert_eq!(stats.total_deletes, 1);
    }

    #[test]
    fn test_flush_threshold() {
        let config = OnlineUpdateConfig {
            max_pending_inserts: 5,
            max_pending_deletes: 3,
            ..Default::default()
        };
        let buffer = UpdateBuffer::new(config);

        // Add inserts below threshold
        for i in 0..4 {
            buffer.buffer_insert(format!("v{}", i), vec![i as f32]);
        }
        assert!(!buffer.should_flush());

        // Add one more to hit threshold
        buffer.buffer_insert("v4".to_string(), vec![4.0]);
        assert!(buffer.should_flush());
    }
}