firecloud_storage/

quota.rs

//! Storage quota management
//!
//! Tracks and enforces storage limits:
//! - Local storage usage tracking
//! - Quota enforcement
//! - Storage statistics

use crate::{StorageError, StorageResult};
use serde::{Deserialize, Serialize};
use sled::Db;
use std::path::Path;
use tracing::{debug, info, warn};

const TREE_QUOTA: &str = "quota";
const KEY_TOTAL_BYTES: &[u8] = b"total_bytes";
const KEY_TOTAL_CHUNKS: &[u8] = b"total_chunks";
const KEY_QUOTA_LIMIT: &[u8] = b"quota_limit";

/// Storage quota and usage statistics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StorageQuota {
    /// Total bytes currently stored
    pub total_bytes: u64,
    /// Total number of chunks stored
    pub total_chunks: u64,
    /// Maximum allowed storage in bytes (0 = unlimited)
    pub quota_limit: u64,
}

impl StorageQuota {
    /// Check if quota is available for storing additional bytes
    pub fn has_space(&self, bytes: u64) -> bool {
        if self.quota_limit == 0 {
            return true; // Unlimited
        }
        self.total_bytes + bytes <= self.quota_limit
    }

    /// Get remaining space in bytes
    pub fn remaining_bytes(&self) -> u64 {
        if self.quota_limit == 0 {
            return u64::MAX; // Unlimited
        }
        self.quota_limit.saturating_sub(self.total_bytes)
    }

    /// Get usage percentage (0-100)
    pub fn usage_percentage(&self) -> f64 {
        if self.quota_limit == 0 {
            return 0.0; // Unlimited
        }
        (self.total_bytes as f64 / self.quota_limit as f64) * 100.0
    }

    /// Check if quota is nearly full (>90%)
    pub fn is_nearly_full(&self) -> bool {
        self.usage_percentage() > 90.0
    }

    /// Check if quota is full
    pub fn is_full(&self) -> bool {
        if self.quota_limit == 0 {
            return false; // Unlimited
        }
        self.total_bytes >= self.quota_limit
    }
}

impl Default for StorageQuota {
    fn default() -> Self {
        Self {
            total_bytes: 0,
            total_chunks: 0,
            quota_limit: 0, // Unlimited by default
        }
    }
}

/// Manages storage quota using Sled
pub struct QuotaManager {
    db: Db,
}

impl QuotaManager {
    /// Open or create a quota manager
    pub fn open<P: AsRef<Path>>(path: P) -> StorageResult<Self> {
        let db = sled::open(path).map_err(|e| StorageError::Database(e.to_string()))?;
        info!("Opened quota manager");
        Ok(Self { db })
    }

    /// Get current quota and usage statistics
    pub fn get_quota(&self) -> StorageResult<StorageQuota> {
        let tree = self
            .db
            .open_tree(TREE_QUOTA)
            .map_err(|e| StorageError::Database(e.to_string()))?;

        let total_bytes = tree
            .get(KEY_TOTAL_BYTES)
            .map_err(|e| StorageError::Database(e.to_string()))?
            .map(|bytes| u64::from_le_bytes(bytes.as_ref().try_into().unwrap_or([0; 8])))
            .unwrap_or(0);

        let total_chunks = tree
            .get(KEY_TOTAL_CHUNKS)
            .map_err(|e| StorageError::Database(e.to_string()))?
            .map(|bytes| u64::from_le_bytes(bytes.as_ref().try_into().unwrap_or([0; 8])))
            .unwrap_or(0);

        let quota_limit = tree
            .get(KEY_QUOTA_LIMIT)
            .map_err(|e| StorageError::Database(e.to_string()))?
            .map(|bytes| u64::from_le_bytes(bytes.as_ref().try_into().unwrap_or([0; 8])))
            .unwrap_or(0);

        Ok(StorageQuota {
            total_bytes,
            total_chunks,
            quota_limit,
        })
    }

    /// Set quota limit in bytes (0 = unlimited)
    pub fn set_quota_limit(&self, limit: u64) -> StorageResult<()> {
        let tree = self
            .db
            .open_tree(TREE_QUOTA)
            .map_err(|e| StorageError::Database(e.to_string()))?;

        tree.insert(KEY_QUOTA_LIMIT, &limit.to_le_bytes())
            .map_err(|e| StorageError::Database(e.to_string()))?;

        info!("Set quota limit to {} bytes", limit);
        Ok(())
    }

    /// Add storage usage (when storing a new chunk)
    pub fn add_usage(&self, bytes: u64) -> StorageResult<()> {
        let tree = self
            .db
            .open_tree(TREE_QUOTA)
            .map_err(|e| StorageError::Database(e.to_string()))?;

        // Check if we have space
        let quota = self.get_quota()?;
        if !quota.has_space(bytes) {
            return Err(StorageError::QuotaExceeded {
                requested: bytes,
                available: quota.remaining_bytes(),
            });
        }

        // Update total bytes
        let new_total = quota.total_bytes + bytes;
        tree.insert(KEY_TOTAL_BYTES, &new_total.to_le_bytes())
            .map_err(|e| StorageError::Database(e.to_string()))?;

        // Update total chunks
        let new_chunks = quota.total_chunks + 1;
        tree.insert(KEY_TOTAL_CHUNKS, &new_chunks.to_le_bytes())
            .map_err(|e| StorageError::Database(e.to_string()))?;

        debug!(
            "Added {} bytes, total: {} / {} chunks: {}",
            bytes, new_total, quota.quota_limit, new_chunks
        );

        // Warn if nearly full
        if new_total as f64 / quota.quota_limit as f64 > 0.9 && quota.quota_limit > 0 {
            warn!(
                "Storage quota nearly full: {:.1}%",
                (new_total as f64 / quota.quota_limit as f64) * 100.0
            );
        }

        Ok(())
    }

