diff --git a/src/api/mod.rs b/src/api/mod.rs index 968ea2d..40bfdc7 100644 --- a/src/api/mod.rs +++ b/src/api/mod.rs @@ -23,7 +23,7 @@ use tracing::{debug, info, warn}; use crate::config::ApiGrayAction; use crate::ip_tracker::UserIpTracker; use crate::maestro::generation::{RuntimeGeneration, RuntimeWatchState}; -use crate::maestro::reload::{ReloadControl, ReloadRequest, ReloadSubmitError}; +use crate::maestro::reload::{ReloadAccepted, ReloadControl, ReloadRequest, ReloadSubmitError}; use crate::proxy::route_mode::RouteRuntimeController; use crate::proxy::shared_state::ProxySharedState; use crate::startup::StartupTracker; @@ -37,6 +37,8 @@ mod events; mod http_utils; mod model; mod patch; +#[cfg(test)] +mod reload_tests; mod runtime_edge; mod runtime_init; mod runtime_min; @@ -182,6 +184,35 @@ fn reload_status_route_id(path: &str) -> Option { .and_then(|id| id.parse().ok()) } +async fn submit_reload_from_disk( + config_path: &std::path::Path, + mutation_lock: &Mutex<()>, + reload_control: &ReloadControl, + expected_revision: Option<&str>, + request: ReloadRequest, +) -> Result<(ReloadAccepted, String), ApiFailure> { + let _guard = mutation_lock.lock().await; + ensure_expected_revision(config_path, expected_revision).await?; + let revision = current_revision(config_path).await?; + let config = Arc::new(load_config_for_reload(config_path).await?); + let accepted = 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((accepted, revision)) +} + fn allowed_methods_for_path(path: &str) -> Option<&'static str> { match path { "/v1/health" @@ -740,26 +771,14 @@ async fn handle( .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", - ), - })?; + let (accepted, revision) = submit_reload_from_disk( + &shared.config_path, + shared.mutation_lock.as_ref(), + &shared.reload_control, + expected_revision.as_deref(), + request, + ) + .await?; Ok(success_response(StatusCode::ACCEPTED, accepted, revision)) } ("PATCH", "/v1/config") => { diff --git a/src/api/reload_tests.rs b/src/api/reload_tests.rs new file mode 100644 index 0000000..d4857ef --- /dev/null +++ b/src/api/reload_tests.rs @@ -0,0 +1,119 @@ +use super::*; +use crate::config::ProxyConfig; + +async fn config_file() -> (tempfile::TempDir, PathBuf, String) { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("config.toml"); + let mut config = ProxyConfig::default(); + config.server.max_connections = 4_242; + let body = toml::to_string_pretty(&config).unwrap(); + tokio::fs::write(&path, &body).await.unwrap(); + let revision = config_store::compute_revision(&body); + (directory, path, revision) +} + +#[tokio::test] +async fn reload_submission_uses_matching_disk_revision_and_snapshot() { + let (_directory, path, revision) = config_file().await; + let mutation_lock = Mutex::new(()); + let (control, mut commands) = ReloadControl::channel(1); + let request = ReloadRequest::default(); + + let (accepted, response_revision) = submit_reload_from_disk( + &path, + &mutation_lock, + &control, + Some(&revision), + request.clone(), + ) + .await + .unwrap(); + let command = commands.recv().await.unwrap(); + + assert_eq!(response_revision, revision); + assert_eq!(accepted.config_revision, revision); + assert_eq!(command.config_revision, revision); + assert_eq!(command.request, request); + assert_eq!(command.config.server.max_connections, 4_242); +} + +#[tokio::test] +async fn revision_conflict_rejects_without_enqueuing_reload() { + let (_directory, path, _revision) = config_file().await; + let mutation_lock = Mutex::new(()); + let (control, _commands) = ReloadControl::channel(1); + + let error = submit_reload_from_disk( + &path, + &mutation_lock, + &control, + Some("stale-revision"), + ReloadRequest::default(), + ) + .await + .unwrap_err(); + + assert_eq!(error.status, StatusCode::CONFLICT); + assert_eq!(error.code, "revision_conflict"); + assert_eq!(control.in_progress().await, None); +} + +#[tokio::test] +async fn reload_conflict_and_closed_coordinator_map_to_http_contract() { + let (_directory, path, _revision) = config_file().await; + let mutation_lock = Mutex::new(()); + let (control, mut commands) = ReloadControl::channel(1); + let _accepted = submit_reload_from_disk( + &path, + &mutation_lock, + &control, + None, + ReloadRequest::default(), + ) + .await + .unwrap(); + let _command = commands.recv().await.unwrap(); + + let conflict = submit_reload_from_disk( + &path, + &mutation_lock, + &control, + None, + ReloadRequest::default(), + ) + .await + .unwrap_err(); + assert_eq!(conflict.status, StatusCode::CONFLICT); + assert_eq!(conflict.code, "reload_in_progress"); + + control.fail(1, "test cleanup").await; + drop(commands); + let unavailable = submit_reload_from_disk( + &path, + &mutation_lock, + &control, + None, + ReloadRequest::default(), + ) + .await + .unwrap_err(); + assert_eq!(unavailable.status, StatusCode::SERVICE_UNAVAILABLE); + assert_eq!(unavailable.code, "maestro_unavailable"); +} + +#[test] +fn reload_routes_expose_only_documented_methods_and_ids() { + assert_eq!( + allowed_methods_for_path("/v1/system/reload"), + Some(ALLOW_POST) + ); + assert_eq!( + allowed_methods_for_path("/v1/system/reload/42"), + Some(ALLOW_GET) + ); + assert_eq!(reload_status_route_id("/v1/system/reload/42"), Some(42)); + assert_eq!( + reload_status_route_id("/v1/system/reload/not-a-number"), + None + ); +} diff --git a/src/api/runtime_watch.rs b/src/api/runtime_watch.rs index 71a352f..50a602d 100644 --- a/src/api/runtime_watch.rs +++ b/src/api/runtime_watch.rs @@ -14,19 +14,20 @@ pub(super) fn spawn_runtime_watchers( runtime_state: Arc, runtime_events: Arc, ) { - spawn_config_watcher( + let _config_watcher = spawn_config_watcher( runtime_watch_rx.clone(), runtime_state.clone(), runtime_events.clone(), ); - spawn_admission_watcher(runtime_watch_rx, runtime_state, runtime_events); + let _admission_watcher = + spawn_admission_watcher(runtime_watch_rx, runtime_state, runtime_events); } fn spawn_config_watcher( mut runtime_watch_rx: watch::Receiver>, runtime_state: Arc, runtime_events: Arc, -) { +) -> tokio::task::JoinHandle<()> { tokio::spawn(async move { let Some(mut current) = runtime_watch_rx.borrow().clone() else { return; @@ -77,14 +78,14 @@ fn spawn_config_watcher( } } } - }); + }) } fn spawn_admission_watcher( mut runtime_watch_rx: watch::Receiver>, runtime_state: Arc, runtime_events: Arc, -) { +) -> tokio::task::JoinHandle<()> { tokio::spawn(async move { let Some(mut current) = runtime_watch_rx.borrow().clone() else { return; @@ -123,7 +124,7 @@ fn spawn_admission_watcher( } } } - }); + }) } fn active_generation_id( @@ -283,4 +284,36 @@ mod tests { 3 ); } + + #[tokio::test] + async fn watcher_recovers_from_closed_generation_and_exits_with_process_channel() { + let (initial, initial_config_tx, _initial_admission_tx) = state(1); + let (runtime_watch_tx, runtime_watch_rx) = watch::channel(Some(initial)); + let runtime_state = runtime_state(); + let events = Arc::new(ApiEventStore::new(16)); + let watcher = spawn_config_watcher(runtime_watch_rx, runtime_state.clone(), events.clone()); + drop(initial_config_tx); + tokio::task::yield_now().await; + + let (next, next_config_tx, _next_admission_tx) = state(2); + runtime_watch_tx.send_replace(Some(next)); + wait_for_count(&runtime_state, 1).await; + next_config_tx.send_replace(Arc::new(ProxyConfig::default())); + wait_for_count(&runtime_state, 2).await; + + drop(runtime_watch_tx); + tokio::time::timeout(Duration::from_secs(1), watcher) + .await + .unwrap() + .unwrap(); + assert_eq!( + events + .snapshot(16) + .events + .iter() + .filter(|event| event.event_type == "config.reload.applied") + .count(), + 2 + ); + } } diff --git a/src/maestro/generation.rs b/src/maestro/generation.rs index 365a5bb..70ca21b 100644 --- a/src/maestro/generation.rs +++ b/src/maestro/generation.rs @@ -1,6 +1,6 @@ use std::future::Future; use std::sync::Arc; -use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; use tokio::sync::{RwLock, Semaphore, watch}; @@ -10,6 +10,8 @@ use tokio_util::task::TaskTracker; use crate::config::ProxyConfig; use crate::crypto::SecureRandom; use crate::ip_tracker::UserIpTracker; +#[cfg(test)] +use crate::proxy::route_mode::RelayRouteMode; use crate::proxy::route_mode::RouteRuntimeController; use crate::proxy::shared_state::ProxySharedState; use crate::stats::beobachten::BeobachtenStore; @@ -21,6 +23,66 @@ use crate::transport::middle_proxy::MePool; const SESSION_STOP_TIMEOUT: Duration = Duration::from_secs(5); const BACKGROUND_STOP_TIMEOUT: Duration = Duration::from_secs(5); +const SESSION_ADMISSION_CLOSED: usize = 1 << (usize::BITS - 1); +const SESSION_REGISTRATION_COUNT: usize = SESSION_ADMISSION_CLOSED - 1; + +struct SessionAdmission { + state: AtomicUsize, +} + +struct SessionRegistration<'a> { + admission: &'a SessionAdmission, +} + +impl SessionAdmission { + fn new() -> Self { + Self { + state: AtomicUsize::new(0), + } + } + + fn try_register(&self) -> Option> { + let mut state = self.state.load(Ordering::Acquire); + loop { + if state & SESSION_ADMISSION_CLOSED != 0 + || state & SESSION_REGISTRATION_COUNT == SESSION_REGISTRATION_COUNT + { + return None; + } + match self.state.compare_exchange_weak( + state, + state + 1, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => return Some(SessionRegistration { admission: self }), + Err(observed) => state = observed, + } + } + } + + fn close(&self) { + self.state + .fetch_or(SESSION_ADMISSION_CLOSED, Ordering::AcqRel); + } + + fn reopen(&self) { + self.state + .fetch_and(!SESSION_ADMISSION_CLOSED, Ordering::AcqRel); + } + + async fn wait_for_registrations(&self) { + while self.state.load(Ordering::Acquire) & SESSION_REGISTRATION_COUNT != 0 { + tokio::task::yield_now().await; + } + } +} + +impl Drop for SessionRegistration<'_> { + fn drop(&mut self) { + self.admission.state.fetch_sub(1, Ordering::Release); + } +} /// Process-visible control-plane receivers for one active runtime generation. #[derive(Clone)] @@ -38,6 +100,7 @@ pub(crate) struct RuntimeTaskScope { } impl RuntimeTaskScope { + /// Creates an open generation-owned task scope. pub(crate) fn new() -> Self { Self { tracker: TaskTracker::new(), @@ -45,6 +108,7 @@ impl RuntimeTaskScope { } } + /// Spawns one task that is cancelled when the generation stops. pub(crate) fn spawn(&self, future: F) where F: Future + Send + 'static, @@ -58,10 +122,12 @@ impl RuntimeTaskScope { }); } + /// Returns the cancellation signal shared by generation-owned controllers. pub(crate) fn cancellation_token(&self) -> CancellationToken { self.cancel.clone() } + /// Cancels the scope and waits within the bounded background-task budget. pub(crate) async fn stop(&self) { self.cancel.cancel(); self.tracker.close(); @@ -90,11 +156,12 @@ pub(crate) struct RuntimeGeneration { background_tasks: RuntimeTaskScope, sessions: TaskTracker, session_cancel: CancellationToken, - accepting_sessions: AtomicBool, + session_admission: SessionAdmission, } impl RuntimeGeneration { #[allow(clippy::too_many_arguments)] + /// Builds one fully owned runtime generation. pub(crate) fn new( id: u64, config_rx: watch::Receiver>, @@ -134,10 +201,11 @@ impl RuntimeGeneration { background_tasks, sessions: TaskTracker::new(), session_cancel: CancellationToken::new(), - accepting_sessions: AtomicBool::new(true), + session_admission: SessionAdmission::new(), }) } + /// Returns the latest hot-reloaded configuration for this generation. pub(crate) fn config(&self) -> Arc { self.config_rx.borrow().clone() } @@ -151,6 +219,7 @@ impl RuntimeGeneration { } } + /// Returns the initial or asynchronously published Middle-End pool. pub(crate) async fn current_me_pool(&self) -> Option> { if let Some(pool) = &self.me_pool { return Some(pool.clone()); @@ -158,13 +227,14 @@ impl RuntimeGeneration { self.me_pool_runtime.read().await.clone() } + /// Registers a session only while admission remains open. pub(crate) fn spawn_session(&self, future: F) -> bool where F: Future + Send + 'static, { - if !self.accepting_sessions.load(Ordering::Acquire) { + let Some(_registration) = self.session_admission.try_register() else { return false; - } + }; let cancel = self.session_cancel.clone(); self.sessions.spawn(async move { tokio::select! { @@ -175,16 +245,20 @@ impl RuntimeGeneration { true } + /// Closes admission while preserving already registered sessions. pub(crate) fn stop_accepting_sessions(&self) { - self.accepting_sessions.store(false, Ordering::Release); + self.session_admission.close(); } + /// Reopens admission after a candidate activation rolls back. pub(crate) fn resume_accepting_sessions(&self) { - self.accepting_sessions.store(true, Ordering::Release); + self.session_admission.reopen(); } + /// Waits for registered sessions and cancels them when the deadline expires. pub(crate) async fn drain_sessions(&self, timeout: Duration) -> bool { self.stop_accepting_sessions(); + self.session_admission.wait_for_registrations().await; self.sessions.close(); if tokio::time::timeout(timeout, self.sessions.wait()) .await @@ -196,35 +270,90 @@ impl RuntimeGeneration { false } + /// Cancels all sessions and waits within the bounded session-stop budget. pub(crate) async fn stop_sessions(&self) { self.stop_accepting_sessions(); + self.session_admission.wait_for_registrations().await; self.session_cancel.cancel(); self.sessions.close(); let _ = tokio::time::timeout(SESSION_STOP_TIMEOUT, self.sessions.wait()).await; } + /// Stops all background tasks owned by this generation. pub(crate) async fn stop_background_tasks(&self) { self.background_tasks.stop().await; } } +#[cfg(test)] +/// Builds a lightweight runtime generation without network startup tasks. +pub(super) fn test_runtime_generation(id: u64, config: ProxyConfig) -> Arc { + let (config_tx, config_rx) = watch::channel(Arc::new(config.clone())); + let (_admission_tx, admission_rx) = watch::channel(true); + let stats = Arc::new(Stats::new()); + let upstream_manager = Arc::new(UpstreamManager::new( + config.upstreams, + 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 _config_tx = config_tx; + RuntimeGeneration::new( + id, + config_rx, + admission_rx, + stats, + upstream_manager, + Arc::new(ReplayChecker::new(128, Duration::from_secs(60))), + Arc::new(BufferPool::with_config(4096, 16)), + Arc::new(SecureRandom::new()), + None, + Arc::new(RwLock::new(None)), + Arc::new(RouteRuntimeController::new(RelayRouteMode::Direct)), + None, + Arc::new(UserIpTracker::new()), + Arc::new(BeobachtenStore::new()), + ProxySharedState::new(), + Arc::new(Semaphore::new(64)), + RuntimeTaskScope::new(), + ) +} + #[cfg(test)] mod tests { use super::*; + use tokio::sync::Barrier; #[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()) + let generation = test_runtime_generation(1, ProxyConfig::default()); + let started = Arc::new(tokio::sync::Notify::new()); + let dropped = Arc::new(tokio::sync::Notify::new()); + let started_task = started.clone(); + let dropped_task = dropped.clone(); + assert!(generation.spawn_session(async move { + struct DropSignal(Arc); + impl Drop for DropSignal { + fn drop(&mut self) { + self.0.notify_one(); + } + } + let _drop_signal = DropSignal(dropped_task); + started_task.notify_one(); + std::future::pending::<()>().await; + })); + started.notified().await; + + generation.stop_sessions().await; + + tokio::time::timeout(Duration::from_secs(1), dropped.notified()) .await .unwrap(); + assert!(!generation.spawn_session(async {})); } #[tokio::test] @@ -235,4 +364,72 @@ mod tests { .await .unwrap(); } + + #[tokio::test] + async fn session_admission_waits_for_registration_started_before_cutover() { + let admission = Arc::new(SessionAdmission::new()); + let registration = admission.try_register().unwrap(); + admission.close(); + assert!(admission.try_register().is_none()); + + let wait_admission = admission.clone(); + let waiter = tokio::spawn(async move { + wait_admission.wait_for_registrations().await; + }); + tokio::task::yield_now().await; + assert!(!waiter.is_finished()); + + drop(registration); + tokio::time::timeout(Duration::from_secs(1), waiter) + .await + .unwrap() + .unwrap(); + + admission.reopen(); + assert!(admission.try_register().is_some()); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn cutover_never_leaves_late_session_registrations() { + const ATTEMPTS: usize = 10_000; + + let admission = Arc::new(SessionAdmission::new()); + let tracker = TaskTracker::new(); + let cancel = CancellationToken::new(); + let start = Arc::new(Barrier::new(ATTEMPTS + 1)); + let live = Arc::new(AtomicUsize::new(0)); + let mut attempts = tokio::task::JoinSet::new(); + + for _ in 0..ATTEMPTS { + let admission = admission.clone(); + let tracker = tracker.clone(); + let cancel = cancel.clone(); + let start = start.clone(); + let live = live.clone(); + attempts.spawn(async move { + start.wait().await; + let Some(_registration) = admission.try_register() else { + return; + }; + tracker.spawn(async move { + live.fetch_add(1, Ordering::AcqRel); + cancel.cancelled().await; + live.fetch_sub(1, Ordering::AcqRel); + }); + }); + } + + start.wait().await; + admission.close(); + admission.wait_for_registrations().await; + tracker.close(); + cancel.cancel(); + while attempts.join_next().await.is_some() {} + tokio::time::timeout(Duration::from_secs(1), tracker.wait()) + .await + .unwrap(); + + assert_eq!(live.load(Ordering::Acquire), 0); + assert!(admission.try_register().is_none()); + } } diff --git a/src/maestro/me_startup.rs b/src/maestro/me_startup.rs index 024230c..f6ee98a 100644 --- a/src/maestro/me_startup.rs +++ b/src/maestro/me_startup.rs @@ -585,6 +585,7 @@ pub(crate) async fn initialize_me_pool( #[cfg(test)] mod tests { use super::*; + use std::sync::atomic::{AtomicUsize, Ordering}; use tokio::sync::Notify; struct DropSignal(Arc); @@ -615,4 +616,34 @@ mod tests { .await .unwrap(); } + + #[tokio::test] + async fn supervisor_restarts_exited_child_and_stops_with_runtime_scope() { + let scope = RuntimeTaskScope::new(); + let starts = Arc::new(AtomicUsize::new(0)); + let restarted = Arc::new(Notify::new()); + let starts_task = starts.clone(); + let restarted_task = restarted.clone(); + scope.spawn(supervise_me_task("restart_test", move || { + let starts = starts_task.clone(); + let restarted = restarted_task.clone(); + async move { + if starts.fetch_add(1, Ordering::AcqRel) + 1 >= 3 { + restarted.notify_one(); + } + } + })); + + tokio::time::timeout(Duration::from_secs(1), restarted.notified()) + .await + .unwrap(); + scope.stop().await; + let stopped_at = starts.load(Ordering::Acquire); + for _ in 0..100 { + tokio::task::yield_now().await; + } + + assert!(stopped_at >= 3); + assert_eq!(starts.load(Ordering::Acquire), stopped_at); + } } diff --git a/src/maestro/mod.rs b/src/maestro/mod.rs index 6b2af2a..b00a7b7 100644 --- a/src/maestro/mod.rs +++ b/src/maestro/mod.rs @@ -954,7 +954,7 @@ async fn run_telemt_core( runtime_task_scope.clone(), ); let active_runtime = Arc::new(ArcSwap::from(runtime_generation)); - reload_supervisor::ReloadSupervisor::spawn( + let reload_supervisor = reload_supervisor::ReloadSupervisor::spawn( active_runtime.clone(), reload_control, reload_commands, @@ -1011,6 +1011,7 @@ async fn run_telemt_core( active_runtime, quota_state_path, synlimit_controller, + reload_supervisor, ) .await; diff --git a/src/maestro/reload.rs b/src/maestro/reload.rs index 0a7ab58..375bcb5 100644 --- a/src/maestro/reload.rs +++ b/src/maestro/reload.rs @@ -43,6 +43,7 @@ pub(crate) struct ReloadRequest { } impl ReloadRequest { + /// Validates mode-specific request parameters. pub(crate) fn validate(&self) -> Result<(), &'static str> { match (self.mode, self.timeout_secs) { (ReloadMode::Instant, None) => Ok(()), @@ -53,6 +54,7 @@ impl ReloadRequest { } } + /// Parses optional PATCH query parameters into a reload request. pub(crate) fn from_query(query: Option<&str>) -> Result, String> { let Some(query) = query.filter(|query| !query.is_empty()) else { return Ok(None); @@ -186,11 +188,22 @@ pub(crate) struct ReloadCommandReceiver { command_rx: mpsc::Receiver, } -#[derive(Default)] struct ReloadStatusState { next_reload_id: u64, active_reload_id: Option, statuses: VecDeque, + accepting_commands: bool, +} + +impl Default for ReloadStatusState { + fn default() -> Self { + Self { + next_reload_id: 0, + active_reload_id: None, + statuses: VecDeque::new(), + accepting_commands: true, + } + } } #[derive(Default)] @@ -199,6 +212,7 @@ struct ReloadStatusStore { } impl ReloadControl { + /// Creates the process-scoped coordinator channel and status store. pub(crate) fn channel(initial_generation: u64) -> (Self, ReloadCommandReceiver) { let (command_tx, command_rx) = mpsc::channel(RELOAD_COMMAND_CAPACITY); ( @@ -211,6 +225,7 @@ impl ReloadControl { ) } + /// Atomically reserves and enqueues one reload operation. pub(crate) async fn submit( &self, config: Arc, @@ -232,7 +247,7 @@ impl ReloadControl { config_revision: status.config_revision.clone(), request, }; - if self.command_tx.send(command).await.is_err() { + if self.command_tx.try_send(command).is_err() { self.status_store .finish( status.reload_id, @@ -252,43 +267,55 @@ impl ReloadControl { }) } + /// Returns a retained reload status by identifier. pub(crate) async fn status(&self, reload_id: u64) -> Option { self.status_store.get(reload_id).await } + /// Returns the identifier of the currently active reload. pub(crate) async fn in_progress(&self) -> Option { self.status_store.state.lock().await.active_reload_id } + /// Rejects new commands while preserving an already accepted operation. + pub(crate) async fn begin_shutdown(&self) { + self.status_store.state.lock().await.accepting_commands = false; + } + + /// Records a non-terminal lifecycle phase. pub(crate) async fn mark_phase(&self, reload_id: u64, phase: ReloadPhase) { self.status_store.mark_phase(reload_id, phase).await; } + /// Records process-owned fields deferred until the next process restart. 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; } + /// Commits the active generation and completes the matching reload. 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) + .finish_success(reload_id, generation, &self.active_generation) .await; } + /// Marks the matching reload as failed. pub(crate) async fn fail(&self, reload_id: u64, error: impl Into) { self.status_store .finish(reload_id, ReloadPhase::Failed, Some(error.into())) .await; } + /// Marks the matching reload as rolled back. 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; } + /// Appends a non-fatal warning to the matching reload status. pub(crate) async fn add_warning(&self, reload_id: u64, warning: impl Into) { let warning = warning.into(); self.status_store @@ -298,6 +325,7 @@ impl ReloadControl { } impl ReloadCommandReceiver { + /// Receives the next accepted reload command. pub(crate) async fn recv(&mut self) -> Option { self.command_rx.recv().await } @@ -311,6 +339,9 @@ impl ReloadStatusStore { request: ReloadRequest, ) -> Result { let mut state = self.state.lock().await; + if !state.accepting_commands { + return Err(ReloadSubmitError::MaestroUnavailable); + } if let Some(reload_id) = state.active_reload_id { return Err(ReloadSubmitError::InProgress(reload_id)); } @@ -375,6 +406,25 @@ impl ReloadStatusStore { } } + async fn finish_success(&self, reload_id: u64, generation: u64, active_generation: &AtomicU64) { + let mut state = self.state.lock().await; + if state.active_reload_id != Some(reload_id) { + return; + } + let Some(status) = state + .statuses + .iter_mut() + .find(|status| status.reload_id == reload_id) + else { + return; + }; + status.state = ReloadPhase::Succeeded; + status.error = None; + status.finished_at_epoch_secs = Some(now_epoch_secs()); + active_generation.store(generation, Ordering::Release); + 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 @@ -395,99 +445,5 @@ fn now_epoch_secs() -> u64 { } #[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()); - } - - #[test] - fn status_uses_documented_deferred_process_fields_key() { - let status = ReloadStatus { - reload_id: 1, - target_generation: 2, - config_revision: "revision".to_string(), - state: ReloadPhase::Succeeded, - mode: ReloadMode::Instant, - failure_policy: ReloadFailurePolicy::KeepNew, - requested_at_epoch_secs: 10, - started_at_epoch_secs: Some(11), - finished_at_epoch_secs: Some(12), - deferred_fields: vec!["server.listeners".to_string()], - warnings: Vec::new(), - error: None, - }; - let value = serde_json::to_value(status).unwrap(); - - assert_eq!( - value["deferred_process_fields"], - serde_json::json!(["server.listeners"]) - ); - assert!(value.get("deferred_fields").is_none()); - } - - #[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); - } -} +#[path = "reload_tests.rs"] +mod tests; diff --git a/src/maestro/reload_supervisor.rs b/src/maestro/reload_supervisor.rs index f71f58a..8344218 100644 --- a/src/maestro/reload_supervisor.rs +++ b/src/maestro/reload_supervisor.rs @@ -4,6 +4,7 @@ use std::time::Duration; use arc_swap::ArcSwap; use tokio::sync::watch; +use tokio_util::sync::CancellationToken; use tracing::{info, warn}; use crate::stats::QuotaStore; @@ -13,7 +14,7 @@ use super::reload::{ ReloadCommand, ReloadCommandReceiver, ReloadControl, ReloadFailurePolicy, ReloadMode, ReloadPhase, }; -use super::runtime_build::{deferred_process_fields, prepare_runtime}; +use super::runtime_build::{PreparedRuntime, deferred_process_fields, prepare_runtime}; use super::runtime_tasks::RuntimeLogFilter; pub(crate) struct ReloadSupervisor { @@ -27,6 +28,24 @@ pub(crate) struct ReloadSupervisor { runtime_watch_tx: watch::Sender>, } +/// Process-owned handle that quiesces reloads before shutdown snapshots the runtime. +pub(crate) struct ReloadSupervisorHandle { + control: ReloadControl, + shutdown: CancellationToken, + join: tokio::task::JoinHandle<()>, +} + +impl ReloadSupervisorHandle { + /// Stops new submissions and waits for the accepted reload to finish. + pub(crate) async fn quiesce(self) { + self.control.begin_shutdown().await; + self.shutdown.cancel(); + if let Err(error) = self.join.await { + warn!(error = %error, "Reload supervisor failed while quiescing"); + } + } +} + #[derive(Debug, PartialEq, Eq)] enum RevisionGateAction { Proceed, @@ -70,6 +89,7 @@ async fn cleanup_candidate(generation: &RuntimeGeneration) -> bool { impl ReloadSupervisor { #[allow(clippy::too_many_arguments)] + /// Starts the process-scoped reload supervisor and returns its shutdown owner. pub(crate) fn spawn( active_runtime: Arc>, control: ReloadControl, @@ -79,7 +99,7 @@ impl ReloadSupervisor { detected_ips_tx: watch::Sender<(Option, Option)>, runtime_log_filter: RuntimeLogFilter, runtime_watch_tx: watch::Sender>, - ) { + ) -> ReloadSupervisorHandle { let supervisor = Self { active_runtime, control, @@ -90,12 +110,35 @@ impl ReloadSupervisor { runtime_log_filter, runtime_watch_tx, }; - tokio::spawn(supervisor.run()); + let control = supervisor.control.clone(); + let shutdown = CancellationToken::new(); + let join = tokio::spawn(supervisor.run(shutdown.clone())); + ReloadSupervisorHandle { + control, + shutdown, + join, + } } - async fn run(mut self) { - while let Some(command) = self.commands.recv().await { - self.reload(command).await; + async fn run(mut self, shutdown: CancellationToken) { + loop { + tokio::select! { + biased; + _ = shutdown.cancelled() => { + if self.control.in_progress().await.is_some() + && let Some(command) = self.commands.recv().await + { + self.reload(command).await; + } + break; + } + command = self.commands.recv() => { + let Some(command) = command else { + break; + }; + self.reload(command).await; + } + } } } @@ -130,6 +173,23 @@ impl ReloadSupervisor { crate::api::config_store::current_revision_for_maestro(&self.config_path).await, command.request.failure_policy, ); + self.activate_prepared(command, old_runtime, prepared, revision_action, |entries| { + crate::network::dns_overrides::install_entries(entries) + .map_err(|error| error.to_string()) + }) + .await; + } + + async fn activate_prepared( + &self, + command: ReloadCommand, + old_runtime: Arc, + prepared: PreparedRuntime, + revision_action: RevisionGateAction, + install_dns: InstallDns, + ) where + InstallDns: FnOnce(&[String]) -> Result<(), String>, + { match revision_action { RevisionGateAction::Proceed => {} RevisionGateAction::Warn(warning) => { @@ -149,9 +209,7 @@ impl ReloadSupervisor { .await; let new_runtime = prepared.generation; old_runtime.stop_accepting_sessions(); - if let Err(error) = crate::network::dns_overrides::install_entries( - &new_runtime.config().network.dns_overrides, - ) { + if let Err(error) = install_dns(&new_runtime.config().network.dns_overrides) { let message = format!("runtime DNS activation failed: {}", error); if command.request.failure_policy == ReloadFailurePolicy::Rollback { old_runtime.resume_accepting_sessions(); @@ -218,39 +276,5 @@ impl ReloadSupervisor { } #[cfg(test)] -mod tests { - use super::*; - - #[test] - fn revision_gate_proceeds_only_on_verified_match() { - assert_eq!( - revision_gate_action( - "accepted", - Ok("accepted".to_string()), - ReloadFailurePolicy::Rollback, - ), - RevisionGateAction::Proceed - ); - } - - #[test] - fn revision_gate_applies_failure_policy_to_mismatch_and_read_error() { - for result in [ - Ok("changed".to_string()), - Err("read failed".to_string()), - ] { - assert!(matches!( - revision_gate_action( - "accepted", - result.clone(), - ReloadFailurePolicy::KeepNew, - ), - RevisionGateAction::Warn(_) - )); - assert!(matches!( - revision_gate_action("accepted", result, ReloadFailurePolicy::Rollback), - RevisionGateAction::Rollback(_) - )); - } - } -} +#[path = "reload_supervisor_tests.rs"] +mod tests; diff --git a/src/maestro/reload_supervisor_tests.rs b/src/maestro/reload_supervisor_tests.rs new file mode 100644 index 0000000..d934016 --- /dev/null +++ b/src/maestro/reload_supervisor_tests.rs @@ -0,0 +1,320 @@ +use super::*; +use crate::config::ProxyConfig; +use crate::maestro::generation::test_runtime_generation; +use crate::maestro::reload::{ReloadRequest, ReloadSubmitError}; +use crate::stats::QuotaStore; +use tokio::sync::Notify; +use tracing_subscriber::{EnvFilter, Registry}; + +struct ReloadFixture { + supervisor: Arc, + control: ReloadControl, + command: ReloadCommand, + old_runtime: Arc, + new_runtime: Arc, + runtime_watch_rx: watch::Receiver>, +} + +fn runtime_log_filter() -> RuntimeLogFilter { + let (_layer, handle) = + tracing_subscriber::reload::Layer::::new(EnvFilter::new("info")); + RuntimeLogFilter::new(handle) +} + +async fn fixture(request: ReloadRequest) -> ReloadFixture { + let old_runtime = test_runtime_generation(1, ProxyConfig::default()); + let new_config = Arc::new(ProxyConfig::default()); + let new_runtime = test_runtime_generation(2, new_config.as_ref().clone()); + let active_runtime = Arc::new(ArcSwap::from(old_runtime.clone())); + let (control, commands) = ReloadControl::channel(old_runtime.id); + let accepted = control + .submit(new_config.clone(), "revision".to_string(), request.clone()) + .await + .unwrap(); + let (detected_ips_tx, _detected_ips_rx) = watch::channel((None, None)); + let (runtime_watch_tx, runtime_watch_rx) = watch::channel(Some(old_runtime.watch_state())); + let supervisor = Arc::new(ReloadSupervisor { + active_runtime, + control: control.clone(), + commands, + config_path: PathBuf::new(), + quota_store: Arc::new(QuotaStore::default()), + detected_ips_tx, + runtime_log_filter: runtime_log_filter(), + runtime_watch_tx, + }); + let command = ReloadCommand { + reload_id: accepted.reload_id, + target_generation: accepted.target_generation, + config: new_config, + config_revision: accepted.config_revision, + request, + }; + ReloadFixture { + supervisor, + control, + command, + old_runtime, + new_runtime, + runtime_watch_rx, + } +} + +struct DropSignal(Arc); + +impl Drop for DropSignal { + fn drop(&mut self) { + self.0.notify_one(); + } +} + +#[test] +fn revision_gate_proceeds_only_on_verified_match() { + assert_eq!( + revision_gate_action( + "accepted", + Ok("accepted".to_string()), + ReloadFailurePolicy::Rollback, + ), + RevisionGateAction::Proceed + ); +} + +#[test] +fn revision_gate_applies_failure_policy_to_mismatch_and_read_error() { + for result in [Ok("changed".to_string()), Err("read failed".to_string())] { + assert!(matches!( + revision_gate_action("accepted", result.clone(), ReloadFailurePolicy::KeepNew,), + RevisionGateAction::Warn(_) + )); + assert!(matches!( + revision_gate_action("accepted", result, ReloadFailurePolicy::Rollback), + RevisionGateAction::Rollback(_) + )); + } +} +#[tokio::test] +async fn revision_rollback_keeps_old_generation_and_cleans_candidate() { + let fixture = fixture(ReloadRequest { + failure_policy: ReloadFailurePolicy::Rollback, + ..ReloadRequest::default() + }) + .await; + let candidate_dropped = Arc::new(Notify::new()); + let candidate_drop = candidate_dropped.clone(); + assert!(fixture.new_runtime.spawn_session(async move { + let _drop_signal = DropSignal(candidate_drop); + std::future::pending::<()>().await; + })); + tokio::task::yield_now().await; + + fixture + .supervisor + .activate_prepared( + fixture.command, + fixture.old_runtime.clone(), + PreparedRuntime { + generation: fixture.new_runtime, + detected_ips: (None, None), + }, + RevisionGateAction::Rollback("revision changed".to_string()), + |_| -> Result<(), String> { panic!("DNS activation must not run on rollback") }, + ) + .await; + + tokio::time::timeout(Duration::from_secs(1), candidate_dropped.notified()) + .await + .unwrap(); + assert_eq!(fixture.supervisor.active_runtime.load().id, 1); + assert_eq!( + fixture + .runtime_watch_rx + .borrow() + .as_ref() + .unwrap() + .generation_id, + 1 + ); + assert!(fixture.old_runtime.spawn_session(async {})); + let status = fixture.control.status(1).await.unwrap(); + assert_eq!(status.state, ReloadPhase::RolledBack); + fixture.old_runtime.stop_sessions().await; +} + +#[tokio::test] +async fn dns_failure_policy_controls_rollback_or_keep_new() { + for policy in [ReloadFailurePolicy::Rollback, ReloadFailurePolicy::KeepNew] { + let fixture = fixture(ReloadRequest { + failure_policy: policy, + ..ReloadRequest::default() + }) + .await; + fixture + .supervisor + .activate_prepared( + fixture.command, + fixture.old_runtime.clone(), + PreparedRuntime { + generation: fixture.new_runtime.clone(), + detected_ips: (None, None), + }, + RevisionGateAction::Proceed, + |_| Err("invalid DNS entry".to_string()), + ) + .await; + + let status = fixture.control.status(1).await.unwrap(); + match policy { + ReloadFailurePolicy::Rollback => { + assert_eq!(fixture.supervisor.active_runtime.load().id, 1); + assert_eq!(status.state, ReloadPhase::RolledBack); + assert!(fixture.old_runtime.spawn_session(async {})); + fixture.old_runtime.stop_sessions().await; + } + ReloadFailurePolicy::KeepNew => { + assert_eq!(fixture.supervisor.active_runtime.load().id, 2); + assert_eq!(status.state, ReloadPhase::Succeeded); + assert_eq!(status.warnings.len(), 1); + assert!(!fixture.old_runtime.spawn_session(async {})); + fixture.new_runtime.stop_sessions().await; + } + } + } +} + +#[tokio::test] +async fn drain_publishes_new_generation_before_old_sessions_finish() { + let mut fixture = fixture(ReloadRequest { + mode: ReloadMode::Drain, + timeout_secs: Some(30), + ..ReloadRequest::default() + }) + .await; + let old_started = Arc::new(Notify::new()); + let old_release = Arc::new(Notify::new()); + let started = old_started.clone(); + let release = old_release.clone(); + assert!(fixture.old_runtime.spawn_session(async move { + started.notify_one(); + release.notified().await; + })); + old_started.notified().await; + + let supervisor = fixture.supervisor.clone(); + let old_runtime = fixture.old_runtime.clone(); + let new_runtime = fixture.new_runtime.clone(); + let activation = tokio::spawn(async move { + supervisor + .activate_prepared( + fixture.command, + old_runtime, + PreparedRuntime { + generation: new_runtime, + detected_ips: (None, None), + }, + RevisionGateAction::Proceed, + |_| Ok(()), + ) + .await; + }); + + fixture.runtime_watch_rx.changed().await.unwrap(); + assert_eq!( + fixture + .runtime_watch_rx + .borrow() + .as_ref() + .unwrap() + .generation_id, + 2 + ); + assert!(!activation.is_finished()); + assert!(!fixture.old_runtime.spawn_session(async {})); + + old_release.notify_one(); + activation.await.unwrap(); + assert_eq!( + fixture.control.status(1).await.unwrap().state, + ReloadPhase::Succeeded + ); + fixture.new_runtime.stop_sessions().await; +} + +#[tokio::test(start_paused = true)] +async fn drain_timeout_cancels_old_sessions_and_records_one_warning() { + let mut fixture = fixture(ReloadRequest { + mode: ReloadMode::Drain, + timeout_secs: Some(1), + ..ReloadRequest::default() + }) + .await; + let dropped = Arc::new(Notify::new()); + let drop_signal = dropped.clone(); + assert!(fixture.old_runtime.spawn_session(async move { + let _drop_signal = DropSignal(drop_signal); + std::future::pending::<()>().await; + })); + tokio::task::yield_now().await; + + let supervisor = fixture.supervisor.clone(); + let old_runtime = fixture.old_runtime.clone(); + let new_runtime = fixture.new_runtime.clone(); + let activation = tokio::spawn(async move { + supervisor + .activate_prepared( + fixture.command, + old_runtime, + PreparedRuntime { + generation: new_runtime, + detected_ips: (None, None), + }, + RevisionGateAction::Proceed, + |_| Ok(()), + ) + .await; + }); + fixture.runtime_watch_rx.changed().await.unwrap(); + tokio::task::yield_now().await; + tokio::time::advance(Duration::from_secs(1)).await; + activation.await.unwrap(); + + dropped.notified().await; + let status = fixture.control.status(1).await.unwrap(); + assert_eq!(status.state, ReloadPhase::Succeeded); + assert_eq!(status.warnings.len(), 1); + assert!(status.warnings[0].contains("exceeded drain timeout")); + fixture.new_runtime.stop_sessions().await; +} + +#[tokio::test] +async fn quiesce_joins_idle_supervisor_and_rejects_later_submissions() { + let runtime = test_runtime_generation(1, ProxyConfig::default()); + let active_runtime = Arc::new(ArcSwap::from(runtime.clone())); + let (control, commands) = ReloadControl::channel(runtime.id); + let (detected_ips_tx, _detected_ips_rx) = watch::channel((None, None)); + let (runtime_watch_tx, _runtime_watch_rx) = watch::channel(Some(runtime.watch_state())); + let handle = ReloadSupervisor::spawn( + active_runtime, + control.clone(), + commands, + PathBuf::new(), + Arc::new(QuotaStore::default()), + detected_ips_tx, + runtime_log_filter(), + runtime_watch_tx, + ); + + tokio::time::timeout(Duration::from_secs(1), handle.quiesce()) + .await + .unwrap(); + let result = control + .submit( + Arc::new(ProxyConfig::default()), + "revision".to_string(), + ReloadRequest::default(), + ) + .await; + + assert_eq!(result, Err(ReloadSubmitError::MaestroUnavailable)); + runtime.stop_sessions().await; +} diff --git a/src/maestro/reload_tests.rs b/src/maestro/reload_tests.rs new file mode 100644 index 0000000..f302214 --- /dev/null +++ b/src/maestro/reload_tests.rs @@ -0,0 +1,256 @@ +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()); +} + +#[test] +fn status_uses_documented_deferred_process_fields_key() { + let status = ReloadStatus { + reload_id: 1, + target_generation: 2, + config_revision: "revision".to_string(), + state: ReloadPhase::Succeeded, + mode: ReloadMode::Instant, + failure_policy: ReloadFailurePolicy::KeepNew, + requested_at_epoch_secs: 10, + started_at_epoch_secs: Some(11), + finished_at_epoch_secs: Some(12), + deferred_fields: vec!["server.listeners".to_string()], + warnings: Vec::new(), + error: None, + }; + let value = serde_json::to_value(status).unwrap(); + + assert_eq!( + value["deferred_process_fields"], + serde_json::json!(["server.listeners"]) + ); + assert!(value.get("deferred_fields").is_none()); +} + +#[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); +} + +#[tokio::test] +async fn terminal_outcomes_release_slot_and_only_success_advances_generation() { + let (control, mut receiver) = ReloadControl::channel(7); + + let failed = control + .submit( + Arc::new(ProxyConfig::default()), + "rev-failed".to_string(), + ReloadRequest::default(), + ) + .await + .unwrap(); + let _command = receiver.recv().await.unwrap(); + control + .mark_phase(failed.reload_id, ReloadPhase::Preparing) + .await; + control.fail(failed.reload_id, "prepare failed").await; + let failed_status = control.status(failed.reload_id).await.unwrap(); + assert_eq!(failed_status.state, ReloadPhase::Failed); + assert_eq!(failed_status.error.as_deref(), Some("prepare failed")); + assert!(failed_status.started_at_epoch_secs.is_some()); + assert!(failed_status.finished_at_epoch_secs.is_some()); + + let rolled_back = control + .submit( + Arc::new(ProxyConfig::default()), + "rev-rollback".to_string(), + ReloadRequest::default(), + ) + .await + .unwrap(); + let _command = receiver.recv().await.unwrap(); + assert_eq!(rolled_back.target_generation, 8); + control + .rolled_back(rolled_back.reload_id, "revision changed") + .await; + + let succeeded = control + .submit( + Arc::new(ProxyConfig::default()), + "rev-success".to_string(), + ReloadRequest::default(), + ) + .await + .unwrap(); + let _command = receiver.recv().await.unwrap(); + assert_eq!(succeeded.target_generation, 8); + control + .succeed(succeeded.reload_id, succeeded.target_generation) + .await; + + let next = control + .submit( + Arc::new(ProxyConfig::default()), + "rev-next".to_string(), + ReloadRequest::default(), + ) + .await + .unwrap(); + assert_eq!(next.target_generation, 9); +} + +#[tokio::test] +async fn stale_success_cannot_advance_generation_or_release_active_reload() { + let (control, mut receiver) = ReloadControl::channel(3); + let active = control + .submit( + Arc::new(ProxyConfig::default()), + "rev-active".to_string(), + ReloadRequest::default(), + ) + .await + .unwrap(); + let _command = receiver.recv().await.unwrap(); + + control.succeed(active.reload_id + 100, 99).await; + + assert_eq!(control.in_progress().await, Some(active.reload_id)); + control.fail(active.reload_id, "expected failure").await; + let next = control + .submit( + Arc::new(ProxyConfig::default()), + "rev-next".to_string(), + ReloadRequest::default(), + ) + .await + .unwrap(); + assert_eq!(next.target_generation, 4); +} + +#[tokio::test] +async fn status_history_retains_only_the_latest_entries() { + let (control, mut receiver) = ReloadControl::channel(1); + let mut reload_ids = Vec::new(); + for index in 0..=RELOAD_HISTORY_CAPACITY { + let accepted = control + .submit( + Arc::new(ProxyConfig::default()), + format!("rev-{index}"), + ReloadRequest::default(), + ) + .await + .unwrap(); + let _command = receiver.recv().await.unwrap(); + reload_ids.push(accepted.reload_id); + control.fail(accepted.reload_id, "expected failure").await; + } + + assert!(control.status(reload_ids[0]).await.is_none()); + assert!(control.status(reload_ids[1]).await.is_some()); + assert!(control.status(*reload_ids.last().unwrap()).await.is_some()); +} + +#[tokio::test] +async fn closed_command_channel_marks_reload_failed_and_releases_slot() { + let (control, receiver) = ReloadControl::channel(1); + drop(receiver); + + let result = control + .submit( + Arc::new(ProxyConfig::default()), + "rev-closed".to_string(), + ReloadRequest::default(), + ) + .await; + + assert_eq!(result, Err(ReloadSubmitError::MaestroUnavailable)); + assert_eq!(control.in_progress().await, None); + let status = control.status(1).await.unwrap(); + assert_eq!(status.state, ReloadPhase::Failed); + assert_eq!( + status.error.as_deref(), + Some("maestro command channel is closed") + ); +} + +#[tokio::test] +async fn shutdown_gate_rejects_new_commands_without_disturbing_active_status() { + let (control, mut receiver) = ReloadControl::channel(4); + let active = control + .submit( + Arc::new(ProxyConfig::default()), + "rev-active".to_string(), + ReloadRequest::default(), + ) + .await + .unwrap(); + let _command = receiver.recv().await.unwrap(); + + control.begin_shutdown().await; + let rejected = control + .submit( + Arc::new(ProxyConfig::default()), + "rev-rejected".to_string(), + ReloadRequest::default(), + ) + .await; + + assert_eq!(rejected, Err(ReloadSubmitError::MaestroUnavailable)); + assert_eq!(control.in_progress().await, Some(active.reload_id)); + control.fail(active.reload_id, "shutdown test").await; +} diff --git a/src/maestro/shutdown.rs b/src/maestro/shutdown.rs index 922a69e..1688bf4 100644 --- a/src/maestro/shutdown.rs +++ b/src/maestro/shutdown.rs @@ -21,6 +21,7 @@ use tracing::{info, warn}; use super::generation::RuntimeGeneration; use super::helpers::{format_uptime, unit_label}; +use super::reload_supervisor::ReloadSupervisorHandle; use crate::stats::Stats; use crate::synlimit_control; @@ -50,7 +51,8 @@ pub(crate) async fn wait_for_shutdown( process_started_at: Instant, active_runtime: Arc>, quota_state_path: PathBuf, - synlimit_controller: tokio::task::JoinHandle<()>, + synlimit_controller: synlimit_control::SynlimitController, + reload_supervisor: ReloadSupervisorHandle, ) { let signal = wait_for_shutdown_signal().await; perform_shutdown( @@ -59,6 +61,7 @@ pub(crate) async fn wait_for_shutdown( active_runtime, quota_state_path, synlimit_controller, + reload_supervisor, ) .await; } @@ -89,13 +92,16 @@ async fn perform_shutdown( process_started_at: Instant, active_runtime: Arc>, quota_state_path: PathBuf, - synlimit_controller: tokio::task::JoinHandle<()>, + synlimit_controller: synlimit_control::SynlimitController, + reload_supervisor: ReloadSupervisorHandle, ) { - let runtime = active_runtime.load_full(); - let stats = runtime.stats.as_ref(); let shutdown_started_at = Instant::now(); info!(signal = %signal, "Received shutdown signal"); + reload_supervisor.quiesce().await; + let runtime = active_runtime.load_full(); + let stats = runtime.stats.as_ref(); + // Dump stats if SIGQUIT if signal == ShutdownSignal::Quit { dump_stats(stats, process_started_at); @@ -124,8 +130,7 @@ async fn perform_shutdown( } } - synlimit_controller.abort(); - let _ = synlimit_controller.await; + synlimit_controller.shutdown().await; if let Err(error) = synlimit_control::clear_synlimit_rules_all_backends().await { warn!(error = %error, "Failed to clear SYN limiter rules during shutdown"); } diff --git a/src/maestro/tls_bootstrap.rs b/src/maestro/tls_bootstrap.rs index 58bd442..123f6ee 100644 --- a/src/maestro/tls_bootstrap.rs +++ b/src/maestro/tls_bootstrap.rs @@ -225,7 +225,46 @@ pub(crate) async fn bootstrap_tls_front( #[cfg(test)] mod tests { - use super::{readiness_error, tls_fetch_host_for_domain}; + use super::*; + use crate::startup::StartupComponentStatus; + use crate::stats::Stats; + + fn test_config(cache_dir: &std::path::Path) -> ProxyConfig { + let mut config = ProxyConfig::default(); + config.censorship.tls_emulation = true; + config.censorship.tls_domain = "front.example".to_string(); + config.censorship.mask_host = Some("127.0.0.1".to_string()); + config.censorship.mask_port = 1; + config.censorship.tls_front_dir = cache_dir.display().to_string(); + config.censorship.tls_fetch.profiles.truncate(1); + config.censorship.tls_fetch.attempt_timeout_ms = 10; + config.censorship.tls_fetch.total_budget_ms = 20; + config + } + + fn upstream_manager(config: &ProxyConfig) -> Arc { + Arc::new(UpstreamManager::new( + Vec::new(), + 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, + Arc::new(Stats::new()), + )) + } + + async fn tls_component_status(tracker: &StartupTracker) -> StartupComponentStatus { + tracker + .snapshot() + .await + .components + .into_iter() + .find(|component| component.id == COMPONENT_TLS_FRONT_BOOTSTRAP) + .unwrap() + .status + } #[test] fn tls_fetch_host_uses_each_domain_when_mask_host_is_primary_default() { @@ -251,4 +290,96 @@ mod tests { Some("TLS-front profiles are not ready for domains: front.example".to_string()) ); } + + #[tokio::test] + async fn require_ready_rejects_default_cache_after_bounded_fetch_failure() { + let cache_dir = tempfile::tempdir().unwrap(); + let config = test_config(cache_dir.path()); + let domains = vec![config.censorship.tls_domain.clone()]; + let tracker = Arc::new(StartupTracker::new(1)); + let scope = RuntimeTaskScope::new(); + + let result = bootstrap_tls_front( + &config, + &domains, + upstream_manager(&config), + &tracker, + scope.clone(), + TlsBootstrapPolicy::RequireReady, + ) + .await; + + assert!(result.is_err()); + assert_eq!( + tls_component_status(&tracker).await, + StartupComponentStatus::Failed + ); + scope.stop().await; + } + + #[tokio::test] + async fn require_ready_accepts_non_default_disk_cache_when_refresh_fails() { + let cache_dir = tempfile::tempdir().unwrap(); + let config = test_config(cache_dir.path()); + let domains = vec![config.censorship.tls_domain.clone()]; + let seed = TlsFrontCache::new(&domains, config.censorship.fake_cert_len, cache_dir.path()); + let mut cached = seed.default_entry().as_ref().clone(); + cached.domain = domains[0].clone(); + tokio::fs::write( + cache_dir.path().join("front.example.json"), + serde_json::to_vec(&cached).unwrap(), + ) + .await + .unwrap(); + let tracker = Arc::new(StartupTracker::new(1)); + let scope = RuntimeTaskScope::new(); + + let cache = bootstrap_tls_front( + &config, + &domains, + upstream_manager(&config), + &tracker, + scope.clone(), + TlsBootstrapPolicy::RequireReady, + ) + .await + .unwrap() + .unwrap(); + + assert!(cache.default_profile_domains(&domains).await.is_empty()); + assert_eq!( + tls_component_status(&tracker).await, + StartupComponentStatus::Ready + ); + scope.stop().await; + } + + #[tokio::test] + async fn best_effort_returns_ready_and_refresh_tasks_are_scope_owned() { + let cache_dir = tempfile::tempdir().unwrap(); + let config = test_config(cache_dir.path()); + let domains = vec![config.censorship.tls_domain.clone()]; + let tracker = Arc::new(StartupTracker::new(1)); + let scope = RuntimeTaskScope::new(); + + let cache = bootstrap_tls_front( + &config, + &domains, + upstream_manager(&config), + &tracker, + scope.clone(), + TlsBootstrapPolicy::BestEffort, + ) + .await + .unwrap(); + + assert!(cache.is_some()); + assert_eq!( + tls_component_status(&tracker).await, + StartupComponentStatus::Ready + ); + tokio::time::timeout(Duration::from_secs(1), scope.stop()) + .await + .unwrap(); + } } diff --git a/src/synlimit_control/mod.rs b/src/synlimit_control/mod.rs index b87f2b2..e6fdb34 100644 --- a/src/synlimit_control/mod.rs +++ b/src/synlimit_control/mod.rs @@ -1,6 +1,7 @@ use std::sync::{Arc, Mutex}; use tokio::sync::watch; +use tokio_util::sync::CancellationToken; use tracing::warn; use crate::config::{ProxyConfig, SynLimitMode}; @@ -16,13 +17,29 @@ use self::model::{SynLimitNamespace, synlimit_namespace, synlimit_targets}; static ACTIVE_SYNLIMIT_NAMESPACE: Mutex> = Mutex::new(None); +/// Process-owned lifecycle handle for the SYN limiter reconciler. +pub(crate) struct SynlimitController { + shutdown: CancellationToken, + join: tokio::task::JoinHandle<()>, +} + +impl SynlimitController { + /// Stops config observation after any in-flight reconcile completes. + pub(crate) async fn shutdown(self) { + self.shutdown.cancel(); + let _ = self.join.await; + } +} + /// Spawns the process-scoped SYN limiter reconciler for active generations. pub(crate) fn spawn_synlimit_controller( runtime_watch_rx: watch::Receiver>, -) -> tokio::task::JoinHandle<()> { - if !cfg!(target_os = "linux") { - return tokio::spawn(watch_active_runtime_configs( +) -> SynlimitController { + let shutdown = CancellationToken::new(); + let join = if !cfg!(target_os = "linux") { + tokio::spawn(watch_active_runtime_configs( runtime_watch_rx, + shutdown.clone(), |_generation_id, cfg| async move { if has_synlimit_config(&cfg) { warn!( @@ -30,22 +47,24 @@ pub(crate) fn spawn_synlimit_controller( ); } }, - )); - } - - tokio::spawn(watch_active_runtime_configs( - runtime_watch_rx, - |_generation_id, cfg| async move { - reconcile_synlimit_rules(&cfg).await; - }, - )) + )) + } else { + tokio::spawn(watch_active_runtime_configs( + runtime_watch_rx, + shutdown.clone(), + |_generation_id, cfg| async move { + reconcile_synlimit_rules(&cfg).await; + }, + )) + }; + SynlimitController { shutdown, join } } async fn watch_active_runtime_configs( mut runtime_watch_rx: watch::Receiver>, + shutdown: CancellationToken, mut on_config: F, -) -where +) where F: FnMut(u64, Arc) -> Fut, Fut: std::future::Future, { @@ -53,16 +72,26 @@ where if let Some(state) = runtime_watch_rx.borrow().clone() { break state; } - if runtime_watch_rx.changed().await.is_err() { - return; + tokio::select! { + biased; + _ = shutdown.cancelled() => return, + changed = runtime_watch_rx.changed() => { + if changed.is_err() { + return; + } + } } }; + if shutdown.is_cancelled() { + return; + } let initial_config = current.config_rx.borrow().clone(); on_config(current.generation_id, initial_config).await; loop { tokio::select! { biased; + _ = shutdown.cancelled() => break, changed = runtime_watch_rx.changed() => { if changed.is_err() { break; @@ -78,10 +107,11 @@ where } changed = current.config_rx.changed() => { if changed.is_err() { - let Some(next) = wait_for_new_runtime( - &mut runtime_watch_rx, - current.generation_id, - ).await else { + let Some(next) = wait_for_new_runtime( + &mut runtime_watch_rx, + current.generation_id, + &shutdown, + ).await else { break; }; current = next; @@ -105,6 +135,7 @@ where async fn wait_for_new_runtime( runtime_watch_rx: &mut watch::Receiver>, previous_generation_id: u64, + shutdown: &CancellationToken, ) -> Option { loop { if let Some(state) = runtime_watch_rx.borrow().clone() @@ -112,8 +143,14 @@ async fn wait_for_new_runtime( { return Some(state); } - if runtime_watch_rx.changed().await.is_err() { - return None; + tokio::select! { + biased; + _ = shutdown.cancelled() => return None, + changed = runtime_watch_rx.changed() => { + if changed.is_err() { + return None; + } + } } } } @@ -250,8 +287,9 @@ fn has_synlimit_config(cfg: &ProxyConfig) -> bool { #[cfg(test)] mod tests { use super::*; + use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; - use tokio::sync::mpsc; + use tokio::sync::{Notify, mpsc}; fn runtime_state( generation_id: u64, @@ -283,6 +321,7 @@ mod tests { let (observed_tx, mut observed_rx) = mpsc::unbounded_channel(); let watcher = tokio::spawn(watch_active_runtime_configs( runtime_rx, + CancellationToken::new(), move |generation_id, cfg| { let observed_tx = observed_tx.clone(); async move { @@ -312,4 +351,68 @@ mod tests { watcher.abort(); } + + #[tokio::test] + async fn shutdown_waits_for_inflight_reconcile_and_stops_future_updates() { + let (initial, config_tx, _admission_tx) = runtime_state(1, 10); + let (_runtime_tx, runtime_rx) = watch::channel(Some(initial)); + let shutdown = CancellationToken::new(); + let started = Arc::new(Notify::new()); + let release = Arc::new(Notify::new()); + let calls = Arc::new(AtomicUsize::new(0)); + let started_callback = started.clone(); + let release_callback = release.clone(); + let calls_callback = calls.clone(); + let watcher_shutdown = shutdown.clone(); + let watcher = tokio::spawn(watch_active_runtime_configs( + runtime_rx, + watcher_shutdown, + move |_generation_id, _cfg| { + let started = started_callback.clone(); + let release = release_callback.clone(); + let calls = calls_callback.clone(); + async move { + calls.fetch_add(1, Ordering::AcqRel); + started.notify_one(); + release.notified().await; + } + }, + )); + started.notified().await; + + shutdown.cancel(); + tokio::task::yield_now().await; + assert!(!watcher.is_finished()); + release.notify_one(); + tokio::time::timeout(Duration::from_secs(1), watcher) + .await + .unwrap() + .unwrap(); + + drop(_runtime_tx); + let mut updated = ProxyConfig::default(); + updated.server.max_connections = 20; + assert!(config_tx.send(Arc::new(updated)).is_err()); + assert_eq!(calls.load(Ordering::Acquire), 1); + } + + #[tokio::test] + async fn shutdown_before_start_skips_initial_reconcile() { + let (initial, _config_tx, _admission_tx) = runtime_state(1, 10); + let (_runtime_tx, runtime_rx) = watch::channel(Some(initial)); + let shutdown = CancellationToken::new(); + shutdown.cancel(); + let calls = Arc::new(AtomicUsize::new(0)); + let calls_callback = calls.clone(); + + watch_active_runtime_configs(runtime_rx, shutdown, move |_generation_id, _cfg| { + let calls = calls_callback.clone(); + async move { + calls.fetch_add(1, Ordering::AcqRel); + } + }) + .await; + + assert_eq!(calls.load(Ordering::Acquire), 0); + } } diff --git a/src/transport/upstream.rs b/src/transport/upstream.rs index 372d4d9..ad1b389 100644 --- a/src/transport/upstream.rs +++ b/src/transport/upstream.rs @@ -2052,6 +2052,45 @@ mod tests { const TEST_SHADOWSOCKS_URL: &str = "ss://2022-blake3-aes-256-gcm:MDEyMzQ1Njc4OTAxMjM0NTY3ODkwMTIzNDU2Nzg5MDE=@127.0.0.1:8388"; + fn manager_with_dns(entries: &[String]) -> UpstreamManager { + UpstreamManager::new(Vec::new(), 1, 1, 1, 1, 1, false, Arc::new(Stats::new())) + .with_dns_overrides(entries) + .unwrap() + } + + #[tokio::test] + async fn generation_local_dns_overrides_are_isolated_and_case_insensitive() { + let active = manager_with_dns(&["Front.Example:443:192.0.2.10".to_string()]); + let candidate = manager_with_dns(&["front.example:443:[2001:db8::10]".to_string()]); + + assert_eq!( + active.resolve_hostname("front.example", 443).await.unwrap(), + "192.0.2.10:443".parse::().unwrap() + ); + assert_eq!( + candidate + .resolve_hostname("FRONT.EXAMPLE", 443) + .await + .unwrap(), + "[2001:db8::10]:443".parse::().unwrap() + ); + + candidate + .update_dns_overrides(&["front.example:443:192.0.2.20".to_string()]) + .unwrap(); + assert_eq!( + active.resolve_hostname("FRONT.EXAMPLE", 443).await.unwrap(), + "192.0.2.10:443".parse::().unwrap() + ); + assert_eq!( + candidate + .resolve_hostname("front.example", 443) + .await + .unwrap(), + "192.0.2.20:443".parse::().unwrap() + ); + } + #[test] fn required_healthy_group_count_applies_three_group_threshold() { assert_eq!(UpstreamManager::required_healthy_group_count(0), 0);