use crate::error::{AppError, ErrorCode}; use crate::state::AppState; use chrono::{DateTime, Duration, Utc}; use serde_json::Value as JsonValue; use sha2::{Digest, Sha256}; use sqlx::FromRow; use uuid::Uuid; #[derive(Debug, Clone, Copy)] pub enum Scope { User(Uuid), ApiKey(Uuid), } #[derive(Debug)] pub enum BeginResult { Acquired, Replay { response_body: JsonValue }, InProgress, } #[derive(Debug, FromRow)] struct IdemRow { request_hash: String, response_status: i32, response_body: Option, } pub fn sha256_hex(parts: &[&[u8]]) -> String { let mut hasher = Sha256::new(); for p in parts { hasher.update(p); hasher.update([0u8]); // separator } hex::encode(hasher.finalize()) } pub async fn begin( state: &AppState, scope: Scope, idempotency_key: &str, request_hash: &str, ttl_hours: i64, ) -> Result { 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)); 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 ) VALUES ( $1, $2, $3, 0, NULL, $4 ) ON CONFLICT DO NOTHING "#, ) .bind(user_id) .bind(idempotency_key) .bind(request_hash) .bind(expires_at) .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 ) VALUES ( $1, $2, $3, 0, NULL, $4 ) ON CONFLICT DO NOTHING "#, ) .bind(api_key_id) .bind(idempotency_key) .bind(request_hash) .bind(expires_at) .execute(&state.db) .await } } .map_err(|err| AppError::new(ErrorCode::Internal, "写入幂等记录失败").with_source(err))?; if inserted.rows_affected() > 0 { return Ok(BeginResult::Acquired); } let row = get_row(state, scope, idempotency_key, now).await?; let Some(row) = row else { return Ok(BeginResult::Acquired); }; if row.request_hash != request_hash { return Err(AppError::new( ErrorCode::IdempotencyConflict, "同一个 Idempotency-Key 的请求参数不一致", )); } if row.response_status == 0 || row.response_body.is_none() { return Ok(BeginResult::InProgress); } Ok(BeginResult::Replay { response_body: row.response_body.unwrap_or(JsonValue::Null), }) } pub async fn wait_for_replay( state: &AppState, scope: Scope, idempotency_key: &str, request_hash: &str, max_wait_ms: u64, ) -> Result, 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( state: &AppState, scope: Scope, idempotency_key: &str, request_hash: &str, 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 WHERE user_id = $1 AND idempotency_key = $2 AND request_hash = $3 AND response_status = 0 "#, ) .bind(user_id) .bind(idempotency_key) .bind(request_hash) .bind(response_status) .bind(response_body) .execute(&state.db) .await } Scope::ApiKey(api_key_id) => { sqlx::query( r#" UPDATE idempotency_keys SET response_status = $4, response_body = $5 WHERE api_key_id = $1 AND idempotency_key = $2 AND request_hash = $3 AND response_status = 0 "#, ) .bind(api_key_id) .bind(idempotency_key) .bind(request_hash) .bind(response_status) .bind(response_body) .execute(&state.db) .await } } .map_err(|err| AppError::new(ErrorCode::Internal, "写入幂等结果失败").with_source(err))?; if updated.rows_affected() == 0 { tracing::warn!("idempotency record not updated (maybe already completed?)"); } Ok(()) } pub async fn abort( state: &AppState, scope: Scope, idempotency_key: &str, request_hash: &str, ) -> 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", ) .bind(user_id) .bind(idempotency_key) .bind(request_hash) .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", ) .bind(api_key_id) .bind(idempotency_key) .bind(request_hash) .execute(&state.db) .await; } } Ok(()) } async fn cleanup_expired_for_key( state: &AppState, scope: Scope, idempotency_key: &str, now: DateTime, ) -> 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, ) -> Result, 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) }