lcpfs 2026.1.102

// Copyright 2025 LunaOS Contributors
// SPDX-License-Identifier: Apache-2.0
//
// Reservation Enforcement
// Guarantee minimum space for datasets.

use crate::FsError;
use crate::mgmt::dataset::DatasetPhys;
use alloc::collections::BTreeMap;
use lazy_static::lazy_static;
use spin::Mutex;

lazy_static! {
    /// Global reservation tracker per pool
    static ref RESERVATION_TRACKER: Mutex<BTreeMap<u64, PoolReservations>> = Mutex::new(BTreeMap::new());
}

/// Pool-wide reservation tracking
#[derive(Debug, Clone)]
pub struct PoolReservations {
    /// Total pool capacity in bytes
    pub total_capacity: u64,
    /// Currently reserved bytes across all datasets
    pub total_reserved: u64,
    /// Current pool usage
    pub total_used: u64,
    /// Per-dataset reservations (dataset_id -> reserved_bytes)
    pub dataset_reservations: BTreeMap<u64, u64>,
}

impl PoolReservations {
    /// Create new pool reservation tracker
    pub fn new(total_capacity: u64) -> Self {
        Self {
            total_capacity,
            total_reserved: 0,
            total_used: 0,
            dataset_reservations: BTreeMap::new(),
        }
    }

    /// Get available space for new reservations
    ///
    /// Available = capacity - max(reserved, used)
    /// We must honor both reserved space and actual usage
    pub fn available_for_reservation(&self) -> u64 {
        let committed = self.total_reserved.max(self.total_used);
        self.total_capacity.saturating_sub(committed)
    }

    /// Get free space available for writes (unreserved)
    ///
    /// Free = capacity - reserved - used
    pub fn free_unreserved(&self) -> u64 {
        self.total_capacity
            .saturating_sub(self.total_reserved)
            .saturating_sub(self.total_used)
    }
}

/// Reservation enforcement engine
pub struct ReservationEngine;

impl ReservationEngine {
    /// Initialize reservation tracking for a pool
    ///
    /// # Arguments
    /// * `pool_id` - Pool identifier
    /// * `total_capacity` - Total pool capacity in bytes
    pub fn init_pool(pool_id: u64, total_capacity: u64) {
        let mut tracker = RESERVATION_TRACKER.lock();
        tracker.insert(pool_id, PoolReservations::new(total_capacity));
    }

    /// Reserve space for a dataset
    ///
    /// # Arguments
    /// * `pool_id` - Pool identifier
    /// * `dataset_id` - Dataset identifier
    /// * `reservation_bytes` - Bytes to reserve
    ///
    /// # Returns
    /// * `Ok(())` - Reservation successful
    /// * `Err(FsError::DiskFull { .. })` - Insufficient space for reservation
    ///
    /// # Algorithm
    /// 1. Check if pool has enough unreserved capacity
    /// 2. Add reservation to pool's total
    /// 3. Track per-dataset reservation
    pub fn reserve_space(
        pool_id: u64,
        dataset_id: u64,
        reservation_bytes: u64,
    ) -> Result<(), FsError> {
        if reservation_bytes == 0 {
            return Ok(()); // No reservation needed
        }

        let mut tracker = RESERVATION_TRACKER.lock();
        let pool = tracker.get_mut(&pool_id).ok_or(FsError::NoDevice)?;

        // Check if we have enough space for this reservation
        if reservation_bytes > pool.available_for_reservation() {
            return Err(FsError::DiskFull {
                needed_bytes: reservation_bytes,
            });
        }

        // Add reservation
        pool.total_reserved = pool.total_reserved.saturating_add(reservation_bytes);
        pool.dataset_reservations
            .insert(dataset_id, reservation_bytes);

        Ok(())
    }

    /// Release a dataset's reservation
    ///
    /// # Arguments
    /// * `pool_id` - Pool identifier
    /// * `dataset_id` - Dataset identifier
    pub fn release_reservation(pool_id: u64, dataset_id: u64) {
        let mut tracker = RESERVATION_TRACKER.lock();
        if let Some(pool) = tracker.get_mut(&pool_id) {
            if let Some(reserved) = pool.dataset_reservations.remove(&dataset_id) {
                pool.total_reserved = pool.total_reserved.saturating_sub(reserved);
            }
        }
    }

