mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-07 18:07:47 +08:00
Merge remote-tracking branch 'origin/main' into codex/fix-antigravity-quota
This commit is contained in:
@@ -1,7 +1,9 @@
|
||||
use self::selection::{
|
||||
collect_selectable_candidates, collect_selectable_candidates_with_skip_reasons,
|
||||
collect_selectable_candidates, collect_selectable_candidates_with_skip_reasons_and_ordering,
|
||||
collect_selectable_enumerated_candidates_with_skip_reasons,
|
||||
resolve_preselection_ordering_config,
|
||||
};
|
||||
use super::config::SchedulerOrderingConfig;
|
||||
use super::state::SchedulerRuntimeState;
|
||||
|
||||
mod affinity;
|
||||
@@ -53,6 +55,10 @@ enum RequiredCapabilityMatchMode {
|
||||
Exclusive,
|
||||
}
|
||||
|
||||
/// `ordering_config` carries the request's routing-policy derived scheduler
|
||||
/// config. `None` falls back to the runtime default (system-default routing
|
||||
/// group, then legacy system-config keys).
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn list_selectable_candidates(
|
||||
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
|
||||
runtime_state: &impl SchedulerRuntimeState,
|
||||
@@ -64,6 +70,7 @@ pub(crate) async fn list_selectable_candidates(
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
now_unix_secs: u64,
|
||||
enable_model_directives: bool,
|
||||
ordering_config: Option<SchedulerOrderingConfig>,
|
||||
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
|
||||
collect_selectable_candidates(
|
||||
selection_row_source,
|
||||
@@ -76,6 +83,7 @@ pub(crate) async fn list_selectable_candidates(
|
||||
client_session_affinity,
|
||||
now_unix_secs,
|
||||
enable_model_directives,
|
||||
ordering_config,
|
||||
)
|
||||
.await
|
||||
}
|
||||
@@ -87,6 +95,7 @@ pub(crate) fn is_exact_all_skipped_by_auth_limit(
|
||||
selection::is_exact_all_skipped_by_auth_limit(selected, skipped)
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn list_selectable_candidates_with_skip_reasons(
|
||||
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
|
||||
runtime_state: &impl SchedulerRuntimeState,
|
||||
@@ -98,6 +107,7 @@ pub(crate) async fn list_selectable_candidates_with_skip_reasons(
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
now_unix_secs: u64,
|
||||
enable_model_directives: bool,
|
||||
ordering_config: Option<SchedulerOrderingConfig>,
|
||||
) -> Result<
|
||||
(
|
||||
Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
@@ -105,7 +115,7 @@ pub(crate) async fn list_selectable_candidates_with_skip_reasons(
|
||||
),
|
||||
GatewayError,
|
||||
> {
|
||||
collect_selectable_candidates_with_skip_reasons(
|
||||
collect_selectable_candidates_with_skip_reasons_and_ordering(
|
||||
selection_row_source,
|
||||
runtime_state,
|
||||
api_format,
|
||||
@@ -117,10 +127,12 @@ pub(crate) async fn list_selectable_candidates_with_skip_reasons(
|
||||
now_unix_secs,
|
||||
enable_model_directives,
|
||||
None,
|
||||
ordering_config,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn list_selectable_candidates_with_skip_reasons_for_request_operation(
|
||||
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
|
||||
runtime_state: &impl SchedulerRuntimeState,
|
||||
@@ -133,6 +145,7 @@ pub(crate) async fn list_selectable_candidates_with_skip_reasons_for_request_ope
|
||||
now_unix_secs: u64,
|
||||
enable_model_directives: bool,
|
||||
request_operation: Option<&str>,
|
||||
ordering_config: Option<SchedulerOrderingConfig>,
|
||||
) -> Result<
|
||||
(
|
||||
Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
@@ -140,7 +153,7 @@ pub(crate) async fn list_selectable_candidates_with_skip_reasons_for_request_ope
|
||||
),
|
||||
GatewayError,
|
||||
> {
|
||||
collect_selectable_candidates_with_skip_reasons(
|
||||
collect_selectable_candidates_with_skip_reasons_and_ordering(
|
||||
selection_row_source,
|
||||
runtime_state,
|
||||
api_format,
|
||||
@@ -152,6 +165,7 @@ pub(crate) async fn list_selectable_candidates_with_skip_reasons_for_request_ope
|
||||
now_unix_secs,
|
||||
enable_model_directives,
|
||||
request_operation,
|
||||
ordering_config,
|
||||
)
|
||||
.await
|
||||
}
|
||||
@@ -166,6 +180,7 @@ pub(crate) async fn list_selectable_enumerated_candidates_with_skip_reasons(
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
now_unix_secs: u64,
|
||||
ordering_config: Option<SchedulerOrderingConfig>,
|
||||
) -> Result<
|
||||
(
|
||||
Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
@@ -173,7 +188,8 @@ pub(crate) async fn list_selectable_enumerated_candidates_with_skip_reasons(
|
||||
),
|
||||
GatewayError,
|
||||
> {
|
||||
let ordering_config = runtime_state.read_scheduler_ordering_config().await?;
|
||||
let ordering_config =
|
||||
resolve_preselection_ordering_config(runtime_state, ordering_config).await?;
|
||||
let priority_affinity_key = selection::scheduling_priority_affinity_key(
|
||||
auth_snapshot,
|
||||
client_session_affinity,
|
||||
@@ -194,6 +210,7 @@ pub(crate) async fn list_selectable_enumerated_candidates_with_skip_reasons(
|
||||
.await
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn list_selectable_candidates_for_required_capability_without_requested_model(
|
||||
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
|
||||
runtime_state: &impl SchedulerRuntimeState,
|
||||
@@ -203,6 +220,7 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
now_unix_secs: u64,
|
||||
ordering_config: Option<SchedulerOrderingConfig>,
|
||||
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
|
||||
Ok(
|
||||
list_selectable_candidates_for_required_capability_without_requested_model_with_auth_limit_signal(
|
||||
@@ -214,12 +232,14 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
|
||||
auth_snapshot,
|
||||
client_session_affinity,
|
||||
now_unix_secs,
|
||||
ordering_config,
|
||||
)
|
||||
.await?
|
||||
.0,
|
||||
)
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn list_selectable_candidates_for_required_capability_without_requested_model_with_auth_limit_signal(
|
||||
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
|
||||
runtime_state: &impl SchedulerRuntimeState,
|
||||
@@ -229,6 +249,7 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
now_unix_secs: u64,
|
||||
ordering_config: Option<SchedulerOrderingConfig>,
|
||||
) -> Result<(Vec<SchedulerMinimalCandidateSelectionCandidate>, bool), GatewayError> {
|
||||
let normalized_api_format = normalize_api_format(candidate_api_format);
|
||||
if normalized_api_format.is_empty() {
|
||||
@@ -262,20 +283,22 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
|
||||
let mut all_attempts_blocked_by_auth_limit = !model_names.is_empty();
|
||||
|
||||
for global_model_name in model_names {
|
||||
let (candidates, skipped_candidates) = collect_selectable_candidates_with_skip_reasons(
|
||||
selection_row_source,
|
||||
runtime_state,
|
||||
&normalized_api_format,
|
||||
&global_model_name,
|
||||
require_streaming,
|
||||
required_capabilities.as_ref(),
|
||||
auth_snapshot,
|
||||
client_session_affinity,
|
||||
now_unix_secs,
|
||||
false,
|
||||
None,
|
||||
)
|
||||
.await?;
|
||||
let (candidates, skipped_candidates) =
|
||||
collect_selectable_candidates_with_skip_reasons_and_ordering(
|
||||
selection_row_source,
|
||||
runtime_state,
|
||||
&normalized_api_format,
|
||||
&global_model_name,
|
||||
require_streaming,
|
||||
required_capabilities.as_ref(),
|
||||
auth_snapshot,
|
||||
client_session_affinity,
|
||||
now_unix_secs,
|
||||
false,
|
||||
None,
|
||||
ordering_config,
|
||||
)
|
||||
.await?;
|
||||
all_attempts_blocked_by_auth_limit &=
|
||||
is_exact_all_skipped_by_auth_limit(&candidates, &skipped_candidates);
|
||||
match capability_mode {
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
use crate::data::auth::GatewayAuthApiKeySnapshot;
|
||||
use crate::data::candidate_selection::MinimalCandidateSelectionRowSource;
|
||||
use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL;
|
||||
use crate::scheduler::config::SchedulerSchedulingMode;
|
||||
use crate::scheduler::config::{SchedulerOrderingConfig, SchedulerSchedulingMode};
|
||||
use crate::GatewayError;
|
||||
use aether_scheduler_core::ClientSessionAffinity;
|
||||
|
||||
@@ -115,6 +115,7 @@ pub(super) async fn select_minimal_candidate(
|
||||
Ok(selected)
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(super) async fn collect_selectable_candidates(
|
||||
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
|
||||
runtime_state: &impl SchedulerRuntimeState,
|
||||
@@ -126,24 +127,32 @@ pub(super) async fn collect_selectable_candidates(
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
now_unix_secs: u64,
|
||||
enable_model_directives: bool,
|
||||
ordering_config: Option<SchedulerOrderingConfig>,
|
||||
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
|
||||
Ok(collect_selectable_candidates_with_skip_reasons(
|
||||
selection_row_source,
|
||||
runtime_state,
|
||||
api_format,
|
||||
global_model_name,
|
||||
require_streaming,
|
||||
required_capabilities,
|
||||
auth_snapshot,
|
||||
client_session_affinity,
|
||||
now_unix_secs,
|
||||
enable_model_directives,
|
||||
None,
|
||||
Ok(
|
||||
collect_selectable_candidates_with_skip_reasons_and_ordering(
|
||||
selection_row_source,
|
||||
runtime_state,
|
||||
api_format,
|
||||
global_model_name,
|
||||
require_streaming,
|
||||
required_capabilities,
|
||||
auth_snapshot,
|
||||
client_session_affinity,
|
||||
now_unix_secs,
|
||||
enable_model_directives,
|
||||
None,
|
||||
ordering_config,
|
||||
)
|
||||
.await?
|
||||
.0,
|
||||
)
|
||||
.await?
|
||||
.0)
|
||||
}
|
||||
|
||||
/// Legacy-shaped entrypoint that resolves the ordering config from the
|
||||
/// runtime state. Prefer `collect_selectable_candidates_with_skip_reasons_and_ordering`
|
||||
/// and pass the request's routing-policy config explicitly.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(super) async fn collect_selectable_candidates_with_skip_reasons(
|
||||
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
|
||||
runtime_state: &impl SchedulerRuntimeState,
|
||||
@@ -163,7 +172,59 @@ pub(super) async fn collect_selectable_candidates_with_skip_reasons(
|
||||
),
|
||||
GatewayError,
|
||||
> {
|
||||
let ordering_config = runtime_state.read_scheduler_ordering_config().await?;
|
||||
collect_selectable_candidates_with_skip_reasons_and_ordering(
|
||||
selection_row_source,
|
||||
runtime_state,
|
||||
api_format,
|
||||
global_model_name,
|
||||
require_streaming,
|
||||
required_capabilities,
|
||||
auth_snapshot,
|
||||
client_session_affinity,
|
||||
now_unix_secs,
|
||||
enable_model_directives,
|
||||
request_operation,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Resolve the ordering config for a preselection pass: the routing-policy
|
||||
/// derived config wins when the caller has one; otherwise fall back to the
|
||||
/// runtime default (system-default routing group, then legacy keys).
|
||||
pub(super) async fn resolve_preselection_ordering_config(
|
||||
runtime_state: &impl SchedulerRuntimeState,
|
||||
ordering_config: Option<SchedulerOrderingConfig>,
|
||||
) -> Result<SchedulerOrderingConfig, GatewayError> {
|
||||
match ordering_config {
|
||||
Some(config) => Ok(config),
|
||||
None => runtime_state.read_scheduler_ordering_config().await,
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(super) async fn collect_selectable_candidates_with_skip_reasons_and_ordering(
|
||||
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
|
||||
runtime_state: &impl SchedulerRuntimeState,
|
||||
api_format: &str,
|
||||
global_model_name: &str,
|
||||
require_streaming: bool,
|
||||
required_capabilities: Option<&serde_json::Value>,
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
now_unix_secs: u64,
|
||||
enable_model_directives: bool,
|
||||
request_operation: Option<&str>,
|
||||
ordering_config: Option<SchedulerOrderingConfig>,
|
||||
) -> Result<
|
||||
(
|
||||
Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
Vec<SchedulerSkippedCandidate>,
|
||||
),
|
||||
GatewayError,
|
||||
> {
|
||||
let ordering_config =
|
||||
resolve_preselection_ordering_config(runtime_state, ordering_config).await?;
|
||||
let priority_affinity_key = scheduling_priority_affinity_key(
|
||||
auth_snapshot,
|
||||
client_session_affinity,
|
||||
@@ -205,7 +266,7 @@ pub(super) async fn collect_selectable_enumerated_candidates_with_skip_reasons(
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
now_unix_secs: u64,
|
||||
ordering_config: crate::scheduler::config::SchedulerOrderingConfig,
|
||||
ordering_config: SchedulerOrderingConfig,
|
||||
priority_affinity_key: Option<&str>,
|
||||
) -> Result<
|
||||
(
|
||||
|
||||
@@ -65,6 +65,7 @@ async fn compatible_required_capability_prefers_matching_keys_without_hard_filte
|
||||
None,
|
||||
None,
|
||||
100,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.expect("selection should succeed");
|
||||
@@ -120,6 +121,7 @@ async fn exclusive_required_capability_keeps_hard_filtering_only_matching_keys()
|
||||
None,
|
||||
None,
|
||||
100,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.expect("selection should succeed");
|
||||
@@ -196,6 +198,7 @@ async fn required_capability_without_model_uses_session_scoped_affinity() {
|
||||
Some(&auth_snapshot),
|
||||
Some(&client_session_affinity),
|
||||
100,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.expect("selection should succeed");
|
||||
@@ -273,6 +276,7 @@ async fn required_capability_reports_auth_limit_signal_when_every_model_is_block
|
||||
Some(&auth_snapshot),
|
||||
None,
|
||||
100,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.expect("selection should succeed");
|
||||
|
||||
@@ -75,6 +75,7 @@ async fn collect_selectable_candidates(
|
||||
None,
|
||||
now_unix_secs,
|
||||
false,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
@@ -1,4 +1,10 @@
|
||||
use aether_data_contracts::repository::routing_profiles::RoutingGroupLookupKey;
|
||||
use aether_routing_core::{
|
||||
ResolvedRoutingPolicy, RoutingDefaultPolicy, RoutingSchedulingMode, RoutingSetPriorityMode,
|
||||
DEFAULT_STICKY_KEY_ATTEMPTS,
|
||||
};
|
||||
use aether_scheduler_core::SchedulerPriorityMode;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
@@ -10,11 +16,30 @@ pub(crate) enum SchedulerSchedulingMode {
|
||||
LoadBalance,
|
||||
}
|
||||
|
||||
impl SchedulerSchedulingMode {
|
||||
pub(crate) fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::FixedOrder => "fixed_order",
|
||||
Self::CacheAffinity => "cache_affinity",
|
||||
Self::LoadBalance => "load_balance",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn scheduler_priority_mode_as_str(mode: SchedulerPriorityMode) -> &'static str {
|
||||
match mode {
|
||||
SchedulerPriorityMode::Provider => "provider",
|
||||
SchedulerPriorityMode::GlobalKey => "global_key",
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) struct SchedulerOrderingConfig {
|
||||
pub(crate) priority_mode: SchedulerPriorityMode,
|
||||
pub(crate) scheduling_mode: SchedulerSchedulingMode,
|
||||
pub(crate) keep_priority_on_conversion: bool,
|
||||
/// Total attempts on the first-ranked (sticky) candidate before failover.
|
||||
pub(crate) sticky_key_attempts: u32,
|
||||
}
|
||||
|
||||
impl Default for SchedulerOrderingConfig {
|
||||
@@ -23,10 +48,72 @@ impl Default for SchedulerOrderingConfig {
|
||||
priority_mode: SchedulerPriorityMode::Provider,
|
||||
scheduling_mode: SchedulerSchedulingMode::CacheAffinity,
|
||||
keep_priority_on_conversion: false,
|
||||
sticky_key_attempts: DEFAULT_STICKY_KEY_ATTEMPTS,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl SchedulerOrderingConfig {
|
||||
/// Ordering config derived from a resolved routing policy. The policy is
|
||||
/// the single source of truth: no legacy system-config value is merged in.
|
||||
pub(crate) fn from_routing_policy(policy: &ResolvedRoutingPolicy) -> Self {
|
||||
Self {
|
||||
priority_mode: scheduler_priority_mode_from_routing(policy.priority_mode),
|
||||
scheduling_mode: scheduler_scheduling_mode_from_routing(policy.scheduling_mode),
|
||||
keep_priority_on_conversion: policy.keep_priority_on_conversion,
|
||||
sticky_key_attempts: policy.sticky_key_attempts,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn from_routing_default_policy(policy: &RoutingDefaultPolicy) -> Self {
|
||||
Self {
|
||||
priority_mode: scheduler_priority_mode_from_routing(policy.priority_mode),
|
||||
scheduling_mode: scheduler_scheduling_mode_from_routing(policy.scheduling_mode),
|
||||
keep_priority_on_conversion: policy.keep_priority_on_conversion,
|
||||
sticky_key_attempts: policy.sticky_key_attempts,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn to_routing_default_policy(self) -> RoutingDefaultPolicy {
|
||||
RoutingDefaultPolicy {
|
||||
priority_mode: match self.priority_mode {
|
||||
SchedulerPriorityMode::Provider => RoutingSetPriorityMode::Provider,
|
||||
SchedulerPriorityMode::GlobalKey => RoutingSetPriorityMode::GlobalKey,
|
||||
},
|
||||
scheduling_mode: match self.scheduling_mode {
|
||||
SchedulerSchedulingMode::FixedOrder => RoutingSchedulingMode::FixedOrder,
|
||||
SchedulerSchedulingMode::CacheAffinity => RoutingSchedulingMode::CacheAffinity,
|
||||
SchedulerSchedulingMode::LoadBalance => RoutingSchedulingMode::LoadBalance,
|
||||
},
|
||||
keep_priority_on_conversion: self.keep_priority_on_conversion,
|
||||
sticky_key_attempts: self.sticky_key_attempts,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn priority_mode_str(self) -> &'static str {
|
||||
scheduler_priority_mode_as_str(self.priority_mode)
|
||||
}
|
||||
|
||||
pub(crate) fn scheduling_mode_str(self) -> &'static str {
|
||||
self.scheduling_mode.as_str()
|
||||
}
|
||||
}
|
||||
|
||||
fn scheduler_priority_mode_from_routing(mode: RoutingSetPriorityMode) -> SchedulerPriorityMode {
|
||||
match mode {
|
||||
RoutingSetPriorityMode::Provider => SchedulerPriorityMode::Provider,
|
||||
RoutingSetPriorityMode::GlobalKey => SchedulerPriorityMode::GlobalKey,
|
||||
}
|
||||
}
|
||||
|
||||
fn scheduler_scheduling_mode_from_routing(mode: RoutingSchedulingMode) -> SchedulerSchedulingMode {
|
||||
match mode {
|
||||
RoutingSchedulingMode::FixedOrder => SchedulerSchedulingMode::FixedOrder,
|
||||
RoutingSchedulingMode::CacheAffinity => SchedulerSchedulingMode::CacheAffinity,
|
||||
RoutingSchedulingMode::LoadBalance => SchedulerSchedulingMode::LoadBalance,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn parse_scheduler_priority_mode(
|
||||
value: Option<&serde_json::Value>,
|
||||
) -> SchedulerPriorityMode {
|
||||
@@ -62,8 +149,62 @@ pub(crate) fn parse_scheduler_scheduling_mode(
|
||||
}
|
||||
}
|
||||
|
||||
/// Effective scheduler ordering config for requests that carry no resolved
|
||||
/// routing policy.
|
||||
///
|
||||
/// Resolution order:
|
||||
/// 1. the enabled system-default routing group's `default_policy`;
|
||||
/// 2. the legacy system-config keys (`provider_priority_mode`,
|
||||
/// `scheduling_mode`, `keep_priority_on_conversion`).
|
||||
///
|
||||
/// Step 2 only exists so deployments that never created a routing group keep
|
||||
/// their behaviour; once the legacy keys are removed this function collapses
|
||||
/// to step 1 plus `SchedulerOrderingConfig::default()`.
|
||||
pub(crate) async fn read_scheduler_ordering_config(
|
||||
state: &AppState,
|
||||
) -> Result<SchedulerOrderingConfig, GatewayError> {
|
||||
if let Some(config) = read_system_default_routing_ordering_config(state).await? {
|
||||
return Ok(config);
|
||||
}
|
||||
read_legacy_scheduler_ordering_config(state).await
|
||||
}
|
||||
|
||||
/// Ordering config from the enabled system-default routing group, if any.
|
||||
pub(crate) async fn read_system_default_routing_ordering_config(
|
||||
state: &AppState,
|
||||
) -> Result<Option<SchedulerOrderingConfig>, GatewayError> {
|
||||
let Some(group) = state
|
||||
.find_routing_group(RoutingGroupLookupKey::SystemDefault)
|
||||
.await?
|
||||
.filter(|group| group.enabled)
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let default_policy = match group.config_json.get("default_policy") {
|
||||
None | Some(serde_json::Value::Null) => RoutingDefaultPolicy::default(),
|
||||
Some(value) => match serde_json::from_value::<RoutingDefaultPolicy>(value.clone()) {
|
||||
Ok(policy) => policy,
|
||||
Err(error) => {
|
||||
warn!(
|
||||
event_name = "scheduler_system_default_routing_policy_invalid",
|
||||
log_type = "event",
|
||||
group_id = %group.id,
|
||||
error = %error,
|
||||
"system default routing group has an invalid default_policy; ignoring it"
|
||||
);
|
||||
return Ok(None);
|
||||
}
|
||||
},
|
||||
};
|
||||
Ok(Some(SchedulerOrderingConfig::from_routing_default_policy(
|
||||
&default_policy,
|
||||
)))
|
||||
}
|
||||
|
||||
/// Legacy system-config based ordering config. Kept only as a migration
|
||||
/// fallback; see `read_scheduler_ordering_config`.
|
||||
pub(crate) async fn read_legacy_scheduler_ordering_config(
|
||||
state: &AppState,
|
||||
) -> Result<SchedulerOrderingConfig, GatewayError> {
|
||||
let priority_mode = parse_scheduler_priority_mode(
|
||||
state
|
||||
@@ -87,5 +228,167 @@ pub(crate) async fn read_scheduler_ordering_config(
|
||||
priority_mode,
|
||||
scheduling_mode,
|
||||
keep_priority_on_conversion,
|
||||
// Legacy config never carried a sticky-key setting; use the routing default.
|
||||
sticky_key_attempts: DEFAULT_STICKY_KEY_ATTEMPTS,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use aether_data::repository::routing_profiles::InMemoryRoutingGroupRepository;
|
||||
use aether_data_contracts::repository::routing_profiles::{
|
||||
CreateRoutingGroupRecord, RoutingGroupLookupKey, RoutingGroupReadRepository,
|
||||
RoutingGroupWriteRepository,
|
||||
};
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
use crate::data::GatewayDataState;
|
||||
|
||||
fn legacy_values() -> [(String, serde_json::Value); 3] {
|
||||
[
|
||||
("provider_priority_mode".to_string(), json!("global_key")),
|
||||
("scheduling_mode".to_string(), json!("load_balance")),
|
||||
("keep_priority_on_conversion".to_string(), json!(true)),
|
||||
]
|
||||
}
|
||||
|
||||
async fn create_system_default(
|
||||
repository: &InMemoryRoutingGroupRepository,
|
||||
enabled: bool,
|
||||
config_json: serde_json::Value,
|
||||
) {
|
||||
repository
|
||||
.create_routing_group(CreateRoutingGroupRecord {
|
||||
id: "system-default".to_string(),
|
||||
name: "system-default".to_string(),
|
||||
description: None,
|
||||
enabled,
|
||||
is_system_default: true,
|
||||
config_json,
|
||||
version: 1,
|
||||
created_at: 1,
|
||||
updated_at: 1,
|
||||
published_at: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn system_default_routing_group_overrides_legacy_keys() {
|
||||
let repository = Arc::new(InMemoryRoutingGroupRepository::default());
|
||||
create_system_default(
|
||||
&repository,
|
||||
true,
|
||||
json!({
|
||||
"default_policy": {
|
||||
"priority_mode": "provider",
|
||||
"scheduling_mode": "fixed_order",
|
||||
"keep_priority_on_conversion": false
|
||||
}
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
let state = AppState::new().unwrap().with_data_state_for_tests(
|
||||
GatewayDataState::disabled()
|
||||
.with_system_config_values_for_tests(legacy_values())
|
||||
.with_routing_group_repository_for_tests(repository),
|
||||
);
|
||||
|
||||
let config = read_scheduler_ordering_config(&state).await.unwrap();
|
||||
|
||||
assert_eq!(config.priority_mode, SchedulerPriorityMode::Provider);
|
||||
assert_eq!(config.scheduling_mode, SchedulerSchedulingMode::FixedOrder);
|
||||
assert!(!config.keep_priority_on_conversion);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn missing_default_policy_in_system_default_group_uses_routing_defaults() {
|
||||
let repository = Arc::new(InMemoryRoutingGroupRepository::default());
|
||||
create_system_default(&repository, true, json!({})).await;
|
||||
let state = AppState::new().unwrap().with_data_state_for_tests(
|
||||
GatewayDataState::disabled()
|
||||
.with_system_config_values_for_tests(legacy_values())
|
||||
.with_routing_group_repository_for_tests(repository),
|
||||
);
|
||||
|
||||
let config = read_scheduler_ordering_config(&state).await.unwrap();
|
||||
|
||||
assert_eq!(config, SchedulerOrderingConfig::default());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn disabled_or_missing_system_default_group_falls_back_to_legacy_keys() {
|
||||
let repository = Arc::new(InMemoryRoutingGroupRepository::default());
|
||||
create_system_default(
|
||||
&repository,
|
||||
false,
|
||||
json!({"default_policy": {"scheduling_mode": "fixed_order"}}),
|
||||
)
|
||||
.await;
|
||||
let with_disabled_group = AppState::new().unwrap().with_data_state_for_tests(
|
||||
GatewayDataState::disabled()
|
||||
.with_system_config_values_for_tests(legacy_values())
|
||||
.with_routing_group_repository_for_tests(repository),
|
||||
);
|
||||
let without_repository = AppState::new().unwrap().with_data_state_for_tests(
|
||||
GatewayDataState::disabled().with_system_config_values_for_tests(legacy_values()),
|
||||
);
|
||||
|
||||
for state in [with_disabled_group, without_repository] {
|
||||
let config = read_scheduler_ordering_config(&state).await.unwrap();
|
||||
assert_eq!(config.priority_mode, SchedulerPriorityMode::GlobalKey);
|
||||
assert_eq!(config.scheduling_mode, SchedulerSchedulingMode::LoadBalance);
|
||||
assert!(config.keep_priority_on_conversion);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bootstrap_creates_system_default_group_from_legacy_keys_once() {
|
||||
let repository = Arc::new(InMemoryRoutingGroupRepository::default());
|
||||
let state = AppState::new().unwrap().with_data_state_for_tests(
|
||||
GatewayDataState::disabled()
|
||||
.with_system_config_values_for_tests(legacy_values())
|
||||
.with_routing_group_repository_for_tests(repository.clone()),
|
||||
);
|
||||
|
||||
let created = state
|
||||
.ensure_system_default_routing_group_inner()
|
||||
.await
|
||||
.unwrap()
|
||||
.expect("first bootstrap should create the system default group");
|
||||
assert!(created.enabled);
|
||||
assert!(created.is_system_default);
|
||||
assert_eq!(
|
||||
created.config_json["default_policy"],
|
||||
json!({
|
||||
"priority_mode": "global_key",
|
||||
"scheduling_mode": "load_balance",
|
||||
"keep_priority_on_conversion": true,
|
||||
"sticky_key_attempts": DEFAULT_STICKY_KEY_ATTEMPTS
|
||||
})
|
||||
);
|
||||
|
||||
let second = state
|
||||
.ensure_system_default_routing_group_inner()
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(second.is_none(), "bootstrap must be idempotent");
|
||||
assert_eq!(
|
||||
repository
|
||||
.find_routing_group(RoutingGroupLookupKey::SystemDefault)
|
||||
.await
|
||||
.unwrap()
|
||||
.map(|group| group.id),
|
||||
Some(created.id)
|
||||
);
|
||||
|
||||
let config = read_scheduler_ordering_config(&state).await.unwrap();
|
||||
assert_eq!(config.priority_mode, SchedulerPriorityMode::GlobalKey);
|
||||
assert_eq!(config.scheduling_mode, SchedulerSchedulingMode::LoadBalance);
|
||||
assert!(config.keep_priority_on_conversion);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user