Implement compression quota refunds and admin manual subscription
This commit is contained in:
341
src/services/idempotency.rs
Normal file
341
src/services/idempotency.rs
Normal file
@@ -0,0 +1,341 @@
|
||||
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 { expires_at: DateTime<Utc> },
|
||||
Replay { response_status: i32, response_body: JsonValue },
|
||||
InProgress,
|
||||
}
|
||||
|
||||
#[derive(Debug, FromRow)]
|
||||
struct IdemRow {
|
||||
request_hash: String,
|
||||
response_status: i32,
|
||||
response_body: Option<JsonValue>,
|
||||
expires_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
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<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));
|
||||
|
||||
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 { expires_at });
|
||||
}
|
||||
|
||||
let row = get_row(state, scope, idempotency_key, now).await?;
|
||||
let Some(row) = row else {
|
||||
return Ok(BeginResult::Acquired { expires_at });
|
||||
};
|
||||
|
||||
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_status: row.response_status,
|
||||
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<Option<(i32, JsonValue)>, AppError> {
|
||||
let started = tokio::time::Instant::now();
|
||||
let now = Utc::now();
|
||||
|
||||
loop {
|
||||
let row = get_row(state, scope, idempotency_key, 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<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, expires_at
|
||||
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, expires_at
|
||||
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)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user