Harden Maestro reload lifecycle and readiness barriers

Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
This commit is contained in:
Alexey
2026-07-18 14:27:24 +03:00
parent 91e05265be
commit c6f40e3717
15 changed files with 1053 additions and 506 deletions
+15 -6
View File
@@ -20,9 +20,9 @@ use tokio::sync::{Mutex, RwLock, Semaphore, watch};
use tokio::time::timeout; use tokio::time::timeout;
use tracing::{debug, info, warn}; use tracing::{debug, info, warn};
use crate::config::{ApiGrayAction, ProxyConfig}; use crate::config::ApiGrayAction;
use crate::ip_tracker::UserIpTracker; use crate::ip_tracker::UserIpTracker;
use crate::maestro::generation::RuntimeGeneration; use crate::maestro::generation::{RuntimeGeneration, RuntimeWatchState};
use crate::maestro::reload::{ReloadControl, ReloadRequest, ReloadSubmitError}; use crate::maestro::reload::{ReloadControl, ReloadRequest, ReloadSubmitError};
use crate::proxy::route_mode::RouteRuntimeController; use crate::proxy::route_mode::RouteRuntimeController;
use crate::proxy::shared_state::ProxySharedState; use crate::proxy::shared_state::ProxySharedState;
@@ -240,8 +240,6 @@ pub async fn serve(
route_runtime: Arc<RouteRuntimeController>, route_runtime: Arc<RouteRuntimeController>,
proxy_shared: Arc<ProxySharedState>, proxy_shared: Arc<ProxySharedState>,
upstream_manager: Arc<UpstreamManager>, upstream_manager: Arc<UpstreamManager>,
config_rx: watch::Receiver<Arc<ProxyConfig>>,
admission_rx: watch::Receiver<bool>,
config_path: PathBuf, config_path: PathBuf,
quota_state_path: PathBuf, quota_state_path: PathBuf,
detected_ips_rx: watch::Receiver<(Option<IpAddr>, Option<IpAddr>)>, detected_ips_rx: watch::Receiver<(Option<IpAddr>, Option<IpAddr>)>,
@@ -249,6 +247,7 @@ pub async fn serve(
startup_tracker: Arc<StartupTracker>, startup_tracker: Arc<StartupTracker>,
reload_control: ReloadControl, reload_control: ReloadControl,
mut active_runtime_rx: watch::Receiver<Option<Arc<ArcSwap<RuntimeGeneration>>>>, mut active_runtime_rx: watch::Receiver<Option<Arc<ArcSwap<RuntimeGeneration>>>>,
mut runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
) { ) {
let active_runtime = loop { let active_runtime = loop {
if let Some(active_runtime) = active_runtime_rx.borrow().clone() { if let Some(active_runtime) = active_runtime_rx.borrow().clone() {
@@ -259,6 +258,17 @@ pub async fn serve(
return; return;
} }
}; };
let initial_watch_state = loop {
if let Some(watch_state) = runtime_watch_rx.borrow().clone() {
break watch_state;
}
if runtime_watch_rx.changed().await.is_err() {
warn!("Runtime watch channel closed before API bootstrap");
return;
}
};
let config_rx = initial_watch_state.config_rx.clone();
let admission_rx = initial_watch_state.admission_rx.clone();
let listener = match TcpListener::bind(listen).await { let listener = match TcpListener::bind(listen).await {
Ok(listener) => listener, Ok(listener) => listener,
Err(error) => { Err(error) => {
@@ -306,8 +316,7 @@ pub async fn serve(
}); });
spawn_runtime_watchers( spawn_runtime_watchers(
config_rx.clone(), runtime_watch_rx,
admission_rx.clone(),
runtime_state.clone(), runtime_state.clone(),
shared.runtime_events.clone(), shared.runtime_events.clone(),
); );
+258 -38
View File
@@ -4,63 +4,283 @@ use std::time::{SystemTime, UNIX_EPOCH};
use tokio::sync::watch; use tokio::sync::watch;
use crate::config::ProxyConfig; use crate::maestro::generation::RuntimeWatchState;
use super::ApiRuntimeState; use super::ApiRuntimeState;
use super::events::ApiEventStore; use super::events::ApiEventStore;
pub(super) fn spawn_runtime_watchers( pub(super) fn spawn_runtime_watchers(
config_rx: watch::Receiver<Arc<ProxyConfig>>, runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
admission_rx: watch::Receiver<bool>,
runtime_state: Arc<ApiRuntimeState>, runtime_state: Arc<ApiRuntimeState>,
runtime_events: Arc<ApiEventStore>, runtime_events: Arc<ApiEventStore>,
) { ) {
let mut config_rx_reload = config_rx; spawn_config_watcher(
let runtime_state_reload = runtime_state.clone(); runtime_watch_rx.clone(),
let runtime_events_reload = runtime_events.clone(); runtime_state.clone(),
tokio::spawn(async move { runtime_events.clone(),
loop { );
if config_rx_reload.changed().await.is_err() { spawn_admission_watcher(runtime_watch_rx, runtime_state, runtime_events);
break; }
}
runtime_state_reload
.config_reload_count
.fetch_add(1, Ordering::Relaxed);
runtime_state_reload
.last_config_reload_epoch_secs
.store(now_epoch_secs(), Ordering::Relaxed);
runtime_events_reload.record("config.reload.applied", "config receiver updated");
}
});
let mut admission_rx_watch = admission_rx; fn spawn_config_watcher(
mut runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
runtime_state: Arc<ApiRuntimeState>,
runtime_events: Arc<ApiEventStore>,
) {
tokio::spawn(async move { tokio::spawn(async move {
runtime_state let Some(mut current) = runtime_watch_rx.borrow().clone() else {
.admission_open return;
.store(*admission_rx_watch.borrow(), Ordering::Relaxed); };
runtime_events.record(
"admission.state",
format!("accepting_new_connections={}", *admission_rx_watch.borrow()),
);
loop { loop {
if admission_rx_watch.changed().await.is_err() { tokio::select! {
break; biased;
changed = runtime_watch_rx.changed() => {
if changed.is_err() {
break;
}
let Some(next) = runtime_watch_rx.borrow().clone() else {
continue;
};
if next.generation_id != current.generation_id {
current = next;
record_config_reload(
&runtime_state,
&runtime_events,
format!("runtime generation {} activated", current.generation_id),
);
}
}
changed = current.config_rx.changed() => {
if changed.is_err() {
let Some(next) = wait_for_new_generation(
&mut runtime_watch_rx,
current.generation_id,
).await else {
break;
};
current = next;
record_config_reload(
&runtime_state,
&runtime_events,
format!("runtime generation {} activated", current.generation_id),
);
continue;
}
if active_generation_id(&runtime_watch_rx) != Some(current.generation_id) {
continue;
}
record_config_reload(
&runtime_state,
&runtime_events,
format!("generation {} config receiver updated", current.generation_id),
);
}
} }
let admission_open = *admission_rx_watch.borrow();
runtime_state
.admission_open
.store(admission_open, Ordering::Relaxed);
runtime_events.record(
"admission.state",
format!("accepting_new_connections={}", admission_open),
);
} }
}); });
} }
fn spawn_admission_watcher(
mut runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
runtime_state: Arc<ApiRuntimeState>,
runtime_events: Arc<ApiEventStore>,
) {
tokio::spawn(async move {
let Some(mut current) = runtime_watch_rx.borrow().clone() else {
return;
};
record_admission_state(&runtime_state, &runtime_events, &current);
loop {
tokio::select! {
biased;
changed = runtime_watch_rx.changed() => {
if changed.is_err() {
break;
}
let Some(next) = runtime_watch_rx.borrow().clone() else {
continue;
};
if next.generation_id != current.generation_id {
current = next;
record_admission_state(&runtime_state, &runtime_events, &current);
}
}
changed = current.admission_rx.changed() => {
if changed.is_err() {
let Some(next) = wait_for_new_generation(
&mut runtime_watch_rx,
current.generation_id,
).await else {
break;
};
current = next;
record_admission_state(&runtime_state, &runtime_events, &current);
continue;
}
if active_generation_id(&runtime_watch_rx) == Some(current.generation_id) {
record_admission_state(&runtime_state, &runtime_events, &current);
}
}
}
}
});
}
fn active_generation_id(
runtime_watch_rx: &watch::Receiver<Option<RuntimeWatchState>>,
) -> Option<u64> {
runtime_watch_rx
.borrow()
.as_ref()
.map(|state| state.generation_id)
}
async fn wait_for_new_generation(
runtime_watch_rx: &mut watch::Receiver<Option<RuntimeWatchState>>,
previous_generation_id: u64,
) -> Option<RuntimeWatchState> {
loop {
if let Some(state) = runtime_watch_rx.borrow().clone()
&& state.generation_id != previous_generation_id
{
return Some(state);
}
if runtime_watch_rx.changed().await.is_err() {
return None;
}
}
}
fn record_config_reload(
runtime_state: &ApiRuntimeState,
runtime_events: &ApiEventStore,
context: String,
) {
runtime_state
.config_reload_count
.fetch_add(1, Ordering::Relaxed);
runtime_state
.last_config_reload_epoch_secs
.store(now_epoch_secs(), Ordering::Relaxed);
runtime_events.record("config.reload.applied", context);
}
fn record_admission_state(
runtime_state: &ApiRuntimeState,
runtime_events: &ApiEventStore,
current: &RuntimeWatchState,
) {
let admission_open = *current.admission_rx.borrow();
runtime_state
.admission_open
.store(admission_open, Ordering::Relaxed);
runtime_events.record(
"admission.state",
format!(
"generation={} accepting_new_connections={}",
current.generation_id, admission_open
),
);
}
fn now_epoch_secs() -> u64 { fn now_epoch_secs() -> u64 {
SystemTime::now() SystemTime::now()
.duration_since(UNIX_EPOCH) .duration_since(UNIX_EPOCH)
.unwrap_or_default() .unwrap_or_default()
.as_secs() .as_secs()
} }
#[cfg(test)]
mod tests {
use super::*;
use crate::config::ProxyConfig;
use std::sync::atomic::{AtomicBool, AtomicU64};
use std::time::Duration;
fn state(
generation_id: u64,
) -> (
RuntimeWatchState,
watch::Sender<Arc<ProxyConfig>>,
watch::Sender<bool>,
) {
let (config_tx, config_rx) = watch::channel(Arc::new(ProxyConfig::default()));
let (admission_tx, admission_rx) = watch::channel(true);
(
RuntimeWatchState {
generation_id,
config_rx,
admission_rx,
},
config_tx,
admission_tx,
)
}
fn runtime_state() -> Arc<ApiRuntimeState> {
Arc::new(ApiRuntimeState {
process_started_at_epoch_secs: 1,
config_reload_count: AtomicU64::new(0),
last_config_reload_epoch_secs: AtomicU64::new(0),
admission_open: AtomicBool::new(false),
})
}
async fn wait_for_count(runtime_state: &ApiRuntimeState, expected: u64) {
tokio::time::timeout(Duration::from_secs(1), async {
loop {
if runtime_state.config_reload_count.load(Ordering::Relaxed) == expected {
break;
}
tokio::task::yield_now().await;
}
})
.await
.unwrap();
}
#[tokio::test]
async fn watchers_follow_only_the_active_generation() {
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));
spawn_runtime_watchers(runtime_watch_rx, runtime_state.clone(), events.clone());
tokio::task::yield_now().await;
assert_eq!(runtime_state.config_reload_count.load(Ordering::Relaxed), 0);
initial_config_tx.send_replace(Arc::new(ProxyConfig::default()));
wait_for_count(&runtime_state, 1).await;
let (next, next_config_tx, next_admission_tx) = state(2);
runtime_watch_tx.send_replace(Some(next));
wait_for_count(&runtime_state, 2).await;
initial_config_tx.send_replace(Arc::new(ProxyConfig::default()));
initial_admission_tx.send_replace(false);
tokio::task::yield_now().await;
assert_eq!(runtime_state.config_reload_count.load(Ordering::Relaxed), 2);
assert!(runtime_state.admission_open.load(Ordering::Relaxed));
next_config_tx.send_replace(Arc::new(ProxyConfig::default()));
next_admission_tx.send_replace(false);
wait_for_count(&runtime_state, 3).await;
tokio::time::timeout(Duration::from_secs(1), async {
while runtime_state.admission_open.load(Ordering::Relaxed) {
tokio::task::yield_now().await;
}
})
.await
.unwrap();
let snapshot = events.snapshot(16);
assert_eq!(
snapshot
.events
.iter()
.filter(|event| event.event_type == "config.reload.applied")
.count(),
3
);
}
}
+25 -1
View File
@@ -1,4 +1,5 @@
use std::sync::atomic::Ordering; use std::sync::atomic::Ordering;
use std::time::{SystemTime, UNIX_EPOCH};
use serde::Serialize; use serde::Serialize;
@@ -162,7 +163,7 @@ pub(super) fn build_system_info_data(
build_time_utc, build_time_utc,
rustc_version, rustc_version,
process_started_at_epoch_secs: shared.runtime_state.process_started_at_epoch_secs, process_started_at_epoch_secs: shared.runtime_state.process_started_at_epoch_secs,
uptime_seconds: shared.stats.uptime_secs(), uptime_seconds: process_uptime_seconds(shared.runtime_state.process_started_at_epoch_secs),
config_path: shared.config_path.display().to_string(), config_path: shared.config_path.display().to_string(),
config_hash: revision.to_string(), config_hash: revision.to_string(),
config_reload_count: shared config_reload_count: shared
@@ -173,6 +174,18 @@ pub(super) fn build_system_info_data(
} }
} }
fn process_uptime_seconds(process_started_at_epoch_secs: u64) -> f64 {
let now_epoch_secs = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
process_uptime_seconds_at(process_started_at_epoch_secs, now_epoch_secs)
}
fn process_uptime_seconds_at(process_started_at_epoch_secs: u64, now_epoch_secs: u64) -> f64 {
now_epoch_secs.saturating_sub(process_started_at_epoch_secs) as f64
}
pub(super) async fn build_runtime_gates_data( pub(super) async fn build_runtime_gates_data(
shared: &ApiShared, shared: &ApiShared,
cfg: &ProxyConfig, cfg: &ProxyConfig,
@@ -339,3 +352,14 @@ fn me_writer_pick_mode_label(mode: MeWriterPickMode) -> &'static str {
MeWriterPickMode::P2c => "p2c", MeWriterPickMode::P2c => "p2c",
} }
} }
#[cfg(test)]
mod tests {
use super::process_uptime_seconds_at;
#[test]
fn process_uptime_is_monotonic_and_saturating() {
assert_eq!(process_uptime_seconds_at(100, 135), 35.0);
assert_eq!(process_uptime_seconds_at(135, 100), 0.0);
}
}
+17
View File
@@ -22,6 +22,14 @@ use crate::transport::middle_proxy::MePool;
const SESSION_STOP_TIMEOUT: Duration = Duration::from_secs(5); const SESSION_STOP_TIMEOUT: Duration = Duration::from_secs(5);
const BACKGROUND_STOP_TIMEOUT: Duration = Duration::from_secs(5); const BACKGROUND_STOP_TIMEOUT: Duration = Duration::from_secs(5);
/// Process-visible control-plane receivers for one active runtime generation.
#[derive(Clone)]
pub(crate) struct RuntimeWatchState {
pub(crate) generation_id: u64,
pub(crate) config_rx: watch::Receiver<Arc<ProxyConfig>>,
pub(crate) admission_rx: watch::Receiver<bool>,
}
/// Cancellation and join ownership for one generation's background tasks. /// Cancellation and join ownership for one generation's background tasks.
#[derive(Clone)] #[derive(Clone)]
pub(crate) struct RuntimeTaskScope { pub(crate) struct RuntimeTaskScope {
@@ -134,6 +142,15 @@ impl RuntimeGeneration {
self.config_rx.borrow().clone() self.config_rx.borrow().clone()
} }
/// Returns receivers used by process-scoped observers of this generation.
pub(crate) fn watch_state(&self) -> RuntimeWatchState {
RuntimeWatchState {
generation_id: self.id,
config_rx: self.config_rx.clone(),
admission_rx: self.admission_rx.clone(),
}
}
pub(crate) async fn current_me_pool(&self) -> Option<Arc<MePool>> { pub(crate) async fn current_me_pool(&self) -> Option<Arc<MePool>> {
if let Some(pool) = &self.me_pool { if let Some(pool) = &self.me_pool {
return Some(pool.clone()); return Some(pool.clone());
+150 -200
View File
@@ -1,9 +1,11 @@
#![allow(clippy::too_many_arguments)] #![allow(clippy::too_many_arguments)]
use std::future::Future;
use std::sync::Arc; use std::sync::Arc;
use std::time::Duration; use std::time::Duration;
use tokio::sync::{RwLock, watch}; use tokio::sync::{RwLock, watch};
use tokio_util::task::AbortOnDropHandle;
use tracing::{error, info, warn}; use tracing::{error, info, warn};
use crate::config::ProxyConfig; use crate::config::ProxyConfig;
@@ -17,8 +19,58 @@ use crate::stats::Stats;
use crate::transport::UpstreamManager; use crate::transport::UpstreamManager;
use crate::transport::middle_proxy::MePool; use crate::transport::middle_proxy::MePool;
use super::generation::RuntimeTaskScope;
use super::helpers::load_startup_proxy_config_snapshot; use super::helpers::load_startup_proxy_config_snapshot;
async fn supervise_me_task<F, Fut>(task_name: &'static str, mut task: F)
where
F: FnMut() -> Fut,
Fut: Future<Output = ()> + Send + 'static,
{
loop {
let result = AbortOnDropHandle::new(tokio::spawn(task())).await;
match result {
Ok(()) => warn!(task = task_name, "Middle-End supervisor task exited unexpectedly, restarting"),
Err(error) => {
error!(task = task_name, error = %error, "Middle-End supervisor task panicked, restarting in 1s");
tokio::time::sleep(Duration::from_secs(1)).await;
}
}
}
}
fn spawn_me_supervisors(
task_scope: RuntimeTaskScope,
pool: Arc<MePool>,
rng: Arc<SecureRandom>,
min_connections: usize,
) {
let health_pool = pool.clone();
let health_rng = rng;
task_scope.spawn(supervise_me_task("health_monitor", move || {
let pool = health_pool.clone();
let rng = health_rng.clone();
async move {
crate::transport::middle_proxy::me_health_monitor(pool, rng, min_connections).await;
}
}));
let drain_pool = pool.clone();
task_scope.spawn(supervise_me_task("drain_timeout_enforcer", move || {
let pool = drain_pool.clone();
async move {
crate::transport::middle_proxy::me_drain_timeout_enforcer(pool).await;
}
}));
task_scope.spawn(supervise_me_task("zombie_writer_watchdog", move || {
let pool = pool.clone();
async move {
crate::transport::middle_proxy::me_zombie_writer_watchdog(pool).await;
}
}));
}
pub(crate) async fn initialize_me_pool( pub(crate) async fn initialize_me_pool(
use_middle_proxy: bool, use_middle_proxy: bool,
config: &ProxyConfig, config: &ProxyConfig,
@@ -30,6 +82,7 @@ pub(crate) async fn initialize_me_pool(
stats: Arc<Stats>, stats: Arc<Stats>,
api_me_pool: Arc<RwLock<Option<Arc<MePool>>>>, api_me_pool: Arc<RwLock<Option<Arc<MePool>>>>,
me_ready_tx: watch::Sender<u64>, me_ready_tx: watch::Sender<u64>,
task_scope: RuntimeTaskScope,
) -> Option<Arc<MePool>> { ) -> Option<Arc<MePool>> {
if !use_middle_proxy { if !use_middle_proxy {
return None; return None;
@@ -52,15 +105,8 @@ pub(crate) async fn initialize_me_pool(
.as_ref() .as_ref()
.map(|tag| hex::decode(tag).expect("general.ad_tag must be validated before startup")); .map(|tag| hex::decode(tag).expect("general.ad_tag must be validated before startup"));
// ============================================================= // The Telegram proxy-secret authenticates ME RPC and is distinct from client secrets.
// CRITICAL: Download Telegram proxy-secret (NOT user secret!) // It corresponds to the C MTProxy --aes-pwd input and may be fetched from Telegram.
//
// C MTProxy uses TWO separate secrets:
// -S flag = 16-byte user secret for client obfuscation
// --aes-pwd = 32-512 byte binary file for ME RPC auth
//
// proxy-secret is from: https://core.telegram.org/getProxySecret
// =============================================================
let proxy_secret_path = config.general.proxy_secret_path.as_deref(); let proxy_secret_path = config.general.proxy_secret_path.as_deref();
let pool_size = config.general.middle_proxy_pool_size.max(1); let pool_size = config.general.middle_proxy_pool_size.max(1);
let proxy_secret = loop { let proxy_secret = loop {
@@ -319,143 +365,70 @@ pub(crate) async fn initialize_me_pool(
let rng_bg = rng.clone(); let rng_bg = rng.clone();
let startup_tracker_bg = startup_tracker.clone(); let startup_tracker_bg = startup_tracker.clone();
let me_ready_tx_bg = me_ready_tx.clone(); let me_ready_tx_bg = me_ready_tx.clone();
let task_scope_bg = task_scope.clone();
let retry_limit = if me_init_retry_attempts == 0 { let retry_limit = if me_init_retry_attempts == 0 {
String::from("unlimited") String::from("unlimited")
} else { } else {
me_init_retry_attempts.to_string() me_init_retry_attempts.to_string()
}; };
std::thread::spawn(move || { task_scope.spawn(async move {
let runtime = match tokio::runtime::Builder::new_current_thread() let mut init_attempt: u32 = 0;
.enable_all() loop {
.build() init_attempt = init_attempt.saturating_add(1);
{ startup_tracker_bg.set_me_init_attempt(init_attempt).await;
Ok(runtime) => runtime, match pool_bg.init(pool_size, &rng_bg).await {
Err(error) => { Ok(()) => {
error!(error = %error, "Failed to build background runtime for ME initialization"); startup_tracker_bg.set_me_last_error(None).await;
return; startup_tracker_bg
} .complete_component(
}; COMPONENT_ME_POOL_INIT_STAGE1,
runtime.block_on(async move { Some("ME pool initialized".to_string()),
let mut init_attempt: u32 = 0; )
loop { .await;
init_attempt = init_attempt.saturating_add(1); startup_tracker_bg
startup_tracker_bg.set_me_init_attempt(init_attempt).await; .set_me_status(StartupMeStatus::Ready, "ready")
match pool_bg.init(pool_size, &rng_bg).await { .await;
Ok(()) => { me_ready_tx_bg.send_modify(|version| {
startup_tracker_bg.set_me_last_error(None).await; *version = version.saturating_add(1);
startup_tracker_bg });
.complete_component( info!(
COMPONENT_ME_POOL_INIT_STAGE1, attempt = init_attempt,
Some("ME pool initialized".to_string()), "Middle-End pool initialized successfully"
) );
.await; spawn_me_supervisors(
startup_tracker_bg task_scope_bg,
.set_me_status(StartupMeStatus::Ready, "ready") pool_bg.clone(),
.await; rng_bg.clone(),
me_ready_tx_bg.send_modify(|version| { pool_size,
*version = version.saturating_add(1); );
}); break;
info!( }
Err(e) => {
startup_tracker_bg
.set_me_last_error(Some(e.to_string()))
.await;
if init_attempt >= me_init_warn_after_attempts {
warn!(
error = %e,
attempt = init_attempt, attempt = init_attempt,
"Middle-End pool initialized successfully" retry_limit = %retry_limit,
retry_in_secs = 2,
"ME pool is not ready yet; retrying background initialization"
);
} else {
info!(
error = %e,
attempt = init_attempt,
retry_limit = %retry_limit,
retry_in_secs = 2,
"ME pool startup warmup: retrying background initialization"
); );
// ── Supervised background tasks ──────────────────
// Each task runs inside a nested tokio::spawn so
// that a panic is caught via JoinHandle and the
// outer loop restarts the task automatically.
let pool_health = pool_bg.clone();
let rng_health = rng_bg.clone();
let min_conns = pool_size;
tokio::spawn(async move {
loop {
let p = pool_health.clone();
let r = rng_health.clone();
let res = tokio::spawn(async move {
crate::transport::middle_proxy::me_health_monitor(
p, r, min_conns,
)
.await;
})
.await;
match res {
Ok(()) => warn!("me_health_monitor exited unexpectedly, restarting"),
Err(e) => {
error!(error = %e, "me_health_monitor panicked, restarting in 1s");
tokio::time::sleep(Duration::from_secs(1)).await;
}
}
}
});
let pool_drain_enforcer = pool_bg.clone();
tokio::spawn(async move {
loop {
let p = pool_drain_enforcer.clone();
let res = tokio::spawn(async move {
crate::transport::middle_proxy::me_drain_timeout_enforcer(p).await;
})
.await;
match res {
Ok(()) => warn!("me_drain_timeout_enforcer exited unexpectedly, restarting"),
Err(e) => {
error!(error = %e, "me_drain_timeout_enforcer panicked, restarting in 1s");
tokio::time::sleep(Duration::from_secs(1)).await;
}
}
}
});
let pool_watchdog = pool_bg.clone();
tokio::spawn(async move {
loop {
let p = pool_watchdog.clone();
let res = tokio::spawn(async move {
crate::transport::middle_proxy::me_zombie_writer_watchdog(p).await;
})
.await;
match res {
Ok(()) => warn!("me_zombie_writer_watchdog exited unexpectedly, restarting"),
Err(e) => {
error!(error = %e, "me_zombie_writer_watchdog panicked, restarting in 1s");
tokio::time::sleep(Duration::from_secs(1)).await;
}
}
}
});
// CRITICAL: keep the current-thread runtime
// alive. Without this, block_on() returns,
// the Runtime is dropped, and ALL spawned
// background tasks (health monitor, drain
// enforcer, zombie watchdog) are silently
// cancelled — causing the draining-writer
// leak that brought us here.
std::future::pending::<()>().await;
unreachable!();
}
Err(e) => {
startup_tracker_bg.set_me_last_error(Some(e.to_string())).await;
if init_attempt >= me_init_warn_after_attempts {
warn!(
error = %e,
attempt = init_attempt,
retry_limit = %retry_limit,
retry_in_secs = 2,
"ME pool is not ready yet; retrying background initialization"
);
} else {
info!(
error = %e,
attempt = init_attempt,
retry_limit = %retry_limit,
retry_in_secs = 2,
"ME pool startup warmup: retrying background initialization"
);
}
pool_bg.reset_stun_state();
tokio::time::sleep(Duration::from_secs(2)).await;
} }
pool_bg.reset_stun_state();
tokio::time::sleep(Duration::from_secs(2)).await;
} }
} }
}); }
}); });
startup_tracker startup_tracker
.set_me_status(StartupMeStatus::Initializing, "background_init") .set_me_status(StartupMeStatus::Initializing, "background_init")
@@ -490,70 +463,12 @@ pub(crate) async fn initialize_me_pool(
"Middle-End pool initialized successfully" "Middle-End pool initialized successfully"
); );
// ── Supervised background tasks ────────────────── spawn_me_supervisors(
let pool_clone = pool.clone(); task_scope.clone(),
let rng_clone = rng.clone(); pool.clone(),
let min_conns = pool_size; rng.clone(),
tokio::spawn(async move { pool_size,
loop { );
let p = pool_clone.clone();
let r = rng_clone.clone();
let res = tokio::spawn(async move {
crate::transport::middle_proxy::me_health_monitor(
p, r, min_conns,
)
.await;
})
.await;
match res {
Ok(()) => warn!(
"me_health_monitor exited unexpectedly, restarting"
),
Err(e) => {
error!(error = %e, "me_health_monitor panicked, restarting in 1s");
tokio::time::sleep(Duration::from_secs(1)).await;
}
}
}
});
let pool_drain_enforcer = pool.clone();
tokio::spawn(async move {
loop {
let p = pool_drain_enforcer.clone();
let res = tokio::spawn(async move {
crate::transport::middle_proxy::me_drain_timeout_enforcer(p).await;
})
.await;
match res {
Ok(()) => warn!(
"me_drain_timeout_enforcer exited unexpectedly, restarting"
),
Err(e) => {
error!(error = %e, "me_drain_timeout_enforcer panicked, restarting in 1s");
tokio::time::sleep(Duration::from_secs(1)).await;
}
}
}
});
let pool_watchdog = pool.clone();
tokio::spawn(async move {
loop {
let p = pool_watchdog.clone();
let res = tokio::spawn(async move {
crate::transport::middle_proxy::me_zombie_writer_watchdog(p).await;
})
.await;
match res {
Ok(()) => warn!(
"me_zombie_writer_watchdog exited unexpectedly, restarting"
),
Err(e) => {
error!(error = %e, "me_zombie_writer_watchdog panicked, restarting in 1s");
tokio::time::sleep(Duration::from_secs(1)).await;
}
}
}
});
break Some(pool); break Some(pool);
} }
@@ -666,3 +581,38 @@ pub(crate) async fn initialize_me_pool(
} }
} }
} }
#[cfg(test)]
mod tests {
use super::*;
use tokio::sync::Notify;
struct DropSignal(Arc<Notify>);
impl Drop for DropSignal {
fn drop(&mut self) {
self.0.notify_one();
}
}
#[tokio::test]
async fn scoped_supervisor_aborts_its_current_child() {
let scope = RuntimeTaskScope::new();
let dropped = Arc::new(Notify::new());
let dropped_for_task = dropped.clone();
scope.spawn(supervise_me_task("test", move || {
let dropped = dropped_for_task.clone();
async move {
let _signal = DropSignal(dropped);
std::future::pending::<()>().await;
}
}));
tokio::task::yield_now().await;
scope.stop().await;
tokio::time::timeout(Duration::from_secs(1), dropped.notified())
.await
.unwrap();
}
}
+18 -10
View File
@@ -503,7 +503,6 @@ async fn run_telemt_core(
config.access.cidr_rate_limits.clone(), config.access.cidr_rate_limits.clone(),
); );
let (api_config_tx, api_config_rx) = watch::channel(Arc::new(config.clone()));
let (detected_ips_tx, detected_ips_rx) = watch::channel((None::<IpAddr>, None::<IpAddr>)); let (detected_ips_tx, detected_ips_rx) = watch::channel((None::<IpAddr>, None::<IpAddr>));
let initial_direct_first = config.general.use_middle_proxy && config.general.me2dc_fallback; let initial_direct_first = config.general.use_middle_proxy && config.general.me2dc_fallback;
let initial_admission_open = !config.general.use_middle_proxy || initial_direct_first; let initial_admission_open = !config.general.use_middle_proxy || initial_direct_first;
@@ -511,6 +510,8 @@ async fn run_telemt_core(
let (reload_control, reload_commands) = reload::ReloadControl::channel(1); let (reload_control, reload_commands) = reload::ReloadControl::channel(1);
let (active_runtime_tx, active_runtime_rx) = let (active_runtime_tx, active_runtime_rx) =
watch::channel(None::<Arc<ArcSwap<generation::RuntimeGeneration>>>); watch::channel(None::<Arc<ArcSwap<generation::RuntimeGeneration>>>);
let (runtime_watch_tx, runtime_watch_rx) =
watch::channel(None::<generation::RuntimeWatchState>);
let initial_route_mode = if !config.general.use_middle_proxy || initial_direct_first { let initial_route_mode = if !config.general.use_middle_proxy || initial_direct_first {
RelayRouteMode::Direct RelayRouteMode::Direct
} else { } else {
@@ -544,14 +545,13 @@ async fn run_telemt_core(
let upstream_manager_api = upstream_manager.clone(); let upstream_manager_api = upstream_manager.clone();
let route_runtime_api = route_runtime.clone(); let route_runtime_api = route_runtime.clone();
let proxy_shared_api = shared_state.clone(); let proxy_shared_api = shared_state.clone();
let config_rx_api = api_config_rx.clone();
let admission_rx_api = admission_rx.clone();
let config_path_api = config_path.clone(); let config_path_api = config_path.clone();
let quota_state_path_api = quota_state_path.clone(); let quota_state_path_api = quota_state_path.clone();
let startup_tracker_api = startup_tracker.clone(); let startup_tracker_api = startup_tracker.clone();
let detected_ips_rx_api = detected_ips_rx.clone(); let detected_ips_rx_api = detected_ips_rx.clone();
let reload_control_api = reload_control.clone(); let reload_control_api = reload_control.clone();
let active_runtime_rx_api = active_runtime_rx.clone(); let active_runtime_rx_api = active_runtime_rx.clone();
let runtime_watch_rx_api = runtime_watch_rx.clone();
tokio::spawn(async move { tokio::spawn(async move {
api::serve( api::serve(
listen, listen,
@@ -561,8 +561,6 @@ async fn run_telemt_core(
route_runtime_api, route_runtime_api,
proxy_shared_api, proxy_shared_api,
upstream_manager_api, upstream_manager_api,
config_rx_api,
admission_rx_api,
config_path_api, config_path_api,
quota_state_path_api, quota_state_path_api,
detected_ips_rx_api, detected_ips_rx_api,
@@ -570,6 +568,7 @@ async fn run_telemt_core(
startup_tracker_api, startup_tracker_api,
reload_control_api, reload_control_api,
active_runtime_rx_api, active_runtime_rx_api,
runtime_watch_rx_api,
) )
.await; .await;
}); });
@@ -610,8 +609,9 @@ async fn run_telemt_core(
upstream_manager.clone(), upstream_manager.clone(),
&startup_tracker, &startup_tracker,
runtime_task_scope.clone(), runtime_task_scope.clone(),
tls_bootstrap::TlsBootstrapPolicy::BestEffort,
) )
.await; .await?;
startup_tracker startup_tracker
.start_component( .start_component(
@@ -737,6 +737,7 @@ async fn run_telemt_core(
stats.clone(), stats.clone(),
api_me_pool.clone(), api_me_pool.clone(),
me_ready_tx.clone(), me_ready_tx.clone(),
runtime_task_scope.clone(),
) )
.await .await
}; };
@@ -824,7 +825,6 @@ async fn run_telemt_core(
rng.clone(), rng.clone(),
ip_tracker.clone(), ip_tracker.clone(),
beobachten.clone(), beobachten.clone(),
api_config_tx.clone(),
me_pool.clone(), me_pool.clone(),
shared_state.clone(), shared_state.clone(),
me_ready_tx.clone(), me_ready_tx.clone(),
@@ -869,6 +869,7 @@ async fn run_telemt_core(
stats_bg.clone(), stats_bg.clone(),
api_me_pool_bg.clone(), api_me_pool_bg.clone(),
me_ready_tx_bg.clone(), me_ready_tx_bg.clone(),
task_scope_bg.clone(),
) )
.await; .await;
if let Some(pool) = pool { if let Some(pool) = pool {
@@ -961,6 +962,7 @@ async fn run_telemt_core(
quota_store, quota_store,
detected_ips_tx, detected_ips_tx,
runtime_log_filter, runtime_log_filter,
runtime_watch_tx.clone(),
); );
let bound = listeners::bind_listeners( let bound = listeners::bind_listeners(
@@ -988,12 +990,12 @@ async fn run_telemt_core(
// On Unix, caller supplies privilege drop after bind (may require root for port < 1024). // On Unix, caller supplies privilege drop after bind (may require root for port < 1024).
drop_after_bind(); drop_after_bind();
synlimit_control::reconcile_synlimit_rules(&config).await; let synlimit_controller = synlimit_control::spawn_synlimit_controller(runtime_watch_rx);
synlimit_control::spawn_synlimit_controller(config_rx.clone());
runtime_tasks::spawn_metrics_if_configured(&config, &startup_tracker, active_runtime.clone()) runtime_tasks::spawn_metrics_if_configured(&config, &startup_tracker, active_runtime.clone())
.await; .await;
runtime_watch_tx.send_replace(Some(active_runtime.load_full().watch_state()));
active_runtime_tx.send_replace(Some(active_runtime.clone())); active_runtime_tx.send_replace(Some(active_runtime.clone()));
runtime_tasks::mark_runtime_ready(&startup_tracker).await; runtime_tasks::mark_runtime_ready(&startup_tracker).await;
@@ -1004,7 +1006,13 @@ async fn run_telemt_core(
#[cfg(unix)] #[cfg(unix)]
listeners::spawn_unix_accept_loop(unix_listener, active_runtime.clone()); listeners::spawn_unix_accept_loop(unix_listener, active_runtime.clone());
shutdown::wait_for_shutdown(process_started_at, active_runtime, quota_state_path).await; shutdown::wait_for_shutdown(
process_started_at,
active_runtime,
quota_state_path,
synlimit_controller,
)
.await;
Ok(()) Ok(())
} }
+30 -1
View File
@@ -138,7 +138,11 @@ pub(crate) struct ReloadStatus {
pub(crate) started_at_epoch_secs: Option<u64>, pub(crate) started_at_epoch_secs: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")] #[serde(skip_serializing_if = "Option::is_none")]
pub(crate) finished_at_epoch_secs: Option<u64>, pub(crate) finished_at_epoch_secs: Option<u64>,
#[serde(default, skip_serializing_if = "Vec::is_empty")] #[serde(
rename = "deferred_process_fields",
default,
skip_serializing_if = "Vec::is_empty"
)]
pub(crate) deferred_fields: Vec<String>, pub(crate) deferred_fields: Vec<String>,
#[serde(default, skip_serializing_if = "Vec::is_empty")] #[serde(default, skip_serializing_if = "Vec::is_empty")]
pub(crate) warnings: Vec<String>, pub(crate) warnings: Vec<String>,
@@ -428,6 +432,31 @@ mod tests {
assert!(ReloadRequest::from_query(Some("timeout_secs=30")).is_err()); 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] #[tokio::test]
async fn coordinator_rejects_concurrent_reload_and_releases_terminal_slot() { async fn coordinator_rejects_concurrent_reload_and_releases_terminal_slot() {
let (control, mut receiver) = ReloadControl::channel(1); let (control, mut receiver) = ReloadControl::channel(1);
+102 -43
View File
@@ -8,7 +8,7 @@ use tracing::{info, warn};
use crate::stats::QuotaStore; use crate::stats::QuotaStore;
use super::generation::RuntimeGeneration; use super::generation::{RuntimeGeneration, RuntimeWatchState};
use super::reload::{ use super::reload::{
ReloadCommand, ReloadCommandReceiver, ReloadControl, ReloadFailurePolicy, ReloadMode, ReloadCommand, ReloadCommandReceiver, ReloadControl, ReloadFailurePolicy, ReloadMode,
ReloadPhase, ReloadPhase,
@@ -24,6 +24,48 @@ pub(crate) struct ReloadSupervisor {
quota_store: Arc<QuotaStore>, quota_store: Arc<QuotaStore>,
detected_ips_tx: watch::Sender<(Option<std::net::IpAddr>, Option<std::net::IpAddr>)>, detected_ips_tx: watch::Sender<(Option<std::net::IpAddr>, Option<std::net::IpAddr>)>,
runtime_log_filter: RuntimeLogFilter, runtime_log_filter: RuntimeLogFilter,
runtime_watch_tx: watch::Sender<Option<RuntimeWatchState>>,
}
#[derive(Debug, PartialEq, Eq)]
enum RevisionGateAction {
Proceed,
Warn(String),
Rollback(String),
}
fn revision_gate_action(
accepted_revision: &str,
current_revision: Result<String, String>,
failure_policy: ReloadFailurePolicy,
) -> RevisionGateAction {
let warning = match current_revision {
Ok(current) if current == accepted_revision => return RevisionGateAction::Proceed,
Ok(current) => format!(
"config revision changed during preparation: accepted={} current={}",
accepted_revision, current
),
Err(error) => format!("config revision verification failed: {}", error),
};
match failure_policy {
ReloadFailurePolicy::KeepNew => RevisionGateAction::Warn(warning),
ReloadFailurePolicy::Rollback => RevisionGateAction::Rollback(warning),
}
}
async fn stop_background_and_middle_end(generation: &RuntimeGeneration) -> bool {
generation.stop_background_tasks().await;
let Some(pool) = generation.current_me_pool().await else {
return false;
};
tokio::time::timeout(Duration::from_secs(2), pool.shutdown_send_close_conn_all())
.await
.is_err()
}
async fn cleanup_candidate(generation: &RuntimeGeneration) -> bool {
generation.stop_sessions().await;
stop_background_and_middle_end(generation).await
} }
impl ReloadSupervisor { impl ReloadSupervisor {
@@ -36,6 +78,7 @@ impl ReloadSupervisor {
quota_store: Arc<QuotaStore>, quota_store: Arc<QuotaStore>,
detected_ips_tx: watch::Sender<(Option<std::net::IpAddr>, Option<std::net::IpAddr>)>, detected_ips_tx: watch::Sender<(Option<std::net::IpAddr>, Option<std::net::IpAddr>)>,
runtime_log_filter: RuntimeLogFilter, runtime_log_filter: RuntimeLogFilter,
runtime_watch_tx: watch::Sender<Option<RuntimeWatchState>>,
) { ) {
let supervisor = Self { let supervisor = Self {
active_runtime, active_runtime,
@@ -45,6 +88,7 @@ impl ReloadSupervisor {
quota_store, quota_store,
detected_ips_tx, detected_ips_tx,
runtime_log_filter, runtime_log_filter,
runtime_watch_tx,
}; };
tokio::spawn(supervisor.run()); tokio::spawn(supervisor.run());
} }
@@ -81,30 +125,23 @@ impl ReloadSupervisor {
} }
}; };
if let Ok(current_revision) = let revision_action = revision_gate_action(
crate::api::config_store::current_revision_for_maestro(&self.config_path).await &command.config_revision,
&& current_revision != command.config_revision crate::api::config_store::current_revision_for_maestro(&self.config_path).await,
{ command.request.failure_policy,
let warning = format!( );
"config revision changed during preparation: accepted={} current={}", match revision_action {
command.config_revision, current_revision RevisionGateAction::Proceed => {}
); RevisionGateAction::Warn(warning) => {
if command.request.failure_policy == ReloadFailurePolicy::Rollback { self.control.add_warning(command.reload_id, warning).await;
}
RevisionGateAction::Rollback(warning) => {
let _ = cleanup_candidate(&prepared.generation).await;
self.runtime_log_filter self.runtime_log_filter
.apply_reload(&old_runtime.config().general.log_level); .apply_reload(&old_runtime.config().general.log_level);
prepared.generation.stop_sessions().await;
if let Some(pool) = prepared.generation.current_me_pool().await {
let _ = tokio::time::timeout(
Duration::from_secs(2),
pool.shutdown_send_close_conn_all(),
)
.await;
}
prepared.generation.stop_background_tasks().await;
self.control.rolled_back(command.reload_id, warning).await; self.control.rolled_back(command.reload_id, warning).await;
return; return;
} }
self.control.add_warning(command.reload_id, warning).await;
} }
self.control self.control
@@ -112,37 +149,26 @@ impl ReloadSupervisor {
.await; .await;
let new_runtime = prepared.generation; let new_runtime = prepared.generation;
old_runtime.stop_accepting_sessions(); old_runtime.stop_accepting_sessions();
let replaced = self.active_runtime.swap(new_runtime.clone());
if let Err(error) = crate::network::dns_overrides::install_entries( if let Err(error) = crate::network::dns_overrides::install_entries(
&new_runtime.config().network.dns_overrides, &new_runtime.config().network.dns_overrides,
) { ) {
let message = format!("runtime DNS activation failed: {}", error); let message = format!("runtime DNS activation failed: {}", error);
if command.request.failure_policy == ReloadFailurePolicy::Rollback { if command.request.failure_policy == ReloadFailurePolicy::Rollback {
let candidate = self.active_runtime.swap(replaced.clone()); old_runtime.resume_accepting_sessions();
replaced.resume_accepting_sessions(); let _ = cleanup_candidate(&new_runtime).await;
self.runtime_log_filter self.runtime_log_filter
.apply_reload(&replaced.config().general.log_level); .apply_reload(&old_runtime.config().general.log_level);
let _ = crate::network::dns_overrides::install_entries(
&replaced.config().network.dns_overrides,
);
candidate.stop_sessions().await;
if let Some(pool) = candidate.current_me_pool().await {
let _ = tokio::time::timeout(
Duration::from_secs(2),
pool.shutdown_send_close_conn_all(),
)
.await;
}
candidate.stop_background_tasks().await;
self.control.rolled_back(command.reload_id, message).await; self.control.rolled_back(command.reload_id, message).await;
return; return;
} }
self.control.add_warning(command.reload_id, message).await; self.control.add_warning(command.reload_id, message).await;
} }
let replaced = self.active_runtime.swap(new_runtime.clone());
self.detected_ips_tx.send_replace(prepared.detected_ips); self.detected_ips_tx.send_replace(prepared.detected_ips);
self.runtime_log_filter self.runtime_log_filter
.apply_reload(&new_runtime.config().general.log_level); .apply_reload(&new_runtime.config().general.log_level);
crate::synlimit_control::reconcile_synlimit_rules(&new_runtime.config()).await; self.runtime_watch_tx
.send_replace(Some(new_runtime.watch_state()));
info!( info!(
reload_id = command.reload_id, reload_id = command.reload_id,
@@ -177,11 +203,7 @@ impl ReloadSupervisor {
} }
} }
if let Some(pool) = replaced.current_me_pool().await if stop_background_and_middle_end(&replaced).await {
&& tokio::time::timeout(Duration::from_secs(2), pool.shutdown_send_close_conn_all())
.await
.is_err()
{
let warning = format!( let warning = format!(
"generation {} Middle-End close broadcast timed out", "generation {} Middle-End close broadcast timed out",
replaced.id replaced.id
@@ -189,9 +211,46 @@ impl ReloadSupervisor {
warn!(reload_id = command.reload_id, warning = %warning); warn!(reload_id = command.reload_id, warning = %warning);
self.control.add_warning(command.reload_id, warning).await; self.control.add_warning(command.reload_id, warning).await;
} }
replaced.stop_background_tasks().await;
self.control self.control
.succeed(command.reload_id, new_runtime.id) .succeed(command.reload_id, new_runtime.id)
.await; .await;
} }
} }
#[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(_)
));
}
}
}
+32 -6
View File
@@ -34,7 +34,7 @@ pub(crate) struct PreparedRuntime {
pub(crate) async fn prepare_runtime( pub(crate) async fn prepare_runtime(
generation_id: u64, generation_id: u64,
mut config: ProxyConfig, config: ProxyConfig,
config_path: &Path, config_path: &Path,
quota_store: Arc<QuotaStore>, quota_store: Arc<QuotaStore>,
runtime_log_filter: RuntimeLogFilter, runtime_log_filter: RuntimeLogFilter,
@@ -112,8 +112,10 @@ pub(crate) async fn prepare_runtime(
upstream_manager.clone(), upstream_manager.clone(),
&startup_tracker, &startup_tracker,
task_scope.clone(), task_scope.clone(),
tls_bootstrap::TlsBootstrapPolicy::RequireReady,
) )
.await; .await
.map_err(|error| error.to_string())?;
let beobachten = Arc::new(BeobachtenStore::new()); let beobachten = Arc::new(BeobachtenStore::new());
let rng = Arc::new(SecureRandom::new()); let rng = Arc::new(SecureRandom::new());
@@ -140,11 +142,20 @@ pub(crate) async fn prepare_runtime(
stats.clone(), stats.clone(),
me_pool_runtime.clone(), me_pool_runtime.clone(),
me_ready_tx.clone(), me_ready_tx.clone(),
task_scope.clone(),
) )
.await .await
}; };
if config.general.use_middle_proxy && me_pool.is_none() && !direct_first_startup { if strict_middle_proxy_unavailable(
config.general.use_middle_proxy = false; config.general.use_middle_proxy,
direct_first_startup,
me_pool.is_some(),
) {
task_scope.stop().await;
return Err(
"Middle-End pool is required but did not become ready during reload preparation"
.to_string(),
);
} }
let config = Arc::new(config); let config = Arc::new(config);
@@ -159,7 +170,6 @@ pub(crate) async fn prepare_runtime(
config.server.max_connections as usize config.server.max_connections as usize
}; };
let max_connections = Arc::new(Semaphore::new(max_connections_limit)); let max_connections = Arc::new(Semaphore::new(max_connections_limit));
let (api_config_tx, _) = watch::channel(config.clone());
let watches = runtime_tasks::spawn_runtime_tasks( let watches = runtime_tasks::spawn_runtime_tasks(
&config, &config,
config_path, config_path,
@@ -175,7 +185,6 @@ pub(crate) async fn prepare_runtime(
rng.clone(), rng.clone(),
ip_tracker.clone(), ip_tracker.clone(),
beobachten.clone(), beobachten.clone(),
api_config_tx,
me_pool.clone(), me_pool.clone(),
proxy_shared.clone(), proxy_shared.clone(),
me_ready_tx.clone(), me_ready_tx.clone(),
@@ -226,6 +235,7 @@ pub(crate) async fn prepare_runtime(
stats_bg.clone(), stats_bg.clone(),
me_pool_runtime_bg.clone(), me_pool_runtime_bg.clone(),
me_ready_tx_bg.clone(), me_ready_tx_bg.clone(),
task_scope_bg.clone(),
) )
.await; .await;
if let Some(pool) = pool { if let Some(pool) = pool {
@@ -291,6 +301,14 @@ pub(crate) async fn prepare_runtime(
}) })
} }
fn strict_middle_proxy_unavailable(
use_middle_proxy: bool,
direct_first_startup: bool,
pool_available: bool,
) -> bool {
use_middle_proxy && !direct_first_startup && !pool_available
}
pub(crate) fn deferred_process_fields(old: &ProxyConfig, new: &ProxyConfig) -> Vec<String> { pub(crate) fn deferred_process_fields(old: &ProxyConfig, new: &ProxyConfig) -> Vec<String> {
let mut fields = Vec::new(); let mut fields = Vec::new();
if old.server.port != new.server.port if old.server.port != new.server.port
@@ -354,4 +372,12 @@ mod tests {
new.censorship.tls_domain = "reload.example".to_string(); new.censorship.tls_domain = "reload.example".to_string();
assert!(deferred_process_fields(&old, &new).is_empty()); assert!(deferred_process_fields(&old, &new).is_empty());
} }
#[test]
fn strict_middle_proxy_requires_a_prepared_pool() {
assert!(strict_middle_proxy_unavailable(true, false, false));
assert!(!strict_middle_proxy_unavailable(true, false, true));
assert!(!strict_middle_proxy_unavailable(true, true, false));
assert!(!strict_middle_proxy_unavailable(false, false, false));
}
} }
-13
View File
@@ -103,7 +103,6 @@ pub(crate) async fn spawn_runtime_tasks(
rng: Arc<SecureRandom>, rng: Arc<SecureRandom>,
ip_tracker: Arc<UserIpTracker>, ip_tracker: Arc<UserIpTracker>,
beobachten: Arc<BeobachtenStore>, beobachten: Arc<BeobachtenStore>,
api_config_tx: watch::Sender<Arc<ProxyConfig>>,
me_pool_for_policy: Option<Arc<MePool>>, me_pool_for_policy: Option<Arc<MePool>>,
shared_state: Arc<ProxySharedState>, shared_state: Arc<ProxySharedState>,
me_ready_tx: watch::Sender<u64>, me_ready_tx: watch::Sender<u64>,
@@ -166,18 +165,6 @@ pub(crate) async fn spawn_runtime_tasks(
Some("config hot-reload watcher started".to_string()), Some("config hot-reload watcher started".to_string()),
) )
.await; .await;
let mut config_rx_api_bridge = config_rx.clone();
let api_config_tx_bridge = api_config_tx.clone();
task_scope.spawn(async move {
loop {
if config_rx_api_bridge.changed().await.is_err() {
break;
}
let cfg = config_rx_api_bridge.borrow_and_update().clone();
api_config_tx_bridge.send_replace(cfg);
}
});
let stats_policy = stats.clone(); let stats_policy = stats.clone();
let upstream_policy = upstream_manager.clone(); let upstream_policy = upstream_manager.clone();
let mut config_rx_policy = config_rx.clone(); let mut config_rx_policy = config_rx.clone();
+17 -6
View File
@@ -50,9 +50,17 @@ pub(crate) async fn wait_for_shutdown(
process_started_at: Instant, process_started_at: Instant,
active_runtime: Arc<ArcSwap<RuntimeGeneration>>, active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
quota_state_path: PathBuf, quota_state_path: PathBuf,
synlimit_controller: tokio::task::JoinHandle<()>,
) { ) {
let signal = wait_for_shutdown_signal().await; let signal = wait_for_shutdown_signal().await;
perform_shutdown(signal, process_started_at, active_runtime, quota_state_path).await; perform_shutdown(
signal,
process_started_at,
active_runtime,
quota_state_path,
synlimit_controller,
)
.await;
} }
/// Waits for any shutdown signal (SIGINT, SIGTERM, SIGQUIT). /// Waits for any shutdown signal (SIGINT, SIGTERM, SIGQUIT).
@@ -81,6 +89,7 @@ async fn perform_shutdown(
process_started_at: Instant, process_started_at: Instant,
active_runtime: Arc<ArcSwap<RuntimeGeneration>>, active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
quota_state_path: PathBuf, quota_state_path: PathBuf,
synlimit_controller: tokio::task::JoinHandle<()>,
) { ) {
let runtime = active_runtime.load_full(); let runtime = active_runtime.load_full();
let stats = runtime.stats.as_ref(); let stats = runtime.stats.as_ref();
@@ -96,12 +105,9 @@ async fn perform_shutdown(
let uptime_secs = process_started_at.elapsed().as_secs(); let uptime_secs = process_started_at.elapsed().as_secs();
info!("Uptime: {}", format_uptime(uptime_secs)); info!("Uptime: {}", format_uptime(uptime_secs));
if let Err(error) = synlimit_control::clear_synlimit_rules_all_backends().await {
warn!(error = %error, "Failed to clear SYN limiter rules during shutdown");
}
// Graceful ME pool shutdown // Graceful ME pool shutdown
runtime.stop_sessions().await; runtime.stop_sessions().await;
runtime.stop_background_tasks().await;
if let Some(pool) = runtime.current_me_pool().await { if let Some(pool) = runtime.current_me_pool().await {
match tokio::time::timeout(Duration::from_secs(2), pool.shutdown_send_close_conn_all()) match tokio::time::timeout(Duration::from_secs(2), pool.shutdown_send_close_conn_all())
.await .await
@@ -117,7 +123,12 @@ async fn perform_shutdown(
} }
} }
} }
runtime.stop_background_tasks().await;
synlimit_controller.abort();
let _ = synlimit_controller.await;
if let Err(error) = synlimit_control::clear_synlimit_rules_all_backends().await {
warn!(error = %error, "Failed to clear SYN limiter rules during shutdown");
}
match crate::quota_state::save_quota_state(&quota_state_path, stats).await { match crate::quota_state::save_quota_state(&quota_state_path, stats).await {
Ok(()) => { Ok(()) => {
+176 -146
View File
@@ -5,6 +5,7 @@ use rand::RngExt;
use tracing::warn; use tracing::warn;
use crate::config::ProxyConfig; use crate::config::ProxyConfig;
use crate::error::{ProxyError, Result};
use crate::startup::{COMPONENT_TLS_FRONT_BOOTSTRAP, StartupTracker}; use crate::startup::{COMPONENT_TLS_FRONT_BOOTSTRAP, StartupTracker};
use crate::tls_front::TlsFrontCache; use crate::tls_front::TlsFrontCache;
use crate::tls_front::fetcher::TlsFetchStrategy; use crate::tls_front::fetcher::TlsFetchStrategy;
@@ -12,6 +13,82 @@ use crate::transport::UpstreamManager;
use super::generation::RuntimeTaskScope; use super::generation::RuntimeTaskScope;
/// Readiness requirement for TLS-front cache initialization.
#[derive(Clone, Copy)]
pub(crate) enum TlsBootstrapPolicy {
BestEffort,
RequireReady,
}
#[derive(Clone)]
struct TlsFetchContext {
cache: Arc<TlsFrontCache>,
domains: Vec<String>,
mask_host: String,
primary_domain: String,
mask_unix_sock: Option<String>,
tls_fetch_scope: Option<String>,
upstream_manager: Arc<UpstreamManager>,
strategy: TlsFetchStrategy,
port: u16,
proxy_protocol: u8,
}
impl TlsFetchContext {
async fn fetch_all(&self, failure_message: &'static str) {
let mut join = tokio::task::JoinSet::new();
for domain in self.domains.clone() {
let cache = self.cache.clone();
let host = tls_fetch_host_for_domain(
&self.mask_host,
&self.primary_domain,
&domain,
);
let unix_sock = self.mask_unix_sock.clone();
let scope = self.tls_fetch_scope.clone();
let upstream = self.upstream_manager.clone();
let strategy = self.strategy.clone();
let port = self.port;
let proxy_protocol = self.proxy_protocol;
join.spawn(async move {
match crate::tls_front::fetcher::fetch_real_tls_with_strategy(
&host,
port,
&domain,
&strategy,
Some(upstream),
scope.as_deref(),
proxy_protocol,
unix_sock.as_deref(),
)
.await
{
Ok(result) => cache.update_from_fetch(&domain, result).await,
Err(error) => warn!(domain = %domain, error = %error, failure_message),
}
});
}
while let Some(result) = join.join_next().await {
if let Err(error) = result {
warn!(error = %error, "TLS emulation fetch task join failed");
}
}
}
async fn fetch_all_with_budget(&self, phase: &'static str) {
if tokio::time::timeout(self.strategy.total_budget, self.fetch_all(phase))
.await
.is_err()
{
warn!(
phase,
timeout_ms = self.strategy.total_budget.as_millis(),
"TLS emulation fetch budget exhausted"
);
}
}
}
fn tls_fetch_host_for_domain(mask_host: &str, primary_tls_domain: &str, domain: &str) -> String { fn tls_fetch_host_for_domain(mask_host: &str, primary_tls_domain: &str, domain: &str) -> String {
if mask_host.eq_ignore_ascii_case(primary_tls_domain) { if mask_host.eq_ignore_ascii_case(primary_tls_domain) {
domain.to_string() domain.to_string()
@@ -20,13 +97,24 @@ fn tls_fetch_host_for_domain(mask_host: &str, primary_tls_domain: &str, domain:
} }
} }
fn readiness_error(default_domains: &[String]) -> Option<String> {
(!default_domains.is_empty()).then(|| {
format!(
"TLS-front profiles are not ready for domains: {}",
default_domains.join(", ")
)
})
}
/// Initializes the TLS-front cache and generation-owned refresh tasks.
pub(crate) async fn bootstrap_tls_front( pub(crate) async fn bootstrap_tls_front(
config: &ProxyConfig, config: &ProxyConfig,
tls_domains: &[String], tls_domains: &[String],
upstream_manager: Arc<UpstreamManager>, upstream_manager: Arc<UpstreamManager>,
startup_tracker: &Arc<StartupTracker>, startup_tracker: &Arc<StartupTracker>,
task_scope: RuntimeTaskScope, task_scope: RuntimeTaskScope,
) -> Option<Arc<TlsFrontCache>> { policy: TlsBootstrapPolicy,
) -> Result<Option<Arc<TlsFrontCache>>> {
startup_tracker startup_tracker
.start_component( .start_component(
COMPONENT_TLS_FRONT_BOOTSTRAP, COMPONENT_TLS_FRONT_BOOTSTRAP,
@@ -34,26 +122,38 @@ pub(crate) async fn bootstrap_tls_front(
) )
.await; .await;
let tls_cache: Option<Arc<TlsFrontCache>> = if config.censorship.tls_emulation { if !config.censorship.tls_emulation {
let cache = Arc::new(TlsFrontCache::new( startup_tracker
tls_domains, .skip_component(
config.censorship.fake_cert_len, COMPONENT_TLS_FRONT_BOOTSTRAP,
&config.censorship.tls_front_dir, Some("censorship.tls_emulation is false".to_string()),
)); )
cache.load_from_disk().await; .await;
return Ok(None);
}
let port = config.censorship.mask_port; let cache = Arc::new(TlsFrontCache::new(
let proxy_protocol = config.censorship.mask_proxy_protocol; tls_domains,
let mask_host = config config.censorship.fake_cert_len,
&config.censorship.tls_front_dir,
));
cache.load_from_disk().await;
let tls_fetch = config.censorship.tls_fetch.clone();
let fetch_context = TlsFetchContext {
cache: cache.clone(),
domains: tls_domains.to_vec(),
mask_host: config
.censorship .censorship
.mask_host .mask_host
.clone() .clone()
.unwrap_or_else(|| config.censorship.tls_domain.clone()); .unwrap_or_else(|| config.censorship.tls_domain.clone()),
let mask_unix_sock = config.censorship.mask_unix_sock.clone(); primary_domain: config.censorship.tls_domain.clone(),
let tls_fetch_scope = (!config.censorship.tls_fetch_scope.is_empty()) mask_unix_sock: config.censorship.mask_unix_sock.clone(),
.then(|| config.censorship.tls_fetch_scope.clone()); tls_fetch_scope: (!config.censorship.tls_fetch_scope.is_empty())
let tls_fetch = config.censorship.tls_fetch.clone(); .then(|| config.censorship.tls_fetch_scope.clone()),
let fetch_strategy = TlsFetchStrategy { upstream_manager,
strategy: TlsFetchStrategy {
profiles: tls_fetch.profiles, profiles: tls_fetch.profiles,
strict_route: tls_fetch.strict_route, strict_route: tls_fetch.strict_route,
attempt_timeout: Duration::from_millis(tls_fetch.attempt_timeout_ms.max(1)), attempt_timeout: Duration::from_millis(tls_fetch.attempt_timeout_ms.max(1)),
@@ -61,150 +161,71 @@ pub(crate) async fn bootstrap_tls_front(
grease_enabled: tls_fetch.grease_enabled, grease_enabled: tls_fetch.grease_enabled,
deterministic: tls_fetch.deterministic, deterministic: tls_fetch.deterministic,
profile_cache_ttl: Duration::from_secs(tls_fetch.profile_cache_ttl_secs), profile_cache_ttl: Duration::from_secs(tls_fetch.profile_cache_ttl_secs),
}; },
let fetch_timeout = fetch_strategy.total_budget; port: config.censorship.mask_port,
proxy_protocol: config.censorship.mask_proxy_protocol,
};
let cache_initial = cache.clone(); match policy {
let domains_initial = tls_domains.to_vec(); TlsBootstrapPolicy::BestEffort => {
let host_initial = mask_host.clone(); let initial_fetch = fetch_context.clone();
let primary_initial = config.censorship.tls_domain.clone(); let fake_cert_len = config.censorship.fake_cert_len;
let unix_sock_initial = mask_unix_sock.clone(); task_scope.spawn(async move {
let scope_initial = tls_fetch_scope.clone(); initial_fetch
let upstream_initial = upstream_manager.clone(); .fetch_all_with_budget("TLS emulation initial fetch failed")
let strategy_initial = fetch_strategy.clone(); .await;
task_scope.spawn(async move { for domain in initial_fetch
let mut join = tokio::task::JoinSet::new(); .cache
for domain in domains_initial { .default_profile_domains(&initial_fetch.domains)
let cache_domain = cache_initial.clone();
let host_domain =
tls_fetch_host_for_domain(&host_initial, &primary_initial, &domain);
let unix_sock_domain = unix_sock_initial.clone();
let scope_domain = scope_initial.clone();
let upstream_domain = upstream_initial.clone();
let strategy_domain = strategy_initial.clone();
join.spawn(async move {
match crate::tls_front::fetcher::fetch_real_tls_with_strategy(
&host_domain,
port,
&domain,
&strategy_domain,
Some(upstream_domain),
scope_domain.as_deref(),
proxy_protocol,
unix_sock_domain.as_deref(),
)
.await .await
{ {
Ok(res) => cache_domain.update_from_fetch(&domain, res).await,
Err(e) => {
warn!(domain = %domain, error = %e, "TLS emulation initial fetch failed")
}
}
});
}
while let Some(res) = join.join_next().await {
if let Err(e) = res {
warn!(error = %e, "TLS emulation initial fetch task join failed");
}
}
});
let cache_timeout = cache.clone();
let domains_timeout = tls_domains.to_vec();
let fake_cert_len = config.censorship.fake_cert_len;
task_scope.spawn(async move {
tokio::time::sleep(fetch_timeout).await;
for domain in domains_timeout {
let cached = cache_timeout.get(&domain).await;
if cached.domain == "default" {
warn!( warn!(
domain = %domain, domain = %domain,
timeout_secs = fetch_timeout.as_secs(), timeout_ms = initial_fetch.strategy.total_budget.as_millis(),
fake_cert_len, fake_cert_len,
"TLS-front fetch not ready within timeout; using cache/default fake cert fallback" "TLS-front fetch not ready within timeout; using cache/default fake cert fallback"
); );
} }
});
}
TlsBootstrapPolicy::RequireReady => {
fetch_context
.fetch_all_with_budget("TLS emulation initial fetch failed")
.await;
let default_domains = cache.default_profile_domains(tls_domains).await;
if let Some(error) = readiness_error(&default_domains) {
startup_tracker
.fail_component(COMPONENT_TLS_FRONT_BOOTSTRAP, Some(error.clone()))
.await;
return Err(ProxyError::Proxy(error));
} }
}); }
let cache_refresh = cache.clone();
let domains_refresh = tls_domains.to_vec();
let host_refresh = mask_host.clone();
let primary_refresh = config.censorship.tls_domain.clone();
let unix_sock_refresh = mask_unix_sock.clone();
let scope_refresh = tls_fetch_scope.clone();
let upstream_refresh = upstream_manager.clone();
let strategy_refresh = fetch_strategy.clone();
task_scope.spawn(async move {
loop {
let base_secs = rand::rng().random_range(4 * 3600..=6 * 3600);
let jitter_secs = rand::rng().random_range(0..=7200);
tokio::time::sleep(Duration::from_secs(base_secs + jitter_secs)).await;
let mut join = tokio::task::JoinSet::new();
for domain in domains_refresh.clone() {
let cache_domain = cache_refresh.clone();
let host_domain =
tls_fetch_host_for_domain(&host_refresh, &primary_refresh, &domain);
let unix_sock_domain = unix_sock_refresh.clone();
let scope_domain = scope_refresh.clone();
let upstream_domain = upstream_refresh.clone();
let strategy_domain = strategy_refresh.clone();
join.spawn(async move {
match crate::tls_front::fetcher::fetch_real_tls_with_strategy(
&host_domain,
port,
&domain,
&strategy_domain,
Some(upstream_domain),
scope_domain.as_deref(),
proxy_protocol,
unix_sock_domain.as_deref(),
)
.await
{
Ok(res) => cache_domain.update_from_fetch(&domain, res).await,
Err(e) => {
warn!(domain = %domain, error = %e, "TLS emulation refresh failed")
}
}
});
}
while let Some(res) = join.join_next().await {
if let Err(e) = res {
warn!(error = %e, "TLS emulation refresh task join failed");
}
}
}
});
Some(cache)
} else {
startup_tracker
.skip_component(
COMPONENT_TLS_FRONT_BOOTSTRAP,
Some("censorship.tls_emulation is false".to_string()),
)
.await;
None
};
if tls_cache.is_some() {
startup_tracker
.complete_component(
COMPONENT_TLS_FRONT_BOOTSTRAP,
Some("tls front cache is initialized".to_string()),
)
.await;
} }
tls_cache let refresh_context = fetch_context;
task_scope.spawn(async move {
loop {
let base_secs = rand::rng().random_range(4 * 3600..=6 * 3600);
let jitter_secs = rand::rng().random_range(0..=7200);
tokio::time::sleep(Duration::from_secs(base_secs + jitter_secs)).await;
refresh_context
.fetch_all_with_budget("TLS emulation refresh failed")
.await;
}
});
startup_tracker
.complete_component(
COMPONENT_TLS_FRONT_BOOTSTRAP,
Some("tls front cache is initialized".to_string()),
)
.await;
Ok(Some(cache))
} }
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::tls_fetch_host_for_domain; use super::{readiness_error, tls_fetch_host_for_domain};
#[test] #[test]
fn tls_fetch_host_uses_each_domain_when_mask_host_is_primary_default() { fn tls_fetch_host_uses_each_domain_when_mask_host_is_primary_default() {
@@ -221,4 +242,13 @@ mod tests {
"origin.example" "origin.example"
); );
} }
#[test]
fn readiness_rejects_only_default_profiles() {
assert!(readiness_error(&[]).is_none());
assert_eq!(
readiness_error(&["front.example".to_string()]),
Some("TLS-front profiles are not ready for domains: front.example".to_string())
);
}
} }
+162 -17
View File
@@ -4,6 +4,7 @@ use tokio::sync::watch;
use tracing::warn; use tracing::warn;
use crate::config::{ProxyConfig, SynLimitMode}; use crate::config::{ProxyConfig, SynLimitMode};
use crate::maestro::generation::RuntimeWatchState;
mod command; mod command;
mod iptables; mod iptables;
@@ -15,28 +16,105 @@ use self::model::{SynLimitNamespace, synlimit_namespace, synlimit_targets};
static ACTIVE_SYNLIMIT_NAMESPACE: Mutex<Option<SynLimitNamespace>> = Mutex::new(None); static ACTIVE_SYNLIMIT_NAMESPACE: Mutex<Option<SynLimitNamespace>> = Mutex::new(None);
pub(crate) fn spawn_synlimit_controller(config_rx: watch::Receiver<Arc<ProxyConfig>>) { /// 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") { if !cfg!(target_os = "linux") {
if has_synlimit_config(&config_rx.borrow()) { return tokio::spawn(watch_active_runtime_configs(
warn!("SYN limiter is configured but unsupported on this OS; skipping netfilter rules"); runtime_watch_rx,
} |_generation_id, cfg| async move {
return; if has_synlimit_config(&cfg) {
warn!(
"SYN limiter is configured but unsupported on this OS; skipping netfilter rules"
);
}
},
));
} }
tokio::spawn(async move { tokio::spawn(watch_active_runtime_configs(
wait_for_config_channel_close_and_reconcile(config_rx).await; runtime_watch_rx,
if let Err(error) = clear_synlimit_rules_all_backends().await { |_generation_id, cfg| async move {
warn!(error = %error, "Failed to clear SYN limiter rules after config channel close"); reconcile_synlimit_rules(&cfg).await;
} },
}); ))
} }
async fn wait_for_config_channel_close_and_reconcile( async fn watch_active_runtime_configs<F, Fut>(
mut config_rx: watch::Receiver<Arc<ProxyConfig>>, mut runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
) { mut on_config: F,
while config_rx.changed().await.is_ok() { )
let cfg = config_rx.borrow_and_update().clone(); where
reconcile_synlimit_rules(&cfg).await; F: FnMut(u64, Arc<ProxyConfig>) -> Fut,
Fut: std::future::Future<Output = ()>,
{
let mut current = loop {
if let Some(state) = runtime_watch_rx.borrow().clone() {
break state;
}
if runtime_watch_rx.changed().await.is_err() {
return;
}
};
let initial_config = current.config_rx.borrow().clone();
on_config(current.generation_id, initial_config).await;
loop {
tokio::select! {
biased;
changed = runtime_watch_rx.changed() => {
if changed.is_err() {
break;
}
let Some(next) = runtime_watch_rx.borrow().clone() else {
continue;
};
if next.generation_id != current.generation_id {
current = next;
let config = current.config_rx.borrow().clone();
on_config(current.generation_id, config).await;
}
}
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 {
break;
};
current = next;
let config = current.config_rx.borrow().clone();
on_config(current.generation_id, config).await;
continue;
}
let active_generation_id = runtime_watch_rx
.borrow()
.as_ref()
.map(|state| state.generation_id);
if active_generation_id == Some(current.generation_id) {
let cfg = current.config_rx.borrow_and_update().clone();
on_config(current.generation_id, cfg).await;
}
}
}
}
}
async fn wait_for_new_runtime(
runtime_watch_rx: &mut watch::Receiver<Option<RuntimeWatchState>>,
previous_generation_id: u64,
) -> Option<RuntimeWatchState> {
loop {
if let Some(state) = runtime_watch_rx.borrow().clone()
&& state.generation_id != previous_generation_id
{
return Some(state);
}
if runtime_watch_rx.changed().await.is_err() {
return None;
}
} }
} }
@@ -168,3 +246,70 @@ fn has_synlimit_config(cfg: &ProxyConfig) -> bool {
.iter() .iter()
.any(|listener| !matches!(listener.synlimit, SynLimitMode::Off)) .any(|listener| !matches!(listener.synlimit, SynLimitMode::Off))
} }
#[cfg(test)]
mod tests {
use super::*;
use std::time::Duration;
use tokio::sync::mpsc;
fn runtime_state(
generation_id: u64,
max_connections: u32,
) -> (
RuntimeWatchState,
watch::Sender<Arc<ProxyConfig>>,
watch::Sender<bool>,
) {
let mut config = ProxyConfig::default();
config.server.max_connections = max_connections;
let (config_tx, config_rx) = watch::channel(Arc::new(config));
let (admission_tx, admission_rx) = watch::channel(true);
(
RuntimeWatchState {
generation_id,
config_rx,
admission_rx,
},
config_tx,
admission_tx,
)
}
#[tokio::test]
async fn config_watcher_ignores_retired_generation_updates() {
let (initial, initial_config_tx, _initial_admission_tx) = runtime_state(1, 10);
let (runtime_tx, runtime_rx) = watch::channel(Some(initial));
let (observed_tx, mut observed_rx) = mpsc::unbounded_channel();
let watcher = tokio::spawn(watch_active_runtime_configs(
runtime_rx,
move |generation_id, cfg| {
let observed_tx = observed_tx.clone();
async move {
let _ = observed_tx.send((generation_id, cfg.server.max_connections));
}
},
));
assert_eq!(observed_rx.recv().await, Some((1, 10)));
let (next, next_config_tx, _next_admission_tx) = runtime_state(2, 20);
runtime_tx.send_replace(Some(next));
assert_eq!(observed_rx.recv().await, Some((2, 20)));
let mut stale = ProxyConfig::default();
stale.server.max_connections = 30;
initial_config_tx.send_replace(Arc::new(stale));
assert!(
tokio::time::timeout(Duration::from_millis(50), observed_rx.recv())
.await
.is_err()
);
let mut active = ProxyConfig::default();
active.server.max_connections = 40;
next_config_tx.send_replace(Arc::new(active));
assert_eq!(observed_rx.recv().await, Some((2, 40)));
watcher.abort();
}
}
+33
View File
@@ -190,6 +190,22 @@ impl TlsFrontCache {
(snapshot, suppressed) (snapshot, suppressed)
} }
/// Returns configured domains that still resolve to the synthetic default profile.
pub(crate) async fn default_profile_domains(&self, domains: &[String]) -> Vec<String> {
let guard = self.memory.read().await;
domains
.iter()
.filter(|domain| {
guard
.get(domain.as_str())
.unwrap_or(&self.default)
.domain
== "default"
})
.cloned()
.collect()
}
fn full_cert_sent_shard_index(client_ip: IpAddr) -> usize { fn full_cert_sent_shard_index(client_ip: IpAddr) -> usize {
let mut hasher = DefaultHasher::new(); let mut hasher = DefaultHasher::new();
client_ip.hash(&mut hasher); client_ip.hash(&mut hasher);
@@ -546,6 +562,23 @@ mod tests {
assert!(!cert_info_matches_domain(&cached)); assert!(!cert_info_matches_domain(&cached));
} }
#[tokio::test]
async fn default_profile_domains_reports_only_unprepared_entries() {
let domains = vec!["ready.example".to_string(), "pending.example".to_string()];
let cache = TlsFrontCache::new(&domains, 1024, "tlsfront-test-cache");
cache
.set(
"ready.example",
cached_with_cert_info("ready.example", None, Vec::new()),
)
.await;
assert_eq!(
cache.default_profile_domains(&domains).await,
vec!["pending.example".to_string()]
);
}
#[tokio::test] #[tokio::test]
async fn test_take_full_cert_budget_for_ip_uses_ttl() { async fn test_take_full_cert_budget_for_ip_uses_ttl() {
let cache = TlsFrontCache::new(&["example.com".to_string()], 1024, "tlsfront-test-cache"); let cache = TlsFrontCache::new(&["example.com".to_string()], 1024, "tlsfront-test-cache");
+18 -19
View File
@@ -916,26 +916,22 @@ async fn connect_tcp_with_upstream(
strict_route: bool, strict_route: bool,
) -> Result<UpstreamStream> { ) -> Result<UpstreamStream> {
if let Some(manager) = upstream { if let Some(manager) = upstream {
let resolved = if let Some(addr) = resolve_socket_addr(host, port) { let resolved = match manager.resolve_hostname(host, port).await {
Some(addr) Ok(addr) => Some(addr),
} else { Err(e) => {
match tokio::net::lookup_host((host, port)).await { if strict_route {
Ok(mut addrs) => addrs.find(|a| a.is_ipv4()), return Err(anyhow!(
Err(e) => { "upstream route DNS resolution failed for {host}:{port}: {e}"
if strict_route { ));
return Err(anyhow!(
"upstream route DNS resolution failed for {host}:{port}: {e}"
));
}
warn!(
host = %host,
port = port,
scope = ?scope,
error = %e,
"Upstream DNS resolution failed, using direct connect"
);
None
} }
warn!(
host = %host,
port = port,
scope = ?scope,
error = %e,
"Upstream DNS resolution failed, using direct connect"
);
None
} }
}; };
@@ -955,6 +951,9 @@ async fn connect_tcp_with_upstream(
error = %e, error = %e,
"Upstream connect failed, using direct connect" "Upstream connect failed, using direct connect"
); );
return Ok(UpstreamStream::Tcp(
timeout(connect_timeout, TcpStream::connect(addr)).await??,
));
} }
} }
} else if strict_route { } else if strict_route {