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

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.

apps/aether-gateway/src/handlers/admin/provider/endpoint_keys/mutations/reset_cycle_stats.rs

@@ -2,6 +2,7 @@ use crate::handlers::admin::provider::shared::paths::admin_reset_cycle_stats_key
 use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
 use crate::handlers::admin::shared::provider_key_status_snapshot_payload;
 use crate::GatewayError;
+use aether_data_contracts::repository::provider_catalog::ProviderCatalogKeyStatusSnapshotUpdate;
 use axum::{
     body::{Body, Bytes},
     http,
@@ -33,7 +34,7 @@ pub(super) async fn maybe_handle(
     let Some(key_id) = admin_reset_cycle_stats_key_id(request_context.path()) else {
         return Ok(Some(not_found_response("Key 不存在")));
     };
-    let Some(mut key) = state
+    let Some(key) = state
         .read_provider_catalog_keys_by_ids(std::slice::from_ref(&key_id))
         .await?
         .into_iter()
@@ -65,11 +66,17 @@ pub(super) async fn maybe_handle(
         return Ok(Some(bad_request_response("当前账号没有可重置的周期窗口")));
     }
 
-    key.status_snapshot = Some(status_snapshot);
-    key.updated_at_unix_secs = Some(now_unix_secs);
-    let Some(_) = state.update_provider_catalog_key(&key).await? else {
+    let quota = status_snapshot.get("quota").cloned().unwrap_or(Value::Null);
+    if !state
+        .update_provider_catalog_key_status_snapshot(&ProviderCatalogKeyStatusSnapshotUpdate {
+            key_id: key.id.clone(),
+            status_snapshot_patch: json!({"quota":quota}),
+            updated_at_unix_secs: Some(now_unix_secs),
+        })
+        .await?
+    {
         return Ok(None);
-    };
+    }
 
     Ok(Some(
         Json(json!({

apps/aether-gateway/src/handlers/admin/provider/endpoint_keys/mutations/update.rs

@@ -82,9 +82,22 @@ pub(super) async fn maybe_handle(
         Ok(record) => record,
         Err(detail) => return Ok(Some(bad_request_response(detail))),
     };
-    let Some(updated) = state.update_provider_catalog_key(&updated_record).await? else {
+    let Some(mut updated) = state.update_provider_catalog_key(&updated_record).await? else {
         return Ok(None);
     };
+    if updated_record.learned_rpm_limit != existing_key.learned_rpm_limit {
+        let Some(reloaded) = state
+            .set_provider_catalog_key_learned_rpm_limit(
+                &key_id,
+                updated_record.learned_rpm_limit,
+                updated_record.updated_at_unix_secs,
+            )
+            .await?
+        else {
+            return Ok(None);
+        };
+        updated = reloaded;
+    }
     let should_overwrite_allowed_models_immediately =
         admin_provider_key_update_requires_immediate_model_fetch(&existing_key, &updated);
     let updated = if should_overwrite_allowed_models_immediately {

apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/complete/key.rs

@@ -1,5 +1,7 @@
 use super::super::super::errors::build_internal_control_error_response;
-use super::super::super::provisioning::provider_oauth_token_payload_expires_at_unix_secs;
+use super::super::super::provisioning::{
+    provider_oauth_token_payload_expires_at_unix_secs, seed_provider_oauth_pool_score,
+};
 use super::super::super::runtime::{
     resolve_provider_oauth_runtime_endpoints,
     spawn_provider_oauth_account_state_refresh_after_update,
@@ -220,6 +222,25 @@ pub(super) async fn handle_admin_provider_oauth_complete_key(
             "Key 不存在",
         ));
     }
+    if !state
+        .clear_provider_catalog_key_oauth_invalid_marker(&key_id)
+        .await?
+    {
+        return Ok(build_internal_control_error_response(
+            http::StatusCode::NOT_FOUND,
+            "Key 不存在",
+        ));
+    }
+    let Some(recovered_key) = state
+        .reset_provider_catalog_key_recovery_state(&key_id)
+        .await?
+    else {
+        return Ok(build_internal_control_error_response(
+            http::StatusCode::NOT_FOUND,
+            "Key 不存在",
+        ));
+    };
+    seed_provider_oauth_pool_score(state, &provider.id, &recovered_key, now_unix_secs).await;
 
     spawn_provider_oauth_account_state_refresh_after_update(
         state.cloned_app(),

apps/aether-gateway/src/handlers/admin/provider/oauth/provisioning.rs

@@ -2,11 +2,16 @@ use super::state::{
     decode_jwt_claims, enrich_admin_provider_oauth_auth_config, json_non_empty_string,
     json_u64_value,
 };
+use crate::ai_serving::{
+    build_provider_key_pool_score_upsert, provider_key_pool_score_id, provider_key_pool_score_scope,
+};
 use crate::handlers::admin::admin_provider_pool_config;
 use crate::handlers::admin::request::AdminAppState;
-use crate::maintenance::ensure_provider_key_pool_scores_for_keys;
 use crate::provider_key_auth::provider_active_api_formats;
 use crate::GatewayError;
+use aether_data_contracts::repository::pool_scores::{
+    GetPoolMemberScoresByIdsQuery, PoolMemberIdentity,
+};
 use aether_data_contracts::repository::provider_catalog::{
     StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
 };
@@ -211,14 +216,22 @@ pub(crate) async fn update_existing_provider_oauth_catalog_key(
     if updated.fingerprint.is_none() {
         updated.fingerprint = grok_oauth_catalog_key_fingerprint(provider_type, auth_config);
     }
-    updated.health_by_format = Some(json!({}));
-    updated.circuit_breaker_by_format = Some(json!({}));
-    updated.error_count = Some(0);
     if let Some(proxy) = proxy {
         updated.proxy = Some(proxy);
     }
     updated.updated_at_unix_secs = Some(now_unix_secs);
-    let persisted = state.update_provider_catalog_key(&updated).await?;
+    if state.update_provider_catalog_key(&updated).await?.is_none() {
+        return Ok(None);
+    }
+    if !state
+        .clear_provider_catalog_key_oauth_invalid_marker(&updated.id)
+        .await?
+    {
+        return Ok(None);
+    }
+    let persisted = state
+        .reset_provider_catalog_key_recovery_state(&updated.id)
+        .await?;
     if let Some(key) = persisted.as_ref() {
         let _ = state
             .app()
@@ -229,7 +242,7 @@ pub(crate) async fn update_existing_provider_oauth_catalog_key(
     Ok(persisted)
 }
 
-async fn seed_provider_oauth_pool_score(
+pub(super) async fn seed_provider_oauth_pool_score(
     state: &AdminAppState<'_>,
     provider_id: &str,
     key: &StoredProviderCatalogKey,
@@ -257,38 +270,45 @@ async fn seed_provider_oauth_pool_score(
     let Some(pool_config) = admin_provider_pool_config(&provider) else {
         return;
     };
-    let endpoints = match state
-        .list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider_id))
+    if !key.is_active || key.provider_id != provider.id {
+        return;
+    }
+
+    let identity = PoolMemberIdentity::provider_api_key(provider.id.clone(), key.id.clone());
+    let scope = provider_key_pool_score_scope();
+    let score_id = provider_key_pool_score_id(&identity, &scope);
+    let existing = match state
+        .app()
+        .data
+        .get_pool_member_scores_by_ids(&GetPoolMemberScoresByIdsQuery {
+            ids: vec![score_id],
+        })
         .await
     {
-        Ok(endpoints) => endpoints,
+        Ok(mut scores) => scores.pop(),
         Err(err) => {
             tracing::debug!(
                 provider_id = %provider_id,
                 key_id = %key.id,
                 error = ?err,
-                "gateway provider oauth provisioning: failed to read endpoints for pool score seed"
+                "gateway provider oauth provisioning: failed to read existing pool score"
             );
             return;
         }
     };
-    let score_ensure_budget = (pool_config.score_fallback_scan_limit as usize).clamp(1, 50_000);
-    if let Err(err) = ensure_provider_key_pool_scores_for_keys(
-        state.as_ref(),
-        &provider,
-        &pool_config,
-        &endpoints,
-        std::slice::from_ref(key),
+    let upsert = build_provider_key_pool_score_upsert(
+        key,
+        provider.provider_type.as_str(),
+        existing.as_ref(),
         now_unix_secs,
-        score_ensure_budget,
-    )
-    .await
-    {
+        pool_config.score_rules,
+    );
+    if let Err(err) = state.app().data.upsert_pool_member_score(upsert).await {
         tracing::debug!(
             provider_id = %provider_id,
             key_id = %key.id,
             error = ?err,
-            "gateway provider oauth provisioning: failed to seed pool score row"
+            "gateway provider oauth provisioning: failed to refresh pool score row"
         );
     }
 }

apps/aether-gateway/src/handlers/admin/provider/oauth/quota/shared.rs

@@ -14,6 +14,7 @@ use aether_contracts::{
     ResolvedTransportProfile, EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER,
 };
 use aether_data_contracts::repository::provider_catalog::{
+    ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate,
     StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
 };
 use aether_provider_pool::{ProviderPoolQuotaRequestSpec, ProviderPoolService};
@@ -288,6 +289,30 @@ pub(crate) async fn persist_provider_quota_refresh_state(
     oauth_invalid_reason: Option<String>,
     encrypted_auth_config: Option<String>,
 ) -> Result<bool, GatewayError> {
+    persist_provider_quota_refresh_state_after_read(
+        state,
+        key_id,
+        metadata_update,
+        oauth_invalid_at_unix_secs,
+        oauth_invalid_reason,
+        encrypted_auth_config,
+        std::future::ready(()),
+    )
+    .await
+}
+
+async fn persist_provider_quota_refresh_state_after_read<F>(
+    state: &AdminAppState<'_>,
+    key_id: &str,
+    metadata_update: Option<&serde_json::Value>,
+    oauth_invalid_at_unix_secs: Option<u64>,
+    oauth_invalid_reason: Option<String>,
+    encrypted_auth_config: Option<String>,
+    after_read: F,
+) -> Result<bool, GatewayError>
+where
+    F: std::future::Future<Output = ()>,
+{
     let Some(mut latest_key) = state
         .read_provider_catalog_keys_by_ids(&[key_id.to_string()])
         .await?
@@ -296,7 +321,11 @@ pub(crate) async fn persist_provider_quota_refresh_state(
     else {
         return Ok(false);
     };
+    after_read.await;
 
+    // Keep the namespace values observed before applying the refresh response;
+    // each runtime metadata write uses them as its CAS expectation.
+    let observed_upstream_metadata = latest_key.upstream_metadata.clone();
     let mut quota_snapshot_provider_type = None::<String>;
     if let Some(metadata_update) = metadata_update {
         latest_key.upstream_metadata = Some(merge_upstream_metadata(
@@ -306,8 +335,8 @@ pub(crate) async fn persist_provider_quota_refresh_state(
         quota_snapshot_provider_type =
             aether_provider_pool::provider_pool_quota_metadata_provider_type(metadata_update);
     }
-    if let Some(encrypted_auth_config) = encrypted_auth_config {
-        latest_key.encrypted_auth_config = Some(encrypted_auth_config);
+    if let Some(encrypted_auth_config) = encrypted_auth_config.as_ref() {
+        latest_key.encrypted_auth_config = Some(encrypted_auth_config.clone());
     }
     latest_key.oauth_invalid_at_unix_secs = oauth_invalid_at_unix_secs;
     latest_key.oauth_invalid_reason = oauth_invalid_reason;
@@ -325,10 +354,91 @@ pub(crate) async fn persist_provider_quota_refresh_state(
         .duration_since(UNIX_EPOCH)
         .ok()
         .map(|duration| duration.as_secs());
-    Ok(state
-        .update_provider_catalog_key(&latest_key)
-        .await?
-        .is_some())
+    let status_patch = provider_quota_refresh_status_patch(latest_key.status_snapshot.as_ref());
+    let metadata_updates = metadata_update
+        .and_then(serde_json::Value::as_object)
+        .map(|updates| {
+            updates
+                .iter()
+                .map(|(namespace, value)| (namespace.clone(), value.clone()))
+                .collect::<Vec<_>>()
+        })
+        .unwrap_or_default();
+    if metadata_updates.is_empty() {
+        if !state
+            .update_provider_catalog_key_oauth_runtime_state(
+                key_id,
+                latest_key.oauth_invalid_at_unix_secs,
+                latest_key.oauth_invalid_reason.as_deref(),
+                encrypted_auth_config.as_deref(),
+                latest_key.updated_at_unix_secs,
+            )
+            .await?
+        {
+            return Ok(false);
+        }
+        return state
+            .update_provider_catalog_key_status_snapshot(&ProviderCatalogKeyStatusSnapshotUpdate {
+                key_id: key_id.to_string(),
+                status_snapshot_patch: status_patch,
+                updated_at_unix_secs: latest_key.updated_at_unix_secs,
+            })
+            .await;
+    }
+
+    for (index, (namespace, value)) in metadata_updates.iter().enumerate() {
+        let patch = if index + 1 == metadata_updates.len() {
+            status_patch.clone()
+        } else {
+            serde_json::json!({})
+        };
+        let mut expected = observed_upstream_metadata
+            .as_ref()
+            .and_then(serde_json::Value::as_object)
+            .and_then(|metadata| metadata.get(namespace))
+            .cloned();
+        let persisted = state
+            .app()
+            .update_provider_catalog_key_runtime_metadata(
+                &ProviderCatalogKeyRuntimeMetadataUpdate {
+                    key_id: key_id.to_string(),
+                    namespace: namespace.clone(),
+                    expected_upstream_metadata_value: expected.clone(),
+                    upstream_metadata_value: value.clone(),
+                    status_snapshot_patch: patch.clone(),
+                    updated_at_unix_secs: latest_key.updated_at_unix_secs,
+                },
+            )
+            .await?;
+        if !persisted {
+            // The refresh response is an authoritative snapshot.  Do not
+            // replay it over a newer namespace after a CAS conflict.
+            return Ok(false);
+        }
+    }
+    state
+        .update_provider_catalog_key_oauth_runtime_state(
+            key_id,
+            latest_key.oauth_invalid_at_unix_secs,
+            latest_key.oauth_invalid_reason.as_deref(),
+            encrypted_auth_config.as_deref(),
+            latest_key.updated_at_unix_secs,
+        )
+        .await
+}
+
+fn provider_quota_refresh_status_patch(
+    status_snapshot: Option<&serde_json::Value>,
+) -> serde_json::Value {
+    let mut patch = serde_json::Map::new();
+    if let Some(snapshot) = status_snapshot.and_then(serde_json::Value::as_object) {
+        for field in ["quota", "oauth"] {
+            if let Some(value) = snapshot.get(field) {
+                patch.insert(field.to_string(), value.clone());
+            }
+        }
+    }
+    serde_json::Value::Object(patch)
 }
 
 pub(super) async fn execute_provider_quota_plan(
@@ -372,3 +482,95 @@ pub(super) async fn execute_provider_quota_plan(
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::data::GatewayDataState;
+    use crate::AppState;
+    use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
+    use aether_data_contracts::repository::provider_catalog::{
+        ProviderCatalogReadRepository, ProviderCatalogWriteRepository, StoredProviderCatalogKey,
+    };
+    use serde_json::json;
+    use std::sync::Arc;
+
+    #[tokio::test]
+    async fn metadata_cas_conflict_does_not_persist_stale_oauth_runtime_state() {
+        let mut key = StoredProviderCatalogKey::new(
+            "key-codex-cas".to_string(),
+            "provider-codex-cas".to_string(),
+            "Codex CAS".to_string(),
+            "oauth".to_string(),
+            None,
+            true,
+        )
+        .expect("key should build");
+        key.encrypted_auth_config = Some("old-auth-config".to_string());
+        key.oauth_invalid_at_unix_secs = Some(100);
+        key.oauth_invalid_reason = Some("old-invalid-reason".to_string());
+        key.upstream_metadata = Some(json!({"codex":{"remaining":5}}));
+        key.status_snapshot = Some(json!({"oauth":{"invalid":true}}));
+
+        let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
+            vec![],
+            vec![],
+            vec![key],
+        ));
+        let app = AppState::new()
+            .expect("app should build")
+            .with_data_state_for_tests(
+                GatewayDataState::with_provider_catalog_repository_for_tests(Arc::clone(
+                    &repository,
+                )),
+            );
+        let admin_state = AdminAppState::new(&app);
+        let concurrent_repository = Arc::clone(&repository);
+        let metadata_update = json!({"codex":{"remaining":3}});
+
+        let persisted = persist_provider_quota_refresh_state_after_read(
+            &admin_state,
+            "key-codex-cas",
+            Some(&metadata_update),
+            Some(200),
+            Some("new-invalid-reason".to_string()),
+            Some("new-auth-config".to_string()),
+            async move {
+                assert!(concurrent_repository
+                    .update_key_runtime_metadata(&ProviderCatalogKeyRuntimeMetadataUpdate {
+                        key_id: "key-codex-cas".to_string(),
+                        namespace: "codex".to_string(),
+                        expected_upstream_metadata_value: Some(json!({"remaining":5})),
+                        upstream_metadata_value: json!({"remaining":4}),
+                        status_snapshot_patch: json!({}),
+                        updated_at_unix_secs: Some(150),
+                    })
+                    .await
+                    .expect("concurrent metadata update should execute"));
+            },
+        )
+        .await
+        .expect("quota refresh persistence should not error");
+
+        assert!(!persisted, "stale namespace should report a CAS conflict");
+        let stored = repository
+            .list_keys_by_ids(&["key-codex-cas".to_string()])
+            .await
+            .expect("key should reload")
+            .pop()
+            .expect("key should remain");
+        assert_eq!(
+            stored.encrypted_auth_config.as_deref(),
+            Some("old-auth-config")
+        );
+        assert_eq!(stored.oauth_invalid_at_unix_secs, Some(100));
+        assert_eq!(
+            stored.oauth_invalid_reason.as_deref(),
+            Some("old-invalid-reason")
+        );
+        assert_eq!(
+            stored.upstream_metadata.as_ref().unwrap()["codex"],
+            json!({"remaining":4})
+        );
+    }
+}