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

Restrict scheduler affinity to cache affinity mode

Browse Source
This commit is contained in:
fawney19
2026-05-11 14:06:49 +08:00
parent e91c874863
commit 247ea9d1bd
12 changed files with 515 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files

52
apps/aether-gateway/src/handlers/admin/provider/pool/config.rs
View File

@@ -215,6 +215,28 @@ fn admin_provider_pool_lru_enabled(
.any(|item| item.enabled && item.preset.eq_ignore_ascii_case("lru"))
}
pub(crate) fn admin_provider_pool_cache_affinity_enabled(
pool_config: &AdminProviderPoolConfig,
) -> bool {
let mut seen = std::collections::BTreeSet::new();
for item in &pool_config.scheduling_presets {
let preset = item.preset.trim().to_ascii_lowercase();
if preset.is_empty() || !seen.insert(preset.clone()) {
continue;
}
if !item.enabled {
continue;
}
if matches!(
preset.as_str(),
"lru" | "cache_affinity" | "load_balance" | "single_account"
) {
return preset == "cache_affinity";
}
}
false
}
pub(crate) fn admin_provider_pool_config(
provider: &aether_data_contracts::repository::provider_catalog::StoredProviderCatalogProvider,
) -> Option<AdminProviderPoolConfig> {
@@ -331,7 +353,10 @@ pub(crate) fn admin_provider_pool_config_from_config_value(
#[cfg(test)]
mod tests {
use super::{admin_provider_pool_config, admin_provider_pool_config_from_config_value};
use super::{
admin_provider_pool_cache_affinity_enabled, admin_provider_pool_config,
admin_provider_pool_config_from_config_value,
};
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogProvider;
use serde_json::json;
@@ -551,4 +576,29 @@ mod tests {
assert_eq!(config.unschedulable_rules[1].keyword, "review_required");
assert_eq!(config.unschedulable_rules[1].duration_minutes, 5);
}
#[test]
fn cache_affinity_enabled_only_when_it_is_distribution_mode() {
let cache_affinity = admin_provider_pool_config_from_config_value(Some(&json!({
"pool_advanced": {
"scheduling_presets": [
{"preset": "cache_affinity", "enabled": true},
{"preset": "priority_first", "enabled": true}
]
}
})))
.expect("pool config should parse");
assert!(admin_provider_pool_cache_affinity_enabled(&cache_affinity));
let load_balance = admin_provider_pool_config_from_config_value(Some(&json!({
"pool_advanced": {
"scheduling_presets": [
{"preset": "load_balance", "enabled": true},
{"preset": "cache_affinity", "enabled": true}
]
}
})))
.expect("pool config should parse");
assert!(!admin_provider_pool_cache_affinity_enabled(&load_balance));
}
}

39
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/reads.rs
View File

@@ -3,6 +3,7 @@ use super::keys::{
pool_cooldown_keys, pool_cost_keys, pool_latency_keys, pool_lru_key, pool_sticky_key,
pool_sticky_pattern,
};
use crate::handlers::admin::provider::pool::config::admin_provider_pool_cache_affinity_enabled;
use crate::handlers::admin::provider::shared::support::{
AdminProviderPoolConfig, AdminProviderPoolRuntimeState,
};
@@ -44,11 +45,13 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
let cooldown_keys = pool_cooldown_keys(provider_id, key_ids);
let cost_keys = pool_cost_keys(provider_id, key_ids);
let latency_keys = pool_latency_keys(provider_id, key_ids);
let sticky_sessions_enabled = pool_config.sticky_session_ttl_seconds > 0
&& admin_provider_pool_cache_affinity_enabled(pool_config);
if let Some(sticky_session_token) = sticky_session_token
.map(str::trim)
.filter(|value| !value.is_empty())
.filter(|_| pool_config.sticky_session_ttl_seconds > 0)
.filter(|_| sticky_sessions_enabled)
{
let sticky_key = pool_sticky_key(provider_id, sticky_session_token);
if let Ok(Some(bound_key_id)) = runtime.kv_get(&sticky_key).await {
@@ -77,22 +80,24 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
}
}
let sticky_keys = runtime
.scan_keys(&pool_sticky_pattern(provider_id), 200)
.await
.unwrap_or_default();
state.total_sticky_sessions = sticky_keys.len();
if !sticky_keys.is_empty() {
let raw_keys = sticky_keys
.iter()
.map(|key| runtime.strip_namespace(key).to_string())
.collect::<Vec<_>>();
if let Ok(values) = runtime.kv_get_many(&raw_keys).await {
for bound_key_id in values.into_iter().flatten() {
*state
.sticky_sessions_by_key
.entry(bound_key_id)
.or_insert(0) += 1;
if sticky_sessions_enabled {
let sticky_keys = runtime
.scan_keys(&pool_sticky_pattern(provider_id), 200)
.await
.unwrap_or_default();
state.total_sticky_sessions = sticky_keys.len();
if !sticky_keys.is_empty() {
let raw_keys = sticky_keys
.iter()
.map(|key| runtime.strip_namespace(key).to_string())
.collect::<Vec<_>>();
if let Ok(values) = runtime.kv_get_many(&raw_keys).await {
for bound_key_id in values.into_iter().flatten() {
*state
.sticky_sessions_by_key
.entry(bound_key_id)
.or_insert(0) += 1;
}
}
}
}

46
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/writes.rs
View File

@@ -2,6 +2,7 @@ use super::keys::{
pool_cooldown_index_key, pool_cooldown_key, pool_cost_key, pool_latency_key, pool_lru_key,
pool_sticky_key, pool_stream_timeout_key,
};
use crate::handlers::admin::provider::pool::config::admin_provider_pool_cache_affinity_enabled;
use crate::handlers::admin::provider::shared::support::{
AdminProviderPoolConfig, AdminProviderPoolUnschedulableRule,
};
@@ -346,6 +347,7 @@ pub(crate) async fn record_admin_provider_pool_success(
.map(str::trim)
.filter(|value| !value.is_empty())
.filter(|_| pool_config.sticky_session_ttl_seconds > 0)
.filter(|_| admin_provider_pool_cache_affinity_enabled(pool_config))
{
let _ = runtime
.kv_set(
@@ -810,6 +812,50 @@ mod tests {
assert!(runtime.lru_score_by_key.contains_key("key-1"));
}
#[tokio::test]
async fn success_feedback_does_not_write_sticky_without_cache_affinity() {
let Some(redis) = start_managed_redis_or_skip().await else {
return;
};
let app = build_runner_app(redis.redis_url(), "pool_runtime_no_sticky_load_balance").await;
let runtime = app.runtime_state.as_ref();
let mut pool_config = sample_pool_config();
pool_config.scheduling_presets = vec![AdminProviderPoolSchedulingPreset {
preset: "load_balance".to_string(),
enabled: true,
mode: None,
}];
pool_config.lru_enabled = false;
let key_ids = vec!["key-1".to_string()];
record_admin_provider_pool_success(
runtime,
"provider-1",
"key-1",
&pool_config,
Some("session-1"),
120,
Some(80),
)
.await;
let runtime = read_admin_provider_pool_runtime_state(
runtime,
"provider-1",
&key_ids,
&pool_config,
Some("session-1"),
)
.await;
assert_eq!(runtime.total_sticky_sessions, 0);
assert_eq!(runtime.sticky_bound_key_id, None);
assert_eq!(runtime.sticky_sessions_by_key.get("key-1"), None);
assert_eq!(runtime.cost_window_usage_by_key.get("key-1"), Some(&120));
assert_eq!(runtime.latency_avg_ms_by_key.get("key-1"), Some(&80.0));
assert!(runtime.lru_score_by_key.contains_key("key-1"));
}
#[tokio::test]
async fn error_feedback_respects_retry_after_for_rate_limits() {
let Some(redis) = start_managed_redis_or_skip().await else {

15
apps/aether-gateway/src/handlers/proxy/mod.rs
View File

@@ -54,6 +54,7 @@ use crate::headers::{
should_skip_request_header,
};
use crate::router::RequestAdmissionError;
use crate::scheduler::config::{read_scheduler_ordering_config, SchedulerSchedulingMode};
use crate::{
AppState, FrontdoorUserRpmOutcome, GatewayError, GatewayFallbackMetricKind,
GatewayFallbackReason, LocalExecutionRuntimeMissDiagnostic,
@@ -297,6 +298,20 @@ async fn maybe_forward_public_request_to_tunnel_owner(
}) else {
return Ok(None);
};
let cache_affinity_enabled = match read_scheduler_ordering_config(state).await {
Ok(config) => config.scheduling_mode == SchedulerSchedulingMode::CacheAffinity,
Err(err) => {
warn!(
trace_id = %request_context.trace_id,
error = ?err,
"gateway failed to load scheduler config while checking tunnel affinity forwarding mode"
);
SchedulerSchedulingMode::default() == SchedulerSchedulingMode::CacheAffinity
}
};
if !cache_affinity_enabled {
return Ok(None);
}
let Some(api_format) = decision
.auth_endpoint_signature
.as_deref()

4
apps/aether-gateway/src/handlers/shared/provider_pool.rs
View File

@@ -1,4 +1,6 @@
pub(crate) use super::super::admin::provider::pool::config::admin_provider_pool_config_from_config_value;
pub(crate) use super::super::admin::provider::pool::config::{
admin_provider_pool_cache_affinity_enabled, admin_provider_pool_config_from_config_value,
};
pub(crate) use super::super::admin::provider::pool::runtime::{
admin_provider_pool_key_circuit_breaker_reason, read_admin_provider_pool_runtime_state,
record_admin_provider_pool_error, record_admin_provider_pool_stream_timeout,