    /// Update a dataset's reservation
    ///
    /// # Arguments
    /// * `pool_id` - Pool identifier
    /// * `dataset_id` - Dataset identifier
    /// * `new_reservation` - New reservation size in bytes
    pub fn update_reservation(
        pool_id: u64,
        dataset_id: u64,
        new_reservation: u64,
    ) -> Result<(), FsError> {
        let mut tracker = RESERVATION_TRACKER.lock();
        let pool = tracker.get_mut(&pool_id).ok_or(FsError::NoDevice)?;

        let old_reservation = pool
            .dataset_reservations
            .get(&dataset_id)
            .copied()
            .unwrap_or(0);

        if new_reservation > old_reservation {
            // Increasing reservation - check if space available
            let additional = new_reservation - old_reservation;
            if additional > pool.available_for_reservation() {
                return Err(FsError::DiskFull {
                    needed_bytes: additional,
                });
            }
            pool.total_reserved = pool.total_reserved.saturating_add(additional);
        } else {
            // Decreasing reservation - free up space
            let decrease = old_reservation - new_reservation;
            pool.total_reserved = pool.total_reserved.saturating_sub(decrease);
        }

        if new_reservation > 0 {
            pool.dataset_reservations
                .insert(dataset_id, new_reservation);
        } else {
            pool.dataset_reservations.remove(&dataset_id);
        }

        Ok(())
    }

    /// Check if a write is allowed (doesn't violate other datasets' reservations)
    ///
    /// # Arguments
    /// * `pool_id` - Pool identifier
    /// * `dataset` - Dataset physical metadata
    /// * `write_bytes` - Bytes to be written
    ///
    /// # Returns
    /// * `Ok(())` - Write is allowed
    /// * `Err(FsError::DiskFull { .. })` - Write would violate reservations
    ///
    /// # Algorithm
    /// 1. If dataset has reservation, allow (guaranteed space)
    /// 2. Otherwise, check if write would consume reserved space
    /// 3. Only allow if write fits in unreserved free space
    pub fn check_write_allowed(
        pool_id: u64,
        dataset: &DatasetPhys,
        write_bytes: u64,
    ) -> Result<(), FsError> {
        let tracker = RESERVATION_TRACKER.lock();
        let pool = tracker.get(&pool_id).ok_or(FsError::NoDevice)?;

        // If this dataset has a reservation, it has guaranteed space
        if dataset.reservation > 0 {
            // Check if write stays within reservation
            let new_usage = dataset.used_bytes.saturating_add(write_bytes);
            if new_usage <= dataset.reservation {
                return Ok(()); // Within reservation, always allowed
            }
        }

        // For non-reserved or over-reservation writes, check unreserved space
        let new_total_used = pool.total_used.saturating_add(write_bytes);
        let max_allowed_usage = pool.total_capacity.saturating_sub(pool.total_reserved);

        if new_total_used > max_allowed_usage {
            return Err(FsError::DiskFull {
                needed_bytes: write_bytes,
            });
        }

        Ok(())
    }

    /// Update pool usage after a write
    ///
    /// # Arguments
    /// * `pool_id` - Pool identifier
    /// * `bytes_written` - Bytes written to pool
    pub fn update_pool_usage(pool_id: u64, bytes_written: u64) {
        let mut tracker = RESERVATION_TRACKER.lock();
        if let Some(pool) = tracker.get_mut(&pool_id) {
            pool.total_used = pool.total_used.saturating_add(bytes_written);
        }
    }

    /// Decrease pool usage after deletion/truncation
    ///
    /// # Arguments
    /// * `pool_id` - Pool identifier
    /// * `bytes_freed` - Bytes freed from pool
    pub fn decrease_pool_usage(pool_id: u64, bytes_freed: u64) {
        let mut tracker = RESERVATION_TRACKER.lock();
        if let Some(pool) = tracker.get_mut(&pool_id) {
            pool.total_used = pool.total_used.saturating_sub(bytes_freed);
        }
    }

