mirror of
https://github.com/telemt/telemt.git
synced 2026-08-02 08:35:55 +03:00
Harden Maestro reload lifecycle and readiness barriers
Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
This commit is contained in:
+15
-6
@@ -20,9 +20,9 @@ use tokio::sync::{Mutex, RwLock, Semaphore, watch};
|
||||
use tokio::time::timeout;
|
||||
use tracing::{debug, info, warn};
|
||||
|
||||
use crate::config::{ApiGrayAction, ProxyConfig};
|
||||
use crate::config::ApiGrayAction;
|
||||
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::proxy::route_mode::RouteRuntimeController;
|
||||
use crate::proxy::shared_state::ProxySharedState;
|
||||
@@ -240,8 +240,6 @@ pub async fn serve(
|
||||
route_runtime: Arc<RouteRuntimeController>,
|
||||
proxy_shared: Arc<ProxySharedState>,
|
||||
upstream_manager: Arc<UpstreamManager>,
|
||||
config_rx: watch::Receiver<Arc<ProxyConfig>>,
|
||||
admission_rx: watch::Receiver<bool>,
|
||||
config_path: PathBuf,
|
||||
quota_state_path: PathBuf,
|
||||
detected_ips_rx: watch::Receiver<(Option<IpAddr>, Option<IpAddr>)>,
|
||||
@@ -249,6 +247,7 @@ pub async fn serve(
|
||||
startup_tracker: Arc<StartupTracker>,
|
||||
reload_control: ReloadControl,
|
||||
mut active_runtime_rx: watch::Receiver<Option<Arc<ArcSwap<RuntimeGeneration>>>>,
|
||||
mut runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
|
||||
) {
|
||||
let active_runtime = loop {
|
||||
if let Some(active_runtime) = active_runtime_rx.borrow().clone() {
|
||||
@@ -259,6 +258,17 @@ pub async fn serve(
|
||||
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 {
|
||||
Ok(listener) => listener,
|
||||
Err(error) => {
|
||||
@@ -306,8 +316,7 @@ pub async fn serve(
|
||||
});
|
||||
|
||||
spawn_runtime_watchers(
|
||||
config_rx.clone(),
|
||||
admission_rx.clone(),
|
||||
runtime_watch_rx,
|
||||
runtime_state.clone(),
|
||||
shared.runtime_events.clone(),
|
||||
);
|
||||
|
||||
+258
-38
@@ -4,63 +4,283 @@ use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use tokio::sync::watch;
|
||||
|
||||
use crate::config::ProxyConfig;
|
||||
use crate::maestro::generation::RuntimeWatchState;
|
||||
|
||||
use super::ApiRuntimeState;
|
||||
use super::events::ApiEventStore;
|
||||
|
||||
pub(super) fn spawn_runtime_watchers(
|
||||
config_rx: watch::Receiver<Arc<ProxyConfig>>,
|
||||
admission_rx: watch::Receiver<bool>,
|
||||
runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
|
||||
runtime_state: Arc<ApiRuntimeState>,
|
||||
runtime_events: Arc<ApiEventStore>,
|
||||
) {
|
||||
let mut config_rx_reload = config_rx;
|
||||
let runtime_state_reload = runtime_state.clone();
|
||||
let runtime_events_reload = runtime_events.clone();
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
if config_rx_reload.changed().await.is_err() {
|
||||
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");
|
||||
}
|
||||
});
|
||||
spawn_config_watcher(
|
||||
runtime_watch_rx.clone(),
|
||||
runtime_state.clone(),
|
||||
runtime_events.clone(),
|
||||
);
|
||||
spawn_admission_watcher(runtime_watch_rx, runtime_state, runtime_events);
|
||||
}
|
||||
|
||||
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 {
|
||||
runtime_state
|
||||
.admission_open
|
||||
.store(*admission_rx_watch.borrow(), Ordering::Relaxed);
|
||||
runtime_events.record(
|
||||
"admission.state",
|
||||
format!("accepting_new_connections={}", *admission_rx_watch.borrow()),
|
||||
);
|
||||
let Some(mut current) = runtime_watch_rx.borrow().clone() else {
|
||||
return;
|
||||
};
|
||||
loop {
|
||||
if admission_rx_watch.changed().await.is_err() {
|
||||
break;
|
||||
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_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, ¤t);
|
||||
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, ¤t);
|
||||
}
|
||||
}
|
||||
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, ¤t);
|
||||
continue;
|
||||
}
|
||||
if active_generation_id(&runtime_watch_rx) == Some(current.generation_id) {
|
||||
record_admission_state(&runtime_state, &runtime_events, ¤t);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
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 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.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
@@ -1,4 +1,5 @@
|
||||
use std::sync::atomic::Ordering;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use serde::Serialize;
|
||||
|
||||
@@ -162,7 +163,7 @@ pub(super) fn build_system_info_data(
|
||||
build_time_utc,
|
||||
rustc_version,
|
||||
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_hash: revision.to_string(),
|
||||
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(
|
||||
shared: &ApiShared,
|
||||
cfg: &ProxyConfig,
|
||||
@@ -339,3 +352,14 @@ fn me_writer_pick_mode_label(mode: MeWriterPickMode) -> &'static str {
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -22,6 +22,14 @@ use crate::transport::middle_proxy::MePool;
|
||||
const SESSION_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.
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct RuntimeTaskScope {
|
||||
@@ -134,6 +142,15 @@ impl RuntimeGeneration {
|
||||
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>> {
|
||||
if let Some(pool) = &self.me_pool {
|
||||
return Some(pool.clone());
|
||||
|
||||
+150
-200
@@ -1,9 +1,11 @@
|
||||
#![allow(clippy::too_many_arguments)]
|
||||
|
||||
use std::future::Future;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use tokio::sync::{RwLock, watch};
|
||||
use tokio_util::task::AbortOnDropHandle;
|
||||
use tracing::{error, info, warn};
|
||||
|
||||
use crate::config::ProxyConfig;
|
||||
@@ -17,8 +19,58 @@ use crate::stats::Stats;
|
||||
use crate::transport::UpstreamManager;
|
||||
use crate::transport::middle_proxy::MePool;
|
||||
|
||||
use super::generation::RuntimeTaskScope;
|
||||
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(
|
||||
use_middle_proxy: bool,
|
||||
config: &ProxyConfig,
|
||||
@@ -30,6 +82,7 @@ pub(crate) async fn initialize_me_pool(
|
||||
stats: Arc<Stats>,
|
||||
api_me_pool: Arc<RwLock<Option<Arc<MePool>>>>,
|
||||
me_ready_tx: watch::Sender<u64>,
|
||||
task_scope: RuntimeTaskScope,
|
||||
) -> Option<Arc<MePool>> {
|
||||
if !use_middle_proxy {
|
||||
return None;
|
||||
@@ -52,15 +105,8 @@ pub(crate) async fn initialize_me_pool(
|
||||
.as_ref()
|
||||
.map(|tag| hex::decode(tag).expect("general.ad_tag must be validated before startup"));
|
||||
|
||||
// =============================================================
|
||||
// CRITICAL: Download Telegram proxy-secret (NOT user secret!)
|
||||
//
|
||||
// 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
|
||||
// =============================================================
|
||||
// The Telegram proxy-secret authenticates ME RPC and is distinct from client secrets.
|
||||
// It corresponds to the C MTProxy --aes-pwd input and may be fetched from Telegram.
|
||||
let proxy_secret_path = config.general.proxy_secret_path.as_deref();
|
||||
let pool_size = config.general.middle_proxy_pool_size.max(1);
|
||||
let proxy_secret = loop {
|
||||
@@ -319,143 +365,70 @@ pub(crate) async fn initialize_me_pool(
|
||||
let rng_bg = rng.clone();
|
||||
let startup_tracker_bg = startup_tracker.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 {
|
||||
String::from("unlimited")
|
||||
} else {
|
||||
me_init_retry_attempts.to_string()
|
||||
};
|
||||
std::thread::spawn(move || {
|
||||
let runtime = match tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
{
|
||||
Ok(runtime) => runtime,
|
||||
Err(error) => {
|
||||
error!(error = %error, "Failed to build background runtime for ME initialization");
|
||||
return;
|
||||
}
|
||||
};
|
||||
runtime.block_on(async move {
|
||||
let mut init_attempt: u32 = 0;
|
||||
loop {
|
||||
init_attempt = init_attempt.saturating_add(1);
|
||||
startup_tracker_bg.set_me_init_attempt(init_attempt).await;
|
||||
match pool_bg.init(pool_size, &rng_bg).await {
|
||||
Ok(()) => {
|
||||
startup_tracker_bg.set_me_last_error(None).await;
|
||||
startup_tracker_bg
|
||||
.complete_component(
|
||||
COMPONENT_ME_POOL_INIT_STAGE1,
|
||||
Some("ME pool initialized".to_string()),
|
||||
)
|
||||
.await;
|
||||
startup_tracker_bg
|
||||
.set_me_status(StartupMeStatus::Ready, "ready")
|
||||
.await;
|
||||
me_ready_tx_bg.send_modify(|version| {
|
||||
*version = version.saturating_add(1);
|
||||
});
|
||||
info!(
|
||||
task_scope.spawn(async move {
|
||||
let mut init_attempt: u32 = 0;
|
||||
loop {
|
||||
init_attempt = init_attempt.saturating_add(1);
|
||||
startup_tracker_bg.set_me_init_attempt(init_attempt).await;
|
||||
match pool_bg.init(pool_size, &rng_bg).await {
|
||||
Ok(()) => {
|
||||
startup_tracker_bg.set_me_last_error(None).await;
|
||||
startup_tracker_bg
|
||||
.complete_component(
|
||||
COMPONENT_ME_POOL_INIT_STAGE1,
|
||||
Some("ME pool initialized".to_string()),
|
||||
)
|
||||
.await;
|
||||
startup_tracker_bg
|
||||
.set_me_status(StartupMeStatus::Ready, "ready")
|
||||
.await;
|
||||
me_ready_tx_bg.send_modify(|version| {
|
||||
*version = version.saturating_add(1);
|
||||
});
|
||||
info!(
|
||||
attempt = init_attempt,
|
||||
"Middle-End pool initialized successfully"
|
||||
);
|
||||
spawn_me_supervisors(
|
||||
task_scope_bg,
|
||||
pool_bg.clone(),
|
||||
rng_bg.clone(),
|
||||
pool_size,
|
||||
);
|
||||
break;
|
||||
}
|
||||
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,
|
||||
"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
|
||||
.set_me_status(StartupMeStatus::Initializing, "background_init")
|
||||
@@ -490,70 +463,12 @@ pub(crate) async fn initialize_me_pool(
|
||||
"Middle-End pool initialized successfully"
|
||||
);
|
||||
|
||||
// ── Supervised background tasks ──────────────────
|
||||
let pool_clone = pool.clone();
|
||||
let rng_clone = rng.clone();
|
||||
let min_conns = pool_size;
|
||||
tokio::spawn(async move {
|
||||
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;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
spawn_me_supervisors(
|
||||
task_scope.clone(),
|
||||
pool.clone(),
|
||||
rng.clone(),
|
||||
pool_size,
|
||||
);
|
||||
|
||||
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
@@ -503,7 +503,6 @@ async fn run_telemt_core(
|
||||
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 initial_direct_first = config.general.use_middle_proxy && config.general.me2dc_fallback;
|
||||
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 (active_runtime_tx, active_runtime_rx) =
|
||||
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 {
|
||||
RelayRouteMode::Direct
|
||||
} else {
|
||||
@@ -544,14 +545,13 @@ async fn run_telemt_core(
|
||||
let upstream_manager_api = upstream_manager.clone();
|
||||
let route_runtime_api = route_runtime.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 quota_state_path_api = quota_state_path.clone();
|
||||
let startup_tracker_api = startup_tracker.clone();
|
||||
let detected_ips_rx_api = detected_ips_rx.clone();
|
||||
let reload_control_api = reload_control.clone();
|
||||
let active_runtime_rx_api = active_runtime_rx.clone();
|
||||
let runtime_watch_rx_api = runtime_watch_rx.clone();
|
||||
tokio::spawn(async move {
|
||||
api::serve(
|
||||
listen,
|
||||
@@ -561,8 +561,6 @@ async fn run_telemt_core(
|
||||
route_runtime_api,
|
||||
proxy_shared_api,
|
||||
upstream_manager_api,
|
||||
config_rx_api,
|
||||
admission_rx_api,
|
||||
config_path_api,
|
||||
quota_state_path_api,
|
||||
detected_ips_rx_api,
|
||||
@@ -570,6 +568,7 @@ async fn run_telemt_core(
|
||||
startup_tracker_api,
|
||||
reload_control_api,
|
||||
active_runtime_rx_api,
|
||||
runtime_watch_rx_api,
|
||||
)
|
||||
.await;
|
||||
});
|
||||
@@ -610,8 +609,9 @@ async fn run_telemt_core(
|
||||
upstream_manager.clone(),
|
||||
&startup_tracker,
|
||||
runtime_task_scope.clone(),
|
||||
tls_bootstrap::TlsBootstrapPolicy::BestEffort,
|
||||
)
|
||||
.await;
|
||||
.await?;
|
||||
|
||||
startup_tracker
|
||||
.start_component(
|
||||
@@ -737,6 +737,7 @@ async fn run_telemt_core(
|
||||
stats.clone(),
|
||||
api_me_pool.clone(),
|
||||
me_ready_tx.clone(),
|
||||
runtime_task_scope.clone(),
|
||||
)
|
||||
.await
|
||||
};
|
||||
@@ -824,7 +825,6 @@ async fn run_telemt_core(
|
||||
rng.clone(),
|
||||
ip_tracker.clone(),
|
||||
beobachten.clone(),
|
||||
api_config_tx.clone(),
|
||||
me_pool.clone(),
|
||||
shared_state.clone(),
|
||||
me_ready_tx.clone(),
|
||||
@@ -869,6 +869,7 @@ async fn run_telemt_core(
|
||||
stats_bg.clone(),
|
||||
api_me_pool_bg.clone(),
|
||||
me_ready_tx_bg.clone(),
|
||||
task_scope_bg.clone(),
|
||||
)
|
||||
.await;
|
||||
if let Some(pool) = pool {
|
||||
@@ -961,6 +962,7 @@ async fn run_telemt_core(
|
||||
quota_store,
|
||||
detected_ips_tx,
|
||||
runtime_log_filter,
|
||||
runtime_watch_tx.clone(),
|
||||
);
|
||||
|
||||
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).
|
||||
drop_after_bind();
|
||||
|
||||
synlimit_control::reconcile_synlimit_rules(&config).await;
|
||||
synlimit_control::spawn_synlimit_controller(config_rx.clone());
|
||||
let synlimit_controller = synlimit_control::spawn_synlimit_controller(runtime_watch_rx);
|
||||
|
||||
runtime_tasks::spawn_metrics_if_configured(&config, &startup_tracker, active_runtime.clone())
|
||||
.await;
|
||||
|
||||
runtime_watch_tx.send_replace(Some(active_runtime.load_full().watch_state()));
|
||||
active_runtime_tx.send_replace(Some(active_runtime.clone()));
|
||||
runtime_tasks::mark_runtime_ready(&startup_tracker).await;
|
||||
|
||||
@@ -1004,7 +1006,13 @@ async fn run_telemt_core(
|
||||
#[cfg(unix)]
|
||||
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(())
|
||||
}
|
||||
|
||||
+30
-1
@@ -138,7 +138,11 @@ pub(crate) struct ReloadStatus {
|
||||
pub(crate) started_at_epoch_secs: Option<u64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
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>,
|
||||
#[serde(default, skip_serializing_if = "Vec::is_empty")]
|
||||
pub(crate) warnings: Vec<String>,
|
||||
@@ -428,6 +432,31 @@ mod tests {
|
||||
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);
|
||||
|
||||
@@ -8,7 +8,7 @@ use tracing::{info, warn};
|
||||
|
||||
use crate::stats::QuotaStore;
|
||||
|
||||
use super::generation::RuntimeGeneration;
|
||||
use super::generation::{RuntimeGeneration, RuntimeWatchState};
|
||||
use super::reload::{
|
||||
ReloadCommand, ReloadCommandReceiver, ReloadControl, ReloadFailurePolicy, ReloadMode,
|
||||
ReloadPhase,
|
||||
@@ -24,6 +24,48 @@ pub(crate) struct ReloadSupervisor {
|
||||
quota_store: Arc<QuotaStore>,
|
||||
detected_ips_tx: watch::Sender<(Option<std::net::IpAddr>, Option<std::net::IpAddr>)>,
|
||||
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 {
|
||||
@@ -36,6 +78,7 @@ impl ReloadSupervisor {
|
||||
quota_store: Arc<QuotaStore>,
|
||||
detected_ips_tx: watch::Sender<(Option<std::net::IpAddr>, Option<std::net::IpAddr>)>,
|
||||
runtime_log_filter: RuntimeLogFilter,
|
||||
runtime_watch_tx: watch::Sender<Option<RuntimeWatchState>>,
|
||||
) {
|
||||
let supervisor = Self {
|
||||
active_runtime,
|
||||
@@ -45,6 +88,7 @@ impl ReloadSupervisor {
|
||||
quota_store,
|
||||
detected_ips_tx,
|
||||
runtime_log_filter,
|
||||
runtime_watch_tx,
|
||||
};
|
||||
tokio::spawn(supervisor.run());
|
||||
}
|
||||
@@ -81,30 +125,23 @@ impl ReloadSupervisor {
|
||||
}
|
||||
};
|
||||
|
||||
if let Ok(current_revision) =
|
||||
crate::api::config_store::current_revision_for_maestro(&self.config_path).await
|
||||
&& current_revision != command.config_revision
|
||||
{
|
||||
let warning = format!(
|
||||
"config revision changed during preparation: accepted={} current={}",
|
||||
command.config_revision, current_revision
|
||||
);
|
||||
if command.request.failure_policy == ReloadFailurePolicy::Rollback {
|
||||
let revision_action = revision_gate_action(
|
||||
&command.config_revision,
|
||||
crate::api::config_store::current_revision_for_maestro(&self.config_path).await,
|
||||
command.request.failure_policy,
|
||||
);
|
||||
match revision_action {
|
||||
RevisionGateAction::Proceed => {}
|
||||
RevisionGateAction::Warn(warning) => {
|
||||
self.control.add_warning(command.reload_id, warning).await;
|
||||
}
|
||||
RevisionGateAction::Rollback(warning) => {
|
||||
let _ = cleanup_candidate(&prepared.generation).await;
|
||||
self.runtime_log_filter
|
||||
.apply_reload(&old_runtime.config().general.log_level);
|
||||
prepared.generation.stop_sessions().await;
|
||||
if let Some(pool) = prepared.generation.current_me_pool().await {
|
||||
let _ = tokio::time::timeout(
|
||||
Duration::from_secs(2),
|
||||
pool.shutdown_send_close_conn_all(),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
prepared.generation.stop_background_tasks().await;
|
||||
self.control.rolled_back(command.reload_id, warning).await;
|
||||
return;
|
||||
}
|
||||
self.control.add_warning(command.reload_id, warning).await;
|
||||
}
|
||||
|
||||
self.control
|
||||
@@ -112,37 +149,26 @@ impl ReloadSupervisor {
|
||||
.await;
|
||||
let new_runtime = prepared.generation;
|
||||
old_runtime.stop_accepting_sessions();
|
||||
let replaced = self.active_runtime.swap(new_runtime.clone());
|
||||
if let Err(error) = crate::network::dns_overrides::install_entries(
|
||||
&new_runtime.config().network.dns_overrides,
|
||||
) {
|
||||
let message = format!("runtime DNS activation failed: {}", error);
|
||||
if command.request.failure_policy == ReloadFailurePolicy::Rollback {
|
||||
let candidate = self.active_runtime.swap(replaced.clone());
|
||||
replaced.resume_accepting_sessions();
|
||||
old_runtime.resume_accepting_sessions();
|
||||
let _ = cleanup_candidate(&new_runtime).await;
|
||||
self.runtime_log_filter
|
||||
.apply_reload(&replaced.config().general.log_level);
|
||||
let _ = crate::network::dns_overrides::install_entries(
|
||||
&replaced.config().network.dns_overrides,
|
||||
);
|
||||
candidate.stop_sessions().await;
|
||||
if let Some(pool) = candidate.current_me_pool().await {
|
||||
let _ = tokio::time::timeout(
|
||||
Duration::from_secs(2),
|
||||
pool.shutdown_send_close_conn_all(),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
candidate.stop_background_tasks().await;
|
||||
.apply_reload(&old_runtime.config().general.log_level);
|
||||
self.control.rolled_back(command.reload_id, message).await;
|
||||
return;
|
||||
}
|
||||
self.control.add_warning(command.reload_id, message).await;
|
||||
}
|
||||
let replaced = self.active_runtime.swap(new_runtime.clone());
|
||||
self.detected_ips_tx.send_replace(prepared.detected_ips);
|
||||
self.runtime_log_filter
|
||||
.apply_reload(&new_runtime.config().general.log_level);
|
||||
crate::synlimit_control::reconcile_synlimit_rules(&new_runtime.config()).await;
|
||||
self.runtime_watch_tx
|
||||
.send_replace(Some(new_runtime.watch_state()));
|
||||
|
||||
info!(
|
||||
reload_id = command.reload_id,
|
||||
@@ -177,11 +203,7 @@ impl ReloadSupervisor {
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(pool) = replaced.current_me_pool().await
|
||||
&& tokio::time::timeout(Duration::from_secs(2), pool.shutdown_send_close_conn_all())
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
if stop_background_and_middle_end(&replaced).await {
|
||||
let warning = format!(
|
||||
"generation {} Middle-End close broadcast timed out",
|
||||
replaced.id
|
||||
@@ -189,9 +211,46 @@ impl ReloadSupervisor {
|
||||
warn!(reload_id = command.reload_id, warning = %warning);
|
||||
self.control.add_warning(command.reload_id, warning).await;
|
||||
}
|
||||
replaced.stop_background_tasks().await;
|
||||
self.control
|
||||
.succeed(command.reload_id, new_runtime.id)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
#[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(_)
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -34,7 +34,7 @@ pub(crate) struct PreparedRuntime {
|
||||
|
||||
pub(crate) async fn prepare_runtime(
|
||||
generation_id: u64,
|
||||
mut config: ProxyConfig,
|
||||
config: ProxyConfig,
|
||||
config_path: &Path,
|
||||
quota_store: Arc<QuotaStore>,
|
||||
runtime_log_filter: RuntimeLogFilter,
|
||||
@@ -112,8 +112,10 @@ pub(crate) async fn prepare_runtime(
|
||||
upstream_manager.clone(),
|
||||
&startup_tracker,
|
||||
task_scope.clone(),
|
||||
tls_bootstrap::TlsBootstrapPolicy::RequireReady,
|
||||
)
|
||||
.await;
|
||||
.await
|
||||
.map_err(|error| error.to_string())?;
|
||||
|
||||
let beobachten = Arc::new(BeobachtenStore::new());
|
||||
let rng = Arc::new(SecureRandom::new());
|
||||
@@ -140,11 +142,20 @@ pub(crate) async fn prepare_runtime(
|
||||
stats.clone(),
|
||||
me_pool_runtime.clone(),
|
||||
me_ready_tx.clone(),
|
||||
task_scope.clone(),
|
||||
)
|
||||
.await
|
||||
};
|
||||
if config.general.use_middle_proxy && me_pool.is_none() && !direct_first_startup {
|
||||
config.general.use_middle_proxy = false;
|
||||
if strict_middle_proxy_unavailable(
|
||||
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);
|
||||
@@ -159,7 +170,6 @@ pub(crate) async fn prepare_runtime(
|
||||
config.server.max_connections as usize
|
||||
};
|
||||
let max_connections = Arc::new(Semaphore::new(max_connections_limit));
|
||||
let (api_config_tx, _) = watch::channel(config.clone());
|
||||
let watches = runtime_tasks::spawn_runtime_tasks(
|
||||
&config,
|
||||
config_path,
|
||||
@@ -175,7 +185,6 @@ pub(crate) async fn prepare_runtime(
|
||||
rng.clone(),
|
||||
ip_tracker.clone(),
|
||||
beobachten.clone(),
|
||||
api_config_tx,
|
||||
me_pool.clone(),
|
||||
proxy_shared.clone(),
|
||||
me_ready_tx.clone(),
|
||||
@@ -226,6 +235,7 @@ pub(crate) async fn prepare_runtime(
|
||||
stats_bg.clone(),
|
||||
me_pool_runtime_bg.clone(),
|
||||
me_ready_tx_bg.clone(),
|
||||
task_scope_bg.clone(),
|
||||
)
|
||||
.await;
|
||||
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> {
|
||||
let mut fields = Vec::new();
|
||||
if old.server.port != new.server.port
|
||||
@@ -354,4 +372,12 @@ mod tests {
|
||||
new.censorship.tls_domain = "reload.example".to_string();
|
||||
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));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -103,7 +103,6 @@ pub(crate) async fn spawn_runtime_tasks(
|
||||
rng: Arc<SecureRandom>,
|
||||
ip_tracker: Arc<UserIpTracker>,
|
||||
beobachten: Arc<BeobachtenStore>,
|
||||
api_config_tx: watch::Sender<Arc<ProxyConfig>>,
|
||||
me_pool_for_policy: Option<Arc<MePool>>,
|
||||
shared_state: Arc<ProxySharedState>,
|
||||
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()),
|
||||
)
|
||||
.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 upstream_policy = upstream_manager.clone();
|
||||
let mut config_rx_policy = config_rx.clone();
|
||||
|
||||
+17
-6
@@ -50,9 +50,17 @@ 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<()>,
|
||||
) {
|
||||
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).
|
||||
@@ -81,6 +89,7 @@ async fn perform_shutdown(
|
||||
process_started_at: Instant,
|
||||
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
|
||||
quota_state_path: PathBuf,
|
||||
synlimit_controller: tokio::task::JoinHandle<()>,
|
||||
) {
|
||||
let runtime = active_runtime.load_full();
|
||||
let stats = runtime.stats.as_ref();
|
||||
@@ -96,12 +105,9 @@ async fn perform_shutdown(
|
||||
let uptime_secs = process_started_at.elapsed().as_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
|
||||
runtime.stop_sessions().await;
|
||||
runtime.stop_background_tasks().await;
|
||||
if let Some(pool) = runtime.current_me_pool().await {
|
||||
match tokio::time::timeout(Duration::from_secs(2), pool.shutdown_send_close_conn_all())
|
||||
.await
|
||||
@@ -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("a_state_path, stats).await {
|
||||
Ok(()) => {
|
||||
|
||||
+176
-146
@@ -5,6 +5,7 @@ use rand::RngExt;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::config::ProxyConfig;
|
||||
use crate::error::{ProxyError, Result};
|
||||
use crate::startup::{COMPONENT_TLS_FRONT_BOOTSTRAP, StartupTracker};
|
||||
use crate::tls_front::TlsFrontCache;
|
||||
use crate::tls_front::fetcher::TlsFetchStrategy;
|
||||
@@ -12,6 +13,82 @@ use crate::transport::UpstreamManager;
|
||||
|
||||
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 {
|
||||
if mask_host.eq_ignore_ascii_case(primary_tls_domain) {
|
||||
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(
|
||||
config: &ProxyConfig,
|
||||
tls_domains: &[String],
|
||||
upstream_manager: Arc<UpstreamManager>,
|
||||
startup_tracker: &Arc<StartupTracker>,
|
||||
task_scope: RuntimeTaskScope,
|
||||
) -> Option<Arc<TlsFrontCache>> {
|
||||
policy: TlsBootstrapPolicy,
|
||||
) -> Result<Option<Arc<TlsFrontCache>>> {
|
||||
startup_tracker
|
||||
.start_component(
|
||||
COMPONENT_TLS_FRONT_BOOTSTRAP,
|
||||
@@ -34,26 +122,38 @@ pub(crate) async fn bootstrap_tls_front(
|
||||
)
|
||||
.await;
|
||||
|
||||
let tls_cache: Option<Arc<TlsFrontCache>> = if config.censorship.tls_emulation {
|
||||
let cache = Arc::new(TlsFrontCache::new(
|
||||
tls_domains,
|
||||
config.censorship.fake_cert_len,
|
||||
&config.censorship.tls_front_dir,
|
||||
));
|
||||
cache.load_from_disk().await;
|
||||
if !config.censorship.tls_emulation {
|
||||
startup_tracker
|
||||
.skip_component(
|
||||
COMPONENT_TLS_FRONT_BOOTSTRAP,
|
||||
Some("censorship.tls_emulation is false".to_string()),
|
||||
)
|
||||
.await;
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let port = config.censorship.mask_port;
|
||||
let proxy_protocol = config.censorship.mask_proxy_protocol;
|
||||
let mask_host = config
|
||||
let cache = Arc::new(TlsFrontCache::new(
|
||||
tls_domains,
|
||||
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
|
||||
.mask_host
|
||||
.clone()
|
||||
.unwrap_or_else(|| config.censorship.tls_domain.clone());
|
||||
let mask_unix_sock = config.censorship.mask_unix_sock.clone();
|
||||
let tls_fetch_scope = (!config.censorship.tls_fetch_scope.is_empty())
|
||||
.then(|| config.censorship.tls_fetch_scope.clone());
|
||||
let tls_fetch = config.censorship.tls_fetch.clone();
|
||||
let fetch_strategy = TlsFetchStrategy {
|
||||
.unwrap_or_else(|| config.censorship.tls_domain.clone()),
|
||||
primary_domain: config.censorship.tls_domain.clone(),
|
||||
mask_unix_sock: config.censorship.mask_unix_sock.clone(),
|
||||
tls_fetch_scope: (!config.censorship.tls_fetch_scope.is_empty())
|
||||
.then(|| config.censorship.tls_fetch_scope.clone()),
|
||||
upstream_manager,
|
||||
strategy: TlsFetchStrategy {
|
||||
profiles: tls_fetch.profiles,
|
||||
strict_route: tls_fetch.strict_route,
|
||||
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,
|
||||
deterministic: tls_fetch.deterministic,
|
||||
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();
|
||||
let domains_initial = tls_domains.to_vec();
|
||||
let host_initial = mask_host.clone();
|
||||
let primary_initial = config.censorship.tls_domain.clone();
|
||||
let unix_sock_initial = mask_unix_sock.clone();
|
||||
let scope_initial = tls_fetch_scope.clone();
|
||||
let upstream_initial = upstream_manager.clone();
|
||||
let strategy_initial = fetch_strategy.clone();
|
||||
task_scope.spawn(async move {
|
||||
let mut join = tokio::task::JoinSet::new();
|
||||
for domain in domains_initial {
|
||||
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(),
|
||||
)
|
||||
match policy {
|
||||
TlsBootstrapPolicy::BestEffort => {
|
||||
let initial_fetch = fetch_context.clone();
|
||||
let fake_cert_len = config.censorship.fake_cert_len;
|
||||
task_scope.spawn(async move {
|
||||
initial_fetch
|
||||
.fetch_all_with_budget("TLS emulation initial fetch failed")
|
||||
.await;
|
||||
for domain in initial_fetch
|
||||
.cache
|
||||
.default_profile_domains(&initial_fetch.domains)
|
||||
.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!(
|
||||
domain = %domain,
|
||||
timeout_secs = fetch_timeout.as_secs(),
|
||||
timeout_ms = initial_fetch.strategy.total_budget.as_millis(),
|
||||
fake_cert_len,
|
||||
"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)]
|
||||
mod tests {
|
||||
use super::tls_fetch_host_for_domain;
|
||||
use super::{readiness_error, tls_fetch_host_for_domain};
|
||||
|
||||
#[test]
|
||||
fn tls_fetch_host_uses_each_domain_when_mask_host_is_primary_default() {
|
||||
@@ -221,4 +242,13 @@ mod tests {
|
||||
"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
@@ -4,6 +4,7 @@ use tokio::sync::watch;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::config::{ProxyConfig, SynLimitMode};
|
||||
use crate::maestro::generation::RuntimeWatchState;
|
||||
|
||||
mod command;
|
||||
mod iptables;
|
||||
@@ -15,28 +16,105 @@ use self::model::{SynLimitNamespace, synlimit_namespace, synlimit_targets};
|
||||
|
||||
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 has_synlimit_config(&config_rx.borrow()) {
|
||||
warn!("SYN limiter is configured but unsupported on this OS; skipping netfilter rules");
|
||||
}
|
||||
return;
|
||||
return tokio::spawn(watch_active_runtime_configs(
|
||||
runtime_watch_rx,
|
||||
|_generation_id, cfg| async move {
|
||||
if has_synlimit_config(&cfg) {
|
||||
warn!(
|
||||
"SYN limiter is configured but unsupported on this OS; skipping netfilter rules"
|
||||
);
|
||||
}
|
||||
},
|
||||
));
|
||||
}
|
||||
|
||||
tokio::spawn(async move {
|
||||
wait_for_config_channel_close_and_reconcile(config_rx).await;
|
||||
if let Err(error) = clear_synlimit_rules_all_backends().await {
|
||||
warn!(error = %error, "Failed to clear SYN limiter rules after config channel close");
|
||||
}
|
||||
});
|
||||
tokio::spawn(watch_active_runtime_configs(
|
||||
runtime_watch_rx,
|
||||
|_generation_id, cfg| async move {
|
||||
reconcile_synlimit_rules(&cfg).await;
|
||||
},
|
||||
))
|
||||
}
|
||||
|
||||
async fn wait_for_config_channel_close_and_reconcile(
|
||||
mut config_rx: watch::Receiver<Arc<ProxyConfig>>,
|
||||
) {
|
||||
while config_rx.changed().await.is_ok() {
|
||||
let cfg = config_rx.borrow_and_update().clone();
|
||||
reconcile_synlimit_rules(&cfg).await;
|
||||
async fn watch_active_runtime_configs<F, Fut>(
|
||||
mut runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
|
||||
mut on_config: F,
|
||||
)
|
||||
where
|
||||
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()
|
||||
.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();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -190,6 +190,22 @@ impl TlsFrontCache {
|
||||
(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 {
|
||||
let mut hasher = DefaultHasher::new();
|
||||
client_ip.hash(&mut hasher);
|
||||
@@ -546,6 +562,23 @@ mod tests {
|
||||
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]
|
||||
async fn test_take_full_cert_budget_for_ip_uses_ttl() {
|
||||
let cache = TlsFrontCache::new(&["example.com".to_string()], 1024, "tlsfront-test-cache");
|
||||
|
||||
+18
-19
@@ -916,26 +916,22 @@ async fn connect_tcp_with_upstream(
|
||||
strict_route: bool,
|
||||
) -> Result<UpstreamStream> {
|
||||
if let Some(manager) = upstream {
|
||||
let resolved = if let Some(addr) = resolve_socket_addr(host, port) {
|
||||
Some(addr)
|
||||
} else {
|
||||
match tokio::net::lookup_host((host, port)).await {
|
||||
Ok(mut addrs) => addrs.find(|a| a.is_ipv4()),
|
||||
Err(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
|
||||
let resolved = match manager.resolve_hostname(host, port).await {
|
||||
Ok(addr) => Some(addr),
|
||||
Err(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
|
||||
}
|
||||
};
|
||||
|
||||
@@ -955,6 +951,9 @@ async fn connect_tcp_with_upstream(
|
||||
error = %e,
|
||||
"Upstream connect failed, using direct connect"
|
||||
);
|
||||
return Ok(UpstreamStream::Tcp(
|
||||
timeout(connect_timeout, TcpStream::connect(addr)).await??,
|
||||
));
|
||||
}
|
||||
}
|
||||
} else if strict_route {
|
||||
|
||||
Reference in New Issue
Block a user