237899745/ystp
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(storage): make queue and object ownership durable

Browse Source
This commit is contained in:
237899745
2026-07-26 10:54:13 +08:00
parent f2d490edce
commit 408e09cda8
14 changed files with 2855 additions and 608 deletions
Download Patch File Download Diff File Expand all files Collapse all files

368
src/services/idempotency.rs
View File

@@ -6,6 +6,8 @@ use serde_json::Value as JsonValue;
use sqlx::FromRow;
use uuid::Uuid;
const OPERATION_LEASE_MINUTES: i64 = 30;
#[derive(Debug, Clone, Copy)]
pub enum Scope {
User(Uuid),
@@ -14,7 +16,7 @@ pub enum Scope {
#[derive(Debug)]
pub enum BeginResult {
Acquired,
Acquired { owner: Uuid },
Replay { response_body: JsonValue },
InProgress,
}
@@ -54,6 +56,8 @@ pub async fn begin(
let now = Utc::now();
let expires_at = now + Duration::hours(ttl_hours.max(1));
let owner = Uuid::new_v4();
let lease_until = now + Duration::minutes(OPERATION_LEASE_MINUTES);
cleanup_expired_for_key(state, scope, idempotency_key, now).await?;
@@ -64,11 +68,11 @@ pub async fn begin(
INSERT INTO idempotency_keys (
user_id, idempotency_key, request_hash,
response_status, response_body,
expires_at
expires_at, lease_owner, lease_until
) VALUES (
$1, $2, $3,
0, NULL,
$4
$4, $5, $6
)
ON CONFLICT DO NOTHING
"#,
@@ -77,6 +81,8 @@ pub async fn begin(
.bind(idempotency_key)
.bind(request_hash)
.bind(expires_at)
.bind(owner)
.bind(lease_until)
.execute(&state.db)
.await
}
@@ -86,11 +92,11 @@ pub async fn begin(
INSERT INTO idempotency_keys (
api_key_id, idempotency_key, request_hash,
response_status, response_body,
expires_at
expires_at, lease_owner, lease_until
) VALUES (
$1, $2, $3,
0, NULL,
$4
$4, $5, $6
)
ON CONFLICT DO NOTHING
"#,
@@ -99,6 +105,8 @@ pub async fn begin(
.bind(idempotency_key)
.bind(request_hash)
.bind(expires_at)
.bind(owner)
.bind(lease_until)
.execute(&state.db)
.await
}
@@ -106,12 +114,15 @@ pub async fn begin(
.map_err(|err| AppError::new(ErrorCode::Internal, "写入幂等记录失败").with_source(err))?;
if inserted.rows_affected() > 0 {
return Ok(BeginResult::Acquired);
return Ok(BeginResult::Acquired { owner });
}
let row = get_row(state, scope, idempotency_key, now).await?;
let Some(row) = row else {
return Ok(BeginResult::Acquired);
return Err(AppError::new(
ErrorCode::StorageUnavailable,
"幂等记录状态已变化，请重试",
));
};
if row.request_hash != request_hash {
@@ -122,6 +133,19 @@ pub async fn begin(
}
if row.response_status == 0 || row.response_body.is_none() {
if take_over_stale_operation(
state,
scope,
idempotency_key,
request_hash,
owner,
lease_until,
now,
)
.await?
{
return Ok(BeginResult::Acquired { owner });
}
return Ok(BeginResult::InProgress);
}
@@ -130,6 +154,152 @@ pub async fn begin(
})
}
async fn take_over_stale_operation(
state: &AppState,
scope: Scope,
idempotency_key: &str,
request_hash: &str,
owner: Uuid,
lease_until: DateTime<Utc>,
now: DateTime<Utc>,
) -> Result<bool, AppError> {
let updated = match scope {
Scope::User(user_id) => {
sqlx::query(
r#"
UPDATE idempotency_keys
SET lease_owner = $4, lease_until = $5
WHERE user_id = $1 AND idempotency_key = $2
AND request_hash = $3 AND response_status = 0
AND (lease_until IS NULL OR lease_until <= $6)
"#,
)
.bind(user_id)
.bind(idempotency_key)
.bind(request_hash)
.bind(owner)
.bind(lease_until)
.bind(now)
.execute(&state.db)
.await
}
Scope::ApiKey(api_key_id) => {
sqlx::query(
r#"
UPDATE idempotency_keys
SET lease_owner = $4, lease_until = $5
WHERE api_key_id = $1 AND idempotency_key = $2
AND request_hash = $3 AND response_status = 0
AND (lease_until IS NULL OR lease_until <= $6)
"#,
)
.bind(api_key_id)
.bind(idempotency_key)
.bind(request_hash)
.bind(owner)
.bind(lease_until)
.bind(now)
.execute(&state.db)
.await
}
}
.map_err(|err| AppError::new(ErrorCode::Internal, "接管过期幂等操作失败").with_source(err))?;
Ok(updated.rows_affected() == 1)
}
pub struct LeaseHeartbeat(Option<tokio::sync::oneshot::Sender<()>>);
impl Drop for LeaseHeartbeat {
fn drop(&mut self) {
if let Some(stop) = self.0.take() {
let _ = stop.send(());
}
}
}
pub fn start_lease_heartbeat(
state: AppState,
scope: Scope,
idempotency_key: String,
request_hash: String,
owner: Uuid,
) -> LeaseHeartbeat {
let (stop_tx, mut stop_rx) = tokio::sync::oneshot::channel();
tokio::spawn(async move {
let mut interval = tokio::time::interval(std::time::Duration::from_secs(60));
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
interval.tick().await;
loop {
tokio::select! {
_ = &mut stop_rx => break,
_ = interval.tick() => {
match renew_lease(
&state,
scope,
&idempotency_key,
&request_hash,
owner,
).await {
Ok(true) => {}
Ok(false) => break,
Err(err) => tracing::warn!(error = %err, "failed to renew idempotency operation lease"),
}
}
}
}
});
LeaseHeartbeat(Some(stop_tx))
}
async fn renew_lease(
state: &AppState,
scope: Scope,
idempotency_key: &str,
request_hash: &str,
owner: Uuid,
) -> Result<bool, AppError> {
let updated = match scope {
Scope::User(user_id) => {
sqlx::query(
r#"
UPDATE idempotency_keys
SET lease_until = NOW() + ($5 * INTERVAL '1 minute')
WHERE user_id = $1 AND idempotency_key = $2
AND request_hash = $3 AND lease_owner = $4
AND response_status = 0
"#,
)
.bind(user_id)
.bind(idempotency_key)
.bind(request_hash)
.bind(owner)
.bind(OPERATION_LEASE_MINUTES)
.execute(&state.db)
.await
}
Scope::ApiKey(api_key_id) => {
sqlx::query(
r#"
UPDATE idempotency_keys
SET lease_until = NOW() + ($5 * INTERVAL '1 minute')
WHERE api_key_id = $1 AND idempotency_key = $2
AND request_hash = $3 AND lease_owner = $4
AND response_status = 0
"#,
)
.bind(api_key_id)
.bind(idempotency_key)
.bind(request_hash)
.bind(owner)
.bind(OPERATION_LEASE_MINUTES)
.execute(&state.db)
.await
}
}
.map_err(|err| AppError::new(ErrorCode::Internal, "续租幂等操作失败").with_source(err))?;
Ok(updated.rows_affected() == 1)
}
pub async fn wait_for_replay(
state: &AppState,
scope: Scope,
@@ -165,11 +335,12 @@ pub async fn wait_for_replay(
}
}
pub async fn complete(
state: &AppState,
pub async fn complete_in_tx(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
scope: Scope,
idempotency_key: &str,
request_hash: &str,
owner: Uuid,
response_status: i32,
response_body: JsonValue,
) -> Result<(), AppError> {
@@ -179,10 +350,13 @@ pub async fn complete(
r#"
UPDATE idempotency_keys
SET response_status = $4,
response_body = $5
response_body = $5,
lease_owner = NULL,
lease_until = NULL
WHERE user_id = $1
AND idempotency_key = $2
AND request_hash = $3
AND lease_owner = $6
AND response_status = 0
"#,
)
@@ -191,7 +365,8 @@ pub async fn complete(
.bind(request_hash)
.bind(response_status)
.bind(response_body)
.execute(&state.db)
.bind(owner)
.execute(&mut **tx)
.await
}
Scope::ApiKey(api_key_id) => {
@@ -199,10 +374,13 @@ pub async fn complete(
r#"
UPDATE idempotency_keys
SET response_status = $4,
response_body = $5
response_body = $5,
lease_owner = NULL,
lease_until = NULL
WHERE api_key_id = $1
AND idempotency_key = $2
AND request_hash = $3
AND lease_owner = $6
AND response_status = 0
"#,
)
@@ -211,16 +389,19 @@ pub async fn complete(
.bind(request_hash)
.bind(response_status)
.bind(response_body)
.execute(&state.db)
.bind(owner)
.execute(&mut **tx)
.await
}
}
.map_err(|err| AppError::new(ErrorCode::Internal, "写入幂等结果失败").with_source(err))?;
.map_err(|err| AppError::new(ErrorCode::Internal, "事务内写入幂等结果失败").with_source(err))?;
if updated.rows_affected() == 0 {
tracing::warn!("idempotency record not updated (maybe already completed?)");
if updated.rows_affected() != 1 {
return Err(AppError::new(
ErrorCode::IdempotencyConflict,
"幂等请求所有权已变化，请重试",
));
}
Ok(())
}
@@ -229,25 +410,28 @@ pub async fn abort(
scope: Scope,
idempotency_key: &str,
request_hash: &str,
owner: Uuid,
) -> Result<(), AppError> {
match scope {
Scope::User(user_id) => {
let _ = sqlx::query(
"DELETE FROM idempotency_keys WHERE user_id = $1 AND idempotency_key = $2 AND request_hash = $3 AND response_status = 0",
"DELETE FROM idempotency_keys WHERE user_id = $1 AND idempotency_key = $2 AND request_hash = $3 AND lease_owner = $4 AND response_status = 0",
)
.bind(user_id)
.bind(idempotency_key)
.bind(request_hash)
.bind(owner)
.execute(&state.db)
.await;
}
Scope::ApiKey(api_key_id) => {
let _ = sqlx::query(
"DELETE FROM idempotency_keys WHERE api_key_id = $1 AND idempotency_key = $2 AND request_hash = $3 AND response_status = 0",
"DELETE FROM idempotency_keys WHERE api_key_id = $1 AND idempotency_key = $2 AND request_hash = $3 AND lease_owner = $4 AND response_status = 0",
)
.bind(api_key_id)
.bind(idempotency_key)
.bind(request_hash)
.bind(owner)
.execute(&state.db)
.await;
}
@@ -334,3 +518,149 @@ async fn get_row(
Ok(row)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::Config;
use crate::services::mail::Mailer;
use crate::services::settings::RuntimePolicyCache;
use crate::services::storage::StorageCache;
use sqlx::postgres::PgPoolOptions;
use std::sync::Arc;
use tokio::sync::Semaphore;
async fn test_state(database_url: String, redis_url: String) -> AppState {
let mut config = Config::from_env().expect("load idempotency test config");
config.database_url = database_url.clone();
config.redis_url = redis_url;
config.mail_enabled = false;
config.mail_log_links_when_disabled = false;
let db = PgPoolOptions::new()
.max_connections(8)
.connect(&database_url)
.await
.expect("connect idempotency test database");
sqlx::migrate!().run(&db).await.expect("run migrations");
let redis = redis::Client::open(config.redis_url.clone())
.expect("create idempotency test Redis client")
.get_connection_manager()
.await
.expect("connect idempotency test Redis");
AppState {
mailer: Arc::new(Mailer::new(&config).expect("create disabled test mailer")),
image_processing_semaphore: Arc::new(Semaphore::new(1)),
zip_build_semaphore: Arc::new(Semaphore::new(1)),
runtime_policy_cache: RuntimePolicyCache::new(),
storage_cache: StorageCache::new(),
config,
db,
redis,
}
}
#[tokio::test]
#[ignore = "requires IMAGEFORGE_TEST_DATABASE_URL and IMAGEFORGE_TEST_REDIS_URL"]
async fn stale_operation_is_fenced_and_replay_remains_atomic() {
let database_url = std::env::var("IMAGEFORGE_TEST_DATABASE_URL")
.expect("IMAGEFORGE_TEST_DATABASE_URL is required");
let redis_url = std::env::var("IMAGEFORGE_TEST_REDIS_URL")
.expect("IMAGEFORGE_TEST_REDIS_URL is required");
let state = test_state(database_url, redis_url).await;
let marker = Uuid::new_v4().simple().to_string();
let user_id: Uuid = sqlx::query_scalar(
r#"
INSERT INTO users (email, username, password_hash, email_verified_at)
VALUES ($1, $2, 'test', NOW())
RETURNING id
"#,
)
.bind(format!("idem-{marker}@example.test"))
.bind(format!("idem-{marker}"))
.fetch_one(&state.db)
.await
.expect("insert idempotency test user");
let scope = Scope::User(user_id);
let key = format!("idem-{marker}");
let request_hash = "a".repeat(64);
let owner_one = match begin(&state, scope, &key, &request_hash, 24)
.await
.expect("acquire first operation")
{
BeginResult::Acquired { owner } => owner,
other => panic!("unexpected first begin result: {other:?}"),
};
assert!(matches!(
begin(&state, scope, &key, &request_hash, 24)
.await
.expect("probe live operation"),
BeginResult::InProgress
));
sqlx::query(
"UPDATE idempotency_keys SET lease_until = NOW() - INTERVAL '1 second' WHERE user_id = $1 AND idempotency_key = $2",
)
.bind(user_id)
.bind(&key)
.execute(&state.db)
.await
.expect("expire first operation lease");
let owner_two = match begin(&state, scope, &key, &request_hash, 24)
.await
.expect("take over stale operation")
{
BeginResult::Acquired { owner } => owner,
other => panic!("unexpected takeover result: {other:?}"),
};
assert_ne!(owner_one, owner_two);
let mut stale_tx = state.db.begin().await.expect("begin stale completion tx");
let stale_error = complete_in_tx(
&mut stale_tx,
scope,
&key,
&request_hash,
owner_one,
200,
serde_json::json!({"owner": "stale"}),
)
.await
.expect_err("stale operation completed after takeover");
assert_eq!(stale_error.code, ErrorCode::IdempotencyConflict);
stale_tx
.rollback()
.await
.expect("rollback stale completion");
let expected = serde_json::json!({"owner": "current"});
let mut current_tx = state.db.begin().await.expect("begin current completion tx");
complete_in_tx(
&mut current_tx,
scope,
&key,
&request_hash,
owner_two,
200,
expected.clone(),
)
.await
.expect("complete current operation");
current_tx
.commit()
.await
.expect("commit current completion");
match begin(&state, scope, &key, &request_hash, 24)
.await
.expect("replay completed operation")
{
BeginResult::Replay { response_body } => assert_eq!(response_body, expected),
other => panic!("unexpected replay result: {other:?}"),
}
sqlx::query("DELETE FROM users WHERE id = $1")
.bind(user_id)
.execute(&state.db)
.await
.expect("clean idempotency test user");
}
}