BlueArchiveToolkit/crates/bat-cas-engine/src/repository.rs

//! CAS V1 repository，组合对象存储与引用计数元数据。

use crate::error::{CasError, Result};
use crate::hash::{compute_hash, Hash};
use crate::refcount::SqliteRefCounter;
use crate::storage::{FileSystemStorage, Storage, StorageStats};
use std::path::Path;

/// 文件系统 CAS repository。
///
/// 该类型是 CAS V1 的核心实现。对象文件由 `FileSystemStorage` 管理，
/// 引用计数和对象元数据由 `SqliteRefCounter` 管理。
#[derive(Debug, Clone)]
pub struct FileSystemCasRepository {
    storage: FileSystemStorage,
    ref_counter: SqliteRefCounter,
}

impl FileSystemCasRepository {
    /// 打开或初始化 CAS repository。
    pub async fn new(root: impl AsRef<Path>) -> Result<Self> {
        let root = root.as_ref();
        let storage = FileSystemStorage::new(root).await?;
        let ref_counter = SqliteRefCounter::new(root.join("metadata.sqlite")).await?;

        Ok(Self {
            storage,
            ref_counter,
        })
    }

    /// 返回底层文件存储。
    pub fn storage(&self) -> &FileSystemStorage {
        &self.storage
    }

    /// 存储对象并增加引用计数。
    pub async fn store(&self, data: &[u8]) -> Result<Hash> {
        let hash = compute_hash(data);
        let existed = self.storage.exists(&hash).await?;
        let stored_hash = self.storage.put(data).await?;

        // 单条 UPSERT 原子建行并 +1：对象行不会在 store 期间以 ref_count=0
        // 暴露给并发 gc，消除“已存对象、尚未加引用”的删除窗口。
        match self
            .ref_counter
            .store_reference(&stored_hash, data.len() as u64)
            .await
        {
            Ok(count) => {
                debug_assert!(count > 0);
                Ok(stored_hash)
            }
            Err(error) => {
                if !existed {
                    self.delete_new_object_after_metadata_failure(&stored_hash)
                        .await?;
                }
                Err(error)
            }
        }
    }

    async fn delete_new_object_after_metadata_failure(&self, hash: &Hash) -> Result<()> {
        match self.storage.delete(hash).await {
            Ok(()) | Err(CasError::ObjectNotFound(_)) => Ok(()),
            Err(error) => Err(error),
        }
    }

    /// 读取对象并验证 Hash。
    pub async fn get(&self, hash: &Hash) -> Result<Vec<u8>> {
        let data = self.storage.get(hash).await?;
        Ok(data)
    }

    /// 检查对象是否存在。
    pub async fn exists(&self, hash: &Hash) -> Result<bool> {
        self.storage.exists(hash).await
    }

    /// 增加引用计数。
    pub async fn add_reference(&self, hash: &Hash) -> Result<u64> {
        if !self.storage.exists(hash).await? {
            return Err(CasError::ObjectNotFound(hash.to_string()));
        }

        // 同样走原子 UPSERT 递增，避免 ensure_object 与 add_reference 之间
        // 出现可被并发 gc 删除的 ref_count=0 窗口。
        let size = self.storage.size(hash).await?;
        let count = self.ref_counter.store_reference(hash, size).await?;
        Ok(count)
    }

    /// 减少引用计数。
    pub async fn remove_reference(&self, hash: &Hash) -> Result<u64> {
        self.ref_counter.remove_reference(hash).await
    }

    /// 获取引用计数。
    pub async fn get_reference_count(&self, hash: &Hash) -> Result<u64> {
        self.ref_counter.get_reference_count(hash).await
    }

    /// 返回当前 GC 候选对象。
    pub async fn gc_candidates(&self) -> Result<Vec<Hash>> {
        self.ref_counter.zero_ref_objects().await
    }

    /// 删除引用计数为 0 的对象。
    pub async fn gc(&self) -> Result<u64> {
        let candidates = self.gc_candidates().await?;
        let mut deleted = 0u64;

        for hash in candidates {
            // 先原子删除 ref_count=0 的元数据行；若被并发递增抢先，rows_affected=0，
            // 跳过，绝不删除仍被引用对象的文件。删元数据成功后再删文件——最坏只留下
            // 无元数据的孤儿文件（可被后续覆盖，无数据丢失），而非删掉被引用的内容。
            if !self.ref_counter.delete_zero_ref_metadata(&hash).await? {
                continue;
            }

            match self.storage.delete(&hash).await {
                Ok(()) | Err(CasError::ObjectNotFound(_)) => {}
                Err(error) => return Err(error),
            }
            deleted += 1;
        }

        Ok(deleted)
    }

