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

//! CAS 引用计数和对象元数据。

use crate::error::{CasError, Result};
use crate::hash::Hash;
use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions, SqliteQueryResult};
use sqlx::SqlitePool;
use std::path::Path;
use std::str::FromStr;
use std::time::{Duration, SystemTime, UNIX_EPOCH};

/// CAS 对象元数据。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ObjectMetadata {
    /// 对象 Hash。
    pub hash: Hash,
    /// 对象大小，单位字节。
    pub size: u64,
    /// 当前引用计数。
    pub ref_count: u64,
    /// 创建时间，Unix 秒。
    pub created_at: i64,
    /// 更新时间，Unix 秒。
    pub updated_at: i64,
    /// 引用计数首次归零时间，Unix 秒。
    pub zero_ref_at: Option<i64>,
}

/// SQLite 引用计数存储。
#[derive(Debug, Clone)]
pub struct SqliteRefCounter {
    pool: SqlitePool,
}

impl SqliteRefCounter {
    /// 打开或创建 SQLite 元数据库。
    pub async fn new(path: impl AsRef<Path>) -> Result<Self> {
        if let Some(parent) = path.as_ref().parent() {
            tokio::fs::create_dir_all(parent).await?;
        }

        let options =
            SqliteConnectOptions::from_str(&format!("sqlite://{}", path.as_ref().display()))
                .map_err(|error| CasError::Database(error.to_string()))?
                .create_if_missing(true)
                .journal_mode(SqliteJournalMode::Wal)
                .busy_timeout(Duration::from_secs(30));

        let pool = SqlitePoolOptions::new()
            .max_connections(1)
            .connect_with(options)
            .await?;

        let counter = Self { pool };
        counter.init_schema().await?;
        Ok(counter)
    }

    async fn init_schema(&self) -> Result<()> {
        let _schema_result = Self::execute_query(
            &self.pool,
            sqlx::query(
                r#"
                CREATE TABLE IF NOT EXISTS cas_objects (
                    hash TEXT PRIMARY KEY NOT NULL,
                    size INTEGER NOT NULL CHECK(size >= 0),
                    ref_count INTEGER NOT NULL CHECK(ref_count >= 0),
                    created_at INTEGER NOT NULL,
                    updated_at INTEGER NOT NULL,
                    zero_ref_at INTEGER
                )
                "#,
            ),
        )
        .await?;

        let _index_result = Self::execute_query(
            &self.pool,
            sqlx::query(
                r#"
                CREATE INDEX IF NOT EXISTS idx_cas_objects_ref_count
                ON cas_objects(ref_count)
                "#,
            ),
        )
        .await?;

        Ok(())
    }

