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

fix: harden concurrency limits and high-RPM runtime paths

Browse Source 
Bound request, stream, queue, and shutdown resource lifetimes. Reduce scheduler and Redis hot-path work and isolate database maintenance. Include regression coverage, load probes, and concurrency audit results.
This commit is contained in:
elky
2026-09-10 08:14:58 +08:00
parent 361952ada9
commit ecc16673eb
149 changed files with 27963 additions and 1926 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/ai_serving/planner/candidate_materialization.rs
+109 -9
View File
@@ -996,7 +996,7 @@ impl<'a> RequestedModelAttemptPageCursor<'a> {
);
if page_is_exact_auth_api_key_concurrency_limited(&page) {
if self.wait_for_auth_api_key_concurrency_retry().await {
if self.wait_for_auth_api_key_concurrency_retry().await? {
continue;
}
self.persist_final_auth_api_key_concurrency_skips(page.skipped_candidates)
@@ -1087,20 +1087,23 @@ impl<'a> RequestedModelAttemptPageCursor<'a> {
}
}
async fn wait_for_auth_api_key_concurrency_retry(&mut self) -> bool {
async fn wait_for_auth_api_key_concurrency_retry(&mut self) -> Result<bool, GatewayError> {
let now = Instant::now();
let deadline = *self
.auth_api_key_concurrency_wait_deadline
.get_or_insert(now + AUTH_API_KEY_CONCURRENCY_WAIT_BUDGET);
if now >= deadline {
return false;
if !crate::scheduler::candidate::wait_for_auth_api_key_concurrency_retry(
self.state.app(),
Some(&self.auth_snapshot),
deadline,
AUTH_API_KEY_CONCURRENCY_RETRY_DELAY,
)
.await?
{
return Ok(false);
}
let sleep_duration =
AUTH_API_KEY_CONCURRENCY_RETRY_DELAY.min(deadline.saturating_duration_since(now));
tokio::time::sleep(sleep_duration).await;
self.page_cursor.restart_scan();
true
Ok(true)
}
async fn persist_final_auth_api_key_concurrency_skips(
@@ -2289,6 +2292,103 @@ mod tests {
candidate
}
#[tokio::test]
async fn auth_concurrency_wait_paged_scan_retries_once_at_original_deadline() {
let now = current_unix_ms();
let active = serde_json::from_value(json!({
"id": "active-candidate",
"request_id": "active-request",
"api_key_id": "api-key-1",
"candidate_index": 0,
"retry_index": 0,
"status": "pending",
"is_cached": false,
"created_at_unix_ms": now,
"started_at_unix_ms": now
}))
.expect("active candidate should build");
let repository = Arc::new(InMemoryRequestCandidateRepository::seed([active]));
let app = AppState::new()
.expect("state should build")
.with_data_state_for_tests(
GatewayDataState::with_request_candidate_repository_for_tests(repository),
);
let mut auth_snapshot = sample_auth_snapshot();
auth_snapshot.api_key_concurrent_limit = Some(1);
let page_cursor = LocalCandidatePreselectionPageCursor::new(
PlannerAppState::new(&app),
&crate::system_features::ModelDirectivePolicySnapshot::default(),
"openai:chat",
"gpt-5",
None,
true,
None,
&auth_snapshot,
None,
None,
None,
false,
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModel,
true,
Some("trace-auth-wait"),
)
.await;
let mut cursor = RequestedModelAttemptPageCursor {
state: PlannerAppState::new(&app),
trace_id: "trace-auth-wait".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 auth wait",
record_runtime_miss_diagnostic: false,
resolution_mode: LocalCandidateResolutionMode::Standard,
decorate_skipped_candidate: Arc::new(identity_skipped_candidate),
page_cursor,
pending_items: VecDeque::new(),
skipped_provider_ids: BTreeSet::new(),
skipped_endpoint_ids: BTreeSet::new(),
skipped_credential_ids: BTreeSet::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,
};
let started = Instant::now();
let mut scan_restarts = 0;
while cursor
.wait_for_auth_api_key_concurrency_retry()
.await
.expect("auth wait should succeed")
{
scan_restarts += 1;
}
assert_eq!(
scan_restarts, 1,
"blocked polls must not restart page scans"
);
assert!(started.elapsed() >= AUTH_API_KEY_CONCURRENCY_WAIT_BUDGET);
let original_deadline = cursor.auth_api_key_concurrency_wait_deadline;
assert!(!cursor
.wait_for_auth_api_key_concurrency_retry()
.await
.expect("expired auth wait should succeed"));
assert_eq!(
cursor.auth_api_key_concurrency_wait_deadline,
original_deadline
);
}
#[tokio::test]
async fn pool_group_keys_are_not_persisted_as_available_before_attempt() {
let repository = Arc::new(InMemoryRequestCandidateRepository::default());
apps/aether-gateway/src/ai_serving/planner/state/scheduler.rs
+24 -43
View File
@@ -1,9 +1,7 @@
use aether_scheduler_core::{ClientSessionAffinity, SchedulerMinimalCandidateSelectionCandidate};
use std::time::Duration;
use tokio::time::Instant;
use super::{GatewayAuthApiKeySnapshot, PlannerAppState};
use crate::clock::current_unix_secs;
use crate::constants::{
API_KEY_CONCURRENCY_WAIT_POLL_INTERVAL_MS, API_KEY_CONCURRENCY_WAIT_TIMEOUT_MS,
};
@@ -97,11 +95,13 @@ impl<'a> PlannerAppState<'a> {
),
GatewayError,
> {
let wait_timeout = Duration::from_millis(API_KEY_CONCURRENCY_WAIT_TIMEOUT_MS);
let wait_interval = Duration::from_millis(API_KEY_CONCURRENCY_WAIT_POLL_INTERVAL_MS.max(1));
let wait_deadline = Instant::now() + wait_timeout;
let mut attempt_now_unix_secs = now_unix_secs;
loop {
crate::scheduler::candidate::select_with_auth_concurrency_wait(
self.app(),
auth_snapshot,
now_unix_secs,
Duration::from_millis(API_KEY_CONCURRENCY_WAIT_TIMEOUT_MS),
Duration::from_millis(API_KEY_CONCURRENCY_WAIT_POLL_INTERVAL_MS),
|attempt_now_unix_secs| async move {
let result = crate::scheduler::candidate::list_selectable_candidates_with_skip_reasons_for_request_operation(
self.app().data.as_ref(),
self.app(),
@@ -118,21 +118,13 @@ impl<'a> PlannerAppState<'a> {
)
.await?;
if !crate::scheduler::candidate::is_exact_all_skipped_by_auth_limit(
let auth_limit_blocked = crate::scheduler::candidate::is_exact_all_skipped_by_auth_limit(
&result.0, &result.1,
) {
return Ok(result);
}
let now = Instant::now();
if now >= wait_deadline {
return Ok(result);
}
let remaining = wait_deadline.duration_since(now);
tokio::time::sleep(wait_interval.min(remaining)).await;
attempt_now_unix_secs = current_unix_secs();
}
);
Ok((result, auth_limit_blocked))
},
)
.await
}
#[allow(clippy::too_many_arguments)]
@@ -178,13 +170,14 @@ impl<'a> PlannerAppState<'a> {
now_unix_secs: u64,
ordering_config: SchedulerOrderingConfig,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
let wait_timeout = Duration::from_millis(API_KEY_CONCURRENCY_WAIT_TIMEOUT_MS);
let wait_interval = Duration::from_millis(API_KEY_CONCURRENCY_WAIT_POLL_INTERVAL_MS.max(1));
let wait_deadline = Instant::now() + wait_timeout;
let mut attempt_now_unix_secs = now_unix_secs;
loop {
let (result, auth_limit_blocked) = crate::scheduler::candidate::list_selectable_candidates_for_required_capability_without_requested_model_with_auth_limit_signal(
crate::scheduler::candidate::select_with_auth_concurrency_wait(
self.app(),
auth_snapshot,
now_unix_secs,
Duration::from_millis(API_KEY_CONCURRENCY_WAIT_TIMEOUT_MS),
Duration::from_millis(API_KEY_CONCURRENCY_WAIT_POLL_INTERVAL_MS),
|attempt_now_unix_secs| {
crate::scheduler::candidate::list_selectable_candidates_for_required_capability_without_requested_model_with_auth_limit_signal(
self.app().data.as_ref(),
self.app(),
candidate_api_format,
@@ -195,20 +188,8 @@ impl<'a> PlannerAppState<'a> {
attempt_now_unix_secs,
ordering_config,
)
.await?;
if !auth_limit_blocked {
return Ok(result);
}
let now = Instant::now();
if now >= wait_deadline {
return Ok(result);
}
let remaining = wait_deadline.duration_since(now);
tokio::time::sleep(wait_interval.min(remaining)).await;
attempt_now_unix_secs = current_unix_secs();
}
},
)
.await
}
}