From 991d5b2c38f3a51c187b9f75c9543204d4533d9d Mon Sep 17 00:00:00 2001 From: Alexey <247128645+axkurcom@users.noreply.github.com> Date: Fri, 17 Jul 2026 23:56:00 +0300 Subject: [PATCH] Maestro: add in-process runtime generation reload Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com> --- src/api/config_edit.rs | 53 ++- src/api/config_store.rs | 15 + src/api/mod.rs | 145 ++++++- src/config/hot_reload.rs | 10 +- src/conntrack_control.rs | 12 +- src/maestro/admission.rs | 5 +- src/maestro/generation.rs | 221 +++++++++++ src/maestro/listeners.rs | 239 +++--------- src/maestro/listeners/unix.rs | 117 ++++++ src/maestro/mod.rs | 156 ++++---- src/maestro/reload.rs | 464 +++++++++++++++++++++++ src/maestro/reload_supervisor.rs | 197 ++++++++++ src/maestro/runtime_build.rs | 357 +++++++++++++++++ src/maestro/runtime_tasks.rs | 151 ++++---- src/maestro/shutdown.rs | 40 +- src/maestro/tls_bootstrap.rs | 9 +- src/metrics.rs | 110 ++---- src/network/dns_overrides.rs | 43 +++ src/proxy/direct_buffer_budget.rs | 190 +++++----- src/proxy/relay/io/quota.rs | 15 +- src/stats/helpers.rs | 11 +- src/stats/mod.rs | 49 +-- src/stats/quota_store.rs | 147 +++++++ src/stats/users.rs | 40 +- src/transport/middle_proxy/http_fetch.rs | 22 +- src/transport/upstream.rs | 41 +- 26 files changed, 2207 insertions(+), 652 deletions(-) create mode 100644 src/maestro/generation.rs create mode 100644 src/maestro/listeners/unix.rs create mode 100644 src/maestro/reload.rs create mode 100644 src/maestro/reload_supervisor.rs create mode 100644 src/maestro/runtime_build.rs create mode 100644 src/stats/quota_store.rs diff --git a/src/api/config_edit.rs b/src/api/config_edit.rs index 8f25cfc..7b432fd 100644 --- a/src/api/config_edit.rs +++ b/src/api/config_edit.rs @@ -6,19 +6,28 @@ use toml::Value as Toml; use super::ApiShared; use super::config_store::{ - EDITABLE_SECTIONS, compute_revision, current_revision, save_sections_to_disk, + EDITABLE_SECTIONS, compute_revision, current_revision, load_config_from_disk, + save_sections_to_disk, }; use super::model::ApiFailure; use crate::config::ProxyConfig; use crate::config::hot_reload::classify_config_changes; +use crate::maestro::reload::{ReloadAccepted, ReloadRequest, ReloadSubmitError}; +use crate::maestro::runtime_build::deferred_process_fields; use serde::Serialize; use std::path::Path; +use std::sync::Arc; #[derive(Debug, Serialize)] pub(super) struct PatchConfigResponse { pub revision: String, pub restart_required: bool, + pub runtime_reload_required: bool, + pub process_restart_required: bool, + pub deferred_process_fields: Vec, pub changed: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + pub reload: Option, } /// Shared-state wrapper around [`apply_patch_to_path`]: serializes config @@ -27,10 +36,40 @@ pub(super) struct PatchConfigResponse { pub(super) async fn patch_config( patch_json: Json, expected_revision: Option, + reload_request: Option, shared: &ApiShared, ) -> Result { let _guard = shared.mutation_lock.lock().await; - let resp = apply_patch_to_path(&shared.config_path, &patch_json, expected_revision).await?; + if reload_request.is_some() + && let Some(reload_id) = shared.reload_control.in_progress().await + { + return Err(ApiFailure::new( + hyper::StatusCode::CONFLICT, + "reload_in_progress", + format!("Reload {} is already in progress", reload_id), + )); + } + let mut resp = apply_patch_to_path(&shared.config_path, &patch_json, expected_revision).await?; + if let Some(request) = reload_request { + let config = Arc::new(load_config_from_disk(&shared.config_path).await?); + let accepted = shared + .reload_control + .submit(config, resp.revision.clone(), request) + .await + .map_err(|error| match error { + ReloadSubmitError::InProgress(reload_id) => ApiFailure::new( + hyper::StatusCode::CONFLICT, + "reload_in_progress", + format!("Reload {} is already in progress", reload_id), + ), + ReloadSubmitError::MaestroUnavailable => ApiFailure::new( + hyper::StatusCode::SERVICE_UNAVAILABLE, + "maestro_unavailable", + "Maestro reload coordinator is unavailable", + ), + })?; + resp.reload = Some(accepted); + } drop(_guard); shared .runtime_events @@ -114,6 +153,7 @@ pub(super) async fn apply_patch_to_path( // 4. classify changes (Telemt's own hot/restart rule) let class = classify_config_changes(&old_cfg, &new_cfg); + let deferred_process_fields = deferred_process_fields(&old_cfg, &new_cfg); // 5. write only the touched top-level sections let revision = save_sections_to_disk(config_path, &new_cfg, &touched).await?; @@ -121,7 +161,11 @@ pub(super) async fn apply_patch_to_path( Ok(PatchConfigResponse { revision, restart_required: class.restart_required, + runtime_reload_required: class.restart_required, + process_restart_required: !deferred_process_fields.is_empty(), + deferred_process_fields, changed: class.changed, + reload: None, }) } @@ -266,6 +310,9 @@ mod tests { let patch: Json = serde_json::json!({"censorship": {"tls_domain": "b.com"}}); let resp = apply_patch_to_path(&path, &patch, None).await.unwrap(); assert!(resp.restart_required); + assert!(resp.runtime_reload_required); + assert!(!resp.process_restart_required); + assert!(resp.deferred_process_fields.is_empty()); assert!(resp.changed.iter().any(|c| c == "censorship")); let written = std::fs::read_to_string(&path).unwrap(); assert!(written.contains("tls_domain = \"b.com\"")); @@ -406,6 +453,8 @@ mod tests { let patch: Json = serde_json::json!({"general": {"log_level": "debug"}}); let resp = apply_patch_to_path(&path, &patch, None).await.unwrap(); assert!(!resp.restart_required); + assert!(!resp.runtime_reload_required); + assert!(!resp.process_restart_required); assert!(resp.changed.iter().any(|c| c == "general")); } } diff --git a/src/api/config_store.rs b/src/api/config_store.rs index cc5dc25..9a28537 100644 --- a/src/api/config_store.rs +++ b/src/api/config_store.rs @@ -72,6 +72,13 @@ pub(super) async fn current_revision(config_path: &Path) -> Result Result { + let content = tokio::fs::read_to_string(config_path) + .await + .map_err(|error| format!("failed to read config: {}", error))?; + Ok(compute_revision(&content)) +} + pub(super) fn compute_revision(content: &str) -> String { let mut hasher = Sha256::new(); hasher.update(content.as_bytes()); @@ -86,6 +93,14 @@ pub(super) async fn load_config_from_disk(config_path: &Path) -> Result Result { + let config_path = config_path.to_path_buf(); + tokio::task::spawn_blocking(move || ProxyConfig::load(config_path)) + .await + .map_err(|error| ApiFailure::internal(format!("failed to join config loader: {}", error)))? + .map_err(|error| ApiFailure::bad_request(format!("invalid runtime config: {}", error))) +} + #[allow(dead_code)] pub(super) async fn save_config_to_disk( config_path: &Path, diff --git a/src/api/mod.rs b/src/api/mod.rs index 8724416..1677a0c 100644 --- a/src/api/mod.rs +++ b/src/api/mod.rs @@ -7,6 +7,7 @@ use std::sync::Arc; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::time::Duration; +use arc_swap::ArcSwap; use http_body_util::Full; use hyper::body::{Bytes, Incoming}; use hyper::header::AUTHORIZATION; @@ -21,6 +22,8 @@ use tracing::{debug, info, warn}; use crate::config::{ApiGrayAction, ProxyConfig}; use crate::ip_tracker::UserIpTracker; +use crate::maestro::generation::RuntimeGeneration; +use crate::maestro::reload::{ReloadControl, ReloadRequest, ReloadSubmitError}; use crate::proxy::route_mode::RouteRuntimeController; use crate::proxy::shared_state::ProxySharedState; use crate::startup::StartupTracker; @@ -29,7 +32,7 @@ use crate::transport::UpstreamManager; use crate::transport::middle_proxy::MePool; mod config_edit; -mod config_store; +pub(crate) mod config_store; mod events; mod http_utils; mod model; @@ -44,7 +47,8 @@ mod runtime_zero; mod users; use config_store::{ - current_revision, ensure_expected_revision, load_config_from_disk, parse_if_match, + current_revision, ensure_expected_revision, load_config_for_reload, load_config_from_disk, + parse_if_match, }; use events::ApiEventStore; use http_utils::{error_response, read_json, read_optional_json, success_response}; @@ -107,12 +111,15 @@ pub(super) struct ApiShared { pub(super) minimal_cache: Arc>>, pub(super) runtime_edge_connections_cache: Arc>>, pub(super) runtime_edge_recompute_lock: Arc>, + pub(super) cache_generation: Arc, pub(super) runtime_events: Arc, pub(super) request_id: Arc, pub(super) runtime_state: Arc, pub(super) startup_tracker: Arc, pub(super) route_runtime: Arc, pub(super) proxy_shared: Arc, + pub(super) reload_control: ReloadControl, + pub(super) active_runtime: Arc>, } impl ApiShared { @@ -123,6 +130,31 @@ impl ApiShared { fn detected_link_ips(&self) -> (Option, Option) { *self.detected_ips_rx.borrow() } + + fn for_runtime(&self, runtime: &RuntimeGeneration) -> Self { + Self { + stats: runtime.stats.clone(), + ip_tracker: runtime.ip_tracker.clone(), + 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(), + detected_ips_rx: self.detected_ips_rx.clone(), + mutation_lock: self.mutation_lock.clone(), + minimal_cache: self.minimal_cache.clone(), + runtime_edge_connections_cache: self.runtime_edge_connections_cache.clone(), + runtime_edge_recompute_lock: self.runtime_edge_recompute_lock.clone(), + cache_generation: self.cache_generation.clone(), + runtime_events: self.runtime_events.clone(), + request_id: self.request_id.clone(), + runtime_state: self.runtime_state.clone(), + startup_tracker: self.startup_tracker.clone(), + route_runtime: runtime.route_runtime.clone(), + proxy_shared: runtime.proxy_shared.clone(), + reload_control: self.reload_control.clone(), + active_runtime: self.active_runtime.clone(), + } + } } fn auth_header_matches(actual: &str, expected: &str) -> bool { @@ -144,6 +176,12 @@ fn user_action_route_matches(path: &str, suffix: &str) -> bool { .unwrap_or(false) } +fn reload_status_route_id(path: &str) -> Option { + path.strip_prefix("/v1/system/reload/") + .filter(|id| !id.is_empty() && !id.contains('/')) + .and_then(|id| id.parse().ok()) +} + fn allowed_methods_for_path(path: &str) -> Option<&'static str> { match path { "/v1/health" @@ -175,12 +213,14 @@ fn allowed_methods_for_path(path: &str) -> Option<&'static str> { | "/v1/stats/users/active-ips" | "/v1/stats/users/quota" | "/v1/stats/users" => Some(ALLOW_GET), + "/v1/system/reload" => Some(ALLOW_POST), "/v1/users" => Some(ALLOW_GET_POST), "/v1/config" => Some(ALLOW_GET_PATCH), _ if user_action_route_matches(path, "/reset-quota") => Some(ALLOW_POST), _ if user_action_route_matches(path, "/rotate-secret") => Some(ALLOW_POST), _ if user_action_route_matches(path, "/enable") => Some(ALLOW_POST), _ if user_action_route_matches(path, "/disable") => Some(ALLOW_POST), + _ if reload_status_route_id(path).is_some() => Some(ALLOW_GET), _ if path .strip_prefix("/v1/users/") .map(|user| !user.is_empty() && !user.contains('/')) @@ -207,7 +247,18 @@ pub async fn serve( detected_ips_rx: watch::Receiver<(Option, Option)>, process_started_at_epoch_secs: u64, startup_tracker: Arc, + reload_control: ReloadControl, + mut active_runtime_rx: watch::Receiver>>>, ) { + let active_runtime = loop { + if let Some(active_runtime) = active_runtime_rx.borrow().clone() { + break active_runtime; + } + if active_runtime_rx.changed().await.is_err() { + warn!("Runtime generation channel closed before API bootstrap"); + return; + } + }; let listener = match TcpListener::bind(listen).await { Ok(listener) => listener, Err(error) => { @@ -241,6 +292,7 @@ pub async fn serve( minimal_cache: Arc::new(Mutex::new(None)), runtime_edge_connections_cache: Arc::new(Mutex::new(None)), runtime_edge_recompute_lock: Arc::new(Mutex::new(())), + cache_generation: Arc::new(AtomicU64::new(1)), runtime_events: Arc::new(ApiEventStore::new( config_rx.borrow().server.api.runtime_edge_events_capacity, )), @@ -249,6 +301,8 @@ pub async fn serve( startup_tracker, route_runtime, proxy_shared, + reload_control, + active_runtime, }); spawn_runtime_watchers( @@ -282,13 +336,11 @@ pub async fn serve( }; let shared_conn = shared.clone(); - let config_rx_conn = config_rx.clone(); tokio::spawn(async move { let _connection_permit = connection_permit; let svc = service_fn(move |req: Request| { let shared_req = shared_conn.clone(); - let config_rx_req = config_rx_conn.clone(); - async move { handle(req, peer, shared_req, config_rx_req).await } + async move { handle(req, peer, shared_req).await } }); match timeout( API_HTTP_CONNECTION_TIMEOUT, @@ -318,8 +370,19 @@ async fn handle( req: Request, peer: SocketAddr, shared: Arc, - config_rx: watch::Receiver>, ) -> Result>, IoError> { + let runtime = shared.active_runtime.load_full(); + let previous_cache_generation = shared.cache_generation.swap(runtime.id, Ordering::AcqRel); + if previous_cache_generation != runtime.id { + *shared.minimal_cache.lock().await = None; + *shared.runtime_edge_connections_cache.lock().await = None; + } + let shared = Arc::new(shared.for_runtime(runtime.as_ref())); + let config_rx = runtime.config_rx.clone(); + shared + .runtime_state + .admission_open + .store(*runtime.admission_rx.borrow(), Ordering::Relaxed); let request_id = shared.next_request_id(); let cfg = config_rx.borrow().clone(); let api_cfg = &cfg.server.api; @@ -651,6 +714,45 @@ async fn handle( config_edit::read_managed_config(&shared.config_path).await?; Ok(success_response(StatusCode::OK, value, revision)) } + ("POST", "/v1/system/reload") => { + if api_cfg.read_only { + return Ok(error_response( + request_id, + ApiFailure::new( + StatusCode::FORBIDDEN, + "read_only", + "API runs in read-only mode", + ), + )); + } + let expected_revision = parse_if_match(req.headers()); + let request = read_optional_json::(req.into_body(), body_limit) + .await? + .unwrap_or_default(); + request.validate().map_err(ApiFailure::bad_request)?; + + let _guard = shared.mutation_lock.lock().await; + ensure_expected_revision(&shared.config_path, expected_revision.as_deref()).await?; + let revision = current_revision(&shared.config_path).await?; + let config = Arc::new(load_config_for_reload(&shared.config_path).await?); + let accepted = shared + .reload_control + .submit(config, revision.clone(), request) + .await + .map_err(|error| match error { + ReloadSubmitError::InProgress(reload_id) => ApiFailure::new( + StatusCode::CONFLICT, + "reload_in_progress", + format!("Reload {} is already in progress", reload_id), + ), + ReloadSubmitError::MaestroUnavailable => ApiFailure::new( + StatusCode::SERVICE_UNAVAILABLE, + "maestro_unavailable", + "Maestro reload coordinator is unavailable", + ), + })?; + Ok(success_response(StatusCode::ACCEPTED, accepted, revision)) + } ("PATCH", "/v1/config") => { if api_cfg.read_only { return Ok(error_response( @@ -663,11 +765,20 @@ async fn handle( )); } let expected_revision = parse_if_match(req.headers()); + let reload_request = + ReloadRequest::from_query(query.as_deref()).map_err(ApiFailure::bad_request)?; let body = read_json::(req.into_body(), body_limit).await?; - match config_edit::patch_config(body, expected_revision, &shared).await { + match config_edit::patch_config(body, expected_revision, reload_request, &shared) + .await + { Ok(resp) => { let revision = resp.revision.clone(); - Ok(success_response(StatusCode::OK, resp, revision)) + let status = if resp.reload.is_some() { + StatusCode::ACCEPTED + } else { + StatusCode::OK + }; + Ok(success_response(status, resp, revision)) } Err(error) => { shared @@ -678,6 +789,24 @@ async fn handle( } } _ => { + if method == Method::GET + && let Some(reload_id) = reload_status_route_id(normalized_path) + { + let revision = current_revision(&shared.config_path).await?; + let status = + shared + .reload_control + .status(reload_id) + .await + .ok_or_else(|| { + ApiFailure::new( + StatusCode::NOT_FOUND, + "reload_not_found", + format!("Reload {} was not found", reload_id), + ) + })?; + return Ok(success_response(StatusCode::OK, status, revision)); + } if method == Method::POST && let Some(base_user) = normalized_path .strip_prefix("/v1/users/") diff --git a/src/config/hot_reload.rs b/src/config/hot_reload.rs index e89ff95..1a4ab2c 100644 --- a/src/config/hot_reload.rs +++ b/src/config/hot_reload.rs @@ -1403,11 +1403,13 @@ fn reload_config( /// `detected_ip_v4` / `detected_ip_v6` are the IPs discovered during the /// startup probe — used when generating proxy links for newly added users, /// matching the same logic as the startup output. +/// The watcher releases its notify and signal resources when `cancellation` fires. pub fn spawn_config_watcher( config_path: PathBuf, initial: Arc, detected_ip_v4: Option, detected_ip_v6: Option, + cancellation: tokio_util::sync::CancellationToken, ) -> (watch::Receiver>, watch::Receiver) { let initial_level = initial.general.log_level.clone(); let (config_tx, config_rx) = watch::channel(initial); @@ -1515,10 +1517,14 @@ pub fn spawn_config_watcher( _ = sighup.recv() => { info!("SIGHUP received — reloading {:?}", config_path); } + _ = cancellation.cancelled() => break, } #[cfg(not(unix))] - if notify_rx.recv().await.is_none() { - break; + tokio::select! { + msg = notify_rx.recv() => { + if msg.is_none() { break; } + } + _ = cancellation.cancelled() => break, } // Debounce: drain extra events that arrive within a short quiet window. diff --git a/src/conntrack_control.rs b/src/conntrack_control.rs index dfbf5fc..1cb5b7d 100644 --- a/src/conntrack_control.rs +++ b/src/conntrack_control.rs @@ -7,6 +7,7 @@ use std::time::Duration; use tokio::io::AsyncWriteExt; use tokio::process::Command; use tokio::sync::{mpsc, watch}; +use tokio_util::sync::CancellationToken; use tracing::{debug, info, warn}; use crate::config::{ConntrackBackend, ConntrackMode, ProxyConfig}; @@ -57,10 +58,11 @@ impl PressureState { } } -pub(crate) fn spawn_conntrack_controller( +pub(crate) async fn run_conntrack_controller( config_rx: watch::Receiver>, stats: Arc, shared: Arc, + cancellation: CancellationToken, ) { if !cfg!(target_os = "linux") { let cfg = config_rx.borrow(); @@ -87,16 +89,15 @@ pub(crate) fn spawn_conntrack_controller( let (tx, rx) = mpsc::channel(CONNTRACK_EVENT_QUEUE_CAPACITY); shared.set_conntrack_close_sender(tx); - tokio::spawn(async move { - run_conntrack_controller(config_rx, stats, shared, rx).await; - }); + run_conntrack_controller_worker(config_rx, stats, shared, rx, cancellation).await; } -async fn run_conntrack_controller( +async fn run_conntrack_controller_worker( mut config_rx: watch::Receiver>, stats: Arc, shared: Arc, mut close_rx: mpsc::Receiver, + cancellation: CancellationToken, ) { let mut cfg = config_rx.borrow().clone(); let mut pressure_state = PressureState::new(stats.as_ref()); @@ -115,6 +116,7 @@ async fn run_conntrack_controller( loop { tokio::select! { + _ = cancellation.cancelled() => break, changed = config_rx.changed() => { if changed.is_err() { break; diff --git a/src/maestro/admission.rs b/src/maestro/admission.rs index fb25c45..b6aadc3 100644 --- a/src/maestro/admission.rs +++ b/src/maestro/admission.rs @@ -8,6 +8,8 @@ use crate::config::ProxyConfig; use crate::proxy::route_mode::{RelayRouteMode, RouteRuntimeController}; use crate::transport::middle_proxy::MePool; +use super::generation::RuntimeTaskScope; + const STARTUP_FALLBACK_AFTER: Duration = Duration::from_secs(80); const RUNTIME_FALLBACK_AFTER: Duration = Duration::from_secs(6); @@ -19,6 +21,7 @@ pub(crate) async fn configure_admission_gate( admission_tx: &watch::Sender, config_rx: watch::Receiver>, me_ready_rx: watch::Receiver, + task_scope: RuntimeTaskScope, ) { if config.general.use_middle_proxy { if me_pool.is_some() || config.general.me2dc_fallback { @@ -64,7 +67,7 @@ pub(crate) async fn configure_admission_gate( let mut config_rx_gate = config_rx.clone(); let mut me_ready_rx_gate = me_ready_rx; let mut admission_poll_ms = config.general.me_admission_poll_ms.max(1); - tokio::spawn(async move { + task_scope.spawn(async move { let mut gate_open = initial_gate_open; let mut route_mode = initial_route_mode; let mut ready_observed = initial_ready; diff --git a/src/maestro/generation.rs b/src/maestro/generation.rs new file mode 100644 index 0000000..f4f52de --- /dev/null +++ b/src/maestro/generation.rs @@ -0,0 +1,221 @@ +use std::future::Future; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::Duration; + +use tokio::sync::{RwLock, Semaphore, watch}; +use tokio_util::sync::CancellationToken; +use tokio_util::task::TaskTracker; + +use crate::config::ProxyConfig; +use crate::crypto::SecureRandom; +use crate::ip_tracker::UserIpTracker; +use crate::proxy::route_mode::RouteRuntimeController; +use crate::proxy::shared_state::ProxySharedState; +use crate::stats::beobachten::BeobachtenStore; +use crate::stats::{ReplayChecker, Stats}; +use crate::stream::BufferPool; +use crate::tls_front::TlsFrontCache; +use crate::transport::UpstreamManager; +use crate::transport::middle_proxy::MePool; + +const SESSION_STOP_TIMEOUT: Duration = Duration::from_secs(5); +const BACKGROUND_STOP_TIMEOUT: Duration = Duration::from_secs(5); + +/// Cancellation and join ownership for one generation's background tasks. +#[derive(Clone)] +pub(crate) struct RuntimeTaskScope { + tracker: TaskTracker, + cancel: CancellationToken, +} + +impl RuntimeTaskScope { + pub(crate) fn new() -> Self { + Self { + tracker: TaskTracker::new(), + cancel: CancellationToken::new(), + } + } + + pub(crate) fn spawn(&self, future: F) + where + F: Future + Send + 'static, + { + let cancel = self.cancel.clone(); + self.tracker.spawn(async move { + tokio::select! { + _ = cancel.cancelled() => {} + _ = future => {} + } + }); + } + + pub(crate) fn cancellation_token(&self) -> CancellationToken { + self.cancel.clone() + } + + pub(crate) async fn stop(&self) { + self.cancel.cancel(); + self.tracker.close(); + let _ = tokio::time::timeout(BACKGROUND_STOP_TIMEOUT, self.tracker.wait()).await; + } +} + +/// Runtime-owned data plane and control-plane dependencies for one generation. +pub(crate) struct RuntimeGeneration { + pub(crate) id: u64, + pub(crate) config_rx: watch::Receiver>, + pub(crate) admission_rx: watch::Receiver, + pub(crate) stats: Arc, + pub(crate) upstream_manager: Arc, + pub(crate) replay_checker: Arc, + pub(crate) buffer_pool: Arc, + pub(crate) rng: Arc, + pub(crate) me_pool: Option>, + pub(crate) me_pool_runtime: Arc>>>, + pub(crate) route_runtime: Arc, + pub(crate) tls_cache: Option>, + pub(crate) ip_tracker: Arc, + pub(crate) beobachten: Arc, + pub(crate) proxy_shared: Arc, + pub(crate) max_connections: Arc, + background_tasks: RuntimeTaskScope, + sessions: TaskTracker, + session_cancel: CancellationToken, + accepting_sessions: AtomicBool, +} + +impl RuntimeGeneration { + #[allow(clippy::too_many_arguments)] + pub(crate) fn new( + id: u64, + config_rx: watch::Receiver>, + admission_rx: watch::Receiver, + stats: Arc, + upstream_manager: Arc, + replay_checker: Arc, + buffer_pool: Arc, + rng: Arc, + me_pool: Option>, + me_pool_runtime: Arc>>>, + route_runtime: Arc, + tls_cache: Option>, + ip_tracker: Arc, + beobachten: Arc, + proxy_shared: Arc, + max_connections: Arc, + background_tasks: RuntimeTaskScope, + ) -> Arc { + Arc::new(Self { + id, + config_rx, + admission_rx, + stats, + upstream_manager, + replay_checker, + buffer_pool, + rng, + me_pool, + me_pool_runtime, + route_runtime, + tls_cache, + ip_tracker, + beobachten, + proxy_shared, + max_connections, + background_tasks, + sessions: TaskTracker::new(), + session_cancel: CancellationToken::new(), + accepting_sessions: AtomicBool::new(true), + }) + } + + pub(crate) fn config(&self) -> Arc { + self.config_rx.borrow().clone() + } + + pub(crate) async fn current_me_pool(&self) -> Option> { + if let Some(pool) = &self.me_pool { + return Some(pool.clone()); + } + self.me_pool_runtime.read().await.clone() + } + + pub(crate) fn spawn_session(&self, future: F) -> bool + where + F: Future + Send + 'static, + { + if !self.accepting_sessions.load(Ordering::Acquire) { + return false; + } + let cancel = self.session_cancel.clone(); + self.sessions.spawn(async move { + tokio::select! { + _ = cancel.cancelled() => {} + _ = future => {} + } + }); + true + } + + pub(crate) fn stop_accepting_sessions(&self) { + self.accepting_sessions.store(false, Ordering::Release); + } + + pub(crate) fn resume_accepting_sessions(&self) { + self.accepting_sessions.store(true, Ordering::Release); + } + + pub(crate) async fn drain_sessions(&self, timeout: Duration) -> bool { + self.stop_accepting_sessions(); + self.sessions.close(); + if tokio::time::timeout(timeout, self.sessions.wait()) + .await + .is_ok() + { + return true; + } + self.stop_sessions().await; + false + } + + pub(crate) async fn stop_sessions(&self) { + self.stop_accepting_sessions(); + self.session_cancel.cancel(); + self.sessions.close(); + let _ = tokio::time::timeout(SESSION_STOP_TIMEOUT, self.sessions.wait()).await; + } + + pub(crate) async fn stop_background_tasks(&self) { + self.background_tasks.stop().await; + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn stop_sessions_cancels_tracked_future() { + let tracker = TaskTracker::new(); + let cancel = CancellationToken::new(); + let child_cancel = cancel.clone(); + tracker.spawn(async move { + child_cancel.cancelled().await; + }); + cancel.cancel(); + tracker.close(); + tokio::time::timeout(Duration::from_secs(1), tracker.wait()) + .await + .unwrap(); + } + + #[tokio::test] + async fn runtime_task_scope_joins_cancelled_background_task() { + let scope = RuntimeTaskScope::new(); + scope.spawn(std::future::pending()); + tokio::time::timeout(Duration::from_secs(1), scope.stop()) + .await + .unwrap(); + } +} diff --git a/src/maestro/listeners.rs b/src/maestro/listeners.rs index 15cd31f..4e3e545 100644 --- a/src/maestro/listeners.rs +++ b/src/maestro/listeners.rs @@ -3,35 +3,33 @@ use std::net::{IpAddr, SocketAddr}; use std::sync::Arc; use std::time::Duration; +use arc_swap::ArcSwap; use tokio::net::TcpListener; #[cfg(unix)] use tokio::net::UnixListener; -use tokio::sync::{RwLock, Semaphore, watch}; use tracing::{debug, error, info, warn}; use crate::config::{ProxyConfig, RstOnCloseMode}; -use crate::crypto::SecureRandom; -use crate::ip_tracker::UserIpTracker; use crate::proxy::ClientHandler; -use crate::proxy::route_mode::RouteRuntimeController; -use crate::proxy::shared_state::ProxySharedState; use crate::startup::{COMPONENT_LISTENERS_BIND, StartupTracker}; -use crate::stats::beobachten::BeobachtenStore; -use crate::stats::{ReplayChecker, Stats}; -use crate::stream::BufferPool; -use crate::tls_front::TlsFrontCache; -use crate::transport::middle_proxy::MePool; use crate::transport::socket::set_linger_zero; -use crate::transport::{ListenOptions, UpstreamManager, create_listener, find_listener_processes}; +use crate::transport::{ListenOptions, create_listener, find_listener_processes}; +use super::generation::RuntimeGeneration; use super::helpers::{ expected_handshake_close_description, is_expected_handshake_eof, peer_close_description, print_proxy_links, }; +#[cfg(unix)] +mod unix; +#[cfg(unix)] +pub(crate) use unix::spawn_unix_accept_loop; + pub(crate) struct BoundListeners { pub(crate) listeners: Vec<(TcpListener, bool)>, - pub(crate) has_unix_listener: bool, + #[cfg(unix)] + pub(crate) unix_listener: Option, } fn listener_port_or_legacy(listener: &crate::config::ListenerConfig, config: &ProxyConfig) -> u16 { @@ -59,21 +57,6 @@ pub(crate) async fn bind_listeners( detected_ip_v4: Option, detected_ip_v6: Option, startup_tracker: &Arc, - config_rx: watch::Receiver>, - admission_rx: watch::Receiver, - stats: Arc, - upstream_manager: Arc, - replay_checker: Arc, - buffer_pool: Arc, - rng: Arc, - me_pool: Option>, - me_pool_runtime: Arc>>>, - route_runtime: Arc, - tls_cache: Option>, - ip_tracker: Arc, - beobachten: Arc, - shared: Arc, - max_connections: Arc, ) -> Result> { startup_tracker .start_component( @@ -218,7 +201,8 @@ pub(crate) async fn bind_listeners( print_proxy_links(&host, port, config); } - let mut has_unix_listener = false; + #[cfg(unix)] + let mut unix_listener_out = None; #[cfg(unix)] if let Some(ref unix_path) = config.server.listen_unix_sock { let _ = tokio::fs::remove_file(unix_path).await; @@ -251,123 +235,14 @@ pub(crate) async fn bind_listeners( info!("Listening on unix:{}", unix_path); } - has_unix_listener = true; - - let mut config_rx_unix: watch::Receiver> = config_rx.clone(); - let admission_rx_unix = admission_rx.clone(); - let stats = stats.clone(); - let upstream_manager = upstream_manager.clone(); - let replay_checker = replay_checker.clone(); - let buffer_pool = buffer_pool.clone(); - let rng = rng.clone(); - let me_pool = me_pool.clone(); - let me_pool_runtime = me_pool_runtime.clone(); - let route_runtime = route_runtime.clone(); - let tls_cache = tls_cache.clone(); - let ip_tracker = ip_tracker.clone(); - let beobachten = beobachten.clone(); - let shared = shared.clone(); - let max_connections_unix = max_connections.clone(); - - tokio::spawn(async move { - let unix_conn_counter = Arc::new(std::sync::atomic::AtomicU64::new(1)); - - loop { - match unix_listener.accept().await { - Ok((stream, _)) => { - if !*admission_rx_unix.borrow() { - drop(stream); - continue; - } - let accept_permit_timeout_ms = - config_rx_unix.borrow().server.accept_permit_timeout_ms; - let permit = if accept_permit_timeout_ms == 0 { - match max_connections_unix.clone().acquire_owned().await { - Ok(permit) => permit, - Err(_) => { - error!("Connection limiter is closed"); - break; - } - } - } else { - match tokio::time::timeout( - Duration::from_millis(accept_permit_timeout_ms), - max_connections_unix.clone().acquire_owned(), - ) - .await - { - Ok(Ok(permit)) => permit, - Ok(Err(_)) => { - error!("Connection limiter is closed"); - break; - } - Err(_) => { - stats.increment_accept_permit_timeout_total(); - debug!( - timeout_ms = accept_permit_timeout_ms, - "Dropping accepted unix connection: permit wait timeout" - ); - drop(stream); - continue; - } - } - }; - let conn_id = - unix_conn_counter.fetch_add(1, std::sync::atomic::Ordering::Relaxed); - let fake_peer = - SocketAddr::from(([127, 0, 0, 1], (conn_id % 65535) as u16)); - - let config = config_rx_unix.borrow_and_update().clone(); - let stats = stats.clone(); - let upstream_manager = upstream_manager.clone(); - let replay_checker = replay_checker.clone(); - let buffer_pool = buffer_pool.clone(); - let rng = rng.clone(); - let me_pool = me_pool.clone(); - let me_pool_runtime = me_pool_runtime.clone(); - let route_runtime = route_runtime.clone(); - let tls_cache = tls_cache.clone(); - let ip_tracker = ip_tracker.clone(); - let beobachten = beobachten.clone(); - let shared = shared.clone(); - let proxy_protocol_enabled = config.server.proxy_protocol; - - tokio::spawn(async move { - let _permit = permit; - if let Err(e) = - crate::proxy::client::handle_client_stream_with_shared_and_pool_runtime( - stream, - fake_peer, - config, - stats, - upstream_manager, - replay_checker, - buffer_pool, - rng, - me_pool, - Some(me_pool_runtime), - route_runtime, - tls_cache, - ip_tracker, - beobachten, - shared, - proxy_protocol_enabled, - ) - .await - { - debug!(error = %e, "Unix socket connection error"); - } - }); - } - Err(e) => { - error!("Unix socket accept error: {}", e); - tokio::time::sleep(Duration::from_millis(100)).await; - } - } - } - }); + unix_listener_out = Some(unix_listener); } + #[cfg(unix)] + let has_unix_listener = unix_listener_out.is_some(); + #[cfg(not(unix))] + let has_unix_listener = false; + startup_tracker .complete_component( COMPONENT_LISTENERS_BIND, @@ -381,51 +256,25 @@ pub(crate) async fn bind_listeners( Ok(BoundListeners { listeners, - has_unix_listener, + #[cfg(unix)] + unix_listener: unix_listener_out, }) } -#[allow(clippy::too_many_arguments)] pub(crate) fn spawn_tcp_accept_loops( listeners: Vec<(TcpListener, bool)>, - config_rx: watch::Receiver>, - admission_rx: watch::Receiver, - stats: Arc, - upstream_manager: Arc, - replay_checker: Arc, - buffer_pool: Arc, - rng: Arc, - me_pool: Option>, - me_pool_runtime: Arc>>>, - route_runtime: Arc, - tls_cache: Option>, - ip_tracker: Arc, - beobachten: Arc, - shared: Arc, - max_connections: Arc, + active_runtime: Arc>, ) { for (listener, listener_proxy_protocol) in listeners { - let mut config_rx: watch::Receiver> = config_rx.clone(); - let admission_rx_tcp = admission_rx.clone(); - let stats = stats.clone(); - let upstream_manager = upstream_manager.clone(); - let replay_checker = replay_checker.clone(); - let buffer_pool = buffer_pool.clone(); - let rng = rng.clone(); - let me_pool = me_pool.clone(); - let me_pool_runtime = me_pool_runtime.clone(); - let route_runtime = route_runtime.clone(); - let tls_cache = tls_cache.clone(); - let ip_tracker = ip_tracker.clone(); - let beobachten = beobachten.clone(); - let shared = shared.clone(); - let max_connections_tcp = max_connections.clone(); + let active_runtime = active_runtime.clone(); tokio::spawn(async move { loop { match listener.accept().await { Ok((stream, peer_addr)) => { - let rst_mode = config_rx.borrow().general.rst_on_close; + let runtime = active_runtime.load_full(); + let config = runtime.config(); + let rst_mode = config.general.rst_on_close; #[cfg(unix)] let raw_fd = { use std::os::unix::io::AsRawFd; @@ -434,15 +283,14 @@ pub(crate) fn spawn_tcp_accept_loops( if matches!(rst_mode, RstOnCloseMode::Errors | RstOnCloseMode::Always) { let _ = set_linger_zero(&stream); } - if !*admission_rx_tcp.borrow() { + if !*runtime.admission_rx.borrow() { debug!(peer = %peer_addr, "Admission gate closed, dropping connection"); drop(stream); continue; } - let accept_permit_timeout_ms = - config_rx.borrow().server.accept_permit_timeout_ms; + let accept_permit_timeout_ms = config.server.accept_permit_timeout_ms; let permit = if accept_permit_timeout_ms == 0 { - match max_connections_tcp.clone().acquire_owned().await { + match runtime.max_connections.clone().acquire_owned().await { Ok(permit) => permit, Err(_) => { error!("Connection limiter is closed"); @@ -452,7 +300,7 @@ pub(crate) fn spawn_tcp_accept_loops( } else { match tokio::time::timeout( Duration::from_millis(accept_permit_timeout_ms), - max_connections_tcp.clone().acquire_owned(), + runtime.max_connections.clone().acquire_owned(), ) .await { @@ -462,7 +310,7 @@ pub(crate) fn spawn_tcp_accept_loops( break; } Err(_) => { - stats.increment_accept_permit_timeout_total(); + runtime.stats.increment_accept_permit_timeout_total(); debug!( peer = %peer_addr, timeout_ms = accept_permit_timeout_ms, @@ -473,24 +321,23 @@ pub(crate) fn spawn_tcp_accept_loops( } } }; - let config = config_rx.borrow_and_update().clone(); - let stats = stats.clone(); - let upstream_manager = upstream_manager.clone(); - let replay_checker = replay_checker.clone(); - let buffer_pool = buffer_pool.clone(); - let rng = rng.clone(); - let me_pool = me_pool.clone(); - let me_pool_runtime = me_pool_runtime.clone(); - let route_runtime = route_runtime.clone(); - let tls_cache = tls_cache.clone(); - let ip_tracker = ip_tracker.clone(); - let beobachten = beobachten.clone(); - let shared = shared.clone(); + let stats = runtime.stats.clone(); + let upstream_manager = runtime.upstream_manager.clone(); + let replay_checker = runtime.replay_checker.clone(); + let buffer_pool = runtime.buffer_pool.clone(); + let rng = runtime.rng.clone(); + let me_pool = runtime.me_pool.clone(); + let me_pool_runtime = runtime.me_pool_runtime.clone(); + let route_runtime = runtime.route_runtime.clone(); + let tls_cache = runtime.tls_cache.clone(); + let ip_tracker = runtime.ip_tracker.clone(); + let beobachten = runtime.beobachten.clone(); + let shared = runtime.proxy_shared.clone(); let proxy_protocol_enabled = listener_proxy_protocol; let real_peer_report = Arc::new(std::sync::Mutex::new(None)); let real_peer_report_for_handler = real_peer_report.clone(); - tokio::spawn(async move { + let _ = runtime.spawn_session(async move { let _permit = permit; if let Err(e) = ClientHandler::new_with_shared( stream, diff --git a/src/maestro/listeners/unix.rs b/src/maestro/listeners/unix.rs new file mode 100644 index 0000000..ba5091a --- /dev/null +++ b/src/maestro/listeners/unix.rs @@ -0,0 +1,117 @@ +use std::net::SocketAddr; +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::Duration; + +use arc_swap::ArcSwap; +use tokio::net::UnixListener; +use tracing::{debug, error}; + +use super::RuntimeGeneration; + +pub(crate) fn spawn_unix_accept_loop( + listener: Option, + active_runtime: Arc>, +) { + let Some(listener) = listener else { + return; + }; + + tokio::spawn(async move { + let connection_counter = AtomicU64::new(1); + + loop { + match listener.accept().await { + Ok((stream, _)) => { + let runtime = active_runtime.load_full(); + if !*runtime.admission_rx.borrow() { + drop(stream); + continue; + } + + let config = runtime.config(); + let timeout_ms = config.server.accept_permit_timeout_ms; + let permit = if timeout_ms == 0 { + match runtime.max_connections.clone().acquire_owned().await { + Ok(permit) => permit, + Err(_) => { + error!("Connection limiter is closed"); + break; + } + } + } else { + match tokio::time::timeout( + Duration::from_millis(timeout_ms), + runtime.max_connections.clone().acquire_owned(), + ) + .await + { + Ok(Ok(permit)) => permit, + Ok(Err(_)) => { + error!("Connection limiter is closed"); + break; + } + Err(_) => { + runtime.stats.increment_accept_permit_timeout_total(); + debug!( + timeout_ms, + "Dropping accepted unix connection: permit wait timeout" + ); + drop(stream); + continue; + } + } + }; + + let connection_id = connection_counter.fetch_add(1, Ordering::Relaxed); + let fake_peer = + SocketAddr::from(([127, 0, 0, 1], (connection_id % 65535) as u16)); + let stats = runtime.stats.clone(); + let upstream_manager = runtime.upstream_manager.clone(); + let replay_checker = runtime.replay_checker.clone(); + let buffer_pool = runtime.buffer_pool.clone(); + let rng = runtime.rng.clone(); + let me_pool = runtime.me_pool.clone(); + let me_pool_runtime = runtime.me_pool_runtime.clone(); + let route_runtime = runtime.route_runtime.clone(); + let tls_cache = runtime.tls_cache.clone(); + let ip_tracker = runtime.ip_tracker.clone(); + let beobachten = runtime.beobachten.clone(); + let shared = runtime.proxy_shared.clone(); + let proxy_protocol_enabled = config.server.proxy_protocol; + + let _ = runtime.spawn_session(async move { + let _permit = permit; + if let Err(error) = + crate::proxy::client::handle_client_stream_with_shared_and_pool_runtime( + stream, + fake_peer, + config, + stats, + upstream_manager, + replay_checker, + buffer_pool, + rng, + me_pool, + Some(me_pool_runtime), + route_runtime, + tls_cache, + ip_tracker, + beobachten, + shared, + proxy_protocol_enabled, + ) + .await + { + debug!(error = %error, "Unix socket connection error"); + } + }); + } + Err(error) => { + error!(error = %error, "Unix socket accept error"); + tokio::time::sleep(Duration::from_millis(100)).await; + } + } + } + }); +} diff --git a/src/maestro/mod.rs b/src/maestro/mod.rs index e01670b..dd2cadf 100644 --- a/src/maestro/mod.rs +++ b/src/maestro/mod.rs @@ -13,19 +13,24 @@ // - shutdown: graceful shutdown sequence and uptime logging. mod admission; mod connectivity; +pub(crate) mod generation; mod helpers; mod listeners; mod me_startup; +pub(crate) mod reload; +mod reload_supervisor; +pub(crate) mod runtime_build; mod runtime_tasks; mod shutdown; mod tls_bootstrap; +use arc_swap::ArcSwap; use std::net::{IpAddr, SocketAddr}; use std::sync::Arc; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; use tokio::sync::{RwLock, Semaphore, watch}; use tracing::{error, info, warn}; -use tracing_subscriber::{EnvFilter, fmt, prelude::*, reload}; +use tracing_subscriber::{EnvFilter, fmt, prelude::*, reload as tracing_reload}; use crate::api; use crate::config::{LogLevel, ProxyConfig}; @@ -34,7 +39,7 @@ use crate::crypto::SecureRandom; use crate::ip_tracker::UserIpTracker; use crate::network::probe::{decide_network_capabilities, log_probe_result, run_probe}; use crate::proxy::direct_buffer_budget::{ - DirectBufferBudget, resolve_direct_buffer_hard_limit, spawn_direct_buffer_budget_controller, + DirectBufferBudget, resolve_direct_buffer_hard_limit, run_direct_buffer_budget_controller, }; use crate::proxy::route_mode::{RelayRouteMode, RouteRuntimeController}; use crate::proxy::shared_state::ProxySharedState; @@ -46,7 +51,7 @@ use crate::startup::{ }; use crate::stats::beobachten::BeobachtenStore; use crate::stats::telemetry::TelemetryPolicy; -use crate::stats::{ReplayChecker, Stats}; +use crate::stats::{QuotaStore, ReplayChecker, Stats}; use crate::stream::BufferPool; use crate::synlimit_control; use crate::transport::UpstreamManager; @@ -343,7 +348,7 @@ async fn run_telemt_core( } }; let (filter_layer, filter_handle) = - reload::Layer::new(EnvFilter::new(initial_filter_spec.clone())); + tracing_reload::Layer::new(EnvFilter::new(initial_filter_spec.clone())); startup_tracker .start_component( COMPONENT_TRACING_INIT, @@ -387,6 +392,7 @@ async fn run_telemt_core( _logging_guard = Some(guard); } } + let runtime_log_filter = runtime_tasks::RuntimeLogFilter::new(filter_handle); startup_tracker .complete_component( @@ -433,21 +439,26 @@ async fn run_telemt_core( warn!("Using default tls_domain. Consider setting a custom domain."); } - let stats = Arc::new(Stats::new()); + let quota_store = Arc::new(QuotaStore::default()); + let stats = Arc::new(Stats::with_quota_store(quota_store.clone())); + 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 upstream_manager = Arc::new(UpstreamManager::new( - config.upstreams.clone(), - config.general.upstream_connect_retry_attempts, - config.general.upstream_connect_retry_backoff_ms, - config.general.upstream_connect_budget_ms, - config.general.tg_connect, - config.general.upstream_unhealthy_fail_threshold, - config.general.upstream_connect_failfast_hard_errors, - stats.clone(), - )); + let upstream_manager = Arc::new( + UpstreamManager::new( + config.upstreams.clone(), + config.general.upstream_connect_retry_attempts, + config.general.upstream_connect_retry_backoff_ms, + config.general.upstream_connect_budget_ms, + config.general.tg_connect, + config.general.upstream_unhealthy_fail_threshold, + config.general.upstream_connect_failfast_hard_errors, + stats.clone(), + ) + .with_dns_overrides(&config.network.dns_overrides)?, + ); let ip_tracker = Arc::new(UserIpTracker::new()); ip_tracker .load_limits( @@ -497,6 +508,9 @@ async fn run_telemt_core( let initial_direct_first = config.general.use_middle_proxy && config.general.me2dc_fallback; let initial_admission_open = !config.general.use_middle_proxy || initial_direct_first; let (admission_tx, admission_rx) = watch::channel(initial_admission_open); + let (reload_control, reload_commands) = reload::ReloadControl::channel(1); + let (active_runtime_tx, active_runtime_rx) = + watch::channel(None::>>); let initial_route_mode = if !config.general.use_middle_proxy || initial_direct_first { RelayRouteMode::Direct } else { @@ -536,6 +550,8 @@ async fn run_telemt_core( let quota_state_path_api = quota_state_path.clone(); let startup_tracker_api = startup_tracker.clone(); let detected_ips_rx_api = detected_ips_rx.clone(); + let reload_control_api = reload_control.clone(); + let active_runtime_rx_api = active_runtime_rx.clone(); tokio::spawn(async move { api::serve( listen, @@ -552,6 +568,8 @@ async fn run_telemt_core( detected_ips_rx_api, process_started_at_epoch_secs, startup_tracker_api, + reload_control_api, + active_runtime_rx_api, ) .await; }); @@ -591,6 +609,7 @@ async fn run_telemt_core( &tls_domains, upstream_manager.clone(), &startup_tracker, + runtime_task_scope.clone(), ) .await; @@ -809,12 +828,19 @@ async fn run_telemt_core( me_pool.clone(), shared_state.clone(), me_ready_tx.clone(), + runtime_task_scope.clone(), ) .await; let config_rx = runtime_watches.config_rx; let log_level_rx = runtime_watches.log_level_rx; let detected_ip_v4 = runtime_watches.detected_ip_v4; let detected_ip_v6 = runtime_watches.detected_ip_v6; + runtime_log_filter.start( + has_rust_log, + &effective_log_level, + log_level_rx, + runtime_task_scope.clone(), + ); if direct_first_startup { let config_bg = config.clone(); @@ -827,7 +853,8 @@ async fn run_telemt_core( let api_me_pool_bg = api_me_pool.clone(); let me_ready_tx_bg = me_ready_tx.clone(); let config_rx_bg = config_rx.clone(); - tokio::spawn(async move { + let task_scope_bg = runtime_task_scope.clone(); + runtime_task_scope.spawn(async move { let mut bootstrap_attempt: u32 = 0; loop { bootstrap_attempt = bootstrap_attempt.saturating_add(1); @@ -851,6 +878,7 @@ async fn run_telemt_core( pool, rng_bg, me_ready_tx_bg, + task_scope_bg, ); break; } @@ -864,7 +892,7 @@ async fn run_telemt_core( let startup_tracker_ready = startup_tracker.clone(); let api_me_pool_ready = api_me_pool.clone(); let mut me_ready_rx_transport = me_ready_tx.subscribe(); - tokio::spawn(async move { + runtime_task_scope.spawn(async move { if me_ready_rx_transport.changed().await.is_ok() { if let Some(pool) = api_me_pool_ready.read().await.as_ref() { pool.set_runtime_ready(true); @@ -886,29 +914,27 @@ async fn run_telemt_core( &admission_tx, config_rx.clone(), me_ready_rx, + runtime_task_scope.clone(), ) .await; let _admission_tx_hold = admission_tx; - conntrack_control::spawn_conntrack_controller( + let conntrack_scope = runtime_task_scope.clone(); + runtime_task_scope.spawn(conntrack_control::run_conntrack_controller( config_rx.clone(), stats.clone(), shared_state.clone(), - ); - spawn_direct_buffer_budget_controller( + conntrack_scope.cancellation_token(), + )); + runtime_task_scope.spawn(run_direct_buffer_budget_controller( direct_buffer_budget, buffer_pool.clone(), stats.clone(), shared_state.clone(), config.server.max_connections, - ); + )); - let bound = listeners::bind_listeners( - &config, - decision.ipv4_dc, - decision.ipv6_dc, - detected_ip_v4, - detected_ip_v6, - &startup_tracker, + let runtime_generation = generation::RuntimeGeneration::new( + 1, config_rx.clone(), admission_rx.clone(), stats.clone(), @@ -924,10 +950,35 @@ async fn run_telemt_core( beobachten.clone(), shared_state.clone(), max_connections.clone(), + runtime_task_scope.clone(), + ); + let active_runtime = Arc::new(ArcSwap::from(runtime_generation)); + reload_supervisor::ReloadSupervisor::spawn( + active_runtime.clone(), + reload_control, + reload_commands, + config_path.clone(), + quota_store, + detected_ips_tx, + runtime_log_filter, + ); + + let bound = listeners::bind_listeners( + &config, + decision.ipv4_dc, + decision.ipv6_dc, + detected_ip_v4, + detected_ip_v6, + &startup_tracker, ) .await?; let listeners = bound.listeners; - let has_unix_listener = bound.has_unix_listener; + #[cfg(unix)] + let unix_listener = bound.unix_listener; + #[cfg(unix)] + let has_unix_listener = unix_listener.is_some(); + #[cfg(not(unix))] + let has_unix_listener = false; if listeners.is_empty() && !has_unix_listener { error!("No listeners. Exiting."); @@ -940,51 +991,20 @@ async fn run_telemt_core( synlimit_control::reconcile_synlimit_rules(&config).await; synlimit_control::spawn_synlimit_controller(config_rx.clone()); - runtime_tasks::apply_runtime_log_filter( - has_rust_log, - &effective_log_level, - filter_handle, - log_level_rx, - ) - .await; - - runtime_tasks::spawn_metrics_if_configured( - &config, - &startup_tracker, - stats.clone(), - beobachten.clone(), - shared_state.clone(), - ip_tracker.clone(), - tls_cache.clone(), - config_rx.clone(), - ) - .await; + runtime_tasks::spawn_metrics_if_configured(&config, &startup_tracker, active_runtime.clone()) + .await; + active_runtime_tx.send_replace(Some(active_runtime.clone())); runtime_tasks::mark_runtime_ready(&startup_tracker).await; // Spawn signal handlers for SIGUSR1/SIGUSR2 (non-shutdown signals) - shutdown::spawn_signal_handlers(stats.clone(), process_started_at); + shutdown::spawn_signal_handlers(active_runtime.clone(), process_started_at); - listeners::spawn_tcp_accept_loops( - listeners, - config_rx.clone(), - admission_rx.clone(), - stats.clone(), - upstream_manager.clone(), - replay_checker.clone(), - buffer_pool.clone(), - rng.clone(), - me_pool.clone(), - api_me_pool.clone(), - route_runtime.clone(), - tls_cache.clone(), - ip_tracker.clone(), - beobachten.clone(), - shared_state, - max_connections.clone(), - ); + listeners::spawn_tcp_accept_loops(listeners, active_runtime.clone()); + #[cfg(unix)] + listeners::spawn_unix_accept_loop(unix_listener, active_runtime.clone()); - shutdown::wait_for_shutdown(process_started_at, me_pool, stats, quota_state_path).await; + shutdown::wait_for_shutdown(process_started_at, active_runtime, quota_state_path).await; Ok(()) } diff --git a/src/maestro/reload.rs b/src/maestro/reload.rs new file mode 100644 index 0000000..b2b3bbc --- /dev/null +++ b/src/maestro/reload.rs @@ -0,0 +1,464 @@ +use std::collections::VecDeque; +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{SystemTime, UNIX_EPOCH}; + +use serde::{Deserialize, Serialize}; +use tokio::sync::{Mutex, mpsc}; + +use crate::config::ProxyConfig; + +const RELOAD_HISTORY_CAPACITY: usize = 32; +const RELOAD_COMMAND_CAPACITY: usize = 1; +const MAX_DRAIN_TIMEOUT_SECS: u64 = 3_600; + +/// Session handling policy for an in-process runtime reload. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Deserialize, Serialize)] +#[serde(rename_all = "snake_case")] +pub(crate) enum ReloadMode { + #[default] + Instant, + Drain, +} + +/// Failure policy applied during the activation barrier. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Deserialize, Serialize)] +#[serde(rename_all = "snake_case")] +pub(crate) enum ReloadFailurePolicy { + #[default] + KeepNew, + Rollback, +} + +/// Request body accepted by the maestro reload endpoint. +#[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize, Serialize)] +#[serde(deny_unknown_fields)] +pub(crate) struct ReloadRequest { + #[serde(default)] + pub(crate) mode: ReloadMode, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub(crate) timeout_secs: Option, + #[serde(default)] + pub(crate) failure_policy: ReloadFailurePolicy, +} + +impl ReloadRequest { + pub(crate) fn validate(&self) -> Result<(), &'static str> { + match (self.mode, self.timeout_secs) { + (ReloadMode::Instant, None) => Ok(()), + (ReloadMode::Instant, Some(_)) => Err("timeout_secs is only valid when mode is drain"), + (ReloadMode::Drain, Some(1..=MAX_DRAIN_TIMEOUT_SECS)) => Ok(()), + (ReloadMode::Drain, Some(_)) => Err("timeout_secs must be within 1..=3600"), + (ReloadMode::Drain, None) => Err("timeout_secs is required when mode is drain"), + } + } + + pub(crate) fn from_query(query: Option<&str>) -> Result, String> { + let Some(query) = query.filter(|query| !query.is_empty()) else { + return Ok(None); + }; + let mut mode = None; + let mut timeout_secs = None; + let mut failure_policy = None; + for (key, value) in url::form_urlencoded::parse(query.as_bytes()) { + match key.as_ref() { + "reload" if mode.is_none() => { + mode = Some(match value.as_ref() { + "instant" => ReloadMode::Instant, + "drain" => ReloadMode::Drain, + _ => return Err("reload must be instant or drain".to_string()), + }); + } + "timeout_secs" if timeout_secs.is_none() => { + timeout_secs = Some( + value + .parse::() + .map_err(|_| "timeout_secs must be an integer".to_string())?, + ); + } + "failure_policy" if failure_policy.is_none() => { + failure_policy = Some(match value.as_ref() { + "keep_new" => ReloadFailurePolicy::KeepNew, + "rollback" => ReloadFailurePolicy::Rollback, + _ => { + return Err("failure_policy must be keep_new or rollback".to_string()); + } + }); + } + "reload" | "timeout_secs" | "failure_policy" => { + return Err(format!("duplicate query parameter: {}", key)); + } + _ => return Err(format!("unknown query parameter: {}", key)), + } + } + let mode = mode.ok_or_else(|| "reload query parameter is required".to_string())?; + let request = Self { + mode, + timeout_secs, + failure_policy: failure_policy.unwrap_or_default(), + }; + request.validate().map_err(str::to_string)?; + Ok(Some(request)) + } +} + +/// Observable phase of one reload operation. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] +pub(crate) enum ReloadPhase { + Accepted, + Preparing, + Activating, + Draining, + Succeeded, + RolledBack, + Failed, +} + +impl ReloadPhase { + fn is_terminal(self) -> bool { + matches!( + self, + ReloadPhase::Succeeded | ReloadPhase::RolledBack | ReloadPhase::Failed + ) + } +} + +/// Bounded public status for one reload operation. +#[derive(Debug, Clone, Serialize)] +pub(crate) struct ReloadStatus { + pub(crate) reload_id: u64, + pub(crate) target_generation: u64, + pub(crate) config_revision: String, + pub(crate) state: ReloadPhase, + pub(crate) mode: ReloadMode, + pub(crate) failure_policy: ReloadFailurePolicy, + pub(crate) requested_at_epoch_secs: u64, + #[serde(skip_serializing_if = "Option::is_none")] + pub(crate) started_at_epoch_secs: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub(crate) finished_at_epoch_secs: Option, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub(crate) deferred_fields: Vec, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub(crate) warnings: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + pub(crate) error: Option, +} + +/// Accepted operation metadata returned before asynchronous preparation starts. +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub(crate) struct ReloadAccepted { + pub(crate) reload_id: u64, + pub(crate) target_generation: u64, + pub(crate) config_revision: String, + pub(crate) state: ReloadPhase, + pub(crate) mode: ReloadMode, + pub(crate) failure_policy: ReloadFailurePolicy, +} + +pub(crate) struct ReloadCommand { + pub(crate) reload_id: u64, + pub(crate) target_generation: u64, + pub(crate) config: Arc, + pub(crate) config_revision: String, + pub(crate) request: ReloadRequest, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum ReloadSubmitError { + InProgress(u64), + MaestroUnavailable, +} + +#[derive(Clone)] +pub(crate) struct ReloadControl { + command_tx: mpsc::Sender, + status_store: Arc, + active_generation: Arc, +} + +pub(crate) struct ReloadCommandReceiver { + command_rx: mpsc::Receiver, +} + +#[derive(Default)] +struct ReloadStatusState { + next_reload_id: u64, + active_reload_id: Option, + statuses: VecDeque, +} + +#[derive(Default)] +struct ReloadStatusStore { + state: Mutex, +} + +impl ReloadControl { + pub(crate) fn channel(initial_generation: u64) -> (Self, ReloadCommandReceiver) { + let (command_tx, command_rx) = mpsc::channel(RELOAD_COMMAND_CAPACITY); + ( + Self { + command_tx, + status_store: Arc::new(ReloadStatusStore::default()), + active_generation: Arc::new(AtomicU64::new(initial_generation)), + }, + ReloadCommandReceiver { command_rx }, + ) + } + + pub(crate) async fn submit( + &self, + config: Arc, + config_revision: String, + request: ReloadRequest, + ) -> Result { + let target_generation = self + .active_generation + .load(Ordering::Acquire) + .saturating_add(1); + let status = self + .status_store + .reserve(target_generation, config_revision, request.clone()) + .await?; + let command = ReloadCommand { + reload_id: status.reload_id, + target_generation, + config, + config_revision: status.config_revision.clone(), + request, + }; + if self.command_tx.send(command).await.is_err() { + self.status_store + .finish( + status.reload_id, + ReloadPhase::Failed, + Some("maestro command channel is closed".to_string()), + ) + .await; + return Err(ReloadSubmitError::MaestroUnavailable); + } + Ok(ReloadAccepted { + reload_id: status.reload_id, + target_generation, + config_revision: status.config_revision, + state: ReloadPhase::Accepted, + mode: status.mode, + failure_policy: status.failure_policy, + }) + } + + pub(crate) async fn status(&self, reload_id: u64) -> Option { + self.status_store.get(reload_id).await + } + + pub(crate) async fn in_progress(&self) -> Option { + self.status_store.state.lock().await.active_reload_id + } + + pub(crate) async fn mark_phase(&self, reload_id: u64, phase: ReloadPhase) { + self.status_store.mark_phase(reload_id, phase).await; + } + + pub(crate) async fn set_deferred_fields(&self, reload_id: u64, fields: Vec) { + self.status_store + .update(reload_id, |status| status.deferred_fields = fields) + .await; + } + + pub(crate) async fn succeed(&self, reload_id: u64, generation: u64) { + self.active_generation.store(generation, Ordering::Release); + self.status_store + .finish(reload_id, ReloadPhase::Succeeded, None) + .await; + } + + pub(crate) async fn fail(&self, reload_id: u64, error: impl Into) { + self.status_store + .finish(reload_id, ReloadPhase::Failed, Some(error.into())) + .await; + } + + pub(crate) async fn rolled_back(&self, reload_id: u64, error: impl Into) { + self.status_store + .finish(reload_id, ReloadPhase::RolledBack, Some(error.into())) + .await; + } + + pub(crate) async fn add_warning(&self, reload_id: u64, warning: impl Into) { + let warning = warning.into(); + self.status_store + .update(reload_id, |status| status.warnings.push(warning)) + .await; + } +} + +impl ReloadCommandReceiver { + pub(crate) async fn recv(&mut self) -> Option { + self.command_rx.recv().await + } +} + +impl ReloadStatusStore { + async fn reserve( + &self, + target_generation: u64, + config_revision: String, + request: ReloadRequest, + ) -> Result { + let mut state = self.state.lock().await; + if let Some(reload_id) = state.active_reload_id { + return Err(ReloadSubmitError::InProgress(reload_id)); + } + state.next_reload_id = state.next_reload_id.saturating_add(1).max(1); + let reload_id = state.next_reload_id; + let status = ReloadStatus { + reload_id, + target_generation, + config_revision, + state: ReloadPhase::Accepted, + mode: request.mode, + failure_policy: request.failure_policy, + requested_at_epoch_secs: now_epoch_secs(), + started_at_epoch_secs: None, + finished_at_epoch_secs: None, + deferred_fields: Vec::new(), + warnings: Vec::new(), + error: None, + }; + state.active_reload_id = Some(reload_id); + state.statuses.push_back(status.clone()); + while state.statuses.len() > RELOAD_HISTORY_CAPACITY { + state.statuses.pop_front(); + } + Ok(status) + } + + async fn get(&self, reload_id: u64) -> Option { + self.state + .lock() + .await + .statuses + .iter() + .find(|status| status.reload_id == reload_id) + .cloned() + } + + async fn mark_phase(&self, reload_id: u64, phase: ReloadPhase) { + self.update(reload_id, |status| { + status.state = phase; + if status.started_at_epoch_secs.is_none() && phase != ReloadPhase::Accepted { + status.started_at_epoch_secs = Some(now_epoch_secs()); + } + }) + .await; + } + + async fn finish(&self, reload_id: u64, phase: ReloadPhase, error: Option) { + debug_assert!(phase.is_terminal()); + let mut state = self.state.lock().await; + if let Some(status) = state + .statuses + .iter_mut() + .find(|status| status.reload_id == reload_id) + { + status.state = phase; + status.error = error; + status.finished_at_epoch_secs = Some(now_epoch_secs()); + } + if state.active_reload_id == Some(reload_id) { + state.active_reload_id = None; + } + } + + async fn update(&self, reload_id: u64, update: impl FnOnce(&mut ReloadStatus)) { + let mut state = self.state.lock().await; + if let Some(status) = state + .statuses + .iter_mut() + .find(|status| status.reload_id == reload_id) + { + update(status); + } + } +} + +fn now_epoch_secs() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn request_defaults_to_instant_keep_new() { + let request: ReloadRequest = serde_json::from_str("{}").unwrap(); + assert_eq!(request, ReloadRequest::default()); + assert_eq!(request.validate(), Ok(())); + } + + #[test] + fn drain_requires_bounded_timeout() { + let missing = ReloadRequest { + mode: ReloadMode::Drain, + ..ReloadRequest::default() + }; + assert!(missing.validate().is_err()); + let valid = ReloadRequest { + mode: ReloadMode::Drain, + timeout_secs: Some(30), + ..ReloadRequest::default() + }; + assert_eq!(valid.validate(), Ok(())); + } + + #[test] + fn patch_query_parses_reload_policy() { + let request = + ReloadRequest::from_query(Some("reload=drain&timeout_secs=30&failure_policy=rollback")) + .unwrap() + .unwrap(); + assert_eq!(request.mode, ReloadMode::Drain); + assert_eq!(request.timeout_secs, Some(30)); + assert_eq!(request.failure_policy, ReloadFailurePolicy::Rollback); + assert!(ReloadRequest::from_query(Some("timeout_secs=30")).is_err()); + } + + #[tokio::test] + async fn coordinator_rejects_concurrent_reload_and_releases_terminal_slot() { + let (control, mut receiver) = ReloadControl::channel(1); + let first = control + .submit( + Arc::new(ProxyConfig::default()), + "rev-1".to_string(), + ReloadRequest::default(), + ) + .await + .unwrap(); + let _command = receiver.recv().await.unwrap(); + let second = control + .submit( + Arc::new(ProxyConfig::default()), + "rev-2".to_string(), + ReloadRequest::default(), + ) + .await; + assert_eq!(second, Err(ReloadSubmitError::InProgress(first.reload_id))); + control + .succeed(first.reload_id, first.target_generation) + .await; + let third = control + .submit( + Arc::new(ProxyConfig::default()), + "rev-3".to_string(), + ReloadRequest::default(), + ) + .await + .unwrap(); + assert_eq!(third.reload_id, first.reload_id + 1); + } +} diff --git a/src/maestro/reload_supervisor.rs b/src/maestro/reload_supervisor.rs new file mode 100644 index 0000000..aa88cce --- /dev/null +++ b/src/maestro/reload_supervisor.rs @@ -0,0 +1,197 @@ +use std::path::PathBuf; +use std::sync::Arc; +use std::time::Duration; + +use arc_swap::ArcSwap; +use tokio::sync::watch; +use tracing::{info, warn}; + +use crate::stats::QuotaStore; + +use super::generation::RuntimeGeneration; +use super::reload::{ + ReloadCommand, ReloadCommandReceiver, ReloadControl, ReloadFailurePolicy, ReloadMode, + ReloadPhase, +}; +use super::runtime_build::{deferred_process_fields, prepare_runtime}; +use super::runtime_tasks::RuntimeLogFilter; + +pub(crate) struct ReloadSupervisor { + active_runtime: Arc>, + control: ReloadControl, + commands: ReloadCommandReceiver, + config_path: PathBuf, + quota_store: Arc, + detected_ips_tx: watch::Sender<(Option, Option)>, + runtime_log_filter: RuntimeLogFilter, +} + +impl ReloadSupervisor { + #[allow(clippy::too_many_arguments)] + pub(crate) fn spawn( + active_runtime: Arc>, + control: ReloadControl, + commands: ReloadCommandReceiver, + config_path: PathBuf, + quota_store: Arc, + detected_ips_tx: watch::Sender<(Option, Option)>, + runtime_log_filter: RuntimeLogFilter, + ) { + let supervisor = Self { + active_runtime, + control, + commands, + config_path, + quota_store, + detected_ips_tx, + runtime_log_filter, + }; + tokio::spawn(supervisor.run()); + } + + async fn run(mut self) { + while let Some(command) = self.commands.recv().await { + self.reload(command).await; + } + } + + async fn reload(&self, command: ReloadCommand) { + self.control + .mark_phase(command.reload_id, ReloadPhase::Preparing) + .await; + let old_runtime = self.active_runtime.load_full(); + let deferred = deferred_process_fields(&old_runtime.config(), &command.config); + self.control + .set_deferred_fields(command.reload_id, deferred) + .await; + + let prepared = match prepare_runtime( + command.target_generation, + command.config.as_ref().clone(), + &self.config_path, + self.quota_store.clone(), + self.runtime_log_filter.clone(), + ) + .await + { + Ok(prepared) => prepared, + Err(error) => { + self.control.fail(command.reload_id, error).await; + return; + } + }; + + if let Ok(current_revision) = + crate::api::config_store::current_revision_for_maestro(&self.config_path).await + && current_revision != command.config_revision + { + let warning = format!( + "config revision changed during preparation: accepted={} current={}", + command.config_revision, current_revision + ); + if command.request.failure_policy == ReloadFailurePolicy::Rollback { + self.runtime_log_filter + .apply_reload(&old_runtime.config().general.log_level); + prepared.generation.stop_sessions().await; + if let Some(pool) = prepared.generation.current_me_pool().await { + let _ = tokio::time::timeout( + Duration::from_secs(2), + pool.shutdown_send_close_conn_all(), + ) + .await; + } + prepared.generation.stop_background_tasks().await; + self.control.rolled_back(command.reload_id, warning).await; + return; + } + self.control.add_warning(command.reload_id, warning).await; + } + + self.control + .mark_phase(command.reload_id, ReloadPhase::Activating) + .await; + let new_runtime = prepared.generation; + old_runtime.stop_accepting_sessions(); + let replaced = self.active_runtime.swap(new_runtime.clone()); + if let Err(error) = crate::network::dns_overrides::install_entries( + &new_runtime.config().network.dns_overrides, + ) { + let message = format!("runtime DNS activation failed: {}", error); + if command.request.failure_policy == ReloadFailurePolicy::Rollback { + let candidate = self.active_runtime.swap(replaced.clone()); + replaced.resume_accepting_sessions(); + self.runtime_log_filter + .apply_reload(&replaced.config().general.log_level); + let _ = crate::network::dns_overrides::install_entries( + &replaced.config().network.dns_overrides, + ); + candidate.stop_sessions().await; + if let Some(pool) = candidate.current_me_pool().await { + let _ = tokio::time::timeout( + Duration::from_secs(2), + pool.shutdown_send_close_conn_all(), + ) + .await; + } + candidate.stop_background_tasks().await; + self.control.rolled_back(command.reload_id, message).await; + return; + } + self.control.add_warning(command.reload_id, message).await; + } + self.detected_ips_tx.send_replace(prepared.detected_ips); + self.runtime_log_filter + .apply_reload(&new_runtime.config().general.log_level); + crate::synlimit_control::reconcile_synlimit_rules(&new_runtime.config()).await; + + info!( + reload_id = command.reload_id, + old_generation = replaced.id, + new_generation = new_runtime.id, + config_revision = %command.config_revision, + "Runtime generation activated" + ); + + match command.request.mode { + ReloadMode::Instant => { + replaced.stop_sessions().await; + } + ReloadMode::Drain => { + self.control + .mark_phase(command.reload_id, ReloadPhase::Draining) + .await; + let timeout = Duration::from_secs( + command + .request + .timeout_secs + .expect("validated drain request must carry timeout_secs"), + ); + if !replaced.drain_sessions(timeout).await { + let warning = format!( + "generation {} exceeded drain timeout; remaining sessions were cancelled", + replaced.id + ); + warn!(reload_id = command.reload_id, warning = %warning); + self.control.add_warning(command.reload_id, warning).await; + } + } + } + + if let Some(pool) = replaced.current_me_pool().await + && tokio::time::timeout(Duration::from_secs(2), pool.shutdown_send_close_conn_all()) + .await + .is_err() + { + let warning = format!( + "generation {} Middle-End close broadcast timed out", + replaced.id + ); + warn!(reload_id = command.reload_id, warning = %warning); + self.control.add_warning(command.reload_id, warning).await; + } + replaced.stop_background_tasks().await; + self.control + .succeed(command.reload_id, new_runtime.id) + .await; + } +} diff --git a/src/maestro/runtime_build.rs b/src/maestro/runtime_build.rs new file mode 100644 index 0000000..4ec1062 --- /dev/null +++ b/src/maestro/runtime_build.rs @@ -0,0 +1,357 @@ +use std::net::IpAddr; +use std::path::Path; +use std::sync::Arc; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use tokio::sync::{RwLock, Semaphore, watch}; + +use crate::config::ProxyConfig; +use crate::crypto::SecureRandom; +use crate::ip_tracker::UserIpTracker; +use crate::network::probe::{decide_network_capabilities, run_probe}; +use crate::proxy::direct_buffer_budget::{ + DirectBufferBudget, resolve_direct_buffer_hard_limit, run_direct_buffer_budget_controller, +}; +use crate::proxy::route_mode::{RelayRouteMode, RouteRuntimeController}; +use crate::proxy::shared_state::ProxySharedState; +use crate::startup::StartupTracker; +use crate::stats::beobachten::BeobachtenStore; +use crate::stats::telemetry::TelemetryPolicy; +use crate::stats::{QuotaStore, ReplayChecker, Stats}; +use crate::stream::BufferPool; +use crate::transport::UpstreamManager; +use crate::transport::middle_proxy::MePool; + +use super::admission; +use super::generation::{RuntimeGeneration, RuntimeTaskScope}; +use super::runtime_tasks::RuntimeLogFilter; +use super::{me_startup, runtime_tasks, tls_bootstrap}; + +pub(crate) struct PreparedRuntime { + pub(crate) generation: Arc, + pub(crate) detected_ips: (Option, Option), +} + +pub(crate) async fn prepare_runtime( + generation_id: u64, + mut config: ProxyConfig, + config_path: &Path, + quota_store: Arc, + runtime_log_filter: RuntimeLogFilter, +) -> Result { + let started_at_epoch_secs = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs(); + let startup_tracker = Arc::new(StartupTracker::new(started_at_epoch_secs)); + let task_scope = RuntimeTaskScope::new(); + let stats = Arc::new(Stats::with_quota_store(quota_store)); + stats.apply_telemetry_policy(TelemetryPolicy::from_config(&config.general.telemetry)); + + let upstream_manager = Arc::new( + UpstreamManager::new( + config.upstreams.clone(), + config.general.upstream_connect_retry_attempts, + config.general.upstream_connect_retry_backoff_ms, + config.general.upstream_connect_budget_ms, + config.general.tg_connect, + config.general.upstream_unhealthy_fail_threshold, + config.general.upstream_connect_failfast_hard_errors, + stats.clone(), + ) + .with_dns_overrides(&config.network.dns_overrides) + .map_err(|error| format!("DNS override preparation failed: {}", error))?, + ); + let ip_tracker = Arc::new(UserIpTracker::new()); + ip_tracker + .load_limits( + config.access.user_max_unique_ips_global_each, + &config.access.user_max_unique_ips, + ) + .await; + ip_tracker + .set_limit_policy( + config.access.user_max_unique_ips_mode, + config.access.user_max_unique_ips_window_secs, + ) + .await; + + let hard_limit = + resolve_direct_buffer_hard_limit(config.general.direct_relay_buffer_budget_max_bytes).await; + let direct_buffer_budget = DirectBufferBudget::new(hard_limit); + let proxy_shared = + ProxySharedState::new_with_direct_buffer_budget(direct_buffer_budget.clone()); + proxy_shared.apply_user_enabled_config(&config.access.user_enabled); + proxy_shared.traffic_limiter.apply_policy( + config.access.user_rate_limits.clone(), + config.access.cidr_rate_limits.clone(), + ); + + let probe = run_probe( + &config.network, + &config.upstreams, + config.general.middle_proxy_nat_probe, + config.general.stun_nat_probe_concurrency, + ) + .await + .map_err(|error| format!("network probe failed: {}", error))?; + let decision = + decide_network_capabilities(&config.network, &probe, config.general.middle_proxy_nat_ip); + let prefer_ipv6 = decision.prefer_ipv6(); + + let mut tls_domains = Vec::with_capacity(1 + config.censorship.tls_domains.len()); + tls_domains.push(config.censorship.tls_domain.clone()); + for domain in &config.censorship.tls_domains { + if !tls_domains.contains(domain) { + tls_domains.push(domain.clone()); + } + } + let tls_cache = tls_bootstrap::bootstrap_tls_front( + &config, + &tls_domains, + upstream_manager.clone(), + &startup_tracker, + task_scope.clone(), + ) + .await; + + let beobachten = Arc::new(BeobachtenStore::new()); + let rng = Arc::new(SecureRandom::new()); + let route_mode = if !config.general.use_middle_proxy || config.general.me2dc_fallback { + RelayRouteMode::Direct + } else { + RelayRouteMode::Middle + }; + let route_runtime = Arc::new(RouteRuntimeController::new(route_mode)); + let me_pool_runtime = Arc::new(RwLock::new(None::>)); + let (me_ready_tx, me_ready_rx) = watch::channel(0_u64); + let direct_first_startup = config.general.use_middle_proxy && config.general.me2dc_fallback; + let me_pool = if direct_first_startup { + None + } else { + me_startup::initialize_me_pool( + config.general.use_middle_proxy, + &config, + &decision, + &probe, + &startup_tracker, + upstream_manager.clone(), + rng.clone(), + stats.clone(), + me_pool_runtime.clone(), + me_ready_tx.clone(), + ) + .await + }; + if config.general.use_middle_proxy && me_pool.is_none() && !direct_first_startup { + config.general.use_middle_proxy = false; + } + + let config = Arc::new(config); + let replay_checker = Arc::new(ReplayChecker::new( + config.access.replay_check_len, + Duration::from_secs(config.access.replay_window_secs), + )); + let buffer_pool = Arc::new(BufferPool::with_config(64 * 1024, 4096)); + let max_connections_limit = if config.server.max_connections == 0 { + Semaphore::MAX_PERMITS + } else { + config.server.max_connections as usize + }; + let max_connections = Arc::new(Semaphore::new(max_connections_limit)); + let (api_config_tx, _) = watch::channel(config.clone()); + let watches = runtime_tasks::spawn_runtime_tasks( + &config, + config_path, + &probe, + prefer_ipv6, + decision.ipv4_dc, + decision.ipv6_dc, + &startup_tracker, + stats.clone(), + upstream_manager.clone(), + replay_checker.clone(), + me_pool.clone(), + rng.clone(), + ip_tracker.clone(), + beobachten.clone(), + api_config_tx, + me_pool.clone(), + proxy_shared.clone(), + me_ready_tx.clone(), + task_scope.clone(), + ) + .await; + let config_rx = watches.config_rx; + runtime_log_filter.spawn_watcher(watches.log_level_rx, task_scope.clone()); + let initial_admission_open = !config.general.use_middle_proxy || me_pool.is_some(); + let (admission_tx, admission_rx) = watch::channel(initial_admission_open); + admission::configure_admission_gate( + &config, + me_pool.clone(), + me_pool_runtime.clone(), + route_runtime.clone(), + &admission_tx, + config_rx.clone(), + me_ready_rx, + task_scope.clone(), + ) + .await; + + if direct_first_startup { + let config_bg = config.clone(); + let decision_bg = decision.clone(); + let probe_bg = probe.clone(); + let startup_tracker_bg = startup_tracker.clone(); + let upstream_manager_bg = upstream_manager.clone(); + let rng_bg = rng.clone(); + let stats_bg = stats.clone(); + let me_pool_runtime_bg = me_pool_runtime.clone(); + let me_ready_tx_bg = me_ready_tx.clone(); + let config_rx_bg = config_rx.clone(); + let task_scope_bg = task_scope.clone(); + let retry_limit = config.general.me_init_retry_attempts; + task_scope.spawn(async move { + let mut attempt = 0_u32; + loop { + attempt = attempt.saturating_add(1); + let pool = me_startup::initialize_me_pool( + true, + config_bg.as_ref(), + &decision_bg, + &probe_bg, + &startup_tracker_bg, + upstream_manager_bg.clone(), + rng_bg.clone(), + stats_bg.clone(), + me_pool_runtime_bg.clone(), + me_ready_tx_bg.clone(), + ) + .await; + if let Some(pool) = pool { + runtime_tasks::spawn_middle_proxy_runtime_tasks( + config_bg.as_ref(), + config_rx_bg, + pool, + rng_bg, + me_ready_tx_bg, + task_scope_bg, + ); + break; + } + if retry_limit > 0 && attempt >= retry_limit { + break; + } + tokio::time::sleep(Duration::from_secs(2)).await; + } + }); + } + + let conntrack_scope = task_scope.clone(); + task_scope.spawn(crate::conntrack_control::run_conntrack_controller( + config_rx.clone(), + stats.clone(), + proxy_shared.clone(), + conntrack_scope.cancellation_token(), + )); + task_scope.spawn(run_direct_buffer_budget_controller( + direct_buffer_budget, + buffer_pool.clone(), + stats.clone(), + proxy_shared.clone(), + config.server.max_connections, + )); + let generation = RuntimeGeneration::new( + generation_id, + config_rx, + admission_rx, + stats, + upstream_manager, + replay_checker, + buffer_pool, + rng, + me_pool, + me_pool_runtime, + route_runtime, + tls_cache, + ip_tracker, + beobachten, + proxy_shared, + max_connections, + task_scope, + ); + drop(admission_tx); + + Ok(PreparedRuntime { + generation, + detected_ips: ( + probe.detected_ipv4.map(IpAddr::V4), + probe.detected_ipv6.map(IpAddr::V6), + ), + }) +} + +pub(crate) fn deferred_process_fields(old: &ProxyConfig, new: &ProxyConfig) -> Vec { + let mut fields = Vec::new(); + if old.server.port != new.server.port + || old.server.proxy_protocol != new.server.proxy_protocol + || old.server.listen_backlog != new.server.listen_backlog + || serde_json::to_value(&old.server.listeners).ok() + != serde_json::to_value(&new.server.listeners).ok() + { + fields.push("server.listeners".to_string()); + } + if old.server.listen_unix_sock != new.server.listen_unix_sock + || old.server.listen_unix_sock_perm != new.server.listen_unix_sock_perm + { + fields.push("server.listen_unix_sock".to_string()); + } + if old.server.api.listen != new.server.api.listen + || old.server.api.enabled != new.server.api.enabled + { + fields.push("server.api.listen".to_string()); + } + if old.server.metrics_listen != new.server.metrics_listen + || old.server.metrics_port != new.server.metrics_port + { + fields.push("server.metrics_listen".to_string()); + } + if old.general.quota_state_path != new.general.quota_state_path { + fields.push("general.quota_state_path".to_string()); + } + if old.general.disable_colors != new.general.disable_colors { + fields.push("general.disable_colors".to_string()); + } + if old.general.data_path != new.general.data_path { + fields.push("general.data_path".to_string()); + } + if serde_json::to_value(&old.logging).ok() != serde_json::to_value(&new.logging).ok() { + fields.push("logging".to_string()); + } + fields +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn process_socket_and_logging_changes_are_deferred() { + let old = ProxyConfig::default(); + let mut new = old.clone(); + new.server.listen_backlog = new.server.listen_backlog.saturating_add(1); + new.general.disable_colors = !new.general.disable_colors; + + let fields = deferred_process_fields(&old, &new); + assert!(fields.contains(&"server.listeners".to_string())); + assert!(fields.contains(&"general.disable_colors".to_string())); + } + + #[test] + fn runtime_only_change_does_not_require_process_rebind() { + let old = ProxyConfig::default(); + let mut new = old.clone(); + new.censorship.tls_domain = "reload.example".to_string(); + assert!(deferred_process_fields(&old, &new).is_empty()); + } +} diff --git a/src/maestro/runtime_tasks.rs b/src/maestro/runtime_tasks.rs index 6099014..7bd3976 100644 --- a/src/maestro/runtime_tasks.rs +++ b/src/maestro/runtime_tasks.rs @@ -2,9 +2,11 @@ use std::net::IpAddr; use std::path::Path; use std::sync::Arc; +use arc_swap::ArcSwap; use tokio::sync::{mpsc, watch}; use tracing::{debug, info, warn}; use tracing_subscriber::EnvFilter; +use tracing_subscriber::Registry; use tracing_subscriber::reload; use crate::config::hot_reload::spawn_config_watcher; @@ -21,10 +23,11 @@ use crate::startup::{ use crate::stats::beobachten::BeobachtenStore; use crate::stats::telemetry::TelemetryPolicy; use crate::stats::{ReplayChecker, Stats}; -use crate::tls_front::TlsFrontCache; use crate::transport::UpstreamManager; use crate::transport::middle_proxy::{MePool, MeReinitTrigger}; +use super::generation::RuntimeGeneration; +use super::generation::RuntimeTaskScope; use super::helpers::write_beobachten_snapshot; pub(crate) struct RuntimeWatches { @@ -34,6 +37,56 @@ pub(crate) struct RuntimeWatches { pub(crate) detected_ip_v6: Option, } +#[derive(Clone)] +pub(crate) struct RuntimeLogFilter { + handle: reload::Handle, +} + +impl RuntimeLogFilter { + pub(crate) fn new(handle: reload::Handle) -> Self { + Self { handle } + } + + pub(crate) fn start( + &self, + has_rust_log: bool, + effective_log_level: &LogLevel, + log_level_rx: watch::Receiver, + task_scope: RuntimeTaskScope, + ) { + self.apply(effective_log_level, has_rust_log); + self.spawn_watcher(log_level_rx, task_scope); + } + + pub(crate) fn apply_reload(&self, level: &LogLevel) { + self.apply(level, false); + } + + pub(crate) fn spawn_watcher( + &self, + mut log_level_rx: watch::Receiver, + task_scope: RuntimeTaskScope, + ) { + let filter = self.clone(); + task_scope.spawn(async move { + loop { + if log_level_rx.changed().await.is_err() { + break; + } + let level = log_level_rx.borrow_and_update().clone(); + filter.apply_reload(&level); + } + }); + } + + fn apply(&self, level: &LogLevel, has_rust_log: bool) { + let runtime_filter = EnvFilter::new(log_filter_spec(has_rust_log, level)); + if let Err(error) = self.handle.reload(runtime_filter) { + tracing::error!(error = %error, "Failed to update runtime log filter"); + } + } +} + #[allow(clippy::too_many_arguments)] pub(crate) async fn spawn_runtime_tasks( config: &Arc, @@ -54,10 +107,11 @@ pub(crate) async fn spawn_runtime_tasks( me_pool_for_policy: Option>, shared_state: Arc, me_ready_tx: watch::Sender, + task_scope: RuntimeTaskScope, ) -> RuntimeWatches { let um_clone = upstream_manager.clone(); let dc_overrides_for_health = config.dc_overrides.clone(); - tokio::spawn(async move { + task_scope.spawn(async move { um_clone .run_health_checks( prefer_ipv6, @@ -69,19 +123,19 @@ pub(crate) async fn spawn_runtime_tasks( }); let rc_clone = replay_checker.clone(); - tokio::spawn(async move { + task_scope.spawn(async move { rc_clone.run_periodic_cleanup().await; }); let stats_maintenance = stats.clone(); - tokio::spawn(async move { + task_scope.spawn(async move { stats_maintenance .run_periodic_user_stats_maintenance() .await; }); let ip_tracker_maintenance = ip_tracker.clone(); - tokio::spawn(async move { + task_scope.spawn(async move { ip_tracker_maintenance.run_periodic_maintenance().await; }); @@ -104,6 +158,7 @@ pub(crate) async fn spawn_runtime_tasks( config.clone(), detected_ip_v4, detected_ip_v6, + task_scope.cancellation_token(), ); startup_tracker .complete_component( @@ -113,7 +168,7 @@ pub(crate) async fn spawn_runtime_tasks( .await; let mut config_rx_api_bridge = config_rx.clone(); let api_config_tx_bridge = api_config_tx.clone(); - tokio::spawn(async move { + task_scope.spawn(async move { loop { if config_rx_api_bridge.changed().await.is_err() { break; @@ -124,8 +179,9 @@ pub(crate) async fn spawn_runtime_tasks( }); let stats_policy = stats.clone(); + let upstream_policy = upstream_manager.clone(); let mut config_rx_policy = config_rx.clone(); - tokio::spawn(async move { + task_scope.spawn(async move { loop { if config_rx_policy.changed().await.is_err() { break; @@ -133,6 +189,9 @@ 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, @@ -149,7 +208,7 @@ pub(crate) async fn spawn_runtime_tasks( let ip_tracker_policy = ip_tracker.clone(); let mut config_rx_ip_limits = config_rx.clone(); - tokio::spawn(async move { + task_scope.spawn(async move { let mut prev_limits = config_rx_ip_limits .borrow() .access @@ -205,7 +264,7 @@ pub(crate) async fn spawn_runtime_tasks( config.access.cidr_rate_limits.clone(), ); let mut config_rx_rate_limits = config_rx.clone(); - tokio::spawn(async move { + task_scope.spawn(async move { let mut prev_user_limits = config_rx_rate_limits .borrow() .access @@ -236,7 +295,7 @@ pub(crate) async fn spawn_runtime_tasks( let shared_user_enabled = shared_state.clone(); let mut config_rx_user_enabled = config_rx.clone(); - tokio::spawn(async move { + task_scope.spawn(async move { loop { if config_rx_user_enabled.changed().await.is_err() { break; @@ -257,7 +316,7 @@ pub(crate) async fn spawn_runtime_tasks( let beobachten_writer = beobachten.clone(); let config_rx_beobachten = config_rx.clone(); - tokio::spawn(async move { + task_scope.spawn(async move { loop { let cfg = config_rx_beobachten.borrow().clone(); let sleep_secs = cfg.general.beobachten_flush_secs.max(1); @@ -278,7 +337,14 @@ pub(crate) async fn spawn_runtime_tasks( }); if let Some(pool) = me_pool { - spawn_middle_proxy_runtime_tasks(config, config_rx.clone(), pool, rng, me_ready_tx); + spawn_middle_proxy_runtime_tasks( + config, + config_rx.clone(), + pool, + rng, + me_ready_tx, + task_scope, + ); } RuntimeWatches { @@ -295,6 +361,7 @@ pub(crate) fn spawn_middle_proxy_runtime_tasks( pool: Arc, rng: Arc, me_ready_tx: watch::Sender, + task_scope: RuntimeTaskScope, ) { let reinit_trigger_capacity = config.general.me_reinit_trigger_channel.max(1); let (reinit_tx, reinit_rx) = mpsc::channel::(reinit_trigger_capacity); @@ -303,7 +370,7 @@ pub(crate) fn spawn_middle_proxy_runtime_tasks( let rng_clone_sched = rng.clone(); let config_rx_clone_sched = config_rx.clone(); let me_ready_tx_sched = me_ready_tx.clone(); - tokio::spawn(async move { + task_scope.spawn(async move { crate::transport::middle_proxy::me_reinit_scheduler( pool_clone_sched, rng_clone_sched, @@ -317,7 +384,7 @@ pub(crate) fn spawn_middle_proxy_runtime_tasks( let pool_clone = pool.clone(); let config_rx_clone = config_rx.clone(); let reinit_tx_updater = reinit_tx.clone(); - tokio::spawn(async move { + task_scope.spawn(async move { crate::transport::middle_proxy::me_config_updater( pool_clone, config_rx_clone, @@ -328,37 +395,12 @@ pub(crate) fn spawn_middle_proxy_runtime_tasks( let config_rx_clone_rot = config_rx.clone(); let reinit_tx_rotation = reinit_tx.clone(); - tokio::spawn(async move { + task_scope.spawn(async move { crate::transport::middle_proxy::me_rotation_task(config_rx_clone_rot, reinit_tx_rotation) .await; }); } -pub(crate) async fn apply_runtime_log_filter( - has_rust_log: bool, - effective_log_level: &LogLevel, - filter_handle: reload::Handle, - mut log_level_rx: watch::Receiver, -) { - let runtime_filter = EnvFilter::new(log_filter_spec(has_rust_log, effective_log_level)); - filter_handle - .reload(runtime_filter) - .expect("Failed to switch log filter"); - - tokio::spawn(async move { - loop { - if log_level_rx.changed().await.is_err() { - break; - } - let level = log_level_rx.borrow_and_update().clone(); - let new_filter = tracing_subscriber::EnvFilter::new(log_filter_spec(false, &level)); - if let Err(e) = filter_handle.reload(new_filter) { - tracing::error!("config reload: failed to update log filter: {}", e); - } - } - }); -} - pub(crate) fn log_filter_spec(has_rust_log: bool, effective_log_level: &LogLevel) -> String { if has_rust_log { std::env::var("RUST_LOG") @@ -373,12 +415,7 @@ pub(crate) fn log_filter_spec(has_rust_log: bool, effective_log_level: &LogLevel pub(crate) async fn spawn_metrics_if_configured( config: &Arc, startup_tracker: &Arc, - stats: Arc, - beobachten: Arc, - shared_state: Arc, - ip_tracker: Arc, - tls_cache: Option>, - config_rx: watch::Receiver>, + active_runtime: Arc>, ) { // metrics_listen takes precedence; fall back to metrics_port for backward compat. let metrics_target: Option<(u16, Option)> = @@ -408,28 +445,10 @@ pub(crate) async fn spawn_metrics_if_configured( Some(format!("spawn metrics endpoint on {}", label)), ) .await; - let stats = stats.clone(); - let beobachten = beobachten.clone(); - let shared_state = shared_state.clone(); - let config_rx_metrics = config_rx.clone(); - let ip_tracker_metrics = ip_tracker.clone(); - let tls_cache_metrics = tls_cache.clone(); - let whitelist = config.server.metrics_whitelist.clone(); + let active_runtime = active_runtime.clone(); let listen_backlog = config.server.listen_backlog; tokio::spawn(async move { - metrics::serve( - port, - listen, - listen_backlog, - stats, - beobachten, - shared_state, - ip_tracker_metrics, - tls_cache_metrics, - config_rx_metrics, - whitelist, - ) - .await; + metrics::serve(port, listen, listen_backlog, active_runtime).await; }); startup_tracker .complete_component( diff --git a/src/maestro/shutdown.rs b/src/maestro/shutdown.rs index 4665337..fcb739a 100644 --- a/src/maestro/shutdown.rs +++ b/src/maestro/shutdown.rs @@ -12,17 +12,17 @@ use std::path::PathBuf; use std::sync::Arc; use std::time::{Duration, Instant}; +use arc_swap::ArcSwap; #[cfg(not(unix))] use tokio::signal; #[cfg(unix)] use tokio::signal::unix::{SignalKind, signal}; use tracing::{info, warn}; +use super::generation::RuntimeGeneration; +use super::helpers::{format_uptime, unit_label}; use crate::stats::Stats; use crate::synlimit_control; -use crate::transport::middle_proxy::MePool; - -use super::helpers::{format_uptime, unit_label}; /// Signal that triggered shutdown. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -48,19 +48,11 @@ impl std::fmt::Display for ShutdownSignal { /// Waits for a shutdown signal and performs graceful shutdown. pub(crate) async fn wait_for_shutdown( process_started_at: Instant, - me_pool: Option>, - stats: Arc, + active_runtime: Arc>, quota_state_path: PathBuf, ) { let signal = wait_for_shutdown_signal().await; - perform_shutdown( - signal, - process_started_at, - me_pool, - &stats, - quota_state_path, - ) - .await; + perform_shutdown(signal, process_started_at, active_runtime, quota_state_path).await; } /// Waits for any shutdown signal (SIGINT, SIGTERM, SIGQUIT). @@ -87,10 +79,11 @@ async fn wait_for_shutdown_signal() -> ShutdownSignal { async fn perform_shutdown( signal: ShutdownSignal, process_started_at: Instant, - me_pool: Option>, - stats: &Stats, + active_runtime: Arc>, quota_state_path: PathBuf, ) { + let runtime = active_runtime.load_full(); + let stats = runtime.stats.as_ref(); let shutdown_started_at = Instant::now(); info!(signal = %signal, "Received shutdown signal"); @@ -108,7 +101,8 @@ async fn perform_shutdown( } // Graceful ME pool shutdown - if let Some(pool) = &me_pool { + runtime.stop_sessions().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 { @@ -123,6 +117,7 @@ async fn perform_shutdown( } } } + runtime.stop_background_tasks().await; match crate::quota_state::save_quota_state("a_state_path, stats).await { Ok(()) => { @@ -191,7 +186,10 @@ fn dump_stats(stats: &Stats, process_started_at: Instant) { /// - SIGUSR1: Log rotation acknowledgment (for external log rotation tools) /// - SIGUSR2: Dump runtime status to log #[cfg(unix)] -pub(crate) fn spawn_signal_handlers(stats: Arc, process_started_at: Instant) { +pub(crate) fn spawn_signal_handlers( + active_runtime: Arc>, + process_started_at: Instant, +) { tokio::spawn(async move { let mut sigusr1 = signal(SignalKind::user_defined1()).expect("Failed to register SIGUSR1 handler"); @@ -204,7 +202,8 @@ pub(crate) fn spawn_signal_handlers(stats: Arc, process_started_at: Insta handle_sigusr1(); } _ = sigusr2.recv() => { - handle_sigusr2(&stats, process_started_at); + let runtime = active_runtime.load_full(); + handle_sigusr2(runtime.stats.as_ref(), process_started_at); } } } @@ -213,7 +212,10 @@ pub(crate) fn spawn_signal_handlers(stats: Arc, process_started_at: Insta /// No-op on non-Unix platforms. #[cfg(not(unix))] -pub(crate) fn spawn_signal_handlers(_stats: Arc, _process_started_at: Instant) { +pub(crate) fn spawn_signal_handlers( + _active_runtime: Arc>, + _process_started_at: Instant, +) { // No SIGUSR1/SIGUSR2 on non-Unix } diff --git a/src/maestro/tls_bootstrap.rs b/src/maestro/tls_bootstrap.rs index 4412723..7b7075f 100644 --- a/src/maestro/tls_bootstrap.rs +++ b/src/maestro/tls_bootstrap.rs @@ -10,6 +10,8 @@ use crate::tls_front::TlsFrontCache; use crate::tls_front::fetcher::TlsFetchStrategy; use crate::transport::UpstreamManager; +use super::generation::RuntimeTaskScope; + fn tls_fetch_host_for_domain(mask_host: &str, primary_tls_domain: &str, domain: &str) -> String { if mask_host.eq_ignore_ascii_case(primary_tls_domain) { domain.to_string() @@ -23,6 +25,7 @@ pub(crate) async fn bootstrap_tls_front( tls_domains: &[String], upstream_manager: Arc, startup_tracker: &Arc, + task_scope: RuntimeTaskScope, ) -> Option> { startup_tracker .start_component( @@ -69,7 +72,7 @@ pub(crate) async fn bootstrap_tls_front( let scope_initial = tls_fetch_scope.clone(); let upstream_initial = upstream_manager.clone(); let strategy_initial = fetch_strategy.clone(); - tokio::spawn(async move { + task_scope.spawn(async move { let mut join = tokio::task::JoinSet::new(); for domain in domains_initial { let cache_domain = cache_initial.clone(); @@ -109,7 +112,7 @@ pub(crate) async fn bootstrap_tls_front( let cache_timeout = cache.clone(); let domains_timeout = tls_domains.to_vec(); let fake_cert_len = config.censorship.fake_cert_len; - tokio::spawn(async move { + task_scope.spawn(async move { tokio::time::sleep(fetch_timeout).await; for domain in domains_timeout { let cached = cache_timeout.get(&domain).await; @@ -132,7 +135,7 @@ pub(crate) async fn bootstrap_tls_front( let scope_refresh = tls_fetch_scope.clone(); let upstream_refresh = upstream_manager.clone(); let strategy_refresh = fetch_strategy.clone(); - tokio::spawn(async move { + task_scope.spawn(async move { loop { let base_secs = rand::rng().random_range(4 * 3600..=6 * 3600); let jitter_secs = rand::rng().random_range(0..=7200); diff --git a/src/metrics.rs b/src/metrics.rs index f2816e2..fb35b84 100644 --- a/src/metrics.rs +++ b/src/metrics.rs @@ -4,12 +4,12 @@ use std::net::SocketAddr; use std::sync::Arc; use std::time::Duration; +use arc_swap::ArcSwap; use http_body_util::Full; use hyper::body::Bytes; use hyper::server::conn::http1; use hyper::service::service_fn; use hyper::{Request, Response, StatusCode}; -use ipnetwork::IpNetwork; use tokio::net::TcpListener; use tokio::sync::Semaphore; use tokio::time::timeout; @@ -17,6 +17,7 @@ use tracing::{debug, info, warn}; use crate::config::ProxyConfig; use crate::ip_tracker::UserIpTracker; +use crate::maestro::generation::RuntimeGeneration; use crate::proxy::shared_state::ProxySharedState; use crate::stats::Stats; use crate::stats::beobachten::BeobachtenStore; @@ -36,16 +37,8 @@ pub async fn serve( port: u16, listen: Option, listen_backlog: u32, - stats: Arc, - beobachten: Arc, - shared_state: Arc, - ip_tracker: Arc, - tls_cache: Option>, - config_rx: tokio::sync::watch::Receiver>, - whitelist: Vec, + active_runtime: Arc>, ) { - let whitelist = Arc::new(whitelist); - // If `metrics_listen` is set, bind on that single address only. if let Some(ref listen_addr) = listen { let addr: SocketAddr = match listen_addr.parse() { @@ -61,17 +54,7 @@ pub async fn serve( match bind_metrics_listener(addr, ipv6_only, listen_backlog) { Ok(listener) => { info!("Metrics endpoint: http://{}/metrics and /beobachten", addr); - serve_listener( - listener, - stats, - beobachten, - shared_state, - ip_tracker, - tls_cache, - config_rx, - whitelist, - ) - .await; + serve_listener(listener, active_runtime).await; } Err(e) => { warn!(error = %e, "Failed to bind metrics on {}", addr); @@ -117,50 +100,14 @@ pub async fn serve( warn!("Metrics listener is unavailable on both IPv4 and IPv6"); } (Some(listener), None) | (None, Some(listener)) => { - serve_listener( - listener, - stats, - beobachten, - shared_state, - ip_tracker, - tls_cache, - config_rx, - whitelist, - ) - .await; + serve_listener(listener, active_runtime).await; } (Some(listener4), Some(listener6)) => { - let stats_v6 = stats.clone(); - let beobachten_v6 = beobachten.clone(); - let shared_state_v6 = shared_state.clone(); - let ip_tracker_v6 = ip_tracker.clone(); - let tls_cache_v6 = tls_cache.clone(); - let config_rx_v6 = config_rx.clone(); - let whitelist_v6 = whitelist.clone(); + let active_runtime_v6 = active_runtime.clone(); tokio::spawn(async move { - serve_listener( - listener6, - stats_v6, - beobachten_v6, - shared_state_v6, - ip_tracker_v6, - tls_cache_v6, - config_rx_v6, - whitelist_v6, - ) - .await; + serve_listener(listener6, active_runtime_v6).await; }); - serve_listener( - listener4, - stats, - beobachten, - shared_state, - ip_tracker, - tls_cache, - config_rx, - whitelist, - ) - .await; + serve_listener(listener4, active_runtime).await; } } } @@ -180,16 +127,7 @@ fn bind_metrics_listener( TcpListener::from_std(socket.into()) } -async fn serve_listener( - listener: TcpListener, - stats: Arc, - beobachten: Arc, - shared_state: Arc, - ip_tracker: Arc, - tls_cache: Option>, - config_rx: tokio::sync::watch::Receiver>, - whitelist: Arc>, -) { +async fn serve_listener(listener: TcpListener, active_runtime: Arc>) { let connection_permits = Arc::new(Semaphore::new(METRICS_MAX_CONTROL_CONNECTIONS)); loop { @@ -201,7 +139,15 @@ async fn serve_listener( } }; - if !whitelist.is_empty() && !whitelist.iter().any(|net| net.contains(peer.ip())) { + let runtime = active_runtime.load_full(); + let config = runtime.config(); + if !config.server.metrics_whitelist.is_empty() + && !config + .server + .metrics_whitelist + .iter() + .any(|net| net.contains(peer.ip())) + { debug!(peer = %peer, "Metrics request denied by whitelist"); continue; } @@ -218,21 +164,17 @@ async fn serve_listener( } }; - let stats = stats.clone(); - let beobachten = beobachten.clone(); - let shared_state = shared_state.clone(); - let ip_tracker = ip_tracker.clone(); - let tls_cache = tls_cache.clone(); - let config_rx_conn = config_rx.clone(); + let active_runtime = active_runtime.clone(); tokio::spawn(async move { let _connection_permit = connection_permit; let svc = service_fn(move |req| { - let stats = stats.clone(); - let beobachten = beobachten.clone(); - let shared_state = shared_state.clone(); - let ip_tracker = ip_tracker.clone(); - let tls_cache = tls_cache.clone(); - let config = config_rx_conn.borrow().clone(); + let runtime = active_runtime.load_full(); + let stats = runtime.stats.clone(); + let beobachten = runtime.beobachten.clone(); + let shared_state = runtime.proxy_shared.clone(); + let ip_tracker = runtime.ip_tracker.clone(); + let tls_cache = runtime.tls_cache.clone(); + let config = runtime.config(); async move { handle( req, diff --git a/src/network/dns_overrides.rs b/src/network/dns_overrides.rs index 86fb325..710fa26 100644 --- a/src/network/dns_overrides.rs +++ b/src/network/dns_overrides.rs @@ -8,6 +8,33 @@ use crate::error::{ProxyError, Result}; type OverrideMap = HashMap<(String, u16), IpAddr>; +/// Immutable DNS override snapshot owned by one runtime generation. +#[derive(Debug, Clone, Default)] +pub struct DnsOverrides { + entries: std::sync::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)?), + }) + } + + /// Resolves a generation-local hostname override. + pub fn resolve(&self, host: &str, port: u16) -> Option { + self.entries + .get(&(host.to_ascii_lowercase(), port)) + .copied() + } + + /// Resolves a generation-local override as a socket address. + pub fn resolve_socket_addr(&self, host: &str, port: u16) -> Option { + self.resolve(host, port).map(|ip| SocketAddr::new(ip, port)) + } +} + static DNS_OVERRIDES: OnceLock> = OnceLock::new(); fn overrides_store() -> &'static RwLock { @@ -180,6 +207,22 @@ mod tests { assert_eq!(resolved, Some("127.0.0.1".parse().unwrap())); } + #[test] + fn generation_snapshots_do_not_observe_each_other() { + let first = DnsOverrides::from_entries(&["example.com:443:127.0.0.1".to_string()]).unwrap(); + let second = + DnsOverrides::from_entries(&["example.com:443:127.0.0.2".to_string()]).unwrap(); + + assert_eq!( + first.resolve("example.com", 443), + Some("127.0.0.1".parse().unwrap()) + ); + assert_eq!( + second.resolve("example.com", 443), + Some("127.0.0.2".parse().unwrap()) + ); + } + #[test] fn split_host_port_parses_supported_shapes() { assert_eq!( diff --git a/src/proxy/direct_buffer_budget.rs b/src/proxy/direct_buffer_budget.rs index eea3c0d..0242294 100644 --- a/src/proxy/direct_buffer_budget.rs +++ b/src/proxy/direct_buffer_budget.rs @@ -337,113 +337,111 @@ pub(crate) async fn resolve_direct_buffer_hard_limit(configured: usize) -> usize align_down(derived as usize).max(DIRECT_BUFFER_UNIT_BYTES) } -/// Starts the single control-plane task for Direct budget and shared pool pressure. -pub(crate) fn spawn_direct_buffer_budget_controller( +/// Runs the control-plane loop for Direct budget and shared pool pressure. +pub(crate) async fn run_direct_buffer_budget_controller( budget: Arc, buffer_pool: Arc, stats: Arc, shared: Arc, max_connections: u32, ) { - tokio::spawn(async move { - let mut interval = tokio::time::interval(CONTROL_INTERVAL); - interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); - let mut healthy_streak = 0u8; - let mut previous_denied = 0u64; - let mut previous_fallback = 0u64; - let mut previous_rejected = 0u64; - let pool_trim_low = buffer_pool - .max_buffers() - .min(BUFFER_POOL_TRIM_LOW_WATERMARK); - let pool_trim_high = buffer_pool - .max_buffers() - .min(BUFFER_POOL_TRIM_HIGH_WATERMARK); - let mut pool_trim_armed = true; + let mut interval = tokio::time::interval(CONTROL_INTERVAL); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + let mut healthy_streak = 0u8; + let mut previous_denied = 0u64; + let mut previous_fallback = 0u64; + let mut previous_rejected = 0u64; + let pool_trim_low = buffer_pool + .max_buffers() + .min(BUFFER_POOL_TRIM_LOW_WATERMARK); + let pool_trim_high = buffer_pool + .max_buffers() + .min(BUFFER_POOL_TRIM_HIGH_WATERMARK); + let mut pool_trim_armed = true; - loop { - interval.tick().await; - let sample = read_system_memory_sample().await; - budget.update_system_sample(sample); + loop { + interval.tick().await; + let sample = read_system_memory_sample().await; + budget.update_system_sample(sample); - let snapshot = budget.snapshot(); - let denied_delta = snapshot - .promotion_denied_total - .saturating_sub(previous_denied); - previous_denied = snapshot.promotion_denied_total; - let fallback_delta = snapshot - .minimum_fallback_total - .saturating_sub(previous_fallback); - previous_fallback = snapshot.minimum_fallback_total; - let rejected_delta = snapshot - .admission_rejected_total - .saturating_sub(previous_rejected); - previous_rejected = snapshot.admission_rejected_total; + let snapshot = budget.snapshot(); + let denied_delta = snapshot + .promotion_denied_total + .saturating_sub(previous_denied); + previous_denied = snapshot.promotion_denied_total; + let fallback_delta = snapshot + .minimum_fallback_total + .saturating_sub(previous_fallback); + previous_fallback = snapshot.minimum_fallback_total; + let rejected_delta = snapshot + .admission_rejected_total + .saturating_sub(previous_rejected); + previous_rejected = snapshot.admission_rejected_total; - let connection_pct = connection_fill_pct(stats.as_ref(), max_connections); - let memory_available_pct = percentage(sample.available_bytes, sample.total_bytes); - let target_utilization_pct = percentage(snapshot.reserved_bytes, snapshot.target_bytes); - let pressure = shared.conntrack_pressure_active() - || connection_pct.is_some_and(|value| value >= 85) - || memory_available_pct.is_some_and(|value| value <= 15) - || target_utilization_pct.is_some_and(|value| value >= 90) - || denied_delta > 0 - || fallback_delta > 0 - || rejected_delta > 0; + let connection_pct = connection_fill_pct(stats.as_ref(), max_connections); + let memory_available_pct = percentage(sample.available_bytes, sample.total_bytes); + let target_utilization_pct = percentage(snapshot.reserved_bytes, snapshot.target_bytes); + let pressure = shared.conntrack_pressure_active() + || connection_pct.is_some_and(|value| value >= 85) + || memory_available_pct.is_some_and(|value| value <= 15) + || target_utilization_pct.is_some_and(|value| value >= 90) + || denied_delta > 0 + || fallback_delta > 0 + || rejected_delta > 0; - if !pressure { - pool_trim_armed = true; - } else if pool_trim_armed && buffer_pool.pooled() > pool_trim_high { - buffer_pool.trim_to(pool_trim_low); - pool_trim_armed = false; - } - - let pool_snapshot = buffer_pool.stats(); - stats.set_buffer_pool_gauges( - pool_snapshot.pooled, - pool_snapshot.allocated, - pool_snapshot.allocated.saturating_sub(pool_snapshot.pooled), - ); - stats.set_buffer_pool_replaced_nonstandard_total(pool_snapshot.replaced_nonstandard); - - let headroom_target = if sample.total_bytes == 0 { - snapshot.hard_limit_bytes - } else { - snapshot - .reserved_bytes - .saturating_add(sample.available_bytes / 4) - .min(snapshot.hard_limit_bytes) - }; - - if pressure { - healthy_streak = 0; - let reduced = snapshot.target_bytes.saturating_mul(3) / 4; - budget.set_target_bytes(reduced.min(headroom_target)); - continue; - } - - let healthy = memory_available_pct.is_none_or(|value| value >= 30) - && connection_pct.is_none_or(|value| value <= 70); - if !healthy { - healthy_streak = 0; - if headroom_target < snapshot.target_bytes { - budget.set_target_bytes(headroom_target); - } - continue; - } - - healthy_streak = healthy_streak.saturating_add(1); - if healthy_streak >= HEALTHY_RECOVERY_SAMPLES { - healthy_streak = 0; - let increment = (snapshot.target_bytes / 16).max(4 * 1024 * 1024); - budget.set_target_bytes( - snapshot - .target_bytes - .saturating_add(increment) - .min(headroom_target), - ); - } + if !pressure { + pool_trim_armed = true; + } else if pool_trim_armed && buffer_pool.pooled() > pool_trim_high { + buffer_pool.trim_to(pool_trim_low); + pool_trim_armed = false; } - }); + + let pool_snapshot = buffer_pool.stats(); + stats.set_buffer_pool_gauges( + pool_snapshot.pooled, + pool_snapshot.allocated, + pool_snapshot.allocated.saturating_sub(pool_snapshot.pooled), + ); + stats.set_buffer_pool_replaced_nonstandard_total(pool_snapshot.replaced_nonstandard); + + let headroom_target = if sample.total_bytes == 0 { + snapshot.hard_limit_bytes + } else { + snapshot + .reserved_bytes + .saturating_add(sample.available_bytes / 4) + .min(snapshot.hard_limit_bytes) + }; + + if pressure { + healthy_streak = 0; + let reduced = snapshot.target_bytes.saturating_mul(3) / 4; + budget.set_target_bytes(reduced.min(headroom_target)); + continue; + } + + let healthy = memory_available_pct.is_none_or(|value| value >= 30) + && connection_pct.is_none_or(|value| value <= 70); + if !healthy { + healthy_streak = 0; + if headroom_target < snapshot.target_bytes { + budget.set_target_bytes(headroom_target); + } + continue; + } + + healthy_streak = healthy_streak.saturating_add(1); + if healthy_streak >= HEALTHY_RECOVERY_SAMPLES { + healthy_streak = 0; + let increment = (snapshot.target_bytes / 16).max(4 * 1024 * 1024); + budget.set_target_bytes( + snapshot + .target_bytes + .saturating_add(increment) + .min(headroom_target), + ); + } + } } fn connection_fill_pct(stats: &Stats, max_connections: u32) -> Option { diff --git a/src/proxy/relay/io/quota.rs b/src/proxy/relay/io/quota.rs index 1faf0b5..82e83f5 100644 --- a/src/proxy/relay/io/quota.rs +++ b/src/proxy/relay/io/quota.rs @@ -1,6 +1,5 @@ use crate::stats::UserStats; use std::io; -use std::sync::atomic::Ordering; #[derive(Debug)] struct QuotaIoSentinel; @@ -52,17 +51,5 @@ pub(super) fn refund_reserved_quota_bytes(user_stats: &UserStats, reserved_bytes if reserved_bytes == 0 { return; } - let mut current = user_stats.quota_used.load(Ordering::Relaxed); - loop { - let next = current.saturating_sub(reserved_bytes); - match user_stats.quota_used.compare_exchange_weak( - current, - next, - Ordering::Relaxed, - Ordering::Relaxed, - ) { - Ok(_) => return, - Err(observed) => current = observed, - } - } + user_stats.refund_quota(reserved_bytes); } diff --git a/src/stats/helpers.rs b/src/stats/helpers.rs index 7a398d4..2db8cfa 100644 --- a/src/stats/helpers.rs +++ b/src/stats/helpers.rs @@ -80,7 +80,11 @@ impl Stats { return handle; } - let entry = self.user_stats.entry(user.to_string()).or_default(); + let quota = self.quota_store.user(user); + let entry = self + .user_stats + .entry(user.to_string()) + .or_insert_with(|| Arc::new(UserStats::with_quota(quota))); if entry.last_seen_epoch_secs.load(Ordering::Relaxed) == 0 { self.touch_user_stats(entry.value().as_ref()); } @@ -166,10 +170,7 @@ impl Stats { #[inline] pub(crate) fn quota_charge_post_write(&self, user_stats: &UserStats, bytes: u64) -> u64 { self.touch_user_stats(user_stats); - user_stats - .quota_used - .fetch_add(bytes, Ordering::Relaxed) - .saturating_add(bytes) + user_stats.quota.charge(bytes) } pub(super) fn maybe_cleanup_user_stats(&self) { diff --git a/src/stats/mod.rs b/src/stats/mod.rs index ba71113..a16f8eb 100644 --- a/src/stats/mod.rs +++ b/src/stats/mod.rs @@ -8,6 +8,7 @@ mod core_getters; mod helpers; mod me_counters; mod me_getters; +mod quota_store; mod replay; pub mod telemetry; pub mod tls_fingerprints; @@ -20,6 +21,7 @@ use std::sync::Arc; use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU64, Ordering}; use std::time::Instant; +pub(crate) use self::quota_store::QuotaStore; #[allow(unused_imports)] pub use self::replay::{ReplayChecker, ReplayStats}; use self::telemetry::TelemetryPolicy; @@ -344,6 +346,7 @@ pub struct Stats { cached_epoch_secs: AtomicU64, tls_fingerprints: tls_fingerprints::TlsFingerprintCollector, user_stats: DashMap>, + quota_store: Arc, user_stats_last_cleanup_epoch_secs: AtomicU64, start_time: parking_lot::RwLock>, } @@ -356,12 +359,7 @@ pub struct UserStats { pub octets_to_client: AtomicU64, pub msgs_from_client: AtomicU64, pub msgs_to_client: AtomicU64, - /// Total bytes charged against per-user quota admission. - /// - /// This counter is the single source of truth for quota enforcement and - /// intentionally tracks attempted traffic, not guaranteed delivery. - pub quota_used: AtomicU64, - pub quota_last_reset_epoch_secs: AtomicU64, + quota: Arc, pub last_seen_epoch_secs: AtomicU64, } @@ -378,9 +376,21 @@ pub enum QuotaReserveError { } impl UserStats { + fn with_quota(quota: Arc) -> Self { + Self { + quota, + ..Self::default() + } + } + #[inline] pub fn quota_used(&self) -> u64 { - self.quota_used.load(Ordering::Relaxed) + self.quota.used() + } + + #[inline] + pub(crate) fn refund_quota(&self, bytes: u64) { + self.quota.refund(bytes); } /// Attempts one CAS reservation step against the quota counter. @@ -390,27 +400,20 @@ impl UserStats { /// with their own contention strategy. #[inline] pub fn quota_try_reserve(&self, bytes: u64, limit: u64) -> Result { - let current = self.quota_used.load(Ordering::Relaxed); - if bytes > limit.saturating_sub(current) { - return Err(QuotaReserveError::LimitExceeded); - } - - let next = current.saturating_add(bytes); - match self.quota_used.compare_exchange_weak( - current, - next, - Ordering::Relaxed, - Ordering::Relaxed, - ) { - Ok(_) => Ok(next), - Err(_) => Err(QuotaReserveError::Contended), - } + self.quota.try_reserve(bytes, limit) } } impl Stats { pub fn new() -> Self { - let stats = Self::default(); + Self::with_quota_store(Arc::new(QuotaStore::default())) + } + + pub(crate) fn with_quota_store(quota_store: Arc) -> Self { + let stats = Self { + quota_store, + ..Self::default() + }; stats.apply_telemetry_policy(TelemetryPolicy::default()); stats.refresh_cached_epoch_secs(); *stats.start_time.write() = Some(Instant::now()); diff --git a/src/stats/quota_store.rs b/src/stats/quota_store.rs new file mode 100644 index 0000000..79f4c99 --- /dev/null +++ b/src/stats/quota_store.rs @@ -0,0 +1,147 @@ +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; + +use dashmap::DashMap; + +use super::{QuotaReserveError, UserQuotaSnapshot}; + +/// Process-scoped per-user quota accounting shared by runtime generations. +#[derive(Default)] +pub struct QuotaStore { + users: DashMap>, +} + +/// Atomic quota state for one configured user. +#[derive(Default)] +pub(crate) struct UserQuotaCounters { + used_bytes: AtomicU64, + last_reset_epoch_secs: AtomicU64, +} + +impl QuotaStore { + pub(crate) fn user(&self, user: &str) -> Arc { + if let Some(existing) = self.users.get(user) { + return Arc::clone(existing.value()); + } + Arc::clone( + self.users + .entry(user.to_string()) + .or_insert_with(|| Arc::new(UserQuotaCounters::default())) + .value(), + ) + } + + pub(crate) fn used(&self, user: &str) -> u64 { + self.users.get(user).map(|state| state.used()).unwrap_or(0) + } + + pub(crate) fn load(&self, user: &str, used_bytes: u64, last_reset_epoch_secs: u64) { + let state = self.user(user); + state.used_bytes.store(used_bytes, Ordering::Relaxed); + state + .last_reset_epoch_secs + .store(last_reset_epoch_secs, Ordering::Relaxed); + } + + pub(crate) fn reset(&self, user: &str, now_epoch_secs: u64) -> UserQuotaSnapshot { + let state = self.user(user); + state.used_bytes.store(0, Ordering::Relaxed); + state + .last_reset_epoch_secs + .store(now_epoch_secs, Ordering::Relaxed); + UserQuotaSnapshot { + used_bytes: 0, + last_reset_epoch_secs: now_epoch_secs, + } + } + + pub(crate) fn snapshot(&self) -> HashMap { + let mut out = HashMap::new(); + for entry in self.users.iter() { + let state = entry.value(); + let used_bytes = state.used(); + let last_reset_epoch_secs = state.last_reset_epoch_secs.load(Ordering::Relaxed); + if used_bytes == 0 && last_reset_epoch_secs == 0 { + continue; + } + out.insert( + entry.key().clone(), + UserQuotaSnapshot { + used_bytes, + last_reset_epoch_secs, + }, + ); + } + out + } +} + +impl UserQuotaCounters { + #[inline] + pub(crate) fn used(&self) -> u64 { + self.used_bytes.load(Ordering::Relaxed) + } + + #[inline] + pub(crate) fn charge(&self, bytes: u64) -> u64 { + self.used_bytes + .fetch_add(bytes, Ordering::Relaxed) + .saturating_add(bytes) + } + + #[inline] + pub(crate) fn refund(&self, bytes: u64) { + let mut current = self.used_bytes.load(Ordering::Relaxed); + loop { + let next = current.saturating_sub(bytes); + match self.used_bytes.compare_exchange_weak( + current, + next, + Ordering::Relaxed, + Ordering::Relaxed, + ) { + Ok(_) => return, + Err(observed) => current = observed, + } + } + } + + #[inline] + pub(crate) fn try_reserve(&self, bytes: u64, limit: u64) -> Result { + let current = self.used_bytes.load(Ordering::Relaxed); + if bytes > limit.saturating_sub(current) { + return Err(QuotaReserveError::LimitExceeded); + } + + let next = current.saturating_add(bytes); + match self.used_bytes.compare_exchange_weak( + current, + next, + Ordering::Relaxed, + Ordering::Relaxed, + ) { + Ok(_) => Ok(next), + Err(_) => Err(QuotaReserveError::Contended), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::stats::Stats; + + #[test] + fn quota_counters_are_shared_across_stats_generations() { + let store = Arc::new(QuotaStore::default()); + let first = Stats::with_quota_store(store.clone()); + store.user("alice").charge(512); + assert_eq!(first.get_user_quota_used("alice"), 512); + + let second = Stats::with_quota_store(store); + assert_eq!(second.get_user_quota_used("alice"), 512); + second.reset_user_quota("alice"); + assert_eq!(first.get_user_quota_used("alice"), 0); + } +} diff --git a/src/stats/users.rs b/src/stats/users.rs index 72849cf..2c9d4af 100644 --- a/src/stats/users.rs +++ b/src/stats/users.rs @@ -119,51 +119,21 @@ impl Stats { } pub fn get_user_quota_used(&self, user: &str) -> u64 { - self.user_stats - .get(user) - .map(|s| s.quota_used.load(Ordering::Relaxed)) - .unwrap_or(0) + self.quota_store.used(user) } pub fn load_user_quota_state(&self, user: &str, used_bytes: u64, last_reset_epoch_secs: u64) { - let stats = self.get_or_create_user_stats_handle(user); - stats.quota_used.store(used_bytes, Ordering::Relaxed); - stats - .quota_last_reset_epoch_secs - .store(last_reset_epoch_secs, Ordering::Relaxed); + self.quota_store + .load(user, used_bytes, last_reset_epoch_secs); } pub fn reset_user_quota(&self, user: &str) -> UserQuotaSnapshot { - let stats = self.get_or_create_user_stats_handle(user); let last_reset_epoch_secs = Self::now_epoch_secs(); - stats.quota_used.store(0, Ordering::Relaxed); - stats - .quota_last_reset_epoch_secs - .store(last_reset_epoch_secs, Ordering::Relaxed); - UserQuotaSnapshot { - used_bytes: 0, - last_reset_epoch_secs, - } + self.quota_store.reset(user, last_reset_epoch_secs) } pub fn user_quota_snapshot(&self) -> HashMap { - let mut out = HashMap::new(); - for entry in self.user_stats.iter() { - let stats = entry.value(); - let used_bytes = stats.quota_used.load(Ordering::Relaxed); - let last_reset_epoch_secs = stats.quota_last_reset_epoch_secs.load(Ordering::Relaxed); - if used_bytes == 0 && last_reset_epoch_secs == 0 { - continue; - } - out.insert( - entry.key().clone(), - UserQuotaSnapshot { - used_bytes, - last_reset_epoch_secs, - }, - ); - } - out + self.quota_store.snapshot() } pub fn get_handshake_timeouts(&self) -> u64 { diff --git a/src/transport/middle_proxy/http_fetch.rs b/src/transport/middle_proxy/http_fetch.rs index 5be601e..33e48e2 100644 --- a/src/transport/middle_proxy/http_fetch.rs +++ b/src/transport/middle_proxy/http_fetch.rs @@ -66,33 +66,13 @@ fn extract_host_port_path(url: &str) -> Result<(String, u16, String)> { Ok((host, port, path)) } -async fn resolve_target_addr(host: &str, port: u16) -> Result { - if let Some(addr) = resolve_socket_addr(host, port) { - return Ok(addr); - } - - let addrs: Vec = tokio::net::lookup_host((host, port)) - .await - .map_err(|e| ProxyError::Proxy(format!("DNS resolve failed for {host}:{port}: {e}")))? - .collect(); - - if let Some(addr) = addrs.iter().copied().find(|addr| addr.is_ipv4()) { - return Ok(addr); - } - - addrs - .first() - .copied() - .ok_or_else(|| ProxyError::Proxy(format!("DNS returned no addresses for {host}:{port}"))) -} - async fn connect_https_transport( host: &str, port: u16, upstream: Option>, ) -> Result { if let Some(manager) = upstream { - let target = resolve_target_addr(host, port).await?; + let target = manager.resolve_hostname(host, port).await?; return timeout(HTTP_CONNECT_TIMEOUT, manager.connect(target, None, None)) .await .map_err(|_| ProxyError::Proxy(format!("upstream connect timeout for {host}:{port}")))? diff --git a/src/transport/upstream.rs b/src/transport/upstream.rs index f741d02..372d4d9 100644 --- a/src/transport/upstream.rs +++ b/src/transport/upstream.rs @@ -4,6 +4,7 @@ #![allow(deprecated)] +use arc_swap::ArcSwap; use rand::RngExt; use std::collections::{BTreeSet, HashMap}; use std::net::{IpAddr, SocketAddr}; @@ -20,7 +21,7 @@ use tracing::{debug, info, trace, warn}; use crate::config::{UpstreamConfig, UpstreamType}; use crate::error::{ProxyError, Result}; -use crate::network::dns_overrides::{resolve_socket_addr, split_host_port}; +use crate::network::dns_overrides::{DnsOverrides, split_host_port}; use crate::protocol::constants::{TG_DATACENTER_PORT, TG_DATACENTERS_V4, TG_DATACENTERS_V6}; use crate::stats::Stats; use crate::transport::shadowsocks::{ @@ -333,6 +334,7 @@ pub struct UpstreamManager { no_upstreams_warn_epoch_ms: Arc, no_healthy_warn_epoch_ms: Arc, stats: Arc, + dns_overrides: Arc>, } impl UpstreamManager { @@ -374,9 +376,37 @@ 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())), } } + pub(crate) fn with_dns_overrides(mut self, entries: &[String]) -> Result { + self.dns_overrides = Arc::new(ArcSwap::from_pointee(DnsOverrides::from_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(()) + } + + 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(); + if let Some(addr) = addrs.iter().copied().find(SocketAddr::is_ipv4) { + return Ok(addr); + } + addrs.first().copied().ok_or_else(|| { + ProxyError::Proxy(format!("DNS returned no addresses for {host}:{port}")) + }) + } + fn now_epoch_ms() -> u64 { std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) @@ -709,11 +739,12 @@ impl UpstreamManager { } async fn connect_hostname_with_dns_override( + &self, address: &str, connect_timeout: Duration, ) -> Result { if let Some((host, port)) = split_host_port(address) - && let Some(addr) = resolve_socket_addr(&host, port) + && let Some(addr) = self.dns_overrides.load().resolve_socket_addr(&host, port) { return match tokio::time::timeout(connect_timeout, TcpStream::connect(addr)).await { Ok(Ok(stream)) => Ok(stream), @@ -1203,7 +1234,8 @@ impl UpstreamManager { "SOCKS4 interface binding is not supported for hostname addresses, ignoring" ); } - Self::connect_hostname_with_dns_override(address, connect_timeout).await? + self.connect_hostname_with_dns_override(address, connect_timeout) + .await? }; // replace socks user_id with config.selected_scope, if set @@ -1291,7 +1323,8 @@ impl UpstreamManager { "SOCKS5 interface binding is not supported for hostname addresses, ignoring" ); } - Self::connect_hostname_with_dns_override(address, connect_timeout).await? + self.connect_hostname_with_dns_override(address, connect_timeout) + .await? }; debug!(config = ?config, "Socks5 connection");