diff --git a/src/api/mod.rs b/src/api/mod.rs index af70f42..ecda81e 100644 --- a/src/api/mod.rs +++ b/src/api/mod.rs @@ -1,5 +1,6 @@ #![allow(clippy::too_many_arguments)] +use std::collections::BTreeSet; use std::io::{Error as IoError, ErrorKind}; use std::net::{IpAddr, SocketAddr}; use std::path::PathBuf; @@ -22,10 +23,12 @@ use tracing::{debug, info, warn}; use crate::config::ApiGrayAction; use crate::ip_tracker::UserIpTracker; +use crate::maestro::control_plane::ProcessControlPlane; use crate::maestro::generation::{RuntimeGeneration, RuntimeWatchState}; use crate::maestro::reload::{ReloadAccepted, ReloadControl, ReloadRequest, ReloadSubmitError}; use crate::proxy::route_mode::RouteRuntimeController; use crate::proxy::shared_state::ProxySharedState; +use crate::quota_state::QuotaStateOwner; use crate::startup::StartupTracker; use crate::stats::Stats; use crate::transport::UpstreamManager; @@ -112,7 +115,7 @@ pub(super) struct ApiShared { pub(super) me_pool: Arc>>>, pub(super) upstream_manager: Arc, pub(super) config_path: PathBuf, - pub(super) quota_state_path: PathBuf, + pub(super) quota_state: Arc, pub(super) detected_ips_rx: watch::Receiver<(Option, Option)>, pub(super) mutation_lock: Arc>, pub(super) minimal_cache: Arc>>, @@ -147,7 +150,7 @@ impl ApiShared { me_pool: runtime.me_pool_runtime.clone(), upstream_manager: runtime.upstream_manager.clone(), config_path: self.config_path.clone(), - quota_state_path: self.quota_state_path.clone(), + quota_state: self.quota_state.clone(), detected_ips_rx: self.detected_ips_rx.clone(), mutation_lock: self.mutation_lock.clone(), minimal_cache: self.minimal_cache.clone(), @@ -282,8 +285,9 @@ fn allowed_methods_for_path(path: &str) -> Option<&'static str> { } } -pub async fn serve( - listen: SocketAddr, +/// Serves the API on a process-owned listener and task scope. +pub(crate) async fn serve( + listener: TcpListener, stats: Arc, ip_tracker: Arc, me_pool: Arc>>>, @@ -291,7 +295,7 @@ pub async fn serve( proxy_shared: Arc, upstream_manager: Arc, config_path: PathBuf, - quota_state_path: PathBuf, + quota_state: Arc, detected_ips_rx: watch::Receiver<(Option, Option)>, process_started_at_epoch_secs: u64, startup_tracker: Arc, @@ -300,6 +304,7 @@ pub async fn serve( mut runtime_watch_rx: watch::Receiver>, web_trace: Arc, web_runtime_rx: watch::Receiver, + control_plane: ProcessControlPlane, ) { let active_runtime = loop { if let Some(active_runtime) = active_runtime_rx.borrow().clone() { @@ -321,19 +326,9 @@ pub async fn serve( }; let config_rx = initial_watch_state.config_rx.clone(); let admission_rx = initial_watch_state.admission_rx.clone(); - let listener = match TcpListener::bind(listen).await { - Ok(listener) => listener, - Err(error) => { - warn!( - error = %error, - listen = %listen, - "Failed to bind API listener" - ); - return; - } - }; + let listen = listener.local_addr().ok(); - info!("API endpoint: http://{}/v1/* and /web-status", listen); + info!(listen = ?listen, "API endpoint ready at /v1/* and /web-status"); let runtime_state = Arc::new(ApiRuntimeState { process_started_at_epoch_secs, @@ -348,7 +343,7 @@ pub async fn serve( me_pool, upstream_manager, config_path, - quota_state_path, + quota_state, detected_ips_rx, mutation_lock: Arc::new(Mutex::new(())), minimal_cache: Arc::new(Mutex::new(None)), @@ -373,6 +368,7 @@ pub async fn serve( runtime_watch_rx, runtime_state.clone(), shared.runtime_events.clone(), + &control_plane, ); let connection_permits = Arc::new(Semaphore::new(API_MAX_CONTROL_CONNECTIONS)); @@ -399,7 +395,7 @@ pub async fn serve( }; let shared_conn = shared.clone(); - tokio::spawn(async move { + let _ = control_plane.spawn(async move { let _connection_permit = connection_permit; let svc = service_fn(move |req: Request| { let shared_req = shared_conn.clone(); @@ -994,6 +990,7 @@ async fn handle( )); } let expected_revision = parse_if_match(req.headers()); + let _mutation_guard = shared.mutation_lock.lock().await; let disk_cfg = load_config_from_disk(&shared.config_path).await?; ensure_expected_revision(&shared.config_path, expected_revision.as_deref()) .await?; @@ -1003,12 +1000,16 @@ async fn handle( ApiFailure::new(StatusCode::NOT_FOUND, "not_found", "User not found"), )); } - let snapshot = match crate::quota_state::reset_user_quota( - &shared.quota_state_path, - shared.stats.as_ref(), - user, - ) - .await + let configured_users = disk_cfg + .access + .users + .keys() + .cloned() + .collect::>(); + let snapshot = match shared + .quota_state + .reset_user(&configured_users, user) + .await { Ok(snapshot) => snapshot, Err(error) => { diff --git a/src/api/model.rs b/src/api/model.rs index 5a183d5..2bae5fe 100644 --- a/src/api/model.rs +++ b/src/api/model.rs @@ -211,6 +211,7 @@ pub(super) struct ZeroMiddleProxyData { pub(super) reconnect_success_total: u64, pub(super) handshake_reject_total: u64, pub(super) handshake_error_codes: Vec, + pub(super) handshake_error_code_overflow_total: u64, pub(super) reader_eof_total: u64, pub(super) idle_close_by_peer_total: u64, pub(super) route_drop_no_conn_total: u64, @@ -388,8 +389,11 @@ pub(super) struct MinimalDcPathData { pub(super) struct MinimalMeRuntimeData { pub(super) active_generation: u64, pub(super) warm_generation: u64, + pub(super) warm_generations: Vec, pub(super) pending_hardswap_generation: u64, pub(super) pending_hardswap_age_secs: Option, + pub(super) reinit_inflight: usize, + pub(super) reinit_max_concurrency_effective: usize, pub(super) hardswap_enabled: bool, pub(super) floor_mode: &'static str, pub(super) adaptive_floor_idle_secs: u64, diff --git a/src/api/runtime_min.rs b/src/api/runtime_min.rs index 986f138..794fd42 100644 --- a/src/api/runtime_min.rs +++ b/src/api/runtime_min.rs @@ -21,8 +21,11 @@ pub(super) struct SecurityWhitelistData { pub(super) struct RuntimeMePoolStateGenerationData { pub(super) active_generation: u64, pub(super) warm_generation: u64, + pub(super) warm_generations: Vec, pub(super) pending_hardswap_generation: u64, pub(super) pending_hardswap_age_secs: Option, + pub(super) reinit_inflight: usize, + pub(super) reinit_max_concurrency_effective: usize, pub(super) draining_generations: Vec, } @@ -67,6 +70,8 @@ pub(super) struct RuntimeMePoolStateRefillDcData { pub(super) struct RuntimeMePoolStateRefillData { pub(super) inflight_endpoints_total: usize, pub(super) inflight_dc_total: usize, + pub(super) running_dc_total: usize, + pub(super) pending_dc_total: usize, pub(super) by_dc: Vec, } @@ -291,8 +296,7 @@ pub(super) async fn build_runtime_me_pool_state_data(shared: &ApiShared) -> Runt }; }; - let status = pool.api_status_snapshot().await; - let runtime = pool.api_runtime_snapshot().await; + let (status, runtime) = pool.api_coherent_snapshots().await; let refill = pool.api_refill_snapshot().await; let mut draining_generations = BTreeSet::::new(); @@ -329,8 +333,11 @@ pub(super) async fn build_runtime_me_pool_state_data(shared: &ApiShared) -> Runt generations: RuntimeMePoolStateGenerationData { active_generation: runtime.active_generation, warm_generation: runtime.warm_generation, + warm_generations: runtime.warm_generations, pending_hardswap_generation: runtime.pending_hardswap_generation, pending_hardswap_age_secs: runtime.pending_hardswap_age_secs, + reinit_inflight: runtime.reinit_inflight, + reinit_max_concurrency_effective: runtime.reinit_max_concurrency_effective, draining_generations: draining_generations.into_iter().collect(), }, hardswap: RuntimeMePoolStateHardswapData { @@ -356,6 +363,8 @@ pub(super) async fn build_runtime_me_pool_state_data(shared: &ApiShared) -> Runt refill: RuntimeMePoolStateRefillData { inflight_endpoints_total: refill.inflight_endpoints_total, inflight_dc_total: refill.inflight_dc_total, + running_dc_total: refill.running_dc_total, + pending_dc_total: refill.pending_dc_total, by_dc: refill .by_dc .into_iter() diff --git a/src/api/runtime_stats.rs b/src/api/runtime_stats.rs index 6fd9bb6..298e5a0 100644 --- a/src/api/runtime_stats.rs +++ b/src/api/runtime_stats.rs @@ -84,6 +84,8 @@ pub(super) fn build_zero_all_data(stats: &Stats, configured_users: usize) -> Zer reconnect_success_total: stats.get_me_reconnect_success(), handshake_reject_total: stats.get_me_handshake_reject_total(), handshake_error_codes, + handshake_error_code_overflow_total: stats + .get_me_handshake_error_code_overflow_total(), reader_eof_total: stats.get_me_reader_eof_total(), idle_close_by_peer_total: stats.get_me_idle_close_by_peer_total(), route_drop_no_conn_total: stats.get_me_route_drop_no_conn(), @@ -342,8 +344,7 @@ async fn get_minimal_payload_cached( } let pool = shared.me_pool.read().await.clone()?; - let status = pool.api_status_snapshot().await; - let runtime = pool.api_runtime_snapshot().await; + let (status, runtime) = pool.api_coherent_snapshots().await; let generated_at_epoch_secs = status.generated_at_epoch_secs; let me_writers = MeWritersData { @@ -425,8 +426,11 @@ async fn get_minimal_payload_cached( let me_runtime = MinimalMeRuntimeData { active_generation: runtime.active_generation, warm_generation: runtime.warm_generation, + warm_generations: runtime.warm_generations, pending_hardswap_generation: runtime.pending_hardswap_generation, pending_hardswap_age_secs: runtime.pending_hardswap_age_secs, + reinit_inflight: runtime.reinit_inflight, + reinit_max_concurrency_effective: runtime.reinit_max_concurrency_effective, hardswap_enabled: runtime.hardswap_enabled, floor_mode: runtime.floor_mode, adaptive_floor_idle_secs: runtime.adaptive_floor_idle_secs, diff --git a/src/api/runtime_watch.rs b/src/api/runtime_watch.rs index 50a602d..4788a96 100644 --- a/src/api/runtime_watch.rs +++ b/src/api/runtime_watch.rs @@ -5,6 +5,7 @@ use std::time::{SystemTime, UNIX_EPOCH}; use tokio::sync::watch; use crate::maestro::generation::RuntimeWatchState; +use crate::maestro::control_plane::ProcessControlPlane; use super::ApiRuntimeState; use super::events::ApiEventStore; @@ -13,22 +14,29 @@ pub(super) fn spawn_runtime_watchers( runtime_watch_rx: watch::Receiver>, runtime_state: Arc, runtime_events: Arc, + control_plane: &ProcessControlPlane, ) { - let _config_watcher = spawn_config_watcher( + spawn_config_watcher( runtime_watch_rx.clone(), runtime_state.clone(), runtime_events.clone(), + control_plane, + ); + spawn_admission_watcher( + runtime_watch_rx, + runtime_state, + runtime_events, + control_plane, ); - let _admission_watcher = - spawn_admission_watcher(runtime_watch_rx, runtime_state, runtime_events); } fn spawn_config_watcher( mut runtime_watch_rx: watch::Receiver>, runtime_state: Arc, runtime_events: Arc, -) -> tokio::task::JoinHandle<()> { - tokio::spawn(async move { + control_plane: &ProcessControlPlane, +) { + let _ = control_plane.spawn(async move { let Some(mut current) = runtime_watch_rx.borrow().clone() else { return; }; @@ -78,15 +86,16 @@ fn spawn_config_watcher( } } } - }) + }); } fn spawn_admission_watcher( mut runtime_watch_rx: watch::Receiver>, runtime_state: Arc, runtime_events: Arc, -) -> tokio::task::JoinHandle<()> { - tokio::spawn(async move { + control_plane: &ProcessControlPlane, +) { + let _ = control_plane.spawn(async move { let Some(mut current) = runtime_watch_rx.borrow().clone() else { return; }; @@ -124,7 +133,7 @@ fn spawn_admission_watcher( } } } - }) + }); } fn active_generation_id( @@ -246,7 +255,13 @@ mod tests { let (runtime_watch_tx, runtime_watch_rx) = watch::channel(Some(initial)); let runtime_state = runtime_state(); let events = Arc::new(ApiEventStore::new(16)); - spawn_runtime_watchers(runtime_watch_rx, runtime_state.clone(), events.clone()); + let control_plane = ProcessControlPlane::new(); + spawn_runtime_watchers( + runtime_watch_rx, + runtime_state.clone(), + events.clone(), + &control_plane, + ); tokio::task::yield_now().await; assert_eq!(runtime_state.config_reload_count.load(Ordering::Relaxed), 0); @@ -283,6 +298,7 @@ mod tests { .count(), 3 ); + assert!(control_plane.shutdown(Duration::from_secs(1)).await); } #[tokio::test] @@ -291,7 +307,13 @@ mod tests { let (runtime_watch_tx, runtime_watch_rx) = watch::channel(Some(initial)); let runtime_state = runtime_state(); let events = Arc::new(ApiEventStore::new(16)); - let watcher = spawn_config_watcher(runtime_watch_rx, runtime_state.clone(), events.clone()); + let control_plane = ProcessControlPlane::new(); + spawn_config_watcher( + runtime_watch_rx, + runtime_state.clone(), + events.clone(), + &control_plane, + ); drop(initial_config_tx); tokio::task::yield_now().await; @@ -302,10 +324,7 @@ mod tests { wait_for_count(&runtime_state, 2).await; drop(runtime_watch_tx); - tokio::time::timeout(Duration::from_secs(1), watcher) - .await - .unwrap() - .unwrap(); + assert!(control_plane.shutdown(Duration::from_secs(1)).await); assert_eq!( events .snapshot(16) diff --git a/src/api/users/lifecycle.rs b/src/api/users/lifecycle.rs index fd0a281..a68850e 100644 --- a/src/api/users/lifecycle.rs +++ b/src/api/users/lifecycle.rs @@ -1,4 +1,5 @@ use super::*; +use tracing::warn; pub(in crate::api) async fn rotate_secret( user: &str, @@ -108,6 +109,18 @@ pub(in crate::api) async fn delete_user( .map_err(|e| ApiFailure::bad_request(format!("config validation failed: {}", e)))?; let revision = save_access_sections_to_disk(&shared.config_path, &cfg, &touched_sections).await?; + let configured_users = cfg.access.users.keys().cloned().collect(); + if let Err(error) = shared + .quota_state + .remove_user(&configured_users, user) + .await + { + warn!( + user, + error = %error, + "Deleted user quota checkpoint cleanup will be reconciled on restart" + ); + } drop(_guard); shared.ip_tracker.remove_user_limit(user).await; shared.ip_tracker.clear_user_ips(user).await; diff --git a/src/api/web_runtime/observability.rs b/src/api/web_runtime/observability.rs index f327f93..9dad37c 100644 --- a/src/api/web_runtime/observability.rs +++ b/src/api/web_runtime/observability.rs @@ -18,6 +18,7 @@ pub(super) struct WebIngressStatus { } impl WebIngressStatus { + /// Builds a process-ingress snapshot without probing external TLS termination. pub(super) fn new(publication: &WebRuntimePublication, runtime_available: bool) -> Self { let configured_listeners = publication.listeners.len(); let live_acceptors = publication.telemetry.live_acceptors(); @@ -64,6 +65,7 @@ pub(super) struct WebCapacityStatus { } impl WebCapacityStatus { + /// Builds a bounded capacity snapshot from non-blocking runtime observations. pub(super) fn new( publication: &WebRuntimePublication, runtime: Option<&WebProcessRuntime>, @@ -104,6 +106,7 @@ pub(super) struct WebDecoyUpstreamStatus { } impl WebDecoyUpstreamStatus { + /// Builds the fixed internal decoy-origin outcome snapshot. pub(super) fn new(publication: &WebRuntimePublication) -> Self { let last = publication.telemetry.last_decoy(); Self { diff --git a/src/config/defaults.rs b/src/config/defaults.rs index 8d3ad17..3236b74 100644 --- a/src/config/defaults.rs +++ b/src/config/defaults.rs @@ -780,6 +780,10 @@ pub(crate) fn default_me_reinit_singleflight() -> bool { true } +pub(crate) fn default_me_reinit_max_concurrency() -> usize { + 2 +} + pub(crate) fn default_me_reinit_trigger_channel() -> usize { 64 } diff --git a/src/config/hot_reload/fields.rs b/src/config/hot_reload/fields.rs index dd83903..372681c 100644 --- a/src/config/hot_reload/fields.rs +++ b/src/config/hot_reload/fields.rs @@ -10,6 +10,7 @@ pub struct HotFields { pub update_every_secs: u64, pub me_reinit_every_secs: u64, pub me_reinit_singleflight: bool, + pub me_reinit_max_concurrency: usize, pub me_reinit_coalesce_window_ms: u64, pub hardswap: bool, pub me_pool_drain_ttl_secs: u64, @@ -102,6 +103,7 @@ impl HotFields { update_every_secs: cfg.general.effective_update_every_secs(), me_reinit_every_secs: cfg.general.me_reinit_every_secs, me_reinit_singleflight: cfg.general.me_reinit_singleflight, + me_reinit_max_concurrency: cfg.general.me_reinit_max_concurrency, me_reinit_coalesce_window_ms: cfg.general.me_reinit_coalesce_window_ms, hardswap: cfg.general.hardswap, me_pool_drain_ttl_secs: cfg.general.me_pool_drain_ttl_secs, @@ -236,6 +238,7 @@ pub(super) fn overlay_hot_fields(old: &ProxyConfig, new: &ProxyConfig) -> ProxyC cfg.general.proxy_config_auto_reload_secs = new.general.proxy_config_auto_reload_secs; cfg.general.me_reinit_every_secs = new.general.me_reinit_every_secs; cfg.general.me_reinit_singleflight = new.general.me_reinit_singleflight; + cfg.general.me_reinit_max_concurrency = new.general.me_reinit_max_concurrency; cfg.general.me_reinit_coalesce_window_ms = new.general.me_reinit_coalesce_window_ms; cfg.general.hardswap = new.general.hardswap; cfg.general.me_pool_drain_ttl_secs = new.general.me_pool_drain_ttl_secs; diff --git a/src/config/hot_reload/reporting.rs b/src/config/hot_reload/reporting.rs index f93b8dd..f532ee1 100644 --- a/src/config/hot_reload/reporting.rs +++ b/src/config/hot_reload/reporting.rs @@ -114,12 +114,14 @@ pub(super) fn log_changes( } if old_hot.me_reinit_every_secs != new_hot.me_reinit_every_secs || old_hot.me_reinit_singleflight != new_hot.me_reinit_singleflight + || old_hot.me_reinit_max_concurrency != new_hot.me_reinit_max_concurrency || old_hot.me_reinit_coalesce_window_ms != new_hot.me_reinit_coalesce_window_ms { info!( - "config reload: me_reinit: interval={}s singleflight={} coalesce={}ms", + "config reload: me_reinit: interval={}s singleflight={} max_concurrency={} coalesce={}ms", new_hot.me_reinit_every_secs, new_hot.me_reinit_singleflight, + new_hot.me_reinit_max_concurrency, new_hot.me_reinit_coalesce_window_ms ); } diff --git a/src/config/hot_reload/tests.rs b/src/config/hot_reload/tests.rs index 2d95608..606ecd7 100644 --- a/src/config/hot_reload/tests.rs +++ b/src/config/hot_reload/tests.rs @@ -272,6 +272,7 @@ async fn candidate_watcher_waits_for_activation_and_reconciles_disk() { None, None, cancellation.clone(), + None, Some(activation_rx), ); let watcher = tokio::spawn(watcher); diff --git a/src/config/hot_reload/watcher.rs b/src/config/hot_reload/watcher.rs index 0770f2d..b0c0538 100644 --- a/src/config/hot_reload/watcher.rs +++ b/src/config/hot_reload/watcher.rs @@ -124,13 +124,14 @@ fn apply_watch_manifest( } /// Load config, validate, diff against current, and broadcast if changed. -pub(super) fn reload_config( +fn reload_config_with_resolver( config_path: &PathBuf, config_tx: &watch::Sender>, log_tx: &watch::Sender, detected_ip_v4: Option, detected_ip_v6: Option, reload_state: &mut ReloadState, + dns_resolver: Option<&crate::network::dns_overrides::GenerationDnsResolver>, ) -> Option { let loaded = match ProxyConfig::load_with_metadata(config_path) { Ok(loaded) => loaded, @@ -176,7 +177,8 @@ pub(super) fn reload_config( } if old_hot.dns_overrides != applied_hot.dns_overrides - && let Err(e) = crate::network::dns_overrides::install_entries(&applied_hot.dns_overrides) + && let Some(dns_resolver) = dns_resolver + && let Err(e) = dns_resolver.apply_entries(&applied_hot.dns_overrides) { error!( "config reload: invalid network.dns_overrides: {}; keeping old config", @@ -198,6 +200,26 @@ pub(super) fn reload_config( Some(next_manifest) } +#[cfg(test)] +pub(super) fn reload_config( + config_path: &PathBuf, + config_tx: &watch::Sender>, + log_tx: &watch::Sender, + detected_ip_v4: Option, + detected_ip_v6: Option, + reload_state: &mut ReloadState, +) -> Option { + reload_config_with_resolver( + config_path, + config_tx, + log_tx, + detected_ip_v4, + detected_ip_v6, + reload_state, + None, + ) +} + /// Spawn the hot-reload watcher task. /// /// Uses `notify` (inotify on Linux) to detect file changes instantly. @@ -213,6 +235,7 @@ pub fn spawn_config_watcher( detected_ip_v4: Option, detected_ip_v6: Option, cancellation: tokio_util::sync::CancellationToken, + dns_resolver: Option>, mut activation: Option>, ) -> ( watch::Receiver>, @@ -364,24 +387,26 @@ pub fn spawn_config_watcher( tokio::time::sleep(HOT_RELOAD_DEBOUNCE).await; while notify_rx.try_recv().is_ok() {} - let mut next_manifest = reload_config( + let mut next_manifest = reload_config_with_resolver( &config_path, &config_tx, &log_tx, detected_ip_v4, detected_ip_v6, &mut reload_state, + dns_resolver.as_deref(), ); if next_manifest.is_none() { tokio::time::sleep(HOT_RELOAD_DEBOUNCE).await; while notify_rx.try_recv().is_ok() {} - next_manifest = reload_config( + next_manifest = reload_config_with_resolver( &config_path, &config_tx, &log_tx, detected_ip_v4, detected_ip_v6, &mut reload_state, + dns_resolver.as_deref(), ); } diff --git a/src/config/load/strict_keys.rs b/src/config/load/strict_keys.rs index 08b9ee5..ffb3d4b 100644 --- a/src/config/load/strict_keys.rs +++ b/src/config/load/strict_keys.rs @@ -161,6 +161,7 @@ const GENERAL_CONFIG_KEYS: &[&str] = &[ "proxy_secret_auto_reload_secs", "proxy_config_auto_reload_secs", "me_reinit_singleflight", + "me_reinit_max_concurrency", "me_reinit_trigger_channel", "me_reinit_coalesce_window_ms", "me_deterministic_writer_sort", diff --git a/src/config/load/validate_me.rs b/src/config/load/validate_me.rs index 8c24fa6..73ca3e2 100644 --- a/src/config/load/validate_me.rs +++ b/src/config/load/validate_me.rs @@ -143,9 +143,15 @@ pub(super) fn validate(config: &mut ProxyConfig) -> Result<()> { )); } - if config.general.me_reinit_trigger_channel == 0 { + if !(1..=8).contains(&config.general.me_reinit_max_concurrency) { return Err(ProxyError::Config( - "general.me_reinit_trigger_channel must be > 0".to_string(), + "general.me_reinit_max_concurrency must be within [1, 8]".to_string(), + )); + } + + if !(1..=4096).contains(&config.general.me_reinit_trigger_channel) { + return Err(ProxyError::Config( + "general.me_reinit_trigger_channel must be within [1, 4096]".to_string(), )); } diff --git a/src/config/types/general.rs b/src/config/types/general.rs index 2218132..8ec14ea 100644 --- a/src/config/types/general.rs +++ b/src/config/types/general.rs @@ -441,6 +441,9 @@ pub struct GeneralConfig { /// Serialize ME reinit cycles across all trigger sources. #[serde(default = "default_me_reinit_singleflight")] pub me_reinit_singleflight: bool, + /// Maximum concurrent ME reinit warmups when single-flight mode is disabled. + #[serde(default = "default_me_reinit_max_concurrency")] + pub me_reinit_max_concurrency: usize, /// Trigger queue capacity for reinit scheduler. #[serde(default = "default_me_reinit_trigger_channel")] pub me_reinit_trigger_channel: usize, diff --git a/src/config/types/general_impl.rs b/src/config/types/general_impl.rs index a0afb36..36de4da 100644 --- a/src/config/types/general_impl.rs +++ b/src/config/types/general_impl.rs @@ -159,6 +159,7 @@ impl Default for GeneralConfig { proxy_secret_auto_reload_secs: default_proxy_secret_reload_secs(), proxy_config_auto_reload_secs: default_proxy_config_reload_secs(), me_reinit_singleflight: default_me_reinit_singleflight(), + me_reinit_max_concurrency: default_me_reinit_max_concurrency(), me_reinit_trigger_channel: default_me_reinit_trigger_channel(), me_reinit_coalesce_window_ms: default_me_reinit_coalesce_window_ms(), me_deterministic_writer_sort: default_me_deterministic_writer_sort(), diff --git a/src/healthcheck.rs b/src/healthcheck.rs index c9fa610..ddbced1 100644 --- a/src/healthcheck.rs +++ b/src/healthcheck.rs @@ -6,6 +6,8 @@ use serde_json::Value; use crate::config::ProxyConfig; +const HEALTHCHECK_RESPONSE_MAX_BYTES: u64 = 64 * 1024; + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) enum HealthcheckMode { Liveness, @@ -73,10 +75,7 @@ fn run_inner(config_path: &str, mode: HealthcheckMode) -> Result<(), String> { .flush() .map_err(|error| format!("request flush failed: {error}"))?; - let mut raw_response = Vec::new(); - stream - .read_to_end(&mut raw_response) - .map_err(|error| format!("response read failed: {error}"))?; + let raw_response = read_response_bounded(&mut stream)?; let response = String::from_utf8(raw_response).map_err(|_| "response is not valid UTF-8".to_string())?; @@ -89,6 +88,18 @@ fn run_inner(config_path: &str, mode: HealthcheckMode) -> Result<(), String> { Ok(()) } +fn read_response_bounded(reader: &mut impl Read) -> Result, String> { + let mut raw_response = Vec::new(); + reader + .take(HEALTHCHECK_RESPONSE_MAX_BYTES.saturating_add(1)) + .read_to_end(&mut raw_response) + .map_err(|error| format!("response read failed: {error}"))?; + if raw_response.len() as u64 > HEALTHCHECK_RESPONSE_MAX_BYTES { + return Err("response exceeds the 64 KiB healthcheck limit".to_string()); + } + Ok(raw_response) +} + fn probe_target(listen: SocketAddr) -> SocketAddr { match listen { SocketAddr::V4(addr) => { @@ -180,7 +191,10 @@ fn validate_payload(mode: HealthcheckMode, body: &str) -> Result<(), String> { #[cfg(test)] mod tests { - use super::{HealthcheckMode, parse_status_code, split_response, validate_payload}; + use super::{ + HEALTHCHECK_RESPONSE_MAX_BYTES, HealthcheckMode, parse_status_code, + read_response_bounded, split_response, validate_payload, + }; #[test] fn parse_status_code_reads_http_200() { @@ -208,4 +222,11 @@ mod tests { let result = validate_payload(HealthcheckMode::Ready, body); assert!(result.is_err()); } + + #[test] + fn bounded_reader_rejects_oversized_health_response() { + let payload = vec![b'x'; HEALTHCHECK_RESPONSE_MAX_BYTES as usize + 1]; + + assert!(read_response_bounded(&mut payload.as_slice()).is_err()); + } } diff --git a/src/ip_tracker.rs b/src/ip_tracker.rs index e35ad73..a54410f 100644 --- a/src/ip_tracker.rs +++ b/src/ip_tracker.rs @@ -8,10 +8,10 @@ use std::hash::{Hash, Hasher}; use std::net::IpAddr; use std::sync::Arc; use std::sync::Mutex; -use std::sync::atomic::{AtomicU8, AtomicU64, AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{Duration, Instant}; -use dashmap::DashMap; +use arc_swap::ArcSwap; use tokio::sync::{Mutex as AsyncMutex, RwLock}; use crate::config::UserMaxUniqueIpsMode; @@ -39,6 +39,25 @@ struct CleanupShard { queue: Mutex>>, } +#[derive(Debug, Clone)] +struct UserIpLimitPolicy { + max_ips: Arc>, + default_max_ips: usize, + mode: UserMaxUniqueIpsMode, + window_secs: u64, +} + +impl Default for UserIpLimitPolicy { + fn default() -> Self { + Self { + max_ips: Arc::new(HashMap::new()), + default_max_ips: 0, + mode: UserMaxUniqueIpsMode::ActiveWindow, + window_secs: 30, + } + } +} + /// Tracks active and recent client IPs for per-user admission control. #[derive(Debug, Clone)] pub struct UserIpTracker { @@ -48,10 +67,7 @@ pub struct UserIpTracker { active_cap_rejects: Arc, recent_cap_rejects: Arc, cleanup_deferred_releases: Arc, - max_ips: Arc>, - default_max_ips: Arc, - limit_mode: Arc, - limit_window_secs: Arc, + limit_policy: Arc>, last_compact_epoch_secs: Arc, cleanup_queue_len: Arc, cleanup_shards: Arc>, @@ -102,12 +118,7 @@ impl UserIpTracker { active_cap_rejects: Arc::new(AtomicU64::new(0)), recent_cap_rejects: Arc::new(AtomicU64::new(0)), cleanup_deferred_releases: Arc::new(AtomicU64::new(0)), - max_ips: Arc::new(DashMap::new()), - default_max_ips: Arc::new(AtomicUsize::new(0)), - limit_mode: Arc::new(AtomicU8::new(Self::mode_to_u8( - UserMaxUniqueIpsMode::ActiveWindow, - ))), - limit_window_secs: Arc::new(AtomicU64::new(30)), + limit_policy: Arc::new(ArcSwap::from_pointee(UserIpLimitPolicy::default())), last_compact_epoch_secs: Arc::new(AtomicU64::new(0)), cleanup_queue_len: Arc::new(AtomicU64::new(0)), cleanup_shards: Arc::new(cleanup_shards), @@ -117,41 +128,23 @@ impl UserIpTracker { } } - pub(super) fn mode_to_u8(mode: UserMaxUniqueIpsMode) -> u8 { - match mode { - UserMaxUniqueIpsMode::ActiveWindow => 0, - UserMaxUniqueIpsMode::TimeWindow => 1, - UserMaxUniqueIpsMode::Combined => 2, - } - } - - pub(super) fn mode_from_u8(raw: u8) -> UserMaxUniqueIpsMode { - match raw { - 1 => UserMaxUniqueIpsMode::TimeWindow, - 2 => UserMaxUniqueIpsMode::Combined, - _ => UserMaxUniqueIpsMode::ActiveWindow, - } - } - pub(super) fn shard_idx(username: &str) -> usize { let mut hasher = DefaultHasher::new(); username.hash(&mut hasher); (hasher.finish() as usize) & USER_IP_TRACKER_SHARD_MASK } - pub(super) fn limit_window(&self) -> Duration { - Duration::from_secs(self.limit_window_secs.load(Ordering::Relaxed).max(1)) + fn limit_window(policy: &UserIpLimitPolicy) -> Duration { + Duration::from_secs(policy.window_secs) } - pub(super) fn user_limit(&self, username: &str) -> Option { - self.max_ips + fn user_limit(policy: &UserIpLimitPolicy, username: &str) -> Option { + policy + .max_ips .get(username) - .map(|limit| *limit) + .copied() .filter(|limit| *limit > 0) - .or_else(|| { - let default_limit = self.default_max_ips.load(Ordering::Relaxed); - (default_limit > 0).then_some(default_limit) - }) + .or_else(|| (policy.default_max_ips > 0).then_some(policy.default_max_ips)) } pub(super) fn decrement_counter(counter: &AtomicU64, amount: usize) { diff --git a/src/ip_tracker/admission.rs b/src/ip_tracker/admission.rs index 314b273..df3bd98 100644 --- a/src/ip_tracker/admission.rs +++ b/src/ip_tracker/admission.rs @@ -2,26 +2,47 @@ use super::*; impl UserIpTracker { pub async fn set_limit_policy(&self, mode: UserMaxUniqueIpsMode, window_secs: u64) { - self.limit_mode - .store(Self::mode_to_u8(mode), Ordering::Relaxed); - self.limit_window_secs - .store(window_secs.max(1), Ordering::Relaxed); + self.limit_policy.rcu(|current| { + Arc::new(UserIpLimitPolicy { + mode, + window_secs: window_secs.max(1), + ..(**current).clone() + }) + }); } pub async fn set_user_limit(&self, username: &str, max_ips: usize) { - self.max_ips.insert(username.to_string(), max_ips); + let username = username.to_string(); + self.limit_policy.rcu(|current| { + let mut limits = current.max_ips.as_ref().clone(); + limits.insert(username.clone(), max_ips); + Arc::new(UserIpLimitPolicy { + max_ips: Arc::new(limits), + ..(**current).clone() + }) + }); } pub async fn remove_user_limit(&self, username: &str) { - self.max_ips.remove(username); + self.limit_policy.rcu(|current| { + let mut limits = current.max_ips.as_ref().clone(); + limits.remove(username); + Arc::new(UserIpLimitPolicy { + max_ips: Arc::new(limits), + ..(**current).clone() + }) + }); } pub async fn load_limits(&self, default_limit: usize, limits: &HashMap) { - self.default_max_ips.store(default_limit, Ordering::Relaxed); - self.max_ips.clear(); - for (username, limit) in limits { - self.max_ips.insert(username.clone(), *limit); - } + let limits = Arc::new(limits.clone()); + self.limit_policy.rcu(|current| { + Arc::new(UserIpLimitPolicy { + max_ips: Arc::clone(&limits), + default_max_ips: default_limit, + ..(**current).clone() + }) + }); } pub(super) fn prune_recent( @@ -40,9 +61,10 @@ impl UserIpTracker { pub async fn check_and_add(&self, username: &str, ip: IpAddr) -> Result<(), String> { self.drain_cleanup_for_user(username).await; self.maybe_compact_empty_users().await; - let limit = self.user_limit(username); - let mode = Self::mode_from_u8(self.limit_mode.load(Ordering::Relaxed)); - let window = self.limit_window(); + let policy = self.limit_policy.load(); + let limit = Self::user_limit(&policy, username); + let mode = policy.mode; + let window = Self::limit_window(&policy); let now = Instant::now(); let shard_idx = Self::shard_idx(username); diff --git a/src/ip_tracker/snapshot.rs b/src/ip_tracker/snapshot.rs index 0d96b4b..f696c4e 100644 --- a/src/ip_tracker/snapshot.rs +++ b/src/ip_tracker/snapshot.rs @@ -21,7 +21,8 @@ impl UserIpTracker { return; } - let window = self.limit_window(); + let policy = self.limit_policy.load(); + let window = Self::limit_window(&policy); let now = Instant::now(); for shard_lock in self.shards.iter() { let mut shard = shard_lock.write().await; @@ -113,7 +114,8 @@ impl UserIpTracker { &self, users: &[String], ) -> HashMap { - let window = self.limit_window(); + let policy = self.limit_policy.load(); + let window = Self::limit_window(&policy); let now = Instant::now(); let mut counts = HashMap::with_capacity(users.len()); @@ -152,7 +154,8 @@ impl UserIpTracker { pub async fn get_recent_ips_for_users(&self, users: &[String]) -> HashMap> { self.drain_cleanup_queue().await; - let window = self.limit_window(); + let policy = self.limit_policy.load(); + let window = Self::limit_window(&policy); let now = Instant::now(); let mut out = HashMap::with_capacity(users.len()); @@ -202,6 +205,7 @@ impl UserIpTracker { } pub(crate) async fn get_stats_snapshot(&self) -> Vec<(String, usize, usize)> { + let policy = self.limit_policy.load(); let mut active_counts = Vec::new(); for shard_lock in self.shards.iter() { let shard = shard_lock.read().await; @@ -215,7 +219,7 @@ impl UserIpTracker { let mut stats = Vec::with_capacity(active_counts.len()); for (username, active_count) in active_counts { - let limit = self.user_limit(&username).unwrap_or(0); + let limit = Self::user_limit(&policy, &username).unwrap_or(0); stats.push((username, active_count, limit)); } @@ -273,7 +277,8 @@ impl UserIpTracker { } pub async fn get_user_limit(&self, username: &str) -> Option { - self.user_limit(username) + let policy = self.limit_policy.load(); + Self::user_limit(&policy, username) } pub async fn format_stats(&self) -> String { diff --git a/src/ip_tracker/tests.rs b/src/ip_tracker/tests.rs index ba1b1fe..bf2925b 100644 --- a/src/ip_tracker/tests.rs +++ b/src/ip_tracker/tests.rs @@ -1,5 +1,6 @@ use super::*; use std::net::{IpAddr, Ipv4Addr, Ipv6Addr}; +use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; fn test_ipv4(oct1: u8, oct2: u8, oct3: u8, oct4: u8) -> IpAddr { @@ -232,6 +233,55 @@ async fn test_load_limits_replaces_previous_map() { assert_eq!(tracker.get_user_limit("user2").await, Some(5)); } +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn concurrent_policy_replacement_never_exposes_partial_limit_map() { + const USER_COUNT: usize = 4_096; + const REPLACEMENTS: usize = 32; + + let tracker = Arc::new(UserIpTracker::new()); + let first = (0..USER_COUNT) + .map(|index| (format!("user-{index}"), 3usize)) + .collect::>(); + let second = (0..USER_COUNT) + .map(|index| (format!("user-{index}"), 5usize)) + .collect::>(); + tracker.load_limits(7, &first).await; + + let running = Arc::new(AtomicBool::new(true)); + let writer_tracker = Arc::clone(&tracker); + let writer_running = Arc::clone(&running); + let writer = tokio::spawn(async move { + for _ in 0..REPLACEMENTS { + writer_tracker.load_limits(7, &second).await; + tokio::task::yield_now().await; + writer_tracker.load_limits(7, &first).await; + tokio::task::yield_now().await; + } + writer_running.store(false, Ordering::Release); + }); + + let mut readers = Vec::new(); + for reader in 0..3usize { + let reader_tracker = Arc::clone(&tracker); + let reader_running = Arc::clone(&running); + readers.push(tokio::spawn(async move { + let mut index = reader; + while reader_running.load(Ordering::Acquire) { + let username = format!("user-{}", index % USER_COUNT); + let limit = reader_tracker.get_user_limit(&username).await; + assert!(matches!(limit, Some(3 | 5)), "partial policy: {limit:?}"); + index = index.wrapping_add(17); + tokio::task::yield_now().await; + } + })); + } + + writer.await.unwrap(); + for reader in readers { + reader.await.unwrap(); + } +} + #[tokio::test] async fn test_global_each_limit_applies_without_user_override() { let tracker = UserIpTracker::new(); diff --git a/src/maestro/bootstrap.rs b/src/maestro/bootstrap.rs index 9fe8a58..a89ee39 100644 --- a/src/maestro/bootstrap.rs +++ b/src/maestro/bootstrap.rs @@ -253,10 +253,6 @@ pub(super) async fn bootstrap( } } - if let Err(e) = crate::network::dns_overrides::install_entries(&config.network.dns_overrides) { - eprintln!("[telemt] Invalid network.dns_overrides: {}", e); - std::process::exit(1); - } set_maestro_colors_enabled(!config.general.disable_colors); startup_tracker .complete_component(COMPONENT_CONFIG_LOAD, Some("config is ready".to_string())) diff --git a/src/maestro/control_plane.rs b/src/maestro/control_plane.rs new file mode 100644 index 0000000..a3100fb --- /dev/null +++ b/src/maestro/control_plane.rs @@ -0,0 +1,217 @@ +use std::future::Future; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::time::Duration; + +use tokio::sync::Notify; +use tokio_util::sync::CancellationToken; +use tokio_util::task::TaskTracker; + +const CONTROL_TASK_ADMISSION_CLOSED: usize = 1 << (usize::BITS - 1); +const CONTROL_TASK_REGISTRATION_COUNT: usize = CONTROL_TASK_ADMISSION_CLOSED - 1; + +struct ControlTaskAdmission { + state: AtomicUsize, + registrations_drained: Notify, +} + +struct ControlTaskRegistration<'a> { + admission: &'a ControlTaskAdmission, +} + +impl ControlTaskAdmission { + fn new() -> Self { + Self { + state: AtomicUsize::new(0), + registrations_drained: Notify::new(), + } + } + + fn try_register(&self) -> Option> { + let mut state = self.state.load(Ordering::Acquire); + loop { + if state & CONTROL_TASK_ADMISSION_CLOSED != 0 + || state & CONTROL_TASK_REGISTRATION_COUNT == CONTROL_TASK_REGISTRATION_COUNT + { + return None; + } + match self.state.compare_exchange_weak( + state, + state + 1, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => return Some(ControlTaskRegistration { admission: self }), + Err(observed) => state = observed, + } + } + } + + fn close(&self) { + self.state + .fetch_or(CONTROL_TASK_ADMISSION_CLOSED, Ordering::AcqRel); + } + + async fn wait_for_registrations(&self) { + loop { + let notified = self.registrations_drained.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); + if self.state.load(Ordering::Acquire) & CONTROL_TASK_REGISTRATION_COUNT == 0 { + return; + } + notified.await; + } + } +} + +impl Drop for ControlTaskRegistration<'_> { + fn drop(&mut self) { + let previous = self.admission.state.fetch_sub(1, Ordering::AcqRel); + if previous & CONTROL_TASK_REGISTRATION_COUNT == 1 { + self.admission.registrations_drained.notify_waiters(); + } + } +} + +struct ProcessControlPlaneInner { + admission: ControlTaskAdmission, + cancellation: CancellationToken, + tasks: TaskTracker, + shutdown_completed: AtomicBool, +} + +/// Process-owned cancellation and join scope for API, metrics, and signal tasks. +#[derive(Clone)] +pub(crate) struct ProcessControlPlane { + inner: Arc, +} + +impl ProcessControlPlane { + /// Creates an open process control-plane scope. + pub(crate) fn new() -> Self { + Self { + inner: Arc::new(ProcessControlPlaneInner { + admission: ControlTaskAdmission::new(), + cancellation: CancellationToken::new(), + tasks: TaskTracker::new(), + shutdown_completed: AtomicBool::new(false), + }), + } + } + + /// Registers a cancellable process control-plane task before it can be unpolled. + pub(crate) fn spawn(&self, future: F) -> Result<(), F> + where + F: Future + Send + 'static, + { + let Some(registration) = self.inner.admission.try_register() else { + return Err(future); + }; + let cancellation = self.inner.cancellation.clone(); + self.inner.tasks.spawn(async move { + tokio::select! { + biased; + _ = cancellation.cancelled() => {} + _ = future => {} + } + }); + drop(registration); + Ok(()) + } + + /// Closes task admission, cancels all owned work, and joins it within the deadline. + pub(crate) async fn shutdown(&self, timeout: Duration) -> bool { + let deadline = tokio::time::Instant::now() + timeout; + self.inner.admission.close(); + self.inner.cancellation.cancel(); + self.inner.tasks.close(); + if self.inner.shutdown_completed.load(Ordering::Acquire) { + return true; + } + let registrations_stopped = tokio::time::timeout_at( + deadline, + self.inner.admission.wait_for_registrations(), + ) + .await + .is_ok(); + let tasks_stopped = tokio::time::timeout_at(deadline, self.inner.tasks.wait()) + .await + .is_ok(); + let outcome = registrations_stopped && tasks_stopped; + if outcome { + self.inner.shutdown_completed.store(true, Ordering::Release); + } + outcome + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::sync::atomic::{AtomicBool, Ordering}; + + use super::*; + + #[tokio::test] + async fn shutdown_cancels_owned_tasks_and_rejects_late_registration() { + struct DropSignal(Arc); + + impl Drop for DropSignal { + fn drop(&mut self) { + self.0.store(true, Ordering::Release); + } + } + + let scope = ProcessControlPlane::new(); + let dropped = Arc::new(AtomicBool::new(false)); + let drop_signal = DropSignal(dropped.clone()); + assert!( + scope + .spawn(async move { + let _drop_signal = drop_signal; + std::future::pending::<()>().await; + }) + .is_ok() + ); + + assert!(scope.shutdown(Duration::from_secs(1)).await); + assert!(dropped.load(Ordering::Acquire)); + assert!(scope.spawn(async {}).is_err()); + } + + #[tokio::test] + async fn concurrent_shutdown_callers_wait_for_completion() { + let scope = ProcessControlPlane::new(); + let registration = scope.inner.admission.try_register().unwrap(); + let first_scope = scope.clone(); + let first = tokio::spawn(async move { first_scope.shutdown(Duration::from_secs(1)).await }); + tokio::task::yield_now().await; + let second_scope = scope.clone(); + let second = tokio::spawn(async move { second_scope.shutdown(Duration::from_secs(1)).await }); + + tokio::task::yield_now().await; + assert!(!first.is_finished()); + assert!(!second.is_finished()); + drop(registration); + + assert!(first.await.unwrap()); + assert!(second.await.unwrap()); + } + + #[tokio::test] + async fn cancelled_shutdown_caller_cannot_orphan_the_control_plane() { + let scope = ProcessControlPlane::new(); + let registration = scope.inner.admission.try_register().unwrap(); + let first_scope = scope.clone(); + let first = tokio::spawn(async move { first_scope.shutdown(Duration::from_secs(30)).await }); + tokio::task::yield_now().await; + + first.abort(); + assert!(first.await.unwrap_err().is_cancelled()); + assert!(scope.spawn(async {}).is_err()); + drop(registration); + + assert!(scope.shutdown(Duration::from_secs(1)).await); + } +} diff --git a/src/maestro/generation.rs b/src/maestro/generation.rs index a81c786..870b683 100644 --- a/src/maestro/generation.rs +++ b/src/maestro/generation.rs @@ -3,7 +3,7 @@ use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; -use tokio::sync::{RwLock, Semaphore, watch}; +use tokio::sync::{Notify, RwLock, Semaphore, watch}; use tokio_util::sync::CancellationToken; use tokio_util::task::TaskTracker; @@ -29,6 +29,7 @@ const SESSION_REGISTRATION_COUNT: usize = SESSION_ADMISSION_CLOSED - 1; struct SessionAdmission { state: AtomicUsize, + registrations_drained: Notify, } struct SessionRegistration<'a> { @@ -39,6 +40,7 @@ impl SessionAdmission { fn new() -> Self { Self { state: AtomicUsize::new(0), + registrations_drained: Notify::new(), } } @@ -73,15 +75,24 @@ impl SessionAdmission { } async fn wait_for_registrations(&self) { - while self.state.load(Ordering::Acquire) & SESSION_REGISTRATION_COUNT != 0 { - tokio::task::yield_now().await; + loop { + let notified = self.registrations_drained.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); + if self.state.load(Ordering::Acquire) & SESSION_REGISTRATION_COUNT == 0 { + return; + } + notified.await; } } } impl Drop for SessionRegistration<'_> { fn drop(&mut self) { - self.admission.state.fetch_sub(1, Ordering::Release); + let previous = self.admission.state.fetch_sub(1, Ordering::AcqRel); + if previous & SESSION_REGISTRATION_COUNT == 1 { + self.admission.registrations_drained.notify_waiters(); + } } } @@ -98,6 +109,7 @@ pub(crate) struct RuntimeWatchState { pub(crate) struct RuntimeTaskScope { tracker: TaskTracker, cancel: CancellationToken, + admission: Arc, } impl RuntimeTaskScope { @@ -106,6 +118,7 @@ impl RuntimeTaskScope { Self { tracker: TaskTracker::new(), cancel: CancellationToken::new(), + admission: Arc::new(SessionAdmission::new()), } } @@ -114,9 +127,13 @@ impl RuntimeTaskScope { where F: Future + Send + 'static, { + let Some(_registration) = self.admission.try_register() else { + return; + }; let cancel = self.cancel.clone(); self.tracker.spawn(async move { tokio::select! { + biased; _ = cancel.cancelled() => {} _ = future => {} } @@ -130,6 +147,8 @@ impl RuntimeTaskScope { /// Cancels the scope and waits within the bounded background-task budget. pub(crate) async fn stop(&self) { + self.admission.close(); + self.admission.wait_for_registrations().await; self.cancel.cancel(); self.tracker.close(); let _ = tokio::time::timeout(BACKGROUND_STOP_TIMEOUT, self.tracker.wait()).await; @@ -263,6 +282,7 @@ impl RuntimeGeneration { let cancel = self.session_cancel.clone(); self.sessions.spawn(async move { tokio::select! { + biased; _ = cancel.cancelled() => {} _ = future => {} } @@ -275,11 +295,6 @@ impl RuntimeGeneration { self.session_admission.close(); } - /// Reopens admission after a candidate activation rolls back. - pub(crate) fn resume_accepting_sessions(&self) { - self.session_admission.reopen(); - } - /// Waits for registered sessions and cancels them when the deadline expires. pub(crate) async fn drain_sessions(&self, timeout: Duration) -> bool { self.stop_accepting_sessions(); @@ -308,6 +323,27 @@ impl RuntimeGeneration { pub(crate) async fn stop_background_tasks(&self) { self.background_tasks.stop().await; } + + /// Terminally stops the generation's Middle-End task and writer scope. + pub(crate) async fn stop_middle_end(&self, timeout: Duration) -> bool { + let Some(pool) = self.current_me_pool().await else { + return true; + }; + pool.shutdown_until(timeout).await + } +} + +impl Drop for RuntimeGeneration { + fn drop(&mut self) { + if let Some(pool) = self.me_pool.as_ref() { + pool.begin_shutdown(); + } + if let Ok(pool) = self.me_pool_runtime.try_read() + && let Some(pool) = pool.as_ref() + { + pool.begin_shutdown(); + } + } } #[cfg(test)] @@ -393,11 +429,31 @@ mod tests { #[tokio::test] async fn runtime_task_scope_joins_cancelled_background_task() { + struct DropSignal(Arc); + + impl Drop for DropSignal { + fn drop(&mut self) { + self.0.fetch_add(1, Ordering::AcqRel); + } + } + let scope = RuntimeTaskScope::new(); - scope.spawn(std::future::pending()); + let dropped = Arc::new(AtomicUsize::new(0)); + let drop_signal = DropSignal(dropped.clone()); + scope.spawn(async move { + let _drop_signal = drop_signal; + std::future::pending::<()>().await; + }); tokio::time::timeout(Duration::from_secs(1), scope.stop()) .await .unwrap(); + assert_eq!(dropped.load(Ordering::Acquire), 1); + + let late_drop_signal = DropSignal(dropped.clone()); + scope.spawn(async move { + let _late_drop_signal = late_drop_signal; + }); + assert_eq!(dropped.load(Ordering::Acquire), 2); } #[tokio::test] diff --git a/src/maestro/listeners/accept.rs b/src/maestro/listeners/accept.rs index 9385b3e..c847c77 100644 --- a/src/maestro/listeners/accept.rs +++ b/src/maestro/listeners/accept.rs @@ -12,7 +12,7 @@ use tracing::{debug, error, info, warn}; use crate::config::{ListenerTransport, RstOnCloseMode}; use crate::proxy::ClientHandler; use crate::transport::socket::set_linger_zero; -use crate::web::manager::WebProcessRuntime; +use crate::web::manager::{HttpConnectionAdmissionError, WebProcessRuntime}; use crate::web::telemetry::{WebAcceptorGuard, WebHttpConnectionOverloadOutcome}; use super::bind::BoundTcpListener; @@ -208,45 +208,82 @@ async fn run_accept_loop( return; }; web_runtime.telemetry().record_accept(); - let Some(connection_permit) = web_runtime.try_http_connection() else { - let config = web_runtime.active_generation().config(); - let action = config.web.http_connection_capacity_action; - let phase_timeout = - Duration::from_millis(config.web.timeouts.http_overload_timeout_ms); - drop(config); - if action == crate::config::WebHttpConnectionCapacityAction::Drop { - web_runtime.telemetry().record_rejection( - crate::web::telemetry::WebRejectionReason::HttpConnectionCapacity, + if cancellation.is_cancelled() { + drop(stream); + continue; + } + if web_runtime.is_shutdown() { + web_runtime + .telemetry() + .record_rejection( + crate::web::telemetry::WebRejectionReason::RuntimeClosed, + ); + drop(stream); + continue; + } + let connection_permit = match web_runtime.try_http_connection() { + Ok(permit) => permit, + Err(HttpConnectionAdmissionError::Closed) => { + web_runtime.telemetry().record_rejection( + crate::web::telemetry::WebRejectionReason::RuntimeClosed, ); - web_runtime - .telemetry() - .record_overload(WebHttpConnectionOverloadOutcome::Dropped); drop(stream); continue; } - let Some(overload_permit) = web_runtime.try_http_overload_connection() - else { - web_runtime.telemetry().record_rejection( - crate::web::telemetry::WebRejectionReason::HttpConnectionCapacity, + Err(HttpConnectionAdmissionError::AtCapacity) => { + let config = web_runtime.active_generation().config(); + let action = config.web.http_connection_capacity_action; + let phase_timeout = Duration::from_millis( + config.web.timeouts.http_overload_timeout_ms, ); - web_runtime.telemetry().record_overload( - WebHttpConnectionOverloadOutcome::OverflowCapacityDrop, - ); - drop(stream); + drop(config); + if action == crate::config::WebHttpConnectionCapacityAction::Drop { + web_runtime.telemetry().record_rejection( + crate::web::telemetry::WebRejectionReason::HttpConnectionCapacity, + ); + web_runtime + .telemetry() + .record_overload(WebHttpConnectionOverloadOutcome::Dropped); + drop(stream); + continue; + } + let overload_permit = + match web_runtime.try_http_overload_connection() { + Ok(permit) => permit, + Err(HttpConnectionAdmissionError::Closed) => { + web_runtime.telemetry().record_rejection( + crate::web::telemetry::WebRejectionReason::RuntimeClosed, + ); + web_runtime.telemetry().record_overload( + WebHttpConnectionOverloadOutcome::ShutdownDrop, + ); + drop(stream); + continue; + } + Err(HttpConnectionAdmissionError::AtCapacity) => { + web_runtime.telemetry().record_rejection( + crate::web::telemetry::WebRejectionReason::HttpConnectionCapacity, + ); + web_runtime.telemetry().record_overload( + WebHttpConnectionOverloadOutcome::OverflowCapacityDrop, + ); + drop(stream); + continue; + } + }; + connections.spawn(web_overload::serve( + stream, + peer_addr, + spec.web_client_ip_source, + Arc::clone(&spec.web_trusted_proxy_cidrs), + Arc::clone(web_runtime), + cancellation.clone(), + overload_permit, + action, + phase_timeout, + )); continue; - }; - connections.spawn(web_overload::serve( - stream, - peer_addr, - spec.web_client_ip_source, - Arc::clone(&spec.web_trusted_proxy_cidrs), - Arc::clone(web_runtime), - cancellation.clone(), - overload_permit, - action, - phase_timeout, - )); - continue; + } }; connections.spawn(crate::web::http::serve_connection( stream, diff --git a/src/maestro/listeners/web_overload.rs b/src/maestro/listeners/web_overload.rs index 240739f..22a7b58 100644 --- a/src/maestro/listeners/web_overload.rs +++ b/src/maestro/listeners/web_overload.rs @@ -9,9 +9,10 @@ use tokio::sync::OwnedSemaphorePermit; use tokio_util::sync::CancellationToken; use crate::config::{WebClientIpSource, WebHttpConnectionCapacityAction}; -use crate::web::manager::WebProcessRuntime; +use crate::web::manager::{HttpConnectionAdmissionError, WebProcessRuntime}; use crate::web::telemetry::{WebHttpConnectionOverloadOutcome, WebRejectionReason}; +/// Exact bounded retryable response emitted before HTTP request parsing. pub(super) const SERVICE_UNAVAILABLE_RESPONSE: &[u8] = b"HTTP/1.1 503 Service Unavailable\r\nContent-Length: 0\r\nCache-Control: no-store\r\nRetry-After: 1\r\nConnection: close\r\n\r\n"; /// Handles one accepted WEB socket outside ordinary connection capacity. @@ -44,25 +45,31 @@ pub(super) async fn serve( return; } permit = tokio::time::timeout(phase_timeout, runtime.acquire_http_connection()) => { - permit.ok().flatten() + permit } }; - let Some(connection_permit) = connection_permit else { - if runtime.is_shutdown() { + let connection_permit = match connection_permit { + Ok(Ok(permit)) => permit, + Ok(Err(HttpConnectionAdmissionError::Closed)) => { + runtime + .telemetry() + .record_rejection(WebRejectionReason::RuntimeClosed); runtime .telemetry() .record_overload(WebHttpConnectionOverloadOutcome::ShutdownDrop); return; } - let outcome = match respond(stream, &cancellation, phase_timeout).await { - WebHttpConnectionOverloadOutcome::Responded503 => { - WebHttpConnectionOverloadOutcome::WaitTimeout503 - } - other => other, - }; - record_final_capacity_rejection(&runtime, outcome); - runtime.telemetry().record_overload(outcome); - return; + Ok(Err(HttpConnectionAdmissionError::AtCapacity)) | Err(_) => { + let outcome = match respond(stream, &cancellation, phase_timeout).await { + WebHttpConnectionOverloadOutcome::Responded503 => { + WebHttpConnectionOverloadOutcome::WaitTimeout503 + } + other => other, + }; + record_final_capacity_rejection(&runtime, outcome); + runtime.telemetry().record_overload(outcome); + return; + } }; runtime .telemetry() diff --git a/src/maestro/me_startup.rs b/src/maestro/me_startup.rs index 9ac9de4..e96f8ec 100644 --- a/src/maestro/me_startup.rs +++ b/src/maestro/me_startup.rs @@ -345,6 +345,7 @@ pub(crate) async fn initialize_me_pool( config.general.me_route_blocking_send_timeout_ms, config.general.me_route_inline_recovery_attempts, config.general.me_route_inline_recovery_wait_ms, + (config.server.max_connections as usize).saturating_add(128), ); startup_tracker .complete_component( diff --git a/src/maestro/mod.rs b/src/maestro/mod.rs index 0afbed3..07565ae 100644 --- a/src/maestro/mod.rs +++ b/src/maestro/mod.rs @@ -6,6 +6,7 @@ // - admission: conditional-cast gate and route mode switching. // - bootstrap: configuration and tracing initialization. // - connectivity: startup ME/DC connectivity diagnostics. +// - control_plane: process-owned API, metrics, and signal task lifecycle. // - generation: runtime generation state and task ownership. // - helpers: CLI and shared startup/runtime helper routines. // - listeners: TCP/Unix listener planning, binding, and lifecycle control. @@ -21,6 +22,7 @@ mod admission; mod bootstrap; mod connectivity; +pub(crate) mod control_plane; pub(crate) mod generation; mod helpers; mod listeners; diff --git a/src/maestro/orchestrator.rs b/src/maestro/orchestrator.rs index a862ae8..c49c9fc 100644 --- a/src/maestro/orchestrator.rs +++ b/src/maestro/orchestrator.rs @@ -1,9 +1,10 @@ +use std::collections::BTreeSet; use std::net::{IpAddr, SocketAddr}; use std::sync::Arc; use arc_swap::ArcSwap; use tokio::sync::{RwLock, watch}; -use tracing::{error, info, warn}; +use tracing::{error, info}; use crate::api; use crate::ip_tracker::UserIpTracker; @@ -15,14 +16,15 @@ use crate::startup::{COMPONENT_API_BOOTSTRAP, COMPONENT_NETWORK_PROBE}; use crate::stats::telemetry::TelemetryPolicy; use crate::stats::{QuotaStore, Stats}; use crate::synlimit_control; +use crate::tls_front::cache::TlsFullCertBudget; use crate::transport::UpstreamManager; use crate::transport::middle_proxy::MePool; use crate::web::control::WebRuntimeControl; use crate::web::trace::WebTraceStore; use super::{ - bootstrap, generation, listeners, reload, reload_supervisor, runtime_startup, runtime_tasks, - shutdown, tls_bootstrap, + bootstrap, control_plane, generation, listeners, reload, reload_supervisor, runtime_startup, + runtime_tasks, shutdown, tls_bootstrap, }; // Shared maestro startup and main loop. `drop_after_bind` runs on Unix after listeners are bound @@ -45,10 +47,22 @@ pub(super) async fn run_telemt_core( let quota_store = Arc::new(QuotaStore::default()); let stats = Arc::new(Stats::with_quota_store(quota_store.clone())); + let tls_full_cert_budget = Arc::new(TlsFullCertBudget::new()); + let process_control_plane = control_plane::ProcessControlPlane::new(); let runtime_task_scope = generation::RuntimeTaskScope::new(); stats.apply_telemetry_policy(TelemetryPolicy::from_config(&config.general.telemetry)); let quota_state_path = config.general.quota_state_path.clone(); - crate::quota_state::load_quota_state("a_state_path, stats.as_ref()).await; + let quota_state = crate::quota_state::QuotaStateOwner::new( + quota_state_path, + quota_store.clone(), + ); + let configured_quota_users = config + .access + .users + .keys() + .cloned() + .collect::>(); + quota_state.load(&configured_quota_users).await; let upstream_manager = Arc::new( UpstreamManager::new( @@ -136,15 +150,29 @@ pub(super) async fn run_telemt_core( let listen = match config.server.api.listen.parse::() { Ok(listen) => listen, Err(error) => { - warn!( - error = %error, - listen = %config.server.api.listen, - "Invalid server.api.listen; API is disabled" + let message = format!( + "invalid server.api.listen \"{}\": {}", + config.server.api.listen, error ); - SocketAddr::from(([127, 0, 0, 1], 0)) + startup_tracker + .fail_component(COMPONENT_API_BOOTSTRAP, Some(message.clone())) + .await; + return Err(std::io::Error::new(std::io::ErrorKind::InvalidInput, message).into()); } }; if listen.port() != 0 { + let api_listener = match tokio::net::TcpListener::bind(listen).await { + Ok(listener) => listener, + Err(error) => { + startup_tracker + .fail_component( + COMPONENT_API_BOOTSTRAP, + Some(format!("API listener bind failed on {listen}: {error}")), + ) + .await; + return Err(error.into()); + } + }; let stats_api = stats.clone(); let ip_tracker_api = ip_tracker.clone(); let me_pool_api = api_me_pool.clone(); @@ -152,7 +180,7 @@ pub(super) async fn run_telemt_core( let route_runtime_api = route_runtime.clone(); let proxy_shared_api = shared_state.clone(); let config_path_api = config_path.clone(); - let quota_state_path_api = quota_state_path.clone(); + let quota_state_api = quota_state.clone(); let startup_tracker_api = startup_tracker.clone(); let detected_ips_rx_api = detected_ips_rx.clone(); let reload_control_api = reload_control.clone(); @@ -160,9 +188,11 @@ pub(super) async fn run_telemt_core( let runtime_watch_rx_api = runtime_watch_rx.clone(); let web_trace_api = web_trace.clone(); let web_runtime_rx_api = web_runtime_control.subscribe(); - tokio::spawn(async move { + let api_control_plane = process_control_plane.clone(); + let api_task_control_plane = process_control_plane.clone(); + let api_task = async move { api::serve( - listen, + api_listener, stats_api, ip_tracker_api, me_pool_api, @@ -170,7 +200,7 @@ pub(super) async fn run_telemt_core( proxy_shared_api, upstream_manager_api, config_path_api, - quota_state_path_api, + quota_state_api, detected_ips_rx_api, process_started_at_epoch_secs, startup_tracker_api, @@ -179,13 +209,21 @@ pub(super) async fn run_telemt_core( runtime_watch_rx_api, web_trace_api, web_runtime_rx_api, + api_task_control_plane, ) .await; - }); + }; + if api_control_plane.spawn(api_task).is_err() { + let message = "process control-plane task admission closed during API startup"; + startup_tracker + .fail_component(COMPONENT_API_BOOTSTRAP, Some(message.to_string())) + .await; + return Err(std::io::Error::other(message).into()); + } startup_tracker .complete_component( COMPONENT_API_BOOTSTRAP, - Some(format!("api task spawned on {}", listen)), + Some(format!("API listener bound and supervised on {}", listen)), ) .await; } else { @@ -219,6 +257,7 @@ pub(super) async fn run_telemt_core( upstream_manager.clone(), &startup_tracker, runtime_task_scope.clone(), + tls_full_cert_budget.clone(), tls_bootstrap::TlsBootstrapPolicy::BestEffort, ) .await?; @@ -316,8 +355,10 @@ pub(super) async fn run_telemt_core( &startup_tracker, active_runtime.clone(), web_runtime_control.subscribe(), + tls_full_cert_budget.clone(), + process_control_plane.clone(), ) - .await; + .await?; runtime_watch_tx.send_replace(Some(active_runtime.load_full().watch_state())); active_runtime_tx.send_replace(Some(active_runtime.clone())); @@ -335,6 +376,7 @@ pub(super) async fn run_telemt_core( reload_commands, config_path, quota_store, + tls_full_cert_budget, detected_ips_tx, runtime_log_filter, runtime_watch_tx, @@ -342,12 +384,17 @@ pub(super) async fn run_telemt_core( web_trace, ); - shutdown::spawn_signal_handlers(active_runtime.clone(), process_started_at); + shutdown::spawn_signal_handlers( + active_runtime.clone(), + process_started_at, + process_control_plane.clone(), + ); shutdown::wait_for_shutdown( process_started_at, active_runtime, - quota_state_path, + quota_state, reload_supervisor, + process_control_plane, ) .await; diff --git a/src/maestro/reload_supervisor.rs b/src/maestro/reload_supervisor.rs index 21c2ea5..77dcc96 100644 --- a/src/maestro/reload_supervisor.rs +++ b/src/maestro/reload_supervisor.rs @@ -9,6 +9,7 @@ use tokio_util::sync::CancellationToken; use tracing::{info, warn}; use crate::stats::QuotaStore; +use crate::tls_front::cache::TlsFullCertBudget; use crate::web::trace::WebTraceStore; use super::generation::{RuntimeGeneration, RuntimeWatchState}; @@ -26,6 +27,7 @@ pub(crate) struct ReloadSupervisor { commands: ReloadCommandReceiver, config_path: PathBuf, quota_store: Arc, + tls_full_cert_budget: Arc, detected_ips_tx: watch::Sender<(Option, Option)>, runtime_log_filter: RuntimeLogFilter, runtime_watch_tx: watch::Sender>, @@ -81,12 +83,7 @@ fn revision_gate_action( async fn stop_background_and_middle_end(generation: &RuntimeGeneration) -> bool { generation.stop_background_tasks().await; - let Some(pool) = generation.current_me_pool().await else { - return false; - }; - tokio::time::timeout(Duration::from_secs(2), pool.shutdown_send_close_conn_all()) - .await - .is_err() + !generation.stop_middle_end(Duration::from_secs(5)).await } async fn cleanup_candidate(generation: &RuntimeGeneration) -> bool { @@ -103,6 +100,7 @@ impl ReloadSupervisor { commands: ReloadCommandReceiver, config_path: PathBuf, quota_store: Arc, + tls_full_cert_budget: Arc, detected_ips_tx: watch::Sender<(Option, Option)>, runtime_log_filter: RuntimeLogFilter, runtime_watch_tx: watch::Sender>, @@ -116,6 +114,7 @@ impl ReloadSupervisor { commands, config_path, quota_store, + tls_full_cert_budget, detected_ips_tx, runtime_log_filter, runtime_watch_tx, @@ -171,6 +170,7 @@ impl ReloadSupervisor { &self.config_path, self.quota_store.clone(), self.runtime_log_filter.clone(), + self.tls_full_cert_budget.clone(), ) .await { @@ -207,25 +207,18 @@ impl ReloadSupervisor { prepared, listener_transition, revision_action, - |entries| { - crate::network::dns_overrides::install_entries(entries) - .map_err(|error| error.to_string()) - }, ) .await; } #[cfg(test)] - async fn activate_prepared( + async fn activate_prepared( &self, command: ReloadCommand, old_runtime: Arc, prepared: PreparedRuntime, revision_action: RevisionGateAction, - install_dns: InstallDns, - ) where - InstallDns: FnOnce(&[String]) -> Result<(), String>, - { + ) { let listener_transition = match self .listener_manager .lock() @@ -245,22 +238,18 @@ impl ReloadSupervisor { prepared, listener_transition, revision_action, - install_dns, ) .await; } - async fn activate_prepared_with_transition( + async fn activate_prepared_with_transition( &self, command: ReloadCommand, old_runtime: Arc, prepared: PreparedRuntime, listener_transition: Option, revision_action: RevisionGateAction, - install_dns: InstallDns, - ) where - InstallDns: FnOnce(&[String]) -> Result<(), String>, - { + ) { match revision_action { RevisionGateAction::Proceed => {} RevisionGateAction::Warn(warning) => { @@ -283,18 +272,6 @@ impl ReloadSupervisor { detected_ips, config_watcher_activation, } = prepared; - if let Err(error) = install_dns(&new_runtime.config().network.dns_overrides) { - let message = format!("runtime DNS activation failed: {}", error); - if command.request.failure_policy == ReloadFailurePolicy::Rollback { - old_runtime.resume_accepting_sessions(); - let _ = cleanup_candidate(&new_runtime).await; - self.runtime_log_filter - .apply_reload(&old_runtime.config().general.log_level); - self.control.rolled_back(command.reload_id, message).await; - return; - } - self.control.add_warning(command.reload_id, message).await; - } let pending_listener_transition = if let Some(listener_transition) = listener_transition { match self .listener_manager @@ -367,7 +344,7 @@ impl ReloadSupervisor { if stop_background_and_middle_end(&replaced).await { let warning = format!( - "generation {} Middle-End close broadcast timed out", + "generation {} Middle-End lifecycle shutdown timed out", replaced.id ); warn!(reload_id = command.reload_id, warning = %warning); diff --git a/src/maestro/reload_supervisor_tests.rs b/src/maestro/reload_supervisor_tests.rs index d742760..27b6e4b 100644 --- a/src/maestro/reload_supervisor_tests.rs +++ b/src/maestro/reload_supervisor_tests.rs @@ -53,6 +53,7 @@ async fn fixture(request: ReloadRequest) -> ReloadFixture { commands, config_path: PathBuf::new(), quota_store: Arc::new(QuotaStore::default()), + tls_full_cert_budget: Arc::new(crate::tls_front::cache::TlsFullCertBudget::new()), detected_ips_tx, runtime_log_filter: runtime_log_filter(), runtime_watch_tx, @@ -131,7 +132,6 @@ async fn revision_rollback_keeps_old_generation_and_cleans_candidate() { fixture.old_runtime.clone(), prepared_runtime(fixture.new_runtime), RevisionGateAction::Rollback("revision changed".to_string()), - |_| -> Result<(), String> { panic!("DNS activation must not run on rollback") }, ) .await; @@ -154,44 +154,6 @@ async fn revision_rollback_keeps_old_generation_and_cleans_candidate() { fixture.old_runtime.stop_sessions().await; } -#[tokio::test] -async fn dns_failure_policy_controls_rollback_or_keep_new() { - for policy in [ReloadFailurePolicy::Rollback, ReloadFailurePolicy::KeepNew] { - let fixture = fixture(ReloadRequest { - failure_policy: policy, - ..ReloadRequest::default() - }) - .await; - fixture - .supervisor - .activate_prepared( - fixture.command, - fixture.old_runtime.clone(), - prepared_runtime(fixture.new_runtime.clone()), - RevisionGateAction::Proceed, - |_| Err("invalid DNS entry".to_string()), - ) - .await; - - let status = fixture.control.status(1).await.unwrap(); - match policy { - ReloadFailurePolicy::Rollback => { - assert_eq!(fixture.supervisor.active_runtime.load().id, 1); - assert_eq!(status.state, ReloadPhase::RolledBack); - assert!(fixture.old_runtime.spawn_session(async {})); - fixture.old_runtime.stop_sessions().await; - } - ReloadFailurePolicy::KeepNew => { - assert_eq!(fixture.supervisor.active_runtime.load().id, 2); - assert_eq!(status.state, ReloadPhase::Succeeded); - assert_eq!(status.warnings.len(), 1); - assert!(!fixture.old_runtime.spawn_session(async {})); - fixture.new_runtime.stop_sessions().await; - } - } - } -} - #[tokio::test] async fn drain_publishes_new_generation_before_old_sessions_finish() { let mut fixture = fixture(ReloadRequest { @@ -220,7 +182,6 @@ async fn drain_publishes_new_generation_before_old_sessions_finish() { old_runtime, prepared_runtime(new_runtime), RevisionGateAction::Proceed, - |_| Ok(()), ) .await; }); @@ -273,7 +234,6 @@ async fn drain_timeout_cancels_old_sessions_and_records_one_warning() { old_runtime, prepared_runtime(new_runtime), RevisionGateAction::Proceed, - |_| Ok(()), ) .await; }); @@ -304,6 +264,7 @@ async fn quiesce_joins_idle_supervisor_and_rejects_later_submissions() { commands, PathBuf::new(), Arc::new(QuotaStore::default()), + Arc::new(crate::tls_front::cache::TlsFullCertBudget::new()), detected_ips_tx, runtime_log_filter(), runtime_watch_tx, diff --git a/src/maestro/runtime_build.rs b/src/maestro/runtime_build.rs index 14af642..12f9c83 100644 --- a/src/maestro/runtime_build.rs +++ b/src/maestro/runtime_build.rs @@ -19,6 +19,7 @@ use crate::stats::beobachten::BeobachtenStore; use crate::stats::telemetry::TelemetryPolicy; use crate::stats::{QuotaStore, ReplayChecker, Stats}; use crate::stream::BufferPool; +use crate::tls_front::cache::TlsFullCertBudget; use crate::transport::UpstreamManager; use crate::transport::middle_proxy::MePool; @@ -44,6 +45,7 @@ pub(crate) async fn prepare_runtime( config_path: &Path, quota_store: Arc, runtime_log_filter: RuntimeLogFilter, + tls_full_cert_budget: Arc, ) -> Result { config .validate_web_decoy_listener_separation() @@ -121,6 +123,7 @@ pub(crate) async fn prepare_runtime( upstream_manager.clone(), &startup_tracker, task_scope.clone(), + tls_full_cert_budget, tls_bootstrap::TlsBootstrapPolicy::RequireReady, ) .await diff --git a/src/maestro/runtime_tasks.rs b/src/maestro/runtime_tasks.rs index 826335d..cc2a124 100644 --- a/src/maestro/runtime_tasks.rs +++ b/src/maestro/runtime_tasks.rs @@ -26,6 +26,7 @@ use crate::stats::{ReplayChecker, Stats}; use crate::transport::UpstreamManager; use crate::transport::middle_proxy::{MePool, MeReinitTrigger}; +use super::control_plane::ProcessControlPlane; use super::generation::RuntimeGeneration; use super::generation::RuntimeTaskScope; use super::helpers::write_beobachten_snapshot; @@ -158,6 +159,7 @@ pub(crate) async fn spawn_runtime_tasks( detected_ip_v4, detected_ip_v6, task_scope.cancellation_token(), + Some(upstream_manager.dns_resolver()), config_watcher_activation, ); task_scope.spawn(config_watcher_task); @@ -168,7 +170,6 @@ pub(crate) async fn spawn_runtime_tasks( ) .await; let stats_policy = stats.clone(); - let upstream_policy = upstream_manager.clone(); let mut config_rx_policy = config_rx.clone(); task_scope.spawn(async move { loop { @@ -178,9 +179,6 @@ pub(crate) async fn spawn_runtime_tasks( let cfg = config_rx_policy.borrow_and_update().clone(); stats_policy .apply_telemetry_policy(TelemetryPolicy::from_config(&cfg.general.telemetry)); - if let Err(error) = upstream_policy.update_dns_overrides(&cfg.network.dns_overrides) { - warn!(error = %error, "Failed to update generation DNS overrides"); - } if let Some(pool) = &me_pool_for_policy { pool.update_runtime_transport_policy( cfg.general.me_socks_kdf_policy, @@ -406,7 +404,9 @@ pub(crate) async fn spawn_metrics_if_configured( startup_tracker: &Arc, active_runtime: Arc>, web_runtime_rx: tokio::sync::watch::Receiver, -) { + tls_full_cert_budget: Arc, + control_plane: ProcessControlPlane, +) -> std::io::Result<()> { // metrics_listen takes precedence; fall back to metrics_port for backward compat. let metrics_target: Option<(u16, Option)> = if let Some(ref listen) = config.server.metrics_listen { @@ -414,12 +414,15 @@ pub(crate) async fn spawn_metrics_if_configured( Ok(addr) => Some((addr.port(), Some(listen.clone()))), Err(e) => { startup_tracker - .skip_component( + .fail_component( COMPONENT_METRICS_START, Some(format!("invalid metrics_listen \"{}\": {}", listen, e)), ) .await; - None + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!("invalid metrics_listen \"{}\": {}", listen, e), + )); } } } else { @@ -435,15 +438,30 @@ pub(crate) async fn spawn_metrics_if_configured( Some(format!("spawn metrics endpoint on {}", label)), ) .await; - let active_runtime = active_runtime.clone(); let listen_backlog = config.server.listen_backlog; - tokio::spawn(async move { - metrics::serve(port, listen, listen_backlog, active_runtime, web_runtime_rx).await; - }); + let bound = match metrics::bind(port, listen, listen_backlog) { + Ok(bound) => bound, + Err(error) => { + startup_tracker + .fail_component( + COMPONENT_METRICS_START, + Some(format!("metrics listener bind failed: {error}")), + ) + .await; + return Err(error); + } + }; + metrics::serve( + bound, + active_runtime, + web_runtime_rx, + tls_full_cert_budget, + control_plane, + ); startup_tracker .complete_component( COMPONENT_METRICS_START, - Some("metrics task spawned".to_string()), + Some("metrics listeners bound and supervised".to_string()), ) .await; } else if config.server.metrics_listen.is_none() { @@ -454,6 +472,7 @@ pub(crate) async fn spawn_metrics_if_configured( ) .await; } + Ok(()) } pub(crate) async fn mark_runtime_ready(startup_tracker: &Arc) { diff --git a/src/maestro/shutdown.rs b/src/maestro/shutdown.rs index d061518..5ef3cee 100644 --- a/src/maestro/shutdown.rs +++ b/src/maestro/shutdown.rs @@ -8,7 +8,7 @@ //! //! SIGHUP is handled separately in config/hot_reload.rs for config reload. -use std::path::PathBuf; +use std::collections::BTreeSet; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -19,11 +19,13 @@ use tokio::signal; use tokio::signal::unix::{SignalKind, signal}; use tracing::{info, warn}; +use super::control_plane::ProcessControlPlane; use super::generation::RuntimeGeneration; use super::helpers::{format_uptime, unit_label}; use super::reload_supervisor::ReloadSupervisorHandle; use crate::stats::Stats; use crate::synlimit_control; +use crate::quota_state::QuotaStateOwner; /// Signal that triggered shutdown. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -50,16 +52,18 @@ impl std::fmt::Display for ShutdownSignal { pub(crate) async fn wait_for_shutdown( process_started_at: Instant, active_runtime: Arc>, - quota_state_path: PathBuf, + quota_state: Arc, reload_supervisor: ReloadSupervisorHandle, + process_control_plane: ProcessControlPlane, ) { let signal = wait_for_shutdown_signal().await; perform_shutdown( signal, process_started_at, active_runtime, - quota_state_path, + quota_state, reload_supervisor, + process_control_plane, ) .await; } @@ -89,8 +93,9 @@ async fn perform_shutdown( signal: ShutdownSignal, process_started_at: Instant, active_runtime: Arc>, - quota_state_path: PathBuf, + quota_state: Arc, reload_supervisor: ReloadSupervisorHandle, + process_control_plane: ProcessControlPlane, ) { let shutdown_started_at = Instant::now(); info!(signal = %signal, "Received shutdown signal"); @@ -115,37 +120,41 @@ async fn perform_shutdown( // Graceful ME pool shutdown runtime.stop_sessions().await; runtime.stop_background_tasks().await; - if let Some(pool) = runtime.current_me_pool().await { - match tokio::time::timeout(Duration::from_secs(2), pool.shutdown_send_close_conn_all()) - .await - { - Ok(total) => { - info!( - close_conn_sent = total, - "ME shutdown: RPC_CLOSE_CONN broadcast completed" - ); - } - Err(_) => { - warn!("ME shutdown: RPC_CLOSE_CONN broadcast timed out"); - } - } + if runtime.stop_middle_end(Duration::from_secs(5)).await { + info!("ME shutdown: pool lifecycle completed"); + } else { + warn!("ME shutdown: pool lifecycle deadline expired"); } if let Err(error) = synlimit_control::clear_synlimit_rules_all_backends().await { warn!(error = %error, "Failed to clear SYN limiter rules during shutdown"); } - match crate::quota_state::save_quota_state("a_state_path, stats).await { + if !process_control_plane + .shutdown(Duration::from_secs(5)) + .await + { + warn!("Process control-plane task shutdown deadline expired"); + } + + let configured_quota_users = runtime + .config() + .access + .users + .keys() + .cloned() + .collect::>(); + match quota_state.save(&configured_quota_users).await { Ok(()) => { info!( - path = %quota_state_path.display(), + path = %quota_state.path().display(), "Persisted per-user quota state" ); } Err(error) => { warn!( error = %error, - path = %quota_state_path.display(), + path = %quota_state.path().display(), "Failed to persist per-user quota state" ); } @@ -205,8 +214,9 @@ fn dump_stats(stats: &Stats, process_started_at: Instant) { pub(crate) fn spawn_signal_handlers( active_runtime: Arc>, process_started_at: Instant, + process_control_plane: ProcessControlPlane, ) { - tokio::spawn(async move { + let _ = process_control_plane.spawn(async move { let mut sigusr1 = signal(SignalKind::user_defined1()).expect("Failed to register SIGUSR1 handler"); let mut sigusr2 = @@ -231,6 +241,7 @@ pub(crate) fn spawn_signal_handlers( pub(crate) fn spawn_signal_handlers( _active_runtime: Arc>, _process_started_at: Instant, + _process_control_plane: ProcessControlPlane, ) { // No SIGUSR1/SIGUSR2 on non-Unix } diff --git a/src/maestro/tls_bootstrap.rs b/src/maestro/tls_bootstrap.rs index 94f7751..1a670b6 100644 --- a/src/maestro/tls_bootstrap.rs +++ b/src/maestro/tls_bootstrap.rs @@ -8,6 +8,7 @@ use crate::config::ProxyConfig; use crate::error::{ProxyError, Result}; use crate::startup::{COMPONENT_TLS_FRONT_BOOTSTRAP, StartupTracker}; use crate::tls_front::TlsFrontCache; +use crate::tls_front::cache::TlsFullCertBudget; use crate::tls_front::fetcher::TlsFetchStrategy; use crate::transport::UpstreamManager; @@ -109,6 +110,7 @@ pub(crate) async fn bootstrap_tls_front( upstream_manager: Arc, startup_tracker: &Arc, task_scope: RuntimeTaskScope, + full_cert_budget: Arc, policy: TlsBootstrapPolicy, ) -> Result>> { startup_tracker @@ -128,10 +130,11 @@ pub(crate) async fn bootstrap_tls_front( return Ok(None); } - let cache = Arc::new(TlsFrontCache::new( + let cache = Arc::new(TlsFrontCache::new_with_full_cert_budget( tls_domains, config.censorship.fake_cert_len, &config.censorship.tls_front_dir, + full_cert_budget, )); cache.load_from_disk().await; @@ -301,6 +304,7 @@ mod tests { upstream_manager(&config), &tracker, scope.clone(), + Arc::new(TlsFullCertBudget::new()), TlsBootstrapPolicy::RequireReady, ) .await; @@ -336,6 +340,7 @@ mod tests { upstream_manager(&config), &tracker, scope.clone(), + Arc::new(TlsFullCertBudget::new()), TlsBootstrapPolicy::RequireReady, ) .await @@ -364,6 +369,7 @@ mod tests { upstream_manager(&config), &tracker, scope.clone(), + Arc::new(TlsFullCertBudget::new()), TlsBootstrapPolicy::BestEffort, ) .await diff --git a/src/metrics.rs b/src/metrics.rs index d5e3500..87b18b1 100644 --- a/src/metrics.rs +++ b/src/metrics.rs @@ -17,12 +17,13 @@ use tracing::{debug, info, warn}; use crate::config::ProxyConfig; use crate::ip_tracker::UserIpTracker; +use crate::maestro::control_plane::ProcessControlPlane; use crate::maestro::generation::RuntimeGeneration; use crate::proxy::shared_state::ProxySharedState; use crate::stats::Stats; use crate::stats::beobachten::BeobachtenStore; use crate::tls_front::TlsFrontCache; -use crate::tls_front::cache; +use crate::tls_front::cache::TlsFullCertBudget; use crate::tls_front::fetcher; use crate::transport::{ListenOptions, create_listener}; @@ -36,84 +37,89 @@ const TLS_FRONT_PROFILE_HEALTH_MAX_DOMAINS: usize = 256; const METRICS_MAX_CONTROL_CONNECTIONS: usize = 512; const METRICS_HTTP_CONNECTION_TIMEOUT: Duration = Duration::from_secs(15); -pub async fn serve( +/// Bound process-owned metrics listeners ready for supervised serving. +pub(crate) struct BoundMetricsListeners { + listeners: Vec<(TcpListener, SocketAddr)>, +} + +/// Binds every configured metrics socket before process readiness is published. +pub(crate) fn bind( port: u16, listen: Option, listen_backlog: u32, - active_runtime: Arc>, - web_runtime_rx: tokio::sync::watch::Receiver, -) { - // If `metrics_listen` is set, bind on that single address only. +) -> std::io::Result { if let Some(ref listen_addr) = listen { - let addr: SocketAddr = match listen_addr.parse() { - Ok(a) => a, - Err(e) => { - warn!(error = %e, "Invalid metrics_listen address: {}", listen_addr); - return; - } - }; + let addr: SocketAddr = listen_addr.parse().map_err(|error| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!("invalid metrics_listen address {listen_addr}: {error}"), + ) + })?; // Match `server.api.listen`: `[::]:port` is a dual-stack wildcard // on Linux when `net.ipv6.bindv6only=0`. let ipv6_only = addr.is_ipv6() && !addr.ip().is_unspecified(); - match bind_metrics_listener(addr, ipv6_only, listen_backlog) { - Ok(listener) => { - info!("Metrics endpoint: http://{}/metrics and /beobachten", addr); - serve_listener(listener, active_runtime, web_runtime_rx).await; - } - Err(e) => { - warn!(error = %e, "Failed to bind metrics on {}", addr); - } - } - return; + let listener = bind_metrics_listener(addr, ipv6_only, listen_backlog)?; + return Ok(BoundMetricsListeners { + listeners: vec![(listener, addr)], + }); } - // Fallback: keep metrics local unless an explicit metrics_listen is configured. - let mut listener_v4 = None; - let mut listener_v6 = None; + let mut listeners = Vec::with_capacity(2); + let mut last_error = None; let addr_v4 = SocketAddr::from(([127, 0, 0, 1], port)); match bind_metrics_listener(addr_v4, false, listen_backlog) { - Ok(listener) => { - info!( - "Metrics endpoint: http://{}/metrics and /beobachten", - addr_v4 - ); - listener_v4 = Some(listener); - } + Ok(listener) => listeners.push((listener, addr_v4)), Err(e) => { warn!(error = %e, "Failed to bind metrics on {}", addr_v4); + last_error = Some(e); } } let addr_v6 = SocketAddr::from(([0, 0, 0, 0, 0, 0, 0, 1], port)); match bind_metrics_listener(addr_v6, true, listen_backlog) { - Ok(listener) => { - info!( - "Metrics endpoint: http://[::1]:{}/metrics and /beobachten", - port - ); - listener_v6 = Some(listener); - } + Ok(listener) => listeners.push((listener, addr_v6)), Err(e) => { warn!(error = %e, "Failed to bind metrics on {}", addr_v6); + last_error = Some(e); } } - match (listener_v4, listener_v6) { - (None, None) => { - warn!("Metrics listener is unavailable on both IPv4 and IPv6"); - } - (Some(listener), None) | (None, Some(listener)) => { - serve_listener(listener, active_runtime, web_runtime_rx).await; - } - (Some(listener4), Some(listener6)) => { - let active_runtime_v6 = active_runtime.clone(); - let web_runtime_rx_v6 = web_runtime_rx.clone(); - tokio::spawn(async move { - serve_listener(listener6, active_runtime_v6, web_runtime_rx_v6).await; - }); - serve_listener(listener4, active_runtime, web_runtime_rx).await; - } + if listeners.is_empty() { + return Err(last_error.unwrap_or_else(|| { + std::io::Error::new( + std::io::ErrorKind::AddrNotAvailable, + "metrics listener is unavailable on both IPv4 and IPv6", + ) + })); + } + Ok(BoundMetricsListeners { listeners }) +} + +/// Starts supervised accept loops for previously bound metrics sockets. +pub(crate) fn serve( + bound: BoundMetricsListeners, + active_runtime: Arc>, + web_runtime_rx: tokio::sync::watch::Receiver, + tls_full_cert_budget: Arc, + control_plane: ProcessControlPlane, +) { + for (listener, addr) in bound.listeners { + info!("Metrics endpoint: http://{}/metrics and /beobachten", addr); + let active_runtime = active_runtime.clone(); + let web_runtime_rx = web_runtime_rx.clone(); + let tls_full_cert_budget = Arc::clone(&tls_full_cert_budget); + let listener_scope = control_plane.clone(); + let _ = control_plane.spawn(async move { + serve_listener( + listener, + active_runtime, + web_runtime_rx, + tls_full_cert_budget, + listener_scope, + ) + .await; + }); } } @@ -136,6 +142,8 @@ async fn serve_listener( listener: TcpListener, active_runtime: Arc>, web_runtime_rx: tokio::sync::watch::Receiver, + tls_full_cert_budget: Arc, + control_plane: ProcessControlPlane, ) { let connection_permits = Arc::new(Semaphore::new(METRICS_MAX_CONTROL_CONNECTIONS)); @@ -175,12 +183,22 @@ async fn serve_listener( let active_runtime = active_runtime.clone(); let web_runtime_rx = web_runtime_rx.clone(); - tokio::spawn(async move { + let tls_full_cert_budget = Arc::clone(&tls_full_cert_budget); + let _ = control_plane.spawn(async move { let _connection_permit = connection_permit; let svc = service_fn(move |req| { let runtime = active_runtime.load_full(); let web_publication = web_runtime_rx.borrow().clone(); - async move { handle(req, &runtime, &web_publication).await } + let tls_full_cert_budget = Arc::clone(&tls_full_cert_budget); + async move { + handle( + req, + &runtime, + &web_publication, + tls_full_cert_budget.as_ref(), + ) + .await + } }); match timeout( METRICS_HTTP_CONNECTION_TIMEOUT, @@ -208,6 +226,7 @@ async fn handle( req: Request, runtime: &RuntimeGeneration, web_publication: &crate::web::control::WebRuntimePublication, + tls_full_cert_budget: &TlsFullCertBudget, ) -> Result>, Infallible> { let stats = &runtime.stats; let beobachten = &runtime.beobachten; @@ -223,6 +242,7 @@ async fn handle( &config, ip_tracker, tls_cache, + tls_full_cert_budget, web_publication, ) .await; @@ -437,6 +457,7 @@ async fn render_metrics( config: &ProxyConfig, ip_tracker: &UserIpTracker, tls_cache: Option<&TlsFrontCache>, + tls_full_cert_budget: &TlsFullCertBudget, web_publication: &crate::web::control::WebRuntimePublication, ) -> String { use std::fmt::Write; @@ -644,17 +665,17 @@ async fn render_metrics( ); let _ = writeln!( out, - "# HELP telemt_tls_front_full_cert_budget_ips Current IP entries tracked by TLS full-cert budget" + "# HELP telemt_tls_front_full_cert_budget_entries Current domain and IP entries tracked by the process-owned TLS full-cert budget" ); - let _ = writeln!(out, "# TYPE telemt_tls_front_full_cert_budget_ips gauge"); + let _ = writeln!(out, "# TYPE telemt_tls_front_full_cert_budget_entries gauge"); let _ = writeln!( out, - "telemt_tls_front_full_cert_budget_ips {}", - cache::full_cert_sent_ips_for_metrics() + "telemt_tls_front_full_cert_budget_entries {}", + tls_full_cert_budget.entries_for_metrics() ); let _ = writeln!( out, - "# HELP telemt_tls_front_full_cert_budget_cap_drops_total New IPs denied full-cert budget tracking because the cap was reached" + "# HELP telemt_tls_front_full_cert_budget_cap_drops_total New domain and IP entries denied full-cert budget tracking because a bound was reached" ); let _ = writeln!( out, @@ -663,7 +684,7 @@ async fn render_metrics( let _ = writeln!( out, "telemt_tls_front_full_cert_budget_cap_drops_total {}", - cache::full_cert_sent_cap_drops_for_metrics() + tls_full_cert_budget.cap_drops_for_metrics() ); render_tls_front_profile_health(&mut out, config, tls_cache).await; @@ -1559,6 +1580,11 @@ async fn render_metrics( error_code, count ); } + let _ = writeln!( + out, + "telemt_me_handshake_error_code_total{{error_code=\"overflow\"}} {}", + stats.get_me_handshake_error_code_overflow_total() + ); } let _ = writeln!( @@ -3955,6 +3981,7 @@ mod tests { &config, &tracker, None, + &TlsFullCertBudget::new(), &test_web_publication(), ) .await; @@ -4091,6 +4118,7 @@ mod tests { &config, &tracker, Some(&cache), + &TlsFullCertBudget::new(), &test_web_publication(), ) .await; @@ -4135,6 +4163,43 @@ mod tests { ); } + #[tokio::test] + async fn process_tls_budget_metrics_survive_a_generation_without_tls_cache() { + let stats = Stats::new(); + let shared_state = ProxySharedState::new(); + let tracker = UserIpTracker::new(); + let config = ProxyConfig::default(); + let budget = Arc::new(TlsFullCertBudget::new()); + let cache = TlsFrontCache::new_with_full_cert_budget( + &["example.com".to_string()], + 1024, + "tlsfront-test-cache", + Arc::clone(&budget), + ); + assert!( + cache + .take_full_cert_budget_for_ip( + "example.com", + "127.0.0.1".parse().unwrap(), + Duration::from_secs(60), + ) + .await + ); + + let output = render_metrics( + &stats, + &shared_state, + &config, + &tracker, + None, + budget.as_ref(), + &test_web_publication(), + ) + .await; + + assert!(output.contains("telemt_tls_front_full_cert_budget_entries 1")); + } + #[tokio::test] async fn test_render_empty_stats() { let stats = Stats::new(); @@ -4147,6 +4212,7 @@ mod tests { &config, &tracker, None, + &TlsFullCertBudget::new(), &test_web_publication(), ) .await; @@ -4179,6 +4245,7 @@ mod tests { &config, &tracker, None, + &TlsFullCertBudget::new(), &test_web_publication(), ) .await; @@ -4199,6 +4266,7 @@ mod tests { &config, &tracker, None, + &TlsFullCertBudget::new(), &test_web_publication(), ) .await; @@ -4246,7 +4314,7 @@ mod tests { assert!(output.contains("# TYPE telemt_ip_tracker_cap_rejects_total counter")); assert!(output.contains("# TYPE telemt_tls_fetch_profile_cache_entries gauge")); assert!(output.contains("# TYPE telemt_tls_fetch_profile_cache_cap_drops_total counter")); - assert!(output.contains("# TYPE telemt_tls_front_full_cert_budget_ips gauge")); + assert!(output.contains("# TYPE telemt_tls_front_full_cert_budget_entries gauge")); assert!( output.contains("# TYPE telemt_tls_front_full_cert_budget_cap_drops_total counter") ); @@ -4271,12 +4339,15 @@ mod tests { config.general.beobachten_minutes = 10; let runtime = crate::maestro::generation::test_runtime_generation(1, config); let web_publication = test_web_publication(); + let tls_full_cert_budget = TlsFullCertBudget::new(); runtime.stats.increment_connects_all(); runtime.stats.increment_connects_all(); runtime.stats.increment_connects_all(); let req = Request::builder().uri("/metrics").body(()).unwrap(); - let resp = handle(req, &runtime, &web_publication).await.unwrap(); + let resp = handle(req, &runtime, &web_publication, &tls_full_cert_budget) + .await + .unwrap(); assert_eq!(resp.status(), StatusCode::OK); let body = resp.into_body().collect().await.unwrap().to_bytes(); assert!( @@ -4299,7 +4370,14 @@ mod tests { Duration::from_secs(600), ); let req_beob = Request::builder().uri("/beobachten").body(()).unwrap(); - let resp_beob = handle(req_beob, &runtime, &web_publication).await.unwrap(); + let resp_beob = handle( + req_beob, + &runtime, + &web_publication, + &tls_full_cert_budget, + ) + .await + .unwrap(); assert_eq!(resp_beob.status(), StatusCode::OK); let body_beob = resp_beob.into_body().collect().await.unwrap().to_bytes(); let beob_text = std::str::from_utf8(body_beob.as_ref()).unwrap(); @@ -4307,7 +4385,9 @@ mod tests { assert!(beob_text.contains("203.0.113.10-1")); let req404 = Request::builder().uri("/other").body(()).unwrap(); - let resp404 = handle(req404, &runtime, &web_publication).await.unwrap(); + let resp404 = handle(req404, &runtime, &web_publication, &tls_full_cert_budget) + .await + .unwrap(); assert_eq!(resp404.status(), StatusCode::NOT_FOUND); } } diff --git a/src/network/dns_overrides.rs b/src/network/dns_overrides.rs index 710fa26..a262a87 100644 --- a/src/network/dns_overrides.rs +++ b/src/network/dns_overrides.rs @@ -2,23 +2,26 @@ use std::collections::HashMap; use std::net::{IpAddr, Ipv6Addr, SocketAddr}; -use std::sync::{OnceLock, RwLock}; +use std::sync::Arc; + +use arc_swap::ArcSwap; use crate::error::{ProxyError, Result}; type OverrideMap = HashMap<(String, u16), IpAddr>; +const DNS_OVERRIDE_MAX_ENTRIES: usize = 4096; /// Immutable DNS override snapshot owned by one runtime generation. #[derive(Debug, Clone, Default)] pub struct DnsOverrides { - entries: std::sync::Arc, + entries: Arc, } impl DnsOverrides { /// Parses a validated generation-local override snapshot. pub fn from_entries(entries: &[String]) -> Result { Ok(Self { - entries: std::sync::Arc::new(parse_entries(entries)?), + entries: Arc::new(parse_entries(entries)?), }) } @@ -35,10 +38,31 @@ impl DnsOverrides { } } -static DNS_OVERRIDES: OnceLock> = OnceLock::new(); +/// Atomically published DNS override snapshot owned by one runtime generation. +#[derive(Debug, Default)] +pub struct GenerationDnsResolver { + snapshot: ArcSwap, +} -fn overrides_store() -> &'static RwLock { - DNS_OVERRIDES.get_or_init(|| RwLock::new(HashMap::new())) +impl GenerationDnsResolver { + /// Creates one resolver from a validated immutable entry set. + pub fn from_entries(entries: &[String]) -> Result { + Ok(Self { + snapshot: ArcSwap::from_pointee(DnsOverrides::from_entries(entries)?), + }) + } + + /// Validates and atomically publishes a new generation-local snapshot. + pub fn apply_entries(&self, entries: &[String]) -> Result<()> { + let snapshot = DnsOverrides::from_entries(entries)?; + self.snapshot.store(Arc::new(snapshot)); + Ok(()) + } + + /// Resolves one configured override without consulting system DNS. + pub fn resolve_socket_addr(&self, host: &str, port: u16) -> Option { + self.snapshot.load().resolve_socket_addr(host, port) + } } fn parse_ip_spec(ip_spec: &str) -> Result { @@ -111,6 +135,11 @@ fn parse_entry(entry: &str) -> Result<((String, u16), IpAddr)> { } fn parse_entries(entries: &[String]) -> Result { + if entries.len() > DNS_OVERRIDE_MAX_ENTRIES { + return Err(ProxyError::Config(format!( + "network.dns_overrides exceeds maximum entry count {DNS_OVERRIDE_MAX_ENTRIES}" + ))); + } let mut parsed = HashMap::new(); for entry in entries { let (key, ip) = parse_entry(entry)?; @@ -125,30 +154,6 @@ pub fn validate_entries(entries: &[String]) -> Result<()> { Ok(()) } -/// Replace runtime DNS overrides with a new validated snapshot. -pub fn install_entries(entries: &[String]) -> Result<()> { - let parsed = parse_entries(entries)?; - let mut guard = overrides_store().write().map_err(|_| { - ProxyError::Config("network.dns_overrides runtime lock is poisoned".to_string()) - })?; - *guard = parsed; - Ok(()) -} - -/// Resolve a hostname override for `(host, port)` if present. -pub fn resolve(host: &str, port: u16) -> Option { - let key = (host.to_ascii_lowercase(), port); - overrides_store() - .read() - .ok() - .and_then(|guard| guard.get(&key).copied()) -} - -/// Resolve a hostname override and construct a socket address when present. -pub fn resolve_socket_addr(host: &str, port: u16) -> Option { - resolve(host, port).map(|ip| SocketAddr::new(ip, port)) -} - /// Parse a runtime endpoint in `host:port` format. /// /// Supports: @@ -199,12 +204,14 @@ mod tests { } #[test] - fn install_and_resolve_are_case_insensitive_for_host() { + fn generation_resolver_updates_are_case_insensitive_for_host() { let entries = vec!["MyPetrovich.ru:8443:127.0.0.1".to_string()]; - install_entries(&entries).unwrap(); + let resolver = GenerationDnsResolver::from_entries(&entries).unwrap(); - let resolved = resolve("mypetrovich.ru", 8443); - assert_eq!(resolved, Some("127.0.0.1".parse().unwrap())); + assert_eq!( + resolver.resolve_socket_addr("mypetrovich.ru", 8443), + Some("127.0.0.1:8443".parse().unwrap()) + ); } #[test] diff --git a/src/network/probe.rs b/src/network/probe.rs index 5c3cbb8..bf18029 100644 --- a/src/network/probe.rs +++ b/src/network/probe.rs @@ -3,6 +3,7 @@ use std::collections::HashMap; use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, UdpSocket}; +use std::sync::Arc; use std::time::Duration; use tokio::task::JoinSet; @@ -12,7 +13,8 @@ use tracing::{debug, info, warn}; use crate::config::{NetworkConfig, UpstreamConfig, UpstreamType}; use crate::error::Result; use crate::network::stun::{ - DualStunResult, IpFamily, StunProbeResult, stun_probe_family_with_bind_and_tcp_fallback, + DualStunResult, IpFamily, StunProbeResult, + stun_probe_family_with_bind_tcp_fallback_and_resolver, }; use crate::transport::UpstreamManager; @@ -67,6 +69,11 @@ pub async fn run_probe( stun_nat_probe_concurrency: usize, ) -> Result { let mut probe = NetworkProbe::default(); + let dns_resolver = Arc::new( + crate::network::dns_overrides::GenerationDnsResolver::from_entries( + &config.dns_overrides, + )?, + ); let servers = collect_stun_servers(config); let mut detected_ipv4 = detect_local_ip_v4(); let mut detected_ipv6 = detect_local_ip_v6(); @@ -88,6 +95,7 @@ pub async fn run_probe( None, None, config.stun_tcp_fallback, + Arc::clone(&dns_resolver), ) .await } @@ -171,6 +179,7 @@ pub async fn run_probe( bind_v4, bind_v6, config.stun_tcp_fallback, + Arc::clone(&dns_resolver), ) .await; if let Some(reflected) = direct_stun_res.v4.map(|r| r.reflected_addr) { @@ -286,6 +295,7 @@ async fn probe_stun_servers_parallel( bind_v4: Option, bind_v6: Option, tcp_fallback: bool, + dns_resolver: Arc, ) -> DualStunResult { let mut join_set = JoinSet::new(); let mut next_idx = 0usize; @@ -295,6 +305,7 @@ async fn probe_stun_servers_parallel( while next_idx < servers.len() || !join_set.is_empty() { while next_idx < servers.len() && join_set.len() < concurrency { let stun_addr = servers[next_idx].clone(); + let dns_resolver = Arc::clone(&dns_resolver); next_idx += 1; join_set.spawn(async move { let batch_timeout = if tcp_fallback { @@ -303,18 +314,20 @@ async fn probe_stun_servers_parallel( STUN_BATCH_TIMEOUT }; let res = timeout(batch_timeout, async { - let v4 = stun_probe_family_with_bind_and_tcp_fallback( + let v4 = stun_probe_family_with_bind_tcp_fallback_and_resolver( &stun_addr, IpFamily::V4, bind_v4, tcp_fallback, + Some(dns_resolver.as_ref()), ) .await?; - let v6 = stun_probe_family_with_bind_and_tcp_fallback( + let v6 = stun_probe_family_with_bind_tcp_fallback_and_resolver( &stun_addr, IpFamily::V6, bind_v6, tcp_fallback, + Some(dns_resolver.as_ref()), ) .await?; Ok::(DualStunResult { v4, v6 }) diff --git a/src/network/stun.rs b/src/network/stun.rs index ca4a8cb..bda8eae 100644 --- a/src/network/stun.rs +++ b/src/network/stun.rs @@ -10,7 +10,7 @@ use tokio::time::{Duration, sleep, timeout}; use crate::crypto::SecureRandom; use crate::error::{ProxyError, Result}; -use crate::network::dns_overrides::{resolve, split_host_port}; +use crate::network::dns_overrides::{GenerationDnsResolver, split_host_port}; fn stun_rng() -> &'static SecureRandom { static STUN_RNG: OnceLock = OnceLock::new(); @@ -80,13 +80,32 @@ pub async fn stun_probe_family_with_bind_and_tcp_fallback( family: IpFamily, bind_ip: Option, tcp_fallback: bool, +) -> Result> { + stun_probe_family_with_bind_tcp_fallback_and_resolver( + stun_addr, + family, + bind_ip, + tcp_fallback, + None, + ) + .await +} + +/// Probes one STUN family with an optional generation-owned DNS resolver. +pub async fn stun_probe_family_with_bind_tcp_fallback_and_resolver( + stun_addr: &str, + family: IpFamily, + bind_ip: Option, + tcp_fallback: bool, + dns_resolver: Option<&GenerationDnsResolver>, ) -> Result> { let udp_attempts = if tcp_fallback { 1 } else { 3 }; - let udp_result = stun_probe_family_udp(stun_addr, family, bind_ip, udp_attempts).await?; + let udp_result = + stun_probe_family_udp(stun_addr, family, bind_ip, udp_attempts, dns_resolver).await?; if udp_result.is_some() || !tcp_fallback { return Ok(udp_result); } - stun_probe_family_tcp(stun_addr, family, bind_ip).await + stun_probe_family_tcp(stun_addr, family, bind_ip, dns_resolver).await } async fn stun_probe_family_udp( @@ -94,6 +113,7 @@ async fn stun_probe_family_udp( family: IpFamily, bind_ip: Option, max_attempts: u8, + dns_resolver: Option<&GenerationDnsResolver>, ) -> Result> { let bind_addr = match (family, bind_ip) { (IpFamily::V4, Some(IpAddr::V4(ip))) => SocketAddr::new(IpAddr::V4(ip), 0), @@ -111,7 +131,7 @@ async fn stun_probe_family_udp( Err(e) => return Err(ProxyError::Proxy(format!("STUN bind failed: {e}"))), }; - let target_addr = resolve_stun_addr(stun_addr, family).await?; + let target_addr = resolve_stun_addr(stun_addr, family, dns_resolver).await?; if let Some(addr) = target_addr { match socket.connect(addr).await { Ok(()) => {} @@ -182,8 +202,9 @@ async fn stun_probe_family_tcp( stun_addr: &str, family: IpFamily, bind_ip: Option, + dns_resolver: Option<&GenerationDnsResolver>, ) -> Result> { - let target_addr = match resolve_stun_addr(stun_addr, family).await? { + let target_addr = match resolve_stun_addr(stun_addr, family, dns_resolver).await? { Some(addr) => addr, None => return Ok(None), }; @@ -360,7 +381,11 @@ fn parse_reflected_addr(buf: &[u8], txid: &[u8]) -> Option { None } -async fn resolve_stun_addr(stun_addr: &str, family: IpFamily) -> Result> { +async fn resolve_stun_addr( + stun_addr: &str, + family: IpFamily, + dns_resolver: Option<&GenerationDnsResolver>, +) -> Result> { if let Ok(addr) = stun_addr.parse::() { return Ok(match (addr.is_ipv4(), family) { (true, IpFamily::V4) | (false, IpFamily::V6) => Some(addr), @@ -369,9 +394,9 @@ async fn resolve_stun_addr(stun_addr: &str, family: IpFamily) -> Result Some(addr), _ => None, diff --git a/src/proxy/adaptive_buffers.rs b/src/proxy/adaptive_buffers.rs index 9870522..6830838 100644 --- a/src/proxy/adaptive_buffers.rs +++ b/src/proxy/adaptive_buffers.rs @@ -2,7 +2,7 @@ use dashmap::DashMap; use std::cmp::max; -use std::sync::OnceLock; +use std::sync::{Mutex, OnceLock}; use std::time::{Duration, Instant}; const EMA_ALPHA: f64 = 0.2; @@ -294,6 +294,11 @@ fn profiles() -> &'static DashMap { USER_PROFILES.get_or_init(DashMap::new) } +fn profile_insert_guard() -> &'static Mutex<()> { + static PROFILE_INSERT_GUARD: OnceLock> = OnceLock::new(); + PROFILE_INSERT_GUARD.get_or_init(|| Mutex::new(())) +} + /// Returns a fresh user's recent successful Direct tier, or `Base` when stale. #[allow(dead_code)] pub fn seed_tier_for_user(user: &str) -> AdaptiveTier { @@ -320,29 +325,58 @@ pub fn record_user_tier(user: &str, tier: AdaptiveTier) { if user.len() > MAX_USER_KEY_BYTES { return; } + record_user_tier_with_cap( + profiles(), + profile_insert_guard(), + user, + tier, + MAX_USER_PROFILES_ENTRIES, + ); +} + +fn record_user_tier_with_cap( + profiles: &DashMap, + insert_guard: &Mutex<()>, + user: &str, + tier: AdaptiveTier, + max_entries: usize, +) { let now = Instant::now(); - let mut was_vacant = false; - match profiles().entry(user.to_string()) { - dashmap::mapref::entry::Entry::Occupied(mut entry) => { - let existing = *entry.get(); - let effective = if now.saturating_duration_since(existing.seen_at) > PROFILE_TTL { - tier - } else { - max(existing.tier, tier) - }; - entry.insert(UserAdaptiveProfile { - tier: effective, - seen_at: now, - }); - } - dashmap::mapref::entry::Entry::Vacant(slot) => { - slot.insert(UserAdaptiveProfile { tier, seen_at: now }); - was_vacant = true; - } + if let Some(mut entry) = profiles.get_mut(user) { + let effective = if now.saturating_duration_since(entry.seen_at) > PROFILE_TTL { + tier + } else { + max(entry.tier, tier) + }; + *entry = UserAdaptiveProfile { + tier: effective, + seen_at: now, + }; + return; } - if was_vacant && profiles().len() > MAX_USER_PROFILES_ENTRIES { - profiles().retain(|_, v| now.saturating_duration_since(v.seen_at) <= PROFILE_TTL); + + let _guard = insert_guard + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if let Some(mut entry) = profiles.get_mut(user) { + let effective = if now.saturating_duration_since(entry.seen_at) > PROFILE_TTL { + tier + } else { + max(entry.tier, tier) + }; + *entry = UserAdaptiveProfile { + tier: effective, + seen_at: now, + }; + return; } + if profiles.len() >= max_entries { + profiles.retain(|_, value| now.saturating_duration_since(value.seen_at) <= PROFILE_TTL); + } + if profiles.len() >= max_entries { + return; + } + profiles.insert(user.to_string(), UserAdaptiveProfile { tier, seen_at: now }); } #[cfg(test)] @@ -438,6 +472,17 @@ mod adaptive_direct_budget_policy_tests; mod tests { use super::*; + #[test] + fn fresh_profile_cardinality_is_hard_bounded() { + let profiles = DashMap::new(); + let insert_guard = Mutex::new(()); + for index in 0..64 { + let user = format!("user-{index}"); + record_user_tier_with_cap(&profiles, &insert_guard, &user, AdaptiveTier::Base, 16); + } + assert_eq!(profiles.len(), 16); + } + fn sample( c2s_bytes: u64, s2c_requested_bytes: u64, diff --git a/src/proxy/client.rs b/src/proxy/client.rs index c5108ab..2e30370 100644 --- a/src/proxy/client.rs +++ b/src/proxy/client.rs @@ -50,7 +50,6 @@ use crate::proxy::handshake::{ }; #[cfg(test)] use crate::proxy::handshake::{handle_mtproto_handshake, handle_tls_handshake}; -use crate::proxy::masking::handle_bad_client_with_shared; #[cfg(test)] use crate::proxy::route_mode::RelayRouteMode; use crate::proxy::route_mode::RouteRuntimeController; @@ -247,6 +246,7 @@ fn masking_outcome( peer: SocketAddr, local_addr: SocketAddr, config: Arc, + upstream_manager: Arc, beobachten: Arc, shared: Arc, ) -> HandshakeOutcome @@ -264,7 +264,7 @@ where ) .await; - handle_bad_client_with_shared( + crate::proxy::masking::handle_bad_client_with_shared_resolver( reader, writer, &initial_data, @@ -273,6 +273,7 @@ where &config, &beobachten, shared.as_ref(), + Some(upstream_manager.as_ref()), ) .await; Ok(()) @@ -657,6 +658,7 @@ where real_peer, local_addr, config.clone(), + upstream_manager.clone(), beobachten.clone(), shared.clone(), )); @@ -679,6 +681,7 @@ where real_peer, local_addr, config.clone(), + upstream_manager.clone(), beobachten.clone(), shared.clone(), )); @@ -698,6 +701,7 @@ where real_peer, local_addr, config.clone(), + upstream_manager.clone(), beobachten.clone(), shared.clone(), )); @@ -729,6 +733,7 @@ where real_peer, local_addr, config.clone(), + upstream_manager.clone(), beobachten.clone(), shared.clone(), )); @@ -787,6 +792,7 @@ where real_peer, local_addr, config.clone(), + upstream_manager.clone(), beobachten.clone(), shared.clone(), )); @@ -817,6 +823,7 @@ where real_peer, local_addr, config.clone(), + upstream_manager.clone(), beobachten.clone(), shared.clone(), )); @@ -843,6 +850,7 @@ where real_peer, local_addr, config.clone(), + upstream_manager.clone(), beobachten.clone(), shared.clone(), )); @@ -1279,6 +1287,7 @@ impl RunningClientHandler { peer, local_addr, self.config.clone(), + self.upstream_manager.clone(), self.beobachten.clone(), self.shared.clone(), )); @@ -1301,6 +1310,7 @@ impl RunningClientHandler { peer, local_addr, self.config.clone(), + self.upstream_manager.clone(), self.beobachten.clone(), self.shared.clone(), )); @@ -1321,6 +1331,7 @@ impl RunningClientHandler { peer, local_addr, self.config.clone(), + self.upstream_manager.clone(), self.beobachten.clone(), self.shared.clone(), )); @@ -1377,6 +1388,7 @@ impl RunningClientHandler { peer, local_addr, config.clone(), + self.upstream_manager.clone(), self.beobachten.clone(), self.shared.clone(), )); @@ -1445,6 +1457,7 @@ impl RunningClientHandler { peer, local_addr, config.clone(), + self.upstream_manager.clone(), self.beobachten.clone(), self.shared.clone(), )); @@ -1493,6 +1506,7 @@ impl RunningClientHandler { peer, local_addr, self.config.clone(), + self.upstream_manager.clone(), self.beobachten.clone(), self.shared.clone(), )); @@ -1532,6 +1546,7 @@ impl RunningClientHandler { peer, local_addr, config.clone(), + self.upstream_manager.clone(), self.beobachten.clone(), self.shared.clone(), )); diff --git a/src/proxy/handshake/auth_probe.rs b/src/proxy/handshake/auth_probe.rs index 91f4e94..73e6aac 100644 --- a/src/proxy/handshake/auth_probe.rs +++ b/src/proxy/handshake/auth_probe.rs @@ -98,7 +98,7 @@ pub(super) fn auth_probe_is_throttled_in( }; if auth_probe_state_expired(&entry, now) { drop(entry); - state.remove(&peer_ip); + state.remove_if(&peer_ip, |_, current| auth_probe_state_expired(current, now)); return false; } now < entry.blocked_until @@ -116,7 +116,7 @@ pub(super) fn auth_probe_saturation_grace_exhausted_in( }; if auth_probe_state_expired(&entry, now) { drop(entry); - state.remove(&peer_ip); + state.remove_if(&peer_ip, |_, current| auth_probe_state_expired(current, now)); return false; } @@ -264,11 +264,19 @@ pub(super) fn auth_probe_record_failure_with_state_in( } } - let Some((evict_key, _, _)) = eviction_candidate else { + let Some((evict_key, evict_fail_streak, evict_last_seen)) = eviction_candidate else { return; }; - state.remove(&evict_key); - break; + if state + .remove_if(&evict_key, |_, current| { + current.fail_streak == evict_fail_streak + && current.last_seen == evict_last_seen + }) + .is_some() + { + break; + } + continue; } let mut stale_keys = Vec::new(); @@ -334,18 +342,22 @@ pub(super) fn auth_probe_record_failure_with_state_in( } for stale_key in stale_keys { - state.remove(&stale_key); + state.remove_if(&stale_key, |_, current| { + auth_probe_state_expired(current, now) + }); } if state.len() < AUTH_PROBE_TRACK_MAX_ENTRIES { break; } - let Some((evict_key, _, _)) = eviction_candidate else { + let Some((evict_key, evict_fail_streak, evict_last_seen)) = eviction_candidate else { auth_probe_note_saturation_in(shared, now); return; }; - state.remove(&evict_key); + state.remove_if(&evict_key, |_, current| { + current.fail_streak == evict_fail_streak && current.last_seen == evict_last_seen + }); auth_probe_note_saturation_in(shared, now); } } diff --git a/src/proxy/handshake/tls_handshake.rs b/src/proxy/handshake/tls_handshake.rs index 0e4e2bb..ccc877a 100644 --- a/src/proxy/handshake/tls_handshake.rs +++ b/src/proxy/handshake/tls_handshake.rs @@ -266,11 +266,11 @@ where return HandshakeResult::BadClient { reader, writer }; } + let selected_tls_domain = + matched_tls_domain.unwrap_or(config.censorship.tls_domain.as_str()); let cached_entry = if config.censorship.tls_emulation { if let Some(cache) = tls_cache.as_ref() { - let selected_domain = - matched_tls_domain.unwrap_or(config.censorship.tls_domain.as_str()); - let cached_entry = cache.get(selected_domain).await; + let cached_entry = cache.get(selected_tls_domain).await; Some(cached_entry) } else { None @@ -322,6 +322,7 @@ where if let Some(cache) = tls_cache.as_ref() { cache .take_full_cert_budget_for_ip( + selected_tls_domain, peer.ip(), Duration::from_secs(config.censorship.tls_full_cert_ttl_secs), ) diff --git a/src/proxy/masking.rs b/src/proxy/masking.rs index 7e99deb..cb60d21 100644 --- a/src/proxy/masking.rs +++ b/src/proxy/masking.rs @@ -1,7 +1,6 @@ //! Masking - forward unrecognized traffic to mask host use crate::config::ProxyConfig; -use crate::network::dns_overrides::resolve_socket_addr; use crate::protocol::tls; use crate::proxy::shared_state::ProxySharedState; use crate::stats::beobachten::BeobachtenStore; @@ -41,6 +40,7 @@ const MASK_RELAY_IDLE_TIMEOUT: Duration = Duration::from_millis(100); const MASK_BUFFER_SIZE: usize = 8192; const MASK_BUFFER_GROW_AFTER_BYTES: usize = 256 * 1024; const MASK_BUFFER_MAX_SIZE: usize = 64 * 1024; +const MASK_DNS_RESULT_MAX_ADDRESSES: usize = 64; #[cfg(unix)] #[cfg(not(test))] const LOCAL_INTERFACE_CACHE_TTL: Duration = Duration::from_secs(300); @@ -532,19 +532,25 @@ fn parse_mask_host_ip_literal(host: &str) -> Option { async fn resolve_mask_target_addrs( mask_host: &str, mask_port: u16, + upstream_manager: Option<&crate::transport::UpstreamManager>, ) -> std::io::Result> { - if let Some(addr) = resolve_socket_addr(mask_host, mask_port) { - return Ok(vec![addr]); - } - if let Some(ip) = parse_mask_host_ip_literal(mask_host) { return Ok(vec![SocketAddr::new(ip, mask_port)]); } + if let Some(upstream_manager) = upstream_manager { + return upstream_manager + .resolve_all(mask_host, mask_port) + .await + .map_err(|error| IoError::new(ErrorKind::NotFound, error.to_string())); + } + let addrs = timeout(MASK_TIMEOUT, lookup_host((mask_host, mask_port))) .await .map_err(|_| IoError::new(ErrorKind::TimedOut, "mask target DNS lookup timed out"))??; - let addrs = addrs.collect::>(); + let addrs = addrs + .take(MASK_DNS_RESULT_MAX_ADDRESSES) + .collect::>(); if addrs.is_empty() { return Err(IoError::new( ErrorKind::NotFound, @@ -999,6 +1005,34 @@ pub(crate) async fn handle_bad_client_with_shared( ) where R: AsyncRead + Unpin + Send + 'static, W: AsyncWrite + Unpin + Send + 'static, +{ + handle_bad_client_with_shared_resolver( + reader, + writer, + initial_data, + peer, + local_addr, + config, + beobachten, + shared, + None, + ) + .await; +} + +pub(super) async fn handle_bad_client_with_shared_resolver( + reader: R, + writer: W, + initial_data: &[u8], + peer: SocketAddr, + local_addr: SocketAddr, + config: &ProxyConfig, + beobachten: &BeobachtenStore, + shared: &ProxySharedState, + upstream_manager: Option<&crate::transport::UpstreamManager>, +) where + R: AsyncRead + Unpin + Send + 'static, + W: AsyncWrite + Unpin + Send + 'static, { let client_type = detect_client_type(initial_data); if config.general.beobachten { @@ -1112,7 +1146,8 @@ pub(crate) async fn handle_bad_client_with_shared( let mask_host = mask_target.host; let mask_port = mask_target.port; - let resolved_mask_addrs = match resolve_mask_target_addrs(mask_host, mask_port).await { + let resolved_mask_addrs = + match resolve_mask_target_addrs(mask_host, mask_port, upstream_manager).await { Ok(addrs) => addrs, Err(e) => { let outcome_started = Instant::now(); diff --git a/src/proxy/middle_relay.rs b/src/proxy/middle_relay.rs index 3049675..8182c03 100644 --- a/src/proxy/middle_relay.rs +++ b/src/proxy/middle_relay.rs @@ -34,7 +34,7 @@ use crate::stats::{ MeD2cFlushReason, MeD2cQuotaRejectStage, MeD2cWriteMode, QuotaReserveError, Stats, UserStats, }; use crate::stream::{BufferPool, CryptoReader, CryptoWriter, PooledBuffer}; -use crate::transport::middle_proxy::{MePool, MeResponse, proto_flags_for_tag}; +use crate::transport::middle_proxy::{ConnLease, MePool, MeResponse, proto_flags_for_tag}; mod c2me; mod d2c; diff --git a/src/proxy/middle_relay/session.rs b/src/proxy/middle_relay/session.rs index 47a731a..7c704f1 100644 --- a/src/proxy/middle_relay/session.rs +++ b/src/proxy/middle_relay/session.rs @@ -1,5 +1,42 @@ use super::*; +struct RelayConnLease { + connection: Option, + conn_id: u64, + shared: Arc, +} + +impl RelayConnLease { + fn new(connection: ConnLease, shared: Arc) -> Self { + let conn_id = connection.conn_id(); + Self { + connection: Some(connection), + conn_id, + shared, + } + } + + fn conn_id(&self) -> u64 { + self.conn_id + } + + async fn unregister(mut self) { + let Some(connection) = self.connection.take() else { + return; + }; + clear_relay_idle_candidate_in(self.shared.as_ref(), connection.conn_id()); + connection.unregister().await; + } +} + +impl Drop for RelayConnLease { + fn drop(&mut self) { + if let Some(connection) = self.connection.as_ref() { + clear_relay_idle_candidate_in(self.shared.as_ref(), connection.conn_id()); + } + } +} + /// Runs Middle-End relay with explicit kernel-conntrack close publication policy. pub(crate) async fn handle_via_middle_proxy_with_conntrack( mut crypto_reader: CryptoReader, @@ -44,7 +81,11 @@ where "Routing via Middle-End" ); - let (conn_id, me_rx) = me_pool.registry().register().await; + let Some((connection, me_rx)) = me_pool.register_connection().await else { + return Err(ProxyError::MiddleConnectionLost); + }; + let relay_connection = RelayConnLease::new(connection, Arc::clone(&shared)); + let conn_id = relay_connection.conn_id(); let trace_id = session_id; let bytes_me2c = Arc::new(AtomicU64::new(0)); let mut forensics = RelayForensicsState { @@ -76,7 +117,7 @@ where let _cutover_park_lease = stats.acquire_middle_cutover_park_lease(); tokio::time::sleep(delay).await; let _ = me_pool.send_close(conn_id).await; - me_pool.registry().unregister(conn_id).await; + relay_connection.unregister().await; return Err(ProxyError::RouteSwitched); } @@ -825,8 +866,7 @@ where } } - clear_relay_idle_candidate_in(shared.as_ref(), conn_id); - me_pool.registry().unregister(conn_id).await; + relay_connection.unregister().await; let pool_snapshot = buffer_pool.stats(); stats.set_buffer_pool_gauges( pool_snapshot.pooled, diff --git a/src/proxy/tests/masking_adversarial_tests.rs b/src/proxy/tests/masking_adversarial_tests.rs index ce2807a..6d930b6 100644 --- a/src/proxy/tests/masking_adversarial_tests.rs +++ b/src/proxy/tests/masking_adversarial_tests.rs @@ -229,7 +229,7 @@ async fn masking_fallback_down_mimics_timeout() { #[tokio::test] async fn masking_ssrf_resolve_internal_ranges_blocked() { - use crate::network::dns_overrides::resolve_socket_addr; + use crate::network::dns_overrides::DnsOverrides; let blocked_ips = [ "127.0.0.1", @@ -238,10 +238,11 @@ async fn masking_ssrf_resolve_internal_ranges_blocked() { "192.168.1.1", "0.0.0.0", ]; + let resolver = DnsOverrides::default(); for ip in blocked_ips { assert!( - resolve_socket_addr(ip, 80).is_none(), + resolver.resolve_socket_addr(ip, 80).is_none(), "runtime DNS overrides must not resolve unconfigured literal host targets" ); } diff --git a/src/proxy/tests/masking_connect_failure_close_matrix_security_tests.rs b/src/proxy/tests/masking_connect_failure_close_matrix_security_tests.rs index 718189c..af3f118 100644 --- a/src/proxy/tests/masking_connect_failure_close_matrix_security_tests.rs +++ b/src/proxy/tests/masking_connect_failure_close_matrix_security_tests.rs @@ -1,5 +1,5 @@ use super::*; -use crate::network::dns_overrides::install_entries; +use std::sync::Arc; use tokio::io::{AsyncReadExt, AsyncWriteExt, duplex}; use tokio::time::{Duration, Instant, timeout}; @@ -8,6 +8,7 @@ async fn run_connect_failure_case( port: u16, timing_normalization_enabled: bool, peer: SocketAddr, + dns_overrides: Vec, ) -> Duration { let mut config = ProxyConfig::default(); config.general.beobachten = false; @@ -21,6 +22,19 @@ async fn run_connect_failure_case( let local_addr: SocketAddr = "127.0.0.1:443".parse().unwrap(); let beobachten = BeobachtenStore::new(); + let upstream_manager = crate::transport::UpstreamManager::new( + Vec::new(), + 1, + 1, + 100, + 1, + 1, + true, + Arc::new(crate::stats::Stats::new()), + ) + .with_dns_overrides(&dns_overrides) + .unwrap(); + let shared = ProxySharedState::new(); let probe = b"CONNECT example.org:443 HTTP/1.1\r\nHost: example.org\r\n\r\n"; let (mut client_writer, client_reader) = duplex(1024); @@ -28,7 +42,7 @@ async fn run_connect_failure_case( let started = Instant::now(); let task = tokio::spawn(async move { - handle_bad_client( + handle_bad_client_with_shared_resolver( client_reader, client_visible_writer, probe, @@ -36,6 +50,8 @@ async fn run_connect_failure_case( local_addr, &config, &beobachten, + shared.as_ref(), + Some(&upstream_manager), ) .await; }); @@ -71,8 +87,14 @@ async fn connect_failure_refusal_close_behavior_matrix() { .parse() .unwrap(); let elapsed = - run_connect_failure_case("127.0.0.1", unused_port, timing_normalization_enabled, peer) - .await; + run_connect_failure_case( + "127.0.0.1", + unused_port, + timing_normalization_enabled, + peer, + Vec::new(), + ) + .await; if timing_normalization_enabled { assert!( @@ -94,9 +116,6 @@ async fn connect_failure_overridden_hostname_close_behavior_matrix() { let unused_port = temp_listener.local_addr().unwrap().port(); drop(temp_listener); - // Make hostname resolution deterministic in tests so timing ceilings are meaningful. - install_entries(&[format!("mask.invalid:{}:127.0.0.1", unused_port)]).unwrap(); - for (idx, timing_normalization_enabled) in [false, true].into_iter().enumerate() { let peer: SocketAddr = format!("203.0.113.220:{}", 54200 + idx as u16) .parse() @@ -106,6 +125,7 @@ async fn connect_failure_overridden_hostname_close_behavior_matrix() { unused_port, timing_normalization_enabled, peer, + vec![format!("mask.invalid:{}:127.0.0.1", unused_port)], ) .await; @@ -121,6 +141,4 @@ async fn connect_failure_overridden_hostname_close_behavior_matrix() { ); } } - - install_entries(&[]).unwrap(); } diff --git a/src/proxy/traffic_limiter.rs b/src/proxy/traffic_limiter.rs index 9d5d977..4785075 100644 --- a/src/proxy/traffic_limiter.rs +++ b/src/proxy/traffic_limiter.rs @@ -355,9 +355,10 @@ impl CidrBucket { } fn acquire_user_share(&self, user: &str) -> Arc { - let share = self.users.get_or_insert_with(user, CidrUserShare::new); - share.active_conns.fetch_add(1, Ordering::Relaxed); - share + self.users + .get_or_insert_with(user, CidrUserShare::new, |share| { + share.active_conns.fetch_add(1, Ordering::Relaxed); + }) } fn release_user_share(&self, user: &str, share: &Arc) { @@ -487,15 +488,20 @@ impl ShardedRegistry { (hasher.finish() as usize) & self.mask } - fn get_or_insert_with(&self, key: &str, make: F) -> Arc + fn get_or_insert_with(&self, key: &str, make: F, activate: A) -> Arc where F: FnOnce() -> T, + A: FnOnce(&Arc), { let shard = &self.shards[self.shard_index(key)]; match shard.entry(key.to_string()) { - dashmap::mapref::entry::Entry::Occupied(entry) => Arc::clone(entry.get()), + dashmap::mapref::entry::Entry::Occupied(entry) => { + activate(entry.get()); + Arc::clone(entry.get()) + } dashmap::mapref::entry::Entry::Vacant(slot) => { let value = Arc::new(make()); + activate(&value); slot.insert(Arc::clone(&value)); value } @@ -516,14 +522,7 @@ impl ShardedRegistry { F: Fn(&Arc) -> bool, { let shard = &self.shards[self.shard_index(key)]; - let should_remove = match shard.get(key) { - Some(entry) => predicate(entry.value()), - None => false, - }; - if !should_remove { - return false; - } - shard.remove(key).is_some() + shard.remove_if(key, |_, value| predicate(value)).is_some() } } @@ -743,9 +742,10 @@ impl TrafficLimiter { if let Some(limit) = policy.user_limits.get(user).copied() { let bucket = self .user_buckets - .get_or_insert_with(user, || UserBucket::new(limit)); + .get_or_insert_with(user, || UserBucket::new(limit), |bucket| { + bucket.active_leases.fetch_add(1, Ordering::Relaxed); + }); bucket.set_rates(limit); - bucket.active_leases.fetch_add(1, Ordering::Relaxed); self.user_scope .active_leases .fetch_add(1, Ordering::Relaxed); @@ -762,9 +762,10 @@ impl TrafficLimiter { }; let bucket = self .cidr_buckets - .get_or_insert_with(key, || CidrBucket::new(limits)); + .get_or_insert_with(key, || CidrBucket::new(limits), |bucket| { + bucket.active_leases.fetch_add(1, Ordering::Relaxed); + }); bucket.set_rates(limits); - bucket.active_leases.fetch_add(1, Ordering::Relaxed); self.cidr_scope .active_leases .fetch_add(1, Ordering::Relaxed); diff --git a/src/quota_state.rs b/src/quota_state.rs index 74abce8..96727d7 100644 --- a/src/quota_state.rs +++ b/src/quota_state.rs @@ -1,12 +1,18 @@ -use std::collections::BTreeMap; -use std::path::Path; +use std::collections::{BTreeMap, BTreeSet}; +use std::io::Write; +use std::path::{Path, PathBuf}; +use std::sync::Arc; use std::time::{SystemTime, UNIX_EPOCH}; use serde::{Deserialize, Serialize}; -use tokio::io::AsyncWriteExt; +use tokio::io::AsyncReadExt; +use tokio::sync::Mutex; use tracing::{info, warn}; -use crate::stats::{Stats, UserQuotaSnapshot}; +use crate::stats::{QuotaStore, UserQuotaSnapshot}; + +const QUOTA_STATE_MAX_BYTES: u64 = 16 * 1024 * 1024; +const QUOTA_STATE_MAX_USERS: usize = 65_536; #[derive(Debug, Default, Serialize, Deserialize)] pub(crate) struct QuotaStateFile { @@ -20,6 +26,148 @@ pub(crate) struct QuotaUserState { pub(crate) last_reset_epoch_secs: u64, } +/// Serialized process owner for bounded durable quota checkpoints and resets. +pub(crate) struct QuotaStateOwner { + path: PathBuf, + store: Arc, + mutation: Arc>, +} + +impl QuotaStateOwner { + /// Creates a process-owned quota persistence coordinator. + pub(crate) fn new(path: PathBuf, store: Arc) -> Arc { + Arc::new(Self { + path, + store, + mutation: Arc::new(Mutex::new(())), + }) + } + + /// Returns the process-owned quota checkpoint path. + pub(crate) fn path(&self) -> &Path { + &self.path + } + + /// Loads only quota entries owned by currently configured users. + pub(crate) async fn load(&self, configured_users: &BTreeSet) { + let _guard = self.mutation.lock().await; + let state = match read_state_file(&self.path).await { + Ok(Some(state)) => state, + Ok(None) => return, + Err(error) => { + warn!( + error = %error, + path = %self.path.display(), + "Failed to load quota state file" + ); + return; + } + }; + + let persisted_users = state.users.len(); + let mut loaded_users = 0usize; + for (user, quota) in state.users { + if loaded_users >= QUOTA_STATE_MAX_USERS || !configured_users.contains(&user) { + continue; + } + self.store + .load(&user, quota.used_bytes, quota.last_reset_epoch_secs); + loaded_users = loaded_users.saturating_add(1); + } + info!( + path = %self.path.display(), + loaded_users, + skipped_users = persisted_users.saturating_sub(loaded_users), + "Loaded bounded per-user quota state" + ); + } + + /// Persists a checkpoint filtered to the active configured user set. + pub(crate) async fn save( + &self, + configured_users: &BTreeSet, + ) -> std::io::Result<()> { + let guard = Arc::clone(&self.mutation).lock_owned().await; + let state = self.state_for_users(configured_users, None); + let path = self.path.clone(); + let task = tokio::task::spawn_blocking(move || { + let _guard = guard; + write_state_file_blocking(&path, &state) + }); + wait_for_blocking_io(task).await + } + + /// Durably commits a prospective reset before publishing it to live counters. + pub(crate) async fn reset_user( + &self, + configured_users: &BTreeSet, + user: &str, + ) -> std::io::Result { + let guard = Arc::clone(&self.mutation).lock_owned().await; + let last_reset_epoch_secs = now_epoch_secs(); + let prospective = UserQuotaSnapshot { + used_bytes: 0, + last_reset_epoch_secs, + }; + let state = self.state_for_users(configured_users, Some((user, prospective.clone()))); + let path = self.path.clone(); + let store = Arc::clone(&self.store); + let user = user.to_string(); + let task = tokio::task::spawn_blocking(move || { + let _guard = guard; + write_state_file_blocking(&path, &state)?; + Ok(store.reset(&user, last_reset_epoch_secs)) + }); + wait_for_blocking_io(task).await + } + + /// Removes a deleted user's persisted and in-memory quota ownership. + pub(crate) async fn remove_user( + &self, + configured_users: &BTreeSet, + user: &str, + ) -> std::io::Result<()> { + let guard = Arc::clone(&self.mutation).lock_owned().await; + let state = self.state_for_users(configured_users, None); + let path = self.path.clone(); + let store = Arc::clone(&self.store); + let user = user.to_string(); + let task = tokio::task::spawn_blocking(move || { + let _guard = guard; + let persisted = write_state_file_blocking(&path, &state); + store.remove(&user); + persisted + }); + wait_for_blocking_io(task).await + } + + fn state_for_users( + &self, + configured_users: &BTreeSet, + override_user: Option<(&str, UserQuotaSnapshot)>, + ) -> QuotaStateFile { + let snapshot = self.store.snapshot(); + let mut users = BTreeMap::new(); + let mut last_reset_epoch_secs = 0; + for user in configured_users.iter().take(QUOTA_STATE_MAX_USERS) { + let quota = override_user + .as_ref() + .filter(|(override_name, _)| *override_name == user.as_str()) + .map(|(_, quota)| quota.clone()) + .or_else(|| snapshot.get(user).cloned()); + let Some(quota) = quota else { + continue; + }; + last_reset_epoch_secs = last_reset_epoch_secs.max(quota.last_reset_epoch_secs); + users.insert(user.clone(), quota_user_state(quota)); + } + QuotaStateFile { + last_reset_epoch_secs, + users, + } + } +} + fn now_epoch_secs() -> u64 { SystemTime::now() .duration_since(UNIX_EPOCH) @@ -27,83 +175,104 @@ fn now_epoch_secs() -> u64 { .as_secs() } -pub(crate) async fn load_quota_state(path: &Path, stats: &Stats) { - let bytes = match tokio::fs::read(path).await { - Ok(bytes) => bytes, - Err(error) if error.kind() == std::io::ErrorKind::NotFound => return, - Err(error) => { - warn!( - error = %error, - path = %path.display(), - "Failed to read quota state file" - ); - return; +async fn read_state_file(path: &Path) -> std::io::Result> { + let file = match tokio::fs::File::open(path).await { + Ok(file) => file, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None), + Err(error) => return Err(error), + }; + if file.metadata().await?.len() > QUOTA_STATE_MAX_BYTES { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "quota state file exceeds the 16 MiB limit", + )); + } + let mut payload = Vec::new(); + file.take(QUOTA_STATE_MAX_BYTES.saturating_add(1)) + .read_to_end(&mut payload) + .await?; + if payload.len() as u64 > QUOTA_STATE_MAX_BYTES { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "quota state file grew beyond the 16 MiB limit while reading", + )); + } + let state = serde_json::from_slice(&payload).map_err(|error| { + std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("failed to parse quota state: {error}"), + ) + })?; + Ok(Some(state)) +} + +async fn write_state_file(path: &Path, state: QuotaStateFile) -> std::io::Result<()> { + let path = path.to_path_buf(); + let task = tokio::task::spawn_blocking(move || write_state_file_blocking(&path, &state)); + wait_for_blocking_io(task).await +} + +async fn wait_for_blocking_io( + task: tokio::task::JoinHandle>, +) -> std::io::Result { + task.await + .map_err(|error| std::io::Error::other(format!("quota checkpoint join failed: {error}")))? +} + +fn write_state_file_blocking(path: &Path, state: &QuotaStateFile) -> std::io::Result<()> { + let parent = path + .parent() + .filter(|parent| !parent.as_os_str().is_empty()) + .unwrap_or_else(|| Path::new(".")); + std::fs::create_dir_all(parent)?; + let mut payload = serde_json::to_vec_pretty(state)?; + payload.push(b'\n'); + if payload.len() as u64 > QUOTA_STATE_MAX_BYTES { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "quota state payload exceeds the 16 MiB limit", + )); + } + + let mut last_collision = None; + for _ in 0..8 { + let tmp_path = path.with_extension(format!( + "tmp.{}.{}", + std::process::id(), + rand::random::() + )); + let mut file = match std::fs::OpenOptions::new() + .write(true) + .create_new(true) + .open(&tmp_path) + { + Ok(file) => file, + Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => { + last_collision = Some(error); + continue; + } + Err(error) => return Err(error), + }; + let result = (|| { + file.write_all(&payload)?; + file.sync_all()?; + drop(file); + std::fs::rename(&tmp_path, path)?; + #[cfg(unix)] + std::fs::File::open(parent)?.sync_all()?; + Ok(()) + })(); + if result.is_err() { + let _ = std::fs::remove_file(&tmp_path); } - }; - - let state = match serde_json::from_slice::(&bytes) { - Ok(state) => state, - Err(error) => { - warn!( - error = %error, - path = %path.display(), - "Failed to parse quota state file" - ); - return; - } - }; - - let loaded_users = state.users.len(); - for (user, quota) in state.users { - stats.load_user_quota_state(&user, quota.used_bytes, quota.last_reset_epoch_secs); + return result; } - info!( - path = %path.display(), - loaded_users, - "Loaded per-user quota state" - ); -} - -pub(crate) async fn save_quota_state(path: &Path, stats: &Stats) -> std::io::Result<()> { - let mut users = BTreeMap::new(); - let mut last_reset_epoch_secs = 0; - for (user, quota) in stats.user_quota_snapshot() { - last_reset_epoch_secs = last_reset_epoch_secs.max(quota.last_reset_epoch_secs); - users.insert(user, quota_user_state(quota)); - } - - let state = QuotaStateFile { - last_reset_epoch_secs, - users, - }; - write_state_file(path, &state).await -} - -pub(crate) async fn reset_user_quota( - path: &Path, - stats: &Stats, - user: &str, -) -> std::io::Result { - let snapshot = stats.reset_user_quota(user); - save_quota_state(path, stats).await?; - Ok(snapshot) -} - -async fn write_state_file(path: &Path, state: &QuotaStateFile) -> std::io::Result<()> { - if let Some(parent) = path.parent() - && !parent.as_os_str().is_empty() - { - tokio::fs::create_dir_all(parent).await?; - } - - let tmp_path = path.with_extension(format!("tmp.{}", now_epoch_secs())); - let payload = serde_json::to_vec_pretty(state)?; - let mut file = tokio::fs::File::create(&tmp_path).await?; - file.write_all(&payload).await?; - file.write_all(b"\n").await?; - file.sync_all().await?; - drop(file); - tokio::fs::rename(&tmp_path, path).await + Err(last_collision.unwrap_or_else(|| { + std::io::Error::new( + std::io::ErrorKind::AlreadyExists, + "failed to allocate a unique quota checkpoint temporary file", + ) + })) } fn quota_user_state(quota: UserQuotaSnapshot) -> QuotaUserState { @@ -112,3 +281,131 @@ fn quota_user_state(quota: UserQuotaSnapshot) -> QuotaUserState { last_reset_epoch_secs: quota.last_reset_epoch_secs, } } + +#[cfg(test)] +mod tests { + use super::*; + + fn users(names: &[&str]) -> BTreeSet { + names.iter().map(|name| (*name).to_string()).collect() + } + + #[tokio::test] + async fn failed_durable_reset_does_not_publish_live_reset() { + let directory = tempfile::tempdir().unwrap(); + let store = Arc::new(QuotaStore::default()); + store.user("alice").charge(512); + let owner = QuotaStateOwner::new(directory.path().to_path_buf(), store.clone()); + + assert!(owner.reset_user(&users(&["alice"]), "alice").await.is_err()); + assert_eq!(store.used("alice"), 512); + } + + #[tokio::test] + async fn concurrent_resets_serialize_without_losing_committed_users() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("quota.json"); + let store = Arc::new(QuotaStore::default()); + store.user("alice").charge(512); + store.user("bob").charge(1024); + let owner = QuotaStateOwner::new(path.clone(), store.clone()); + let configured = Arc::new(users(&["alice", "bob"])); + + let alice_owner = owner.clone(); + let alice_users = configured.clone(); + let alice = tokio::spawn(async move { + alice_owner.reset_user(&alice_users, "alice").await + }); + let bob_owner = owner.clone(); + let bob_users = configured.clone(); + let bob = tokio::spawn(async move { bob_owner.reset_user(&bob_users, "bob").await }); + alice.await.unwrap().unwrap(); + bob.await.unwrap().unwrap(); + + let state = read_state_file(&path).await.unwrap().unwrap(); + assert_eq!(state.users["alice"].used_bytes, 0); + assert_eq!(state.users["bob"].used_bytes, 0); + assert_eq!(store.used("alice"), 0); + assert_eq!(store.used("bob"), 0); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn cancelled_waiter_does_not_release_blocking_mutation_ownership() { + let mutation = Arc::new(Mutex::new(())); + let guard = Arc::clone(&mutation).lock_owned().await; + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let (release_tx, release_rx) = std::sync::mpsc::channel(); + let operation = tokio::task::spawn_blocking(move || { + let _guard = guard; + let _ = started_tx.send(()); + release_rx.recv().unwrap(); + Ok::<_, std::io::Error>(()) + }); + let waiter = tokio::spawn(wait_for_blocking_io(operation)); + started_rx.await.unwrap(); + + waiter.abort(); + assert!(waiter.await.unwrap_err().is_cancelled()); + assert!(mutation.try_lock().is_err()); + + release_tx.send(()).unwrap(); + tokio::time::timeout(std::time::Duration::from_secs(1), async { + loop { + if mutation.try_lock().is_ok() { + return; + } + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + } + + #[tokio::test] + async fn load_rejects_oversized_state_before_parsing() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("quota.json"); + let file = std::fs::File::create(&path).unwrap(); + file.set_len(QUOTA_STATE_MAX_BYTES + 1).unwrap(); + let store = Arc::new(QuotaStore::default()); + let owner = QuotaStateOwner::new(path, store.clone()); + + owner.load(&users(&["alice"])).await; + + assert_eq!(store.used("alice"), 0); + } + + #[tokio::test] + async fn load_filters_entries_without_configured_ownership() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("quota.json"); + let state = QuotaStateFile { + last_reset_epoch_secs: 1, + users: BTreeMap::from([ + ( + "alice".to_string(), + QuotaUserState { + used_bytes: 512, + last_reset_epoch_secs: 1, + }, + ), + ( + "orphan".to_string(), + QuotaUserState { + used_bytes: 1024, + last_reset_epoch_secs: 1, + }, + ), + ]), + }; + write_state_file(&path, state).await.unwrap(); + let store = Arc::new(QuotaStore::default()); + let owner = QuotaStateOwner::new(path, store.clone()); + + owner.load(&users(&["alice"])).await; + + assert_eq!(store.used("alice"), 512); + assert_eq!(store.used("orphan"), 0); + assert!(!store.snapshot().contains_key("orphan")); + } +} diff --git a/src/stats/beobachten.rs b/src/stats/beobachten.rs index 5684455..046111c 100644 --- a/src/stats/beobachten.rs +++ b/src/stats/beobachten.rs @@ -1,17 +1,21 @@ //! Per-IP forensic buckets for scanner and handshake failure observation. +use std::collections::hash_map::RandomState; use std::collections::{BTreeMap, HashMap}; +use std::hash::{BuildHasher, Hash, Hasher}; use std::net::IpAddr; use std::time::{Duration, Instant}; use parking_lot::Mutex; const CLEANUP_INTERVAL: Duration = Duration::from_secs(30); -const MAX_BEOBACHTEN_ENTRIES: usize = 65_536; +const BEOBACHTEN_SHARDS: usize = 64; +const BEOBACHTEN_ENTRIES_PER_SHARD: usize = 1024; #[derive(Default)] -struct BeobachtenInner { - entries: HashMap<(String, IpAddr), BeobachtenEntry>, +struct BeobachtenShard { + classes: HashMap>, + entries: usize, last_cleanup: Option, } @@ -23,7 +27,8 @@ struct BeobachtenEntry { /// In-memory, TTL-scoped per-IP counters keyed by source class. pub struct BeobachtenStore { - inner: Mutex, + shards: Vec>, + hash_builder: RandomState, } impl Default for BeobachtenStore { @@ -35,7 +40,10 @@ impl Default for BeobachtenStore { impl BeobachtenStore { pub fn new() -> Self { Self { - inner: Mutex::new(BeobachtenInner::default()), + shards: (0..BEOBACHTEN_SHARDS) + .map(|_| Mutex::new(BeobachtenShard::default())) + .collect(), + hash_builder: RandomState::new(), } } @@ -45,27 +53,31 @@ impl BeobachtenStore { } let now = Instant::now(); - let mut guard = self.inner.lock(); - Self::cleanup_if_needed(&mut guard, now, ttl); + let shard_index = self.shard_index(class, ip); + let mut shard = self.shards[shard_index].lock(); + Self::cleanup_if_needed(&mut shard, now, ttl); - let key = (class.to_string(), ip); - if let Some(entry) = guard.entries.get_mut(&key) { + if let Some(entry) = shard + .classes + .get_mut(class) + .and_then(|entries| entries.get_mut(&ip)) + { entry.tries = entry.tries.saturating_add(1); entry.last_seen = now; return; } - if guard.entries.len() >= MAX_BEOBACHTEN_ENTRIES { + if shard.entries >= BEOBACHTEN_ENTRIES_PER_SHARD { return; } - - guard.entries.insert( - key, + shard.classes.entry(class.to_string()).or_default().insert( + ip, BeobachtenEntry { tries: 1, last_seen: now, }, ); + shard.entries = shard.entries.saturating_add(1); } pub fn snapshot_text(&self, ttl: Duration) -> String { @@ -74,21 +86,15 @@ impl BeobachtenStore { } let now = Instant::now(); - let entries = { - let mut guard = self.inner.lock(); - Self::cleanup(&mut guard, now, ttl); - guard.last_cleanup = Some(now); - - guard - .entries - .iter() - .map(|((class, ip), entry)| (class.clone(), *ip, entry.tries)) - .collect::>() - }; - let mut grouped = BTreeMap::>::new(); - for (class, ip, tries) in entries { - grouped.entry(class).or_default().push((ip, tries)); + for shard in &self.shards { + let mut shard = shard.lock(); + Self::cleanup(&mut shard, now, ttl); + shard.last_cleanup = Some(now); + for (class, entries) in &shard.classes { + let output = grouped.entry(class.clone()).or_default(); + output.extend(entries.iter().map(|(ip, entry)| (*ip, entry.tries))); + } } if grouped.is_empty() { @@ -111,24 +117,61 @@ impl BeobachtenStore { out.push_str(&format!("{ip}-{tries}\n")); } } - out } - fn cleanup_if_needed(inner: &mut BeobachtenInner, now: Instant, ttl: Duration) { - let should_cleanup = match inner.last_cleanup { + fn shard_index(&self, class: &str, ip: IpAddr) -> usize { + let mut hasher = self.hash_builder.build_hasher(); + class.hash(&mut hasher); + ip.hash(&mut hasher); + (hasher.finish() as usize) % BEOBACHTEN_SHARDS + } + + fn cleanup_if_needed(shard: &mut BeobachtenShard, now: Instant, ttl: Duration) { + let should_cleanup = match shard.last_cleanup { Some(last) => now.saturating_duration_since(last) >= CLEANUP_INTERVAL, None => true, }; if should_cleanup { - Self::cleanup(inner, now, ttl); - inner.last_cleanup = Some(now); + Self::cleanup(shard, now, ttl); + shard.last_cleanup = Some(now); } } - fn cleanup(inner: &mut BeobachtenInner, now: Instant, ttl: Duration) { - inner - .entries - .retain(|_, entry| now.saturating_duration_since(entry.last_seen) <= ttl); + fn cleanup(shard: &mut BeobachtenShard, now: Instant, ttl: Duration) { + for entries in shard.classes.values_mut() { + entries.retain(|_, entry| now.saturating_duration_since(entry.last_seen) <= ttl); + } + shard.classes.retain(|_, entries| !entries.is_empty()); + shard.entries = shard.classes.values().map(HashMap::len).sum(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn one_shard_never_exceeds_its_allocation_bound() { + let store = BeobachtenStore::new(); + let ttl = Duration::from_secs(60); + let target_shard = 0; + let mut inserted = 0usize; + for suffix in 0..u32::MAX { + let ip = IpAddr::V4(std::net::Ipv4Addr::from(suffix)); + if store.shard_index("scanner", ip) != target_shard { + continue; + } + store.record("scanner", ip, ttl); + inserted = inserted.saturating_add(1); + if inserted > BEOBACHTEN_ENTRIES_PER_SHARD + 64 { + break; + } + } + + assert_eq!( + store.shards[target_shard].lock().entries, + BEOBACHTEN_ENTRIES_PER_SHARD + ); } } diff --git a/src/stats/me_counters.rs b/src/stats/me_counters.rs index a6c9fb4..e4c588c 100644 --- a/src/stats/me_counters.rs +++ b/src/stats/me_counters.rs @@ -77,11 +77,45 @@ impl Stats { if !self.telemetry_me_allows_normal() { return; } - let entry = self - .me_handshake_error_codes - .entry(code) - .or_insert_with(|| AtomicU64::new(0)); - entry.fetch_add(1, Ordering::Relaxed); + if let Some(entry) = self.me_handshake_error_codes.get(&code) { + entry.fetch_add(1, Ordering::Relaxed); + return; + } + + let mut slots = self + .me_handshake_error_code_slots + .load(Ordering::Acquire); + loop { + if slots >= ME_HANDSHAKE_ERROR_CODE_MAX { + if let Some(entry) = self.me_handshake_error_codes.get(&code) { + entry.fetch_add(1, Ordering::Relaxed); + } else { + self.me_handshake_error_code_overflow_total + .fetch_add(1, Ordering::Relaxed); + } + return; + } + match self.me_handshake_error_code_slots.compare_exchange_weak( + slots, + slots + 1, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => break, + Err(observed) => slots = observed, + } + } + + match self.me_handshake_error_codes.entry(code) { + dashmap::mapref::entry::Entry::Occupied(entry) => { + self.me_handshake_error_code_slots + .fetch_sub(1, Ordering::AcqRel); + entry.get().fetch_add(1, Ordering::Relaxed); + } + dashmap::mapref::entry::Entry::Vacant(entry) => { + entry.insert(AtomicU64::new(1)); + } + } } pub fn increment_me_reader_eof_total(&self) { if self.telemetry_me_allows_normal() { @@ -440,3 +474,44 @@ impl Stats { } } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use super::*; + + #[test] + fn handshake_error_code_cardinality_is_bounded_under_concurrency() { + const WORKERS: usize = 8; + const CODES_PER_WORKER: usize = 128; + + let stats = Arc::new(Stats::new()); + std::thread::scope(|scope| { + for worker in 0..WORKERS { + let stats = stats.clone(); + scope.spawn(move || { + for code in 0..CODES_PER_WORKER { + stats.increment_me_handshake_error_code( + (worker * CODES_PER_WORKER + code) as i32, + ); + } + }); + } + }); + + let counts = stats.get_me_handshake_error_code_counts(); + assert!(counts.len() <= ME_HANDSHAKE_ERROR_CODE_MAX); + assert_eq!( + counts.iter().map(|(_, total)| *total).sum::() + + stats.get_me_handshake_error_code_overflow_total(), + (WORKERS * CODES_PER_WORKER) as u64 + ); + assert_eq!( + stats + .me_handshake_error_code_slots + .load(Ordering::Acquire), + counts.len() + ); + } +} diff --git a/src/stats/me_getters.rs b/src/stats/me_getters.rs index ce04423..78c3dff 100644 --- a/src/stats/me_getters.rs +++ b/src/stats/me_getters.rs @@ -10,6 +10,10 @@ impl Stats { out.sort_by_key(|(code, _)| *code); out } + pub fn get_me_handshake_error_code_overflow_total(&self) -> u64 { + self.me_handshake_error_code_overflow_total + .load(Ordering::Relaxed) + } pub fn get_me_route_drop_no_conn(&self) -> u64 { self.me_route_drop_no_conn.load(Ordering::Relaxed) } diff --git a/src/stats/mod.rs b/src/stats/mod.rs index a16f8eb..170ad08 100644 --- a/src/stats/mod.rs +++ b/src/stats/mod.rs @@ -18,7 +18,7 @@ mod writer_counters; use dashmap::DashMap; use std::collections::HashMap; use std::sync::Arc; -use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU64, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU64, AtomicUsize, Ordering}; use std::time::Instant; pub(crate) use self::quota_store::QuotaStore; @@ -28,6 +28,8 @@ use self::telemetry::TelemetryPolicy; pub use self::tls_fingerprints::TlsFingerprintSnapshotRow; use crate::config::MeWriterPickMode; +const ME_HANDSHAKE_ERROR_CODE_MAX: usize = 64; + #[derive(Clone, Copy)] enum RouteConnectionGauge { Direct, @@ -219,6 +221,8 @@ pub struct Stats { me_floor_swap_idle_total: AtomicU64, me_floor_swap_idle_failed_total: AtomicU64, me_handshake_error_codes: DashMap, + me_handshake_error_code_slots: AtomicUsize, + me_handshake_error_code_overflow_total: AtomicU64, me_route_drop_no_conn: AtomicU64, me_route_drop_channel_closed: AtomicU64, me_route_drop_queue_full: AtomicU64, diff --git a/src/stats/quota_store.rs b/src/stats/quota_store.rs index 79f4c99..e338b76 100644 --- a/src/stats/quota_store.rs +++ b/src/stats/quota_store.rs @@ -56,6 +56,10 @@ impl QuotaStore { } } + pub(crate) fn remove(&self, user: &str) { + self.users.remove(user); + } + pub(crate) fn snapshot(&self) -> HashMap { let mut out = HashMap::new(); for entry in self.users.iter() { diff --git a/src/tls_front/cache.rs b/src/tls_front/cache.rs index 711b981..a5516e6 100644 --- a/src/tls_front/cache.rs +++ b/src/tls_front/cache.rs @@ -1,13 +1,14 @@ -use std::collections::HashMap; -use std::collections::hash_map::DefaultHasher; -use std::hash::{Hash, Hasher}; +use std::collections::{HashMap, HashSet}; +use std::collections::hash_map::RandomState; +use std::hash::{BuildHasher, Hash, Hasher}; use std::net::IpAddr; use std::path::{Path, PathBuf}; use std::sync::Arc; -use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; use tokio::sync::RwLock; +use tokio::io::AsyncReadExt; use tokio::time::sleep; use tracing::{debug, info, warn}; @@ -16,21 +17,160 @@ use crate::tls_front::types::{ TlsProfileSource, }; -const FULL_CERT_SENT_SWEEP_INTERVAL_SECS: u64 = 30; -const FULL_CERT_SENT_MAX_IPS: usize = 65_536; +const FULL_CERT_SENT_SWEEP_INTERVAL_SECS: u64 = 1; +const FULL_CERT_SENT_MAX_ENTRIES: usize = 65_536; const FULL_CERT_SENT_SHARDS: usize = 64; +const FULL_CERT_SENT_MAX_ENTRIES_PER_SHARD: usize = + (FULL_CERT_SENT_MAX_ENTRIES / FULL_CERT_SENT_SHARDS) * 2; +const TLS_FRONT_DISK_ENTRY_MAX_BYTES: u64 = 1024 * 1024; -static FULL_CERT_SENT_IPS_GAUGE: AtomicU64 = AtomicU64::new(0); -static FULL_CERT_SENT_CAP_DROPS: AtomicU64 = AtomicU64::new(0); - -/// Current number of IPs tracked by the TLS full-cert budget gate. -pub(crate) fn full_cert_sent_ips_for_metrics() -> u64 { - FULL_CERT_SENT_IPS_GAUGE.load(Ordering::Relaxed) +#[derive(Debug, Clone, Hash, PartialEq, Eq)] +struct FullCertBudgetKey { + domain: Arc, + client_ip: IpAddr, } -/// Number of new IPs denied a full-cert budget slot because the cap was reached. -pub(crate) fn full_cert_sent_cap_drops_for_metrics() -> u64 { - FULL_CERT_SENT_CAP_DROPS.load(Ordering::Relaxed) +#[derive(Debug)] +struct FullCertBudgetEntry { + expires_at: Option, +} + +/// Process-owned, bounded TLS full-certificate admission history. +#[derive(Debug)] +pub(crate) struct TlsFullCertBudget { + shards: Vec>>, + hash_builder: RandomState, + entries: AtomicUsize, + cap_drops: AtomicU64, + last_sweep_epoch_secs: AtomicU64, + sweep_cursor: AtomicUsize, +} + +impl TlsFullCertBudget { + /// Creates an empty process-owned full-certificate budget. + pub(crate) fn new() -> Self { + Self { + shards: (0..FULL_CERT_SENT_SHARDS) + .map(|_| RwLock::new(HashMap::new())) + .collect(), + hash_builder: RandomState::new(), + entries: AtomicUsize::new(0), + cap_drops: AtomicU64::new(0), + last_sweep_epoch_secs: AtomicU64::new(0), + sweep_cursor: AtomicUsize::new(0), + } + } + + fn shard_index(&self, key: &FullCertBudgetKey) -> usize { + let mut hasher = self.hash_builder.build_hasher(); + key.hash(&mut hasher); + (hasher.finish() as usize) % FULL_CERT_SENT_SHARDS + } + + fn decrement_entries(&self, amount: usize) { + if amount == 0 { + return; + } + let _ = self + .entries + .fetch_update(Ordering::AcqRel, Ordering::Relaxed, |current| { + Some(current.saturating_sub(amount)) + }); + } + + fn try_reserve_entry(&self) -> bool { + let mut current = self.entries.load(Ordering::Relaxed); + loop { + if current >= FULL_CERT_SENT_MAX_ENTRIES { + return false; + } + match self.entries.compare_exchange_weak( + current, + current.saturating_add(1), + Ordering::AcqRel, + Ordering::Relaxed, + ) { + Ok(_) => return true, + Err(actual) => current = actual, + } + } + } + + async fn sweep_one_shard(&self, now: Instant) { + let shard_index = self.sweep_cursor.fetch_add(1, Ordering::Relaxed) + % FULL_CERT_SENT_SHARDS; + let mut guard = self.shards[shard_index].write().await; + let before = guard.len(); + guard.retain(|_, entry| entry.expires_at.map_or(true, |expires_at| expires_at > now)); + self.decrement_entries(before.saturating_sub(guard.len())); + } + + async fn maybe_sweep(&self, now: Instant) { + let now_epoch_secs = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs(); + let should_sweep = self + .last_sweep_epoch_secs + .fetch_update(Ordering::AcqRel, Ordering::Relaxed, |last_sweep| { + if now_epoch_secs.saturating_sub(last_sweep) + >= FULL_CERT_SENT_SWEEP_INTERVAL_SECS + { + Some(now_epoch_secs) + } else { + None + } + }) + .is_ok(); + if should_sweep { + self.sweep_one_shard(now).await; + } + } + + async fn take(&self, domain: Arc, client_ip: IpAddr, ttl: Duration) -> bool { + if ttl.is_zero() { + return true; + } + + let now = Instant::now(); + self.maybe_sweep(now).await; + let expires_at = now.checked_add(ttl); + let key = FullCertBudgetKey { domain, client_ip }; + let shard_index = self.shard_index(&key); + let mut guard = self.shards[shard_index].write().await; + + if let Some(entry) = guard.get_mut(&key) { + if entry.expires_at.is_some_and(|expires_at| expires_at <= now) { + entry.expires_at = expires_at; + return true; + } + return false; + } + + if guard.len() >= FULL_CERT_SENT_MAX_ENTRIES_PER_SHARD { + let before = guard.len(); + guard.retain(|_, entry| { + entry.expires_at.map_or(true, |expires_at| expires_at > now) + }); + self.decrement_entries(before.saturating_sub(guard.len())); + } + if guard.len() >= FULL_CERT_SENT_MAX_ENTRIES_PER_SHARD || !self.try_reserve_entry() { + self.cap_drops.fetch_add(1, Ordering::Relaxed); + return false; + } + guard.insert(key, FullCertBudgetEntry { expires_at }); + true + } + + /// Returns the current number of retained domain and client IP entries. + pub(crate) fn entries_for_metrics(&self) -> u64 { + self.entries.load(Ordering::Relaxed) as u64 + } + + /// Returns the cumulative number of entries rejected by hard bounds. + pub(crate) fn cap_drops_for_metrics(&self) -> u64 { + self.cap_drops.load(Ordering::Relaxed) + } } /// Lightweight in-memory + optional on-disk cache for TLS fronting data. @@ -38,8 +178,9 @@ pub(crate) fn full_cert_sent_cap_drops_for_metrics() -> u64 { pub struct TlsFrontCache { memory: RwLock>>, default: Arc, - full_cert_sent_shards: Vec>>, - full_cert_sent_last_sweep_epoch_secs: AtomicU64, + full_cert_budget: Arc, + full_cert_domain_keys: HashMap>, + disk_entry_names: HashSet, disk_path: PathBuf, } @@ -91,6 +232,21 @@ fn key_share_group_label(group: Option) -> &'static str { #[allow(dead_code)] impl TlsFrontCache { pub fn new(domains: &[String], default_len: usize, disk_path: impl AsRef) -> Self { + Self::new_with_full_cert_budget( + domains, + default_len, + disk_path, + Arc::new(TlsFullCertBudget::new()), + ) + } + + /// Creates a generation-local cache backed by the process-owned full-cert budget. + pub(crate) fn new_with_full_cert_budget( + domains: &[String], + default_len: usize, + disk_path: impl AsRef, + full_cert_budget: Arc, + ) -> Self { let default_template = ParsedServerHello { version: [0x03, 0x03], random: [0u8; 32], @@ -112,17 +268,24 @@ impl TlsFrontCache { }); let mut map = HashMap::new(); + let mut full_cert_domain_keys = HashMap::new(); + let mut disk_entry_names = HashSet::new(); for d in domains { map.insert(d.clone(), default.clone()); + disk_entry_names.insert(format!("{}.json", d.replace(['/', '\\'], "_"))); + let canonical: Arc = Arc::from(normalize_dns_name(d)); + full_cert_domain_keys.insert(d.clone(), canonical.clone()); + full_cert_domain_keys + .entry(canonical.to_string()) + .or_insert(canonical); } Self { memory: RwLock::new(map), default, - full_cert_sent_shards: (0..FULL_CERT_SENT_SHARDS) - .map(|_| RwLock::new(HashMap::new())) - .collect(), - full_cert_sent_last_sweep_epoch_secs: AtomicU64::new(0), + full_cert_budget, + full_cert_domain_keys, + disk_entry_names, disk_path: disk_path.as_ref().to_path_buf(), } } @@ -202,114 +365,66 @@ impl TlsFrontCache { .collect() } - fn full_cert_sent_shard_index(client_ip: IpAddr) -> usize { - let mut hasher = DefaultHasher::new(); - client_ip.hash(&mut hasher); - (hasher.finish() as usize) % FULL_CERT_SENT_SHARDS + fn full_cert_domain_key(&self, domain: &str) -> Arc { + self.full_cert_domain_keys + .get(domain) + .cloned() + .unwrap_or_else(|| Arc::from(normalize_dns_name(domain))) } - fn full_cert_sent_shard(&self, client_ip: IpAddr) -> &RwLock> { - &self.full_cert_sent_shards[Self::full_cert_sent_shard_index(client_ip)] + /// Returns true when the selected domain and client IP may receive a full cert payload. + pub async fn take_full_cert_budget_for_ip( + &self, + domain: &str, + client_ip: IpAddr, + ttl: Duration, + ) -> bool { + self.full_cert_budget + .take(self.full_cert_domain_key(domain), client_ip, ttl) + .await } - fn decrement_full_cert_sent_entries(amount: usize) { - if amount == 0 { - return; - } - let amount = amount as u64; - let _ = - FULL_CERT_SENT_IPS_GAUGE.fetch_update(Ordering::AcqRel, Ordering::Relaxed, |current| { - Some(current.saturating_sub(amount)) - }); + /// Returns the current process-owned full-cert budget entry count. + pub(crate) fn full_cert_budget_entries_for_metrics(&self) -> u64 { + self.full_cert_budget.entries_for_metrics() } - fn try_reserve_full_cert_sent_entry() -> bool { - let mut current = FULL_CERT_SENT_IPS_GAUGE.load(Ordering::Relaxed); - loop { - if current >= FULL_CERT_SENT_MAX_IPS as u64 { - return false; - } - match FULL_CERT_SENT_IPS_GAUGE.compare_exchange_weak( - current, - current.saturating_add(1), - Ordering::AcqRel, - Ordering::Relaxed, - ) { - Ok(_) => return true, - Err(actual) => current = actual, - } - } - } - - async fn sweep_full_cert_sent_shards(&self, now: Instant, ttl: Duration) { - for shard in &self.full_cert_sent_shards { - let mut guard = shard.write().await; - let before = guard.len(); - guard.retain(|_, seen_at| now.duration_since(*seen_at) < ttl); - Self::decrement_full_cert_sent_entries(before.saturating_sub(guard.len())); - } - } - - /// Returns true when full cert payload should be sent for client_ip - /// according to TTL policy. - pub async fn take_full_cert_budget_for_ip(&self, client_ip: IpAddr, ttl: Duration) -> bool { - if ttl.is_zero() { - return true; - } - - let now = Instant::now(); - let now_epoch_secs = SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - .as_secs(); - let should_sweep = self - .full_cert_sent_last_sweep_epoch_secs - .fetch_update(Ordering::AcqRel, Ordering::Relaxed, |last_sweep| { - if now_epoch_secs.saturating_sub(last_sweep) >= FULL_CERT_SENT_SWEEP_INTERVAL_SECS { - Some(now_epoch_secs) - } else { - None - } - }) - .is_ok(); - - if should_sweep { - self.sweep_full_cert_sent_shards(now, ttl).await; - } - - let mut guard = self.full_cert_sent_shard(client_ip).write().await; - let allowed = match guard.get_mut(&client_ip) { - Some(seen_at) => { - if now.duration_since(*seen_at) >= ttl { - *seen_at = now; - true - } else { - false - } - } - None => { - if !Self::try_reserve_full_cert_sent_entry() { - FULL_CERT_SENT_CAP_DROPS.fetch_add(1, Ordering::Relaxed); - return false; - } - guard.insert(client_ip, now); - true - } - }; - allowed + /// Returns the cumulative process-owned full-cert budget cap drops. + pub(crate) fn full_cert_budget_cap_drops_for_metrics(&self) -> u64 { + self.full_cert_budget.cap_drops_for_metrics() } #[cfg(test)] - async fn insert_full_cert_sent_for_tests(&self, client_ip: IpAddr, seen_at: Instant) { - let mut guard = self.full_cert_sent_shard(client_ip).write().await; - if guard.insert(client_ip, seen_at).is_none() { - FULL_CERT_SENT_IPS_GAUGE.fetch_add(1, Ordering::Relaxed); + async fn insert_full_cert_sent_for_tests( + &self, + domain: &str, + client_ip: IpAddr, + expires_at: Instant, + ) { + let key = FullCertBudgetKey { + domain: self.full_cert_domain_key(domain), + client_ip, + }; + let shard_index = self.full_cert_budget.shard_index(&key); + let mut guard = self.full_cert_budget.shards[shard_index].write().await; + if guard + .insert( + key, + FullCertBudgetEntry { + expires_at: Some(expires_at), + }, + ) + .is_none() + { + self.full_cert_budget + .entries + .fetch_add(1, Ordering::Relaxed); } } #[cfg(test)] async fn full_cert_sent_is_empty_for_tests(&self) -> bool { - for shard in &self.full_cert_sent_shards { + for shard in &self.full_cert_budget.shards { if !shard.read().await.is_empty() { return false; } @@ -318,11 +433,16 @@ impl TlsFrontCache { } #[cfg(test)] - async fn full_cert_sent_contains_for_tests(&self, client_ip: IpAddr) -> bool { - self.full_cert_sent_shard(client_ip) + async fn full_cert_sent_contains_for_tests(&self, domain: &str, client_ip: IpAddr) -> bool { + let key = FullCertBudgetKey { + domain: self.full_cert_domain_key(domain), + client_ip, + }; + let shard_index = self.full_cert_budget.shard_index(&key); + self.full_cert_budget.shards[shard_index] .read() .await - .contains_key(&client_ip) + .contains_key(&key) } pub async fn set(&self, domain: &str, data: CachedTlsData) { @@ -336,54 +456,60 @@ impl TlsFrontCache { return; } let mut loaded = 0usize; - if let Ok(mut dir) = tokio::fs::read_dir(&path).await { - while let Ok(Some(entry)) = dir.next_entry().await { - if let Ok(name) = entry.file_name().into_string() { - if !name.ends_with(".json") { - continue; - } - if let Ok(data) = tokio::fs::read(entry.path()).await - && let Ok(mut cached) = serde_json::from_slice::(&data) - { - if cached.domain.is_empty() - || cached.domain.len() > 255 - || !cached - .domain - .chars() - .all(|c| c.is_ascii_alphanumeric() || c == '.' || c == '-') - { - warn!(file = %name, "Skipping TLS cache entry with invalid domain"); - continue; - } - if !cert_info_matches_domain(&cached) { - warn!( - file = %name, - domain = %cached.domain, - "Skipping TLS cache entry with mismatched certificate metadata" - ); - continue; - } - // fetched_at is skipped during deserialization; approximate with file mtime if available. - if let Ok(meta) = entry.metadata().await - && let Ok(modified) = meta.modified() - { - cached.fetched_at = modified; - } - // Drop entries older than 72h - if let Ok(age) = cached.fetched_at.elapsed() - && age > Duration::from_secs(72 * 3600) - { - warn!(domain = %cached.domain, "Skipping stale TLS cache entry (>72h)"); - continue; - } - cached - .behavior_profile - .refresh_server_hello_summary(&cached.server_hello_template); - let domain = cached.domain.clone(); - self.set(&domain, cached).await; - loaded += 1; - } + for name in &self.disk_entry_names { + let entry_path = path.join(name); + let Ok(metadata) = tokio::fs::symlink_metadata(&entry_path).await else { + continue; + }; + if !metadata.file_type().is_file() { + continue; + } + if let Ok(data) = read_disk_entry_bounded(&entry_path).await + && let Ok(mut cached) = serde_json::from_slice::(&data) + { + if cached.domain.is_empty() + || cached.domain.len() > 255 + || !cached + .domain + .chars() + .all(|c| c.is_ascii_alphanumeric() || c == '.' || c == '-') + { + warn!(file = %name, "Skipping TLS cache entry with invalid domain"); + continue; } + if !self.full_cert_domain_keys.contains_key(&cached.domain) { + warn!( + file = %name, + domain = %cached.domain, + "Skipping TLS cache entry outside configured domains" + ); + continue; + } + if !cert_info_matches_domain(&cached) { + warn!( + file = %name, + domain = %cached.domain, + "Skipping TLS cache entry with mismatched certificate metadata" + ); + continue; + } + // fetched_at is skipped during deserialization; approximate with file mtime if available. + if let Ok(modified) = metadata.modified() { + cached.fetched_at = modified; + } + // Drop entries older than 72h + if let Ok(age) = cached.fetched_at.elapsed() + && age > Duration::from_secs(72 * 3600) + { + warn!(domain = %cached.domain, "Skipping stale TLS cache entry (>72h)"); + continue; + } + cached + .behavior_profile + .refresh_server_hello_summary(&cached.server_hello_template); + let domain = cached.domain.clone(); + self.set(&domain, cached).await; + loaded += 1; } } if loaded > 0 { @@ -398,6 +524,14 @@ impl TlsFrontCache { let fname = format!("{}.json", domain.replace(['/', '\\'], "_")); let path = self.disk_path.join(fname); if let Ok(json) = serde_json::to_vec_pretty(data) { + if json.len() as u64 > TLS_FRONT_DISK_ENTRY_MAX_BYTES { + warn!( + domain, + bytes = json.len(), + "Skipping oversized TLS cache persistence" + ); + return; + } // best-effort write let _ = tokio::fs::write(path, json).await; } @@ -464,6 +598,27 @@ impl TlsFrontCache { } } +async fn read_disk_entry_bounded(path: &Path) -> std::io::Result> { + let file = tokio::fs::File::open(path).await?; + if file.metadata().await?.len() > TLS_FRONT_DISK_ENTRY_MAX_BYTES { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "TLS cache entry exceeds the 1 MiB limit", + )); + } + let mut bytes = Vec::new(); + file.take(TLS_FRONT_DISK_ENTRY_MAX_BYTES.saturating_add(1)) + .read_to_end(&mut bytes) + .await?; + if bytes.len() as u64 > TLS_FRONT_DISK_ENTRY_MAX_BYTES { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "TLS cache entry grew beyond the 1 MiB limit while reading", + )); + } + Ok(bytes) +} + fn cert_info_matches_domain(cached: &CachedTlsData) -> bool { let Some(cert_info) = cached.cert_info.as_ref() else { return true; @@ -581,12 +736,24 @@ mod tests { let ip: IpAddr = "127.0.0.1".parse().expect("ip"); let ttl = Duration::from_millis(80); - assert!(cache.take_full_cert_budget_for_ip(ip, ttl).await); - assert!(!cache.take_full_cert_budget_for_ip(ip, ttl).await); + assert!( + cache + .take_full_cert_budget_for_ip("example.com", ip, ttl) + .await + ); + assert!( + !cache + .take_full_cert_budget_for_ip("example.com", ip, ttl) + .await + ); tokio::time::sleep(Duration::from_millis(90)).await; - assert!(cache.take_full_cert_budget_for_ip(ip, ttl).await); + assert!( + cache + .take_full_cert_budget_for_ip("example.com", ip, ttl) + .await + ); } #[tokio::test] @@ -601,7 +768,11 @@ mod tests { ((idx >> 8) & 0xff) as u8, (idx & 0xff) as u8, )); - assert!(cache.take_full_cert_budget_for_ip(ip, ttl).await); + assert!( + cache + .take_full_cert_budget_for_ip("example.com", ip, ttl) + .await + ); } assert!(cache.full_cert_sent_is_empty_for_tests().await); @@ -613,21 +784,42 @@ mod tests { let stale_ip: IpAddr = "127.0.0.1".parse().expect("ip"); let new_ip: IpAddr = "127.0.0.2".parse().expect("ip"); let ttl = Duration::from_secs(1); - let stale_seen_at = Instant::now() - .checked_sub(Duration::from_secs(10)) + let stale_expires_at = Instant::now() + .checked_sub(Duration::from_secs(1)) .unwrap_or_else(Instant::now); cache - .insert_full_cert_sent_for_tests(stale_ip, stale_seen_at) + .insert_full_cert_sent_for_tests("example.com", stale_ip, stale_expires_at) .await; + let stale_key = FullCertBudgetKey { + domain: cache.full_cert_domain_key("example.com"), + client_ip: stale_ip, + }; + cache.full_cert_budget.sweep_cursor.store( + cache.full_cert_budget.shard_index(&stale_key), + Ordering::Relaxed, + ); cache - .full_cert_sent_last_sweep_epoch_secs + .full_cert_budget + .last_sweep_epoch_secs .store(0, Ordering::Relaxed); - assert!(cache.take_full_cert_budget_for_ip(new_ip, ttl).await); + assert!( + cache + .take_full_cert_budget_for_ip("example.com", new_ip, ttl) + .await + ); - assert!(!cache.full_cert_sent_contains_for_tests(stale_ip).await); - assert!(cache.full_cert_sent_contains_for_tests(new_ip).await); + assert!( + !cache + .full_cert_sent_contains_for_tests("example.com", stale_ip) + .await + ); + assert!( + cache + .full_cert_sent_contains_for_tests("example.com", new_ip) + .await + ); } #[tokio::test] @@ -636,8 +828,8 @@ mod tests { let stale_ip: IpAddr = "127.0.0.1".parse().expect("ip"); let new_ip: IpAddr = "127.0.0.2".parse().expect("ip"); let ttl = Duration::from_secs(1); - let stale_seen_at = Instant::now() - .checked_sub(Duration::from_secs(10)) + let stale_expires_at = Instant::now() + .checked_sub(Duration::from_secs(1)) .unwrap_or_else(Instant::now); let now_epoch_secs = SystemTime::now() .duration_since(UNIX_EPOCH) @@ -645,15 +837,126 @@ mod tests { .as_secs(); cache - .insert_full_cert_sent_for_tests(stale_ip, stale_seen_at) + .insert_full_cert_sent_for_tests("example.com", stale_ip, stale_expires_at) .await; cache - .full_cert_sent_last_sweep_epoch_secs + .full_cert_budget + .last_sweep_epoch_secs .store(now_epoch_secs, Ordering::Relaxed); - assert!(cache.take_full_cert_budget_for_ip(new_ip, ttl).await); + assert!( + cache + .take_full_cert_budget_for_ip("example.com", new_ip, ttl) + .await + ); - assert!(cache.full_cert_sent_contains_for_tests(stale_ip).await); - assert!(cache.full_cert_sent_contains_for_tests(new_ip).await); + assert!( + cache + .full_cert_sent_contains_for_tests("example.com", stale_ip) + .await + ); + assert!( + cache + .full_cert_sent_contains_for_tests("example.com", new_ip) + .await + ); + } + + #[tokio::test] + async fn full_cert_budget_is_shared_across_cache_generations_and_scoped_by_domain() { + let budget = Arc::new(TlsFullCertBudget::new()); + let domains = ["one.example".to_string(), "two.example".to_string()]; + let first = TlsFrontCache::new_with_full_cert_budget( + &domains, + 1024, + "tlsfront-test-cache", + budget.clone(), + ); + let second = TlsFrontCache::new_with_full_cert_budget( + &domains, + 1024, + "tlsfront-test-cache", + budget, + ); + let ip: IpAddr = "127.0.0.1".parse().expect("ip"); + let ttl = Duration::from_secs(60); + + assert!( + first + .take_full_cert_budget_for_ip("one.example", ip, ttl) + .await + ); + assert!( + !second + .take_full_cert_budget_for_ip("one.example", ip, ttl) + .await + ); + assert!( + second + .take_full_cert_budget_for_ip("two.example", ip, ttl) + .await + ); + assert_eq!(second.full_cert_budget_entries_for_metrics(), 2); + } + + #[tokio::test] + async fn existing_full_cert_entry_keeps_its_own_expiry_after_ttl_change() { + let cache = TlsFrontCache::new(&["example.com".to_string()], 1024, "tlsfront-test-cache"); + let ip: IpAddr = "127.0.0.1".parse().expect("ip"); + + assert!( + cache + .take_full_cert_budget_for_ip("example.com", ip, Duration::from_millis(80)) + .await + ); + tokio::time::sleep(Duration::from_millis(20)).await; + assert!( + !cache + .take_full_cert_budget_for_ip("example.com", ip, Duration::from_millis(1)) + .await + ); + tokio::time::sleep(Duration::from_millis(70)).await; + assert!( + cache + .take_full_cert_budget_for_ip("example.com", ip, Duration::from_millis(1)) + .await + ); + } + + #[tokio::test] + async fn disk_reader_rejects_an_entry_above_the_hard_limit() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("oversized.json"); + tokio::fs::write( + &path, + vec![0u8; TLS_FRONT_DISK_ENTRY_MAX_BYTES as usize + 1], + ) + .await + .unwrap(); + + let error = read_disk_entry_bounded(&path).await.unwrap_err(); + + assert_eq!(error.kind(), std::io::ErrorKind::InvalidData); + } + + #[cfg(unix)] + #[tokio::test] + async fn disk_loader_does_not_follow_a_configured_name_symlink() { + let directory = tempfile::tempdir().unwrap(); + let target = directory.path().join("outside.json"); + let cached = cached_with_cert_info("example.com", None, Vec::new()); + tokio::fs::write(&target, serde_json::to_vec(&cached).unwrap()) + .await + .unwrap(); + std::os::unix::fs::symlink(&target, directory.path().join("example.com.json")).unwrap(); + let cache = TlsFrontCache::new( + &["example.com".to_string()], + 1024, + directory.path(), + ); + + cache.load_from_disk().await; + + assert_eq!(cache.get("example.com").await.domain, "default"); } } diff --git a/src/tls_front/fetcher.rs b/src/tls_front/fetcher.rs index 3bcc3c5..e6bf81c 100644 --- a/src/tls_front/fetcher.rs +++ b/src/tls_front/fetcher.rs @@ -30,7 +30,6 @@ use x509_parser::prelude::FromDer; use crate::config::TlsFetchProfile; use crate::crypto::{SecureRandom, sha256}; -use crate::network::dns_overrides::resolve_socket_addr; use crate::protocol::constants::{ TLS_RECORD_APPLICATION, TLS_RECORD_CHANGE_CIPHER, TLS_RECORD_HANDSHAKE, }; @@ -901,9 +900,6 @@ async fn connect_with_dns_override( port: u16, connect_timeout: Duration, ) -> Result { - if let Some(addr) = resolve_socket_addr(host, port) { - return Ok(timeout(connect_timeout, TcpStream::connect(addr)).await??); - } Ok(timeout(connect_timeout, TcpStream::connect((host, port))).await??) } diff --git a/src/transport/middle_proxy/config_updater.rs b/src/transport/middle_proxy/config_updater.rs index b777445..070e020 100644 --- a/src/transport/middle_proxy/config_updater.rs +++ b/src/transport/middle_proxy/config_updater.rs @@ -14,7 +14,7 @@ use crate::error::Result; use crate::transport::UpstreamManager; use super::MePool; -use super::http_fetch::https_get; +use super::http_fetch::{HTTPS_RESPONSE_BODY_MAX_BYTES, https_get}; use super::rotation::{MeReinitTrigger, enqueue_reinit_trigger}; use super::secret::download_proxy_secret_with_max_len_via_upstream; use super::selftest::record_timeskew_sample; @@ -106,7 +106,7 @@ pub async fn fetch_proxy_config_with_raw_via_upstream( url: &str, upstream: Option>, ) -> Result<(ProxyConfigData, String)> { - let resp = https_get(url, upstream).await?; + let resp = https_get(url, upstream, HTTPS_RESPONSE_BODY_MAX_BYTES).await?; let http_status = resp.status; if let Some(date_str) = resp.date_header.as_deref() diff --git a/src/transport/middle_proxy/health.rs b/src/transport/middle_proxy/health.rs index b288f5b..d40430c 100644 --- a/src/transport/middle_proxy/health.rs +++ b/src/transport/middle_proxy/health.rs @@ -71,6 +71,40 @@ struct FamilyReconnectOutcome { endpoint_count: usize, } +struct ScheduledReconnects<'a> { + inflight: &'a mut HashMap<(i32, IpFamily), usize>, + keys: Vec<(i32, IpFamily)>, +} + +impl ScheduledReconnects<'_> { + fn current(&self, key: &(i32, IpFamily)) -> usize { + self.inflight.get(key).copied().unwrap_or(0) + } + + fn reserve(&mut self, key: (i32, IpFamily)) { + self.keys.push(key); + *self.inflight.entry(key).or_insert(0) += 1; + } +} + +impl Drop for ScheduledReconnects<'_> { + fn drop(&mut self) { + for key in self.keys.drain(..) { + let std::collections::hash_map::Entry::Occupied(mut entry) = + self.inflight.entry(key) + else { + continue; + }; + let remaining = entry.get().saturating_sub(1); + if remaining == 0 { + entry.remove(); + } else { + *entry.get_mut() = remaining; + } + } + } +} + pub async fn me_health_monitor(pool: Arc, rng: Arc, _min_connections: usize) { let mut backoff: HashMap<(i32, IpFamily), u64> = HashMap::new(); let mut next_attempt: HashMap<(i32, IpFamily), Instant> = HashMap::new(); @@ -437,6 +471,10 @@ async fn check_family( let writer_idle_since = Arc::new(writer_idle_since); let bound_clients_by_writer = Arc::new(bound_clients_by_writer); let mut reconnect_set = JoinSet::::new(); + let mut scheduled_reconnects = ScheduledReconnects { + inflight, + keys: Vec::new(), + }; for (dc, endpoints) in dc_endpoints { if endpoints.is_empty() { @@ -562,7 +600,7 @@ async fn check_family( .reconnect_runtime .me_reconnect_max_concurrent_per_dc .max(1) as usize; - if *inflight.get(&key).unwrap_or(&0) >= max_concurrent { + if scheduled_reconnects.current(&key) >= max_concurrent { continue; } if pool @@ -579,7 +617,7 @@ async fn check_family( ); continue; } - *inflight.entry(key).or_insert(0) += 1; + scheduled_reconnects.reserve(key); let pool_for_reconnect = pool.clone(); let rng_for_reconnect = rng.clone(); let reconnect_sem_for_dc = reconnect_sem.clone(); @@ -740,9 +778,6 @@ async fn check_family( ); } } - if let Some(v) = inflight.get_mut(&outcome.key) { - *v = v.saturating_sub(1); - } } family_degraded @@ -1701,15 +1736,35 @@ mod tests { use tokio::sync::mpsc; use tokio_util::sync::CancellationToken; - use super::reap_draining_writers; + use super::{ScheduledReconnects, reap_draining_writers}; use crate::config::{GeneralConfig, MeRouteNoWriterMode, MeSocksKdfPolicy, MeWriterPickMode}; use crate::crypto::SecureRandom; use crate::network::probe::NetworkDecision; + use crate::network::IpFamily; use crate::stats::Stats; use crate::transport::middle_proxy::codec::WriterCommand; use crate::transport::middle_proxy::pool::{MePool, MeWriter, WriterContour}; use crate::transport::middle_proxy::registry::ConnMeta; + #[test] + fn reconnect_batch_releases_every_reserved_key_after_join_failures() { + let retained = (1, IpFamily::V4); + let removed = (2, IpFamily::V6); + let mut inflight = HashMap::from([(retained, 1)]); + + { + let mut scheduled = ScheduledReconnects { + inflight: &mut inflight, + keys: Vec::new(), + }; + scheduled.reserve(retained); + scheduled.reserve(removed); + } + + assert_eq!(inflight.get(&retained), Some(&1)); + assert!(!inflight.contains_key(&removed)); + } + async fn make_pool(me_pool_drain_threshold: u64) -> Arc { let general = GeneralConfig { me_pool_drain_threshold, @@ -1812,6 +1867,7 @@ mod tests { general.me_route_blocking_send_timeout_ms, general.me_route_inline_recovery_attempts, general.me_route_inline_recovery_wait_ms, + 16_384, ) } diff --git a/src/transport/middle_proxy/http_fetch.rs b/src/transport/middle_proxy/http_fetch.rs index 33e48e2..938dfac 100644 --- a/src/transport/middle_proxy/http_fetch.rs +++ b/src/transport/middle_proxy/http_fetch.rs @@ -1,7 +1,7 @@ use std::sync::Arc; use std::time::Duration; -use http_body_util::{BodyExt, Empty}; +use http_body_util::{BodyExt, Empty, Limited}; use hyper::header::{CONNECTION, DATE, HOST, USER_AGENT}; use hyper::{Method, Request}; use hyper_util::rt::TokioIo; @@ -12,11 +12,19 @@ use tokio_rustls::TlsConnector; use tracing::debug; use crate::error::{ProxyError, Result}; -use crate::network::dns_overrides::resolve_socket_addr; use crate::transport::{UpstreamManager, UpstreamStream}; const HTTP_CONNECT_TIMEOUT: Duration = Duration::from_secs(10); const HTTP_REQUEST_TIMEOUT: Duration = Duration::from_secs(15); +pub(super) const HTTPS_RESPONSE_BODY_MAX_BYTES: usize = 1024 * 1024; + +struct HttpsConnectionDriver(tokio::task::JoinHandle<()>); + +impl Drop for HttpsConnectionDriver { + fn drop(&mut self) { + self.0.abort(); + } +} pub(crate) struct HttpsGetResponse { pub(crate) status: u16, @@ -81,14 +89,6 @@ async fn connect_https_transport( }); } - if let Some(addr) = resolve_socket_addr(host, port) { - let stream = timeout(HTTP_CONNECT_TIMEOUT, TcpStream::connect(addr)) - .await - .map_err(|_| ProxyError::Proxy(format!("connect timeout for {host}:{port}")))? - .map_err(|e| ProxyError::Proxy(format!("connect failed for {host}:{port}: {e}")))?; - return Ok(UpstreamStream::Tcp(stream)); - } - let stream = timeout(HTTP_CONNECT_TIMEOUT, TcpStream::connect((host, port))) .await .map_err(|_| ProxyError::Proxy(format!("connect timeout for {host}:{port}")))? @@ -99,7 +99,14 @@ async fn connect_https_transport( pub(crate) async fn https_get( url: &str, upstream: Option>, + max_body_bytes: usize, ) -> Result { + if max_body_bytes == 0 { + return Err(ProxyError::Proxy( + "HTTPS response body limit must be greater than zero".to_string(), + )); + } + let max_body_bytes = max_body_bytes.min(HTTPS_RESPONSE_BODY_MAX_BYTES); let (host, port, path_and_query) = extract_host_port_path(url)?; let stream = connect_https_transport(&host, port, upstream).await?; @@ -115,11 +122,11 @@ pub(crate) async fn https_get( .await .map_err(|e| ProxyError::Proxy(format!("HTTP handshake failed for {host}:{port}: {e}")))?; - tokio::spawn(async move { + let _connection_driver = HttpsConnectionDriver(tokio::spawn(async move { if let Err(e) = connection.await { debug!(error = %e, "HTTPS fetch connection task failed"); } - }); + })); let host_header = if port == 443 { host.clone() @@ -148,10 +155,17 @@ pub(crate) async fn https_get( .and_then(|value| value.to_str().ok()) .map(|value| value.to_string()); - let body = timeout(HTTP_REQUEST_TIMEOUT, response.into_body().collect()) + let body = timeout( + HTTP_REQUEST_TIMEOUT, + Limited::new(response.into_body(), max_body_bytes).collect(), + ) .await .map_err(|_| ProxyError::Proxy(format!("HTTP body read timeout for {url}")))? - .map_err(|e| ProxyError::Proxy(format!("HTTP body read failed for {url}: {e}")))? + .map_err(|e| { + ProxyError::Proxy(format!( + "HTTP body read failed or exceeded {max_body_bytes} bytes for {url}: {e}" + )) + })? .to_bytes() .to_vec(); @@ -161,3 +175,16 @@ pub(crate) async fn https_get( body, }) } + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn limited_body_rejects_payload_above_caller_bound() { + let body = http_body_util::Full::new(bytes::Bytes::from_static(b"12345")); + let result = Limited::new(body, 4).collect().await; + + assert!(result.is_err()); + } +} diff --git a/src/transport/middle_proxy/mod.rs b/src/transport/middle_proxy/mod.rs index 3f46a80..feb658d 100644 --- a/src/transport/middle_proxy/mod.rs +++ b/src/transport/middle_proxy/mod.rs @@ -22,6 +22,7 @@ mod ping; mod pool; mod pool_config; mod pool_init; +mod pool_lifecycle; mod pool_nat; mod pool_refill; #[cfg(test)] @@ -60,6 +61,7 @@ pub use ping::{ MePingFamily, MePingReport, MePingSample, format_me_route, format_sample_line, run_me_ping, }; pub use pool::MePool; +pub(crate) use registry::ConnLease; #[allow(unused_imports)] pub use pool_nat::{detect_public_ip, stun_probe}; pub use registry::ConnRegistry; diff --git a/src/transport/middle_proxy/pool.rs b/src/transport/middle_proxy/pool.rs index 69754de..30e6d94 100644 --- a/src/transport/middle_proxy/pool.rs +++ b/src/transport/middle_proxy/pool.rs @@ -10,6 +10,7 @@ use std::sync::atomic::{ use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; use arc_swap::ArcSwap; +use parking_lot::Mutex as ParkingMutex; use tokio::sync::{Mutex, RwLock, Semaphore, mpsc, watch}; use tokio_util::sync::CancellationToken; @@ -23,6 +24,7 @@ use crate::transport::UpstreamManager; use super::ConnRegistry; use super::codec::WriterCommand; +use super::pool_lifecycle::MePoolLifecycle; const ME_FORCE_CLOSE_SAFETY_FALLBACK_SECS: u64 = 300; @@ -32,12 +34,6 @@ pub(super) struct RefillDcKey { pub family: IpFamily, } -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] -pub(super) struct RefillEndpointKey { - pub dc: i32, - pub addr: SocketAddr, -} - #[derive(Clone)] pub struct MeWriter { pub id: u64, @@ -148,6 +144,18 @@ pub(super) enum WriterContour { Draining = 2, } +pub(super) struct WriterOpenReservation<'a> { + counter: Option<&'a AtomicUsize>, +} + +impl Drop for WriterOpenReservation<'_> { + fn drop(&mut self) { + if let Some(counter) = self.counter { + counter.fetch_sub(1, Ordering::AcqRel); + } + } +} + impl WriterContour { pub(super) fn as_u8(self) -> u8 { self as u8 @@ -265,6 +273,10 @@ pub(super) struct ReinitCore { pub(super) pending_hardswap_generation: AtomicU64, pub(super) pending_hardswap_started_at_epoch_secs: AtomicU64, pub(super) pending_hardswap_map_hash: AtomicU64, + pub(super) scheduler_inflight: AtomicUsize, + pub(super) max_concurrency_effective: AtomicUsize, + pub(super) coordinator: ParkingMutex, + pub(super) status: ArcSwap, pub(super) hardswap: AtomicBool, pub(super) me_hardswap_warmup_delay_min_ms: AtomicU64, pub(super) me_hardswap_warmup_delay_max_ms: AtomicU64, @@ -272,6 +284,39 @@ pub(super) struct ReinitCore { pub(super) me_hardswap_warmup_pass_backoff_base_ms: AtomicU64, } +#[derive(Clone, Debug)] +pub(super) struct ReinitStatusSnapshot { + pub(super) active_generation: u64, + pub(super) warm_generations: Vec, + pub(super) pending_hardswap_generation: u64, + pub(super) pending_hardswap_started_at_epoch_secs: u64, + pub(super) pending_hardswap_map_hash: u64, + pub(super) inflight: usize, +} + +#[derive(Clone, Copy)] +pub(super) struct ReinitPendingState { + pub(super) generation: u64, + pub(super) started_at_epoch_secs: u64, + pub(super) map_hash: u64, +} + +#[derive(Clone, Copy)] +pub(super) struct ReinitAttemptState { + pub(super) generation: u64, + pub(super) map_hash: u64, + pub(super) hardswap: bool, + pub(super) committed: bool, +} + +pub(super) struct ReinitCoordinatorState { + pub(super) next_attempt_id: u64, + pub(super) active_generation: u64, + pub(super) desired_map_hash: u64, + pub(super) pending: Option, + pub(super) attempts: HashMap, +} + pub(super) struct WriterLifecycleCore { pub(super) me_keepalive_enabled: bool, pub(super) me_keepalive_interval: Duration, @@ -421,6 +466,7 @@ pub struct MePool { pub(super) floor_runtime: Arc, pub(super) writer_selection_policy: Arc, pub(super) transport_policy: Arc, + pub(super) lifecycle: MePoolLifecycle, pub(super) decision: NetworkDecision, pub(super) upstream: Option>, pub(super) rng: Arc, @@ -431,9 +477,12 @@ pub struct MePool { pub(super) endpoint_dc_map: Arc>>>, pub(super) default_dc: AtomicI32, pub(super) next_writer_id: AtomicU64, + pub(super) writer_connect_active_reserved: AtomicUsize, + pub(super) writer_connect_warm_reserved: AtomicUsize, pub(super) rtt_stats: Arc>>, - pub(super) refill_inflight: Arc>>, - pub(super) refill_inflight_dc: Arc>>, + pub(super) refill_states: Arc>>>, + pub(super) refill_running: AtomicUsize, + pub(super) refill_pending: AtomicUsize, pub(super) conn_count: AtomicUsize, pub(super) draining_active_runtime: AtomicU64, pub(super) stats: Arc, @@ -573,12 +622,14 @@ impl MePool { me_route_blocking_send_timeout_ms: u64, me_route_inline_recovery_attempts: u32, me_route_inline_recovery_wait_ms: u64, + me_connection_cleanup_capacity: usize, ) -> Arc { let endpoint_dc_map = Self::build_endpoint_dc_map_from_maps(&proxy_map_v4, &proxy_map_v6); let preferred_endpoints_by_dc = Self::build_preferred_endpoints_by_dc(&decision, &proxy_map_v4, &proxy_map_v6); - let registry = Arc::new(ConnRegistry::with_route_channel_capacity( + let registry = Arc::new(ConnRegistry::with_route_and_cleanup_capacity( me_route_channel_capacity, + me_connection_cleanup_capacity, )); registry.update_route_backpressure_policy( me_route_backpressure_base_timeout_ms, @@ -587,6 +638,14 @@ impl MePool { ); let (writer_epoch, _) = watch::channel(0u64); let now_epoch_secs = Self::now_epoch_secs(); + let reinit_status = ReinitStatusSnapshot { + active_generation: 1, + warm_generations: Vec::new(), + pending_hardswap_generation: 0, + pending_hardswap_started_at_epoch_secs: 0, + pending_hardswap_map_hash: 0, + inflight: 0, + }; stats.set_me_writer_byte_budget_limit_bytes(me_writer_byte_budget_bytes); Arc::new(Self { routing: Arc::new(RoutingCore { @@ -603,6 +662,16 @@ impl MePool { pending_hardswap_generation: AtomicU64::new(0), pending_hardswap_started_at_epoch_secs: AtomicU64::new(0), pending_hardswap_map_hash: AtomicU64::new(0), + scheduler_inflight: AtomicUsize::new(0), + max_concurrency_effective: AtomicUsize::new(1), + coordinator: ParkingMutex::new(ReinitCoordinatorState { + next_attempt_id: 1, + active_generation: 1, + desired_map_hash: 0, + pending: None, + attempts: HashMap::new(), + }), + status: ArcSwap::from_pointee(reinit_status), hardswap: AtomicBool::new(hardswap), me_hardswap_warmup_delay_min_ms: AtomicU64::new(me_hardswap_warmup_delay_min_ms), me_hardswap_warmup_delay_max_ms: AtomicU64::new(me_hardswap_warmup_delay_max_ms), @@ -805,6 +874,7 @@ impl MePool { me_reader_route_data_wait_ms, )), }), + lifecycle: MePoolLifecycle::new(), decision, upstream, rng, @@ -830,9 +900,12 @@ impl MePool { endpoint_dc_map: Arc::new(RwLock::new(endpoint_dc_map)), default_dc: AtomicI32::new(default_dc.unwrap_or(2)), next_writer_id: AtomicU64::new(1), + writer_connect_active_reserved: AtomicUsize::new(0), + writer_connect_warm_reserved: AtomicUsize::new(0), rtt_stats: Arc::new(Mutex::new(HashMap::new())), - refill_inflight: Arc::new(Mutex::new(HashSet::new())), - refill_inflight_dc: Arc::new(Mutex::new(HashSet::new())), + refill_states: Arc::new(ParkingMutex::new(HashMap::new())), + refill_running: AtomicUsize::new(0), + refill_pending: AtomicUsize::new(0), conn_count: AtomicUsize::new(0), draining_active_runtime: AtomicU64::new(0), endpoint_quarantine: Arc::new(Mutex::new(HashMap::new())), @@ -1263,6 +1336,7 @@ impl MePool { self.translate_our_addr_with_reflection(addr, None) } + #[allow(dead_code)] pub fn registry(&self) -> &Arc { &self.registry } @@ -1733,6 +1807,69 @@ impl MePool { } } + pub(super) async fn reserve_writer_open( + &self, + contour: WriterContour, + allow_coverage_override: bool, + writer_dc: i32, + ) -> Option> { + let counter = match contour { + WriterContour::Active => &self.writer_connect_active_reserved, + WriterContour::Warm => &self.writer_connect_warm_reserved, + WriterContour::Draining => { + return Some(WriterOpenReservation { counter: None }); + } + }; + + loop { + if !self + .can_open_writer_for_contour(contour, allow_coverage_override, writer_dc) + .await + { + return None; + } + let (active_writers, warm_writers, _) = + self.non_draining_writer_counts_by_contour().await; + let live = match contour { + WriterContour::Active => active_writers, + WriterContour::Warm => warm_writers, + WriterContour::Draining => 0, + }; + let mut limit = match contour { + WriterContour::Active => self.adaptive_floor_active_cap_configured_total(), + WriterContour::Warm => self.adaptive_floor_warm_cap_configured_total(), + WriterContour::Draining => usize::MAX, + }; + if contour == WriterContour::Active && allow_coverage_override { + limit = limit + .max(self.active_coverage_required_total().await) + .saturating_add( + self.reconnect_runtime + .me_reconnect_max_concurrent_per_dc + .max(1) as usize, + ); + } + + let reserved = counter.load(Ordering::Acquire); + if live.saturating_add(reserved) >= limit { + return None; + } + if counter + .compare_exchange_weak( + reserved, + reserved + 1, + Ordering::AcqRel, + Ordering::Acquire, + ) + .is_ok() + { + return Some(WriterOpenReservation { + counter: Some(counter), + }); + } + } + } + pub(super) fn required_writers_for_dc_with_floor_mode( &self, endpoint_count: usize, diff --git a/src/transport/middle_proxy/pool_init.rs b/src/transport/middle_proxy/pool_init.rs index 3eee903..0936246 100644 --- a/src/transport/middle_proxy/pool_init.rs +++ b/src/transport/middle_proxy/pool_init.rs @@ -107,7 +107,7 @@ impl MePool { let pool = Arc::clone(self); let rng_clone = Arc::clone(rng); let dc_addrs_bg = dc_addrs.clone(); - tokio::spawn(async move { + let saturation = async move { let mut join_bg = tokio::task::JoinSet::new(); for (dc, addrs) in dc_addrs_bg { if addrs.len() <= 1 { @@ -135,7 +135,10 @@ impl MePool { current_pool_size = pool.connection_count(), "Background ME saturation warmup finished" ); - }); + }; + if self.lifecycle.spawn_producer(saturation).is_err() { + debug!("Background ME saturation skipped: pool lifecycle closed"); + } if !self.decision.effective_multipath && self.connection_count() > 0 { break; diff --git a/src/transport/middle_proxy/pool_lifecycle.rs b/src/transport/middle_proxy/pool_lifecycle.rs new file mode 100644 index 0000000..c4eb73a --- /dev/null +++ b/src/transport/middle_proxy/pool_lifecycle.rs @@ -0,0 +1,326 @@ +use std::future::Future; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::time::Duration; + +use tokio::sync::Notify; +use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; +use tokio_util::task::TaskTracker; + +use super::pool::MePool; +use super::registry::ConnLease; + +const ME_TASK_ADMISSION_CLOSED: usize = 1 << (usize::BITS - 1); +const ME_TASK_REGISTRATION_COUNT: usize = ME_TASK_ADMISSION_CLOSED - 1; + +struct MeTaskAdmission { + state: AtomicUsize, + registrations_drained: Notify, +} + +/// RAII ownership of one ME task-publication section. +pub(super) struct MeTaskRegistration<'a> { + admission: &'a MeTaskAdmission, +} + +impl MeTaskAdmission { + fn new() -> Self { + Self { + state: AtomicUsize::new(0), + registrations_drained: Notify::new(), + } + } + + fn try_register(&self) -> Option> { + let mut state = self.state.load(Ordering::Acquire); + loop { + if state & ME_TASK_ADMISSION_CLOSED != 0 + || state & ME_TASK_REGISTRATION_COUNT == ME_TASK_REGISTRATION_COUNT + { + return None; + } + match self.state.compare_exchange_weak( + state, + state + 1, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => return Some(MeTaskRegistration { admission: self }), + Err(observed) => state = observed, + } + } + } + + fn close(&self) { + self.state + .fetch_or(ME_TASK_ADMISSION_CLOSED, Ordering::AcqRel); + } + + async fn wait_for_registrations(&self) { + loop { + let notified = self.registrations_drained.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); + if self.state.load(Ordering::Acquire) & ME_TASK_REGISTRATION_COUNT == 0 { + return; + } + notified.await; + } + } +} + +impl Drop for MeTaskRegistration<'_> { + fn drop(&mut self) { + let previous = self.admission.state.fetch_sub(1, Ordering::AcqRel); + if previous & ME_TASK_REGISTRATION_COUNT == 1 { + self.admission.registrations_drained.notify_waiters(); + } + } +} + +/// Pool-owned admission, cancellation, and join authority for ME tasks. +pub(super) struct MePoolLifecycle { + admission: MeTaskAdmission, + producer_cancel: CancellationToken, + producer_tasks: TaskTracker, + writer_tasks: TaskTracker, + cleanup_cancel: CancellationToken, + cleanup_tasks: TaskTracker, + cleanup_started: AtomicBool, + shutdown_started: AtomicBool, +} + +impl MePoolLifecycle { + /// Creates an open ME task lifecycle. + pub(super) fn new() -> Self { + Self { + admission: MeTaskAdmission::new(), + producer_cancel: CancellationToken::new(), + producer_tasks: TaskTracker::new(), + writer_tasks: TaskTracker::new(), + cleanup_cancel: CancellationToken::new(), + cleanup_tasks: TaskTracker::new(), + cleanup_started: AtomicBool::new(false), + shutdown_started: AtomicBool::new(false), + } + } + + /// Registers one task-publication section while lifecycle admission is open. + pub(super) fn try_register(&self) -> Option> { + self.admission.try_register() + } + + /// Registers and spawns one cancellation-aware ME producer. + pub(super) fn spawn_producer(&self, future: F) -> Result<(), F> + where + F: Future + Send + 'static, + { + let Some(registration) = self.try_register() else { + return Err(future); + }; + self.spawn_registered_producer(registration, future); + Ok(()) + } + + /// Spawns a producer after its caller published cancellation cleanup ownership. + pub(super) fn spawn_registered_producer( + &self, + registration: MeTaskRegistration<'_>, + future: F, + ) where + F: Future + Send + 'static, + { + let cancel = self.producer_cancel.clone(); + self.producer_tasks.spawn(async move { + tokio::select! { + biased; + _ = cancel.cancelled() => {} + _ = future => {} + } + }); + drop(registration); + } + + /// Spawns a writer after its caller completed task registration. + pub(super) fn spawn_registered_writer( + &self, + registration: MeTaskRegistration<'_>, + future: F, + ) where + F: Future + Send + 'static, + { + self.writer_tasks.spawn(future); + drop(registration); + } + + fn start_cleanup_worker(&self, pool: &Arc) -> bool { + let Some(registration) = self.try_register() else { + return false; + }; + if self + .cleanup_started + .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) + .is_err() + { + return true; + } + let Some(mut cleanup_rx) = pool.registry.take_cleanup_receiver() else { + self.cleanup_started.store(false, Ordering::Release); + return false; + }; + let registry = Arc::clone(&pool.registry); + let cancel = self.cleanup_cancel.clone(); + self.cleanup_tasks.spawn(async move { + loop { + tokio::select! { + biased; + cleanup = cleanup_rx.recv() => { + let Some(conn_id) = cleanup else { + return; + }; + registry.unregister(conn_id).await; + } + _ = cancel.cancelled() => { + while let Ok(conn_id) = cleanup_rx.try_recv() { + registry.unregister(conn_id).await; + } + return; + } + } + } + }); + drop(registration); + true + } + + /// Idempotently closes task admission and cancels ME producers. + pub(super) fn begin_shutdown(&self) { + if self.shutdown_started.swap(true, Ordering::AcqRel) { + return; + } + self.admission.close(); + self.producer_cancel.cancel(); + } + + async fn wait_until(deadline: tokio::time::Instant, future: F) -> bool + where + F: Future, + { + let now = tokio::time::Instant::now(); + if now >= deadline { + return false; + } + tokio::time::timeout_at(deadline, future).await.is_ok() + } + + /// Joins producers, writers, and cleanup ownership under one deadline. + pub(super) async fn shutdown_pool( + &self, + pool: &Arc, + timeout: Duration, + ) -> bool { + let deadline = tokio::time::Instant::now() + timeout; + self.begin_shutdown(); + pool.set_runtime_ready(false); + + let registrations_stopped = + Self::wait_until(deadline, self.admission.wait_for_registrations()).await; + + self.producer_tasks.close(); + let producers_stopped = Self::wait_until(deadline, self.producer_tasks.wait()).await; + + let close_signals_sent = if tokio::time::Instant::now() < deadline { + tokio::time::timeout_at(deadline, pool.shutdown_send_close_conn_all()) + .await + .is_ok() + } else { + false + }; + + let writers = pool.writers.snapshot(); + for writer in writers.iter() { + writer.cancel.cancel(); + } + self.writer_tasks.close(); + let writers_stopped = Self::wait_until(deadline, self.writer_tasks.wait()).await; + + self.cleanup_cancel.cancel(); + self.cleanup_tasks.close(); + let cleanup_stopped = Self::wait_until(deadline, self.cleanup_tasks.wait()).await; + + registrations_stopped + && producers_stopped + && close_signals_sent + && writers_stopped + && cleanup_stopped + } +} + +impl MePool { + /// Registers one cancellation-safe client route in the bounded cleanup plane. + pub(crate) async fn register_connection( + self: &Arc, + ) -> Option<(ConnLease, mpsc::Receiver)> { + if !self.lifecycle.start_cleanup_worker(self) { + return None; + } + let registration = self.lifecycle.try_register()?; + let registered = tokio::select! { + biased; + _ = self.lifecycle.producer_cancel.cancelled() => None, + registered = self.registry.register_leased() => registered, + }; + drop(registration); + registered + } + + /// Closes ME task admission and joins pool-owned producers and writers. + pub(crate) async fn shutdown_until(self: &Arc, timeout: Duration) -> bool { + self.lifecycle.shutdown_pool(self, timeout).await + } + + /// Terminally closes ME task admission without waiting for asynchronous teardown. + pub(crate) fn begin_shutdown(&self) { + self.lifecycle.begin_shutdown(); + self.set_runtime_ready(false); + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::sync::atomic::{AtomicBool, Ordering}; + use std::time::Duration; + + use super::MePoolLifecycle; + + #[tokio::test] + async fn shutdown_fence_rejects_late_producer_registration() { + struct DropSignal(Arc); + + impl Drop for DropSignal { + fn drop(&mut self) { + self.0.store(true, Ordering::Release); + } + } + + let lifecycle = Arc::new(MePoolLifecycle::new()); + let dropped = Arc::new(AtomicBool::new(false)); + let drop_signal = DropSignal(dropped.clone()); + let spawned = lifecycle.spawn_producer(async move { + let _drop_signal = drop_signal; + std::future::pending::<()>().await; + }); + assert!(spawned.is_ok()); + + lifecycle.begin_shutdown(); + lifecycle.producer_tasks.close(); + tokio::time::timeout(Duration::from_secs(1), lifecycle.producer_tasks.wait()) + .await + .unwrap(); + + assert!(dropped.load(Ordering::Acquire)); + assert!(lifecycle.spawn_producer(async {}).is_err()); + } +} diff --git a/src/transport/middle_proxy/pool_nat.rs b/src/transport/middle_proxy/pool_nat.rs index d9b7973..9c1a15c 100644 --- a/src/transport/middle_proxy/pool_nat.rs +++ b/src/transport/middle_proxy/pool_nat.rs @@ -9,7 +9,8 @@ use tracing::{debug, info}; use crate::error::{ProxyError, Result}; use crate::network::probe::{detect_public_ipv4_http, is_bogon}; use crate::network::stun::{ - IpFamily, stun_probe_dual_with_tcp_fallback, stun_probe_family_with_bind_and_tcp_fallback, + IpFamily, stun_probe_dual_with_tcp_fallback, + stun_probe_family_with_bind_tcp_fallback_and_resolver, }; use super::MePool; @@ -66,10 +67,12 @@ impl MePool { let mut best_by_ip: HashMap = HashMap::new(); let concurrency = self.nat_runtime.nat_probe_concurrency.max(1); let tcp_fallback = self.nat_runtime.stun_tcp_fallback; + let dns_resolver = self.upstream.as_ref().map(|manager| manager.dns_resolver()); while next_idx < servers.len() || !join_set.is_empty() { while next_idx < servers.len() && join_set.len() < concurrency { let stun_addr = servers[next_idx].clone(); + let dns_resolver = dns_resolver.clone(); next_idx += 1; join_set.spawn(async move { let batch_timeout = if tcp_fallback { @@ -79,11 +82,12 @@ impl MePool { }; let res = timeout( batch_timeout, - stun_probe_family_with_bind_and_tcp_fallback( + stun_probe_family_with_bind_tcp_fallback_and_resolver( &stun_addr, family, bind_ip, tcp_fallback, + dns_resolver.as_deref(), ), ) .await; diff --git a/src/transport/middle_proxy/pool_refill.rs b/src/transport/middle_proxy/pool_refill.rs index bb62604..4a8c914 100644 --- a/src/transport/middle_proxy/pool_refill.rs +++ b/src/transport/middle_proxy/pool_refill.rs @@ -9,13 +9,48 @@ use tracing::{debug, info, warn}; use crate::crypto::SecureRandom; use crate::network::IpFamily; -use super::pool::{MePool, RefillDcKey, RefillEndpointKey, WriterContour}; +use super::pool::{MePool, RefillDcKey, WriterContour}; const ME_FLAP_UPTIME_THRESHOLD_SECS: u64 = 20; const ME_FLAP_QUARANTINE_SECS: u64 = 25; const ME_FLAP_MIN_UPTIME_MILLIS: u64 = 500; const ME_REFILL_TOTAL_ATTEMPT_CAP: u32 = 20; +struct RefillRunGuard { + pool: Arc, + key: RefillDcKey, + active: bool, +} + +impl RefillRunGuard { + fn next_or_finish(&mut self) -> Option { + let mut states = self.pool.refill_states.lock(); + let next = states.get_mut(&self.key).and_then(Option::take); + if next.is_some() { + self.pool.refill_pending.fetch_sub(1, Ordering::AcqRel); + return next; + } + states.remove(&self.key); + self.pool.refill_running.fetch_sub(1, Ordering::AcqRel); + self.active = false; + None + } +} + +impl Drop for RefillRunGuard { + fn drop(&mut self) { + if !self.active { + return; + } + if let Some(pending) = self.pool.refill_states.lock().remove(&self.key) + && pending.is_some() + { + self.pool.refill_pending.fetch_sub(1, Ordering::AcqRel); + } + self.pool.refill_running.fetch_sub(1, Ordering::AcqRel); + } +} + impl MePool { pub(super) async fn sweep_endpoint_quarantine(&self) { let configured = self @@ -131,8 +166,7 @@ impl MePool { } pub(super) async fn has_refill_inflight_for_dc_key(&self, key: RefillDcKey) -> bool { - let guard = self.refill_inflight_dc.lock().await; - guard.contains(&key) + self.refill_states.lock().contains_key(&key) } pub(super) async fn connect_endpoints_round_robin( @@ -327,62 +361,53 @@ impl MePool { addr: SocketAddr, writer_dc: i32, ) { - let endpoint_key = RefillEndpointKey { - dc: writer_dc, - addr, + let Some(registration) = self.lifecycle.try_register() else { + return; }; - let pre_inserted = if let Ok(mut guard) = self.refill_inflight.try_lock() { - if !guard.insert(endpoint_key) { + let dc_key = RefillDcKey { + dc: writer_dc, + family: if addr.is_ipv4() { + IpFamily::V4 + } else { + IpFamily::V6 + }, + }; + { + let mut states = self.refill_states.lock(); + if let Some(pending) = states.get_mut(&dc_key) { + if pending.is_none() { + self.refill_pending.fetch_add(1, Ordering::AcqRel); + } + *pending = Some(addr); self.stats.increment_me_refill_skipped_inflight_total(); return; } - true - } else { - false - }; + states.insert(dc_key, None); + self.refill_running.fetch_add(1, Ordering::AcqRel); + } let pool = Arc::clone(self); - tokio::spawn(async move { - let dc_key = RefillDcKey { - dc: writer_dc, - family: if addr.is_ipv4() { - IpFamily::V4 - } else { - IpFamily::V6 - }, - }; - - if !pre_inserted { - let mut guard = pool.refill_inflight.lock().await; - if !guard.insert(endpoint_key) { - pool.stats.increment_me_refill_skipped_inflight_total(); - return; + let mut run_guard = RefillRunGuard { + pool: Arc::clone(&pool), + key: dc_key, + active: true, + }; + self.lifecycle.spawn_registered_producer(registration, async move { + let mut current_addr = addr; + loop { + pool.stats.increment_me_refill_triggered_total(); + let restored = pool + .refill_writer_after_loss(current_addr, writer_dc) + .await; + if !restored { + warn!(%current_addr, dc = writer_dc, "ME immediate refill failed"); } - } - { - let mut dc_guard = pool.refill_inflight_dc.lock().await; - if dc_guard.contains(&dc_key) { - pool.stats.increment_me_refill_skipped_inflight_total(); - drop(dc_guard); - let mut guard = pool.refill_inflight.lock().await; - guard.remove(&endpoint_key); + let Some(next_addr) = run_guard.next_or_finish() else { return; - } - dc_guard.insert(dc_key); + }; + current_addr = next_addr; } - - pool.stats.increment_me_refill_triggered_total(); - let restored = pool.refill_writer_after_loss(addr, writer_dc).await; - if !restored { - warn!(%addr, dc = writer_dc, "ME immediate refill failed"); - } - - let mut guard = pool.refill_inflight.lock().await; - guard.remove(&endpoint_key); - drop(guard); - let mut dc_guard = pool.refill_inflight_dc.lock().await; - dc_guard.remove(&dc_key); }); } } diff --git a/src/transport/middle_proxy/pool_reinit.rs b/src/transport/middle_proxy/pool_reinit.rs index 5b0d7d9..de2cc17 100644 --- a/src/transport/middle_proxy/pool_reinit.rs +++ b/src/transport/middle_proxy/pool_reinit.rs @@ -13,10 +13,106 @@ use tracing::{debug, info, warn}; use crate::crypto::SecureRandom; use crate::network::IpFamily; -use super::pool::{MeDrainGateReason, MePool, WriterContour}; +use super::pool::{ + MeDrainGateReason, MePool, ReinitAttemptState, ReinitCoordinatorState, ReinitCore, + ReinitPendingState, ReinitStatusSnapshot, WriterContour, +}; const ME_HARDSWAP_PENDING_TTL_SECS: u64 = 1800; +struct ReinitAttemptGuard { + reinit: Arc, + attempt_id: u64, + generation: u64, + previous_generation: u64, + map_hash: u64, + hardswap: bool, +} + +impl Drop for ReinitAttemptGuard { + fn drop(&mut self) { + let mut state = self.reinit.coordinator.lock(); + state.attempts.remove(&self.attempt_id); + publish_reinit_state(self.reinit.as_ref(), &state); + } +} + +struct ReinitReservation { + attempt: ReinitAttemptGuard, + pending_reused: bool, + pending_expired: bool, + pending_age_secs: u64, +} + +fn publish_reinit_state(reinit: &ReinitCore, state: &ReinitCoordinatorState) { + let mut warm_generations = state + .attempts + .values() + .filter(|attempt| attempt.hardswap && !attempt.committed) + .map(|attempt| attempt.generation) + .collect::>(); + warm_generations.sort_unstable(); + warm_generations.dedup(); + let pending = state.pending; + let snapshot = ReinitStatusSnapshot { + active_generation: state.active_generation, + warm_generations, + pending_hardswap_generation: pending.map_or(0, |value| value.generation), + pending_hardswap_started_at_epoch_secs: pending + .map_or(0, |value| value.started_at_epoch_secs), + pending_hardswap_map_hash: pending.map_or(0, |value| value.map_hash), + inflight: state.attempts.len(), + }; + reinit + .active_generation + .store(snapshot.active_generation, Ordering::Release); + reinit.warm_generation.store( + snapshot.warm_generations.last().copied().unwrap_or(0), + Ordering::Release, + ); + reinit.pending_hardswap_generation.store( + snapshot.pending_hardswap_generation, + Ordering::Release, + ); + reinit.pending_hardswap_started_at_epoch_secs.store( + snapshot.pending_hardswap_started_at_epoch_secs, + Ordering::Release, + ); + reinit + .pending_hardswap_map_hash + .store(snapshot.pending_hardswap_map_hash, Ordering::Release); + reinit.status.store(Arc::new(snapshot)); +} + +fn commit_reinit_state( + state: &mut ReinitCoordinatorState, + attempt_id: u64, + generation: u64, + map_hash: u64, + hardswap: bool, +) -> bool { + let Some(record) = state.attempts.get(&attempt_id).copied() else { + return false; + }; + if record.map_hash != state.desired_map_hash || record.map_hash != map_hash { + return false; + } + if hardswap { + let pending_matches = state.pending.is_some_and(|pending| { + pending.generation == generation && pending.map_hash == map_hash + }); + if !pending_matches || generation < state.active_generation { + return false; + } + state.active_generation = generation; + state.pending = None; + } + if let Some(record) = state.attempts.get_mut(&attempt_id) { + record.committed = true; + } + true +} + impl MePool { fn desired_map_hash(desired_by_dc: &HashMap>) -> u64 { let mut hasher = DefaultHasher::new(); @@ -36,36 +132,97 @@ impl MePool { hasher.finish() } - fn clear_pending_hardswap_state(&self) { - self.reinit - .pending_hardswap_generation - .store(0, Ordering::Relaxed); - self.reinit - .pending_hardswap_started_at_epoch_secs - .store(0, Ordering::Relaxed); - self.reinit - .pending_hardswap_map_hash - .store(0, Ordering::Relaxed); - self.reinit.warm_generation.store(0, Ordering::Relaxed); + fn reserve_reinit_attempt( + self: &Arc, + hardswap: bool, + map_hash: u64, + now_epoch_secs: u64, + ) -> ReinitReservation { + let mut state = self.reinit.coordinator.lock(); + state.desired_map_hash = map_hash; + let previous_generation = state.active_generation; + let mut pending_reused = false; + let mut pending_expired = false; + let mut pending_age_secs = 0; + + let generation = if hardswap { + let reusable = state.pending.filter(|pending| { + pending_age_secs = now_epoch_secs.saturating_sub(pending.started_at_epoch_secs); + pending_expired = pending.started_at_epoch_secs > 0 + && pending_age_secs > ME_HARDSWAP_PENDING_TTL_SECS; + pending.generation >= previous_generation + && pending.map_hash == map_hash + && !pending_expired + }); + if let Some(pending) = reusable { + pending_reused = true; + pending.generation + } else { + let generation = self.reinit.generation.fetch_add(1, Ordering::AcqRel) + 1; + state.pending = Some(ReinitPendingState { + generation, + started_at_epoch_secs: now_epoch_secs, + map_hash, + }); + generation + } + } else { + state.pending = None; + self.reinit.generation.fetch_add(1, Ordering::AcqRel) + 1 + }; + + let attempt_id = state.next_attempt_id; + state.next_attempt_id = state.next_attempt_id.saturating_add(1); + state.attempts.insert( + attempt_id, + ReinitAttemptState { + generation, + map_hash, + hardswap, + committed: false, + }, + ); + publish_reinit_state(self.reinit.as_ref(), &state); + ReinitReservation { + attempt: ReinitAttemptGuard { + reinit: Arc::clone(&self.reinit), + attempt_id, + generation, + previous_generation, + map_hash, + hardswap, + }, + pending_reused, + pending_expired, + pending_age_secs, + } } - async fn promote_warm_generation_to_active(&self, generation: u64) { - self.reinit - .active_generation - .store(generation, Ordering::Relaxed); - self.reinit.warm_generation.store(0, Ordering::Relaxed); - - let ws = self.writers.read().await; - for writer in ws.iter() { - if writer.draining.load(Ordering::Relaxed) { - continue; - } - if writer.generation == generation { - writer - .contour - .store(WriterContour::Active.as_u8(), Ordering::Relaxed); + fn commit_reinit_attempt(&self, attempt: &ReinitAttemptGuard) -> bool { + let mut state = self.reinit.coordinator.lock(); + if !commit_reinit_state( + &mut state, + attempt.attempt_id, + attempt.generation, + attempt.map_hash, + attempt.hardswap, + ) { + return false; + } + if attempt.hardswap { + let writers = self.writers.snapshot(); + for writer in writers.iter() { + if !writer.draining.load(Ordering::Relaxed) + && writer.generation == attempt.generation + { + writer + .contour + .store(WriterContour::Active.as_u8(), Ordering::Release); + } } } + publish_reinit_state(self.reinit.as_ref(), &state); + true } fn coverage_ratio( @@ -387,74 +544,32 @@ impl MePool { } let desired_map_hash = Self::desired_map_hash(&desired_by_dc); - let previous_generation = self.current_generation(); let hardswap = self.reinit.hardswap.load(Ordering::Relaxed); - let generation = if hardswap { - let pending_generation = self - .reinit - .pending_hardswap_generation - .load(Ordering::Relaxed); - let pending_started_at = self - .reinit - .pending_hardswap_started_at_epoch_secs - .load(Ordering::Relaxed); - let pending_map_hash = self - .reinit - .pending_hardswap_map_hash - .load(Ordering::Relaxed); - let pending_age_secs = now_epoch_secs.saturating_sub(pending_started_at); - let pending_ttl_expired = - pending_started_at > 0 && pending_age_secs > ME_HARDSWAP_PENDING_TTL_SECS; - let pending_matches_map = pending_map_hash != 0 && pending_map_hash == desired_map_hash; - - if pending_generation != 0 - && pending_generation >= previous_generation - && pending_matches_map - && !pending_ttl_expired - { - self.stats.increment_me_hardswap_pending_reuse_total(); - debug!( - previous_generation, - generation = pending_generation, - pending_age_secs, - "ME hardswap continues with pending generation" - ); - pending_generation - } else { - if pending_generation != 0 && pending_ttl_expired { - self.stats.increment_me_hardswap_pending_ttl_expired_total(); - warn!( - previous_generation, - generation = pending_generation, - pending_age_secs, - pending_ttl_secs = ME_HARDSWAP_PENDING_TTL_SECS, - "ME hardswap pending generation expired by TTL; starting fresh generation" - ); - } - let next_generation = self.reinit.generation.fetch_add(1, Ordering::Relaxed) + 1; - self.reinit - .pending_hardswap_generation - .store(next_generation, Ordering::Relaxed); - self.reinit - .pending_hardswap_started_at_epoch_secs - .store(now_epoch_secs, Ordering::Relaxed); - self.reinit - .pending_hardswap_map_hash - .store(desired_map_hash, Ordering::Relaxed); - self.reinit - .warm_generation - .store(next_generation, Ordering::Relaxed); - next_generation - } - } else { - self.clear_pending_hardswap_state(); - self.reinit.generation.fetch_add(1, Ordering::Relaxed) + 1 - }; + let reservation = + self.reserve_reinit_attempt(hardswap, desired_map_hash, now_epoch_secs); + let attempt = reservation.attempt; + let previous_generation = attempt.previous_generation; + let generation = attempt.generation; + if reservation.pending_reused { + self.stats.increment_me_hardswap_pending_reuse_total(); + debug!( + previous_generation, + generation, + pending_age_secs = reservation.pending_age_secs, + "ME hardswap continues with pending generation" + ); + } else if reservation.pending_expired { + self.stats.increment_me_hardswap_pending_ttl_expired_total(); + warn!( + previous_generation, + generation, + pending_age_secs = reservation.pending_age_secs, + pending_ttl_secs = ME_HARDSWAP_PENDING_TTL_SECS, + "ME hardswap pending generation expired by TTL; starting fresh generation" + ); + } if hardswap { - self.reinit - .warm_generation - .store(generation, Ordering::Relaxed); self.warmup_generation_for_all_dcs(rng, generation, &desired_by_dc) .await; } else { @@ -542,8 +657,13 @@ impl MePool { ); } - if hardswap { - self.promote_warm_generation_to_active(generation).await; + if !self.commit_reinit_attempt(&attempt) { + debug!( + previous_generation, + generation, + "ME reinit result discarded after a newer desired-map attempt" + ); + return false; } let desired_addrs: HashSet<(i32, SocketAddr)> = desired_by_dc @@ -566,9 +686,6 @@ impl MePool { drop(writers); if stale_writer_ids.is_empty() { - if hardswap { - self.clear_pending_hardswap_state(); - } debug!("ME reinit cycle completed with no stale writers"); return true; } @@ -606,9 +723,6 @@ impl MePool { self.remove_writer_and_close_clients(writer_id).await; } } - if hardswap { - self.clear_pending_hardswap_state(); - } true } @@ -622,7 +736,10 @@ mod tests { use std::collections::{HashMap, HashSet}; use std::net::{IpAddr, Ipv4Addr, SocketAddr}; - use super::MePool; + use super::{MePool, commit_reinit_state}; + use crate::transport::middle_proxy::pool::{ + ReinitAttemptState, ReinitCoordinatorState, ReinitPendingState, + }; fn addr(octet: u8, port: u16) -> SocketAddr { SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, octet)), port) @@ -673,4 +790,44 @@ mod tests { assert_eq!(ratio, 0.0); assert_eq!(missing_dc, vec![1, 2]); } + + #[test] + fn stale_concurrent_attempt_cannot_regress_active_generation() { + let mut state = ReinitCoordinatorState { + next_attempt_id: 3, + active_generation: 1, + desired_map_hash: 22, + pending: Some(ReinitPendingState { + generation: 3, + started_at_epoch_secs: 1, + map_hash: 22, + }), + attempts: HashMap::from([ + ( + 1, + ReinitAttemptState { + generation: 2, + map_hash: 11, + hardswap: true, + committed: false, + }, + ), + ( + 2, + ReinitAttemptState { + generation: 3, + map_hash: 22, + hardswap: true, + committed: false, + }, + ), + ]), + }; + + assert!(commit_reinit_state(&mut state, 2, 3, 22, true)); + assert_eq!(state.active_generation, 3); + assert!(!commit_reinit_state(&mut state, 1, 2, 11, true)); + assert_eq!(state.active_generation, 3); + assert!(state.pending.is_none()); + } } diff --git a/src/transport/middle_proxy/pool_runtime_api.rs b/src/transport/middle_proxy/pool_runtime_api.rs index 539f397..7b26583 100644 --- a/src/transport/middle_proxy/pool_runtime_api.rs +++ b/src/transport/middle_proxy/pool_runtime_api.rs @@ -15,6 +15,8 @@ pub(crate) struct MeApiRefillDcSnapshot { pub(crate) struct MeApiRefillSnapshot { pub inflight_endpoints_total: usize, pub inflight_dc_total: usize, + pub running_dc_total: usize, + pub pending_dc_total: usize, pub by_dc: Vec, } @@ -56,14 +58,21 @@ pub(crate) struct MeApiDrainGateSnapshot { impl MePool { pub(crate) async fn api_refill_snapshot(&self) -> MeApiRefillSnapshot { - let inflight_endpoints_total = self.refill_inflight.lock().await.len(); - let inflight_dc_keys = self - .refill_inflight_dc - .lock() - .await - .iter() + let refill_states = self.refill_states.lock(); + let inflight_endpoints_total = refill_states + .values() + .map(|pending| 1usize + usize::from(pending.is_some())) + .sum(); + let running_dc_total = refill_states.len(); + let pending_dc_total = refill_states + .values() + .filter(|pending| pending.is_some()) + .count(); + let inflight_dc_keys = refill_states + .keys() .copied() .collect::>(); + drop(refill_states); let mut by_dc_map = HashMap::<(i16, &'static str), usize>::new(); for key in inflight_dc_keys { @@ -88,6 +97,8 @@ impl MePool { MeApiRefillSnapshot { inflight_endpoints_total, inflight_dc_total: by_dc.len(), + running_dc_total, + pending_dc_total, by_dc, } } diff --git a/src/transport/middle_proxy/pool_status.rs b/src/transport/middle_proxy/pool_status.rs index ae9038b..8eea3a6 100644 --- a/src/transport/middle_proxy/pool_status.rs +++ b/src/transport/middle_proxy/pool_status.rs @@ -1,9 +1,10 @@ use std::collections::{BTreeMap, BTreeSet, HashMap}; use std::net::{IpAddr, SocketAddr}; +use std::sync::Arc; use std::sync::atomic::Ordering; use std::time::Instant; -use super::pool::{MePool, WriterContour}; +use super::pool::{MePool, ReinitStatusSnapshot, WriterContour}; use crate::config::{MeBindStaleMode, MeFloorMode, MeSocksKdfPolicy}; use crate::transport::upstream::IpPreference; @@ -87,8 +88,11 @@ pub(crate) struct MeApiDcPathSnapshot { pub(crate) struct MeApiRuntimeSnapshot { pub active_generation: u64, pub warm_generation: u64, + pub warm_generations: Vec, pub pending_hardswap_generation: u64, pub pending_hardswap_age_secs: Option, + pub reinit_inflight: usize, + pub reinit_max_concurrency_effective: usize, pub hardswap_enabled: bool, pub floor_mode: &'static str, pub adaptive_floor_idle_secs: u64, @@ -222,8 +226,16 @@ impl MePool { } pub(crate) async fn api_status_snapshot(&self) -> MeApiStatusSnapshot { + let reinit = self.reinit.status.load_full(); + self.api_status_snapshot_for_reinit(reinit.as_ref()).await + } + + async fn api_status_snapshot_for_reinit( + &self, + reinit: &ReinitStatusSnapshot, + ) -> MeApiStatusSnapshot { let now_epoch_secs = Self::now_epoch_secs(); - let active_generation = self.current_generation(); + let active_generation = reinit.active_generation; let drain_ttl_secs = self .drain_runtime .me_pool_drain_ttl_secs @@ -440,13 +452,19 @@ impl MePool { } } + #[allow(dead_code)] pub(crate) async fn api_runtime_snapshot(&self) -> MeApiRuntimeSnapshot { + let reinit = self.reinit.status.load_full(); + self.api_runtime_snapshot_for_reinit(reinit.as_ref()).await + } + + async fn api_runtime_snapshot_for_reinit( + &self, + reinit: &ReinitStatusSnapshot, + ) -> MeApiRuntimeSnapshot { let now = Instant::now(); let now_epoch_secs = Self::now_epoch_secs(); - let pending_started_at = self - .reinit - .pending_hardswap_started_at_epoch_secs - .load(Ordering::Relaxed); + let pending_started_at = reinit.pending_hardswap_started_at_epoch_secs; let pending_hardswap_age_secs = (pending_started_at > 0).then_some(now_epoch_secs.saturating_sub(pending_started_at)); @@ -486,13 +504,16 @@ impl MePool { } MeApiRuntimeSnapshot { - active_generation: self.reinit.active_generation.load(Ordering::Relaxed), - warm_generation: self.reinit.warm_generation.load(Ordering::Relaxed), - pending_hardswap_generation: self - .reinit - .pending_hardswap_generation - .load(Ordering::Relaxed), + active_generation: reinit.active_generation, + warm_generation: reinit.warm_generations.last().copied().unwrap_or(0), + warm_generations: reinit.warm_generations.clone(), + pending_hardswap_generation: reinit.pending_hardswap_generation, pending_hardswap_age_secs, + reinit_inflight: reinit.inflight, + reinit_max_concurrency_effective: self + .reinit + .max_concurrency_effective + .load(Ordering::Acquire), hardswap_enabled: self.reinit.hardswap.load(Ordering::Relaxed), floor_mode: floor_mode_label(self.floor_mode()), adaptive_floor_idle_secs: self @@ -662,6 +683,23 @@ impl MePool { network_path, } } + + pub(crate) async fn api_coherent_snapshots( + &self, + ) -> (MeApiStatusSnapshot, MeApiRuntimeSnapshot) { + let mut attempts = 0usize; + loop { + let reinit = self.reinit.status.load_full(); + let status = self.api_status_snapshot_for_reinit(reinit.as_ref()).await; + let runtime = self + .api_runtime_snapshot_for_reinit(reinit.as_ref()) + .await; + attempts += 1; + if Arc::ptr_eq(&reinit, &self.reinit.status.load_full()) || attempts >= 3 { + return (status, runtime); + } + } + } } fn ratio_pct(part: usize, total: usize) -> f64 { diff --git a/src/transport/middle_proxy/pool_writer.rs b/src/transport/middle_proxy/pool_writer.rs index c87c296..9c54c07 100644 --- a/src/transport/middle_proxy/pool_writer.rs +++ b/src/transport/middle_proxy/pool_writer.rs @@ -269,7 +269,10 @@ async fn rpc_proxy_req_signal_loop( continue; }; - let (conn_id, mut service_rx) = pool.registry.register().await; + let Some((conn_lease, mut service_rx)) = pool.register_connection().await else { + return; + }; + let conn_id = conn_lease.conn_id(); // Service RPC_PROXY_REQ signal path is intentionally route-only: // do not bind synthetic conn_id into regular writer/client accounting. @@ -286,7 +289,7 @@ async fn rpc_proxy_req_signal_loop( send_service_writer_command(&tx_signal, WriterCommand::DataAndFlush(payload)).await { stats_signal.increment_me_rpc_proxy_req_signal_failed_total(); - let _ = pool.registry.unregister(conn_id).await; + conn_lease.unregister().await; match error { ServiceWriterCommandSendError::Closed => return, ServiceWriterCommandSendError::TimedOut => continue, @@ -313,7 +316,7 @@ async fn rpc_proxy_req_signal_loop( .await { stats_signal.increment_me_rpc_proxy_req_signal_failed_total(); - let _ = pool.registry.unregister(conn_id).await; + conn_lease.unregister().await; match error { ServiceWriterCommandSendError::Closed => return, ServiceWriterCommandSendError::TimedOut => continue, @@ -321,7 +324,7 @@ async fn rpc_proxy_req_signal_loop( } stats_signal.increment_me_rpc_proxy_req_signal_close_sent_total(); - let _ = pool.registry.unregister(conn_id).await; + conn_lease.unregister().await; } } @@ -394,14 +397,14 @@ impl MePool { writer_dc: i32, allow_coverage_override: bool, ) -> Result<()> { - if !self - .can_open_writer_for_contour(contour, allow_coverage_override, writer_dc) + let Some(_writer_open_reservation) = self + .reserve_writer_open(contour, allow_coverage_override, writer_dc) .await - { + else { return Err(ProxyError::Proxy(format!( "ME {contour:?} writer cap reached" ))); - } + }; let secret_len = self.proxy_secret.read().await.secret.len(); if secret_len < 32 { @@ -415,6 +418,9 @@ impl MePool { let hs = self .handshake_only(stream, addr, upstream_egress, rng) .await?; + let Some(task_registration) = self.lifecycle.try_register() else { + return Err(ProxyError::Proxy("ME pool lifecycle closed".into())); + }; let writer_id = self.next_writer_id.fetch_add(1, Ordering::Relaxed); let contour = Arc::new(AtomicU8::new(contour.as_u8())); @@ -499,7 +505,7 @@ impl MePool { let route_fairshare_enabled = self.transport_policy.me_route_fairshare_enabled.clone(); let reader_route_data_wait_ms = self.transport_policy.me_reader_route_data_wait_ms.clone(); - tokio::spawn(async move { + self.lifecycle.spawn_registered_writer(task_registration, async move { // Reader MUST be the first branch in biased select! to avoid read starvation. let exit = tokio::select! { biased; diff --git a/src/transport/middle_proxy/registry.rs b/src/transport/middle_proxy/registry.rs index b3d3f4b..e97a233 100644 --- a/src/transport/middle_proxy/registry.rs +++ b/src/transport/middle_proxy/registry.rs @@ -1,6 +1,7 @@ use std::collections::{HashMap, HashSet}; use std::net::SocketAddr; use std::sync::Arc; +use std::sync::Mutex as StdMutex; use std::sync::atomic::{AtomicU8, AtomicU64, Ordering}; use std::time::{SystemTime, UNIX_EPOCH}; @@ -118,6 +119,38 @@ pub struct ConnRegistry { route_backpressure_high_timeout_ms: AtomicU64, route_backpressure_high_watermark_pct: AtomicU8, route_byte_permits_per_conn: usize, + cleanup_tx: mpsc::Sender, + cleanup_rx: StdMutex>>, +} + +/// Cancellation-safe ownership of one registered ME client route. +pub(crate) struct ConnLease { + registry: Arc, + conn_id: u64, + cleanup_permit: Option>, +} + +impl ConnLease { + /// Returns the stable connection identifier owned by this lease. + pub(crate) fn conn_id(&self) -> u64 { + self.conn_id + } + + /// Completes asynchronous registry cleanup and disarms drop cleanup. + pub(crate) async fn unregister(mut self) { + self.registry.unregister(self.conn_id).await; + self.cleanup_permit.take(); + } +} + +impl Drop for ConnLease { + fn drop(&mut self) { + let Some(permit) = self.cleanup_permit.take() else { + return; + }; + self.registry.remove_route_now(self.conn_id); + permit.send(self.conn_id); + } } impl ConnRegistry { @@ -133,15 +166,30 @@ impl ConnRegistry { Self::with_route_limits( route_channel_capacity, Self::route_byte_permit_budget(route_channel_capacity), + 16_384, + ) + } + + pub(crate) fn with_route_and_cleanup_capacity( + route_channel_capacity: usize, + connection_cleanup_capacity: usize, + ) -> Self { + let route_channel_capacity = route_channel_capacity.max(1); + Self::with_route_limits( + route_channel_capacity, + Self::route_byte_permit_budget(route_channel_capacity), + connection_cleanup_capacity, ) } fn with_route_limits( route_channel_capacity: usize, route_byte_permits_per_conn: usize, + connection_cleanup_capacity: usize, ) -> Self { let start = rand::random::() | 1; let route_channel_capacity = route_channel_capacity.max(1); + let (cleanup_tx, cleanup_rx) = mpsc::channel(connection_cleanup_capacity.max(1)); Self { routing: RoutingTable { map: DashMap::new(), @@ -168,6 +216,8 @@ impl ConnRegistry { ROUTE_BACKPRESSURE_HIGH_WATERMARK_PCT, ), route_byte_permits_per_conn: route_byte_permits_per_conn.max(1), + cleanup_tx, + cleanup_rx: StdMutex::new(Some(cleanup_rx)), } } @@ -199,7 +249,7 @@ impl ConnRegistry { route_channel_capacity: usize, route_byte_permits_per_conn: usize, ) -> Self { - Self::with_route_limits(route_channel_capacity, route_byte_permits_per_conn) + Self::with_route_limits(route_channel_capacity, route_byte_permits_per_conn, 16_384) } pub fn update_route_backpressure_policy( @@ -220,6 +270,29 @@ impl ConnRegistry { } pub async fn register(&self) -> (u64, mpsc::Receiver) { + self.register_route() + } + + pub(crate) async fn register_leased( + self: &Arc, + ) -> Option<(ConnLease, mpsc::Receiver)> { + let cleanup_permit = self.cleanup_tx.clone().reserve_owned().await.ok()?; + let (conn_id, rx) = self.register_route(); + Some(( + ConnLease { + registry: Arc::clone(self), + conn_id, + cleanup_permit: Some(cleanup_permit), + }, + rx, + )) + } + + pub(super) fn take_cleanup_receiver(&self) -> Option> { + self.cleanup_rx.lock().ok()?.take() + } + + fn register_route(&self) -> (u64, mpsc::Receiver) { let id = self.next_id.fetch_add(1, Ordering::Relaxed); let (tx, rx) = mpsc::channel(self.route_channel_capacity); self.routing.map.insert(id, tx); @@ -229,6 +302,12 @@ impl ConnRegistry { ); (id, rx) } + + fn remove_route_now(&self, id: u64) { + self.routing.map.remove(&id); + self.routing.byte_budget.remove(&id); + self.hot_binding.map.remove(&id); + } } #[cfg(test)] diff --git a/src/transport/middle_proxy/registry/tests.rs b/src/transport/middle_proxy/registry/tests.rs index fe981a5..c90daf0 100644 --- a/src/transport/middle_proxy/registry/tests.rs +++ b/src/transport/middle_proxy/registry/tests.rs @@ -308,3 +308,21 @@ async fn non_empty_writer_ids_returns_only_writers_with_bound_clients() { assert!(!non_empty.contains(&20)); assert!(!non_empty.contains(&30)); } + +#[tokio::test] +async fn leased_registration_removes_hot_route_before_async_cleanup() { + let registry = Arc::new(ConnRegistry::with_route_and_cleanup_capacity(8, 1)); + let mut cleanup_rx = registry.take_cleanup_receiver().unwrap(); + let (lease, _rx) = registry.register_leased().await.unwrap(); + let conn_id = lease.conn_id(); + + drop(lease); + + assert_eq!( + registry.route_nowait(conn_id, MeResponse::Ack(1)).await, + RouteResult::NoConn + ); + assert_eq!(cleanup_rx.recv().await, Some(conn_id)); + registry.unregister(conn_id).await; + assert!(registry.active_conn_ids().await.is_empty()); +} diff --git a/src/transport/middle_proxy/rotation.rs b/src/transport/middle_proxy/rotation.rs index 9308226..90a08e1 100644 --- a/src/transport/middle_proxy/rotation.rs +++ b/src/transport/middle_proxy/rotation.rs @@ -2,6 +2,7 @@ use std::sync::Arc; use std::time::Duration; use tokio::sync::{mpsc, watch}; +use tokio::task::JoinSet; use tracing::{debug, info, warn}; use crate::config::ProxyConfig; @@ -42,6 +43,45 @@ pub fn enqueue_reinit_trigger(tx: &mpsc::Sender, trigger: MeRei } } +const REINIT_TRIGGER_PERIODIC: u8 = 1; +const REINIT_TRIGGER_MAP_CHANGED: u8 = 2; + +struct ReinitInflightGuard { + pool: Arc, +} + +impl Drop for ReinitInflightGuard { + fn drop(&mut self) { + self.pool + .reinit + .scheduler_inflight + .fetch_sub(1, std::sync::atomic::Ordering::AcqRel); + } +} + +fn effective_reinit_concurrency(config: &ProxyConfig) -> usize { + if config.general.me_reinit_singleflight { + 1 + } else { + config.general.me_reinit_max_concurrency.clamp(1, 8) + } +} + +fn trigger_bit(trigger: MeReinitTrigger) -> u8 { + match trigger { + MeReinitTrigger::Periodic => REINIT_TRIGGER_PERIODIC, + MeReinitTrigger::MapChanged => REINIT_TRIGGER_MAP_CHANGED, + } +} + +fn trigger_reason(pending: u8) -> &'static str { + match pending { + REINIT_TRIGGER_PERIODIC => "periodic", + REINIT_TRIGGER_MAP_CHANGED => "map-change", + _ => "map-change+periodic", + } +} + pub async fn me_reinit_scheduler( pool: Arc, rng: Arc, @@ -50,68 +90,104 @@ pub async fn me_reinit_scheduler( me_ready_tx: watch::Sender, ) { info!("ME reinit scheduler started"); + let mut tasks = JoinSet::::new(); + let mut pending = 0u8; + let mut pending_deadline = None; + let mut trigger_channel_open = true; + loop { - let Some(first_trigger) = trigger_rx.recv().await else { - warn!("ME reinit scheduler stopped: trigger channel closed"); - break; - }; - - let mut map_change_seen = matches!(first_trigger, MeReinitTrigger::MapChanged); - let mut periodic_seen = matches!(first_trigger, MeReinitTrigger::Periodic); let cfg = config_rx.borrow().clone(); - let coalesce_window = Duration::from_millis(cfg.general.me_reinit_coalesce_window_ms); - if !coalesce_window.is_zero() { - let deadline = tokio::time::Instant::now() + coalesce_window; - loop { - let now = tokio::time::Instant::now(); - if now >= deadline { - break; - } - match tokio::time::timeout(deadline - now, trigger_rx.recv()).await { - Ok(Some(next)) => { - if next == MeReinitTrigger::MapChanged { - map_change_seen = true; - } else { - periodic_seen = true; - } - } - Ok(None) => break, - Err(_) => break, - } - } - } + let max_concurrency = effective_reinit_concurrency(&cfg); + pool.reinit + .max_concurrency_effective + .store(max_concurrency, std::sync::atomic::Ordering::Release); - let reason = if map_change_seen && periodic_seen { - "map-change+periodic" - } else if map_change_seen { - "map-change" - } else { - "periodic" - }; - - if cfg.general.me_reinit_singleflight { - debug!(reason, "ME reinit scheduled (single-flight)"); - if pool.zero_downtime_reinit_periodic(rng.as_ref()).await { - me_ready_tx.send_modify(|version| { - *version = version.saturating_add(1); - }); - } - } else { - debug!(reason, "ME reinit scheduled (concurrent mode)"); + let pending_ready = pending != 0 + && pending_deadline.is_none_or(|deadline| deadline <= tokio::time::Instant::now()); + if pending_ready && tasks.len() < max_concurrency { + let reason = trigger_reason(pending); + pending = 0; + pending_deadline = None; + debug!(reason, max_concurrency, "ME reinit scheduled"); let pool_clone = pool.clone(); let rng_clone = rng.clone(); - let me_ready_tx_clone = me_ready_tx.clone(); - tokio::spawn(async move { - if pool_clone + pool.reinit + .scheduler_inflight + .fetch_add(1, std::sync::atomic::Ordering::AcqRel); + let inflight = ReinitInflightGuard { + pool: Arc::clone(&pool_clone), + }; + tasks.spawn(async move { + let _inflight = inflight; + pool_clone .zero_downtime_reinit_periodic(rng_clone.as_ref()) .await - { - me_ready_tx_clone.send_modify(|version| { - *version = version.saturating_add(1); - }); - } }); + continue; } + + if !trigger_channel_open && pending == 0 && tasks.is_empty() { + warn!("ME reinit scheduler stopped: trigger channel closed"); + return; + } + + let timer_enabled = pending != 0 && tasks.len() < max_concurrency; + tokio::select! { + trigger = trigger_rx.recv(), if trigger_channel_open => { + match trigger { + Some(trigger) => { + if pending == 0 { + pending_deadline = Some( + tokio::time::Instant::now() + + Duration::from_millis( + cfg.general.me_reinit_coalesce_window_ms, + ), + ); + } + pending |= trigger_bit(trigger); + } + None => trigger_channel_open = false, + } + } + joined = tasks.join_next(), if !tasks.is_empty() => { + match joined { + Some(Ok(true)) => { + me_ready_tx.send_modify(|version| { + *version = version.saturating_add(1); + }); + } + Some(Ok(false)) => {} + Some(Err(error)) => { + warn!(error = %error, "ME reinit task failed"); + } + None => {} + } + } + _ = async { + if let Some(deadline) = pending_deadline { + tokio::time::sleep_until(deadline).await; + } + }, if timer_enabled => {} + } + } +} + +#[cfg(test)] +mod tests { + use crate::config::ProxyConfig; + + #[test] + fn effective_concurrency_is_one_for_singleflight_and_bounded_otherwise() { + let mut config = ProxyConfig::default(); + config.general.me_reinit_singleflight = true; + config.general.me_reinit_max_concurrency = 8; + assert_eq!(super::effective_reinit_concurrency(&config), 1); + + config.general.me_reinit_singleflight = false; + config.general.me_reinit_max_concurrency = 2; + assert_eq!(super::effective_reinit_concurrency(&config), 2); + config.general.me_reinit_max_concurrency = usize::MAX; + assert_eq!(super::effective_reinit_concurrency(&config), 8); } } diff --git a/src/transport/middle_proxy/secret.rs b/src/transport/middle_proxy/secret.rs index e3495e8..6bae755 100644 --- a/src/transport/middle_proxy/secret.rs +++ b/src/transport/middle_proxy/secret.rs @@ -108,6 +108,7 @@ pub async fn download_proxy_secret_with_max_len_via_upstream( let resp = https_get( proxy_secret_url.unwrap_or("https://core.telegram.org/getProxySecret"), upstream, + max_len, ) .await?; diff --git a/src/transport/middle_proxy/tests/health_adversarial_tests.rs b/src/transport/middle_proxy/tests/health_adversarial_tests.rs index 80cad99..e42151d 100644 --- a/src/transport/middle_proxy/tests/health_adversarial_tests.rs +++ b/src/transport/middle_proxy/tests/health_adversarial_tests.rs @@ -122,6 +122,7 @@ async fn make_pool( general.me_route_blocking_send_timeout_ms, general.me_route_inline_recovery_attempts, general.me_route_inline_recovery_wait_ms, + 16_384, ); (pool, rng) diff --git a/src/transport/middle_proxy/tests/health_integration_tests.rs b/src/transport/middle_proxy/tests/health_integration_tests.rs index 72ac165..9aecfba 100644 --- a/src/transport/middle_proxy/tests/health_integration_tests.rs +++ b/src/transport/middle_proxy/tests/health_integration_tests.rs @@ -120,6 +120,7 @@ async fn make_pool( general.me_route_blocking_send_timeout_ms, general.me_route_inline_recovery_attempts, general.me_route_inline_recovery_wait_ms, + 16_384, ); (pool, rng) } diff --git a/src/transport/middle_proxy/tests/health_regression_tests.rs b/src/transport/middle_proxy/tests/health_regression_tests.rs index 85f0764..04c7773 100644 --- a/src/transport/middle_proxy/tests/health_regression_tests.rs +++ b/src/transport/middle_proxy/tests/health_regression_tests.rs @@ -115,6 +115,7 @@ async fn make_pool(me_pool_drain_threshold: u64) -> Arc { general.me_route_blocking_send_timeout_ms, general.me_route_inline_recovery_attempts, general.me_route_inline_recovery_wait_ms, + 16_384, ) } diff --git a/src/transport/middle_proxy/tests/pool_refill_security_tests.rs b/src/transport/middle_proxy/tests/pool_refill_security_tests.rs index 9d7fb99..a513d1a 100644 --- a/src/transport/middle_proxy/tests/pool_refill_security_tests.rs +++ b/src/transport/middle_proxy/tests/pool_refill_security_tests.rs @@ -1,6 +1,7 @@ use std::collections::HashMap; use std::net::{IpAddr, Ipv4Addr, SocketAddr}; use std::sync::Arc; +use std::sync::atomic::Ordering; use std::time::{Duration, Instant}; use crate::config::{GeneralConfig, MeRouteNoWriterMode, MeSocksKdfPolicy, MeWriterPickMode}; @@ -104,6 +105,7 @@ async fn make_pool() -> Arc { general.me_route_blocking_send_timeout_ms, general.me_route_inline_recovery_attempts, general.me_route_inline_recovery_wait_ms, + 16_384, ) } @@ -159,3 +161,31 @@ async fn connectable_endpoints_releases_quarantine_lock_before_sleep() { .expect("task join failed"); assert_eq!(endpoints, vec![addr]); } + +#[tokio::test(flavor = "current_thread")] +async fn refill_coalesces_one_pending_endpoint_and_cleans_up_before_first_poll() { + let pool = make_pool().await; + let first = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 31, 0, 21)), 443); + let second = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 31, 0, 22)), 443); + let latest = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 31, 0, 23)), 443); + + pool.trigger_immediate_refill_for_dc(first, 2); + pool.trigger_immediate_refill_for_dc(second, 2); + pool.trigger_immediate_refill_for_dc(latest, 2); + + assert_eq!(pool.refill_states.lock().len(), 1); + assert_eq!(pool.refill_running.load(Ordering::Acquire), 1); + assert_eq!(pool.refill_pending.load(Ordering::Acquire), 1); + + pool.begin_shutdown(); + tokio::time::timeout(Duration::from_secs(1), async { + while pool.refill_running.load(Ordering::Acquire) != 0 { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + + assert!(pool.refill_states.lock().is_empty()); + assert_eq!(pool.refill_pending.load(Ordering::Acquire), 0); +} diff --git a/src/transport/middle_proxy/tests/pool_writer_security_tests.rs b/src/transport/middle_proxy/tests/pool_writer_security_tests.rs index f7908dc..1d57257 100644 --- a/src/transport/middle_proxy/tests/pool_writer_security_tests.rs +++ b/src/transport/middle_proxy/tests/pool_writer_security_tests.rs @@ -109,6 +109,7 @@ async fn make_pool() -> Arc { general.me_route_blocking_send_timeout_ms, general.me_route_inline_recovery_attempts, general.me_route_inline_recovery_wait_ms, + 16_384, ) } diff --git a/src/transport/middle_proxy/tests/send_adversarial_tests.rs b/src/transport/middle_proxy/tests/send_adversarial_tests.rs index 7ccd62d..667eb33 100644 --- a/src/transport/middle_proxy/tests/send_adversarial_tests.rs +++ b/src/transport/middle_proxy/tests/send_adversarial_tests.rs @@ -119,6 +119,7 @@ async fn make_pool_with_decision(decision: NetworkDecision) -> (Arc, Arc general.me_route_blocking_send_timeout_ms, general.me_route_inline_recovery_attempts, general.me_route_inline_recovery_wait_ms, + 16_384, ); (pool, rng) diff --git a/src/transport/upstream.rs b/src/transport/upstream.rs index ad1b389..f899b7e 100644 --- a/src/transport/upstream.rs +++ b/src/transport/upstream.rs @@ -4,7 +4,6 @@ #![allow(deprecated)] -use arc_swap::ArcSwap; use rand::RngExt; use std::collections::{BTreeSet, HashMap}; use std::net::{IpAddr, SocketAddr}; @@ -21,7 +20,7 @@ use tracing::{debug, info, trace, warn}; use crate::config::{UpstreamConfig, UpstreamType}; use crate::error::{ProxyError, Result}; -use crate::network::dns_overrides::{DnsOverrides, split_host_port}; +use crate::network::dns_overrides::{GenerationDnsResolver, split_host_port}; use crate::protocol::constants::{TG_DATACENTER_PORT, TG_DATACENTERS_V4, TG_DATACENTERS_V6}; use crate::stats::Stats; use crate::transport::shadowsocks::{ @@ -43,6 +42,7 @@ const HEALTH_CHECK_INTERVAL_SECS: u64 = 30; const HEALTH_CHECK_CONNECT_TIMEOUT_SECS: u64 = 10; /// Upstream is considered healthy when at least this many DC groups are reachable. const MIN_HEALTHY_DC_GROUPS: usize = 3; +const DNS_RESULT_MAX_ADDRESSES: usize = 64; // ============= RTT Tracking ============= @@ -334,7 +334,7 @@ pub struct UpstreamManager { no_upstreams_warn_epoch_ms: Arc, no_healthy_warn_epoch_ms: Arc, stats: Arc, - dns_overrides: Arc>, + dns_resolver: Arc, } impl UpstreamManager { @@ -376,29 +376,42 @@ impl UpstreamManager { no_upstreams_warn_epoch_ms: Arc::new(AtomicU64::new(0)), no_healthy_warn_epoch_ms: Arc::new(AtomicU64::new(0)), stats, - dns_overrides: Arc::new(ArcSwap::from_pointee(DnsOverrides::default())), + dns_resolver: Arc::new(GenerationDnsResolver::default()), } } - pub(crate) fn with_dns_overrides(mut self, entries: &[String]) -> Result { - self.dns_overrides = Arc::new(ArcSwap::from_pointee(DnsOverrides::from_entries(entries)?)); + pub(crate) fn with_dns_overrides(self, entries: &[String]) -> Result { + self.dns_resolver.apply_entries(entries)?; Ok(self) } pub(crate) fn update_dns_overrides(&self, entries: &[String]) -> Result<()> { - let snapshot = DnsOverrides::from_entries(entries)?; - self.dns_overrides.store(Arc::new(snapshot)); - Ok(()) + self.dns_resolver.apply_entries(entries) + } + + pub(crate) fn dns_resolver(&self) -> Arc { + Arc::clone(&self.dns_resolver) + } + + pub(crate) async fn resolve_all(&self, host: &str, port: u16) -> Result> { + if let Some(addr) = self.dns_resolver.resolve_socket_addr(host, port) { + return Ok(vec![addr]); + } + let addrs = tokio::net::lookup_host((host, port)) + .await + .map_err(ProxyError::Io)? + .take(DNS_RESULT_MAX_ADDRESSES) + .collect::>(); + if addrs.is_empty() { + return Err(ProxyError::Proxy(format!( + "DNS returned no addresses for {host}:{port}" + ))); + } + Ok(addrs) } pub(crate) async fn resolve_hostname(&self, host: &str, port: u16) -> Result { - if let Some(addr) = self.dns_overrides.load().resolve_socket_addr(host, port) { - return Ok(addr); - } - let addrs: Vec = tokio::net::lookup_host((host, port)) - .await - .map_err(ProxyError::Io)? - .collect(); + let addrs = self.resolve_all(host, port).await?; if let Some(addr) = addrs.iter().copied().find(SocketAddr::is_ipv4) { return Ok(addr); } @@ -744,7 +757,7 @@ impl UpstreamManager { connect_timeout: Duration, ) -> Result { if let Some((host, port)) = split_host_port(address) - && let Some(addr) = self.dns_overrides.load().resolve_socket_addr(&host, port) + && let Some(addr) = self.dns_resolver.resolve_socket_addr(&host, port) { return match tokio::time::timeout(connect_timeout, TcpStream::connect(addr)).await { Ok(Ok(stream)) => Ok(stream), diff --git a/src/web/http/websocket/driver/io.rs b/src/web/http/websocket/driver/io.rs index 220bed5..20223e8 100644 --- a/src/web/http/websocket/driver/io.rs +++ b/src/web/http/websocket/driver/io.rs @@ -53,6 +53,8 @@ pub(super) async fn reserve_data( } let notify = runtime.budget_notify(); let notified = notify.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); if let Some(budget) = runtime.try_websocket_data_budget(owner, bytes.max(1)) { return Ok(budget); } @@ -104,6 +106,8 @@ pub(super) async fn process_lane( } let notify = runtime.budget_notify(); let notified = notify.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); match session.process_websocket_lane(reservation, sequence, body) { Ok(progressed) => return Ok(progressed), Err(ManagerError::Backpressure) => {} @@ -135,6 +139,8 @@ where } let notify = runtime.budget_notify(); let notified = notify.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); match operation() { Ok(value) => return Ok(value), Err(ManagerError::Backpressure) => {} diff --git a/src/web/manager.rs b/src/web/manager.rs index 028813f..584a0ef 100644 --- a/src/web/manager.rs +++ b/src/web/manager.rs @@ -6,7 +6,7 @@ use std::time::Duration; use arc_swap::ArcSwap; use parking_lot::Mutex; -use tokio::sync::{Notify, OwnedSemaphorePermit, Semaphore}; +use tokio::sync::{Notify, OwnedSemaphorePermit, Semaphore, TryAcquireError}; use tokio_util::sync::CancellationToken; use tokio_util::task::TaskTracker; @@ -71,6 +71,15 @@ pub(crate) type TokenHash = [u8; TOKEN_BYTES]; /// Stable non-allocating key used for per-profile quotas. pub(crate) type ProfileKey = [u8; TOKEN_BYTES]; +/// Stable failure category for accepted-socket capacity admission. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum HttpConnectionAdmissionError { + /// The bounded connection plane currently has no free permit. + AtCapacity, + /// Terminal runtime shutdown closed the connection plane. + Closed, +} + /// WEB manager operation failure category. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub(crate) enum ManagerError { @@ -294,27 +303,38 @@ impl WebProcessRuntime { } /// Reserves one accepted HTTP connection. - pub(crate) fn try_http_connection(&self) -> Option { - let permit = Arc::clone(&self.http_connections).try_acquire_owned().ok(); - if permit.is_none() { - self.record_limit_hit(); + pub(crate) fn try_http_connection( + &self, + ) -> Result { + match Arc::clone(&self.http_connections).try_acquire_owned() { + Ok(permit) => Ok(permit), + Err(TryAcquireError::NoPermits) => { + self.record_limit_hit(); + Err(HttpConnectionAdmissionError::AtCapacity) + } + Err(TryAcquireError::Closed) => Err(HttpConnectionAdmissionError::Closed), } - permit } /// Waits for one accepted HTTP connection slot after bounded overload admission. - pub(crate) async fn acquire_http_connection(&self) -> Option { + pub(crate) async fn acquire_http_connection( + &self, + ) -> Result { Arc::clone(&self.http_connections) .acquire_owned() .await - .ok() + .map_err(|_| HttpConnectionAdmissionError::Closed) } /// Reserves one accepted socket outside ordinary HTTP connection capacity. - pub(crate) fn try_http_overload_connection(&self) -> Option { - Arc::clone(&self.http_overload_connections) - .try_acquire_owned() - .ok() + pub(crate) fn try_http_overload_connection( + &self, + ) -> Result { + match Arc::clone(&self.http_overload_connections).try_acquire_owned() { + Ok(permit) => Ok(permit), + Err(TryAcquireError::NoPermits) => Err(HttpConnectionAdmissionError::AtCapacity), + Err(TryAcquireError::Closed) => Err(HttpConnectionAdmissionError::Closed), + } } /// Reserves one concurrently executing HTTP request handler. diff --git a/src/web/manager/lifecycle.rs b/src/web/manager/lifecycle.rs index 236df4b..6f1cdec 100644 --- a/src/web/manager/lifecycle.rs +++ b/src/web/manager/lifecycle.rs @@ -448,7 +448,10 @@ mod tests { tokio::task::yield_now().await; let drain = runtime.begin_shutdown(); - assert!(runtime.try_http_connection().is_none()); + assert_eq!( + runtime.try_http_connection().unwrap_err(), + super::super::HttpConnectionAdmissionError::Closed + ); assert!(runtime.try_http_handler().is_none()); assert!(runtime.try_lane_poll(false).is_none()); assert_eq!( diff --git a/src/web/manager/operator_lifecycle.rs b/src/web/manager/operator_lifecycle.rs index d2b9fe8..06df139 100644 --- a/src/web/manager/operator_lifecycle.rs +++ b/src/web/manager/operator_lifecycle.rs @@ -1,5 +1,5 @@ use std::sync::Arc; -use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{Duration, Instant}; use arc_swap::ArcSwap; @@ -16,12 +16,14 @@ pub(crate) use status::{ OperatorDrainOutcome, OperatorDrainState, OperatorDrainStatus, OperatorLifecycleState, OperatorLifecycleStatus, }; +// Admission fencing and registration-drain synchronization. +mod admission; +use admission::{OperatorAdmission, OperatorAdmissionRejection}; +pub(super) use admission::OperatorRegistration; // Mutable and published lifecycle state storage. mod state; use state::{ActiveDrain, OperatorLifecycleInner, OperatorSnapshot, WorkCounts}; -const OPERATOR_ADMISSION_CLOSED: usize = 1 << (usize::BITS - 1); -const OPERATOR_REGISTRATION_COUNT: usize = OPERATOR_ADMISSION_CLOSED - 1; const DRAIN_REF_VERSION: &str = "wd1"; /// Stable operator-control rejection category. @@ -33,77 +35,7 @@ pub(crate) enum OperatorLifecycleError { OperationInProgress, } -struct OperatorAdmission { - state: AtomicUsize, - registrations_drained: Notify, -} - -pub(super) struct OperatorRegistration<'a> { - admission: &'a OperatorAdmission, -} - -impl OperatorAdmission { - fn new() -> Self { - Self { - state: AtomicUsize::new(0), - registrations_drained: Notify::new(), - } - } - - fn try_register(&self) -> Option> { - let mut state = self.state.load(Ordering::Acquire); - loop { - if state & OPERATOR_ADMISSION_CLOSED != 0 - || state & OPERATOR_REGISTRATION_COUNT == OPERATOR_REGISTRATION_COUNT - { - return None; - } - match self.state.compare_exchange_weak( - state, - state + 1, - Ordering::AcqRel, - Ordering::Acquire, - ) { - Ok(_) => return Some(OperatorRegistration { admission: self }), - Err(observed) => state = observed, - } - } - } - - fn close(&self) { - self.state - .fetch_or(OPERATOR_ADMISSION_CLOSED, Ordering::AcqRel); - } - - fn reopen(&self) { - self.state - .fetch_and(!OPERATOR_ADMISSION_CLOSED, Ordering::AcqRel); - } - - fn is_closed(&self) -> bool { - self.state.load(Ordering::Acquire) & OPERATOR_ADMISSION_CLOSED != 0 - } - - async fn wait_for_registrations(&self) { - loop { - let notified = self.registrations_drained.notified(); - if self.state.load(Ordering::Acquire) & OPERATOR_REGISTRATION_COUNT == 0 { - return; - } - notified.await; - } - } -} - -impl Drop for OperatorRegistration<'_> { - fn drop(&mut self) { - let previous = self.admission.state.fetch_sub(1, Ordering::AcqRel); - if previous & OPERATOR_REGISTRATION_COUNT == 1 { - self.admission.registrations_drained.notify_waiters(); - } - } -} - +/// Process-owned reversible lifecycle and admission authority. pub(super) struct OperatorLifecycle { runtime_instance: Arc, admission: OperatorAdmission, @@ -115,6 +47,7 @@ pub(super) struct OperatorLifecycle { } impl OperatorLifecycle { + /// Creates a running lifecycle for one immutable process instance. pub(super) fn new(runtime_instance: Arc) -> Self { let since = Instant::now(); let snapshot = OperatorSnapshot { @@ -142,16 +75,21 @@ impl OperatorLifecycle { } } - pub(super) fn try_register(&self) -> Option> { + /// Registers one synchronous operator-fenced admission section. + pub(super) fn try_register( + &self, + ) -> Result, crate::web::telemetry::WebRejectionReason> { self.admission.try_register() } + /// Wakes an active drain after tracked work ownership changes. pub(super) fn notify_work_changed(&self) { if self.admission.is_closed() { self.work_changed.notify_waiters(); } } + /// Returns the lock-free lifecycle snapshot with effective config admission. pub(super) fn status(&self, config_enabled: bool) -> OperatorLifecycleStatus { let snapshot = self.published.load(); let admission_open = @@ -170,30 +108,6 @@ impl OperatorLifecycle { self.published.load().terminal } - fn rejection_reason(&self) -> crate::web::telemetry::WebRejectionReason { - let inner = self.inner.lock(); - if inner.terminal { - return crate::web::telemetry::WebRejectionReason::RuntimeClosed; - } - match inner.state { - OperatorLifecycleState::Paused => { - crate::web::telemetry::WebRejectionReason::OperatorPaused - } - OperatorLifecycleState::Draining => { - crate::web::telemetry::WebRejectionReason::OperatorDraining - } - OperatorLifecycleState::ForceClosing => { - crate::web::telemetry::WebRejectionReason::OperatorForceClosing - } - OperatorLifecycleState::Drained => { - crate::web::telemetry::WebRejectionReason::OperatorDrained - } - OperatorLifecycleState::Running => { - crate::web::telemetry::WebRejectionReason::RuntimeClosed - } - } - } - fn publish_locked(&self, inner: &OperatorLifecycleInner) { self.published.store(Arc::new(OperatorSnapshot { state: inner.state, @@ -246,6 +160,8 @@ impl OperatorLifecycle { return false; } self.transition_locked(&mut inner, OperatorLifecycleState::ForceClosing); + self.admission + .close(OperatorAdmissionRejection::ForceClosing); let Some(drain) = inner.drain.as_mut() else { return false; }; @@ -269,6 +185,7 @@ impl OperatorLifecycle { } inner.active = None; self.transition_locked(&mut inner, OperatorLifecycleState::Drained); + self.admission.close(OperatorAdmissionRejection::Drained); if let Some(drain) = inner.drain.as_mut() { drain.state = OperatorDrainState::Completed; drain.outcome = Some(if forced { @@ -286,7 +203,8 @@ impl OperatorLifecycle { fn close_terminal(&self) { let mut inner = self.inner.lock(); - self.admission.close(); + self.admission + .close(OperatorAdmissionRejection::RuntimeClosed); if inner.terminal { return; } @@ -321,7 +239,6 @@ impl WebProcessRuntime { if inner.terminal || self.shutdown.is_cancelled() { return Err(OperatorLifecycleError::Closed); } - self.operator_lifecycle.admission.close(); if matches!( inner.state, OperatorLifecycleState::Running | OperatorLifecycleState::Drained @@ -329,6 +246,9 @@ impl WebProcessRuntime { self.operator_lifecycle .transition_locked(&mut inner, OperatorLifecycleState::Paused); } + self.operator_lifecycle.admission.close( + OperatorAdmissionRejection::for_state(inner.state), + ); self.operator_lifecycle.publish_locked(&inner); } self.operator_lifecycle @@ -355,7 +275,7 @@ impl WebProcessRuntime { let started = TokioInstant::now(); let deadline = started + timeout; let started_epoch_millis = crate::web::trace::store_epoch_millis(); - { + let accepted = { let mut inner = self.operator_lifecycle.inner.lock(); if inner.terminal || self.shutdown.is_cancelled() { return Err(OperatorLifecycleError::Closed); @@ -368,7 +288,6 @@ impl WebProcessRuntime { { return Err(OperatorLifecycleError::OperationInProgress); } - self.operator_lifecycle.admission.close(); inner.active = Some(ActiveDrain { sequence, cancellation: cancellation.clone(), @@ -392,26 +311,20 @@ impl WebProcessRuntime { }); self.operator_lifecycle .transition_locked(&mut inner, OperatorLifecycleState::Draining); + self.operator_lifecycle + .admission + .close(OperatorAdmissionRejection::Draining); + let runtime = Arc::clone(self); + self.spawn_auxiliary(async move { + runtime + .run_operator_drain(sequence, deadline, cancellation) + .await; + }); self.operator_lifecycle.publish_locked(&inner); - } - self.operator_lifecycle - .admission - .wait_for_registrations() - .await; - if self.shutdown.is_cancelled() || self.operator_lifecycle.is_terminal() { - return Err(OperatorLifecycleError::Closed); - } - let counts = self.operator_work_counts(); - if !self.operator_lifecycle.update_counts(sequence, counts) { - return Err(OperatorLifecycleError::Closed); - } - let runtime = Arc::clone(self); - self.spawn_auxiliary(async move { - runtime - .run_operator_drain(sequence, deadline, cancellation) - .await; - }); - Ok(self.operator_lifecycle_status()) + let config_enabled = self.active_generation().config().web.enabled; + self.operator_lifecycle.status(config_enabled) + }; + Ok(accepted) } /// Resumes operator admission and invalidates any active drain waiter. @@ -444,23 +357,25 @@ impl WebProcessRuntime { Ok(self.operator_lifecycle_status()) } + /// Registers one WEB manager admission commit under the operator fence. pub(super) fn try_operator_admission( &self, ) -> Result, super::ManagerError> { match self.operator_lifecycle.try_register() { - Some(registration) => Ok(registration), - None => { - self.telemetry - .record_rejection(self.operator_lifecycle.rejection_reason()); + Ok(registration) => Ok(registration), + Err(reason) => { + self.telemetry.record_rejection(reason); Err(super::ManagerError::AdmissionPaused) } } } + /// Notifies an active drain that tracked WEB work changed. pub(super) fn notify_operator_work_changed(&self) { self.operator_lifecycle.notify_work_changed(); } + /// Terminally closes operator lifecycle during process shutdown. pub(super) fn close_operator_lifecycle(&self) { self.operator_lifecycle.close_terminal(); } @@ -482,9 +397,16 @@ impl WebProcessRuntime { deadline: TokioInstant, cancellation: CancellationToken, ) { + tokio::select! { + biased; + _ = cancellation.cancelled() => return, + _ = self.operator_lifecycle.admission.wait_for_registrations() => {} + } let mut forced = false; loop { let notified = self.operator_lifecycle.work_changed.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); let counts = self.operator_work_counts(); if !self.operator_lifecycle.update_counts(sequence, counts) { return; @@ -497,14 +419,14 @@ impl WebProcessRuntime { tokio::select! { biased; _ = cancellation.cancelled() => return, - _ = notified => {} + _ = notified.as_mut() => {} } continue; } tokio::select! { biased; _ = cancellation.cancelled() => return, - _ = notified => {}, + _ = notified.as_mut() => {}, _ = tokio::time::sleep_until(deadline) => { let counts = self.operator_work_counts(); if counts.is_zero() { diff --git a/src/web/manager/operator_lifecycle/admission.rs b/src/web/manager/operator_lifecycle/admission.rs new file mode 100644 index 0000000..5462241 --- /dev/null +++ b/src/web/manager/operator_lifecycle/admission.rs @@ -0,0 +1,165 @@ +use std::sync::atomic::{AtomicUsize, Ordering}; + +use tokio::sync::Notify; + +use super::OperatorLifecycleState; + +const OPERATOR_ADMISSION_CLOSED: usize = 1 << (usize::BITS - 1); +const OPERATOR_REJECTION_BITS: u32 = 3; +const OPERATOR_REJECTION_SHIFT: u32 = usize::BITS - 1 - OPERATOR_REJECTION_BITS; +const OPERATOR_REJECTION_MASK: usize = + ((1usize << OPERATOR_REJECTION_BITS) - 1) << OPERATOR_REJECTION_SHIFT; +const OPERATOR_REGISTRATION_COUNT: usize = (1usize << OPERATOR_REJECTION_SHIFT) - 1; + +#[derive(Clone, Copy)] +#[repr(usize)] +/// Stable rejection encoded into the closed admission-fence word. +pub(super) enum OperatorAdmissionRejection { + /// Operator pause closed admission. + Paused = 1, + /// Graceful drain closed admission. + Draining = 2, + /// Deadline-triggered forced closure remains in progress. + ForceClosing = 3, + /// The latest drain reached confirmed zero. + Drained = 4, + /// Terminal process shutdown closed admission. + RuntimeClosed = 5, +} + +impl OperatorAdmissionRejection { + /// Maps a closed lifecycle state to its admission rejection. + pub(super) fn for_state(state: OperatorLifecycleState) -> Self { + match state { + OperatorLifecycleState::Paused => Self::Paused, + OperatorLifecycleState::Draining => Self::Draining, + OperatorLifecycleState::ForceClosing => Self::ForceClosing, + OperatorLifecycleState::Drained => Self::Drained, + OperatorLifecycleState::Running => Self::RuntimeClosed, + } + } + + fn from_admission_state(state: usize) -> Self { + match (state & OPERATOR_REJECTION_MASK) >> OPERATOR_REJECTION_SHIFT { + value if value == Self::Paused as usize => Self::Paused, + value if value == Self::Draining as usize => Self::Draining, + value if value == Self::ForceClosing as usize => Self::ForceClosing, + value if value == Self::Drained as usize => Self::Drained, + _ => Self::RuntimeClosed, + } + } + + fn telemetry_reason(self) -> crate::web::telemetry::WebRejectionReason { + match self { + Self::Paused => crate::web::telemetry::WebRejectionReason::OperatorPaused, + Self::Draining => crate::web::telemetry::WebRejectionReason::OperatorDraining, + Self::ForceClosing => { + crate::web::telemetry::WebRejectionReason::OperatorForceClosing + } + Self::Drained => crate::web::telemetry::WebRejectionReason::OperatorDrained, + Self::RuntimeClosed => crate::web::telemetry::WebRejectionReason::RuntimeClosed, + } + } +} + +/// Lock-free admission fence with bounded pre-cutover registration tracking. +pub(super) struct OperatorAdmission { + state: AtomicUsize, + registrations_drained: Notify, +} + +/// RAII ownership of one synchronous pre-cutover admission section. +pub(in crate::web::manager) struct OperatorRegistration<'a> { + admission: &'a OperatorAdmission, +} + +impl OperatorAdmission { + /// Creates an open admission fence. + pub(super) fn new() -> Self { + Self { + state: AtomicUsize::new(0), + registrations_drained: Notify::new(), + } + } + + /// Registers one pre-cutover section or returns its stable rejection reason. + pub(super) fn try_register( + &self, + ) -> Result, crate::web::telemetry::WebRejectionReason> { + let mut state = self.state.load(Ordering::Acquire); + loop { + if state & OPERATOR_ADMISSION_CLOSED != 0 { + return Err( + OperatorAdmissionRejection::from_admission_state(state).telemetry_reason(), + ); + } + if state & OPERATOR_REGISTRATION_COUNT == OPERATOR_REGISTRATION_COUNT { + return Err(crate::web::telemetry::WebRejectionReason::Concurrent); + } + match self.state.compare_exchange_weak( + state, + state + 1, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => return Ok(OperatorRegistration { admission: self }), + Err(observed) => state = observed, + } + } + } + + /// Atomically closes admission while preserving active registration count. + pub(super) fn close(&self, reason: OperatorAdmissionRejection) { + let reason = (reason as usize) << OPERATOR_REJECTION_SHIFT; + let mut state = self.state.load(Ordering::Acquire); + loop { + let next = (state & OPERATOR_REGISTRATION_COUNT) + | OPERATOR_ADMISSION_CLOSED + | reason; + match self.state.compare_exchange_weak( + state, + next, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => return, + Err(observed) => state = observed, + } + } + } + + /// Reopens admission without altering active registration ownership. + pub(super) fn reopen(&self) { + self.state.fetch_and( + !(OPERATOR_ADMISSION_CLOSED | OPERATOR_REJECTION_MASK), + Ordering::AcqRel, + ); + } + + /// Returns whether the admission fence is closed. + pub(super) fn is_closed(&self) -> bool { + self.state.load(Ordering::Acquire) & OPERATOR_ADMISSION_CLOSED != 0 + } + + /// Waits until every pre-cutover registration has been released. + pub(super) async fn wait_for_registrations(&self) { + loop { + let notified = self.registrations_drained.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); + if self.state.load(Ordering::Acquire) & OPERATOR_REGISTRATION_COUNT == 0 { + return; + } + notified.await; + } + } +} + +impl Drop for OperatorRegistration<'_> { + fn drop(&mut self) { + let previous = self.admission.state.fetch_sub(1, Ordering::AcqRel); + if previous & OPERATOR_REGISTRATION_COUNT == 1 { + self.admission.registrations_drained.notify_waiters(); + } + } +} diff --git a/src/web/manager/operator_lifecycle/tests.rs b/src/web/manager/operator_lifecycle/tests.rs index a4d7dce..2ff32b7 100644 --- a/src/web/manager/operator_lifecycle/tests.rs +++ b/src/web/manager/operator_lifecycle/tests.rs @@ -133,6 +133,43 @@ async fn empty_drain_completes_gracefully_and_stays_closed_until_resume() { stop_runtime(runtime, generation).await; } +#[tokio::test] +async fn drain_request_returns_after_registering_its_worker() { + let (runtime, generation) = test_runtime(); + let registration = runtime.try_operator_admission().unwrap(); + let drain_runtime = Arc::clone(&runtime); + let drain = tokio::spawn(async move { + drain_runtime + .drain_operator(Duration::from_secs(30)) + .await + }); + + tokio::time::timeout(Duration::from_secs(1), async { + while runtime.operator_lifecycle_status().state != OperatorLifecycleState::Draining { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + let accepted = tokio::time::timeout(Duration::from_secs(1), drain) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!(accepted.state, OperatorLifecycleState::Draining); + drop(registration); + + tokio::time::timeout(Duration::from_secs(1), async { + while runtime.operator_lifecycle_status().state != OperatorLifecycleState::Drained { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + + stop_runtime(runtime, generation).await; +} + #[tokio::test] async fn resume_after_force_commit_cancels_wait_but_preserves_force_evidence() { let (runtime, generation) = test_runtime(); @@ -140,7 +177,10 @@ async fn resume_after_force_commit_cancels_wait_but_preserves_force_evidence() { let cancellation = CancellationToken::new(); { let mut inner = runtime.operator_lifecycle.inner.lock(); - runtime.operator_lifecycle.admission.close(); + runtime + .operator_lifecycle + .admission + .close(OperatorAdmissionRejection::Draining); inner.active = Some(ActiveDrain { sequence, cancellation, diff --git a/src/web/session/downlink.rs b/src/web/session/downlink.rs index 1147296..45df0c5 100644 --- a/src/web/session/downlink.rs +++ b/src/web/session/downlink.rs @@ -64,6 +64,8 @@ impl WebSession { let poll = async { loop { let notified = self.down_notify.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); { let mut state = self.state.lock(); if state.down_epoch != epoch { diff --git a/src/web/session/lanes.rs b/src/web/session/lanes.rs index cd0a184..cc6748a 100644 --- a/src/web/session/lanes.rs +++ b/src/web/session/lanes.rs @@ -151,6 +151,8 @@ impl WebSession { let poll = async { loop { let notified = notify.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); { let mut state = self.state.lock(); if state.closed { @@ -306,6 +308,8 @@ impl WebSession { let opened = tokio::time::timeout(deadline, async { loop { let notified = self.lane_open_notify.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); { let state = self.state.lock(); if state.closed { diff --git a/src/web/session/lifecycle.rs b/src/web/session/lifecycle.rs index 7ed8e99..e2fc2b3 100644 --- a/src/web/session/lifecycle.rs +++ b/src/web/session/lifecycle.rs @@ -77,6 +77,8 @@ impl WebSession { pub(crate) async fn wait(&self) { loop { let notified = self.tasks_done.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); if self.tasks_live.load(Ordering::Acquire) == 0 { return; } diff --git a/src/web/telemetry.rs b/src/web/telemetry.rs index c23b08c..d06d6c2 100644 --- a/src/web/telemetry.rs +++ b/src/web/telemetry.rs @@ -4,6 +4,10 @@ use std::time::Instant; use serde::Serialize; +const LAST_DECOY_OUTCOME_BITS: u32 = 4; +const LAST_DECOY_OUTCOME_MASK: u64 = (1 << LAST_DECOY_OUTCOME_BITS) - 1; +const LAST_DECOY_ELAPSED_MAX: u64 = u64::MAX >> LAST_DECOY_OUTCOME_BITS; + /// Stable operational rejection reason recorded at the decision point. #[derive(Clone, Copy, Debug, PartialEq, Eq)] #[repr(usize)] @@ -296,8 +300,7 @@ pub(crate) struct WebTelemetry { rejections: [AtomicU64; WebRejectionReason::ALL.len()], overload_outcomes: [AtomicU64; WebHttpConnectionOverloadOutcome::ALL.len()], decoy_outcomes: [AtomicU64; WebDecoyUpstreamOutcome::ALL.len()], - last_decoy_outcome: AtomicUsize, - last_decoy_elapsed_ms: AtomicU64, + last_decoy: AtomicU64, sessions_created: AtomicU64, sessions_closed: AtomicU64, streams_opened: AtomicU64, @@ -318,8 +321,7 @@ impl WebTelemetry { rejections: std::array::from_fn(|_| AtomicU64::new(0)), overload_outcomes: std::array::from_fn(|_| AtomicU64::new(0)), decoy_outcomes: std::array::from_fn(|_| AtomicU64::new(0)), - last_decoy_outcome: AtomicUsize::new(usize::MAX), - last_decoy_elapsed_ms: AtomicU64::new(0), + last_decoy: AtomicU64::new(0), sessions_created: AtomicU64::new(0), sessions_closed: AtomicU64::new(0), streams_opened: AtomicU64::new(0), @@ -408,11 +410,13 @@ impl WebTelemetry { /// Records one internal plain-HTTP decoy origin outcome. pub(crate) fn record_decoy(&self, outcome: WebDecoyUpstreamOutcome) { self.decoy_outcomes[outcome as usize].fetch_add(1, Ordering::Relaxed); - let elapsed_ms = self.started.elapsed().as_millis().min(u128::from(u64::MAX)) as u64; - self.last_decoy_elapsed_ms - .store(elapsed_ms.saturating_add(1), Ordering::Relaxed); - self.last_decoy_outcome - .store(outcome as usize, Ordering::Release); + let elapsed_ms = self + .started + .elapsed() + .as_millis() + .min(u128::from(LAST_DECOY_ELAPSED_MAX)) as u64; + let packed = (elapsed_ms << LAST_DECOY_OUTCOME_BITS) | (outcome as u64 + 1); + self.last_decoy.store(packed, Ordering::Release); } /// Returns one fixed internal decoy origin counter. @@ -433,14 +437,16 @@ impl WebTelemetry { /// Returns the last decoy outcome and its monotonic age in milliseconds. pub(crate) fn last_decoy(&self) -> Option<(&'static str, u64)> { - let raw = self.last_decoy_outcome.load(Ordering::Acquire); - let outcome = WebDecoyUpstreamOutcome::ALL.get(raw).copied()?; - let recorded = self.last_decoy_elapsed_ms.load(Ordering::Relaxed); - let now = self.started.elapsed().as_millis().min(u128::from(u64::MAX)) as u64; - Some(( - outcome.as_str(), - now.saturating_sub(recorded.saturating_sub(1)), - )) + let packed = self.last_decoy.load(Ordering::Acquire); + let outcome_index = (packed & LAST_DECOY_OUTCOME_MASK).checked_sub(1)? as usize; + let outcome = WebDecoyUpstreamOutcome::ALL.get(outcome_index).copied()?; + let recorded = packed >> LAST_DECOY_OUTCOME_BITS; + let now = self + .started + .elapsed() + .as_millis() + .min(u128::from(LAST_DECOY_ELAPSED_MAX)) as u64; + Some((outcome.as_str(), now.saturating_sub(recorded))) } /// Records one created session incarnation.