    /// 获取存储统计信息。
    pub async fn stats(&self) -> Result<StorageStats> {
        self.storage.stats().await
    }
}

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

    async fn temp_repo() -> (tempfile::TempDir, FileSystemCasRepository) {
        let temp_dir = tempfile::tempdir().unwrap();
        let repo = FileSystemCasRepository::new(temp_dir.path()).await.unwrap();
        (temp_dir, repo)
    }

    #[tokio::test]
    async fn store_increments_reference_count() {
        let (_temp_dir, repo) = temp_repo().await;

        let hash = repo.store(b"same data").await.unwrap();
        let second = repo.store(b"same data").await.unwrap();

        assert_eq!(hash, second);
        assert_eq!(repo.get_reference_count(&hash).await.unwrap(), 2);
        assert_eq!(repo.stats().await.unwrap().object_count, 1);
    }

    #[tokio::test]
    async fn gc_deletes_only_zero_reference_objects() {
        let (_temp_dir, repo) = temp_repo().await;

        let keep = repo.store(b"keep").await.unwrap();
        let delete = repo.store(b"delete").await.unwrap();

        repo.remove_reference(&delete).await.unwrap();

        assert_eq!(repo.gc().await.unwrap(), 1);
        assert!(repo.exists(&keep).await.unwrap());
        assert!(!repo.exists(&delete).await.unwrap());
        assert!(matches!(
            repo.get_reference_count(&delete).await,
            Err(CasError::ObjectNotFound(_))
        ));
    }

    #[tokio::test]
    async fn concurrent_store_is_safe_and_counts_references() {
        let (_temp_dir, repo) = temp_repo().await;
        let repo = Arc::new(repo);
        let mut tasks = Vec::new();

        for _ in 0..32 {
            let repo = Arc::clone(&repo);
            tasks.push(tokio::spawn(async move { repo.store(b"shared").await }));
        }

        let mut hashes = Vec::new();
        for task in tasks {
            hashes.push(task.await.unwrap().unwrap());
        }

        let first = hashes[0];
        assert!(hashes.iter().all(|hash| *hash == first));
        assert_eq!(repo.get_reference_count(&first).await.unwrap(), 32);
        assert_eq!(repo.stats().await.unwrap().object_count, 1);
        assert_eq!(repo.get(&first).await.unwrap(), b"shared");
    }

    #[tokio::test]
    async fn store_never_exposes_zero_reference_window() {
        let (_temp_dir, repo) = temp_repo().await;

        let hash = repo.store(b"payload").await.unwrap();
        // store 结束后引用计数为 1，绝不会成为 gc 候选。
        assert_eq!(repo.get_reference_count(&hash).await.unwrap(), 1);
        assert!(repo.gc_candidates().await.unwrap().is_empty());
        assert_eq!(repo.gc().await.unwrap(), 0);
        assert!(repo.exists(&hash).await.unwrap());
    }

    #[tokio::test]
    async fn gc_skips_object_reacquired_before_delete() {
        let (_temp_dir, repo) = temp_repo().await;

        let hash = repo.store(b"reacquired").await.unwrap();
        assert_eq!(repo.remove_reference(&hash).await.unwrap(), 0);
        // 归零后成为 gc 候选。
        assert_eq!(repo.gc_candidates().await.unwrap(), vec![hash]);
        // 在删除前被重新引用（模拟 store/gc 竞态中的重新获取）。
        assert_eq!(repo.add_reference(&hash).await.unwrap(), 1);
        // gc 的原子闸门应跳过它，对象文件保留。
        assert_eq!(repo.gc().await.unwrap(), 0);
        assert!(repo.exists(&hash).await.unwrap());
        assert_eq!(repo.get_reference_count(&hash).await.unwrap(), 1);
    }

    #[tokio::test]
    async fn corrupted_object_is_detected_through_repository() {
        let (_temp_dir, repo) = temp_repo().await;

        let hash = repo.store(b"valid").await.unwrap();
        let write_result = tokio::fs::write(repo.storage().object_path(&hash), b"invalid").await;
        write_result.unwrap();

        assert!(matches!(
            repo.get(&hash).await,
            Err(CasError::HashMismatch { .. })
        ));
    }

    #[tokio::test]
    async fn add_reference_requires_existing_object() {
        let (_temp_dir, repo) = temp_repo().await;
        let missing = compute_hash(b"missing");

        assert!(matches!(
            repo.add_reference(&missing).await,
            Err(CasError::ObjectNotFound(_))
        ));
    }

    #[tokio::test]
    async fn remove_reference_does_not_go_below_zero() {
        let (_temp_dir, repo) = temp_repo().await;

        let hash = repo.store(b"underflow").await.unwrap();
        assert_eq!(repo.remove_reference(&hash).await.unwrap(), 0);

        assert!(matches!(
            repo.remove_reference(&hash).await,
            Err(CasError::ReferenceUnderflow(_))
        ));
    }
}