astelm/telemt
1
0
Fork 0
mirror of https://github.com/telemt/telemt.git synced 2026-07-21 13:10:19 +03:00
Code Issues Packages Projects Releases Wiki Activity

Atomic Maestro Sessions + Shutdown gate

Browse Source 
Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
This commit is contained in:
Alexey
2026-07-19 16:10:31 +03:00
parent fabd98ce89
commit 4679bdcfd5
14 changed files with 1453 additions and 219 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/api/mod.rs
+40 -21
View File
@@ -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<u64> {
.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") => {
src/api/reload_tests.rs
+119
View File
@@ -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
);
}
src/api/runtime_watch.rs
+39 -6
View File
@@ -14,19 +14,20 @@ pub(super) fn spawn_runtime_watchers(
runtime_state: Arc<ApiRuntimeState>,
runtime_events: Arc<ApiEventStore>,
) {
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<Option<RuntimeWatchState>>,
runtime_state: Arc<ApiRuntimeState>,
runtime_events: Arc<ApiEventStore>,
) {
) -> 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<Option<RuntimeWatchState>>,
runtime_state: Arc<ApiRuntimeState>,
runtime_events: Arc<ApiEventStore>,
) {
) -> 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
);
}
}
src/maestro/generation.rs
+213 -16
View File
@@ -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<SessionRegistration<'_>> {
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<F>(&self, future: F)
where
F: Future<Output = ()> + 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<Arc<ProxyConfig>>,
@@ -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<ProxyConfig> {
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<Arc<MePool>> {
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<F>(&self, future: F) -> bool
where
F: Future<Output = ()> + 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<RuntimeGeneration> {
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<tokio::sync::Notify>);
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());
}
}
src/maestro/me_startup.rs
+31
View File
@@ -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<Notify>);
@@ -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);
}
}
src/maestro/mod.rs
+2 -1
View File
@@ -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;
src/maestro/reload.rs
+56 -100
View File
@@ -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<Option<Self>, 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<ReloadCommand>,
}
#[derive(Default)]
struct ReloadStatusState {
next_reload_id: u64,
active_reload_id: Option<u64>,
statuses: VecDeque<ReloadStatus>,
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<ProxyConfig>,
@@ -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<ReloadStatus> {
self.status_store.get(reload_id).await
}
/// Returns the identifier of the currently active reload.
pub(crate) async fn in_progress(&self) -> Option<u64> {
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<String>) {
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<String>) {
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<String>) {
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<String>) {
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<ReloadCommand> {
self.command_rx.recv().await
}
@@ -311,6 +339,9 @@ impl ReloadStatusStore {
request: ReloadRequest,
) -> Result<ReloadStatus, ReloadSubmitError> {
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;
src/maestro/reload_supervisor.rs
+69 -45
View File
@@ -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<Option<RuntimeWatchState>>,
}
/// 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<ArcSwap<RuntimeGeneration>>,
control: ReloadControl,
@@ -79,7 +99,7 @@ impl ReloadSupervisor {
detected_ips_tx: watch::Sender<(Option<std::net::IpAddr>, Option<std::net::IpAddr>)>,
runtime_log_filter: RuntimeLogFilter,
runtime_watch_tx: watch::Sender<Option<RuntimeWatchState>>,
) {
) -> 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<InstallDns>(
&self,
command: ReloadCommand,
old_runtime: Arc<RuntimeGeneration>,
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;
src/maestro/reload_supervisor_tests.rs
+320
View File
@@ -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<ReloadSupervisor>,
control: ReloadControl,
command: ReloadCommand,
old_runtime: Arc<RuntimeGeneration>,
new_runtime: Arc<RuntimeGeneration>,
runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
}
fn runtime_log_filter() -> RuntimeLogFilter {
let (_layer, handle) =
tracing_subscriber::reload::Layer::<EnvFilter, Registry>::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<Notify>);
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;
}
src/maestro/reload_tests.rs
+256
View File
@@ -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;
}
src/maestro/shutdown.rs
+11 -6
View File
@@ -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<ArcSwap<RuntimeGeneration>>,
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<ArcSwap<RuntimeGeneration>>,
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");
}
src/maestro/tls_bootstrap.rs
+132 -1
View File
@@ -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<UpstreamManager> {
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();
}
}
src/synlimit_control/mod.rs
+126 -23
View File
@@ -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<Option<SynLimitNamespace>> = 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<Option<RuntimeWatchState>>,
) -> 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<F, Fut>(
mut runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
shutdown: CancellationToken,
mut on_config: F,
)
where
) where
F: FnMut(u64, Arc<ProxyConfig>) -> Fut,
Fut: std::future::Future<Output = ()>,
{
@@ -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<Option<RuntimeWatchState>>,
previous_generation_id: u64,
shutdown: &CancellationToken,
) -> Option<RuntimeWatchState> {
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);
}
}
src/transport/upstream.rs
+39
View File
@@ -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::<SocketAddr>().unwrap()
);
assert_eq!(
candidate
.resolve_hostname("FRONT.EXAMPLE", 443)
.await
.unwrap(),
"[2001:db8::10]:443".parse::<SocketAddr>().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::<SocketAddr>().unwrap()
);
assert_eq!(
candidate
.resolve_hostname("front.example", 443)
.await
.unwrap(),
"192.0.2.20:443".parse::<SocketAddr>().unwrap()
);
}
#[test]
fn required_healthy_group_count_applies_three_group_threshold() {
assert_eq!(UpstreamManager::required_healthy_group_count(0), 0);