    /// Remove storage usage (when deleting a chunk)
    pub fn remove_usage(&self, bytes: u64) -> StorageResult<()> {
        let tree = self
            .db
            .open_tree(TREE_QUOTA)
            .map_err(|e| StorageError::Database(e.to_string()))?;

        let quota = self.get_quota()?;

        // Update total bytes
        let new_total = quota.total_bytes.saturating_sub(bytes);
        tree.insert(KEY_TOTAL_BYTES, &new_total.to_le_bytes())
            .map_err(|e| StorageError::Database(e.to_string()))?;

        // Update total chunks
        let new_chunks = quota.total_chunks.saturating_sub(1);
        tree.insert(KEY_TOTAL_CHUNKS, &new_chunks.to_le_bytes())
            .map_err(|e| StorageError::Database(e.to_string()))?;

        debug!(
            "Removed {} bytes, total: {} chunks: {}",
            bytes, new_total, new_chunks
        );

        Ok(())
    }

    /// Check if we can store additional bytes
    pub fn can_store(&self, bytes: u64) -> StorageResult<bool> {
        let quota = self.get_quota()?;
        Ok(quota.has_space(bytes))
    }

    /// Get usage statistics
    pub fn get_stats(&self) -> StorageResult<QuotaStats> {
        let quota = self.get_quota()?;

        Ok(QuotaStats {
            total_bytes: quota.total_bytes,
            total_chunks: quota.total_chunks,
            quota_limit: quota.quota_limit,
            remaining_bytes: quota.remaining_bytes(),
            usage_percentage: quota.usage_percentage(),
            is_nearly_full: quota.is_nearly_full(),
            is_full: quota.is_full(),
        })
    }

    /// Reset all usage statistics (dangerous!)
    pub fn reset(&self) -> StorageResult<()> {
        let tree = self
            .db
            .open_tree(TREE_QUOTA)
            .map_err(|e| StorageError::Database(e.to_string()))?;

        tree.clear().map_err(|e| StorageError::Database(e.to_string()))?;

        warn!("Storage quota statistics reset!");
        Ok(())
    }
}

/// Detailed quota statistics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QuotaStats {
    pub total_bytes: u64,
    pub total_chunks: u64,
    pub quota_limit: u64,
    pub remaining_bytes: u64,
    pub usage_percentage: f64,
    pub is_nearly_full: bool,
    pub is_full: bool,
}

impl QuotaStats {
    /// Format bytes in human-readable format
    pub fn format_bytes(bytes: u64) -> String {
        const UNITS: &[&str] = &["B", "KB", "MB", "GB", "TB"];
        let mut size = bytes as f64;
        let mut unit_idx = 0;

        while size >= 1024.0 && unit_idx < UNITS.len() - 1 {
            size /= 1024.0;
            unit_idx += 1;
        }

        format!("{:.2} {}", size, UNITS[unit_idx])
    }

    /// Get a human-readable summary
    pub fn summary(&self) -> String {
        if self.quota_limit == 0 {
            format!(
                "{} used in {} chunks (unlimited)",
                Self::format_bytes(self.total_bytes),
                self.total_chunks
            )
        } else {
            format!(
                "{} / {} used ({:.1}%) in {} chunks",
                Self::format_bytes(self.total_bytes),
                Self::format_bytes(self.quota_limit),
                self.usage_percentage,
                self.total_chunks
            )
        }
    }
}

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

    #[test]
    fn test_quota_basic() {
        let dir = tempdir().unwrap();
        let manager = QuotaManager::open(dir.path().join("quota")).unwrap();

        // Set quota to 1GB
        manager.set_quota_limit(1_000_000_000).unwrap();

        // Add some usage
        manager.add_usage(500_000_000).unwrap();

        let stats = manager.get_stats().unwrap();
        assert_eq!(stats.total_bytes, 500_000_000);
        assert_eq!(stats.total_chunks, 1);
        assert_eq!(stats.usage_percentage, 50.0);

        // Check we can store more
        assert!(manager.can_store(400_000_000).unwrap());

        // Check we can't exceed quota
        assert!(!manager.can_store(600_000_000).unwrap());
    }

    #[test]
    fn test_quota_exceeded() {
        let dir = tempdir().unwrap();
        let manager = QuotaManager::open(dir.path().join("quota")).unwrap();

        manager.set_quota_limit(1000).unwrap();
        manager.add_usage(800).unwrap();

        // Should fail when exceeding quota
        let result = manager.add_usage(300);
        assert!(matches!(result, Err(StorageError::QuotaExceeded { .. })));
    }

    #[test]
    fn test_remove_usage() {
        let dir = tempdir().unwrap();
        let manager = QuotaManager::open(dir.path().join("quota")).unwrap();

        manager.add_usage(1000).unwrap();
        assert_eq!(manager.get_stats().unwrap().total_bytes, 1000);

        manager.remove_usage(400).unwrap();
        assert_eq!(manager.get_stats().unwrap().total_bytes, 600);
        assert_eq!(manager.get_stats().unwrap().total_chunks, 0);
    }

    #[test]
    fn test_unlimited_quota() {
        let dir = tempdir().unwrap();
        let manager = QuotaManager::open(dir.path().join("quota")).unwrap();

        // No limit set (0 = unlimited)
        manager.add_usage(u64::MAX / 2).unwrap();
        assert!(manager.can_store(u64::MAX / 2).unwrap());
    }
}