i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-04 08:27: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/handlers/admin/provider/pool/runtime/keys.rs
-14
View File
@@ -40,20 +40,6 @@ pub(super) fn pool_stream_timeout_key(provider_id: &str, key_id: &str) -> String
format!("ap:{provider_id}:stream_timeout:{key_id}")
}
pub(super) fn parse_pool_cost_member(member: &str) -> u64 {
member
.rsplit_once(':')
.and_then(|(_, suffix)| suffix.parse::<u64>().ok())
.unwrap_or(0)
}
pub(super) fn parse_pool_latency_member(member: &str) -> u64 {
member
.rsplit_once(':')
.and_then(|(_, suffix)| suffix.parse::<u64>().ok())
.unwrap_or(0)
}
pub(super) fn pool_cooldown_keys(provider_id: &str, key_ids: &[String]) -> Vec<String> {
key_ids
.iter()
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/mod.rs
+2 -1
View File
@@ -12,7 +12,8 @@ pub(crate) use self::mutations::{
pub(crate) use self::reads::{
read_admin_provider_pool_cooldown_count, read_admin_provider_pool_cooldown_counts,
read_admin_provider_pool_cooldown_key_ids, read_admin_provider_pool_key_cooldown_reason,
read_admin_provider_pool_runtime_state,
read_admin_provider_pool_runtime_state, read_provider_pool_scheduling_runtime_state,
read_provider_pool_sticky_bound_key_id,
};
pub(crate) use self::status::build_admin_provider_pool_status_payload;
pub(crate) use self::writes::{
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/reads.rs
+560 -94
View File
@@ -1,7 +1,6 @@
use super::keys::{
parse_pool_cost_member, parse_pool_latency_member, pool_cooldown_index_key, pool_cooldown_key,
pool_cooldown_keys, pool_cost_keys, pool_latency_keys, pool_lru_key, pool_sticky_key,
pool_sticky_pattern,
pool_cooldown_index_key, pool_cooldown_key, 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::{
@@ -12,7 +11,8 @@ use crate::maintenance::PoolQuotaProbeWorkerConfig;
use crate::provider_pool_demand::{
provider_pool_burst_pending, read_provider_pool_demand_snapshot,
};
use aether_runtime_state::{DataLayerError, RuntimeState};
use aether_pool_core::{normalize_enabled_pool_presets, PoolSchedulingPreset};
use aether_runtime_state::{DataLayerError, RuntimeState, ScoreWindowU64Stats};
use futures_util::future::join_all;
use std::collections::{BTreeMap, BTreeSet};
use std::time::{SystemTime, UNIX_EPOCH};
@@ -48,6 +48,103 @@ fn bounded_runtime_window_metric_key_ids(key_ids: &[String], limit: usize) -> &[
&key_ids[..end]
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum PoolRuntimeReadPurpose {
Admin,
Scheduling,
}
fn scheduling_window_metrics(pool_config: &AdminProviderPoolConfig) -> (bool, bool) {
let presets = pool_config
.scheduling_presets
.iter()
.map(|preset| PoolSchedulingPreset {
preset: preset.preset.clone(),
enabled: preset.enabled,
mode: preset.mode.clone(),
})
.collect::<Vec<_>>();
let active = normalize_enabled_pool_presets(&presets);
let cost = pool_config.cost_limit_per_key_tokens.is_some()
|| active
.iter()
.any(|preset| matches!(preset.as_str(), "cost_first" | "quota_balanced"));
let latency = active.iter().any(|preset| preset == "latency_first");
(cost, latency)
}
async fn read_window_stats(
runtime: &RuntimeState,
keys: &[String],
min_score: f64,
) -> Vec<ScoreWindowU64Stats> {
let aggregates = match runtime.score_window_u64_stats_by_min(keys, min_score).await {
Ok(values) => values,
Err(err) => {
warn!(
"gateway provider pool: bounded window aggregation failed, using exact range reads: {err:?}"
);
vec![None; keys.len()]
}
};
join_all(keys.iter().zip(aggregates).map(|(key, stats)| async move {
match stats {
Some(stats) => stats,
// Large windows and failed aggregation retain the original exact
// read. A missing aggregate must never be treated as zero cost.
None => {
let members = runtime
.score_range_by_min(key, min_score)
.await
.unwrap_or_default();
ScoreWindowU64Stats::from_members(members.iter().map(String::as_str))
}
}
}))
.await
}
pub(crate) async fn read_provider_pool_sticky_bound_key_id(
runtime: &RuntimeState,
provider_id: &str,
pool_config: &AdminProviderPoolConfig,
sticky_session_token: Option<&str>,
) -> Option<String> {
if pool_config.sticky_session_ttl_seconds == 0
|| !admin_provider_pool_cache_affinity_enabled(pool_config)
{
return None;
}
let sticky_session_token = sticky_session_token
.map(str::trim)
.filter(|value| !value.is_empty())?;
let sticky_key = pool_sticky_key(provider_id, sticky_session_token);
let bound_key_id = runtime.kv_get(&sticky_key).await.ok().flatten()?;
let cooldown_key = pool_cooldown_key(provider_id, &bound_key_id);
match runtime.kv_exists(&cooldown_key).await {
Ok(false) => {
let _ = runtime
.key_expire(
&sticky_key,
std::time::Duration::from_secs(pool_config.sticky_session_ttl_seconds),
)
.await;
Some(bound_key_id)
}
Ok(true) => {
let _ = runtime.kv_delete(&sticky_key).await;
None
}
Err(err) => {
warn!(
"gateway admin provider pool: failed to validate sticky cooldown for provider {provider_id}: {:?}",
err
);
Some(bound_key_id)
}
}
}
pub(crate) async fn read_admin_provider_pool_cooldown_counts(
runtime: &RuntimeState,
provider_ids: &[String],
@@ -71,9 +168,52 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
pool_config: &AdminProviderPoolConfig,
sticky_session_token: Option<&str>,
) -> AdminProviderPoolRuntimeState {
read_provider_pool_runtime_state(
runtime,
provider_id,
key_ids,
pool_config,
sticky_session_token,
PoolRuntimeReadPurpose::Admin,
)
.await
}
pub(crate) async fn read_provider_pool_scheduling_runtime_state(
runtime: &RuntimeState,
provider_id: &str,
key_ids: &[String],
pool_config: &AdminProviderPoolConfig,
sticky_session_token: Option<&str>,
) -> AdminProviderPoolRuntimeState {
read_provider_pool_runtime_state(
runtime,
provider_id,
key_ids,
pool_config,
sticky_session_token,
PoolRuntimeReadPurpose::Scheduling,
)
.await
}
async fn read_provider_pool_runtime_state(
runtime: &RuntimeState,
provider_id: &str,
key_ids: &[String],
pool_config: &AdminProviderPoolConfig,
sticky_session_token: Option<&str>,
purpose: PoolRuntimeReadPurpose,
) -> AdminProviderPoolRuntimeState {
let include_admin_metrics = purpose == PoolRuntimeReadPurpose::Admin;
let mut state = AdminProviderPoolRuntimeState::default();
let cooldown_keys = pool_cooldown_keys(provider_id, key_ids);
let metric_key_limit = pool_runtime_window_metric_key_limit();
let metric_key_limit = if include_admin_metrics {
pool_runtime_window_metric_key_limit()
} else {
key_ids.len()
};
// The admin display cap must not hide a candidate's strict cost limit.
let metric_key_ids = bounded_runtime_window_metric_key_ids(key_ids, metric_key_limit);
if metric_key_ids.len() < key_ids.len() {
info!(
@@ -86,44 +226,33 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
"gateway limited admin pool runtime cost/latency window reads"
);
}
let cost_keys = pool_cost_keys(provider_id, metric_key_ids);
let latency_keys = pool_latency_keys(provider_id, metric_key_ids);
let (load_cost, load_latency) = if include_admin_metrics {
(true, true)
} else {
scheduling_window_metrics(pool_config)
};
let cost_keys = if load_cost {
pool_cost_keys(provider_id, metric_key_ids)
} else {
Vec::new()
};
let latency_keys = if load_latency {
pool_latency_keys(provider_id, metric_key_ids)
} else {
Vec::new()
};
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(|_| 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 {
let cooldown_key = pool_cooldown_key(provider_id, &bound_key_id);
match runtime.kv_exists(&cooldown_key).await {
Ok(false) => {
let _ = runtime
.key_expire(
&sticky_key,
std::time::Duration::from_secs(pool_config.sticky_session_ttl_seconds),
)
.await;
state.sticky_bound_key_id = Some(bound_key_id);
}
Ok(true) => {
let _ = runtime.kv_delete(&sticky_key).await;
}
Err(err) => {
warn!(
"gateway admin provider pool: failed to validate sticky cooldown for provider {provider_id}: {:?}",
err
);
state.sticky_bound_key_id = Some(bound_key_id);
}
}
}
}
state.sticky_bound_key_id = read_provider_pool_sticky_bound_key_id(
runtime,
provider_id,
pool_config,
sticky_session_token,
)
.await;
if sticky_sessions_enabled {
if include_admin_metrics && sticky_sessions_enabled {
let sticky_keys = runtime
.scan_keys(&pool_sticky_pattern(provider_id), 200)
.await
@@ -161,23 +290,25 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
.unwrap_or_default();
}
let probe_config = PoolQuotaProbeWorkerConfig::from_env();
let demand_snapshot = read_provider_pool_demand_snapshot(
runtime,
provider_id,
key_ids.len(),
probe_config.max_keys_per_provider,
)
.await;
state.provider_in_flight = demand_snapshot.in_flight;
state.provider_ema_in_flight = demand_snapshot.ema_in_flight;
state.provider_desired_hot = if pool_config.probing_enabled {
demand_snapshot.desired_hot
} else {
0
};
state.provider_burst_pending =
pool_config.probing_enabled && provider_pool_burst_pending(runtime, provider_id).await;
if include_admin_metrics || pool_config.probing_enabled {
let probe_config = PoolQuotaProbeWorkerConfig::from_env();
let demand_snapshot = read_provider_pool_demand_snapshot(
runtime,
provider_id,
key_ids.len(),
probe_config.max_keys_per_provider,
)
.await;
state.provider_in_flight = demand_snapshot.in_flight;
state.provider_ema_in_flight = demand_snapshot.ema_in_flight;
state.provider_desired_hot = if pool_config.probing_enabled {
demand_snapshot.desired_hot
} else {
0
};
state.provider_burst_pending =
pool_config.probing_enabled && provider_pool_burst_pending(runtime, provider_id).await;
}
if !cooldown_keys.is_empty() {
let cooldown_reasons = runtime
@@ -190,12 +321,14 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
{
if let Some(reason) = reason {
state.cooldown_reason_by_key.insert(key_id.clone(), reason);
if let Ok(Some(ttl)) = runtime.kv_ttl_seconds(cooldown_key).await {
if let Ok(ttl_seconds) = u64::try_from(ttl) {
if ttl_seconds > 0 {
state
.cooldown_ttl_by_key
.insert(key_id.clone(), ttl_seconds);
if include_admin_metrics {
if let Ok(Some(ttl)) = runtime.kv_ttl_seconds(cooldown_key).await {
if let Ok(ttl_seconds) = u64::try_from(ttl) {
if ttl_seconds > 0 {
state
.cooldown_ttl_by_key
.insert(key_id.clone(), ttl_seconds);
}
}
}
}
@@ -205,42 +338,24 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
let now = current_unix_secs();
let cost_window_start = now.saturating_sub(pool_config.cost_window_seconds) as f64;
let cost_results = join_all(
cost_keys
.iter()
.map(|cost_key| runtime.score_range_by_min(cost_key, cost_window_start)),
)
.await;
for (key_id, members) in metric_key_ids.iter().zip(cost_results) {
let total = members
.unwrap_or_default()
.iter()
.map(|member| parse_pool_cost_member(member))
.sum::<u64>();
if total > 0 {
state.cost_window_usage_by_key.insert(key_id.clone(), total);
let latency_window_start = now.saturating_sub(pool_config.latency_window_seconds) as f64;
let (cost_results, latency_results) = tokio::join!(
read_window_stats(runtime, &cost_keys, cost_window_start),
read_window_stats(runtime, &latency_keys, latency_window_start),
);
for (key_id, stats) in metric_key_ids.iter().zip(cost_results) {
if stats.sum > 0 {
state
.cost_window_usage_by_key
.insert(key_id.clone(), stats.sum);
}
}
let latency_window_start = now.saturating_sub(pool_config.latency_window_seconds) as f64;
let latency_results = join_all(
latency_keys
.iter()
.map(|latency_key| runtime.score_range_by_min(latency_key, latency_window_start)),
)
.await;
for (key_id, members) in metric_key_ids.iter().zip(latency_results) {
let samples = members
.unwrap_or_default()
.iter()
.map(|member| parse_pool_latency_member(member))
.filter(|value| *value > 0)
.collect::<Vec<_>>();
if samples.is_empty() {
for (key_id, stats) in metric_key_ids.iter().zip(latency_results) {
if stats.positive_count == 0 {
continue;
}
let total = samples.iter().sum::<u64>() as f64;
let average = total / samples.len() as f64;
let average = stats.sum as f64 / stats.positive_count as f64;
if average.is_finite() && average >= 0.0 {
state.latency_avg_ms_by_key.insert(key_id.clone(), average);
}
@@ -300,7 +415,358 @@ pub(crate) async fn read_admin_provider_pool_key_cooldown_reason(
#[cfg(test)]
mod tests {
use super::bounded_runtime_window_metric_key_ids;
use super::super::keys::{pool_cooldown_key, pool_cost_key, pool_latency_key, pool_sticky_key};
use super::{
bounded_runtime_window_metric_key_ids, current_unix_secs,
read_admin_provider_pool_runtime_state, read_provider_pool_scheduling_runtime_state,
read_provider_pool_sticky_bound_key_id,
};
use crate::handlers::admin::provider::pool::config::admin_provider_pool_config_from_config_value;
use crate::handlers::admin::provider::shared::support::AdminProviderPoolConfig;
use aether_runtime_state::{MemoryRuntimeStateConfig, RedisClientConfig, RuntimeState};
use aether_test_support::ManagedRedisServer;
use serde_json::json;
use std::time::Duration;
fn config(value: serde_json::Value) -> AdminProviderPoolConfig {
admin_provider_pool_config_from_config_value(Some(&json!({ "pool_advanced": value })))
.expect("pool config")
}
async fn seed_window_metrics(runtime: &RuntimeState, provider_id: &str, key_id: &str) {
let now = current_unix_secs() as f64;
for (key, member, timestamp) in [
(pool_cost_key(provider_id, key_id), "current:70", now),
(pool_cost_key(provider_id, key_id), "earlier:30", now - 1.0),
(
pool_cost_key(provider_id, key_id),
"expired:999",
now - 20_000.0,
),
(pool_latency_key(provider_id, key_id), "first:10", now),
(
pool_latency_key(provider_id, key_id),
"second:30",
now - 1.0,
),
] {
runtime
.score_set(&key, member, timestamp)
.await
.expect("seed window");
}
}
async fn admin_command_count(runtime: &RuntimeState) -> u64 {
runtime
.redis_diagnostics()
.await
.expect("diagnostics")
.expect("Redis runtime")
.lanes
.into_iter()
.find(|lane| lane.lane == "admin")
.expect("admin lane")
.command_count
}
#[tokio::test]
async fn scheduling_runtime_aggregates_bounded_windows_and_falls_back_for_large_windows() {
let redis = match ManagedRedisServer::start().await {
Ok(server) => server,
Err(err) if err.to_string().contains("No such file or directory") => {
eprintln!("skipping redis-backed scheduling runtime test: {err}");
return;
}
Err(err) => panic!("start Redis: {err}"),
};
let runtime = RuntimeState::redis(
RedisClientConfig {
url: redis.redis_url().to_string(),
key_prefix: Some("pool-window-aggregation-test".to_string()),
},
Some(1_000),
)
.await
.expect("runtime Redis");
let keys = vec!["bounded".to_string(), "large".to_string()];
let now = current_unix_secs() as f64;
for (key_id, count) in [(&keys[0], 512), (&keys[1], 2048)] {
let cost_key = pool_cost_key("pool", key_id);
for index in 0..count {
runtime
.score_set(&cost_key, &format!("{index}:100"), now)
.await
.expect("seed cost window");
}
runtime
.score_set(&cost_key, "expired:9999999", now - 20_000.0)
.await
.expect("expired cost");
for (member, score) in [("first:10", now), ("second:30", now), ("zero:0", now)] {
runtime
.score_set(&pool_latency_key("pool", key_id), member, score)
.await
.expect("seed latency");
}
}
let pool_config = config(json!({
"cost_limit_per_key_tokens": 50_000,
"cost_window_seconds": 600,
"latency_window_seconds": 600,
"scheduling_presets": [{"preset": "latency_first", "enabled": true}]
}));
let scheduled = read_provider_pool_scheduling_runtime_state(
&runtime,
"pool",
&keys,
&pool_config,
None,
)
.await;
assert_eq!(
scheduled.cost_window_usage_by_key.get("bounded"),
Some(&51_200)
);
assert_eq!(
scheduled.cost_window_usage_by_key.get("large"),
Some(&204_800)
);
assert_eq!(scheduled.latency_avg_ms_by_key.get("bounded"), Some(&20.0));
assert_eq!(scheduled.latency_avg_ms_by_key.get("large"), Some(&20.0));
runtime
.score_remove_by_score(&pool_cost_key("pool", "large"), f64::INFINITY)
.await
.expect("reset window");
runtime
.score_set(&pool_cost_key("pool", "large"), "after-reset:75", now)
.await
.expect("post-reset cost");
let reset = read_provider_pool_scheduling_runtime_state(
&runtime,
"pool",
&keys,
&pool_config,
None,
)
.await;
assert_eq!(reset.cost_window_usage_by_key.get("large"), Some(&75));
runtime
.kv_set(&pool_cost_key("pool", "large"), "wrong-type", None)
.await
.expect("simulate invalid metric key");
let partial = read_provider_pool_scheduling_runtime_state(
&runtime,
"pool",
&keys,
&pool_config,
None,
)
.await;
assert_eq!(
partial.cost_window_usage_by_key.get("bounded"),
Some(&51_200),
"one failed aggregate must not discard another key's strict cost check"
);
}
#[tokio::test]
async fn scheduling_runtime_skips_admin_scan_and_unused_window_queries() {
let redis = match ManagedRedisServer::start().await {
Ok(server) => server,
Err(err) if err.to_string().contains("No such file or directory") => {
eprintln!("skipping redis-backed scheduling runtime test: {err}");
return;
}
Err(err) => panic!("Redis server should start: {err}"),
};
let runtime = RuntimeState::redis(
RedisClientConfig {
url: redis.redis_url().to_string(),
key_prefix: Some("scheduling-runtime-reads".to_string()),
},
Some(2_000),
)
.await
.expect("Redis runtime");
let pool_config = config(json!({}));
let keys = vec!["ready".to_string(), "cooling".to_string()];
seed_window_metrics(&runtime, "pool", "ready").await;
for session in ["current", "other"] {
runtime
.kv_set(
&pool_sticky_key("pool", session),
"ready".to_string(),
Some(Duration::from_secs(60)),
)
.await
.expect("seed sticky session");
}
runtime
.kv_set(
&pool_cooldown_key("pool", "cooling"),
"rate_limit".to_string(),
Some(Duration::from_secs(60)),
)
.await
.expect("seed cooldown");
let before = admin_command_count(&runtime).await;
let scheduled = read_provider_pool_scheduling_runtime_state(
&runtime,
"pool",
&keys,
&pool_config,
Some("current"),
)
.await;
let after = admin_command_count(&runtime).await;
assert_eq!(
after - before,
1,
"only the diagnostics INFO may use the admin lane"
);
assert_eq!(scheduled.sticky_bound_key_id.as_deref(), Some("ready"));
assert_eq!(
scheduled
.cooldown_reason_by_key
.get("cooling")
.map(String::as_str),
Some("rate_limit")
);
assert!(scheduled.cooldown_ttl_by_key.is_empty());
assert!(scheduled.cost_window_usage_by_key.is_empty());
assert!(scheduled.latency_avg_ms_by_key.is_empty());
assert_eq!(scheduled.total_sticky_sessions, 0);
let admin = read_admin_provider_pool_runtime_state(
&runtime,
"pool",
&keys,
&pool_config,
Some("current"),
)
.await;
assert_eq!(admin.total_sticky_sessions, 2);
assert_eq!(admin.sticky_sessions_by_key.get("ready"), Some(&2));
assert_eq!(admin.cost_window_usage_by_key.get("ready"), Some(&100));
assert_eq!(admin.latency_avg_ms_by_key.get("ready"), Some(&20.0));
assert!(admin
.cooldown_ttl_by_key
.get("cooling")
.is_some_and(|ttl| *ttl > 0));
}
#[tokio::test]
async fn scheduling_runtime_loads_only_metrics_used_by_enabled_strategies() {
let runtime = RuntimeState::memory(MemoryRuntimeStateConfig::default());
let keys = vec!["key".to_string()];
seed_window_metrics(&runtime, "pool", "key").await;
for (value, expected_cost, expected_latency) in [
(json!({}), false, false),
(json!({"cost_limit_per_key_tokens": 100}), true, false),
(json!({"cost_limit_per_key_tokens": 0}), true, false),
(
json!({"scheduling_presets": [{"preset": "cost_first", "enabled": true}]}),
true,
false,
),
(
json!({"scheduling_presets": [{"preset": "quota_balanced", "enabled": true}]}),
true,
false,
),
(
json!({"scheduling_presets": [{"preset": "latency_first", "enabled": true}]}),
false,
true,
),
(
json!({"scheduling_presets": [
{"preset": "cost_first", "enabled": false},
{"preset": "latency_first", "enabled": false}
]}),
false,
false,
),
] {
let pool_config = config(value.clone());
let snapshot = read_provider_pool_scheduling_runtime_state(
&runtime,
"pool",
&keys,
&pool_config,
None,
)
.await;
assert_eq!(
snapshot.cost_window_usage_by_key.get("key").copied(),
expected_cost.then_some(100),
"config: {value}"
);
assert_eq!(
snapshot.latency_avg_ms_by_key.get("key").copied(),
expected_latency.then_some(20.0),
"config: {value}"
);
}
}
#[tokio::test]
async fn scheduling_runtime_checks_cost_for_candidates_beyond_admin_display_limit() {
let runtime = RuntimeState::memory(MemoryRuntimeStateConfig::default());
let keys = (0..513)
.map(|index| format!("key-{index}"))
.collect::<Vec<_>>();
seed_window_metrics(&runtime, "pool", &keys[512]).await;
let pool_config = config(json!({ "cost_limit_per_key_tokens": 100 }));
let snapshot = read_provider_pool_scheduling_runtime_state(
&runtime,
"pool",
&keys,
&pool_config,
None,
)
.await;
assert_eq!(
snapshot.cost_window_usage_by_key.get(&keys[512]),
Some(&100)
);
}
#[tokio::test]
async fn scheduling_sticky_lookup_invalidates_a_cooled_down_binding() {
let runtime = RuntimeState::memory(MemoryRuntimeStateConfig::default());
let pool_config = config(json!({}));
let sticky_key = pool_sticky_key("pool", "session");
runtime
.kv_set(&sticky_key, "key".to_string(), None)
.await
.expect("sticky session");
runtime
.kv_set(
&pool_cooldown_key("pool", "key"),
"rate_limit".to_string(),
None,
)
.await
.expect("cooldown");
assert!(read_provider_pool_sticky_bound_key_id(
&runtime,
"pool",
&pool_config,
Some("session")
)
.await
.is_none());
assert!(!runtime
.kv_exists(&sticky_key)
.await
.expect("sticky existence"));
}
#[test]
fn runtime_window_metric_key_ids_are_bounded() {
apps/aether-gateway/src/handlers/admin/system/shared/update.rs
+1
View File
@@ -1937,6 +1937,7 @@ pub(crate) async fn start_admin_system_rollback_task(
}
fn request_process_restart() -> ! {
let _ = aether_runtime::shutdown_logging(std::time::Duration::from_secs(2));
std::process::exit(RESTART_EXIT_CODE);
}
apps/aether-gateway/src/handlers/proxy/body_buffer.rs
+28 -3
View File
@@ -225,6 +225,9 @@ impl RequestBodyBufferError {
RequestBodyNormalizationError::RequestBodyTooLarge { .. } => {
"request_body_too_large"
}
RequestBodyNormalizationError::BodyBufferOverloaded { .. } => {
"request_body_buffer_overloaded"
}
},
Self::TooLarge { .. } => "request_body_too_large",
Self::Overloaded { .. } => "request_body_buffer_overloaded",
@@ -292,15 +295,37 @@ pub(super) async fn buffer_and_normalize_request_body(
.await
.map_err(RequestBodyBufferError::from)?;
let elapsed_ms = buffered.elapsed().as_millis() as u64;
let retained_input_capacity = buffered
.requested_bytes()
.saturating_sub(buffered.bytes().len());
let normalized = buffered
.try_map(|body| {
crate::headers::normalize_request_body_headers_and_bytes_with_limit(
.try_map_with_budget(|body, memory| {
crate::headers::normalize_request_body_headers_and_bytes_with_budget(
headers,
body,
policy.effective_max_bytes(),
&mut |requested_bytes| {
let requested_bytes = requested_bytes.saturating_add(retained_input_capacity);
memory.try_reserve_bytes(requested_bytes).map_err(|_| {
RequestBodyNormalizationError::BodyBufferOverloaded {
requested_bytes,
budget_bytes: policy.budget_bytes(),
}
})
},
)
})
.map_err(RequestBodyBufferError::Normalization)?;
.map_err(|error| match error {
RequestBodyNormalizationError::BodyBufferOverloaded {
requested_bytes,
budget_bytes,
} => RequestBodyBufferError::Overloaded {
requested_bytes,
budget_bytes,
timeout_ms: 0,
},
error => RequestBodyBufferError::Normalization(error),
})?;
info!(
event_name = "frontdoor_request_body_buffer_completed",
log_type = "event",
apps/aether-gateway/src/handlers/proxy/mod.rs
+111 -5
View File
@@ -1035,11 +1035,10 @@ pub(crate) async fn proxy_request(
ConnectInfo(remote_addr): ConnectInfo<std::net::SocketAddr>,
request: Request,
) -> Result<Response<Body>, GatewayError> {
crate::request_lifecycle::run_request(Box::pin(proxy_request_inner(
state,
remote_addr,
request,
)))
crate::request_lifecycle::run_request_with_usage(
state.usage_runtime.clone(),
Box::pin(proxy_request_inner(state, remote_addr, request)),
)
.await
}
@@ -3269,6 +3268,113 @@ mod tests {
));
}
#[tokio::test]
async fn request_body_buffer_allows_parallel_compressed_uploads() {
let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
encoder.write_all(br#"{"model":"test"}"#).unwrap();
let encoded = Bytes::from(encoder.finish().unwrap());
let budget_bytes = 2 * crate::state::REQUEST_BODY_BUFFER_PERMIT_BYTES;
let budget = Arc::new(Semaphore::new(2));
let policy = RequestBodyBufferPolicy::for_tests_with_budget(
budget_bytes as u64,
Duration::from_secs(1),
Duration::from_millis(50),
budget_bytes,
Arc::clone(&budget),
);
let mut headers = HeaderMap::new();
headers.insert(header::CONTENT_ENCODING, HeaderValue::from_static("gzip"));
headers.insert(header::CONTENT_LENGTH, HeaderValue::from(encoded.len()));
let (started_tx, started_rx) = tokio::sync::oneshot::channel();
let (finish_tx, finish_rx) = tokio::sync::oneshot::channel();
let first_policy = policy.clone();
let mut first_headers = headers.clone();
let first_encoded = encoded.clone();
let first = async move {
let stream = async_stream::stream! {
let middle = first_encoded.len() / 2;
yield Ok::<_, std::io::Error>(first_encoded.slice(..middle));
let _ = started_tx.send(());
let _ = finish_rx.await;
yield Ok(first_encoded.slice(middle..));
};
buffer_and_normalize_request_body(
&mut Some(Body::from_stream(stream)),
&mut first_headers,
"test owns body",
"trace-compressed-first",
&Method::POST,
"/v1/responses",
"test",
first_policy,
)
.await
};
let second = async move {
started_rx.await.unwrap();
let result = buffer_and_normalize_request_body(
&mut Some(Body::from(encoded)),
&mut headers,
"test owns body",
"trace-compressed-second",
&Method::POST,
"/v1/responses",
"test",
policy,
)
.await;
let _ = finish_tx.send(());
result
};
let (first, second) = tokio::time::timeout(Duration::from_secs(2), async {
tokio::join!(first, second)
})
.await
.expect("concurrent compressed requests should finish");
assert_eq!(first.unwrap().as_ref(), br#"{"model":"test"}"#);
assert_eq!(second.unwrap().as_ref(), br#"{"model":"test"}"#);
assert_eq!(budget.available_permits(), 2);
}
#[tokio::test]
async fn request_body_buffer_rejects_decompression_growth_when_budget_is_busy() {
let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
encoder.write_all(&vec![b'a'; 100_000]).unwrap();
let encoded = encoder.finish().unwrap();
let budget_bytes = 2 * crate::state::REQUEST_BODY_BUFFER_PERMIT_BYTES;
let budget = Arc::new(Semaphore::new(2));
let held = Arc::clone(&budget).acquire_owned().await.unwrap();
let mut headers = HeaderMap::new();
headers.insert(header::CONTENT_ENCODING, HeaderValue::from_static("gzip"));
headers.insert(header::CONTENT_LENGTH, HeaderValue::from(encoded.len()));
let result = buffer_and_normalize_request_body(
&mut Some(Body::from(encoded)),
&mut headers,
"test owns body",
"trace-decompression-overload",
&Method::POST,
"/v1/responses",
"test",
RequestBodyBufferPolicy::for_tests_with_budget(
budget_bytes as u64,
Duration::from_secs(1),
Duration::from_secs(1),
budget_bytes,
Arc::clone(&budget),
),
)
.await
.unwrap_err();
assert!(matches!(
result,
RequestBodyBufferError::Overloaded { timeout_ms: 0, .. }
));
assert_eq!(result.http_status(), http::StatusCode::SERVICE_UNAVAILABLE);
assert_eq!(budget.available_permits(), 1);
drop(held);
assert_eq!(budget.available_permits(), 2);
}
#[tokio::test]
async fn request_body_buffer_times_out_instead_of_waiting_forever() {
let stream = async_stream::stream! {
apps/aether-gateway/src/handlers/proxy/websocket/live/audit.rs
+4
View File
@@ -508,7 +508,9 @@ async fn persist_live_audit_event(
let usage_runtime = std::sync::Arc::clone(&state.usage_runtime);
let usage_data = std::sync::Arc::clone(state.usage_lifecycle_data_state());
let write_request_id = request_id.clone();
let usage_producer = usage_runtime.track_producer();
let task = tokio::spawn(async move {
let _usage_producer = usage_producer;
if tokio::time::timeout(
LIVE_AUDIT_WRITE_HARD_TIMEOUT,
usage_runtime.record_terminal_event_direct(usage_data.as_ref(), event),
@@ -572,7 +574,9 @@ fn spawn_live_audit_event_detached(state: &AppState, event: UsageEvent, audit_sc
};
let usage_runtime = std::sync::Arc::clone(&state.usage_runtime);
let usage_data = std::sync::Arc::clone(state.usage_lifecycle_data_state());
let usage_producer = usage_runtime.track_producer();
runtime.spawn(async move {
let _usage_producer = usage_producer;
if tokio::time::timeout(
LIVE_AUDIT_WRITE_HARD_TIMEOUT,
usage_runtime.record_terminal_event_direct(usage_data.as_ref(), event),
apps/aether-gateway/src/handlers/shared/provider_pool.rs
+2 -1
View File
@@ -3,7 +3,8 @@ pub(crate) use super::super::admin::provider::pool::config::{
};
pub(crate) use super::super::admin::provider::pool::runtime::{
admin_provider_pool_key_terminal_error_reason, read_admin_provider_pool_key_cooldown_reason,
read_admin_provider_pool_runtime_state, record_admin_provider_pool_error,
read_admin_provider_pool_runtime_state, read_provider_pool_scheduling_runtime_state,
read_provider_pool_sticky_bound_key_id, record_admin_provider_pool_error,
record_admin_provider_pool_stream_timeout, record_admin_provider_pool_success,
release_admin_provider_pool_key_lease,
};