mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-09 02:47:45 +08:00
Improve gateway transport and usage runtime
This commit is contained in:
@@ -2,13 +2,15 @@ use std::collections::HashMap;
|
||||
use std::sync::atomic::AtomicU64;
|
||||
use std::sync::Arc;
|
||||
use std::sync::Mutex as StdMutex;
|
||||
use std::sync::RwLock as StdRwLock;
|
||||
use std::time::Duration;
|
||||
|
||||
use aether_data::repository::users::StoredUserGroup;
|
||||
use aether_data_contracts::repository::billing::UserDailyQuotaAvailabilityRecord;
|
||||
use aether_data_contracts::repository::quota::StoredProviderQuotaSnapshot;
|
||||
use aether_runtime::ConcurrencyGate;
|
||||
use aether_runtime_state::{RuntimeSemaphore, RuntimeState};
|
||||
use dashmap::DashMap;
|
||||
use tokio::sync::Mutex as TokioMutex;
|
||||
|
||||
use super::super::async_task::{VideoTaskPollerConfig, VideoTaskService};
|
||||
use super::super::cache::{
|
||||
@@ -39,14 +41,25 @@ const MAX_LOCAL_EXECUTION_PLANNING_TIMEOUT_MS: u64 = 120_000;
|
||||
const LOCAL_EXECUTION_PLANNING_TIMEOUT_MS_ENV: &str =
|
||||
"AETHER_GATEWAY_LOCAL_EXECUTION_PLANNING_TIMEOUT_MS";
|
||||
const DEFAULT_CANDIDATE_PLANNING_GATE_LIMIT: usize = 1024;
|
||||
const DEFAULT_UPSTREAM_EXECUTION_GATE_LIMIT: usize = 2000;
|
||||
const DEFAULT_UPSTREAM_TARGET_GATE_LIMIT: usize = 2000;
|
||||
const DEFAULT_UPSTREAM_EXECUTION_GATE_LIMIT: usize = 10_000;
|
||||
const DEFAULT_UPSTREAM_TARGET_GATE_LIMIT: usize = 10_000;
|
||||
const MAX_CANDIDATE_PLANNING_GATE_LIMIT: usize = 8192;
|
||||
const MAX_UPSTREAM_EXECUTION_GATE_LIMIT: usize = 16_384;
|
||||
const MAX_UPSTREAM_TARGET_GATE_LIMIT: usize = 16_384;
|
||||
const CANDIDATE_PLANNING_GATE_LIMIT_PER_CPU: usize = 256;
|
||||
const UPSTREAM_EXECUTION_GATE_LIMIT_PER_CPU: usize = 1024;
|
||||
const UPSTREAM_TARGET_GATE_LIMIT_PER_CPU: usize = 1024;
|
||||
const GATE_LIMIT_FD_RESERVE: usize = 128;
|
||||
const DEFAULT_INTERNAL_GATE_QUEUE_BUDGET_MS: u64 = 250;
|
||||
const MAX_INTERNAL_GATE_QUEUE_BUDGET_MS: u64 = 5_000;
|
||||
const CANDIDATE_PLANNING_GATE_LIMIT_ENV: &str = "AETHER_GATEWAY_CANDIDATE_PLANNING_GATE_LIMIT";
|
||||
const UPSTREAM_EXECUTION_GATE_LIMIT_ENV: &str = "AETHER_GATEWAY_UPSTREAM_EXECUTION_GATE_LIMIT";
|
||||
const UPSTREAM_TARGET_GATE_LIMIT_ENV: &str = "AETHER_GATEWAY_UPSTREAM_TARGET_GATE_LIMIT";
|
||||
const INTERNAL_GATE_QUEUE_BUDGET_MS_ENV: &str = "AETHER_GATEWAY_INTERNAL_GATE_QUEUE_BUDGET_MS";
|
||||
const DEFAULT_AUTH_CAPACITY_CACHE_TTL_MS: u64 = 500;
|
||||
const MIN_AUTH_CAPACITY_CACHE_TTL_MS: u64 = 10;
|
||||
const MAX_AUTH_CAPACITY_CACHE_TTL_MS: u64 = 10_000;
|
||||
const AUTH_CAPACITY_CACHE_TTL_MS_ENV: &str = "AETHER_GATEWAY_AUTH_CAPACITY_CACHE_TTL_MS";
|
||||
|
||||
#[cfg(test)]
|
||||
type TestExecutionRuntimeSyncOverrideFn = dyn Fn(
|
||||
@@ -73,6 +86,7 @@ pub(crate) struct FrontdoorRuntimeGuardConfig {
|
||||
pub(crate) request_body_read_timeout: Duration,
|
||||
pub(crate) local_execution_planning_timeout: Duration,
|
||||
pub(crate) internal_gate_queue_budget: Duration,
|
||||
pub(crate) auth_capacity_cache_ttl: Duration,
|
||||
pub(crate) candidate_planning_gate_limit: Option<usize>,
|
||||
pub(crate) upstream_execution_gate_limit: Option<usize>,
|
||||
pub(crate) upstream_target_gate_limit: Option<usize>,
|
||||
@@ -99,18 +113,15 @@ impl FrontdoorRuntimeGuardConfig {
|
||||
1,
|
||||
MAX_INTERNAL_GATE_QUEUE_BUDGET_MS,
|
||||
),
|
||||
candidate_planning_gate_limit: env_optional_usize(
|
||||
CANDIDATE_PLANNING_GATE_LIMIT_ENV,
|
||||
DEFAULT_CANDIDATE_PLANNING_GATE_LIMIT,
|
||||
),
|
||||
upstream_execution_gate_limit: env_optional_usize(
|
||||
UPSTREAM_EXECUTION_GATE_LIMIT_ENV,
|
||||
DEFAULT_UPSTREAM_EXECUTION_GATE_LIMIT,
|
||||
),
|
||||
upstream_target_gate_limit: env_optional_usize(
|
||||
UPSTREAM_TARGET_GATE_LIMIT_ENV,
|
||||
DEFAULT_UPSTREAM_TARGET_GATE_LIMIT,
|
||||
auth_capacity_cache_ttl: env_cache_duration_ms(
|
||||
AUTH_CAPACITY_CACHE_TTL_MS_ENV,
|
||||
DEFAULT_AUTH_CAPACITY_CACHE_TTL_MS,
|
||||
MIN_AUTH_CAPACITY_CACHE_TTL_MS,
|
||||
MAX_AUTH_CAPACITY_CACHE_TTL_MS,
|
||||
),
|
||||
candidate_planning_gate_limit: candidate_planning_gate_limit_from_env(),
|
||||
upstream_execution_gate_limit: upstream_execution_gate_limit_from_env(),
|
||||
upstream_target_gate_limit: upstream_target_gate_limit_from_env(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -125,6 +136,7 @@ impl FrontdoorRuntimeGuardConfig {
|
||||
internal_gate_queue_budget: Duration::from_millis(
|
||||
DEFAULT_INTERNAL_GATE_QUEUE_BUDGET_MS,
|
||||
),
|
||||
auth_capacity_cache_ttl: Duration::from_millis(DEFAULT_AUTH_CAPACITY_CACHE_TTL_MS),
|
||||
candidate_planning_gate_limit: Some(DEFAULT_CANDIDATE_PLANNING_GATE_LIMIT),
|
||||
upstream_execution_gate_limit: Some(DEFAULT_UPSTREAM_EXECUTION_GATE_LIMIT),
|
||||
upstream_target_gate_limit: Some(DEFAULT_UPSTREAM_TARGET_GATE_LIMIT),
|
||||
@@ -142,15 +154,152 @@ fn env_duration_ms(key: &str, default_ms: u64, min_ms: u64, max_ms: u64) -> Dura
|
||||
Duration::from_millis(ms)
|
||||
}
|
||||
|
||||
fn env_optional_usize(key: &str, default_value: usize) -> Option<usize> {
|
||||
match std::env::var(key)
|
||||
fn env_cache_duration_ms(key: &str, default_ms: u64, min_ms: u64, max_ms: u64) -> Duration {
|
||||
let Some(raw) = std::env::var(key)
|
||||
.ok()
|
||||
.and_then(|value| value.trim().parse::<usize>().ok())
|
||||
{
|
||||
Some(0) => None,
|
||||
Some(value) => Some(value.max(1)),
|
||||
None => Some(default_value.max(1)),
|
||||
.map(|value| value.trim().to_string())
|
||||
.filter(|value| !value.is_empty())
|
||||
else {
|
||||
return Duration::from_millis(default_ms);
|
||||
};
|
||||
let Some(parsed) = raw.parse::<u64>().ok() else {
|
||||
return Duration::from_millis(default_ms);
|
||||
};
|
||||
if parsed == 0 {
|
||||
return Duration::ZERO;
|
||||
}
|
||||
Duration::from_millis(parsed.clamp(min_ms, max_ms))
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
struct GateAutoProfile {
|
||||
floor: usize,
|
||||
cap: usize,
|
||||
per_cpu: usize,
|
||||
fd_divisor: Option<usize>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
struct GateAutoCapacity {
|
||||
cpu_parallelism: usize,
|
||||
fd_soft_limit: usize,
|
||||
}
|
||||
|
||||
const CANDIDATE_PLANNING_GATE_AUTO_PROFILE: GateAutoProfile = GateAutoProfile {
|
||||
floor: DEFAULT_CANDIDATE_PLANNING_GATE_LIMIT,
|
||||
cap: MAX_CANDIDATE_PLANNING_GATE_LIMIT,
|
||||
per_cpu: CANDIDATE_PLANNING_GATE_LIMIT_PER_CPU,
|
||||
fd_divisor: None,
|
||||
};
|
||||
|
||||
const UPSTREAM_EXECUTION_GATE_AUTO_PROFILE: GateAutoProfile = GateAutoProfile {
|
||||
floor: DEFAULT_UPSTREAM_EXECUTION_GATE_LIMIT,
|
||||
cap: MAX_UPSTREAM_EXECUTION_GATE_LIMIT,
|
||||
per_cpu: UPSTREAM_EXECUTION_GATE_LIMIT_PER_CPU,
|
||||
fd_divisor: Some(2),
|
||||
};
|
||||
|
||||
const UPSTREAM_TARGET_GATE_AUTO_PROFILE: GateAutoProfile = GateAutoProfile {
|
||||
floor: DEFAULT_UPSTREAM_TARGET_GATE_LIMIT,
|
||||
cap: MAX_UPSTREAM_TARGET_GATE_LIMIT,
|
||||
per_cpu: UPSTREAM_TARGET_GATE_LIMIT_PER_CPU,
|
||||
fd_divisor: Some(4),
|
||||
};
|
||||
|
||||
fn candidate_planning_gate_limit_from_env() -> Option<usize> {
|
||||
env_gate_limit(
|
||||
CANDIDATE_PLANNING_GATE_LIMIT_ENV,
|
||||
CANDIDATE_PLANNING_GATE_AUTO_PROFILE,
|
||||
)
|
||||
}
|
||||
|
||||
fn upstream_execution_gate_limit_from_env() -> Option<usize> {
|
||||
env_gate_limit(
|
||||
UPSTREAM_EXECUTION_GATE_LIMIT_ENV,
|
||||
UPSTREAM_EXECUTION_GATE_AUTO_PROFILE,
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn upstream_target_gate_limit_from_env() -> Option<usize> {
|
||||
env_gate_limit(
|
||||
UPSTREAM_TARGET_GATE_LIMIT_ENV,
|
||||
UPSTREAM_TARGET_GATE_AUTO_PROFILE,
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn upstream_target_gate_auto_limit() -> usize {
|
||||
auto_gate_limit(
|
||||
UPSTREAM_TARGET_GATE_AUTO_PROFILE,
|
||||
current_gate_auto_capacity(),
|
||||
)
|
||||
}
|
||||
|
||||
fn env_gate_limit(key: &str, profile: GateAutoProfile) -> Option<usize> {
|
||||
let raw = std::env::var(key).ok();
|
||||
parse_gate_limit_value(raw.as_deref(), profile, current_gate_auto_capacity())
|
||||
}
|
||||
|
||||
fn parse_gate_limit_value(
|
||||
raw: Option<&str>,
|
||||
profile: GateAutoProfile,
|
||||
capacity: GateAutoCapacity,
|
||||
) -> Option<usize> {
|
||||
let Some(value) = raw.map(str::trim).filter(|value| !value.is_empty()) else {
|
||||
return Some(auto_gate_limit(profile, capacity));
|
||||
};
|
||||
let normalized = value.to_ascii_lowercase();
|
||||
match normalized.as_str() {
|
||||
"auto" => Some(auto_gate_limit(profile, capacity)),
|
||||
"off" | "none" | "disabled" | "disable" => None,
|
||||
_ => match value.parse::<usize>() {
|
||||
Ok(0) => None,
|
||||
Ok(limit) => Some(limit.max(1)),
|
||||
Err(_) => Some(auto_gate_limit(profile, capacity)),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn auto_gate_limit(profile: GateAutoProfile, capacity: GateAutoCapacity) -> usize {
|
||||
let cpu_limit = capacity
|
||||
.cpu_parallelism
|
||||
.max(1)
|
||||
.saturating_mul(profile.per_cpu.max(1));
|
||||
let mut limit = cpu_limit.max(profile.floor).min(profile.cap);
|
||||
if let Some(fd_divisor) = profile.fd_divisor.filter(|value| *value > 0) {
|
||||
let fd_budget = capacity
|
||||
.fd_soft_limit
|
||||
.saturating_sub(GATE_LIMIT_FD_RESERVE)
|
||||
.checked_div(fd_divisor)
|
||||
.unwrap_or(1)
|
||||
.max(1);
|
||||
limit = limit.min(fd_budget);
|
||||
}
|
||||
limit.max(1)
|
||||
}
|
||||
|
||||
fn current_gate_auto_capacity() -> GateAutoCapacity {
|
||||
GateAutoCapacity {
|
||||
cpu_parallelism: std::thread::available_parallelism()
|
||||
.map(|value| value.get())
|
||||
.unwrap_or(1)
|
||||
.max(1),
|
||||
fd_soft_limit: soft_fd_limit().unwrap_or(1024).max(1),
|
||||
}
|
||||
}
|
||||
|
||||
fn soft_fd_limit() -> Option<usize> {
|
||||
#[cfg(unix)]
|
||||
{
|
||||
let mut limit = libc::rlimit {
|
||||
rlim_cur: 0,
|
||||
rlim_max: 0,
|
||||
};
|
||||
let result = unsafe { libc::getrlimit(libc::RLIMIT_NOFILE, &mut limit) };
|
||||
if result == 0 {
|
||||
return usize::try_from(limit.rlim_cur).ok();
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
@@ -178,8 +327,15 @@ pub struct AppState {
|
||||
pub(crate) auth_api_key_force_capabilities_cache:
|
||||
Arc<JsonValueCache<AuthApiKeyIdentityCacheKey>>,
|
||||
pub(crate) auth_api_key_feature_settings_cache: Arc<JsonValueCache<AuthApiKeyFeatureCacheKey>>,
|
||||
pub(crate) auth_daily_quota_availability_cache:
|
||||
Arc<ValueCache<String, UserDailyQuotaAvailabilityRecord>>,
|
||||
pub(crate) auth_wallet_snapshot_cache:
|
||||
Arc<ValueCache<String, aether_data::repository::wallet::StoredWalletSnapshot>>,
|
||||
pub(crate) auth_request_cost_upper_bound_cache: Arc<ValueCache<String, f64>>,
|
||||
pub(crate) provider_quota_snapshot_cache: Arc<ValueCache<String, StoredProviderQuotaSnapshot>>,
|
||||
pub(crate) user_groups_for_user_cache: Arc<ValueCache<String, Vec<StoredUserGroup>>>,
|
||||
pub(crate) routing_group_selection_cache:
|
||||
Arc<ValueCache<String, crate::routing::GatewayRoutingGroupSelection>>,
|
||||
pub(crate) auth_api_key_last_used_cache: Arc<AuthApiKeyLastUsedCache>,
|
||||
pub(crate) oauth_refresh: Arc<provider_transport::LocalOAuthRefreshCoordinator>,
|
||||
pub(crate) direct_plan_bypass_cache: Arc<DirectPlanBypassCache>,
|
||||
@@ -187,6 +343,7 @@ pub struct AppState {
|
||||
pub(crate) scheduler_affinity_epoch: Arc<AtomicU64>,
|
||||
pub(crate) dashboard_response_cache: Arc<DashboardResponseCache>,
|
||||
pub(crate) system_config_cache: Arc<SystemConfigCache>,
|
||||
pub(crate) candidate_row_page_cache: Arc<super::super::cache::CandidateRowPageCache>,
|
||||
pub(crate) candidate_page_cache: Arc<super::super::cache::CandidatePageCache>,
|
||||
pub(crate) candidate_resolved_page_cache: Arc<super::super::cache::CandidateResolvedPageCache>,
|
||||
pub(crate) chat_pii_redaction_runtime_config_cache:
|
||||
@@ -197,10 +354,12 @@ pub struct AppState {
|
||||
pub(crate) frontdoor_user_rpm: Arc<FrontdoorUserRpmLimiter>,
|
||||
pub(crate) tunnel: crate::tunnel::EmbeddedTunnelState,
|
||||
pub(crate) provider_transport_snapshot_cache:
|
||||
Arc<StdRwLock<HashMap<ProviderTransportSnapshotCacheKey, CachedProviderTransportSnapshot>>>,
|
||||
Arc<DashMap<ProviderTransportSnapshotCacheKey, CachedProviderTransportSnapshot>>,
|
||||
pub(crate) provider_transport_snapshot_inflight:
|
||||
Arc<DashMap<ProviderTransportSnapshotCacheKey, Arc<TokioMutex<()>>>>,
|
||||
pub(crate) provider_key_rpm_resets: Arc<StdMutex<HashMap<String, u64>>>,
|
||||
pub(crate) local_execution_runtime_miss_diagnostics:
|
||||
Arc<StdMutex<HashMap<String, LocalExecutionRuntimeMissDiagnostic>>>,
|
||||
Arc<DashMap<String, LocalExecutionRuntimeMissDiagnostic>>,
|
||||
pub(crate) admin_monitoring_error_stats_reset_at: Arc<StdMutex<Option<u64>>>,
|
||||
pub(crate) provider_delete_tasks: Arc<StdMutex<HashMap<String, LocalProviderDeleteTaskState>>>,
|
||||
#[cfg(test)]
|
||||
@@ -262,3 +421,69 @@ pub struct AppState {
|
||||
#[cfg(test)]
|
||||
pub(crate) provider_oauth_token_url_overrides: Arc<StdMutex<HashMap<String, String>>>,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
const TEST_PROFILE: GateAutoProfile = GateAutoProfile {
|
||||
floor: 10_000,
|
||||
cap: 16_384,
|
||||
per_cpu: 1024,
|
||||
fd_divisor: Some(2),
|
||||
};
|
||||
|
||||
const TEST_CAPACITY: GateAutoCapacity = GateAutoCapacity {
|
||||
cpu_parallelism: 12,
|
||||
fd_soft_limit: 1_048_576,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn gate_limit_parser_defaults_to_auto() {
|
||||
assert_eq!(
|
||||
parse_gate_limit_value(None, TEST_PROFILE, TEST_CAPACITY),
|
||||
Some(12_288)
|
||||
);
|
||||
assert_eq!(
|
||||
parse_gate_limit_value(Some("auto"), TEST_PROFILE, TEST_CAPACITY),
|
||||
Some(12_288)
|
||||
);
|
||||
assert_eq!(
|
||||
parse_gate_limit_value(Some(" AUTO "), TEST_PROFILE, TEST_CAPACITY),
|
||||
Some(12_288)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gate_limit_parser_accepts_fixed_numbers() {
|
||||
assert_eq!(
|
||||
parse_gate_limit_value(Some("4096"), TEST_PROFILE, TEST_CAPACITY),
|
||||
Some(4096)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gate_limit_parser_accepts_off_and_legacy_zero() {
|
||||
for value in ["off", "none", "disabled", "disable", "0"] {
|
||||
assert_eq!(
|
||||
parse_gate_limit_value(Some(value), TEST_PROFILE, TEST_CAPACITY),
|
||||
None
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn auto_gate_limit_respects_fd_budget_when_fd_limit_is_low() {
|
||||
assert_eq!(
|
||||
parse_gate_limit_value(
|
||||
Some("auto"),
|
||||
TEST_PROFILE,
|
||||
GateAutoCapacity {
|
||||
cpu_parallelism: 32,
|
||||
fd_soft_limit: 1024,
|
||||
},
|
||||
),
|
||||
Some(448)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user