graph_d 1.3.2

//! Concurrency control for multi-threaded transactions.

use crate::error::{GraphError, Result};
use crate::graph::Id;
use crate::transaction::{Transaction, TransactionId};
use parking_lot::{Mutex, RwLock};
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use std::time::{Duration, Instant};

/// Lock types for different operations.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum LockType {
    /// Shared read lock
    Read,
    /// Exclusive write lock
    Write,
}

/// Resource that can be locked (nodes, relationships, etc.).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum LockableResource {
    /// Node with given ID
    Node(Id),
    /// Relationship with given ID
    Relationship(Id),
    /// Global graph structure
    Schema,
}

/// Lock information tracking.
#[derive(Debug, Clone)]
pub struct LockInfo {
    /// Unique identifier of the transaction holding this lock
    pub transaction_id: TransactionId,
    /// Type of lock (read or write)
    pub lock_type: LockType,
    /// Timestamp when the lock was acquired
    pub acquired_at: Instant,
    /// The resource that is locked
    pub resource: LockableResource,
}

/// Lock manager for coordinating concurrent access to graph resources.
pub struct LockManager {
    /// Currently held locks: resource -> set of lock info
    active_locks: Arc<RwLock<HashMap<LockableResource, Vec<LockInfo>>>>,
    /// Lock requests waiting for resources
    waiting_locks: Arc<Mutex<Vec<LockRequest>>>,
    /// Deadlock detection interval
    #[allow(dead_code)]
    deadlock_check_interval: Duration,
}

/// A lock request waiting to be granted.
#[derive(Debug, Clone)]
pub struct LockRequest {
    /// Unique identifier of the requesting transaction
    pub transaction_id: TransactionId,
    /// The resource being requested
    pub resource: LockableResource,
    /// Type of lock being requested (read or write)
    pub lock_type: LockType,
    /// Timestamp when the request was made
    pub requested_at: Instant,
}

impl LockManager {
    /// Create a new lock manager.
    pub fn new() -> Self {
        LockManager {
            active_locks: Arc::new(RwLock::new(HashMap::new())),
            waiting_locks: Arc::new(Mutex::new(Vec::new())),
            deadlock_check_interval: Duration::from_millis(100),
        }
    }

    /// Acquire a lock on a resource for a transaction.
    pub fn acquire_lock(
        &self,
        transaction_id: TransactionId,
        resource: LockableResource,
        lock_type: LockType,
        timeout: Option<Duration>,
    ) -> Result<bool> {
        let start_time = Instant::now();
        let timeout = timeout.unwrap_or(Duration::from_secs(30));

        loop {
            // Try to acquire the lock immediately
            if self.try_acquire_lock(transaction_id, resource, lock_type)? {
                return Ok(true);
            }

            // Check timeout
            if start_time.elapsed() > timeout {
                return Err(GraphError::Concurrency(format!(
                    "Lock acquisition timeout for transaction {transaction_id} on resource {resource:?}"
                )));
            }

            // Add to waiting queue
            self.add_to_waiting_queue(transaction_id, resource, lock_type);

            // Check for deadlocks
            if self.detect_deadlock(transaction_id)? {
                return Err(GraphError::Concurrency(format!(
                    "Deadlock detected involving transaction {transaction_id}"
                )));
            }

            // Wait a bit before retrying
            std::thread::sleep(Duration::from_millis(10));
        }
    }

    /// Try to acquire a lock immediately without waiting.
    fn try_acquire_lock(
        &self,
        transaction_id: TransactionId,
        resource: LockableResource,
        lock_type: LockType,
    ) -> Result<bool> {
        let mut locks = self.active_locks.write();
        let existing_locks = locks.get(&resource).cloned().unwrap_or_default();

        // Check if the lock can be acquired
        if self.can_acquire_lock(&existing_locks, lock_type) {
            let lock_info = LockInfo {
                transaction_id,
                lock_type,
                acquired_at: Instant::now(),
                resource,
            };

            locks.entry(resource).or_default().push(lock_info);
            Ok(true)
        } else {
            Ok(false)
        }
    }

    /// Check if a lock can be acquired given existing locks.
    fn can_acquire_lock(&self, existing_locks: &[LockInfo], requested_type: LockType) -> bool {
        if existing_locks.is_empty() {
            return true;
        }

        match requested_type {
            LockType::Read => {
                // Read locks can coexist with other read locks
                existing_locks
                    .iter()
                    .all(|lock| lock.lock_type == LockType::Read)
            }
            LockType::Write => {
                // Write locks are exclusive
                existing_locks.is_empty()
            }
        }
    }

