i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-08 02:17:46 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(gateway): stop candidate persistence retry storms

Browse Source
This commit is contained in:
ZheFox
2026-08-17 13:49:12 +08:00
parent 0b17026eab
commit 32944538e9
5 changed files with 1535 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-data/adapters/postgres/src/candidates.rs
+355 -13
View File
@@ -188,13 +188,30 @@ DO UPDATE SET
concurrent_requests = COALESCE(EXCLUDED.concurrent_requests, request_candidates.concurrent_requests),
extra_data = CASE
WHEN request_candidates.extra_data IS NULL THEN EXCLUDED.extra_data
WHEN EXCLUDED.extra_data IS NULL THEN request_candidates.extra_data
WHEN EXCLUDED.extra_data IS NULL THEN regexp_replace(
request_candidates.extra_data::text,
$aether_nul$(?<!\\)((?:\\\\)*)\\u0000$aether_nul$,
$aether_replacement$\1\\ufffd$aether_replacement$,
'g'
)::json
WHEN json_typeof(request_candidates.extra_data) = 'object'
AND json_typeof(EXCLUDED.extra_data) = 'object'
THEN (request_candidates.extra_data::jsonb || EXCLUDED.extra_data::jsonb)::json
THEN (
regexp_replace(
request_candidates.extra_data::text,
$aether_nul$(?<!\\)((?:\\\\)*)\\u0000$aether_nul$,
$aether_replacement$\1\\ufffd$aether_replacement$,
'g'
)::jsonb || EXCLUDED.extra_data::jsonb
)::json
ELSE EXCLUDED.extra_data
END,
required_capabilities = COALESCE(EXCLUDED.required_capabilities, request_candidates.required_capabilities),
required_capabilities = regexp_replace(
COALESCE(EXCLUDED.required_capabilities, request_candidates.required_capabilities)::text,
$aether_nul$(?<!\\)((?:\\\\)*)\\u0000$aether_nul$,
$aether_replacement$\1\\ufffd$aether_replacement$,
'g'
)::json,
created_at = CASE
WHEN request_candidates.created_at <= TO_TIMESTAMP(1)
THEN EXCLUDED.created_at
@@ -303,13 +320,30 @@ DO UPDATE SET
concurrent_requests = COALESCE(EXCLUDED.concurrent_requests, request_candidates.concurrent_requests),
extra_data = CASE
WHEN request_candidates.extra_data IS NULL THEN EXCLUDED.extra_data
WHEN EXCLUDED.extra_data IS NULL THEN request_candidates.extra_data
WHEN EXCLUDED.extra_data IS NULL THEN regexp_replace(
request_candidates.extra_data::text,
$aether_nul$(?<!\\)((?:\\\\)*)\\u0000$aether_nul$,
$aether_replacement$\1\\ufffd$aether_replacement$,
'g'
)::json
WHEN json_typeof(request_candidates.extra_data) = 'object'
AND json_typeof(EXCLUDED.extra_data) = 'object'
THEN (request_candidates.extra_data::jsonb || EXCLUDED.extra_data::jsonb)::json
THEN (
regexp_replace(
request_candidates.extra_data::text,
$aether_nul$(?<!\\)((?:\\\\)*)\\u0000$aether_nul$,
$aether_replacement$\1\\ufffd$aether_replacement$,
'g'
)::jsonb || EXCLUDED.extra_data::jsonb
)::json
ELSE EXCLUDED.extra_data
END,
required_capabilities = COALESCE(EXCLUDED.required_capabilities, request_candidates.required_capabilities),
required_capabilities = regexp_replace(
COALESCE(EXCLUDED.required_capabilities, request_candidates.required_capabilities)::text,
$aether_nul$(?<!\\)((?:\\\\)*)\\u0000$aether_nul$,
$aether_replacement$\1\\ufffd$aether_replacement$,
'g'
)::json,
created_at = CASE
WHEN request_candidates.created_at <= TO_TIMESTAMP(1)
THEN EXCLUDED.created_at
@@ -393,13 +427,30 @@ DO UPDATE SET
concurrent_requests = COALESCE(EXCLUDED.concurrent_requests, request_candidates.concurrent_requests),
extra_data = CASE
WHEN request_candidates.extra_data IS NULL THEN EXCLUDED.extra_data
WHEN EXCLUDED.extra_data IS NULL THEN request_candidates.extra_data
WHEN EXCLUDED.extra_data IS NULL THEN regexp_replace(
request_candidates.extra_data::text,
$aether_nul$(?<!\\)((?:\\\\)*)\\u0000$aether_nul$,
$aether_replacement$\1\\ufffd$aether_replacement$,
'g'
)::json
WHEN json_typeof(request_candidates.extra_data) = 'object'
AND json_typeof(EXCLUDED.extra_data) = 'object'
THEN (request_candidates.extra_data::jsonb || EXCLUDED.extra_data::jsonb)::json
THEN (
regexp_replace(
request_candidates.extra_data::text,
$aether_nul$(?<!\\)((?:\\\\)*)\\u0000$aether_nul$,
$aether_replacement$\1\\ufffd$aether_replacement$,
'g'
)::jsonb || EXCLUDED.extra_data::jsonb
)::json
ELSE EXCLUDED.extra_data
END,
required_capabilities = COALESCE(EXCLUDED.required_capabilities, request_candidates.required_capabilities),
required_capabilities = regexp_replace(
COALESCE(EXCLUDED.required_capabilities, request_candidates.required_capabilities)::text,
$aether_nul$(?<!\\)((?:\\\\)*)\\u0000$aether_nul$,
$aether_replacement$\1\\ufffd$aether_replacement$,
'g'
)::json,
created_at = CASE
WHEN request_candidates.created_at <= TO_TIMESTAMP(1)
THEN EXCLUDED.created_at
@@ -724,8 +775,9 @@ impl SqlxRequestCandidateReadRepository {
pub async fn upsert(
&self,
candidate: UpsertRequestCandidateRecord,
mut candidate: UpsertRequestCandidateRecord,
) -> Result<StoredRequestCandidate, DataLayerError> {
sanitize_request_candidate_for_postgres(&mut candidate);
candidate.validate()?;
self.tx_runner
.run_read_write(|tx| {
@@ -873,7 +925,8 @@ struct BatchUpsertRequestCandidateRow {
impl TryFrom<UpsertRequestCandidateRecord> for BatchUpsertRequestCandidateRow {
type Error = DataLayerError;
fn try_from(candidate: UpsertRequestCandidateRecord) -> Result<Self, Self::Error> {
fn try_from(mut candidate: UpsertRequestCandidateRecord) -> Result<Self, Self::Error> {
sanitize_request_candidate_for_postgres(&mut candidate);
candidate.validate()?;
Ok(Self {
id: if candidate.id.trim().is_empty() {
@@ -1176,13 +1229,95 @@ fn to_i32_u64(value: u64) -> Result<i32, DataLayerError> {
})
}
fn sanitize_request_candidate_for_postgres(candidate: &mut UpsertRequestCandidateRecord) -> usize {
let mut replacements = 0usize;
for value in [
&mut candidate.username,
&mut candidate.api_key_name,
&mut candidate.skip_reason,
&mut candidate.error_type,
&mut candidate.error_message,
] {
if let Some(value) = value.as_mut() {
replacements = replacements.saturating_add(replace_nul_characters(value));
}
}
for value in [
&mut candidate.extra_data,
&mut candidate.required_capabilities,
] {
if let Some(value) = value.as_mut() {
replacements = replacements.saturating_add(sanitize_json_nul_characters(value));
}
}
if replacements > 0 {
tracing::warn!(
event_name = "request_candidate_postgres_nul_sanitized",
log_type = "event",
candidate_index = candidate.candidate_index,
retry_index = candidate.retry_index,
status = ?candidate.status,
replacements,
"postgres request candidate persistence replaced unsupported NUL characters"
);
}
replacements
}
fn sanitize_json_nul_characters(value: &mut serde_json::Value) -> usize {
match value {
serde_json::Value::String(value) => replace_nul_characters(value),
serde_json::Value::Array(values) => values.iter_mut().fold(0usize, |count, value| {
count.saturating_add(sanitize_json_nul_characters(value))
}),
serde_json::Value::Object(values) => {
let mut replacements = 0usize;
let original = std::mem::take(values);
for (mut key, mut value) in original {
replacements = replacements.saturating_add(replace_nul_characters(&mut key));
replacements =
replacements.saturating_add(sanitize_json_nul_characters(&mut value));
if values.contains_key(&key) {
let base = key.clone();
let mut suffix = 1usize;
loop {
let candidate = format!("{base}#{suffix}");
if !values.contains_key(&candidate) {
key = candidate;
break;
}
suffix = suffix.saturating_add(1);
}
}
values.insert(key, value);
}
replacements
}
serde_json::Value::Null | serde_json::Value::Bool(_) | serde_json::Value::Number(_) => 0,
}
}
fn replace_nul_characters(value: &mut String) -> usize {
let replacements = value.matches('\0').count();
if replacements > 0 {
*value = value.replace('\0', "\u{fffd}");
}
replacements
}
#[cfg(test)]
mod tests {
use super::{
SqlxRequestCandidateReadRepository, UPSERT_CONFLICT_INHERIT_IS_CACHED_SQL,
UPSERT_CONFLICT_SQL, UPSERT_SQL,
sanitize_request_candidate_for_postgres, SqlxRequestCandidateReadRepository,
UPSERT_CONFLICT_INHERIT_IS_CACHED_SQL, UPSERT_CONFLICT_SQL, UPSERT_SQL,
};
use crate::error::SqlxResultExt;
use crate::{PostgresPoolConfig, PostgresPoolFactory};
use aether_data_contracts::repository::candidates::{
RequestCandidateStatus, UpsertRequestCandidateRecord,
};
use serde_json::{json, Map, Value};
#[test]
fn upsert_sql_does_not_default_missing_or_epoch_created_at_to_epoch() {
@@ -1220,6 +1355,71 @@ mod tests {
}
}
#[test]
fn postgres_candidate_sanitizer_replaces_nul_in_text_and_nested_json() {
let mut extra_data = Map::new();
extra_data.insert(
"bad\0key".to_string(),
json!({"nested": ["bad\0value", {"literal": "\\u0000"}]}),
);
let mut required_capabilities = Map::new();
required_capabilities.insert("cap\0key".to_string(), Value::String("cap\0value".into()));
let mut candidate = UpsertRequestCandidateRecord {
id: "candidate-1".to_string(),
request_id: "request-1".to_string(),
user_id: None,
api_key_id: None,
username: Some("user\0name".to_string()),
api_key_name: Some("key\0name".to_string()),
candidate_index: 0,
retry_index: 0,
provider_id: None,
endpoint_id: None,
key_id: None,
status: RequestCandidateStatus::Failed,
skip_reason: Some("skip\0reason".to_string()),
is_cached: None,
status_code: Some(500),
error_type: Some("upstream\0error".to_string()),
error_message: Some("bad\0message".to_string()),
latency_ms: None,
concurrent_requests: None,
extra_data: Some(Value::Object(extra_data)),
required_capabilities: Some(Value::Object(required_capabilities)),
created_at_unix_ms: Some(1),
started_at_unix_ms: None,
finished_at_unix_ms: Some(2),
};
assert_eq!(sanitize_request_candidate_for_postgres(&mut candidate), 9);
assert_eq!(candidate.username.as_deref(), Some("user�name"));
assert_eq!(candidate.api_key_name.as_deref(), Some("key�name"));
assert_eq!(candidate.skip_reason.as_deref(), Some("skip�reason"));
assert_eq!(candidate.error_type.as_deref(), Some("upstream�error"));
assert_eq!(candidate.error_message.as_deref(), Some("bad�message"));
assert_eq!(
candidate.extra_data,
Some(json!({"bad�key": {"nested": ["bad�value", {"literal": "\\u0000"}]}}))
);
assert_eq!(
candidate.required_capabilities,
Some(json!({"cap�key": "cap�value"}))
);
}
#[test]
fn every_postgres_candidate_conflict_path_repairs_legacy_json_nul_escapes() {
for sql in [
UPSERT_SQL,
UPSERT_CONFLICT_SQL,
UPSERT_CONFLICT_INHERIT_IS_CACHED_SQL,
] {
assert!(sql.contains("regexp_replace("));
assert!(sql.contains(r"(?<!\\)((?:\\\\)*)\\u0000"));
assert!(sql.contains(r"\1\\ufffd"));
}
}
#[tokio::test]
async fn repository_constructs_from_lazy_pool() {
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
@@ -1239,4 +1439,146 @@ mod tests {
let _ = repository.pool();
let _ = repository.transaction_runner();
}
#[tokio::test]
#[ignore = "requires AETHER_TEST_DATABASE_URL and PostgreSQL migrations"]
async fn live_postgres_candidate_nul_is_sanitized_and_legacy_json_is_repaired() {
let database_url = std::env::var("AETHER_TEST_DATABASE_URL")
.expect("AETHER_TEST_DATABASE_URL must point at the test database");
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
database_url,
min_connections: 1,
max_connections: 2,
acquire_timeout_ms: 10_000,
idle_timeout_ms: 30_000,
max_lifetime_ms: 60_000,
statement_cache_capacity: 64,
require_ssl: false,
})
.expect("factory should build");
let repository = SqlxRequestCandidateReadRepository::new(
factory.connect_lazy().expect("lazy pool should build"),
);
crate::run_migrations(repository.pool())
.await
.expect("test database migrations should succeed");
let mapped_error = sqlx::query("SELECT $1::jsonb")
.bind(json!("bad\0value"))
.execute(repository.pool())
.await
.map_postgres_err()
.expect_err("PostgreSQL jsonb should reject a NUL string");
assert!(mapped_error.to_string().contains("SQLSTATE 22P05"));
let suffix = uuid::Uuid::new_v4().simple().to_string();
let single_request_id = format!("candidate-nul-single-{suffix}");
let batch_request_id = format!("candidate-nul-batch-{suffix}");
let healthy_request_id = format!("candidate-nul-healthy-{suffix}");
let legacy_extra =
r#"{"old\u0000key":"old\u0000value","literal":"\\u0000","adjacent":"\u0000\u0000"}"#;
let legacy_capabilities = r#"{"cap\u0000key":"cap\u0000value"}"#;
for request_id in [&single_request_id, &batch_request_id] {
sqlx::query(
r#"
INSERT INTO request_candidates (
id, request_id, candidate_index, retry_index, status,
extra_data, required_capabilities, created_at
)
VALUES ($1, $2, 0, 0, 'pending', $3::json, $4::json, NOW())
"#,
)
.bind(uuid::Uuid::new_v4().to_string())
.bind(request_id)
.bind(legacy_extra)
.bind(legacy_capabilities)
.execute(repository.pool())
.await
.expect("legacy JSON poison seed should persist in the json column");
}
let candidate = |request_id: &str, id: String| UpsertRequestCandidateRecord {
id,
request_id: request_id.to_string(),
user_id: None,
api_key_id: None,
username: None,
api_key_name: None,
candidate_index: 0,
retry_index: 0,
provider_id: None,
endpoint_id: None,
key_id: None,
status: RequestCandidateStatus::Success,
skip_reason: None,
is_cached: Some(false),
status_code: Some(200),
error_type: None,
error_message: Some("bad\0message".to_string()),
latency_ms: Some(1),
concurrent_requests: None,
extra_data: Some(json!({"new": true, "nested": "new\0value"})),
required_capabilities: None,
created_at_unix_ms: Some(1_700_000_000_000),
started_at_unix_ms: Some(1_700_000_000_000),
finished_at_unix_ms: Some(1_700_000_000_001),
};
repository
.upsert(candidate(
&single_request_id,
uuid::Uuid::new_v4().to_string(),
))
.await
.expect("single conflict should sanitize incoming and legacy JSON");
repository
.upsert_many(vec![
candidate(&batch_request_id, uuid::Uuid::new_v4().to_string()),
candidate(&healthy_request_id, uuid::Uuid::new_v4().to_string()),
])
.await
.expect("batch conflict should sanitize poison without blocking a healthy peer");
for request_id in [&single_request_id, &batch_request_id] {
let rows = repository
.list_by_request_id(request_id)
.await
.expect("sanitized candidate should be readable");
assert_eq!(rows.len(), 1);
assert_eq!(rows[0].status, RequestCandidateStatus::Success);
assert_eq!(rows[0].error_message.as_deref(), Some("bad�message"));
assert_eq!(
rows[0].extra_data,
Some(json!({
"old�key": "old�value",
"literal": "\\u0000",
"adjacent": "��",
"new": true,
"nested": "new�value"
}))
);
assert_eq!(
rows[0].required_capabilities,
Some(json!({"cap�key": "cap�value"}))
);
}
assert_eq!(
repository
.list_by_request_id(&healthy_request_id)
.await
.expect("healthy batch peer should be readable")
.len(),
1
);
sqlx::query("DELETE FROM request_candidates WHERE request_id = ANY($1)")
.bind(vec![
single_request_id,
batch_request_id,
healthy_request_id,
])
.execute(repository.pool())
.await
.expect("candidate NUL test rows should clean up");
}
}
crates/aether-data/adapters/postgres/src/error.rs
+73 -2
View File
@@ -1,7 +1,14 @@
use crate::DataLayerError;
pub(crate) fn postgres_error(error: impl std::fmt::Display) -> DataLayerError {
DataLayerError::postgres(error)
pub(crate) fn postgres_error(error: sqlx::Error) -> DataLayerError {
let sqlstate = error
.as_database_error()
.and_then(|database_error| database_error.code())
.map(|code| code.into_owned());
match sqlstate {
Some(code) => DataLayerError::Postgres(format!("{error} (SQLSTATE {code})")),
None => DataLayerError::postgres(error),
}
}
pub(crate) trait SqlxResultExt<T> {
@@ -13,3 +20,67 @@ impl<T> SqlxResultExt<T> for Result<T, sqlx::Error> {
self.map_err(postgres_error)
}
}
#[cfg(test)]
mod tests {
use super::postgres_error;
use sqlx::error::{DatabaseError, ErrorKind};
use std::borrow::Cow;
use std::error::Error;
use std::fmt::{Display, Formatter};
#[derive(Debug)]
struct TestDatabaseError;
impl Display for TestDatabaseError {
fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
formatter.write_str("unsupported Unicode escape sequence")
}
}
impl Error for TestDatabaseError {}
impl DatabaseError for TestDatabaseError {
fn message(&self) -> &str {
"unsupported Unicode escape sequence"
}
fn code(&self) -> Option<Cow<'_, str>> {
Some(Cow::Borrowed("22P05"))
}
fn as_error(&self) -> &(dyn Error + Send + Sync + 'static) {
self
}
fn as_error_mut(&mut self) -> &mut (dyn Error + Send + Sync + 'static) {
self
}
fn into_error(self: Box<Self>) -> Box<dyn Error + Send + Sync + 'static> {
self
}
fn kind(&self) -> ErrorKind {
ErrorKind::Other
}
}
#[test]
fn postgres_error_preserves_sqlstate_without_database_detail() {
let mapped = postgres_error(sqlx::Error::Database(Box::new(TestDatabaseError)));
assert_eq!(
mapped.to_string(),
"postgres error: error returned from database: unsupported Unicode escape sequence (SQLSTATE 22P05)"
);
}
#[test]
fn postgres_error_without_database_code_keeps_original_message() {
let mapped = postgres_error(sqlx::Error::PoolTimedOut);
assert_eq!(
mapped.to_string(),
"postgres error: pool timed out while waiting for an open connection"
);
}
}
crates/aether-data/adapters/postgres/src/provider_catalog.rs
+2 -2
View File
@@ -694,7 +694,7 @@ impl SqlxProviderCatalogReadRepository {
)
.await?
}
Err(error) => return Err(postgres_error(error)),
Err(error) => return Err(error),
};
Ok(rows)
}
@@ -733,7 +733,7 @@ impl SqlxProviderCatalogReadRepository {
)
.await?
}
Err(error) => return Err(postgres_error(error)),
Err(error) => return Err(error),
};
Ok(rows)
}