    /// Get pool reservation statistics
    ///
    /// # Arguments
    /// * `pool_id` - Pool identifier
    ///
    /// # Returns
    /// * `Some(stats)` - Reservation statistics
    /// * `None` - Pool not found
    pub fn get_stats(pool_id: u64) -> Option<PoolReservations> {
        let tracker = RESERVATION_TRACKER.lock();
        tracker.get(&pool_id).cloned()
    }
}

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

    #[test]
    fn test_pool_init() {
        ReservationEngine::init_pool(1, 1_000_000);
        let stats = ReservationEngine::get_stats(1).expect("test: operation should succeed");
        assert_eq!(stats.total_capacity, 1_000_000);
        assert_eq!(stats.total_reserved, 0);
        assert_eq!(stats.total_used, 0);
    }

    #[test]
    fn test_reserve_space() {
        ReservationEngine::init_pool(2, 1_000_000);

        // Reserve 100K for dataset 1
        assert!(ReservationEngine::reserve_space(2, 1, 100_000).is_ok());

        let stats = ReservationEngine::get_stats(2).expect("test: operation should succeed");
        assert_eq!(stats.total_reserved, 100_000);
        assert_eq!(stats.available_for_reservation(), 900_000);
    }

    #[test]
    fn test_reservation_exceeds_capacity() {
        ReservationEngine::init_pool(3, 1_000_000);

        // Try to reserve more than capacity
        assert!(matches!(
            ReservationEngine::reserve_space(3, 1, 2_000_000),
            Err(FsError::DiskFull { .. })
        ));
    }

    #[test]
    fn test_multiple_reservations() {
        ReservationEngine::init_pool(4, 1_000_000);

        assert!(ReservationEngine::reserve_space(4, 1, 300_000).is_ok());
        assert!(ReservationEngine::reserve_space(4, 2, 400_000).is_ok());

        let stats = ReservationEngine::get_stats(4).expect("test: operation should succeed");
        assert_eq!(stats.total_reserved, 700_000);
        assert_eq!(stats.available_for_reservation(), 300_000);

        // Should reject reservation that would exceed capacity
        assert!(matches!(
            ReservationEngine::reserve_space(4, 3, 400_000),
            Err(FsError::DiskFull { .. })
        ));
    }

    #[test]
    fn test_release_reservation() {
        ReservationEngine::init_pool(5, 1_000_000);

        ReservationEngine::reserve_space(5, 1, 500_000).expect("test: operation should succeed");
        ReservationEngine::release_reservation(5, 1);

        let stats = ReservationEngine::get_stats(5).expect("test: operation should succeed");
        assert_eq!(stats.total_reserved, 0);
        assert_eq!(stats.available_for_reservation(), 1_000_000);
    }

    #[test]
    fn test_write_with_reservation() {
        ReservationEngine::init_pool(6, 1_000_000);
        ReservationEngine::reserve_space(6, 1, 500_000).expect("test: operation should succeed");

        let dataset = DatasetPhys {
            reservation: 500_000,
            used_bytes: 100_000,
            ..DatasetPhys::new(1)
        };

        // Write within reservation should succeed
        assert!(ReservationEngine::check_write_allowed(6, &dataset, 300_000).is_ok());
    }

    #[test]
    fn test_write_without_reservation() {
        ReservationEngine::init_pool(7, 1_000_000);
        ReservationEngine::reserve_space(7, 1, 800_000).expect("test: operation should succeed"); // Reserve 800K for dataset 1

        let dataset2 = DatasetPhys {
            reservation: 0, // No reservation
            used_bytes: 0,
            ..DatasetPhys::new(2)
        };

        // Dataset 2 can only use unreserved space (200K)
        assert!(ReservationEngine::check_write_allowed(7, &dataset2, 100_000).is_ok());
        assert!(matches!(
            ReservationEngine::check_write_allowed(7, &dataset2, 300_000),
            Err(FsError::DiskFull { .. })
        ));
    }
}