    /// Add a lock request to the waiting queue.
    fn add_to_waiting_queue(
        &self,
        transaction_id: TransactionId,
        resource: LockableResource,
        lock_type: LockType,
    ) {
        let mut waiting = self.waiting_locks.lock();

        // Avoid duplicate requests
        if !waiting.iter().any(|req| {
            req.transaction_id == transaction_id
                && req.resource == resource
                && req.lock_type == lock_type
        }) {
            waiting.push(LockRequest {
                transaction_id,
                resource,
                lock_type,
                requested_at: Instant::now(),
            });
        }
    }

    /// Release all locks held by a transaction.
    pub fn release_all_locks(&self, transaction_id: TransactionId) -> Result<()> {
        let mut locks = self.active_locks.write();

        // Remove all locks held by this transaction
        for (_resource, lock_list) in locks.iter_mut() {
            lock_list.retain(|lock| lock.transaction_id != transaction_id);
        }

        // Remove empty entries
        locks.retain(|_resource, lock_list| !lock_list.is_empty());

        // Remove from waiting queue
        let mut waiting = self.waiting_locks.lock();
        waiting.retain(|req| req.transaction_id != transaction_id);

        Ok(())
    }

    /// Release a specific lock held by a transaction.
    pub fn release_lock(
        &self,
        transaction_id: TransactionId,
        resource: LockableResource,
    ) -> Result<()> {
        let mut locks = self.active_locks.write();

        if let Some(lock_list) = locks.get_mut(&resource) {
            lock_list.retain(|lock| {
                !(lock.transaction_id == transaction_id && lock.resource == resource)
            });

            if lock_list.is_empty() {
                locks.remove(&resource);
            }
        }

        Ok(())
    }

    /// Simple deadlock detection using wait-for graph.
    fn detect_deadlock(&self, transaction_id: TransactionId) -> Result<bool> {
        let locks = self.active_locks.read();
        let waiting = self.waiting_locks.lock();

        // Build wait-for graph
        let mut wait_for: HashMap<TransactionId, HashSet<TransactionId>> = HashMap::new();

        for request in waiting.iter() {
            if let Some(holders) = locks.get(&request.resource) {
                let waiting_tx = request.transaction_id;
                let holder_txs: HashSet<TransactionId> = holders
                    .iter()
                    .map(|lock| lock.transaction_id)
                    .filter(|&id| id != waiting_tx)
                    .collect();

                if !holder_txs.is_empty() {
                    wait_for.insert(waiting_tx, holder_txs);
                }
            }
        }

        // Check for cycles using DFS
        let mut visited = HashSet::new();
        let mut rec_stack = HashSet::new();

        self.has_cycle_dfs(transaction_id, &wait_for, &mut visited, &mut rec_stack)
    }

    /// DFS-based cycle detection.
    #[allow(clippy::only_used_in_recursion)]
    fn has_cycle_dfs(
        &self,
        node: TransactionId,
        graph: &HashMap<TransactionId, HashSet<TransactionId>>,
        visited: &mut HashSet<TransactionId>,
        rec_stack: &mut HashSet<TransactionId>,
    ) -> Result<bool> {
        visited.insert(node);
        rec_stack.insert(node);

        if let Some(neighbors) = graph.get(&node) {
            for &neighbor in neighbors {
                if !visited.contains(&neighbor) {
                    if self.has_cycle_dfs(neighbor, graph, visited, rec_stack)? {
                        return Ok(true);
                    }
                } else if rec_stack.contains(&neighbor) {
                    return Ok(true); // Cycle detected
                }
            }
        }

        rec_stack.remove(&node);
        Ok(false)
    }

    /// Get lock statistics for monitoring.
    pub fn get_lock_statistics(&self) -> LockStatistics {
        let locks = self.active_locks.read();
        let waiting = self.waiting_locks.lock();

        let total_active_locks = locks.values().map(|v| v.len()).sum();
        let waiting_requests = waiting.len();
        let locked_resources = locks.len();

        LockStatistics {
            total_active_locks,
            waiting_requests,
            locked_resources,
        }
    }
}

impl Default for LockManager {
    fn default() -> Self {
        Self::new()
    }
}

/// Statistics about the current locking state.
#[derive(Debug, Clone)]
pub struct LockStatistics {
    /// Total number of locks currently held across all transactions
    pub total_active_locks: usize,
    /// Number of lock requests waiting to be granted
    pub waiting_requests: usize,
    /// Number of distinct resources that have at least one lock
    pub locked_resources: usize,
}

/// Enhanced transaction with concurrency control.
pub struct ConcurrentTransaction {
    /// Base transaction
    pub transaction: Transaction,
    /// Lock manager reference
    lock_manager: Arc<LockManager>,
    /// Locks held by this transaction
    held_locks: HashSet<LockableResource>,
}

