mirror of
https://github.com/Yuyi-Oak/BlueArchiveToolkit.git
synced 2026-07-22 09:15:15 +08:00
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>
367 lines
12 KiB
Rust
367 lines
12 KiB
Rust
//! 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());
|
||
}
|
||
}
|