667 lines
20 KiB
Rust
667 lines
20 KiB
Rust
use crate::error::{AppError, ErrorCode};
|
|
use crate::state::AppState;
|
|
|
|
use chrono::{DateTime, Duration, Utc};
|
|
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),
|
|
ApiKey(Uuid),
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub enum BeginResult {
|
|
Acquired { owner: Uuid },
|
|
Replay { response_body: JsonValue },
|
|
InProgress,
|
|
}
|
|
|
|
#[derive(Debug, FromRow)]
|
|
struct IdemRow {
|
|
request_hash: String,
|
|
response_status: i32,
|
|
response_body: Option<JsonValue>,
|
|
}
|
|
|
|
pub async fn begin(
|
|
state: &AppState,
|
|
scope: Scope,
|
|
idempotency_key: &str,
|
|
request_hash: &str,
|
|
ttl_hours: i64,
|
|
) -> Result<BeginResult, AppError> {
|
|
if idempotency_key.trim().is_empty() {
|
|
return Err(AppError::new(
|
|
ErrorCode::InvalidRequest,
|
|
"Idempotency-Key 不能为空",
|
|
));
|
|
}
|
|
if idempotency_key.len() > 128 {
|
|
return Err(AppError::new(
|
|
ErrorCode::InvalidRequest,
|
|
"Idempotency-Key 过长",
|
|
));
|
|
}
|
|
if request_hash.len() != 64 {
|
|
return Err(AppError::new(
|
|
ErrorCode::InvalidRequest,
|
|
"request_hash 不合法",
|
|
));
|
|
}
|
|
|
|
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?;
|
|
|
|
let inserted = match scope {
|
|
Scope::User(user_id) => {
|
|
sqlx::query(
|
|
r#"
|
|
INSERT INTO idempotency_keys (
|
|
user_id, idempotency_key, request_hash,
|
|
response_status, response_body,
|
|
expires_at, lease_owner, lease_until
|
|
) VALUES (
|
|
$1, $2, $3,
|
|
0, NULL,
|
|
$4, $5, $6
|
|
)
|
|
ON CONFLICT DO NOTHING
|
|
"#,
|
|
)
|
|
.bind(user_id)
|
|
.bind(idempotency_key)
|
|
.bind(request_hash)
|
|
.bind(expires_at)
|
|
.bind(owner)
|
|
.bind(lease_until)
|
|
.execute(&state.db)
|
|
.await
|
|
}
|
|
Scope::ApiKey(api_key_id) => {
|
|
sqlx::query(
|
|
r#"
|
|
INSERT INTO idempotency_keys (
|
|
api_key_id, idempotency_key, request_hash,
|
|
response_status, response_body,
|
|
expires_at, lease_owner, lease_until
|
|
) VALUES (
|
|
$1, $2, $3,
|
|
0, NULL,
|
|
$4, $5, $6
|
|
)
|
|
ON CONFLICT DO NOTHING
|
|
"#,
|
|
)
|
|
.bind(api_key_id)
|
|
.bind(idempotency_key)
|
|
.bind(request_hash)
|
|
.bind(expires_at)
|
|
.bind(owner)
|
|
.bind(lease_until)
|
|
.execute(&state.db)
|
|
.await
|
|
}
|
|
}
|
|
.map_err(|err| AppError::new(ErrorCode::Internal, "写入幂等记录失败").with_source(err))?;
|
|
|
|
if inserted.rows_affected() > 0 {
|
|
return Ok(BeginResult::Acquired { owner });
|
|
}
|
|
|
|
let row = get_row(state, scope, idempotency_key, now).await?;
|
|
let Some(row) = row else {
|
|
return Err(AppError::new(
|
|
ErrorCode::StorageUnavailable,
|
|
"幂等记录状态已变化，请重试",
|
|
));
|
|
};
|
|
|
|
if row.request_hash != request_hash {
|
|
return Err(AppError::new(
|
|
ErrorCode::IdempotencyConflict,
|
|
"同一个 Idempotency-Key 的请求参数不一致",
|
|
));
|
|
}
|
|
|
|
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);
|
|
}
|
|
|
|
Ok(BeginResult::Replay {
|
|
response_body: row.response_body.unwrap_or(JsonValue::Null),
|
|
})
|
|
}
|
|
|
|
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,
|
|
idempotency_key: &str,
|
|
request_hash: &str,
|
|
max_wait_ms: u64,
|
|
) -> Result<Option<(i32, JsonValue)>, AppError> {
|
|
let started = tokio::time::Instant::now();
|
|
|
|
loop {
|
|
let row = get_row(state, scope, idempotency_key, Utc::now()).await?;
|
|
let Some(row) = row else { return Ok(None) };
|
|
|
|
if row.request_hash != request_hash {
|
|
return Err(AppError::new(
|
|
ErrorCode::IdempotencyConflict,
|
|
"同一个 Idempotency-Key 的请求参数不一致",
|
|
));
|
|
}
|
|
|
|
if row.response_status != 0 {
|
|
return Ok(Some((
|
|
row.response_status,
|
|
row.response_body.unwrap_or(JsonValue::Null),
|
|
)));
|
|
}
|
|
|
|
if started.elapsed().as_millis() as u64 >= max_wait_ms {
|
|
return Ok(None);
|
|
}
|
|
|
|
tokio::time::sleep(std::time::Duration::from_millis(200)).await;
|
|
}
|
|
}
|
|
|
|
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> {
|
|
let updated = match scope {
|
|
Scope::User(user_id) => {
|
|
sqlx::query(
|
|
r#"
|
|
UPDATE idempotency_keys
|
|
SET response_status = $4,
|
|
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
|
|
"#,
|
|
)
|
|
.bind(user_id)
|
|
.bind(idempotency_key)
|
|
.bind(request_hash)
|
|
.bind(response_status)
|
|
.bind(response_body)
|
|
.bind(owner)
|
|
.execute(&mut **tx)
|
|
.await
|
|
}
|
|
Scope::ApiKey(api_key_id) => {
|
|
sqlx::query(
|
|
r#"
|
|
UPDATE idempotency_keys
|
|
SET response_status = $4,
|
|
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
|
|
"#,
|
|
)
|
|
.bind(api_key_id)
|
|
.bind(idempotency_key)
|
|
.bind(request_hash)
|
|
.bind(response_status)
|
|
.bind(response_body)
|
|
.bind(owner)
|
|
.execute(&mut **tx)
|
|
.await
|
|
}
|
|
}
|
|
.map_err(|err| AppError::new(ErrorCode::Internal, "事务内写入幂等结果失败").with_source(err))?;
|
|
|
|
if updated.rows_affected() != 1 {
|
|
return Err(AppError::new(
|
|
ErrorCode::IdempotencyConflict,
|
|
"幂等请求所有权已变化，请重试",
|
|
));
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
pub async fn abort(
|
|
state: &AppState,
|
|
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 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 lease_owner = $4 AND response_status = 0",
|
|
)
|
|
.bind(api_key_id)
|
|
.bind(idempotency_key)
|
|
.bind(request_hash)
|
|
.bind(owner)
|
|
.execute(&state.db)
|
|
.await;
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
async fn cleanup_expired_for_key(
|
|
state: &AppState,
|
|
scope: Scope,
|
|
idempotency_key: &str,
|
|
now: DateTime<Utc>,
|
|
) -> Result<(), AppError> {
|
|
match scope {
|
|
Scope::User(user_id) => {
|
|
let _ = sqlx::query(
|
|
"DELETE FROM idempotency_keys WHERE user_id = $1 AND idempotency_key = $2 AND expires_at < $3",
|
|
)
|
|
.bind(user_id)
|
|
.bind(idempotency_key)
|
|
.bind(now)
|
|
.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 expires_at < $3",
|
|
)
|
|
.bind(api_key_id)
|
|
.bind(idempotency_key)
|
|
.bind(now)
|
|
.execute(&state.db)
|
|
.await;
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
async fn get_row(
|
|
state: &AppState,
|
|
scope: Scope,
|
|
idempotency_key: &str,
|
|
now: DateTime<Utc>,
|
|
) -> Result<Option<IdemRow>, AppError> {
|
|
let row = match scope {
|
|
Scope::User(user_id) => {
|
|
sqlx::query_as::<_, IdemRow>(
|
|
r#"
|
|
SELECT request_hash, response_status, response_body
|
|
FROM idempotency_keys
|
|
WHERE user_id = $1
|
|
AND idempotency_key = $2
|
|
AND expires_at > $3
|
|
ORDER BY created_at DESC
|
|
LIMIT 1
|
|
"#,
|
|
)
|
|
.bind(user_id)
|
|
.bind(idempotency_key)
|
|
.bind(now)
|
|
.fetch_optional(&state.db)
|
|
.await
|
|
}
|
|
Scope::ApiKey(api_key_id) => {
|
|
sqlx::query_as::<_, IdemRow>(
|
|
r#"
|
|
SELECT request_hash, response_status, response_body
|
|
FROM idempotency_keys
|
|
WHERE api_key_id = $1
|
|
AND idempotency_key = $2
|
|
AND expires_at > $3
|
|
ORDER BY created_at DESC
|
|
LIMIT 1
|
|
"#,
|
|
)
|
|
.bind(api_key_id)
|
|
.bind(idempotency_key)
|
|
.bind(now)
|
|
.fetch_optional(&state.db)
|
|
.await
|
|
}
|
|
}
|
|
.map_err(|err| AppError::new(ErrorCode::Internal, "查询幂等记录失败").with_source(err))?;
|
|
|
|
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");
|
|
}
|
|
}
|