impl ConcurrentTransaction {
    /// Create a new concurrent transaction.
    pub fn new(transaction: Transaction, lock_manager: Arc<LockManager>) -> Self {
        ConcurrentTransaction {
            transaction,
            lock_manager,
            held_locks: HashSet::new(),
        }
    }

    /// Acquire a read lock on a resource.
    pub fn read_lock(&mut self, resource: LockableResource) -> Result<()> {
        self.lock_manager.acquire_lock(
            self.transaction.id(),
            resource,
            LockType::Read,
            Some(Duration::from_secs(30)),
        )?;
        self.held_locks.insert(resource);
        Ok(())
    }

    /// Acquire a write lock on a resource.
    pub fn write_lock(&mut self, resource: LockableResource) -> Result<()> {
        self.lock_manager.acquire_lock(
            self.transaction.id(),
            resource,
            LockType::Write,
            Some(Duration::from_secs(30)),
        )?;
        self.held_locks.insert(resource);
        Ok(())
    }

    /// Commit the transaction and release all locks.
    pub fn commit(mut self) -> Result<()> {
        self.transaction.commit()?;
        self.lock_manager.release_all_locks(self.transaction.id())?;
        self.held_locks.clear();
        Ok(())
    }

    /// Rollback the transaction and release all locks.
    pub fn rollback(mut self) -> Result<()> {
        self.transaction.rollback()?;
        self.lock_manager.release_all_locks(self.transaction.id())?;
        self.held_locks.clear();
        Ok(())
    }

    /// Get the transaction ID.
    pub fn id(&self) -> TransactionId {
        self.transaction.id()
    }

    /// Check if the transaction is active.
    pub fn is_active(&self) -> bool {
        self.transaction.is_active()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::transaction::{IsolationLevel, TransactionManager};
    use std::time::Duration;

    #[test]
    fn test_lock_acquisition() {
        let lock_manager = LockManager::new();
        let tx_id = 1;
        let resource = LockableResource::Node(100);

        // Acquire read lock
        let result = lock_manager.acquire_lock(
            tx_id,
            resource,
            LockType::Read,
            Some(Duration::from_millis(100)),
        );
        assert!(result.is_ok());
        assert!(result.unwrap());

        // Acquire another read lock on same resource
        let result = lock_manager.acquire_lock(
            2,
            resource,
            LockType::Read,
            Some(Duration::from_millis(100)),
        );
        assert!(result.is_ok());
        assert!(result.unwrap());

        // Try to acquire write lock (should fail due to existing read locks)
        let result = lock_manager.acquire_lock(
            3,
            resource,
            LockType::Write,
            Some(Duration::from_millis(100)),
        );
        assert!(result.is_err());

        // Release locks
        lock_manager.release_all_locks(tx_id).unwrap();
        lock_manager.release_all_locks(2).unwrap();

        // Now write lock should succeed
        let result = lock_manager.acquire_lock(
            3,
            resource,
            LockType::Write,
            Some(Duration::from_millis(100)),
        );
        assert!(result.is_ok());
        assert!(result.unwrap());
    }

    #[test]
    fn test_concurrent_transaction() {
        let lock_manager = Arc::new(LockManager::new());
        let tx_manager = TransactionManager::new(IsolationLevel::ReadCommitted);

        let transaction = tx_manager.begin();
        let mut concurrent_tx = ConcurrentTransaction::new(transaction, lock_manager);

        let resource = LockableResource::Node(200);

        // Test locking
        assert!(concurrent_tx.read_lock(resource).is_ok());
        assert!(concurrent_tx
            .write_lock(LockableResource::Node(201))
            .is_ok());

        // Test commit
        assert!(concurrent_tx.commit().is_ok());
    }

    #[test]
    fn test_lock_statistics() {
        let lock_manager = LockManager::new();

        // Initially no locks
        let stats = lock_manager.get_lock_statistics();
        assert_eq!(stats.total_active_locks, 0);
        assert_eq!(stats.waiting_requests, 0);
        assert_eq!(stats.locked_resources, 0);

        // Acquire some locks
        lock_manager
            .acquire_lock(
                1,
                LockableResource::Node(100),
                LockType::Read,
                Some(Duration::from_millis(100)),
            )
            .unwrap();
        lock_manager
            .acquire_lock(
                2,
                LockableResource::Node(101),
                LockType::Write,
                Some(Duration::from_millis(100)),
            )
            .unwrap();

        let stats = lock_manager.get_lock_statistics();
        assert_eq!(stats.total_active_locks, 2);
        assert_eq!(stats.locked_resources, 2);
    }
}