Files
BlueArchiveToolkit/crates/bat-cas-engine/src/refcount.rs
T
nyaKazuhaandClaude Fable 5 1c4debbaa7 fix(cas): 消除 store/gc 之间的引用计数竞态
store() 原来先 ensure_object(建行 ref_count=0)再 add_reference,两条语句
非事务;即使连接池 max_connections=1,两语句间的 await 也会释放连接,让并发
gc() 在 ref_count=0 窗口删掉刚存的对象文件与元数据(静默丢数据),或使 store
返回 ObjectNotFound。

- 新增 store_reference:单条 UPSERT 原子建行为 ref_count=1 或 +1,对象行不再
  出现 ref_count=0 的可见窗口;store() 与 repository add_reference() 均改用它。
- 新增 delete_zero_ref_metadata:gc 先原子执行 DELETE ... WHERE ref_count=0,
  仅当 rows_affected>0 才删对象文件;被并发递增抢先时跳过,绝不删除仍被引用对象。
  元数据先删、文件后删,最坏只留无元数据的孤儿文件(可覆盖,无数据丢失)。

新增单测:store 后不暴露 ref_count=0(gc_candidates 为空)、候选被重新引用后
gc 跳过、store_reference 原子建行/递增、delete_zero_ref_metadata 守卫。

对应 issue #18 维护清单 1-5。

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-16 07:39:07 -07:00

367 lines
12 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
//! CAS 引用计数和对象元数据。
use crate::error::{CasError, Result};
use crate::hash::Hash;
use sqlx::sqlite::{SqliteConnectOptions, SqlitePoolOptions, SqliteQueryResult};
use sqlx::SqlitePool;
use std::path::Path;
use std::str::FromStr;
use std::time::{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);
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());
}
}