From b8ea1cd076c744744983dec8a353252443eb81c5 Mon Sep 17 00:00:00 2001 From: Zyy Date: Fri, 14 Aug 2026 23:27:06 +0800 Subject: [PATCH] fix(core): normalize account scheduling load --- .../gateway-core/src/engine/credential.rs | 84 ++++++++++++++----- .../gateway-core/tests/engine/credential.rs | 39 +++++++++ 2 files changed, 101 insertions(+), 22 deletions(-) diff --git a/backend/crates/gateway-core/src/engine/credential.rs b/backend/crates/gateway-core/src/engine/credential.rs index cf53ff4e4..9c0f10eb4 100644 --- a/backend/crates/gateway-core/src/engine/credential.rs +++ b/backend/crates/gateway-core/src/engine/credential.rs @@ -1588,14 +1588,26 @@ impl AccountSelector { let candidate = match context.policy.strategy() { RotationStrategy::QuotaResetPriority => { - eligible.sort_by_key(|candidate| { + let default_concurrency = context.policy.max_concurrent_per_account(); + eligible.sort_by(|left, right| { ( - candidate.signals.quota_reset_at.is_none(), - candidate.signals.quota_reset_at, - candidate.signals.in_flight, - candidate.signals.last_started_at, - candidate.account.id().clone(), + left.signals.quota_reset_at.is_none(), + left.signals.quota_reset_at, ) + .cmp(&( + right.signals.quota_reset_at.is_none(), + right.signals.quota_reset_at, + )) + .then_with(|| { + capacity_utilization(left, default_concurrency) + .total_cmp(&capacity_utilization(right, default_concurrency)) + }) + .then_with(|| { + left.signals + .last_started_at + .cmp(&right.signals.last_started_at) + }) + .then_with(|| left.account.id().cmp(right.account.id())) }); eligible.first().copied()? } @@ -1604,9 +1616,11 @@ impl AccountSelector { let index = context.round_robin_cursor as usize % eligible.len(); eligible.get(index).copied()? } - RotationStrategy::Smart => { - select_smart_candidate(&eligible, context.round_robin_cursor)? - } + RotationStrategy::Smart => select_smart_candidate( + &eligible, + context.policy.max_concurrent_per_account(), + context.round_robin_cursor, + )?, RotationStrategy::Sticky => { eligible.sort_by_key(|candidate| { ( @@ -1676,8 +1690,18 @@ const SMART_QUOTA_WEIGHT: f64 = 0.8; const SMART_FAILURE_WEIGHT: f64 = 1.0; const SMART_LATENCY_WEIGHT: f64 = 0.5; +fn capacity_utilization(candidate: &AccountCandidate, default_concurrency: NonZeroU32) -> f64 { + f64::from(candidate.signals.in_flight) + / f64::from( + candidate + .account + .effective_concurrency(default_concurrency) + .get(), + ) +} + struct SmartNormalization { - max_in_flight: u32, + max_capacity_utilization: f64, min_quota: Option, max_quota: Option, min_latency_ms: Option, @@ -1685,11 +1709,11 @@ struct SmartNormalization { } impl SmartNormalization { - fn from_candidates(candidates: &[&AccountCandidate]) -> Self { - let max_in_flight = candidates + fn from_candidates(candidates: &[&AccountCandidate], default_concurrency: NonZeroU32) -> Self { + let max_capacity_utilization = candidates .iter() - .map(|candidate| candidate.signals.in_flight) - .max() + .map(|candidate| capacity_utilization(candidate, default_concurrency)) + .reduce(f64::max) .unwrap_or_default(); let min_quota = candidates .iter() @@ -1710,7 +1734,7 @@ impl SmartNormalization { .filter(|latency| *latency > 0) .max(); Self { - max_in_flight, + max_capacity_utilization, min_quota, max_quota, min_latency_ms, @@ -1721,12 +1745,18 @@ impl SmartNormalization { fn select_smart_candidate<'a>( candidates: &[&'a AccountCandidate], + default_concurrency: NonZeroU32, cursor: u64, ) -> Option<&'a AccountCandidate> { - let normalization = SmartNormalization::from_candidates(candidates); + let normalization = SmartNormalization::from_candidates(candidates, default_concurrency); let mut ranked = candidates .iter() - .map(|candidate| (*candidate, smart_score(candidate, &normalization))) + .map(|candidate| { + ( + *candidate, + smart_score(candidate, default_concurrency, &normalization), + ) + }) .collect::>(); ranked.sort_by(|(left, left_score), (right, right_score)| { right_score @@ -1742,11 +1772,14 @@ fn select_smart_candidate<'a>( Some(ranked[index].0) } -fn smart_score(candidate: &AccountCandidate, normalization: &SmartNormalization) -> f64 { - let load = lower_is_better( - u64::from(candidate.signals.in_flight), - 0, - u64::from(normalization.max_in_flight), +fn smart_score( + candidate: &AccountCandidate, + default_concurrency: NonZeroU32, + normalization: &SmartNormalization, +) -> f64 { + let load = lower_ratio_is_better( + capacity_utilization(candidate, default_concurrency), + normalization.max_capacity_utilization, ); let quota = candidate.signals.quota_remaining_rank.map_or(0.5, |quota| { higher_is_better( @@ -1781,6 +1814,13 @@ fn smart_score(candidate: &AccountCandidate, normalization: &SmartNormalization) + SMART_LATENCY_WEIGHT * latency } +fn lower_ratio_is_better(value: f64, maximum: f64) -> f64 { + if maximum <= 0.0 { + return 1.0; + } + 1.0 - (value / maximum).clamp(0.0, 1.0) +} + fn lower_is_better(value: u64, minimum: u64, maximum: u64) -> f64 { if maximum <= minimum { return 1.0; diff --git a/backend/crates/gateway-core/tests/engine/credential.rs b/backend/crates/gateway-core/tests/engine/credential.rs index 878e3da04..6e8e47717 100644 --- a/backend/crates/gateway-core/tests/engine/credential.rs +++ b/backend/crates/gateway-core/tests/engine/credential.rs @@ -51,6 +51,15 @@ fn candidate(id: &str, in_flight: u32, remaining: Option) -> AccountCandida } } +fn candidate_with_concurrency(id: &str, in_flight: u32, concurrency: u32) -> AccountCandidate { + let mut candidate = candidate(id, in_flight, Some(100)); + candidate.account = candidate.account.with_scheduling( + Some(AccountConcurrencyLimit::new(concurrency).expect("concurrency override")), + AccountWeight::DEFAULT, + ); + candidate +} + fn weighted_candidate(id: &str, weight: u16, in_flight: u32) -> AccountCandidate { let mut candidate = candidate(id, in_flight, None); candidate.account = candidate.account.with_scheduling( @@ -348,6 +357,20 @@ fn smart_selector_should_prefer_lower_inflight_count() { assert_eq!(selected.candidate().account.id().as_str(), "acct_idle"); } +#[test] +fn smart_selector_should_prefer_lower_capacity_utilization_over_lower_inflight_count() { + let candidates = [ + candidate_with_concurrency("acct_smaller", 1, 2), + candidate_with_concurrency("acct_larger", 2, 10), + ]; + + let selected = AccountSelector + .select(&candidates, &context(RotationStrategy::Smart)) + .expect("candidate available"); + + assert_eq!(selected.candidate().account.id().as_str(), "acct_larger"); +} + #[test] fn selector_should_prioritize_the_highest_weight_for_every_rotation_strategy() { let candidates = vec![ @@ -622,6 +645,22 @@ fn quota_reset_selector_should_prefer_known_earliest_window() { assert_eq!(selected.candidate().account.id().as_str(), "acct_earlier"); } +#[test] +fn quota_reset_selector_should_use_capacity_utilization_as_load_tiebreaker() { + let reset_at = SystemTime::now() + Duration::from_secs(60); + let mut smaller = candidate_with_concurrency("acct_smaller", 1, 2); + smaller.signals.quota_reset_at = Some(reset_at); + let mut larger = candidate_with_concurrency("acct_larger", 2, 10); + larger.signals.quota_reset_at = Some(reset_at); + let candidates = [smaller, larger]; + + let selected = AccountSelector + .select(&candidates, &context(RotationStrategy::QuotaResetPriority)) + .expect("candidate available"); + + assert_eq!(selected.candidate().account.id().as_str(), "acct_larger"); +} + #[test] fn smart_selector_should_prefer_known_quota_over_unknown_after_load_is_equal() { let candidates = vec![