siftdb_core/

locking.rs

use anyhow::{Context, Result};
use std::fs::{File, OpenOptions};
use std::io::{Read, Write};
use std::path::{Path, PathBuf};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use serde::{Deserialize, Serialize};

/// Lock types for SiftDB operations
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum LockType {
    Read,
    Write,
}

/// Lock information stored in lock files
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LockInfo {
    pub lock_type: String, // "read" or "write"
    pub process_id: u32,
    pub acquired_at: u64,
    pub expires_at: u64,
    pub holder_info: String, // Additional info about lock holder
}

impl LockInfo {
    pub fn new(lock_type: LockType, duration_secs: u64, holder_info: String) -> Self {
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_secs();
        
        Self {
            lock_type: match lock_type {
                LockType::Read => "read".to_string(),
                LockType::Write => "write".to_string(),
            },
            process_id: std::process::id(),
            acquired_at: now,
            expires_at: now + duration_secs,
            holder_info,
        }
    }

    pub fn is_expired(&self) -> bool {
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_secs();
        now >= self.expires_at
    }

    pub fn is_write_lock(&self) -> bool {
        self.lock_type == "write"
    }

    pub fn is_read_lock(&self) -> bool {
        self.lock_type == "read"
    }
}

/// SWMR lock manager for SiftDB collections
pub struct SWMRLockManager {
    collection_path: PathBuf,
    locks_dir: PathBuf,
    write_lock_path: PathBuf,
}

impl SWMRLockManager {
    pub fn new(collection_path: &Path) -> Self {
        let locks_dir = collection_path.join("locks");
        let write_lock_path = locks_dir.join("write.lock");
        
        Self {
            collection_path: collection_path.to_path_buf(),
            locks_dir,
            write_lock_path,
        }
    }

    /// Initialize the locks directory
    pub fn init(&self) -> Result<()> {
        std::fs::create_dir_all(&self.locks_dir)
            .context("Failed to create locks directory")?;
        Ok(())
    }

    /// Acquire a read lock
    pub fn acquire_read_lock(&self, timeout_secs: u64, holder_info: String) -> Result<ReadLock> {
        self.init()?;
        
        let start_time = SystemTime::now();
        let timeout = Duration::from_secs(timeout_secs);

        loop {
            // Check if there's an active write lock
            if let Ok(write_lock_info) = self.read_write_lock() {
                if !write_lock_info.is_expired() {
                    if start_time.elapsed().unwrap() >= timeout {
                        anyhow::bail!("Timeout waiting for read lock - write lock held by process {}", 
                                    write_lock_info.process_id);
                    }
                    std::thread::sleep(Duration::from_millis(100));
                    continue;
                }
            }

            // No active write lock, acquire read lock
            let lock_info = LockInfo::new(LockType::Read, 3600, holder_info); // 1 hour default
            let read_lock_path = self.locks_dir.join(format!("read_{}.lock", std::process::id()));
            
            self.write_lock_file(&read_lock_path, &lock_info)?;
            
            return Ok(ReadLock {
                manager: self,
                lock_path: read_lock_path,
                lock_info,
            });
        }
    }

    /// Acquire a write lock
    pub fn acquire_write_lock(&self, timeout_secs: u64, holder_info: String) -> Result<WriteLock> {
        self.init()?;
        
        let start_time = SystemTime::now();
        let timeout = Duration::from_secs(timeout_secs);

        loop {
            // Check for existing write lock
            if let Ok(existing_write) = self.read_write_lock() {
                if !existing_write.is_expired() {
                    if start_time.elapsed().unwrap() >= timeout {
                        anyhow::bail!("Timeout waiting for write lock - held by process {}", 
                                    existing_write.process_id);
                    }
                    std::thread::sleep(Duration::from_millis(100));
                    continue;
                }
            }

            // Check for active read locks
            let active_read_locks = self.get_active_read_locks()?;
            if !active_read_locks.is_empty() {
                if start_time.elapsed().unwrap() >= timeout {
                    anyhow::bail!("Timeout waiting for write lock - {} read locks active", 
                                active_read_locks.len());
                }
                std::thread::sleep(Duration::from_millis(100));
                continue;
            }

            // No conflicts, acquire write lock
            let lock_info = LockInfo::new(LockType::Write, 1800, holder_info); // 30 minutes default
            
            self.write_lock_file(&self.write_lock_path, &lock_info)?;
            
            return Ok(WriteLock {
                manager: self,
                lock_info,
            });
        }
    }

    /// Check if collection is currently locked for writing
    pub fn is_write_locked(&self) -> Result<bool> {
        match self.read_write_lock() {
            Ok(lock_info) => Ok(!lock_info.is_expired()),
            Err(_) => Ok(false),
        }
    }

