Improve gateway transport and usage runtime

This commit is contained in:
elky
2026-06-25 22:36:27 +08:00
parent d336d1a7fa
commit 6f00e9fc67
112 changed files with 12456 additions and 1387 deletions
+248 -23
View File
@@ -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)
);
}
}