    async fn execute_query<'q>(
        pool: &SqlitePool,
        query: sqlx::query::Query<'q, sqlx::Sqlite, sqlx::sqlite::SqliteArguments<'q>>,
    ) -> Result<SqliteQueryResult> {
        let result = query.execute(pool).await?;
        Ok(result)
    }

    fn insert_or_update_object_query(
        hash: &Hash,
        size: u64,
        now: i64,
    ) -> sqlx::query::Query<'_, sqlx::Sqlite, sqlx::sqlite::SqliteArguments<'_>> {
        sqlx::query(
            r#"
            INSERT INTO cas_objects(hash, size, ref_count, created_at, updated_at, zero_ref_at)
            VALUES(?1, ?2, 0, ?3, ?3, ?3)
            ON CONFLICT(hash) DO UPDATE SET
                size = excluded.size,
                updated_at = excluded.updated_at
            "#,
        )
        .bind(hash.to_string())
        .bind(size as i64)
        .bind(now)
    }

    fn now() -> i64 {
        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .map(|duration| duration.as_secs() as i64)
            .unwrap_or_default()
    }

    /// 注册对象元数据。如果对象已存在，仅更新大小和更新时间。
    pub async fn ensure_object(&self, hash: &Hash, size: u64) -> Result<()> {
        let now = Self::now();
        let query = Self::insert_or_update_object_query(hash, size, now);
        let _upsert_result = Self::execute_query(&self.pool, query).await?;
        Ok(())
    }

    /// 原子地注册对象并把引用计数加一。
    ///
    /// 新对象直接以 `ref_count = 1` 建行，已存在对象在同一 UPSERT 语句内 `+1`，
    /// 因此对象行不会在“已注册但尚未引用”的瞬间以 `ref_count = 0` 暴露给并发 GC。
    pub async fn store_reference(&self, hash: &Hash, size: u64) -> Result<u64> {
        let now = Self::now();
        let count: i64 = sqlx::query_scalar(
            r#"
            INSERT INTO cas_objects(hash, size, ref_count, created_at, updated_at, zero_ref_at)
            VALUES(?1, ?2, 1, ?3, ?3, NULL)
            ON CONFLICT(hash) DO UPDATE SET
                ref_count = ref_count + 1,
                size = excluded.size,
                updated_at = excluded.updated_at,
                zero_ref_at = NULL
            RETURNING ref_count
            "#,
        )
        .bind(hash.to_string())
        .bind(size as i64)
        .bind(now)
        .fetch_one(&self.pool)
        .await?;

        Ok(count as u64)
    }

    /// 增加对象引用计数。
    pub async fn add_reference(&self, hash: &Hash) -> Result<u64> {
        let now = Self::now();
        let updated: Option<i64> = sqlx::query_scalar(
            r#"
            UPDATE cas_objects
            SET ref_count = ref_count + 1,
                updated_at = ?1,
                zero_ref_at = NULL
            WHERE hash = ?2
            RETURNING ref_count
            "#,
        )
        .bind(now)
        .bind(hash.to_string())
        .fetch_optional(&self.pool)
        .await?;

        updated
            .map(|count| count as u64)
            .ok_or_else(|| CasError::ObjectNotFound(hash.to_string()))
    }

    /// 减少对象引用计数。
    pub async fn remove_reference(&self, hash: &Hash) -> Result<u64> {
        let now = Self::now();
        let updated: Option<i64> = sqlx::query_scalar(
            r#"
            UPDATE cas_objects
            SET ref_count = ref_count - 1,
                updated_at = ?1,
                zero_ref_at = CASE WHEN ref_count = 1 THEN ?1 ELSE zero_ref_at END
            WHERE hash = ?2 AND ref_count > 0
            RETURNING ref_count
            "#,
        )
        .bind(now)
        .bind(hash.to_string())
        .fetch_optional(&self.pool)
        .await?;

        if let Some(count) = updated {
            return Ok(count as u64);
        }

        if self.metadata(hash).await?.is_some() {
            Err(CasError::ReferenceUnderflow(hash.to_string()))
        } else {
            Err(CasError::ObjectNotFound(hash.to_string()))
        }
    }

    /// 获取对象引用计数。
    pub async fn get_reference_count(&self, hash: &Hash) -> Result<u64> {
        let count: Option<i64> =
            sqlx::query_scalar("SELECT ref_count FROM cas_objects WHERE hash = ?1")
                .bind(hash.to_string())
                .fetch_optional(&self.pool)
                .await?;

        count
            .map(|count| count as u64)
            .ok_or_else(|| CasError::ObjectNotFound(hash.to_string()))
    }

    /// 获取对象元数据。
    pub async fn metadata(&self, hash: &Hash) -> Result<Option<ObjectMetadata>> {
        let row: Option<(String, i64, i64, i64, i64, Option<i64>)> = sqlx::query_as(
            r#"
            SELECT hash, size, ref_count, created_at, updated_at, zero_ref_at
            FROM cas_objects
            WHERE hash = ?1
            "#,
        )
        .bind(hash.to_string())
        .fetch_optional(&self.pool)
        .await?;

        row.map(
            |(hash, size, ref_count, created_at, updated_at, zero_ref_at)| {
                Ok(ObjectMetadata {
                    hash: hash.parse()?,
                    size: size as u64,
                    ref_count: ref_count as u64,
                    created_at,
                    updated_at,
                    zero_ref_at,
                })
            },
        )
        .transpose()
    }

    /// 列出引用计数为 0 的对象。
    pub async fn zero_ref_objects(&self) -> Result<Vec<Hash>> {
        let hashes: Vec<String> =
            sqlx::query_scalar("SELECT hash FROM cas_objects WHERE ref_count = 0")
                .fetch_all(&self.pool)
                .await?;

        hashes
            .into_iter()
            .map(|hash| hash.parse())
            .collect::<Result<Vec<_>>>()
    }

    /// 删除对象元数据。
    pub async fn delete_metadata(&self, hash: &Hash) -> Result<bool> {
        let result = sqlx::query("DELETE FROM cas_objects WHERE hash = ?1")
            .bind(hash.to_string())
            .execute(&self.pool)
            .await?;
        Ok(result.rows_affected() > 0)
    }

    /// 原子地删除引用计数为 0 的对象元数据。
    ///
    /// `WHERE ref_count = 0` 是 GC 的原子闸门：若对象被并发 `store_reference` /
    /// `add_reference` 抢先递增，本语句 `rows_affected = 0` 返回 `false`，调用方据此
    /// 跳过删除对象文件，绝不删掉仍被引用的对象。
    pub async fn delete_zero_ref_metadata(&self, hash: &Hash) -> Result<bool> {
        let result = sqlx::query("DELETE FROM cas_objects WHERE hash = ?1 AND ref_count = 0")
            .bind(hash.to_string())
            .execute(&self.pool)
            .await?;
        Ok(result.rows_affected() > 0)
    }
}

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

    #[tokio::test]
    async fn reference_count_lifecycle() {
        let temp_dir = tempfile::tempdir().unwrap();
        let counter = SqliteRefCounter::new(temp_dir.path().join("metadata.sqlite"))
            .await
            .unwrap();
        let hash = compute_hash(b"tracked object");

        counter.ensure_object(&hash, 14).await.unwrap();
        assert_eq!(counter.add_reference(&hash).await.unwrap(), 1);
        assert_eq!(counter.add_reference(&hash).await.unwrap(), 2);
        assert_eq!(counter.remove_reference(&hash).await.unwrap(), 1);
        assert_eq!(counter.remove_reference(&hash).await.unwrap(), 0);
        assert_eq!(counter.get_reference_count(&hash).await.unwrap(), 0);
    }

    #[tokio::test]
    async fn zero_reference_objects_are_listed() {
        let temp_dir = tempfile::tempdir().unwrap();
        let counter = SqliteRefCounter::new(temp_dir.path().join("metadata.sqlite"))
            .await
            .unwrap();
        let hash = compute_hash(b"candidate");

        counter.ensure_object(&hash, 9).await.unwrap();
        counter.add_reference(&hash).await.unwrap();
        counter.remove_reference(&hash).await.unwrap();

        assert_eq!(counter.zero_ref_objects().await.unwrap(), vec![hash]);
    }

    #[tokio::test]
    async fn removing_zero_reference_returns_underflow() {
        let temp_dir = tempfile::tempdir().unwrap();
        let counter = SqliteRefCounter::new(temp_dir.path().join("metadata.sqlite"))
            .await
            .unwrap();
        let hash = compute_hash(b"underflow");

        counter.ensure_object(&hash, 9).await.unwrap();

        let error = counter.remove_reference(&hash).await.unwrap_err();
        assert!(matches!(error, CasError::ReferenceUnderflow(_)));
    }

    #[tokio::test]
    async fn store_reference_creates_at_one_and_increments() {
        let temp_dir = tempfile::tempdir().unwrap();
        let counter = SqliteRefCounter::new(temp_dir.path().join("metadata.sqlite"))
            .await
            .unwrap();
        let hash = compute_hash(b"atomic");

        // 新对象直接建行为 1，不经过 ref_count=0 阶段。
        assert_eq!(counter.store_reference(&hash, 6).await.unwrap(), 1);
        assert_eq!(counter.store_reference(&hash, 6).await.unwrap(), 2);
    }

    #[tokio::test]
    async fn delete_zero_ref_metadata_guards_referenced_objects() {
        let temp_dir = tempfile::tempdir().unwrap();
        let counter = SqliteRefCounter::new(temp_dir.path().join("metadata.sqlite"))
            .await
            .unwrap();
        let hash = compute_hash(b"guarded");

        counter.store_reference(&hash, 7).await.unwrap();
        // ref_count>0：守卫删除返回 false 且元数据保留。
        assert!(!counter.delete_zero_ref_metadata(&hash).await.unwrap());
        assert!(counter.metadata(&hash).await.unwrap().is_some());
        // 归零后可原子删除。
        counter.remove_reference(&hash).await.unwrap();
        assert!(counter.delete_zero_ref_metadata(&hash).await.unwrap());
        assert!(counter.metadata(&hash).await.unwrap().is_none());
    }
}