astelm/telemt
1
0
Fork 0
mirror of https://github.com/telemt/telemt.git synced 2026-10-30 20:40:51 +03:00
Code Issues Packages Projects Releases Wiki Activity

Maestro: add in-process runtime generation reload

Browse Source 
Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
This commit is contained in:
Alexey
2026-07-17 23:56:00 +03:00
parent 1f9c82c924
commit 991d5b2c38
26 changed files with 2207 additions and 652 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/maestro/admission.rs
+4 -1
View File
@@ -8,6 +8,8 @@ use crate::config::ProxyConfig;
use crate::proxy::route_mode::{RelayRouteMode, RouteRuntimeController};
use crate::transport::middle_proxy::MePool;
use super::generation::RuntimeTaskScope;
const STARTUP_FALLBACK_AFTER: Duration = Duration::from_secs(80);
const RUNTIME_FALLBACK_AFTER: Duration = Duration::from_secs(6);
@@ -19,6 +21,7 @@ pub(crate) async fn configure_admission_gate(
admission_tx: &watch::Sender<bool>,
config_rx: watch::Receiver<Arc<ProxyConfig>>,
me_ready_rx: watch::Receiver<u64>,
task_scope: RuntimeTaskScope,
) {
if config.general.use_middle_proxy {
if me_pool.is_some() || config.general.me2dc_fallback {
@@ -64,7 +67,7 @@ pub(crate) async fn configure_admission_gate(
let mut config_rx_gate = config_rx.clone();
let mut me_ready_rx_gate = me_ready_rx;
let mut admission_poll_ms = config.general.me_admission_poll_ms.max(1);
tokio::spawn(async move {
task_scope.spawn(async move {
let mut gate_open = initial_gate_open;
let mut route_mode = initial_route_mode;
let mut ready_observed = initial_ready;
src/maestro/generation.rs
+221
View File
@@ -0,0 +1,221 @@
use std::future::Future;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use tokio::sync::{RwLock, Semaphore, watch};
use tokio_util::sync::CancellationToken;
use tokio_util::task::TaskTracker;
use crate::config::ProxyConfig;
use crate::crypto::SecureRandom;
use crate::ip_tracker::UserIpTracker;
use crate::proxy::route_mode::RouteRuntimeController;
use crate::proxy::shared_state::ProxySharedState;
use crate::stats::beobachten::BeobachtenStore;
use crate::stats::{ReplayChecker, Stats};
use crate::stream::BufferPool;
use crate::tls_front::TlsFrontCache;
use crate::transport::UpstreamManager;
use crate::transport::middle_proxy::MePool;
const SESSION_STOP_TIMEOUT: Duration = Duration::from_secs(5);
const BACKGROUND_STOP_TIMEOUT: Duration = Duration::from_secs(5);
/// Cancellation and join ownership for one generation's background tasks.
#[derive(Clone)]
pub(crate) struct RuntimeTaskScope {
tracker: TaskTracker,
cancel: CancellationToken,
}
impl RuntimeTaskScope {
pub(crate) fn new() -> Self {
Self {
tracker: TaskTracker::new(),
cancel: CancellationToken::new(),
}
}
pub(crate) fn spawn<F>(&self, future: F)
where
F: Future<Output = ()> + Send + 'static,
{
let cancel = self.cancel.clone();
self.tracker.spawn(async move {
tokio::select! {
_ = cancel.cancelled() => {}
_ = future => {}
}
});
}
pub(crate) fn cancellation_token(&self) -> CancellationToken {
self.cancel.clone()
}
pub(crate) async fn stop(&self) {
self.cancel.cancel();
self.tracker.close();
let _ = tokio::time::timeout(BACKGROUND_STOP_TIMEOUT, self.tracker.wait()).await;
}
}
/// Runtime-owned data plane and control-plane dependencies for one generation.
pub(crate) struct RuntimeGeneration {
pub(crate) id: u64,
pub(crate) config_rx: watch::Receiver<Arc<ProxyConfig>>,
pub(crate) admission_rx: watch::Receiver<bool>,
pub(crate) stats: Arc<Stats>,
pub(crate) upstream_manager: Arc<UpstreamManager>,
pub(crate) replay_checker: Arc<ReplayChecker>,
pub(crate) buffer_pool: Arc<BufferPool>,
pub(crate) rng: Arc<SecureRandom>,
pub(crate) me_pool: Option<Arc<MePool>>,
pub(crate) me_pool_runtime: Arc<RwLock<Option<Arc<MePool>>>>,
pub(crate) route_runtime: Arc<RouteRuntimeController>,
pub(crate) tls_cache: Option<Arc<TlsFrontCache>>,
pub(crate) ip_tracker: Arc<UserIpTracker>,
pub(crate) beobachten: Arc<BeobachtenStore>,
pub(crate) proxy_shared: Arc<ProxySharedState>,
pub(crate) max_connections: Arc<Semaphore>,
background_tasks: RuntimeTaskScope,
sessions: TaskTracker,
session_cancel: CancellationToken,
accepting_sessions: AtomicBool,
}
impl RuntimeGeneration {
#[allow(clippy::too_many_arguments)]
pub(crate) fn new(
id: u64,
config_rx: watch::Receiver<Arc<ProxyConfig>>,
admission_rx: watch::Receiver<bool>,
stats: Arc<Stats>,
upstream_manager: Arc<UpstreamManager>,
replay_checker: Arc<ReplayChecker>,
buffer_pool: Arc<BufferPool>,
rng: Arc<SecureRandom>,
me_pool: Option<Arc<MePool>>,
me_pool_runtime: Arc<RwLock<Option<Arc<MePool>>>>,
route_runtime: Arc<RouteRuntimeController>,
tls_cache: Option<Arc<TlsFrontCache>>,
ip_tracker: Arc<UserIpTracker>,
beobachten: Arc<BeobachtenStore>,
proxy_shared: Arc<ProxySharedState>,
max_connections: Arc<Semaphore>,
background_tasks: RuntimeTaskScope,
) -> Arc<Self> {
Arc::new(Self {
id,
config_rx,
admission_rx,
stats,
upstream_manager,
replay_checker,
buffer_pool,
rng,
me_pool,
me_pool_runtime,
route_runtime,
tls_cache,
ip_tracker,
beobachten,
proxy_shared,
max_connections,
background_tasks,
sessions: TaskTracker::new(),
session_cancel: CancellationToken::new(),
accepting_sessions: AtomicBool::new(true),
})
}
pub(crate) fn config(&self) -> Arc<ProxyConfig> {
self.config_rx.borrow().clone()
}
pub(crate) async fn current_me_pool(&self) -> Option<Arc<MePool>> {
if let Some(pool) = &self.me_pool {
return Some(pool.clone());
}
self.me_pool_runtime.read().await.clone()
}
pub(crate) fn spawn_session<F>(&self, future: F) -> bool
where
F: Future<Output = ()> + Send + 'static,
{
if !self.accepting_sessions.load(Ordering::Acquire) {
return false;
}
let cancel = self.session_cancel.clone();
self.sessions.spawn(async move {
tokio::select! {
_ = cancel.cancelled() => {}
_ = future => {}
}
});
true
}
pub(crate) fn stop_accepting_sessions(&self) {
self.accepting_sessions.store(false, Ordering::Release);
}
pub(crate) fn resume_accepting_sessions(&self) {
self.accepting_sessions.store(true, Ordering::Release);
}
pub(crate) async fn drain_sessions(&self, timeout: Duration) -> bool {
self.stop_accepting_sessions();
self.sessions.close();
if tokio::time::timeout(timeout, self.sessions.wait())
.await
.is_ok()
{
return true;
}
self.stop_sessions().await;
false
}
pub(crate) async fn stop_sessions(&self) {
self.stop_accepting_sessions();
self.session_cancel.cancel();
self.sessions.close();
let _ = tokio::time::timeout(SESSION_STOP_TIMEOUT, self.sessions.wait()).await;
}
pub(crate) async fn stop_background_tasks(&self) {
self.background_tasks.stop().await;
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn stop_sessions_cancels_tracked_future() {
let tracker = TaskTracker::new();
let cancel = CancellationToken::new();
let child_cancel = cancel.clone();
tracker.spawn(async move {
child_cancel.cancelled().await;
});
cancel.cancel();
tracker.close();
tokio::time::timeout(Duration::from_secs(1), tracker.wait())
.await
.unwrap();
}
#[tokio::test]
async fn runtime_task_scope_joins_cancelled_background_task() {
let scope = RuntimeTaskScope::new();
scope.spawn(std::future::pending());
tokio::time::timeout(Duration::from_secs(1), scope.stop())
.await
.unwrap();
}
}
src/maestro/listeners.rs
+43 -196
View File
@@ -3,35 +3,33 @@ use std::net::{IpAddr, SocketAddr};
use std::sync::Arc;
use std::time::Duration;
use arc_swap::ArcSwap;
use tokio::net::TcpListener;
#[cfg(unix)]
use tokio::net::UnixListener;
use tokio::sync::{RwLock, Semaphore, watch};
use tracing::{debug, error, info, warn};
use crate::config::{ProxyConfig, RstOnCloseMode};
use crate::crypto::SecureRandom;
use crate::ip_tracker::UserIpTracker;
use crate::proxy::ClientHandler;
use crate::proxy::route_mode::RouteRuntimeController;
use crate::proxy::shared_state::ProxySharedState;
use crate::startup::{COMPONENT_LISTENERS_BIND, StartupTracker};
use crate::stats::beobachten::BeobachtenStore;
use crate::stats::{ReplayChecker, Stats};
use crate::stream::BufferPool;
use crate::tls_front::TlsFrontCache;
use crate::transport::middle_proxy::MePool;
use crate::transport::socket::set_linger_zero;
use crate::transport::{ListenOptions, UpstreamManager, create_listener, find_listener_processes};
use crate::transport::{ListenOptions, create_listener, find_listener_processes};
use super::generation::RuntimeGeneration;
use super::helpers::{
expected_handshake_close_description, is_expected_handshake_eof, peer_close_description,
print_proxy_links,
};
#[cfg(unix)]
mod unix;
#[cfg(unix)]
pub(crate) use unix::spawn_unix_accept_loop;
pub(crate) struct BoundListeners {
pub(crate) listeners: Vec<(TcpListener, bool)>,
pub(crate) has_unix_listener: bool,
#[cfg(unix)]
pub(crate) unix_listener: Option<UnixListener>,
}
fn listener_port_or_legacy(listener: &crate::config::ListenerConfig, config: &ProxyConfig) -> u16 {
@@ -59,21 +57,6 @@ pub(crate) async fn bind_listeners(
detected_ip_v4: Option<IpAddr>,
detected_ip_v6: Option<IpAddr>,
startup_tracker: &Arc<StartupTracker>,
config_rx: watch::Receiver<Arc<ProxyConfig>>,
admission_rx: watch::Receiver<bool>,
stats: Arc<Stats>,
upstream_manager: Arc<UpstreamManager>,
replay_checker: Arc<ReplayChecker>,
buffer_pool: Arc<BufferPool>,
rng: Arc<SecureRandom>,
me_pool: Option<Arc<MePool>>,
me_pool_runtime: Arc<RwLock<Option<Arc<MePool>>>>,
route_runtime: Arc<RouteRuntimeController>,
tls_cache: Option<Arc<TlsFrontCache>>,
ip_tracker: Arc<UserIpTracker>,
beobachten: Arc<BeobachtenStore>,
shared: Arc<ProxySharedState>,
max_connections: Arc<Semaphore>,
) -> Result<BoundListeners, Box<dyn Error>> {
startup_tracker
.start_component(
@@ -218,7 +201,8 @@ pub(crate) async fn bind_listeners(
print_proxy_links(&host, port, config);
}
let mut has_unix_listener = false;
#[cfg(unix)]
let mut unix_listener_out = None;
#[cfg(unix)]
if let Some(ref unix_path) = config.server.listen_unix_sock {
let _ = tokio::fs::remove_file(unix_path).await;
@@ -251,123 +235,14 @@ pub(crate) async fn bind_listeners(
info!("Listening on unix:{}", unix_path);
}
has_unix_listener = true;
let mut config_rx_unix: watch::Receiver<Arc<ProxyConfig>> = config_rx.clone();
let admission_rx_unix = admission_rx.clone();
let stats = stats.clone();
let upstream_manager = upstream_manager.clone();
let replay_checker = replay_checker.clone();
let buffer_pool = buffer_pool.clone();
let rng = rng.clone();
let me_pool = me_pool.clone();
let me_pool_runtime = me_pool_runtime.clone();
let route_runtime = route_runtime.clone();
let tls_cache = tls_cache.clone();
let ip_tracker = ip_tracker.clone();
let beobachten = beobachten.clone();
let shared = shared.clone();
let max_connections_unix = max_connections.clone();
tokio::spawn(async move {
let unix_conn_counter = Arc::new(std::sync::atomic::AtomicU64::new(1));
loop {
match unix_listener.accept().await {
Ok((stream, _)) => {
if !*admission_rx_unix.borrow() {
drop(stream);
continue;
}
let accept_permit_timeout_ms =
config_rx_unix.borrow().server.accept_permit_timeout_ms;
let permit = if accept_permit_timeout_ms == 0 {
match max_connections_unix.clone().acquire_owned().await {
Ok(permit) => permit,
Err(_) => {
error!("Connection limiter is closed");
break;
}
}
} else {
match tokio::time::timeout(
Duration::from_millis(accept_permit_timeout_ms),
max_connections_unix.clone().acquire_owned(),
)
.await
{
Ok(Ok(permit)) => permit,
Ok(Err(_)) => {
error!("Connection limiter is closed");
break;
}
Err(_) => {
stats.increment_accept_permit_timeout_total();
debug!(
timeout_ms = accept_permit_timeout_ms,
"Dropping accepted unix connection: permit wait timeout"
);
drop(stream);
continue;
}
}
};
let conn_id =
unix_conn_counter.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
let fake_peer =
SocketAddr::from(([127, 0, 0, 1], (conn_id % 65535) as u16));
let config = config_rx_unix.borrow_and_update().clone();
let stats = stats.clone();
let upstream_manager = upstream_manager.clone();
let replay_checker = replay_checker.clone();
let buffer_pool = buffer_pool.clone();
let rng = rng.clone();
let me_pool = me_pool.clone();
let me_pool_runtime = me_pool_runtime.clone();
let route_runtime = route_runtime.clone();
let tls_cache = tls_cache.clone();
let ip_tracker = ip_tracker.clone();
let beobachten = beobachten.clone();
let shared = shared.clone();
let proxy_protocol_enabled = config.server.proxy_protocol;
tokio::spawn(async move {
let _permit = permit;
if let Err(e) =
crate::proxy::client::handle_client_stream_with_shared_and_pool_runtime(
stream,
fake_peer,
config,
stats,
upstream_manager,
replay_checker,
buffer_pool,
rng,
me_pool,
Some(me_pool_runtime),
route_runtime,
tls_cache,
ip_tracker,
beobachten,
shared,
proxy_protocol_enabled,
)
.await
{
debug!(error = %e, "Unix socket connection error");
}
});
}
Err(e) => {
error!("Unix socket accept error: {}", e);
tokio::time::sleep(Duration::from_millis(100)).await;
}
}
}
});
unix_listener_out = Some(unix_listener);
}
#[cfg(unix)]
let has_unix_listener = unix_listener_out.is_some();
#[cfg(not(unix))]
let has_unix_listener = false;
startup_tracker
.complete_component(
COMPONENT_LISTENERS_BIND,
@@ -381,51 +256,25 @@ pub(crate) async fn bind_listeners(
Ok(BoundListeners {
listeners,
has_unix_listener,
#[cfg(unix)]
unix_listener: unix_listener_out,
})
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn spawn_tcp_accept_loops(
listeners: Vec<(TcpListener, bool)>,
config_rx: watch::Receiver<Arc<ProxyConfig>>,
admission_rx: watch::Receiver<bool>,
stats: Arc<Stats>,
upstream_manager: Arc<UpstreamManager>,
replay_checker: Arc<ReplayChecker>,
buffer_pool: Arc<BufferPool>,
rng: Arc<SecureRandom>,
me_pool: Option<Arc<MePool>>,
me_pool_runtime: Arc<RwLock<Option<Arc<MePool>>>>,
route_runtime: Arc<RouteRuntimeController>,
tls_cache: Option<Arc<TlsFrontCache>>,
ip_tracker: Arc<UserIpTracker>,
beobachten: Arc<BeobachtenStore>,
shared: Arc<ProxySharedState>,
max_connections: Arc<Semaphore>,
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
) {
for (listener, listener_proxy_protocol) in listeners {
let mut config_rx: watch::Receiver<Arc<ProxyConfig>> = config_rx.clone();
let admission_rx_tcp = admission_rx.clone();
let stats = stats.clone();
let upstream_manager = upstream_manager.clone();
let replay_checker = replay_checker.clone();
let buffer_pool = buffer_pool.clone();
let rng = rng.clone();
let me_pool = me_pool.clone();
let me_pool_runtime = me_pool_runtime.clone();
let route_runtime = route_runtime.clone();
let tls_cache = tls_cache.clone();
let ip_tracker = ip_tracker.clone();
let beobachten = beobachten.clone();
let shared = shared.clone();
let max_connections_tcp = max_connections.clone();
let active_runtime = active_runtime.clone();
tokio::spawn(async move {
loop {
match listener.accept().await {
Ok((stream, peer_addr)) => {
let rst_mode = config_rx.borrow().general.rst_on_close;
let runtime = active_runtime.load_full();
let config = runtime.config();
let rst_mode = config.general.rst_on_close;
#[cfg(unix)]
let raw_fd = {
use std::os::unix::io::AsRawFd;
@@ -434,15 +283,14 @@ pub(crate) fn spawn_tcp_accept_loops(
if matches!(rst_mode, RstOnCloseMode::Errors | RstOnCloseMode::Always) {
let _ = set_linger_zero(&stream);
}
if !*admission_rx_tcp.borrow() {
if !*runtime.admission_rx.borrow() {
debug!(peer = %peer_addr, "Admission gate closed, dropping connection");
drop(stream);
continue;
}
let accept_permit_timeout_ms =
config_rx.borrow().server.accept_permit_timeout_ms;
let accept_permit_timeout_ms = config.server.accept_permit_timeout_ms;
let permit = if accept_permit_timeout_ms == 0 {
match max_connections_tcp.clone().acquire_owned().await {
match runtime.max_connections.clone().acquire_owned().await {
Ok(permit) => permit,
Err(_) => {
error!("Connection limiter is closed");
@@ -452,7 +300,7 @@ pub(crate) fn spawn_tcp_accept_loops(
} else {
match tokio::time::timeout(
Duration::from_millis(accept_permit_timeout_ms),
max_connections_tcp.clone().acquire_owned(),
runtime.max_connections.clone().acquire_owned(),
)
.await
{
@@ -462,7 +310,7 @@ pub(crate) fn spawn_tcp_accept_loops(
break;
}
Err(_) => {
stats.increment_accept_permit_timeout_total();
runtime.stats.increment_accept_permit_timeout_total();
debug!(
peer = %peer_addr,
timeout_ms = accept_permit_timeout_ms,
@@ -473,24 +321,23 @@ pub(crate) fn spawn_tcp_accept_loops(
}
}
};
let config = config_rx.borrow_and_update().clone();
let stats = stats.clone();
let upstream_manager = upstream_manager.clone();
let replay_checker = replay_checker.clone();
let buffer_pool = buffer_pool.clone();
let rng = rng.clone();
let me_pool = me_pool.clone();
let me_pool_runtime = me_pool_runtime.clone();
let route_runtime = route_runtime.clone();
let tls_cache = tls_cache.clone();
let ip_tracker = ip_tracker.clone();
let beobachten = beobachten.clone();
let shared = shared.clone();
let stats = runtime.stats.clone();
let upstream_manager = runtime.upstream_manager.clone();
let replay_checker = runtime.replay_checker.clone();
let buffer_pool = runtime.buffer_pool.clone();
let rng = runtime.rng.clone();
let me_pool = runtime.me_pool.clone();
let me_pool_runtime = runtime.me_pool_runtime.clone();
let route_runtime = runtime.route_runtime.clone();
let tls_cache = runtime.tls_cache.clone();
let ip_tracker = runtime.ip_tracker.clone();
let beobachten = runtime.beobachten.clone();
let shared = runtime.proxy_shared.clone();
let proxy_protocol_enabled = listener_proxy_protocol;
let real_peer_report = Arc::new(std::sync::Mutex::new(None));
let real_peer_report_for_handler = real_peer_report.clone();
tokio::spawn(async move {
let _ = runtime.spawn_session(async move {
let _permit = permit;
if let Err(e) = ClientHandler::new_with_shared(
stream,
src/maestro/listeners/unix.rs
+117
View File
@@ -0,0 +1,117 @@
use std::net::SocketAddr;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
use arc_swap::ArcSwap;
use tokio::net::UnixListener;
use tracing::{debug, error};
use super::RuntimeGeneration;
pub(crate) fn spawn_unix_accept_loop(
listener: Option<UnixListener>,
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
) {
let Some(listener) = listener else {
return;
};
tokio::spawn(async move {
let connection_counter = AtomicU64::new(1);
loop {
match listener.accept().await {
Ok((stream, _)) => {
let runtime = active_runtime.load_full();
if !*runtime.admission_rx.borrow() {
drop(stream);
continue;
}
let config = runtime.config();
let timeout_ms = config.server.accept_permit_timeout_ms;
let permit = if timeout_ms == 0 {
match runtime.max_connections.clone().acquire_owned().await {
Ok(permit) => permit,
Err(_) => {
error!("Connection limiter is closed");
break;
}
}
} else {
match tokio::time::timeout(
Duration::from_millis(timeout_ms),
runtime.max_connections.clone().acquire_owned(),
)
.await
{
Ok(Ok(permit)) => permit,
Ok(Err(_)) => {
error!("Connection limiter is closed");
break;
}
Err(_) => {
runtime.stats.increment_accept_permit_timeout_total();
debug!(
timeout_ms,
"Dropping accepted unix connection: permit wait timeout"
);
drop(stream);
continue;
}
}
};
let connection_id = connection_counter.fetch_add(1, Ordering::Relaxed);
let fake_peer =
SocketAddr::from(([127, 0, 0, 1], (connection_id % 65535) as u16));
let stats = runtime.stats.clone();
let upstream_manager = runtime.upstream_manager.clone();
let replay_checker = runtime.replay_checker.clone();
let buffer_pool = runtime.buffer_pool.clone();
let rng = runtime.rng.clone();
let me_pool = runtime.me_pool.clone();
let me_pool_runtime = runtime.me_pool_runtime.clone();
let route_runtime = runtime.route_runtime.clone();
let tls_cache = runtime.tls_cache.clone();
let ip_tracker = runtime.ip_tracker.clone();
let beobachten = runtime.beobachten.clone();
let shared = runtime.proxy_shared.clone();
let proxy_protocol_enabled = config.server.proxy_protocol;
let _ = runtime.spawn_session(async move {
let _permit = permit;
if let Err(error) =
crate::proxy::client::handle_client_stream_with_shared_and_pool_runtime(
stream,
fake_peer,
config,
stats,
upstream_manager,
replay_checker,
buffer_pool,
rng,
me_pool,
Some(me_pool_runtime),
route_runtime,
tls_cache,
ip_tracker,
beobachten,
shared,
proxy_protocol_enabled,
)
.await
{
debug!(error = %error, "Unix socket connection error");
}
});
}
Err(error) => {
error!(error = %error, "Unix socket accept error");
tokio::time::sleep(Duration::from_millis(100)).await;
}
}
}
});
}
src/maestro/mod.rs
+88 -68
View File
@@ -13,19 +13,24 @@
// - shutdown: graceful shutdown sequence and uptime logging.
mod admission;
mod connectivity;
pub(crate) mod generation;
mod helpers;
mod listeners;
mod me_startup;
pub(crate) mod reload;
mod reload_supervisor;
pub(crate) mod runtime_build;
mod runtime_tasks;
mod shutdown;
mod tls_bootstrap;
use arc_swap::ArcSwap;
use std::net::{IpAddr, SocketAddr};
use std::sync::Arc;
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use tokio::sync::{RwLock, Semaphore, watch};
use tracing::{error, info, warn};
use tracing_subscriber::{EnvFilter, fmt, prelude::*, reload};
use tracing_subscriber::{EnvFilter, fmt, prelude::*, reload as tracing_reload};
use crate::api;
use crate::config::{LogLevel, ProxyConfig};
@@ -34,7 +39,7 @@ use crate::crypto::SecureRandom;
use crate::ip_tracker::UserIpTracker;
use crate::network::probe::{decide_network_capabilities, log_probe_result, run_probe};
use crate::proxy::direct_buffer_budget::{
DirectBufferBudget, resolve_direct_buffer_hard_limit, spawn_direct_buffer_budget_controller,
DirectBufferBudget, resolve_direct_buffer_hard_limit, run_direct_buffer_budget_controller,
};
use crate::proxy::route_mode::{RelayRouteMode, RouteRuntimeController};
use crate::proxy::shared_state::ProxySharedState;
@@ -46,7 +51,7 @@ use crate::startup::{
};
use crate::stats::beobachten::BeobachtenStore;
use crate::stats::telemetry::TelemetryPolicy;
use crate::stats::{ReplayChecker, Stats};
use crate::stats::{QuotaStore, ReplayChecker, Stats};
use crate::stream::BufferPool;
use crate::synlimit_control;
use crate::transport::UpstreamManager;
@@ -343,7 +348,7 @@ async fn run_telemt_core(
}
};
let (filter_layer, filter_handle) =
reload::Layer::new(EnvFilter::new(initial_filter_spec.clone()));
tracing_reload::Layer::new(EnvFilter::new(initial_filter_spec.clone()));
startup_tracker
.start_component(
COMPONENT_TRACING_INIT,
@@ -387,6 +392,7 @@ async fn run_telemt_core(
_logging_guard = Some(guard);
}
}
let runtime_log_filter = runtime_tasks::RuntimeLogFilter::new(filter_handle);
startup_tracker
.complete_component(
@@ -433,21 +439,26 @@ async fn run_telemt_core(
warn!("Using default tls_domain. Consider setting a custom domain.");
}
let stats = Arc::new(Stats::new());
let quota_store = Arc::new(QuotaStore::default());
let stats = Arc::new(Stats::with_quota_store(quota_store.clone()));
let runtime_task_scope = generation::RuntimeTaskScope::new();
stats.apply_telemetry_policy(TelemetryPolicy::from_config(&config.general.telemetry));
let quota_state_path = config.general.quota_state_path.clone();
crate::quota_state::load_quota_state(&quota_state_path, stats.as_ref()).await;
let upstream_manager = Arc::new(UpstreamManager::new(
config.upstreams.clone(),
config.general.upstream_connect_retry_attempts,
config.general.upstream_connect_retry_backoff_ms,
config.general.upstream_connect_budget_ms,
config.general.tg_connect,
config.general.upstream_unhealthy_fail_threshold,
config.general.upstream_connect_failfast_hard_errors,
stats.clone(),
));
let upstream_manager = Arc::new(
UpstreamManager::new(
config.upstreams.clone(),
config.general.upstream_connect_retry_attempts,
config.general.upstream_connect_retry_backoff_ms,
config.general.upstream_connect_budget_ms,
config.general.tg_connect,
config.general.upstream_unhealthy_fail_threshold,
config.general.upstream_connect_failfast_hard_errors,
stats.clone(),
)
.with_dns_overrides(&config.network.dns_overrides)?,
);
let ip_tracker = Arc::new(UserIpTracker::new());
ip_tracker
.load_limits(
@@ -497,6 +508,9 @@ async fn run_telemt_core(
let initial_direct_first = config.general.use_middle_proxy && config.general.me2dc_fallback;
let initial_admission_open = !config.general.use_middle_proxy || initial_direct_first;
let (admission_tx, admission_rx) = watch::channel(initial_admission_open);
let (reload_control, reload_commands) = reload::ReloadControl::channel(1);
let (active_runtime_tx, active_runtime_rx) =
watch::channel(None::<Arc<ArcSwap<generation::RuntimeGeneration>>>);
let initial_route_mode = if !config.general.use_middle_proxy || initial_direct_first {
RelayRouteMode::Direct
} else {
@@ -536,6 +550,8 @@ async fn run_telemt_core(
let quota_state_path_api = quota_state_path.clone();
let startup_tracker_api = startup_tracker.clone();
let detected_ips_rx_api = detected_ips_rx.clone();
let reload_control_api = reload_control.clone();
let active_runtime_rx_api = active_runtime_rx.clone();
tokio::spawn(async move {
api::serve(
listen,
@@ -552,6 +568,8 @@ async fn run_telemt_core(
detected_ips_rx_api,
process_started_at_epoch_secs,
startup_tracker_api,
reload_control_api,
active_runtime_rx_api,
)
.await;
});
@@ -591,6 +609,7 @@ async fn run_telemt_core(
&tls_domains,
upstream_manager.clone(),
&startup_tracker,
runtime_task_scope.clone(),
)
.await;
@@ -809,12 +828,19 @@ async fn run_telemt_core(
me_pool.clone(),
shared_state.clone(),
me_ready_tx.clone(),
runtime_task_scope.clone(),
)
.await;
let config_rx = runtime_watches.config_rx;
let log_level_rx = runtime_watches.log_level_rx;
let detected_ip_v4 = runtime_watches.detected_ip_v4;
let detected_ip_v6 = runtime_watches.detected_ip_v6;
runtime_log_filter.start(
has_rust_log,
&effective_log_level,
log_level_rx,
runtime_task_scope.clone(),
);
if direct_first_startup {
let config_bg = config.clone();
@@ -827,7 +853,8 @@ async fn run_telemt_core(
let api_me_pool_bg = api_me_pool.clone();
let me_ready_tx_bg = me_ready_tx.clone();
let config_rx_bg = config_rx.clone();
tokio::spawn(async move {
let task_scope_bg = runtime_task_scope.clone();
runtime_task_scope.spawn(async move {
let mut bootstrap_attempt: u32 = 0;
loop {
bootstrap_attempt = bootstrap_attempt.saturating_add(1);
@@ -851,6 +878,7 @@ async fn run_telemt_core(
pool,
rng_bg,
me_ready_tx_bg,
task_scope_bg,
);
break;
}
@@ -864,7 +892,7 @@ async fn run_telemt_core(
let startup_tracker_ready = startup_tracker.clone();
let api_me_pool_ready = api_me_pool.clone();
let mut me_ready_rx_transport = me_ready_tx.subscribe();
tokio::spawn(async move {
runtime_task_scope.spawn(async move {
if me_ready_rx_transport.changed().await.is_ok() {
if let Some(pool) = api_me_pool_ready.read().await.as_ref() {
pool.set_runtime_ready(true);
@@ -886,29 +914,27 @@ async fn run_telemt_core(
&admission_tx,
config_rx.clone(),
me_ready_rx,
runtime_task_scope.clone(),
)
.await;
let _admission_tx_hold = admission_tx;
conntrack_control::spawn_conntrack_controller(
let conntrack_scope = runtime_task_scope.clone();
runtime_task_scope.spawn(conntrack_control::run_conntrack_controller(
config_rx.clone(),
stats.clone(),
shared_state.clone(),
);
spawn_direct_buffer_budget_controller(
conntrack_scope.cancellation_token(),
));
runtime_task_scope.spawn(run_direct_buffer_budget_controller(
direct_buffer_budget,
buffer_pool.clone(),
stats.clone(),
shared_state.clone(),
config.server.max_connections,
);
));
let bound = listeners::bind_listeners(
&config,
decision.ipv4_dc,
decision.ipv6_dc,
detected_ip_v4,
detected_ip_v6,
&startup_tracker,
let runtime_generation = generation::RuntimeGeneration::new(
1,
config_rx.clone(),
admission_rx.clone(),
stats.clone(),
@@ -924,10 +950,35 @@ async fn run_telemt_core(
beobachten.clone(),
shared_state.clone(),
max_connections.clone(),
runtime_task_scope.clone(),
);
let active_runtime = Arc::new(ArcSwap::from(runtime_generation));
reload_supervisor::ReloadSupervisor::spawn(
active_runtime.clone(),
reload_control,
reload_commands,
config_path.clone(),
quota_store,
detected_ips_tx,
runtime_log_filter,
);
let bound = listeners::bind_listeners(
&config,
decision.ipv4_dc,
decision.ipv6_dc,
detected_ip_v4,
detected_ip_v6,
&startup_tracker,
)
.await?;
let listeners = bound.listeners;
let has_unix_listener = bound.has_unix_listener;
#[cfg(unix)]
let unix_listener = bound.unix_listener;
#[cfg(unix)]
let has_unix_listener = unix_listener.is_some();
#[cfg(not(unix))]
let has_unix_listener = false;
if listeners.is_empty() && !has_unix_listener {
error!("No listeners. Exiting.");
@@ -940,51 +991,20 @@ async fn run_telemt_core(
synlimit_control::reconcile_synlimit_rules(&config).await;
synlimit_control::spawn_synlimit_controller(config_rx.clone());
runtime_tasks::apply_runtime_log_filter(
has_rust_log,
&effective_log_level,
filter_handle,
log_level_rx,
)
.await;
runtime_tasks::spawn_metrics_if_configured(
&config,
&startup_tracker,
stats.clone(),
beobachten.clone(),
shared_state.clone(),
ip_tracker.clone(),
tls_cache.clone(),
config_rx.clone(),
)
.await;
runtime_tasks::spawn_metrics_if_configured(&config, &startup_tracker, active_runtime.clone())
.await;
active_runtime_tx.send_replace(Some(active_runtime.clone()));
runtime_tasks::mark_runtime_ready(&startup_tracker).await;
// Spawn signal handlers for SIGUSR1/SIGUSR2 (non-shutdown signals)
shutdown::spawn_signal_handlers(stats.clone(), process_started_at);
shutdown::spawn_signal_handlers(active_runtime.clone(), process_started_at);
listeners::spawn_tcp_accept_loops(
listeners,
config_rx.clone(),
admission_rx.clone(),
stats.clone(),
upstream_manager.clone(),
replay_checker.clone(),
buffer_pool.clone(),
rng.clone(),
me_pool.clone(),
api_me_pool.clone(),
route_runtime.clone(),
tls_cache.clone(),
ip_tracker.clone(),
beobachten.clone(),
shared_state,
max_connections.clone(),
);
listeners::spawn_tcp_accept_loops(listeners, active_runtime.clone());
#[cfg(unix)]
listeners::spawn_unix_accept_loop(unix_listener, active_runtime.clone());
shutdown::wait_for_shutdown(process_started_at, me_pool, stats, quota_state_path).await;
shutdown::wait_for_shutdown(process_started_at, active_runtime, quota_state_path).await;
Ok(())
}
src/maestro/reload.rs
+464
View File
@@ -0,0 +1,464 @@
use std::collections::VecDeque;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};
use serde::{Deserialize, Serialize};
use tokio::sync::{Mutex, mpsc};
use crate::config::ProxyConfig;
const RELOAD_HISTORY_CAPACITY: usize = 32;
const RELOAD_COMMAND_CAPACITY: usize = 1;
const MAX_DRAIN_TIMEOUT_SECS: u64 = 3_600;
/// Session handling policy for an in-process runtime reload.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Deserialize, Serialize)]
#[serde(rename_all = "snake_case")]
pub(crate) enum ReloadMode {
#[default]
Instant,
Drain,
}
/// Failure policy applied during the activation barrier.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Deserialize, Serialize)]
#[serde(rename_all = "snake_case")]
pub(crate) enum ReloadFailurePolicy {
#[default]
KeepNew,
Rollback,
}
/// Request body accepted by the maestro reload endpoint.
#[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize, Serialize)]
#[serde(deny_unknown_fields)]
pub(crate) struct ReloadRequest {
#[serde(default)]
pub(crate) mode: ReloadMode,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) timeout_secs: Option<u64>,
#[serde(default)]
pub(crate) failure_policy: ReloadFailurePolicy,
}
impl ReloadRequest {
pub(crate) fn validate(&self) -> Result<(), &'static str> {
match (self.mode, self.timeout_secs) {
(ReloadMode::Instant, None) => Ok(()),
(ReloadMode::Instant, Some(_)) => Err("timeout_secs is only valid when mode is drain"),
(ReloadMode::Drain, Some(1..=MAX_DRAIN_TIMEOUT_SECS)) => Ok(()),
(ReloadMode::Drain, Some(_)) => Err("timeout_secs must be within 1..=3600"),
(ReloadMode::Drain, None) => Err("timeout_secs is required when mode is drain"),
}
}
pub(crate) fn from_query(query: Option<&str>) -> Result<Option<Self>, String> {
let Some(query) = query.filter(|query| !query.is_empty()) else {
return Ok(None);
};
let mut mode = None;
let mut timeout_secs = None;
let mut failure_policy = None;
for (key, value) in url::form_urlencoded::parse(query.as_bytes()) {
match key.as_ref() {
"reload" if mode.is_none() => {
mode = Some(match value.as_ref() {
"instant" => ReloadMode::Instant,
"drain" => ReloadMode::Drain,
_ => return Err("reload must be instant or drain".to_string()),
});
}
"timeout_secs" if timeout_secs.is_none() => {
timeout_secs = Some(
value
.parse::<u64>()
.map_err(|_| "timeout_secs must be an integer".to_string())?,
);
}
"failure_policy" if failure_policy.is_none() => {
failure_policy = Some(match value.as_ref() {
"keep_new" => ReloadFailurePolicy::KeepNew,
"rollback" => ReloadFailurePolicy::Rollback,
_ => {
return Err("failure_policy must be keep_new or rollback".to_string());
}
});
}
"reload" | "timeout_secs" | "failure_policy" => {
return Err(format!("duplicate query parameter: {}", key));
}
_ => return Err(format!("unknown query parameter: {}", key)),
}
}
let mode = mode.ok_or_else(|| "reload query parameter is required".to_string())?;
let request = Self {
mode,
timeout_secs,
failure_policy: failure_policy.unwrap_or_default(),
};
request.validate().map_err(str::to_string)?;
Ok(Some(request))
}
}
/// Observable phase of one reload operation.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
#[serde(rename_all = "snake_case")]
pub(crate) enum ReloadPhase {
Accepted,
Preparing,
Activating,
Draining,
Succeeded,
RolledBack,
Failed,
}
impl ReloadPhase {
fn is_terminal(self) -> bool {
matches!(
self,
ReloadPhase::Succeeded | ReloadPhase::RolledBack | ReloadPhase::Failed
)
}
}
/// Bounded public status for one reload operation.
#[derive(Debug, Clone, Serialize)]
pub(crate) struct ReloadStatus {
pub(crate) reload_id: u64,
pub(crate) target_generation: u64,
pub(crate) config_revision: String,
pub(crate) state: ReloadPhase,
pub(crate) mode: ReloadMode,
pub(crate) failure_policy: ReloadFailurePolicy,
pub(crate) requested_at_epoch_secs: u64,
#[serde(skip_serializing_if = "Option::is_none")]
pub(crate) started_at_epoch_secs: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub(crate) finished_at_epoch_secs: Option<u64>,
#[serde(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>,
#[serde(skip_serializing_if = "Option::is_none")]
pub(crate) error: Option<String>,
}
/// Accepted operation metadata returned before asynchronous preparation starts.
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub(crate) struct ReloadAccepted {
pub(crate) reload_id: u64,
pub(crate) target_generation: u64,
pub(crate) config_revision: String,
pub(crate) state: ReloadPhase,
pub(crate) mode: ReloadMode,
pub(crate) failure_policy: ReloadFailurePolicy,
}
pub(crate) struct ReloadCommand {
pub(crate) reload_id: u64,
pub(crate) target_generation: u64,
pub(crate) config: Arc<ProxyConfig>,
pub(crate) config_revision: String,
pub(crate) request: ReloadRequest,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum ReloadSubmitError {
InProgress(u64),
MaestroUnavailable,
}
#[derive(Clone)]
pub(crate) struct ReloadControl {
command_tx: mpsc::Sender<ReloadCommand>,
status_store: Arc<ReloadStatusStore>,
active_generation: Arc<AtomicU64>,
}
pub(crate) struct ReloadCommandReceiver {
command_rx: mpsc::Receiver<ReloadCommand>,
}
#[derive(Default)]
struct ReloadStatusState {
next_reload_id: u64,
active_reload_id: Option<u64>,
statuses: VecDeque<ReloadStatus>,
}
#[derive(Default)]
struct ReloadStatusStore {
state: Mutex<ReloadStatusState>,
}
impl ReloadControl {
pub(crate) fn channel(initial_generation: u64) -> (Self, ReloadCommandReceiver) {
let (command_tx, command_rx) = mpsc::channel(RELOAD_COMMAND_CAPACITY);
(
Self {
command_tx,
status_store: Arc::new(ReloadStatusStore::default()),
active_generation: Arc::new(AtomicU64::new(initial_generation)),
},
ReloadCommandReceiver { command_rx },
)
}
pub(crate) async fn submit(
&self,
config: Arc<ProxyConfig>,
config_revision: String,
request: ReloadRequest,
) -> Result<ReloadAccepted, ReloadSubmitError> {
let target_generation = self
.active_generation
.load(Ordering::Acquire)
.saturating_add(1);
let status = self
.status_store
.reserve(target_generation, config_revision, request.clone())
.await?;
let command = ReloadCommand {
reload_id: status.reload_id,
target_generation,
config,
config_revision: status.config_revision.clone(),
request,
};
if self.command_tx.send(command).await.is_err() {
self.status_store
.finish(
status.reload_id,
ReloadPhase::Failed,
Some("maestro command channel is closed".to_string()),
)
.await;
return Err(ReloadSubmitError::MaestroUnavailable);
}
Ok(ReloadAccepted {
reload_id: status.reload_id,
target_generation,
config_revision: status.config_revision,
state: ReloadPhase::Accepted,
mode: status.mode,
failure_policy: status.failure_policy,
})
}
pub(crate) async fn status(&self, reload_id: u64) -> Option<ReloadStatus> {
self.status_store.get(reload_id).await
}
pub(crate) async fn in_progress(&self) -> Option<u64> {
self.status_store.state.lock().await.active_reload_id
}
pub(crate) async fn mark_phase(&self, reload_id: u64, phase: ReloadPhase) {
self.status_store.mark_phase(reload_id, phase).await;
}
pub(crate) async fn set_deferred_fields(&self, reload_id: u64, fields: Vec<String>) {
self.status_store
.update(reload_id, |status| status.deferred_fields = fields)
.await;
}
pub(crate) async fn succeed(&self, reload_id: u64, generation: u64) {
self.active_generation.store(generation, Ordering::Release);
self.status_store
.finish(reload_id, ReloadPhase::Succeeded, None)
.await;
}
pub(crate) async fn fail(&self, reload_id: u64, error: impl Into<String>) {
self.status_store
.finish(reload_id, ReloadPhase::Failed, Some(error.into()))
.await;
}
pub(crate) async fn rolled_back(&self, reload_id: u64, error: impl Into<String>) {
self.status_store
.finish(reload_id, ReloadPhase::RolledBack, Some(error.into()))
.await;
}
pub(crate) async fn add_warning(&self, reload_id: u64, warning: impl Into<String>) {
let warning = warning.into();
self.status_store
.update(reload_id, |status| status.warnings.push(warning))
.await;
}
}
impl ReloadCommandReceiver {
pub(crate) async fn recv(&mut self) -> Option<ReloadCommand> {
self.command_rx.recv().await
}
}
impl ReloadStatusStore {
async fn reserve(
&self,
target_generation: u64,
config_revision: String,
request: ReloadRequest,
) -> Result<ReloadStatus, ReloadSubmitError> {
let mut state = self.state.lock().await;
if let Some(reload_id) = state.active_reload_id {
return Err(ReloadSubmitError::InProgress(reload_id));
}
state.next_reload_id = state.next_reload_id.saturating_add(1).max(1);
let reload_id = state.next_reload_id;
let status = ReloadStatus {
reload_id,
target_generation,
config_revision,
state: ReloadPhase::Accepted,
mode: request.mode,
failure_policy: request.failure_policy,
requested_at_epoch_secs: now_epoch_secs(),
started_at_epoch_secs: None,
finished_at_epoch_secs: None,
deferred_fields: Vec::new(),
warnings: Vec::new(),
error: None,
};
state.active_reload_id = Some(reload_id);
state.statuses.push_back(status.clone());
while state.statuses.len() > RELOAD_HISTORY_CAPACITY {
state.statuses.pop_front();
}
Ok(status)
}
async fn get(&self, reload_id: u64) -> Option<ReloadStatus> {
self.state
.lock()
.await
.statuses
.iter()
.find(|status| status.reload_id == reload_id)
.cloned()
}
async fn mark_phase(&self, reload_id: u64, phase: ReloadPhase) {
self.update(reload_id, |status| {
status.state = phase;
if status.started_at_epoch_secs.is_none() && phase != ReloadPhase::Accepted {
status.started_at_epoch_secs = Some(now_epoch_secs());
}
})
.await;
}
async fn finish(&self, reload_id: u64, phase: ReloadPhase, error: Option<String>) {
debug_assert!(phase.is_terminal());
let mut state = self.state.lock().await;
if let Some(status) = state
.statuses
.iter_mut()
.find(|status| status.reload_id == reload_id)
{
status.state = phase;
status.error = error;
status.finished_at_epoch_secs = Some(now_epoch_secs());
}
if state.active_reload_id == Some(reload_id) {
state.active_reload_id = None;
}
}
async fn update(&self, reload_id: u64, update: impl FnOnce(&mut ReloadStatus)) {
let mut state = self.state.lock().await;
if let Some(status) = state
.statuses
.iter_mut()
.find(|status| status.reload_id == reload_id)
{
update(status);
}
}
}
fn now_epoch_secs() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn request_defaults_to_instant_keep_new() {
let request: ReloadRequest = serde_json::from_str("{}").unwrap();
assert_eq!(request, ReloadRequest::default());
assert_eq!(request.validate(), Ok(()));
}
#[test]
fn drain_requires_bounded_timeout() {
let missing = ReloadRequest {
mode: ReloadMode::Drain,
..ReloadRequest::default()
};
assert!(missing.validate().is_err());
let valid = ReloadRequest {
mode: ReloadMode::Drain,
timeout_secs: Some(30),
..ReloadRequest::default()
};
assert_eq!(valid.validate(), Ok(()));
}
#[test]
fn patch_query_parses_reload_policy() {
let request =
ReloadRequest::from_query(Some("reload=drain&timeout_secs=30&failure_policy=rollback"))
.unwrap()
.unwrap();
assert_eq!(request.mode, ReloadMode::Drain);
assert_eq!(request.timeout_secs, Some(30));
assert_eq!(request.failure_policy, ReloadFailurePolicy::Rollback);
assert!(ReloadRequest::from_query(Some("timeout_secs=30")).is_err());
}
#[tokio::test]
async fn coordinator_rejects_concurrent_reload_and_releases_terminal_slot() {
let (control, mut receiver) = ReloadControl::channel(1);
let first = control
.submit(
Arc::new(ProxyConfig::default()),
"rev-1".to_string(),
ReloadRequest::default(),
)
.await
.unwrap();
let _command = receiver.recv().await.unwrap();
let second = control
.submit(
Arc::new(ProxyConfig::default()),
"rev-2".to_string(),
ReloadRequest::default(),
)
.await;
assert_eq!(second, Err(ReloadSubmitError::InProgress(first.reload_id)));
control
.succeed(first.reload_id, first.target_generation)
.await;
let third = control
.submit(
Arc::new(ProxyConfig::default()),
"rev-3".to_string(),
ReloadRequest::default(),
)
.await
.unwrap();
assert_eq!(third.reload_id, first.reload_id + 1);
}
}
src/maestro/reload_supervisor.rs
+197
View File
@@ -0,0 +1,197 @@
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
use arc_swap::ArcSwap;
use tokio::sync::watch;
use tracing::{info, warn};
use crate::stats::QuotaStore;
use super::generation::RuntimeGeneration;
use super::reload::{
ReloadCommand, ReloadCommandReceiver, ReloadControl, ReloadFailurePolicy, ReloadMode,
ReloadPhase,
};
use super::runtime_build::{deferred_process_fields, prepare_runtime};
use super::runtime_tasks::RuntimeLogFilter;
pub(crate) struct ReloadSupervisor {
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
control: ReloadControl,
commands: ReloadCommandReceiver,
config_path: PathBuf,
quota_store: Arc<QuotaStore>,
detected_ips_tx: watch::Sender<(Option<std::net::IpAddr>, Option<std::net::IpAddr>)>,
runtime_log_filter: RuntimeLogFilter,
}
impl ReloadSupervisor {
#[allow(clippy::too_many_arguments)]
pub(crate) fn spawn(
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
control: ReloadControl,
commands: ReloadCommandReceiver,
config_path: PathBuf,
quota_store: Arc<QuotaStore>,
detected_ips_tx: watch::Sender<(Option<std::net::IpAddr>, Option<std::net::IpAddr>)>,
runtime_log_filter: RuntimeLogFilter,
) {
let supervisor = Self {
active_runtime,
control,
commands,
config_path,
quota_store,
detected_ips_tx,
runtime_log_filter,
};
tokio::spawn(supervisor.run());
}
async fn run(mut self) {
while let Some(command) = self.commands.recv().await {
self.reload(command).await;
}
}
async fn reload(&self, command: ReloadCommand) {
self.control
.mark_phase(command.reload_id, ReloadPhase::Preparing)
.await;
let old_runtime = self.active_runtime.load_full();
let deferred = deferred_process_fields(&old_runtime.config(), &command.config);
self.control
.set_deferred_fields(command.reload_id, deferred)
.await;
let prepared = match prepare_runtime(
command.target_generation,
command.config.as_ref().clone(),
&self.config_path,
self.quota_store.clone(),
self.runtime_log_filter.clone(),
)
.await
{
Ok(prepared) => prepared,
Err(error) => {
self.control.fail(command.reload_id, error).await;
return;
}
};
if let Ok(current_revision) =
crate::api::config_store::current_revision_for_maestro(&self.config_path).await
&& current_revision != command.config_revision
{
let warning = format!(
"config revision changed during preparation: accepted={} current={}",
command.config_revision, current_revision
);
if command.request.failure_policy == ReloadFailurePolicy::Rollback {
self.runtime_log_filter
.apply_reload(&old_runtime.config().general.log_level);
prepared.generation.stop_sessions().await;
if let Some(pool) = prepared.generation.current_me_pool().await {
let _ = tokio::time::timeout(
Duration::from_secs(2),
pool.shutdown_send_close_conn_all(),
)
.await;
}
prepared.generation.stop_background_tasks().await;
self.control.rolled_back(command.reload_id, warning).await;
return;
}
self.control.add_warning(command.reload_id, warning).await;
}
self.control
.mark_phase(command.reload_id, ReloadPhase::Activating)
.await;
let new_runtime = prepared.generation;
old_runtime.stop_accepting_sessions();
let replaced = self.active_runtime.swap(new_runtime.clone());
if let Err(error) = crate::network::dns_overrides::install_entries(
&new_runtime.config().network.dns_overrides,
) {
let message = format!("runtime DNS activation failed: {}", error);
if command.request.failure_policy == ReloadFailurePolicy::Rollback {
let candidate = self.active_runtime.swap(replaced.clone());
replaced.resume_accepting_sessions();
self.runtime_log_filter
.apply_reload(&replaced.config().general.log_level);
let _ = crate::network::dns_overrides::install_entries(
&replaced.config().network.dns_overrides,
);
candidate.stop_sessions().await;
if let Some(pool) = candidate.current_me_pool().await {
let _ = tokio::time::timeout(
Duration::from_secs(2),
pool.shutdown_send_close_conn_all(),
)
.await;
}
candidate.stop_background_tasks().await;
self.control.rolled_back(command.reload_id, message).await;
return;
}
self.control.add_warning(command.reload_id, message).await;
}
self.detected_ips_tx.send_replace(prepared.detected_ips);
self.runtime_log_filter
.apply_reload(&new_runtime.config().general.log_level);
crate::synlimit_control::reconcile_synlimit_rules(&new_runtime.config()).await;
info!(
reload_id = command.reload_id,
old_generation = replaced.id,
new_generation = new_runtime.id,
config_revision = %command.config_revision,
"Runtime generation activated"
);
match command.request.mode {
ReloadMode::Instant => {
replaced.stop_sessions().await;
}
ReloadMode::Drain => {
self.control
.mark_phase(command.reload_id, ReloadPhase::Draining)
.await;
let timeout = Duration::from_secs(
command
.request
.timeout_secs
.expect("validated drain request must carry timeout_secs"),
);
if !replaced.drain_sessions(timeout).await {
let warning = format!(
"generation {} exceeded drain timeout; remaining sessions were cancelled",
replaced.id
);
warn!(reload_id = command.reload_id, warning = %warning);
self.control.add_warning(command.reload_id, warning).await;
}
}
}
if let Some(pool) = replaced.current_me_pool().await
&& tokio::time::timeout(Duration::from_secs(2), pool.shutdown_send_close_conn_all())
.await
.is_err()
{
let warning = format!(
"generation {} Middle-End close broadcast timed out",
replaced.id
);
warn!(reload_id = command.reload_id, warning = %warning);
self.control.add_warning(command.reload_id, warning).await;
}
replaced.stop_background_tasks().await;
self.control
.succeed(command.reload_id, new_runtime.id)
.await;
}
}
src/maestro/runtime_build.rs
+357
View File
@@ -0,0 +1,357 @@
use std::net::IpAddr;
use std::path::Path;
use std::sync::Arc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use tokio::sync::{RwLock, Semaphore, watch};
use crate::config::ProxyConfig;
use crate::crypto::SecureRandom;
use crate::ip_tracker::UserIpTracker;
use crate::network::probe::{decide_network_capabilities, run_probe};
use crate::proxy::direct_buffer_budget::{
DirectBufferBudget, resolve_direct_buffer_hard_limit, run_direct_buffer_budget_controller,
};
use crate::proxy::route_mode::{RelayRouteMode, RouteRuntimeController};
use crate::proxy::shared_state::ProxySharedState;
use crate::startup::StartupTracker;
use crate::stats::beobachten::BeobachtenStore;
use crate::stats::telemetry::TelemetryPolicy;
use crate::stats::{QuotaStore, ReplayChecker, Stats};
use crate::stream::BufferPool;
use crate::transport::UpstreamManager;
use crate::transport::middle_proxy::MePool;
use super::admission;
use super::generation::{RuntimeGeneration, RuntimeTaskScope};
use super::runtime_tasks::RuntimeLogFilter;
use super::{me_startup, runtime_tasks, tls_bootstrap};
pub(crate) struct PreparedRuntime {
pub(crate) generation: Arc<RuntimeGeneration>,
pub(crate) detected_ips: (Option<IpAddr>, Option<IpAddr>),
}
pub(crate) async fn prepare_runtime(
generation_id: u64,
mut config: ProxyConfig,
config_path: &Path,
quota_store: Arc<QuotaStore>,
runtime_log_filter: RuntimeLogFilter,
) -> Result<PreparedRuntime, String> {
let started_at_epoch_secs = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
let startup_tracker = Arc::new(StartupTracker::new(started_at_epoch_secs));
let task_scope = RuntimeTaskScope::new();
let stats = Arc::new(Stats::with_quota_store(quota_store));
stats.apply_telemetry_policy(TelemetryPolicy::from_config(&config.general.telemetry));
let upstream_manager = Arc::new(
UpstreamManager::new(
config.upstreams.clone(),
config.general.upstream_connect_retry_attempts,
config.general.upstream_connect_retry_backoff_ms,
config.general.upstream_connect_budget_ms,
config.general.tg_connect,
config.general.upstream_unhealthy_fail_threshold,
config.general.upstream_connect_failfast_hard_errors,
stats.clone(),
)
.with_dns_overrides(&config.network.dns_overrides)
.map_err(|error| format!("DNS override preparation failed: {}", error))?,
);
let ip_tracker = Arc::new(UserIpTracker::new());
ip_tracker
.load_limits(
config.access.user_max_unique_ips_global_each,
&config.access.user_max_unique_ips,
)
.await;
ip_tracker
.set_limit_policy(
config.access.user_max_unique_ips_mode,
config.access.user_max_unique_ips_window_secs,
)
.await;
let hard_limit =
resolve_direct_buffer_hard_limit(config.general.direct_relay_buffer_budget_max_bytes).await;
let direct_buffer_budget = DirectBufferBudget::new(hard_limit);
let proxy_shared =
ProxySharedState::new_with_direct_buffer_budget(direct_buffer_budget.clone());
proxy_shared.apply_user_enabled_config(&config.access.user_enabled);
proxy_shared.traffic_limiter.apply_policy(
config.access.user_rate_limits.clone(),
config.access.cidr_rate_limits.clone(),
);
let probe = run_probe(
&config.network,
&config.upstreams,
config.general.middle_proxy_nat_probe,
config.general.stun_nat_probe_concurrency,
)
.await
.map_err(|error| format!("network probe failed: {}", error))?;
let decision =
decide_network_capabilities(&config.network, &probe, config.general.middle_proxy_nat_ip);
let prefer_ipv6 = decision.prefer_ipv6();
let mut tls_domains = Vec::with_capacity(1 + config.censorship.tls_domains.len());
tls_domains.push(config.censorship.tls_domain.clone());
for domain in &config.censorship.tls_domains {
if !tls_domains.contains(domain) {
tls_domains.push(domain.clone());
}
}
let tls_cache = tls_bootstrap::bootstrap_tls_front(
&config,
&tls_domains,
upstream_manager.clone(),
&startup_tracker,
task_scope.clone(),
)
.await;
let beobachten = Arc::new(BeobachtenStore::new());
let rng = Arc::new(SecureRandom::new());
let route_mode = if !config.general.use_middle_proxy || config.general.me2dc_fallback {
RelayRouteMode::Direct
} else {
RelayRouteMode::Middle
};
let route_runtime = Arc::new(RouteRuntimeController::new(route_mode));
let me_pool_runtime = Arc::new(RwLock::new(None::<Arc<MePool>>));
let (me_ready_tx, me_ready_rx) = watch::channel(0_u64);
let direct_first_startup = config.general.use_middle_proxy && config.general.me2dc_fallback;
let me_pool = if direct_first_startup {
None
} else {
me_startup::initialize_me_pool(
config.general.use_middle_proxy,
&config,
&decision,
&probe,
&startup_tracker,
upstream_manager.clone(),
rng.clone(),
stats.clone(),
me_pool_runtime.clone(),
me_ready_tx.clone(),
)
.await
};
if config.general.use_middle_proxy && me_pool.is_none() && !direct_first_startup {
config.general.use_middle_proxy = false;
}
let config = Arc::new(config);
let replay_checker = Arc::new(ReplayChecker::new(
config.access.replay_check_len,
Duration::from_secs(config.access.replay_window_secs),
));
let buffer_pool = Arc::new(BufferPool::with_config(64 * 1024, 4096));
let max_connections_limit = if config.server.max_connections == 0 {
Semaphore::MAX_PERMITS
} else {
config.server.max_connections as usize
};
let max_connections = Arc::new(Semaphore::new(max_connections_limit));
let (api_config_tx, _) = watch::channel(config.clone());
let watches = runtime_tasks::spawn_runtime_tasks(
&config,
config_path,
&probe,
prefer_ipv6,
decision.ipv4_dc,
decision.ipv6_dc,
&startup_tracker,
stats.clone(),
upstream_manager.clone(),
replay_checker.clone(),
me_pool.clone(),
rng.clone(),
ip_tracker.clone(),
beobachten.clone(),
api_config_tx,
me_pool.clone(),
proxy_shared.clone(),
me_ready_tx.clone(),
task_scope.clone(),
)
.await;
let config_rx = watches.config_rx;
runtime_log_filter.spawn_watcher(watches.log_level_rx, task_scope.clone());
let initial_admission_open = !config.general.use_middle_proxy || me_pool.is_some();
let (admission_tx, admission_rx) = watch::channel(initial_admission_open);
admission::configure_admission_gate(
&config,
me_pool.clone(),
me_pool_runtime.clone(),
route_runtime.clone(),
&admission_tx,
config_rx.clone(),
me_ready_rx,
task_scope.clone(),
)
.await;
if direct_first_startup {
let config_bg = config.clone();
let decision_bg = decision.clone();
let probe_bg = probe.clone();
let startup_tracker_bg = startup_tracker.clone();
let upstream_manager_bg = upstream_manager.clone();
let rng_bg = rng.clone();
let stats_bg = stats.clone();
let me_pool_runtime_bg = me_pool_runtime.clone();
let me_ready_tx_bg = me_ready_tx.clone();
let config_rx_bg = config_rx.clone();
let task_scope_bg = task_scope.clone();
let retry_limit = config.general.me_init_retry_attempts;
task_scope.spawn(async move {
let mut attempt = 0_u32;
loop {
attempt = attempt.saturating_add(1);
let pool = me_startup::initialize_me_pool(
true,
config_bg.as_ref(),
&decision_bg,
&probe_bg,
&startup_tracker_bg,
upstream_manager_bg.clone(),
rng_bg.clone(),
stats_bg.clone(),
me_pool_runtime_bg.clone(),
me_ready_tx_bg.clone(),
)
.await;
if let Some(pool) = pool {
runtime_tasks::spawn_middle_proxy_runtime_tasks(
config_bg.as_ref(),
config_rx_bg,
pool,
rng_bg,
me_ready_tx_bg,
task_scope_bg,
);
break;
}
if retry_limit > 0 && attempt >= retry_limit {
break;
}
tokio::time::sleep(Duration::from_secs(2)).await;
}
});
}
let conntrack_scope = task_scope.clone();
task_scope.spawn(crate::conntrack_control::run_conntrack_controller(
config_rx.clone(),
stats.clone(),
proxy_shared.clone(),
conntrack_scope.cancellation_token(),
));
task_scope.spawn(run_direct_buffer_budget_controller(
direct_buffer_budget,
buffer_pool.clone(),
stats.clone(),
proxy_shared.clone(),
config.server.max_connections,
));
let generation = RuntimeGeneration::new(
generation_id,
config_rx,
admission_rx,
stats,
upstream_manager,
replay_checker,
buffer_pool,
rng,
me_pool,
me_pool_runtime,
route_runtime,
tls_cache,
ip_tracker,
beobachten,
proxy_shared,
max_connections,
task_scope,
);
drop(admission_tx);
Ok(PreparedRuntime {
generation,
detected_ips: (
probe.detected_ipv4.map(IpAddr::V4),
probe.detected_ipv6.map(IpAddr::V6),
),
})
}
pub(crate) fn deferred_process_fields(old: &ProxyConfig, new: &ProxyConfig) -> Vec<String> {
let mut fields = Vec::new();
if old.server.port != new.server.port
|| old.server.proxy_protocol != new.server.proxy_protocol
|| old.server.listen_backlog != new.server.listen_backlog
|| serde_json::to_value(&old.server.listeners).ok()
!= serde_json::to_value(&new.server.listeners).ok()
{
fields.push("server.listeners".to_string());
}
if old.server.listen_unix_sock != new.server.listen_unix_sock
|| old.server.listen_unix_sock_perm != new.server.listen_unix_sock_perm
{
fields.push("server.listen_unix_sock".to_string());
}
if old.server.api.listen != new.server.api.listen
|| old.server.api.enabled != new.server.api.enabled
{
fields.push("server.api.listen".to_string());
}
if old.server.metrics_listen != new.server.metrics_listen
|| old.server.metrics_port != new.server.metrics_port
{
fields.push("server.metrics_listen".to_string());
}
if old.general.quota_state_path != new.general.quota_state_path {
fields.push("general.quota_state_path".to_string());
}
if old.general.disable_colors != new.general.disable_colors {
fields.push("general.disable_colors".to_string());
}
if old.general.data_path != new.general.data_path {
fields.push("general.data_path".to_string());
}
if serde_json::to_value(&old.logging).ok() != serde_json::to_value(&new.logging).ok() {
fields.push("logging".to_string());
}
fields
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn process_socket_and_logging_changes_are_deferred() {
let old = ProxyConfig::default();
let mut new = old.clone();
new.server.listen_backlog = new.server.listen_backlog.saturating_add(1);
new.general.disable_colors = !new.general.disable_colors;
let fields = deferred_process_fields(&old, &new);
assert!(fields.contains(&"server.listeners".to_string()));
assert!(fields.contains(&"general.disable_colors".to_string()));
}
#[test]
fn runtime_only_change_does_not_require_process_rebind() {
let old = ProxyConfig::default();
let mut new = old.clone();
new.censorship.tls_domain = "reload.example".to_string();
assert!(deferred_process_fields(&old, &new).is_empty());
}
}
src/maestro/runtime_tasks.rs
+85 -66
View File
@@ -2,9 +2,11 @@ use std::net::IpAddr;
use std::path::Path;
use std::sync::Arc;
use arc_swap::ArcSwap;
use tokio::sync::{mpsc, watch};
use tracing::{debug, info, warn};
use tracing_subscriber::EnvFilter;
use tracing_subscriber::Registry;
use tracing_subscriber::reload;
use crate::config::hot_reload::spawn_config_watcher;
@@ -21,10 +23,11 @@ use crate::startup::{
use crate::stats::beobachten::BeobachtenStore;
use crate::stats::telemetry::TelemetryPolicy;
use crate::stats::{ReplayChecker, Stats};
use crate::tls_front::TlsFrontCache;
use crate::transport::UpstreamManager;
use crate::transport::middle_proxy::{MePool, MeReinitTrigger};
use super::generation::RuntimeGeneration;
use super::generation::RuntimeTaskScope;
use super::helpers::write_beobachten_snapshot;
pub(crate) struct RuntimeWatches {
@@ -34,6 +37,56 @@ pub(crate) struct RuntimeWatches {
pub(crate) detected_ip_v6: Option<IpAddr>,
}
#[derive(Clone)]
pub(crate) struct RuntimeLogFilter {
handle: reload::Handle<EnvFilter, Registry>,
}
impl RuntimeLogFilter {
pub(crate) fn new(handle: reload::Handle<EnvFilter, Registry>) -> Self {
Self { handle }
}
pub(crate) fn start(
&self,
has_rust_log: bool,
effective_log_level: &LogLevel,
log_level_rx: watch::Receiver<LogLevel>,
task_scope: RuntimeTaskScope,
) {
self.apply(effective_log_level, has_rust_log);
self.spawn_watcher(log_level_rx, task_scope);
}
pub(crate) fn apply_reload(&self, level: &LogLevel) {
self.apply(level, false);
}
pub(crate) fn spawn_watcher(
&self,
mut log_level_rx: watch::Receiver<LogLevel>,
task_scope: RuntimeTaskScope,
) {
let filter = self.clone();
task_scope.spawn(async move {
loop {
if log_level_rx.changed().await.is_err() {
break;
}
let level = log_level_rx.borrow_and_update().clone();
filter.apply_reload(&level);
}
});
}
fn apply(&self, level: &LogLevel, has_rust_log: bool) {
let runtime_filter = EnvFilter::new(log_filter_spec(has_rust_log, level));
if let Err(error) = self.handle.reload(runtime_filter) {
tracing::error!(error = %error, "Failed to update runtime log filter");
}
}
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn spawn_runtime_tasks(
config: &Arc<ProxyConfig>,
@@ -54,10 +107,11 @@ pub(crate) async fn spawn_runtime_tasks(
me_pool_for_policy: Option<Arc<MePool>>,
shared_state: Arc<ProxySharedState>,
me_ready_tx: watch::Sender<u64>,
task_scope: RuntimeTaskScope,
) -> RuntimeWatches {
let um_clone = upstream_manager.clone();
let dc_overrides_for_health = config.dc_overrides.clone();
tokio::spawn(async move {
task_scope.spawn(async move {
um_clone
.run_health_checks(
prefer_ipv6,
@@ -69,19 +123,19 @@ pub(crate) async fn spawn_runtime_tasks(
});
let rc_clone = replay_checker.clone();
tokio::spawn(async move {
task_scope.spawn(async move {
rc_clone.run_periodic_cleanup().await;
});
let stats_maintenance = stats.clone();
tokio::spawn(async move {
task_scope.spawn(async move {
stats_maintenance
.run_periodic_user_stats_maintenance()
.await;
});
let ip_tracker_maintenance = ip_tracker.clone();
tokio::spawn(async move {
task_scope.spawn(async move {
ip_tracker_maintenance.run_periodic_maintenance().await;
});
@@ -104,6 +158,7 @@ pub(crate) async fn spawn_runtime_tasks(
config.clone(),
detected_ip_v4,
detected_ip_v6,
task_scope.cancellation_token(),
);
startup_tracker
.complete_component(
@@ -113,7 +168,7 @@ pub(crate) async fn spawn_runtime_tasks(
.await;
let mut config_rx_api_bridge = config_rx.clone();
let api_config_tx_bridge = api_config_tx.clone();
tokio::spawn(async move {
task_scope.spawn(async move {
loop {
if config_rx_api_bridge.changed().await.is_err() {
break;
@@ -124,8 +179,9 @@ pub(crate) async fn spawn_runtime_tasks(
});
let stats_policy = stats.clone();
let upstream_policy = upstream_manager.clone();
let mut config_rx_policy = config_rx.clone();
tokio::spawn(async move {
task_scope.spawn(async move {
loop {
if config_rx_policy.changed().await.is_err() {
break;
@@ -133,6 +189,9 @@ pub(crate) async fn spawn_runtime_tasks(
let cfg = config_rx_policy.borrow_and_update().clone();
stats_policy
.apply_telemetry_policy(TelemetryPolicy::from_config(&cfg.general.telemetry));
if let Err(error) = upstream_policy.update_dns_overrides(&cfg.network.dns_overrides) {
warn!(error = %error, "Failed to update generation DNS overrides");
}
if let Some(pool) = &me_pool_for_policy {
pool.update_runtime_transport_policy(
cfg.general.me_socks_kdf_policy,
@@ -149,7 +208,7 @@ pub(crate) async fn spawn_runtime_tasks(
let ip_tracker_policy = ip_tracker.clone();
let mut config_rx_ip_limits = config_rx.clone();
tokio::spawn(async move {
task_scope.spawn(async move {
let mut prev_limits = config_rx_ip_limits
.borrow()
.access
@@ -205,7 +264,7 @@ pub(crate) async fn spawn_runtime_tasks(
config.access.cidr_rate_limits.clone(),
);
let mut config_rx_rate_limits = config_rx.clone();
tokio::spawn(async move {
task_scope.spawn(async move {
let mut prev_user_limits = config_rx_rate_limits
.borrow()
.access
@@ -236,7 +295,7 @@ pub(crate) async fn spawn_runtime_tasks(
let shared_user_enabled = shared_state.clone();
let mut config_rx_user_enabled = config_rx.clone();
tokio::spawn(async move {
task_scope.spawn(async move {
loop {
if config_rx_user_enabled.changed().await.is_err() {
break;
@@ -257,7 +316,7 @@ pub(crate) async fn spawn_runtime_tasks(
let beobachten_writer = beobachten.clone();
let config_rx_beobachten = config_rx.clone();
tokio::spawn(async move {
task_scope.spawn(async move {
loop {
let cfg = config_rx_beobachten.borrow().clone();
let sleep_secs = cfg.general.beobachten_flush_secs.max(1);
@@ -278,7 +337,14 @@ pub(crate) async fn spawn_runtime_tasks(
});
if let Some(pool) = me_pool {
spawn_middle_proxy_runtime_tasks(config, config_rx.clone(), pool, rng, me_ready_tx);
spawn_middle_proxy_runtime_tasks(
config,
config_rx.clone(),
pool,
rng,
me_ready_tx,
task_scope,
);
}
RuntimeWatches {
@@ -295,6 +361,7 @@ pub(crate) fn spawn_middle_proxy_runtime_tasks(
pool: Arc<MePool>,
rng: Arc<SecureRandom>,
me_ready_tx: watch::Sender<u64>,
task_scope: RuntimeTaskScope,
) {
let reinit_trigger_capacity = config.general.me_reinit_trigger_channel.max(1);
let (reinit_tx, reinit_rx) = mpsc::channel::<MeReinitTrigger>(reinit_trigger_capacity);
@@ -303,7 +370,7 @@ pub(crate) fn spawn_middle_proxy_runtime_tasks(
let rng_clone_sched = rng.clone();
let config_rx_clone_sched = config_rx.clone();
let me_ready_tx_sched = me_ready_tx.clone();
tokio::spawn(async move {
task_scope.spawn(async move {
crate::transport::middle_proxy::me_reinit_scheduler(
pool_clone_sched,
rng_clone_sched,
@@ -317,7 +384,7 @@ pub(crate) fn spawn_middle_proxy_runtime_tasks(
let pool_clone = pool.clone();
let config_rx_clone = config_rx.clone();
let reinit_tx_updater = reinit_tx.clone();
tokio::spawn(async move {
task_scope.spawn(async move {
crate::transport::middle_proxy::me_config_updater(
pool_clone,
config_rx_clone,
@@ -328,37 +395,12 @@ pub(crate) fn spawn_middle_proxy_runtime_tasks(
let config_rx_clone_rot = config_rx.clone();
let reinit_tx_rotation = reinit_tx.clone();
tokio::spawn(async move {
task_scope.spawn(async move {
crate::transport::middle_proxy::me_rotation_task(config_rx_clone_rot, reinit_tx_rotation)
.await;
});
}
pub(crate) async fn apply_runtime_log_filter(
has_rust_log: bool,
effective_log_level: &LogLevel,
filter_handle: reload::Handle<EnvFilter, tracing_subscriber::Registry>,
mut log_level_rx: watch::Receiver<LogLevel>,
) {
let runtime_filter = EnvFilter::new(log_filter_spec(has_rust_log, effective_log_level));
filter_handle
.reload(runtime_filter)
.expect("Failed to switch log filter");
tokio::spawn(async move {
loop {
if log_level_rx.changed().await.is_err() {
break;
}
let level = log_level_rx.borrow_and_update().clone();
let new_filter = tracing_subscriber::EnvFilter::new(log_filter_spec(false, &level));
if let Err(e) = filter_handle.reload(new_filter) {
tracing::error!("config reload: failed to update log filter: {}", e);
}
}
});
}
pub(crate) fn log_filter_spec(has_rust_log: bool, effective_log_level: &LogLevel) -> String {
if has_rust_log {
std::env::var("RUST_LOG")
@@ -373,12 +415,7 @@ pub(crate) fn log_filter_spec(has_rust_log: bool, effective_log_level: &LogLevel
pub(crate) async fn spawn_metrics_if_configured(
config: &Arc<ProxyConfig>,
startup_tracker: &Arc<StartupTracker>,
stats: Arc<Stats>,
beobachten: Arc<BeobachtenStore>,
shared_state: Arc<ProxySharedState>,
ip_tracker: Arc<UserIpTracker>,
tls_cache: Option<Arc<TlsFrontCache>>,
config_rx: watch::Receiver<Arc<ProxyConfig>>,
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
) {
// metrics_listen takes precedence; fall back to metrics_port for backward compat.
let metrics_target: Option<(u16, Option<String>)> =
@@ -408,28 +445,10 @@ pub(crate) async fn spawn_metrics_if_configured(
Some(format!("spawn metrics endpoint on {}", label)),
)
.await;
let stats = stats.clone();
let beobachten = beobachten.clone();
let shared_state = shared_state.clone();
let config_rx_metrics = config_rx.clone();
let ip_tracker_metrics = ip_tracker.clone();
let tls_cache_metrics = tls_cache.clone();
let whitelist = config.server.metrics_whitelist.clone();
let active_runtime = active_runtime.clone();
let listen_backlog = config.server.listen_backlog;
tokio::spawn(async move {
metrics::serve(
port,
listen,
listen_backlog,
stats,
beobachten,
shared_state,
ip_tracker_metrics,
tls_cache_metrics,
config_rx_metrics,
whitelist,
)
.await;
metrics::serve(port, listen, listen_backlog, active_runtime).await;
});
startup_tracker
.complete_component(
src/maestro/shutdown.rs
+21 -19
View File
@@ -12,17 +12,17 @@ use std::path::PathBuf;
use std::sync::Arc;
use std::time::{Duration, Instant};
use arc_swap::ArcSwap;
#[cfg(not(unix))]
use tokio::signal;
#[cfg(unix)]
use tokio::signal::unix::{SignalKind, signal};
use tracing::{info, warn};
use super::generation::RuntimeGeneration;
use super::helpers::{format_uptime, unit_label};
use crate::stats::Stats;
use crate::synlimit_control;
use crate::transport::middle_proxy::MePool;
use super::helpers::{format_uptime, unit_label};
/// Signal that triggered shutdown.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -48,19 +48,11 @@ impl std::fmt::Display for ShutdownSignal {
/// Waits for a shutdown signal and performs graceful shutdown.
pub(crate) async fn wait_for_shutdown(
process_started_at: Instant,
me_pool: Option<Arc<MePool>>,
stats: Arc<Stats>,
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
quota_state_path: PathBuf,
) {
let signal = wait_for_shutdown_signal().await;
perform_shutdown(
signal,
process_started_at,
me_pool,
&stats,
quota_state_path,
)
.await;
perform_shutdown(signal, process_started_at, active_runtime, quota_state_path).await;
}
/// Waits for any shutdown signal (SIGINT, SIGTERM, SIGQUIT).
@@ -87,10 +79,11 @@ async fn wait_for_shutdown_signal() -> ShutdownSignal {
async fn perform_shutdown(
signal: ShutdownSignal,
process_started_at: Instant,
me_pool: Option<Arc<MePool>>,
stats: &Stats,
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
quota_state_path: PathBuf,
) {
let runtime = active_runtime.load_full();
let stats = runtime.stats.as_ref();
let shutdown_started_at = Instant::now();
info!(signal = %signal, "Received shutdown signal");
@@ -108,7 +101,8 @@ async fn perform_shutdown(
}
// Graceful ME pool shutdown
if let Some(pool) = &me_pool {
runtime.stop_sessions().await;
if let Some(pool) = runtime.current_me_pool().await {
match tokio::time::timeout(Duration::from_secs(2), pool.shutdown_send_close_conn_all())
.await
{
@@ -123,6 +117,7 @@ async fn perform_shutdown(
}
}
}
runtime.stop_background_tasks().await;
match crate::quota_state::save_quota_state(&quota_state_path, stats).await {
Ok(()) => {
@@ -191,7 +186,10 @@ fn dump_stats(stats: &Stats, process_started_at: Instant) {
/// - SIGUSR1: Log rotation acknowledgment (for external log rotation tools)
/// - SIGUSR2: Dump runtime status to log
#[cfg(unix)]
pub(crate) fn spawn_signal_handlers(stats: Arc<Stats>, process_started_at: Instant) {
pub(crate) fn spawn_signal_handlers(
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
process_started_at: Instant,
) {
tokio::spawn(async move {
let mut sigusr1 =
signal(SignalKind::user_defined1()).expect("Failed to register SIGUSR1 handler");
@@ -204,7 +202,8 @@ pub(crate) fn spawn_signal_handlers(stats: Arc<Stats>, process_started_at: Insta
handle_sigusr1();
}
_ = sigusr2.recv() => {
handle_sigusr2(&stats, process_started_at);
let runtime = active_runtime.load_full();
handle_sigusr2(runtime.stats.as_ref(), process_started_at);
}
}
}
@@ -213,7 +212,10 @@ pub(crate) fn spawn_signal_handlers(stats: Arc<Stats>, process_started_at: Insta
/// No-op on non-Unix platforms.
#[cfg(not(unix))]
pub(crate) fn spawn_signal_handlers(_stats: Arc<Stats>, _process_started_at: Instant) {
pub(crate) fn spawn_signal_handlers(
_active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
_process_started_at: Instant,
) {
// No SIGUSR1/SIGUSR2 on non-Unix
}
src/maestro/tls_bootstrap.rs
+6 -3
View File
@@ -10,6 +10,8 @@ use crate::tls_front::TlsFrontCache;
use crate::tls_front::fetcher::TlsFetchStrategy;
use crate::transport::UpstreamManager;
use super::generation::RuntimeTaskScope;
fn tls_fetch_host_for_domain(mask_host: &str, primary_tls_domain: &str, domain: &str) -> String {
if mask_host.eq_ignore_ascii_case(primary_tls_domain) {
domain.to_string()
@@ -23,6 +25,7 @@ pub(crate) async fn bootstrap_tls_front(
tls_domains: &[String],
upstream_manager: Arc<UpstreamManager>,
startup_tracker: &Arc<StartupTracker>,
task_scope: RuntimeTaskScope,
) -> Option<Arc<TlsFrontCache>> {
startup_tracker
.start_component(
@@ -69,7 +72,7 @@ pub(crate) async fn bootstrap_tls_front(
let scope_initial = tls_fetch_scope.clone();
let upstream_initial = upstream_manager.clone();
let strategy_initial = fetch_strategy.clone();
tokio::spawn(async move {
task_scope.spawn(async move {
let mut join = tokio::task::JoinSet::new();
for domain in domains_initial {
let cache_domain = cache_initial.clone();
@@ -109,7 +112,7 @@ pub(crate) async fn bootstrap_tls_front(
let cache_timeout = cache.clone();
let domains_timeout = tls_domains.to_vec();
let fake_cert_len = config.censorship.fake_cert_len;
tokio::spawn(async move {
task_scope.spawn(async move {
tokio::time::sleep(fetch_timeout).await;
for domain in domains_timeout {
let cached = cache_timeout.get(&domain).await;
@@ -132,7 +135,7 @@ pub(crate) async fn bootstrap_tls_front(
let scope_refresh = tls_fetch_scope.clone();
let upstream_refresh = upstream_manager.clone();
let strategy_refresh = fetch_strategy.clone();
tokio::spawn(async move {
task_scope.spawn(async move {
loop {
let base_secs = rand::rng().random_range(4 * 3600..=6 * 3600);
let jitter_secs = rand::rng().random_range(0..=7200);