    /// Get count of active read locks
    pub fn active_read_lock_count(&self) -> Result<usize> {
        Ok(self.get_active_read_locks()?.len())
    }

    fn read_write_lock(&self) -> Result<LockInfo> {
        let content = std::fs::read_to_string(&self.write_lock_path)
            .context("Failed to read write lock file")?;
        let lock_info: LockInfo = serde_json::from_str(&content)
            .context("Failed to parse write lock info")?;
        Ok(lock_info)
    }

    fn get_active_read_locks(&self) -> Result<Vec<LockInfo>> {
        let mut active_locks = Vec::new();
        
        if !self.locks_dir.exists() {
            return Ok(active_locks);
        }

        for entry in std::fs::read_dir(&self.locks_dir)? {
            let entry = entry?;
            let path = entry.path();
            
            if let Some(filename) = path.file_name() {
                if let Some(filename_str) = filename.to_str() {
                    if filename_str.starts_with("read_") && filename_str.ends_with(".lock") {
                        if let Ok(content) = std::fs::read_to_string(&path) {
                            if let Ok(lock_info) = serde_json::from_str::<LockInfo>(&content) {
                                if !lock_info.is_expired() {
                                    active_locks.push(lock_info);
                                } else {
                                    // Clean up expired lock
                                    let _ = std::fs::remove_file(&path);
                                }
                            }
                        }
                    }
                }
            }
        }
        
        Ok(active_locks)
    }

    fn write_lock_file(&self, path: &Path, lock_info: &LockInfo) -> Result<()> {
        let json = serde_json::to_string_pretty(lock_info)
            .context("Failed to serialize lock info")?;
        std::fs::write(path, json)
            .context("Failed to write lock file")?;
        Ok(())
    }

    fn release_read_lock(&self, lock_path: &Path) -> Result<()> {
        if lock_path.exists() {
            std::fs::remove_file(lock_path)
                .context("Failed to remove read lock file")?;
        }
        Ok(())
    }

    fn release_write_lock(&self) -> Result<()> {
        if self.write_lock_path.exists() {
            std::fs::remove_file(&self.write_lock_path)
                .context("Failed to remove write lock file")?;
        }
        Ok(())
    }
}

/// Read lock guard - automatically released when dropped
pub struct ReadLock<'a> {
    manager: &'a SWMRLockManager,
    lock_path: PathBuf,
    lock_info: LockInfo,
}

impl<'a> ReadLock<'a> {
    pub fn lock_info(&self) -> &LockInfo {
        &self.lock_info
    }

    /// Extend the lock duration
    pub fn extend(&mut self, additional_secs: u64) -> Result<()> {
        self.lock_info.expires_at += additional_secs;
        self.manager.write_lock_file(&self.lock_path, &self.lock_info)?;
        Ok(())
    }
}

impl<'a> Drop for ReadLock<'a> {
    fn drop(&mut self) {
        let _ = self.manager.release_read_lock(&self.lock_path);
    }
}

/// Write lock guard - automatically released when dropped  
pub struct WriteLock<'a> {
    manager: &'a SWMRLockManager,
    lock_info: LockInfo,
}

impl<'a> WriteLock<'a> {
    pub fn lock_info(&self) -> &LockInfo {
        &self.lock_info
    }

    /// Extend the lock duration
    pub fn extend(&mut self, additional_secs: u64) -> Result<()> {
        self.lock_info.expires_at += additional_secs;
        self.manager.write_lock_file(&self.manager.write_lock_path, &self.lock_info)?;
        Ok(())
    }
}

impl<'a> Drop for WriteLock<'a> {
    fn drop(&mut self) {
        let _ = self.manager.release_write_lock();
    }
}

/// Helper trait for lock-aware operations
pub trait LockAware {
    fn with_read_lock<F, R>(&self, timeout_secs: u64, operation: F) -> Result<R>
    where
        F: FnOnce() -> Result<R>;

    fn with_write_lock<F, R>(&self, timeout_secs: u64, operation: F) -> Result<R>
    where
        F: FnOnce() -> Result<R>;
}

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

    #[test]
    fn test_read_lock_acquisition() {
        let temp_dir = TempDir::new().unwrap();
        let manager = SWMRLockManager::new(temp_dir.path());
        
        let _lock = manager.acquire_read_lock(5, "test".to_string()).unwrap();
        assert_eq!(manager.active_read_lock_count().unwrap(), 1);
    }

    #[test]
    fn test_write_lock_exclusivity() {
        let temp_dir = TempDir::new().unwrap();
        let manager = SWMRLockManager::new(temp_dir.path());
        
        let _write_lock = manager.acquire_write_lock(5, "test".to_string()).unwrap();
        
        // Should not be able to acquire read lock while write lock is held
        let result = manager.acquire_read_lock(1, "test2".to_string());
        assert!(result.is_err());
    }
}