From d336d1a7fabf3e3b5e3b97adb3fa0b656a5b586b Mon Sep 17 00:00:00 2001 From: elky Date: Wed, 24 Jun 2026 01:47:26 +0800 Subject: [PATCH] Improve gateway scheduling and runtime admission --- Cargo.lock | 1 + Cargo.toml | 1 + apps/aether-gateway/Cargo.toml | 1 + .../planner/candidate_affinity_cache.rs | 15 + .../planner/candidate_materialization.rs | 375 +++++- .../ai_serving/planner/candidate_ranking.rs | 77 +- .../planner/candidate_resolution.rs | 41 +- .../ai_serving/planner/candidate_source.rs | 157 ++- .../candidate_transport_ranking_facts.rs | 173 ++- .../passthrough/provider/family/build.rs | 4 +- .../planner/passthrough/provider/plans.rs | 8 +- .../src/ai_serving/planner/redaction.rs | 116 +- .../ai_serving/planner/specialized/files.rs | 12 +- .../ai_serving/planner/specialized/image.rs | 12 +- .../ai_serving/planner/specialized/video.rs | 6 +- .../planner/standard/family/build.rs | 12 +- .../standard/openai/chat/decision/payload.rs | 25 + .../standard/openai/chat/decision/request.rs | 29 +- .../planner/standard/openai/chat/mod.rs | 4 +- .../standard/openai/chat/plans/stream.rs | 33 +- .../standard/openai/chat/plans/sync.rs | 2 +- .../planner/standard/openai/responses/mod.rs | 4 +- .../standard/openai/responses/plans.rs | 8 +- .../src/ai_serving/planner/state/transport.rs | 12 + apps/aether-gateway/src/cache/auth_context.rs | 76 +- apps/aether-gateway/src/cache/auth_runtime.rs | 310 ++++- .../src/cache/candidate_page.rs | 479 +++++++ apps/aether-gateway/src/cache/mod.rs | 17 +- .../aether-gateway/src/cache/system_config.rs | 72 +- apps/aether-gateway/src/clock.rs | 9 + .../src/control/auth/resolution.rs | 291 +++- apps/aether-gateway/src/data/state/auth.rs | 30 +- apps/aether-gateway/src/data/state/catalog.rs | 15 +- apps/aether-gateway/src/data/state/core.rs | 150 ++- apps/aether-gateway/src/data/state/mod.rs | 13 +- apps/aether-gateway/src/data/state/models.rs | 90 +- apps/aether-gateway/src/data/state/runtime.rs | 368 ++++- .../src/dispatch/pool_scheduler.rs | 53 +- apps/aether-gateway/src/error.rs | 78 ++ .../src/execution_runtime/stream/execution.rs | 1198 ++++++++++++++++- .../stream/execution_failures.rs | 8 +- .../src/execution_runtime/stream_pump.rs | 2 + .../src/execution_runtime/sync/execution.rs | 8 +- .../src/execution_runtime/transport.rs | 814 ++++++++++- .../src/executor/candidate_loop.rs | 121 +- apps/aether-gateway/src/executor/mod.rs | 3 +- apps/aether-gateway/src/executor/outcome.rs | 93 +- .../src/executor/stream_path.rs | 6 + apps/aether-gateway/src/handlers/proxy/mod.rs | 97 +- apps/aether-gateway/src/lib.rs | 6 +- apps/aether-gateway/src/main.rs | 209 ++- .../src/middleware/access_log.rs | 6 + apps/aether-gateway/src/middleware/mod.rs | 3 +- .../src/orchestration/effects.rs | 134 ++ apps/aether-gateway/src/privacy/mod.rs | 314 ++++- .../src/provider_pool_demand.rs | 191 ++- .../src/request_candidate_queue.rs | 295 +++- .../aether-gateway/src/request_diagnostics.rs | 200 +++ .../src/scheduler/candidate/affinity.rs | 9 + .../src/scheduler/candidate/mod.rs | 7 +- .../src/scheduler/candidate/selection.rs | 28 +- .../scheduler/candidate/tests/selection.rs | 104 +- apps/aether-gateway/src/stage_metrics.rs | 335 +++++ apps/aether-gateway/src/state/app.rs | 58 +- apps/aether-gateway/src/state/cache.rs | 4 +- apps/aether-gateway/src/state/core.rs | 132 +- apps/aether-gateway/src/state/mod.rs | 2 +- apps/aether-gateway/src/state/oauth.rs | 103 +- apps/aether-gateway/src/state/testing.rs | 5 + apps/aether-gateway/src/upstream_admission.rs | 269 ++++ apps/aether-tunnel/src/upstream_client.rs | 11 +- crates/aether-cache/src/ttl_map.rs | 101 ++ crates/aether-contracts/src/lib.rs | 3 +- crates/aether-contracts/src/plan.rs | 1 + .../src/repository/candidates/types.rs | 12 + .../aether-data/src/repository/auth/memory.rs | 76 +- .../src/repository/candidates/postgres.rs | 321 +++++ .../aether-provider-transport/src/network.rs | 18 + .../aether-runtime-state/src/redis/client.rs | 35 +- .../src/bin/gateway_pressure_probe.rs | 51 +- .../aether-testkit/src/bin/http_load_probe.rs | 51 +- .../src/bin/mock_openai_upstream.rs | 61 +- crates/aether-testkit/src/load.rs | 529 +++++++- crates/aether-usage-runtime/src/config.rs | 33 + .../src/request_metadata.rs | 70 +- crates/aether-usage-runtime/src/runtime.rs | 1052 ++++++++++++++- crates/aether-usage-runtime/src/worker.rs | 107 +- 87 files changed, 9671 insertions(+), 804 deletions(-) create mode 100644 apps/aether-gateway/src/cache/candidate_page.rs create mode 100644 apps/aether-gateway/src/request_diagnostics.rs create mode 100644 apps/aether-gateway/src/stage_metrics.rs create mode 100644 apps/aether-gateway/src/upstream_admission.rs diff --git a/Cargo.lock b/Cargo.lock index 98390c5f1..af87bf166 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -270,6 +270,7 @@ dependencies = [ "serde_json", "sha1", "sha2", + "socket2 0.6.3", "sqlx", "tar", "thiserror 2.0.18", diff --git a/Cargo.toml b/Cargo.toml index 13b1787f1..6ec414e7f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -92,6 +92,7 @@ serde = { version = "1", features = ["derive"] } serde_json = { version = "1", features = ["preserve_order"] } serde_path_to_error = "0.1" sha2 = "0.10" +socket2 = "0.6" tar = "0.4" sqlx = { version = "0.8", default-features = false, features = ["postgres", "mysql", "sqlite", "runtime-tokio-rustls", "chrono"] } thiserror = "2" diff --git a/apps/aether-gateway/Cargo.toml b/apps/aether-gateway/Cargo.toml index fa0757dfc..df4c60941 100644 --- a/apps/aether-gateway/Cargo.toml +++ b/apps/aether-gateway/Cargo.toml @@ -58,6 +58,7 @@ serde.workspace = true serde_json.workspace = true sha1 = "0.10" sha2 = { workspace = true, features = ["oid"] } +socket2.workspace = true tar.workspace = true sqlx.workspace = true thiserror.workspace = true diff --git a/apps/aether-gateway/src/ai_serving/planner/candidate_affinity_cache.rs b/apps/aether-gateway/src/ai_serving/planner/candidate_affinity_cache.rs index 2c80b84de..3a77c8b95 100644 --- a/apps/aether-gateway/src/ai_serving/planner/candidate_affinity_cache.rs +++ b/apps/aether-gateway/src/ai_serving/planner/candidate_affinity_cache.rs @@ -8,6 +8,12 @@ use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL; const PLANNER_SCHEDULER_AFFINITY_MAX_ENTRIES: usize = 10_000; +pub(crate) fn has_explicit_session_affinity( + client_session_affinity: Option<&ClientSessionAffinity>, +) -> bool { + client_session_affinity.is_some_and(ClientSessionAffinity::has_session_key) +} + pub(crate) fn read_cached_scheduler_affinity_target( state: PlannerAppState<'_>, auth_snapshot: Option<&GatewayAuthApiKeySnapshot>, @@ -15,6 +21,9 @@ pub(crate) fn read_cached_scheduler_affinity_target( client_api_format: &str, requested_model: Option<&str>, ) -> Option { + if !has_explicit_session_affinity(client_session_affinity) { + return None; + } let requested_model = requested_model .map(str::trim) .filter(|value| !value.is_empty())?; @@ -41,6 +50,9 @@ pub(crate) fn remember_scheduler_affinity_for_candidate( requested_model: &str, candidate: &SchedulerMinimalCandidateSelectionCandidate, ) { + if !has_explicit_session_affinity(client_session_affinity) { + return; + } remember_scheduler_affinity_for_candidate_at_epoch( state, auth_snapshot, @@ -61,6 +73,9 @@ pub(crate) fn remember_scheduler_affinity_for_candidate_at_epoch( candidate: &SchedulerMinimalCandidateSelectionCandidate, expected_epoch: Option, ) { + if !has_explicit_session_affinity(client_session_affinity) { + return; + } let Some(api_key_id) = auth_snapshot .map(|snapshot| snapshot.api_key_id.trim()) .filter(|value| !value.is_empty()) diff --git a/apps/aether-gateway/src/ai_serving/planner/candidate_materialization.rs b/apps/aether-gateway/src/ai_serving/planner/candidate_materialization.rs index 54892f127..39a070d30 100644 --- a/apps/aether-gateway/src/ai_serving/planner/candidate_materialization.rs +++ b/apps/aether-gateway/src/ai_serving/planner/candidate_materialization.rs @@ -38,12 +38,19 @@ use crate::ai_serving::planner::pool_scheduler::PoolKeyCursor; use crate::ai_serving::planner::runtime_miss::record_local_runtime_candidate_skip_reason; use crate::ai_serving::planner::CandidateFailureDiagnostic; use crate::ai_serving::{GatewayAuthApiKeySnapshot, PlannerAppState}; +use crate::cache::{ + candidate_page_cache_stale_ttl, candidate_page_cache_ttl_from_env, + record_candidate_page_resolve_cache_follower_wait, record_candidate_page_resolve_cache_hit, + record_candidate_page_resolve_cache_load, record_candidate_page_resolve_cache_miss, + CacheLoadObserver, CandidateResolvedPageCacheKey, CandidateResolvedPageSnapshot, +}; use crate::clock::current_unix_ms; use crate::dispatch::refs::dispatch_ref_for_local_candidate; use crate::handlers::shared::provider_pool::admin_provider_pool_config_from_config_value; use crate::orchestration::{local_attempt_slot_count, ExecutionAttemptIdentity}; use crate::scheduler::candidate::is_auth_api_key_concurrency_limit_skip_reason; use crate::scheduler::config::SchedulerSchedulingMode; +use crate::stage_metrics::observe_gateway_stage_ms; use crate::{AppState, GatewayError}; const POOL_KEY_RETRY_INDEX_STRIDE: u32 = 100; @@ -101,16 +108,20 @@ impl<'a> LocalExecutionCandidateAttemptSource<'a> { Self { items } } - pub(crate) async fn next_attempt(&mut self) -> Option { + pub(crate) async fn next_attempt( + &mut self, + ) -> Result, GatewayError> { loop { - let front = self.items.front_mut()?; + let Some(front) = self.items.front_mut() else { + return Ok(None); + }; match front { LocalExecutionCandidateAttemptSourceItem::Static { attempts } => { if let Some(attempt) = next_attempt_from_dispatch_sequence(attempts) { if dispatch_sequence_exhausted(attempts) { self.items.pop_front(); } - return Some(attempt); + return Ok(Some(attempt)); } self.items.pop_front(); } @@ -121,7 +132,7 @@ impl<'a> LocalExecutionCandidateAttemptSource<'a> { pool_exhaustion_persistence, } => { if let Some(attempt) = next_attempt_from_dispatch_sequence(pending_attempts) { - return Some(attempt); + return Ok(Some(attempt)); } let Some(candidate) = cursor.next_key().await else { if let Some(skipped) = cursor.exhausted_group_skipped_candidate() { @@ -146,11 +157,11 @@ impl<'a> LocalExecutionCandidateAttemptSource<'a> { ); } LocalExecutionCandidateAttemptSourceItem::RequestedModelPage { cursor } => { - let Some(attempt) = cursor.next_attempt().await else { + let Some(attempt) = cursor.next_attempt().await? else { self.items.pop_front(); continue; }; - return Some(attempt); + return Ok(Some(attempt)); } } } @@ -726,8 +737,10 @@ where auth_snapshot, routing_policy, client_session_affinity, + request_auth_channel, use_api_format_alias_match, key_mode, + Some(trace_id), ) .await; let mut cursor = RequestedModelAttemptPageCursor { @@ -755,11 +768,14 @@ where remembered_affinity: false, scheduler_cache_affinity_enabled, auth_api_key_concurrency_wait_deadline: None, + deferred_error: None, }; - cursor.load_next_page().await; + if let Err(error) = cursor.load_next_page().await { + cursor.deferred_error = Some(error); + } let candidate_count = cursor.candidate_count; let mut items = VecDeque::new(); - if !cursor.pending_items.is_empty() { + if !cursor.pending_items.is_empty() || cursor.deferred_error.is_some() { items.push_back( LocalExecutionCandidateAttemptSourceItem::RequestedModelPage { cursor: Box::new(cursor), @@ -797,34 +813,58 @@ struct RequestedModelAttemptPageCursor<'a> { remembered_affinity: bool, scheduler_cache_affinity_enabled: bool, auth_api_key_concurrency_wait_deadline: Option, + deferred_error: Option, } impl<'a> RequestedModelAttemptPageCursor<'a> { - async fn next_attempt(&mut self) -> Option { + async fn next_attempt( + &mut self, + ) -> Result, GatewayError> { + if let Some(error) = self.deferred_error.take() { + return Err(error); + } loop { if let Some(attempt) = pop_attempt_from_items(&mut self.pending_items).await { - return Some(attempt); + return Ok(Some(attempt)); } - if !self.load_next_page().await { - return None; + if !self.load_next_page().await? { + return Ok(None); } } } - async fn load_next_page(&mut self) -> bool { + async fn load_next_page(&mut self) -> Result { loop { + let page_started_at = std::time::Instant::now(); let page = match self.page_cursor.next_page().await { Ok(Some(page)) => page, - Ok(None) => return false, + Ok(None) => { + observe_gateway_stage_ms( + "candidate_page_load", + page_started_at.elapsed().as_millis() as u64, + ); + return Ok(false); + } Err(error) => { + observe_gateway_stage_ms( + "candidate_page_load", + page_started_at.elapsed().as_millis() as u64, + ); + if matches!(error, GatewayError::AdmissionTimeout { .. }) { + return Err(error); + } warn!( trace_id = %self.trace_id, error = ?error, "gateway lazy requested-model candidate page read failed" ); - return false; + return Ok(false); } }; + observe_gateway_stage_ms( + "candidate_page_load", + page_started_at.elapsed().as_millis() as u64, + ); if page_is_exact_auth_api_key_concurrency_limited(&page) { if self.wait_for_auth_api_key_concurrency_retry().await { @@ -832,24 +872,16 @@ impl<'a> RequestedModelAttemptPageCursor<'a> { } self.persist_final_auth_api_key_concurrency_skips(page.skipped_candidates) .await; - return false; + return Ok(false); } + let resolve_started_at = std::time::Instant::now(); let (candidates, resolved_skipped) = - resolve_and_rank_logical_local_execution_candidates( - self.state, - page.candidates, - &self.client_api_format, - Some(&self.requested_model), - Some(&self.auth_snapshot), - self.client_session_affinity.as_ref(), - self.required_capabilities.as_ref(), - self.routing_policy.as_ref(), - self.sticky_session_token.as_deref(), - self.request_auth_channel.as_deref(), - self.resolution_mode, - ) - .await; + resolve_priority_candidate_page_with_cache(self, page.candidates).await; + observe_gateway_stage_ms( + "candidate_page_resolve", + resolve_started_at.elapsed().as_millis() as u64, + ); let skipped_candidates = page .skipped_candidates .into_iter() @@ -899,7 +931,7 @@ impl<'a> RequestedModelAttemptPageCursor<'a> { .saturating_add(u32::try_from(skipped_candidate_count).unwrap_or(u32::MAX)); if !items.is_empty() { self.pending_items = items; - return true; + return Ok(true); } let skipped_starting_candidate_index = next_candidate_index; let skipped_persistence = LocalSkippedCandidatePersistenceContext { @@ -1080,6 +1112,120 @@ pub(crate) fn remember_first_local_candidate_affinity( ); } +async fn resolve_priority_candidate_page_with_cache( + cursor: &RequestedModelAttemptPageCursor<'_>, + page_candidates: Vec, +) -> ( + Vec, + Vec, +) { + if !should_cache_resolved_candidate_page(cursor) { + return resolve_and_rank_logical_local_execution_candidates( + cursor.state, + page_candidates, + &cursor.client_api_format, + Some(&cursor.requested_model), + Some(&cursor.auth_snapshot), + cursor.client_session_affinity.as_ref(), + cursor.required_capabilities.as_ref(), + cursor.routing_policy.as_ref(), + cursor.sticky_session_token.as_deref(), + cursor.request_auth_channel.as_deref(), + cursor.resolution_mode, + ) + .await; + } + + let key = CandidateResolvedPageCacheKey::new( + &cursor.requested_model, + &cursor.client_api_format, + true, + &cursor.auth_snapshot, + cursor.required_capabilities.as_ref(), + cursor.routing_policy.as_ref(), + cursor.request_auth_channel.as_deref(), + cursor.state.app().scheduler_affinity_epoch(), + cursor.page_cursor.resolved_page_cache_preselection_mode(), + cursor + .page_cursor + .resolved_page_cache_use_api_format_alias_match(), + cursor.client_session_affinity.as_ref(), + cursor.resolution_mode, + ); + let page_candidates_for_fallback = page_candidates.clone(); + let page_candidates_for_load = page_candidates; + let cache = cursor.state.app().candidate_resolved_page_cache.clone(); + let ttl = candidate_page_cache_ttl_from_env(); + let stale_ttl = candidate_page_cache_stale_ttl(ttl); + let cached = cache + .get_or_load_once_stale_while_refreshing( + key, + ttl, + stale_ttl, + || async move { + let (candidates, resolved_skipped) = + resolve_and_rank_logical_local_execution_candidates( + cursor.state, + page_candidates_for_load, + &cursor.client_api_format, + Some(&cursor.requested_model), + Some(&cursor.auth_snapshot), + cursor.client_session_affinity.as_ref(), + cursor.required_capabilities.as_ref(), + cursor.routing_policy.as_ref(), + cursor.sticky_session_token.as_deref(), + cursor.request_auth_channel.as_deref(), + cursor.resolution_mode, + ) + .await; + Ok::<_, GatewayError>(Some(Arc::new(CandidateResolvedPageSnapshot { + candidates, + resolved_skipped, + }))) + }, + CacheLoadObserver::new() + .on_hit(record_candidate_page_resolve_cache_hit) + .on_miss(record_candidate_page_resolve_cache_miss) + .on_load(record_candidate_page_resolve_cache_load) + .on_follower_wait(record_candidate_page_resolve_cache_follower_wait), + ) + .await + .unwrap_or(None); + + match cached { + Some(snapshot) => ( + snapshot.candidates.clone(), + snapshot.resolved_skipped.clone(), + ), + None => { + if page_candidates_for_fallback.is_empty() { + return (Vec::new(), Vec::new()); + } + resolve_and_rank_logical_local_execution_candidates( + cursor.state, + page_candidates_for_fallback, + &cursor.client_api_format, + Some(&cursor.requested_model), + Some(&cursor.auth_snapshot), + cursor.client_session_affinity.as_ref(), + cursor.required_capabilities.as_ref(), + cursor.routing_policy.as_ref(), + cursor.sticky_session_token.as_deref(), + cursor.request_auth_channel.as_deref(), + cursor.resolution_mode, + ) + .await + } + } +} + +fn should_cache_resolved_candidate_page(cursor: &RequestedModelAttemptPageCursor<'_>) -> bool { + cursor.sticky_session_token.is_none() + && cursor + .page_cursor + .should_cache_current_priority_resolved_page() +} + fn should_persist_available_local_candidate(eligible: &EligibleLocalExecutionCandidate) -> bool { ai_should_persist_available_candidate_for_pool_key(eligible.orchestration.pool_key_index) } @@ -1937,7 +2083,8 @@ mod tests { GatewayDataState::with_request_candidate_repository_for_tests(Arc::clone( &repository, )), - ); + ) + .without_request_candidate_queue_for_tests(); let attempts = persist_available_local_execution_candidates( PlannerAppState::new(&app), @@ -2025,6 +2172,145 @@ mod tests { .is_none()); } + #[tokio::test] + async fn resolved_candidate_page_cache_requires_fixed_order_or_explicit_affinity() { + let app = AppState::new().expect("state should build"); + let auth_snapshot = sample_auth_snapshot(); + let mut page_cursor = LocalCandidatePreselectionPageCursor::new( + PlannerAppState::new(&app), + "openai:chat", + "gpt-5", + true, + None, + &auth_snapshot, + None, + None, + None, + false, + LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModel, + Some("trace-no-session-affinity"), + ) + .await; + page_cursor.mark_priority_page_emitted_for_tests(); + let cursor = RequestedModelAttemptPageCursor { + state: PlannerAppState::new(&app), + trace_id: "trace-no-session-affinity".to_string(), + client_api_format: "openai:chat".to_string(), + requested_model: "gpt-5".to_string(), + auth_snapshot: auth_snapshot.clone(), + client_session_affinity: None, + required_capabilities: None, + routing_policy: None, + sticky_session_token: None, + request_auth_channel: None, + skipped_user_id: "user-1".to_string(), + skipped_api_key_id: "api-key-1".to_string(), + skipped_required_capabilities: None, + skipped_error_context: "test skipped", + record_runtime_miss_diagnostic: false, + resolution_mode: LocalCandidateResolutionMode::Standard, + decorate_skipped_candidate: Arc::new(identity_skipped_candidate), + page_cursor, + pending_items: VecDeque::new(), + candidate_count: 0, + next_candidate_index: 0, + remembered_affinity: false, + scheduler_cache_affinity_enabled: false, + auth_api_key_concurrency_wait_deadline: None, + deferred_error: None, + }; + + assert!(!should_cache_resolved_candidate_page(&cursor)); + + let sticky_cursor = RequestedModelAttemptPageCursor { + sticky_session_token: Some("sticky-token".to_string()), + ..cursor + }; + + assert!(!should_cache_resolved_candidate_page(&sticky_cursor)); + + let mut page_cursor = LocalCandidatePreselectionPageCursor::new( + PlannerAppState::new(&app), + "openai:chat", + "gpt-5", + true, + None, + &auth_snapshot, + None, + Some(&ClientSessionAffinity::from_session_key("chat-session-1")), + None, + false, + LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModel, + Some("trace-session-affinity"), + ) + .await; + page_cursor.mark_priority_page_emitted_for_tests(); + let cursor = RequestedModelAttemptPageCursor { + client_session_affinity: Some(ClientSessionAffinity::from_session_key( + "chat-session-1", + )), + page_cursor, + sticky_session_token: None, + ..sticky_cursor + }; + + assert!(should_cache_resolved_candidate_page(&cursor)); + + let fixed_order_app = AppState::new() + .expect("state should build") + .with_data_state_for_tests( + GatewayDataState::disabled().with_system_config_values_for_tests([( + "scheduling_mode".to_string(), + json!("fixed_order"), + )]), + ); + let mut page_cursor = LocalCandidatePreselectionPageCursor::new( + PlannerAppState::new(&fixed_order_app), + "openai:chat", + "gpt-5", + true, + None, + &auth_snapshot, + None, + None, + None, + false, + LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModel, + Some("trace-fixed-order"), + ) + .await; + page_cursor.mark_priority_page_emitted_for_tests(); + let cursor = RequestedModelAttemptPageCursor { + state: PlannerAppState::new(&fixed_order_app), + trace_id: "trace-fixed-order".to_string(), + client_api_format: "openai:chat".to_string(), + requested_model: "gpt-5".to_string(), + auth_snapshot, + client_session_affinity: None, + required_capabilities: None, + routing_policy: None, + sticky_session_token: None, + request_auth_channel: None, + skipped_user_id: "user-1".to_string(), + skipped_api_key_id: "api-key-1".to_string(), + skipped_required_capabilities: None, + skipped_error_context: "test skipped", + record_runtime_miss_diagnostic: false, + resolution_mode: LocalCandidateResolutionMode::Standard, + decorate_skipped_candidate: Arc::new(identity_skipped_candidate), + page_cursor, + pending_items: VecDeque::new(), + candidate_count: 0, + next_candidate_index: 0, + remembered_affinity: false, + scheduler_cache_affinity_enabled: false, + auth_api_key_concurrency_wait_deadline: None, + deferred_error: None, + }; + + assert!(should_cache_resolved_candidate_page(&cursor)); + } + #[tokio::test] async fn logical_materialization_does_not_persist_pool_group_representative() { let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); @@ -2042,7 +2328,8 @@ mod tests { Arc::clone(&request_candidate_repository), "test-encryption-key", ), - ); + ) + .without_request_candidate_queue_for_tests(); let mut pool_group = sample_eligible("pool-group", None); pool_group.kind = LocalExecutionCandidateKind::PoolGroup; pool_group.transport = sample_transport( @@ -2140,7 +2427,8 @@ mod tests { GatewayDataState::with_request_candidate_repository_for_tests(Arc::clone( &repository, )), - ); + ) + .without_request_candidate_queue_for_tests(); let mut eligible = sample_eligible("ranked-key", None); eligible.ranking = Some(SchedulerRankingOutcome { original_index: 1, @@ -2214,12 +2502,17 @@ mod tests { let first = source .next_attempt() .await + .expect("first attempt read should succeed") .expect("first attempt should be available"); assert_eq!(first.eligible.candidate.key_id, "normal-key"); let remaining = source.drain_static_attempts(); assert!(remaining.is_empty()); - assert!(source.next_attempt().await.is_none()); + assert!(source + .next_attempt() + .await + .expect("remaining attempt read should succeed") + .is_none()); } #[tokio::test] @@ -2239,7 +2532,8 @@ mod tests { Arc::clone(&request_candidate_repository), "test-encryption-key", ), - ); + ) + .without_request_candidate_queue_for_tests(); let mut pool_group = sample_eligible("pool-group", None); pool_group.kind = LocalExecutionCandidateKind::PoolGroup; pool_group.transport = sample_transport( @@ -2275,7 +2569,11 @@ mod tests { }]), }; - assert!(source.next_attempt().await.is_none()); + assert!(source + .next_attempt() + .await + .expect("pool attempt read should succeed") + .is_none()); let stored = app .read_request_candidates_by_request_id("trace-dynamic-pool") @@ -2308,7 +2606,8 @@ mod tests { GatewayDataState::with_request_candidate_repository_for_tests(Arc::clone( &repository, )), - ); + ) + .without_request_candidate_queue_for_tests(); persist_skipped_local_execution_candidates( &app, diff --git a/apps/aether-gateway/src/ai_serving/planner/candidate_ranking.rs b/apps/aether-gateway/src/ai_serving/planner/candidate_ranking.rs index 1dcc615dd..854ad130b 100644 --- a/apps/aether-gateway/src/ai_serving/planner/candidate_ranking.rs +++ b/apps/aether-gateway/src/ai_serving/planner/candidate_ranking.rs @@ -1,5 +1,3 @@ -use std::collections::BTreeMap; - use aether_ai_serving::{ ai_ranking_context, build_ai_rankable_candidate, run_ai_candidate_ranking, AiCandidateRankingPort, AiRankableCandidateParts, AiRankingContextConfig, @@ -7,6 +5,7 @@ use aether_ai_serving::{ }; use aether_routing_core::{ResolvedRoutingPolicy, RoutingSchedulingMode, RoutingSetPriorityMode}; use async_trait::async_trait; +use tokio::sync::Mutex; use tracing::warn; use crate::ai_serving::{GatewayAuthApiKeySnapshot, PlannerAppState}; @@ -24,7 +23,7 @@ use aether_scheduler_core::{ use super::candidate_affinity_cache::read_cached_scheduler_affinity_target; use super::candidate_resolution::{EligibleLocalExecutionCandidate, LocalExecutionCandidateKind}; use super::candidate_transport_ranking_facts::{ - resolve_cached_transport_ranking_facts, CandidateTransportRankingFacts, + resolve_cached_transport_ranking_facts, CandidateTransportRankingFactsCache, }; struct GatewayLocalCandidateRankingPort<'a> { @@ -35,6 +34,7 @@ struct GatewayLocalCandidateRankingPort<'a> { required_capabilities: Option<&'a serde_json::Value>, ordering_config: SchedulerOrderingConfig, routing_policy: Option<&'a ResolvedRoutingPolicy>, + transport_ranking_facts_cache: Mutex, } #[async_trait] @@ -84,13 +84,17 @@ impl AiCandidateRankingPort for GatewayLocalCandidateRankingPort<'_> { normalized_client_api_format: &str, cached_affinity_match: bool, ) -> Result { - let ranking_facts = resolve_transport_ranking_facts_for_candidate( - self.state, - &candidate.candidate, - candidate.transport.as_ref(), - self.ordering_config, - ) - .await; + let ranking_facts = { + let mut cache = self.transport_ranking_facts_cache.lock().await; + resolve_cached_transport_ranking_facts( + self.state, + &mut cache, + &candidate.candidate, + candidate.transport.as_ref(), + self.ordering_config, + ) + .await + }; let routing_overlaid_candidate = routing_overlaid_candidate(self.routing_policy, candidate.kind, &candidate.candidate); Ok(build_ai_rankable_candidate(AiRankableCandidateParts { @@ -137,6 +141,7 @@ pub(crate) async fn rank_eligible_local_execution_candidates( required_capabilities, ordering_config, routing_policy, + transport_ranking_facts_cache: Mutex::new(CandidateTransportRankingFactsCache::default()), }; match run_ai_candidate_ranking(&port, candidates, normalized_client_api_format).await { @@ -145,23 +150,6 @@ pub(crate) async fn rank_eligible_local_execution_candidates( } } -async fn resolve_transport_ranking_facts_for_candidate( - state: PlannerAppState<'_>, - candidate: &SchedulerMinimalCandidateSelectionCandidate, - transport: &crate::ai_serving::GatewayProviderTransportSnapshot, - ordering_config: SchedulerOrderingConfig, -) -> CandidateTransportRankingFacts { - let mut ordering_cache = BTreeMap::new(); - resolve_cached_transport_ranking_facts( - state, - &mut ordering_cache, - candidate, - transport, - ordering_config, - ) - .await -} - fn cached_affinity_matches_local_execution_scope( eligible: &EligibleLocalExecutionCandidate, target: &SchedulerAffinityTarget, @@ -284,7 +272,9 @@ mod tests { use serde_json::json; use super::super::candidate_affinity_cache::remember_scheduler_affinity_for_candidate; - use super::super::candidate_transport_ranking_facts::resolve_cached_candidate_transport_ranking_facts; + use super::super::candidate_transport_ranking_facts::{ + resolve_cached_candidate_transport_ranking_facts, CandidateTransportRankingFactsCache, + }; use super::{PlannerAppState, SchedulerMinimalCandidateSelectionCandidate}; use crate::ai_serving::planner::candidate_resolution::{ resolve_and_rank_local_execution_candidates, @@ -306,7 +296,7 @@ mod tests { let ordering_config = super::read_scheduler_ordering_config_or_default(state).await; let mut candidates = candidates; let mut rankables = Vec::with_capacity(candidates.len()); - let mut ordering_cache = BTreeMap::new(); + let mut ordering_cache = CandidateTransportRankingFactsCache::default(); for (original_index, candidate) in candidates.iter().enumerate() { let ranking_facts = resolve_cached_candidate_transport_ranking_facts( @@ -1540,6 +1530,7 @@ mod tests { .expect("state should build") .with_data_state_for_tests(data_state); let auth_snapshot = sample_auth_snapshot(); + let client_session_affinity = ClientSessionAffinity::from_session_key("session-1"); let cached_candidate = sample_priority_candidate( "provider-cached", "endpoint-cached", @@ -1551,7 +1542,7 @@ mod tests { remember_scheduler_affinity_for_candidate( PlannerAppState::new(&state), Some(&auth_snapshot), - None, + Some(&client_session_affinity), "openai:chat", "gpt-4.1", &cached_candidate, @@ -1573,7 +1564,7 @@ mod tests { "openai:chat", "gpt-4.1", Some(&auth_snapshot), - None, + Some(&client_session_affinity), None, None, None, @@ -1714,6 +1705,7 @@ mod tests { .expect("state should build") .with_data_state_for_tests(data_state); let auth_snapshot = sample_auth_snapshot(); + let client_session_affinity = ClientSessionAffinity::from_session_key("session-1"); let cached_cross_format = sample_priority_candidate( "provider-shared", "endpoint-openai", @@ -1725,7 +1717,7 @@ mod tests { remember_scheduler_affinity_for_candidate( PlannerAppState::new(&state), Some(&auth_snapshot), - None, + Some(&client_session_affinity), "claude:messages", "gpt-4.1", &cached_cross_format, @@ -1747,7 +1739,7 @@ mod tests { "claude:messages", "gpt-4.1", Some(&auth_snapshot), - None, + Some(&client_session_affinity), None, None, None, @@ -1915,6 +1907,7 @@ mod tests { .expect("state should build") .with_data_state_for_tests(data_state); let auth_snapshot = sample_auth_snapshot(); + let client_session_affinity = ClientSessionAffinity::from_session_key("session-1"); let cached_candidate = sample_priority_candidate( "provider-pool", "endpoint-pool", @@ -1926,7 +1919,7 @@ mod tests { remember_scheduler_affinity_for_candidate( PlannerAppState::new(&state), Some(&auth_snapshot), - None, + Some(&client_session_affinity), "openai:chat", "gpt-4.1", &cached_candidate, @@ -1948,7 +1941,7 @@ mod tests { "openai:chat", Some("gpt-4.1"), Some(&auth_snapshot), - None, + Some(&client_session_affinity), None, None, None, @@ -2008,6 +2001,7 @@ mod tests { .expect("state should build") .with_data_state_for_tests(data_state); let auth_snapshot = sample_auth_snapshot(); + let client_session_affinity = ClientSessionAffinity::from_session_key("session-1"); let cached_candidate = sample_priority_candidate( "provider-pool", "endpoint-pool", @@ -2019,7 +2013,7 @@ mod tests { remember_scheduler_affinity_for_candidate( PlannerAppState::new(&state), Some(&auth_snapshot), - None, + Some(&client_session_affinity), "openai:chat", "gpt-4.1", &cached_candidate, @@ -2041,7 +2035,7 @@ mod tests { "openai:chat", Some("gpt-4.1"), Some(&auth_snapshot), - None, + Some(&client_session_affinity), None, None, None, @@ -2064,7 +2058,7 @@ mod tests { } #[tokio::test] - async fn remembers_scheduler_affinity_for_candidate_using_requested_model_key() { + async fn ignores_scheduler_affinity_without_client_session_scope() { let state = AppState::new().expect("state should build"); let auth_snapshot = sample_auth_snapshot(); let candidate = sample_candidate("endpoint-1", "key-1"); @@ -2078,15 +2072,12 @@ mod tests { &candidate, ); - let remembered = state + assert!(state .read_scheduler_affinity_target( "scheduler_affinity:api-key-1:openai:chat:gpt-5", SCHEDULER_AFFINITY_TTL, ) - .expect("affinity target should be cached"); - assert_eq!(remembered.provider_id, "provider-1"); - assert_eq!(remembered.endpoint_id, "endpoint-1"); - assert_eq!(remembered.key_id, "key-1"); + .is_none()); } #[tokio::test] diff --git a/apps/aether-gateway/src/ai_serving/planner/candidate_resolution.rs b/apps/aether-gateway/src/ai_serving/planner/candidate_resolution.rs index aaca79a21..df1b27317 100644 --- a/apps/aether-gateway/src/ai_serving/planner/candidate_resolution.rs +++ b/apps/aether-gateway/src/ai_serving/planner/candidate_resolution.rs @@ -7,6 +7,7 @@ use aether_ai_serving::{ use aether_routing_core::ResolvedRoutingPolicy; use async_trait::async_trait; use std::convert::Infallible; +use std::time::Instant; use tracing::warn; use aether_scheduler_core::{ @@ -20,6 +21,7 @@ use crate::ai_serving::{ PlannerAppState, }; use crate::orchestration::LocalExecutionCandidateMetadata; +use crate::stage_metrics::observe_gateway_stage_ms; use super::candidate_ranking::rank_eligible_local_execution_candidates; @@ -68,7 +70,7 @@ struct GatewayLocalCandidateResolutionPort<'a> { #[async_trait] impl AiCandidateResolutionPort for GatewayLocalCandidateResolutionPort<'_> { type Candidate = SchedulerMinimalCandidateSelectionCandidate; - type Transport = GatewayProviderTransportSnapshot; + type Transport = Arc; type Eligible = EligibleLocalExecutionCandidate; type Skipped = SkippedLocalExecutionCandidate; type Error = Infallible; @@ -77,7 +79,12 @@ impl AiCandidateResolutionPort for GatewayLocalCandidateResolutionPort<'_> { &self, candidate: &Self::Candidate, ) -> Result, Self::Error> { - Ok(read_candidate_transport_snapshot(self.state, candidate).await) + let started_at = Instant::now(); + let transport = read_candidate_transport_snapshot_arc(self.state, candidate).await; + let elapsed_ms = started_at.elapsed().as_millis() as u64; + observe_gateway_stage_ms("candidate_transport_snapshot", elapsed_ms); + observe_gateway_stage_ms("candidate_resolution_transport_read", elapsed_ms); + Ok(transport) } fn build_missing_transport_skipped_candidate( @@ -139,7 +146,7 @@ impl AiCandidateResolutionPort for GatewayLocalCandidateResolutionPort<'_> { SkippedLocalExecutionCandidate { candidate, skip_reason, - transport: Some(Arc::new(transport)), + transport: Some(transport), ranking: None, extra_data: None, } @@ -159,7 +166,7 @@ impl AiCandidateResolutionPort for GatewayLocalCandidateResolutionPort<'_> { EligibleLocalExecutionCandidate { kind, candidate, - transport: Arc::new(transport), + transport, provider_api_format, orchestration: LocalExecutionCandidateMetadata::default(), ranking: None, @@ -171,7 +178,8 @@ impl AiCandidateResolutionPort for GatewayLocalCandidateResolutionPort<'_> { candidates: Vec, normalized_client_api_format: &str, ) -> Result, Self::Error> { - Ok(rank_eligible_local_execution_candidates( + let started_at = Instant::now(); + let ranked = rank_eligible_local_execution_candidates( self.state, candidates, normalized_client_api_format, @@ -181,7 +189,12 @@ impl AiCandidateResolutionPort for GatewayLocalCandidateResolutionPort<'_> { self.required_capabilities, self.routing_policy, ) - .await) + .await; + observe_gateway_stage_ms( + "candidate_resolution_rank", + started_at.elapsed().as_millis() as u64, + ); + Ok(ranked) } async fn apply_pool_scheduler( @@ -358,8 +371,13 @@ async fn resolve_and_rank_local_execution_candidates_with_pool_expansion( expand_pool_groups, }; + let started_at = Instant::now(); match run_ai_candidate_resolution(&port, candidates, request).await { Ok(mut outcome) => { + observe_gateway_stage_ms( + "candidate_resolution_core", + started_at.elapsed().as_millis() as u64, + ); for candidate in &mut outcome.eligible_candidates { candidate.orchestration.scheduler_affinity_epoch = Some(scheduler_affinity_epoch); } @@ -506,8 +524,17 @@ pub(crate) async fn read_candidate_transport_snapshot( state: PlannerAppState<'_>, candidate: &SchedulerMinimalCandidateSelectionCandidate, ) -> Option { + read_candidate_transport_snapshot_arc(state, candidate) + .await + .map(|transport| (*transport).clone()) +} + +pub(crate) async fn read_candidate_transport_snapshot_arc( + state: PlannerAppState<'_>, + candidate: &SchedulerMinimalCandidateSelectionCandidate, +) -> Option> { match state - .read_provider_transport_snapshot( + .read_provider_transport_snapshot_arc( &candidate.provider_id, &candidate.endpoint_id, &candidate.key_id, diff --git a/apps/aether-gateway/src/ai_serving/planner/candidate_source.rs b/apps/aether-gateway/src/ai_serving/planner/candidate_source.rs index 8128a1e8c..09d1413a6 100644 --- a/apps/aether-gateway/src/ai_serving/planner/candidate_source.rs +++ b/apps/aether-gateway/src/ai_serving/planner/candidate_source.rs @@ -3,6 +3,7 @@ use aether_ai_serving::{ }; use aether_data_contracts::repository::candidate_selection::StoredMinimalCandidateSelectionRow; use aether_routing_core::ResolvedRoutingPolicy; +use aether_runtime::ConcurrencyPermit; use aether_scheduler_core::{ enumerate_minimal_candidate_selection_with_model_directives, normalize_api_format, resolve_requested_global_model_name_with_model_directives, @@ -12,16 +13,17 @@ use aether_scheduler_core::{ use async_trait::async_trait; use std::collections::{BTreeMap, BTreeSet, VecDeque}; +use crate::ai_serving::planner::candidate_affinity_cache::has_explicit_session_affinity; use crate::ai_serving::planner::candidate_resolution::SkippedLocalExecutionCandidate; use crate::ai_serving::{GatewayAuthApiKeySnapshot, PlannerAppState}; -use crate::clock::current_unix_secs; +use crate::clock::request_distribution_seed; use crate::data::candidate_selection::{ read_requested_model_rows_fast_path_page, requested_model_candidate_names, MinimalCandidateSelectionRowSource, REQUESTED_MODEL_CANDIDATE_PAGE_SIZE, REQUESTED_MODEL_MAX_SCANNED_ROWS, }; use crate::scheduler::candidate::SchedulerSkippedCandidate; -use crate::scheduler::config::SchedulerOrderingConfig; +use crate::scheduler::config::{SchedulerOrderingConfig, SchedulerSchedulingMode}; use crate::GatewayError; #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -30,6 +32,15 @@ pub(crate) enum LocalCandidatePreselectionKeyMode { ProviderEndpointKeyModelAndApiFormat, } +impl LocalCandidatePreselectionKeyMode { + pub(crate) fn cache_key_name(self) -> &'static str { + match self { + Self::ProviderEndpointKeyModel => "provider_endpoint_key_model", + Self::ProviderEndpointKeyModelAndApiFormat => "provider_endpoint_key_model_api_format", + } + } +} + struct GatewayLocalCandidatePreselectionPort<'a> { state: PlannerAppState<'a>, client_api_format: &'a str, @@ -43,6 +54,7 @@ struct GatewayLocalCandidatePreselectionPort<'a> { key_mode: LocalCandidatePreselectionKeyMode, candidate_api_formats: Vec, model_directive_enabled_api_formats: BTreeSet, + ranking_seed: u64, } #[async_trait] @@ -81,7 +93,7 @@ impl AiCandidatePreselectionPort for GatewayLocalCandidatePreselectionPort<'_> { self.required_capabilities, auth_snapshot, self.client_session_affinity, - current_unix_secs(), + self.ranking_seed, ) .await?; @@ -227,6 +239,7 @@ pub(crate) async fn preselect_local_execution_candidates_for_api_formats_with_se key_mode, candidate_api_formats, model_directive_enabled_api_formats, + ranking_seed: request_distribution_seed(), }; run_ai_candidate_preselection(&port).await @@ -234,6 +247,7 @@ pub(crate) async fn preselect_local_execution_candidates_for_api_formats_with_se pub(crate) struct LocalCandidatePreselectionPageCursor<'a> { state: PlannerAppState<'a>, + trace_id: String, client_api_format: String, requested_model: String, require_streaming: bool, @@ -241,11 +255,13 @@ pub(crate) struct LocalCandidatePreselectionPageCursor<'a> { auth_snapshot: GatewayAuthApiKeySnapshot, routing_policy: Option, client_session_affinity: Option, + request_auth_channel: Option, use_api_format_alias_match: bool, key_mode: LocalCandidatePreselectionKeyMode, candidate_api_formats: Vec, model_directive_enabled_api_formats: BTreeSet, ordering_config: SchedulerOrderingConfig, + ranking_seed: u64, priority_page_emitted: bool, deferred_pages_by_format: BTreeMap< String, @@ -276,8 +292,10 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> { auth_snapshot: &GatewayAuthApiKeySnapshot, routing_policy: Option<&ResolvedRoutingPolicy>, client_session_affinity: Option<&ClientSessionAffinity>, + request_auth_channel: Option<&str>, use_api_format_alias_match: bool, key_mode: LocalCandidatePreselectionKeyMode, + trace_id: Option<&str>, ) -> Self { let candidate_api_formats = crate::ai_serving::request_candidate_api_formats(client_api_format, require_streaming) @@ -307,6 +325,7 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> { Self { state, + trace_id: trace_id.unwrap_or_default().to_string(), client_api_format: client_api_format.to_string(), requested_model: requested_model.to_string(), require_streaming, @@ -314,11 +333,13 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> { auth_snapshot: auth_snapshot.clone(), routing_policy: routing_policy.cloned(), client_session_affinity: client_session_affinity.cloned(), + request_auth_channel: request_auth_channel.map(str::to_string), use_api_format_alias_match, key_mode, candidate_api_formats, model_directive_enabled_api_formats, ordering_config, + ranking_seed: request_distribution_seed(), priority_page_emitted: false, deferred_pages_by_format: BTreeMap::new(), format_index: 0, @@ -344,7 +365,7 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> { > { if !self.priority_page_emitted { self.priority_page_emitted = true; - let priority_page = self.next_priority_page().await?; + let priority_page = self.cached_next_priority_page().await?; if !priority_page.candidates.is_empty() || !priority_page.skipped_candidates.is_empty() { return Ok(Some(priority_page)); @@ -356,7 +377,10 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> { if let Some(outcome) = self.pop_deferred_page(&candidate_api_format) { return Ok(Some(outcome)); } - let Some(outcome) = self.next_page_for_api_format(&candidate_api_format).await? else { + let Some(outcome) = self + .next_page_for_api_format_with_planning_gate(&candidate_api_format) + .await? + else { self.format_index += 1; continue; }; @@ -380,6 +404,70 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> { self.deferred_pages_by_format.clear(); } + pub(crate) fn resolved_page_cache_preselection_mode(&self) -> &'static str { + self.key_mode.cache_key_name() + } + + pub(crate) fn resolved_page_cache_use_api_format_alias_match(&self) -> bool { + self.use_api_format_alias_match + } + + pub(crate) fn should_cache_current_priority_resolved_page(&self) -> bool { + if !(self.priority_page_emitted + && self.format_index == 0 + && self.deferred_pages_by_format.is_empty()) + { + return false; + } + + match self.ordering_config.scheduling_mode { + SchedulerSchedulingMode::FixedOrder => true, + SchedulerSchedulingMode::CacheAffinity => { + has_explicit_session_affinity(self.client_session_affinity.as_ref()) + } + SchedulerSchedulingMode::LoadBalance => false, + } + } + + #[cfg(test)] + pub(crate) fn mark_priority_page_emitted_for_tests(&mut self) { + self.priority_page_emitted = true; + } + + async fn cached_next_priority_page( + &mut self, + ) -> Result< + AiCandidatePreselectionOutcome< + SchedulerMinimalCandidateSelectionCandidate, + SkippedLocalExecutionCandidate, + >, + GatewayError, + > { + let page = self.next_priority_page_with_planning_gate().await?; + self.remember_seen_candidates_from_page(&page); + Ok(page) + } + + fn remember_seen_candidates_from_page( + &mut self, + page: &AiCandidatePreselectionOutcome< + SchedulerMinimalCandidateSelectionCandidate, + SkippedLocalExecutionCandidate, + >, + ) { + for candidate in &page.candidates { + self.seen_candidate_keys + .insert(local_candidate_preselection_key(candidate, self.key_mode)); + } + for skipped_candidate in &page.skipped_candidates { + self.seen_candidate_keys + .insert(local_candidate_preselection_key( + &skipped_candidate.candidate, + self.key_mode, + )); + } + } + async fn next_priority_page( &mut self, ) -> Result< @@ -423,6 +511,35 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> { Ok(priority_page) } + async fn next_priority_page_with_planning_gate( + &mut self, + ) -> Result< + AiCandidatePreselectionOutcome< + SchedulerMinimalCandidateSelectionCandidate, + SkippedLocalExecutionCandidate, + >, + GatewayError, + > { + let _permit = acquire_candidate_planning_gate(self.state, &self.trace_id).await?; + self.next_priority_page().await + } + + async fn next_page_for_api_format_with_planning_gate( + &mut self, + candidate_api_format: &str, + ) -> Result< + Option< + AiCandidatePreselectionOutcome< + SchedulerMinimalCandidateSelectionCandidate, + SkippedLocalExecutionCandidate, + >, + >, + GatewayError, + > { + let _permit = acquire_candidate_planning_gate(self.state, &self.trace_id).await?; + self.next_page_for_api_format(candidate_api_format).await + } + async fn split_priority_conversion_page( &self, candidate_api_format: &str, @@ -797,7 +914,7 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> { self.required_capabilities.as_ref(), auth_snapshot, self.client_session_affinity.as_ref(), - current_unix_secs(), + self.ranking_seed, ) .await?; let skipped_candidates = skipped_candidates @@ -894,6 +1011,28 @@ fn local_candidate_preselection_key( } } +async fn acquire_candidate_planning_gate( + state: PlannerAppState<'_>, + trace_id: &str, +) -> Result, GatewayError> { + let Some(gate) = state.app().candidate_planning_gate.as_ref() else { + return Ok(None); + }; + let budget = state + .app() + .frontdoor_runtime_guards + .internal_gate_queue_budget; + match tokio::time::timeout(budget, gate.acquire()).await { + Ok(Ok(permit)) => Ok(Some(permit)), + Ok(Err(err)) => Err(GatewayError::Internal(err.to_string())), + Err(_) => Err(GatewayError::AdmissionTimeout { + trace_id: trace_id.to_string(), + gate: "gateway_candidate_planning", + queue_budget_ms: budget.as_millis() as u64, + }), + } +} + fn matches_client_api_format( use_api_format_alias_match: bool, candidate_api_format: &str, @@ -1220,8 +1359,10 @@ mod tests { &auth_snapshot, None, None, + None, true, LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModelAndApiFormat, + None, ) .await; @@ -1277,8 +1418,10 @@ mod tests { &auth_snapshot, None, None, + None, true, LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModelAndApiFormat, + None, ) .await; @@ -1349,8 +1492,10 @@ mod tests { &auth_snapshot, None, None, + None, true, LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModelAndApiFormat, + None, ) .await; diff --git a/apps/aether-gateway/src/ai_serving/planner/candidate_transport_ranking_facts.rs b/apps/aether-gateway/src/ai_serving/planner/candidate_transport_ranking_facts.rs index 1d2991e55..29de32135 100644 --- a/apps/aether-gateway/src/ai_serving/planner/candidate_transport_ranking_facts.rs +++ b/apps/aether-gateway/src/ai_serving/planner/candidate_transport_ranking_facts.rs @@ -1,8 +1,10 @@ use std::collections::BTreeMap; +use aether_contracts::ProxySnapshot; use aether_scheduler_core::{ SchedulerMinimalCandidateSelectionCandidate, SchedulerTunnelAffinityBucket, }; +use serde_json::Value; use tracing::warn; use crate::ai_serving::{GatewayProviderTransportSnapshot, PlannerAppState}; @@ -10,7 +12,9 @@ use crate::scheduler::config::SchedulerOrderingConfig; use super::candidate_resolution::read_candidate_transport_snapshot; -pub(super) type CandidateTransportIdentity<'a> = (&'a str, &'a str, &'a str); +const TUNNEL_OWNER_INSTANCE_ID_EXTRA_KEY: &str = "tunnel_owner_instance_id"; + +pub(super) type CandidateTransportIdentity = (String, String, String); #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(super) struct CandidateTransportRankingFacts { @@ -18,38 +22,51 @@ pub(super) struct CandidateTransportRankingFacts { pub(super) keep_priority_on_conversion: bool, } -pub(super) async fn resolve_cached_candidate_transport_ranking_facts<'a>( - state: PlannerAppState<'_>, - cache: &mut BTreeMap, CandidateTransportRankingFacts>, - candidate: &'a SchedulerMinimalCandidateSelectionCandidate, - ordering_config: SchedulerOrderingConfig, -) -> CandidateTransportRankingFacts { - let identity = candidate_transport_identity(candidate); - if let Some(facts) = cache.get(&identity).copied() { - return facts; - } - - let facts = resolve_candidate_transport_ranking_facts(state, candidate, ordering_config).await; - cache.insert(identity, facts); - facts +#[derive(Debug, Default)] +pub(super) struct CandidateTransportRankingFactsCache { + candidate_facts: BTreeMap, + configured_proxy_snapshots: BTreeMap>, + system_proxy_snapshot: Option>, + tunnel_buckets_by_node_id: BTreeMap, } -pub(super) async fn resolve_cached_transport_ranking_facts<'a>( +pub(super) async fn resolve_cached_candidate_transport_ranking_facts( state: PlannerAppState<'_>, - cache: &mut BTreeMap, CandidateTransportRankingFacts>, - candidate: &'a SchedulerMinimalCandidateSelectionCandidate, - transport: &GatewayProviderTransportSnapshot, + cache: &mut CandidateTransportRankingFactsCache, + candidate: &SchedulerMinimalCandidateSelectionCandidate, ordering_config: SchedulerOrderingConfig, ) -> CandidateTransportRankingFacts { let identity = candidate_transport_identity(candidate); - if let Some(facts) = cache.get(&identity).copied() { + if let Some(facts) = cache.candidate_facts.get(&identity).copied() { return facts; } let facts = - resolve_candidate_transport_ranking_facts_from_transport(state, transport, ordering_config) - .await; - cache.insert(identity, facts); + resolve_candidate_transport_ranking_facts(state, cache, candidate, ordering_config).await; + cache.candidate_facts.insert(identity, facts); + facts +} + +pub(super) async fn resolve_cached_transport_ranking_facts( + state: PlannerAppState<'_>, + cache: &mut CandidateTransportRankingFactsCache, + candidate: &SchedulerMinimalCandidateSelectionCandidate, + transport: &GatewayProviderTransportSnapshot, + ordering_config: SchedulerOrderingConfig, +) -> CandidateTransportRankingFacts { + let identity = candidate_transport_identity(candidate); + if let Some(facts) = cache.candidate_facts.get(&identity).copied() { + return facts; + } + + let facts = resolve_candidate_transport_ranking_facts_from_transport( + state, + cache, + transport, + ordering_config, + ) + .await; + cache.candidate_facts.insert(identity, facts); facts } @@ -58,13 +75,15 @@ pub(super) async fn candidate_keeps_priority_on_conversion( candidate: &SchedulerMinimalCandidateSelectionCandidate, ordering_config: SchedulerOrderingConfig, ) -> bool { - resolve_candidate_transport_ranking_facts(state, candidate, ordering_config) + let mut cache = CandidateTransportRankingFactsCache::default(); + resolve_candidate_transport_ranking_facts(state, &mut cache, candidate, ordering_config) .await .keep_priority_on_conversion } async fn resolve_candidate_transport_ranking_facts( state: PlannerAppState<'_>, + cache: &mut CandidateTransportRankingFactsCache, candidate: &SchedulerMinimalCandidateSelectionCandidate, ordering_config: SchedulerOrderingConfig, ) -> CandidateTransportRankingFacts { @@ -75,17 +94,23 @@ async fn resolve_candidate_transport_ranking_facts( }; }; - resolve_candidate_transport_ranking_facts_from_transport(state, &transport, ordering_config) - .await + resolve_candidate_transport_ranking_facts_from_transport( + state, + cache, + &transport, + ordering_config, + ) + .await } async fn resolve_candidate_transport_ranking_facts_from_transport( state: PlannerAppState<'_>, + cache: &mut CandidateTransportRankingFactsCache, transport: &GatewayProviderTransportSnapshot, ordering_config: SchedulerOrderingConfig, ) -> CandidateTransportRankingFacts { CandidateTransportRankingFacts { - tunnel_bucket: resolve_tunnel_owner_affinity_from_transport(state, transport).await, + tunnel_bucket: resolve_tunnel_owner_affinity_from_transport(state, cache, transport).await, keep_priority_on_conversion: ordering_config.keep_priority_on_conversion || transport.provider.keep_priority_on_conversion, } @@ -93,12 +118,11 @@ async fn resolve_candidate_transport_ranking_facts_from_transport( async fn resolve_tunnel_owner_affinity_from_transport( state: PlannerAppState<'_>, + cache: &mut CandidateTransportRankingFactsCache, transport: &GatewayProviderTransportSnapshot, ) -> SchedulerTunnelAffinityBucket { - let Some(proxy) = state - .app() - .resolve_transport_proxy_snapshot_with_tunnel_affinity(transport) - .await + let Some(proxy) = + resolve_transport_proxy_snapshot_with_tunnel_affinity_cached(state, cache, transport).await else { return SchedulerTunnelAffinityBucket::Neutral; }; @@ -114,10 +138,74 @@ async fn resolve_tunnel_owner_affinity_from_transport( return SchedulerTunnelAffinityBucket::Neutral; }; + if let Some(bucket) = cache.tunnel_buckets_by_node_id.get(node_id).copied() { + return bucket; + } + + let bucket = resolve_tunnel_owner_affinity_from_proxy(state, &proxy, node_id).await; + cache + .tunnel_buckets_by_node_id + .insert(node_id.to_string(), bucket); + bucket +} + +async fn resolve_transport_proxy_snapshot_with_tunnel_affinity_cached( + state: PlannerAppState<'_>, + cache: &mut CandidateTransportRankingFactsCache, + transport: &GatewayProviderTransportSnapshot, +) -> Option { + for raw in [ + transport.key.proxy.as_ref(), + transport.endpoint.proxy.as_ref(), + transport.provider.proxy.as_ref(), + ] + .into_iter() + .flatten() + { + let cache_key = proxy_config_cache_key(raw); + if let Some(snapshot) = cache.configured_proxy_snapshots.get(&cache_key) { + if snapshot.is_some() { + return snapshot.clone(); + } + continue; + } + let snapshot = state + .app() + .resolve_configured_proxy_snapshot_with_tunnel_affinity(Some(raw)) + .await; + cache + .configured_proxy_snapshots + .insert(cache_key, snapshot.clone()); + if snapshot.is_some() { + return snapshot; + } + } + + if let Some(snapshot) = cache.system_proxy_snapshot.as_ref() { + return snapshot.clone(); + } + let snapshot = state.app().resolve_system_proxy_snapshot().await; + cache.system_proxy_snapshot = Some(snapshot.clone()); + snapshot +} + +async fn resolve_tunnel_owner_affinity_from_proxy( + state: PlannerAppState<'_>, + proxy: &ProxySnapshot, + node_id: &str, +) -> SchedulerTunnelAffinityBucket { if state.app().tunnel.has_local_proxy(node_id) { return SchedulerTunnelAffinityBucket::LocalTunnel; } + if let Some(owner_instance_id) = proxy_tunnel_owner_instance_id(proxy) { + return if owner_instance_id == state.app().tunnel.local_instance_id() { + SchedulerTunnelAffinityBucket::LocalTunnel + } else { + SchedulerTunnelAffinityBucket::RemoteTunnel + }; + } + match state .app() .tunnel @@ -144,10 +232,25 @@ async fn resolve_tunnel_owner_affinity_from_transport( fn candidate_transport_identity( candidate: &SchedulerMinimalCandidateSelectionCandidate, -) -> CandidateTransportIdentity<'_> { +) -> CandidateTransportIdentity { ( - candidate.provider_id.as_str(), - candidate.endpoint_id.as_str(), - candidate.key_id.as_str(), + candidate.provider_id.clone(), + candidate.endpoint_id.clone(), + candidate.key_id.clone(), ) } + +fn proxy_config_cache_key(raw: &Value) -> String { + serde_json::to_string(raw).unwrap_or_else(|_| raw.to_string()) +} + +fn proxy_tunnel_owner_instance_id(proxy: &ProxySnapshot) -> Option<&str> { + proxy + .extra + .as_ref() + .and_then(Value::as_object) + .and_then(|extra| extra.get(TUNNEL_OWNER_INSTANCE_ID_EXTRA_KEY)) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) +} diff --git a/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/build.rs b/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/build.rs index 964cc2616..cfdd9bbcf 100644 --- a/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/build.rs +++ b/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/build.rs @@ -66,7 +66,7 @@ pub(crate) async fn maybe_build_sync_local_same_format_provider_decision_payload candidate_count, ); - while let Some(attempt) = source.next_attempt().await { + while let Some(attempt) = source.next_attempt().await? { if let Some(payload) = maybe_build_local_same_format_provider_decision_payload_for_candidate( state, parts, trace_id, body_json, &input, attempt, spec, @@ -134,7 +134,7 @@ pub(crate) async fn maybe_build_stream_local_same_format_provider_decision_paylo candidate_count, ); - while let Some(attempt) = source.next_attempt().await { + while let Some(attempt) = source.next_attempt().await? { if let Some(payload) = maybe_build_local_same_format_provider_decision_payload_for_candidate( state, parts, trace_id, body_json, &input, attempt, spec, diff --git a/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/plans.rs b/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/plans.rs index 8cf5d0269..42c196222 100644 --- a/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/plans.rs +++ b/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/plans.rs @@ -190,7 +190,7 @@ pub(crate) async fn build_local_stream_attempt_source<'a>( #[async_trait] impl LocalExecutionAttemptSource for LocalSameFormatProviderSyncAttemptSource<'_> { async fn next_execution_attempt(&mut self) -> Result, GatewayError> { - while let Some(attempt) = self.candidates.next_attempt().await { + while let Some(attempt) = self.candidates.next_attempt().await? { match self.build_sync_attempt(attempt).await? { Some(attempt) => return Ok(Some(attempt)), None => continue, @@ -220,7 +220,7 @@ impl LocalExecutionAttemptSource for LocalSameFormatProviderStreamAttemptSource<'_> { async fn next_execution_attempt(&mut self) -> Result, GatewayError> { - while let Some(attempt) = self.candidates.next_attempt().await { + while let Some(attempt) = self.candidates.next_attempt().await? { match self.build_stream_attempt(attempt).await? { Some(attempt) => return Ok(Some(attempt)), None => continue, @@ -372,7 +372,7 @@ pub(crate) async fn build_local_sync_plan_and_reports( } let mut plans = Vec::new(); - while let Some(attempt) = source.next_attempt().await { + while let Some(attempt) = source.next_attempt().await? { let Some(payload) = maybe_build_local_same_format_provider_decision_payload_for_candidate( state, parts, trace_id, body_json, &input, attempt, spec, ) @@ -458,7 +458,7 @@ pub(crate) async fn build_local_stream_plan_and_reports( } let mut plans = Vec::new(); - while let Some(attempt) = source.next_attempt().await { + while let Some(attempt) = source.next_attempt().await? { let Some(payload) = maybe_build_local_same_format_provider_decision_payload_for_candidate( state, parts, trace_id, body_json, &input, attempt, spec, ) diff --git a/apps/aether-gateway/src/ai_serving/planner/redaction.rs b/apps/aether-gateway/src/ai_serving/planner/redaction.rs index 5c40306b8..55506929f 100644 --- a/apps/aether-gateway/src/ai_serving/planner/redaction.rs +++ b/apps/aether-gateway/src/ai_serving/planner/redaction.rs @@ -1,5 +1,5 @@ use std::borrow::Cow; -use std::time::{SystemTime, UNIX_EPOCH}; +use std::time::{Instant, SystemTime, UNIX_EPOCH}; use serde_json::Value; use tracing::warn; @@ -7,9 +7,11 @@ use tracing::warn; use crate::ai_serving::ExecutionRuntimeAuthContext; use crate::privacy::{ build_redaction_session_config, read_chat_pii_redaction_runtime_config, - try_mask_chat_pii_request_json_with_cache_options, ChatPiiRedactionRequestFormat, - MaskChatRequestOptions, RedactionMaskError, RedactionSessionSlot, RedisRedactionMappingCache, + try_mask_chat_pii_request_value_with_cache_options, CachedRequestRedaction, + ChatPiiRedactionRequestFormat, MaskChatRequestOptions, RedactionMaskError, + RedactionSessionSlot, RedisRedactionMappingCache, }; +use crate::stage_metrics::observe_gateway_stage_ms; use crate::{AppState, GatewayError}; pub(crate) struct ProviderRequestRedaction<'a> { @@ -73,6 +75,18 @@ pub(crate) async fn resolve_provider_chat_pii_redaction<'a>( let Some(slot) = parts.extensions.get::() else { return Ok(ProviderRequestRedaction::disabled(body_json)); }; + let request_cache_key = request_redaction_cache_key(format, body_json); + if let Some(cached) = slot.cached_request_redaction(&request_cache_key) { + observe_gateway_stage_ms("chat_pii_redaction_request_cache_hit", 0); + return Ok(provider_redaction_from_cached( + slot, + candidate_id, + body_json, + cached, + )); + } + + let runtime_config_started_at = Instant::now(); let runtime_config = read_chat_pii_redaction_runtime_config(state) .await .map_err(|err| { @@ -82,11 +96,22 @@ pub(crate) async fn resolve_provider_chat_pii_redaction<'a>( ); GatewayError::Internal("chat pii redaction setup failed".to_string()) })?; + observe_gateway_stage_ms( + "chat_pii_redaction_runtime_config", + runtime_config_started_at.elapsed().as_millis() as u64, + ); if !runtime_config.enabled { + slot.put_cached_request_redaction(request_cache_key, CachedRequestRedaction::unredacted()); return Ok(ProviderRequestRedaction::disabled(body_json)); } + let feature_settings_started_at = Instant::now(); let feature_settings = resolve_chat_pii_redaction_feature_settings(state, auth_context).await?; + observe_gateway_stage_ms( + "chat_pii_redaction_feature_settings", + feature_settings_started_at.elapsed().as_millis() as u64, + ); if !feature_settings.effective_enabled() { + slot.put_cached_request_redaction(request_cache_key, CachedRequestRedaction::unredacted()); return Ok(ProviderRequestRedaction::disabled(body_json)); } let Some(hmac_key) = state.encryption_key().map(str::as_bytes).map(Vec::from) else { @@ -95,20 +120,14 @@ pub(crate) async fn resolve_provider_chat_pii_redaction<'a>( "chat pii redaction setup failed".to_string(), )); }; - let body_bytes = serde_json::to_vec(body_json).map_err(|err| { - warn!( - error = ?err, - "gateway failed to serialize provider chat pii redaction body" - ); - GatewayError::Internal("chat pii redaction setup failed".to_string()) - })?; let now_unix_secs = SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap_or_default() .as_secs(); let cache = RedisRedactionMappingCache::new(state.runtime_state.as_ref()); - let masked = try_mask_chat_pii_request_json_with_cache_options( - &body_bytes, + let mask_started_at = Instant::now(); + let masked = try_mask_chat_pii_request_value_with_cache_options( + body_json, format, build_redaction_session_config(hmac_key, &runtime_config, now_unix_secs), MaskChatRequestOptions::runtime(), @@ -116,19 +135,27 @@ pub(crate) async fn resolve_provider_chat_pii_redaction<'a>( ) .await .map_err(redaction_mask_error_to_gateway_error)?; + observe_gateway_stage_ms( + "chat_pii_redaction_mask_body", + mask_started_at.elapsed().as_millis() as u64, + ); if !masked.redacted { + slot.put_cached_request_redaction(request_cache_key, CachedRequestRedaction::unredacted()); return Ok(ProviderRequestRedaction { body_json: Cow::Borrowed(body_json), redacted: false, }); } - let masked_body_json = serde_json::from_slice::(&masked.body).map_err(|err| { - warn!( - error = ?err, - "gateway failed to decode redacted provider chat pii body" - ); - GatewayError::Internal("chat pii redaction setup failed".to_string()) - })?; + let Some(masked_body_json) = masked.body_json else { + warn!("gateway pii redaction reported redacted without masked body"); + return Err(GatewayError::Internal( + "chat pii redaction setup failed".to_string(), + )); + }; + slot.put_cached_request_redaction( + request_cache_key, + CachedRequestRedaction::redacted(masked_body_json.clone(), masked.session.clone()), + ); slot.put_for_candidate(candidate_id, masked.session); Ok(ProviderRequestRedaction { body_json: Cow::Owned(masked_body_json), @@ -136,31 +163,46 @@ pub(crate) async fn resolve_provider_chat_pii_redaction<'a>( }) } +fn request_redaction_cache_key(format: ChatPiiRedactionRequestFormat, body_json: &Value) -> String { + format!("{format:?}:{:p}", body_json) +} + +fn provider_redaction_from_cached<'a>( + slot: &RedactionSessionSlot, + candidate_id: &str, + body_json: &'a Value, + cached: CachedRequestRedaction, +) -> ProviderRequestRedaction<'a> { + if !cached.redacted { + return ProviderRequestRedaction::disabled(body_json); + } + let Some(masked_body_json) = cached.body_json else { + return ProviderRequestRedaction::disabled(body_json); + }; + if let Some(session) = cached.session { + slot.put_for_candidate(candidate_id, session); + } + ProviderRequestRedaction { + body_json: Cow::Owned(masked_body_json), + redacted: true, + } +} + async fn resolve_chat_pii_redaction_feature_settings( state: &AppState, auth_context: &ExecutionRuntimeAuthContext, ) -> Result { - let user_settings = state - .read_user_feature_settings(&auth_context.user_id) - .await + let user_settings_fut = state.read_user_feature_settings(&auth_context.user_id); + let key_settings_fut = state.read_auth_api_key_feature_settings( + &auth_context.user_id, + &auth_context.api_key_id, + auth_context.api_key_is_standalone, + ); + let (user_settings, key_settings) = tokio::try_join!(user_settings_fut, key_settings_fut) .map_err(|err| { warn!( error = ?err, - "gateway failed to read user chat pii redaction feature settings" - ); - GatewayError::Internal("chat pii redaction setup failed".to_string()) - })?; - let key_settings = state - .read_auth_api_key_feature_settings( - &auth_context.user_id, - &auth_context.api_key_id, - auth_context.api_key_is_standalone, - ) - .await - .map_err(|err| { - warn!( - error = ?err, - "gateway failed to read api key chat pii redaction feature settings" + "gateway failed to read chat pii redaction feature settings" ); GatewayError::Internal("chat pii redaction setup failed".to_string()) })?; diff --git a/apps/aether-gateway/src/ai_serving/planner/specialized/files.rs b/apps/aether-gateway/src/ai_serving/planner/specialized/files.rs index cc5c8335a..d08aab594 100644 --- a/apps/aether-gateway/src/ai_serving/planner/specialized/files.rs +++ b/apps/aether-gateway/src/ai_serving/planner/specialized/files.rs @@ -175,7 +175,7 @@ pub(crate) async fn build_local_gemini_files_stream_attempt_source_for_kind<'a>( #[async_trait] impl LocalExecutionAttemptSource for LocalGeminiFilesSyncAttemptSource<'_> { async fn next_execution_attempt(&mut self) -> Result, GatewayError> { - while let Some(attempt) = self.candidates.next_attempt().await { + while let Some(attempt) = self.candidates.next_attempt().await? { match self.build_sync_attempt(attempt).await? { Some(attempt) => return Ok(Some(attempt)), None => continue, @@ -198,7 +198,7 @@ impl LocalExecutionAttemptSource for LocalGeminiFilesSyncAttemptS #[async_trait] impl LocalExecutionAttemptSource for LocalGeminiFilesStreamAttemptSource<'_> { async fn next_execution_attempt(&mut self) -> Result, GatewayError> { - while let Some(attempt) = self.candidates.next_attempt().await { + while let Some(attempt) = self.candidates.next_attempt().await? { match self.build_stream_attempt(attempt).await? { Some(attempt) => return Ok(Some(attempt)), None => continue, @@ -323,7 +323,7 @@ pub(crate) async fn maybe_build_sync_local_gemini_files_decision_payload( let (mut source, _) = build_local_gemini_files_candidate_attempt_source(state, trace_id, &input).await?; - while let Some(attempt) = source.next_attempt().await { + while let Some(attempt) = source.next_attempt().await? { if let Some(payload) = maybe_build_local_gemini_files_decision_payload_for_candidate( state, parts, @@ -365,7 +365,7 @@ pub(crate) async fn maybe_build_stream_local_gemini_files_decision_payload( build_local_gemini_files_candidate_attempt_source(state, trace_id, &input).await?; let empty_body_json = serde_json::Value::Null; - while let Some(attempt) = source.next_attempt().await { + while let Some(attempt) = source.next_attempt().await? { if let Some(payload) = maybe_build_local_gemini_files_decision_payload_for_candidate( state, parts, @@ -414,7 +414,7 @@ async fn build_local_sync_plan_and_reports( build_local_gemini_files_candidate_attempt_source(state, trace_id, &input).await?; let mut plans = Vec::new(); - while let Some(attempt) = source.next_attempt().await { + while let Some(attempt) = source.next_attempt().await? { let Some(payload) = maybe_build_local_gemini_files_decision_payload_for_candidate( state, parts, @@ -467,7 +467,7 @@ async fn build_local_stream_plan_and_reports( let mut plans = Vec::new(); let empty_body_json = serde_json::Value::Null; - while let Some(attempt) = source.next_attempt().await { + while let Some(attempt) = source.next_attempt().await? { let Some(payload) = maybe_build_local_gemini_files_decision_payload_for_candidate( state, parts, diff --git a/apps/aether-gateway/src/ai_serving/planner/specialized/image.rs b/apps/aether-gateway/src/ai_serving/planner/specialized/image.rs index 12a0fdb77..12d504057 100644 --- a/apps/aether-gateway/src/ai_serving/planner/specialized/image.rs +++ b/apps/aether-gateway/src/ai_serving/planner/specialized/image.rs @@ -253,7 +253,7 @@ pub(crate) async fn build_local_image_stream_attempt_source_for_kind<'a>( #[async_trait] impl LocalExecutionAttemptSource for LocalOpenAiImageSyncAttemptSource<'_> { async fn next_execution_attempt(&mut self) -> Result, GatewayError> { - while let Some(attempt) = self.candidates.next_attempt().await { + while let Some(attempt) = self.candidates.next_attempt().await? { match self.build_sync_attempt(attempt).await? { Some(attempt) => return Ok(Some(attempt)), None => continue, @@ -276,7 +276,7 @@ impl LocalExecutionAttemptSource for LocalOpenAiImageSyncAttemptS #[async_trait] impl LocalExecutionAttemptSource for LocalOpenAiImageStreamAttemptSource<'_> { async fn next_execution_attempt(&mut self) -> Result, GatewayError> { - while let Some(attempt) = self.candidates.next_attempt().await { + while let Some(attempt) = self.candidates.next_attempt().await? { match self.build_stream_attempt(attempt).await? { Some(attempt) => return Ok(Some(attempt)), None => continue, @@ -421,7 +421,7 @@ pub(crate) async fn maybe_build_sync_local_image_decision_payload( return Ok(None); }; - while let Some(attempt) = source.next_attempt().await { + while let Some(attempt) = source.next_attempt().await? { if let Some(payload) = maybe_build_local_openai_image_decision_payload_for_candidate( state, parts, @@ -482,7 +482,7 @@ pub(crate) async fn maybe_build_stream_local_image_decision_payload( return Ok(None); }; - while let Some(attempt) = source.next_attempt().await { + while let Some(attempt) = source.next_attempt().await? { if let Some(payload) = maybe_build_local_openai_image_decision_payload_for_candidate( state, parts, @@ -540,7 +540,7 @@ async fn build_local_sync_plan_and_reports( }; let mut plans = Vec::new(); - while let Some(attempt) = source.next_attempt().await { + while let Some(attempt) = source.next_attempt().await? { let Some(payload) = maybe_build_local_openai_image_decision_payload_for_candidate( state, parts, @@ -617,7 +617,7 @@ async fn build_local_stream_plan_and_reports( }; let mut plans = Vec::new(); - while let Some(attempt) = source.next_attempt().await { + while let Some(attempt) = source.next_attempt().await? { let Some(payload) = maybe_build_local_openai_image_decision_payload_for_candidate( state, parts, diff --git a/apps/aether-gateway/src/ai_serving/planner/specialized/video.rs b/apps/aether-gateway/src/ai_serving/planner/specialized/video.rs index f4669c09f..70184da79 100644 --- a/apps/aether-gateway/src/ai_serving/planner/specialized/video.rs +++ b/apps/aether-gateway/src/ai_serving/planner/specialized/video.rs @@ -105,7 +105,7 @@ pub(crate) async fn build_local_video_sync_attempt_source_for_kind<'a>( #[async_trait] impl LocalExecutionAttemptSource for LocalVideoCreateSyncAttemptSource<'_> { async fn next_execution_attempt(&mut self) -> Result, GatewayError> { - while let Some(attempt) = self.candidates.next_attempt().await { + while let Some(attempt) = self.candidates.next_attempt().await? { match self.build_sync_attempt(attempt).await? { Some(attempt) => return Ok(Some(attempt)), None => continue, @@ -195,7 +195,7 @@ pub(crate) async fn maybe_build_sync_local_video_decision_payload( return Ok(None); }; - while let Some(attempt) = source.next_attempt().await { + while let Some(attempt) = source.next_attempt().await? { if let Some(payload) = maybe_build_local_video_create_decision_payload_for_candidate( state, parts, body_json, trace_id, &input, attempt, spec, ) @@ -240,7 +240,7 @@ async fn build_local_sync_plan_and_reports( }; let mut plans = Vec::new(); - while let Some(attempt) = source.next_attempt().await { + while let Some(attempt) = source.next_attempt().await? { let Some(payload) = maybe_build_local_video_create_decision_payload_for_candidate( state, parts, body_json, trace_id, &input, attempt, spec, ) diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/family/build.rs b/apps/aether-gateway/src/ai_serving/planner/standard/family/build.rs index 7997f4760..a52a73c2c 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/family/build.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/family/build.rs @@ -178,7 +178,7 @@ pub(crate) async fn build_local_stream_attempt_source<'a>( #[async_trait] impl LocalExecutionAttemptSource for LocalStandardSyncAttemptSource<'_> { async fn next_execution_attempt(&mut self) -> Result, GatewayError> { - while let Some(attempt) = self.candidates.next_attempt().await { + while let Some(attempt) = self.candidates.next_attempt().await? { match self.build_sync_attempt(attempt).await? { Some(attempt) => return Ok(Some(attempt)), None => continue, @@ -206,7 +206,7 @@ impl LocalExecutionAttemptSource for LocalStandardSyncAttemptSour #[async_trait] impl LocalExecutionAttemptSource for LocalStandardStreamAttemptSource<'_> { async fn next_execution_attempt(&mut self) -> Result, GatewayError> { - while let Some(attempt) = self.candidates.next_attempt().await { + while let Some(attempt) = self.candidates.next_attempt().await? { match self.build_stream_attempt(attempt).await? { Some(attempt) => return Ok(Some(attempt)), None => continue, @@ -340,7 +340,7 @@ pub(crate) async fn maybe_build_sync_via_standard_family_payload( .await?; apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count); - while let Some(attempt) = source.next_attempt().await { + while let Some(attempt) = source.next_attempt().await? { if let Some(payload) = maybe_build_local_standard_decision_payload_for_candidate( state, parts, trace_id, body_json, &input, attempt, spec, ) @@ -390,7 +390,7 @@ pub(crate) async fn maybe_build_stream_via_standard_family_payload( .await?; apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count); - while let Some(attempt) = source.next_attempt().await { + while let Some(attempt) = source.next_attempt().await? { if let Some(payload) = maybe_build_local_standard_decision_payload_for_candidate( state, parts, trace_id, body_json, &input, attempt, spec, ) @@ -449,7 +449,7 @@ pub(crate) async fn build_local_sync_plan_and_reports( return Ok(Vec::new()); } let mut plans = Vec::new(); - while let Some(attempt) = source.next_attempt().await { + while let Some(attempt) = source.next_attempt().await? { let Some(payload) = maybe_build_local_standard_decision_payload_for_candidate( state, parts, trace_id, body_json, &input, attempt, spec, ) @@ -524,7 +524,7 @@ pub(crate) async fn build_local_stream_plan_and_reports( return Ok(Vec::new()); } let mut plans = Vec::new(); - while let Some(attempt) = source.next_attempt().await { + while let Some(attempt) = source.next_attempt().await? { let Some(payload) = maybe_build_local_standard_decision_payload_for_candidate( state, parts, trace_id, body_json, &input, attempt, spec, ) diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/payload.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/payload.rs index 4efda97be..cdda7a5f7 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/payload.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/payload.rs @@ -12,6 +12,7 @@ use crate::ai_serving::planner::{ use crate::ai_serving::transport::{ resolve_transport_execution_timeouts, resolve_transport_profile, }; +use crate::stage_metrics::observe_gateway_stage_ms; use crate::{ append_execution_contract_fields_to_value, append_local_failover_policy_to_value, AiExecutionDecision, AppState, GatewayError, @@ -40,6 +41,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate candidate_id, .. } = attempt; + let payload_started_at = std::time::Instant::now(); let Some(resolved) = resolve_local_openai_chat_candidate_payload_parts( state, parts, @@ -55,8 +57,16 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate ) .await? else { + observe_gateway_stage_ms( + "stream_candidate_payload_parts", + payload_started_at.elapsed().as_millis() as u64, + ); return Ok(None); }; + observe_gateway_stage_ms( + "stream_candidate_payload_parts", + payload_started_at.elapsed().as_millis() as u64, + ); let candidate = &eligible.candidate; let prompt_cache_key = resolved @@ -66,9 +76,14 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate .map(str::trim) .filter(|value| !value.is_empty()) .map(ToOwned::to_owned); + let proxy_started_at = std::time::Instant::now(); let proxy = state .resolve_transport_proxy_snapshot_with_tunnel_affinity(&resolved.transport) .await; + observe_gateway_stage_ms( + "stream_candidate_proxy", + proxy_started_at.elapsed().as_millis() as u64, + ); let transport_profile = resolved .transport_profile .clone() @@ -130,6 +145,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate Some(body_json) }; let effective_headers = input.effective_headers(&parts.headers); + let report_context_started_at = std::time::Instant::now(); let report_context = append_local_failover_policy_to_value( append_execution_contract_fields_to_value( build_local_execution_report_context(LocalExecutionReportContextParts { @@ -184,8 +200,13 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate ), &transport, ); + observe_gateway_stage_ms( + "stream_candidate_report_context", + report_context_started_at.elapsed().as_millis() as u64, + ); let request_gzip = resolve_transport_request_gzip_policy(&transport); + let decision_started_at = std::time::Instant::now(); let mut decision = build_ai_execution_decision_response(AiExecutionDecisionResponseParts { decision_is_stream, decision_kind: decision_kind.to_string(), @@ -222,5 +243,9 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate auth_context: input.auth_context.clone(), }); apply_provider_request_routing_policy_to_decision(input, &mut decision)?; + observe_gateway_stage_ms( + "stream_candidate_decision_build", + decision_started_at.elapsed().as_millis() as u64, + ); Ok(Some(decision)) } diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs index 13c7e2af9..1eaba9dd5 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs @@ -54,6 +54,7 @@ use crate::ai_serving::{ LocalResolvedOAuthRequestAuth, }; use crate::ai_serving::{ConversionMode, ExecutionStrategy}; +use crate::stage_metrics::observe_gateway_stage_ms; use crate::{AppState, GatewayError}; use super::support::{ @@ -102,6 +103,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts( report_kind: &str, upstream_is_stream: bool, ) -> Result, GatewayError> { + let prepare_started_at = std::time::Instant::now(); let planner_state = crate::ai_serving::PlannerAppState::new(state); let candidate = &eligible.candidate; let provider_api_format = eligible.provider_api_format.as_str(); @@ -109,6 +111,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts( let transport_profile = crate::ai_serving::transport::resolve_transport_profile(transport); let force_body_stream_field = endpoint_config_forces_body_stream_field(transport.endpoint.config.as_ref()); + let model_directives_started_at = std::time::Instant::now(); let enable_model_directives = crate::system_features::reasoning_model_directive_enabled_for_api_format_and_model( state, @@ -116,6 +119,11 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts( Some(&input.requested_model), ) .await; + observe_gateway_stage_ms( + "openai_chat_payload_model_directives", + model_directives_started_at.elapsed().as_millis() as u64, + ); + let redaction_started_at = std::time::Instant::now(); let redaction = resolve_provider_chat_pii_redaction( state, parts, @@ -125,6 +133,10 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts( candidate_id, ) .await?; + observe_gateway_stage_ms( + "openai_chat_payload_redaction", + redaction_started_at.elapsed().as_millis() as u64, + ); let body_json = redaction.body_json.as_ref(); let effective_headers = input.effective_headers(&parts.headers); let is_grok = transport @@ -234,7 +246,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts( redaction.redacted, ); - return Ok(Some(LocalOpenAiChatCandidatePayloadParts { + let result = Ok(Some(LocalOpenAiChatCandidatePayloadParts { client_api_format: "openai:chat".to_string(), auth_header: prepared_candidate.auth_header, auth_value: prepared_candidate.auth_value, @@ -252,6 +264,11 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts( transport_profile, image_request_summary: None, })); + observe_gateway_stage_ms( + "openai_chat_payload_parts_prepare", + prepare_started_at.elapsed().as_millis() as u64, + ); + return result; } if provider_api_format == "openai:chat" && is_windsurf_provider_transport(transport) { @@ -288,6 +305,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts( return Ok(None); }; + let auth_prepare_started_at = std::time::Instant::now(); let prepared_candidate = match prepare_header_authenticated_candidate( planner_state, transport, @@ -316,7 +334,12 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts( return Ok(None); } }; + observe_gateway_stage_ms( + "openai_chat_payload_auth_prepare", + auth_prepare_started_at.elapsed().as_millis() as u64, + ); + let body_build_started_at = std::time::Instant::now(); let Some(mut provider_request_body) = build_local_openai_chat_request_body( body_json, &prepared_candidate.mapped_model, @@ -343,6 +366,10 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts( .await; return Ok(None); }; + observe_gateway_stage_ms( + "openai_chat_payload_body_build", + body_build_started_at.elapsed().as_millis() as u64, + ); apply_deepseek_tool_call_thinking_compat( &mut provider_request_body, transport.provider.provider_type.as_str(), diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/mod.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/mod.rs index a53ec1a40..16b8cfbe0 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/mod.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/mod.rs @@ -147,7 +147,7 @@ pub(crate) async fn maybe_build_sync_local_decision_payload( ) .await; - while let Some(attempt) = source.next_attempt().await { + while let Some(attempt) = source.next_attempt().await? { let upstream_is_stream = self::plans::openai_chat_upstream_is_stream_for_candidate( &attempt.eligible.transport, attempt.eligible.provider_api_format.as_str(), @@ -199,7 +199,7 @@ pub(crate) async fn maybe_build_stream_local_decision_payload( ) .await; - while let Some(attempt) = source.next_attempt().await { + while let Some(attempt) = source.next_attempt().await? { let upstream_is_stream = self::plans::openai_chat_upstream_is_stream_for_candidate( &attempt.eligible.transport, attempt.eligible.provider_api_format.as_str(), diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/stream.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/stream.rs index bbd7f58f5..01dd1e668 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/stream.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/stream.rs @@ -18,6 +18,7 @@ use crate::ai_serving::planner::plan_builders::{ build_openai_chat_stream_plan_from_decision, AiStreamAttempt, }; use crate::ai_serving::planner::runtime_miss::apply_local_runtime_candidate_terminal_reason; +use crate::stage_metrics::observe_gateway_stage_ms; pub(crate) struct LocalOpenAiChatStreamAttemptSource<'a> { state: &'a AppState, @@ -93,10 +94,36 @@ pub(crate) async fn build_local_openai_chat_stream_attempt_source<'a>( #[async_trait] impl LocalExecutionAttemptSource for LocalOpenAiChatStreamAttemptSource<'_> { async fn next_execution_attempt(&mut self) -> Result, GatewayError> { - while let Some(attempt) = self.candidates.next_attempt().await { + loop { + let source_started_at = std::time::Instant::now(); + let Some(attempt) = self.candidates.next_attempt().await? else { + observe_gateway_stage_ms( + "stream_candidate_source_next", + source_started_at.elapsed().as_millis() as u64, + ); + break; + }; + observe_gateway_stage_ms( + "stream_candidate_source_next", + source_started_at.elapsed().as_millis() as u64, + ); + + let plan_started_at = std::time::Instant::now(); match self.build_stream_attempt(attempt).await? { - Some(attempt) => return Ok(Some(attempt)), - None => continue, + Some(attempt) => { + observe_gateway_stage_ms( + "stream_candidate_plan_build", + plan_started_at.elapsed().as_millis() as u64, + ); + return Ok(Some(attempt)); + } + None => { + observe_gateway_stage_ms( + "stream_candidate_plan_build", + plan_started_at.elapsed().as_millis() as u64, + ); + continue; + } } } apply_local_runtime_candidate_terminal_reason( diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/sync.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/sync.rs index f1bfd6186..8dfdbd479 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/sync.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/sync.rs @@ -93,7 +93,7 @@ pub(crate) async fn build_local_openai_chat_sync_attempt_source<'a>( #[async_trait] impl LocalExecutionAttemptSource for LocalOpenAiChatSyncAttemptSource<'_> { async fn next_execution_attempt(&mut self) -> Result, GatewayError> { - while let Some(attempt) = self.candidates.next_attempt().await { + while let Some(attempt) = self.candidates.next_attempt().await? { match self.build_sync_attempt(attempt).await? { Some(attempt) => return Ok(Some(attempt)), None => continue, diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/mod.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/mod.rs index fc00c3549..3ec54f2af 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/mod.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/mod.rs @@ -114,7 +114,7 @@ pub(crate) async fn maybe_build_sync_local_openai_responses_decision_payload( ) .await?; - while let Some(attempt) = source.next_attempt().await { + while let Some(attempt) = source.next_attempt().await? { if let Some(payload) = maybe_build_local_openai_responses_decision_payload_for_candidate( state, parts, trace_id, body_json, &input, attempt, spec, ) @@ -153,7 +153,7 @@ pub(crate) async fn maybe_build_stream_local_openai_responses_decision_payload( ) .await?; - while let Some(attempt) = source.next_attempt().await { + while let Some(attempt) = source.next_attempt().await? { if let Some(payload) = maybe_build_local_openai_responses_decision_payload_for_candidate( state, parts, trace_id, body_json, &input, attempt, spec, ) diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/plans.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/plans.rs index 48282fb85..48aebaf36 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/plans.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/plans.rs @@ -162,7 +162,7 @@ pub(super) async fn build_local_stream_attempt_source<'a>( #[async_trait] impl LocalExecutionAttemptSource for LocalOpenAiResponsesSyncAttemptSource<'_> { async fn next_execution_attempt(&mut self) -> Result, GatewayError> { - while let Some(attempt) = self.candidates.next_attempt().await { + while let Some(attempt) = self.candidates.next_attempt().await? { match self.build_sync_attempt(attempt).await? { Some(attempt) => return Ok(Some(attempt)), None => continue, @@ -190,7 +190,7 @@ impl LocalExecutionAttemptSource for LocalOpenAiResponsesSyncAtte #[async_trait] impl LocalExecutionAttemptSource for LocalOpenAiResponsesStreamAttemptSource<'_> { async fn next_execution_attempt(&mut self) -> Result, GatewayError> { - while let Some(attempt) = self.candidates.next_attempt().await { + while let Some(attempt) = self.candidates.next_attempt().await? { match self.build_stream_attempt(attempt).await? { Some(attempt) => return Ok(Some(attempt)), None => continue, @@ -331,7 +331,7 @@ pub(super) async fn build_local_sync_plan_and_reports( } let mut plans = Vec::new(); - while let Some(attempt) = source.next_attempt().await { + while let Some(attempt) = source.next_attempt().await? { let Some(payload) = maybe_build_local_openai_responses_decision_payload_for_candidate( state, parts, trace_id, body_json, &input, attempt, spec, ) @@ -403,7 +403,7 @@ pub(super) async fn build_local_stream_plan_and_reports( } let mut plans = Vec::new(); - while let Some(attempt) = source.next_attempt().await { + while let Some(attempt) = source.next_attempt().await? { let Some(payload) = maybe_build_local_openai_responses_decision_payload_for_candidate( state, parts, trace_id, body_json, &input, attempt, spec, ) diff --git a/apps/aether-gateway/src/ai_serving/planner/state/transport.rs b/apps/aether-gateway/src/ai_serving/planner/state/transport.rs index c71822607..cc6a45d8b 100644 --- a/apps/aether-gateway/src/ai_serving/planner/state/transport.rs +++ b/apps/aether-gateway/src/ai_serving/planner/state/transport.rs @@ -3,8 +3,20 @@ pub(crate) use crate::ai_serving::transport::{ GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth, }; use crate::GatewayError; +use std::sync::Arc; impl<'a> PlannerAppState<'a> { + pub(crate) async fn read_provider_transport_snapshot_arc( + self, + provider_id: &str, + endpoint_id: &str, + key_id: &str, + ) -> Result>, GatewayError> { + self.app() + .read_provider_transport_snapshot_arc(provider_id, endpoint_id, key_id) + .await + } + pub(crate) async fn read_provider_transport_snapshot( self, provider_id: &str, diff --git a/apps/aether-gateway/src/cache/auth_context.rs b/apps/aether-gateway/src/cache/auth_context.rs index ddbd70af3..16110515e 100644 --- a/apps/aether-gateway/src/cache/auth_context.rs +++ b/apps/aether-gateway/src/cache/auth_context.rs @@ -1,12 +1,54 @@ +use std::collections::HashSet; use std::time::Duration; use aether_cache::ExpiringMap; +use tokio::sync::Notify; use crate::control::GatewayControlAuthContext; -#[derive(Debug, Default)] +#[derive(Debug)] pub(crate) struct AuthContextCache { entries: ExpiringMap, + inflight: std::sync::Mutex>, + notify: Notify, +} + +impl Default for AuthContextCache { + fn default() -> Self { + Self { + entries: ExpiringMap::default(), + inflight: std::sync::Mutex::new(HashSet::new()), + notify: Notify::new(), + } + } +} + +pub(crate) enum AuthContextInflightRegistration<'a> { + Leader(AuthContextInflightGuard<'a>), + Follower, + Bypass, +} + +pub(crate) struct AuthContextInflightGuard<'a> { + cache: &'a AuthContextCache, + cache_key: Option, +} + +impl Drop for AuthContextInflightGuard<'_> { + fn drop(&mut self) { + let Some(cache_key) = self.cache_key.take() else { + return; + }; + let removed = self + .cache + .inflight + .lock() + .map(|mut inflight| inflight.remove(&cache_key)) + .unwrap_or(false); + if removed { + self.cache.notify.notify_waiters(); + } + } } impl AuthContextCache { @@ -29,7 +71,39 @@ impl AuthContextCache { .insert(cache_key, auth_context, ttl, max_entries); } + pub(crate) fn notified(&self) -> tokio::sync::futures::Notified<'_> { + self.notify.notified() + } + + pub(crate) fn register_inflight(&self, cache_key: &str) -> AuthContextInflightRegistration<'_> { + let cache_key = cache_key.trim(); + if cache_key.is_empty() { + return AuthContextInflightRegistration::Bypass; + } + match self.inflight.lock() { + Ok(mut inflight) => { + if inflight.contains(cache_key) { + AuthContextInflightRegistration::Follower + } else { + inflight.insert(cache_key.to_string()); + AuthContextInflightRegistration::Leader(AuthContextInflightGuard { + cache: self, + cache_key: Some(cache_key.to_string()), + }) + } + } + Err(_) => AuthContextInflightRegistration::Bypass, + } + } + pub(crate) fn clear(&self) { self.entries.clear(); + if let Ok(mut inflight) = self.inflight.lock() { + let had_inflight = !inflight.is_empty(); + inflight.clear(); + if had_inflight { + self.notify.notify_waiters(); + } + } } } diff --git a/apps/aether-gateway/src/cache/auth_runtime.rs b/apps/aether-gateway/src/cache/auth_runtime.rs index 8b6ac7808..15eb3e5fb 100644 --- a/apps/aether-gateway/src/cache/auth_runtime.rs +++ b/apps/aether-gateway/src/cache/auth_runtime.rs @@ -198,10 +198,10 @@ impl AuthSnapshotCache { &self, key: AuthSnapshotCacheKey, ttl: Duration, - load: F, + mut load: F, ) -> Result, E> where - F: Fn() -> Fut, + F: FnMut() -> Fut, Fut: Future, E>>, { if let Some(value) = self.get(&key, ttl) { @@ -269,10 +269,10 @@ where &self, key: K, ttl: Duration, - load: F, + mut load: F, ) -> Result, E> where - F: Fn() -> Fut, + F: FnMut() -> Fut, Fut: Future, E>>, { if let Some(value) = self.get(&key, ttl) { @@ -341,10 +341,10 @@ where &self, key: K, ttl: Duration, - load: F, + mut load: F, ) -> Result, E> where - F: Fn() -> Fut, + F: FnMut() -> Fut, Fut: Future, E>>, { if let Some(value) = self.get(&key, ttl) { @@ -374,15 +374,200 @@ where } } + pub(crate) async fn get_or_load_once( + &self, + key: K, + ttl: Duration, + load: F, + ) -> Result, E> + where + F: FnOnce() -> Fut, + Fut: Future, E>>, + { + self.get_or_load_once_with_observer(key, ttl, load, CacheLoadObserver::default()) + .await + } + + pub(crate) async fn get_or_load_once_with_observer( + &self, + key: K, + ttl: Duration, + load: F, + observer: CacheLoadObserver, + ) -> Result, E> + where + F: FnOnce() -> Fut, + Fut: Future, E>>, + { + if let Some(value) = self.get(&key, ttl) { + observer.hit(); + return Ok(value); + } + observer.miss(); + + let mut load = Some(load); + loop { + let notified = self.singleflight.notified(); + match self.singleflight.register(&key) { + CacheInflightRegistration::Bypass => { + observer.load(); + let value = + load.take().expect("cache load closure should be available")().await?; + self.insert(key, value.clone(), ttl); + return Ok(value); + } + CacheInflightRegistration::Follower => { + observer.follower_wait(); + notified.await; + if let Some(value) = self.get(&key, ttl) { + observer.hit(); + return Ok(value); + } + } + CacheInflightRegistration::Leader(_guard) => { + observer.load(); + let value = + load.take().expect("cache load closure should be available")().await?; + self.insert(key, value.clone(), ttl); + return Ok(value); + } + } + } + } + + pub(crate) async fn get_or_load_once_stale_while_refreshing( + &self, + key: K, + ttl: Duration, + stale_ttl: Duration, + load: F, + observer: CacheLoadObserver, + ) -> Result, E> + where + F: FnOnce() -> Fut, + Fut: Future, E>>, + { + if let Some((value, age)) = self.entries.get_with_age(&key, stale_ttl) { + if age <= ttl { + observer.hit(); + return Ok(value); + } + + // Keep stale snapshots off the request critical path. The caller's + // invalidation path clears entries when provider/catalog/routing + // state changes, and the bounded stale TTL limits passive drift. + observer.hit(); + return Ok(value); + } + + observer.miss(); + let mut load = Some(load); + loop { + let notified = self.singleflight.notified(); + match self.singleflight.register(&key) { + CacheInflightRegistration::Bypass => { + observer.load(); + let value = + load.take().expect("cache load closure should be available")().await?; + self.entries.insert( + key, + value.clone(), + stale_ttl, + AUTH_RUNTIME_CACHE_MAX_ENTRIES, + ); + return Ok(value); + } + CacheInflightRegistration::Follower => { + observer.follower_wait(); + notified.await; + if let Some((value, _age)) = self.entries.get_with_age(&key, stale_ttl) { + observer.hit(); + return Ok(value); + } + } + CacheInflightRegistration::Leader(_guard) => { + observer.load(); + let value = + load.take().expect("cache load closure should be available")().await?; + self.entries.insert( + key, + value.clone(), + stale_ttl, + AUTH_RUNTIME_CACHE_MAX_ENTRIES, + ); + return Ok(value); + } + } + } + } + pub(crate) fn clear(&self) { self.entries.clear(); self.singleflight.clear(); } } +#[derive(Clone, Copy, Default)] +pub(crate) struct CacheLoadObserver { + on_hit: Option, + on_miss: Option, + on_load: Option, + on_follower_wait: Option, +} + +impl CacheLoadObserver { + pub(crate) fn new() -> Self { + Self::default() + } + + pub(crate) fn on_hit(mut self, callback: fn()) -> Self { + self.on_hit = Some(callback); + self + } + + pub(crate) fn on_miss(mut self, callback: fn()) -> Self { + self.on_miss = Some(callback); + self + } + + pub(crate) fn on_load(mut self, callback: fn()) -> Self { + self.on_load = Some(callback); + self + } + + pub(crate) fn on_follower_wait(mut self, callback: fn()) -> Self { + self.on_follower_wait = Some(callback); + self + } + + fn hit(self) { + if let Some(callback) = self.on_hit { + callback(); + } + } + + fn miss(self) { + if let Some(callback) = self.on_miss { + callback(); + } + } + + fn load(self) { + if let Some(callback) = self.on_load { + callback(); + } + } + + fn follower_wait(self) { + if let Some(callback) = self.on_follower_wait { + callback(); + } + } +} + #[cfg(test)] mod tests { - use super::ValueCache; + use super::{CacheLoadObserver, ValueCache}; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -472,6 +657,117 @@ mod tests { assert_eq!(max_active.load(Ordering::Acquire), 2); } + #[tokio::test] + async fn value_cache_returns_stale_without_refreshing_on_request_path() { + let cache = Arc::new(ValueCache::::default()); + let key = "hot-key".to_string(); + let calls = Arc::new(AtomicUsize::new(0)); + cache.insert(key.clone(), Some(1), Duration::from_millis(10)); + tokio::time::sleep(Duration::from_millis(20)).await; + + let first_cache = Arc::clone(&cache); + let first_key = key.clone(); + let first_calls = Arc::clone(&calls); + let first_started = Instant::now(); + let first = tokio::spawn(async move { + first_cache + .get_or_load_once_stale_while_refreshing::<(), _, _>( + first_key, + Duration::from_millis(10), + Duration::from_secs(1), + || async move { + first_calls.fetch_add(1, Ordering::AcqRel); + tokio::time::sleep(Duration::from_millis(100)).await; + Ok(Some(2)) + }, + CacheLoadObserver::default(), + ) + .await + }); + + let follower_cache = Arc::clone(&cache); + let follower_started = Instant::now(); + let follower_calls = Arc::clone(&calls); + let follower = tokio::spawn(async move { + follower_cache + .get_or_load_once_stale_while_refreshing::<(), _, _>( + key, + Duration::from_millis(10), + Duration::from_secs(1), + || async move { + follower_calls.fetch_add(1, Ordering::AcqRel); + Ok(Some(3)) + }, + CacheLoadObserver::default(), + ) + .await + }); + + assert_eq!(first.await.unwrap().unwrap(), Some(1)); + assert!( + first_started.elapsed() < Duration::from_millis(80), + "stale value should not wait for request-path refresh" + ); + assert_eq!(follower.await.unwrap().unwrap(), Some(1)); + assert!( + follower_started.elapsed() < Duration::from_millis(80), + "follower should return stale value without waiting for refresh" + ); + assert_eq!(calls.load(Ordering::Acquire), 0); + } + + #[tokio::test] + async fn value_cache_cold_stale_followers_do_not_reload_after_fresh_ttl() { + let cache = Arc::new(ValueCache::::default()); + let key = "cold-hot-key".to_string(); + let calls = Arc::new(AtomicUsize::new(0)); + + let leader_cache = Arc::clone(&cache); + let leader_key = key.clone(); + let leader_calls = Arc::clone(&calls); + let leader = tokio::spawn(async move { + leader_cache + .get_or_load_once_stale_while_refreshing::<(), _, _>( + leader_key, + Duration::from_millis(10), + Duration::from_secs(1), + || async move { + leader_calls.fetch_add(1, Ordering::AcqRel); + Ok(Some(1)) + }, + CacheLoadObserver::default(), + ) + .await + }); + assert_eq!(leader.await.unwrap().unwrap(), Some(1)); + tokio::time::sleep(Duration::from_millis(25)).await; + + let follower_started = Instant::now(); + let follower_cache = Arc::clone(&cache); + let follower_calls = Arc::clone(&calls); + let follower = tokio::spawn(async move { + follower_cache + .get_or_load_once_stale_while_refreshing::<(), _, _>( + key, + Duration::from_millis(10), + Duration::from_secs(1), + || async move { + follower_calls.fetch_add(1, Ordering::AcqRel); + Ok(Some(2)) + }, + CacheLoadObserver::default(), + ) + .await + }); + + assert_eq!(follower.await.unwrap().unwrap(), Some(1)); + assert_eq!(calls.load(Ordering::Acquire), 1); + assert!( + follower_started.elapsed() < Duration::from_millis(50), + "follower should reuse cold-loaded stale value without reloading" + ); + } + #[tokio::test] async fn value_cache_clear_releases_same_key_followers() { let cache = Arc::new(ValueCache::::default()); diff --git a/apps/aether-gateway/src/cache/candidate_page.rs b/apps/aether-gateway/src/cache/candidate_page.rs new file mode 100644 index 000000000..ff8997aa7 --- /dev/null +++ b/apps/aether-gateway/src/cache/candidate_page.rs @@ -0,0 +1,479 @@ +use std::collections::hash_map::DefaultHasher; +use std::hash::{Hash, Hasher}; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, LazyLock}; +use std::time::Duration; + +use aether_ai_serving::AiCandidatePreselectionOutcome; +use aether_ai_serving::AiCandidateResolutionMode; +use aether_routing_core::ResolvedRoutingPolicy; +use aether_runtime::{MetricKind, MetricSample}; +use aether_scheduler_core::{ + normalize_api_format, ClientSessionAffinity, SchedulerMinimalCandidateSelectionCandidate, +}; +use serde_json::Value; +use sha2::Digest as _; + +use crate::ai_serving::{ + EligibleLocalExecutionCandidate, GatewayAuthApiKeySnapshot, SkippedLocalExecutionCandidate, +}; + +const DEFAULT_CANDIDATE_PAGE_CACHE_TTL_MS: u64 = 250; +const MIN_CANDIDATE_PAGE_CACHE_TTL_MS: u64 = 50; +const MAX_CANDIDATE_PAGE_CACHE_TTL_MS: u64 = 1_000; +const CANDIDATE_PAGE_CACHE_TTL_ENV: &str = "AETHER_GATEWAY_CANDIDATE_PAGE_CACHE_TTL_MS"; + +pub(crate) type CandidatePageSnapshot = AiCandidatePreselectionOutcome< + SchedulerMinimalCandidateSelectionCandidate, + SkippedLocalExecutionCandidate, +>; + +pub(crate) type CandidatePageCache = + super::ValueCache>; + +#[derive(Debug, Clone)] +pub(crate) struct CandidateResolvedPageSnapshot { + pub(crate) candidates: Vec, + pub(crate) resolved_skipped: Vec, +} + +pub(crate) type CandidateResolvedPageCache = + super::ValueCache>; + +static CANDIDATE_PAGE_CACHE_METRICS: LazyLock = + LazyLock::new(CandidatePageCacheMetrics::default); + +#[derive(Debug, Default)] +struct CandidatePageCacheMetrics { + hit_total: AtomicU64, + load_total: AtomicU64, + follower_wait_total: AtomicU64, + miss_total: AtomicU64, + none_total: AtomicU64, + resolve_hit_total: AtomicU64, + resolve_load_total: AtomicU64, + resolve_follower_wait_total: AtomicU64, + resolve_miss_total: AtomicU64, +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub(crate) struct CandidatePageCacheKey { + requested_model: String, + client_api_format: String, + auth_identity: CandidatePageAuthIdentity, + require_streaming: bool, + required_capabilities_hash: String, + routing_policy_hash: String, + request_auth_channel: String, + scheduler_affinity_epoch: u64, + preselection_mode: &'static str, + use_api_format_alias_match: bool, + client_session_affinity_hash: String, +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +enum CandidatePageAuthIdentity { + Standalone { api_key_id: String }, + UserApiKey { user_id: String, api_key_id: String }, +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub(crate) struct CandidateResolvedPageCacheKey { + page_key: CandidatePageCacheKey, + resolution_mode: &'static str, +} + +impl CandidatePageCacheKey { + #[allow(clippy::too_many_arguments)] + pub(crate) fn new( + requested_model: &str, + client_api_format: &str, + require_streaming: bool, + auth_snapshot: &GatewayAuthApiKeySnapshot, + required_capabilities: Option<&Value>, + routing_policy: Option<&ResolvedRoutingPolicy>, + request_auth_channel: Option<&str>, + scheduler_affinity_epoch: u64, + preselection_mode: &'static str, + use_api_format_alias_match: bool, + client_session_affinity: Option<&ClientSessionAffinity>, + ) -> Self { + Self { + requested_model: normalize_text_key(requested_model), + client_api_format: normalize_api_format(client_api_format), + auth_identity: CandidatePageAuthIdentity::from_auth_snapshot(auth_snapshot), + require_streaming, + required_capabilities_hash: stable_json_hash(required_capabilities), + routing_policy_hash: stable_json_hash(routing_policy), + request_auth_channel: normalize_text_key(request_auth_channel.unwrap_or_default()), + scheduler_affinity_epoch, + preselection_mode, + use_api_format_alias_match, + client_session_affinity_hash: client_session_affinity_key(client_session_affinity), + } + } +} + +impl CandidateResolvedPageCacheKey { + #[allow(clippy::too_many_arguments)] + pub(crate) fn new( + requested_model: &str, + client_api_format: &str, + require_streaming: bool, + auth_snapshot: &GatewayAuthApiKeySnapshot, + required_capabilities: Option<&Value>, + routing_policy: Option<&ResolvedRoutingPolicy>, + request_auth_channel: Option<&str>, + scheduler_affinity_epoch: u64, + preselection_mode: &'static str, + use_api_format_alias_match: bool, + client_session_affinity: Option<&ClientSessionAffinity>, + resolution_mode: AiCandidateResolutionMode, + ) -> Self { + Self { + page_key: CandidatePageCacheKey::new( + requested_model, + client_api_format, + require_streaming, + auth_snapshot, + required_capabilities, + routing_policy, + request_auth_channel, + scheduler_affinity_epoch, + preselection_mode, + use_api_format_alias_match, + client_session_affinity, + ), + resolution_mode: resolution_mode_name(resolution_mode), + } + } +} + +impl CandidatePageAuthIdentity { + fn from_auth_snapshot(auth_snapshot: &GatewayAuthApiKeySnapshot) -> Self { + let api_key_id = normalize_text_key(&auth_snapshot.api_key_id); + if auth_snapshot.api_key_is_standalone { + Self::Standalone { api_key_id } + } else { + Self::UserApiKey { + user_id: normalize_text_key(&auth_snapshot.user_id), + api_key_id, + } + } + } +} + +pub(crate) fn candidate_page_cache_ttl_from_env() -> Duration { + let ttl_ms = std::env::var(CANDIDATE_PAGE_CACHE_TTL_ENV) + .ok() + .and_then(|value| value.trim().parse::().ok()) + .filter(|value| *value > 0) + .unwrap_or(DEFAULT_CANDIDATE_PAGE_CACHE_TTL_MS) + .clamp( + MIN_CANDIDATE_PAGE_CACHE_TTL_MS, + MAX_CANDIDATE_PAGE_CACHE_TTL_MS, + ); + Duration::from_millis(ttl_ms) +} + +pub(crate) fn candidate_page_cache_stale_ttl(ttl: Duration) -> Duration { + let stale_ttl = ttl.saturating_mul(8); + stale_ttl.min(Duration::from_secs(2)).max(ttl) +} + +pub(crate) fn record_candidate_page_cache_hit() { + CANDIDATE_PAGE_CACHE_METRICS + .hit_total + .fetch_add(1, Ordering::Relaxed); +} + +pub(crate) fn record_candidate_page_cache_miss() { + CANDIDATE_PAGE_CACHE_METRICS + .miss_total + .fetch_add(1, Ordering::Relaxed); +} + +pub(crate) fn record_candidate_page_cache_load() { + CANDIDATE_PAGE_CACHE_METRICS + .load_total + .fetch_add(1, Ordering::Relaxed); +} + +pub(crate) fn record_candidate_page_cache_follower_wait() { + CANDIDATE_PAGE_CACHE_METRICS + .follower_wait_total + .fetch_add(1, Ordering::Relaxed); +} + +pub(crate) fn record_candidate_page_cache_none() { + CANDIDATE_PAGE_CACHE_METRICS + .none_total + .fetch_add(1, Ordering::Relaxed); +} + +pub(crate) fn record_candidate_page_resolve_cache_hit() { + CANDIDATE_PAGE_CACHE_METRICS + .resolve_hit_total + .fetch_add(1, Ordering::Relaxed); +} + +pub(crate) fn record_candidate_page_resolve_cache_miss() { + CANDIDATE_PAGE_CACHE_METRICS + .resolve_miss_total + .fetch_add(1, Ordering::Relaxed); +} + +pub(crate) fn record_candidate_page_resolve_cache_load() { + CANDIDATE_PAGE_CACHE_METRICS + .resolve_load_total + .fetch_add(1, Ordering::Relaxed); +} + +pub(crate) fn record_candidate_page_resolve_cache_follower_wait() { + CANDIDATE_PAGE_CACHE_METRICS + .resolve_follower_wait_total + .fetch_add(1, Ordering::Relaxed); +} + +pub(crate) fn candidate_page_cache_metric_samples() -> Vec { + vec![ + MetricSample::new( + "candidate_page_cache_hit_total", + "Total candidate page cache hits.", + MetricKind::Counter, + CANDIDATE_PAGE_CACHE_METRICS + .hit_total + .load(Ordering::Relaxed), + ), + MetricSample::new( + "candidate_page_cache_miss_total", + "Total candidate page cache misses before singleflight registration.", + MetricKind::Counter, + CANDIDATE_PAGE_CACHE_METRICS + .miss_total + .load(Ordering::Relaxed), + ), + MetricSample::new( + "candidate_page_cache_load_total", + "Total candidate page cache loader executions.", + MetricKind::Counter, + CANDIDATE_PAGE_CACHE_METRICS + .load_total + .load(Ordering::Relaxed), + ), + MetricSample::new( + "candidate_page_cache_follower_wait_total", + "Total candidate page cache requests that waited for another loader.", + MetricKind::Counter, + CANDIDATE_PAGE_CACHE_METRICS + .follower_wait_total + .load(Ordering::Relaxed), + ), + MetricSample::new( + "candidate_page_cache_none_total", + "Total candidate page cache lookups that resolved to no page.", + MetricKind::Counter, + CANDIDATE_PAGE_CACHE_METRICS + .none_total + .load(Ordering::Relaxed), + ), + MetricSample::new( + "candidate_page_resolve_cache_hit_total", + "Total resolved candidate page cache hits.", + MetricKind::Counter, + CANDIDATE_PAGE_CACHE_METRICS + .resolve_hit_total + .load(Ordering::Relaxed), + ), + MetricSample::new( + "candidate_page_resolve_cache_miss_total", + "Total resolved candidate page cache misses before singleflight registration.", + MetricKind::Counter, + CANDIDATE_PAGE_CACHE_METRICS + .resolve_miss_total + .load(Ordering::Relaxed), + ), + MetricSample::new( + "candidate_page_resolve_cache_load_total", + "Total resolved candidate page cache loader executions.", + MetricKind::Counter, + CANDIDATE_PAGE_CACHE_METRICS + .resolve_load_total + .load(Ordering::Relaxed), + ), + MetricSample::new( + "candidate_page_resolve_cache_follower_wait_total", + "Total resolved candidate page cache requests that waited for another loader.", + MetricKind::Counter, + CANDIDATE_PAGE_CACHE_METRICS + .resolve_follower_wait_total + .load(Ordering::Relaxed), + ), + ] +} + +fn normalize_text_key(value: &str) -> String { + value.trim().to_string() +} + +fn client_session_affinity_key(affinity: Option<&ClientSessionAffinity>) -> String { + let Some(affinity) = affinity else { + return String::new(); + }; + let family = affinity + .client_family + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_ascii_lowercase) + .unwrap_or_default(); + let session = affinity + .session_key + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(sha256_hex) + .unwrap_or_default(); + format!("{family}:{session}") +} + +fn stable_json_hash(value: Option<&T>) -> String +where + T: serde::Serialize, +{ + let Some(value) = value else { + return String::new(); + }; + match serde_json::to_vec(value) { + Ok(serialized) => sha256_hex(&serialized), + Err(_) => { + let mut hasher = DefaultHasher::new(); + std::any::type_name::().hash(&mut hasher); + format!("fallback:{:016x}", hasher.finish()) + } + } +} + +fn sha256_hex(value: impl AsRef<[u8]>) -> String { + let digest = sha2::Sha256::digest(value.as_ref()); + digest.iter().map(|byte| format!("{byte:02x}")).collect() +} + +fn resolution_mode_name(mode: AiCandidateResolutionMode) -> &'static str { + match mode { + AiCandidateResolutionMode::Standard => "standard", + AiCandidateResolutionMode::WithoutTransportPairGate => "without_transport_pair_gate", + } +} + +#[cfg(test)] +mod tests { + use super::*; + use aether_data::repository::auth::ResolvedAuthApiKeySnapshot; + use serde_json::json; + + fn auth_snapshot(user_id: &str, api_key_id: &str) -> ResolvedAuthApiKeySnapshot { + ResolvedAuthApiKeySnapshot { + user_id: user_id.to_string(), + username: "user".to_string(), + email: None, + user_role: "user".to_string(), + user_auth_source: "local".to_string(), + user_is_active: true, + user_is_deleted: false, + user_rate_limit: None, + user_allowed_providers: None, + user_allowed_api_formats: None, + user_allowed_models: None, + api_key_id: api_key_id.to_string(), + api_key_name: None, + api_key_is_active: true, + api_key_is_locked: false, + api_key_is_standalone: false, + api_key_rate_limit: None, + api_key_concurrent_limit: None, + api_key_expires_at_unix_secs: None, + api_key_allowed_providers: None, + api_key_allowed_api_formats: None, + api_key_allowed_models: None, + api_key_ip_rules: None, + currently_usable: true, + } + } + + #[test] + fn candidate_page_cache_key_isolates_auth_model_format_and_capabilities() { + let auth_a = auth_snapshot("user-a", "key-a"); + let auth_b = auth_snapshot("user-b", "key-a"); + let base = CandidatePageCacheKey::new( + "gpt-4o", + "openai:chat", + true, + &auth_a, + Some(&json!({"vision": true})), + None, + Some("bearer"), + 7, + "provider_endpoint_key_model", + true, + None, + ); + let different_user = CandidatePageCacheKey::new( + "gpt-4o", + "openai:chat", + true, + &auth_b, + Some(&json!({"vision": true})), + None, + Some("bearer"), + 7, + "provider_endpoint_key_model", + true, + None, + ); + let different_model = CandidatePageCacheKey::new( + "gpt-4.1", + "openai:chat", + true, + &auth_a, + Some(&json!({"vision": true})), + None, + Some("bearer"), + 7, + "provider_endpoint_key_model", + true, + None, + ); + let different_format = CandidatePageCacheKey::new( + "gpt-4o", + "openai:responses", + true, + &auth_a, + Some(&json!({"vision": true})), + None, + Some("bearer"), + 7, + "provider_endpoint_key_model", + true, + None, + ); + let different_capabilities = CandidatePageCacheKey::new( + "gpt-4o", + "openai:chat", + true, + &auth_a, + Some(&json!({"vision": false})), + None, + Some("bearer"), + 7, + "provider_endpoint_key_model", + true, + None, + ); + + assert_ne!(base, different_user); + assert_ne!(base, different_model); + assert_ne!(base, different_format); + assert_ne!(base, different_capabilities); + } +} diff --git a/apps/aether-gateway/src/cache/mod.rs b/apps/aether-gateway/src/cache/mod.rs index 832f22398..73362e203 100644 --- a/apps/aether-gateway/src/cache/mod.rs +++ b/apps/aether-gateway/src/cache/mod.rs @@ -1,20 +1,31 @@ mod auth_api_key_last_used; mod auth_context; mod auth_runtime; +mod candidate_page; mod dashboard_response; mod direct_plan_bypass; mod scheduler_affinity; mod system_config; pub(crate) use auth_api_key_last_used::AuthApiKeyLastUsedCache; -pub(crate) use auth_context::AuthContextCache; +pub(crate) use auth_context::{AuthContextCache, AuthContextInflightRegistration}; pub(crate) use auth_runtime::{ AuthApiKeyFeatureCacheKey, AuthApiKeyIdentityCacheKey, AuthSnapshotCache, AuthSnapshotCacheKey, - JsonValueCache, ValueCache, + CacheLoadObserver, JsonValueCache, ValueCache, +}; +pub(crate) use candidate_page::{ + candidate_page_cache_metric_samples, candidate_page_cache_stale_ttl, + candidate_page_cache_ttl_from_env, record_candidate_page_cache_follower_wait, + record_candidate_page_cache_hit, record_candidate_page_cache_load, + record_candidate_page_cache_miss, record_candidate_page_cache_none, + record_candidate_page_resolve_cache_follower_wait, record_candidate_page_resolve_cache_hit, + record_candidate_page_resolve_cache_load, record_candidate_page_resolve_cache_miss, + CandidatePageCache, CandidatePageCacheKey, CandidatePageSnapshot, CandidateResolvedPageCache, + CandidateResolvedPageCacheKey, CandidateResolvedPageSnapshot, }; pub(crate) use dashboard_response::DashboardResponseCache; pub(crate) use direct_plan_bypass::DirectPlanBypassCache; pub(crate) use scheduler_affinity::{ SchedulerAffinityCache, SchedulerAffinitySnapshotEntry, SchedulerAffinityTarget, }; -pub(crate) use system_config::SystemConfigCache; +pub(crate) use system_config::{SystemConfigCache, SystemConfigInflightRegistration}; diff --git a/apps/aether-gateway/src/cache/system_config.rs b/apps/aether-gateway/src/cache/system_config.rs index ec25cb7a9..20f75ea68 100644 --- a/apps/aether-gateway/src/cache/system_config.rs +++ b/apps/aether-gateway/src/cache/system_config.rs @@ -1,21 +1,43 @@ +use std::collections::HashSet; use std::time::Duration; use aether_cache::ExpiringMap; -use tokio::sync::{Mutex, MutexGuard}; +use tokio::sync::Notify; const MAX_ENTRIES: usize = 512; #[derive(Debug)] pub(crate) struct SystemConfigCache { entries: ExpiringMap>, - load_guard: Mutex<()>, + inflight: std::sync::Mutex>, + notify: Notify, } impl Default for SystemConfigCache { fn default() -> Self { Self { entries: ExpiringMap::new(), - load_guard: Mutex::new(()), + inflight: std::sync::Mutex::new(HashSet::new()), + notify: Notify::new(), + } + } +} + +pub(crate) enum SystemConfigInflightRegistration<'a> { + Leader(SystemConfigInflightGuard<'a>), + Follower, + Bypass, +} + +pub(crate) struct SystemConfigInflightGuard<'a> { + cache: &'a SystemConfigCache, + key: Option, +} + +impl Drop for SystemConfigInflightGuard<'_> { + fn drop(&mut self) { + if let Some(key) = self.key.take() { + self.cache.finish_load(&key); } } } @@ -29,11 +51,51 @@ impl SystemConfigCache { self.entries.insert(key, value, ttl, MAX_ENTRIES); } - pub(crate) async fn load_guard(&self) -> MutexGuard<'_, ()> { - self.load_guard.lock().await + pub(crate) fn register_load(&self, key: &str) -> SystemConfigInflightRegistration<'_> { + match self.inflight.lock() { + Ok(mut inflight) => { + if inflight.contains(key) { + SystemConfigInflightRegistration::Follower + } else { + inflight.insert(key.to_string()); + SystemConfigInflightRegistration::Leader(SystemConfigInflightGuard { + cache: self, + key: Some(key.to_string()), + }) + } + } + Err(_) => SystemConfigInflightRegistration::Bypass, + } + } + + pub(crate) fn notified(&self) -> tokio::sync::futures::Notified<'_> { + self.notify.notified() + } + + fn finish_load(&self, key: &str) { + let removed = self + .inflight + .lock() + .map(|mut inflight| inflight.remove(key)) + .unwrap_or(false); + if removed { + self.notify.notify_waiters(); + } } pub(crate) fn clear(&self) { self.entries.clear(); + let cleared = self + .inflight + .lock() + .map(|mut inflight| { + let had_entries = !inflight.is_empty(); + inflight.clear(); + had_entries + }) + .unwrap_or(false); + if cleared { + self.notify.notify_waiters(); + } } } diff --git a/apps/aether-gateway/src/clock.rs b/apps/aether-gateway/src/clock.rs index 7615f29db..84616ed7d 100644 --- a/apps/aether-gateway/src/clock.rs +++ b/apps/aether-gateway/src/clock.rs @@ -1,5 +1,8 @@ +use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{SystemTime, UNIX_EPOCH}; +static REQUEST_DISTRIBUTION_COUNTER: AtomicU64 = AtomicU64::new(0); + pub(crate) fn current_unix_secs() -> u64 { SystemTime::now() .duration_since(UNIX_EPOCH) @@ -13,3 +16,9 @@ pub(crate) fn current_unix_ms() -> u64 { .unwrap_or_default() .as_millis() as u64 } + +pub(crate) fn request_distribution_seed() -> u64 { + let now_ms = current_unix_ms(); + let counter = REQUEST_DISTRIBUTION_COUNTER.fetch_add(1, Ordering::Relaxed); + now_ms.rotate_left(21) ^ counter +} diff --git a/apps/aether-gateway/src/control/auth/resolution.rs b/apps/aether-gateway/src/control/auth/resolution.rs index f72841238..37514d6f0 100644 --- a/apps/aether-gateway/src/control/auth/resolution.rs +++ b/apps/aether-gateway/src/control/auth/resolution.rs @@ -23,6 +23,7 @@ use super::principal::derive_principal_candidate; use super::types::{ GatewayCredentialCarrier, GatewayPrincipalCandidate, GatewayTrustedAuthHeaders, }; +use crate::cache::AuthContextInflightRegistration; use crate::headers::header_value_str; const AUTH_CONTEXT_CACHE_TTL: Duration = Duration::from_secs(60); @@ -137,14 +138,15 @@ pub(in super::super) async fn resolve_control_decision_auth( if let Some(cache_key) = auth_context_cache_key.as_deref() { if let Some(auth_context) = get_cached_auth_context(state, cache_key) { resolved_auth_context = if auth_context_cache_refresh_on_hit() { - let refreshed = refresh_execution_runtime_auth_context( - state, - auth_context, - decision.auth_endpoint_signature.as_deref(), + Some( + refresh_cached_auth_context_or_reuse( + state, + cache_key, + auth_context, + decision.auth_endpoint_signature.as_deref(), + ) + .await?, ) - .await?; - put_cached_auth_context(state, cache_key.to_string(), refreshed.clone()); - Some(refreshed) } else { Some(auth_context) }; @@ -152,11 +154,13 @@ pub(in super::super) async fn resolve_control_decision_auth( } if resolved_auth_context.is_none() { - resolved_auth_context = resolve_data_backed_auth_context( + resolved_auth_context = resolve_data_backed_auth_context_cached( state, + auth_context_cache_key.as_deref(), headers, uri, decision.auth_endpoint_signature.as_deref(), + true, ) .await?; if let (Some(cache_key), Some(auth_context)) = ( @@ -501,14 +505,15 @@ pub(crate) async fn resolve_execution_runtime_auth_context( if !auth_context_cache_refresh_on_hit() { return Ok(Some(auth_context)); } - return Ok(Some( - refresh_execution_runtime_auth_context( - state, - auth_context, - decision.auth_endpoint_signature.as_deref(), - ) - .await?, - )); + return refresh_decision_auth_context_on_hit( + state, + headers, + uri, + decision.auth_endpoint_signature.as_deref(), + auth_context, + ) + .await + .map(Some); } let Some(auth_endpoint_signature) = decision.auth_endpoint_signature.as_deref() else { @@ -524,18 +529,25 @@ pub(crate) async fn resolve_execution_runtime_auth_context( return Ok(Some(auth_context)); } - let refreshed = refresh_execution_runtime_auth_context( + let refreshed = refresh_cached_auth_context_or_reuse( state, + &cache_key, auth_context, Some(auth_endpoint_signature), ) .await?; - put_cached_auth_context(state, cache_key, refreshed.clone()); return Ok(Some(refreshed)); } - if let Some(auth_context) = - resolve_data_backed_auth_context(state, headers, uri, Some(auth_endpoint_signature)).await? + if let Some(auth_context) = resolve_data_backed_auth_context_cached( + state, + Some(cache_key.as_str()), + headers, + uri, + Some(auth_endpoint_signature), + true, + ) + .await? { if auth_context.user_id.is_empty() || auth_context.api_key_id.is_empty() { return Ok(None); @@ -547,6 +559,109 @@ pub(crate) async fn resolve_execution_runtime_auth_context( Ok(None) } +async fn refresh_decision_auth_context_on_hit( + state: &AppState, + headers: &http::HeaderMap, + uri: &Uri, + auth_endpoint_signature: Option<&str>, + auth_context: GatewayControlAuthContext, +) -> Result { + let Some(auth_endpoint_signature) = auth_endpoint_signature else { + return Ok(auth_context); + }; + let Some(cache_key) = build_auth_context_cache_key(headers, uri, auth_endpoint_signature) + else { + return refresh_execution_runtime_auth_context( + state, + auth_context, + Some(auth_endpoint_signature), + ) + .await; + }; + refresh_cached_auth_context_or_reuse( + state, + &cache_key, + auth_context, + Some(auth_endpoint_signature), + ) + .await +} + +async fn refresh_cached_auth_context_or_reuse( + state: &AppState, + cache_key: &str, + auth_context: GatewayControlAuthContext, + auth_endpoint_signature: Option<&str>, +) -> Result { + match state.auth_context_cache.register_inflight(cache_key) { + AuthContextInflightRegistration::Leader(_guard) => { + let refreshed = refresh_execution_runtime_auth_context( + state, + auth_context, + auth_endpoint_signature, + ) + .await?; + put_cached_auth_context(state, cache_key.to_string(), refreshed.clone()); + Ok(refreshed) + } + AuthContextInflightRegistration::Follower => Ok(auth_context), + AuthContextInflightRegistration::Bypass => { + refresh_execution_runtime_auth_context(state, auth_context, auth_endpoint_signature) + .await + } + } +} + +async fn resolve_data_backed_auth_context_cached( + state: &AppState, + cache_key: Option<&str>, + headers: &http::HeaderMap, + uri: &Uri, + auth_endpoint_signature: Option<&str>, + cache_negative: bool, +) -> Result, GatewayError> { + let Some(cache_key) = cache_key else { + return resolve_data_backed_auth_context(state, headers, uri, auth_endpoint_signature) + .await; + }; + loop { + let notified = state.auth_context_cache.notified(); + match state.auth_context_cache.register_inflight(cache_key) { + AuthContextInflightRegistration::Leader(_guard) => { + let resolved = + resolve_data_backed_auth_context(state, headers, uri, auth_endpoint_signature) + .await?; + if let Some(auth_context) = resolved.as_ref() { + if cache_negative + || (!auth_context.user_id.is_empty() && !auth_context.api_key_id.is_empty()) + { + put_cached_auth_context(state, cache_key.to_string(), auth_context.clone()); + } + } + return Ok(resolved); + } + AuthContextInflightRegistration::Follower => { + notified.await; + if let Some(auth_context) = get_cached_auth_context(state, cache_key) { + return Ok(Some(auth_context)); + } + if !cache_negative { + return Ok(None); + } + } + AuthContextInflightRegistration::Bypass => { + return resolve_data_backed_auth_context( + state, + headers, + uri, + auth_endpoint_signature, + ) + .await; + } + } + } +} + pub(crate) async fn refresh_execution_runtime_auth_context( state: &AppState, auth_context: GatewayControlAuthContext, @@ -1176,6 +1291,7 @@ fn get_cached_auth_context(state: &AppState, cache_key: &str) -> Option Result, DataLayerError> { - let snapshot = read_resolved_auth_api_key_snapshot_by_user_api_key_ids( - self, - user_id, - api_key_id, - now_unix_secs, + let snapshot = crate::request_diagnostics::observe_db_operation( + "auth_api_key_snapshot", + self.database_pool_summary(), + read_resolved_auth_api_key_snapshot_by_user_api_key_ids( + self, + user_id, + api_key_id, + now_unix_secs, + ), ) .await?; self.apply_user_group_effective_policies(snapshot).await @@ -1763,8 +1767,12 @@ impl GatewayDataState { key_hash: &str, now_unix_secs: u64, ) -> Result, DataLayerError> { - let snapshot = - read_resolved_auth_api_key_snapshot_by_key_hash(self, key_hash, now_unix_secs).await?; + let snapshot = crate::request_diagnostics::observe_db_operation( + "auth_api_key_snapshot_by_hash", + self.database_pool_summary(), + read_resolved_auth_api_key_snapshot_by_key_hash(self, key_hash, now_unix_secs), + ) + .await?; self.apply_user_group_effective_policies(snapshot).await } @@ -1782,7 +1790,13 @@ impl GatewayDataState { let Some(repository) = self.user_reader.as_ref() else { return Ok(Some(snapshot)); }; - let Some(user) = repository.find_user_auth_by_id(&snapshot.user_id).await? else { + let Some(user) = crate::request_diagnostics::observe_db_operation( + "auth_user_policy", + self.database_pool_summary(), + repository.find_user_auth_by_id(&snapshot.user_id), + ) + .await? + else { return Ok(Some(snapshot)); }; if user.role.eq_ignore_ascii_case("admin") && !snapshot.api_key_is_standalone { diff --git a/apps/aether-gateway/src/data/state/catalog.rs b/apps/aether-gateway/src/data/state/catalog.rs index 5661c54c1..064d66fc5 100644 --- a/apps/aether-gateway/src/data/state/catalog.rs +++ b/apps/aether-gateway/src/data/state/catalog.rs @@ -107,10 +107,17 @@ impl GatewayDataState { &self, candidate: UpsertRequestCandidateRecord, ) -> Result, DataLayerError> { - match &self.request_candidate_writer { - Some(repository) => repository.upsert(candidate).await.map(Some), - None => Ok(None), - } + crate::request_diagnostics::observe_db_operation( + "request_candidate_upsert", + self.database_pool_summary(), + async { + match &self.request_candidate_writer { + Some(repository) => repository.upsert(candidate).await.map(Some), + None => Ok(None), + } + }, + ) + .await } pub(crate) async fn delete_request_candidates_created_before( diff --git a/apps/aether-gateway/src/data/state/core.rs b/apps/aether-gateway/src/data/state/core.rs index 8a0a67c86..e129ee0fc 100644 --- a/apps/aether-gateway/src/data/state/core.rs +++ b/apps/aether-gateway/src/data/state/core.rs @@ -3,17 +3,77 @@ use aether_data_contracts::repository::candidate_selection::MinimalCandidateSele use aether_data_contracts::repository::candidates::RequestCandidateReadRepository; use aether_data_contracts::repository::provider_catalog::ProviderCatalogReadRepository; use aether_runtime_state::RuntimeQueueStore; -use std::sync::{Arc, OnceLock}; +use std::collections::HashSet; +use std::sync::Arc; use std::time::{Duration, Instant}; -use tokio::sync::Mutex; +use tokio::sync::Notify; use super::{GatewayDataConfig, GatewayDataState, StoredSystemConfigEntry}; const SYSTEM_CONFIG_VALUE_CACHE_TTL: Duration = Duration::from_secs(30); -fn system_config_value_load_guard() -> &'static Mutex<()> { - static GUARD: OnceLock> = OnceLock::new(); - GUARD.get_or_init(|| Mutex::new(())) +fn system_config_value_load_state() -> &'static SystemConfigValueLoadState { + static STATE: std::sync::OnceLock = std::sync::OnceLock::new(); + STATE.get_or_init(SystemConfigValueLoadState::default) +} + +#[derive(Debug, Default)] +struct SystemConfigValueLoadState { + inflight: std::sync::Mutex>, + notify: Notify, +} + +enum SystemConfigValueLoadRegistration<'a> { + Leader(SystemConfigValueLoadGuard<'a>), + Follower, + Bypass, +} + +struct SystemConfigValueLoadGuard<'a> { + state: &'a SystemConfigValueLoadState, + key: Option, +} + +impl Drop for SystemConfigValueLoadGuard<'_> { + fn drop(&mut self) { + if let Some(key) = self.key.take() { + self.state.finish(&key); + } + } +} + +impl SystemConfigValueLoadState { + fn register(&self, key: &str) -> SystemConfigValueLoadRegistration<'_> { + match self.inflight.lock() { + Ok(mut inflight) => { + if inflight.contains(key) { + SystemConfigValueLoadRegistration::Follower + } else { + inflight.insert(key.to_string()); + SystemConfigValueLoadRegistration::Leader(SystemConfigValueLoadGuard { + state: self, + key: Some(key.to_string()), + }) + } + } + Err(_) => SystemConfigValueLoadRegistration::Bypass, + } + } + + fn notified(&self) -> tokio::sync::futures::Notified<'_> { + self.notify.notified() + } + + fn finish(&self, key: &str) { + let removed = self + .inflight + .lock() + .map(|mut inflight| inflight.remove(key)) + .unwrap_or(false); + if removed { + self.notify.notify_waiters(); + } + } } fn current_system_config_updated_at_unix_secs() -> u64 { @@ -447,27 +507,69 @@ impl GatewayDataState { return Ok(value); } } - let _guard = system_config_value_load_guard().lock().await; - let cached_value = self - .system_config_value_cache - .read() - .expect("system config value cache lock") - .get(key) - .cloned(); - if let Some((cached_at, value)) = cached_value { - if cached_at.elapsed() <= SYSTEM_CONFIG_VALUE_CACHE_TTL { - return Ok(value); + let load_state = system_config_value_load_state(); + loop { + let notified = load_state.notified(); + match load_state.register(key) { + SystemConfigValueLoadRegistration::Bypass => { + let Some(backends) = self.backends.as_ref() else { + return Ok(None); + }; + let value = crate::request_diagnostics::observe_db_operation( + "system_config_value", + self.database_pool_summary(), + backends.find_system_config_value(key), + ) + .await?; + self.system_config_value_cache + .write() + .expect("system config value cache lock") + .insert(key.to_string(), (Instant::now(), value.clone())); + return Ok(value); + } + SystemConfigValueLoadRegistration::Follower => { + notified.await; + let cached_value = self + .system_config_value_cache + .read() + .expect("system config value cache lock") + .get(key) + .cloned(); + if let Some((cached_at, value)) = cached_value { + if cached_at.elapsed() <= SYSTEM_CONFIG_VALUE_CACHE_TTL { + return Ok(value); + } + } + } + SystemConfigValueLoadRegistration::Leader(_guard) => { + let cached_value = self + .system_config_value_cache + .read() + .expect("system config value cache lock") + .get(key) + .cloned(); + if let Some((cached_at, value)) = cached_value { + if cached_at.elapsed() <= SYSTEM_CONFIG_VALUE_CACHE_TTL { + return Ok(value); + } + } + let Some(backends) = self.backends.as_ref() else { + return Ok(None); + }; + let value = crate::request_diagnostics::observe_db_operation( + "system_config_value", + self.database_pool_summary(), + backends.find_system_config_value(key), + ) + .await?; + self.system_config_value_cache + .write() + .expect("system config value cache lock") + .insert(key.to_string(), (Instant::now(), value.clone())); + return Ok(value); + } } } - let Some(backends) = self.backends.as_ref() else { - return Ok(None); - }; - let value = backends.find_system_config_value(key).await?; - self.system_config_value_cache - .write() - .expect("system config value cache lock") - .insert(key.to_string(), (Instant::now(), value.clone())); - Ok(value) } pub(crate) async fn upsert_system_config_value( diff --git a/apps/aether-gateway/src/data/state/mod.rs b/apps/aether-gateway/src/data/state/mod.rs index 7f4c777c0..59dd21887 100644 --- a/apps/aether-gateway/src/data/state/mod.rs +++ b/apps/aether-gateway/src/data/state/mod.rs @@ -197,9 +197,7 @@ pub(crate) struct GatewayDataState { settlement_writer: Option>, system_config_values: Option>>>, system_config_value_cache: Arc)>>>, - billing_model_context_cache: Arc< - RwLock)>>, - >, + billing_model_context_cache: Arc, } #[derive(Clone, Debug, Eq, Hash, PartialEq)] @@ -216,6 +214,15 @@ pub(super) enum BillingModelContextCacheKey { }, } +#[derive(Default)] +pub(super) struct BillingModelContextCacheState { + pub(super) entries: + RwLock)>>, + pub(super) inflight: std::sync::Mutex>, + pub(super) inflight_notify: tokio::sync::Notify, + pub(super) next_inflight_token: std::sync::atomic::AtomicU64, +} + impl fmt::Debug for GatewayDataState { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("GatewayDataState") diff --git a/apps/aether-gateway/src/data/state/models.rs b/apps/aether-gateway/src/data/state/models.rs index f8618c5ae..c979510c2 100644 --- a/apps/aether-gateway/src/data/state/models.rs +++ b/apps/aether-gateway/src/data/state/models.rs @@ -15,14 +15,24 @@ impl GatewayDataState { api_format: &str, global_model_name: &str, ) -> Result, DataLayerError> { - match &self.minimal_candidate_selection_reader { - Some(repository) => { - repository - .list_for_exact_api_format_and_global_model(api_format, global_model_name) - .await - } - None => Ok(Vec::new()), - } + crate::request_diagnostics::observe_db_operation( + "candidate_selection", + self.database_pool_summary(), + async { + match &self.minimal_candidate_selection_reader { + Some(repository) => { + repository + .list_for_exact_api_format_and_global_model( + api_format, + global_model_name, + ) + .await + } + None => Ok(Vec::new()), + } + }, + ) + .await } pub(crate) async fn list_minimal_candidate_selection_rows_for_requested_model( @@ -30,38 +40,62 @@ impl GatewayDataState { api_format: &str, requested_model_name: &str, ) -> Result, DataLayerError> { - match &self.minimal_candidate_selection_reader { - Some(repository) => { - repository - .list_for_exact_api_format_and_requested_model(api_format, requested_model_name) - .await - } - None => Ok(Vec::new()), - } + crate::request_diagnostics::observe_db_operation( + "candidate_selection", + self.database_pool_summary(), + async { + match &self.minimal_candidate_selection_reader { + Some(repository) => { + repository + .list_for_exact_api_format_and_requested_model( + api_format, + requested_model_name, + ) + .await + } + None => Ok(Vec::new()), + } + }, + ) + .await } pub(crate) async fn list_minimal_candidate_selection_rows_for_requested_model_page( &self, query: &StoredRequestedModelCandidateRowsQuery, ) -> Result, DataLayerError> { - match &self.minimal_candidate_selection_reader { - Some(repository) => { - repository - .list_for_exact_api_format_and_requested_model_page(query) - .await - } - None => Ok(Vec::new()), - } + crate::request_diagnostics::observe_db_operation( + "candidate_selection", + self.database_pool_summary(), + async { + match &self.minimal_candidate_selection_reader { + Some(repository) => { + repository + .list_for_exact_api_format_and_requested_model_page(query) + .await + } + None => Ok(Vec::new()), + } + }, + ) + .await } pub(crate) async fn list_minimal_candidate_selection_rows_for_api_format( &self, api_format: &str, ) -> Result, DataLayerError> { - match &self.minimal_candidate_selection_reader { - Some(repository) => repository.list_for_exact_api_format(api_format).await, - None => Ok(Vec::new()), - } + crate::request_diagnostics::observe_db_operation( + "candidate_selection", + self.database_pool_summary(), + async { + match &self.minimal_candidate_selection_reader { + Some(repository) => repository.list_for_exact_api_format(api_format).await, + None => Ok(Vec::new()), + } + }, + ) + .await } pub(crate) async fn list_pool_key_candidate_rows_for_group( diff --git a/apps/aether-gateway/src/data/state/runtime.rs b/apps/aether-gateway/src/data/state/runtime.rs index d43bdd3fb..9cebfc6cf 100644 --- a/apps/aether-gateway/src/data/state/runtime.rs +++ b/apps/aether-gateway/src/data/state/runtime.rs @@ -43,6 +43,7 @@ use aether_data_contracts::repository::usage::{ use aether_runtime_state::RuntimeQueueStore; use aether_video_tasks_core::read_data_backed_video_task_response; use std::time::{Duration, Instant}; +use tokio::time::timeout; fn normalize_billing_context_cache_part(value: &str) -> String { value.trim().to_string() @@ -55,12 +56,51 @@ fn normalize_optional_billing_context_cache_part(value: Option<&str>) -> Option< .map(ToOwned::to_owned) } +enum BillingModelContextInflightRegistration { + Leader(u64), + Follower, + Bypass, +} + +struct BillingModelContextInflightGuard<'a> { + state: &'a GatewayDataState, + key: Option, + token: u64, +} + +impl<'a> BillingModelContextInflightGuard<'a> { + fn new(state: &'a GatewayDataState, key: BillingModelContextCacheKey, token: u64) -> Self { + Self { + state, + key: Some(key), + token, + } + } + + fn finish(&mut self) { + if let Some(key) = self.key.take() { + self.state + .finish_billing_model_context_inflight(&key, self.token); + } + } +} + +impl Drop for BillingModelContextInflightGuard<'_> { + fn drop(&mut self) { + self.finish(); + } +} + impl GatewayDataState { const MAINTENANCE_POOL_IDLE_RESERVE_ENV: &'static str = "AETHER_GATEWAY_MAINTENANCE_POOL_IDLE_RESERVE"; const MAINTENANCE_POOL_PRESSURE_MAX_DEFER: Duration = Duration::from_secs(30); const BILLING_MODEL_CONTEXT_CACHE_TTL: Duration = Duration::from_secs(30); const BILLING_MODEL_CONTEXT_CACHE_MAX_ENTRIES: usize = 4096; + #[cfg(not(test))] + const BILLING_MODEL_CONTEXT_CACHE_INFLIGHT_WAIT_TIMEOUT: Duration = Duration::from_secs(10); + #[cfg(test)] + const BILLING_MODEL_CONTEXT_CACHE_INFLIGHT_WAIT_TIMEOUT: Duration = Duration::from_millis(100); pub(crate) async fn run_database_maintenance( &self, @@ -1041,10 +1081,17 @@ impl GatewayDataState { &self, usage: UpsertUsageRecord, ) -> Result, DataLayerError> { - match &self.usage_writer { - Some(repository) => repository.upsert(usage).await.map(Some), - None => Ok(None), - } + crate::request_diagnostics::observe_db_operation( + "usage_upsert", + self.database_pool_summary(), + async { + match &self.usage_writer { + Some(repository) => repository.upsert(usage).await.map(Some), + None => Ok(None), + } + }, + ) + .await } #[allow(dead_code)] @@ -1710,17 +1757,47 @@ impl GatewayDataState { if let Some(value) = self.cached_billing_model_context(&key) { return Ok(value); } - match &self.billing_reader { - Some(repository) => { - let value = repository - .find_model_context(provider_id, provider_api_key_id, global_model_name) - .await?; - self.remember_billing_model_context(key, value.clone()); - Ok(value) - } - None => { - self.remember_billing_model_context(key, None); - Ok(None) + loop { + let notified = self.billing_model_context_cache.inflight_notify.notified(); + match self.register_billing_model_context_inflight(&key) { + BillingModelContextInflightRegistration::Bypass => { + return self + .load_billing_model_context_by_name( + key, + provider_id, + provider_api_key_id, + global_model_name, + ) + .await; + } + BillingModelContextInflightRegistration::Follower => { + if timeout( + Self::BILLING_MODEL_CONTEXT_CACHE_INFLIGHT_WAIT_TIMEOUT, + notified, + ) + .await + .is_err() + { + self.expire_billing_model_context_inflight(&key); + } + if let Some(value) = self.cached_billing_model_context(&key) { + return Ok(value); + } + continue; + } + BillingModelContextInflightRegistration::Leader(token) => { + let mut guard = BillingModelContextInflightGuard::new(self, key.clone(), token); + let result = self + .load_billing_model_context_by_name( + key, + provider_id, + provider_api_key_id, + global_model_name, + ) + .await; + guard.finish(); + return result; + } } } } @@ -1739,18 +1816,164 @@ impl GatewayDataState { if let Some(value) = self.cached_billing_model_context(&key) { return Ok(value); } - match &self.billing_reader { - Some(repository) => { - let value = repository - .find_model_context_by_model_id(provider_id, provider_api_key_id, model_id) - .await?; - self.remember_billing_model_context(key, value.clone()); - Ok(value) + loop { + let notified = self.billing_model_context_cache.inflight_notify.notified(); + match self.register_billing_model_context_inflight(&key) { + BillingModelContextInflightRegistration::Bypass => { + return self + .load_billing_model_context_by_model_id( + key, + provider_id, + provider_api_key_id, + model_id, + ) + .await; + } + BillingModelContextInflightRegistration::Follower => { + if timeout( + Self::BILLING_MODEL_CONTEXT_CACHE_INFLIGHT_WAIT_TIMEOUT, + notified, + ) + .await + .is_err() + { + self.expire_billing_model_context_inflight(&key); + } + if let Some(value) = self.cached_billing_model_context(&key) { + return Ok(value); + } + continue; + } + BillingModelContextInflightRegistration::Leader(token) => { + let mut guard = BillingModelContextInflightGuard::new(self, key.clone(), token); + let result = self + .load_billing_model_context_by_model_id( + key, + provider_id, + provider_api_key_id, + model_id, + ) + .await; + guard.finish(); + return result; + } } - None => { - self.remember_billing_model_context(key, None); - Ok(None) + } + } + + async fn load_billing_model_context_by_name( + &self, + key: BillingModelContextCacheKey, + provider_id: &str, + provider_api_key_id: Option<&str>, + global_model_name: &str, + ) -> Result, DataLayerError> { + crate::request_diagnostics::observe_db_operation( + "billing_model_context", + self.database_pool_summary(), + async { + match &self.billing_reader { + Some(repository) => { + let value = repository + .find_model_context(provider_id, provider_api_key_id, global_model_name) + .await?; + self.remember_billing_model_context(key, value.clone()); + Ok(value) + } + None => { + self.remember_billing_model_context(key, None); + Ok(None) + } + } + }, + ) + .await + } + + async fn load_billing_model_context_by_model_id( + &self, + key: BillingModelContextCacheKey, + provider_id: &str, + provider_api_key_id: Option<&str>, + model_id: &str, + ) -> Result, DataLayerError> { + crate::request_diagnostics::observe_db_operation( + "billing_model_context", + self.database_pool_summary(), + async { + match &self.billing_reader { + Some(repository) => { + let value = repository + .find_model_context_by_model_id( + provider_id, + provider_api_key_id, + model_id, + ) + .await?; + self.remember_billing_model_context(key, value.clone()); + Ok(value) + } + None => { + self.remember_billing_model_context(key, None); + Ok(None) + } + } + }, + ) + .await + } + + fn register_billing_model_context_inflight( + &self, + key: &BillingModelContextCacheKey, + ) -> BillingModelContextInflightRegistration { + match self.billing_model_context_cache.inflight.lock() { + Ok(mut inflight) => { + if inflight.contains_key(key) { + return BillingModelContextInflightRegistration::Follower; + } + let token = self + .billing_model_context_cache + .next_inflight_token + .fetch_add(1, std::sync::atomic::Ordering::AcqRel); + inflight.insert(key.clone(), token); + BillingModelContextInflightRegistration::Leader(token) } + Err(_) => BillingModelContextInflightRegistration::Bypass, + } + } + + fn finish_billing_model_context_inflight(&self, key: &BillingModelContextCacheKey, token: u64) { + let mut removed = false; + if let Ok(mut inflight) = self.billing_model_context_cache.inflight.lock() { + if inflight.get(key).copied() == Some(token) { + inflight.remove(key); + removed = true; + } + } + if removed { + self.billing_model_context_cache + .inflight_notify + .notify_waiters(); + } + } + + fn expire_billing_model_context_inflight(&self, key: &BillingModelContextCacheKey) { + let mut removed = false; + if let Ok(mut inflight) = self.billing_model_context_cache.inflight.lock() { + removed = inflight.remove(key).is_some(); + } + if removed { + tracing::warn!( + event_name = "billing_model_context_cache_inflight_expired", + log_type = "ops", + cache_key = ?key, + wait_timeout_ms = Self::BILLING_MODEL_CONTEXT_CACHE_INFLIGHT_WAIT_TIMEOUT.as_millis() as u64, + "gateway billing model context cache expired stale inflight load" + ); + self.billing_model_context_cache + .inflight_notify + .notify_waiters(); } } @@ -1759,6 +1982,7 @@ impl GatewayDataState { key: &BillingModelContextCacheKey, ) -> Option> { self.billing_model_context_cache + .entries .read() .expect("billing model context cache lock") .get(key) @@ -1775,6 +1999,7 @@ impl GatewayDataState { ) { let mut cache = self .billing_model_context_cache + .entries .write() .expect("billing model context cache lock"); cache.retain(|_, (cached_at, _)| { @@ -1794,9 +2019,25 @@ impl GatewayDataState { fn clear_billing_model_context_cache(&self) { self.billing_model_context_cache + .entries .write() .expect("billing model context cache lock") .clear(); + let mut cleared_inflight = false; + if let Ok(mut inflight) = self.billing_model_context_cache.inflight.lock() { + cleared_inflight = !inflight.is_empty(); + inflight.clear(); + } + if cleared_inflight { + tracing::warn!( + event_name = "billing_model_context_cache_inflight_cleared", + log_type = "ops", + "gateway billing model context cache cleared in-flight loads" + ); + self.billing_model_context_cache + .inflight_notify + .notify_waiters(); + } } pub(crate) async fn admin_billing_enabled_default_value_exists( @@ -2211,12 +2452,89 @@ impl GatewayDataState { #[cfg(test)] mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; + use std::time::Duration; use aether_data::repository::users::{InMemoryUserReadRepository, StoredUserExportRow}; + use aether_data_contracts::repository::billing::{ + BillingReadRepository, StoredBillingModelContext, + }; + use async_trait::async_trait; + use serde_json::json; use super::GatewayDataState; + struct SlowBillingContextRepository { + calls: AtomicUsize, + context: StoredBillingModelContext, + } + + #[async_trait] + impl BillingReadRepository for SlowBillingContextRepository { + async fn find_model_context( + &self, + _provider_id: &str, + _provider_api_key_id: Option<&str>, + _global_model_name: &str, + ) -> Result, aether_data_contracts::DataLayerError> + { + self.calls.fetch_add(1, Ordering::AcqRel); + tokio::time::sleep(Duration::from_millis(25)).await; + Ok(Some(self.context.clone())) + } + } + + fn billing_context() -> StoredBillingModelContext { + StoredBillingModelContext::new( + "provider-1".to_string(), + Some("pay_as_you_go".to_string()), + Some("key-1".to_string()), + None, + None, + "global-model-1".to_string(), + "gpt-5".to_string(), + None, + Some(0.02), + Some(json!({"tiers":[{"up_to":null,"input_price_per_1m":3.0,"output_price_per_1m":15.0}]})), + Some("model-1".to_string()), + Some("gpt-5-upstream".to_string()), + None, + None, + None, + ) + .expect("billing context should build") + } + + #[tokio::test] + async fn billing_model_context_cache_coalesces_concurrent_loads() { + let repository = Arc::new(SlowBillingContextRepository { + calls: AtomicUsize::new(0), + context: billing_context(), + }); + let state = Arc::new(GatewayDataState::with_billing_reader_for_tests( + repository.clone(), + )); + let mut tasks = Vec::new(); + + for _ in 0..16 { + let state = Arc::clone(&state); + tasks.push(tokio::spawn(async move { + state + .find_billing_model_context("provider-1", Some("key-1"), "gpt-5") + .await + .expect("billing context lookup should succeed") + .expect("billing context should exist"); + })); + } + + for task in tasks { + task.await.expect("lookup task should complete"); + } + + assert_eq!(repository.calls.load(Ordering::Acquire), 1); + } + #[tokio::test] async fn lists_non_admin_export_users_from_user_reader() { let repository = Arc::new(InMemoryUserReadRepository::seed_export_users(vec![ diff --git a/apps/aether-gateway/src/dispatch/pool_scheduler.rs b/apps/aether-gateway/src/dispatch/pool_scheduler.rs index f8268d950..4900a6b49 100644 --- a/apps/aether-gateway/src/dispatch/pool_scheduler.rs +++ b/apps/aether-gateway/src/dispatch/pool_scheduler.rs @@ -47,6 +47,7 @@ use crate::handlers::shared::provider_pool::{ use crate::handlers::shared::{parse_catalog_auth_config_json, provider_key_health_summary}; use crate::maintenance::spawn_pool_quota_probe_replenish_for_request; use crate::orchestration::LocalExecutionCandidateMetadata; +use crate::stage_metrics::observe_gateway_stage_ms; static LOAD_BALANCE_SEQUENCE: AtomicU64 = AtomicU64::new(0); static POOL_SCORE_SCHEDULE_INTEREST_SEMAPHORE: LazyLock> = @@ -133,14 +134,20 @@ async fn schedule_pool_page_candidates( let runtime = if key_ids.is_empty() { AdminProviderPoolRuntimeState::default() } else { - read_admin_provider_pool_runtime_state( + let runtime_started_at = std::time::Instant::now(); + let runtime = read_admin_provider_pool_runtime_state( state.app().runtime_state.as_ref(), provider_id.as_str(), &key_ids, &pool_config, sticky_session_token, ) - .await + .await; + observe_gateway_stage_ms( + "pool_runtime_state", + runtime_started_at.elapsed().as_millis() as u64, + ); + runtime }; pool_config_by_provider.insert(provider_id.clone(), pool_config); runtime_by_provider.insert(provider_id, runtime); @@ -449,9 +456,17 @@ impl<'a> PoolKeyCursor<'a> { } pub(crate) async fn next_key(&mut self) -> Option { + let started_at = std::time::Instant::now(); + let mut observed = false; loop { if let Some(candidate) = self.next_queued_candidate().await { self.returned_key_count = self.returned_key_count.saturating_add(1); + if !observed { + observe_gateway_stage_ms( + "pool_cursor_next_key", + started_at.elapsed().as_millis() as u64, + ); + } return Some(candidate); } @@ -464,6 +479,13 @@ impl<'a> PoolKeyCursor<'a> { } if !self.refill_queued_candidates().await { + if !observed { + observe_gateway_stage_ms( + "pool_cursor_next_key", + started_at.elapsed().as_millis() as u64, + ); + observed = true; + } return None; } } @@ -634,9 +656,14 @@ impl<'a> PoolKeyCursor<'a> { offset: 0, limit: limit as usize, }; + let score_started_at = std::time::Instant::now(); let scores = match self.state.app().data.list_ranked_pool_members(&query).await { Ok(scores) => scores, Err(err) => { + observe_gateway_stage_ms( + "pool_score_load", + score_started_at.elapsed().as_millis() as u64, + ); warn!( event_name = "pool_group_score_load_failed", log_type = "event", @@ -650,6 +677,10 @@ impl<'a> PoolKeyCursor<'a> { return None; } }; + observe_gateway_stage_ms( + "pool_score_load", + score_started_at.elapsed().as_millis() as u64, + ); if scores.is_empty() { return None; } @@ -668,6 +699,7 @@ impl<'a> PoolKeyCursor<'a> { selected_provider_model_name: self.group.candidate.selected_provider_model_name.clone(), key_ids, }; + let rows_started_at = std::time::Instant::now(); let rows = match self .state .app() @@ -676,6 +708,10 @@ impl<'a> PoolKeyCursor<'a> { { Ok(rows) => rows, Err(err) => { + observe_gateway_stage_ms( + "pool_score_key_rows", + rows_started_at.elapsed().as_millis() as u64, + ); warn!( event_name = "pool_group_score_key_load_failed", log_type = "event", @@ -690,6 +726,10 @@ impl<'a> PoolKeyCursor<'a> { return None; } }; + observe_gateway_stage_ms( + "pool_score_key_rows", + rows_started_at.elapsed().as_millis() as u64, + ); let materialized_row_count = rows.len() as u32; let missing_score_count = scores.len().saturating_sub(rows.len()); if missing_score_count > 0 { @@ -1012,11 +1052,20 @@ impl<'a> PoolKeyCursor<'a> { return None; } + let transport_started_at = std::time::Instant::now(); let Some(transport) = read_candidate_transport_snapshot(self.state, &candidate).await else { + observe_gateway_stage_ms( + "candidate_transport_snapshot", + transport_started_at.elapsed().as_millis() as u64, + ); self.record_skip_reason("transport_snapshot_missing"); return None; }; + observe_gateway_stage_ms( + "candidate_transport_snapshot", + transport_started_at.elapsed().as_millis() as u64, + ); if let Some(skip_reason) = candidate_auth_channel_skip_reason(&transport, self.request_auth_channel.as_deref()) { diff --git a/apps/aether-gateway/src/error.rs b/apps/aether-gateway/src/error.rs index 2cc789013..70b346ccc 100644 --- a/apps/aether-gateway/src/error.rs +++ b/apps/aether-gateway/src/error.rs @@ -24,6 +24,11 @@ pub(crate) enum GatewayError { phase: &'static str, timeout_ms: u64, }, + AdmissionTimeout { + trace_id: String, + gate: &'static str, + queue_budget_ms: u64, + }, Client { status: StatusCode, message: String, @@ -43,6 +48,13 @@ impl GatewayError { } => { format!("local execution planning timed out in {phase} after {timeout_ms}ms") } + Self::AdmissionTimeout { + gate, + queue_budget_ms, + .. + } => { + format!("gateway admission gate {gate} timed out after {queue_budget_ms}ms") + } } } } @@ -113,6 +125,34 @@ impl IntoResponse for GatewayError { ); response } + Self::AdmissionTimeout { + trace_id, + gate, + queue_budget_ms, + } => { + tracing::debug!( + trace_id = %trace_id, + gate, + queue_budget_ms, + "gateway admission gate timed out" + ); + let body = Json(json!({ + "error": { + "message": "gateway admission queue timed out", + "trace_id": trace_id, + } + })); + let mut response = (StatusCode::TOO_MANY_REQUESTS, body).into_response(); + let _ = + insert_header_if_missing(response.headers_mut(), TRACE_ID_HEADER, &trace_id); + let _ = insert_header_if_missing( + response.headers_mut(), + GATEWAY_HEADER, + "rust-phase3b", + ); + let _ = insert_header_if_missing(response.headers_mut(), "Retry-After", "1"); + response + } Self::Client { status, message } => ( status, Json(json!({ @@ -140,3 +180,41 @@ impl From for GatewayError { GatewayError::Internal(error.0) } } + +#[cfg(test)] +mod tests { + use axum::http::{header::RETRY_AFTER, StatusCode}; + use axum::response::IntoResponse; + + use crate::constants::TRACE_ID_HEADER; + + use super::GatewayError; + + #[test] + fn admission_timeout_returns_429_with_retry_after_without_panicking() { + let trace_id = "trace-admission-timeout".to_string(); + + let response = GatewayError::AdmissionTimeout { + trace_id: trace_id.clone(), + gate: "gateway_upstream_execution", + queue_budget_ms: 250, + } + .into_response(); + + assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS); + assert_eq!( + response + .headers() + .get(RETRY_AFTER) + .and_then(|v| v.to_str().ok()), + Some("1") + ); + assert_eq!( + response + .headers() + .get(TRACE_ID_HEADER) + .and_then(|v| v.to_str().ok()), + Some(trace_id.as_str()) + ); + } +} diff --git a/apps/aether-gateway/src/execution_runtime/stream/execution.rs b/apps/aether-gateway/src/execution_runtime/stream/execution.rs index a6c34695d..db903f439 100644 --- a/apps/aether-gateway/src/execution_runtime/stream/execution.rs +++ b/apps/aether-gateway/src/execution_runtime/stream/execution.rs @@ -1,4 +1,5 @@ use std::collections::{BTreeMap, VecDeque}; +use std::future::Future; use std::io::Error as IoError; use std::sync::{ atomic::{AtomicBool, AtomicU64, Ordering}, @@ -17,8 +18,8 @@ use aether_scheduler_core::{ }; use aether_usage_runtime::{ build_lifecycle_usage_seed, build_stream_terminal_usage_payload_seed, - build_sync_terminal_usage_payload_seed, build_terminal_usage_context_seed, LifecycleUsageSeed, - UsageBodyCapturePolicy, UsageRequestRecordLevel, + build_sync_terminal_usage_payload_seed, build_terminal_usage_context_seed, + build_usage_event_data_seed, LifecycleUsageSeed, UsageEvent, UsageEventType, DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES, }; use async_stream::stream; @@ -79,9 +80,11 @@ use crate::execution_runtime::submission::{ strip_utf8_bom_and_ws, submit_local_core_error_or_sync_finalize, }; use crate::execution_runtime::transport::{ - execute_stream_plan_via_local_tunnel, record_manual_proxy_request_failure, + execute_stream_plan_via_local_tunnel, format_upstream_request_error, + format_wreq_upstream_request_error, record_manual_proxy_request_failure, record_manual_proxy_request_success, record_manual_proxy_stream_error, - DirectSyncExecutionRuntime, DirectUpstreamStreamExecution, ExecutionRuntimeTransportError, + stream_first_byte_timeout_message, DirectSyncExecutionRuntime, DirectUpstreamResponse, + DirectUpstreamStreamExecution, ExecutionRuntimeTransportError, }; use crate::execution_runtime::windsurf::maybe_execute_windsurf_stream; use crate::execution_runtime::{ @@ -107,6 +110,14 @@ use crate::request_candidate_runtime::{ ensure_execution_request_candidate_slot, record_local_request_candidate_status, record_local_request_candidate_status_snapshot, snapshot_local_request_candidate_status, }; +use crate::request_diagnostics::{ + attach_current_request_diagnostics_to_report_context, + attach_request_diagnostics_to_report_context, current_request_diagnostics, RequestDiagnostics, +}; +use crate::stage_metrics::{ + attach_stage_trace_to_report_context, observe_gateway_stage_ms, observe_gateway_stage_trace_ms, + RequestStageTrace, +}; use crate::usage::submit_stream_report; use crate::usage::{GatewayStreamReportRequest, GatewaySyncReportRequest}; use crate::{ @@ -122,13 +133,74 @@ const STREAM_IDLE_LOG_INTERVAL: Duration = Duration::from_secs(60); const STREAM_IDLE_LOG_INTERVAL_MS: u64 = 60_000; const REWRITTEN_STREAM_PREFETCH_TIMEOUT: Duration = Duration::from_millis(750); +struct StageElapsedGuard { + stage: &'static str, + started_at: Instant, +} + +#[derive(Debug)] +enum InProcessStreamExecutionError { + Transport(ExecutionRuntimeTransportError), + Gateway(GatewayError), +} + +impl From for InProcessStreamExecutionError { + fn from(error: ExecutionRuntimeTransportError) -> Self { + Self::Transport(error) + } +} + +impl From for InProcessStreamExecutionError { + fn from(error: GatewayError) -> Self { + Self::Gateway(error) + } +} + +impl StageElapsedGuard { + fn from_started_at(stage: &'static str, started_at: Instant) -> Self { + Self { stage, started_at } + } +} + +fn report_context_with_stage_trace( + report_context: Option, + mut stage_trace: RequestStageTrace, + stream_started_at: Instant, + terminal_telemetry: Option<&ExecutionTelemetry>, +) -> Option { + stage_trace.observe("stream_total", stream_elapsed_ms_since(stream_started_at)); + let fallback_elapsed_ms = terminal_telemetry.and_then(|telemetry| telemetry.ttfb_ms); + attach_stage_trace_to_report_context( + report_context, + stage_trace.into_metadata_value(fallback_elapsed_ms), + ) +} + +fn report_context_with_request_diagnostics( + report_context: Option, + diagnostics: Option<&Arc>, +) -> Option { + attach_request_diagnostics_to_report_context(report_context, diagnostics) +} + +impl Drop for StageElapsedGuard { + fn drop(&mut self) { + observe_gateway_stage_ms(self.stage, self.started_at.elapsed().as_millis() as u64); + } +} + fn record_sync_terminal_usage( state: &AppState, plan: &ExecutionPlan, report_context: Option<&serde_json::Value>, payload: &GatewaySyncReportRequest, ) { - let context_seed = build_terminal_usage_context_seed(plan, report_context); + let report_context_with_diagnostics = + attach_current_request_diagnostics_to_report_context(report_context); + let context_seed = build_terminal_usage_context_seed( + plan, + report_context_with_diagnostics.as_ref().or(report_context), + ); let payload_seed = build_sync_terminal_usage_payload_seed(payload); state .usage_runtime @@ -221,6 +293,68 @@ fn record_stream_terminal_usage( ); } +async fn record_stream_admission_timeout_terminal_state( + state: &AppState, + plan: &ExecutionPlan, + report_context: Option<&Value>, + candidate_started_unix_ms: u64, + error: &GatewayError, +) { + let status_code = 429; + let error_type = "gateway_admission_timeout"; + let error_message = match error { + GatewayError::AdmissionTimeout { + gate, + queue_budget_ms, + .. + } => format!("gateway admission gate {gate} timed out after {queue_budget_ms}ms"), + other => format!("{other:?}"), + }; + let terminal_unix_ms = current_request_candidate_unix_ms(); + let latency_ms = terminal_unix_ms.saturating_sub(candidate_started_unix_ms); + record_local_request_candidate_status( + state, + plan, + report_context, + SchedulerRequestCandidateStatusUpdate { + status: RequestCandidateStatus::Failed, + status_code: Some(status_code), + error_type: Some(error_type.to_string()), + error_message: Some(error_message.clone()), + latency_ms: Some(latency_ms), + started_at_unix_ms: Some(candidate_started_unix_ms), + finished_at_unix_ms: Some(terminal_unix_ms), + }, + ) + .await; + + if !state.usage_runtime.is_enabled() { + return; + } + + let mut usage_data = build_usage_event_data_seed(plan, report_context); + usage_data.status_code = Some(status_code); + usage_data.error_message = Some(error_message.clone()); + usage_data.error_category = Some("client_error".to_string()); + usage_data.response_time_ms = Some(latency_ms); + let error_body = json!({ + "error": { + "type": error_type, + "message": error_message, + "code": status_code, + } + }); + usage_data.response_headers = Some(json!({"content-type": "application/json"})); + usage_data.response_body = Some(error_body.clone()); + usage_data.client_response_headers = Some(json!({"content-type": "application/json"})); + usage_data.client_response_body = Some(error_body); + + state.usage_runtime.submit_terminal_event( + state.data.as_ref(), + UsageEvent::new(UsageEventType::Failed, plan.request_id.clone(), usage_data), + ); +} + fn build_stream_body_capture( body: &[u8], truncated: bool, @@ -794,19 +928,25 @@ fn stream_terminal_summary_represents_failure_with_requirement( async fn execute_in_process_stream( state: &AppState, plan: &ExecutionPlan, -) -> Result { + trace_id: &str, +) -> Result { if let Some(execution) = execute_stream_plan_via_local_tunnel(state, plan).await? { return Ok(execution); } + let upstream_target_permit = state + .upstream_target_admission + .acquire(plan, trace_id) + .await?; match DirectSyncExecutionRuntime::new().execute_stream(plan).await { - Ok(execution) => { + Ok(mut execution) => { + execution.upstream_target_permit = upstream_target_permit; record_manual_proxy_request_success(state, plan).await; Ok(execution) } Err(error) => { record_manual_proxy_request_failure(state, plan).await; - Err(error) + Err(error.into()) } } } @@ -816,19 +956,825 @@ async fn execute_in_process_stream_with_oauth_retry( plan: &mut ExecutionPlan, trace_id: &str, report_context: Option<&Value>, -) -> Result { - let mut execution = execute_in_process_stream(state, plan).await?; +) -> Result { + let mut execution = execute_in_process_stream(state, plan, trace_id).await?; apply_stream_summary_report_context(&mut execution, report_context); if execution.status_code >= 400 && refresh_oauth_plan_auth_for_retry(state, plan, execution.status_code, None, trace_id) .await { - execution = execute_in_process_stream(state, plan).await?; + drop(execution); + execution = execute_in_process_stream(state, plan, trace_id).await?; apply_stream_summary_report_context(&mut execution, report_context); } Ok(execution) } +fn should_use_direct_sse_passthrough( + plan: &ExecutionPlan, + plan_kind: &str, + report_context: Option<&Value>, + execution: &DirectUpstreamStreamExecution, +) -> bool { + if !(200..300).contains(&execution.status_code) { + return false; + } + if plan_kind == OPENAI_IMAGE_STREAM_PLAN_KIND { + return false; + } + if !response_headers_indicate_sse(&execution.headers) { + return false; + } + if !plan + .provider_api_format + .eq_ignore_ascii_case(plan.client_api_format.as_str()) + { + return false; + } + if client_format_allows_proxy_generated_sse_control_blocks(plan) { + return false; + } + if maybe_build_provider_private_stream_normalizer(report_context).is_some() { + return false; + } + let normalized_stream_report_context = + normalize_provider_private_report_context(report_context); + if maybe_build_stream_response_rewriter(normalized_stream_report_context.as_ref()).is_some() { + return false; + } + + let direct_stream_finalize_kind = resolve_core_stream_direct_finalize_report_kind(plan_kind); + should_skip_direct_finalize_prefetch( + direct_stream_finalize_kind.as_deref(), + execution.headers.get("content-type").map(String::as_str), + plan.provider_api_format.as_str(), + plan.client_api_format.as_str(), + false, + false, + ) +} + +fn direct_upstream_response_byte_stream( + response: DirectUpstreamResponse, +) -> BoxStream<'static, Result> { + match response { + DirectUpstreamResponse::Reqwest(response) => response + .bytes_stream() + .map(|item| item.map_err(|err| format_upstream_request_error(&err))) + .boxed(), + DirectUpstreamResponse::BrowserWreq(response) => response + .bytes_stream() + .map(|item| item.map_err(|err| format_wreq_upstream_request_error(&err))) + .boxed(), + DirectUpstreamResponse::LocalTunnel(mut response) => stream! { + loop { + match response.next_chunk().await { + Ok(Some(chunk)) => yield Ok(chunk), + Ok(None) => break, + Err(err) => { + yield Err(err); + break; + } + } + } + } + .boxed(), + } +} + +async fn await_direct_passthrough_first_item( + future: F, + started_at: Instant, + timeout: Option, +) -> Result +where + F: Future, +{ + let Some(timeout) = timeout else { + return Ok(future.await); + }; + let Some(remaining) = timeout.checked_sub(started_at.elapsed()) else { + return Err(timeout); + }; + if remaining.is_zero() { + return Err(timeout); + } + tokio::time::timeout(remaining, future) + .await + .map_err(|_| timeout) +} + +#[allow(clippy::too_many_arguments)] +async fn forward_direct_passthrough_client_chunk( + tx: &mpsc::Sender>, + chunk: Bytes, + downstream_dropped: &mut bool, + client_visible_stream_completed: &mut bool, + client_stream_completion_tracker: &mut ClientVisibleStreamCompletionTracker, + client_stream_bytes: &mut u64, + buffered_body: &mut Vec, + client_body_truncated: &mut bool, + max_stream_body_buffer_bytes: usize, + stream_started_at: Instant, + last_client_chunk_elapsed_ms: &mut u64, + trace_id: &str, + request_id_for_log: &str, + candidate_id: Option<&str>, +) -> bool { + if chunk.is_empty() { + return false; + } + append_stream_capture_bytes( + buffered_body, + chunk.as_ref(), + max_stream_body_buffer_bytes, + client_body_truncated, + ); + if *downstream_dropped { + return false; + } + let chunk_len = u64::try_from(chunk.len()).unwrap_or(u64::MAX); + let send_started_at = Instant::now(); + if tx.send(Ok(chunk.clone())).await.is_err() { + debug!( + event_name = "direct_passthrough_downstream_disconnected", + log_type = "ops", + trace_id = %trace_id, + request_id = %request_id_for_log, + candidate_id = ?candidate_id, + "gateway direct passthrough downstream dropped; cancelling upstream stream" + ); + *downstream_dropped = true; + return false; + } + observe_gateway_stage_ms( + "direct_passthrough_body_send_wait", + send_started_at.elapsed().as_millis() as u64, + ); + + *client_visible_stream_completed |= + client_stream_completion_tracker.observe_chunk(chunk.as_ref()); + *client_stream_bytes = client_stream_bytes.saturating_add(chunk_len); + *last_client_chunk_elapsed_ms = stream_started_at + .elapsed() + .as_millis() + .min(u128::from(u64::MAX)) as u64; + true +} + +#[allow(clippy::too_many_arguments)] +async fn execute_stream_from_direct_passthrough( + state: &AppState, + plan: ExecutionPlan, + trace_id: &str, + decision: &GatewayControlDecision, + plan_kind: &str, + report_kind: Option, + report_context: Option, + candidate_started_unix_secs: u64, + stream_started_at: Instant, + mut stage_trace: RequestStageTrace, + execution: DirectUpstreamStreamExecution, + in_flight_guard: Option, +) -> Result>, GatewayError> { + let DirectUpstreamStreamExecution { + request_id: _, + candidate_id: _, + status_code, + mut headers, + provider_api_format: _, + stream_summary_report_context: _, + response, + started_at: upstream_started_at, + stream_first_byte_timeout, + upstream_target_permit, + } = execution; + + let request_id = plan.request_id.clone(); + let candidate_id = plan.candidate_id.clone(); + let request_id_for_log = short_request_id(request_id.as_str()); + let mut report_context = + attach_provider_response_headers_to_report_context(report_context, &headers); + if status_code == 200 { + seed_kiro_simulated_cache_enabled(state, &plan, &mut report_context).await; + if kiro_simulated_cache_enabled_from_report_context(report_context.as_ref()) { + seed_kiro_report_context_input_tokens(&plan, &mut report_context); + } + seed_kiro_report_context_prompt_cache_usage(state, &plan, &mut report_context).await; + } + + let lifecycle_seed = build_lifecycle_usage_seed(&plan, report_context.as_ref()); + state.usage_runtime.record_stream_started( + state.data.as_ref(), + &lifecycle_seed, + status_code, + None, + ); + + let request_candidate_status_snapshot = + snapshot_local_request_candidate_status(&plan, report_context.as_ref()); + if let Some(snapshot) = request_candidate_status_snapshot { + let state_bg = state.clone(); + tokio::spawn(async move { + record_local_request_candidate_status_snapshot( + &state_bg, + &snapshot, + SchedulerRequestCandidateStatusUpdate { + status: RequestCandidateStatus::Streaming, + status_code: Some(status_code), + error_type: None, + error_message: None, + latency_ms: None, + started_at_unix_ms: Some(candidate_started_unix_secs), + finished_at_unix_ms: None, + }, + ) + .await; + }); + } + + apply_endpoint_response_header_rules(state, &plan, &mut headers, None).await?; + let headers_for_report = headers.clone(); + headers.insert(CONTROL_REQUEST_ID_HEADER.to_string(), request_id.clone()); + if let Some(candidate_id) = candidate_id + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + { + headers.insert( + CONTROL_CANDIDATE_ID_HEADER.to_string(), + candidate_id.to_string(), + ); + } + headers.remove("content-length"); + + let (tx, rx) = mpsc::channel::>(128); + let state_for_report = state.clone(); + let plan_for_report = plan; + let trace_id_owned = trace_id.to_string(); + let report_kind_owned = report_kind; + let report_context_owned = report_context; + let lifecycle_seed_for_report = lifecycle_seed; + let direct_stream_finalize_kind_owned = + resolve_core_stream_direct_finalize_report_kind(plan_kind); + let normalized_stream_report_context_owned = + normalize_provider_private_report_context(report_context_owned.as_ref()); + let stream_started_at_for_report = stream_started_at; + observe_gateway_stage_trace_ms( + &mut stage_trace, + "stream_response_ready", + stream_elapsed_ms_since(stream_started_at), + ); + let stage_trace_for_report = stage_trace; + let request_diagnostics_for_report = current_request_diagnostics(); + let request_id_for_report = request_id.clone(); + let request_id_for_report_log = request_id_for_log.clone(); + let candidate_id_for_report = candidate_id.clone(); + let provider_pool_in_flight_guard_for_report = in_flight_guard; + tokio::spawn(async move { + let mut stage_trace_for_report = stage_trace_for_report; + let _stream_total_guard = + StageElapsedGuard::from_started_at("stream_total", stream_started_at_for_report); + let _provider_pool_in_flight_guard = provider_pool_in_flight_guard_for_report; + let _upstream_target_permit = upstream_target_permit; + let max_stream_body_buffer_bytes = DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES; + let stream_usage_report_context = + normalized_stream_report_context_owned.clone().or_else(|| { + Some(serde_json::json!({ + "provider_api_format": plan_for_report.provider_api_format.as_str(), + "client_api_format": plan_for_report.client_api_format.as_str(), + })) + }); + let mut stream_usage_observer = stream_usage_report_context + .as_ref() + .map(|_| StreamingStandardTerminalObserver::default()); + let mut stream_usage_observer_buffered = Vec::new(); + let mut provider_buffered_body = Vec::new(); + let mut buffered_body = Vec::new(); + let mut provider_body_truncated = false; + let mut client_body_truncated = false; + let mut upstream_control_filter = Some(SseControlBlockFilter::default()); + let mut client_stream_completion_tracker = ClientVisibleStreamCompletionTracker::default(); + let mut client_visible_stream_completed = false; + let mut usage_stream_telemetry: Option = None; + let telemetry: Option = None; + let mut provider_stream_bytes = 0u64; + let mut client_stream_bytes = 0u64; + let mut last_client_chunk_elapsed_ms = 0u64; + let mut downstream_dropped = false; + let mut terminal_failure: Option = None; + let mut upstream = direct_upstream_response_byte_stream(response); + let mut observed_first_upstream_body = false; + let mut observed_first_client_send = false; + + loop { + if downstream_dropped { + break; + } + let item = if usage_stream_telemetry + .as_ref() + .and_then(|telemetry| telemetry.ttfb_ms) + .is_none() + { + tokio::select! { + biased; + _ = tx.closed(), if !downstream_dropped => { + downstream_dropped = true; + break; + } + result = await_direct_passthrough_first_item( + upstream.next(), + upstream_started_at, + stream_first_byte_timeout, + ) => { + match result { + Ok(item) => item, + Err(timeout) => { + terminal_failure = Some(build_stream_failure_report( + "first_byte_timeout", + stream_first_byte_timeout_message(timeout), + 504, + )); + break; + } + } + } + } + } else { + tokio::select! { + biased; + _ = tx.closed(), if !downstream_dropped => { + downstream_dropped = true; + break; + } + item = upstream.next() => item, + } + }; + + let Some(item) = item else { + break; + }; + let chunk = match item { + Ok(chunk) => chunk, + Err(message) => { + warn!( + event_name = "direct_passthrough_body_read_error", + log_type = "ops", + trace_id = %trace_id_owned, + request_id = %request_id_for_report_log, + candidate_id = ?candidate_id_for_report.as_deref(), + upstream_bytes = provider_stream_bytes, + error = %message, + "gateway direct passthrough upstream body read failed" + ); + terminal_failure = Some(build_stream_failure_report( + "execution_runtime_stream_read_error", + message, + 502, + )); + break; + } + }; + if chunk.is_empty() { + continue; + } + + let observed_at = Instant::now(); + if !observed_first_upstream_body { + observed_first_upstream_body = true; + observe_gateway_stage_trace_ms( + &mut stage_trace_for_report, + "direct_passthrough_upstream_body_first", + stream_elapsed_ms_at(stream_started_at_for_report, observed_at), + ); + } + maybe_record_first_stream_event_started( + &state_for_report, + &lifecycle_seed_for_report, + status_code, + stream_started_at_for_report, + observed_at, + telemetry.as_ref(), + &mut usage_stream_telemetry, + ); + if usage_stream_telemetry + .as_ref() + .and_then(|telemetry| telemetry.ttfb_ms) + .is_some() + && provider_stream_bytes == 0 + { + observe_gateway_stage_trace_ms( + &mut stage_trace_for_report, + "stream_first_data", + stream_elapsed_ms_at(stream_started_at_for_report, observed_at), + ); + } + + let provider_chunk = chunk.clone(); + if let Some(client_chunk) = + filter_upstream_sse_control_chunk(&mut upstream_control_filter, chunk) + { + let sent_client_chunk = forward_direct_passthrough_client_chunk( + &tx, + client_chunk, + &mut downstream_dropped, + &mut client_visible_stream_completed, + &mut client_stream_completion_tracker, + &mut client_stream_bytes, + &mut buffered_body, + &mut client_body_truncated, + max_stream_body_buffer_bytes, + stream_started_at_for_report, + &mut last_client_chunk_elapsed_ms, + trace_id_owned.as_str(), + request_id_for_report_log.as_str(), + candidate_id_for_report.as_deref(), + ) + .await; + if sent_client_chunk && !observed_first_client_send { + observed_first_client_send = true; + observe_gateway_stage_trace_ms( + &mut stage_trace_for_report, + "direct_passthrough_first_client_send", + stream_elapsed_ms_since(stream_started_at_for_report), + ); + } + } + + provider_stream_bytes = provider_stream_bytes + .saturating_add(u64::try_from(provider_chunk.len()).unwrap_or(u64::MAX)); + append_stream_capture_bytes( + &mut provider_buffered_body, + provider_chunk.as_ref(), + max_stream_body_buffer_bytes, + &mut provider_body_truncated, + ); + if let (Some(observer), Some(report_context)) = ( + stream_usage_observer.as_mut(), + stream_usage_report_context.as_ref(), + ) { + observe_stream_usage_bytes( + observer, + report_context, + &mut stream_usage_observer_buffered, + provider_chunk.as_ref(), + ); + } + let provider_private_error_body_json = extract_provider_private_stream_error_body( + stream_usage_report_context.as_ref(), + provider_chunk.as_ref(), + ); + if let Some(error_body_json) = provider_private_error_body_json { + let error_status_code = + resolve_local_sync_error_status_code(status_code, &error_body_json); + terminal_failure = Some(build_stream_failure_from_provider_error_body( + error_status_code, + &error_body_json, + )); + break; + } + } + + if terminal_failure.is_none() { + if let Some(client_chunk) = + flush_upstream_sse_control_filter(&mut upstream_control_filter) + { + let _ = forward_direct_passthrough_client_chunk( + &tx, + client_chunk, + &mut downstream_dropped, + &mut client_visible_stream_completed, + &mut client_stream_completion_tracker, + &mut client_stream_bytes, + &mut buffered_body, + &mut client_body_truncated, + max_stream_body_buffer_bytes, + stream_started_at_for_report, + &mut last_client_chunk_elapsed_ms, + trace_id_owned.as_str(), + request_id_for_report_log.as_str(), + candidate_id_for_report.as_deref(), + ) + .await; + } + } + + if let Some(failure) = terminal_failure.as_ref().filter(|_| !downstream_dropped) { + match encode_terminal_sse_error_event(failure) { + Ok(error_event) => { + let _ = forward_direct_passthrough_client_chunk( + &tx, + error_event, + &mut downstream_dropped, + &mut client_visible_stream_completed, + &mut client_stream_completion_tracker, + &mut client_stream_bytes, + &mut buffered_body, + &mut client_body_truncated, + max_stream_body_buffer_bytes, + stream_started_at_for_report, + &mut last_client_chunk_elapsed_ms, + trace_id_owned.as_str(), + request_id_for_report_log.as_str(), + candidate_id_for_report.as_deref(), + ) + .await; + } + Err(err) => { + warn!( + event_name = "direct_passthrough_terminal_error_event_encode_failed", + log_type = "ops", + trace_id = %trace_id_owned, + request_id = %request_id_for_report_log, + candidate_id = ?candidate_id_for_report.as_deref(), + error = ?err, + "gateway direct passthrough failed to encode terminal SSE error event" + ); + } + } + } + drop(tx); + + let mut stream_terminal_summary = finalize_stream_usage_observer( + &mut stream_usage_observer, + stream_usage_report_context.as_ref(), + &mut stream_usage_observer_buffered, + ); + + if downstream_dropped && client_visible_stream_completed && terminal_failure.is_none() { + debug!( + event_name = "direct_passthrough_downstream_closed_after_done", + log_type = "debug", + trace_id = %trace_id_owned, + request_id = %request_id_for_report_log, + candidate_id = ?candidate_id_for_report.as_deref(), + "gateway treats direct passthrough downstream close after terminal SSE event as completed" + ); + downstream_dropped = false; + } + + if downstream_dropped { + let terminal_telemetry = Some(build_terminal_stream_telemetry( + stream_started_at_for_report, + telemetry.as_ref(), + usage_stream_telemetry.as_ref(), + provider_stream_bytes, + )); + let report_context_for_payload = report_context_with_stage_trace( + report_context_owned, + stage_trace_for_report, + stream_started_at_for_report, + terminal_telemetry.as_ref(), + ); + let report_context_for_payload = report_context_with_request_diagnostics( + report_context_for_payload, + request_diagnostics_for_report.as_ref(), + ); + let usage_payload = build_stream_usage_payload( + trace_id_owned, + report_kind_owned.unwrap_or_default(), + report_context_for_payload, + 499, + headers_for_report, + &provider_buffered_body, + provider_body_truncated, + &buffered_body, + client_body_truncated, + stream_terminal_summary, + terminal_telemetry, + ); + record_stream_terminal_usage( + &state_for_report, + &plan_for_report, + usage_payload.report_context.as_ref(), + &usage_payload, + true, + ); + record_local_request_candidate_status( + &state_for_report, + &plan_for_report, + usage_payload.report_context.as_ref(), + SchedulerRequestCandidateStatusUpdate { + status: RequestCandidateStatus::Cancelled, + status_code: Some(499), + error_type: Some("downstream_disconnect".to_string()), + error_message: Some("client disconnected before stream completion".to_string()), + latency_ms: usage_payload + .telemetry + .as_ref() + .and_then(|value| value.elapsed_ms), + started_at_unix_ms: Some(candidate_started_unix_secs), + finished_at_unix_ms: Some(current_request_candidate_unix_ms()), + }, + ) + .await; + return; + } + + if let Some(failure) = terminal_failure { + record_manual_proxy_stream_error(&state_for_report, &plan_for_report).await; + let terminal_telemetry = Some(build_terminal_stream_telemetry( + stream_started_at_for_report, + telemetry.as_ref(), + usage_stream_telemetry.as_ref(), + provider_stream_bytes, + )); + let report_context_for_payload = report_context_with_stage_trace( + report_context_owned, + stage_trace_for_report, + stream_started_at_for_report, + terminal_telemetry.as_ref(), + ); + let report_context_for_payload = report_context_with_request_diagnostics( + report_context_for_payload, + request_diagnostics_for_report.as_ref(), + ); + submit_midstream_stream_failure( + &state_for_report, + &trace_id_owned, + &plan_for_report, + direct_stream_finalize_kind_owned.as_deref(), + report_context_for_payload, + headers_for_report, + terminal_telemetry, + &provider_buffered_body, + candidate_started_unix_secs, + failure, + ) + .await; + return; + } + + maybe_apply_kiro_prompt_cache_usage_to_stream_summary( + &state_for_report, + &plan_for_report, + report_context_owned.as_ref(), + &mut stream_terminal_summary, + ) + .await; + let requires_observed_terminal_event = stream_requires_observed_terminal_event( + plan_for_report.provider_api_format.as_str(), + stream_usage_report_context.as_ref(), + ); + ensure_stream_terminal_summary_for_missing_observed_finish( + &mut stream_terminal_summary, + requires_observed_terminal_event, + ); + let missing_observed_finish = + stream_terminal_summary_missing_observed_finish_with_requirement( + stream_terminal_summary.as_ref(), + requires_observed_terminal_event, + ); + let stream_failed = stream_terminal_summary_represents_failure_with_requirement( + stream_terminal_summary.as_ref(), + requires_observed_terminal_event, + ); + let stream_terminal_error_message = stream_terminal_summary + .as_ref() + .and_then(|summary| summary.parser_error.clone()) + .or_else(|| { + missing_observed_finish.then(|| { + "execution runtime stream ended before provider terminal event".to_string() + }) + }); + let should_submit_report = report_kind_owned.is_some(); + let terminal_telemetry = Some(build_terminal_stream_telemetry( + stream_started_at_for_report, + telemetry.as_ref(), + usage_stream_telemetry.as_ref(), + provider_stream_bytes, + )); + let report_context_for_payload = report_context_with_stage_trace( + report_context_owned, + stage_trace_for_report, + stream_started_at_for_report, + terminal_telemetry.as_ref(), + ); + let report_context_for_payload = report_context_with_request_diagnostics( + report_context_for_payload, + request_diagnostics_for_report.as_ref(), + ); + let usage_payload = build_stream_usage_payload( + trace_id_owned.clone(), + report_kind_owned.unwrap_or_default(), + report_context_for_payload, + status_code, + headers_for_report, + &provider_buffered_body, + provider_body_truncated, + &buffered_body, + client_body_truncated, + stream_terminal_summary, + terminal_telemetry, + ); + if stream_failed { + warn!( + event_name = "direct_passthrough_stream_failed", + log_type = "ops", + trace_id = %trace_id_owned, + request_id = %request_id_for_report_log, + candidate_id = ?candidate_id_for_report.as_deref(), + status_code, + error_message = stream_terminal_error_message.as_deref().unwrap_or_default(), + "gateway direct passthrough stream ended with a failed terminal state" + ); + } else { + apply_local_execution_effect( + &state_for_report, + LocalExecutionEffectContext { + plan: &plan_for_report, + report_context: usage_payload.report_context.as_ref(), + }, + LocalExecutionEffect::HealthSuccess(LocalHealthSuccessEffect), + ) + .await; + apply_local_execution_effect( + &state_for_report, + LocalExecutionEffectContext { + plan: &plan_for_report, + report_context: usage_payload.report_context.as_ref(), + }, + LocalExecutionEffect::AdaptiveSuccess(LocalAdaptiveSuccessEffect), + ) + .await; + apply_local_execution_effect( + &state_for_report, + LocalExecutionEffectContext { + plan: &plan_for_report, + report_context: usage_payload.report_context.as_ref(), + }, + LocalExecutionEffect::PoolSuccessStream { + payload: &usage_payload, + }, + ) + .await; + } + record_stream_terminal_usage( + &state_for_report, + &plan_for_report, + usage_payload.report_context.as_ref(), + &usage_payload, + false, + ); + record_local_request_candidate_status( + &state_for_report, + &plan_for_report, + usage_payload.report_context.as_ref(), + SchedulerRequestCandidateStatusUpdate { + status: if stream_failed { + RequestCandidateStatus::Failed + } else { + RequestCandidateStatus::Success + }, + status_code: Some(status_code), + error_type: if stream_failed { + if missing_observed_finish { + Some("stream_missing_terminal_event".to_string()) + } else { + Some("stream_terminal_error".to_string()) + } + } else { + None + }, + error_message: stream_failed + .then_some(stream_terminal_error_message) + .flatten(), + latency_ms: usage_payload + .telemetry + .as_ref() + .and_then(|value| value.elapsed_ms), + started_at_unix_ms: Some(candidate_started_unix_secs), + finished_at_unix_ms: Some(current_request_candidate_unix_ms()), + }, + ) + .await; + + if should_submit_report { + if let Err(err) = submit_stream_report(&state_for_report, usage_payload).await { + warn!( + event_name = "execution_report_submit_failed", + log_type = "ops", + trace_id = %trace_id_owned, + request_id = %request_id_for_report_log, + candidate_id = ?candidate_id_for_report.as_deref(), + report_scope = "direct_passthrough_stream", + error = ?err, + "gateway failed to submit direct passthrough stream execution report" + ); + } + } + }); + + let body_stream = build_sse_body_stream(Vec::new(), rx, false, false, SSE_KEEPALIVE_INTERVAL); + Ok(Some(build_client_response_from_parts( + status_code, + &headers, + Body::from_stream(body_stream), + trace_id, + Some(decision), + )?)) +} + #[allow(clippy::too_many_arguments)] // internal function, grouping would add unnecessary indirection pub(crate) async fn execute_execution_runtime_stream( state: &AppState, @@ -840,13 +1786,26 @@ pub(crate) async fn execute_execution_runtime_stream( mut report_context: Option, ) -> Result>, GatewayError> { let stream_started_at = Instant::now(); + let mut stage_trace = RequestStageTrace::from_env(); + let candidate_slot_started_at = Instant::now(); ensure_execution_request_candidate_slot(state, &mut plan, &mut report_context).await; + observe_gateway_stage_trace_ms( + &mut stage_trace, + "stream_candidate_slot", + candidate_slot_started_at.elapsed().as_millis() as u64, + ); let lifecycle_seed = build_lifecycle_usage_seed(&plan, report_context.as_ref()); let request_candidate_status_snapshot = snapshot_local_request_candidate_status(&plan, report_context.as_ref()); + let usage_pending_started_at = Instant::now(); state .usage_runtime .record_pending(state.data.as_ref(), lifecycle_seed.clone()); + observe_gateway_stage_trace_ms( + &mut stage_trace, + "stream_usage_pending", + usage_pending_started_at.elapsed().as_millis() as u64, + ); let candidate_started_unix_secs = current_request_candidate_unix_ms(); if let Some(snapshot) = request_candidate_status_snapshot.clone() { let state_bg = state.clone(); @@ -879,6 +1838,7 @@ pub(crate) async fn execute_execution_runtime_stream( .and_then(|context| context.candidate_index) .map(|value| value.to_string()) .unwrap_or_else(|| "-".to_string()); + let provider_in_flight_started_at = Instant::now(); let mut provider_pool_in_flight_guard = acquire_provider_pool_in_flight_guard( state.runtime_state.clone(), &plan.provider_id, @@ -887,6 +1847,11 @@ pub(crate) async fn execute_execution_runtime_stream( key_id.as_str(), ) .await; + observe_gateway_stage_trace_ms( + &mut stage_trace, + "stream_provider_in_flight", + provider_in_flight_started_at.elapsed().as_millis() as u64, + ); match maybe_execute_grok_stream(&plan, report_context.as_ref()).await { Ok(Some(grok_stream)) => { return execute_stream_from_frame_stream( @@ -899,6 +1864,7 @@ pub(crate) async fn execute_execution_runtime_stream( grok_stream.report_context.or(report_context), candidate_started_unix_secs, stream_started_at, + stage_trace, grok_stream.frame_stream, provider_pool_in_flight_guard.take(), ) @@ -951,6 +1917,7 @@ pub(crate) async fn execute_execution_runtime_stream( windsurf_stream.report_context.or(report_context), candidate_started_unix_secs, stream_started_at, + stage_trace, windsurf_stream.frame_stream, provider_pool_in_flight_guard.take(), ) @@ -1003,6 +1970,7 @@ pub(crate) async fn execute_execution_runtime_stream( kiro_web_search.report_context.or(report_context), candidate_started_unix_secs, stream_started_at, + stage_trace, kiro_web_search.frame_stream, provider_pool_in_flight_guard.take(), ) @@ -1055,6 +2023,7 @@ pub(crate) async fn execute_execution_runtime_stream( chatgpt_web_image.report_context.or(report_context), candidate_started_unix_secs, stream_started_at, + stage_trace, chatgpt_web_image.frame_stream, provider_pool_in_flight_guard.take(), ) @@ -1097,6 +2066,7 @@ pub(crate) async fn execute_execution_runtime_stream( } #[cfg(not(test))] { + let upstream_headers_started_at = Instant::now(); let execution = match execute_in_process_stream_with_oauth_retry( state, &mut plan, @@ -1106,7 +2076,20 @@ pub(crate) async fn execute_execution_runtime_stream( .await { Ok(execution) => execution, - Err(err) => { + Err(InProcessStreamExecutionError::Gateway(err)) => { + if matches!(err, GatewayError::AdmissionTimeout { .. }) { + record_stream_admission_timeout_terminal_state( + state, + &plan, + report_context.as_ref(), + candidate_started_unix_secs, + &err, + ) + .await; + } + return Err(err); + } + Err(InProcessStreamExecutionError::Transport(err)) => { info!( event_name = "stream_execution_runtime_unavailable", log_type = "ops", @@ -1140,6 +2123,29 @@ pub(crate) async fn execute_execution_runtime_stream( return Ok(None); } }; + observe_gateway_stage_trace_ms( + &mut stage_trace, + "stream_upstream_headers", + upstream_headers_started_at.elapsed().as_millis() as u64, + ); + if should_use_direct_sse_passthrough(&plan, plan_kind, report_context.as_ref(), &execution) + { + return execute_stream_from_direct_passthrough( + state, + plan, + trace_id, + decision, + plan_kind, + report_kind, + report_context, + candidate_started_unix_secs, + stream_started_at, + stage_trace, + execution, + provider_pool_in_flight_guard.take(), + ) + .await; + } let frame_stream = build_direct_execution_frame_stream(execution).boxed(); return execute_stream_from_frame_stream( state, @@ -1151,6 +2157,7 @@ pub(crate) async fn execute_execution_runtime_stream( report_context, candidate_started_unix_secs, stream_started_at, + stage_trace, frame_stream, provider_pool_in_flight_guard.take(), ) @@ -1162,6 +2169,7 @@ pub(crate) async fn execute_execution_runtime_stream( .execution_runtime_override_base_url() .unwrap_or_default(); if remote_execution_runtime_base_url.trim().is_empty() { + let upstream_headers_started_at = Instant::now(); let execution = match execute_in_process_stream_with_oauth_retry( state, &mut plan, @@ -1171,7 +2179,20 @@ pub(crate) async fn execute_execution_runtime_stream( .await { Ok(execution) => execution, - Err(err) => { + Err(InProcessStreamExecutionError::Gateway(err)) => { + if matches!(err, GatewayError::AdmissionTimeout { .. }) { + record_stream_admission_timeout_terminal_state( + state, + &plan, + report_context.as_ref(), + candidate_started_unix_secs, + &err, + ) + .await; + } + return Err(err); + } + Err(InProcessStreamExecutionError::Transport(err)) => { info!( event_name = "stream_execution_runtime_unavailable", log_type = "ops", @@ -1205,6 +2226,33 @@ pub(crate) async fn execute_execution_runtime_stream( return Ok(None); } }; + observe_gateway_stage_trace_ms( + &mut stage_trace, + "stream_upstream_headers", + upstream_headers_started_at.elapsed().as_millis() as u64, + ); + if should_use_direct_sse_passthrough( + &plan, + plan_kind, + report_context.as_ref(), + &execution, + ) { + return execute_stream_from_direct_passthrough( + state, + plan, + trace_id, + decision, + plan_kind, + report_kind, + report_context, + candidate_started_unix_secs, + stream_started_at, + stage_trace, + execution, + provider_pool_in_flight_guard.take(), + ) + .await; + } let frame_stream = build_direct_execution_frame_stream(execution).boxed(); return execute_stream_from_frame_stream( state, @@ -1216,6 +2264,7 @@ pub(crate) async fn execute_execution_runtime_stream( report_context, candidate_started_unix_secs, stream_started_at, + stage_trace, frame_stream, provider_pool_in_flight_guard.take(), ) @@ -1302,6 +2351,7 @@ pub(crate) async fn execute_execution_runtime_stream( report_context, candidate_started_unix_secs, stream_started_at, + stage_trace, frame_stream, provider_pool_in_flight_guard.take(), ) @@ -1958,6 +3008,7 @@ async fn execute_stream_from_frame_stream( report_context: Option, candidate_started_unix_secs: u64, stream_started_at: Instant, + mut stage_trace: RequestStageTrace, frame_stream: BoxStream<'static, Result>, in_flight_guard: Option, ) -> Result>, GatewayError> { @@ -1976,9 +3027,15 @@ async fn execute_stream_from_frame_stream( let reader = StreamReader::new(frame_stream); let mut lines = FramedRead::new(reader, LinesCodec::new()); + let first_frame_started_at = Instant::now(); let first_frame = read_next_frame(&mut lines).await?.ok_or_else(|| { GatewayError::Internal("execution runtime stream ended before headers frame".to_string()) })?; + observe_gateway_stage_trace_ms( + &mut stage_trace, + "stream_first_frame", + first_frame_started_at.elapsed().as_millis() as u64, + ); let StreamFramePayload::Headers { status_code, mut headers, @@ -2480,12 +3537,18 @@ async fn execute_stream_from_frame_stream( let frame_observed_at = observed_frame.observed_at; match observed_frame.frame.payload { StreamFramePayload::Data { chunk_b64, text } => { - maybe_capture_first_stream_event_telemetry( + if maybe_capture_first_stream_event_telemetry( stream_started_at, frame_observed_at, prefetched_telemetry.as_ref(), &mut prefetched_usage_telemetry, - ); + ) { + observe_gateway_stage_trace_ms( + &mut stage_trace, + "stream_first_data", + stream_elapsed_ms_at(stream_started_at, frame_observed_at), + ); + } let chunk = match decode_stream_data_chunk(chunk_b64.as_deref(), text.as_deref()) { Ok(chunk) => chunk, @@ -2836,41 +3899,22 @@ async fn execute_stream_from_frame_stream( let emit_passthrough_sse_terminal_error = skip_direct_finalize_prefetch && response_headers_indicate_sse(&upstream_headers) && !is_openai_image_stream_for_report; - let body_capture_policy = match state - .usage_runtime - .body_capture_policy_for(state.data.as_ref()) - .await - { - Ok(policy) => policy, - Err(err) => { - warn!( - event_name = "stream_execution_body_capture_policy_read_failed", - log_type = "ops", - trace_id = %trace_id, - request_id = %request_id_for_report_log, - candidate_id = ?candidate_id_for_report.as_deref(), - error = %err, - fallback_request_body_bytes = DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES, - "gateway failed to read body capture policy; falling back to default stream capture limits" - ); - UsageBodyCapturePolicy::default() - } - }; - let max_stream_body_buffer_bytes = if matches!( - body_capture_policy.record_level, - UsageRequestRecordLevel::Basic - ) { - 0 - } else { - body_capture_policy - .max_response_body_bytes - .unwrap_or(DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES) - }; let plan_kind_for_report = plan_kind.to_string(); let stream_started_at_for_report = stream_started_at; + observe_gateway_stage_trace_ms( + &mut stage_trace, + "stream_response_ready", + stream_elapsed_ms_since(stream_started_at), + ); + let stage_trace_for_report = stage_trace; + let request_diagnostics_for_report = current_request_diagnostics(); let provider_pool_in_flight_guard_for_report = in_flight_guard; tokio::spawn(async move { + let mut stage_trace_for_report = stage_trace_for_report; + let _stream_total_guard = + StageElapsedGuard::from_started_at("stream_total", stream_started_at_for_report); let _provider_pool_in_flight_guard = provider_pool_in_flight_guard_for_report; + let max_stream_body_buffer_bytes = DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES; let mut provider_buffered_body = Vec::new(); let mut buffered_body = Vec::new(); let mut provider_body_truncated = false; @@ -3135,6 +4179,10 @@ async fn execute_stream_from_frame_stream( last_upstream_frame_elapsed_ms.store(frame_elapsed_ms, Ordering::Relaxed); match observed_frame.frame.payload { StreamFramePayload::Data { chunk_b64, text } => { + let first_data_before = usage_stream_telemetry + .as_ref() + .and_then(|telemetry| telemetry.ttfb_ms) + .is_some(); maybe_record_first_stream_event_started( &state_for_report, &lifecycle_seed_for_report, @@ -3144,6 +4192,20 @@ async fn execute_stream_from_frame_stream( telemetry.as_ref(), &mut usage_stream_telemetry, ); + let first_data_after = usage_stream_telemetry + .as_ref() + .and_then(|telemetry| telemetry.ttfb_ms) + .is_some(); + if !first_data_before && first_data_after { + observe_gateway_stage_trace_ms( + &mut stage_trace_for_report, + "stream_first_data", + stream_elapsed_ms_at( + stream_started_at_for_report, + frame_observed_at, + ), + ); + } if sync_json_stream_bridge_active_for_report { continue; } @@ -3282,7 +4344,7 @@ async fn execute_stream_from_frame_stream( } let rewritten_chunk = Bytes::from(rewritten_chunk); if tx.send(Ok(rewritten_chunk.clone())).await.is_err() { - warn!( + debug!( event_name = "stream_execution_downstream_disconnected", log_type = "ops", trace_id = %trace_id_owned, @@ -3652,10 +4714,20 @@ async fn execute_stream_from_frame_stream( usage_stream_telemetry.as_ref(), provider_stream_bytes.load(Ordering::Relaxed), )); + let report_context_for_payload = report_context_with_stage_trace( + report_context_owned, + stage_trace_for_report, + stream_started_at_for_report, + terminal_telemetry.as_ref(), + ); + let report_context_for_payload = report_context_with_request_diagnostics( + report_context_for_payload, + request_diagnostics_for_report.as_ref(), + ); let usage_payload = build_stream_usage_payload( trace_id_owned, report_kind_owned.unwrap_or_default(), - report_context_owned, + report_context_for_payload, 499, headers_for_report, &provider_buffered_body, @@ -3701,12 +4773,22 @@ async fn execute_stream_from_frame_stream( usage_stream_telemetry.as_ref(), provider_stream_bytes.load(Ordering::Relaxed), )); + let report_context_for_payload = report_context_with_stage_trace( + report_context_owned, + stage_trace_for_report, + stream_started_at_for_report, + terminal_telemetry.as_ref(), + ); + let report_context_for_payload = report_context_with_request_diagnostics( + report_context_for_payload, + request_diagnostics_for_report.as_ref(), + ); submit_midstream_stream_failure( &state_for_report, &trace_id_owned, &plan_for_report, direct_stream_finalize_kind_owned.as_deref(), - report_context_owned, + report_context_for_payload, headers_for_report, terminal_telemetry, &provider_buffered_body, @@ -3757,10 +4839,20 @@ async fn execute_stream_from_frame_stream( "execution runtime stream ended before provider terminal event".to_string() }) }); + let report_context_for_payload = report_context_with_stage_trace( + report_context_owned, + stage_trace_for_report, + stream_started_at_for_report, + terminal_telemetry.as_ref(), + ); + let report_context_for_payload = report_context_with_request_diagnostics( + report_context_for_payload, + request_diagnostics_for_report.as_ref(), + ); let usage_payload = build_stream_usage_payload( trace_id_owned.clone(), report_kind_owned.unwrap_or_default(), - report_context_owned, + report_context_for_payload, status_code, headers_for_report, &provider_buffered_body, @@ -3957,6 +5049,7 @@ mod tests { ClientVisibleStreamCompletionTracker, }; use crate::control::GatewayControlDecision; + use crate::stage_metrics::RequestStageTrace; use crate::tunnel::{tunnel_protocol, TunnelProxyConn}; use crate::AppState; @@ -5367,6 +6460,7 @@ mod tests { })), crate::clock::current_unix_ms(), Instant::now(), + RequestStageTrace::from_env(), frame_stream, None, ) @@ -5476,6 +6570,7 @@ mod tests { })), crate::clock::current_unix_ms(), Instant::now(), + RequestStageTrace::from_env(), frame_stream, None, ) @@ -5582,6 +6677,7 @@ mod tests { })), crate::clock::current_unix_ms(), Instant::now(), + RequestStageTrace::from_env(), frame_stream, None, ) @@ -5718,6 +6814,7 @@ mod tests { })), crate::clock::current_unix_ms(), Instant::now(), + RequestStageTrace::from_env(), frame_stream, None, ) @@ -5843,6 +6940,7 @@ mod tests { })), crate::clock::current_unix_ms(), Instant::now(), + RequestStageTrace::from_env(), frame_stream, None, ) @@ -6016,6 +7114,7 @@ mod tests { })), crate::clock::current_unix_ms(), Instant::now(), + RequestStageTrace::from_env(), frame_stream, None, ) @@ -6155,6 +7254,7 @@ mod tests { })), crate::clock::current_unix_ms(), Instant::now(), + RequestStageTrace::from_env(), frame_stream, None, ) diff --git a/apps/aether-gateway/src/execution_runtime/stream/execution_failures.rs b/apps/aether-gateway/src/execution_runtime/stream/execution_failures.rs index b513d78f2..0e9702420 100644 --- a/apps/aether-gateway/src/execution_runtime/stream/execution_failures.rs +++ b/apps/aether-gateway/src/execution_runtime/stream/execution_failures.rs @@ -26,6 +26,7 @@ use crate::orchestration::{ LocalPoolErrorEffect, }; use crate::request_candidate_runtime::record_report_request_candidate_status; +use crate::request_diagnostics::attach_current_request_diagnostics_to_report_context; use crate::usage::submit_sync_report; use crate::{usage::GatewaySyncReportRequest, AppState, GatewayError}; @@ -319,7 +320,12 @@ async fn record_stream_sync_failure( }), ) .await; - let context_seed = build_terminal_usage_context_seed(plan, report_context); + let report_context_with_diagnostics = + attach_current_request_diagnostics_to_report_context(report_context); + let context_seed = build_terminal_usage_context_seed( + plan, + report_context_with_diagnostics.as_ref().or(report_context), + ); let payload_seed = build_sync_terminal_usage_payload_seed(payload); state .usage_runtime diff --git a/apps/aether-gateway/src/execution_runtime/stream_pump.rs b/apps/aether-gateway/src/execution_runtime/stream_pump.rs index 6c9bcee9c..1f065096a 100644 --- a/apps/aether-gateway/src/execution_runtime/stream_pump.rs +++ b/apps/aether-gateway/src/execution_runtime/stream_pump.rs @@ -39,7 +39,9 @@ pub(crate) fn build_direct_execution_frame_stream( response, started_at, stream_first_byte_timeout, + upstream_target_permit, } = execution; + let _upstream_target_permit = upstream_target_permit; let mut observer_context = stream_summary_report_context; if observer_context diff --git a/apps/aether-gateway/src/execution_runtime/sync/execution.rs b/apps/aether-gateway/src/execution_runtime/sync/execution.rs index 8666e9aea..3b1245753 100644 --- a/apps/aether-gateway/src/execution_runtime/sync/execution.rs +++ b/apps/aether-gateway/src/execution_runtime/sync/execution.rs @@ -79,6 +79,7 @@ use crate::request_candidate_runtime::{ ensure_execution_request_candidate_slot, record_local_request_candidate_extra_data, record_local_request_candidate_status, }; +use crate::request_diagnostics::attach_current_request_diagnostics_to_report_context; use crate::usage::{spawn_sync_report, submit_sync_report}; use crate::video_tasks::VideoTaskSyncReportMode; use crate::{usage::GatewaySyncReportRequest, AppState, GatewayError}; @@ -296,7 +297,12 @@ fn record_sync_terminal_usage( report_context: Option<&serde_json::Value>, payload: &GatewaySyncReportRequest, ) { - let context_seed = build_terminal_usage_context_seed(plan, report_context); + let report_context_with_diagnostics = + attach_current_request_diagnostics_to_report_context(report_context); + let context_seed = build_terminal_usage_context_seed( + plan, + report_context_with_diagnostics.as_ref().or(report_context), + ); let payload_seed = build_sync_terminal_usage_payload_seed(payload); state .usage_runtime diff --git a/apps/aether-gateway/src/execution_runtime/transport.rs b/apps/aether-gateway/src/execution_runtime/transport.rs index 893805800..09ddd1c4a 100644 --- a/apps/aether-gateway/src/execution_runtime/transport.rs +++ b/apps/aether-gateway/src/execution_runtime/transport.rs @@ -3,6 +3,7 @@ use std::error::Error as _; use std::future::Future; use std::io::Read; use std::io::Write; +use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{LazyLock, Mutex as StdMutex}; use std::time::{Duration, Instant}; @@ -11,10 +12,11 @@ use aether_contracts::{ ResponseBody, EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER, EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER, EXECUTION_REQUEST_HTTP1_ONLY_HEADER, TRANSPORT_BACKEND_BROWSER_WREQ, TRANSPORT_BACKEND_REQWEST_RUSTLS, - TRANSPORT_HTTP_MODE_HTTP1_ONLY, + TRANSPORT_HTTP_MODE_H2C_PRIOR_KNOWLEDGE, TRANSPORT_HTTP_MODE_HTTP1_ONLY, }; use aether_data::repository::proxy_nodes::ProxyNodeTrafficMutation; use aether_http::{apply_http_client_config, HttpClientConfig}; +use aether_runtime::{MetricKind, MetricSample}; use axum::body::Bytes; use base64::Engine as _; use flate2::read::{DeflateDecoder, GzDecoder}; @@ -35,7 +37,9 @@ use crate::execution_runtime::windsurf::maybe_execute_windsurf_sync; use crate::frontdoor_loop_guard::{ configured_gateway_frontdoor_base_url, gateway_frontdoor_self_loop_guard_error, }; +use crate::stage_metrics::observe_gateway_stage_ms; use crate::tunnel::{self, tunnel_protocol}; +use crate::upstream_admission::UpstreamTargetAdmissionPermit; use crate::{AppState, GatewayError}; const HUB_RELAY_CONTENT_TYPE: &str = "application/vnd.aether.tunnel-envelope"; @@ -46,18 +50,88 @@ const DEFAULT_STREAM_FIRST_BYTE_TIMEOUT_MS: u64 = 30_000; const DEFAULT_NON_STREAM_TOTAL_TIMEOUT_MS: u64 = 300_000; const MIN_TUNNEL_TIMEOUT_SECS: u64 = 1; const MAX_TUNNEL_TIMEOUT_SECS: u64 = 300; +const DIRECT_REQWEST_H2_CLIENT_SHARDS_ENV: &str = "AETHER_GATEWAY_DIRECT_REQWEST_H2_CLIENT_SHARDS"; +const DIRECT_REQWEST_H2_TARGET_STREAMS_PER_CLIENT_ENV: &str = + "AETHER_GATEWAY_DIRECT_REQWEST_H2_TARGET_STREAMS_PER_CLIENT"; +const DIRECT_REQWEST_SYNC_WARM_CLIENTS_ENV: &str = + "AETHER_GATEWAY_DIRECT_REQWEST_SYNC_WARM_CLIENTS"; +const UPSTREAM_TARGET_GATE_LIMIT_ENV: &str = "AETHER_GATEWAY_UPSTREAM_TARGET_GATE_LIMIT"; +const DEFAULT_UPSTREAM_TARGET_GATE_LIMIT: usize = 2_000; +const DEFAULT_H2_TARGET_STREAMS_PER_CLIENT: usize = 200; +const DEFAULT_DIRECT_REQWEST_SYNC_WARM_CLIENTS: usize = 4; +const MAX_DIRECT_REQWEST_H2_CLIENT_SHARDS: usize = 128; -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +#[derive(Debug, Clone, PartialEq, Eq, Hash)] struct DirectReqwestClientCacheKey { connect_timeout_ms: Option, + proxy_url: Option, follow_redirects: bool, http1_only: bool, accept_invalid_certs: bool, + transport_profile: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct DirectReqwestTransportProfileCacheKey { + profile_id: String, + backend: String, + http_mode: String, + pool_scope: String, + header_fingerprint: Option, + extra: Option, +} + +struct DirectReqwestClientCacheEntry { + clients: Vec, + next: AtomicU64, + target_len: usize, + warming: bool, +} + +impl DirectReqwestClientCacheEntry { + fn new(clients: Vec, target_len: usize, warming: bool) -> Self { + Self { + clients, + next: AtomicU64::new(0), + target_len: target_len.max(1), + warming, + } + } + + fn select(&self) -> reqwest::Client { + if self.clients.len() <= 1 { + return self + .clients + .first() + .expect("direct reqwest client cache entry should contain a client") + .clone(); + } + let index = self.next.fetch_add(1, Ordering::Relaxed) as usize % self.clients.len(); + self.clients[index].clone() + } + + fn len(&self) -> usize { + self.clients.len() + } + + fn should_warm(&self) -> bool { + self.clients.len() < self.target_len && !self.warming + } } static DIRECT_REQWEST_CLIENT_CACHE: LazyLock< - StdMutex>, + StdMutex>, > = LazyLock::new(|| StdMutex::new(HashMap::new())); + +#[derive(Debug, Default)] +struct DirectReqwestClientCacheMetrics { + hits: AtomicU64, + misses: AtomicU64, + builds: AtomicU64, +} + +static DIRECT_REQWEST_CLIENT_CACHE_METRICS: LazyLock = + LazyLock::new(DirectReqwestClientCacheMetrics::default); pub(crate) fn format_upstream_request_error(err: &reqwest::Error) -> String { let mut kinds = Vec::new(); if err.is_connect() { @@ -244,6 +318,7 @@ pub(crate) struct DirectUpstreamStreamExecution { pub(crate) response: DirectUpstreamResponse, pub(crate) started_at: Instant, pub(crate) stream_first_byte_timeout: Option, + pub(crate) upstream_target_permit: Option, } impl DirectSyncExecutionRuntime { @@ -302,10 +377,19 @@ impl DirectSyncExecutionRuntime { return Err(ExecutionRuntimeTransportError::StreamUnsupported); } + let build_body_started_at = Instant::now(); let body_bytes = build_request_body(plan)?; + observe_gateway_stage_ms( + "direct_build_body", + build_body_started_at.elapsed().as_millis() as u64, + ); let started_at = Instant::now(); let response = send_request(plan, body_bytes).await?; + observe_gateway_stage_ms( + "direct_send_headers", + started_at.elapsed().as_millis() as u64, + ); let status_code = response.status_code(); let headers = response.headers(); @@ -321,6 +405,7 @@ impl DirectSyncExecutionRuntime { response: response.into_direct_upstream_response(), started_at, stream_first_byte_timeout: resolve_stream_first_byte_timeout(plan), + upstream_target_permit: None, }) } } @@ -433,6 +518,7 @@ pub(crate) async fn execute_stream_plan_via_local_tunnel( response: DirectUpstreamResponse::LocalTunnel(response), started_at, stream_first_byte_timeout: resolve_stream_first_byte_timeout(plan), + upstream_target_permit: None, })) } @@ -1366,64 +1452,367 @@ fn build_client( ) -> Result { validate_reqwest_transport_profile(transport_profile)?; let resolved_proxy_url = resolve_proxy_url(proxy)?; - if resolved_proxy_url.is_none() && transport_profile.is_none() { - let cache_key = DirectReqwestClientCacheKey { - connect_timeout_ms: timeouts.and_then(|timeouts| timeouts.connect_ms), - follow_redirects: transport_controls.follow_redirects == Some(true), - http1_only: transport_controls.http1_only, - accept_invalid_certs: transport_controls.accept_invalid_certs, - }; - if let Ok(cache) = DIRECT_REQWEST_CLIENT_CACHE.lock() { - if let Some(client) = cache.get(&cache_key) { - return Ok(client.clone()); + let cache_key = direct_reqwest_client_cache_key( + timeouts, + resolved_proxy_url, + transport_profile, + transport_controls, + ); + cached_direct_reqwest_client(cache_key) +} + +pub(crate) fn prewarm_direct_reqwest_client_cache_for_plan(plan: &ExecutionPlan) { + match try_prewarm_direct_reqwest_client_cache_for_plan(plan) { + Ok(true) => {} + Ok(false) => {} + Err(err) => { + tracing::debug!( + error = ?err, + request_id = %plan.request_id, + candidate_id = ?plan.candidate_id, + provider_id = %plan.provider_id, + endpoint_id = %plan.endpoint_id, + key_id = %plan.key_id, + "gateway direct reqwest client prewarm skipped" + ); + } + } +} + +fn try_prewarm_direct_reqwest_client_cache_for_plan( + plan: &ExecutionPlan, +) -> Result { + if transport_profile_uses_browser_wreq(plan.transport_profile.as_ref()) { + return Ok(false); + } + if resolve_tunnel_node_id(plan.proxy.as_ref()).is_some() { + return Ok(false); + } + + let transport_controls = resolve_execution_transport_controls(&plan.headers); + validate_reqwest_transport_profile(plan.transport_profile.as_ref())?; + let resolved_proxy_url = resolve_proxy_url(plan.proxy.as_ref())?; + let cache_key = direct_reqwest_client_cache_key( + plan.timeouts.as_ref(), + resolved_proxy_url, + plan.transport_profile.as_ref(), + transport_controls, + ); + prewarm_direct_reqwest_client_cache(cache_key)?; + Ok(true) +} + +fn prewarm_direct_reqwest_client_cache( + cache_key: DirectReqwestClientCacheKey, +) -> Result<(), ExecutionRuntimeTransportError> { + let mut warm_after_unlock = None; + if let Ok(mut cache) = DIRECT_REQWEST_CLIENT_CACHE.lock() { + if let Some(entry) = cache.get_mut(&cache_key) { + if entry.should_warm() { + entry.warming = true; + warm_after_unlock = Some((cache_key.clone(), entry.len(), entry.target_len)); } + drop(cache); + if let Some((cache_key, existing_len, target_len)) = warm_after_unlock { + let spawned = spawn_direct_reqwest_client_cache_warm( + cache_key.clone(), + existing_len, + target_len, + ); + if !spawned { + mark_direct_reqwest_client_cache_not_warming(&cache_key); + } + } + return Ok(()); } - let client = build_plain_direct_reqwest_client(cache_key)?; - if let Ok(mut cache) = DIRECT_REQWEST_CLIENT_CACHE.lock() { - let client = cache.entry(cache_key).or_insert_with(|| client.clone()); - return Ok(client.clone()); + let target_len = direct_reqwest_client_shard_count(&cache_key); + let initial_len = direct_reqwest_initial_client_shard_count(target_len); + let mut clients = Vec::with_capacity(initial_len); + for _ in 0..initial_len { + clients.push(build_direct_reqwest_client_from_cache_key(&cache_key)?); + DIRECT_REQWEST_CLIENT_CACHE_METRICS + .builds + .fetch_add(1, Ordering::Relaxed); + } + let entry = + DirectReqwestClientCacheEntry::new(clients, target_len, target_len > initial_len); + let warm_key = (target_len > initial_len).then(|| cache_key.clone()); + cache.insert(cache_key, entry); + if let Some(warm_key) = warm_key { + warm_after_unlock = Some((warm_key, initial_len, target_len)); + } + drop(cache); + if let Some((cache_key, existing_len, target_len)) = warm_after_unlock { + let spawned = + spawn_direct_reqwest_client_cache_warm(cache_key.clone(), existing_len, target_len); + if !spawned { + mark_direct_reqwest_client_cache_not_warming(&cache_key); + } + } + } + Ok(()) +} + +fn cached_direct_reqwest_client( + cache_key: DirectReqwestClientCacheKey, +) -> Result { + let mut warm_after_unlock = None; + if let Ok(mut cache) = DIRECT_REQWEST_CLIENT_CACHE.lock() { + if let Some(entry) = cache.get_mut(&cache_key) { + DIRECT_REQWEST_CLIENT_CACHE_METRICS + .hits + .fetch_add(1, Ordering::Relaxed); + let client = entry.select(); + if entry.should_warm() { + entry.warming = true; + warm_after_unlock = Some((cache_key.clone(), entry.len(), entry.target_len)); + } + drop(cache); + if let Some((cache_key, existing_len, target_len)) = warm_after_unlock { + let spawned = spawn_direct_reqwest_client_cache_warm( + cache_key.clone(), + existing_len, + target_len, + ); + if !spawned { + mark_direct_reqwest_client_cache_not_warming(&cache_key); + } + } + return Ok(client); + } + DIRECT_REQWEST_CLIENT_CACHE_METRICS + .misses + .fetch_add(1, Ordering::Relaxed); + let target_len = direct_reqwest_client_shard_count(&cache_key); + let initial_len = direct_reqwest_initial_client_shard_count(target_len); + let mut clients = Vec::with_capacity(initial_len); + for _ in 0..initial_len { + clients.push(build_direct_reqwest_client_from_cache_key(&cache_key)?); + DIRECT_REQWEST_CLIENT_CACHE_METRICS + .builds + .fetch_add(1, Ordering::Relaxed); + } + let entry = + DirectReqwestClientCacheEntry::new(clients, target_len, target_len > initial_len); + let client = entry.select(); + let warm_key = (target_len > initial_len).then(|| cache_key.clone()); + cache.insert(cache_key, entry); + if let Some(warm_key) = warm_key { + warm_after_unlock = Some((warm_key, initial_len, target_len)); + } + drop(cache); + if let Some((cache_key, existing_len, target_len)) = warm_after_unlock { + let spawned = + spawn_direct_reqwest_client_cache_warm(cache_key.clone(), existing_len, target_len); + if !spawned { + mark_direct_reqwest_client_cache_not_warming(&cache_key); + } } return Ok(client); } - let mut builder = reqwest::Client::builder(); - if transport_controls.follow_redirects != Some(true) { - builder = builder.redirect(Policy::none()); - } - if transport_controls.http1_only || transport_profile_http1_only(transport_profile) { - builder = builder.http1_only(); - } - let mut builder = apply_http_client_config( - builder, - &HttpClientConfig { - connect_timeout_ms: timeouts.and_then(|timeouts| timeouts.connect_ms), - ..HttpClientConfig::default() - }, - ); - builder = apply_transport_profile(builder, transport_profile); - if transport_controls.accept_invalid_certs { - builder = builder.danger_accept_invalid_certs(true); - } - if let Some(proxy_url) = resolved_proxy_url { - let proxy = reqwest::Proxy::all(&proxy_url) - .map_err(ExecutionRuntimeTransportError::InvalidProxy)?; - builder = builder.proxy(proxy); - } - builder - .build() - .map_err(ExecutionRuntimeTransportError::ClientBuild) + DIRECT_REQWEST_CLIENT_CACHE_METRICS + .misses + .fetch_add(1, Ordering::Relaxed); + let client = build_direct_reqwest_client_from_cache_key(&cache_key)?; + DIRECT_REQWEST_CLIENT_CACHE_METRICS + .builds + .fetch_add(1, Ordering::Relaxed); + Ok(client) } -fn build_plain_direct_reqwest_client( +fn spawn_direct_reqwest_client_cache_warm( cache_key: DirectReqwestClientCacheKey, + existing_len: usize, + target_len: usize, +) -> bool { + if target_len <= existing_len { + return false; + } + let Ok(handle) = tokio::runtime::Handle::try_current() else { + return false; + }; + handle.spawn_blocking(move || { + for _ in existing_len..target_len { + match build_direct_reqwest_client_from_cache_key(&cache_key) { + Ok(client) => { + DIRECT_REQWEST_CLIENT_CACHE_METRICS + .builds + .fetch_add(1, Ordering::Relaxed); + let Ok(mut cache) = DIRECT_REQWEST_CLIENT_CACHE.lock() else { + return; + }; + let Some(entry) = cache.get_mut(&cache_key) else { + return; + }; + if entry.clients.len() >= entry.target_len { + entry.warming = false; + return; + } + entry.clients.push(client); + if entry.clients.len() >= entry.target_len { + entry.warming = false; + return; + } + } + Err(err) => { + tracing::debug!( + error = ?err, + "gateway direct reqwest client cache warm failed" + ); + mark_direct_reqwest_client_cache_not_warming(&cache_key); + break; + } + } + } + + let Ok(mut cache) = DIRECT_REQWEST_CLIENT_CACHE.lock() else { + return; + }; + let Some(entry) = cache.get_mut(&cache_key) else { + return; + }; + entry.warming = false; + }); + true +} + +fn mark_direct_reqwest_client_cache_warming(cache_key: &DirectReqwestClientCacheKey) { + if let Ok(mut cache) = DIRECT_REQWEST_CLIENT_CACHE.lock() { + if let Some(entry) = cache.get_mut(cache_key) { + entry.warming = true; + } + } +} + +fn mark_direct_reqwest_client_cache_not_warming(cache_key: &DirectReqwestClientCacheKey) { + if let Ok(mut cache) = DIRECT_REQWEST_CLIENT_CACHE.lock() { + if let Some(entry) = cache.get_mut(cache_key) { + entry.warming = false; + } + } +} + +fn direct_reqwest_client_cache_key( + timeouts: Option<&aether_contracts::ExecutionTimeouts>, + proxy_url: Option, + transport_profile: Option<&ResolvedTransportProfile>, + transport_controls: ExecutionTransportControls, +) -> DirectReqwestClientCacheKey { + DirectReqwestClientCacheKey { + connect_timeout_ms: timeouts.and_then(|timeouts| timeouts.connect_ms), + proxy_url, + follow_redirects: transport_controls.follow_redirects == Some(true), + http1_only: transport_controls.http1_only, + accept_invalid_certs: transport_controls.accept_invalid_certs, + transport_profile: transport_profile.map(direct_reqwest_transport_profile_cache_key), + } +} + +fn direct_reqwest_transport_profile_cache_key( + profile: &ResolvedTransportProfile, +) -> DirectReqwestTransportProfileCacheKey { + DirectReqwestTransportProfileCacheKey { + profile_id: profile.profile_id.trim().to_string(), + backend: profile.backend.trim().to_ascii_lowercase(), + http_mode: profile.http_mode.trim().to_ascii_lowercase(), + pool_scope: profile.pool_scope.trim().to_ascii_lowercase(), + header_fingerprint: stable_json_cache_key(profile.header_fingerprint.as_ref()), + extra: stable_json_cache_key(profile.extra.as_ref()), + } +} + +fn stable_json_cache_key(value: Option<&Value>) -> Option { + value.and_then(|value| serde_json::to_string(value).ok()) +} + +fn build_direct_reqwest_client_cache_entry_from_cache_key( + cache_key: &DirectReqwestClientCacheKey, +) -> Result { + let shard_count = direct_reqwest_client_shard_count(cache_key); + let mut clients = Vec::with_capacity(shard_count); + for _ in 0..shard_count { + clients.push(build_direct_reqwest_client_from_cache_key(cache_key)?); + } + Ok(DirectReqwestClientCacheEntry::new( + clients, + shard_count, + false, + )) +} + +fn direct_reqwest_client_shard_count(cache_key: &DirectReqwestClientCacheKey) -> usize { + if !direct_reqwest_client_cache_key_uses_http2(cache_key) { + return 1; + } + direct_reqwest_h2_client_shards_from_config( + env_positive_usize(DIRECT_REQWEST_H2_CLIENT_SHARDS_ENV), + env_positive_usize(UPSTREAM_TARGET_GATE_LIMIT_ENV) + .unwrap_or(DEFAULT_UPSTREAM_TARGET_GATE_LIMIT), + env_positive_usize(DIRECT_REQWEST_H2_TARGET_STREAMS_PER_CLIENT_ENV) + .unwrap_or(DEFAULT_H2_TARGET_STREAMS_PER_CLIENT), + ) +} + +fn direct_reqwest_client_cache_key_uses_http2(cache_key: &DirectReqwestClientCacheKey) -> bool { + if cache_key.http1_only { + return false; + } + let Some(profile) = cache_key.transport_profile.as_ref() else { + return false; + }; + profile.http_mode != TRANSPORT_HTTP_MODE_HTTP1_ONLY +} + +fn direct_reqwest_h2_client_shards_from_config( + explicit_shards: Option, + target_gate_limit: usize, + target_streams_per_client: usize, +) -> usize { + if let Some(shards) = explicit_shards { + return shards.clamp(1, MAX_DIRECT_REQWEST_H2_CLIENT_SHARDS); + } + let streams_per_client = target_streams_per_client.max(1); + target_gate_limit + .max(1) + .div_ceil(streams_per_client) + .clamp(1, MAX_DIRECT_REQWEST_H2_CLIENT_SHARDS) +} + +fn direct_reqwest_initial_client_shard_count(target_len: usize) -> usize { + env_positive_usize(DIRECT_REQWEST_SYNC_WARM_CLIENTS_ENV) + .unwrap_or(DEFAULT_DIRECT_REQWEST_SYNC_WARM_CLIENTS) + .clamp(1, target_len.max(1)) +} + +fn env_positive_usize(name: &str) -> Option { + std::env::var(name) + .ok() + .and_then(|value| value.trim().parse::().ok()) + .filter(|value| *value > 0) +} + +fn build_direct_reqwest_client_from_cache_key( + cache_key: &DirectReqwestClientCacheKey, ) -> Result { let mut builder = reqwest::Client::builder(); if !cache_key.follow_redirects { builder = builder.redirect(Policy::none()); } - if cache_key.http1_only { + if cache_key.http1_only + || cache_key + .transport_profile + .as_ref() + .is_some_and(|profile| profile.http_mode == TRANSPORT_HTTP_MODE_HTTP1_ONLY) + { builder = builder.http1_only(); + } else if cache_key + .transport_profile + .as_ref() + .is_some_and(|profile| profile.http_mode == TRANSPORT_HTTP_MODE_H2C_PRIOR_KNOWLEDGE) + { + builder = builder.http2_prior_knowledge(); } let mut builder = apply_http_client_config( builder, @@ -1433,9 +1822,15 @@ fn build_plain_direct_reqwest_client( ..HttpClientConfig::default() }, ); + builder = apply_transport_profile_cache_key(builder, cache_key.transport_profile.as_ref()); if cache_key.accept_invalid_certs { builder = builder.danger_accept_invalid_certs(true); } + if let Some(proxy_url) = cache_key.proxy_url.as_deref() { + let proxy = + reqwest::Proxy::all(proxy_url).map_err(ExecutionRuntimeTransportError::InvalidProxy)?; + builder = builder.proxy(proxy); + } builder .build() .map_err(ExecutionRuntimeTransportError::ClientBuild) @@ -1450,6 +1845,55 @@ fn direct_reqwest_pool_max_idle_per_host() -> usize { .unwrap_or(DEFAULT_MAX_IDLE_PER_HOST) } +pub(crate) fn direct_reqwest_client_cache_metric_samples() -> Vec { + let (entries, clients) = DIRECT_REQWEST_CLIENT_CACHE + .lock() + .map(|cache| { + let entries = cache.len() as u64; + let clients = cache.values().map(|entry| entry.len() as u64).sum(); + (entries, clients) + }) + .unwrap_or((0, 0)); + vec![ + MetricSample::new( + "direct_reqwest_client_cache_entries", + "Number of cached direct reqwest clients.", + MetricKind::Gauge, + entries, + ), + MetricSample::new( + "direct_reqwest_client_cache_clients", + "Number of direct reqwest clients across all cache entries.", + MetricKind::Gauge, + clients, + ), + MetricSample::new( + "direct_reqwest_client_cache_hits_total", + "Number of direct reqwest client cache hits.", + MetricKind::Counter, + DIRECT_REQWEST_CLIENT_CACHE_METRICS + .hits + .load(Ordering::Relaxed), + ), + MetricSample::new( + "direct_reqwest_client_cache_misses_total", + "Number of direct reqwest client cache misses.", + MetricKind::Counter, + DIRECT_REQWEST_CLIENT_CACHE_METRICS + .misses + .load(Ordering::Relaxed), + ), + MetricSample::new( + "direct_reqwest_client_cache_builds_total", + "Number of direct reqwest clients built after cache misses.", + MetricKind::Counter, + DIRECT_REQWEST_CLIENT_CACHE_METRICS + .builds + .load(Ordering::Relaxed), + ), + ] +} + pub(crate) fn build_browser_wreq_client( timeouts: Option<&aether_contracts::ExecutionTimeouts>, proxy: Option<&ProxySnapshot>, @@ -1622,6 +2066,19 @@ fn transport_profile_http1_only(transport_profile: Option<&ResolvedTransportProf .unwrap_or(false) } +fn transport_profile_h2c_prior_knowledge( + transport_profile: Option<&ResolvedTransportProfile>, +) -> bool { + transport_profile + .map(|profile| { + profile + .http_mode + .trim() + .eq_ignore_ascii_case(TRANSPORT_HTTP_MODE_H2C_PRIOR_KNOWLEDGE) + }) + .unwrap_or(false) +} + fn apply_transport_profile( builder: reqwest::ClientBuilder, transport_profile: Option<&ResolvedTransportProfile>, @@ -1630,7 +2087,24 @@ fn apply_transport_profile( return builder; }; let profile_id = profile.profile_id.trim(); - if profile_id.is_empty() { + if profile_id.is_empty() || transport_profile_h2c_prior_knowledge(Some(profile)) { + return builder; + } + + let _ = rustls::crypto::ring::default_provider().install_default(); + + builder.use_preconfigured_tls(build_best_effort_transport_tls_config()) +} + +fn apply_transport_profile_cache_key( + builder: reqwest::ClientBuilder, + transport_profile: Option<&DirectReqwestTransportProfileCacheKey>, +) -> reqwest::ClientBuilder { + let Some(profile) = transport_profile else { + return builder; + }; + if profile.profile_id.is_empty() || profile.http_mode == TRANSPORT_HTTP_MODE_H2C_PRIOR_KNOWLEDGE + { return builder; } @@ -1915,6 +2389,7 @@ mod tests { ExecutionPlan, ExecutionTimeouts, ProxySnapshot, RequestBody, ResolvedTransportProfile, EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER, EXECUTION_REQUEST_HTTP1_ONLY_HEADER, TRANSPORT_BACKEND_BROWSER_WREQ, TRANSPORT_BACKEND_REQWEST_RUSTLS, + TRANSPORT_HTTP_MODE_H2C_PRIOR_KNOWLEDGE, TRANSPORT_HTTP_MODE_HTTP1_ONLY, }; use aether_data::repository::proxy_nodes::{ InMemoryProxyNodeRepository, ProxyNodeReadRepository, StoredProxyNode, @@ -2024,6 +2499,186 @@ mod tests { } } + #[test] + fn direct_reqwest_client_cache_key_includes_transport_profile() { + let timeouts = ExecutionTimeouts { + connect_ms: Some(5_000), + ..ExecutionTimeouts::default() + }; + let h2c_profile = ResolvedTransportProfile { + profile_id: "mock-h2c".into(), + backend: TRANSPORT_BACKEND_REQWEST_RUSTLS.into(), + http_mode: TRANSPORT_HTTP_MODE_H2C_PRIOR_KNOWLEDGE.into(), + pool_scope: "key".into(), + header_fingerprint: None, + extra: Some(json!({"pool": "a"})), + }; + let same_h2c_profile = ResolvedTransportProfile { + extra: Some(json!({"pool": "a"})), + ..h2c_profile.clone() + }; + let http1_profile = ResolvedTransportProfile { + http_mode: TRANSPORT_HTTP_MODE_HTTP1_ONLY.into(), + ..h2c_profile.clone() + }; + + let left = super::direct_reqwest_client_cache_key( + Some(&timeouts), + None, + Some(&h2c_profile), + ExecutionTransportControls::default(), + ); + let right = super::direct_reqwest_client_cache_key( + Some(&timeouts), + None, + Some(&same_h2c_profile), + ExecutionTransportControls::default(), + ); + let different_mode = super::direct_reqwest_client_cache_key( + Some(&timeouts), + None, + Some(&http1_profile), + ExecutionTransportControls::default(), + ); + let different_proxy = super::direct_reqwest_client_cache_key( + Some(&timeouts), + Some("http://127.0.0.1:8080".into()), + Some(&h2c_profile), + ExecutionTransportControls::default(), + ); + + assert_eq!(left, right); + assert_ne!(left, different_mode); + assert_ne!(left, different_proxy); + assert!(super::direct_reqwest_client_cache_key_uses_http2(&left)); + assert!(!super::direct_reqwest_client_cache_key_uses_http2( + &different_mode + )); + } + + #[test] + fn direct_reqwest_h2_client_shards_scale_from_target_gate() { + assert_eq!( + super::direct_reqwest_h2_client_shards_from_config(None, 12_000, 200), + 60 + ); + assert_eq!( + super::direct_reqwest_h2_client_shards_from_config(None, 2_000, 200), + 10 + ); + assert_eq!( + super::direct_reqwest_h2_client_shards_from_config(Some(4), 12_000, 200), + 4 + ); + assert_eq!( + super::direct_reqwest_h2_client_shards_from_config(None, 100_000, 100), + super::MAX_DIRECT_REQWEST_H2_CLIENT_SHARDS + ); + } + + #[test] + fn direct_reqwest_initial_client_shards_are_bounded_by_target() { + assert_eq!(super::direct_reqwest_initial_client_shard_count(1), 1); + assert_eq!(super::direct_reqwest_initial_client_shard_count(2), 2); + assert_eq!( + super::direct_reqwest_initial_client_shard_count(21), + super::DEFAULT_DIRECT_REQWEST_SYNC_WARM_CLIENTS + ); + } + + #[test] + fn direct_reqwest_prewarm_populates_cache_for_plan() { + let profile = ResolvedTransportProfile { + profile_id: "mock-h2c-prewarm".into(), + backend: TRANSPORT_BACKEND_REQWEST_RUSTLS.into(), + http_mode: TRANSPORT_HTTP_MODE_H2C_PRIOR_KNOWLEDGE.into(), + pool_scope: "key".into(), + header_fingerprint: None, + extra: None, + }; + let plan = ExecutionPlan { + request_id: "req-prewarm".into(), + candidate_id: Some("candidate-prewarm".into()), + provider_name: Some("mock".into()), + provider_id: "provider-1".into(), + endpoint_id: "endpoint-1".into(), + key_id: "key-1".into(), + method: "POST".into(), + url: "http://127.0.0.1:18184/v1/chat/completions".into(), + headers: BTreeMap::new(), + content_type: Some("application/json".into()), + content_encoding: None, + body: RequestBody::from_json(json!({"stream": true})), + stream: true, + client_api_format: "openai:chat".into(), + provider_api_format: "openai:chat".into(), + model_name: Some("mock-model".into()), + proxy: None, + transport_profile: Some(profile.clone()), + timeouts: Some(ExecutionTimeouts { + connect_ms: Some(5_000), + ..ExecutionTimeouts::default() + }), + }; + + assert!( + super::try_prewarm_direct_reqwest_client_cache_for_plan(&plan) + .expect("prewarm should succeed") + ); + + let cache_key = super::direct_reqwest_client_cache_key( + plan.timeouts.as_ref(), + None, + Some(&profile), + super::ExecutionTransportControls::default(), + ); + let target_len = super::direct_reqwest_client_shard_count(&cache_key); + let expected_initial_len = super::direct_reqwest_initial_client_shard_count(target_len); + let cache = super::DIRECT_REQWEST_CLIENT_CACHE + .lock() + .expect("cache lock"); + let entry = cache.get(&cache_key).expect("cache entry"); + assert_eq!(entry.len(), expected_initial_len); + assert_eq!(entry.target_len, target_len); + } + + #[test] + fn direct_reqwest_prewarm_skips_browser_transport() { + let plan = ExecutionPlan { + request_id: "req-browser".into(), + candidate_id: None, + provider_name: Some("browser".into()), + provider_id: "provider-1".into(), + endpoint_id: "endpoint-1".into(), + key_id: "key-1".into(), + method: "POST".into(), + url: "https://example.com/v1/chat/completions".into(), + headers: BTreeMap::new(), + content_type: Some("application/json".into()), + content_encoding: None, + body: RequestBody::from_json(json!({"stream": true})), + stream: true, + client_api_format: "openai:chat".into(), + provider_api_format: "openai:chat".into(), + model_name: Some("mock-model".into()), + proxy: None, + transport_profile: Some(ResolvedTransportProfile { + profile_id: "chrome_136".into(), + backend: TRANSPORT_BACKEND_BROWSER_WREQ.into(), + http_mode: "auto".into(), + pool_scope: "key".into(), + header_fingerprint: None, + extra: None, + }), + timeouts: None, + }; + + assert!( + !super::try_prewarm_direct_reqwest_client_cache_for_plan(&plan) + .expect("browser transport should skip prewarm") + ); + } + #[test] fn direct_sync_execution_runtime_strips_accept_invalid_certs_control_header() { let headers = BTreeMap::from([ @@ -3667,6 +4322,73 @@ mod tests { ); } + #[tokio::test] + async fn direct_sync_execution_runtime_supports_h2c_prior_knowledge_profile() { + let listener = crate::test_support::bind_loopback_listener() + .await + .expect("listener should bind"); + let addr = listener.local_addr().expect("local addr should resolve"); + let app = Router::new().route( + "/chat", + post(|| async { + ( + axum::http::StatusCode::OK, + Json(json!({"transport_profile": "h2c"})), + ) + }), + ); + let server = tokio::spawn(async move { + axum::serve(listener, app) + .await + .expect("test server should run"); + }); + + let execution_runtime = DirectSyncExecutionRuntime::new(); + let result = execution_runtime + .execute_sync(&ExecutionPlan { + request_id: "req-h2c-1".into(), + candidate_id: Some("cand-h2c-1".into()), + provider_name: Some("mock".into()), + provider_id: "prov-1".into(), + endpoint_id: "ep-1".into(), + key_id: "key-1".into(), + method: "POST".into(), + url: format!("http://{addr}/chat"), + headers: BTreeMap::from([("content-type".into(), "application/json".into())]), + content_type: Some("application/json".into()), + content_encoding: None, + body: RequestBody::from_json(json!({"model": "mock-model"})), + stream: false, + client_api_format: "openai:chat".into(), + provider_api_format: "openai:chat".into(), + model_name: Some("mock-model".into()), + proxy: None, + transport_profile: Some(ResolvedTransportProfile { + profile_id: "mock-h2c".into(), + backend: TRANSPORT_BACKEND_REQWEST_RUSTLS.into(), + http_mode: TRANSPORT_HTTP_MODE_H2C_PRIOR_KNOWLEDGE.into(), + pool_scope: "key".into(), + header_fingerprint: None, + extra: None, + }), + timeouts: Some(ExecutionTimeouts { + connect_ms: Some(5_000), + total_ms: Some(LOCAL_HTTP_SUCCESS_TIMEOUT_MS), + ..ExecutionTimeouts::default() + }), + }) + .await + .expect("h2c prior-knowledge execution should succeed"); + + server.abort(); + + assert_eq!(result.status_code, 200); + assert_eq!( + result.body.and_then(|body| body.json_body), + Some(json!({"transport_profile": "h2c"})) + ); + } + #[test] fn direct_sync_execution_runtime_rejects_unsupported_transport_backend() { let profile = ResolvedTransportProfile { diff --git a/apps/aether-gateway/src/executor/candidate_loop.rs b/apps/aether-gateway/src/executor/candidate_loop.rs index 7f756b51c..e23bf6fe9 100644 --- a/apps/aether-gateway/src/executor/candidate_loop.rs +++ b/apps/aether-gateway/src/executor/candidate_loop.rs @@ -2,12 +2,14 @@ use aether_ai_serving::{ run_ai_attempt_loop, AiAttemptLoopOutcome, AiAttemptLoopPort, AiExecutionAttempt, }; use aether_data_contracts::repository::candidates::RequestCandidateStatus; +use aether_runtime::ConcurrencyPermit; use aether_scheduler_core::{ parse_request_candidate_report_context, SchedulerRequestCandidateStatusUpdate, }; use async_trait::async_trait; use axum::body::Body; use axum::http::Response; +use futures_util::StreamExt; use tokio::time::{timeout, Duration}; use tracing::{debug, warn, Instrument}; @@ -23,6 +25,7 @@ use crate::privacy::RedactionExecutionCandidateId; use crate::request_candidate_runtime::{ record_local_request_candidate_status, RequestCandidateRuntimeWriter, }; +use crate::stage_metrics::observe_gateway_stage_ms; use crate::{AppState, GatewayError}; const DEFAULT_STREAM_FIRST_BYTE_WATCHDOG_TIMEOUT_MS: u64 = 30_000; @@ -157,6 +160,8 @@ where type Error = GatewayError; async fn execute_attempt(&self, attempt: &T) -> Result, Self::Error> { + prewarm_direct_reqwest_candidate_client(attempt.execution_plan()); + let _permit = acquire_upstream_execution_gate(self.state, self.trace_id).await?; let mut response = execute_execution_runtime_sync( self.state, self.parts.uri.path(), @@ -323,13 +328,33 @@ where { let mut last_attempted = None; - while let Some(attempt) = - next_execution_attempt_with_timeout(source, trace_id, plan_kind, planning_timeout).await? - { + loop { + let next_started_at = std::time::Instant::now(); + let next_attempt = + next_execution_attempt_with_timeout(source, trace_id, plan_kind, planning_timeout) + .await?; + observe_gateway_stage_ms( + "stream_candidate_next", + next_started_at.elapsed().as_millis() as u64, + ); + let Some(attempt) = next_attempt else { + break; + }; last_attempted = Some((attempt.execution_plan().clone(), attempt.report_context())); - if let Some(response) = port.execute_attempt(&attempt).await? { + let execute_started_at = std::time::Instant::now(); + let response = port.execute_attempt(&attempt).await?; + observe_gateway_stage_ms( + "stream_candidate_execute", + execute_started_at.elapsed().as_millis() as u64, + ); + if let Some(response) = response { let remaining = source.drain_execution_attempts().await?; + let unused_started_at = std::time::Instant::now(); port.mark_unused_attempts(remaining).await?; + observe_gateway_stage_ms( + "stream_candidate_unused", + unused_started_at.elapsed().as_millis() as u64, + ); return Ok(LocalExecutionRequestOutcome::responded(response)); } } @@ -412,6 +437,7 @@ where candidate_index = candidate_index.as_str(), "candidate loop attempting stream execution candidate" ); + prewarm_direct_reqwest_candidate_client(&plan); let watchdog_plan = plan.clone(); let watchdog_report_context = report_context.clone(); let execution_state = self.state.clone(); @@ -475,6 +501,15 @@ where } } +fn prewarm_direct_reqwest_candidate_client(plan: &aether_contracts::ExecutionPlan) { + let started_at = std::time::Instant::now(); + crate::execution_runtime::transport::prewarm_direct_reqwest_client_cache_for_plan(plan); + observe_gateway_stage_ms( + "direct_reqwest_client_prewarm", + started_at.elapsed().as_millis() as u64, + ); +} + pub(crate) async fn mark_unused_local_candidates(state: &AppState, remaining: Vec) where T: AiExecutionAttempt, @@ -543,7 +578,7 @@ fn stream_candidate_watchdog_timeout_message() -> &'static str { } async fn execute_stream_candidate_with_watchdog( - state: &(impl RequestCandidateRuntimeWriter + ?Sized), + state: &(impl RequestCandidateRuntimeWriter + UpstreamExecutionGateProvider + ?Sized), trace_id: &str, plan_kind: &str, plan: &aether_contracts::ExecutionPlan, @@ -556,9 +591,12 @@ where { let timeout_duration = resolve_stream_candidate_watchdog_timeout(plan, report_context); let candidate_started_unix_ms = current_unix_ms(); + let permit = acquire_upstream_execution_gate(state, trace_id).await?; let mut join_handle = tokio::spawn(execute()); match timeout(timeout_duration, &mut join_handle).await { - Ok(Ok(result)) => result, + Ok(Ok(result)) => { + result.map(|response| maybe_hold_upstream_execution_permit(response, permit)) + } Ok(Err(join_error)) => Err(GatewayError::Internal(format!( "local stream candidate task join failed: {join_error}" ))), @@ -608,6 +646,67 @@ where } } +fn maybe_hold_upstream_execution_permit( + response: Option>, + permit: Option, +) -> Option> { + match (response, permit) { + (Some(response), Some(permit)) => { + Some(hold_response_upstream_execution_permit(response, permit)) + } + (response, _) => response, + } +} + +fn hold_response_upstream_execution_permit( + response: Response, + permit: ConcurrencyPermit, +) -> Response { + let (parts, body) = response.into_parts(); + let stream = async_stream::stream! { + let _permit = permit; + let mut body_stream = body.into_data_stream(); + while let Some(item) = body_stream.next().await { + yield item; + } + }; + Response::from_parts(parts, Body::from_stream(stream)) +} + +trait UpstreamExecutionGateProvider { + fn upstream_execution_gate(&self) -> Option<&aether_runtime::ConcurrencyGate>; + fn upstream_execution_gate_queue_budget(&self) -> Duration; +} + +impl UpstreamExecutionGateProvider for AppState { + fn upstream_execution_gate(&self) -> Option<&aether_runtime::ConcurrencyGate> { + self.upstream_execution_gate.as_deref() + } + + fn upstream_execution_gate_queue_budget(&self) -> Duration { + self.frontdoor_runtime_guards.internal_gate_queue_budget + } +} + +async fn acquire_upstream_execution_gate( + state: &(impl UpstreamExecutionGateProvider + ?Sized), + trace_id: &str, +) -> Result, GatewayError> { + let Some(gate) = state.upstream_execution_gate() else { + return Ok(None); + }; + let budget = state.upstream_execution_gate_queue_budget(); + match timeout(budget, gate.acquire()).await { + Ok(Ok(permit)) => Ok(Some(permit)), + Ok(Err(err)) => Err(GatewayError::Internal(err.to_string())), + Err(_) => Err(GatewayError::AdmissionTimeout { + trace_id: trace_id.to_string(), + gate: "gateway_upstream_execution", + queue_budget_ms: budget.as_millis() as u64, + }), + } +} + pub(crate) async fn mark_unused_local_candidate_items( state: &AppState, remaining: Vec, @@ -677,6 +776,16 @@ mod tests { } } + impl UpstreamExecutionGateProvider for TestRequestCandidateWriter { + fn upstream_execution_gate(&self) -> Option<&aether_runtime::ConcurrencyGate> { + None + } + + fn upstream_execution_gate_queue_budget(&self) -> Duration { + Duration::from_millis(250) + } + } + struct PendingAttemptSource; #[async_trait] diff --git a/apps/aether-gateway/src/executor/mod.rs b/apps/aether-gateway/src/executor/mod.rs index dc6aaa44f..8d72a4a58 100644 --- a/apps/aether-gateway/src/executor/mod.rs +++ b/apps/aether-gateway/src/executor/mod.rs @@ -16,7 +16,8 @@ pub(crate) use candidate_loop::{ }; pub(crate) use orchestration::*; pub(crate) use outcome::{ - beautify_local_execution_client_error_message, build_local_execution_exhaustion, + beautify_local_execution_client_error_message, build_fast_local_execution_exhaustion, + build_fast_local_execution_runtime_miss_context, build_local_execution_exhaustion, build_local_execution_runtime_miss_context, record_failed_usage_for_exhausted_request, record_failed_usage_for_runtime_miss_request, LocalExecutionExhaustion, LocalExecutionRequestOutcome, LocalExecutionRuntimeMissContext, diff --git a/apps/aether-gateway/src/executor/outcome.rs b/apps/aether-gateway/src/executor/outcome.rs index 889749fbe..b31ccb9bc 100644 --- a/apps/aether-gateway/src/executor/outcome.rs +++ b/apps/aether-gateway/src/executor/outcome.rs @@ -150,6 +150,7 @@ pub(crate) async fn build_local_execution_exhaustion( plan: &ExecutionPlan, report_context: Option<&Value>, ) -> LocalExecutionExhaustion { + let mut exhaustion = build_fast_local_execution_exhaustion(plan, report_context); let mut data = build_usage_event_data_seed(plan, report_context); let last_failed_candidate = match state .read_request_candidates_by_request_id(plan.request_id.as_str()) @@ -180,28 +181,46 @@ pub(crate) async fn build_local_execution_exhaustion( .or_else(|| candidate.key_id.clone()); } + exhaustion.data = data; + exhaustion.candidate_id = last_failed_candidate + .as_ref() + .map(|candidate| candidate.id.clone()); + exhaustion.candidate_index = last_failed_candidate + .as_ref() + .map(|candidate| candidate.candidate_index); + exhaustion.upstream_status_code = last_failed_candidate + .as_ref() + .and_then(|candidate| candidate.status_code); + exhaustion.upstream_error_type = last_failed_candidate + .as_ref() + .and_then(|candidate| candidate.error_type.clone()) + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()); + exhaustion.upstream_error_message = last_failed_candidate + .as_ref() + .and_then(|candidate| candidate.error_message.clone()) + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()); + exhaustion +} + +pub(crate) fn build_fast_local_execution_exhaustion( + plan: &ExecutionPlan, + report_context: Option<&Value>, +) -> LocalExecutionExhaustion { + let data = build_usage_event_data_seed(plan, report_context); LocalExecutionExhaustion { request_id: plan.request_id.clone(), + candidate_id: plan.candidate_id.clone(), + candidate_index: report_context + .and_then(Value::as_object) + .and_then(|value| value.get("candidate_index")) + .and_then(Value::as_u64) + .and_then(|value| u32::try_from(value).ok()), data, - candidate_id: last_failed_candidate - .as_ref() - .map(|candidate| candidate.id.clone()), - candidate_index: last_failed_candidate - .as_ref() - .map(|candidate| candidate.candidate_index), - upstream_status_code: last_failed_candidate - .as_ref() - .and_then(|candidate| candidate.status_code), - upstream_error_type: last_failed_candidate - .as_ref() - .and_then(|candidate| candidate.error_type.clone()) - .map(|value| value.trim().to_string()) - .filter(|value| !value.is_empty()), - upstream_error_message: last_failed_candidate - .as_ref() - .and_then(|candidate| candidate.error_message.clone()) - .map(|value| value.trim().to_string()) - .filter(|value| !value.is_empty()), + upstream_status_code: None, + upstream_error_type: None, + upstream_error_message: None, } } @@ -224,6 +243,20 @@ pub(crate) async fn build_local_execution_runtime_miss_context( } } +pub(crate) fn build_fast_local_execution_runtime_miss_context( + decision: Option<&GatewayControlDecision>, +) -> LocalExecutionRuntimeMissContext { + let auth_context = decision.and_then(|value| value.auth_context.as_ref()); + + LocalExecutionRuntimeMissContext { + auth_user_id: auth_context.map(|value| value.user_id.clone()), + auth_api_key_id: auth_context.map(|value| value.api_key_id.clone()), + auth_username: auth_context.and_then(|value| value.username.clone()), + auth_api_key_name: auth_context.and_then(|value| value.api_key_name.clone()), + candidate_contexts: Vec::new(), + } +} + pub(crate) async fn record_failed_usage_for_exhausted_request( state: &AppState, exhaustion: LocalExecutionExhaustion, @@ -309,13 +342,10 @@ pub(crate) async fn record_failed_usage_for_exhausted_request( ); data.request_metadata = Some(Value::Object(request_metadata)); - state - .usage_runtime - .record_terminal_event_direct( - state.data.as_ref(), - UsageEvent::new(UsageEventType::Failed, request_id, data), - ) - .await; + state.usage_runtime.submit_terminal_event( + state.data.as_ref(), + UsageEvent::new(UsageEventType::Failed, request_id, data), + ); } pub(crate) async fn record_failed_usage_for_runtime_miss_request( @@ -448,13 +478,10 @@ pub(crate) async fn record_failed_usage_for_runtime_miss_request( data.request_metadata = (!request_metadata.is_empty()).then_some(Value::Object(request_metadata)); - state - .usage_runtime - .record_terminal_event_direct( - state.data.as_ref(), - UsageEvent::new(UsageEventType::Failed, request_id, data), - ) - .await; + state.usage_runtime.submit_terminal_event( + state.data.as_ref(), + UsageEvent::new(UsageEventType::Failed, request_id, data), + ); } pub(crate) fn beautify_local_execution_client_error_message(message: &str) -> String { diff --git a/apps/aether-gateway/src/executor/stream_path.rs b/apps/aether-gateway/src/executor/stream_path.rs index 35962b500..0be7fc885 100644 --- a/apps/aether-gateway/src/executor/stream_path.rs +++ b/apps/aether-gateway/src/executor/stream_path.rs @@ -13,6 +13,7 @@ use crate::ai_serving::api::{ }; use crate::api::response::build_client_response_from_parts; use crate::control::GatewayControlDecision; +use crate::stage_metrics::observe_gateway_stage_ms; use crate::{AppState, GatewayError, GatewayFallbackReason}; use super::{ @@ -94,6 +95,7 @@ impl AiStreamExecutionPathPort for GatewayStreamExecutionPathPort<'_> { &self, step: AiStreamExecutionStep, ) -> Result, Self::Error> { + let step_started_at = std::time::Instant::now(); let outcome = match step { AiStreamExecutionStep::LocalVideoContent => { maybe_execute_local_video_task_content_stream( @@ -188,6 +190,10 @@ impl AiStreamExecutionPathPort for GatewayStreamExecutionPathPort<'_> { } } }; + observe_gateway_stage_ms( + "stream_path_step", + step_started_at.elapsed().as_millis() as u64, + ); Ok(to_ai_serving_outcome(outcome)) } diff --git a/apps/aether-gateway/src/handlers/proxy/mod.rs b/apps/aether-gateway/src/handlers/proxy/mod.rs index 164659c30..8a3b24387 100644 --- a/apps/aether-gateway/src/handlers/proxy/mod.rs +++ b/apps/aether-gateway/src/handlers/proxy/mod.rs @@ -60,6 +60,7 @@ use crate::scheduler::candidate::{ LEGACY_API_KEY_CONCURRENCY_LIMIT_SKIP_REASON, }; use crate::scheduler::config::{read_scheduler_ordering_config, SchedulerSchedulingMode}; +use crate::stage_metrics::observe_gateway_stage_ms; use crate::{ AppState, FrontdoorUserRpmOutcome, GatewayError, GatewayFallbackMetricKind, GatewayFallbackReason, LocalExecutionRuntimeMissDiagnostic, @@ -1023,8 +1024,30 @@ pub(crate) async fn proxy_request( State(state): State, ConnectInfo(remote_addr): ConnectInfo, request: Request, +) -> Result, GatewayError> { + crate::request_diagnostics::scope_request_diagnostics(proxy_request_inner( + state, + remote_addr, + request, + )) + .await +} + +async fn proxy_request_inner( + state: AppState, + remote_addr: std::net::SocketAddr, + request: Request, ) -> Result, GatewayError> { let started_at = Instant::now(); + if let Some(accepted_at) = request + .extensions() + .get::() + { + observe_gateway_stage_ms( + "frontdoor_handler_queue", + started_at.duration_since(accepted_at.0).as_millis() as u64, + ); + } let mut request_permit = match state.try_acquire_request_permit().await { Ok(permit) => permit, Err(RequestAdmissionError::Local(aether_runtime::ConcurrencyError::Saturated { @@ -1085,6 +1108,7 @@ pub(crate) async fn proxy_request( )) => return Err(GatewayError::Internal(message)), }; let request_admission_ms = started_at.elapsed().as_millis() as u64; + observe_gateway_stage_ms("frontdoor_admission", request_admission_ms); let (mut parts, body) = request.into_parts(); let redaction_slot = crate::privacy::RedactionSessionSlot::default(); parts.extensions.insert(redaction_slot.clone()); @@ -1176,6 +1200,7 @@ pub(crate) async fn proxy_request( } } let request_context_ms = request_context_started_at.elapsed().as_millis() as u64; + observe_gateway_stage_ms("frontdoor_context", request_context_ms); if request_context .control_decision .as_ref() @@ -1202,6 +1227,7 @@ pub(crate) async fn proxy_request( let mut request_body = Some(body); let local_proxy_body = if local_proxy_route_requires_buffered_body(&request_context) { let body_buffer_policy = RequestBodyBufferPolicy::from_state(&state); + let stage_started_at = Instant::now(); let body = buffer_and_normalize_request_body( &mut request_body, &mut parts.headers, @@ -1213,6 +1239,10 @@ pub(crate) async fn proxy_request( body_buffer_policy, ) .await; + observe_gateway_stage_ms( + "frontdoor_body_buffer", + stage_started_at.elapsed().as_millis() as u64, + ); match body { Ok(body) => Some(body), Err(err) => { @@ -1388,6 +1418,7 @@ pub(crate) async fn proxy_request( let buffered_body = if should_buffer_body { let body_buffer_policy = RequestBodyBufferPolicy::from_state(&state); + let stage_started_at = Instant::now(); let body = buffer_and_normalize_request_body( &mut request_body, &mut parts.headers, @@ -1399,6 +1430,10 @@ pub(crate) async fn proxy_request( body_buffer_policy, ) .await; + observe_gateway_stage_ms( + "frontdoor_body_buffer", + stage_started_at.elapsed().as_millis() as u64, + ); match body { Ok(body) => Some(body), Err(err) => { @@ -1417,15 +1452,20 @@ pub(crate) async fn proxy_request( None }; - if let Some(response) = maybe_forward_public_request_to_tunnel_owner( + let owner_forward_started_at = Instant::now(); + let owner_forward_response = maybe_forward_public_request_to_tunnel_owner( &state, &remote_addr, &request_context, &parts, buffered_body.as_ref(), ) - .await? - { + .await?; + observe_gateway_stage_ms( + "frontdoor_owner_forward", + owner_forward_started_at.elapsed().as_millis() as u64, + ); + if let Some(response) = owner_forward_response { return Ok(finalize_gateway_response_with_context( &state, response, @@ -1452,15 +1492,20 @@ pub(crate) async fn proxy_request( } if let Some(buffered_body) = buffered_body.as_ref() { - if let Some(rejection) = request_model_local_rejection( + let auth_model_started_at = Instant::now(); + let model_rejection = request_model_local_rejection( &state, control_decision, &parts.uri, &parts.headers, buffered_body, ) - .await? - { + .await?; + observe_gateway_stage_ms( + "frontdoor_auth_model", + auth_model_started_at.elapsed().as_millis() as u64, + ); + if let Some(rejection) = model_rejection { let response = build_local_auth_rejection_response(&trace_id, control_decision, &rejection)?; return Ok(finalize_gateway_response_with_context( @@ -1475,10 +1520,12 @@ pub(crate) async fn proxy_request( } } + let rpm_started_at = Instant::now(); let rate_limit_outcome = state .frontdoor_user_rpm() .check_and_consume(&state, control_decision) .await?; + observe_gateway_stage_ms("frontdoor_rpm", rpm_started_at.elapsed().as_millis() as u64); if let FrontdoorUserRpmOutcome::Rejected(rejection) = &rate_limit_outcome { let auth_context = control_decision.and_then(|decision| decision.auth_context.as_ref()); let user_id = auth_context @@ -1514,13 +1561,18 @@ pub(crate) async fn proxy_request( )); } - if let Some(response) = super::public::maybe_build_local_ai_public_response( + let local_ai_public_started_at = Instant::now(); + let local_ai_public_response = super::public::maybe_build_local_ai_public_response( &state, &request_context, buffered_body.as_ref(), ) - .await - { + .await; + observe_gateway_stage_ms( + "frontdoor_local_ai_public", + local_ai_public_started_at.elapsed().as_millis() as u64, + ); + if let Some(response) = local_ai_public_response { return Ok(finalize_gateway_response_with_context( &state, response, @@ -1557,6 +1609,7 @@ pub(crate) async fn proxy_request( let stream_request = request_wants_stream(&request_context, &parts.headers, buffered_body); let mut local_execution_exhaustion = None; if stream_request { + let execute_stream_started_at = Instant::now(); let stream_outcome = match maybe_execute_stream_request( &state, &parts, @@ -1568,6 +1621,10 @@ pub(crate) async fn proxy_request( { Ok(outcome) => outcome, Err(err) => { + observe_gateway_stage_ms( + "frontdoor_execute_stream", + execute_stream_started_at.elapsed().as_millis() as u64, + ); if let Some((phase, timeout_ms)) = local_execution_planning_timeout_parts(&err) { return finalize_local_execution_planning_timeout( @@ -1585,6 +1642,10 @@ pub(crate) async fn proxy_request( return Err(err); } }; + observe_gateway_stage_ms( + "frontdoor_execute_stream", + execute_stream_started_at.elapsed().as_millis() as u64, + ); debug!( event_name = "proxy_stream_local_execute_outcome", log_type = "debug", @@ -1622,6 +1683,7 @@ pub(crate) async fn proxy_request( LocalExecutionRequestOutcome::NoPath => {} } } + let execute_sync_started_at = Instant::now(); let sync_outcome = match maybe_execute_sync_request( &state, &parts, @@ -1633,6 +1695,10 @@ pub(crate) async fn proxy_request( { Ok(outcome) => outcome, Err(err) => { + observe_gateway_stage_ms( + "frontdoor_execute_sync", + execute_sync_started_at.elapsed().as_millis() as u64, + ); if let Some((phase, timeout_ms)) = local_execution_planning_timeout_parts(&err) { return finalize_local_execution_planning_timeout( &state, @@ -1649,6 +1715,10 @@ pub(crate) async fn proxy_request( return Err(err); } }; + observe_gateway_stage_ms( + "frontdoor_execute_sync", + execute_sync_started_at.elapsed().as_millis() as u64, + ); match sync_outcome { LocalExecutionRequestOutcome::Responded(execution_runtime_response) => { let execution_runtime_response = restore_redacted_sync_execution_response( @@ -1673,6 +1743,7 @@ pub(crate) async fn proxy_request( LocalExecutionRequestOutcome::NoPath => {} } if parts.method != http::Method::POST { + let execute_stream_started_at = Instant::now(); let stream_outcome = match maybe_execute_stream_request( &state, &parts, @@ -1684,6 +1755,10 @@ pub(crate) async fn proxy_request( { Ok(outcome) => outcome, Err(err) => { + observe_gateway_stage_ms( + "frontdoor_execute_stream", + execute_stream_started_at.elapsed().as_millis() as u64, + ); if let Some((phase, timeout_ms)) = local_execution_planning_timeout_parts(&err) { return finalize_local_execution_planning_timeout( @@ -1701,6 +1776,10 @@ pub(crate) async fn proxy_request( return Err(err); } }; + observe_gateway_stage_ms( + "frontdoor_execute_stream", + execute_stream_started_at.elapsed().as_millis() as u64, + ); match stream_outcome { LocalExecutionRequestOutcome::Responded(execution_runtime_response) => { let execution_runtime_response = restore_redacted_stream_execution_response( diff --git a/apps/aether-gateway/src/lib.rs b/apps/aether-gateway/src/lib.rs index 13fc509ea..893d3b787 100644 --- a/apps/aether-gateway/src/lib.rs +++ b/apps/aether-gateway/src/lib.rs @@ -63,15 +63,18 @@ pub(crate) use aether_provider_transport as provider_transport; mod rate_limit; mod request_candidate_queue; mod request_candidate_runtime; +mod request_diagnostics; mod roles; mod router; mod routing; mod scheduler; mod server_chan_push; +mod stage_metrics; mod state; mod system_features; mod task_runtime; mod tunnel; +mod upstream_admission; mod usage; mod video_tasks; mod wallet_runtime; @@ -126,7 +129,8 @@ fn insert_header_if_missing( if headers.contains_key(key) { return Ok(()); } - let name = HeaderName::from_static(key); + let name = HeaderName::from_bytes(key.as_bytes()) + .map_err(|err| GatewayError::Internal(err.to_string()))?; let value = HeaderValue::from_str(value).map_err(|err| GatewayError::Internal(err.to_string()))?; headers.insert(name, value); diff --git a/apps/aether-gateway/src/main.rs b/apps/aether-gateway/src/main.rs index 226d3175e..719eaeda7 100644 --- a/apps/aether-gateway/src/main.rs +++ b/apps/aether-gateway/src/main.rs @@ -236,6 +236,11 @@ const AUTO_SERVER_SQL_POOL_MIN_CONNECTIONS_FLOOR: u32 = 4; const AUTO_SERVER_SQL_POOL_MIN_CONNECTIONS_CAP: u32 = 16; const AUTO_SERVER_SQL_POOL_MAX_CONNECTIONS_FLOOR: u32 = 20; const AUTO_SERVER_SQL_POOL_MAX_CONNECTIONS_CAP: u32 = 100; +const DEFAULT_GATEWAY_LISTEN_BACKLOG: i32 = 8192; +const MIN_GATEWAY_LISTEN_BACKLOG: i32 = 128; +const MAX_GATEWAY_LISTEN_BACKLOG: i32 = 65_535; +const DEFAULT_GATEWAY_LISTENER_SHARDS: usize = 1; +const MAX_GATEWAY_LISTENER_SHARDS: usize = 64; fn env_var_trimmed(name: &str) -> Option { std::env::var(name) .ok() @@ -570,6 +575,34 @@ struct GatewayUsageArgs { default_value_t = 5_000 )] queue_reclaim_interval_ms: u64, + + #[arg( + long, + env = "AETHER_GATEWAY_USAGE_ENQUEUE_RETRY_BUFFER_CAPACITY", + default_value_t = 131_072 + )] + enqueue_retry_buffer_capacity: usize, + + #[arg( + long, + env = "AETHER_GATEWAY_USAGE_ENQUEUE_RETRY_WORKERS", + default_value_t = 4 + )] + enqueue_retry_workers: usize, + + #[arg( + long, + env = "AETHER_GATEWAY_USAGE_ENQUEUE_RETRY_INITIAL_BACKOFF_MS", + default_value_t = 10 + )] + enqueue_retry_initial_backoff_ms: u64, + + #[arg( + long, + env = "AETHER_GATEWAY_USAGE_ENQUEUE_RETRY_MAX_BACKOFF_MS", + default_value_t = 1_000 + )] + enqueue_retry_max_backoff_ms: u64, } impl GatewayUsageArgs { @@ -587,6 +620,12 @@ impl GatewayUsageArgs { reclaim_idle_ms: self.queue_reclaim_idle_ms.max(1), reclaim_count: self.queue_reclaim_count.max(1), reclaim_interval_ms: self.queue_reclaim_interval_ms.max(1), + enqueue_retry_buffer_capacity: self.enqueue_retry_buffer_capacity.max(1), + enqueue_retry_workers: self.enqueue_retry_workers.clamp(1, 64), + enqueue_retry_initial_backoff_ms: self.enqueue_retry_initial_backoff_ms.max(1), + enqueue_retry_max_backoff_ms: self + .enqueue_retry_max_backoff_ms + .max(self.enqueue_retry_initial_backoff_ms.max(1)), } } } @@ -755,6 +794,20 @@ struct Args { #[arg(long, env = "APP_PORT", default_value_t = 8084)] app_port: u16, + #[arg( + long, + env = "AETHER_GATEWAY_LISTEN_BACKLOG", + default_value_t = DEFAULT_GATEWAY_LISTEN_BACKLOG + )] + listen_backlog: i32, + + #[arg( + long, + env = "AETHER_GATEWAY_LISTENER_SHARDS", + default_value_t = DEFAULT_GATEWAY_LISTENER_SHARDS + )] + listener_shards: usize, + /// 容器内健康检查入口：根据当前 bind 端口探测本地 /health。 #[arg(long, hide = true, default_value_t = false)] healthcheck: bool, @@ -1004,6 +1057,98 @@ fn gateway_bind_addr(app_port: u16) -> Result i32 { + backlog.clamp(MIN_GATEWAY_LISTEN_BACKLOG, MAX_GATEWAY_LISTEN_BACKLOG) +} + +fn gateway_listener_shards(shards: usize) -> usize { + shards.clamp(1, MAX_GATEWAY_LISTENER_SHARDS) +} + +fn gateway_listener( + bind_addr: std::net::SocketAddr, + backlog: i32, + reuse_port: bool, +) -> Result { + let domain = match bind_addr { + std::net::SocketAddr::V4(_) => socket2::Domain::IPV4, + std::net::SocketAddr::V6(_) => socket2::Domain::IPV6, + }; + let socket = socket2::Socket::new(domain, socket2::Type::STREAM, Some(socket2::Protocol::TCP))?; + socket.set_reuse_address(true)?; + if reuse_port { + set_gateway_listener_reuse_port(&socket)?; + } + socket.set_nonblocking(true)?; + socket.set_tcp_nodelay(true)?; + socket.bind(&bind_addr.into())?; + socket.listen(gateway_listen_backlog(backlog))?; + tokio::net::TcpListener::from_std(socket.into()) +} + +#[cfg(unix)] +fn set_gateway_listener_reuse_port(socket: &socket2::Socket) -> Result<(), std::io::Error> { + socket.set_reuse_port(true) +} + +#[cfg(not(unix))] +fn set_gateway_listener_reuse_port(_socket: &socket2::Socket) -> Result<(), std::io::Error> { + Err(std::io::Error::new( + std::io::ErrorKind::Unsupported, + "AETHER_GATEWAY_LISTENER_SHARDS > 1 requires SO_REUSEPORT support", + )) +} + +fn gateway_listeners( + bind_addr: std::net::SocketAddr, + backlog: i32, + shards: usize, +) -> Result, std::io::Error> { + let shards = gateway_listener_shards(shards); + let mut listeners = Vec::with_capacity(shards); + for _ in 0..shards { + listeners.push(gateway_listener(bind_addr, backlog, shards > 1)?); + } + Ok(listeners) +} + +async fn serve_gateway_router( + listeners: Vec, + router: axum::Router, +) -> Result<(), Box> { + if listeners.len() == 1 { + let listener = listeners + .into_iter() + .next() + .ok_or_else(|| std::io::Error::other("gateway listener set is empty"))?; + axum::serve( + listener, + router.into_make_service_with_connect_info::(), + ) + .await?; + return Ok(()); + } + + let mut servers = tokio::task::JoinSet::new(); + for listener in listeners { + let router = router.clone(); + servers.spawn(async move { + axum::serve( + listener, + router.into_make_service_with_connect_info::(), + ) + .await + }); + } + if let Some(result) = servers.join_next().await { + servers.abort_all(); + let serve_result = result + .map_err(|err| std::io::Error::other(format!("gateway listener task failed: {err}")))?; + serve_result?; + } + Ok(()) +} + fn resolve_local_http_base_url(app_port: u16) -> Result { Ok(format!("http://127.0.0.1:{}", validate_app_port(app_port)?)) } @@ -1323,6 +1468,20 @@ async fn run() -> Result<(), Box> { ); } state.bootstrap_admin_from_env().await?; + match state.prewarm_chat_pii_redaction_runtime_config().await { + Ok(enabled) => { + info!( + chat_pii_redaction_enabled = enabled, + "prewarmed chat pii redaction runtime config" + ); + } + Err(err) => { + warn!( + error = %err, + "failed to prewarm chat pii redaction runtime config" + ); + } + } let background_tasks = if args.node_role.spawns_background_tasks() { Some(state.spawn_background_tasks()) @@ -1333,7 +1492,9 @@ async fn run() -> Result<(), Box> { ); None }; - let listener = tokio::net::TcpListener::bind(bind_addr).await?; + let listen_backlog = gateway_listen_backlog(args.listen_backlog); + let listener_shards = gateway_listener_shards(args.listener_shards); + let listeners = gateway_listeners(bind_addr, listen_backlog, listener_shards)?; let public_base_url = resolve_local_http_base_url(app_port)?; let frontdoor_health_url = format!("{public_base_url}/_gateway/health"); let api_router = build_router_with_state(state); @@ -1353,17 +1514,15 @@ async fn run() -> Result<(), Box> { log_type = "ops", bind = %bind_addr, app_port, + listen_backlog, + listener_shards, public_url = %public_base_url, healthcheck_url = %frontdoor_health_url, legacy_route_policy = "fail_closed", "aether-gateway ready" ); - axum::serve( - listener, - router.into_make_service_with_connect_info::(), - ) - .await?; + serve_gateway_router(listeners, router).await?; if let Some(background_tasks) = background_tasks { background_tasks.shutdown().await; } @@ -1730,7 +1889,8 @@ mod tests { DatabaseDriverArg, DeploymentTopologyArg, GatewayDataArgs, GatewayFrontdoorArgs, GatewayLogDestinationArg, GatewayLogFormatArg, GatewayLogRotationArg, GatewayLoggingArgs, GatewayRateLimitArgs, GatewayUsageArgs, NodeRoleArg, RuntimeBackendArg, - VideoTaskTruthSourceArg, + VideoTaskTruthSourceArg, DEFAULT_GATEWAY_LISTENER_SHARDS, DEFAULT_GATEWAY_LISTEN_BACKLOG, + MAX_GATEWAY_LISTENER_SHARDS, MAX_GATEWAY_LISTEN_BACKLOG, MIN_GATEWAY_LISTEN_BACKLOG, }; use aether_data::{DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig}; use aether_gateway::AppState; @@ -1739,6 +1899,8 @@ mod tests { Args { command: None, app_port: 8084, + listen_backlog: DEFAULT_GATEWAY_LISTEN_BACKLOG, + listener_shards: DEFAULT_GATEWAY_LISTENER_SHARDS, healthcheck: false, healthcheck_timeout_ms: 3_000, deployment_topology: DeploymentTopologyArg::SingleNode, @@ -1789,6 +1951,10 @@ mod tests { queue_reclaim_idle_ms: 30_000, queue_reclaim_count: 500, queue_reclaim_interval_ms: 5_000, + enqueue_retry_buffer_capacity: 131_072, + enqueue_retry_workers: 4, + enqueue_retry_initial_backoff_ms: 10, + enqueue_retry_max_backoff_ms: 1_000, }, frontdoor: GatewayFrontdoorArgs { environment: "development".to_string(), @@ -1825,6 +1991,35 @@ mod tests { assert_eq!(error.kind(), std::io::ErrorKind::InvalidInput); } + #[test] + fn clamps_gateway_listen_backlog() { + assert_eq!( + super::gateway_listen_backlog(MIN_GATEWAY_LISTEN_BACKLOG - 1), + MIN_GATEWAY_LISTEN_BACKLOG + ); + assert_eq!( + super::gateway_listen_backlog(DEFAULT_GATEWAY_LISTEN_BACKLOG), + DEFAULT_GATEWAY_LISTEN_BACKLOG + ); + assert_eq!( + super::gateway_listen_backlog(MAX_GATEWAY_LISTEN_BACKLOG + 1), + MAX_GATEWAY_LISTEN_BACKLOG + ); + } + + #[test] + fn clamps_gateway_listener_shards() { + assert_eq!(super::gateway_listener_shards(0), 1); + assert_eq!( + super::gateway_listener_shards(DEFAULT_GATEWAY_LISTENER_SHARDS), + DEFAULT_GATEWAY_LISTENER_SHARDS + ); + assert_eq!( + super::gateway_listener_shards(MAX_GATEWAY_LISTENER_SHARDS + 1), + MAX_GATEWAY_LISTENER_SHARDS + ); + } + #[test] fn explicit_migrate_runtime_config_enables_data_logs() { let mut args = test_args(); diff --git a/apps/aether-gateway/src/middleware/access_log.rs b/apps/aether-gateway/src/middleware/access_log.rs index 43d9f4905..bbf4cd6d5 100644 --- a/apps/aether-gateway/src/middleware/access_log.rs +++ b/apps/aether-gateway/src/middleware/access_log.rs @@ -18,6 +18,9 @@ use crate::log_ids::short_request_id; #[derive(Debug, Clone, Copy)] pub(crate) struct RequestLogEmitted; +#[derive(Debug, Clone, Copy)] +pub(crate) struct GatewayRequestAcceptedAt(pub(crate) Instant); + fn is_usage_detail_path(path: &str) -> bool { let Some(detail_id) = path.strip_prefix("/api/admin/usage/") else { return false; @@ -59,6 +62,9 @@ pub(crate) fn sanitize_access_log_path(path: &str) -> String { pub(crate) async fn access_log_middleware(mut request: Request, next: Next) -> Response { let started_at = Instant::now(); + request + .extensions_mut() + .insert(GatewayRequestAcceptedAt(started_at)); let method = request.method().clone(); let raw_path = request .uri() diff --git a/apps/aether-gateway/src/middleware/mod.rs b/apps/aether-gateway/src/middleware/mod.rs index 779c6f129..2aa6137d9 100644 --- a/apps/aether-gateway/src/middleware/mod.rs +++ b/apps/aether-gateway/src/middleware/mod.rs @@ -3,7 +3,8 @@ mod frontdoor_cors; mod strip_cf_headers; pub(crate) use access_log::{ - access_log_middleware, sanitize_access_log_path, should_downgrade_access_log, RequestLogEmitted, + access_log_middleware, sanitize_access_log_path, should_downgrade_access_log, + GatewayRequestAcceptedAt, RequestLogEmitted, }; pub(crate) use frontdoor_cors::frontdoor_cors_middleware; pub use strip_cf_headers::strip_cf_headers_middleware; diff --git a/apps/aether-gateway/src/orchestration/effects.rs b/apps/aether-gateway/src/orchestration/effects.rs index 9a1c03f89..1484f386d 100644 --- a/apps/aether-gateway/src/orchestration/effects.rs +++ b/apps/aether-gateway/src/orchestration/effects.rs @@ -43,16 +43,22 @@ use crate::{ }; const POOL_SCORE_FEEDBACK_GATE_MAX_ENTRIES: usize = 50_000; +const HEALTH_SUCCESS_PERSIST_GATE_MAX_ENTRIES: usize = 50_000; const POOL_SCORE_SUCCESS_FEEDBACK_MIN_INTERVAL_ENV: &str = "AETHER_GATEWAY_POOL_SCORE_SUCCESS_FEEDBACK_MIN_INTERVAL_SECS"; const POOL_SCORE_FAILURE_FEEDBACK_MIN_INTERVAL_ENV: &str = "AETHER_GATEWAY_POOL_SCORE_FAILURE_FEEDBACK_MIN_INTERVAL_SECS"; +const HEALTH_SUCCESS_PERSIST_MIN_INTERVAL_ENV: &str = + "AETHER_GATEWAY_PROVIDER_KEY_HEALTH_SUCCESS_PERSIST_MIN_INTERVAL_SECS"; const DEFAULT_POOL_SCORE_SUCCESS_FEEDBACK_MIN_INTERVAL_SECS: u64 = 5; const DEFAULT_POOL_SCORE_FAILURE_FEEDBACK_MIN_INTERVAL_SECS: u64 = 1; +const DEFAULT_HEALTH_SUCCESS_PERSIST_MIN_INTERVAL_SECS: u64 = 5; const MAX_POOL_SCORE_FEEDBACK_MIN_INTERVAL_SECS: u64 = 300; static POOL_SCORE_FEEDBACK_GATE: LazyLock> = LazyLock::new(ExpiringMap::new); +static HEALTH_SUCCESS_PERSIST_GATE: LazyLock> = + LazyLock::new(ExpiringMap::new); static POOL_SCORE_SUCCESS_FEEDBACK_MIN_INTERVAL: LazyLock = LazyLock::new(|| { pool_score_feedback_interval_from_env( POOL_SCORE_SUCCESS_FEEDBACK_MIN_INTERVAL_ENV, @@ -65,6 +71,12 @@ static POOL_SCORE_FAILURE_FEEDBACK_MIN_INTERVAL: LazyLock = LazyLock:: DEFAULT_POOL_SCORE_FAILURE_FEEDBACK_MIN_INTERVAL_SECS, ) }); +static HEALTH_SUCCESS_PERSIST_MIN_INTERVAL: LazyLock = LazyLock::new(|| { + pool_score_feedback_interval_from_env( + HEALTH_SUCCESS_PERSIST_MIN_INTERVAL_ENV, + DEFAULT_HEALTH_SUCCESS_PERSIST_MIN_INTERVAL_SECS, + ) +}); #[derive(Debug, Clone, Copy)] pub(crate) struct LocalExecutionEffectContext<'a> { @@ -645,6 +657,14 @@ async fn record_health_failure_effect( .or(current_key.circuit_breaker_by_format.as_ref()) }; + if !provider_key_health_success_persist_gate_allows( + &context.plan.key_id, + api_format, + circuit_breaker_update_owned.is_some(), + ) { + return; + } + if let Err(err) = state .update_provider_catalog_key_health_state( &context.plan.key_id, @@ -711,6 +731,8 @@ async fn record_health_success_effect( .or(current_key.circuit_breaker_by_format.as_ref()) }; + clear_provider_key_health_success_persist_gate(&context.plan.key_id, api_format); + if let Err(err) = state .update_provider_catalog_key_health_state( &context.plan.key_id, @@ -727,6 +749,36 @@ async fn record_health_success_effect( } } +fn provider_key_health_success_persist_gate_allows( + key_id: &str, + api_format: &str, + closes_circuit: bool, +) -> bool { + if closes_circuit { + return true; + } + let min_interval = *HEALTH_SUCCESS_PERSIST_MIN_INTERVAL; + if min_interval.is_zero() { + return true; + } + let key = provider_key_health_success_persist_gate_key(key_id, api_format); + HEALTH_SUCCESS_PERSIST_GATE.insert_if_absent_fresh( + key, + (), + min_interval, + HEALTH_SUCCESS_PERSIST_GATE_MAX_ENTRIES, + ) +} + +fn clear_provider_key_health_success_persist_gate(key_id: &str, api_format: &str) { + let key = provider_key_health_success_persist_gate_key(key_id, api_format); + HEALTH_SUCCESS_PERSIST_GATE.remove(&key); +} + +fn provider_key_health_success_persist_gate_key(key_id: &str, api_format: &str) -> String { + format!("success:{key_id}:{api_format}") +} + async fn record_stream_pool_success_effect( state: &AppState, context: LocalExecutionEffectContext<'_>, @@ -2612,6 +2664,88 @@ mod tests { ); } + #[tokio::test] + async fn health_success_projection_is_rate_limited_until_failure_resets_gate() { + let state = health_state(); + let plan = sample_plan(); + + apply_local_execution_effect( + &state, + LocalExecutionEffectContext { + plan: &plan, + report_context: None, + }, + LocalExecutionEffect::HealthSuccess(LocalHealthSuccessEffect), + ) + .await; + let first_updated_at = state + .read_provider_catalog_keys_by_ids(std::slice::from_ref(&plan.key_id)) + .await + .expect("provider catalog keys should load") + .into_iter() + .next() + .expect("stored key should exist") + .updated_at_unix_secs; + + apply_local_execution_effect( + &state, + LocalExecutionEffectContext { + plan: &plan, + report_context: None, + }, + LocalExecutionEffect::HealthSuccess(LocalHealthSuccessEffect), + ) + .await; + let second_updated_at = state + .read_provider_catalog_keys_by_ids(std::slice::from_ref(&plan.key_id)) + .await + .expect("provider catalog keys should load") + .into_iter() + .next() + .expect("stored key should exist") + .updated_at_unix_secs; + assert_eq!(second_updated_at, first_updated_at); + + apply_local_execution_effect( + &state, + LocalExecutionEffectContext { + plan: &plan, + report_context: None, + }, + LocalExecutionEffect::HealthFailure(LocalHealthFailureEffect { + status_code: 503, + classification: LocalFailoverClassification::RetryUpstreamFailure, + }), + ) + .await; + apply_local_execution_effect( + &state, + LocalExecutionEffectContext { + plan: &plan, + report_context: None, + }, + LocalExecutionEffect::HealthSuccess(LocalHealthSuccessEffect), + ) + .await; + + let stored_key = state + .read_provider_catalog_keys_by_ids(std::slice::from_ref(&plan.key_id)) + .await + .expect("provider catalog keys should load") + .into_iter() + .next() + .expect("stored key should exist"); + assert_eq!( + stored_key + .health_by_format + .as_ref() + .and_then(|value| value.get("openai:chat")) + .and_then(|value| value.get("consecutive_failures")) + .and_then(Value::as_u64), + Some(0) + ); + } + #[tokio::test] async fn health_success_projection_closes_key_circuit_for_format() { let mut key = sample_health_key(); diff --git a/apps/aether-gateway/src/privacy/mod.rs b/apps/aether-gateway/src/privacy/mod.rs index 5ea94ea91..42c2caad4 100644 --- a/apps/aether-gateway/src/privacy/mod.rs +++ b/apps/aether-gateway/src/privacy/mod.rs @@ -1,8 +1,9 @@ use std::collections::{BTreeMap, HashMap, HashSet}; use std::fmt; use std::net::{Ipv4Addr, Ipv6Addr}; +use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, LazyLock, Mutex}; -use std::time::Duration; +use std::time::{Duration, Instant}; use aether_data_contracts::DataLayerError; use aether_runtime_state::RuntimeState; @@ -26,6 +27,7 @@ const MAX_SENTINEL_NAMESPACE_LEN: usize = 32; const DIRECT_RESTORE_SENTINEL_LIMIT: usize = 32; const MAX_CACHE_SENTINEL_BYTES: usize = 128; const MAX_CACHE_RECORD_BYTES: usize = 512; +const CHAT_PII_REDACTION_RUNTIME_CONFIG_CACHE_TTL: Duration = Duration::from_secs(5); static EMAIL_REGEX: LazyLock = LazyLock::new(|| { Regex::new(r"(?i)[A-Z0-9._%+-]{1,64}@[A-Z0-9.-]{1,253}\.[A-Z]{2,63}") @@ -368,6 +370,7 @@ struct MappingKey { original: String, } +#[derive(Clone)] pub(crate) struct RedactionSession { config: RedactionSessionConfig, mappings: HashMap, @@ -717,10 +720,53 @@ impl fmt::Debug for MaskedChatRequest { } } +pub(crate) struct MaskedChatRequestValue { + pub(crate) body_json: Option, + pub(crate) session: RedactionSession, + pub(crate) redacted: bool, +} + +impl fmt::Debug for MaskedChatRequestValue { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("MaskedChatRequestValue") + .field("body_json_owned", &self.body_json.is_some()) + .field("session", &self.session) + .field("redacted", &self.redacted) + .finish() + } +} + +#[derive(Clone)] +pub(crate) struct CachedRequestRedaction { + pub(crate) body_json: Option, + pub(crate) session: Option, + pub(crate) redacted: bool, +} + +impl CachedRequestRedaction { + pub(crate) fn unredacted() -> Self { + Self { + body_json: None, + session: None, + redacted: false, + } + } + + pub(crate) fn redacted(body_json: Value, session: RedactionSession) -> Self { + Self { + body_json: Some(body_json), + session: Some(session), + redacted: true, + } + } +} + #[derive(Default, Clone)] pub(crate) struct RedactionSessionSlot { session: Arc>>, sessions_by_candidate: Arc>>, + request_redactions: Arc>>, } impl RedactionSessionSlot { @@ -747,6 +793,10 @@ impl RedactionSessionSlot { .lock() .expect("redaction session candidate slot should lock") .clear(); + self.request_redactions + .lock() + .expect("redaction request cache slot should lock") + .clear(); } pub(crate) fn take(&self) -> Option { @@ -794,6 +844,25 @@ impl RedactionSessionSlot { _ => None, } } + + pub(crate) fn cached_request_redaction(&self, key: &str) -> Option { + self.request_redactions + .lock() + .expect("redaction request cache slot should lock") + .get(key) + .cloned() + } + + pub(crate) fn put_cached_request_redaction( + &self, + key: impl Into, + redaction: CachedRequestRedaction, + ) { + self.request_redactions + .lock() + .expect("redaction request cache slot should lock") + .insert(key.into(), redaction); + } } #[derive(Clone, Debug, Eq, PartialEq)] @@ -828,6 +897,150 @@ impl Default for ChatPiiRedactionRuntimeConfig { } } +impl ChatPiiRedactionRuntimeConfig { + fn disabled() -> Self { + Self { + enabled: false, + rules: Vec::new(), + ttl_seconds: DEFAULT_REDACTION_TTL_SECONDS, + placeholder_prefix: DEFAULT_SENTINEL_NAMESPACE.to_string(), + } + } +} + +#[derive(Debug, Default)] +pub(crate) struct ChatPiiRedactionRuntimeConfigCache { + value: Mutex>, + loading_generation: Mutex>, + generation: AtomicU64, + notify: tokio::sync::Notify, +} + +enum ChatPiiRedactionRuntimeConfigLoadRegistration { + Leader(ChatPiiRedactionRuntimeConfigLoadGuard), + Follower, + Bypass, +} + +struct ChatPiiRedactionRuntimeConfigLoadGuard { + cache: ChatPiiRedactionRuntimeConfigCacheHandle, + generation: u64, + active: bool, +} + +impl ChatPiiRedactionRuntimeConfigLoadGuard { + fn generation(&self) -> u64 { + self.generation + } +} + +impl Drop for ChatPiiRedactionRuntimeConfigLoadGuard { + fn drop(&mut self) { + if self.active { + self.cache.finish_load(self.generation); + } + } +} + +impl ChatPiiRedactionRuntimeConfigCache { + fn get(&self) -> Option { + self.value.lock().ok().and_then(|guard| { + guard.as_ref().and_then(|(loaded_at, value)| { + (loaded_at.elapsed() <= CHAT_PII_REDACTION_RUNTIME_CONFIG_CACHE_TTL) + .then(|| value.clone()) + }) + }) + } + + fn get_stale(&self) -> Option { + self.value + .lock() + .ok() + .and_then(|guard| guard.as_ref().map(|(_, value)| value.clone())) + } + + fn insert(&self, value: ChatPiiRedactionRuntimeConfig) { + if let Ok(mut guard) = self.value.lock() { + *guard = Some((Instant::now(), value)); + } + } + + fn insert_if_generation(&self, generation: u64, value: ChatPiiRedactionRuntimeConfig) { + if self.generation.load(Ordering::Acquire) == generation { + self.insert(value); + } + } + + fn register_load(self: &Arc) -> ChatPiiRedactionRuntimeConfigLoadRegistration { + let generation = self.generation.load(Ordering::Acquire); + match self.loading_generation.lock() { + Ok(mut loading_generation) => { + if loading_generation.is_some() { + ChatPiiRedactionRuntimeConfigLoadRegistration::Follower + } else { + *loading_generation = Some(generation); + ChatPiiRedactionRuntimeConfigLoadRegistration::Leader( + ChatPiiRedactionRuntimeConfigLoadGuard { + cache: Arc::clone(self), + generation, + active: true, + }, + ) + } + } + Err(_) => ChatPiiRedactionRuntimeConfigLoadRegistration::Bypass, + } + } + + fn notified(&self) -> tokio::sync::futures::Notified<'_> { + self.notify.notified() + } + + fn finish_load(&self, generation: u64) { + let finished = self + .loading_generation + .lock() + .map(|mut loading_generation| { + if *loading_generation == Some(generation) { + *loading_generation = None; + true + } else { + false + } + }) + .unwrap_or(false); + if finished { + self.notify.notify_waiters(); + } + } +} + +pub(crate) type ChatPiiRedactionRuntimeConfigCacheHandle = Arc; + +pub(crate) fn new_chat_pii_redaction_runtime_config_cache( +) -> ChatPiiRedactionRuntimeConfigCacheHandle { + Arc::new(ChatPiiRedactionRuntimeConfigCache::default()) +} + +pub(crate) fn clear_chat_pii_redaction_runtime_config_cache( + cache: &ChatPiiRedactionRuntimeConfigCacheHandle, +) { + cache.clear(); +} + +impl ChatPiiRedactionRuntimeConfigCache { + fn clear(&self) { + self.generation.fetch_add(1, Ordering::AcqRel); + if let Ok(mut value) = self.value.lock() { + *value = None; + } + if let Ok(mut loading_generation) = self.loading_generation.lock() { + *loading_generation = None; + } + self.notify.notify_waiters(); + } +} + pub(crate) struct MaskChatRequestOptions { pub(crate) scan_limits: RedactionScanLimits, } @@ -1127,13 +1340,76 @@ fn parse_chat_pii_redaction_rules( pub(crate) async fn read_chat_pii_redaction_runtime_config( state: &crate::AppState, ) -> Result { - let mut config = ChatPiiRedactionRuntimeConfig::default(); - config.enabled = state + let cache = Arc::clone(&state.chat_pii_redaction_runtime_config_cache); + if let Some(value) = cache.get() { + return Ok(value); + } + if let Some(value) = cache.get_stale() { + if let ChatPiiRedactionRuntimeConfigLoadRegistration::Leader(guard) = cache.register_load() + { + spawn_chat_pii_redaction_runtime_config_refresh(state.clone(), cache, guard); + } + return Ok(value); + } + + loop { + let notified = cache.notified(); + match cache.register_load() { + ChatPiiRedactionRuntimeConfigLoadRegistration::Bypass => { + let value = load_chat_pii_redaction_runtime_config(state).await?; + cache.insert(value.clone()); + return Ok(value); + } + ChatPiiRedactionRuntimeConfigLoadRegistration::Follower => { + notified.await; + if let Some(value) = cache.get() { + return Ok(value); + } + } + ChatPiiRedactionRuntimeConfigLoadRegistration::Leader(_guard) => { + let generation = _guard.generation(); + let value = load_chat_pii_redaction_runtime_config(state).await?; + cache.insert_if_generation(generation, value.clone()); + return Ok(value); + } + } + } +} + +fn spawn_chat_pii_redaction_runtime_config_refresh( + state: crate::AppState, + cache: ChatPiiRedactionRuntimeConfigCacheHandle, + guard: ChatPiiRedactionRuntimeConfigLoadGuard, +) { + tokio::spawn(async move { + let generation = guard.generation(); + match load_chat_pii_redaction_runtime_config(&state).await { + Ok(value) => cache.insert_if_generation(generation, value), + Err(err) => { + tracing::warn!( + error = ?err, + "gateway failed to refresh chat pii redaction runtime config" + ); + } + } + drop(guard); + }); +} + +async fn load_chat_pii_redaction_runtime_config( + state: &crate::AppState, +) -> Result { + let enabled = state .read_system_config_json_value("module.chat_pii_redaction.enabled") .await? .as_ref() .and_then(Value::as_bool) - .unwrap_or(config.enabled); + .unwrap_or(false); + if !enabled { + return Ok(ChatPiiRedactionRuntimeConfig::disabled()); + } + let mut config = ChatPiiRedactionRuntimeConfig::default(); + config.enabled = true; config.rules = parse_chat_pii_redaction_rules( state .read_system_config_json_value("module.chat_pii_redaction.rules") @@ -1291,6 +1567,35 @@ pub(crate) async fn try_mask_chat_pii_request_json_with_cache_options( }) } +pub(crate) async fn try_mask_chat_pii_request_value_with_cache_options( + body_json: &Value, + format: ChatPiiRedactionRequestFormat, + config: RedactionSessionConfig, + options: MaskChatRequestOptions, + cache: Option<&RedisRedactionMappingCache<'_>>, +) -> Result { + let mut session = RedactionSession::new(config); + let mut value = body_json.clone(); + + session.set_collision_corpus(request_collision_corpus(format, &value)); + let mut scan_state = RedactionScanState::new(options.scan_limits); + let redacted = mask_request_value_async( + format, + &mut value, + &mut session, + &mut scan_state, + options, + cache, + ) + .await?; + + Ok(MaskedChatRequestValue { + body_json: redacted.then_some(value), + session, + redacted, + }) +} + fn request_collision_corpus(format: ChatPiiRedactionRequestFormat, value: &Value) -> Vec { match format { ChatPiiRedactionRequestFormat::OpenAiChat => value @@ -2752,6 +3057,7 @@ impl fmt::Debug for RedactionMatch { } } +#[derive(Clone)] pub(crate) struct RedactionMapping { pub(crate) rule_label: String, pub(crate) kind: Option, diff --git a/apps/aether-gateway/src/provider_pool_demand.rs b/apps/aether-gateway/src/provider_pool_demand.rs index 535733008..00346cb03 100644 --- a/apps/aether-gateway/src/provider_pool_demand.rs +++ b/apps/aether-gateway/src/provider_pool_demand.rs @@ -1,10 +1,11 @@ use std::sync::{ - atomic::{AtomicBool, Ordering}, + atomic::{AtomicBool, AtomicUsize, Ordering}, Arc, }; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use aether_runtime_state::RuntimeState; +use dashmap::DashMap; use serde::{Deserialize, Serialize}; use tokio::task::JoinHandle; use tracing::debug; @@ -32,14 +33,24 @@ pub(crate) struct ProviderPoolDemandSnapshot { } pub(crate) struct ProviderPoolInFlightGuard { - runtime: Arc, - tokens_key: String, - token: String, - stop_renewal: Arc, - renew_handle: Option>, + kind: ProviderPoolInFlightGuardKind, released: bool, } +enum ProviderPoolInFlightGuardKind { + Local { + provider_id: String, + counter: Arc, + }, + Runtime { + runtime: Arc, + tokens_key: String, + token: String, + stop_renewal: Arc, + renew_handle: Option>, + }, +} + impl ProviderPoolInFlightGuard { pub(crate) async fn release(mut self) { self.release_inner().await; @@ -50,19 +61,29 @@ impl ProviderPoolInFlightGuard { return; } self.released = true; - self.stop_renewal.store(true, Ordering::Release); - if let Some(handle) = self.renew_handle.take() { - handle.abort(); - } - if let Err(err) = self - .runtime - .score_remove(&self.tokens_key, &self.token) - .await - { - debug!( - error = ?err, - "gateway provider pool demand: failed to release in-flight token" - ); + match &mut self.kind { + ProviderPoolInFlightGuardKind::Local { + provider_id, + counter, + } => decrement_local_provider_in_flight(provider_id, counter), + ProviderPoolInFlightGuardKind::Runtime { + runtime, + tokens_key, + token, + stop_renewal, + renew_handle, + } => { + stop_renewal.store(true, Ordering::Release); + if let Some(handle) = renew_handle.take() { + handle.abort(); + } + if let Err(err) = runtime.score_remove(tokens_key, token).await { + debug!( + error = ?err, + "gateway provider pool demand: failed to release in-flight token" + ); + } + } } } } @@ -73,23 +94,37 @@ impl Drop for ProviderPoolInFlightGuard { return; } self.released = true; - self.stop_renewal.store(true, Ordering::Release); - if let Some(handle) = self.renew_handle.take() { - handle.abort(); - } - - let runtime = self.runtime.clone(); - let tokens_key = self.tokens_key.clone(); - let token = self.token.clone(); - if let Ok(handle) = tokio::runtime::Handle::try_current() { - handle.spawn(async move { - if let Err(err) = runtime.score_remove(&tokens_key, &token).await { - debug!( - error = ?err, - "gateway provider pool demand: failed to release dropped in-flight token" - ); + match &mut self.kind { + ProviderPoolInFlightGuardKind::Local { + provider_id, + counter, + } => decrement_local_provider_in_flight(provider_id, counter), + ProviderPoolInFlightGuardKind::Runtime { + runtime, + tokens_key, + token, + stop_renewal, + renew_handle, + } => { + stop_renewal.store(true, Ordering::Release); + if let Some(handle) = renew_handle.take() { + handle.abort(); } - }); + + let runtime = runtime.clone(); + let tokens_key = tokens_key.clone(); + let token = token.clone(); + if let Ok(handle) = tokio::runtime::Handle::try_current() { + handle.spawn(async move { + if let Err(err) = runtime.score_remove(&tokens_key, &token).await { + debug!( + error = ?err, + "gateway provider pool demand: failed to release dropped in-flight token" + ); + } + }); + } + } } } } @@ -106,6 +141,46 @@ fn in_flight_tokens_key(provider_id: &str) -> String { format!("{PROVIDER_POOL_IN_FLIGHT_TOKENS_PREFIX}:{provider_id}") } +fn local_provider_in_flight_counts() -> &'static DashMap> { + static COUNTS: std::sync::OnceLock>> = + std::sync::OnceLock::new(); + COUNTS.get_or_init(DashMap::new) +} + +fn increment_local_provider_in_flight(provider_id: &str) -> Arc { + let counter_ref = local_provider_in_flight_counts() + .entry(provider_id.to_string()) + .or_insert_with(|| Arc::new(AtomicUsize::new(0))); + counter_ref.fetch_add(1, Ordering::AcqRel); + counter_ref.clone() +} + +fn decrement_local_provider_in_flight(provider_id: &str, counter: &AtomicUsize) { + let mut current = counter.load(Ordering::Acquire); + while current > 0 { + match counter.compare_exchange_weak( + current, + current - 1, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => break, + Err(next) => current = next, + } + } + if counter.load(Ordering::Acquire) == 0 { + let _ = local_provider_in_flight_counts() + .remove_if(provider_id, |_, stored| stored.load(Ordering::Acquire) == 0); + } +} + +fn local_provider_live_in_flight_count(provider_id: &str) -> usize { + local_provider_in_flight_counts() + .get(provider_id) + .map(|counter| counter.load(Ordering::Acquire)) + .unwrap_or(0) +} + fn demand_snapshot_key(provider_id: &str) -> String { format!("{PROVIDER_POOL_DEMAND_SNAPSHOT_PREFIX}:{provider_id}") } @@ -185,6 +260,17 @@ pub(crate) async fn acquire_provider_pool_in_flight_guard( return None; } + if runtime.is_memory() { + let counter = increment_local_provider_in_flight(provider_id); + return Some(ProviderPoolInFlightGuard { + kind: ProviderPoolInFlightGuardKind::Local { + provider_id: provider_id.to_string(), + counter, + }, + released: false, + }); + } + let tokens_key = in_flight_tokens_key(provider_id); let token = build_in_flight_token(request_id, candidate_id, key_id); match tokio::time::timeout( @@ -221,11 +307,13 @@ pub(crate) async fn acquire_provider_pool_in_flight_guard( ); Some(ProviderPoolInFlightGuard { - runtime, - tokens_key, - token, - stop_renewal, - renew_handle: Some(renew_handle), + kind: ProviderPoolInFlightGuardKind::Runtime { + runtime, + tokens_key, + token, + stop_renewal, + renew_handle: Some(renew_handle), + }, released: false, }) } @@ -238,6 +326,9 @@ pub(crate) async fn provider_pool_live_in_flight_count( if provider_id.is_empty() { return 0; } + if runtime.is_memory() { + return local_provider_live_in_flight_count(provider_id); + } let key = in_flight_tokens_key(provider_id); let now_ms = current_unix_ms() as f64; if let Err(err) = runtime.score_remove_by_score(&key, now_ms).await { @@ -389,9 +480,10 @@ mod tests { #[tokio::test] async fn in_flight_guard_tracks_and_releases_provider_tokens() { let runtime = Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default())); + let provider_id = "provider-guard-release"; let guard = acquire_provider_pool_in_flight_guard( runtime.clone(), - "provider-1", + provider_id, "request-1", Some("candidate-1"), "key-1", @@ -400,14 +492,14 @@ mod tests { .expect("guard should be acquired"); assert_eq!( - provider_pool_live_in_flight_count(runtime.as_ref(), "provider-1").await, + provider_pool_live_in_flight_count(runtime.as_ref(), provider_id).await, 1 ); guard.release().await; assert_eq!( - provider_pool_live_in_flight_count(runtime.as_ref(), "provider-1").await, + provider_pool_live_in_flight_count(runtime.as_ref(), provider_id).await, 0 ); } @@ -415,26 +507,27 @@ mod tests { #[tokio::test] async fn demand_snapshot_uses_instant_in_flight_for_fast_rise_and_ema_for_fall() { let runtime = RuntimeState::memory(MemoryRuntimeStateConfig::default()); + let provider_id = "provider-demand-snapshot"; + let mut guards = Vec::new(); for idx in 0..10 { let guard = acquire_provider_pool_in_flight_guard( Arc::new(runtime.clone()), - "provider-1", + provider_id, "request-1", Some(&format!("candidate-{idx}")), "key-1", ) .await .expect("guard"); - std::mem::forget(guard); + guards.push(guard); } - let high = sample_provider_pool_demand(&runtime, "provider-1", 100, 50).await; + let high = sample_provider_pool_demand(&runtime, provider_id, 100, 50).await; assert_eq!(high.in_flight, 10); assert_eq!(high.desired_hot, 12); - let key = in_flight_tokens_key("provider-1"); - let _ = runtime.score_remove_by_score(&key, f64::INFINITY).await; - let low = sample_provider_pool_demand(&runtime, "provider-1", 100, 50).await; + drop(guards); + let low = sample_provider_pool_demand(&runtime, provider_id, 100, 50).await; assert_eq!(low.in_flight, 0); assert!(low.ema_in_flight > 0.0); assert!(low.desired_hot >= PROVIDER_POOL_DEMAND_FLOOR); diff --git a/apps/aether-gateway/src/request_candidate_queue.rs b/apps/aether-gateway/src/request_candidate_queue.rs index 8fdbcdd02..c02aeaae8 100644 --- a/apps/aether-gateway/src/request_candidate_queue.rs +++ b/apps/aether-gateway/src/request_candidate_queue.rs @@ -22,6 +22,7 @@ const DEFAULT_QUEUE_CAPACITY: usize = 65_536; const DEFAULT_BATCH_SIZE: usize = 512; const DEFAULT_FLUSH_INTERVAL_MS: u64 = 50; const DEFAULT_WORKERS: usize = 2; +const FAILED_FLUSH_RETRY_DELAY_MS: u64 = 25; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) enum RequestCandidateWriteMode { @@ -53,7 +54,7 @@ impl Default for RequestCandidateQueueConfig { batch_size: DEFAULT_BATCH_SIZE, flush_interval: Duration::from_millis(DEFAULT_FLUSH_INTERVAL_MS), workers: DEFAULT_WORKERS, - full_policy: RequestCandidateQueueFullPolicy::Drop, + full_policy: RequestCandidateQueueFullPolicy::Sync, } } } @@ -62,8 +63,9 @@ impl RequestCandidateQueueConfig { pub(crate) fn from_env() -> Self { let mut config = Self::default(); config.mode = match env_string(MODE_ENV).as_deref() { + Some("sync") | Some("inline") => RequestCandidateWriteMode::Sync, Some("async") | Some("queued") | Some("queue") => RequestCandidateWriteMode::Async, - _ => RequestCandidateWriteMode::Sync, + _ => RequestCandidateWriteMode::Async, }; config.capacity = env_usize(QUEUE_CAPACITY_ENV, DEFAULT_QUEUE_CAPACITY).max(1); config.batch_size = env_usize(BATCH_SIZE_ENV, DEFAULT_BATCH_SIZE).max(1); @@ -71,10 +73,13 @@ impl RequestCandidateQueueConfig { Duration::from_millis(env_u64(FLUSH_INTERVAL_MS_ENV, DEFAULT_FLUSH_INTERVAL_MS).max(1)); config.workers = env_usize(WORKERS_ENV, DEFAULT_WORKERS).clamp(1, 32); config.full_policy = match env_string(QUEUE_FULL_ENV).as_deref() { + Some("drop") | Some("best_effort") | Some("best-effort") => { + RequestCandidateQueueFullPolicy::Drop + } Some("sync") | Some("fallback_sync") | Some("fallback-sync") => { RequestCandidateQueueFullPolicy::Sync } - _ => RequestCandidateQueueFullPolicy::Drop, + _ => RequestCandidateQueueFullPolicy::Sync, }; config } @@ -94,6 +99,7 @@ struct RequestCandidateQueueMetrics { flush_failed_total: AtomicU64, flush_batches_total: AtomicU64, flush_sql_ops_total: AtomicU64, + flush_sql_records_total: AtomicU64, compacted_total: AtomicU64, sync_fallback_total: AtomicU64, } @@ -248,13 +254,19 @@ impl RequestCandidateQueueRuntime { ), MetricSample::new( "request_candidate_queue_flush_sql_ops_total", - "Total repository upsert operations issued by async request candidate persistence workers after compaction.", + "Total repository batch upsert operations issued by async request candidate persistence workers after compaction.", MetricKind::Counter, self.metrics.flush_sql_ops_total.load(Ordering::Acquire), ), + MetricSample::new( + "request_candidate_queue_flush_sql_records_total", + "Total request candidate records submitted to repository batch upsert operations after compaction.", + MetricKind::Counter, + self.metrics.flush_sql_records_total.load(Ordering::Acquire), + ), MetricSample::new( "request_candidate_queue_compacted_total", - "Total request candidate records compacted before async persistence because a later queued record covered the same slot and status.", + "Total request candidate records compacted before async persistence because a later queued record covered the same request candidate slot.", MetricKind::Counter, self.metrics.compacted_total.load(Ordering::Acquire), ), @@ -340,7 +352,7 @@ async fn flush_batch( return; } let source_count = records.len(); - let records = compact_same_status_records(records); + let records = compact_records_for_flush(records); let compacted = source_count.saturating_sub(records.len()); if compacted > 0 { metrics @@ -348,30 +360,51 @@ async fn flush_batch( .fetch_add(compacted as u64, Ordering::AcqRel); } metrics.flush_batches_total.fetch_add(1, Ordering::AcqRel); + let source_count = records + .iter() + .map(|record| record.source_count) + .sum::(); + let record_count = records.len(); + let upsert_records = records + .into_iter() + .map(|record| record.record) + .collect::>(); + metrics.flush_sql_ops_total.fetch_add(1, Ordering::AcqRel); + metrics + .flush_sql_records_total + .fetch_add(record_count as u64, Ordering::AcqRel); let mut failed = 0_u64; - for record in records { - metrics.flush_sql_ops_total.fetch_add(1, Ordering::AcqRel); - if let Err(err) = repository.upsert(record.record).await { - failed = failed.saturating_add(record.source_count as u64); - decrement_atomic_usize_by(&metrics.pending_current, record.source_count); - warn!( - event_name = "request_candidate_async_flush_failed", - log_type = "event", - worker_index, - error = ?err, - "gateway failed to asynchronously persist request candidate" - ); - } else { - metrics - .flushed_total - .fetch_add(record.source_count as u64, Ordering::AcqRel); - decrement_atomic_usize_by(&metrics.pending_current, record.source_count); - } + let mut retry_records = Vec::new(); + if let Err(err) = repository.upsert_many(upsert_records.clone()).await { + failed = source_count as u64; + decrement_atomic_usize_by( + &metrics.pending_current, + source_count.saturating_sub(record_count), + ); + warn!( + event_name = "request_candidate_async_flush_failed", + log_type = "event", + worker_index, + record_count, + source_count, + error = ?err, + "gateway failed to asynchronously persist request candidate batch" + ); + retry_records = upsert_records; + } else { + metrics + .flushed_total + .fetch_add(source_count as u64, Ordering::AcqRel); + decrement_atomic_usize_by(&metrics.pending_current, source_count); } if failed > 0 { metrics .flush_failed_total .fetch_add(failed, Ordering::AcqRel); + tokio::time::sleep(Duration::from_millis(FAILED_FLUSH_RETRY_DELAY_MS)).await; + for record in retry_records { + batch.push(record); + } } debug!( event_name = "request_candidate_async_flush_completed", @@ -388,10 +421,10 @@ struct CompactedRequestCandidateRecord { source_count: usize, } -fn compact_same_status_records( +fn compact_records_for_flush( records: Vec, ) -> Vec { - let mut latest_slot_status = HashMap::<(String, u32, u32), (u8, usize)>::new(); + let mut latest_slot = HashMap::<(String, u32, u32), usize>::new(); let mut compacted = Vec::::with_capacity(records.len()); for record in records { let slot = ( @@ -399,14 +432,13 @@ fn compact_same_status_records( record.candidate_index, record.retry_index, ); - let status = request_candidate_status_discriminant(record.status); - match latest_slot_status.get(&slot).copied() { - Some((latest_status, index)) if latest_status == status => { - merge_request_candidate_record(&mut compacted[index].record, record); + match latest_slot.get(&slot).copied() { + Some(index) => { + merge_request_candidate_record_for_flush(&mut compacted[index].record, record); compacted[index].source_count = compacted[index].source_count.saturating_add(1); } _ => { - latest_slot_status.insert(slot, (status, compacted.len())); + latest_slot.insert(slot, compacted.len()); compacted.push(CompactedRequestCandidateRecord { record, source_count: 1, @@ -489,6 +521,41 @@ fn merge_request_candidate_record( ); } +fn merge_request_candidate_record_for_flush( + target: &mut UpsertRequestCandidateRecord, + incoming: UpsertRequestCandidateRecord, +) { + let target_status = target.status; + let incoming_status = incoming.status; + let next_status = merged_request_candidate_status(target_status, incoming_status); + merge_request_candidate_record(target, incoming); + target.status = next_status; +} + +fn merged_request_candidate_status( + current: RequestCandidateStatus, + incoming: RequestCandidateStatus, +) -> RequestCandidateStatus { + match ( + request_candidate_status_is_terminal(current), + request_candidate_status_is_terminal(incoming), + ) { + (_, true) => incoming, + (true, false) => current, + (false, false) => incoming, + } +} + +fn request_candidate_status_is_terminal(status: RequestCandidateStatus) -> bool { + matches!( + status, + RequestCandidateStatus::Success + | RequestCandidateStatus::Failed + | RequestCandidateStatus::Cancelled + ) +} + +#[cfg(test)] fn request_candidate_status_discriminant(status: RequestCandidateStatus) -> u8 { match status { RequestCandidateStatus::Available => 0, @@ -554,7 +621,7 @@ fn decrement_atomic_usize_by(value: &AtomicUsize, amount: usize) { #[cfg(test)] mod tests { use super::{ - compact_same_status_records, RequestCandidateQueueConfig, RequestCandidateQueueRuntime, + compact_records_for_flush, RequestCandidateQueueConfig, RequestCandidateQueueRuntime, }; use aether_data::repository::candidates::InMemoryRequestCandidateRepository; use aether_data::DataLayerError; @@ -562,7 +629,7 @@ mod tests { RequestCandidateReadRepository, RequestCandidateStatus, RequestCandidateWriteRepository, StoredRequestCandidate, UpsertRequestCandidateRecord, }; - use std::sync::atomic::Ordering; + use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; use std::time::Duration; @@ -594,6 +661,46 @@ mod tests { } } + #[derive(Default)] + struct CountingBatchRequestCandidateRepository { + inner: InMemoryRequestCandidateRepository, + upsert_calls: AtomicUsize, + upsert_many_calls: AtomicUsize, + } + + #[async_trait::async_trait] + impl RequestCandidateWriteRepository for CountingBatchRequestCandidateRepository { + async fn upsert( + &self, + candidate: UpsertRequestCandidateRecord, + ) -> Result { + self.upsert_calls.fetch_add(1, Ordering::AcqRel); + self.inner.upsert(candidate).await + } + + async fn upsert_many( + &self, + candidates: Vec, + ) -> Result { + self.upsert_many_calls.fetch_add(1, Ordering::AcqRel); + let count = candidates.len(); + for candidate in candidates { + self.inner.upsert(candidate).await?; + } + Ok(count) + } + + async fn delete_created_before( + &self, + created_before_unix_secs: u64, + limit: usize, + ) -> Result { + self.inner + .delete_created_before(created_before_unix_secs, limit) + .await + } + } + fn record( request_id: &str, candidate_index: u32, @@ -628,8 +735,45 @@ mod tests { } } + struct EnvGuard { + key: &'static str, + previous: Option, + } + + impl EnvGuard { + fn unset(key: &'static str) -> Self { + let previous = std::env::var(key).ok(); + std::env::remove_var(key); + Self { key, previous } + } + } + + impl Drop for EnvGuard { + fn drop(&mut self) { + if let Some(previous) = self.previous.as_ref() { + std::env::set_var(self.key, previous); + } else { + std::env::remove_var(self.key); + } + } + } + #[test] - fn compact_merges_same_slot_and_status_without_dropping_state_transitions() { + fn from_env_defaults_to_async_with_sync_full_fallback() { + let _mode = EnvGuard::unset(super::MODE_ENV); + let _full = EnvGuard::unset(super::QUEUE_FULL_ENV); + + let config = RequestCandidateQueueConfig::from_env(); + + assert_eq!(config.mode, super::RequestCandidateWriteMode::Async); + assert_eq!( + config.full_policy, + super::RequestCandidateQueueFullPolicy::Sync + ); + } + + #[test] + fn compact_merges_same_slot_without_losing_terminal_fields() { let mut first_success = record("req", 0, 0, RequestCandidateStatus::Success); first_success.provider_id = Some("provider-a".to_string()); first_success.extra_data = Some(serde_json::json!({"first": true})); @@ -637,43 +781,40 @@ mod tests { second_success.latency_ms = Some(123); second_success.extra_data = Some(serde_json::json!({"second": true})); - let compacted = compact_same_status_records(vec![ + let compacted = compact_records_for_flush(vec![ record("req", 0, 0, RequestCandidateStatus::Pending), first_success, record("req", 0, 1, RequestCandidateStatus::Failed), second_success, ]); - assert_eq!(compacted.len(), 3); - assert_eq!(compacted[0].record.status, RequestCandidateStatus::Pending); - assert_eq!(compacted[0].source_count, 1); - assert_eq!(compacted[1].record.status, RequestCandidateStatus::Success); - assert_eq!(compacted[1].source_count, 2); + assert_eq!(compacted.len(), 2); + assert_eq!(compacted[0].record.status, RequestCandidateStatus::Success); + assert_eq!(compacted[0].source_count, 3); assert_eq!( - compacted[1].record.provider_id.as_deref(), + compacted[0].record.provider_id.as_deref(), Some("provider-a") ); - assert_eq!(compacted[1].record.latency_ms, Some(123)); + assert_eq!(compacted[0].record.latency_ms, Some(123)); assert_eq!( - compacted[1].record.extra_data, + compacted[0].record.extra_data, Some(serde_json::json!({"first": true, "second": true})) ); - assert_eq!(compacted[2].record.status, RequestCandidateStatus::Failed); - assert_eq!(compacted[2].source_count, 1); + assert_eq!(compacted[1].record.status, RequestCandidateStatus::Failed); + assert_eq!(compacted[1].source_count, 1); } #[test] - fn compact_preserves_same_slot_status_order_across_transitions() { - let compacted = compact_same_status_records(vec![ - record("req", 0, 0, RequestCandidateStatus::Success), - record("req", 0, 0, RequestCandidateStatus::Failed), + fn compact_keeps_terminal_status_when_later_intermediate_status_arrives() { + let compacted = compact_records_for_flush(vec![ record("req", 0, 0, RequestCandidateStatus::Success), + record("req", 0, 0, RequestCandidateStatus::Streaming), + record("req", 0, 0, RequestCandidateStatus::Unused), ]); - assert_eq!(compacted.len(), 3); + assert_eq!(compacted.len(), 1); assert_eq!(compacted[0].record.status, RequestCandidateStatus::Success); - assert_eq!(compacted[1].record.status, RequestCandidateStatus::Failed); - assert_eq!(compacted[2].record.status, RequestCandidateStatus::Success); + assert_eq!(compacted[0].source_count, 3); } #[tokio::test] @@ -708,6 +849,56 @@ mod tests { panic!("async request candidate queue did not flush record in time"); } + #[tokio::test] + async fn async_queue_flushes_records_with_batch_repository_call() { + let repository = Arc::new(CountingBatchRequestCandidateRepository::default()); + let runtime = RequestCandidateQueueRuntime::spawn( + repository.clone(), + RequestCandidateQueueConfig { + mode: super::RequestCandidateWriteMode::Async, + capacity: 16, + batch_size: 4, + flush_interval: Duration::from_millis(100), + workers: 1, + full_policy: super::RequestCandidateQueueFullPolicy::Drop, + }, + ); + + for index in 0..4 { + runtime + .enqueue_or_fallback(record( + "req-batch", + index, + 0, + RequestCandidateStatus::Success, + )) + .await + .unwrap(); + } + + for _ in 0..50 { + if runtime.metrics.pending_current.load(Ordering::Acquire) == 0 { + assert_eq!(repository.upsert_many_calls.load(Ordering::Acquire), 1); + assert_eq!(repository.upsert_calls.load(Ordering::Acquire), 0); + assert_eq!( + runtime.metrics.flush_sql_ops_total.load(Ordering::Acquire), + 1 + ); + assert_eq!( + runtime + .metrics + .flush_sql_records_total + .load(Ordering::Acquire), + 4 + ); + return; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + + panic!("async request candidate queue did not finish batch flush in time"); + } + #[tokio::test] async fn async_queue_preserves_same_slot_order_with_multiple_workers() { let repository = Arc::new(DelayedPendingRequestCandidateRepository::default()); diff --git a/apps/aether-gateway/src/request_diagnostics.rs b/apps/aether-gateway/src/request_diagnostics.rs new file mode 100644 index 000000000..49c33d334 --- /dev/null +++ b/apps/aether-gateway/src/request_diagnostics.rs @@ -0,0 +1,200 @@ +use std::collections::BTreeMap; +use std::future::Future; +use std::sync::{Arc, Mutex}; +use std::time::Instant; + +use aether_data::DatabasePoolSummary; +use serde_json::{Map, Value}; + +tokio::task_local! { + static REQUEST_DIAGNOSTICS: Arc; +} + +#[derive(Debug, Default)] +pub(crate) struct RequestDiagnostics { + inner: Mutex, +} + +#[derive(Debug, Default)] +struct RequestDiagnosticsInner { + db_operations: BTreeMap<&'static str, DbOperationTiming>, + db_pool: Option, +} + +#[derive(Debug, Clone, Copy, Default)] +struct DbOperationTiming { + count: u64, + sum_ms: u64, + max_ms: u64, +} + +#[derive(Debug, Clone, Copy)] +struct DbPoolObservation { + max_checked_out: u64, + max_pool_size: u64, + min_idle: u64, + max_connections: u64, + max_usage_rate_x100: u64, +} + +impl RequestDiagnostics { + fn record_db_timing_ms(&self, operation: &'static str, elapsed_ms: u64) { + let Ok(mut inner) = self.inner.lock() else { + return; + }; + let timing = inner.db_operations.entry(operation).or_default(); + timing.count = timing.count.saturating_add(1); + timing.sum_ms = timing.sum_ms.saturating_add(elapsed_ms); + timing.max_ms = timing.max_ms.max(elapsed_ms); + } + + fn record_db_pool_summary(&self, summary: DatabasePoolSummary) { + let observation = DbPoolObservation { + max_checked_out: summary.checked_out as u64, + max_pool_size: summary.pool_size as u64, + min_idle: summary.idle as u64, + max_connections: u64::from(summary.max_connections), + max_usage_rate_x100: (summary.usage_rate * 100.0).max(0.0).round() as u64, + }; + let Ok(mut inner) = self.inner.lock() else { + return; + }; + inner.db_pool = Some(match inner.db_pool { + Some(existing) => DbPoolObservation { + max_checked_out: existing.max_checked_out.max(observation.max_checked_out), + max_pool_size: existing.max_pool_size.max(observation.max_pool_size), + min_idle: existing.min_idle.min(observation.min_idle), + max_connections: existing.max_connections.max(observation.max_connections), + max_usage_rate_x100: existing + .max_usage_rate_x100 + .max(observation.max_usage_rate_x100), + }, + None => observation, + }); + } + + pub(crate) fn db_timings_metadata(&self) -> Option { + let Ok(inner) = self.inner.lock() else { + return None; + }; + if inner.db_operations.is_empty() && inner.db_pool.is_none() { + return None; + } + + let mut total_count = 0_u64; + let mut query_total_ms = 0_u64; + let mut query_max_ms = 0_u64; + let mut operations = Map::new(); + for (operation, timing) in &inner.db_operations { + total_count = total_count.saturating_add(timing.count); + query_total_ms = query_total_ms.saturating_add(timing.sum_ms); + query_max_ms = query_max_ms.max(timing.max_ms); + operations.insert( + (*operation).to_string(), + Value::Object(Map::from_iter([ + ("count".to_string(), Value::from(timing.count)), + ("sum".to_string(), Value::from(timing.sum_ms)), + ("max".to_string(), Value::from(timing.max_ms)), + ])), + ); + } + + let mut metadata = Map::new(); + if !operations.is_empty() { + metadata.insert("query_count".to_string(), Value::from(total_count)); + metadata.insert("query_total".to_string(), Value::from(query_total_ms)); + metadata.insert("query_max".to_string(), Value::from(query_max_ms)); + metadata.insert("operations".to_string(), Value::Object(operations)); + } + if let Some(pool) = inner.db_pool { + metadata.insert( + "pool".to_string(), + Value::Object(Map::from_iter([ + ( + "max_checked_out".to_string(), + Value::from(pool.max_checked_out), + ), + ("max_pool_size".to_string(), Value::from(pool.max_pool_size)), + ("min_idle".to_string(), Value::from(pool.min_idle)), + ( + "max_connections".to_string(), + Value::from(pool.max_connections), + ), + ( + "max_usage_rate".to_string(), + Value::from(pool.max_usage_rate_x100 as f64 / 100.0), + ), + ])), + ); + } + + Some(Value::Object(metadata)) + } +} + +pub(crate) async fn scope_request_diagnostics(future: F) -> F::Output +where + F: Future, +{ + REQUEST_DIAGNOSTICS + .scope(Arc::new(RequestDiagnostics::default()), future) + .await +} + +pub(crate) fn current_request_diagnostics() -> Option> { + REQUEST_DIAGNOSTICS.try_with(Arc::clone).ok() +} + +pub(crate) async fn observe_db_operation( + operation: &'static str, + pool_summary: Option, + future: F, +) -> F::Output +where + F: Future, +{ + if let Some(summary) = pool_summary { + record_db_pool_summary(summary); + } + let started_at = Instant::now(); + let output = future.await; + record_db_timing_ms(operation, started_at.elapsed().as_millis() as u64); + output +} + +pub(crate) fn record_db_timing_ms(operation: &'static str, elapsed_ms: u64) { + if let Some(diagnostics) = current_request_diagnostics() { + diagnostics.record_db_timing_ms(operation, elapsed_ms); + } +} + +pub(crate) fn record_db_pool_summary(summary: DatabasePoolSummary) { + if let Some(diagnostics) = current_request_diagnostics() { + diagnostics.record_db_pool_summary(summary); + } +} + +pub(crate) fn attach_request_diagnostics_to_report_context( + report_context: Option, + diagnostics: Option<&Arc>, +) -> Option { + let Some(db_timings_ms) = diagnostics.and_then(|diagnostics| diagnostics.db_timings_metadata()) + else { + return report_context; + }; + + let mut object = match report_context { + Some(Value::Object(object)) => object, + Some(other) => Map::from_iter([("seed".to_string(), other)]), + None => Map::new(), + }; + object.insert("db_timings_ms".to_string(), db_timings_ms); + Some(Value::Object(object)) +} + +pub(crate) fn attach_current_request_diagnostics_to_report_context( + report_context: Option<&Value>, +) -> Option { + let diagnostics = current_request_diagnostics()?; + attach_request_diagnostics_to_report_context(report_context.cloned(), Some(&diagnostics)) +} diff --git a/apps/aether-gateway/src/scheduler/candidate/affinity.rs b/apps/aether-gateway/src/scheduler/candidate/affinity.rs index 635a28f75..b6aee19d1 100644 --- a/apps/aether-gateway/src/scheduler/candidate/affinity.rs +++ b/apps/aether-gateway/src/scheduler/candidate/affinity.rs @@ -17,6 +17,9 @@ pub(super) fn build_scheduler_affinity_cache_key( global_model_name: &str, client_session_affinity: Option<&ClientSessionAffinity>, ) -> Option { + if !has_explicit_session_affinity(client_session_affinity) { + return None; + } let api_key_id = auth_snapshot .map(|snapshot| snapshot.api_key_id.trim()) .filter(|value| !value.is_empty())?; @@ -28,6 +31,12 @@ pub(super) fn build_scheduler_affinity_cache_key( ) } +pub(super) fn has_explicit_session_affinity( + client_session_affinity: Option<&ClientSessionAffinity>, +) -> bool { + client_session_affinity.is_some_and(ClientSessionAffinity::has_session_key) +} + pub(super) fn scheduler_candidate_affinity_hash( affinity_key: &str, candidate: &SchedulerMinimalCandidateSelectionCandidate, diff --git a/apps/aether-gateway/src/scheduler/candidate/mod.rs b/apps/aether-gateway/src/scheduler/candidate/mod.rs index b8311de0e..f793cef21 100644 --- a/apps/aether-gateway/src/scheduler/candidate/mod.rs +++ b/apps/aether-gateway/src/scheduler/candidate/mod.rs @@ -138,8 +138,11 @@ pub(crate) async fn list_selectable_enumerated_candidates_with_skip_reasons( GatewayError, > { let ordering_config = runtime_state.read_scheduler_ordering_config().await?; - let priority_affinity_key = - selection::scheduling_priority_affinity_key(auth_snapshot, ordering_config.scheduling_mode); + let priority_affinity_key = selection::scheduling_priority_affinity_key( + auth_snapshot, + client_session_affinity, + ordering_config.scheduling_mode, + ); collect_selectable_enumerated_candidates_with_skip_reasons( runtime_state, api_format, diff --git a/apps/aether-gateway/src/scheduler/candidate/selection.rs b/apps/aether-gateway/src/scheduler/candidate/selection.rs index cfcfb7729..458d00ac1 100644 --- a/apps/aether-gateway/src/scheduler/candidate/selection.rs +++ b/apps/aether-gateway/src/scheduler/candidate/selection.rs @@ -5,7 +5,9 @@ use crate::scheduler::config::SchedulerSchedulingMode; use crate::GatewayError; use aether_scheduler_core::ClientSessionAffinity; -use super::affinity::{build_scheduler_affinity_cache_key, remember_scheduler_affinity}; +use super::affinity::{ + build_scheduler_affinity_cache_key, has_explicit_session_affinity, remember_scheduler_affinity, +}; use super::enumeration::enumerate_scheduler_candidates; use super::ranking::rank_scheduler_candidates; use super::resolution::resolve_scheduler_candidate_selectability; @@ -66,8 +68,11 @@ pub(super) async fn select_minimal_candidate( global_model_name, client_session_affinity, ); - let priority_affinity_key = - scheduling_priority_affinity_key(auth_snapshot, ordering_config.scheduling_mode); + let priority_affinity_key = scheduling_priority_affinity_key( + auth_snapshot, + client_session_affinity, + ordering_config.scheduling_mode, + ); let candidates = enumerate_scheduler_candidates( selection_row_source, api_format, @@ -94,7 +99,9 @@ pub(super) async fn select_minimal_candidate( .0 .into_iter() .next(); - if ordering_config.scheduling_mode == SchedulerSchedulingMode::CacheAffinity { + if ordering_config.scheduling_mode == SchedulerSchedulingMode::CacheAffinity + && has_explicit_session_affinity(client_session_affinity) + { if let Some(candidate) = selected.as_ref() { remember_scheduler_affinity( affinity_cache_key.as_deref(), @@ -154,8 +161,11 @@ pub(super) async fn collect_selectable_candidates_with_skip_reasons( GatewayError, > { let ordering_config = runtime_state.read_scheduler_ordering_config().await?; - let priority_affinity_key = - scheduling_priority_affinity_key(auth_snapshot, ordering_config.scheduling_mode); + let priority_affinity_key = scheduling_priority_affinity_key( + auth_snapshot, + client_session_affinity, + ordering_config.scheduling_mode, + ); let candidates = enumerate_scheduler_candidates( selection_row_source, api_format, @@ -262,11 +272,17 @@ pub(super) async fn collect_selectable_enumerated_candidates_with_skip_reasons( pub(super) fn scheduling_priority_affinity_key<'a>( auth_snapshot: Option<&'a GatewayAuthApiKeySnapshot>, + client_session_affinity: Option<&ClientSessionAffinity>, scheduling_mode: SchedulerSchedulingMode, ) -> Option<&'a str> { if scheduling_mode == SchedulerSchedulingMode::FixedOrder { return None; } + if scheduling_mode == SchedulerSchedulingMode::CacheAffinity + && !has_explicit_session_affinity(client_session_affinity) + { + return None; + } auth_snapshot .map(|snapshot| snapshot.api_key_id.trim()) diff --git a/apps/aether-gateway/src/scheduler/candidate/tests/selection.rs b/apps/aether-gateway/src/scheduler/candidate/tests/selection.rs index 9a7bcbf15..824964199 100644 --- a/apps/aether-gateway/src/scheduler/candidate/tests/selection.rs +++ b/apps/aether-gateway/src/scheduler/candidate/tests/selection.rs @@ -13,7 +13,7 @@ use aether_data_contracts::repository::candidates::{ }; use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey; use aether_data_contracts::repository::quota::StoredProviderQuotaSnapshot; -use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate; +use aether_scheduler_core::{ClientSessionAffinity, SchedulerMinimalCandidateSelectionCandidate}; use serde_json::json; use crate::cache::SchedulerAffinityTarget; @@ -559,9 +559,85 @@ async fn cache_affinity_promotes_cached_scheduler_affinity_candidate_when_enable second.endpoint_id = "endpoint-b".to_string(); second.key_id = "key-b".to_string(); second.key_name = "beta".to_string(); - second.provider_priority = 0; - second.key_internal_priority = 0; - second.key_global_priority_by_format = Some(json!({"openai:chat": 0})); + second.provider_priority = 10; + second.key_internal_priority = 10; + second.key_global_priority_by_format = Some(json!({"openai:chat": 10})); + + let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![ + first, second, + ])); + let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![])); + let state = AppState::new() + .expect("state should build") + .with_data_state_for_tests( + GatewayDataState::with_candidate_selection_and_quota_for_tests(candidates, quotas) + .with_system_config_values_for_tests(vec![( + "scheduling_mode".to_string(), + json!("cache_affinity"), + )]), + ); + + let auth_snapshot = sample_auth_snapshot("affinity-key-1"); + let client_session_affinity = ClientSessionAffinity::from_session_key("session-1"); + let cache_key = build_scheduler_affinity_cache_key( + Some(&auth_snapshot), + "openai:chat", + "gpt-4.1", + Some(&client_session_affinity), + ) + .expect("scheduler affinity cache key should build"); + state.remember_scheduler_affinity_target( + &cache_key, + SchedulerAffinityTarget { + provider_id: "provider-b".to_string(), + endpoint_id: "endpoint-b".to_string(), + key_id: "key-b".to_string(), + }, + Duration::from_secs(300), + 100, + ); + + let selected = select_candidate_impl( + state.data.as_ref(), + &state, + "openai:chat", + "gpt-4.1", + false, + None, + Some(&auth_snapshot), + Some(&client_session_affinity), + 100, + false, + ) + .await + .expect("selection should succeed") + .expect("candidate should exist"); + + assert_eq!(selected.provider_id, "provider-b"); + assert_eq!(selected.key_id, "key-b"); +} + +#[tokio::test] +async fn cache_affinity_ignores_cached_scheduler_affinity_without_client_session() { + let mut first = sample_row(); + first.provider_id = "provider-a".to_string(); + first.provider_name = "provider-a".to_string(); + first.endpoint_id = "endpoint-a".to_string(); + first.key_id = "key-a".to_string(); + first.key_name = "alpha".to_string(); + first.provider_priority = 0; + first.key_internal_priority = 0; + first.key_global_priority_by_format = Some(json!({"openai:chat": 0})); + + let mut second = sample_row(); + second.provider_id = "provider-b".to_string(); + second.provider_name = "provider-b".to_string(); + second.endpoint_id = "endpoint-b".to_string(); + second.key_id = "key-b".to_string(); + second.key_name = "beta".to_string(); + second.provider_priority = 10; + second.key_internal_priority = 10; + second.key_global_priority_by_format = Some(json!({"openai:chat": 10})); let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![ first, second, @@ -602,8 +678,8 @@ async fn cache_affinity_promotes_cached_scheduler_affinity_candidate_when_enable .expect("selection should succeed") .expect("candidate should exist"); - assert_eq!(selected.provider_id, "provider-b"); - assert_eq!(selected.key_id, "key-b"); + assert_eq!(selected.provider_id, "provider-a"); + assert_eq!(selected.key_id, "key-a"); } #[tokio::test] @@ -623,18 +699,26 @@ async fn load_balance_selection_does_not_remember_scheduler_affinity() { )]), ); let auth_snapshot = sample_auth_snapshot("affinity-key-1"); - let cache_key = - build_scheduler_affinity_cache_key(Some(&auth_snapshot), "openai:chat", "gpt-4.1", None) - .expect("scheduler affinity cache key should build"); + let client_session_affinity = ClientSessionAffinity::from_session_key("session-1"); + let cache_key = build_scheduler_affinity_cache_key( + Some(&auth_snapshot), + "openai:chat", + "gpt-4.1", + Some(&client_session_affinity), + ) + .expect("scheduler affinity cache key should build"); - let selected = select_candidate( + let selected = select_candidate_impl( state.data.as_ref(), &state, "openai:chat", "gpt-4.1", false, + None, Some(&auth_snapshot), + Some(&client_session_affinity), 100, + false, ) .await .expect("selection should succeed") diff --git a/apps/aether-gateway/src/stage_metrics.rs b/apps/aether-gateway/src/stage_metrics.rs new file mode 100644 index 000000000..451357e1e --- /dev/null +++ b/apps/aether-gateway/src/stage_metrics.rs @@ -0,0 +1,335 @@ +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::LazyLock; + +use aether_runtime::{MetricKind, MetricLabel, MetricSample}; +use serde_json::{Map, Value}; + +const BUCKETS_MS: [u64; 12] = [1, 5, 10, 25, 50, 100, 250, 500, 1_000, 2_500, 5_000, 10_000]; + +const STAGES: [&str; 56] = [ + "frontdoor_handler_queue", + "frontdoor_admission", + "frontdoor_context", + "frontdoor_body_buffer", + "frontdoor_owner_forward", + "frontdoor_auth_model", + "frontdoor_rpm", + "frontdoor_local_ai_public", + "frontdoor_execute_stream", + "frontdoor_execute_sync", + "stream_candidate_slot", + "stream_path_step", + "stream_candidate_next", + "stream_candidate_source_next", + "stream_candidate_plan_build", + "stream_candidate_payload_parts", + "stream_candidate_proxy", + "stream_candidate_report_context", + "stream_candidate_decision_build", + "openai_chat_payload_parts_prepare", + "openai_chat_payload_model_directives", + "openai_chat_payload_redaction", + "chat_pii_redaction_request_cache_hit", + "chat_pii_redaction_runtime_config", + "chat_pii_redaction_feature_settings", + "chat_pii_redaction_mask_body", + "openai_chat_payload_auth_prepare", + "openai_chat_payload_body_build", + "candidate_page_load", + "candidate_page_resolve", + "pool_cursor_next_key", + "pool_score_load", + "pool_score_key_rows", + "pool_runtime_state", + "candidate_transport_snapshot", + "candidate_resolution_core", + "candidate_resolution_transport_read", + "candidate_resolution_rank", + "direct_reqwest_client_prewarm", + "stream_candidate_execute", + "stream_candidate_unused", + "stream_usage_pending", + "stream_provider_in_flight", + "stream_upstream_target_admission", + "stream_upstream_headers", + "stream_first_frame", + "stream_first_data", + "stream_response_policy", + "stream_response_ready", + "stream_total", + "direct_passthrough_upstream_body_first", + "direct_passthrough_first_client_send", + "direct_passthrough_body_send_wait", + "direct_passthrough_body_recv_first", + "direct_build_body", + "direct_send_headers", +]; + +const TRACE_STAGE_CAPACITY: usize = 16; +const STAGE_TRACE_MODE_ENV: &str = "AETHER_GATEWAY_STAGE_TRACE_MODE"; +const STAGE_TRACE_SLOW_MS_ENV: &str = "AETHER_GATEWAY_STAGE_TRACE_SLOW_MS"; +const STAGE_TRACE_SAMPLE_RATE_ENV: &str = "AETHER_GATEWAY_STAGE_TRACE_SAMPLE_RATE"; +const DEFAULT_STAGE_TRACE_SLOW_MS: u64 = 1_000; + +static METRICS: LazyLock> = + LazyLock::new(|| STAGES.iter().map(|stage| StageMetric::new(stage)).collect()); +static STAGE_TRACE_CONFIG: LazyLock = + LazyLock::new(read_stage_trace_config); +static STAGE_TRACE_SAMPLE_COUNTER: AtomicU64 = AtomicU64::new(0); + +struct StageMetric { + stage: &'static str, + count: AtomicU64, + sum_ms: AtomicU64, + max_ms: AtomicU64, + buckets: Vec, +} + +impl StageMetric { + fn new(stage: &'static str) -> Self { + Self { + stage, + count: AtomicU64::new(0), + sum_ms: AtomicU64::new(0), + max_ms: AtomicU64::new(0), + buckets: BUCKETS_MS.iter().map(|_| AtomicU64::new(0)).collect(), + } + } + + fn observe(&self, elapsed_ms: u64) { + self.count.fetch_add(1, Ordering::Relaxed); + self.sum_ms.fetch_add(elapsed_ms, Ordering::Relaxed); + update_max(&self.max_ms, elapsed_ms); + for (index, bucket) in BUCKETS_MS.iter().enumerate() { + if elapsed_ms <= *bucket { + self.buckets[index].fetch_add(1, Ordering::Relaxed); + } + } + } + + fn samples(&self) -> Vec { + let stage_label = vec![MetricLabel::new("stage", self.stage)]; + let mut samples = vec![ + MetricSample::new( + "gateway_stage_latency_count", + "Number of gateway stage latency observations.", + MetricKind::Counter, + self.count.load(Ordering::Relaxed), + ) + .with_labels(stage_label.clone()), + MetricSample::new( + "gateway_stage_latency_sum_ms", + "Total gateway stage latency in milliseconds.", + MetricKind::Counter, + self.sum_ms.load(Ordering::Relaxed), + ) + .with_labels(stage_label.clone()), + MetricSample::new( + "gateway_stage_latency_max_ms", + "Maximum observed gateway stage latency in milliseconds since process start.", + MetricKind::Gauge, + self.max_ms.load(Ordering::Relaxed), + ) + .with_labels(stage_label.clone()), + ]; + for (index, upper_bound_ms) in BUCKETS_MS.iter().enumerate() { + samples.push( + MetricSample::new( + "gateway_stage_latency_bucket", + "Cumulative gateway stage latency observations less than or equal to the bucket upper bound.", + MetricKind::Counter, + self.buckets[index].load(Ordering::Relaxed), + ) + .with_labels(vec![ + MetricLabel::new("stage", self.stage), + MetricLabel::new("le_ms", upper_bound_ms.to_string()), + ]), + ); + } + samples + } +} + +pub(crate) fn observe_gateway_stage_ms(stage: &'static str, elapsed_ms: u64) { + if let Some(metric) = METRICS.iter().find(|metric| metric.stage == stage) { + metric.observe(elapsed_ms); + } +} + +pub(crate) fn gateway_stage_metric_samples() -> Vec { + METRICS.iter().flat_map(StageMetric::samples).collect() +} + +fn update_max(max: &AtomicU64, value: u64) { + let mut current = max.load(Ordering::Relaxed); + while value > current { + match max.compare_exchange_weak(current, value, Ordering::Relaxed, Ordering::Relaxed) { + Ok(_) => break, + Err(next) => current = next, + } + } +} + +#[derive(Debug, Clone)] +pub(crate) struct RequestStageTrace { + mode: RequestStageTraceMode, + slow_ms: u64, + sampled: bool, + stages: Vec<(&'static str, u64)>, +} + +#[derive(Debug, Clone, Copy)] +struct RequestStageTraceConfig { + mode: RequestStageTraceMode, + slow_ms: u64, + sample_rate: f64, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum RequestStageTraceMode { + Off, + Slow, + Sample, + All, +} + +impl RequestStageTrace { + pub(crate) fn from_env() -> Self { + let config = *STAGE_TRACE_CONFIG; + let sampled = config.sample_rate > 0.0 && random_unit_sample() < config.sample_rate; + Self { + mode: config.mode, + slow_ms: config.slow_ms, + sampled, + stages: Vec::with_capacity(TRACE_STAGE_CAPACITY), + } + } + + pub(crate) fn observe(&mut self, stage: &'static str, elapsed_ms: u64) { + if self.mode == RequestStageTraceMode::Off { + return; + } + if let Some((_, existing)) = self + .stages + .iter_mut() + .find(|(existing_stage, _)| *existing_stage == stage) + { + *existing = elapsed_ms; + return; + } + if self.stages.len() < TRACE_STAGE_CAPACITY { + self.stages.push((stage, elapsed_ms)); + } + } + + pub(crate) fn into_metadata_value(self, fallback_elapsed_ms: Option) -> Option { + if self.mode == RequestStageTraceMode::Off || self.stages.is_empty() { + return None; + } + + let max_observed_ms = self + .stages + .iter() + .map(|(_, elapsed_ms)| *elapsed_ms) + .max() + .unwrap_or(0); + let fallback_elapsed_ms = fallback_elapsed_ms.unwrap_or(0); + let slow = max_observed_ms.max(fallback_elapsed_ms) >= self.slow_ms; + let should_emit = match self.mode { + RequestStageTraceMode::Off => false, + RequestStageTraceMode::Slow => slow || self.sampled, + RequestStageTraceMode::Sample => self.sampled, + RequestStageTraceMode::All => true, + }; + if !should_emit { + return None; + } + + let mut object = Map::new(); + for (stage, elapsed_ms) in self.stages { + object.insert(stage.to_string(), Value::from(elapsed_ms)); + } + Some(Value::Object(object)) + } +} + +pub(crate) fn observe_gateway_stage_trace_ms( + trace: &mut RequestStageTrace, + stage: &'static str, + elapsed_ms: u64, +) { + observe_gateway_stage_ms(stage, elapsed_ms); + trace.observe(stage, elapsed_ms); +} + +pub(crate) fn attach_stage_trace_to_report_context( + report_context: Option, + stage_timings_ms: Option, +) -> Option { + let Some(stage_timings_ms) = stage_timings_ms else { + return report_context; + }; + + let mut object = match report_context { + Some(Value::Object(object)) => object, + Some(other) => Map::from_iter([("seed".to_string(), other)]), + None => Map::new(), + }; + object.insert("stage_timings_ms".to_string(), stage_timings_ms); + Some(Value::Object(object)) +} + +fn read_stage_trace_config() -> RequestStageTraceConfig { + RequestStageTraceConfig { + mode: read_stage_trace_mode(), + slow_ms: read_stage_trace_slow_ms(), + sample_rate: read_stage_trace_sample_rate(), + } +} + +fn read_stage_trace_mode() -> RequestStageTraceMode { + match std::env::var(STAGE_TRACE_MODE_ENV) + .ok() + .as_deref() + .map(str::trim) + .map(str::to_ascii_lowercase) + .as_deref() + { + Some("all") => RequestStageTraceMode::All, + Some("sample") => RequestStageTraceMode::Sample, + Some("off") | Some("none") | Some("disabled") | Some("0") => RequestStageTraceMode::Off, + _ => RequestStageTraceMode::Slow, + } +} + +fn read_stage_trace_slow_ms() -> u64 { + std::env::var(STAGE_TRACE_SLOW_MS_ENV) + .ok() + .as_deref() + .map(str::trim) + .and_then(|value| value.parse::().ok()) + .filter(|value| *value > 0) + .unwrap_or(DEFAULT_STAGE_TRACE_SLOW_MS) +} + +fn read_stage_trace_sample_rate() -> f64 { + std::env::var(STAGE_TRACE_SAMPLE_RATE_ENV) + .ok() + .as_deref() + .map(str::trim) + .and_then(|value| value.parse::().ok()) + .filter(|value| value.is_finite()) + .map(|value| value.clamp(0.0, 1.0)) + .unwrap_or(0.0) +} + +fn random_unit_sample() -> f64 { + let mut value = STAGE_TRACE_SAMPLE_COUNTER + .fetch_add(1, Ordering::Relaxed) + .wrapping_add(0x9e37_79b9_7f4a_7c15); + value ^= value >> 12; + value ^= value << 25; + value ^= value >> 27; + let mixed = value.wrapping_mul(0x2545_f491_4f6c_dd1d); + (mixed as f64) / (u64::MAX as f64) +} diff --git a/apps/aether-gateway/src/state/app.rs b/apps/aether-gateway/src/state/app.rs index e318ba12a..d802d9351 100644 --- a/apps/aether-gateway/src/state/app.rs +++ b/apps/aether-gateway/src/state/app.rs @@ -2,6 +2,7 @@ use std::collections::HashMap; use std::sync::atomic::AtomicU64; use std::sync::Arc; use std::sync::Mutex as StdMutex; +use std::sync::RwLock as StdRwLock; use std::time::Duration; use aether_data::repository::users::StoredUserGroup; @@ -37,6 +38,15 @@ const MIN_LOCAL_EXECUTION_PLANNING_TIMEOUT_MS: u64 = 500; const MAX_LOCAL_EXECUTION_PLANNING_TIMEOUT_MS: u64 = 120_000; const LOCAL_EXECUTION_PLANNING_TIMEOUT_MS_ENV: &str = "AETHER_GATEWAY_LOCAL_EXECUTION_PLANNING_TIMEOUT_MS"; +const DEFAULT_CANDIDATE_PLANNING_GATE_LIMIT: usize = 1024; +const DEFAULT_UPSTREAM_EXECUTION_GATE_LIMIT: usize = 2000; +const DEFAULT_UPSTREAM_TARGET_GATE_LIMIT: usize = 2000; +const DEFAULT_INTERNAL_GATE_QUEUE_BUDGET_MS: u64 = 250; +const MAX_INTERNAL_GATE_QUEUE_BUDGET_MS: u64 = 5_000; +const CANDIDATE_PLANNING_GATE_LIMIT_ENV: &str = "AETHER_GATEWAY_CANDIDATE_PLANNING_GATE_LIMIT"; +const UPSTREAM_EXECUTION_GATE_LIMIT_ENV: &str = "AETHER_GATEWAY_UPSTREAM_EXECUTION_GATE_LIMIT"; +const UPSTREAM_TARGET_GATE_LIMIT_ENV: &str = "AETHER_GATEWAY_UPSTREAM_TARGET_GATE_LIMIT"; +const INTERNAL_GATE_QUEUE_BUDGET_MS_ENV: &str = "AETHER_GATEWAY_INTERNAL_GATE_QUEUE_BUDGET_MS"; #[cfg(test)] type TestExecutionRuntimeSyncOverrideFn = dyn Fn( @@ -62,6 +72,10 @@ impl std::fmt::Debug for TestExecutionRuntimeSyncOverride { pub(crate) struct FrontdoorRuntimeGuardConfig { pub(crate) request_body_read_timeout: Duration, pub(crate) local_execution_planning_timeout: Duration, + pub(crate) internal_gate_queue_budget: Duration, + pub(crate) candidate_planning_gate_limit: Option, + pub(crate) upstream_execution_gate_limit: Option, + pub(crate) upstream_target_gate_limit: Option, } impl FrontdoorRuntimeGuardConfig { @@ -79,6 +93,24 @@ impl FrontdoorRuntimeGuardConfig { MIN_LOCAL_EXECUTION_PLANNING_TIMEOUT_MS, MAX_LOCAL_EXECUTION_PLANNING_TIMEOUT_MS, ), + internal_gate_queue_budget: env_duration_ms( + INTERNAL_GATE_QUEUE_BUDGET_MS_ENV, + DEFAULT_INTERNAL_GATE_QUEUE_BUDGET_MS, + 1, + MAX_INTERNAL_GATE_QUEUE_BUDGET_MS, + ), + candidate_planning_gate_limit: env_optional_usize( + CANDIDATE_PLANNING_GATE_LIMIT_ENV, + DEFAULT_CANDIDATE_PLANNING_GATE_LIMIT, + ), + upstream_execution_gate_limit: env_optional_usize( + UPSTREAM_EXECUTION_GATE_LIMIT_ENV, + DEFAULT_UPSTREAM_EXECUTION_GATE_LIMIT, + ), + upstream_target_gate_limit: env_optional_usize( + UPSTREAM_TARGET_GATE_LIMIT_ENV, + DEFAULT_UPSTREAM_TARGET_GATE_LIMIT, + ), } } @@ -90,6 +122,12 @@ impl FrontdoorRuntimeGuardConfig { Self { request_body_read_timeout, local_execution_planning_timeout, + internal_gate_queue_budget: Duration::from_millis( + DEFAULT_INTERNAL_GATE_QUEUE_BUDGET_MS, + ), + candidate_planning_gate_limit: Some(DEFAULT_CANDIDATE_PLANNING_GATE_LIMIT), + upstream_execution_gate_limit: Some(DEFAULT_UPSTREAM_EXECUTION_GATE_LIMIT), + upstream_target_gate_limit: Some(DEFAULT_UPSTREAM_TARGET_GATE_LIMIT), } } } @@ -104,6 +142,17 @@ fn env_duration_ms(key: &str, default_ms: u64, min_ms: u64, max_ms: u64) -> Dura Duration::from_millis(ms) } +fn env_optional_usize(key: &str, default_value: usize) -> Option { + match std::env::var(key) + .ok() + .and_then(|value| value.trim().parse::().ok()) + { + Some(0) => None, + Some(value) => Some(value.max(1)), + None => Some(default_value.max(1)), + } +} + #[derive(Debug, Clone)] pub struct AppState { #[cfg(test)] @@ -117,6 +166,9 @@ pub struct AppState { pub(crate) video_task_poller: Option, pub(crate) frontdoor_runtime_guards: Arc, pub(crate) request_gate: Option>, + pub(crate) candidate_planning_gate: Option>, + pub(crate) upstream_execution_gate: Option>, + pub(crate) upstream_target_admission: Arc, pub(crate) distributed_request_gate: Option>, pub(crate) client: reqwest::Client, pub(crate) auth_context_cache: Arc, @@ -135,13 +187,17 @@ pub struct AppState { pub(crate) scheduler_affinity_epoch: Arc, pub(crate) dashboard_response_cache: Arc, pub(crate) system_config_cache: Arc, + pub(crate) candidate_page_cache: Arc, + pub(crate) candidate_resolved_page_cache: Arc, + pub(crate) chat_pii_redaction_runtime_config_cache: + crate::privacy::ChatPiiRedactionRuntimeConfigCacheHandle, pub(crate) fallback_metrics: Arc, pub(crate) request_candidate_queue: Option>, pub(crate) frontdoor_cors: Option>, pub(crate) frontdoor_user_rpm: Arc, pub(crate) tunnel: crate::tunnel::EmbeddedTunnelState, pub(crate) provider_transport_snapshot_cache: - Arc>>, + Arc>>, pub(crate) provider_key_rpm_resets: Arc>>, pub(crate) local_execution_runtime_miss_diagnostics: Arc>>, diff --git a/apps/aether-gateway/src/state/cache.rs b/apps/aether-gateway/src/state/cache.rs index b4281b62d..b464fb4ac 100644 --- a/apps/aether-gateway/src/state/cache.rs +++ b/apps/aether-gateway/src/state/cache.rs @@ -1,3 +1,4 @@ +use std::sync::Arc; use std::time::Duration; use super::super::provider_transport; @@ -5,10 +6,11 @@ use super::super::provider_transport; pub(crate) const AUTH_API_KEY_LAST_USED_TTL: Duration = Duration::from_secs(60); pub(crate) const AUTH_API_KEY_LAST_USED_MAX_ENTRIES: usize = 10_000; pub(crate) const PROVIDER_TRANSPORT_SNAPSHOT_CACHE_TTL: Duration = Duration::from_secs(1); +pub(crate) const PROVIDER_TRANSPORT_SNAPSHOT_CACHE_STALE_TTL: Duration = Duration::from_secs(30); pub(crate) const PROVIDER_TRANSPORT_SNAPSHOT_CACHE_MAX_ENTRIES: usize = 1_024; #[derive(Debug, Clone)] pub(crate) struct CachedProviderTransportSnapshot { pub(crate) loaded_at: std::time::Instant, - pub(crate) snapshot: provider_transport::GatewayProviderTransportSnapshot, + pub(crate) snapshot: Arc, } diff --git a/apps/aether-gateway/src/state/core.rs b/apps/aether-gateway/src/state/core.rs index 6eb7e7404..82e1e16aa 100644 --- a/apps/aether-gateway/src/state/core.rs +++ b/apps/aether-gateway/src/state/core.rs @@ -32,7 +32,7 @@ use super::super::async_task::{ use super::super::cache::{ AuthApiKeyLastUsedCache, AuthContextCache, AuthSnapshotCache, DashboardResponseCache, DirectPlanBypassCache, JsonValueCache, SchedulerAffinityCache, SchedulerAffinitySnapshotEntry, - SchedulerAffinityTarget, SystemConfigCache, ValueCache, + SchedulerAffinityTarget, SystemConfigCache, SystemConfigInflightRegistration, ValueCache, }; use super::super::data::{GatewayDataConfig, GatewayDataState}; use super::super::fallback_metrics; @@ -78,6 +78,7 @@ const AUTH_AFFECTING_SYSTEM_CONFIG_KEYS: &[&str] = &[ crate::constants::ANTIGRAVITY_BEARER_BRIDGE_CONFIG_KEY, ]; const FRONTDOOR_RPM_AFFECTING_SYSTEM_CONFIG_KEYS: &[&str] = &["rate_limit_per_minute"]; +const CHAT_PII_REDACTION_SYSTEM_CONFIG_PREFIX: &str = "module.chat_pii_redaction."; fn system_config_key_affects_scheduler(key: &str) -> bool { let key = key.trim(); @@ -94,7 +95,19 @@ fn system_config_key_affects_frontdoor_rpm(key: &str) -> bool { FRONTDOOR_RPM_AFFECTING_SYSTEM_CONFIG_KEYS.contains(&key) } +fn system_config_key_affects_chat_pii_redaction(key: &str) -> bool { + key.trim() + .starts_with(CHAT_PII_REDACTION_SYSTEM_CONFIG_PREFIX) +} + impl AppState { + pub async fn prewarm_chat_pii_redaction_runtime_config(&self) -> Result { + crate::privacy::read_chat_pii_redaction_runtime_config(self) + .await + .map(|config| config.enabled) + .map_err(|err| format!("{err:?}")) + } + fn usage_worker_queue_for( runtime_state: &Arc, ) -> Option> { @@ -172,6 +185,7 @@ impl AppState { self.clear_provider_transport_snapshot_cache(); self.invalidate_scheduler_affinity_cache(); self.invalidate_auth_context_cache(); + self.candidate_resolved_page_cache.clear(); self.system_config_cache.clear(); self.frontdoor_user_rpm.clear_system_default_cache(); let data = Arc::new( @@ -179,6 +193,8 @@ impl AppState { .clone() .with_usage_worker_queue(Self::usage_worker_queue_for(&self.runtime_state)), ); + self.candidate_page_cache.clear(); + self.candidate_resolved_page_cache.clear(); self.tunnel = crate::tunnel::EmbeddedTunnelState::with_data_and_runtime_state( Arc::clone(&data), self.runtime_state.clone(), @@ -222,6 +238,7 @@ impl AppState { http2_adaptive_window: true, ..HttpClientConfig::default() })?; + let frontdoor_runtime_guards = Arc::new(FrontdoorRuntimeGuardConfig::from_env()); Ok(Self { #[cfg(test)] execution_runtime_override_base_url: execution_runtime_override_base_url @@ -236,8 +253,20 @@ impl AppState { VideoTaskTruthSourceMode::PythonSyncReport, )), video_task_poller: None, - frontdoor_runtime_guards: Arc::new(FrontdoorRuntimeGuardConfig::from_env()), + frontdoor_runtime_guards: Arc::clone(&frontdoor_runtime_guards), request_gate: None, + candidate_planning_gate: frontdoor_runtime_guards + .candidate_planning_gate_limit + .map(|limit| Arc::new(ConcurrencyGate::new("gateway_candidate_planning", limit))), + upstream_execution_gate: frontdoor_runtime_guards + .upstream_execution_gate_limit + .map(|limit| Arc::new(ConcurrencyGate::new("gateway_upstream_execution", limit))), + upstream_target_admission: Arc::new( + crate::upstream_admission::UpstreamTargetAdmission::new( + frontdoor_runtime_guards.upstream_target_gate_limit, + frontdoor_runtime_guards.internal_gate_queue_budget, + ), + ), distributed_request_gate: None, client, auth_context_cache: Arc::new(AuthContextCache::default()), @@ -255,6 +284,12 @@ impl AppState { scheduler_affinity_epoch: Arc::new(AtomicU64::new(0)), dashboard_response_cache: Arc::new(DashboardResponseCache::default()), system_config_cache: Arc::new(SystemConfigCache::default()), + candidate_page_cache: Arc::new(crate::cache::CandidatePageCache::default()), + candidate_resolved_page_cache: Arc::new( + crate::cache::CandidateResolvedPageCache::default(), + ), + chat_pii_redaction_runtime_config_cache: + crate::privacy::new_chat_pii_redaction_runtime_config_cache(), fallback_metrics: Arc::new(fallback_metrics::GatewayFallbackMetrics::default()), request_candidate_queue: None, frontdoor_cors: None, @@ -265,7 +300,7 @@ impl AppState { data, runtime_state.clone(), ), - provider_transport_snapshot_cache: Arc::new(StdMutex::new(HashMap::new())), + provider_transport_snapshot_cache: Arc::new(std::sync::RwLock::new(HashMap::new())), provider_key_rpm_resets: Arc::new(StdMutex::new(HashMap::new())), local_execution_runtime_miss_diagnostics: Arc::new(StdMutex::new(HashMap::new())), admin_monitoring_error_stats_reset_at: Arc::new(StdMutex::new(None)), @@ -535,19 +570,44 @@ impl AppState { return Ok(value); } - let _guard = self.system_config_cache.load_guard().await; - if let Some(value) = self.system_config_cache.get(key, SYSTEM_CONFIG_CACHE_TTL) { - return Ok(value); + loop { + let notified = self.system_config_cache.notified(); + match self.system_config_cache.register_load(key) { + SystemConfigInflightRegistration::Bypass => { + let value = self + .data + .find_system_config_value(key) + .await + .map_err(|err| GatewayError::Internal(err.to_string()))?; + self.system_config_cache.insert( + key.to_string(), + value.clone(), + SYSTEM_CONFIG_CACHE_TTL, + ); + return Ok(value); + } + SystemConfigInflightRegistration::Follower => { + notified.await; + if let Some(value) = self.system_config_cache.get(key, SYSTEM_CONFIG_CACHE_TTL) + { + return Ok(value); + } + } + SystemConfigInflightRegistration::Leader(_guard) => { + let value = self + .data + .find_system_config_value(key) + .await + .map_err(|err| GatewayError::Internal(err.to_string()))?; + self.system_config_cache.insert( + key.to_string(), + value.clone(), + SYSTEM_CONFIG_CACHE_TTL, + ); + return Ok(value); + } + } } - - let value = self - .data - .find_system_config_value(key) - .await - .map_err(|err| GatewayError::Internal(err.to_string()))?; - self.system_config_cache - .insert(key.to_string(), value.clone(), SYSTEM_CONFIG_CACHE_TTL); - Ok(value) } pub(crate) async fn upsert_system_config_json_value( @@ -606,12 +666,19 @@ impl AppState { if deleted && system_config_key_affects_frontdoor_rpm(key) { self.frontdoor_user_rpm.clear_system_default_cache(); } + if deleted && system_config_key_affects_chat_pii_redaction(key) { + crate::privacy::clear_chat_pii_redaction_runtime_config_cache( + &self.chat_pii_redaction_runtime_config_cache, + ); + } Ok(deleted) } pub(crate) fn invalidate_provider_routing_caches(&self) { self.data.clear_minimal_candidate_selection_cache(); self.data.clear_provider_catalog_cache(); + self.candidate_page_cache.clear(); + self.candidate_resolved_page_cache.clear(); self.clear_provider_transport_snapshot_cache(); self.invalidate_scheduler_affinity_cache(); } @@ -619,6 +686,8 @@ impl AppState { pub(crate) fn invalidate_provider_health_routing_caches(&self) { self.data.clear_minimal_candidate_selection_cache(); self.data.clear_provider_catalog_cache(); + self.candidate_page_cache.clear(); + self.candidate_resolved_page_cache.clear(); self.clear_provider_transport_snapshot_cache(); } @@ -631,6 +700,8 @@ impl AppState { self.auth_api_key_feature_settings_cache.clear(); self.provider_quota_snapshot_cache.clear(); self.user_groups_for_user_cache.clear(); + self.candidate_page_cache.clear(); + self.candidate_resolved_page_cache.clear(); } fn remember_system_config_write(&self, key: &str, value: Option) { @@ -645,6 +716,11 @@ impl AppState { if system_config_key_affects_frontdoor_rpm(key) { self.frontdoor_user_rpm.clear_system_default_cache(); } + if system_config_key_affects_chat_pii_redaction(key) { + crate::privacy::clear_chat_pii_redaction_runtime_config_cache( + &self.chat_pii_redaction_runtime_config_cache, + ); + } } pub(crate) async fn read_admin_system_stats( @@ -908,6 +984,18 @@ impl AppState { self.request_gate.as_ref().map(|gate| gate.snapshot()) } + pub(crate) fn candidate_planning_concurrency_snapshot(&self) -> Option { + self.candidate_planning_gate + .as_ref() + .map(|gate| gate.snapshot()) + } + + pub(crate) fn upstream_execution_concurrency_snapshot(&self) -> Option { + self.upstream_execution_gate + .as_ref() + .map(|gate| gate.snapshot()) + } + pub(crate) async fn distributed_request_concurrency_snapshot( &self, ) -> Result, RuntimeSemaphoreError> { @@ -922,6 +1010,12 @@ impl AppState { if let Some(snapshot) = self.request_concurrency_snapshot() { samples.extend(snapshot.to_metric_samples("gateway_requests")); } + if let Some(snapshot) = self.candidate_planning_concurrency_snapshot() { + samples.extend(snapshot.to_metric_samples("gateway_candidate_planning")); + } + if let Some(snapshot) = self.upstream_execution_concurrency_snapshot() { + samples.extend(snapshot.to_metric_samples("gateway_upstream_execution")); + } if let Some(gate) = self.distributed_request_gate.as_ref() { match gate.snapshot().await { Ok(snapshot) => { @@ -947,6 +1041,12 @@ impl AppState { if let Some(queue) = self.request_candidate_queue.as_ref() { samples.extend(queue.metric_samples()); } + samples.extend( + crate::execution_runtime::transport::direct_reqwest_client_cache_metric_samples(), + ); + samples.extend(self.upstream_target_admission.metric_samples()); + samples.extend(crate::cache::candidate_page_cache_metric_samples()); + samples.extend(crate::stage_metrics::gateway_stage_metric_samples()); samples.extend(self.tunnel.metric_samples()); samples.extend(self.fallback_metrics.metric_samples()); samples @@ -1131,6 +1231,8 @@ impl AppState { .fetch_add(1, Ordering::AcqRel) .saturating_add(1); self.scheduler_affinity_cache.clear(); + self.candidate_page_cache.clear(); + self.candidate_resolved_page_cache.clear(); next_epoch } diff --git a/apps/aether-gateway/src/state/mod.rs b/apps/aether-gateway/src/state/mod.rs index 940af360d..0e859ca8d 100644 --- a/apps/aether-gateway/src/state/mod.rs +++ b/apps/aether-gateway/src/state/mod.rs @@ -31,7 +31,7 @@ pub(crate) use self::app::FrontdoorRuntimeGuardConfig; pub(crate) use self::cache::{ CachedProviderTransportSnapshot, AUTH_API_KEY_LAST_USED_MAX_ENTRIES, AUTH_API_KEY_LAST_USED_TTL, PROVIDER_TRANSPORT_SNAPSHOT_CACHE_MAX_ENTRIES, - PROVIDER_TRANSPORT_SNAPSHOT_CACHE_TTL, + PROVIDER_TRANSPORT_SNAPSHOT_CACHE_STALE_TTL, PROVIDER_TRANSPORT_SNAPSHOT_CACHE_TTL, }; pub use self::cors::FrontdoorCorsConfig; pub(crate) use self::types::{ diff --git a/apps/aether-gateway/src/state/oauth.rs b/apps/aether-gateway/src/state/oauth.rs index 2f620b47d..0a621b2c6 100644 --- a/apps/aether-gateway/src/state/oauth.rs +++ b/apps/aether-gateway/src/state/oauth.rs @@ -1,7 +1,7 @@ use super::{ provider_transport_snapshot_looks_refreshed, AppState, CachedProviderTransportSnapshot, GatewayError, ProviderTransportSnapshotCacheKey, PROVIDER_TRANSPORT_SNAPSHOT_CACHE_MAX_ENTRIES, - PROVIDER_TRANSPORT_SNAPSHOT_CACHE_TTL, + PROVIDER_TRANSPORT_SNAPSHOT_CACHE_STALE_TTL, }; use crate::handlers::shared::default_provider_key_status_snapshot; use crate::provider_transport::LocalOAuthHttpExecutor; @@ -20,6 +20,7 @@ use serde_json::{json, Map, Value}; use sha2::{Digest, Sha256}; use std::collections::BTreeMap; use std::io::Read; +use std::sync::Arc; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use aether_crypto::encrypt_python_fernet_plaintext; @@ -479,39 +480,49 @@ impl<'a> provider_transport::LocalOAuthHttpExecutor for GatewayLocalOAuthHttpExe impl AppState { pub(crate) fn clear_provider_transport_snapshot_cache(&self) { self.provider_transport_snapshot_cache - .lock() + .write() .expect("provider transport snapshot cache should lock") .clear(); } - fn get_cached_provider_transport_snapshot( + fn get_cached_provider_transport_snapshot_arc( &self, cache_key: &ProviderTransportSnapshotCacheKey, - ) -> Option { - let mut cache = self - .provider_transport_snapshot_cache - .lock() - .expect("provider transport snapshot cache should lock"); - let cached = cache.get(cache_key).cloned()?; - if cached.loaded_at.elapsed() <= PROVIDER_TRANSPORT_SNAPSHOT_CACHE_TTL { + ) -> Option> { + let cached = { + let cache = self + .provider_transport_snapshot_cache + .read() + .expect("provider transport snapshot cache should lock"); + cache.get(cache_key).cloned() + }?; + if cached.loaded_at.elapsed() <= PROVIDER_TRANSPORT_SNAPSHOT_CACHE_STALE_TTL { return Some(cached.snapshot); } - cache.remove(cache_key); + let mut cache = self + .provider_transport_snapshot_cache + .write() + .expect("provider transport snapshot cache should lock"); + if cache.get(cache_key).is_some_and(|entry| { + entry.loaded_at.elapsed() > PROVIDER_TRANSPORT_SNAPSHOT_CACHE_STALE_TTL + }) { + cache.remove(cache_key); + } None } fn put_cached_provider_transport_snapshot( &self, cache_key: ProviderTransportSnapshotCacheKey, - snapshot: provider_transport::GatewayProviderTransportSnapshot, + snapshot: Arc, ) { let mut cache = self .provider_transport_snapshot_cache - .lock() + .write() .expect("provider transport snapshot cache should lock"); if cache.len() >= PROVIDER_TRANSPORT_SNAPSHOT_CACHE_MAX_ENTRIES { cache.retain(|_, entry| { - entry.loaded_at.elapsed() <= PROVIDER_TRANSPORT_SNAPSHOT_CACHE_TTL + entry.loaded_at.elapsed() <= PROVIDER_TRANSPORT_SNAPSHOT_CACHE_STALE_TTL }); if cache.len() >= PROVIDER_TRANSPORT_SNAPSHOT_CACHE_MAX_ENTRIES { cache.clear(); @@ -796,6 +807,41 @@ impl AppState { None } + pub(crate) async fn read_provider_transport_snapshot_arc( + &self, + provider_id: &str, + endpoint_id: &str, + key_id: &str, + ) -> Result< + Option>, + GatewayError, + > { + let Some(cache_key) = + ProviderTransportSnapshotCacheKey::new(provider_id, endpoint_id, key_id) + else { + return Ok(self + .read_provider_transport_snapshot_uncached(provider_id, endpoint_id, key_id) + .await? + .map(Arc::new)); + }; + if let Some(snapshot) = self.get_cached_provider_transport_snapshot_arc(&cache_key) { + return Ok(Some(snapshot)); + } + + let snapshot = self + .read_provider_transport_snapshot_uncached(provider_id, endpoint_id, key_id) + .await?; + match snapshot { + Some(snapshot) => { + let snapshot = self.apply_global_format_conversion_override(snapshot).await; + let snapshot = Arc::new(snapshot); + self.put_cached_provider_transport_snapshot(cache_key, Arc::clone(&snapshot)); + Ok(Some(snapshot)) + } + None => Ok(None), + } + } + pub(crate) async fn read_provider_transport_snapshot( &self, provider_id: &str, @@ -803,31 +849,10 @@ impl AppState { key_id: &str, ) -> Result, GatewayError> { - let Some(cache_key) = - ProviderTransportSnapshotCacheKey::new(provider_id, endpoint_id, key_id) - else { - return self - .read_provider_transport_snapshot_uncached(provider_id, endpoint_id, key_id) - .await; - }; - if let Some(snapshot) = self.get_cached_provider_transport_snapshot(&cache_key) { - return Ok(Some( - self.apply_global_format_conversion_override(snapshot).await, - )); - } - - let snapshot = self - .read_provider_transport_snapshot_uncached(provider_id, endpoint_id, key_id) - .await?; - if let Some(snapshot) = snapshot.as_ref() { - self.put_cached_provider_transport_snapshot(cache_key, snapshot.clone()); - } - match snapshot { - Some(snapshot) => Ok(Some( - self.apply_global_format_conversion_override(snapshot).await, - )), - None => Ok(None), - } + Ok(self + .read_provider_transport_snapshot_arc(provider_id, endpoint_id, key_id) + .await? + .map(|snapshot| (*snapshot).clone())) } pub(crate) async fn update_provider_catalog_key_oauth_credentials( diff --git a/apps/aether-gateway/src/state/testing.rs b/apps/aether-gateway/src/state/testing.rs index 697521027..ceb3c874c 100644 --- a/apps/aether-gateway/src/state/testing.rs +++ b/apps/aether-gateway/src/state/testing.rs @@ -21,6 +21,11 @@ impl AppState { self } + pub(crate) fn without_request_candidate_queue_for_tests(mut self) -> Self { + self.request_candidate_queue = None; + self + } + pub(crate) fn with_turnstile_siteverify_url_for_tests(mut self, url: &str) -> Self { self.turnstile_siteverify_url_override = Some(url.trim().to_string()); self diff --git a/apps/aether-gateway/src/upstream_admission.rs b/apps/aether-gateway/src/upstream_admission.rs new file mode 100644 index 000000000..82d9eb3cd --- /dev/null +++ b/apps/aether-gateway/src/upstream_admission.rs @@ -0,0 +1,269 @@ +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use aether_contracts::ExecutionPlan; +use aether_runtime::{ConcurrencyGate, ConcurrencyPermit, MetricKind, MetricLabel, MetricSample}; +use dashmap::DashMap; +use tokio::time::timeout; +use url::Url; + +use crate::stage_metrics::observe_gateway_stage_ms; +use crate::GatewayError; + +const GATE_NAME: &str = "gateway_upstream_target"; +const DEFAULT_METRIC_TARGET_LIMIT: usize = 32; +const METRIC_TARGET_LIMIT_ENV: &str = "AETHER_GATEWAY_UPSTREAM_TARGET_GATE_METRIC_LIMIT"; + +#[derive(Debug)] +pub(crate) struct UpstreamTargetAdmission { + limit: Option, + queue_budget: Duration, + gates: DashMap>, +} + +#[derive(Debug)] +pub(crate) struct UpstreamTargetAdmissionPermit { + _permit: ConcurrencyPermit, +} + +impl UpstreamTargetAdmission { + pub(crate) fn new(limit: Option, queue_budget: Duration) -> Self { + Self { + limit, + queue_budget, + gates: DashMap::new(), + } + } + + pub(crate) async fn acquire( + &self, + plan: &ExecutionPlan, + trace_id: &str, + ) -> Result, GatewayError> { + let Some(limit) = self.limit else { + return Ok(None); + }; + let key = upstream_target_key(plan); + let gate = self + .gates + .entry(key.clone()) + .or_insert_with(|| Arc::new(ConcurrencyGate::new(GATE_NAME, limit))) + .clone(); + let started_at = Instant::now(); + let permit = match timeout(self.queue_budget, gate.acquire()).await { + Ok(Ok(permit)) => permit, + Ok(Err(err)) => return Err(GatewayError::Internal(err.to_string())), + Err(_) => { + tracing::debug!( + event_name = "gateway_upstream_target_admission_timeout", + log_type = "ops", + trace_id, + target = key.as_str(), + limit, + queue_budget_ms = self.queue_budget.as_millis() as u64, + "gateway upstream target admission gate timed out" + ); + return Err(GatewayError::AdmissionTimeout { + trace_id: trace_id.to_string(), + gate: GATE_NAME, + queue_budget_ms: self.queue_budget.as_millis() as u64, + }); + } + }; + observe_gateway_stage_ms( + "stream_upstream_target_admission", + started_at.elapsed().as_millis() as u64, + ); + Ok(Some(UpstreamTargetAdmissionPermit { _permit: permit })) + } + + pub(crate) fn metric_samples(&self) -> Vec { + let mut samples = vec![MetricSample::new( + "upstream_target_gate_active_targets", + "Number of upstream targets currently tracked by the gateway upstream target admission gates.", + MetricKind::Gauge, + self.gates.len() as u64, + )]; + + let Some(limit) = self.limit else { + return samples; + }; + + samples.push(MetricSample::new( + "upstream_target_gate_limit", + "Configured per-upstream-target admission gate limit.", + MetricKind::Gauge, + limit as u64, + )); + + let mut snapshots = self + .gates + .iter() + .map(|entry| { + let snapshot = entry.value().snapshot(); + ( + entry.key().clone(), + snapshot.in_flight, + snapshot.available_permits, + snapshot.high_watermark, + snapshot.rejected, + ) + }) + .collect::>(); + snapshots.sort_by(|left, right| { + right + .1 + .cmp(&left.1) + .then_with(|| right.3.cmp(&left.3)) + .then_with(|| right.4.cmp(&left.4)) + }); + + let metric_target_limit = upstream_target_metric_limit(); + for (target, in_flight, available, high_watermark, rejected) in + snapshots.into_iter().take(metric_target_limit) + { + let labels = vec![MetricLabel::new("target", target)]; + samples.push( + MetricSample::new( + "upstream_target_gate_in_flight", + "Current number of in-flight operations for an upstream target admission gate.", + MetricKind::Gauge, + in_flight as u64, + ) + .with_labels(labels.clone()), + ); + samples.push( + MetricSample::new( + "upstream_target_gate_available_permits", + "Currently available permits for an upstream target admission gate.", + MetricKind::Gauge, + available as u64, + ) + .with_labels(labels.clone()), + ); + samples.push( + MetricSample::new( + "upstream_target_gate_high_watermark", + "Highest observed in-flight count for an upstream target admission gate.", + MetricKind::Gauge, + high_watermark as u64, + ) + .with_labels(labels.clone()), + ); + samples.push( + MetricSample::new( + "upstream_target_gate_rejected_total", + "Number of operations rejected by an upstream target admission gate.", + MetricKind::Counter, + rejected, + ) + .with_labels(labels), + ); + } + + samples + } +} + +pub(crate) fn upstream_target_key(plan: &ExecutionPlan) -> String { + let parsed = Url::parse(plan.url.as_str()).ok(); + let Some(url) = parsed else { + return fallback_target_key(plan); + }; + let scheme = url.scheme().to_ascii_lowercase(); + let Some(host) = url.host_str().map(|host| host.to_ascii_lowercase()) else { + return fallback_target_key(plan); + }; + let port = url + .port_or_known_default() + .map(|port| port.to_string()) + .unwrap_or_else(|| "-".to_string()); + let proxy = plan + .proxy + .as_ref() + .and_then(|proxy| proxy.url.as_deref()) + .map(str::trim) + .filter(|value| !value.is_empty()) + .unwrap_or("-"); + format!("{scheme}://{host}:{port}|proxy={proxy}") +} + +fn fallback_target_key(plan: &ExecutionPlan) -> String { + format!( + "unparsed|provider={}|endpoint={}|url={}", + plan.provider_id, plan.endpoint_id, plan.url + ) +} + +fn upstream_target_metric_limit() -> usize { + std::env::var(METRIC_TARGET_LIMIT_ENV) + .ok() + .and_then(|value| value.trim().parse::().ok()) + .unwrap_or(DEFAULT_METRIC_TARGET_LIMIT) +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use aether_contracts::{ExecutionPlan, RequestBody}; + use serde_json::json; + + use super::*; + + fn test_plan(url: &str) -> ExecutionPlan { + ExecutionPlan { + request_id: "req-upstream-target".to_string(), + candidate_id: Some("cand-upstream-target".to_string()), + provider_name: Some("provider".to_string()), + provider_id: "provider_id".to_string(), + endpoint_id: "endpoint_id".to_string(), + key_id: "key_id".to_string(), + method: "POST".to_string(), + url: url.to_string(), + headers: Default::default(), + content_type: Some("application/json".to_string()), + content_encoding: None, + body: RequestBody::from_json(json!({"stream": true})), + stream: true, + client_api_format: "openai".to_string(), + provider_api_format: "openai".to_string(), + model_name: Some("model".to_string()), + proxy: None, + transport_profile: None, + timeouts: None, + } + } + + #[test] + fn upstream_target_key_ignores_path_and_query() { + let left = test_plan("http://127.0.0.1:18181/v1/chat/completions?x=1"); + let right = test_plan("http://127.0.0.1:18181/v1/responses"); + + assert_eq!(upstream_target_key(&left), upstream_target_key(&right)); + } + + #[tokio::test] + async fn acquire_times_out_when_target_gate_is_saturated() { + let admission = UpstreamTargetAdmission::new(Some(1), Duration::from_millis(1)); + let plan = test_plan("http://127.0.0.1:18181/v1/chat/completions"); + let _first = admission + .acquire(&plan, "trace-upstream-target") + .await + .expect("first acquire should succeed") + .expect("gate enabled"); + + let err = admission + .acquire(&plan, "trace-upstream-target") + .await + .expect_err("second acquire should time out"); + + assert!(matches!( + err, + GatewayError::AdmissionTimeout { + gate: "gateway_upstream_target", + .. + } + )); + } +} diff --git a/apps/aether-tunnel/src/upstream_client.rs b/apps/aether-tunnel/src/upstream_client.rs index 4ee2b0ce6..c31ae16d8 100644 --- a/apps/aether-tunnel/src/upstream_client.rs +++ b/apps/aether-tunnel/src/upstream_client.rs @@ -12,7 +12,7 @@ use std::time::Duration; use aether_contracts::{ ResolvedTransportProfile, TRANSPORT_BACKEND_HYPER_RUSTLS, TRANSPORT_BACKEND_REQWEST_RUSTLS, - TRANSPORT_HTTP_MODE_HTTP1_ONLY, + TRANSPORT_HTTP_MODE_H2C_PRIOR_KNOWLEDGE, TRANSPORT_HTTP_MODE_HTTP1_ONLY, }; use bytes::Bytes; use futures_util::Stream; @@ -99,10 +99,15 @@ impl UpstreamClientPool { let http1_only = key .http_mode .eq_ignore_ascii_case(TRANSPORT_HTTP_MODE_HTTP1_ONLY); + let h2c_prior_knowledge = !http1_only + && key + .http_mode + .eq_ignore_ascii_case(TRANSPORT_HTTP_MODE_H2C_PRIOR_KNOWLEDGE); let client = build_upstream_client_with_protocol( &self.config, Arc::clone(&self.dns_cache), http1_only, + h2c_prior_knowledge, )?; let mut clients = self.clients.lock().expect("client pool lock"); if let Some(entry) = clients.get_mut(&key) { @@ -472,6 +477,7 @@ fn build_upstream_client_with_protocol( config: &Config, dns_cache: Arc, http1_only: bool, + h2c_prior_knowledge: bool, ) -> Result { let mut http = HttpConnector::new_with_resolver(ValidatedResolver::new( dns_cache, @@ -507,6 +513,9 @@ fn build_upstream_client_with_protocol( }; let mut builder = Client::builder(TokioExecutor::new()); + if h2c_prior_knowledge { + builder.http2_only(true); + } builder.pool_max_idle_per_host(config.upstream_pool_max_idle_per_host); builder.pool_idle_timeout(Duration::from_secs(config.upstream_pool_idle_timeout_secs)); builder.pool_timer(TokioTimer::new()); diff --git a/crates/aether-cache/src/ttl_map.rs b/crates/aether-cache/src/ttl_map.rs index 250f163ac..fe553e256 100644 --- a/crates/aether-cache/src/ttl_map.rs +++ b/crates/aether-cache/src/ttl_map.rs @@ -63,6 +63,46 @@ where ); } + pub fn insert_if_absent_fresh( + &self, + key: K, + value: V, + ttl: Duration, + max_entries: usize, + ) -> bool { + if ttl.is_zero() { + return true; + } + + let Ok(mut entries) = self.entries.lock() else { + return false; + }; + + prune_expired(&mut entries, ttl); + if entries.contains_key(&key) { + return false; + } + while max_entries > 0 && entries.len() >= max_entries { + let Some(oldest_key) = entries + .iter() + .min_by_key(|(_, entry)| entry.inserted_at) + .map(|(key, _)| key.clone()) + else { + break; + }; + entries.remove(&oldest_key); + } + + entries.insert( + key, + TimedEntry { + value, + inserted_at: Instant::now(), + }, + ); + true + } + pub fn remove(&self, key: &K) -> Option { let Ok(mut entries) = self.entries.lock() else { return None; @@ -94,6 +134,22 @@ where K: Eq + Hash + Clone, V: Clone, { + pub fn get_with_age(&self, key: &K, max_age: Duration) -> Option<(V, Duration)> { + let Ok(mut entries) = self.entries.lock() else { + return None; + }; + + let entry = entries.get(key).cloned()?; + let age = entry.inserted_at.elapsed(); + + if age > max_age { + entries.remove(key); + return None; + } + + Some((entry.value, age)) + } + pub fn get_fresh(&self, key: &K, ttl: Duration) -> Option { let Ok(mut entries) = self.entries.lock() else { return None; @@ -204,4 +260,49 @@ mod tests { Some(3) ); } + + #[test] + fn insert_if_absent_fresh_rejects_fresh_duplicate() { + let cache = ExpiringMap::new(); + + assert!(cache.insert_if_absent_fresh( + "hello".to_string(), + 1_u32, + std::time::Duration::from_secs(60), + 16, + )); + assert!(!cache.insert_if_absent_fresh( + "hello".to_string(), + 2_u32, + std::time::Duration::from_secs(60), + 16, + )); + assert_eq!( + cache.get_fresh(&"hello".to_string(), std::time::Duration::from_secs(60)), + Some(1) + ); + } + + #[test] + fn insert_if_absent_fresh_allows_after_expiry() { + let cache = ExpiringMap::new(); + + assert!(cache.insert_if_absent_fresh( + "hello".to_string(), + 1_u32, + std::time::Duration::from_millis(10), + 16, + )); + sleep(std::time::Duration::from_millis(20)); + assert!(cache.insert_if_absent_fresh( + "hello".to_string(), + 2_u32, + std::time::Duration::from_millis(10), + 16, + )); + assert_eq!( + cache.get_fresh(&"hello".to_string(), std::time::Duration::from_millis(10)), + Some(2) + ); + } } diff --git a/crates/aether-contracts/src/lib.rs b/crates/aether-contracts/src/lib.rs index 69b331025..a84b35e04 100644 --- a/crates/aether-contracts/src/lib.rs +++ b/crates/aether-contracts/src/lib.rs @@ -13,7 +13,8 @@ pub use plan::{ EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER, EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER, EXECUTION_REQUEST_HTTP1_ONLY_HEADER, TRANSPORT_BACKEND_BROWSER_WREQ, TRANSPORT_BACKEND_HYPER_RUSTLS, TRANSPORT_BACKEND_REQWEST_RUSTLS, TRANSPORT_HTTP_MODE_AUTO, - TRANSPORT_HTTP_MODE_HTTP1_ONLY, TRANSPORT_POOL_SCOPE_KEY, + TRANSPORT_HTTP_MODE_H2C_PRIOR_KNOWLEDGE, TRANSPORT_HTTP_MODE_HTTP1_ONLY, + TRANSPORT_POOL_SCOPE_KEY, }; pub use result::{ExecutionResult, ExecutionTelemetry, ResponseBody}; pub use usage::{ diff --git a/crates/aether-contracts/src/plan.rs b/crates/aether-contracts/src/plan.rs index 740bc09ad..2e58736a6 100644 --- a/crates/aether-contracts/src/plan.rs +++ b/crates/aether-contracts/src/plan.rs @@ -66,6 +66,7 @@ pub const TRANSPORT_BACKEND_HYPER_RUSTLS: &str = "hyper_rustls"; pub const TRANSPORT_BACKEND_BROWSER_WREQ: &str = "browser_wreq"; pub const TRANSPORT_HTTP_MODE_AUTO: &str = "auto"; pub const TRANSPORT_HTTP_MODE_HTTP1_ONLY: &str = "http1_only"; +pub const TRANSPORT_HTTP_MODE_H2C_PRIOR_KNOWLEDGE: &str = "h2c_prior_knowledge"; pub const TRANSPORT_POOL_SCOPE_KEY: &str = "key"; #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] diff --git a/crates/aether-data-contracts/src/repository/candidates/types.rs b/crates/aether-data-contracts/src/repository/candidates/types.rs index 49d89d9b6..81243428f 100644 --- a/crates/aether-data-contracts/src/repository/candidates/types.rs +++ b/crates/aether-data-contracts/src/repository/candidates/types.rs @@ -538,6 +538,18 @@ pub trait RequestCandidateWriteRepository: Send + Sync { candidate: UpsertRequestCandidateRecord, ) -> Result; + async fn upsert_many( + &self, + candidates: Vec, + ) -> Result { + let mut persisted = 0usize; + for candidate in candidates { + self.upsert(candidate).await?; + persisted = persisted.saturating_add(1); + } + Ok(persisted) + } + async fn delete_created_before( &self, created_before_unix_secs: u64, diff --git a/crates/aether-data/src/repository/auth/memory.rs b/crates/aether-data/src/repository/auth/memory.rs index 6eeaa833b..4cdefdf48 100644 --- a/crates/aether-data/src/repository/auth/memory.rs +++ b/crates/aether-data/src/repository/auth/memory.rs @@ -1,6 +1,6 @@ use std::collections::BTreeMap; use std::sync::RwLock; -use std::time::{SystemTime, UNIX_EPOCH}; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; use async_trait::async_trait; @@ -26,11 +26,14 @@ struct MemoryAuthApiKeyIndex { export_by_api_key_id: BTreeMap, by_key_hash: BTreeMap, touch_counts: BTreeMap, + snapshot_lookup_counts: BTreeMap, + key_hash_lookup_counts: BTreeMap, } #[derive(Debug, Default)] pub struct InMemoryAuthApiKeySnapshotRepository { index: RwLock, + lookup_delay: Option, } impl InMemoryAuthApiKeySnapshotRepository { @@ -99,10 +102,18 @@ impl InMemoryAuthApiKeySnapshotRepository { export_by_api_key_id, by_key_hash, touch_counts: BTreeMap::new(), + snapshot_lookup_counts: BTreeMap::new(), + key_hash_lookup_counts: BTreeMap::new(), }), + lookup_delay: None, } } + pub fn with_lookup_delay_for_tests(mut self, delay: Duration) -> Self { + self.lookup_delay = Some(delay); + self + } + pub fn with_export_records(mut self, items: I) -> Self where I: IntoIterator, @@ -129,6 +140,26 @@ impl InMemoryAuthApiKeySnapshotRepository { .unwrap_or(0) } + pub fn snapshot_lookup_count(&self, api_key_id: &str) -> usize { + self.index + .read() + .expect("auth api key snapshot repository lock") + .snapshot_lookup_counts + .get(api_key_id) + .copied() + .unwrap_or(0) + } + + pub fn key_hash_lookup_count(&self, key_hash: &str) -> usize { + self.index + .read() + .expect("auth api key snapshot repository lock") + .key_hash_lookup_counts + .get(key_hash) + .copied() + .unwrap_or(0) + } + pub(crate) fn apply_usage_stats_delta( &self, api_key_id: &str, @@ -200,27 +231,46 @@ impl AuthApiKeyReadRepository for InMemoryAuthApiKeySnapshotRepository { &self, key: AuthApiKeyLookupKey<'_>, ) -> Result, DataLayerError> { - let index = self + if let Some(delay) = self.lookup_delay { + tokio::time::sleep(delay).await; + } + let mut index = self .index - .read() + .write() .expect("auth api key snapshot repository lock"); Ok(match key { - AuthApiKeyLookupKey::KeyHash(key_hash) => index - .by_key_hash - .get(key_hash) - .and_then(|api_key_id| index.by_api_key_id.get(api_key_id)) - .cloned(), + AuthApiKeyLookupKey::KeyHash(key_hash) => { + *index + .key_hash_lookup_counts + .entry(key_hash.to_string()) + .or_insert(0) += 1; + index + .by_key_hash + .get(key_hash) + .and_then(|api_key_id| index.by_api_key_id.get(api_key_id)) + .cloned() + } AuthApiKeyLookupKey::ApiKeyId(api_key_id) => { + *index + .snapshot_lookup_counts + .entry(api_key_id.to_string()) + .or_insert(0) += 1; index.by_api_key_id.get(api_key_id).cloned() } AuthApiKeyLookupKey::UserApiKeyIds { user_id, api_key_id, - } => index - .by_api_key_id - .get(api_key_id) - .filter(|snapshot| snapshot.user_id == user_id) - .cloned(), + } => { + *index + .snapshot_lookup_counts + .entry(api_key_id.to_string()) + .or_insert(0) += 1; + index + .by_api_key_id + .get(api_key_id) + .filter(|snapshot| snapshot.user_id == user_id) + .cloned() + } }) } diff --git a/crates/aether-data/src/repository/candidates/postgres.rs b/crates/aether-data/src/repository/candidates/postgres.rs index e357e5021..79d88c87f 100644 --- a/crates/aether-data/src/repository/candidates/postgres.rs +++ b/crates/aether-data/src/repository/candidates/postgres.rs @@ -8,6 +8,7 @@ use super::{ RequestCandidateStatus, RequestCandidateWriteRepository, StoredRequestCandidate, UpsertRequestCandidateRecord, }; +use crate::driver::postgres::PostgresTransaction; use crate::driver::postgres::PostgresTransactionRunner; use crate::{error::SqlxResultExt, DataLayerError}; use aether_data_query::{push_eq, push_in, push_limit, WhereClause}; @@ -182,6 +183,109 @@ RETURNING CAST(EXTRACT(EPOCH FROM finished_at) * 1000 AS BIGINT) AS finished_at_unix_ms "#; +const UPSERT_CONFLICT_SQL: &str = r#" +ON CONFLICT (request_id, candidate_index, retry_index) +DO UPDATE SET + user_id = COALESCE(EXCLUDED.user_id, request_candidates.user_id), + api_key_id = COALESCE(EXCLUDED.api_key_id, request_candidates.api_key_id), + username = COALESCE(EXCLUDED.username, request_candidates.username), + api_key_name = COALESCE(EXCLUDED.api_key_name, request_candidates.api_key_name), + provider_id = COALESCE(EXCLUDED.provider_id, request_candidates.provider_id), + endpoint_id = COALESCE(EXCLUDED.endpoint_id, request_candidates.endpoint_id), + key_id = COALESCE(EXCLUDED.key_id, request_candidates.key_id), + status = EXCLUDED.status, + skip_reason = COALESCE(EXCLUDED.skip_reason, request_candidates.skip_reason), + is_cached = COALESCE(EXCLUDED.is_cached, request_candidates.is_cached), + status_code = COALESCE(EXCLUDED.status_code, request_candidates.status_code), + error_type = COALESCE(EXCLUDED.error_type, request_candidates.error_type), + error_message = COALESCE(EXCLUDED.error_message, request_candidates.error_message), + latency_ms = COALESCE(EXCLUDED.latency_ms, request_candidates.latency_ms), + 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 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 + ELSE EXCLUDED.extra_data + END, + required_capabilities = COALESCE(EXCLUDED.required_capabilities, request_candidates.required_capabilities), + created_at = CASE + WHEN request_candidates.created_at <= TO_TIMESTAMP(1) + THEN EXCLUDED.created_at + ELSE request_candidates.created_at + END, + started_at = COALESCE(EXCLUDED.started_at, request_candidates.started_at), + finished_at = COALESCE(EXCLUDED.finished_at, request_candidates.finished_at) +"#; + +const UPSERT_CONFLICT_INHERIT_IS_CACHED_SQL: &str = r#" +ON CONFLICT (request_id, candidate_index, retry_index) +DO UPDATE SET + user_id = COALESCE(EXCLUDED.user_id, request_candidates.user_id), + api_key_id = COALESCE(EXCLUDED.api_key_id, request_candidates.api_key_id), + username = COALESCE(EXCLUDED.username, request_candidates.username), + api_key_name = COALESCE(EXCLUDED.api_key_name, request_candidates.api_key_name), + provider_id = COALESCE(EXCLUDED.provider_id, request_candidates.provider_id), + endpoint_id = COALESCE(EXCLUDED.endpoint_id, request_candidates.endpoint_id), + key_id = COALESCE(EXCLUDED.key_id, request_candidates.key_id), + status = EXCLUDED.status, + skip_reason = COALESCE(EXCLUDED.skip_reason, request_candidates.skip_reason), + is_cached = request_candidates.is_cached, + status_code = COALESCE(EXCLUDED.status_code, request_candidates.status_code), + error_type = COALESCE(EXCLUDED.error_type, request_candidates.error_type), + error_message = COALESCE(EXCLUDED.error_message, request_candidates.error_message), + latency_ms = COALESCE(EXCLUDED.latency_ms, request_candidates.latency_ms), + 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 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 + ELSE EXCLUDED.extra_data + END, + required_capabilities = COALESCE(EXCLUDED.required_capabilities, request_candidates.required_capabilities), + created_at = CASE + WHEN request_candidates.created_at <= TO_TIMESTAMP(1) + THEN EXCLUDED.created_at + ELSE request_candidates.created_at + END, + started_at = COALESCE(EXCLUDED.started_at, request_candidates.started_at), + finished_at = COALESCE(EXCLUDED.finished_at, request_candidates.finished_at) +"#; + +const UPSERT_MANY_PREFIX_SQL: &str = r#" +INSERT INTO request_candidates ( + id, + request_id, + user_id, + api_key_id, + username, + api_key_name, + candidate_index, + retry_index, + provider_id, + endpoint_id, + key_id, + status, + skip_reason, + is_cached, + status_code, + error_type, + error_message, + latency_ms, + concurrent_requests, + extra_data, + required_capabilities, + created_at, + started_at, + finished_at +) +"#; + +const MAX_POSTGRES_REQUEST_CANDIDATE_UPSERT_ROWS: usize = 1_000; + const DELETE_CREATED_BEFORE_SQL: &str = r#" DELETE FROM request_candidates WHERE id IN ( @@ -501,6 +605,44 @@ impl SqlxRequestCandidateReadRepository { .await } + pub async fn upsert_many( + &self, + candidates: Vec, + ) -> Result { + if candidates.is_empty() { + return Ok(0); + } + let rows = candidates + .into_iter() + .map(BatchUpsertRequestCandidateRow::try_from) + .collect::, _>>()?; + + self.tx_runner + .run_read_write(|tx| { + Box::pin(async move { + let mut persisted = 0usize; + for ordered_batch in split_request_candidate_upsert_batches(rows) { + let (explicit_is_cached, inherited_is_cached): (Vec<_>, Vec<_>) = + ordered_batch + .into_iter() + .partition(|row| row.is_cached.is_some()); + + persisted = persisted.saturating_add( + execute_partitioned_upsert_many_batch(tx, &explicit_is_cached, true) + .await?, + ); + persisted = persisted.saturating_add( + execute_partitioned_upsert_many_batch(tx, &inherited_is_cached, false) + .await?, + ); + } + + Ok(persisted) + }) as BoxFuture<'_, Result> + }) + .await + } + pub async fn delete_created_before( &self, created_before_unix_secs: u64, @@ -524,6 +666,178 @@ impl SqlxRequestCandidateReadRepository { } } +async fn execute_partitioned_upsert_many_batch( + tx: &mut PostgresTransaction, + rows: &[BatchUpsertRequestCandidateRow], + overwrite_is_cached: bool, +) -> Result { + let mut persisted = 0usize; + for chunk in rows.chunks(MAX_POSTGRES_REQUEST_CANDIDATE_UPSERT_ROWS) { + persisted = persisted + .saturating_add(execute_upsert_many_batch(tx, chunk, overwrite_is_cached).await?); + } + Ok(persisted) +} + +#[derive(Debug)] +struct BatchUpsertRequestCandidateRow { + id: String, + request_id: String, + user_id: Option, + api_key_id: Option, + username: Option, + api_key_name: Option, + candidate_index: i32, + retry_index: i32, + provider_id: Option, + endpoint_id: Option, + key_id: Option, + status: &'static str, + skip_reason: Option, + is_cached: Option, + status_code: Option, + error_type: Option, + error_message: Option, + latency_ms: Option, + concurrent_requests: Option, + extra_data: Option, + required_capabilities: Option, + created_at_unix_ms: Option, + started_at_unix_ms: Option, + finished_at_unix_ms: Option, +} + +impl TryFrom for BatchUpsertRequestCandidateRow { + type Error = DataLayerError; + + fn try_from(candidate: UpsertRequestCandidateRecord) -> Result { + candidate.validate()?; + Ok(Self { + id: if candidate.id.trim().is_empty() { + Uuid::new_v4().to_string() + } else { + candidate.id + }, + request_id: candidate.request_id, + user_id: candidate.user_id, + api_key_id: candidate.api_key_id, + username: candidate.username, + api_key_name: candidate.api_key_name, + candidate_index: to_i32(candidate.candidate_index)?, + retry_index: to_i32(candidate.retry_index)?, + provider_id: candidate.provider_id, + endpoint_id: candidate.endpoint_id, + key_id: candidate.key_id, + status: status_to_database(candidate.status), + skip_reason: candidate.skip_reason, + is_cached: candidate.is_cached, + status_code: candidate.status_code.map(i32::from), + error_type: candidate.error_type, + error_message: candidate.error_message, + latency_ms: candidate.latency_ms.map(to_i32_u64).transpose()?, + concurrent_requests: candidate.concurrent_requests.map(to_i32).transpose()?, + extra_data: candidate.extra_data, + required_capabilities: candidate.required_capabilities, + created_at_unix_ms: candidate.created_at_unix_ms.map(|value| value as f64), + started_at_unix_ms: candidate.started_at_unix_ms.map(|value| value as f64), + finished_at_unix_ms: candidate.finished_at_unix_ms.map(|value| value as f64), + }) + } +} + +async fn execute_upsert_many_batch( + tx: &mut PostgresTransaction, + rows: &[BatchUpsertRequestCandidateRow], + overwrite_is_cached: bool, +) -> Result { + if rows.is_empty() { + return Ok(0); + } + + let mut builder = QueryBuilder::::new(UPSERT_MANY_PREFIX_SQL); + builder.push_values(rows, |mut values, row| { + values + .push_bind(row.id.clone()) + .push_bind(row.request_id.clone()) + .push_bind(row.user_id.clone()) + .push_bind(row.api_key_id.clone()) + .push_bind(row.username.clone()) + .push_bind(row.api_key_name.clone()) + .push_bind(row.candidate_index) + .push_bind(row.retry_index) + .push_bind(row.provider_id.clone()) + .push_bind(row.endpoint_id.clone()) + .push_bind(row.key_id.clone()) + .push_bind(row.status) + .push_bind(row.skip_reason.clone()) + .push_bind(row.is_cached.unwrap_or(false)) + .push_bind(row.status_code) + .push_bind(row.error_type.clone()) + .push_bind(row.error_message.clone()) + .push_bind(row.latency_ms) + .push_bind(row.concurrent_requests) + .push_bind(row.extra_data.clone()) + .push_bind(row.required_capabilities.clone()) + .push("COALESCE(CASE WHEN ") + .push_bind_unseparated(row.created_at_unix_ms) + .push_unseparated(" IS NOT NULL AND ") + .push_bind_unseparated(row.created_at_unix_ms) + .push_unseparated(" > 1000.0 THEN TO_TIMESTAMP(") + .push_bind_unseparated(row.created_at_unix_ms) + .push_unseparated(" / 1000.0) END, TO_TIMESTAMP(") + .push_bind_unseparated(row.started_at_unix_ms) + .push_unseparated(" / 1000.0), TO_TIMESTAMP(") + .push_bind_unseparated(row.finished_at_unix_ms) + .push_unseparated(" / 1000.0), NOW())") + .push("TO_TIMESTAMP(") + .push_bind_unseparated(row.started_at_unix_ms) + .push_unseparated(" / 1000.0)") + .push("TO_TIMESTAMP(") + .push_bind_unseparated(row.finished_at_unix_ms) + .push_unseparated(" / 1000.0)"); + }); + builder.push(upsert_many_conflict_sql(overwrite_is_cached)); + let result = builder + .build() + .execute(&mut **tx) + .await + .map_postgres_err()?; + Ok(usize::try_from(result.rows_affected()).unwrap_or(rows.len())) +} + +fn split_request_candidate_upsert_batches( + rows: Vec, +) -> Vec> { + let mut batches = Vec::new(); + let mut current = Vec::new(); + let mut seen = std::collections::HashSet::<(String, i32, i32)>::new(); + + for row in rows { + let key = (row.request_id.clone(), row.candidate_index, row.retry_index); + if seen.contains(&key) && !current.is_empty() { + batches.push(current); + current = Vec::new(); + seen.clear(); + } + seen.insert(key); + current.push(row); + } + + if !current.is_empty() { + batches.push(current); + } + + batches +} + +fn upsert_many_conflict_sql(overwrite_is_cached: bool) -> &'static str { + if overwrite_is_cached { + UPSERT_CONFLICT_SQL + } else { + UPSERT_CONFLICT_INHERIT_IS_CACHED_SQL + } +} + #[async_trait] impl RequestCandidateReadRepository for SqlxRequestCandidateReadRepository { async fn list_by_request_id( @@ -600,6 +914,13 @@ impl RequestCandidateWriteRepository for SqlxRequestCandidateReadRepository { Self::upsert(self, candidate).await } + async fn upsert_many( + &self, + candidates: Vec, + ) -> Result { + Self::upsert_many(self, candidates).await + } + async fn delete_created_before( &self, created_before_unix_secs: u64, diff --git a/crates/aether-provider-transport/src/network.rs b/crates/aether-provider-transport/src/network.rs index d2a4278fa..c2678df85 100644 --- a/crates/aether-provider-transport/src/network.rs +++ b/crates/aether-provider-transport/src/network.rs @@ -344,6 +344,7 @@ mod tests { transport_proxy_is_locally_supported, TransportTunnelAffinityLookup, TransportTunnelAttachmentOwner, }; + use aether_contracts::TRANSPORT_HTTP_MODE_H2C_PRIOR_KNOWLEDGE; #[derive(Default)] struct TestTunnelAffinityLookup { @@ -589,6 +590,23 @@ mod tests { assert_eq!(profile.pool_scope, "key"); } + #[test] + fn resolves_h2c_prior_knowledge_transport_profile() { + let mut transport = sample_transport(); + transport.key.fingerprint = Some(json!({ + "transport_profile": { + "profile_id": "mock-h2c", + "backend": "reqwest_rustls", + "http_mode": "h2c_prior_knowledge" + } + })); + + let profile = resolve_transport_profile(&transport).expect("profile"); + + assert_eq!(profile.profile_id, "mock-h2c"); + assert_eq!(profile.http_mode, TRANSPORT_HTTP_MODE_H2C_PRIOR_KNOWLEDGE); + } + #[test] fn resolves_no_transport_profile_without_fingerprint_configuration() { let mut transport = sample_transport(); diff --git a/crates/aether-runtime-state/src/redis/client.rs b/crates/aether-runtime-state/src/redis/client.rs index f69da5cf9..2f51044cd 100644 --- a/crates/aether-runtime-state/src/redis/client.rs +++ b/crates/aether-runtime-state/src/redis/client.rs @@ -106,24 +106,28 @@ impl RedisConnectionRouter { &client, connection_manager_config(command_timeout_ms), RedisConnectionLane::Fast, + command_timeout_ms, ) .await?; let stream = connect_lane( &client, connection_manager_config(command_timeout_ms), RedisConnectionLane::Stream, + command_timeout_ms, ) .await?; let blocking_stream = connect_lane( &client, connection_manager_config(command_timeout_ms), RedisConnectionLane::BlockingStream, + command_timeout_ms, ) .await?; let admin = connect_lane( &client, connection_manager_config(command_timeout_ms), RedisConnectionLane::Admin, + command_timeout_ms, ) .await?; info!( @@ -228,16 +232,29 @@ async fn connect_lane( client: &RedisClient, config: redis::aio::ConnectionManagerConfig, lane: RedisConnectionLane, + command_timeout_ms: Option, ) -> Result { - client - .get_connection_manager_with_config(config) - .await - .map_err(|err| { - DataLayerError::Redis(format!( - "failed to initialize runtime redis {} lane: {err}", - lane.as_str() - )) - }) + let connect = client.get_connection_manager_with_config(config); + let result = if let Some(timeout_ms) = command_timeout_ms { + match tokio::time::timeout(Duration::from_millis(timeout_ms), connect).await { + Ok(result) => result, + Err(_) => { + return Err(DataLayerError::TimedOut(format!( + "runtime redis {} lane connection exceeded {}ms timeout", + lane.as_str(), + timeout_ms + ))); + } + } + } else { + connect.await + }; + result.map_err(|err| { + DataLayerError::Redis(format!( + "failed to initialize runtime redis {} lane: {err}", + lane.as_str() + )) + }) } #[cfg(test)] diff --git a/crates/aether-testkit/src/bin/gateway_pressure_probe.rs b/crates/aether-testkit/src/bin/gateway_pressure_probe.rs index 845ce0d5e..594cf1c77 100644 --- a/crates/aether-testkit/src/bin/gateway_pressure_probe.rs +++ b/crates/aether-testkit/src/bin/gateway_pressure_probe.rs @@ -176,25 +176,56 @@ fn spawn_metrics_sampler( fn parse_args(args: Vec) -> Result> { let mut target_url: Option = None; + let mut warmup_url: Option = None; let mut metrics_url: Option = None; let mut total_requests: Option = None; let mut concurrency: Option = None; + let mut warmup_connections: usize = 0; let mut timeout_ms: Option = None; + let mut connect_timeout_ms: Option = None; + let mut client_shards: Option = None; + let mut pool_max_idle_per_host: Option = None; + let mut start_ramp_ms: u64 = 0; + let mut first_body_hold_ms: u64 = 0; let mut sample_interval_ms: u64 = 500; let mut method = Method::GET; let mut headers = BTreeMap::new(); let mut body: Option> = None; let mut response_mode = HttpLoadProbeResponseMode::HeadersOnly; + let mut http1_only = false; + let mut http2_prior_knowledge = false; let mut output_path = None; let mut iter = args.into_iter(); while let Some(arg) = iter.next() { match arg.as_str() { "--url" => target_url = Some(next_value(&mut iter, "--url")?), + "--warmup-url" => warmup_url = Some(next_value(&mut iter, "--warmup-url")?), "--metrics-url" => metrics_url = Some(next_value(&mut iter, "--metrics-url")?), "--requests" => total_requests = Some(next_value(&mut iter, "--requests")?.parse()?), "--concurrency" => concurrency = Some(next_value(&mut iter, "--concurrency")?.parse()?), + "--warmup-connections" => { + warmup_connections = next_value(&mut iter, "--warmup-connections")?.parse()? + } "--timeout-ms" => timeout_ms = Some(next_value(&mut iter, "--timeout-ms")?.parse()?), + "--connect-timeout-ms" => { + connect_timeout_ms = Some(next_value(&mut iter, "--connect-timeout-ms")?.parse()?) + } + "--client-shards" => { + client_shards = Some(next_value(&mut iter, "--client-shards")?.parse()?) + } + "--pool-max-idle-per-host" => { + pool_max_idle_per_host = + Some(next_value(&mut iter, "--pool-max-idle-per-host")?.parse()?) + } + "--start-ramp-ms" => { + start_ramp_ms = next_value(&mut iter, "--start-ramp-ms")?.parse()? + } + "--first-body-hold-ms" => { + first_body_hold_ms = next_value(&mut iter, "--first-body-hold-ms")?.parse()? + } + "--http1-only" => http1_only = true, + "--http2-prior-knowledge" => http2_prior_knowledge = true, "--sample-interval-ms" => { sample_interval_ms = next_value(&mut iter, "--sample-interval-ms")?.parse()? } @@ -247,9 +278,20 @@ fn parse_args(args: Vec) -> Result> { response_mode, ..HttpLoadProbeConfig::default() }; + load.warmup_url = warmup_url; + load.warmup_connections = warmup_connections; if let Some(timeout_ms) = timeout_ms { load.timeout = Duration::from_millis(timeout_ms); } + load.connect_timeout = connect_timeout_ms.map(Duration::from_millis); + if let Some(client_shards) = client_shards { + load.client_shards = client_shards; + } + load.pool_max_idle_per_host = pool_max_idle_per_host; + load.start_ramp = Duration::from_millis(start_ramp_ms); + load.first_body_hold = Duration::from_millis(first_body_hold_ms); + load.http1_only = http1_only; + load.http2_prior_knowledge = http2_prior_knowledge; load.validate() .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidInput, err))?; if sample_interval_ms == 0 { @@ -298,10 +340,15 @@ fn parse_response_mode( ) -> Result> { match value.trim().to_ascii_lowercase().as_str() { "headers" | "headers-only" | "header" => Ok(HttpLoadProbeResponseMode::HeadersOnly), + "first-body-byte" | "first-body" | "first-byte" | "first-chunk" => { + Ok(HttpLoadProbeResponseMode::FirstBodyByte) + } "full" | "full-body" | "body" => Ok(HttpLoadProbeResponseMode::FullBody), other => Err(std::io::Error::new( std::io::ErrorKind::InvalidInput, - format!("unsupported --response-mode {other}; expected headers or full"), + format!( + "unsupported --response-mode {other}; expected headers, first-body-byte, or full" + ), ) .into()), } @@ -348,6 +395,6 @@ fn metric_name_matches(actual: &str, expected: &str) -> bool { fn print_usage() { eprintln!( - "usage: cargo run -p aether-testkit --bin gateway_pressure_probe -- --url --metrics-url --requests --concurrency [--method GET] [--timeout-ms 30000] [--sample-interval-ms 500] [-H 'Name: value'] [--body JSON | --body-file path] [--response-mode headers|full] [--output /tmp/gateway_pressure.json]" + "usage: cargo run -p aether-testkit --bin gateway_pressure_probe -- --url --metrics-url --requests --concurrency [--warmup-url ] [--warmup-connections N] [--method GET] [--timeout-ms 30000] [--connect-timeout-ms 10000] [--client-shards 1] [--pool-max-idle-per-host N] [--start-ramp-ms 0] [--first-body-hold-ms 0] [--http1-only | --http2-prior-knowledge] [--sample-interval-ms 500] [-H 'Name: value'] [--body JSON | --body-file path] [--response-mode headers|first-body-byte|full] [--output /tmp/gateway_pressure.json]" ); } diff --git a/crates/aether-testkit/src/bin/http_load_probe.rs b/crates/aether-testkit/src/bin/http_load_probe.rs index ff74b954c..105efbf2d 100644 --- a/crates/aether-testkit/src/bin/http_load_probe.rs +++ b/crates/aether-testkit/src/bin/http_load_probe.rs @@ -15,21 +15,52 @@ async fn main() -> Result<(), Box> { fn parse_args(args: Vec) -> Result> { let mut url: Option = None; + let mut warmup_url: Option = None; let mut total_requests: Option = None; let mut concurrency: Option = None; + let mut warmup_connections: usize = 0; let mut timeout_ms: Option = None; + let mut connect_timeout_ms: Option = None; + let mut client_shards: Option = None; + let mut pool_max_idle_per_host: Option = None; + let mut start_ramp_ms: u64 = 0; + let mut first_body_hold_ms: u64 = 0; let mut method = Method::GET; let mut headers = std::collections::BTreeMap::new(); let mut body: Option> = None; let mut response_mode = aether_testkit::HttpLoadProbeResponseMode::HeadersOnly; + let mut http1_only = false; + let mut http2_prior_knowledge = false; let mut iter = args.into_iter(); while let Some(arg) = iter.next() { match arg.as_str() { "--url" => url = Some(next_value(&mut iter, "--url")?), + "--warmup-url" => warmup_url = Some(next_value(&mut iter, "--warmup-url")?), "--requests" => total_requests = Some(next_value(&mut iter, "--requests")?.parse()?), "--concurrency" => concurrency = Some(next_value(&mut iter, "--concurrency")?.parse()?), + "--warmup-connections" => { + warmup_connections = next_value(&mut iter, "--warmup-connections")?.parse()? + } "--timeout-ms" => timeout_ms = Some(next_value(&mut iter, "--timeout-ms")?.parse()?), + "--connect-timeout-ms" => { + connect_timeout_ms = Some(next_value(&mut iter, "--connect-timeout-ms")?.parse()?) + } + "--client-shards" => { + client_shards = Some(next_value(&mut iter, "--client-shards")?.parse()?) + } + "--pool-max-idle-per-host" => { + pool_max_idle_per_host = + Some(next_value(&mut iter, "--pool-max-idle-per-host")?.parse()?) + } + "--start-ramp-ms" => { + start_ramp_ms = next_value(&mut iter, "--start-ramp-ms")?.parse()? + } + "--first-body-hold-ms" => { + first_body_hold_ms = next_value(&mut iter, "--first-body-hold-ms")?.parse()? + } + "--http1-only" => http1_only = true, + "--http2-prior-knowledge" => http2_prior_knowledge = true, "--method" => { method = Method::from_bytes(next_value(&mut iter, "--method")?.as_bytes())? } @@ -78,9 +109,20 @@ fn parse_args(args: Vec) -> Result { Ok(aether_testkit::HttpLoadProbeResponseMode::HeadersOnly) } + "first-body-byte" | "first-body" | "first-byte" | "first-chunk" => { + Ok(aether_testkit::HttpLoadProbeResponseMode::FirstBodyByte) + } "full" | "full-body" | "body" => Ok(aether_testkit::HttpLoadProbeResponseMode::FullBody), other => Err(std::io::Error::new( std::io::ErrorKind::InvalidInput, - format!("unsupported --response-mode {other}; expected headers or full"), + format!( + "unsupported --response-mode {other}; expected headers, first-body-byte, or full" + ), ) .into()), } @@ -139,6 +186,6 @@ fn next_value( fn print_usage() { eprintln!( - "usage: cargo run -p aether-testkit --bin http_load_probe -- --url --requests --concurrency [--method GET] [--timeout-ms 30000] [-H 'Name: value'] [--body JSON | --body-file path] [--response-mode headers|full]" + "usage: cargo run -p aether-testkit --bin http_load_probe -- --url --requests --concurrency [--warmup-url ] [--warmup-connections N] [--method GET] [--timeout-ms 30000] [--connect-timeout-ms 10000] [--client-shards 1] [--pool-max-idle-per-host N] [--start-ramp-ms 0] [--first-body-hold-ms 0] [--http1-only | --http2-prior-knowledge] [-H 'Name: value'] [--body JSON | --body-file path] [--response-mode headers|first-body-byte|full]" ); } diff --git a/crates/aether-testkit/src/bin/mock_openai_upstream.rs b/crates/aether-testkit/src/bin/mock_openai_upstream.rs index 8e0d1a4bf..f79019919 100644 --- a/crates/aether-testkit/src/bin/mock_openai_upstream.rs +++ b/crates/aether-testkit/src/bin/mock_openai_upstream.rs @@ -14,7 +14,7 @@ use serde_json::json; #[derive(Debug, Clone)] struct Config { - bind: SocketAddr, + binds: Vec, chunks: u64, first_byte_delay: Duration, chunk_delay: Duration, @@ -25,9 +25,9 @@ struct Config { impl Default for Config { fn default() -> Self { Self { - bind: "127.0.0.1:18181" + binds: vec!["127.0.0.1:18181" .parse() - .expect("default bind address should parse"), + .expect("default bind address should parse")], chunks: 8, first_byte_delay: Duration::from_millis(0), chunk_delay: Duration::from_millis(20), @@ -68,9 +68,40 @@ async fn main() -> Result<(), Box> { .route("/responses", post(responses)) .with_state(app_state); - let listener = tokio::net::TcpListener::bind(config.bind).await?; - eprintln!("mock OpenAI upstream listening on http://{}", config.bind); - axum::serve(listener, app).await?; + serve_listeners(&config.binds, app).await?; + Ok(()) +} + +async fn serve_listeners( + binds: &[SocketAddr], + app: Router, +) -> Result<(), Box> { + let mut listeners = Vec::with_capacity(binds.len()); + for bind in binds { + listeners.push((*bind, tokio::net::TcpListener::bind(bind).await?)); + } + if listeners.len() == 1 { + let (bind, listener) = listeners + .into_iter() + .next() + .ok_or_else(|| std::io::Error::other("mock upstream listener set is empty"))?; + eprintln!("mock OpenAI upstream listening on http://{bind}"); + axum::serve(listener, app).await?; + return Ok(()); + } + + let mut servers = tokio::task::JoinSet::new(); + for (bind, listener) in listeners { + let app = app.clone(); + eprintln!("mock OpenAI upstream listening on http://{bind}"); + servers.spawn(async move { axum::serve(listener, app).await }); + } + if let Some(result) = servers.join_next().await { + servers.abort_all(); + let serve_result = result + .map_err(|err| std::io::Error::other(format!("mock listener task failed: {err}")))?; + serve_result?; + } Ok(()) } @@ -292,10 +323,17 @@ fn current_unix_secs() -> u64 { fn parse_args(args: Vec) -> Result> { let mut config = Config::default(); + let mut binds_overridden = false; let mut iter = args.into_iter(); while let Some(arg) = iter.next() { match arg.as_str() { - "--bind" => config.bind = next_value(&mut iter, "--bind")?.parse()?, + "--bind" => { + if !binds_overridden { + config.binds.clear(); + binds_overridden = true; + } + config.binds.push(next_value(&mut iter, "--bind")?.parse()?); + } "--chunks" => config.chunks = next_value(&mut iter, "--chunks")?.parse()?, "--first-byte-delay-ms" => { config.first_byte_delay = @@ -325,6 +363,13 @@ fn parse_args(args: Vec) -> Result> { } } } + if config.binds.is_empty() { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "at least one --bind is required", + ) + .into()); + } Ok(config) } @@ -343,6 +388,6 @@ fn next_value( fn print_usage() { eprintln!( - "usage: cargo run -p aether-testkit --bin mock_openai_upstream -- [--bind 127.0.0.1:18181] [--chunks 8] [--first-byte-delay-ms 0] [--chunk-delay-ms 20] [--payload-bytes 32] [--status 200]" + "usage: cargo run -p aether-testkit --bin mock_openai_upstream -- [--bind 127.0.0.1:18181]... [--chunks 8] [--first-byte-delay-ms 0] [--chunk-delay-ms 20] [--payload-bytes 32] [--status 200]" ); } diff --git a/crates/aether-testkit/src/load.rs b/crates/aether-testkit/src/load.rs index 20a5a50cb..e9ee34c26 100644 --- a/crates/aether-testkit/src/load.rs +++ b/crates/aether-testkit/src/load.rs @@ -1,4 +1,5 @@ use std::collections::BTreeMap; +use std::error::Error as StdError; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -9,36 +10,57 @@ use tokio::sync::Mutex; use crate::runtime::{BenchmarkRuntimeSampler, BenchmarkRuntimeSnapshot}; +const MAX_ERROR_SAMPLES: usize = 32; + #[derive(Debug, Clone, Copy, Default, serde::Serialize, PartialEq, Eq)] pub enum HttpLoadProbeResponseMode { #[default] HeadersOnly, + FirstBodyByte, FullBody, } #[derive(Debug, Clone)] pub struct HttpLoadProbeConfig { pub url: String, + pub warmup_url: Option, pub method: Method, pub headers: BTreeMap, pub body: Option>, pub total_requests: usize, pub concurrency: usize, + pub warmup_connections: usize, pub timeout: Duration, + pub connect_timeout: Option, pub response_mode: HttpLoadProbeResponseMode, + pub client_shards: usize, + pub pool_max_idle_per_host: Option, + pub start_ramp: Duration, + pub http1_only: bool, + pub http2_prior_knowledge: bool, + pub first_body_hold: Duration, } impl Default for HttpLoadProbeConfig { fn default() -> Self { Self { url: String::new(), + warmup_url: None, method: Method::GET, headers: BTreeMap::new(), body: None, total_requests: 100, concurrency: 10, + warmup_connections: 0, timeout: Duration::from_secs(30), + connect_timeout: None, response_mode: HttpLoadProbeResponseMode::HeadersOnly, + client_shards: 1, + pool_max_idle_per_host: None, + start_ramp: Duration::ZERO, + http1_only: false, + http2_prior_knowledge: false, + first_body_hold: Duration::ZERO, } } } @@ -57,10 +79,33 @@ impl HttpLoadProbeConfig { if self.timeout.is_zero() { return Err("load probe timeout must be positive".to_string()); } + if matches!(self.connect_timeout, Some(timeout) if timeout.is_zero()) { + return Err("load probe connect_timeout must be positive when set".to_string()); + } + if self.client_shards == 0 { + return Err("load probe client_shards must be positive".to_string()); + } + if self.http1_only && self.http2_prior_knowledge { + return Err( + "load probe cannot enable both http1_only and http2_prior_knowledge".to_string(), + ); + } Ok(()) } } +#[derive(Debug, Clone, serde::Serialize, PartialEq, Eq)] +pub struct HttpLoadProbeErrorSample { + pub request_index: usize, + pub url: String, + pub phase: String, + pub kind: String, + pub elapsed_ms: u64, + pub message: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub source: Option, +} + #[derive(Debug, Clone, serde::Serialize, PartialEq, Eq)] pub struct HttpLoadProbeResult { pub url: String, @@ -68,6 +113,18 @@ pub struct HttpLoadProbeResult { pub response_mode: HttpLoadProbeResponseMode, pub total_requests: usize, pub concurrency: usize, + pub warmup_connections: usize, + #[serde(skip_serializing_if = "Option::is_none")] + pub warmup_url: Option, + pub client_shards: usize, + pub start_ramp_ms: u64, + #[serde(skip_serializing_if = "Option::is_none")] + pub pool_max_idle_per_host: Option, + pub http1_only: bool, + pub http2_prior_knowledge: bool, + #[serde(skip_serializing_if = "Option::is_none")] + pub connect_timeout_ms: Option, + pub first_body_hold_ms: u64, pub duration_ms: u64, pub throughput_rps: u64, pub p99_ms: u64, @@ -77,9 +134,23 @@ pub struct HttpLoadProbeResult { pub p95_ms: u64, pub max_ms: u64, pub mean_ms: u64, + #[serde(skip_serializing_if = "Option::is_none")] + pub headers_p50_ms: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub headers_p95_ms: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub headers_p99_ms: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub first_body_p50_ms: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub first_body_p95_ms: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub first_body_p99_ms: Option, pub runtime: BenchmarkRuntimeSnapshot, pub status_counts: BTreeMap, pub error_counts: BTreeMap, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub error_samples: Vec, } #[derive(Debug, Clone, serde::Serialize, PartialEq, Eq)] @@ -90,6 +161,18 @@ pub struct MultiUrlHttpLoadProbeResult { pub response_mode: HttpLoadProbeResponseMode, pub total_requests: usize, pub concurrency: usize, + pub warmup_connections: usize, + #[serde(skip_serializing_if = "Option::is_none")] + pub warmup_url: Option, + pub client_shards: usize, + pub start_ramp_ms: u64, + #[serde(skip_serializing_if = "Option::is_none")] + pub pool_max_idle_per_host: Option, + pub http1_only: bool, + pub http2_prior_knowledge: bool, + #[serde(skip_serializing_if = "Option::is_none")] + pub connect_timeout_ms: Option, + pub first_body_hold_ms: u64, pub duration_ms: u64, pub throughput_rps: u64, pub p99_ms: u64, @@ -99,9 +182,23 @@ pub struct MultiUrlHttpLoadProbeResult { pub p95_ms: u64, pub max_ms: u64, pub mean_ms: u64, + #[serde(skip_serializing_if = "Option::is_none")] + pub headers_p50_ms: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub headers_p95_ms: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub headers_p99_ms: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub first_body_p50_ms: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub first_body_p95_ms: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub first_body_p99_ms: Option, pub runtime: BenchmarkRuntimeSnapshot, pub status_counts: BTreeMap, pub error_counts: BTreeMap, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub error_samples: Vec, } pub async fn run_http_load_probe( @@ -120,6 +217,15 @@ pub async fn run_http_load_probe( response_mode: result.response_mode, total_requests: result.total_requests, concurrency: result.concurrency, + warmup_connections: result.warmup_connections, + warmup_url: result.warmup_url, + client_shards: result.client_shards, + start_ramp_ms: result.start_ramp_ms, + pool_max_idle_per_host: result.pool_max_idle_per_host, + http1_only: result.http1_only, + http2_prior_knowledge: result.http2_prior_knowledge, + connect_timeout_ms: result.connect_timeout_ms, + first_body_hold_ms: result.first_body_hold_ms, duration_ms: result.duration_ms, throughput_rps: result.throughput_rps, p99_ms: result.p99_ms, @@ -129,9 +235,16 @@ pub async fn run_http_load_probe( p95_ms: result.p95_ms, max_ms: result.max_ms, mean_ms: result.mean_ms, + headers_p50_ms: result.headers_p50_ms, + headers_p95_ms: result.headers_p95_ms, + headers_p99_ms: result.headers_p99_ms, + first_body_p50_ms: result.first_body_p50_ms, + first_body_p95_ms: result.first_body_p95_ms, + first_body_p99_ms: result.first_body_p99_ms, runtime: result.runtime, status_counts: result.status_counts, error_counts: result.error_counts, + error_samples: result.error_samples, }) } @@ -150,32 +263,38 @@ async fn run_http_load_probe_against_urls( config: &HttpLoadProbeConfig, urls: &[String], ) -> Result { - let client = Client::builder() - .timeout(config.timeout) - .build() - .map_err(|err| format!("failed to build load probe http client: {err}"))?; + let clients = Arc::new(build_probe_clients(config)?); let total_requests = config.total_requests; let request_headers = build_headers(&config.headers)?; let request_body = config.body.clone().map(Arc::new); let response_mode = config.response_mode; + let first_body_hold = config.first_body_hold; + let start_ramp = config.start_ramp; + warmup_probe_connections(config, Arc::clone(&clients)).await?; let mut runtime_sampler = BenchmarkRuntimeSampler::new(); let started_at = Instant::now(); let next_request = Arc::new(AtomicUsize::new(0)); let latencies_ms = Arc::new(Mutex::new(Vec::with_capacity(config.total_requests))); + let header_latencies_ms = Arc::new(Mutex::new(Vec::with_capacity(config.total_requests))); + let first_body_latencies_ms = Arc::new(Mutex::new(Vec::with_capacity(config.total_requests))); let status_counts = Arc::new(Mutex::new(BTreeMap::::new())); let error_counts = Arc::new(Mutex::new(BTreeMap::::new())); + let error_samples = Arc::new(Mutex::new(Vec::::new())); let target_request_counts = Arc::new(Mutex::new(BTreeMap::::new())); let failed_requests = Arc::new(AtomicUsize::new(0)); let completed_requests = Arc::new(AtomicUsize::new(0)); let mut workers = tokio::task::JoinSet::new(); - for _ in 0..config.concurrency { - let client = client.clone(); + for worker_index in 0..config.concurrency { + let client = clients[worker_index % clients.len()].clone(); let next_request = Arc::clone(&next_request); let latencies_ms = Arc::clone(&latencies_ms); + let header_latencies_ms = Arc::clone(&header_latencies_ms); + let first_body_latencies_ms = Arc::clone(&first_body_latencies_ms); let status_counts = Arc::clone(&status_counts); let error_counts = Arc::clone(&error_counts); + let error_samples = Arc::clone(&error_samples); let target_request_counts = Arc::clone(&target_request_counts); let failed_requests = Arc::clone(&failed_requests); let completed_requests = Arc::clone(&completed_requests); @@ -183,8 +302,12 @@ async fn run_http_load_probe_against_urls( let urls = urls.to_vec(); let request_headers = request_headers.clone(); let request_body = request_body.clone(); + let start_delay = worker_start_delay(start_ramp, worker_index, config.concurrency); workers.spawn(async move { + if !start_delay.is_zero() { + tokio::time::sleep(start_delay).await; + } loop { let current = next_request.fetch_add(1, Ordering::AcqRel); if current >= total_requests { @@ -202,36 +325,62 @@ async fn run_http_load_probe_against_urls( } match request.send().await { Ok(response) => { + let headers_latency_ms = started_at.elapsed().as_millis() as u64; let status = response.status().as_u16(); - let body_result = match response_mode { - HttpLoadProbeResponseMode::HeadersOnly => Ok(()), - HttpLoadProbeResponseMode::FullBody => response - .bytes() - .await - .map(|_| ()) - .map_err(|err| classify_reqwest_error(&err)), - }; - if let Err(error_kind) = body_result { - failed_requests.fetch_add(1, Ordering::AcqRel); - let mut counts = error_counts.lock().await; - *counts.entry(error_kind).or_insert(0) += 1; - } else { - let mut counts = status_counts.lock().await; - *counts.entry(status).or_insert(0) += 1; - drop(counts); - let mut target_counts = target_request_counts.lock().await; - *target_counts.entry(url).or_insert(0) += 1; + let body_result = observe_response_body( + response, + response_mode, + started_at, + first_body_hold, + ) + .await; + match body_result { + Err(error) => { + failed_requests.fetch_add(1, Ordering::AcqRel); + record_load_error( + &error_counts, + &error_samples, + current, + &url, + started_at.elapsed().as_millis() as u64, + error, + ) + .await; + } + Ok(observation) => { + let mut counts = status_counts.lock().await; + *counts.entry(status).or_insert(0) += 1; + drop(counts); + let mut target_counts = target_request_counts.lock().await; + *target_counts.entry(url).or_insert(0) += 1; + if let Some(first_body_latency_ms) = + observation.first_body_latency_ms + { + first_body_latencies_ms + .lock() + .await + .push(first_body_latency_ms); + } + } } let latency_ms = started_at.elapsed().as_millis() as u64; latencies_ms.lock().await.push(latency_ms); + header_latencies_ms.lock().await.push(headers_latency_ms); completed_requests.fetch_add(1, Ordering::AcqRel); } Err(err) => { let latency_ms = started_at.elapsed().as_millis() as u64; latencies_ms.lock().await.push(latency_ms); failed_requests.fetch_add(1, Ordering::AcqRel); - let mut counts = error_counts.lock().await; - *counts.entry(classify_reqwest_error(&err)).or_insert(0) += 1; + record_load_error( + &error_counts, + &error_samples, + current, + &url, + latency_ms, + classify_reqwest_error("send", &err), + ) + .await; completed_requests.fetch_add(1, Ordering::AcqRel); } } @@ -245,10 +394,19 @@ async fn run_http_load_probe_against_urls( let status_counts = status_counts.lock().await.clone(); let error_counts = error_counts.lock().await.clone(); + let error_samples = error_samples.lock().await.clone(); let target_request_counts = target_request_counts.lock().await.clone(); let mut latencies = latencies_ms.lock().await.clone(); + let mut header_latencies = header_latencies_ms.lock().await.clone(); + let mut first_body_latencies = first_body_latencies_ms.lock().await.clone(); latencies.sort_unstable(); + header_latencies.sort_unstable(); + first_body_latencies.sort_unstable(); let (p50_ms, p95_ms, p99_ms, max_ms, mean_ms) = summarize_latencies(&latencies); + let (headers_p50_ms, headers_p95_ms, headers_p99_ms, _, _) = + summarize_latencies(&header_latencies); + let (first_body_p50_ms, first_body_p95_ms, first_body_p99_ms, _, _) = + summarize_latencies(&first_body_latencies); let duration_ms = started_at.elapsed().as_millis() as u64; let throughput_rps = if duration_ms == 0 { completed_requests.load(Ordering::Acquire) as u64 @@ -263,6 +421,17 @@ async fn run_http_load_probe_against_urls( response_mode: config.response_mode, total_requests: config.total_requests, concurrency: config.concurrency, + warmup_connections: config.warmup_connections, + warmup_url: config.warmup_url.clone(), + client_shards: config.client_shards, + start_ramp_ms: config.start_ramp.as_millis() as u64, + pool_max_idle_per_host: config.pool_max_idle_per_host, + http1_only: config.http1_only, + http2_prior_knowledge: config.http2_prior_knowledge, + connect_timeout_ms: config + .connect_timeout + .map(|timeout| timeout.as_millis() as u64), + first_body_hold_ms: config.first_body_hold.as_millis() as u64, duration_ms, throughput_rps, p99_ms, @@ -272,32 +441,261 @@ async fn run_http_load_probe_against_urls( p95_ms, max_ms, mean_ms, + headers_p50_ms: (!header_latencies.is_empty()).then_some(headers_p50_ms), + headers_p95_ms: (!header_latencies.is_empty()).then_some(headers_p95_ms), + headers_p99_ms: (!header_latencies.is_empty()).then_some(headers_p99_ms), + first_body_p50_ms: (!first_body_latencies.is_empty()).then_some(first_body_p50_ms), + first_body_p95_ms: (!first_body_latencies.is_empty()).then_some(first_body_p95_ms), + first_body_p99_ms: (!first_body_latencies.is_empty()).then_some(first_body_p99_ms), runtime: runtime_sampler.snapshot(), status_counts, error_counts, + error_samples, }) } -fn classify_reqwest_error(err: &reqwest::Error) -> String { - if err.is_timeout() { - return "timeout".to_string(); +fn build_probe_clients(config: &HttpLoadProbeConfig) -> Result, String> { + let mut clients = Vec::with_capacity(config.client_shards); + for _ in 0..config.client_shards { + let mut builder = Client::builder().timeout(config.timeout); + if let Some(connect_timeout) = config.connect_timeout { + builder = builder.connect_timeout(connect_timeout); + } + if let Some(pool_max_idle_per_host) = config.pool_max_idle_per_host { + builder = builder.pool_max_idle_per_host(pool_max_idle_per_host); + } + if config.http1_only { + builder = builder.http1_only(); + } + if config.http2_prior_knowledge { + builder = builder.http2_prior_knowledge(); + } + clients.push( + builder + .build() + .map_err(|err| format!("failed to build load probe http client: {err}"))?, + ); } - if err.is_connect() { - return "connect".to_string(); + Ok(clients) +} + +fn worker_start_delay(start_ramp: Duration, worker_index: usize, concurrency: usize) -> Duration { + if start_ramp.is_zero() || concurrency <= 1 || worker_index == 0 { + return Duration::ZERO; } - if err.is_body() { - return "body".to_string(); + let ramp_nanos = start_ramp.as_nanos(); + let offset_nanos = ramp_nanos + .saturating_mul(worker_index as u128) + .checked_div((concurrency - 1) as u128) + .unwrap_or_default(); + Duration::from_nanos(offset_nanos.min(u64::MAX as u128) as u64) +} + +async fn warmup_probe_connections( + config: &HttpLoadProbeConfig, + clients: Arc>, +) -> Result<(), String> { + if config.warmup_connections == 0 { + return Ok(()); } - if err.is_request() { - return "request".to_string(); + let warmup_url = config.warmup_url.as_deref().unwrap_or(config.url.as_str()); + let next_request = Arc::new(AtomicUsize::new(0)); + let mut workers = tokio::task::JoinSet::new(); + let concurrency = config.concurrency.min(config.warmup_connections).max(1); + for worker_index in 0..concurrency { + let client = clients[worker_index % clients.len()].clone(); + let next_request = Arc::clone(&next_request); + let warmup_url = warmup_url.to_string(); + let total = config.warmup_connections; + workers.spawn(async move { + loop { + let current = next_request.fetch_add(1, Ordering::AcqRel); + if current >= total { + break; + } + let response = client + .get(&warmup_url) + .send() + .await + .map_err(|err| format!("warmup request failed: {err}"))?; + let response = response + .error_for_status() + .map_err(|err| format!("warmup request returned error status: {err}"))?; + response + .bytes() + .await + .map_err(|err| format!("warmup response body failed: {err}"))?; + } + Ok::<(), String>(()) + }); } - if err.is_decode() { - return "decode".to_string(); + while let Some(result) = workers.join_next().await { + result + .map_err(|err| format!("warmup worker task failed: {err}"))? + .map_err(|err| format!("failed to warm load probe connections: {err}"))?; } - if err.is_redirect() { - return "redirect".to_string(); + Ok(()) +} + +#[derive(Debug, Clone, PartialEq, Eq)] +struct ClassifiedLoadError { + key: String, + phase: String, + kind: String, + message: String, + source: Option, +} + +impl ClassifiedLoadError { + fn static_body(kind: &str, message: &str) -> Self { + Self { + key: format!("body:{kind}"), + phase: "body".to_string(), + kind: kind.to_string(), + message: message.to_string(), + source: None, + } } - "other".to_string() +} + +async fn record_load_error( + error_counts: &Arc>>, + error_samples: &Arc>>, + request_index: usize, + url: &str, + elapsed_ms: u64, + error: ClassifiedLoadError, +) { + let mut counts = error_counts.lock().await; + *counts.entry(error.key).or_insert(0) += 1; + drop(counts); + + let mut samples = error_samples.lock().await; + if samples.len() < MAX_ERROR_SAMPLES { + samples.push(HttpLoadProbeErrorSample { + request_index, + url: url.to_string(), + phase: error.phase, + kind: error.kind, + elapsed_ms, + message: error.message, + source: error.source, + }); + } +} + +fn classify_reqwest_error(phase: &str, err: &reqwest::Error) -> ClassifiedLoadError { + let kind = if err.is_timeout() && err.is_connect() { + "connect_timeout" + } else if err.is_timeout() && err.is_body() { + "body_timeout" + } else if err.is_timeout() { + "timeout" + } else if err.is_connect() { + "connect" + } else if err.is_body() { + "body" + } else if err.is_request() { + "request" + } else if err.is_decode() { + "decode" + } else if err.is_redirect() { + "redirect" + } else { + "other" + }; + + ClassifiedLoadError { + key: format!("{phase}:{kind}"), + phase: phase.to_string(), + kind: kind.to_string(), + message: compact_error_text(err.to_string(), 240), + source: error_source_chain(err, 240), + } +} + +fn error_source_chain(err: &(dyn StdError + 'static), max_chars: usize) -> Option { + let mut sources = Vec::new(); + let mut next = err.source(); + while let Some(source) = next { + sources.push(compact_error_text(source.to_string(), max_chars)); + if sources.len() >= 4 { + break; + } + next = source.source(); + } + (!sources.is_empty()).then(|| compact_error_text(sources.join(" | "), max_chars)) +} + +fn compact_error_text(value: impl AsRef, max_chars: usize) -> String { + let mut compact = value + .as_ref() + .split_whitespace() + .collect::>() + .join(" "); + if compact.chars().count() > max_chars { + compact = compact.chars().take(max_chars.saturating_sub(1)).collect(); + compact.push_str("..."); + } + compact +} + +async fn observe_response_body( + mut response: reqwest::Response, + response_mode: HttpLoadProbeResponseMode, + started_at: Instant, + first_body_hold: Duration, +) -> Result { + match response_mode { + HttpLoadProbeResponseMode::HeadersOnly => Ok(BodyObservation::default()), + HttpLoadProbeResponseMode::FirstBodyByte => { + let first = response + .chunk() + .await + .map_err(|err| classify_reqwest_error("body", &err))? + .ok_or_else(|| { + ClassifiedLoadError::static_body( + "empty_body", + "response body ended before the first chunk", + ) + })?; + let first_body_latency_ms = started_at.elapsed().as_millis() as u64; + drop(first); + if !first_body_hold.is_zero() { + tokio::time::sleep(first_body_hold).await; + } + Ok(BodyObservation { + first_body_latency_ms: Some(first_body_latency_ms), + }) + } + HttpLoadProbeResponseMode::FullBody => { + let mut first_body_latency_ms = None; + while let Some(chunk) = response + .chunk() + .await + .map_err(|err| classify_reqwest_error("body", &err))? + { + if first_body_latency_ms.is_none() { + first_body_latency_ms = Some(started_at.elapsed().as_millis() as u64); + } + drop(chunk); + } + if first_body_latency_ms.is_none() { + return Err(ClassifiedLoadError::static_body( + "empty_body", + "response body ended before the first chunk", + )); + } + Ok(BodyObservation { + first_body_latency_ms, + }) + } + } +} + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +struct BodyObservation { + first_body_latency_ms: Option, } fn build_headers(headers: &BTreeMap) -> Result { @@ -337,7 +735,8 @@ fn percentile(latencies: &[u64], percentile: u8) -> u64 { #[cfg(test)] mod tests { use super::{ - build_headers, summarize_latencies, HttpLoadProbeConfig, HttpLoadProbeResponseMode, + build_headers, summarize_latencies, worker_start_delay, HttpLoadProbeConfig, + HttpLoadProbeResponseMode, }; use reqwest::Method; use std::collections::BTreeMap; @@ -369,6 +768,25 @@ mod tests { } .validate() .is_err()); + assert!(HttpLoadProbeConfig { + connect_timeout: Some(Duration::ZERO), + ..HttpLoadProbeConfig::default() + } + .validate() + .is_err()); + assert!(HttpLoadProbeConfig { + client_shards: 0, + ..HttpLoadProbeConfig::default() + } + .validate() + .is_err()); + assert!(HttpLoadProbeConfig { + http1_only: true, + http2_prior_knowledge: true, + ..HttpLoadProbeConfig::default() + } + .validate() + .is_err()); } #[test] @@ -386,12 +804,21 @@ mod tests { fn default_probe_config_is_reasonable() { let config = HttpLoadProbeConfig::default(); assert_eq!(config.method, Method::GET); + assert!(config.warmup_url.is_none()); assert!(config.headers.is_empty()); assert!(config.body.is_none()); assert_eq!(config.total_requests, 100); assert_eq!(config.concurrency, 10); + assert_eq!(config.warmup_connections, 0); assert_eq!(config.timeout, Duration::from_secs(30)); + assert_eq!(config.connect_timeout, None); assert_eq!(config.response_mode, HttpLoadProbeResponseMode::HeadersOnly); + assert_eq!(config.client_shards, 1); + assert_eq!(config.pool_max_idle_per_host, None); + assert_eq!(config.start_ramp, Duration::ZERO); + assert!(!config.http1_only); + assert!(!config.http2_prior_knowledge); + assert_eq!(config.first_body_hold, Duration::ZERO); } #[test] @@ -408,4 +835,24 @@ mod tests { let invalid = BTreeMap::from([("bad header".to_string(), "ok".to_string())]); assert!(build_headers(&invalid).is_err()); } + + #[test] + fn spreads_worker_start_delay_across_ramp() { + assert_eq!( + worker_start_delay(Duration::from_millis(900), 0, 10), + Duration::ZERO + ); + assert_eq!( + worker_start_delay(Duration::from_millis(900), 5, 10), + Duration::from_millis(500) + ); + assert_eq!( + worker_start_delay(Duration::from_millis(900), 9, 10), + Duration::from_millis(900) + ); + assert_eq!( + worker_start_delay(Duration::from_millis(900), 3, 1), + Duration::ZERO + ); + } } diff --git a/crates/aether-usage-runtime/src/config.rs b/crates/aether-usage-runtime/src/config.rs index 1cb326242..9bb2d3bb8 100644 --- a/crates/aether-usage-runtime/src/config.rs +++ b/crates/aether-usage-runtime/src/config.rs @@ -14,6 +14,10 @@ pub struct UsageRuntimeConfig { pub reclaim_idle_ms: u64, pub reclaim_count: usize, pub reclaim_interval_ms: u64, + pub enqueue_retry_buffer_capacity: usize, + pub enqueue_retry_workers: usize, + pub enqueue_retry_initial_backoff_ms: u64, + pub enqueue_retry_max_backoff_ms: u64, } impl Default for UsageRuntimeConfig { @@ -31,6 +35,10 @@ impl Default for UsageRuntimeConfig { reclaim_idle_ms: 30_000, reclaim_count: 500, reclaim_interval_ms: 5_000, + enqueue_retry_buffer_capacity: 131_072, + enqueue_retry_workers: 4, + enqueue_retry_initial_backoff_ms: 10, + enqueue_retry_max_backoff_ms: 1_000, } } } @@ -90,6 +98,31 @@ impl UsageRuntimeConfig { "usage runtime reclaim_interval_ms must be positive".to_string(), )); } + if self.enqueue_retry_buffer_capacity == 0 { + return Err(DataLayerError::InvalidConfiguration( + "usage runtime enqueue_retry_buffer_capacity must be positive".to_string(), + )); + } + if self.enqueue_retry_workers == 0 { + return Err(DataLayerError::InvalidConfiguration( + "usage runtime enqueue_retry_workers must be positive".to_string(), + )); + } + if self.enqueue_retry_initial_backoff_ms == 0 { + return Err(DataLayerError::InvalidConfiguration( + "usage runtime enqueue_retry_initial_backoff_ms must be positive".to_string(), + )); + } + if self.enqueue_retry_max_backoff_ms == 0 { + return Err(DataLayerError::InvalidConfiguration( + "usage runtime enqueue_retry_max_backoff_ms must be positive".to_string(), + )); + } + if self.enqueue_retry_initial_backoff_ms > self.enqueue_retry_max_backoff_ms { + return Err(DataLayerError::InvalidConfiguration( + "usage runtime enqueue retry initial backoff cannot exceed max backoff".to_string(), + )); + } Ok(()) } diff --git a/crates/aether-usage-runtime/src/request_metadata.rs b/crates/aether-usage-runtime/src/request_metadata.rs index 554572be3..14d6d3c63 100644 --- a/crates/aether-usage-runtime/src/request_metadata.rs +++ b/crates/aether-usage-runtime/src/request_metadata.rs @@ -146,6 +146,8 @@ fn copy_allowed_metadata_fields(source: &Map, target: &mut Map, target: &mut Map remove_number(&mut source, target, "cache_read_price_per_1m"); remove_number(&mut source, target, "price_per_request"); remove_non_null_value(&mut source, target, "proxy"); + remove_non_null_value(&mut source, target, "stage_timings_ms"); + remove_non_null_value(&mut source, target, "db_timings_ms"); sanitize_request_path_metadata_fields(target); } @@ -476,8 +480,38 @@ mod tests { } } + fn sample_stage_timings_metadata() -> Value { + json!({ + "stream_candidate_slot": 1, + "stream_provider_in_flight": 2, + "stream_upstream_headers": 180, + "stream_first_data": 8210 + }) + } + + fn sample_db_timings_metadata() -> Value { + json!({ + "query_count": 2, + "query_total": 950, + "query_max": 650, + "operations": { + "request_candidate_upsert": {"count": 1, "sum": 650, "max": 650}, + "usage_upsert": {"count": 1, "sum": 300, "max": 300} + }, + "pool": { + "max_checked_out": 20, + "max_pool_size": 20, + "min_idle": 0, + "max_connections": 20, + "max_usage_rate": 100.0 + } + }) + } + #[test] fn sanitizes_request_metadata_to_allowlist() { + let stage_timings_ms = sample_stage_timings_metadata(); + let db_timings_ms = sample_db_timings_metadata(); let metadata = sanitize_usage_request_metadata(Some(json!({ "request_id": "req-1", "provider_id": "provider-1", @@ -512,6 +546,8 @@ mod tests { "cache_creation_price_per_1m": 3.75, "cache_read_price_per_1m": 0.3, "price_per_request": 0.02, + "stage_timings_ms": stage_timings_ms.clone(), + "db_timings_ms": db_timings_ms.clone(), "original_headers": {"authorization": "Bearer secret"}, "original_request_body": {"messages": []}, "provider_request_headers": {"authorization": "Bearer secret"}, @@ -549,7 +585,9 @@ mod tests { "output_price_per_1m": 15.0, "cache_creation_price_per_1m": 3.75, "cache_read_price_per_1m": 0.3, - "price_per_request": 0.02 + "price_per_request": 0.02, + "stage_timings_ms": stage_timings_ms, + "db_timings_ms": db_timings_ms }) ); } @@ -657,7 +695,20 @@ mod tests { "global_model_name": "gpt-5", "client_ip": "203.0.113.8", "user_agent": "Claude-Code/1.0", - "billing_snapshot": {"status": "complete"} + "billing_snapshot": {"status": "complete"}, + "stage_timings_ms": { + "stream_candidate_slot": 0, + "stream_upstream_headers": 180, + "stream_first_data": 8210 + }, + "db_timings_ms": { + "query_count": 1, + "query_total": 42, + "query_max": 42, + "operations": { + "auth_api_key_snapshot": {"count": 1, "sum": 42, "max": 42} + } + } }) .as_object() .expect("object"), @@ -676,7 +727,20 @@ mod tests { "global_model_name": "gpt-5", "client_ip": "203.0.113.8", "user_agent": "Claude-Code/1.0", - "billing_snapshot": {"status": "complete"} + "billing_snapshot": {"status": "complete"}, + "stage_timings_ms": { + "stream_candidate_slot": 0, + "stream_upstream_headers": 180, + "stream_first_data": 8210 + }, + "db_timings_ms": { + "query_count": 1, + "query_total": 42, + "query_max": 42, + "operations": { + "auth_api_key_snapshot": {"count": 1, "sum": 42, "max": 42} + } + } }) ); } diff --git a/crates/aether-usage-runtime/src/runtime.rs b/crates/aether-usage-runtime/src/runtime.rs index 9dcd91f16..7633f857c 100644 --- a/crates/aether-usage-runtime/src/runtime.rs +++ b/crates/aether-usage-runtime/src/runtime.rs @@ -1,5 +1,6 @@ use std::future::Future; use std::pin::Pin; +use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -7,6 +8,7 @@ use aether_contracts::ExecutionTelemetry; use aether_data_contracts::DataLayerError; use aether_runtime_state::RuntimeQueueStore; use async_trait::async_trait; +use tokio::sync::mpsc; use tracing::warn; use crate::executor::spawn_on_usage_background_runtime; @@ -76,10 +78,15 @@ pub trait UsageRuntimeAccess: #[derive(Debug, Clone)] pub struct UsageRuntime { config: UsageRuntimeConfig, - body_policy_cache: Arc>>, + body_policy_cache: Arc>>, + enqueue_retry: Arc, + lifecycle_enqueue_state: Arc, } const USAGE_BODY_CAPTURE_POLICY_CACHE_TTL: Duration = Duration::from_secs(30); +const USAGE_BODY_CAPTURE_POLICY_ERROR_CACHE_TTL: Duration = Duration::from_secs(1); +const LIFECYCLE_ENQUEUE_MAX_IN_FLIGHT: u64 = 128; +const LIFECYCLE_ENQUEUE_CIRCUIT_OPEN_MS: u64 = 1_000; impl Default for UsageRuntime { fn default() -> Self { @@ -92,14 +99,19 @@ impl UsageRuntime { Self { config: UsageRuntimeConfig::disabled(), body_policy_cache: Arc::new(tokio::sync::Mutex::new(None)), + enqueue_retry: UsageEnqueueRetryDispatcher::disabled(), + lifecycle_enqueue_state: Arc::new(LifecycleEnqueueState::default()), } } pub fn new(config: UsageRuntimeConfig) -> Result { config.validate()?; + let enqueue_retry = UsageEnqueueRetryDispatcher::spawn(config.clone()); Ok(Self { config, body_policy_cache: Arc::new(tokio::sync::Mutex::new(None)), + enqueue_retry, + lifecycle_enqueue_state: Arc::new(LifecycleEnqueueState::default()), }) } @@ -222,15 +234,6 @@ impl UsageRuntime { runtime .apply_body_capture_policy_from_data(&data, &mut event) .await; - if let Err(err) = data.enrich_usage_event(&mut event).await { - warn!( - event_name = "usage_sync_terminal_billing_enrichment_failed", - log_type = "event", - request_id = %request_id, - error = %err, - "usage runtime failed to enrich sync usage event with billing" - ); - } runtime.enqueue_or_write_terminal(&data, event).await } Err(err) => { @@ -269,15 +272,6 @@ impl UsageRuntime { runtime .apply_body_capture_policy_from_data(&data, &mut event) .await; - if let Err(err) = data.enrich_usage_event(&mut event).await { - warn!( - event_name = "usage_stream_terminal_billing_enrichment_failed", - log_type = "event", - request_id = %request_id, - error = %err, - "usage runtime failed to enrich stream usage event with billing" - ); - } runtime.enqueue_or_write_terminal(&data, event).await } Err(err) => { @@ -316,15 +310,6 @@ impl UsageRuntime { } self.apply_body_capture_policy_from_data(data, &mut event) .await; - if let Err(err) = data.enrich_usage_event(&mut event).await { - warn!( - event_name = "usage_terminal_billing_enrichment_failed", - log_type = "event", - request_id = %event.request_id, - error = %err, - "usage runtime failed to enrich terminal usage event with billing" - ); - } self.enqueue_or_write_terminal(data, event).await; } @@ -387,14 +372,26 @@ impl UsageRuntime { T: UsageRuntimeAccess, { let mut cache = self.body_policy_cache.lock().await; - if let Some((cached_at, policy)) = cache.as_ref() { - if cached_at.elapsed() <= USAGE_BODY_CAPTURE_POLICY_CACHE_TTL { - return Ok(*policy); + if let Some(entry) = cache.as_ref() { + if entry.cached_at.elapsed() <= entry.ttl { + return match entry.source { + UsageBodyCapturePolicyCacheSource::Loaded => Ok(entry.policy), + UsageBodyCapturePolicyCacheSource::FallbackAfterError => { + Ok(UsageBodyCapturePolicy::default()) + } + }; + } + } + match data.body_capture_policy().await { + Ok(policy) => { + *cache = Some(UsageBodyCapturePolicyCacheEntry::loaded(policy)); + Ok(policy) + } + Err(err) => { + *cache = Some(UsageBodyCapturePolicyCacheEntry::fallback_after_error()); + Err(err) } } - let policy = data.body_capture_policy().await?; - *cache = Some((Instant::now(), policy)); - Ok(policy) } async fn enqueue_or_write_terminal(&self, data: &T, event: UsageEvent) @@ -409,14 +406,92 @@ impl UsageRuntime { where T: UsageRuntimeAccess, { - self.enqueue_or_write_event(data, event, "lifecycle", self.config.queue_lifecycle_events) - .await; + self.enqueue_lifecycle_event(data, event).await; + } + + async fn enqueue_lifecycle_event(&self, data: &T, event: UsageEvent) + where + T: UsageRuntimeAccess, + { + if !self.config.queue_lifecycle_events { + warn!( + event_name = "usage_lifecycle_event_not_queued", + log_type = "event", + usage_event_type = ?event.event_type, + request_id = %event.request_id, + fallback = "none", + "usage runtime lifecycle queue is disabled; lifecycle event will not be written directly" + ); + return; + } + + let now_ms = now_unix_ms(); + if self.lifecycle_enqueue_state.is_circuit_open(now_ms) { + self.lifecycle_enqueue_state + .record_skip("circuit_open", &event); + return; + } + + let Some(_guard) = self.lifecycle_enqueue_state.try_acquire_in_flight(&event) else { + return; + }; + + let Some(runner) = data.usage_worker_queue() else { + warn!( + event_name = "usage_lifecycle_event_queue_unavailable", + log_type = "event", + usage_event_type = ?event.event_type, + request_id = %event.request_id, + fallback = "none", + "usage runtime lifecycle queue is unavailable; lifecycle event will not be written directly" + ); + return; + }; + + let queue = match UsageQueue::new(runner, self.config.clone()) { + Ok(queue) => queue, + Err(err) => { + warn!( + event_name = "usage_lifecycle_event_queue_init_failed", + log_type = "event", + usage_event_type = ?event.event_type, + request_id = %event.request_id, + fallback = "none", + error = %err, + "usage runtime failed to build lifecycle queue; lifecycle event will not be written directly" + ); + return; + } + }; + + if let Err(err) = queue.enqueue(&event).await { + self.lifecycle_enqueue_state + .open_circuit(now_unix_ms().saturating_add(LIFECYCLE_ENQUEUE_CIRCUIT_OPEN_MS)); + let failures = self + .lifecycle_enqueue_state + .failed_total + .fetch_add(1, Ordering::AcqRel) + + 1; + if should_log_usage_retry_counter(failures) { + warn!( + event_name = "usage_lifecycle_event_enqueue_failed", + log_type = "event", + usage_event_type = ?event.event_type, + request_id = %event.request_id, + fallback = "none", + failure_total = failures, + circuit_open_ms = LIFECYCLE_ENQUEUE_CIRCUIT_OPEN_MS, + error = %err, + "usage runtime failed to enqueue lifecycle event; lifecycle enqueue circuit opened" + ); + } + } } async fn enqueue_or_write_event( &self, data: &T, - event: UsageEvent, + mut event: UsageEvent, event_phase: &'static str, queue_enabled: bool, ) where @@ -428,16 +503,10 @@ impl UsageRuntime { Ok(queue) => match queue.enqueue(&event).await { Ok(_) => return, Err(err) => { - warn!( - event_name = "usage_event_enqueue_failed", - log_type = "event", - event_phase, - usage_event_type = ?event.event_type, - request_id = %event.request_id, - fallback = "direct_write", - error = %err, - "usage runtime failed to enqueue usage event; falling back to direct write" - ) + self.enqueue_retry + .schedule(queue, event, event_phase, err) + .await; + return; } }, Err(err) => { @@ -456,6 +525,9 @@ impl UsageRuntime { } } + if event_phase == "terminal" { + enrich_terminal_event(data, &mut event).await; + } self.write_event_direct(data, &event).await; } @@ -502,6 +574,363 @@ impl UsageRuntime { } } +async fn enrich_terminal_event(data: &T, event: &mut UsageEvent) +where + T: UsageBillingEventEnricher + Send + Sync, +{ + if let Err(err) = data.enrich_usage_event(event).await { + warn!( + event_name = "usage_terminal_billing_enrichment_failed", + log_type = "event", + request_id = %event.request_id, + error = %err, + "usage runtime failed to enrich terminal usage event with billing" + ); + } +} + +#[derive(Debug, Clone, Copy)] +struct UsageBodyCapturePolicyCacheEntry { + cached_at: Instant, + ttl: Duration, + policy: UsageBodyCapturePolicy, + source: UsageBodyCapturePolicyCacheSource, +} + +impl UsageBodyCapturePolicyCacheEntry { + fn loaded(policy: UsageBodyCapturePolicy) -> Self { + Self { + cached_at: Instant::now(), + ttl: USAGE_BODY_CAPTURE_POLICY_CACHE_TTL, + policy, + source: UsageBodyCapturePolicyCacheSource::Loaded, + } + } + + fn fallback_after_error() -> Self { + Self { + cached_at: Instant::now(), + ttl: USAGE_BODY_CAPTURE_POLICY_ERROR_CACHE_TTL, + policy: UsageBodyCapturePolicy::default(), + source: UsageBodyCapturePolicyCacheSource::FallbackAfterError, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum UsageBodyCapturePolicyCacheSource { + Loaded, + FallbackAfterError, +} + +#[derive(Debug, Default)] +struct LifecycleEnqueueState { + in_flight: AtomicU64, + circuit_open_until_unix_ms: AtomicU64, + skipped_total: AtomicU64, + failed_total: AtomicU64, +} + +impl LifecycleEnqueueState { + fn is_circuit_open(&self, now_unix_ms: u64) -> bool { + self.circuit_open_until_unix_ms.load(Ordering::Acquire) > now_unix_ms + } + + fn open_circuit(&self, open_until_unix_ms: u64) { + let mut current = self.circuit_open_until_unix_ms.load(Ordering::Acquire); + while open_until_unix_ms > current { + match self.circuit_open_until_unix_ms.compare_exchange( + current, + open_until_unix_ms, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => break, + Err(next) => current = next, + } + } + } + + fn try_acquire_in_flight<'a>( + &'a self, + event: &UsageEvent, + ) -> Option> { + let mut current = self.in_flight.load(Ordering::Acquire); + loop { + if current >= LIFECYCLE_ENQUEUE_MAX_IN_FLIGHT { + self.record_skip("in_flight_limit", event); + return None; + } + match self.in_flight.compare_exchange_weak( + current, + current + 1, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => { + return Some(LifecycleEnqueueInFlightGuard { state: self }); + } + Err(next) => current = next, + } + } + } + + fn record_skip(&self, reason: &'static str, event: &UsageEvent) { + let skipped = self.skipped_total.fetch_add(1, Ordering::AcqRel) + 1; + if should_log_usage_retry_counter(skipped) { + warn!( + event_name = "usage_lifecycle_event_enqueue_skipped", + log_type = "event", + usage_event_type = ?event.event_type, + request_id = %event.request_id, + reason, + skipped_total = skipped, + fallback = "none", + "usage runtime skipped lifecycle enqueue" + ); + } + } +} + +struct LifecycleEnqueueInFlightGuard<'a> { + state: &'a LifecycleEnqueueState, +} + +impl Drop for LifecycleEnqueueInFlightGuard<'_> { + fn drop(&mut self) { + self.state.in_flight.fetch_sub(1, Ordering::AcqRel); + } +} + +#[derive(Debug)] +struct UsageEnqueueRetryDispatcher { + senders: Vec>, + scheduled_total: Arc, +} + +struct UsageEnqueueRetryItem { + queue: UsageQueue, + event: UsageEvent, + event_phase: &'static str, + attempts: u64, +} + +impl UsageEnqueueRetryDispatcher { + fn disabled() -> Arc { + Arc::new(Self { + senders: Vec::new(), + scheduled_total: Arc::new(AtomicU64::new(0)), + }) + } + + fn spawn(config: UsageRuntimeConfig) -> Arc { + if !config.enabled || !(config.queue_terminal_events || config.queue_lifecycle_events) { + return Self::disabled(); + } + + let workers = config + .enqueue_retry_workers + .min(config.enqueue_retry_buffer_capacity) + .max(1); + let mut senders = Vec::with_capacity(workers); + let scheduled_total = Arc::new(AtomicU64::new(0)); + let recovered_total = Arc::new(AtomicU64::new(0)); + for worker_index in 0..workers { + let capacity = + retry_worker_capacity(config.enqueue_retry_buffer_capacity, workers, worker_index); + let (sender, receiver) = mpsc::channel(capacity); + senders.push(sender); + let worker_config = config.clone(); + let worker_recovered_total = Arc::clone(&recovered_total); + spawn_on_usage_background_runtime(async move { + run_usage_enqueue_retry_worker( + worker_index, + worker_config, + receiver, + worker_recovered_total, + ) + .await; + }); + } + + Arc::new(Self { + senders, + scheduled_total, + }) + } + + async fn schedule( + &self, + queue: UsageQueue, + event: UsageEvent, + event_phase: &'static str, + cause: DataLayerError, + ) { + let scheduled = self.scheduled_total.fetch_add(1, Ordering::AcqRel) + 1; + if should_log_usage_retry_counter(scheduled) { + warn!( + event_name = "usage_event_enqueue_failed_retry_scheduled", + log_type = "event", + event_phase, + usage_event_type = ?event.event_type, + request_id = %event.request_id, + retry_scheduled_total = scheduled, + fallback = "local_enqueue_retry", + error = %cause, + "usage runtime failed to enqueue usage event; scheduled local retry" + ); + } + + let worker_index = retry_worker_index(&event.request_id, self.senders.len()); + let Some(sender) = self.senders.get(worker_index) else { + warn!( + event_name = "usage_event_enqueue_retry_unavailable", + log_type = "event", + event_phase, + usage_event_type = ?event.event_type, + request_id = %event.request_id, + error = %cause, + "usage runtime local enqueue retry dispatcher is unavailable" + ); + return; + }; + + let item = UsageEnqueueRetryItem { + queue, + event, + event_phase, + attempts: 0, + }; + match sender.try_send(item) { + Ok(()) => {} + Err(mpsc::error::TrySendError::Full(item)) => { + warn!( + event_name = "usage_event_enqueue_retry_buffer_full", + log_type = "event", + event_phase = item.event_phase, + usage_event_type = ?item.event.event_type, + request_id = %item.event.request_id, + worker_index, + "usage runtime local enqueue retry buffer is full; waiting for capacity" + ); + if let Err(err) = sender.send(item).await { + let item = err.0; + warn!( + event_name = "usage_event_enqueue_retry_closed", + log_type = "event", + event_phase = item.event_phase, + usage_event_type = ?item.event.event_type, + request_id = %item.event.request_id, + worker_index, + "usage runtime local enqueue retry dispatcher closed before accepting event" + ); + } + } + Err(mpsc::error::TrySendError::Closed(item)) => { + warn!( + event_name = "usage_event_enqueue_retry_closed", + log_type = "event", + event_phase = item.event_phase, + usage_event_type = ?item.event.event_type, + request_id = %item.event.request_id, + worker_index, + "usage runtime local enqueue retry dispatcher is closed" + ); + } + } + } +} + +async fn run_usage_enqueue_retry_worker( + worker_index: usize, + config: UsageRuntimeConfig, + mut receiver: mpsc::Receiver, + recovered_total: Arc, +) { + while let Some(mut item) = receiver.recv().await { + loop { + match item.queue.enqueue(&item.event).await { + Ok(_) => { + let recovered = recovered_total.fetch_add(1, Ordering::AcqRel) + 1; + if item.attempts > 0 && should_log_usage_retry_counter(recovered) { + warn!( + event_name = "usage_event_enqueue_retry_recovered", + log_type = "event", + event_phase = item.event_phase, + usage_event_type = ?item.event.event_type, + request_id = %item.event.request_id, + worker_index, + retry_attempts = item.attempts, + retry_recovered_total = recovered, + "usage runtime local enqueue retry recovered" + ); + } + break; + } + Err(err) => { + item.attempts = item.attempts.saturating_add(1); + let delay = usage_enqueue_retry_delay(&config, item.attempts); + if should_log_usage_retry_counter(item.attempts) { + warn!( + event_name = "usage_event_enqueue_retry_failed", + log_type = "event", + event_phase = item.event_phase, + usage_event_type = ?item.event.event_type, + request_id = %item.event.request_id, + worker_index, + retry_attempt = item.attempts, + retry_delay_ms = delay.as_millis() as u64, + error = %err, + "usage runtime local enqueue retry failed; will retry" + ); + } + tokio::time::sleep(delay).await; + } + } + } + } +} + +fn usage_enqueue_retry_delay(config: &UsageRuntimeConfig, attempts: u64) -> Duration { + let exponent = attempts.saturating_sub(1).min(16); + let multiplier = 1_u64.checked_shl(exponent as u32).unwrap_or(u64::MAX); + let delay_ms = config + .enqueue_retry_initial_backoff_ms + .saturating_mul(multiplier) + .min(config.enqueue_retry_max_backoff_ms); + Duration::from_millis(delay_ms.max(1)) +} + +fn retry_worker_capacity(total_capacity: usize, workers: usize, worker_index: usize) -> usize { + let workers = workers.max(1); + let base = total_capacity / workers; + let remainder = total_capacity % workers; + (base + usize::from(worker_index < remainder)).max(1) +} + +fn retry_worker_index(request_id: &str, worker_count: usize) -> usize { + if worker_count <= 1 { + return 0; + } + (fnv_hash(request_id.as_bytes()) % worker_count as u64) as usize +} + +fn fnv_hash(bytes: &[u8]) -> u64 { + const FNV_OFFSET: u64 = 14_695_981_039_346_656_037; + const FNV_PRIME: u64 = 1_099_511_628_211; + + let mut hash = FNV_OFFSET; + for byte in bytes { + hash ^= u64::from(*byte); + hash = hash.wrapping_mul(FNV_PRIME); + } + hash +} + +fn should_log_usage_retry_counter(value: u64) -> bool { + value <= 8 || value.is_power_of_two() || value % 1_000 == 0 +} + async fn build_pending_usage_event_offthread( seed: LifecycleUsageSeed, now_unix_secs: u64, @@ -573,15 +1002,20 @@ where } fn now_unix_secs() -> u64 { + now_unix_ms() / 1_000 +} + +fn now_unix_ms() -> u64 { std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .unwrap_or_default() - .as_secs() + .as_millis() as u64 } #[cfg(test)] mod tests { use std::collections::BTreeMap; + use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use aether_contracts::{ExecutionPlan, RequestBody}; @@ -622,6 +1056,24 @@ mod tests { queue: Arc, } + struct EnrichmentCountingQueueStore { + records: Mutex>, + queue: Arc, + enrich_calls: AtomicUsize, + } + + #[derive(Default)] + struct FailingPolicyUsageStore { + inner: NoRedisUsageStore, + policy_reads: AtomicUsize, + } + + struct FlakyAppendQueueStore { + inner: Arc, + remaining_failures: AtomicUsize, + append_attempts: AtomicUsize, + } + #[async_trait] impl UsageRecordWriter for NoRedisUsageStore { async fn upsert_usage_record( @@ -798,6 +1250,207 @@ mod tests { } } + #[async_trait] + impl UsageRecordWriter for EnrichmentCountingQueueStore { + async fn upsert_usage_record( + &self, + record: UpsertUsageRecord, + ) -> Result, DataLayerError> { + self.records.lock().expect("records lock").push(record); + Ok(None) + } + } + + #[async_trait] + impl UsageSettlementWriter for EnrichmentCountingQueueStore { + fn has_usage_settlement_writer(&self) -> bool { + false + } + + async fn settle_usage( + &self, + _input: UsageSettlementInput, + ) -> Result, DataLayerError> { + Ok(None) + } + } + + #[async_trait] + impl UsageBillingEventEnricher for EnrichmentCountingQueueStore { + async fn enrich_usage_event(&self, event: &mut UsageEvent) -> Result<(), DataLayerError> { + self.enrich_calls.fetch_add(1, Ordering::AcqRel); + event.data.total_cost_usd = Some(0.123); + Ok(()) + } + } + + #[async_trait] + impl ManualProxyNodeCounter for EnrichmentCountingQueueStore { + async fn increment_manual_proxy_node_requests( + &self, + _node_id: &str, + _total_delta: i64, + _failed_delta: i64, + _latency_ms: Option, + ) -> Result<(), DataLayerError> { + Ok(()) + } + } + + impl UsageRuntimeAccess for EnrichmentCountingQueueStore { + fn has_usage_writer(&self) -> bool { + true + } + + fn has_usage_worker_queue(&self) -> bool { + true + } + + fn usage_worker_queue(&self) -> Option> { + Some(Arc::clone(&self.queue)) + } + } + + #[async_trait] + impl UsageRecordWriter for FailingPolicyUsageStore { + async fn upsert_usage_record( + &self, + record: UpsertUsageRecord, + ) -> Result, DataLayerError> { + self.inner.upsert_usage_record(record).await + } + } + + #[async_trait] + impl UsageSettlementWriter for FailingPolicyUsageStore { + fn has_usage_settlement_writer(&self) -> bool { + false + } + + async fn settle_usage( + &self, + _input: UsageSettlementInput, + ) -> Result, DataLayerError> { + Ok(None) + } + } + + #[async_trait] + impl UsageBillingEventEnricher for FailingPolicyUsageStore { + async fn enrich_usage_event(&self, _event: &mut UsageEvent) -> Result<(), DataLayerError> { + Ok(()) + } + } + + #[async_trait] + impl ManualProxyNodeCounter for FailingPolicyUsageStore { + async fn increment_manual_proxy_node_requests( + &self, + _node_id: &str, + _total_delta: i64, + _failed_delta: i64, + _latency_ms: Option, + ) -> Result<(), DataLayerError> { + Ok(()) + } + } + + #[async_trait] + impl UsageRuntimeAccess for FailingPolicyUsageStore { + fn has_usage_writer(&self) -> bool { + true + } + + fn has_usage_worker_queue(&self) -> bool { + false + } + + fn usage_worker_queue(&self) -> Option> { + None + } + + async fn body_capture_policy(&self) -> Result { + self.policy_reads.fetch_add(1, Ordering::AcqRel); + Err(DataLayerError::Postgres( + "forced policy read failure".to_string(), + )) + } + } + + #[async_trait] + impl RuntimeQueueStore for FlakyAppendQueueStore { + async fn ensure_consumer_group( + &self, + stream: &str, + group: &str, + start_id: &str, + ) -> Result<(), DataLayerError> { + self.inner + .ensure_consumer_group(stream, group, start_id) + .await + } + + async fn append_fields_with_maxlen( + &self, + stream: &str, + fields: &BTreeMap, + maxlen: Option, + ) -> Result { + self.append_attempts.fetch_add(1, Ordering::AcqRel); + let failed = self + .remaining_failures + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| { + (current > 0).then(|| current - 1) + }) + .is_ok(); + if failed { + return Err(DataLayerError::Redis("forced append failure".to_string())); + } + self.inner + .append_fields_with_maxlen(stream, fields, maxlen) + .await + } + + async fn read_group( + &self, + stream: &str, + group: &str, + consumer: &str, + count: usize, + block_ms: Option, + ) -> Result, DataLayerError> { + self.inner + .read_group(stream, group, consumer, count, block_ms) + .await + } + + async fn claim_stale( + &self, + stream: &str, + group: &str, + consumer: &str, + start_id: &str, + config: aether_runtime_state::RuntimeQueueReclaimConfig, + ) -> Result, DataLayerError> { + self.inner + .claim_stale(stream, group, consumer, start_id, config) + .await + } + + async fn ack( + &self, + stream: &str, + group: &str, + ids: &[String], + ) -> Result { + self.inner.ack(stream, group, ids).await + } + + async fn delete(&self, stream: &str, ids: &[String]) -> Result { + self.inner.delete(stream, ids).await + } + } + #[tokio::test] async fn terminal_usage_without_redis_writes_directly_to_usage_repository() { let runtime = UsageRuntime::new(UsageRuntimeConfig { @@ -933,6 +1586,309 @@ mod tests { panic!("pending lifecycle usage event was not enqueued"); } + #[tokio::test] + async fn lifecycle_queue_append_failure_does_not_write_directly() { + let config = UsageRuntimeConfig { + enabled: true, + queue_lifecycle_events: true, + stream_key: "usage:events:test:lifecycle-failure".to_string(), + consumer_group: "usage_consumers_test_lifecycle_failure".to_string(), + consumer_block_ms: 1, + enqueue_retry_initial_backoff_ms: 1, + enqueue_retry_max_backoff_ms: 5, + enqueue_retry_buffer_capacity: 16, + enqueue_retry_workers: 1, + ..UsageRuntimeConfig::default() + }; + let inner_queue: Arc = + Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default())); + let flaky_queue: Arc = Arc::new(FlakyAppendQueueStore { + inner: Arc::clone(&inner_queue), + remaining_failures: AtomicUsize::new(usize::MAX), + append_attempts: AtomicUsize::new(0), + }); + let store = CloneQueueConfiguredUsageStore { + records: Arc::new(Mutex::new(Vec::new())), + queue: flaky_queue, + }; + let runtime = UsageRuntime::new(config).expect("usage runtime should build"); + let plan = ExecutionPlan { + request_id: "req-lifecycle-queue-failure-1".to_string(), + candidate_id: Some("cand-lifecycle-queue-failure-1".to_string()), + provider_name: Some("openai".to_string()), + provider_id: "provider-1".to_string(), + endpoint_id: "endpoint-1".to_string(), + key_id: "key-1".to_string(), + method: "POST".to_string(), + url: "https://example.com/v1/responses".to_string(), + headers: BTreeMap::new(), + content_type: Some("application/json".to_string()), + content_encoding: None, + body: RequestBody::from_json(json!({"model": "gpt-5"})), + stream: false, + client_api_format: "openai:responses".to_string(), + provider_api_format: "openai:responses".to_string(), + model_name: Some("gpt-5".to_string()), + proxy: None, + transport_profile: None, + timeouts: None, + }; + + runtime.record_pending(&store, build_lifecycle_usage_seed(&plan, None)); + sleep(Duration::from_millis(50)).await; + + assert!( + store.records.lock().expect("records lock").is_empty(), + "lifecycle enqueue failure must not fall back to direct DB writes" + ); + } + + #[tokio::test] + async fn lifecycle_enqueue_failure_opens_short_circuit() { + let config = UsageRuntimeConfig { + enabled: true, + queue_lifecycle_events: true, + stream_key: "usage:events:test:lifecycle-circuit".to_string(), + consumer_group: "usage_consumers_test_lifecycle_circuit".to_string(), + consumer_block_ms: 1, + ..UsageRuntimeConfig::default() + }; + let inner_queue: Arc = + Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default())); + let flaky_queue = Arc::new(FlakyAppendQueueStore { + inner: Arc::clone(&inner_queue), + remaining_failures: AtomicUsize::new(usize::MAX), + append_attempts: AtomicUsize::new(0), + }); + let store = CloneQueueConfiguredUsageStore { + records: Arc::new(Mutex::new(Vec::new())), + queue: flaky_queue.clone(), + }; + let runtime = UsageRuntime::new(config).expect("usage runtime should build"); + for index in 0..10 { + let event = UsageEvent::new( + UsageEventType::Pending, + format!("req-lifecycle-circuit-{index}"), + UsageEventData { + provider_name: "openai".to_string(), + model: "gpt-5".to_string(), + ..UsageEventData::default() + }, + ); + runtime.enqueue_lifecycle_event(&store, event).await; + } + + assert_eq!( + flaky_queue.append_attempts.load(Ordering::Acquire), + 1, + "lifecycle enqueue circuit should prevent repeated Redis appends after a failure" + ); + assert!( + store.records.lock().expect("records lock").is_empty(), + "lifecycle enqueue circuit must not fall back to direct DB writes" + ); + } + + #[tokio::test] + async fn disabled_lifecycle_queue_does_not_write_directly() { + let config = UsageRuntimeConfig { + enabled: true, + queue_lifecycle_events: false, + stream_key: "usage:events:test:lifecycle-disabled".to_string(), + consumer_group: "usage_consumers_test_lifecycle_disabled".to_string(), + consumer_block_ms: 1, + ..UsageRuntimeConfig::default() + }; + let queue_runner: Arc = + Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default())); + let store = CloneQueueConfiguredUsageStore { + records: Arc::new(Mutex::new(Vec::new())), + queue: queue_runner, + }; + let runtime = UsageRuntime::new(config).expect("usage runtime should build"); + let plan = ExecutionPlan { + request_id: "req-lifecycle-disabled-1".to_string(), + candidate_id: Some("cand-lifecycle-disabled-1".to_string()), + provider_name: Some("openai".to_string()), + provider_id: "provider-1".to_string(), + endpoint_id: "endpoint-1".to_string(), + key_id: "key-1".to_string(), + method: "POST".to_string(), + url: "https://example.com/v1/responses".to_string(), + headers: BTreeMap::new(), + content_type: Some("application/json".to_string()), + content_encoding: None, + body: RequestBody::from_json(json!({"model": "gpt-5"})), + stream: false, + client_api_format: "openai:responses".to_string(), + provider_api_format: "openai:responses".to_string(), + model_name: Some("gpt-5".to_string()), + proxy: None, + transport_profile: None, + timeouts: None, + }; + + runtime.record_pending(&store, build_lifecycle_usage_seed(&plan, None)); + sleep(Duration::from_millis(50)).await; + + assert!( + store.records.lock().expect("records lock").is_empty(), + "disabled lifecycle queue must not fall back to direct DB writes" + ); + } + + #[tokio::test] + async fn queued_terminal_usage_does_not_enrich_or_write_before_enqueue() { + let config = UsageRuntimeConfig { + enabled: true, + queue_terminal_events: true, + stream_key: "usage:events:test:terminal".to_string(), + consumer_group: "usage_consumers_test_terminal".to_string(), + consumer_block_ms: 1, + ..UsageRuntimeConfig::default() + }; + let queue_runner: Arc = + Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default())); + let queue = UsageQueue::new(Arc::clone(&queue_runner), config.clone()) + .expect("usage queue should build"); + queue + .ensure_consumer_group() + .await + .expect("consumer group should initialize"); + let store = EnrichmentCountingQueueStore { + records: Mutex::new(Vec::new()), + queue: queue_runner, + enrich_calls: AtomicUsize::new(0), + }; + let runtime = UsageRuntime::new(config).expect("usage runtime should build"); + let event = UsageEvent::new( + UsageEventType::Completed, + "req-terminal-queue-1", + UsageEventData { + user_id: Some("user-terminal-queue-1".to_string()), + provider_name: "openai".to_string(), + model: "gpt-5".to_string(), + input_tokens: Some(4), + output_tokens: Some(8), + total_tokens: Some(12), + status_code: Some(200), + ..UsageEventData::default() + }, + ); + + runtime.record_terminal_event(&store, event).await; + + assert_eq!(store.enrich_calls.load(Ordering::Acquire), 0); + assert!(store.records.lock().expect("records lock").is_empty()); + let entries = queue + .read_group("usage-test-terminal-consumer") + .await + .expect("queue read should succeed"); + let entry = entries.first().expect("terminal event should be queued"); + let queued = UsageEvent::from_stream_fields(&entry.fields) + .expect("queued terminal event should parse"); + assert_eq!(queued.request_id, "req-terminal-queue-1"); + assert_eq!(queued.event_type, UsageEventType::Completed); + assert_eq!(queued.data.total_cost_usd, None); + } + + #[tokio::test] + async fn queue_append_failure_retries_locally_without_direct_write() { + let config = UsageRuntimeConfig { + enabled: true, + queue_terminal_events: true, + stream_key: "usage:events:test:retry".to_string(), + consumer_group: "usage_consumers_test_retry".to_string(), + consumer_block_ms: 1, + enqueue_retry_initial_backoff_ms: 1, + enqueue_retry_max_backoff_ms: 5, + enqueue_retry_buffer_capacity: 16, + enqueue_retry_workers: 1, + ..UsageRuntimeConfig::default() + }; + let inner_queue: Arc = + Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default())); + let flaky_queue: Arc = Arc::new(FlakyAppendQueueStore { + inner: Arc::clone(&inner_queue), + remaining_failures: AtomicUsize::new(1), + append_attempts: AtomicUsize::new(0), + }); + let queue = UsageQueue::new(Arc::clone(&inner_queue), config.clone()) + .expect("usage queue should build"); + queue + .ensure_consumer_group() + .await + .expect("consumer group should initialize"); + let store = EnrichmentCountingQueueStore { + records: Mutex::new(Vec::new()), + queue: flaky_queue, + enrich_calls: AtomicUsize::new(0), + }; + let runtime = UsageRuntime::new(config).expect("usage runtime should build"); + let event = UsageEvent::new( + UsageEventType::Completed, + "req-terminal-retry-1", + UsageEventData { + user_id: Some("user-terminal-retry-1".to_string()), + provider_name: "openai".to_string(), + model: "gpt-5".to_string(), + total_tokens: Some(12), + status_code: Some(200), + ..UsageEventData::default() + }, + ); + + runtime.record_terminal_event(&store, event).await; + + assert_eq!(store.enrich_calls.load(Ordering::Acquire), 0); + assert!(store.records.lock().expect("records lock").is_empty()); + for _ in 0..50 { + let entries = queue + .read_group("usage-test-retry-consumer") + .await + .expect("queue read should succeed"); + if let Some(entry) = entries.into_iter().next() { + let event = UsageEvent::from_stream_fields(&entry.fields) + .expect("queued retry event should parse"); + assert_eq!(event.request_id, "req-terminal-retry-1"); + return; + } + sleep(Duration::from_millis(10)).await; + } + + panic!("terminal usage event was not retried into queue"); + } + + #[tokio::test] + async fn body_capture_policy_read_failure_is_short_cached_as_default() { + let runtime = UsageRuntime::new(UsageRuntimeConfig { + enabled: true, + ..UsageRuntimeConfig::default() + }) + .expect("usage runtime should build"); + let store = FailingPolicyUsageStore::default(); + + let first = runtime.body_capture_policy_for(&store).await; + assert!(first.is_err(), "first failed read should be surfaced"); + + let second = runtime + .body_capture_policy_for(&store) + .await + .expect("short cached fallback should be returned"); + let third = runtime + .body_capture_policy_for(&store) + .await + .expect("short cached fallback should be reused"); + + assert_eq!(second, UsageBodyCapturePolicy::default()); + assert_eq!(third, UsageBodyCapturePolicy::default()); + assert_eq!( + store.policy_reads.load(Ordering::Acquire), + 1, + "policy read failures should be short cached to avoid concurrent DB storms" + ); + } + #[test] fn basic_request_record_level_strips_body_capture_but_preserves_derived_fields() { let mut event = UsageEvent::new( diff --git a/crates/aether-usage-runtime/src/worker.rs b/crates/aether-usage-runtime/src/worker.rs index 4714f7c9e..fb42c27e5 100644 --- a/crates/aether-usage-runtime/src/worker.rs +++ b/crates/aether-usage-runtime/src/worker.rs @@ -8,9 +8,10 @@ use async_trait::async_trait; use tracing::warn; use crate::executor::spawn_on_usage_background_runtime; +use crate::runtime::UsageBillingEventEnricher; use crate::{ - build_upsert_usage_record_from_event, settle_usage_if_needed, UsageEvent, UsageQueue, - UsageRuntimeConfig, UsageSettlementWriter, + build_upsert_usage_record_from_event, settle_usage_if_needed, UsageEvent, UsageEventType, + UsageQueue, UsageRuntimeConfig, UsageSettlementWriter, }; #[async_trait] @@ -50,10 +51,17 @@ impl UsageDataEventRecorder { #[async_trait] impl UsageEventRecorder for UsageDataEventRecorder where - T: UsageRecordWriter + UsageSettlementWriter + ManualProxyNodeCounter + Send + Sync, + T: UsageRecordWriter + + UsageSettlementWriter + + UsageBillingEventEnricher + + ManualProxyNodeCounter + + Send + + Sync, { async fn record_usage_event(&self, event: &UsageEvent) -> Result<(), DataLayerError> { - write_event_record(self.data.as_ref(), event).await + let mut event = event.clone(); + enrich_terminal_event(self.data.as_ref(), &mut event).await; + write_event_record(self.data.as_ref(), &event).await } } @@ -276,7 +284,13 @@ pub fn build_usage_queue_worker( config: UsageRuntimeConfig, ) -> Result where - T: UsageRecordWriter + UsageSettlementWriter + ManualProxyNodeCounter + Send + Sync + 'static, + T: UsageRecordWriter + + UsageSettlementWriter + + UsageBillingEventEnricher + + ManualProxyNodeCounter + + Send + + Sync + + 'static, { UsageQueueWorker::new(runner, Arc::new(UsageDataEventRecorder::new(data)), config) } @@ -293,6 +307,29 @@ where Ok(()) } +async fn enrich_terminal_event(data: &T, event: &mut UsageEvent) +where + T: UsageBillingEventEnricher + Send + Sync, +{ + if !matches!( + event.event_type, + UsageEventType::Completed | UsageEventType::Failed | UsageEventType::Cancelled + ) { + return; + } + + if let Err(err) = data.enrich_usage_event(event).await { + warn!( + event_name = "usage_worker_billing_enrichment_failed", + log_type = "event", + request_id = %event.request_id, + event_type = ?event.event_type, + error = %err, + "usage worker failed to enrich terminal usage event with billing" + ); + } +} + async fn increment_manual_proxy_node_from_event(data: &T, event: &UsageEvent) where T: ManualProxyNodeCounter + Send + Sync, @@ -368,6 +405,7 @@ mod tests { usage_event_record_error_is_permanent, write_event_record, ManualProxyNodeCounter, UsageEventRecorder, UsageQueueWorker, UsageRecordWriter, }; + use crate::UsageBillingEventEnricher; use crate::{ UsageEvent, UsageEventData, UsageEventType, UsageRuntimeConfig, UsageSettlementWriter, }; @@ -376,6 +414,7 @@ mod tests { struct TestUsageStore { records: Mutex>, settlements: Mutex>, + enrich_calls: Mutex>, } #[derive(Default)] @@ -471,6 +510,18 @@ mod tests { } } + #[async_trait] + impl UsageBillingEventEnricher for TestUsageStore { + async fn enrich_usage_event(&self, event: &mut UsageEvent) -> Result<(), DataLayerError> { + self.enrich_calls + .lock() + .expect("enrich calls lock") + .push(event.request_id.clone()); + event.data.total_cost_usd = Some(0.456); + Ok(()) + } + } + #[async_trait] impl UsageEventRecorder for SelectiveFailingRecorder { async fn record_usage_event(&self, event: &UsageEvent) -> Result<(), DataLayerError> { @@ -532,6 +583,52 @@ mod tests { assert_eq!(settlements[0].request_id, "req-worker-123"); } + #[tokio::test] + async fn data_event_recorder_enriches_terminal_event_before_write() { + let store = Arc::new(TestUsageStore::default()); + let recorder = super::UsageDataEventRecorder::new(Arc::clone(&store)); + let event = sample_event(); + + recorder + .record_usage_event(&event) + .await + .expect("recorder should enrich and write usage"); + + assert_eq!( + store + .enrich_calls + .lock() + .expect("enrich calls lock") + .as_slice(), + ["req-worker-123"] + ); + let records = store.records.lock().expect("records lock"); + assert_eq!(records.len(), 1); + assert_eq!(records[0].total_cost_usd, Some(0.456)); + } + + #[tokio::test] + async fn data_event_recorder_skips_enrichment_for_lifecycle_event() { + let store = Arc::new(TestUsageStore::default()); + let recorder = super::UsageDataEventRecorder::new(Arc::clone(&store)); + let mut event = sample_event(); + event.event_type = UsageEventType::Pending; + + recorder + .record_usage_event(&event) + .await + .expect("recorder should write lifecycle usage"); + + assert!(store + .enrich_calls + .lock() + .expect("enrich calls lock") + .is_empty()); + let records = store.records.lock().expect("records lock"); + assert_eq!(records.len(), 1); + assert_eq!(records[0].total_cost_usd, None); + } + #[test] fn usage_event_record_error_classifies_permanent_failures() { assert!(usage_event_record_error_is_permanent(