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

perf(gateway): scale request hot paths for 20k streams

Browse Source 
Shard and singleflight hot-path caches, batch and prioritize candidate and usage lifecycle persistence, and extend database and pressure-test instrumentation for 20k concurrent streams.
This commit is contained in:
elky
2026-07-22 02:11:08 +08:00
parent 7756c0913f
commit fc92c4f431
124 changed files with 36325 additions and 3217 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-data/adapters/mysql/src/provider_catalog.rs
+445 -39
View File
@@ -1,3 +1,5 @@
use std::collections::BTreeMap;
use async_trait::async_trait;
use sqlx::{
mysql::{MySqlArguments, MySqlRow},
@@ -6,7 +8,9 @@ use sqlx::{
};
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogKeyListQuery, ProviderCatalogReadRepository, ProviderCatalogSnapshot,
ProviderCatalogKeyAdaptiveStateUpdate, ProviderCatalogKeyHealthStateUpdate,
ProviderCatalogKeyListQuery, ProviderCatalogKeyRuntimeMetadataUpdate,
ProviderCatalogKeyStatusSnapshotUpdate, ProviderCatalogReadRepository, ProviderCatalogSnapshot,
ProviderCatalogUpstreamMetadataNamespaceUpdate, ProviderCatalogWriteRepository,
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage,
@@ -716,7 +720,23 @@ WHERE id = ?
}
}
transaction.commit().await.map_sql_err()?;
Ok(keys.to_vec())
let key_ids = keys.iter().map(|key| key.id.clone()).collect::<Vec<_>>();
let mut reloaded = self
.list_keys_by_ids(&key_ids)
.await?
.into_iter()
.map(|key| (key.id.clone(), key))
.collect::<BTreeMap<_, _>>();
keys.iter()
.map(|key| {
reloaded.remove(&key.id).ok_or_else(|| {
DataLayerError::UnexpectedValue(format!(
"updated provider catalog key {} could not be reloaded",
key.id
))
})
})
.collect()
}
pub async fn delete_key(&self, key_id: &str) -> Result<bool, DataLayerError> {
@@ -967,6 +987,38 @@ WHERE id = ?
Ok(rows_affected > 0)
}
pub async fn update_key_oauth_runtime_state(
&self,
key_id: &str,
oauth_invalid_at_unix_secs: Option<u64>,
oauth_invalid_reason: Option<&str>,
encrypted_auth_config_update: Option<&str>,
updated_at_unix_secs: Option<u64>,
) -> Result<bool, DataLayerError> {
validate_non_empty(key_id, "provider catalog key_id")?;
let rows_affected = sqlx::query(
r#"
UPDATE provider_api_keys
SET oauth_invalid_at = ?, oauth_invalid_reason = ?,
auth_config = COALESCE(?, auth_config), updated_at = ?
WHERE id = ?
"#,
)
.bind(optional_i64_from_u64(
oauth_invalid_at_unix_secs,
"provider_api_keys.oauth_invalid_at",
)?)
.bind(oauth_invalid_reason)
.bind(encrypted_auth_config_update)
.bind(updated_at_unix_secs.unwrap_or_else(current_unix_secs) as i64)
.bind(key_id)
.execute(&self.pool)
.await
.map_sql_err()?
.rows_affected();
Ok(rows_affected > 0)
}
pub async fn update_key_health_state(
&self,
key_id: &str,
@@ -1000,6 +1052,248 @@ WHERE id = ?
Ok(rows_affected > 0)
}
pub async fn reset_key_error_count(&self, key_id: &str) -> Result<bool, DataLayerError> {
validate_non_empty(key_id, "provider catalog key_id")?;
let rows_affected = sqlx::query(
r#"
UPDATE provider_api_keys
SET error_count = 0, updated_at = ?
WHERE id = ?
"#,
)
.bind(current_unix_secs() as i64)
.bind(key_id)
.execute(&self.pool)
.await
.map_sql_err()?
.rows_affected();
Ok(rows_affected > 0)
}
pub async fn compare_and_update_key_adaptive_state(
&self,
update: &ProviderCatalogKeyAdaptiveStateUpdate,
) -> Result<bool, DataLayerError> {
validate_non_empty(&update.key_id, "provider catalog key_id")?;
let status_snapshot_patch = adaptive_status_snapshot_patch(&update.status_snapshot_patch)?;
let expected = update.expected.canonicalized();
let next = update.next.canonicalized();
let mut builder =
QueryBuilder::<MySql>::new("UPDATE provider_api_keys SET learned_rpm_limit = ");
builder
.push_bind(optional_i64_from_u32(next.learned_rpm_limit))
.push(", rpm_429_count = ")
.push_bind(optional_i64_from_u32(next.rpm_429_count))
.push(", last_429_at = ")
.push_bind(optional_i64_from_u64(
next.last_429_at_unix_secs,
"provider_api_keys.last_429_at",
)?)
.push(", last_429_type = ")
.push_bind(&next.last_429_type)
.push(", adjustment_history = ")
.push_bind(optional_json_to_string(
&next.adjustment_history,
"provider_api_keys.adjustment_history",
)?)
.push(", utilization_samples = ")
.push_bind(optional_json_to_string(
&next.utilization_samples,
"provider_api_keys.utilization_samples",
)?)
.push(", last_probe_increase_at = ")
.push_bind(optional_i64_from_u64(
next.last_probe_increase_at_unix_secs,
"provider_api_keys.last_probe_increase_at",
)?)
.push(", last_rpm_peak = ")
.push_bind(optional_i64_from_u32(next.last_rpm_peak))
.push(", concurrent_429_count = ")
.push_bind(optional_i64_from_u32(next.concurrent_429_count))
.push(", status_snapshot = ");
push_status_snapshot_shallow_patch(&mut builder, &status_snapshot_patch)?;
builder
.push(", updated_at = ")
.push_bind(
update
.updated_at_unix_secs
.unwrap_or_else(current_unix_secs) as i64,
)
.push(" WHERE id = ")
.push_bind(&update.key_id)
.push(" AND learned_rpm_limit <=> ")
.push_bind(optional_i64_from_u32(expected.learned_rpm_limit))
.push(" AND rpm_429_count <=> ")
.push_bind(optional_i64_from_u32(expected.rpm_429_count))
.push(" AND last_429_at <=> ")
.push_bind(optional_i64_from_u64(
expected.last_429_at_unix_secs,
"provider_api_keys.last_429_at",
)?)
.push(" AND last_429_type <=> ")
.push_bind(&expected.last_429_type)
.push(" AND JSON_EXTRACT(adjustment_history, '$') <=> CAST(")
.push_bind(optional_json_to_string(
&expected.adjustment_history,
"provider_api_keys.adjustment_history",
)?)
.push(" AS JSON)")
.push(" AND JSON_EXTRACT(utilization_samples, '$') <=> CAST(")
.push_bind(optional_json_to_string(
&expected.utilization_samples,
"provider_api_keys.utilization_samples",
)?)
.push(" AS JSON)")
.push(" AND last_probe_increase_at <=> ")
.push_bind(optional_i64_from_u64(
expected.last_probe_increase_at_unix_secs,
"provider_api_keys.last_probe_increase_at",
)?)
.push(" AND last_rpm_peak <=> ")
.push_bind(optional_i64_from_u32(expected.last_rpm_peak))
.push(" AND concurrent_429_count <=> ")
.push_bind(optional_i64_from_u32(expected.concurrent_429_count));
let rows_affected = builder
.build()
.execute(&self.pool)
.await
.map_sql_err()?
.rows_affected();
Ok(rows_affected > 0)
}
pub async fn update_key_runtime_metadata(
&self,
update: &ProviderCatalogKeyRuntimeMetadataUpdate,
) -> Result<bool, DataLayerError> {
validate_runtime_metadata_update(update)?;
let namespace_path = format!(
"$.{}",
serde_json::to_string(&update.namespace).map_err(|err| {
DataLayerError::UnexpectedValue(format!(
"provider_api_keys.upstream_metadata namespace is not serializable: {err}"
))
})?
);
let metadata_value =
serde_json::to_string(&update.upstream_metadata_value).map_err(|err| {
DataLayerError::UnexpectedValue(format!(
"provider_api_keys.upstream_metadata value is not serializable: {err}"
))
})?;
let expected_metadata_value = update
.expected_upstream_metadata_value
.as_ref()
.map(serde_json::to_string)
.transpose()
.map_err(|err| {
DataLayerError::UnexpectedValue(format!(
"provider_api_keys.upstream_metadata expected value is not serializable: {err}"
))
})?;
let mut builder = QueryBuilder::<MySql>::new(
"UPDATE provider_api_keys SET upstream_metadata = JSON_SET(\
COALESCE(NULLIF(upstream_metadata, ''), '{}'), ",
);
builder
.push_bind(namespace_path.clone())
.push(", CAST(")
.push_bind(metadata_value)
.push(" AS JSON)), status_snapshot = ");
push_status_snapshot_shallow_patch(&mut builder, &update.status_snapshot_patch)?;
builder
.push(", updated_at = ")
.push_bind(
update
.updated_at_unix_secs
.unwrap_or_else(current_unix_secs) as i64,
)
.push(" WHERE id = ")
.push_bind(&update.key_id)
.push(" AND JSON_EXTRACT(COALESCE(NULLIF(upstream_metadata, ''), '{}'), ")
.push_bind(namespace_path)
.push(") <=> CAST(")
.push_bind(expected_metadata_value)
.push(" AS JSON)");
let rows_affected = builder
.build()
.execute(&self.pool)
.await
.map_sql_err()?
.rows_affected();
Ok(rows_affected > 0)
}
pub async fn update_key_status_snapshot(
&self,
update: &ProviderCatalogKeyStatusSnapshotUpdate,
) -> Result<bool, DataLayerError> {
validate_non_empty(&update.key_id, "provider catalog key_id")?;
if !update.status_snapshot_patch.is_object() {
return Err(DataLayerError::InvalidInput(
"provider catalog status snapshot patch must be an object".to_string(),
));
}
let mut builder =
QueryBuilder::<MySql>::new("UPDATE provider_api_keys SET status_snapshot = ");
push_status_snapshot_shallow_patch(&mut builder, &update.status_snapshot_patch)?;
builder
.push(", updated_at = ")
.push_bind(
update
.updated_at_unix_secs
.unwrap_or_else(current_unix_secs) as i64,
)
.push(" WHERE id = ")
.push_bind(&update.key_id);
let rows_affected = builder
.build()
.execute(&self.pool)
.await
.map_sql_err()?
.rows_affected();
Ok(rows_affected > 0)
}
pub async fn compare_and_update_key_health_state(
&self,
update: &ProviderCatalogKeyHealthStateUpdate,
) -> Result<bool, DataLayerError> {
validate_non_empty(&update.key_id, "provider catalog key_id")?;
let rows_affected = sqlx::query(
r#"
UPDATE provider_api_keys
SET health_by_format = ?, circuit_breaker_by_format = ?, updated_at = ?
WHERE id = ?
AND JSON_EXTRACT(health_by_format, '$') <=> CAST(? AS JSON)
AND JSON_EXTRACT(circuit_breaker_by_format, '$') <=> CAST(? AS JSON)
"#,
)
.bind(optional_json_to_string(
&update.health_by_format,
"provider_api_keys.health_by_format",
)?)
.bind(optional_json_to_string(
&update.circuit_breaker_by_format,
"provider_api_keys.circuit_breaker_by_format",
)?)
.bind(current_unix_secs() as i64)
.bind(&update.key_id)
.bind(optional_json_to_string(
&update.expected_health_by_format,
"provider_api_keys.health_by_format",
)?)
.bind(optional_json_to_string(
&update.expected_circuit_breaker_by_format,
"provider_api_keys.circuit_breaker_by_format",
)?)
.execute(&self.pool)
.await
.map_sql_err()?
.rows_affected();
Ok(rows_affected > 0)
}
async fn reload_provider(
&self,
provider_id: &str,
@@ -1303,6 +1597,25 @@ impl ProviderCatalogWriteRepository for MysqlProviderCatalogReadRepository {
.await
}
async fn update_key_oauth_runtime_state(
&self,
key_id: &str,
oauth_invalid_at_unix_secs: Option<u64>,
oauth_invalid_reason: Option<&str>,
encrypted_auth_config_update: Option<&str>,
updated_at_unix_secs: Option<u64>,
) -> Result<bool, DataLayerError> {
Self::update_key_oauth_runtime_state(
self,
key_id,
oauth_invalid_at_unix_secs,
oauth_invalid_reason,
encrypted_auth_config_update,
updated_at_unix_secs,
)
.await
}
async fn update_key_health_state(
&self,
key_id: &str,
@@ -1319,6 +1632,38 @@ impl ProviderCatalogWriteRepository for MysqlProviderCatalogReadRepository {
)
.await
}
async fn reset_key_error_count(&self, key_id: &str) -> Result<bool, DataLayerError> {
Self::reset_key_error_count(self, key_id).await
}
async fn compare_and_update_key_adaptive_state(
&self,
update: &ProviderCatalogKeyAdaptiveStateUpdate,
) -> Result<bool, DataLayerError> {
Self::compare_and_update_key_adaptive_state(self, update).await
}
async fn update_key_runtime_metadata(
&self,
update: &ProviderCatalogKeyRuntimeMetadataUpdate,
) -> Result<bool, DataLayerError> {
Self::update_key_runtime_metadata(self, update).await
}
async fn update_key_status_snapshot(
&self,
update: &ProviderCatalogKeyStatusSnapshotUpdate,
) -> Result<bool, DataLayerError> {
Self::update_key_status_snapshot(self, update).await
}
async fn compare_and_update_key_health_state(
&self,
update: &ProviderCatalogKeyHealthStateUpdate,
) -> Result<bool, DataLayerError> {
Self::compare_and_update_key_health_state(self, update).await
}
}
fn current_unix_secs() -> u64 {
@@ -1334,6 +1679,87 @@ fn validate_non_empty(value: &str, field_name: &str) -> Result<(), DataLayerErro
Ok(())
}
fn adaptive_status_snapshot_patch(
patch: &serde_json::Value,
) -> Result<serde_json::Value, DataLayerError> {
const OWNED_FIELDS: [&str; 6] = [
"observation_count",
"header_observation_count",
"latest_upstream_limit",
"learning_confidence",
"enforcement_active",
"known_boundary",
];
let object = patch.as_object().ok_or_else(|| {
DataLayerError::InvalidInput(
"provider catalog adaptive status snapshot patch must be an object".to_string(),
)
})?;
Ok(serde_json::Value::Object(
OWNED_FIELDS
.into_iter()
.filter_map(|field| {
object
.get(field)
.cloned()
.map(|value| (field.to_string(), value))
})
.collect(),
))
}
fn validate_runtime_metadata_update(
update: &ProviderCatalogKeyRuntimeMetadataUpdate,
) -> Result<(), DataLayerError> {
validate_non_empty(&update.key_id, "provider catalog key_id")?;
validate_non_empty(
&update.namespace,
"provider catalog runtime metadata namespace",
)?;
if !update.status_snapshot_patch.is_object() {
return Err(DataLayerError::InvalidInput(
"provider catalog runtime status snapshot patch must be an object".to_string(),
));
}
Ok(())
}
fn push_status_snapshot_shallow_patch<'args>(
builder: &mut QueryBuilder<'args, MySql>,
patch: &serde_json::Value,
) -> Result<(), DataLayerError> {
let object = patch.as_object().ok_or_else(|| {
DataLayerError::InvalidInput(
"provider catalog status snapshot patch must be an object".to_string(),
)
})?;
if object.is_empty() {
builder.push("COALESCE(NULLIF(status_snapshot, ''), '{}')");
return Ok(());
}
builder.push("JSON_SET(COALESCE(NULLIF(status_snapshot, ''), '{}')");
for (field, value) in object {
let path = format!(
"$.{}",
serde_json::to_string(field).map_err(|err| {
DataLayerError::UnexpectedValue(format!(
"provider_api_keys.status_snapshot field is not serializable: {err}"
))
})?
);
let value = serde_json::to_string(value).map_err(|err| {
DataLayerError::UnexpectedValue(format!(
"provider_api_keys.status_snapshot value is not serializable: {err}"
))
})?;
builder.push(", ").push_bind(path).push(", CAST(");
builder.push_bind(value).push(" AS JSON)");
}
builder.push(")");
Ok(())
}
fn validate_provider(provider: &StoredProviderCatalogProvider) -> Result<(), DataLayerError> {
validate_non_empty(&provider.id, "provider catalog provider.id")?;
validate_non_empty(&provider.name, "provider catalog provider.name")?;
@@ -1477,17 +1903,6 @@ SET
fingerprint = ?,
rpm_limit = ?,
concurrent_limit = ?,
learned_rpm_limit = ?,
concurrent_429_count = ?,
rpm_429_count = ?,
last_429_at = ?,
last_429_type = ?,
adjustment_history = ?,
utilization_samples = ?,
last_probe_increase_at = ?,
last_rpm_peak = ?,
last_models_fetch_at = ?,
last_models_fetch_error = ?,
auto_fetch_models = ?,
locked_models = ?,
model_include_patterns = ?,
@@ -1554,32 +1969,6 @@ fn key_update_query(
)?)
.bind(optional_i64_from_u32(key.rpm_limit))
.bind(key.concurrent_limit)
.bind(optional_i64_from_u32(key.learned_rpm_limit))
.bind(optional_i64_from_u32(key.concurrent_429_count).unwrap_or(0))
.bind(optional_i64_from_u32(key.rpm_429_count).unwrap_or(0))
.bind(optional_i64_from_u64(
key.last_429_at_unix_secs,
"provider_api_keys.last_429_at",
)?)
.bind(&key.last_429_type)
.bind(optional_json_to_string(
&key.adjustment_history,
"provider_api_keys.adjustment_history",
)?)
.bind(optional_json_to_string(
&key.utilization_samples,
"provider_api_keys.utilization_samples",
)?)
.bind(optional_i64_from_u64(
key.last_probe_increase_at_unix_secs,
"provider_api_keys.last_probe_increase_at",
)?)
.bind(optional_i64_from_u32(key.last_rpm_peak))
.bind(optional_i64_from_u64(
key.last_models_fetch_at_unix_secs,
"provider_api_keys.last_models_fetch_at",
)?)
.bind(&key.last_models_fetch_error)
.bind(key.auto_fetch_models)
.bind(optional_json_to_string(
&key.locked_models,
@@ -1935,6 +2324,23 @@ mod tests {
StoredProviderCatalogProvider,
};
use serde_json::json;
#[test]
fn ordinary_key_update_does_not_own_adaptive_runtime_fields() {
let sql = super::key_update_sql().to_ascii_lowercase();
for runtime_assignment in [
"learned_rpm_limit =",
"rpm_429_count =",
"last_429_at =",
"last_429_type =",
"adjustment_history =",
"utilization_samples =",
"last_probe_increase_at =",
"last_rpm_peak =",
] {
assert!(!sql.contains(runtime_assignment));
}
}
use sqlx::Execute;
#[tokio::test]