mirror of
https://github.com/telemt/telemt.git
synced 2026-09-05 18:16:06 +03:00
Runtime Ownership hardened
Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
This commit is contained in:
+26
-25
@@ -1,5 +1,6 @@
|
|||||||
#![allow(clippy::too_many_arguments)]
|
#![allow(clippy::too_many_arguments)]
|
||||||
|
|
||||||
|
use std::collections::BTreeSet;
|
||||||
use std::io::{Error as IoError, ErrorKind};
|
use std::io::{Error as IoError, ErrorKind};
|
||||||
use std::net::{IpAddr, SocketAddr};
|
use std::net::{IpAddr, SocketAddr};
|
||||||
use std::path::PathBuf;
|
use std::path::PathBuf;
|
||||||
@@ -22,10 +23,12 @@ use tracing::{debug, info, warn};
|
|||||||
|
|
||||||
use crate::config::ApiGrayAction;
|
use crate::config::ApiGrayAction;
|
||||||
use crate::ip_tracker::UserIpTracker;
|
use crate::ip_tracker::UserIpTracker;
|
||||||
|
use crate::maestro::control_plane::ProcessControlPlane;
|
||||||
use crate::maestro::generation::{RuntimeGeneration, RuntimeWatchState};
|
use crate::maestro::generation::{RuntimeGeneration, RuntimeWatchState};
|
||||||
use crate::maestro::reload::{ReloadAccepted, ReloadControl, ReloadRequest, ReloadSubmitError};
|
use crate::maestro::reload::{ReloadAccepted, ReloadControl, ReloadRequest, ReloadSubmitError};
|
||||||
use crate::proxy::route_mode::RouteRuntimeController;
|
use crate::proxy::route_mode::RouteRuntimeController;
|
||||||
use crate::proxy::shared_state::ProxySharedState;
|
use crate::proxy::shared_state::ProxySharedState;
|
||||||
|
use crate::quota_state::QuotaStateOwner;
|
||||||
use crate::startup::StartupTracker;
|
use crate::startup::StartupTracker;
|
||||||
use crate::stats::Stats;
|
use crate::stats::Stats;
|
||||||
use crate::transport::UpstreamManager;
|
use crate::transport::UpstreamManager;
|
||||||
@@ -112,7 +115,7 @@ pub(super) struct ApiShared {
|
|||||||
pub(super) me_pool: Arc<RwLock<Option<Arc<MePool>>>>,
|
pub(super) me_pool: Arc<RwLock<Option<Arc<MePool>>>>,
|
||||||
pub(super) upstream_manager: Arc<UpstreamManager>,
|
pub(super) upstream_manager: Arc<UpstreamManager>,
|
||||||
pub(super) config_path: PathBuf,
|
pub(super) config_path: PathBuf,
|
||||||
pub(super) quota_state_path: PathBuf,
|
pub(super) quota_state: Arc<QuotaStateOwner>,
|
||||||
pub(super) detected_ips_rx: watch::Receiver<(Option<IpAddr>, Option<IpAddr>)>,
|
pub(super) detected_ips_rx: watch::Receiver<(Option<IpAddr>, Option<IpAddr>)>,
|
||||||
pub(super) mutation_lock: Arc<Mutex<()>>,
|
pub(super) mutation_lock: Arc<Mutex<()>>,
|
||||||
pub(super) minimal_cache: Arc<Mutex<Option<MinimalCacheEntry>>>,
|
pub(super) minimal_cache: Arc<Mutex<Option<MinimalCacheEntry>>>,
|
||||||
@@ -147,7 +150,7 @@ impl ApiShared {
|
|||||||
me_pool: runtime.me_pool_runtime.clone(),
|
me_pool: runtime.me_pool_runtime.clone(),
|
||||||
upstream_manager: runtime.upstream_manager.clone(),
|
upstream_manager: runtime.upstream_manager.clone(),
|
||||||
config_path: self.config_path.clone(),
|
config_path: self.config_path.clone(),
|
||||||
quota_state_path: self.quota_state_path.clone(),
|
quota_state: self.quota_state.clone(),
|
||||||
detected_ips_rx: self.detected_ips_rx.clone(),
|
detected_ips_rx: self.detected_ips_rx.clone(),
|
||||||
mutation_lock: self.mutation_lock.clone(),
|
mutation_lock: self.mutation_lock.clone(),
|
||||||
minimal_cache: self.minimal_cache.clone(),
|
minimal_cache: self.minimal_cache.clone(),
|
||||||
@@ -282,8 +285,9 @@ fn allowed_methods_for_path(path: &str) -> Option<&'static str> {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn serve(
|
/// Serves the API on a process-owned listener and task scope.
|
||||||
listen: SocketAddr,
|
pub(crate) async fn serve(
|
||||||
|
listener: TcpListener,
|
||||||
stats: Arc<Stats>,
|
stats: Arc<Stats>,
|
||||||
ip_tracker: Arc<UserIpTracker>,
|
ip_tracker: Arc<UserIpTracker>,
|
||||||
me_pool: Arc<RwLock<Option<Arc<MePool>>>>,
|
me_pool: Arc<RwLock<Option<Arc<MePool>>>>,
|
||||||
@@ -291,7 +295,7 @@ pub async fn serve(
|
|||||||
proxy_shared: Arc<ProxySharedState>,
|
proxy_shared: Arc<ProxySharedState>,
|
||||||
upstream_manager: Arc<UpstreamManager>,
|
upstream_manager: Arc<UpstreamManager>,
|
||||||
config_path: PathBuf,
|
config_path: PathBuf,
|
||||||
quota_state_path: PathBuf,
|
quota_state: Arc<QuotaStateOwner>,
|
||||||
detected_ips_rx: watch::Receiver<(Option<IpAddr>, Option<IpAddr>)>,
|
detected_ips_rx: watch::Receiver<(Option<IpAddr>, Option<IpAddr>)>,
|
||||||
process_started_at_epoch_secs: u64,
|
process_started_at_epoch_secs: u64,
|
||||||
startup_tracker: Arc<StartupTracker>,
|
startup_tracker: Arc<StartupTracker>,
|
||||||
@@ -300,6 +304,7 @@ pub async fn serve(
|
|||||||
mut runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
|
mut runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
|
||||||
web_trace: Arc<WebTraceStore>,
|
web_trace: Arc<WebTraceStore>,
|
||||||
web_runtime_rx: watch::Receiver<WebRuntimePublication>,
|
web_runtime_rx: watch::Receiver<WebRuntimePublication>,
|
||||||
|
control_plane: ProcessControlPlane,
|
||||||
) {
|
) {
|
||||||
let active_runtime = loop {
|
let active_runtime = loop {
|
||||||
if let Some(active_runtime) = active_runtime_rx.borrow().clone() {
|
if let Some(active_runtime) = active_runtime_rx.borrow().clone() {
|
||||||
@@ -321,19 +326,9 @@ pub async fn serve(
|
|||||||
};
|
};
|
||||||
let config_rx = initial_watch_state.config_rx.clone();
|
let config_rx = initial_watch_state.config_rx.clone();
|
||||||
let admission_rx = initial_watch_state.admission_rx.clone();
|
let admission_rx = initial_watch_state.admission_rx.clone();
|
||||||
let listener = match TcpListener::bind(listen).await {
|
let listen = listener.local_addr().ok();
|
||||||
Ok(listener) => listener,
|
|
||||||
Err(error) => {
|
|
||||||
warn!(
|
|
||||||
error = %error,
|
|
||||||
listen = %listen,
|
|
||||||
"Failed to bind API listener"
|
|
||||||
);
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
info!("API endpoint: http://{}/v1/* and /web-status", listen);
|
info!(listen = ?listen, "API endpoint ready at /v1/* and /web-status");
|
||||||
|
|
||||||
let runtime_state = Arc::new(ApiRuntimeState {
|
let runtime_state = Arc::new(ApiRuntimeState {
|
||||||
process_started_at_epoch_secs,
|
process_started_at_epoch_secs,
|
||||||
@@ -348,7 +343,7 @@ pub async fn serve(
|
|||||||
me_pool,
|
me_pool,
|
||||||
upstream_manager,
|
upstream_manager,
|
||||||
config_path,
|
config_path,
|
||||||
quota_state_path,
|
quota_state,
|
||||||
detected_ips_rx,
|
detected_ips_rx,
|
||||||
mutation_lock: Arc::new(Mutex::new(())),
|
mutation_lock: Arc::new(Mutex::new(())),
|
||||||
minimal_cache: Arc::new(Mutex::new(None)),
|
minimal_cache: Arc::new(Mutex::new(None)),
|
||||||
@@ -373,6 +368,7 @@ pub async fn serve(
|
|||||||
runtime_watch_rx,
|
runtime_watch_rx,
|
||||||
runtime_state.clone(),
|
runtime_state.clone(),
|
||||||
shared.runtime_events.clone(),
|
shared.runtime_events.clone(),
|
||||||
|
&control_plane,
|
||||||
);
|
);
|
||||||
|
|
||||||
let connection_permits = Arc::new(Semaphore::new(API_MAX_CONTROL_CONNECTIONS));
|
let connection_permits = Arc::new(Semaphore::new(API_MAX_CONTROL_CONNECTIONS));
|
||||||
@@ -399,7 +395,7 @@ pub async fn serve(
|
|||||||
};
|
};
|
||||||
|
|
||||||
let shared_conn = shared.clone();
|
let shared_conn = shared.clone();
|
||||||
tokio::spawn(async move {
|
let _ = control_plane.spawn(async move {
|
||||||
let _connection_permit = connection_permit;
|
let _connection_permit = connection_permit;
|
||||||
let svc = service_fn(move |req: Request<Incoming>| {
|
let svc = service_fn(move |req: Request<Incoming>| {
|
||||||
let shared_req = shared_conn.clone();
|
let shared_req = shared_conn.clone();
|
||||||
@@ -994,6 +990,7 @@ async fn handle(
|
|||||||
));
|
));
|
||||||
}
|
}
|
||||||
let expected_revision = parse_if_match(req.headers());
|
let expected_revision = parse_if_match(req.headers());
|
||||||
|
let _mutation_guard = shared.mutation_lock.lock().await;
|
||||||
let disk_cfg = load_config_from_disk(&shared.config_path).await?;
|
let disk_cfg = load_config_from_disk(&shared.config_path).await?;
|
||||||
ensure_expected_revision(&shared.config_path, expected_revision.as_deref())
|
ensure_expected_revision(&shared.config_path, expected_revision.as_deref())
|
||||||
.await?;
|
.await?;
|
||||||
@@ -1003,12 +1000,16 @@ async fn handle(
|
|||||||
ApiFailure::new(StatusCode::NOT_FOUND, "not_found", "User not found"),
|
ApiFailure::new(StatusCode::NOT_FOUND, "not_found", "User not found"),
|
||||||
));
|
));
|
||||||
}
|
}
|
||||||
let snapshot = match crate::quota_state::reset_user_quota(
|
let configured_users = disk_cfg
|
||||||
&shared.quota_state_path,
|
.access
|
||||||
shared.stats.as_ref(),
|
.users
|
||||||
user,
|
.keys()
|
||||||
)
|
.cloned()
|
||||||
.await
|
.collect::<BTreeSet<_>>();
|
||||||
|
let snapshot = match shared
|
||||||
|
.quota_state
|
||||||
|
.reset_user(&configured_users, user)
|
||||||
|
.await
|
||||||
{
|
{
|
||||||
Ok(snapshot) => snapshot,
|
Ok(snapshot) => snapshot,
|
||||||
Err(error) => {
|
Err(error) => {
|
||||||
|
|||||||
@@ -211,6 +211,7 @@ pub(super) struct ZeroMiddleProxyData {
|
|||||||
pub(super) reconnect_success_total: u64,
|
pub(super) reconnect_success_total: u64,
|
||||||
pub(super) handshake_reject_total: u64,
|
pub(super) handshake_reject_total: u64,
|
||||||
pub(super) handshake_error_codes: Vec<ZeroCodeCount>,
|
pub(super) handshake_error_codes: Vec<ZeroCodeCount>,
|
||||||
|
pub(super) handshake_error_code_overflow_total: u64,
|
||||||
pub(super) reader_eof_total: u64,
|
pub(super) reader_eof_total: u64,
|
||||||
pub(super) idle_close_by_peer_total: u64,
|
pub(super) idle_close_by_peer_total: u64,
|
||||||
pub(super) route_drop_no_conn_total: u64,
|
pub(super) route_drop_no_conn_total: u64,
|
||||||
@@ -388,8 +389,11 @@ pub(super) struct MinimalDcPathData {
|
|||||||
pub(super) struct MinimalMeRuntimeData {
|
pub(super) struct MinimalMeRuntimeData {
|
||||||
pub(super) active_generation: u64,
|
pub(super) active_generation: u64,
|
||||||
pub(super) warm_generation: u64,
|
pub(super) warm_generation: u64,
|
||||||
|
pub(super) warm_generations: Vec<u64>,
|
||||||
pub(super) pending_hardswap_generation: u64,
|
pub(super) pending_hardswap_generation: u64,
|
||||||
pub(super) pending_hardswap_age_secs: Option<u64>,
|
pub(super) pending_hardswap_age_secs: Option<u64>,
|
||||||
|
pub(super) reinit_inflight: usize,
|
||||||
|
pub(super) reinit_max_concurrency_effective: usize,
|
||||||
pub(super) hardswap_enabled: bool,
|
pub(super) hardswap_enabled: bool,
|
||||||
pub(super) floor_mode: &'static str,
|
pub(super) floor_mode: &'static str,
|
||||||
pub(super) adaptive_floor_idle_secs: u64,
|
pub(super) adaptive_floor_idle_secs: u64,
|
||||||
|
|||||||
+11
-2
@@ -21,8 +21,11 @@ pub(super) struct SecurityWhitelistData {
|
|||||||
pub(super) struct RuntimeMePoolStateGenerationData {
|
pub(super) struct RuntimeMePoolStateGenerationData {
|
||||||
pub(super) active_generation: u64,
|
pub(super) active_generation: u64,
|
||||||
pub(super) warm_generation: u64,
|
pub(super) warm_generation: u64,
|
||||||
|
pub(super) warm_generations: Vec<u64>,
|
||||||
pub(super) pending_hardswap_generation: u64,
|
pub(super) pending_hardswap_generation: u64,
|
||||||
pub(super) pending_hardswap_age_secs: Option<u64>,
|
pub(super) pending_hardswap_age_secs: Option<u64>,
|
||||||
|
pub(super) reinit_inflight: usize,
|
||||||
|
pub(super) reinit_max_concurrency_effective: usize,
|
||||||
pub(super) draining_generations: Vec<u64>,
|
pub(super) draining_generations: Vec<u64>,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -67,6 +70,8 @@ pub(super) struct RuntimeMePoolStateRefillDcData {
|
|||||||
pub(super) struct RuntimeMePoolStateRefillData {
|
pub(super) struct RuntimeMePoolStateRefillData {
|
||||||
pub(super) inflight_endpoints_total: usize,
|
pub(super) inflight_endpoints_total: usize,
|
||||||
pub(super) inflight_dc_total: usize,
|
pub(super) inflight_dc_total: usize,
|
||||||
|
pub(super) running_dc_total: usize,
|
||||||
|
pub(super) pending_dc_total: usize,
|
||||||
pub(super) by_dc: Vec<RuntimeMePoolStateRefillDcData>,
|
pub(super) by_dc: Vec<RuntimeMePoolStateRefillDcData>,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -291,8 +296,7 @@ pub(super) async fn build_runtime_me_pool_state_data(shared: &ApiShared) -> Runt
|
|||||||
};
|
};
|
||||||
};
|
};
|
||||||
|
|
||||||
let status = pool.api_status_snapshot().await;
|
let (status, runtime) = pool.api_coherent_snapshots().await;
|
||||||
let runtime = pool.api_runtime_snapshot().await;
|
|
||||||
let refill = pool.api_refill_snapshot().await;
|
let refill = pool.api_refill_snapshot().await;
|
||||||
|
|
||||||
let mut draining_generations = BTreeSet::<u64>::new();
|
let mut draining_generations = BTreeSet::<u64>::new();
|
||||||
@@ -329,8 +333,11 @@ pub(super) async fn build_runtime_me_pool_state_data(shared: &ApiShared) -> Runt
|
|||||||
generations: RuntimeMePoolStateGenerationData {
|
generations: RuntimeMePoolStateGenerationData {
|
||||||
active_generation: runtime.active_generation,
|
active_generation: runtime.active_generation,
|
||||||
warm_generation: runtime.warm_generation,
|
warm_generation: runtime.warm_generation,
|
||||||
|
warm_generations: runtime.warm_generations,
|
||||||
pending_hardswap_generation: runtime.pending_hardswap_generation,
|
pending_hardswap_generation: runtime.pending_hardswap_generation,
|
||||||
pending_hardswap_age_secs: runtime.pending_hardswap_age_secs,
|
pending_hardswap_age_secs: runtime.pending_hardswap_age_secs,
|
||||||
|
reinit_inflight: runtime.reinit_inflight,
|
||||||
|
reinit_max_concurrency_effective: runtime.reinit_max_concurrency_effective,
|
||||||
draining_generations: draining_generations.into_iter().collect(),
|
draining_generations: draining_generations.into_iter().collect(),
|
||||||
},
|
},
|
||||||
hardswap: RuntimeMePoolStateHardswapData {
|
hardswap: RuntimeMePoolStateHardswapData {
|
||||||
@@ -356,6 +363,8 @@ pub(super) async fn build_runtime_me_pool_state_data(shared: &ApiShared) -> Runt
|
|||||||
refill: RuntimeMePoolStateRefillData {
|
refill: RuntimeMePoolStateRefillData {
|
||||||
inflight_endpoints_total: refill.inflight_endpoints_total,
|
inflight_endpoints_total: refill.inflight_endpoints_total,
|
||||||
inflight_dc_total: refill.inflight_dc_total,
|
inflight_dc_total: refill.inflight_dc_total,
|
||||||
|
running_dc_total: refill.running_dc_total,
|
||||||
|
pending_dc_total: refill.pending_dc_total,
|
||||||
by_dc: refill
|
by_dc: refill
|
||||||
.by_dc
|
.by_dc
|
||||||
.into_iter()
|
.into_iter()
|
||||||
|
|||||||
@@ -84,6 +84,8 @@ pub(super) fn build_zero_all_data(stats: &Stats, configured_users: usize) -> Zer
|
|||||||
reconnect_success_total: stats.get_me_reconnect_success(),
|
reconnect_success_total: stats.get_me_reconnect_success(),
|
||||||
handshake_reject_total: stats.get_me_handshake_reject_total(),
|
handshake_reject_total: stats.get_me_handshake_reject_total(),
|
||||||
handshake_error_codes,
|
handshake_error_codes,
|
||||||
|
handshake_error_code_overflow_total: stats
|
||||||
|
.get_me_handshake_error_code_overflow_total(),
|
||||||
reader_eof_total: stats.get_me_reader_eof_total(),
|
reader_eof_total: stats.get_me_reader_eof_total(),
|
||||||
idle_close_by_peer_total: stats.get_me_idle_close_by_peer_total(),
|
idle_close_by_peer_total: stats.get_me_idle_close_by_peer_total(),
|
||||||
route_drop_no_conn_total: stats.get_me_route_drop_no_conn(),
|
route_drop_no_conn_total: stats.get_me_route_drop_no_conn(),
|
||||||
@@ -342,8 +344,7 @@ async fn get_minimal_payload_cached(
|
|||||||
}
|
}
|
||||||
|
|
||||||
let pool = shared.me_pool.read().await.clone()?;
|
let pool = shared.me_pool.read().await.clone()?;
|
||||||
let status = pool.api_status_snapshot().await;
|
let (status, runtime) = pool.api_coherent_snapshots().await;
|
||||||
let runtime = pool.api_runtime_snapshot().await;
|
|
||||||
let generated_at_epoch_secs = status.generated_at_epoch_secs;
|
let generated_at_epoch_secs = status.generated_at_epoch_secs;
|
||||||
|
|
||||||
let me_writers = MeWritersData {
|
let me_writers = MeWritersData {
|
||||||
@@ -425,8 +426,11 @@ async fn get_minimal_payload_cached(
|
|||||||
let me_runtime = MinimalMeRuntimeData {
|
let me_runtime = MinimalMeRuntimeData {
|
||||||
active_generation: runtime.active_generation,
|
active_generation: runtime.active_generation,
|
||||||
warm_generation: runtime.warm_generation,
|
warm_generation: runtime.warm_generation,
|
||||||
|
warm_generations: runtime.warm_generations,
|
||||||
pending_hardswap_generation: runtime.pending_hardswap_generation,
|
pending_hardswap_generation: runtime.pending_hardswap_generation,
|
||||||
pending_hardswap_age_secs: runtime.pending_hardswap_age_secs,
|
pending_hardswap_age_secs: runtime.pending_hardswap_age_secs,
|
||||||
|
reinit_inflight: runtime.reinit_inflight,
|
||||||
|
reinit_max_concurrency_effective: runtime.reinit_max_concurrency_effective,
|
||||||
hardswap_enabled: runtime.hardswap_enabled,
|
hardswap_enabled: runtime.hardswap_enabled,
|
||||||
floor_mode: runtime.floor_mode,
|
floor_mode: runtime.floor_mode,
|
||||||
adaptive_floor_idle_secs: runtime.adaptive_floor_idle_secs,
|
adaptive_floor_idle_secs: runtime.adaptive_floor_idle_secs,
|
||||||
|
|||||||
+34
-15
@@ -5,6 +5,7 @@ use std::time::{SystemTime, UNIX_EPOCH};
|
|||||||
use tokio::sync::watch;
|
use tokio::sync::watch;
|
||||||
|
|
||||||
use crate::maestro::generation::RuntimeWatchState;
|
use crate::maestro::generation::RuntimeWatchState;
|
||||||
|
use crate::maestro::control_plane::ProcessControlPlane;
|
||||||
|
|
||||||
use super::ApiRuntimeState;
|
use super::ApiRuntimeState;
|
||||||
use super::events::ApiEventStore;
|
use super::events::ApiEventStore;
|
||||||
@@ -13,22 +14,29 @@ pub(super) fn spawn_runtime_watchers(
|
|||||||
runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
|
runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
|
||||||
runtime_state: Arc<ApiRuntimeState>,
|
runtime_state: Arc<ApiRuntimeState>,
|
||||||
runtime_events: Arc<ApiEventStore>,
|
runtime_events: Arc<ApiEventStore>,
|
||||||
|
control_plane: &ProcessControlPlane,
|
||||||
) {
|
) {
|
||||||
let _config_watcher = spawn_config_watcher(
|
spawn_config_watcher(
|
||||||
runtime_watch_rx.clone(),
|
runtime_watch_rx.clone(),
|
||||||
runtime_state.clone(),
|
runtime_state.clone(),
|
||||||
runtime_events.clone(),
|
runtime_events.clone(),
|
||||||
|
control_plane,
|
||||||
|
);
|
||||||
|
spawn_admission_watcher(
|
||||||
|
runtime_watch_rx,
|
||||||
|
runtime_state,
|
||||||
|
runtime_events,
|
||||||
|
control_plane,
|
||||||
);
|
);
|
||||||
let _admission_watcher =
|
|
||||||
spawn_admission_watcher(runtime_watch_rx, runtime_state, runtime_events);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
fn spawn_config_watcher(
|
fn spawn_config_watcher(
|
||||||
mut runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
|
mut runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
|
||||||
runtime_state: Arc<ApiRuntimeState>,
|
runtime_state: Arc<ApiRuntimeState>,
|
||||||
runtime_events: Arc<ApiEventStore>,
|
runtime_events: Arc<ApiEventStore>,
|
||||||
) -> tokio::task::JoinHandle<()> {
|
control_plane: &ProcessControlPlane,
|
||||||
tokio::spawn(async move {
|
) {
|
||||||
|
let _ = control_plane.spawn(async move {
|
||||||
let Some(mut current) = runtime_watch_rx.borrow().clone() else {
|
let Some(mut current) = runtime_watch_rx.borrow().clone() else {
|
||||||
return;
|
return;
|
||||||
};
|
};
|
||||||
@@ -78,15 +86,16 @@ fn spawn_config_watcher(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
})
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
fn spawn_admission_watcher(
|
fn spawn_admission_watcher(
|
||||||
mut runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
|
mut runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
|
||||||
runtime_state: Arc<ApiRuntimeState>,
|
runtime_state: Arc<ApiRuntimeState>,
|
||||||
runtime_events: Arc<ApiEventStore>,
|
runtime_events: Arc<ApiEventStore>,
|
||||||
) -> tokio::task::JoinHandle<()> {
|
control_plane: &ProcessControlPlane,
|
||||||
tokio::spawn(async move {
|
) {
|
||||||
|
let _ = control_plane.spawn(async move {
|
||||||
let Some(mut current) = runtime_watch_rx.borrow().clone() else {
|
let Some(mut current) = runtime_watch_rx.borrow().clone() else {
|
||||||
return;
|
return;
|
||||||
};
|
};
|
||||||
@@ -124,7 +133,7 @@ fn spawn_admission_watcher(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
})
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
fn active_generation_id(
|
fn active_generation_id(
|
||||||
@@ -246,7 +255,13 @@ mod tests {
|
|||||||
let (runtime_watch_tx, runtime_watch_rx) = watch::channel(Some(initial));
|
let (runtime_watch_tx, runtime_watch_rx) = watch::channel(Some(initial));
|
||||||
let runtime_state = runtime_state();
|
let runtime_state = runtime_state();
|
||||||
let events = Arc::new(ApiEventStore::new(16));
|
let events = Arc::new(ApiEventStore::new(16));
|
||||||
spawn_runtime_watchers(runtime_watch_rx, runtime_state.clone(), events.clone());
|
let control_plane = ProcessControlPlane::new();
|
||||||
|
spawn_runtime_watchers(
|
||||||
|
runtime_watch_rx,
|
||||||
|
runtime_state.clone(),
|
||||||
|
events.clone(),
|
||||||
|
&control_plane,
|
||||||
|
);
|
||||||
tokio::task::yield_now().await;
|
tokio::task::yield_now().await;
|
||||||
|
|
||||||
assert_eq!(runtime_state.config_reload_count.load(Ordering::Relaxed), 0);
|
assert_eq!(runtime_state.config_reload_count.load(Ordering::Relaxed), 0);
|
||||||
@@ -283,6 +298,7 @@ mod tests {
|
|||||||
.count(),
|
.count(),
|
||||||
3
|
3
|
||||||
);
|
);
|
||||||
|
assert!(control_plane.shutdown(Duration::from_secs(1)).await);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
@@ -291,7 +307,13 @@ mod tests {
|
|||||||
let (runtime_watch_tx, runtime_watch_rx) = watch::channel(Some(initial));
|
let (runtime_watch_tx, runtime_watch_rx) = watch::channel(Some(initial));
|
||||||
let runtime_state = runtime_state();
|
let runtime_state = runtime_state();
|
||||||
let events = Arc::new(ApiEventStore::new(16));
|
let events = Arc::new(ApiEventStore::new(16));
|
||||||
let watcher = spawn_config_watcher(runtime_watch_rx, runtime_state.clone(), events.clone());
|
let control_plane = ProcessControlPlane::new();
|
||||||
|
spawn_config_watcher(
|
||||||
|
runtime_watch_rx,
|
||||||
|
runtime_state.clone(),
|
||||||
|
events.clone(),
|
||||||
|
&control_plane,
|
||||||
|
);
|
||||||
drop(initial_config_tx);
|
drop(initial_config_tx);
|
||||||
tokio::task::yield_now().await;
|
tokio::task::yield_now().await;
|
||||||
|
|
||||||
@@ -302,10 +324,7 @@ mod tests {
|
|||||||
wait_for_count(&runtime_state, 2).await;
|
wait_for_count(&runtime_state, 2).await;
|
||||||
|
|
||||||
drop(runtime_watch_tx);
|
drop(runtime_watch_tx);
|
||||||
tokio::time::timeout(Duration::from_secs(1), watcher)
|
assert!(control_plane.shutdown(Duration::from_secs(1)).await);
|
||||||
.await
|
|
||||||
.unwrap()
|
|
||||||
.unwrap();
|
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
events
|
events
|
||||||
.snapshot(16)
|
.snapshot(16)
|
||||||
|
|||||||
@@ -1,4 +1,5 @@
|
|||||||
use super::*;
|
use super::*;
|
||||||
|
use tracing::warn;
|
||||||
|
|
||||||
pub(in crate::api) async fn rotate_secret(
|
pub(in crate::api) async fn rotate_secret(
|
||||||
user: &str,
|
user: &str,
|
||||||
@@ -108,6 +109,18 @@ pub(in crate::api) async fn delete_user(
|
|||||||
.map_err(|e| ApiFailure::bad_request(format!("config validation failed: {}", e)))?;
|
.map_err(|e| ApiFailure::bad_request(format!("config validation failed: {}", e)))?;
|
||||||
let revision =
|
let revision =
|
||||||
save_access_sections_to_disk(&shared.config_path, &cfg, &touched_sections).await?;
|
save_access_sections_to_disk(&shared.config_path, &cfg, &touched_sections).await?;
|
||||||
|
let configured_users = cfg.access.users.keys().cloned().collect();
|
||||||
|
if let Err(error) = shared
|
||||||
|
.quota_state
|
||||||
|
.remove_user(&configured_users, user)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
warn!(
|
||||||
|
user,
|
||||||
|
error = %error,
|
||||||
|
"Deleted user quota checkpoint cleanup will be reconciled on restart"
|
||||||
|
);
|
||||||
|
}
|
||||||
drop(_guard);
|
drop(_guard);
|
||||||
shared.ip_tracker.remove_user_limit(user).await;
|
shared.ip_tracker.remove_user_limit(user).await;
|
||||||
shared.ip_tracker.clear_user_ips(user).await;
|
shared.ip_tracker.clear_user_ips(user).await;
|
||||||
|
|||||||
@@ -18,6 +18,7 @@ pub(super) struct WebIngressStatus {
|
|||||||
}
|
}
|
||||||
|
|
||||||
impl WebIngressStatus {
|
impl WebIngressStatus {
|
||||||
|
/// Builds a process-ingress snapshot without probing external TLS termination.
|
||||||
pub(super) fn new(publication: &WebRuntimePublication, runtime_available: bool) -> Self {
|
pub(super) fn new(publication: &WebRuntimePublication, runtime_available: bool) -> Self {
|
||||||
let configured_listeners = publication.listeners.len();
|
let configured_listeners = publication.listeners.len();
|
||||||
let live_acceptors = publication.telemetry.live_acceptors();
|
let live_acceptors = publication.telemetry.live_acceptors();
|
||||||
@@ -64,6 +65,7 @@ pub(super) struct WebCapacityStatus {
|
|||||||
}
|
}
|
||||||
|
|
||||||
impl WebCapacityStatus {
|
impl WebCapacityStatus {
|
||||||
|
/// Builds a bounded capacity snapshot from non-blocking runtime observations.
|
||||||
pub(super) fn new(
|
pub(super) fn new(
|
||||||
publication: &WebRuntimePublication,
|
publication: &WebRuntimePublication,
|
||||||
runtime: Option<&WebProcessRuntime>,
|
runtime: Option<&WebProcessRuntime>,
|
||||||
@@ -104,6 +106,7 @@ pub(super) struct WebDecoyUpstreamStatus {
|
|||||||
}
|
}
|
||||||
|
|
||||||
impl WebDecoyUpstreamStatus {
|
impl WebDecoyUpstreamStatus {
|
||||||
|
/// Builds the fixed internal decoy-origin outcome snapshot.
|
||||||
pub(super) fn new(publication: &WebRuntimePublication) -> Self {
|
pub(super) fn new(publication: &WebRuntimePublication) -> Self {
|
||||||
let last = publication.telemetry.last_decoy();
|
let last = publication.telemetry.last_decoy();
|
||||||
Self {
|
Self {
|
||||||
|
|||||||
@@ -780,6 +780,10 @@ pub(crate) fn default_me_reinit_singleflight() -> bool {
|
|||||||
true
|
true
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub(crate) fn default_me_reinit_max_concurrency() -> usize {
|
||||||
|
2
|
||||||
|
}
|
||||||
|
|
||||||
pub(crate) fn default_me_reinit_trigger_channel() -> usize {
|
pub(crate) fn default_me_reinit_trigger_channel() -> usize {
|
||||||
64
|
64
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -10,6 +10,7 @@ pub struct HotFields {
|
|||||||
pub update_every_secs: u64,
|
pub update_every_secs: u64,
|
||||||
pub me_reinit_every_secs: u64,
|
pub me_reinit_every_secs: u64,
|
||||||
pub me_reinit_singleflight: bool,
|
pub me_reinit_singleflight: bool,
|
||||||
|
pub me_reinit_max_concurrency: usize,
|
||||||
pub me_reinit_coalesce_window_ms: u64,
|
pub me_reinit_coalesce_window_ms: u64,
|
||||||
pub hardswap: bool,
|
pub hardswap: bool,
|
||||||
pub me_pool_drain_ttl_secs: u64,
|
pub me_pool_drain_ttl_secs: u64,
|
||||||
@@ -102,6 +103,7 @@ impl HotFields {
|
|||||||
update_every_secs: cfg.general.effective_update_every_secs(),
|
update_every_secs: cfg.general.effective_update_every_secs(),
|
||||||
me_reinit_every_secs: cfg.general.me_reinit_every_secs,
|
me_reinit_every_secs: cfg.general.me_reinit_every_secs,
|
||||||
me_reinit_singleflight: cfg.general.me_reinit_singleflight,
|
me_reinit_singleflight: cfg.general.me_reinit_singleflight,
|
||||||
|
me_reinit_max_concurrency: cfg.general.me_reinit_max_concurrency,
|
||||||
me_reinit_coalesce_window_ms: cfg.general.me_reinit_coalesce_window_ms,
|
me_reinit_coalesce_window_ms: cfg.general.me_reinit_coalesce_window_ms,
|
||||||
hardswap: cfg.general.hardswap,
|
hardswap: cfg.general.hardswap,
|
||||||
me_pool_drain_ttl_secs: cfg.general.me_pool_drain_ttl_secs,
|
me_pool_drain_ttl_secs: cfg.general.me_pool_drain_ttl_secs,
|
||||||
@@ -236,6 +238,7 @@ pub(super) fn overlay_hot_fields(old: &ProxyConfig, new: &ProxyConfig) -> ProxyC
|
|||||||
cfg.general.proxy_config_auto_reload_secs = new.general.proxy_config_auto_reload_secs;
|
cfg.general.proxy_config_auto_reload_secs = new.general.proxy_config_auto_reload_secs;
|
||||||
cfg.general.me_reinit_every_secs = new.general.me_reinit_every_secs;
|
cfg.general.me_reinit_every_secs = new.general.me_reinit_every_secs;
|
||||||
cfg.general.me_reinit_singleflight = new.general.me_reinit_singleflight;
|
cfg.general.me_reinit_singleflight = new.general.me_reinit_singleflight;
|
||||||
|
cfg.general.me_reinit_max_concurrency = new.general.me_reinit_max_concurrency;
|
||||||
cfg.general.me_reinit_coalesce_window_ms = new.general.me_reinit_coalesce_window_ms;
|
cfg.general.me_reinit_coalesce_window_ms = new.general.me_reinit_coalesce_window_ms;
|
||||||
cfg.general.hardswap = new.general.hardswap;
|
cfg.general.hardswap = new.general.hardswap;
|
||||||
cfg.general.me_pool_drain_ttl_secs = new.general.me_pool_drain_ttl_secs;
|
cfg.general.me_pool_drain_ttl_secs = new.general.me_pool_drain_ttl_secs;
|
||||||
|
|||||||
@@ -114,12 +114,14 @@ pub(super) fn log_changes(
|
|||||||
}
|
}
|
||||||
if old_hot.me_reinit_every_secs != new_hot.me_reinit_every_secs
|
if old_hot.me_reinit_every_secs != new_hot.me_reinit_every_secs
|
||||||
|| old_hot.me_reinit_singleflight != new_hot.me_reinit_singleflight
|
|| old_hot.me_reinit_singleflight != new_hot.me_reinit_singleflight
|
||||||
|
|| old_hot.me_reinit_max_concurrency != new_hot.me_reinit_max_concurrency
|
||||||
|| old_hot.me_reinit_coalesce_window_ms != new_hot.me_reinit_coalesce_window_ms
|
|| old_hot.me_reinit_coalesce_window_ms != new_hot.me_reinit_coalesce_window_ms
|
||||||
{
|
{
|
||||||
info!(
|
info!(
|
||||||
"config reload: me_reinit: interval={}s singleflight={} coalesce={}ms",
|
"config reload: me_reinit: interval={}s singleflight={} max_concurrency={} coalesce={}ms",
|
||||||
new_hot.me_reinit_every_secs,
|
new_hot.me_reinit_every_secs,
|
||||||
new_hot.me_reinit_singleflight,
|
new_hot.me_reinit_singleflight,
|
||||||
|
new_hot.me_reinit_max_concurrency,
|
||||||
new_hot.me_reinit_coalesce_window_ms
|
new_hot.me_reinit_coalesce_window_ms
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -272,6 +272,7 @@ async fn candidate_watcher_waits_for_activation_and_reconciles_disk() {
|
|||||||
None,
|
None,
|
||||||
None,
|
None,
|
||||||
cancellation.clone(),
|
cancellation.clone(),
|
||||||
|
None,
|
||||||
Some(activation_rx),
|
Some(activation_rx),
|
||||||
);
|
);
|
||||||
let watcher = tokio::spawn(watcher);
|
let watcher = tokio::spawn(watcher);
|
||||||
|
|||||||
@@ -124,13 +124,14 @@ fn apply_watch_manifest<W1: Watcher, W2: Watcher>(
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Load config, validate, diff against current, and broadcast if changed.
|
/// Load config, validate, diff against current, and broadcast if changed.
|
||||||
pub(super) fn reload_config(
|
fn reload_config_with_resolver(
|
||||||
config_path: &PathBuf,
|
config_path: &PathBuf,
|
||||||
config_tx: &watch::Sender<Arc<ProxyConfig>>,
|
config_tx: &watch::Sender<Arc<ProxyConfig>>,
|
||||||
log_tx: &watch::Sender<LogLevel>,
|
log_tx: &watch::Sender<LogLevel>,
|
||||||
detected_ip_v4: Option<IpAddr>,
|
detected_ip_v4: Option<IpAddr>,
|
||||||
detected_ip_v6: Option<IpAddr>,
|
detected_ip_v6: Option<IpAddr>,
|
||||||
reload_state: &mut ReloadState,
|
reload_state: &mut ReloadState,
|
||||||
|
dns_resolver: Option<&crate::network::dns_overrides::GenerationDnsResolver>,
|
||||||
) -> Option<WatchManifest> {
|
) -> Option<WatchManifest> {
|
||||||
let loaded = match ProxyConfig::load_with_metadata(config_path) {
|
let loaded = match ProxyConfig::load_with_metadata(config_path) {
|
||||||
Ok(loaded) => loaded,
|
Ok(loaded) => loaded,
|
||||||
@@ -176,7 +177,8 @@ pub(super) fn reload_config(
|
|||||||
}
|
}
|
||||||
|
|
||||||
if old_hot.dns_overrides != applied_hot.dns_overrides
|
if old_hot.dns_overrides != applied_hot.dns_overrides
|
||||||
&& let Err(e) = crate::network::dns_overrides::install_entries(&applied_hot.dns_overrides)
|
&& let Some(dns_resolver) = dns_resolver
|
||||||
|
&& let Err(e) = dns_resolver.apply_entries(&applied_hot.dns_overrides)
|
||||||
{
|
{
|
||||||
error!(
|
error!(
|
||||||
"config reload: invalid network.dns_overrides: {}; keeping old config",
|
"config reload: invalid network.dns_overrides: {}; keeping old config",
|
||||||
@@ -198,6 +200,26 @@ pub(super) fn reload_config(
|
|||||||
Some(next_manifest)
|
Some(next_manifest)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
pub(super) fn reload_config(
|
||||||
|
config_path: &PathBuf,
|
||||||
|
config_tx: &watch::Sender<Arc<ProxyConfig>>,
|
||||||
|
log_tx: &watch::Sender<LogLevel>,
|
||||||
|
detected_ip_v4: Option<IpAddr>,
|
||||||
|
detected_ip_v6: Option<IpAddr>,
|
||||||
|
reload_state: &mut ReloadState,
|
||||||
|
) -> Option<WatchManifest> {
|
||||||
|
reload_config_with_resolver(
|
||||||
|
config_path,
|
||||||
|
config_tx,
|
||||||
|
log_tx,
|
||||||
|
detected_ip_v4,
|
||||||
|
detected_ip_v6,
|
||||||
|
reload_state,
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
/// Spawn the hot-reload watcher task.
|
/// Spawn the hot-reload watcher task.
|
||||||
///
|
///
|
||||||
/// Uses `notify` (inotify on Linux) to detect file changes instantly.
|
/// Uses `notify` (inotify on Linux) to detect file changes instantly.
|
||||||
@@ -213,6 +235,7 @@ pub fn spawn_config_watcher(
|
|||||||
detected_ip_v4: Option<IpAddr>,
|
detected_ip_v4: Option<IpAddr>,
|
||||||
detected_ip_v6: Option<IpAddr>,
|
detected_ip_v6: Option<IpAddr>,
|
||||||
cancellation: tokio_util::sync::CancellationToken,
|
cancellation: tokio_util::sync::CancellationToken,
|
||||||
|
dns_resolver: Option<Arc<crate::network::dns_overrides::GenerationDnsResolver>>,
|
||||||
mut activation: Option<watch::Receiver<bool>>,
|
mut activation: Option<watch::Receiver<bool>>,
|
||||||
) -> (
|
) -> (
|
||||||
watch::Receiver<Arc<ProxyConfig>>,
|
watch::Receiver<Arc<ProxyConfig>>,
|
||||||
@@ -364,24 +387,26 @@ pub fn spawn_config_watcher(
|
|||||||
tokio::time::sleep(HOT_RELOAD_DEBOUNCE).await;
|
tokio::time::sleep(HOT_RELOAD_DEBOUNCE).await;
|
||||||
while notify_rx.try_recv().is_ok() {}
|
while notify_rx.try_recv().is_ok() {}
|
||||||
|
|
||||||
let mut next_manifest = reload_config(
|
let mut next_manifest = reload_config_with_resolver(
|
||||||
&config_path,
|
&config_path,
|
||||||
&config_tx,
|
&config_tx,
|
||||||
&log_tx,
|
&log_tx,
|
||||||
detected_ip_v4,
|
detected_ip_v4,
|
||||||
detected_ip_v6,
|
detected_ip_v6,
|
||||||
&mut reload_state,
|
&mut reload_state,
|
||||||
|
dns_resolver.as_deref(),
|
||||||
);
|
);
|
||||||
if next_manifest.is_none() {
|
if next_manifest.is_none() {
|
||||||
tokio::time::sleep(HOT_RELOAD_DEBOUNCE).await;
|
tokio::time::sleep(HOT_RELOAD_DEBOUNCE).await;
|
||||||
while notify_rx.try_recv().is_ok() {}
|
while notify_rx.try_recv().is_ok() {}
|
||||||
next_manifest = reload_config(
|
next_manifest = reload_config_with_resolver(
|
||||||
&config_path,
|
&config_path,
|
||||||
&config_tx,
|
&config_tx,
|
||||||
&log_tx,
|
&log_tx,
|
||||||
detected_ip_v4,
|
detected_ip_v4,
|
||||||
detected_ip_v6,
|
detected_ip_v6,
|
||||||
&mut reload_state,
|
&mut reload_state,
|
||||||
|
dns_resolver.as_deref(),
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -161,6 +161,7 @@ const GENERAL_CONFIG_KEYS: &[&str] = &[
|
|||||||
"proxy_secret_auto_reload_secs",
|
"proxy_secret_auto_reload_secs",
|
||||||
"proxy_config_auto_reload_secs",
|
"proxy_config_auto_reload_secs",
|
||||||
"me_reinit_singleflight",
|
"me_reinit_singleflight",
|
||||||
|
"me_reinit_max_concurrency",
|
||||||
"me_reinit_trigger_channel",
|
"me_reinit_trigger_channel",
|
||||||
"me_reinit_coalesce_window_ms",
|
"me_reinit_coalesce_window_ms",
|
||||||
"me_deterministic_writer_sort",
|
"me_deterministic_writer_sort",
|
||||||
|
|||||||
@@ -143,9 +143,15 @@ pub(super) fn validate(config: &mut ProxyConfig) -> Result<()> {
|
|||||||
));
|
));
|
||||||
}
|
}
|
||||||
|
|
||||||
if config.general.me_reinit_trigger_channel == 0 {
|
if !(1..=8).contains(&config.general.me_reinit_max_concurrency) {
|
||||||
return Err(ProxyError::Config(
|
return Err(ProxyError::Config(
|
||||||
"general.me_reinit_trigger_channel must be > 0".to_string(),
|
"general.me_reinit_max_concurrency must be within [1, 8]".to_string(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
if !(1..=4096).contains(&config.general.me_reinit_trigger_channel) {
|
||||||
|
return Err(ProxyError::Config(
|
||||||
|
"general.me_reinit_trigger_channel must be within [1, 4096]".to_string(),
|
||||||
));
|
));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -441,6 +441,9 @@ pub struct GeneralConfig {
|
|||||||
/// Serialize ME reinit cycles across all trigger sources.
|
/// Serialize ME reinit cycles across all trigger sources.
|
||||||
#[serde(default = "default_me_reinit_singleflight")]
|
#[serde(default = "default_me_reinit_singleflight")]
|
||||||
pub me_reinit_singleflight: bool,
|
pub me_reinit_singleflight: bool,
|
||||||
|
/// Maximum concurrent ME reinit warmups when single-flight mode is disabled.
|
||||||
|
#[serde(default = "default_me_reinit_max_concurrency")]
|
||||||
|
pub me_reinit_max_concurrency: usize,
|
||||||
/// Trigger queue capacity for reinit scheduler.
|
/// Trigger queue capacity for reinit scheduler.
|
||||||
#[serde(default = "default_me_reinit_trigger_channel")]
|
#[serde(default = "default_me_reinit_trigger_channel")]
|
||||||
pub me_reinit_trigger_channel: usize,
|
pub me_reinit_trigger_channel: usize,
|
||||||
|
|||||||
@@ -159,6 +159,7 @@ impl Default for GeneralConfig {
|
|||||||
proxy_secret_auto_reload_secs: default_proxy_secret_reload_secs(),
|
proxy_secret_auto_reload_secs: default_proxy_secret_reload_secs(),
|
||||||
proxy_config_auto_reload_secs: default_proxy_config_reload_secs(),
|
proxy_config_auto_reload_secs: default_proxy_config_reload_secs(),
|
||||||
me_reinit_singleflight: default_me_reinit_singleflight(),
|
me_reinit_singleflight: default_me_reinit_singleflight(),
|
||||||
|
me_reinit_max_concurrency: default_me_reinit_max_concurrency(),
|
||||||
me_reinit_trigger_channel: default_me_reinit_trigger_channel(),
|
me_reinit_trigger_channel: default_me_reinit_trigger_channel(),
|
||||||
me_reinit_coalesce_window_ms: default_me_reinit_coalesce_window_ms(),
|
me_reinit_coalesce_window_ms: default_me_reinit_coalesce_window_ms(),
|
||||||
me_deterministic_writer_sort: default_me_deterministic_writer_sort(),
|
me_deterministic_writer_sort: default_me_deterministic_writer_sort(),
|
||||||
|
|||||||
+26
-5
@@ -6,6 +6,8 @@ use serde_json::Value;
|
|||||||
|
|
||||||
use crate::config::ProxyConfig;
|
use crate::config::ProxyConfig;
|
||||||
|
|
||||||
|
const HEALTHCHECK_RESPONSE_MAX_BYTES: u64 = 64 * 1024;
|
||||||
|
|
||||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||||
pub(crate) enum HealthcheckMode {
|
pub(crate) enum HealthcheckMode {
|
||||||
Liveness,
|
Liveness,
|
||||||
@@ -73,10 +75,7 @@ fn run_inner(config_path: &str, mode: HealthcheckMode) -> Result<(), String> {
|
|||||||
.flush()
|
.flush()
|
||||||
.map_err(|error| format!("request flush failed: {error}"))?;
|
.map_err(|error| format!("request flush failed: {error}"))?;
|
||||||
|
|
||||||
let mut raw_response = Vec::new();
|
let raw_response = read_response_bounded(&mut stream)?;
|
||||||
stream
|
|
||||||
.read_to_end(&mut raw_response)
|
|
||||||
.map_err(|error| format!("response read failed: {error}"))?;
|
|
||||||
let response =
|
let response =
|
||||||
String::from_utf8(raw_response).map_err(|_| "response is not valid UTF-8".to_string())?;
|
String::from_utf8(raw_response).map_err(|_| "response is not valid UTF-8".to_string())?;
|
||||||
|
|
||||||
@@ -89,6 +88,18 @@ fn run_inner(config_path: &str, mode: HealthcheckMode) -> Result<(), String> {
|
|||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn read_response_bounded(reader: &mut impl Read) -> Result<Vec<u8>, String> {
|
||||||
|
let mut raw_response = Vec::new();
|
||||||
|
reader
|
||||||
|
.take(HEALTHCHECK_RESPONSE_MAX_BYTES.saturating_add(1))
|
||||||
|
.read_to_end(&mut raw_response)
|
||||||
|
.map_err(|error| format!("response read failed: {error}"))?;
|
||||||
|
if raw_response.len() as u64 > HEALTHCHECK_RESPONSE_MAX_BYTES {
|
||||||
|
return Err("response exceeds the 64 KiB healthcheck limit".to_string());
|
||||||
|
}
|
||||||
|
Ok(raw_response)
|
||||||
|
}
|
||||||
|
|
||||||
fn probe_target(listen: SocketAddr) -> SocketAddr {
|
fn probe_target(listen: SocketAddr) -> SocketAddr {
|
||||||
match listen {
|
match listen {
|
||||||
SocketAddr::V4(addr) => {
|
SocketAddr::V4(addr) => {
|
||||||
@@ -180,7 +191,10 @@ fn validate_payload(mode: HealthcheckMode, body: &str) -> Result<(), String> {
|
|||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::{HealthcheckMode, parse_status_code, split_response, validate_payload};
|
use super::{
|
||||||
|
HEALTHCHECK_RESPONSE_MAX_BYTES, HealthcheckMode, parse_status_code,
|
||||||
|
read_response_bounded, split_response, validate_payload,
|
||||||
|
};
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn parse_status_code_reads_http_200() {
|
fn parse_status_code_reads_http_200() {
|
||||||
@@ -208,4 +222,11 @@ mod tests {
|
|||||||
let result = validate_payload(HealthcheckMode::Ready, body);
|
let result = validate_payload(HealthcheckMode::Ready, body);
|
||||||
assert!(result.is_err());
|
assert!(result.is_err());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn bounded_reader_rejects_oversized_health_response() {
|
||||||
|
let payload = vec![b'x'; HEALTHCHECK_RESPONSE_MAX_BYTES as usize + 1];
|
||||||
|
|
||||||
|
assert!(read_response_bounded(&mut payload.as_slice()).is_err());
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
+30
-37
@@ -8,10 +8,10 @@ use std::hash::{Hash, Hasher};
|
|||||||
use std::net::IpAddr;
|
use std::net::IpAddr;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use std::sync::Mutex;
|
use std::sync::Mutex;
|
||||||
use std::sync::atomic::{AtomicU8, AtomicU64, AtomicUsize, Ordering};
|
use std::sync::atomic::{AtomicU64, Ordering};
|
||||||
use std::time::{Duration, Instant};
|
use std::time::{Duration, Instant};
|
||||||
|
|
||||||
use dashmap::DashMap;
|
use arc_swap::ArcSwap;
|
||||||
use tokio::sync::{Mutex as AsyncMutex, RwLock};
|
use tokio::sync::{Mutex as AsyncMutex, RwLock};
|
||||||
|
|
||||||
use crate::config::UserMaxUniqueIpsMode;
|
use crate::config::UserMaxUniqueIpsMode;
|
||||||
@@ -39,6 +39,25 @@ struct CleanupShard {
|
|||||||
queue: Mutex<HashMap<String, HashMap<IpAddr, usize>>>,
|
queue: Mutex<HashMap<String, HashMap<IpAddr, usize>>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
struct UserIpLimitPolicy {
|
||||||
|
max_ips: Arc<HashMap<String, usize>>,
|
||||||
|
default_max_ips: usize,
|
||||||
|
mode: UserMaxUniqueIpsMode,
|
||||||
|
window_secs: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for UserIpLimitPolicy {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
max_ips: Arc::new(HashMap::new()),
|
||||||
|
default_max_ips: 0,
|
||||||
|
mode: UserMaxUniqueIpsMode::ActiveWindow,
|
||||||
|
window_secs: 30,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// Tracks active and recent client IPs for per-user admission control.
|
/// Tracks active and recent client IPs for per-user admission control.
|
||||||
#[derive(Debug, Clone)]
|
#[derive(Debug, Clone)]
|
||||||
pub struct UserIpTracker {
|
pub struct UserIpTracker {
|
||||||
@@ -48,10 +67,7 @@ pub struct UserIpTracker {
|
|||||||
active_cap_rejects: Arc<AtomicU64>,
|
active_cap_rejects: Arc<AtomicU64>,
|
||||||
recent_cap_rejects: Arc<AtomicU64>,
|
recent_cap_rejects: Arc<AtomicU64>,
|
||||||
cleanup_deferred_releases: Arc<AtomicU64>,
|
cleanup_deferred_releases: Arc<AtomicU64>,
|
||||||
max_ips: Arc<DashMap<String, usize>>,
|
limit_policy: Arc<ArcSwap<UserIpLimitPolicy>>,
|
||||||
default_max_ips: Arc<AtomicUsize>,
|
|
||||||
limit_mode: Arc<AtomicU8>,
|
|
||||||
limit_window_secs: Arc<AtomicU64>,
|
|
||||||
last_compact_epoch_secs: Arc<AtomicU64>,
|
last_compact_epoch_secs: Arc<AtomicU64>,
|
||||||
cleanup_queue_len: Arc<AtomicU64>,
|
cleanup_queue_len: Arc<AtomicU64>,
|
||||||
cleanup_shards: Arc<Box<[CleanupShard]>>,
|
cleanup_shards: Arc<Box<[CleanupShard]>>,
|
||||||
@@ -102,12 +118,7 @@ impl UserIpTracker {
|
|||||||
active_cap_rejects: Arc::new(AtomicU64::new(0)),
|
active_cap_rejects: Arc::new(AtomicU64::new(0)),
|
||||||
recent_cap_rejects: Arc::new(AtomicU64::new(0)),
|
recent_cap_rejects: Arc::new(AtomicU64::new(0)),
|
||||||
cleanup_deferred_releases: Arc::new(AtomicU64::new(0)),
|
cleanup_deferred_releases: Arc::new(AtomicU64::new(0)),
|
||||||
max_ips: Arc::new(DashMap::new()),
|
limit_policy: Arc::new(ArcSwap::from_pointee(UserIpLimitPolicy::default())),
|
||||||
default_max_ips: Arc::new(AtomicUsize::new(0)),
|
|
||||||
limit_mode: Arc::new(AtomicU8::new(Self::mode_to_u8(
|
|
||||||
UserMaxUniqueIpsMode::ActiveWindow,
|
|
||||||
))),
|
|
||||||
limit_window_secs: Arc::new(AtomicU64::new(30)),
|
|
||||||
last_compact_epoch_secs: Arc::new(AtomicU64::new(0)),
|
last_compact_epoch_secs: Arc::new(AtomicU64::new(0)),
|
||||||
cleanup_queue_len: Arc::new(AtomicU64::new(0)),
|
cleanup_queue_len: Arc::new(AtomicU64::new(0)),
|
||||||
cleanup_shards: Arc::new(cleanup_shards),
|
cleanup_shards: Arc::new(cleanup_shards),
|
||||||
@@ -117,41 +128,23 @@ impl UserIpTracker {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(super) fn mode_to_u8(mode: UserMaxUniqueIpsMode) -> u8 {
|
|
||||||
match mode {
|
|
||||||
UserMaxUniqueIpsMode::ActiveWindow => 0,
|
|
||||||
UserMaxUniqueIpsMode::TimeWindow => 1,
|
|
||||||
UserMaxUniqueIpsMode::Combined => 2,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
pub(super) fn mode_from_u8(raw: u8) -> UserMaxUniqueIpsMode {
|
|
||||||
match raw {
|
|
||||||
1 => UserMaxUniqueIpsMode::TimeWindow,
|
|
||||||
2 => UserMaxUniqueIpsMode::Combined,
|
|
||||||
_ => UserMaxUniqueIpsMode::ActiveWindow,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
pub(super) fn shard_idx(username: &str) -> usize {
|
pub(super) fn shard_idx(username: &str) -> usize {
|
||||||
let mut hasher = DefaultHasher::new();
|
let mut hasher = DefaultHasher::new();
|
||||||
username.hash(&mut hasher);
|
username.hash(&mut hasher);
|
||||||
(hasher.finish() as usize) & USER_IP_TRACKER_SHARD_MASK
|
(hasher.finish() as usize) & USER_IP_TRACKER_SHARD_MASK
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(super) fn limit_window(&self) -> Duration {
|
fn limit_window(policy: &UserIpLimitPolicy) -> Duration {
|
||||||
Duration::from_secs(self.limit_window_secs.load(Ordering::Relaxed).max(1))
|
Duration::from_secs(policy.window_secs)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(super) fn user_limit(&self, username: &str) -> Option<usize> {
|
fn user_limit(policy: &UserIpLimitPolicy, username: &str) -> Option<usize> {
|
||||||
self.max_ips
|
policy
|
||||||
|
.max_ips
|
||||||
.get(username)
|
.get(username)
|
||||||
.map(|limit| *limit)
|
.copied()
|
||||||
.filter(|limit| *limit > 0)
|
.filter(|limit| *limit > 0)
|
||||||
.or_else(|| {
|
.or_else(|| (policy.default_max_ips > 0).then_some(policy.default_max_ips))
|
||||||
let default_limit = self.default_max_ips.load(Ordering::Relaxed);
|
|
||||||
(default_limit > 0).then_some(default_limit)
|
|
||||||
})
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(super) fn decrement_counter(counter: &AtomicU64, amount: usize) {
|
pub(super) fn decrement_counter(counter: &AtomicU64, amount: usize) {
|
||||||
|
|||||||
+36
-14
@@ -2,26 +2,47 @@ use super::*;
|
|||||||
|
|
||||||
impl UserIpTracker {
|
impl UserIpTracker {
|
||||||
pub async fn set_limit_policy(&self, mode: UserMaxUniqueIpsMode, window_secs: u64) {
|
pub async fn set_limit_policy(&self, mode: UserMaxUniqueIpsMode, window_secs: u64) {
|
||||||
self.limit_mode
|
self.limit_policy.rcu(|current| {
|
||||||
.store(Self::mode_to_u8(mode), Ordering::Relaxed);
|
Arc::new(UserIpLimitPolicy {
|
||||||
self.limit_window_secs
|
mode,
|
||||||
.store(window_secs.max(1), Ordering::Relaxed);
|
window_secs: window_secs.max(1),
|
||||||
|
..(**current).clone()
|
||||||
|
})
|
||||||
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn set_user_limit(&self, username: &str, max_ips: usize) {
|
pub async fn set_user_limit(&self, username: &str, max_ips: usize) {
|
||||||
self.max_ips.insert(username.to_string(), max_ips);
|
let username = username.to_string();
|
||||||
|
self.limit_policy.rcu(|current| {
|
||||||
|
let mut limits = current.max_ips.as_ref().clone();
|
||||||
|
limits.insert(username.clone(), max_ips);
|
||||||
|
Arc::new(UserIpLimitPolicy {
|
||||||
|
max_ips: Arc::new(limits),
|
||||||
|
..(**current).clone()
|
||||||
|
})
|
||||||
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn remove_user_limit(&self, username: &str) {
|
pub async fn remove_user_limit(&self, username: &str) {
|
||||||
self.max_ips.remove(username);
|
self.limit_policy.rcu(|current| {
|
||||||
|
let mut limits = current.max_ips.as_ref().clone();
|
||||||
|
limits.remove(username);
|
||||||
|
Arc::new(UserIpLimitPolicy {
|
||||||
|
max_ips: Arc::new(limits),
|
||||||
|
..(**current).clone()
|
||||||
|
})
|
||||||
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn load_limits(&self, default_limit: usize, limits: &HashMap<String, usize>) {
|
pub async fn load_limits(&self, default_limit: usize, limits: &HashMap<String, usize>) {
|
||||||
self.default_max_ips.store(default_limit, Ordering::Relaxed);
|
let limits = Arc::new(limits.clone());
|
||||||
self.max_ips.clear();
|
self.limit_policy.rcu(|current| {
|
||||||
for (username, limit) in limits {
|
Arc::new(UserIpLimitPolicy {
|
||||||
self.max_ips.insert(username.clone(), *limit);
|
max_ips: Arc::clone(&limits),
|
||||||
}
|
default_max_ips: default_limit,
|
||||||
|
..(**current).clone()
|
||||||
|
})
|
||||||
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(super) fn prune_recent(
|
pub(super) fn prune_recent(
|
||||||
@@ -40,9 +61,10 @@ impl UserIpTracker {
|
|||||||
pub async fn check_and_add(&self, username: &str, ip: IpAddr) -> Result<(), String> {
|
pub async fn check_and_add(&self, username: &str, ip: IpAddr) -> Result<(), String> {
|
||||||
self.drain_cleanup_for_user(username).await;
|
self.drain_cleanup_for_user(username).await;
|
||||||
self.maybe_compact_empty_users().await;
|
self.maybe_compact_empty_users().await;
|
||||||
let limit = self.user_limit(username);
|
let policy = self.limit_policy.load();
|
||||||
let mode = Self::mode_from_u8(self.limit_mode.load(Ordering::Relaxed));
|
let limit = Self::user_limit(&policy, username);
|
||||||
let window = self.limit_window();
|
let mode = policy.mode;
|
||||||
|
let window = Self::limit_window(&policy);
|
||||||
let now = Instant::now();
|
let now = Instant::now();
|
||||||
|
|
||||||
let shard_idx = Self::shard_idx(username);
|
let shard_idx = Self::shard_idx(username);
|
||||||
|
|||||||
@@ -21,7 +21,8 @@ impl UserIpTracker {
|
|||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
|
||||||
let window = self.limit_window();
|
let policy = self.limit_policy.load();
|
||||||
|
let window = Self::limit_window(&policy);
|
||||||
let now = Instant::now();
|
let now = Instant::now();
|
||||||
for shard_lock in self.shards.iter() {
|
for shard_lock in self.shards.iter() {
|
||||||
let mut shard = shard_lock.write().await;
|
let mut shard = shard_lock.write().await;
|
||||||
@@ -113,7 +114,8 @@ impl UserIpTracker {
|
|||||||
&self,
|
&self,
|
||||||
users: &[String],
|
users: &[String],
|
||||||
) -> HashMap<String, usize> {
|
) -> HashMap<String, usize> {
|
||||||
let window = self.limit_window();
|
let policy = self.limit_policy.load();
|
||||||
|
let window = Self::limit_window(&policy);
|
||||||
let now = Instant::now();
|
let now = Instant::now();
|
||||||
|
|
||||||
let mut counts = HashMap::with_capacity(users.len());
|
let mut counts = HashMap::with_capacity(users.len());
|
||||||
@@ -152,7 +154,8 @@ impl UserIpTracker {
|
|||||||
|
|
||||||
pub async fn get_recent_ips_for_users(&self, users: &[String]) -> HashMap<String, Vec<IpAddr>> {
|
pub async fn get_recent_ips_for_users(&self, users: &[String]) -> HashMap<String, Vec<IpAddr>> {
|
||||||
self.drain_cleanup_queue().await;
|
self.drain_cleanup_queue().await;
|
||||||
let window = self.limit_window();
|
let policy = self.limit_policy.load();
|
||||||
|
let window = Self::limit_window(&policy);
|
||||||
let now = Instant::now();
|
let now = Instant::now();
|
||||||
|
|
||||||
let mut out = HashMap::with_capacity(users.len());
|
let mut out = HashMap::with_capacity(users.len());
|
||||||
@@ -202,6 +205,7 @@ impl UserIpTracker {
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) async fn get_stats_snapshot(&self) -> Vec<(String, usize, usize)> {
|
pub(crate) async fn get_stats_snapshot(&self) -> Vec<(String, usize, usize)> {
|
||||||
|
let policy = self.limit_policy.load();
|
||||||
let mut active_counts = Vec::new();
|
let mut active_counts = Vec::new();
|
||||||
for shard_lock in self.shards.iter() {
|
for shard_lock in self.shards.iter() {
|
||||||
let shard = shard_lock.read().await;
|
let shard = shard_lock.read().await;
|
||||||
@@ -215,7 +219,7 @@ impl UserIpTracker {
|
|||||||
|
|
||||||
let mut stats = Vec::with_capacity(active_counts.len());
|
let mut stats = Vec::with_capacity(active_counts.len());
|
||||||
for (username, active_count) in active_counts {
|
for (username, active_count) in active_counts {
|
||||||
let limit = self.user_limit(&username).unwrap_or(0);
|
let limit = Self::user_limit(&policy, &username).unwrap_or(0);
|
||||||
stats.push((username, active_count, limit));
|
stats.push((username, active_count, limit));
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -273,7 +277,8 @@ impl UserIpTracker {
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub async fn get_user_limit(&self, username: &str) -> Option<usize> {
|
pub async fn get_user_limit(&self, username: &str) -> Option<usize> {
|
||||||
self.user_limit(username)
|
let policy = self.limit_policy.load();
|
||||||
|
Self::user_limit(&policy, username)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn format_stats(&self) -> String {
|
pub async fn format_stats(&self) -> String {
|
||||||
|
|||||||
@@ -1,5 +1,6 @@
|
|||||||
use super::*;
|
use super::*;
|
||||||
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
|
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
|
||||||
|
use std::sync::atomic::AtomicBool;
|
||||||
use std::sync::atomic::Ordering;
|
use std::sync::atomic::Ordering;
|
||||||
|
|
||||||
fn test_ipv4(oct1: u8, oct2: u8, oct3: u8, oct4: u8) -> IpAddr {
|
fn test_ipv4(oct1: u8, oct2: u8, oct3: u8, oct4: u8) -> IpAddr {
|
||||||
@@ -232,6 +233,55 @@ async fn test_load_limits_replaces_previous_map() {
|
|||||||
assert_eq!(tracker.get_user_limit("user2").await, Some(5));
|
assert_eq!(tracker.get_user_limit("user2").await, Some(5));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
||||||
|
async fn concurrent_policy_replacement_never_exposes_partial_limit_map() {
|
||||||
|
const USER_COUNT: usize = 4_096;
|
||||||
|
const REPLACEMENTS: usize = 32;
|
||||||
|
|
||||||
|
let tracker = Arc::new(UserIpTracker::new());
|
||||||
|
let first = (0..USER_COUNT)
|
||||||
|
.map(|index| (format!("user-{index}"), 3usize))
|
||||||
|
.collect::<HashMap<_, _>>();
|
||||||
|
let second = (0..USER_COUNT)
|
||||||
|
.map(|index| (format!("user-{index}"), 5usize))
|
||||||
|
.collect::<HashMap<_, _>>();
|
||||||
|
tracker.load_limits(7, &first).await;
|
||||||
|
|
||||||
|
let running = Arc::new(AtomicBool::new(true));
|
||||||
|
let writer_tracker = Arc::clone(&tracker);
|
||||||
|
let writer_running = Arc::clone(&running);
|
||||||
|
let writer = tokio::spawn(async move {
|
||||||
|
for _ in 0..REPLACEMENTS {
|
||||||
|
writer_tracker.load_limits(7, &second).await;
|
||||||
|
tokio::task::yield_now().await;
|
||||||
|
writer_tracker.load_limits(7, &first).await;
|
||||||
|
tokio::task::yield_now().await;
|
||||||
|
}
|
||||||
|
writer_running.store(false, Ordering::Release);
|
||||||
|
});
|
||||||
|
|
||||||
|
let mut readers = Vec::new();
|
||||||
|
for reader in 0..3usize {
|
||||||
|
let reader_tracker = Arc::clone(&tracker);
|
||||||
|
let reader_running = Arc::clone(&running);
|
||||||
|
readers.push(tokio::spawn(async move {
|
||||||
|
let mut index = reader;
|
||||||
|
while reader_running.load(Ordering::Acquire) {
|
||||||
|
let username = format!("user-{}", index % USER_COUNT);
|
||||||
|
let limit = reader_tracker.get_user_limit(&username).await;
|
||||||
|
assert!(matches!(limit, Some(3 | 5)), "partial policy: {limit:?}");
|
||||||
|
index = index.wrapping_add(17);
|
||||||
|
tokio::task::yield_now().await;
|
||||||
|
}
|
||||||
|
}));
|
||||||
|
}
|
||||||
|
|
||||||
|
writer.await.unwrap();
|
||||||
|
for reader in readers {
|
||||||
|
reader.await.unwrap();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_global_each_limit_applies_without_user_override() {
|
async fn test_global_each_limit_applies_without_user_override() {
|
||||||
let tracker = UserIpTracker::new();
|
let tracker = UserIpTracker::new();
|
||||||
|
|||||||
@@ -253,10 +253,6 @@ pub(super) async fn bootstrap(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
if let Err(e) = crate::network::dns_overrides::install_entries(&config.network.dns_overrides) {
|
|
||||||
eprintln!("[telemt] Invalid network.dns_overrides: {}", e);
|
|
||||||
std::process::exit(1);
|
|
||||||
}
|
|
||||||
set_maestro_colors_enabled(!config.general.disable_colors);
|
set_maestro_colors_enabled(!config.general.disable_colors);
|
||||||
startup_tracker
|
startup_tracker
|
||||||
.complete_component(COMPONENT_CONFIG_LOAD, Some("config is ready".to_string()))
|
.complete_component(COMPONENT_CONFIG_LOAD, Some("config is ready".to_string()))
|
||||||
|
|||||||
@@ -0,0 +1,217 @@
|
|||||||
|
use std::future::Future;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
|
||||||
|
use std::time::Duration;
|
||||||
|
|
||||||
|
use tokio::sync::Notify;
|
||||||
|
use tokio_util::sync::CancellationToken;
|
||||||
|
use tokio_util::task::TaskTracker;
|
||||||
|
|
||||||
|
const CONTROL_TASK_ADMISSION_CLOSED: usize = 1 << (usize::BITS - 1);
|
||||||
|
const CONTROL_TASK_REGISTRATION_COUNT: usize = CONTROL_TASK_ADMISSION_CLOSED - 1;
|
||||||
|
|
||||||
|
struct ControlTaskAdmission {
|
||||||
|
state: AtomicUsize,
|
||||||
|
registrations_drained: Notify,
|
||||||
|
}
|
||||||
|
|
||||||
|
struct ControlTaskRegistration<'a> {
|
||||||
|
admission: &'a ControlTaskAdmission,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ControlTaskAdmission {
|
||||||
|
fn new() -> Self {
|
||||||
|
Self {
|
||||||
|
state: AtomicUsize::new(0),
|
||||||
|
registrations_drained: Notify::new(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn try_register(&self) -> Option<ControlTaskRegistration<'_>> {
|
||||||
|
let mut state = self.state.load(Ordering::Acquire);
|
||||||
|
loop {
|
||||||
|
if state & CONTROL_TASK_ADMISSION_CLOSED != 0
|
||||||
|
|| state & CONTROL_TASK_REGISTRATION_COUNT == CONTROL_TASK_REGISTRATION_COUNT
|
||||||
|
{
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
match self.state.compare_exchange_weak(
|
||||||
|
state,
|
||||||
|
state + 1,
|
||||||
|
Ordering::AcqRel,
|
||||||
|
Ordering::Acquire,
|
||||||
|
) {
|
||||||
|
Ok(_) => return Some(ControlTaskRegistration { admission: self }),
|
||||||
|
Err(observed) => state = observed,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn close(&self) {
|
||||||
|
self.state
|
||||||
|
.fetch_or(CONTROL_TASK_ADMISSION_CLOSED, Ordering::AcqRel);
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn wait_for_registrations(&self) {
|
||||||
|
loop {
|
||||||
|
let notified = self.registrations_drained.notified();
|
||||||
|
tokio::pin!(notified);
|
||||||
|
notified.as_mut().enable();
|
||||||
|
if self.state.load(Ordering::Acquire) & CONTROL_TASK_REGISTRATION_COUNT == 0 {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
notified.await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Drop for ControlTaskRegistration<'_> {
|
||||||
|
fn drop(&mut self) {
|
||||||
|
let previous = self.admission.state.fetch_sub(1, Ordering::AcqRel);
|
||||||
|
if previous & CONTROL_TASK_REGISTRATION_COUNT == 1 {
|
||||||
|
self.admission.registrations_drained.notify_waiters();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
struct ProcessControlPlaneInner {
|
||||||
|
admission: ControlTaskAdmission,
|
||||||
|
cancellation: CancellationToken,
|
||||||
|
tasks: TaskTracker,
|
||||||
|
shutdown_completed: AtomicBool,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Process-owned cancellation and join scope for API, metrics, and signal tasks.
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub(crate) struct ProcessControlPlane {
|
||||||
|
inner: Arc<ProcessControlPlaneInner>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ProcessControlPlane {
|
||||||
|
/// Creates an open process control-plane scope.
|
||||||
|
pub(crate) fn new() -> Self {
|
||||||
|
Self {
|
||||||
|
inner: Arc::new(ProcessControlPlaneInner {
|
||||||
|
admission: ControlTaskAdmission::new(),
|
||||||
|
cancellation: CancellationToken::new(),
|
||||||
|
tasks: TaskTracker::new(),
|
||||||
|
shutdown_completed: AtomicBool::new(false),
|
||||||
|
}),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Registers a cancellable process control-plane task before it can be unpolled.
|
||||||
|
pub(crate) fn spawn<F>(&self, future: F) -> Result<(), F>
|
||||||
|
where
|
||||||
|
F: Future<Output = ()> + Send + 'static,
|
||||||
|
{
|
||||||
|
let Some(registration) = self.inner.admission.try_register() else {
|
||||||
|
return Err(future);
|
||||||
|
};
|
||||||
|
let cancellation = self.inner.cancellation.clone();
|
||||||
|
self.inner.tasks.spawn(async move {
|
||||||
|
tokio::select! {
|
||||||
|
biased;
|
||||||
|
_ = cancellation.cancelled() => {}
|
||||||
|
_ = future => {}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
drop(registration);
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Closes task admission, cancels all owned work, and joins it within the deadline.
|
||||||
|
pub(crate) async fn shutdown(&self, timeout: Duration) -> bool {
|
||||||
|
let deadline = tokio::time::Instant::now() + timeout;
|
||||||
|
self.inner.admission.close();
|
||||||
|
self.inner.cancellation.cancel();
|
||||||
|
self.inner.tasks.close();
|
||||||
|
if self.inner.shutdown_completed.load(Ordering::Acquire) {
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
let registrations_stopped = tokio::time::timeout_at(
|
||||||
|
deadline,
|
||||||
|
self.inner.admission.wait_for_registrations(),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.is_ok();
|
||||||
|
let tasks_stopped = tokio::time::timeout_at(deadline, self.inner.tasks.wait())
|
||||||
|
.await
|
||||||
|
.is_ok();
|
||||||
|
let outcome = registrations_stopped && tasks_stopped;
|
||||||
|
if outcome {
|
||||||
|
self.inner.shutdown_completed.store(true, Ordering::Release);
|
||||||
|
}
|
||||||
|
outcome
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use std::sync::Arc;
|
||||||
|
use std::sync::atomic::{AtomicBool, Ordering};
|
||||||
|
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn shutdown_cancels_owned_tasks_and_rejects_late_registration() {
|
||||||
|
struct DropSignal(Arc<AtomicBool>);
|
||||||
|
|
||||||
|
impl Drop for DropSignal {
|
||||||
|
fn drop(&mut self) {
|
||||||
|
self.0.store(true, Ordering::Release);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let scope = ProcessControlPlane::new();
|
||||||
|
let dropped = Arc::new(AtomicBool::new(false));
|
||||||
|
let drop_signal = DropSignal(dropped.clone());
|
||||||
|
assert!(
|
||||||
|
scope
|
||||||
|
.spawn(async move {
|
||||||
|
let _drop_signal = drop_signal;
|
||||||
|
std::future::pending::<()>().await;
|
||||||
|
})
|
||||||
|
.is_ok()
|
||||||
|
);
|
||||||
|
|
||||||
|
assert!(scope.shutdown(Duration::from_secs(1)).await);
|
||||||
|
assert!(dropped.load(Ordering::Acquire));
|
||||||
|
assert!(scope.spawn(async {}).is_err());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn concurrent_shutdown_callers_wait_for_completion() {
|
||||||
|
let scope = ProcessControlPlane::new();
|
||||||
|
let registration = scope.inner.admission.try_register().unwrap();
|
||||||
|
let first_scope = scope.clone();
|
||||||
|
let first = tokio::spawn(async move { first_scope.shutdown(Duration::from_secs(1)).await });
|
||||||
|
tokio::task::yield_now().await;
|
||||||
|
let second_scope = scope.clone();
|
||||||
|
let second = tokio::spawn(async move { second_scope.shutdown(Duration::from_secs(1)).await });
|
||||||
|
|
||||||
|
tokio::task::yield_now().await;
|
||||||
|
assert!(!first.is_finished());
|
||||||
|
assert!(!second.is_finished());
|
||||||
|
drop(registration);
|
||||||
|
|
||||||
|
assert!(first.await.unwrap());
|
||||||
|
assert!(second.await.unwrap());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn cancelled_shutdown_caller_cannot_orphan_the_control_plane() {
|
||||||
|
let scope = ProcessControlPlane::new();
|
||||||
|
let registration = scope.inner.admission.try_register().unwrap();
|
||||||
|
let first_scope = scope.clone();
|
||||||
|
let first = tokio::spawn(async move { first_scope.shutdown(Duration::from_secs(30)).await });
|
||||||
|
tokio::task::yield_now().await;
|
||||||
|
|
||||||
|
first.abort();
|
||||||
|
assert!(first.await.unwrap_err().is_cancelled());
|
||||||
|
assert!(scope.spawn(async {}).is_err());
|
||||||
|
drop(registration);
|
||||||
|
|
||||||
|
assert!(scope.shutdown(Duration::from_secs(1)).await);
|
||||||
|
}
|
||||||
|
}
|
||||||
+66
-10
@@ -3,7 +3,7 @@ use std::sync::Arc;
|
|||||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
|
|
||||||
use tokio::sync::{RwLock, Semaphore, watch};
|
use tokio::sync::{Notify, RwLock, Semaphore, watch};
|
||||||
use tokio_util::sync::CancellationToken;
|
use tokio_util::sync::CancellationToken;
|
||||||
use tokio_util::task::TaskTracker;
|
use tokio_util::task::TaskTracker;
|
||||||
|
|
||||||
@@ -29,6 +29,7 @@ const SESSION_REGISTRATION_COUNT: usize = SESSION_ADMISSION_CLOSED - 1;
|
|||||||
|
|
||||||
struct SessionAdmission {
|
struct SessionAdmission {
|
||||||
state: AtomicUsize,
|
state: AtomicUsize,
|
||||||
|
registrations_drained: Notify,
|
||||||
}
|
}
|
||||||
|
|
||||||
struct SessionRegistration<'a> {
|
struct SessionRegistration<'a> {
|
||||||
@@ -39,6 +40,7 @@ impl SessionAdmission {
|
|||||||
fn new() -> Self {
|
fn new() -> Self {
|
||||||
Self {
|
Self {
|
||||||
state: AtomicUsize::new(0),
|
state: AtomicUsize::new(0),
|
||||||
|
registrations_drained: Notify::new(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -73,15 +75,24 @@ impl SessionAdmission {
|
|||||||
}
|
}
|
||||||
|
|
||||||
async fn wait_for_registrations(&self) {
|
async fn wait_for_registrations(&self) {
|
||||||
while self.state.load(Ordering::Acquire) & SESSION_REGISTRATION_COUNT != 0 {
|
loop {
|
||||||
tokio::task::yield_now().await;
|
let notified = self.registrations_drained.notified();
|
||||||
|
tokio::pin!(notified);
|
||||||
|
notified.as_mut().enable();
|
||||||
|
if self.state.load(Ordering::Acquire) & SESSION_REGISTRATION_COUNT == 0 {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
notified.await;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Drop for SessionRegistration<'_> {
|
impl Drop for SessionRegistration<'_> {
|
||||||
fn drop(&mut self) {
|
fn drop(&mut self) {
|
||||||
self.admission.state.fetch_sub(1, Ordering::Release);
|
let previous = self.admission.state.fetch_sub(1, Ordering::AcqRel);
|
||||||
|
if previous & SESSION_REGISTRATION_COUNT == 1 {
|
||||||
|
self.admission.registrations_drained.notify_waiters();
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -98,6 +109,7 @@ pub(crate) struct RuntimeWatchState {
|
|||||||
pub(crate) struct RuntimeTaskScope {
|
pub(crate) struct RuntimeTaskScope {
|
||||||
tracker: TaskTracker,
|
tracker: TaskTracker,
|
||||||
cancel: CancellationToken,
|
cancel: CancellationToken,
|
||||||
|
admission: Arc<SessionAdmission>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl RuntimeTaskScope {
|
impl RuntimeTaskScope {
|
||||||
@@ -106,6 +118,7 @@ impl RuntimeTaskScope {
|
|||||||
Self {
|
Self {
|
||||||
tracker: TaskTracker::new(),
|
tracker: TaskTracker::new(),
|
||||||
cancel: CancellationToken::new(),
|
cancel: CancellationToken::new(),
|
||||||
|
admission: Arc::new(SessionAdmission::new()),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -114,9 +127,13 @@ impl RuntimeTaskScope {
|
|||||||
where
|
where
|
||||||
F: Future<Output = ()> + Send + 'static,
|
F: Future<Output = ()> + Send + 'static,
|
||||||
{
|
{
|
||||||
|
let Some(_registration) = self.admission.try_register() else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
let cancel = self.cancel.clone();
|
let cancel = self.cancel.clone();
|
||||||
self.tracker.spawn(async move {
|
self.tracker.spawn(async move {
|
||||||
tokio::select! {
|
tokio::select! {
|
||||||
|
biased;
|
||||||
_ = cancel.cancelled() => {}
|
_ = cancel.cancelled() => {}
|
||||||
_ = future => {}
|
_ = future => {}
|
||||||
}
|
}
|
||||||
@@ -130,6 +147,8 @@ impl RuntimeTaskScope {
|
|||||||
|
|
||||||
/// Cancels the scope and waits within the bounded background-task budget.
|
/// Cancels the scope and waits within the bounded background-task budget.
|
||||||
pub(crate) async fn stop(&self) {
|
pub(crate) async fn stop(&self) {
|
||||||
|
self.admission.close();
|
||||||
|
self.admission.wait_for_registrations().await;
|
||||||
self.cancel.cancel();
|
self.cancel.cancel();
|
||||||
self.tracker.close();
|
self.tracker.close();
|
||||||
let _ = tokio::time::timeout(BACKGROUND_STOP_TIMEOUT, self.tracker.wait()).await;
|
let _ = tokio::time::timeout(BACKGROUND_STOP_TIMEOUT, self.tracker.wait()).await;
|
||||||
@@ -263,6 +282,7 @@ impl RuntimeGeneration {
|
|||||||
let cancel = self.session_cancel.clone();
|
let cancel = self.session_cancel.clone();
|
||||||
self.sessions.spawn(async move {
|
self.sessions.spawn(async move {
|
||||||
tokio::select! {
|
tokio::select! {
|
||||||
|
biased;
|
||||||
_ = cancel.cancelled() => {}
|
_ = cancel.cancelled() => {}
|
||||||
_ = future => {}
|
_ = future => {}
|
||||||
}
|
}
|
||||||
@@ -275,11 +295,6 @@ impl RuntimeGeneration {
|
|||||||
self.session_admission.close();
|
self.session_admission.close();
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Reopens admission after a candidate activation rolls back.
|
|
||||||
pub(crate) fn resume_accepting_sessions(&self) {
|
|
||||||
self.session_admission.reopen();
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Waits for registered sessions and cancels them when the deadline expires.
|
/// Waits for registered sessions and cancels them when the deadline expires.
|
||||||
pub(crate) async fn drain_sessions(&self, timeout: Duration) -> bool {
|
pub(crate) async fn drain_sessions(&self, timeout: Duration) -> bool {
|
||||||
self.stop_accepting_sessions();
|
self.stop_accepting_sessions();
|
||||||
@@ -308,6 +323,27 @@ impl RuntimeGeneration {
|
|||||||
pub(crate) async fn stop_background_tasks(&self) {
|
pub(crate) async fn stop_background_tasks(&self) {
|
||||||
self.background_tasks.stop().await;
|
self.background_tasks.stop().await;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Terminally stops the generation's Middle-End task and writer scope.
|
||||||
|
pub(crate) async fn stop_middle_end(&self, timeout: Duration) -> bool {
|
||||||
|
let Some(pool) = self.current_me_pool().await else {
|
||||||
|
return true;
|
||||||
|
};
|
||||||
|
pool.shutdown_until(timeout).await
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Drop for RuntimeGeneration {
|
||||||
|
fn drop(&mut self) {
|
||||||
|
if let Some(pool) = self.me_pool.as_ref() {
|
||||||
|
pool.begin_shutdown();
|
||||||
|
}
|
||||||
|
if let Ok(pool) = self.me_pool_runtime.try_read()
|
||||||
|
&& let Some(pool) = pool.as_ref()
|
||||||
|
{
|
||||||
|
pool.begin_shutdown();
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
@@ -393,11 +429,31 @@ mod tests {
|
|||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn runtime_task_scope_joins_cancelled_background_task() {
|
async fn runtime_task_scope_joins_cancelled_background_task() {
|
||||||
|
struct DropSignal(Arc<AtomicUsize>);
|
||||||
|
|
||||||
|
impl Drop for DropSignal {
|
||||||
|
fn drop(&mut self) {
|
||||||
|
self.0.fetch_add(1, Ordering::AcqRel);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
let scope = RuntimeTaskScope::new();
|
let scope = RuntimeTaskScope::new();
|
||||||
scope.spawn(std::future::pending());
|
let dropped = Arc::new(AtomicUsize::new(0));
|
||||||
|
let drop_signal = DropSignal(dropped.clone());
|
||||||
|
scope.spawn(async move {
|
||||||
|
let _drop_signal = drop_signal;
|
||||||
|
std::future::pending::<()>().await;
|
||||||
|
});
|
||||||
tokio::time::timeout(Duration::from_secs(1), scope.stop())
|
tokio::time::timeout(Duration::from_secs(1), scope.stop())
|
||||||
.await
|
.await
|
||||||
.unwrap();
|
.unwrap();
|
||||||
|
assert_eq!(dropped.load(Ordering::Acquire), 1);
|
||||||
|
|
||||||
|
let late_drop_signal = DropSignal(dropped.clone());
|
||||||
|
scope.spawn(async move {
|
||||||
|
let _late_drop_signal = late_drop_signal;
|
||||||
|
});
|
||||||
|
assert_eq!(dropped.load(Ordering::Acquire), 2);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
|
|||||||
@@ -12,7 +12,7 @@ use tracing::{debug, error, info, warn};
|
|||||||
use crate::config::{ListenerTransport, RstOnCloseMode};
|
use crate::config::{ListenerTransport, RstOnCloseMode};
|
||||||
use crate::proxy::ClientHandler;
|
use crate::proxy::ClientHandler;
|
||||||
use crate::transport::socket::set_linger_zero;
|
use crate::transport::socket::set_linger_zero;
|
||||||
use crate::web::manager::WebProcessRuntime;
|
use crate::web::manager::{HttpConnectionAdmissionError, WebProcessRuntime};
|
||||||
use crate::web::telemetry::{WebAcceptorGuard, WebHttpConnectionOverloadOutcome};
|
use crate::web::telemetry::{WebAcceptorGuard, WebHttpConnectionOverloadOutcome};
|
||||||
|
|
||||||
use super::bind::BoundTcpListener;
|
use super::bind::BoundTcpListener;
|
||||||
@@ -208,45 +208,82 @@ async fn run_accept_loop(
|
|||||||
return;
|
return;
|
||||||
};
|
};
|
||||||
web_runtime.telemetry().record_accept();
|
web_runtime.telemetry().record_accept();
|
||||||
let Some(connection_permit) = web_runtime.try_http_connection() else {
|
if cancellation.is_cancelled() {
|
||||||
let config = web_runtime.active_generation().config();
|
drop(stream);
|
||||||
let action = config.web.http_connection_capacity_action;
|
continue;
|
||||||
let phase_timeout =
|
}
|
||||||
Duration::from_millis(config.web.timeouts.http_overload_timeout_ms);
|
if web_runtime.is_shutdown() {
|
||||||
drop(config);
|
web_runtime
|
||||||
if action == crate::config::WebHttpConnectionCapacityAction::Drop {
|
.telemetry()
|
||||||
web_runtime.telemetry().record_rejection(
|
.record_rejection(
|
||||||
crate::web::telemetry::WebRejectionReason::HttpConnectionCapacity,
|
crate::web::telemetry::WebRejectionReason::RuntimeClosed,
|
||||||
|
);
|
||||||
|
drop(stream);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let connection_permit = match web_runtime.try_http_connection() {
|
||||||
|
Ok(permit) => permit,
|
||||||
|
Err(HttpConnectionAdmissionError::Closed) => {
|
||||||
|
web_runtime.telemetry().record_rejection(
|
||||||
|
crate::web::telemetry::WebRejectionReason::RuntimeClosed,
|
||||||
);
|
);
|
||||||
web_runtime
|
|
||||||
.telemetry()
|
|
||||||
.record_overload(WebHttpConnectionOverloadOutcome::Dropped);
|
|
||||||
drop(stream);
|
drop(stream);
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
let Some(overload_permit) = web_runtime.try_http_overload_connection()
|
Err(HttpConnectionAdmissionError::AtCapacity) => {
|
||||||
else {
|
let config = web_runtime.active_generation().config();
|
||||||
web_runtime.telemetry().record_rejection(
|
let action = config.web.http_connection_capacity_action;
|
||||||
crate::web::telemetry::WebRejectionReason::HttpConnectionCapacity,
|
let phase_timeout = Duration::from_millis(
|
||||||
|
config.web.timeouts.http_overload_timeout_ms,
|
||||||
);
|
);
|
||||||
web_runtime.telemetry().record_overload(
|
drop(config);
|
||||||
WebHttpConnectionOverloadOutcome::OverflowCapacityDrop,
|
if action == crate::config::WebHttpConnectionCapacityAction::Drop {
|
||||||
);
|
web_runtime.telemetry().record_rejection(
|
||||||
drop(stream);
|
crate::web::telemetry::WebRejectionReason::HttpConnectionCapacity,
|
||||||
|
);
|
||||||
|
web_runtime
|
||||||
|
.telemetry()
|
||||||
|
.record_overload(WebHttpConnectionOverloadOutcome::Dropped);
|
||||||
|
drop(stream);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let overload_permit =
|
||||||
|
match web_runtime.try_http_overload_connection() {
|
||||||
|
Ok(permit) => permit,
|
||||||
|
Err(HttpConnectionAdmissionError::Closed) => {
|
||||||
|
web_runtime.telemetry().record_rejection(
|
||||||
|
crate::web::telemetry::WebRejectionReason::RuntimeClosed,
|
||||||
|
);
|
||||||
|
web_runtime.telemetry().record_overload(
|
||||||
|
WebHttpConnectionOverloadOutcome::ShutdownDrop,
|
||||||
|
);
|
||||||
|
drop(stream);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
Err(HttpConnectionAdmissionError::AtCapacity) => {
|
||||||
|
web_runtime.telemetry().record_rejection(
|
||||||
|
crate::web::telemetry::WebRejectionReason::HttpConnectionCapacity,
|
||||||
|
);
|
||||||
|
web_runtime.telemetry().record_overload(
|
||||||
|
WebHttpConnectionOverloadOutcome::OverflowCapacityDrop,
|
||||||
|
);
|
||||||
|
drop(stream);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
connections.spawn(web_overload::serve(
|
||||||
|
stream,
|
||||||
|
peer_addr,
|
||||||
|
spec.web_client_ip_source,
|
||||||
|
Arc::clone(&spec.web_trusted_proxy_cidrs),
|
||||||
|
Arc::clone(web_runtime),
|
||||||
|
cancellation.clone(),
|
||||||
|
overload_permit,
|
||||||
|
action,
|
||||||
|
phase_timeout,
|
||||||
|
));
|
||||||
continue;
|
continue;
|
||||||
};
|
}
|
||||||
connections.spawn(web_overload::serve(
|
|
||||||
stream,
|
|
||||||
peer_addr,
|
|
||||||
spec.web_client_ip_source,
|
|
||||||
Arc::clone(&spec.web_trusted_proxy_cidrs),
|
|
||||||
Arc::clone(web_runtime),
|
|
||||||
cancellation.clone(),
|
|
||||||
overload_permit,
|
|
||||||
action,
|
|
||||||
phase_timeout,
|
|
||||||
));
|
|
||||||
continue;
|
|
||||||
};
|
};
|
||||||
connections.spawn(crate::web::http::serve_connection(
|
connections.spawn(crate::web::http::serve_connection(
|
||||||
stream,
|
stream,
|
||||||
|
|||||||
@@ -9,9 +9,10 @@ use tokio::sync::OwnedSemaphorePermit;
|
|||||||
use tokio_util::sync::CancellationToken;
|
use tokio_util::sync::CancellationToken;
|
||||||
|
|
||||||
use crate::config::{WebClientIpSource, WebHttpConnectionCapacityAction};
|
use crate::config::{WebClientIpSource, WebHttpConnectionCapacityAction};
|
||||||
use crate::web::manager::WebProcessRuntime;
|
use crate::web::manager::{HttpConnectionAdmissionError, WebProcessRuntime};
|
||||||
use crate::web::telemetry::{WebHttpConnectionOverloadOutcome, WebRejectionReason};
|
use crate::web::telemetry::{WebHttpConnectionOverloadOutcome, WebRejectionReason};
|
||||||
|
|
||||||
|
/// Exact bounded retryable response emitted before HTTP request parsing.
|
||||||
pub(super) const SERVICE_UNAVAILABLE_RESPONSE: &[u8] = b"HTTP/1.1 503 Service Unavailable\r\nContent-Length: 0\r\nCache-Control: no-store\r\nRetry-After: 1\r\nConnection: close\r\n\r\n";
|
pub(super) const SERVICE_UNAVAILABLE_RESPONSE: &[u8] = b"HTTP/1.1 503 Service Unavailable\r\nContent-Length: 0\r\nCache-Control: no-store\r\nRetry-After: 1\r\nConnection: close\r\n\r\n";
|
||||||
|
|
||||||
/// Handles one accepted WEB socket outside ordinary connection capacity.
|
/// Handles one accepted WEB socket outside ordinary connection capacity.
|
||||||
@@ -44,25 +45,31 @@ pub(super) async fn serve(
|
|||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
permit = tokio::time::timeout(phase_timeout, runtime.acquire_http_connection()) => {
|
permit = tokio::time::timeout(phase_timeout, runtime.acquire_http_connection()) => {
|
||||||
permit.ok().flatten()
|
permit
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
let Some(connection_permit) = connection_permit else {
|
let connection_permit = match connection_permit {
|
||||||
if runtime.is_shutdown() {
|
Ok(Ok(permit)) => permit,
|
||||||
|
Ok(Err(HttpConnectionAdmissionError::Closed)) => {
|
||||||
|
runtime
|
||||||
|
.telemetry()
|
||||||
|
.record_rejection(WebRejectionReason::RuntimeClosed);
|
||||||
runtime
|
runtime
|
||||||
.telemetry()
|
.telemetry()
|
||||||
.record_overload(WebHttpConnectionOverloadOutcome::ShutdownDrop);
|
.record_overload(WebHttpConnectionOverloadOutcome::ShutdownDrop);
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
let outcome = match respond(stream, &cancellation, phase_timeout).await {
|
Ok(Err(HttpConnectionAdmissionError::AtCapacity)) | Err(_) => {
|
||||||
WebHttpConnectionOverloadOutcome::Responded503 => {
|
let outcome = match respond(stream, &cancellation, phase_timeout).await {
|
||||||
WebHttpConnectionOverloadOutcome::WaitTimeout503
|
WebHttpConnectionOverloadOutcome::Responded503 => {
|
||||||
}
|
WebHttpConnectionOverloadOutcome::WaitTimeout503
|
||||||
other => other,
|
}
|
||||||
};
|
other => other,
|
||||||
record_final_capacity_rejection(&runtime, outcome);
|
};
|
||||||
runtime.telemetry().record_overload(outcome);
|
record_final_capacity_rejection(&runtime, outcome);
|
||||||
return;
|
runtime.telemetry().record_overload(outcome);
|
||||||
|
return;
|
||||||
|
}
|
||||||
};
|
};
|
||||||
runtime
|
runtime
|
||||||
.telemetry()
|
.telemetry()
|
||||||
|
|||||||
@@ -345,6 +345,7 @@ pub(crate) async fn initialize_me_pool(
|
|||||||
config.general.me_route_blocking_send_timeout_ms,
|
config.general.me_route_blocking_send_timeout_ms,
|
||||||
config.general.me_route_inline_recovery_attempts,
|
config.general.me_route_inline_recovery_attempts,
|
||||||
config.general.me_route_inline_recovery_wait_ms,
|
config.general.me_route_inline_recovery_wait_ms,
|
||||||
|
(config.server.max_connections as usize).saturating_add(128),
|
||||||
);
|
);
|
||||||
startup_tracker
|
startup_tracker
|
||||||
.complete_component(
|
.complete_component(
|
||||||
|
|||||||
@@ -6,6 +6,7 @@
|
|||||||
// - admission: conditional-cast gate and route mode switching.
|
// - admission: conditional-cast gate and route mode switching.
|
||||||
// - bootstrap: configuration and tracing initialization.
|
// - bootstrap: configuration and tracing initialization.
|
||||||
// - connectivity: startup ME/DC connectivity diagnostics.
|
// - connectivity: startup ME/DC connectivity diagnostics.
|
||||||
|
// - control_plane: process-owned API, metrics, and signal task lifecycle.
|
||||||
// - generation: runtime generation state and task ownership.
|
// - generation: runtime generation state and task ownership.
|
||||||
// - helpers: CLI and shared startup/runtime helper routines.
|
// - helpers: CLI and shared startup/runtime helper routines.
|
||||||
// - listeners: TCP/Unix listener planning, binding, and lifecycle control.
|
// - listeners: TCP/Unix listener planning, binding, and lifecycle control.
|
||||||
@@ -21,6 +22,7 @@
|
|||||||
mod admission;
|
mod admission;
|
||||||
mod bootstrap;
|
mod bootstrap;
|
||||||
mod connectivity;
|
mod connectivity;
|
||||||
|
pub(crate) mod control_plane;
|
||||||
pub(crate) mod generation;
|
pub(crate) mod generation;
|
||||||
mod helpers;
|
mod helpers;
|
||||||
mod listeners;
|
mod listeners;
|
||||||
|
|||||||
+65
-18
@@ -1,9 +1,10 @@
|
|||||||
|
use std::collections::BTreeSet;
|
||||||
use std::net::{IpAddr, SocketAddr};
|
use std::net::{IpAddr, SocketAddr};
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
|
|
||||||
use arc_swap::ArcSwap;
|
use arc_swap::ArcSwap;
|
||||||
use tokio::sync::{RwLock, watch};
|
use tokio::sync::{RwLock, watch};
|
||||||
use tracing::{error, info, warn};
|
use tracing::{error, info};
|
||||||
|
|
||||||
use crate::api;
|
use crate::api;
|
||||||
use crate::ip_tracker::UserIpTracker;
|
use crate::ip_tracker::UserIpTracker;
|
||||||
@@ -15,14 +16,15 @@ use crate::startup::{COMPONENT_API_BOOTSTRAP, COMPONENT_NETWORK_PROBE};
|
|||||||
use crate::stats::telemetry::TelemetryPolicy;
|
use crate::stats::telemetry::TelemetryPolicy;
|
||||||
use crate::stats::{QuotaStore, Stats};
|
use crate::stats::{QuotaStore, Stats};
|
||||||
use crate::synlimit_control;
|
use crate::synlimit_control;
|
||||||
|
use crate::tls_front::cache::TlsFullCertBudget;
|
||||||
use crate::transport::UpstreamManager;
|
use crate::transport::UpstreamManager;
|
||||||
use crate::transport::middle_proxy::MePool;
|
use crate::transport::middle_proxy::MePool;
|
||||||
use crate::web::control::WebRuntimeControl;
|
use crate::web::control::WebRuntimeControl;
|
||||||
use crate::web::trace::WebTraceStore;
|
use crate::web::trace::WebTraceStore;
|
||||||
|
|
||||||
use super::{
|
use super::{
|
||||||
bootstrap, generation, listeners, reload, reload_supervisor, runtime_startup, runtime_tasks,
|
bootstrap, control_plane, generation, listeners, reload, reload_supervisor, runtime_startup,
|
||||||
shutdown, tls_bootstrap,
|
runtime_tasks, shutdown, tls_bootstrap,
|
||||||
};
|
};
|
||||||
|
|
||||||
// Shared maestro startup and main loop. `drop_after_bind` runs on Unix after listeners are bound
|
// Shared maestro startup and main loop. `drop_after_bind` runs on Unix after listeners are bound
|
||||||
@@ -45,10 +47,22 @@ pub(super) async fn run_telemt_core(
|
|||||||
|
|
||||||
let quota_store = Arc::new(QuotaStore::default());
|
let quota_store = Arc::new(QuotaStore::default());
|
||||||
let stats = Arc::new(Stats::with_quota_store(quota_store.clone()));
|
let stats = Arc::new(Stats::with_quota_store(quota_store.clone()));
|
||||||
|
let tls_full_cert_budget = Arc::new(TlsFullCertBudget::new());
|
||||||
|
let process_control_plane = control_plane::ProcessControlPlane::new();
|
||||||
let runtime_task_scope = generation::RuntimeTaskScope::new();
|
let runtime_task_scope = generation::RuntimeTaskScope::new();
|
||||||
stats.apply_telemetry_policy(TelemetryPolicy::from_config(&config.general.telemetry));
|
stats.apply_telemetry_policy(TelemetryPolicy::from_config(&config.general.telemetry));
|
||||||
let quota_state_path = config.general.quota_state_path.clone();
|
let quota_state_path = config.general.quota_state_path.clone();
|
||||||
crate::quota_state::load_quota_state("a_state_path, stats.as_ref()).await;
|
let quota_state = crate::quota_state::QuotaStateOwner::new(
|
||||||
|
quota_state_path,
|
||||||
|
quota_store.clone(),
|
||||||
|
);
|
||||||
|
let configured_quota_users = config
|
||||||
|
.access
|
||||||
|
.users
|
||||||
|
.keys()
|
||||||
|
.cloned()
|
||||||
|
.collect::<BTreeSet<_>>();
|
||||||
|
quota_state.load(&configured_quota_users).await;
|
||||||
|
|
||||||
let upstream_manager = Arc::new(
|
let upstream_manager = Arc::new(
|
||||||
UpstreamManager::new(
|
UpstreamManager::new(
|
||||||
@@ -136,15 +150,29 @@ pub(super) async fn run_telemt_core(
|
|||||||
let listen = match config.server.api.listen.parse::<SocketAddr>() {
|
let listen = match config.server.api.listen.parse::<SocketAddr>() {
|
||||||
Ok(listen) => listen,
|
Ok(listen) => listen,
|
||||||
Err(error) => {
|
Err(error) => {
|
||||||
warn!(
|
let message = format!(
|
||||||
error = %error,
|
"invalid server.api.listen \"{}\": {}",
|
||||||
listen = %config.server.api.listen,
|
config.server.api.listen, error
|
||||||
"Invalid server.api.listen; API is disabled"
|
|
||||||
);
|
);
|
||||||
SocketAddr::from(([127, 0, 0, 1], 0))
|
startup_tracker
|
||||||
|
.fail_component(COMPONENT_API_BOOTSTRAP, Some(message.clone()))
|
||||||
|
.await;
|
||||||
|
return Err(std::io::Error::new(std::io::ErrorKind::InvalidInput, message).into());
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
if listen.port() != 0 {
|
if listen.port() != 0 {
|
||||||
|
let api_listener = match tokio::net::TcpListener::bind(listen).await {
|
||||||
|
Ok(listener) => listener,
|
||||||
|
Err(error) => {
|
||||||
|
startup_tracker
|
||||||
|
.fail_component(
|
||||||
|
COMPONENT_API_BOOTSTRAP,
|
||||||
|
Some(format!("API listener bind failed on {listen}: {error}")),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
return Err(error.into());
|
||||||
|
}
|
||||||
|
};
|
||||||
let stats_api = stats.clone();
|
let stats_api = stats.clone();
|
||||||
let ip_tracker_api = ip_tracker.clone();
|
let ip_tracker_api = ip_tracker.clone();
|
||||||
let me_pool_api = api_me_pool.clone();
|
let me_pool_api = api_me_pool.clone();
|
||||||
@@ -152,7 +180,7 @@ pub(super) async fn run_telemt_core(
|
|||||||
let route_runtime_api = route_runtime.clone();
|
let route_runtime_api = route_runtime.clone();
|
||||||
let proxy_shared_api = shared_state.clone();
|
let proxy_shared_api = shared_state.clone();
|
||||||
let config_path_api = config_path.clone();
|
let config_path_api = config_path.clone();
|
||||||
let quota_state_path_api = quota_state_path.clone();
|
let quota_state_api = quota_state.clone();
|
||||||
let startup_tracker_api = startup_tracker.clone();
|
let startup_tracker_api = startup_tracker.clone();
|
||||||
let detected_ips_rx_api = detected_ips_rx.clone();
|
let detected_ips_rx_api = detected_ips_rx.clone();
|
||||||
let reload_control_api = reload_control.clone();
|
let reload_control_api = reload_control.clone();
|
||||||
@@ -160,9 +188,11 @@ pub(super) async fn run_telemt_core(
|
|||||||
let runtime_watch_rx_api = runtime_watch_rx.clone();
|
let runtime_watch_rx_api = runtime_watch_rx.clone();
|
||||||
let web_trace_api = web_trace.clone();
|
let web_trace_api = web_trace.clone();
|
||||||
let web_runtime_rx_api = web_runtime_control.subscribe();
|
let web_runtime_rx_api = web_runtime_control.subscribe();
|
||||||
tokio::spawn(async move {
|
let api_control_plane = process_control_plane.clone();
|
||||||
|
let api_task_control_plane = process_control_plane.clone();
|
||||||
|
let api_task = async move {
|
||||||
api::serve(
|
api::serve(
|
||||||
listen,
|
api_listener,
|
||||||
stats_api,
|
stats_api,
|
||||||
ip_tracker_api,
|
ip_tracker_api,
|
||||||
me_pool_api,
|
me_pool_api,
|
||||||
@@ -170,7 +200,7 @@ pub(super) async fn run_telemt_core(
|
|||||||
proxy_shared_api,
|
proxy_shared_api,
|
||||||
upstream_manager_api,
|
upstream_manager_api,
|
||||||
config_path_api,
|
config_path_api,
|
||||||
quota_state_path_api,
|
quota_state_api,
|
||||||
detected_ips_rx_api,
|
detected_ips_rx_api,
|
||||||
process_started_at_epoch_secs,
|
process_started_at_epoch_secs,
|
||||||
startup_tracker_api,
|
startup_tracker_api,
|
||||||
@@ -179,13 +209,21 @@ pub(super) async fn run_telemt_core(
|
|||||||
runtime_watch_rx_api,
|
runtime_watch_rx_api,
|
||||||
web_trace_api,
|
web_trace_api,
|
||||||
web_runtime_rx_api,
|
web_runtime_rx_api,
|
||||||
|
api_task_control_plane,
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
});
|
};
|
||||||
|
if api_control_plane.spawn(api_task).is_err() {
|
||||||
|
let message = "process control-plane task admission closed during API startup";
|
||||||
|
startup_tracker
|
||||||
|
.fail_component(COMPONENT_API_BOOTSTRAP, Some(message.to_string()))
|
||||||
|
.await;
|
||||||
|
return Err(std::io::Error::other(message).into());
|
||||||
|
}
|
||||||
startup_tracker
|
startup_tracker
|
||||||
.complete_component(
|
.complete_component(
|
||||||
COMPONENT_API_BOOTSTRAP,
|
COMPONENT_API_BOOTSTRAP,
|
||||||
Some(format!("api task spawned on {}", listen)),
|
Some(format!("API listener bound and supervised on {}", listen)),
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
} else {
|
} else {
|
||||||
@@ -219,6 +257,7 @@ pub(super) async fn run_telemt_core(
|
|||||||
upstream_manager.clone(),
|
upstream_manager.clone(),
|
||||||
&startup_tracker,
|
&startup_tracker,
|
||||||
runtime_task_scope.clone(),
|
runtime_task_scope.clone(),
|
||||||
|
tls_full_cert_budget.clone(),
|
||||||
tls_bootstrap::TlsBootstrapPolicy::BestEffort,
|
tls_bootstrap::TlsBootstrapPolicy::BestEffort,
|
||||||
)
|
)
|
||||||
.await?;
|
.await?;
|
||||||
@@ -316,8 +355,10 @@ pub(super) async fn run_telemt_core(
|
|||||||
&startup_tracker,
|
&startup_tracker,
|
||||||
active_runtime.clone(),
|
active_runtime.clone(),
|
||||||
web_runtime_control.subscribe(),
|
web_runtime_control.subscribe(),
|
||||||
|
tls_full_cert_budget.clone(),
|
||||||
|
process_control_plane.clone(),
|
||||||
)
|
)
|
||||||
.await;
|
.await?;
|
||||||
|
|
||||||
runtime_watch_tx.send_replace(Some(active_runtime.load_full().watch_state()));
|
runtime_watch_tx.send_replace(Some(active_runtime.load_full().watch_state()));
|
||||||
active_runtime_tx.send_replace(Some(active_runtime.clone()));
|
active_runtime_tx.send_replace(Some(active_runtime.clone()));
|
||||||
@@ -335,6 +376,7 @@ pub(super) async fn run_telemt_core(
|
|||||||
reload_commands,
|
reload_commands,
|
||||||
config_path,
|
config_path,
|
||||||
quota_store,
|
quota_store,
|
||||||
|
tls_full_cert_budget,
|
||||||
detected_ips_tx,
|
detected_ips_tx,
|
||||||
runtime_log_filter,
|
runtime_log_filter,
|
||||||
runtime_watch_tx,
|
runtime_watch_tx,
|
||||||
@@ -342,12 +384,17 @@ pub(super) async fn run_telemt_core(
|
|||||||
web_trace,
|
web_trace,
|
||||||
);
|
);
|
||||||
|
|
||||||
shutdown::spawn_signal_handlers(active_runtime.clone(), process_started_at);
|
shutdown::spawn_signal_handlers(
|
||||||
|
active_runtime.clone(),
|
||||||
|
process_started_at,
|
||||||
|
process_control_plane.clone(),
|
||||||
|
);
|
||||||
shutdown::wait_for_shutdown(
|
shutdown::wait_for_shutdown(
|
||||||
process_started_at,
|
process_started_at,
|
||||||
active_runtime,
|
active_runtime,
|
||||||
quota_state_path,
|
quota_state,
|
||||||
reload_supervisor,
|
reload_supervisor,
|
||||||
|
process_control_plane,
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
|
|||||||
@@ -9,6 +9,7 @@ use tokio_util::sync::CancellationToken;
|
|||||||
use tracing::{info, warn};
|
use tracing::{info, warn};
|
||||||
|
|
||||||
use crate::stats::QuotaStore;
|
use crate::stats::QuotaStore;
|
||||||
|
use crate::tls_front::cache::TlsFullCertBudget;
|
||||||
use crate::web::trace::WebTraceStore;
|
use crate::web::trace::WebTraceStore;
|
||||||
|
|
||||||
use super::generation::{RuntimeGeneration, RuntimeWatchState};
|
use super::generation::{RuntimeGeneration, RuntimeWatchState};
|
||||||
@@ -26,6 +27,7 @@ pub(crate) struct ReloadSupervisor {
|
|||||||
commands: ReloadCommandReceiver,
|
commands: ReloadCommandReceiver,
|
||||||
config_path: PathBuf,
|
config_path: PathBuf,
|
||||||
quota_store: Arc<QuotaStore>,
|
quota_store: Arc<QuotaStore>,
|
||||||
|
tls_full_cert_budget: Arc<TlsFullCertBudget>,
|
||||||
detected_ips_tx: watch::Sender<(Option<std::net::IpAddr>, Option<std::net::IpAddr>)>,
|
detected_ips_tx: watch::Sender<(Option<std::net::IpAddr>, Option<std::net::IpAddr>)>,
|
||||||
runtime_log_filter: RuntimeLogFilter,
|
runtime_log_filter: RuntimeLogFilter,
|
||||||
runtime_watch_tx: watch::Sender<Option<RuntimeWatchState>>,
|
runtime_watch_tx: watch::Sender<Option<RuntimeWatchState>>,
|
||||||
@@ -81,12 +83,7 @@ fn revision_gate_action(
|
|||||||
|
|
||||||
async fn stop_background_and_middle_end(generation: &RuntimeGeneration) -> bool {
|
async fn stop_background_and_middle_end(generation: &RuntimeGeneration) -> bool {
|
||||||
generation.stop_background_tasks().await;
|
generation.stop_background_tasks().await;
|
||||||
let Some(pool) = generation.current_me_pool().await else {
|
!generation.stop_middle_end(Duration::from_secs(5)).await
|
||||||
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 {
|
async fn cleanup_candidate(generation: &RuntimeGeneration) -> bool {
|
||||||
@@ -103,6 +100,7 @@ impl ReloadSupervisor {
|
|||||||
commands: ReloadCommandReceiver,
|
commands: ReloadCommandReceiver,
|
||||||
config_path: PathBuf,
|
config_path: PathBuf,
|
||||||
quota_store: Arc<QuotaStore>,
|
quota_store: Arc<QuotaStore>,
|
||||||
|
tls_full_cert_budget: Arc<TlsFullCertBudget>,
|
||||||
detected_ips_tx: watch::Sender<(Option<std::net::IpAddr>, Option<std::net::IpAddr>)>,
|
detected_ips_tx: watch::Sender<(Option<std::net::IpAddr>, Option<std::net::IpAddr>)>,
|
||||||
runtime_log_filter: RuntimeLogFilter,
|
runtime_log_filter: RuntimeLogFilter,
|
||||||
runtime_watch_tx: watch::Sender<Option<RuntimeWatchState>>,
|
runtime_watch_tx: watch::Sender<Option<RuntimeWatchState>>,
|
||||||
@@ -116,6 +114,7 @@ impl ReloadSupervisor {
|
|||||||
commands,
|
commands,
|
||||||
config_path,
|
config_path,
|
||||||
quota_store,
|
quota_store,
|
||||||
|
tls_full_cert_budget,
|
||||||
detected_ips_tx,
|
detected_ips_tx,
|
||||||
runtime_log_filter,
|
runtime_log_filter,
|
||||||
runtime_watch_tx,
|
runtime_watch_tx,
|
||||||
@@ -171,6 +170,7 @@ impl ReloadSupervisor {
|
|||||||
&self.config_path,
|
&self.config_path,
|
||||||
self.quota_store.clone(),
|
self.quota_store.clone(),
|
||||||
self.runtime_log_filter.clone(),
|
self.runtime_log_filter.clone(),
|
||||||
|
self.tls_full_cert_budget.clone(),
|
||||||
)
|
)
|
||||||
.await
|
.await
|
||||||
{
|
{
|
||||||
@@ -207,25 +207,18 @@ impl ReloadSupervisor {
|
|||||||
prepared,
|
prepared,
|
||||||
listener_transition,
|
listener_transition,
|
||||||
revision_action,
|
revision_action,
|
||||||
|entries| {
|
|
||||||
crate::network::dns_overrides::install_entries(entries)
|
|
||||||
.map_err(|error| error.to_string())
|
|
||||||
},
|
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
async fn activate_prepared<InstallDns>(
|
async fn activate_prepared(
|
||||||
&self,
|
&self,
|
||||||
command: ReloadCommand,
|
command: ReloadCommand,
|
||||||
old_runtime: Arc<RuntimeGeneration>,
|
old_runtime: Arc<RuntimeGeneration>,
|
||||||
prepared: PreparedRuntime,
|
prepared: PreparedRuntime,
|
||||||
revision_action: RevisionGateAction,
|
revision_action: RevisionGateAction,
|
||||||
install_dns: InstallDns,
|
) {
|
||||||
) where
|
|
||||||
InstallDns: FnOnce(&[String]) -> Result<(), String>,
|
|
||||||
{
|
|
||||||
let listener_transition = match self
|
let listener_transition = match self
|
||||||
.listener_manager
|
.listener_manager
|
||||||
.lock()
|
.lock()
|
||||||
@@ -245,22 +238,18 @@ impl ReloadSupervisor {
|
|||||||
prepared,
|
prepared,
|
||||||
listener_transition,
|
listener_transition,
|
||||||
revision_action,
|
revision_action,
|
||||||
install_dns,
|
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn activate_prepared_with_transition<InstallDns>(
|
async fn activate_prepared_with_transition(
|
||||||
&self,
|
&self,
|
||||||
command: ReloadCommand,
|
command: ReloadCommand,
|
||||||
old_runtime: Arc<RuntimeGeneration>,
|
old_runtime: Arc<RuntimeGeneration>,
|
||||||
prepared: PreparedRuntime,
|
prepared: PreparedRuntime,
|
||||||
listener_transition: Option<PreparedListenerTransition>,
|
listener_transition: Option<PreparedListenerTransition>,
|
||||||
revision_action: RevisionGateAction,
|
revision_action: RevisionGateAction,
|
||||||
install_dns: InstallDns,
|
) {
|
||||||
) where
|
|
||||||
InstallDns: FnOnce(&[String]) -> Result<(), String>,
|
|
||||||
{
|
|
||||||
match revision_action {
|
match revision_action {
|
||||||
RevisionGateAction::Proceed => {}
|
RevisionGateAction::Proceed => {}
|
||||||
RevisionGateAction::Warn(warning) => {
|
RevisionGateAction::Warn(warning) => {
|
||||||
@@ -283,18 +272,6 @@ impl ReloadSupervisor {
|
|||||||
detected_ips,
|
detected_ips,
|
||||||
config_watcher_activation,
|
config_watcher_activation,
|
||||||
} = prepared;
|
} = prepared;
|
||||||
if let Err(error) = install_dns(&new_runtime.config().network.dns_overrides) {
|
|
||||||
let message = format!("runtime DNS activation failed: {}", error);
|
|
||||||
if command.request.failure_policy == ReloadFailurePolicy::Rollback {
|
|
||||||
old_runtime.resume_accepting_sessions();
|
|
||||||
let _ = cleanup_candidate(&new_runtime).await;
|
|
||||||
self.runtime_log_filter
|
|
||||||
.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 pending_listener_transition = if let Some(listener_transition) = listener_transition {
|
let pending_listener_transition = if let Some(listener_transition) = listener_transition {
|
||||||
match self
|
match self
|
||||||
.listener_manager
|
.listener_manager
|
||||||
@@ -367,7 +344,7 @@ impl ReloadSupervisor {
|
|||||||
|
|
||||||
if stop_background_and_middle_end(&replaced).await {
|
if stop_background_and_middle_end(&replaced).await {
|
||||||
let warning = format!(
|
let warning = format!(
|
||||||
"generation {} Middle-End close broadcast timed out",
|
"generation {} Middle-End lifecycle shutdown timed out",
|
||||||
replaced.id
|
replaced.id
|
||||||
);
|
);
|
||||||
warn!(reload_id = command.reload_id, warning = %warning);
|
warn!(reload_id = command.reload_id, warning = %warning);
|
||||||
|
|||||||
@@ -53,6 +53,7 @@ async fn fixture(request: ReloadRequest) -> ReloadFixture {
|
|||||||
commands,
|
commands,
|
||||||
config_path: PathBuf::new(),
|
config_path: PathBuf::new(),
|
||||||
quota_store: Arc::new(QuotaStore::default()),
|
quota_store: Arc::new(QuotaStore::default()),
|
||||||
|
tls_full_cert_budget: Arc::new(crate::tls_front::cache::TlsFullCertBudget::new()),
|
||||||
detected_ips_tx,
|
detected_ips_tx,
|
||||||
runtime_log_filter: runtime_log_filter(),
|
runtime_log_filter: runtime_log_filter(),
|
||||||
runtime_watch_tx,
|
runtime_watch_tx,
|
||||||
@@ -131,7 +132,6 @@ async fn revision_rollback_keeps_old_generation_and_cleans_candidate() {
|
|||||||
fixture.old_runtime.clone(),
|
fixture.old_runtime.clone(),
|
||||||
prepared_runtime(fixture.new_runtime),
|
prepared_runtime(fixture.new_runtime),
|
||||||
RevisionGateAction::Rollback("revision changed".to_string()),
|
RevisionGateAction::Rollback("revision changed".to_string()),
|
||||||
|_| -> Result<(), String> { panic!("DNS activation must not run on rollback") },
|
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
@@ -154,44 +154,6 @@ async fn revision_rollback_keeps_old_generation_and_cleans_candidate() {
|
|||||||
fixture.old_runtime.stop_sessions().await;
|
fixture.old_runtime.stop_sessions().await;
|
||||||
}
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn dns_failure_policy_controls_rollback_or_keep_new() {
|
|
||||||
for policy in [ReloadFailurePolicy::Rollback, ReloadFailurePolicy::KeepNew] {
|
|
||||||
let fixture = fixture(ReloadRequest {
|
|
||||||
failure_policy: policy,
|
|
||||||
..ReloadRequest::default()
|
|
||||||
})
|
|
||||||
.await;
|
|
||||||
fixture
|
|
||||||
.supervisor
|
|
||||||
.activate_prepared(
|
|
||||||
fixture.command,
|
|
||||||
fixture.old_runtime.clone(),
|
|
||||||
prepared_runtime(fixture.new_runtime.clone()),
|
|
||||||
RevisionGateAction::Proceed,
|
|
||||||
|_| Err("invalid DNS entry".to_string()),
|
|
||||||
)
|
|
||||||
.await;
|
|
||||||
|
|
||||||
let status = fixture.control.status(1).await.unwrap();
|
|
||||||
match policy {
|
|
||||||
ReloadFailurePolicy::Rollback => {
|
|
||||||
assert_eq!(fixture.supervisor.active_runtime.load().id, 1);
|
|
||||||
assert_eq!(status.state, ReloadPhase::RolledBack);
|
|
||||||
assert!(fixture.old_runtime.spawn_session(async {}));
|
|
||||||
fixture.old_runtime.stop_sessions().await;
|
|
||||||
}
|
|
||||||
ReloadFailurePolicy::KeepNew => {
|
|
||||||
assert_eq!(fixture.supervisor.active_runtime.load().id, 2);
|
|
||||||
assert_eq!(status.state, ReloadPhase::Succeeded);
|
|
||||||
assert_eq!(status.warnings.len(), 1);
|
|
||||||
assert!(!fixture.old_runtime.spawn_session(async {}));
|
|
||||||
fixture.new_runtime.stop_sessions().await;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn drain_publishes_new_generation_before_old_sessions_finish() {
|
async fn drain_publishes_new_generation_before_old_sessions_finish() {
|
||||||
let mut fixture = fixture(ReloadRequest {
|
let mut fixture = fixture(ReloadRequest {
|
||||||
@@ -220,7 +182,6 @@ async fn drain_publishes_new_generation_before_old_sessions_finish() {
|
|||||||
old_runtime,
|
old_runtime,
|
||||||
prepared_runtime(new_runtime),
|
prepared_runtime(new_runtime),
|
||||||
RevisionGateAction::Proceed,
|
RevisionGateAction::Proceed,
|
||||||
|_| Ok(()),
|
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
});
|
});
|
||||||
@@ -273,7 +234,6 @@ async fn drain_timeout_cancels_old_sessions_and_records_one_warning() {
|
|||||||
old_runtime,
|
old_runtime,
|
||||||
prepared_runtime(new_runtime),
|
prepared_runtime(new_runtime),
|
||||||
RevisionGateAction::Proceed,
|
RevisionGateAction::Proceed,
|
||||||
|_| Ok(()),
|
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
});
|
});
|
||||||
@@ -304,6 +264,7 @@ async fn quiesce_joins_idle_supervisor_and_rejects_later_submissions() {
|
|||||||
commands,
|
commands,
|
||||||
PathBuf::new(),
|
PathBuf::new(),
|
||||||
Arc::new(QuotaStore::default()),
|
Arc::new(QuotaStore::default()),
|
||||||
|
Arc::new(crate::tls_front::cache::TlsFullCertBudget::new()),
|
||||||
detected_ips_tx,
|
detected_ips_tx,
|
||||||
runtime_log_filter(),
|
runtime_log_filter(),
|
||||||
runtime_watch_tx,
|
runtime_watch_tx,
|
||||||
|
|||||||
@@ -19,6 +19,7 @@ use crate::stats::beobachten::BeobachtenStore;
|
|||||||
use crate::stats::telemetry::TelemetryPolicy;
|
use crate::stats::telemetry::TelemetryPolicy;
|
||||||
use crate::stats::{QuotaStore, ReplayChecker, Stats};
|
use crate::stats::{QuotaStore, ReplayChecker, Stats};
|
||||||
use crate::stream::BufferPool;
|
use crate::stream::BufferPool;
|
||||||
|
use crate::tls_front::cache::TlsFullCertBudget;
|
||||||
use crate::transport::UpstreamManager;
|
use crate::transport::UpstreamManager;
|
||||||
use crate::transport::middle_proxy::MePool;
|
use crate::transport::middle_proxy::MePool;
|
||||||
|
|
||||||
@@ -44,6 +45,7 @@ pub(crate) async fn prepare_runtime(
|
|||||||
config_path: &Path,
|
config_path: &Path,
|
||||||
quota_store: Arc<QuotaStore>,
|
quota_store: Arc<QuotaStore>,
|
||||||
runtime_log_filter: RuntimeLogFilter,
|
runtime_log_filter: RuntimeLogFilter,
|
||||||
|
tls_full_cert_budget: Arc<TlsFullCertBudget>,
|
||||||
) -> Result<PreparedRuntime, String> {
|
) -> Result<PreparedRuntime, String> {
|
||||||
config
|
config
|
||||||
.validate_web_decoy_listener_separation()
|
.validate_web_decoy_listener_separation()
|
||||||
@@ -121,6 +123,7 @@ pub(crate) async fn prepare_runtime(
|
|||||||
upstream_manager.clone(),
|
upstream_manager.clone(),
|
||||||
&startup_tracker,
|
&startup_tracker,
|
||||||
task_scope.clone(),
|
task_scope.clone(),
|
||||||
|
tls_full_cert_budget,
|
||||||
tls_bootstrap::TlsBootstrapPolicy::RequireReady,
|
tls_bootstrap::TlsBootstrapPolicy::RequireReady,
|
||||||
)
|
)
|
||||||
.await
|
.await
|
||||||
|
|||||||
@@ -26,6 +26,7 @@ use crate::stats::{ReplayChecker, Stats};
|
|||||||
use crate::transport::UpstreamManager;
|
use crate::transport::UpstreamManager;
|
||||||
use crate::transport::middle_proxy::{MePool, MeReinitTrigger};
|
use crate::transport::middle_proxy::{MePool, MeReinitTrigger};
|
||||||
|
|
||||||
|
use super::control_plane::ProcessControlPlane;
|
||||||
use super::generation::RuntimeGeneration;
|
use super::generation::RuntimeGeneration;
|
||||||
use super::generation::RuntimeTaskScope;
|
use super::generation::RuntimeTaskScope;
|
||||||
use super::helpers::write_beobachten_snapshot;
|
use super::helpers::write_beobachten_snapshot;
|
||||||
@@ -158,6 +159,7 @@ pub(crate) async fn spawn_runtime_tasks(
|
|||||||
detected_ip_v4,
|
detected_ip_v4,
|
||||||
detected_ip_v6,
|
detected_ip_v6,
|
||||||
task_scope.cancellation_token(),
|
task_scope.cancellation_token(),
|
||||||
|
Some(upstream_manager.dns_resolver()),
|
||||||
config_watcher_activation,
|
config_watcher_activation,
|
||||||
);
|
);
|
||||||
task_scope.spawn(config_watcher_task);
|
task_scope.spawn(config_watcher_task);
|
||||||
@@ -168,7 +170,6 @@ pub(crate) async fn spawn_runtime_tasks(
|
|||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
let stats_policy = stats.clone();
|
let stats_policy = stats.clone();
|
||||||
let upstream_policy = upstream_manager.clone();
|
|
||||||
let mut config_rx_policy = config_rx.clone();
|
let mut config_rx_policy = config_rx.clone();
|
||||||
task_scope.spawn(async move {
|
task_scope.spawn(async move {
|
||||||
loop {
|
loop {
|
||||||
@@ -178,9 +179,6 @@ pub(crate) async fn spawn_runtime_tasks(
|
|||||||
let cfg = config_rx_policy.borrow_and_update().clone();
|
let cfg = config_rx_policy.borrow_and_update().clone();
|
||||||
stats_policy
|
stats_policy
|
||||||
.apply_telemetry_policy(TelemetryPolicy::from_config(&cfg.general.telemetry));
|
.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 {
|
if let Some(pool) = &me_pool_for_policy {
|
||||||
pool.update_runtime_transport_policy(
|
pool.update_runtime_transport_policy(
|
||||||
cfg.general.me_socks_kdf_policy,
|
cfg.general.me_socks_kdf_policy,
|
||||||
@@ -406,7 +404,9 @@ pub(crate) async fn spawn_metrics_if_configured(
|
|||||||
startup_tracker: &Arc<StartupTracker>,
|
startup_tracker: &Arc<StartupTracker>,
|
||||||
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
|
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
|
||||||
web_runtime_rx: tokio::sync::watch::Receiver<crate::web::control::WebRuntimePublication>,
|
web_runtime_rx: tokio::sync::watch::Receiver<crate::web::control::WebRuntimePublication>,
|
||||||
) {
|
tls_full_cert_budget: Arc<crate::tls_front::cache::TlsFullCertBudget>,
|
||||||
|
control_plane: ProcessControlPlane,
|
||||||
|
) -> std::io::Result<()> {
|
||||||
// metrics_listen takes precedence; fall back to metrics_port for backward compat.
|
// metrics_listen takes precedence; fall back to metrics_port for backward compat.
|
||||||
let metrics_target: Option<(u16, Option<String>)> =
|
let metrics_target: Option<(u16, Option<String>)> =
|
||||||
if let Some(ref listen) = config.server.metrics_listen {
|
if let Some(ref listen) = config.server.metrics_listen {
|
||||||
@@ -414,12 +414,15 @@ pub(crate) async fn spawn_metrics_if_configured(
|
|||||||
Ok(addr) => Some((addr.port(), Some(listen.clone()))),
|
Ok(addr) => Some((addr.port(), Some(listen.clone()))),
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
startup_tracker
|
startup_tracker
|
||||||
.skip_component(
|
.fail_component(
|
||||||
COMPONENT_METRICS_START,
|
COMPONENT_METRICS_START,
|
||||||
Some(format!("invalid metrics_listen \"{}\": {}", listen, e)),
|
Some(format!("invalid metrics_listen \"{}\": {}", listen, e)),
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
None
|
return Err(std::io::Error::new(
|
||||||
|
std::io::ErrorKind::InvalidInput,
|
||||||
|
format!("invalid metrics_listen \"{}\": {}", listen, e),
|
||||||
|
));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
} else {
|
} else {
|
||||||
@@ -435,15 +438,30 @@ pub(crate) async fn spawn_metrics_if_configured(
|
|||||||
Some(format!("spawn metrics endpoint on {}", label)),
|
Some(format!("spawn metrics endpoint on {}", label)),
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
let active_runtime = active_runtime.clone();
|
|
||||||
let listen_backlog = config.server.listen_backlog;
|
let listen_backlog = config.server.listen_backlog;
|
||||||
tokio::spawn(async move {
|
let bound = match metrics::bind(port, listen, listen_backlog) {
|
||||||
metrics::serve(port, listen, listen_backlog, active_runtime, web_runtime_rx).await;
|
Ok(bound) => bound,
|
||||||
});
|
Err(error) => {
|
||||||
|
startup_tracker
|
||||||
|
.fail_component(
|
||||||
|
COMPONENT_METRICS_START,
|
||||||
|
Some(format!("metrics listener bind failed: {error}")),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
return Err(error);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
metrics::serve(
|
||||||
|
bound,
|
||||||
|
active_runtime,
|
||||||
|
web_runtime_rx,
|
||||||
|
tls_full_cert_budget,
|
||||||
|
control_plane,
|
||||||
|
);
|
||||||
startup_tracker
|
startup_tracker
|
||||||
.complete_component(
|
.complete_component(
|
||||||
COMPONENT_METRICS_START,
|
COMPONENT_METRICS_START,
|
||||||
Some("metrics task spawned".to_string()),
|
Some("metrics listeners bound and supervised".to_string()),
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
} else if config.server.metrics_listen.is_none() {
|
} else if config.server.metrics_listen.is_none() {
|
||||||
@@ -454,6 +472,7 @@ pub(crate) async fn spawn_metrics_if_configured(
|
|||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
}
|
}
|
||||||
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) async fn mark_runtime_ready(startup_tracker: &Arc<StartupTracker>) {
|
pub(crate) async fn mark_runtime_ready(startup_tracker: &Arc<StartupTracker>) {
|
||||||
|
|||||||
+33
-22
@@ -8,7 +8,7 @@
|
|||||||
//!
|
//!
|
||||||
//! SIGHUP is handled separately in config/hot_reload.rs for config reload.
|
//! SIGHUP is handled separately in config/hot_reload.rs for config reload.
|
||||||
|
|
||||||
use std::path::PathBuf;
|
use std::collections::BTreeSet;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use std::time::{Duration, Instant};
|
use std::time::{Duration, Instant};
|
||||||
|
|
||||||
@@ -19,11 +19,13 @@ use tokio::signal;
|
|||||||
use tokio::signal::unix::{SignalKind, signal};
|
use tokio::signal::unix::{SignalKind, signal};
|
||||||
use tracing::{info, warn};
|
use tracing::{info, warn};
|
||||||
|
|
||||||
|
use super::control_plane::ProcessControlPlane;
|
||||||
use super::generation::RuntimeGeneration;
|
use super::generation::RuntimeGeneration;
|
||||||
use super::helpers::{format_uptime, unit_label};
|
use super::helpers::{format_uptime, unit_label};
|
||||||
use super::reload_supervisor::ReloadSupervisorHandle;
|
use super::reload_supervisor::ReloadSupervisorHandle;
|
||||||
use crate::stats::Stats;
|
use crate::stats::Stats;
|
||||||
use crate::synlimit_control;
|
use crate::synlimit_control;
|
||||||
|
use crate::quota_state::QuotaStateOwner;
|
||||||
|
|
||||||
/// Signal that triggered shutdown.
|
/// Signal that triggered shutdown.
|
||||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||||
@@ -50,16 +52,18 @@ impl std::fmt::Display for ShutdownSignal {
|
|||||||
pub(crate) async fn wait_for_shutdown(
|
pub(crate) async fn wait_for_shutdown(
|
||||||
process_started_at: Instant,
|
process_started_at: Instant,
|
||||||
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
|
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
|
||||||
quota_state_path: PathBuf,
|
quota_state: Arc<QuotaStateOwner>,
|
||||||
reload_supervisor: ReloadSupervisorHandle,
|
reload_supervisor: ReloadSupervisorHandle,
|
||||||
|
process_control_plane: ProcessControlPlane,
|
||||||
) {
|
) {
|
||||||
let signal = wait_for_shutdown_signal().await;
|
let signal = wait_for_shutdown_signal().await;
|
||||||
perform_shutdown(
|
perform_shutdown(
|
||||||
signal,
|
signal,
|
||||||
process_started_at,
|
process_started_at,
|
||||||
active_runtime,
|
active_runtime,
|
||||||
quota_state_path,
|
quota_state,
|
||||||
reload_supervisor,
|
reload_supervisor,
|
||||||
|
process_control_plane,
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
}
|
}
|
||||||
@@ -89,8 +93,9 @@ async fn perform_shutdown(
|
|||||||
signal: ShutdownSignal,
|
signal: ShutdownSignal,
|
||||||
process_started_at: Instant,
|
process_started_at: Instant,
|
||||||
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
|
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
|
||||||
quota_state_path: PathBuf,
|
quota_state: Arc<QuotaStateOwner>,
|
||||||
reload_supervisor: ReloadSupervisorHandle,
|
reload_supervisor: ReloadSupervisorHandle,
|
||||||
|
process_control_plane: ProcessControlPlane,
|
||||||
) {
|
) {
|
||||||
let shutdown_started_at = Instant::now();
|
let shutdown_started_at = Instant::now();
|
||||||
info!(signal = %signal, "Received shutdown signal");
|
info!(signal = %signal, "Received shutdown signal");
|
||||||
@@ -115,37 +120,41 @@ async fn perform_shutdown(
|
|||||||
// Graceful ME pool shutdown
|
// Graceful ME pool shutdown
|
||||||
runtime.stop_sessions().await;
|
runtime.stop_sessions().await;
|
||||||
runtime.stop_background_tasks().await;
|
runtime.stop_background_tasks().await;
|
||||||
if let Some(pool) = runtime.current_me_pool().await {
|
if runtime.stop_middle_end(Duration::from_secs(5)).await {
|
||||||
match tokio::time::timeout(Duration::from_secs(2), pool.shutdown_send_close_conn_all())
|
info!("ME shutdown: pool lifecycle completed");
|
||||||
.await
|
} else {
|
||||||
{
|
warn!("ME shutdown: pool lifecycle deadline expired");
|
||||||
Ok(total) => {
|
|
||||||
info!(
|
|
||||||
close_conn_sent = total,
|
|
||||||
"ME shutdown: RPC_CLOSE_CONN broadcast completed"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
Err(_) => {
|
|
||||||
warn!("ME shutdown: RPC_CLOSE_CONN broadcast timed out");
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
if let Err(error) = synlimit_control::clear_synlimit_rules_all_backends().await {
|
if let Err(error) = synlimit_control::clear_synlimit_rules_all_backends().await {
|
||||||
warn!(error = %error, "Failed to clear SYN limiter rules during shutdown");
|
warn!(error = %error, "Failed to clear SYN limiter rules during shutdown");
|
||||||
}
|
}
|
||||||
|
|
||||||
match crate::quota_state::save_quota_state("a_state_path, stats).await {
|
if !process_control_plane
|
||||||
|
.shutdown(Duration::from_secs(5))
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
warn!("Process control-plane task shutdown deadline expired");
|
||||||
|
}
|
||||||
|
|
||||||
|
let configured_quota_users = runtime
|
||||||
|
.config()
|
||||||
|
.access
|
||||||
|
.users
|
||||||
|
.keys()
|
||||||
|
.cloned()
|
||||||
|
.collect::<BTreeSet<_>>();
|
||||||
|
match quota_state.save(&configured_quota_users).await {
|
||||||
Ok(()) => {
|
Ok(()) => {
|
||||||
info!(
|
info!(
|
||||||
path = %quota_state_path.display(),
|
path = %quota_state.path().display(),
|
||||||
"Persisted per-user quota state"
|
"Persisted per-user quota state"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
Err(error) => {
|
Err(error) => {
|
||||||
warn!(
|
warn!(
|
||||||
error = %error,
|
error = %error,
|
||||||
path = %quota_state_path.display(),
|
path = %quota_state.path().display(),
|
||||||
"Failed to persist per-user quota state"
|
"Failed to persist per-user quota state"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
@@ -205,8 +214,9 @@ fn dump_stats(stats: &Stats, process_started_at: Instant) {
|
|||||||
pub(crate) fn spawn_signal_handlers(
|
pub(crate) fn spawn_signal_handlers(
|
||||||
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
|
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
|
||||||
process_started_at: Instant,
|
process_started_at: Instant,
|
||||||
|
process_control_plane: ProcessControlPlane,
|
||||||
) {
|
) {
|
||||||
tokio::spawn(async move {
|
let _ = process_control_plane.spawn(async move {
|
||||||
let mut sigusr1 =
|
let mut sigusr1 =
|
||||||
signal(SignalKind::user_defined1()).expect("Failed to register SIGUSR1 handler");
|
signal(SignalKind::user_defined1()).expect("Failed to register SIGUSR1 handler");
|
||||||
let mut sigusr2 =
|
let mut sigusr2 =
|
||||||
@@ -231,6 +241,7 @@ pub(crate) fn spawn_signal_handlers(
|
|||||||
pub(crate) fn spawn_signal_handlers(
|
pub(crate) fn spawn_signal_handlers(
|
||||||
_active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
|
_active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
|
||||||
_process_started_at: Instant,
|
_process_started_at: Instant,
|
||||||
|
_process_control_plane: ProcessControlPlane,
|
||||||
) {
|
) {
|
||||||
// No SIGUSR1/SIGUSR2 on non-Unix
|
// No SIGUSR1/SIGUSR2 on non-Unix
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -8,6 +8,7 @@ use crate::config::ProxyConfig;
|
|||||||
use crate::error::{ProxyError, Result};
|
use crate::error::{ProxyError, Result};
|
||||||
use crate::startup::{COMPONENT_TLS_FRONT_BOOTSTRAP, StartupTracker};
|
use crate::startup::{COMPONENT_TLS_FRONT_BOOTSTRAP, StartupTracker};
|
||||||
use crate::tls_front::TlsFrontCache;
|
use crate::tls_front::TlsFrontCache;
|
||||||
|
use crate::tls_front::cache::TlsFullCertBudget;
|
||||||
use crate::tls_front::fetcher::TlsFetchStrategy;
|
use crate::tls_front::fetcher::TlsFetchStrategy;
|
||||||
use crate::transport::UpstreamManager;
|
use crate::transport::UpstreamManager;
|
||||||
|
|
||||||
@@ -109,6 +110,7 @@ pub(crate) async fn bootstrap_tls_front(
|
|||||||
upstream_manager: Arc<UpstreamManager>,
|
upstream_manager: Arc<UpstreamManager>,
|
||||||
startup_tracker: &Arc<StartupTracker>,
|
startup_tracker: &Arc<StartupTracker>,
|
||||||
task_scope: RuntimeTaskScope,
|
task_scope: RuntimeTaskScope,
|
||||||
|
full_cert_budget: Arc<TlsFullCertBudget>,
|
||||||
policy: TlsBootstrapPolicy,
|
policy: TlsBootstrapPolicy,
|
||||||
) -> Result<Option<Arc<TlsFrontCache>>> {
|
) -> Result<Option<Arc<TlsFrontCache>>> {
|
||||||
startup_tracker
|
startup_tracker
|
||||||
@@ -128,10 +130,11 @@ pub(crate) async fn bootstrap_tls_front(
|
|||||||
return Ok(None);
|
return Ok(None);
|
||||||
}
|
}
|
||||||
|
|
||||||
let cache = Arc::new(TlsFrontCache::new(
|
let cache = Arc::new(TlsFrontCache::new_with_full_cert_budget(
|
||||||
tls_domains,
|
tls_domains,
|
||||||
config.censorship.fake_cert_len,
|
config.censorship.fake_cert_len,
|
||||||
&config.censorship.tls_front_dir,
|
&config.censorship.tls_front_dir,
|
||||||
|
full_cert_budget,
|
||||||
));
|
));
|
||||||
cache.load_from_disk().await;
|
cache.load_from_disk().await;
|
||||||
|
|
||||||
@@ -301,6 +304,7 @@ mod tests {
|
|||||||
upstream_manager(&config),
|
upstream_manager(&config),
|
||||||
&tracker,
|
&tracker,
|
||||||
scope.clone(),
|
scope.clone(),
|
||||||
|
Arc::new(TlsFullCertBudget::new()),
|
||||||
TlsBootstrapPolicy::RequireReady,
|
TlsBootstrapPolicy::RequireReady,
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
@@ -336,6 +340,7 @@ mod tests {
|
|||||||
upstream_manager(&config),
|
upstream_manager(&config),
|
||||||
&tracker,
|
&tracker,
|
||||||
scope.clone(),
|
scope.clone(),
|
||||||
|
Arc::new(TlsFullCertBudget::new()),
|
||||||
TlsBootstrapPolicy::RequireReady,
|
TlsBootstrapPolicy::RequireReady,
|
||||||
)
|
)
|
||||||
.await
|
.await
|
||||||
@@ -364,6 +369,7 @@ mod tests {
|
|||||||
upstream_manager(&config),
|
upstream_manager(&config),
|
||||||
&tracker,
|
&tracker,
|
||||||
scope.clone(),
|
scope.clone(),
|
||||||
|
Arc::new(TlsFullCertBudget::new()),
|
||||||
TlsBootstrapPolicy::BestEffort,
|
TlsBootstrapPolicy::BestEffort,
|
||||||
)
|
)
|
||||||
.await
|
.await
|
||||||
|
|||||||
+147
-67
@@ -17,12 +17,13 @@ use tracing::{debug, info, warn};
|
|||||||
|
|
||||||
use crate::config::ProxyConfig;
|
use crate::config::ProxyConfig;
|
||||||
use crate::ip_tracker::UserIpTracker;
|
use crate::ip_tracker::UserIpTracker;
|
||||||
|
use crate::maestro::control_plane::ProcessControlPlane;
|
||||||
use crate::maestro::generation::RuntimeGeneration;
|
use crate::maestro::generation::RuntimeGeneration;
|
||||||
use crate::proxy::shared_state::ProxySharedState;
|
use crate::proxy::shared_state::ProxySharedState;
|
||||||
use crate::stats::Stats;
|
use crate::stats::Stats;
|
||||||
use crate::stats::beobachten::BeobachtenStore;
|
use crate::stats::beobachten::BeobachtenStore;
|
||||||
use crate::tls_front::TlsFrontCache;
|
use crate::tls_front::TlsFrontCache;
|
||||||
use crate::tls_front::cache;
|
use crate::tls_front::cache::TlsFullCertBudget;
|
||||||
use crate::tls_front::fetcher;
|
use crate::tls_front::fetcher;
|
||||||
use crate::transport::{ListenOptions, create_listener};
|
use crate::transport::{ListenOptions, create_listener};
|
||||||
|
|
||||||
@@ -36,84 +37,89 @@ const TLS_FRONT_PROFILE_HEALTH_MAX_DOMAINS: usize = 256;
|
|||||||
const METRICS_MAX_CONTROL_CONNECTIONS: usize = 512;
|
const METRICS_MAX_CONTROL_CONNECTIONS: usize = 512;
|
||||||
const METRICS_HTTP_CONNECTION_TIMEOUT: Duration = Duration::from_secs(15);
|
const METRICS_HTTP_CONNECTION_TIMEOUT: Duration = Duration::from_secs(15);
|
||||||
|
|
||||||
pub async fn serve(
|
/// Bound process-owned metrics listeners ready for supervised serving.
|
||||||
|
pub(crate) struct BoundMetricsListeners {
|
||||||
|
listeners: Vec<(TcpListener, SocketAddr)>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Binds every configured metrics socket before process readiness is published.
|
||||||
|
pub(crate) fn bind(
|
||||||
port: u16,
|
port: u16,
|
||||||
listen: Option<String>,
|
listen: Option<String>,
|
||||||
listen_backlog: u32,
|
listen_backlog: u32,
|
||||||
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
|
) -> std::io::Result<BoundMetricsListeners> {
|
||||||
web_runtime_rx: tokio::sync::watch::Receiver<crate::web::control::WebRuntimePublication>,
|
|
||||||
) {
|
|
||||||
// If `metrics_listen` is set, bind on that single address only.
|
|
||||||
if let Some(ref listen_addr) = listen {
|
if let Some(ref listen_addr) = listen {
|
||||||
let addr: SocketAddr = match listen_addr.parse() {
|
let addr: SocketAddr = listen_addr.parse().map_err(|error| {
|
||||||
Ok(a) => a,
|
std::io::Error::new(
|
||||||
Err(e) => {
|
std::io::ErrorKind::InvalidInput,
|
||||||
warn!(error = %e, "Invalid metrics_listen address: {}", listen_addr);
|
format!("invalid metrics_listen address {listen_addr}: {error}"),
|
||||||
return;
|
)
|
||||||
}
|
})?;
|
||||||
};
|
|
||||||
// Match `server.api.listen`: `[::]:port` is a dual-stack wildcard
|
// Match `server.api.listen`: `[::]:port` is a dual-stack wildcard
|
||||||
// on Linux when `net.ipv6.bindv6only=0`.
|
// on Linux when `net.ipv6.bindv6only=0`.
|
||||||
let ipv6_only = addr.is_ipv6() && !addr.ip().is_unspecified();
|
let ipv6_only = addr.is_ipv6() && !addr.ip().is_unspecified();
|
||||||
match bind_metrics_listener(addr, ipv6_only, listen_backlog) {
|
let listener = bind_metrics_listener(addr, ipv6_only, listen_backlog)?;
|
||||||
Ok(listener) => {
|
return Ok(BoundMetricsListeners {
|
||||||
info!("Metrics endpoint: http://{}/metrics and /beobachten", addr);
|
listeners: vec![(listener, addr)],
|
||||||
serve_listener(listener, active_runtime, web_runtime_rx).await;
|
});
|
||||||
}
|
|
||||||
Err(e) => {
|
|
||||||
warn!(error = %e, "Failed to bind metrics on {}", addr);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
return;
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// Fallback: keep metrics local unless an explicit metrics_listen is configured.
|
let mut listeners = Vec::with_capacity(2);
|
||||||
let mut listener_v4 = None;
|
let mut last_error = None;
|
||||||
let mut listener_v6 = None;
|
|
||||||
|
|
||||||
let addr_v4 = SocketAddr::from(([127, 0, 0, 1], port));
|
let addr_v4 = SocketAddr::from(([127, 0, 0, 1], port));
|
||||||
match bind_metrics_listener(addr_v4, false, listen_backlog) {
|
match bind_metrics_listener(addr_v4, false, listen_backlog) {
|
||||||
Ok(listener) => {
|
Ok(listener) => listeners.push((listener, addr_v4)),
|
||||||
info!(
|
|
||||||
"Metrics endpoint: http://{}/metrics and /beobachten",
|
|
||||||
addr_v4
|
|
||||||
);
|
|
||||||
listener_v4 = Some(listener);
|
|
||||||
}
|
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
warn!(error = %e, "Failed to bind metrics on {}", addr_v4);
|
warn!(error = %e, "Failed to bind metrics on {}", addr_v4);
|
||||||
|
last_error = Some(e);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
let addr_v6 = SocketAddr::from(([0, 0, 0, 0, 0, 0, 0, 1], port));
|
let addr_v6 = SocketAddr::from(([0, 0, 0, 0, 0, 0, 0, 1], port));
|
||||||
match bind_metrics_listener(addr_v6, true, listen_backlog) {
|
match bind_metrics_listener(addr_v6, true, listen_backlog) {
|
||||||
Ok(listener) => {
|
Ok(listener) => listeners.push((listener, addr_v6)),
|
||||||
info!(
|
|
||||||
"Metrics endpoint: http://[::1]:{}/metrics and /beobachten",
|
|
||||||
port
|
|
||||||
);
|
|
||||||
listener_v6 = Some(listener);
|
|
||||||
}
|
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
warn!(error = %e, "Failed to bind metrics on {}", addr_v6);
|
warn!(error = %e, "Failed to bind metrics on {}", addr_v6);
|
||||||
|
last_error = Some(e);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
match (listener_v4, listener_v6) {
|
if listeners.is_empty() {
|
||||||
(None, None) => {
|
return Err(last_error.unwrap_or_else(|| {
|
||||||
warn!("Metrics listener is unavailable on both IPv4 and IPv6");
|
std::io::Error::new(
|
||||||
}
|
std::io::ErrorKind::AddrNotAvailable,
|
||||||
(Some(listener), None) | (None, Some(listener)) => {
|
"metrics listener is unavailable on both IPv4 and IPv6",
|
||||||
serve_listener(listener, active_runtime, web_runtime_rx).await;
|
)
|
||||||
}
|
}));
|
||||||
(Some(listener4), Some(listener6)) => {
|
}
|
||||||
let active_runtime_v6 = active_runtime.clone();
|
Ok(BoundMetricsListeners { listeners })
|
||||||
let web_runtime_rx_v6 = web_runtime_rx.clone();
|
}
|
||||||
tokio::spawn(async move {
|
|
||||||
serve_listener(listener6, active_runtime_v6, web_runtime_rx_v6).await;
|
/// Starts supervised accept loops for previously bound metrics sockets.
|
||||||
});
|
pub(crate) fn serve(
|
||||||
serve_listener(listener4, active_runtime, web_runtime_rx).await;
|
bound: BoundMetricsListeners,
|
||||||
}
|
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
|
||||||
|
web_runtime_rx: tokio::sync::watch::Receiver<crate::web::control::WebRuntimePublication>,
|
||||||
|
tls_full_cert_budget: Arc<TlsFullCertBudget>,
|
||||||
|
control_plane: ProcessControlPlane,
|
||||||
|
) {
|
||||||
|
for (listener, addr) in bound.listeners {
|
||||||
|
info!("Metrics endpoint: http://{}/metrics and /beobachten", addr);
|
||||||
|
let active_runtime = active_runtime.clone();
|
||||||
|
let web_runtime_rx = web_runtime_rx.clone();
|
||||||
|
let tls_full_cert_budget = Arc::clone(&tls_full_cert_budget);
|
||||||
|
let listener_scope = control_plane.clone();
|
||||||
|
let _ = control_plane.spawn(async move {
|
||||||
|
serve_listener(
|
||||||
|
listener,
|
||||||
|
active_runtime,
|
||||||
|
web_runtime_rx,
|
||||||
|
tls_full_cert_budget,
|
||||||
|
listener_scope,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
});
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -136,6 +142,8 @@ async fn serve_listener(
|
|||||||
listener: TcpListener,
|
listener: TcpListener,
|
||||||
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
|
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
|
||||||
web_runtime_rx: tokio::sync::watch::Receiver<crate::web::control::WebRuntimePublication>,
|
web_runtime_rx: tokio::sync::watch::Receiver<crate::web::control::WebRuntimePublication>,
|
||||||
|
tls_full_cert_budget: Arc<TlsFullCertBudget>,
|
||||||
|
control_plane: ProcessControlPlane,
|
||||||
) {
|
) {
|
||||||
let connection_permits = Arc::new(Semaphore::new(METRICS_MAX_CONTROL_CONNECTIONS));
|
let connection_permits = Arc::new(Semaphore::new(METRICS_MAX_CONTROL_CONNECTIONS));
|
||||||
|
|
||||||
@@ -175,12 +183,22 @@ async fn serve_listener(
|
|||||||
|
|
||||||
let active_runtime = active_runtime.clone();
|
let active_runtime = active_runtime.clone();
|
||||||
let web_runtime_rx = web_runtime_rx.clone();
|
let web_runtime_rx = web_runtime_rx.clone();
|
||||||
tokio::spawn(async move {
|
let tls_full_cert_budget = Arc::clone(&tls_full_cert_budget);
|
||||||
|
let _ = control_plane.spawn(async move {
|
||||||
let _connection_permit = connection_permit;
|
let _connection_permit = connection_permit;
|
||||||
let svc = service_fn(move |req| {
|
let svc = service_fn(move |req| {
|
||||||
let runtime = active_runtime.load_full();
|
let runtime = active_runtime.load_full();
|
||||||
let web_publication = web_runtime_rx.borrow().clone();
|
let web_publication = web_runtime_rx.borrow().clone();
|
||||||
async move { handle(req, &runtime, &web_publication).await }
|
let tls_full_cert_budget = Arc::clone(&tls_full_cert_budget);
|
||||||
|
async move {
|
||||||
|
handle(
|
||||||
|
req,
|
||||||
|
&runtime,
|
||||||
|
&web_publication,
|
||||||
|
tls_full_cert_budget.as_ref(),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
}
|
||||||
});
|
});
|
||||||
match timeout(
|
match timeout(
|
||||||
METRICS_HTTP_CONNECTION_TIMEOUT,
|
METRICS_HTTP_CONNECTION_TIMEOUT,
|
||||||
@@ -208,6 +226,7 @@ async fn handle<B>(
|
|||||||
req: Request<B>,
|
req: Request<B>,
|
||||||
runtime: &RuntimeGeneration,
|
runtime: &RuntimeGeneration,
|
||||||
web_publication: &crate::web::control::WebRuntimePublication,
|
web_publication: &crate::web::control::WebRuntimePublication,
|
||||||
|
tls_full_cert_budget: &TlsFullCertBudget,
|
||||||
) -> Result<Response<Full<Bytes>>, Infallible> {
|
) -> Result<Response<Full<Bytes>>, Infallible> {
|
||||||
let stats = &runtime.stats;
|
let stats = &runtime.stats;
|
||||||
let beobachten = &runtime.beobachten;
|
let beobachten = &runtime.beobachten;
|
||||||
@@ -223,6 +242,7 @@ async fn handle<B>(
|
|||||||
&config,
|
&config,
|
||||||
ip_tracker,
|
ip_tracker,
|
||||||
tls_cache,
|
tls_cache,
|
||||||
|
tls_full_cert_budget,
|
||||||
web_publication,
|
web_publication,
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
@@ -437,6 +457,7 @@ async fn render_metrics(
|
|||||||
config: &ProxyConfig,
|
config: &ProxyConfig,
|
||||||
ip_tracker: &UserIpTracker,
|
ip_tracker: &UserIpTracker,
|
||||||
tls_cache: Option<&TlsFrontCache>,
|
tls_cache: Option<&TlsFrontCache>,
|
||||||
|
tls_full_cert_budget: &TlsFullCertBudget,
|
||||||
web_publication: &crate::web::control::WebRuntimePublication,
|
web_publication: &crate::web::control::WebRuntimePublication,
|
||||||
) -> String {
|
) -> String {
|
||||||
use std::fmt::Write;
|
use std::fmt::Write;
|
||||||
@@ -644,17 +665,17 @@ async fn render_metrics(
|
|||||||
);
|
);
|
||||||
let _ = writeln!(
|
let _ = writeln!(
|
||||||
out,
|
out,
|
||||||
"# HELP telemt_tls_front_full_cert_budget_ips Current IP entries tracked by TLS full-cert budget"
|
"# HELP telemt_tls_front_full_cert_budget_entries Current domain and IP entries tracked by the process-owned TLS full-cert budget"
|
||||||
);
|
);
|
||||||
let _ = writeln!(out, "# TYPE telemt_tls_front_full_cert_budget_ips gauge");
|
let _ = writeln!(out, "# TYPE telemt_tls_front_full_cert_budget_entries gauge");
|
||||||
let _ = writeln!(
|
let _ = writeln!(
|
||||||
out,
|
out,
|
||||||
"telemt_tls_front_full_cert_budget_ips {}",
|
"telemt_tls_front_full_cert_budget_entries {}",
|
||||||
cache::full_cert_sent_ips_for_metrics()
|
tls_full_cert_budget.entries_for_metrics()
|
||||||
);
|
);
|
||||||
let _ = writeln!(
|
let _ = writeln!(
|
||||||
out,
|
out,
|
||||||
"# HELP telemt_tls_front_full_cert_budget_cap_drops_total New IPs denied full-cert budget tracking because the cap was reached"
|
"# HELP telemt_tls_front_full_cert_budget_cap_drops_total New domain and IP entries denied full-cert budget tracking because a bound was reached"
|
||||||
);
|
);
|
||||||
let _ = writeln!(
|
let _ = writeln!(
|
||||||
out,
|
out,
|
||||||
@@ -663,7 +684,7 @@ async fn render_metrics(
|
|||||||
let _ = writeln!(
|
let _ = writeln!(
|
||||||
out,
|
out,
|
||||||
"telemt_tls_front_full_cert_budget_cap_drops_total {}",
|
"telemt_tls_front_full_cert_budget_cap_drops_total {}",
|
||||||
cache::full_cert_sent_cap_drops_for_metrics()
|
tls_full_cert_budget.cap_drops_for_metrics()
|
||||||
);
|
);
|
||||||
render_tls_front_profile_health(&mut out, config, tls_cache).await;
|
render_tls_front_profile_health(&mut out, config, tls_cache).await;
|
||||||
|
|
||||||
@@ -1559,6 +1580,11 @@ async fn render_metrics(
|
|||||||
error_code, count
|
error_code, count
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
let _ = writeln!(
|
||||||
|
out,
|
||||||
|
"telemt_me_handshake_error_code_total{{error_code=\"overflow\"}} {}",
|
||||||
|
stats.get_me_handshake_error_code_overflow_total()
|
||||||
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
let _ = writeln!(
|
let _ = writeln!(
|
||||||
@@ -3955,6 +3981,7 @@ mod tests {
|
|||||||
&config,
|
&config,
|
||||||
&tracker,
|
&tracker,
|
||||||
None,
|
None,
|
||||||
|
&TlsFullCertBudget::new(),
|
||||||
&test_web_publication(),
|
&test_web_publication(),
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
@@ -4091,6 +4118,7 @@ mod tests {
|
|||||||
&config,
|
&config,
|
||||||
&tracker,
|
&tracker,
|
||||||
Some(&cache),
|
Some(&cache),
|
||||||
|
&TlsFullCertBudget::new(),
|
||||||
&test_web_publication(),
|
&test_web_publication(),
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
@@ -4135,6 +4163,43 @@ mod tests {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn process_tls_budget_metrics_survive_a_generation_without_tls_cache() {
|
||||||
|
let stats = Stats::new();
|
||||||
|
let shared_state = ProxySharedState::new();
|
||||||
|
let tracker = UserIpTracker::new();
|
||||||
|
let config = ProxyConfig::default();
|
||||||
|
let budget = Arc::new(TlsFullCertBudget::new());
|
||||||
|
let cache = TlsFrontCache::new_with_full_cert_budget(
|
||||||
|
&["example.com".to_string()],
|
||||||
|
1024,
|
||||||
|
"tlsfront-test-cache",
|
||||||
|
Arc::clone(&budget),
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
cache
|
||||||
|
.take_full_cert_budget_for_ip(
|
||||||
|
"example.com",
|
||||||
|
"127.0.0.1".parse().unwrap(),
|
||||||
|
Duration::from_secs(60),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
);
|
||||||
|
|
||||||
|
let output = render_metrics(
|
||||||
|
&stats,
|
||||||
|
&shared_state,
|
||||||
|
&config,
|
||||||
|
&tracker,
|
||||||
|
None,
|
||||||
|
budget.as_ref(),
|
||||||
|
&test_web_publication(),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
assert!(output.contains("telemt_tls_front_full_cert_budget_entries 1"));
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_render_empty_stats() {
|
async fn test_render_empty_stats() {
|
||||||
let stats = Stats::new();
|
let stats = Stats::new();
|
||||||
@@ -4147,6 +4212,7 @@ mod tests {
|
|||||||
&config,
|
&config,
|
||||||
&tracker,
|
&tracker,
|
||||||
None,
|
None,
|
||||||
|
&TlsFullCertBudget::new(),
|
||||||
&test_web_publication(),
|
&test_web_publication(),
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
@@ -4179,6 +4245,7 @@ mod tests {
|
|||||||
&config,
|
&config,
|
||||||
&tracker,
|
&tracker,
|
||||||
None,
|
None,
|
||||||
|
&TlsFullCertBudget::new(),
|
||||||
&test_web_publication(),
|
&test_web_publication(),
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
@@ -4199,6 +4266,7 @@ mod tests {
|
|||||||
&config,
|
&config,
|
||||||
&tracker,
|
&tracker,
|
||||||
None,
|
None,
|
||||||
|
&TlsFullCertBudget::new(),
|
||||||
&test_web_publication(),
|
&test_web_publication(),
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
@@ -4246,7 +4314,7 @@ mod tests {
|
|||||||
assert!(output.contains("# TYPE telemt_ip_tracker_cap_rejects_total counter"));
|
assert!(output.contains("# TYPE telemt_ip_tracker_cap_rejects_total counter"));
|
||||||
assert!(output.contains("# TYPE telemt_tls_fetch_profile_cache_entries gauge"));
|
assert!(output.contains("# TYPE telemt_tls_fetch_profile_cache_entries gauge"));
|
||||||
assert!(output.contains("# TYPE telemt_tls_fetch_profile_cache_cap_drops_total counter"));
|
assert!(output.contains("# TYPE telemt_tls_fetch_profile_cache_cap_drops_total counter"));
|
||||||
assert!(output.contains("# TYPE telemt_tls_front_full_cert_budget_ips gauge"));
|
assert!(output.contains("# TYPE telemt_tls_front_full_cert_budget_entries gauge"));
|
||||||
assert!(
|
assert!(
|
||||||
output.contains("# TYPE telemt_tls_front_full_cert_budget_cap_drops_total counter")
|
output.contains("# TYPE telemt_tls_front_full_cert_budget_cap_drops_total counter")
|
||||||
);
|
);
|
||||||
@@ -4271,12 +4339,15 @@ mod tests {
|
|||||||
config.general.beobachten_minutes = 10;
|
config.general.beobachten_minutes = 10;
|
||||||
let runtime = crate::maestro::generation::test_runtime_generation(1, config);
|
let runtime = crate::maestro::generation::test_runtime_generation(1, config);
|
||||||
let web_publication = test_web_publication();
|
let web_publication = test_web_publication();
|
||||||
|
let tls_full_cert_budget = TlsFullCertBudget::new();
|
||||||
runtime.stats.increment_connects_all();
|
runtime.stats.increment_connects_all();
|
||||||
runtime.stats.increment_connects_all();
|
runtime.stats.increment_connects_all();
|
||||||
runtime.stats.increment_connects_all();
|
runtime.stats.increment_connects_all();
|
||||||
|
|
||||||
let req = Request::builder().uri("/metrics").body(()).unwrap();
|
let req = Request::builder().uri("/metrics").body(()).unwrap();
|
||||||
let resp = handle(req, &runtime, &web_publication).await.unwrap();
|
let resp = handle(req, &runtime, &web_publication, &tls_full_cert_budget)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
assert_eq!(resp.status(), StatusCode::OK);
|
assert_eq!(resp.status(), StatusCode::OK);
|
||||||
let body = resp.into_body().collect().await.unwrap().to_bytes();
|
let body = resp.into_body().collect().await.unwrap().to_bytes();
|
||||||
assert!(
|
assert!(
|
||||||
@@ -4299,7 +4370,14 @@ mod tests {
|
|||||||
Duration::from_secs(600),
|
Duration::from_secs(600),
|
||||||
);
|
);
|
||||||
let req_beob = Request::builder().uri("/beobachten").body(()).unwrap();
|
let req_beob = Request::builder().uri("/beobachten").body(()).unwrap();
|
||||||
let resp_beob = handle(req_beob, &runtime, &web_publication).await.unwrap();
|
let resp_beob = handle(
|
||||||
|
req_beob,
|
||||||
|
&runtime,
|
||||||
|
&web_publication,
|
||||||
|
&tls_full_cert_budget,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
assert_eq!(resp_beob.status(), StatusCode::OK);
|
assert_eq!(resp_beob.status(), StatusCode::OK);
|
||||||
let body_beob = resp_beob.into_body().collect().await.unwrap().to_bytes();
|
let body_beob = resp_beob.into_body().collect().await.unwrap().to_bytes();
|
||||||
let beob_text = std::str::from_utf8(body_beob.as_ref()).unwrap();
|
let beob_text = std::str::from_utf8(body_beob.as_ref()).unwrap();
|
||||||
@@ -4307,7 +4385,9 @@ mod tests {
|
|||||||
assert!(beob_text.contains("203.0.113.10-1"));
|
assert!(beob_text.contains("203.0.113.10-1"));
|
||||||
|
|
||||||
let req404 = Request::builder().uri("/other").body(()).unwrap();
|
let req404 = Request::builder().uri("/other").body(()).unwrap();
|
||||||
let resp404 = handle(req404, &runtime, &web_publication).await.unwrap();
|
let resp404 = handle(req404, &runtime, &web_publication, &tls_full_cert_budget)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
assert_eq!(resp404.status(), StatusCode::NOT_FOUND);
|
assert_eq!(resp404.status(), StatusCode::NOT_FOUND);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -2,23 +2,26 @@
|
|||||||
|
|
||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
use std::net::{IpAddr, Ipv6Addr, SocketAddr};
|
use std::net::{IpAddr, Ipv6Addr, SocketAddr};
|
||||||
use std::sync::{OnceLock, RwLock};
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use arc_swap::ArcSwap;
|
||||||
|
|
||||||
use crate::error::{ProxyError, Result};
|
use crate::error::{ProxyError, Result};
|
||||||
|
|
||||||
type OverrideMap = HashMap<(String, u16), IpAddr>;
|
type OverrideMap = HashMap<(String, u16), IpAddr>;
|
||||||
|
const DNS_OVERRIDE_MAX_ENTRIES: usize = 4096;
|
||||||
|
|
||||||
/// Immutable DNS override snapshot owned by one runtime generation.
|
/// Immutable DNS override snapshot owned by one runtime generation.
|
||||||
#[derive(Debug, Clone, Default)]
|
#[derive(Debug, Clone, Default)]
|
||||||
pub struct DnsOverrides {
|
pub struct DnsOverrides {
|
||||||
entries: std::sync::Arc<OverrideMap>,
|
entries: Arc<OverrideMap>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl DnsOverrides {
|
impl DnsOverrides {
|
||||||
/// Parses a validated generation-local override snapshot.
|
/// Parses a validated generation-local override snapshot.
|
||||||
pub fn from_entries(entries: &[String]) -> Result<Self> {
|
pub fn from_entries(entries: &[String]) -> Result<Self> {
|
||||||
Ok(Self {
|
Ok(Self {
|
||||||
entries: std::sync::Arc::new(parse_entries(entries)?),
|
entries: Arc::new(parse_entries(entries)?),
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -35,10 +38,31 @@ impl DnsOverrides {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
static DNS_OVERRIDES: OnceLock<RwLock<OverrideMap>> = OnceLock::new();
|
/// Atomically published DNS override snapshot owned by one runtime generation.
|
||||||
|
#[derive(Debug, Default)]
|
||||||
|
pub struct GenerationDnsResolver {
|
||||||
|
snapshot: ArcSwap<DnsOverrides>,
|
||||||
|
}
|
||||||
|
|
||||||
fn overrides_store() -> &'static RwLock<OverrideMap> {
|
impl GenerationDnsResolver {
|
||||||
DNS_OVERRIDES.get_or_init(|| RwLock::new(HashMap::new()))
|
/// Creates one resolver from a validated immutable entry set.
|
||||||
|
pub fn from_entries(entries: &[String]) -> Result<Self> {
|
||||||
|
Ok(Self {
|
||||||
|
snapshot: ArcSwap::from_pointee(DnsOverrides::from_entries(entries)?),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Validates and atomically publishes a new generation-local snapshot.
|
||||||
|
pub fn apply_entries(&self, entries: &[String]) -> Result<()> {
|
||||||
|
let snapshot = DnsOverrides::from_entries(entries)?;
|
||||||
|
self.snapshot.store(Arc::new(snapshot));
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Resolves one configured override without consulting system DNS.
|
||||||
|
pub fn resolve_socket_addr(&self, host: &str, port: u16) -> Option<SocketAddr> {
|
||||||
|
self.snapshot.load().resolve_socket_addr(host, port)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn parse_ip_spec(ip_spec: &str) -> Result<IpAddr> {
|
fn parse_ip_spec(ip_spec: &str) -> Result<IpAddr> {
|
||||||
@@ -111,6 +135,11 @@ fn parse_entry(entry: &str) -> Result<((String, u16), IpAddr)> {
|
|||||||
}
|
}
|
||||||
|
|
||||||
fn parse_entries(entries: &[String]) -> Result<OverrideMap> {
|
fn parse_entries(entries: &[String]) -> Result<OverrideMap> {
|
||||||
|
if entries.len() > DNS_OVERRIDE_MAX_ENTRIES {
|
||||||
|
return Err(ProxyError::Config(format!(
|
||||||
|
"network.dns_overrides exceeds maximum entry count {DNS_OVERRIDE_MAX_ENTRIES}"
|
||||||
|
)));
|
||||||
|
}
|
||||||
let mut parsed = HashMap::new();
|
let mut parsed = HashMap::new();
|
||||||
for entry in entries {
|
for entry in entries {
|
||||||
let (key, ip) = parse_entry(entry)?;
|
let (key, ip) = parse_entry(entry)?;
|
||||||
@@ -125,30 +154,6 @@ pub fn validate_entries(entries: &[String]) -> Result<()> {
|
|||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Replace runtime DNS overrides with a new validated snapshot.
|
|
||||||
pub fn install_entries(entries: &[String]) -> Result<()> {
|
|
||||||
let parsed = parse_entries(entries)?;
|
|
||||||
let mut guard = overrides_store().write().map_err(|_| {
|
|
||||||
ProxyError::Config("network.dns_overrides runtime lock is poisoned".to_string())
|
|
||||||
})?;
|
|
||||||
*guard = parsed;
|
|
||||||
Ok(())
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Resolve a hostname override for `(host, port)` if present.
|
|
||||||
pub fn resolve(host: &str, port: u16) -> Option<IpAddr> {
|
|
||||||
let key = (host.to_ascii_lowercase(), port);
|
|
||||||
overrides_store()
|
|
||||||
.read()
|
|
||||||
.ok()
|
|
||||||
.and_then(|guard| guard.get(&key).copied())
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Resolve a hostname override and construct a socket address when present.
|
|
||||||
pub fn resolve_socket_addr(host: &str, port: u16) -> Option<SocketAddr> {
|
|
||||||
resolve(host, port).map(|ip| SocketAddr::new(ip, port))
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Parse a runtime endpoint in `host:port` format.
|
/// Parse a runtime endpoint in `host:port` format.
|
||||||
///
|
///
|
||||||
/// Supports:
|
/// Supports:
|
||||||
@@ -199,12 +204,14 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn install_and_resolve_are_case_insensitive_for_host() {
|
fn generation_resolver_updates_are_case_insensitive_for_host() {
|
||||||
let entries = vec!["MyPetrovich.ru:8443:127.0.0.1".to_string()];
|
let entries = vec!["MyPetrovich.ru:8443:127.0.0.1".to_string()];
|
||||||
install_entries(&entries).unwrap();
|
let resolver = GenerationDnsResolver::from_entries(&entries).unwrap();
|
||||||
|
|
||||||
let resolved = resolve("mypetrovich.ru", 8443);
|
assert_eq!(
|
||||||
assert_eq!(resolved, Some("127.0.0.1".parse().unwrap()));
|
resolver.resolve_socket_addr("mypetrovich.ru", 8443),
|
||||||
|
Some("127.0.0.1:8443".parse().unwrap())
|
||||||
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
|
|||||||
+16
-3
@@ -3,6 +3,7 @@
|
|||||||
|
|
||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, UdpSocket};
|
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, UdpSocket};
|
||||||
|
use std::sync::Arc;
|
||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
|
|
||||||
use tokio::task::JoinSet;
|
use tokio::task::JoinSet;
|
||||||
@@ -12,7 +13,8 @@ use tracing::{debug, info, warn};
|
|||||||
use crate::config::{NetworkConfig, UpstreamConfig, UpstreamType};
|
use crate::config::{NetworkConfig, UpstreamConfig, UpstreamType};
|
||||||
use crate::error::Result;
|
use crate::error::Result;
|
||||||
use crate::network::stun::{
|
use crate::network::stun::{
|
||||||
DualStunResult, IpFamily, StunProbeResult, stun_probe_family_with_bind_and_tcp_fallback,
|
DualStunResult, IpFamily, StunProbeResult,
|
||||||
|
stun_probe_family_with_bind_tcp_fallback_and_resolver,
|
||||||
};
|
};
|
||||||
use crate::transport::UpstreamManager;
|
use crate::transport::UpstreamManager;
|
||||||
|
|
||||||
@@ -67,6 +69,11 @@ pub async fn run_probe(
|
|||||||
stun_nat_probe_concurrency: usize,
|
stun_nat_probe_concurrency: usize,
|
||||||
) -> Result<NetworkProbe> {
|
) -> Result<NetworkProbe> {
|
||||||
let mut probe = NetworkProbe::default();
|
let mut probe = NetworkProbe::default();
|
||||||
|
let dns_resolver = Arc::new(
|
||||||
|
crate::network::dns_overrides::GenerationDnsResolver::from_entries(
|
||||||
|
&config.dns_overrides,
|
||||||
|
)?,
|
||||||
|
);
|
||||||
let servers = collect_stun_servers(config);
|
let servers = collect_stun_servers(config);
|
||||||
let mut detected_ipv4 = detect_local_ip_v4();
|
let mut detected_ipv4 = detect_local_ip_v4();
|
||||||
let mut detected_ipv6 = detect_local_ip_v6();
|
let mut detected_ipv6 = detect_local_ip_v6();
|
||||||
@@ -88,6 +95,7 @@ pub async fn run_probe(
|
|||||||
None,
|
None,
|
||||||
None,
|
None,
|
||||||
config.stun_tcp_fallback,
|
config.stun_tcp_fallback,
|
||||||
|
Arc::clone(&dns_resolver),
|
||||||
)
|
)
|
||||||
.await
|
.await
|
||||||
}
|
}
|
||||||
@@ -171,6 +179,7 @@ pub async fn run_probe(
|
|||||||
bind_v4,
|
bind_v4,
|
||||||
bind_v6,
|
bind_v6,
|
||||||
config.stun_tcp_fallback,
|
config.stun_tcp_fallback,
|
||||||
|
Arc::clone(&dns_resolver),
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
if let Some(reflected) = direct_stun_res.v4.map(|r| r.reflected_addr) {
|
if let Some(reflected) = direct_stun_res.v4.map(|r| r.reflected_addr) {
|
||||||
@@ -286,6 +295,7 @@ async fn probe_stun_servers_parallel(
|
|||||||
bind_v4: Option<IpAddr>,
|
bind_v4: Option<IpAddr>,
|
||||||
bind_v6: Option<IpAddr>,
|
bind_v6: Option<IpAddr>,
|
||||||
tcp_fallback: bool,
|
tcp_fallback: bool,
|
||||||
|
dns_resolver: Arc<crate::network::dns_overrides::GenerationDnsResolver>,
|
||||||
) -> DualStunResult {
|
) -> DualStunResult {
|
||||||
let mut join_set = JoinSet::new();
|
let mut join_set = JoinSet::new();
|
||||||
let mut next_idx = 0usize;
|
let mut next_idx = 0usize;
|
||||||
@@ -295,6 +305,7 @@ async fn probe_stun_servers_parallel(
|
|||||||
while next_idx < servers.len() || !join_set.is_empty() {
|
while next_idx < servers.len() || !join_set.is_empty() {
|
||||||
while next_idx < servers.len() && join_set.len() < concurrency {
|
while next_idx < servers.len() && join_set.len() < concurrency {
|
||||||
let stun_addr = servers[next_idx].clone();
|
let stun_addr = servers[next_idx].clone();
|
||||||
|
let dns_resolver = Arc::clone(&dns_resolver);
|
||||||
next_idx += 1;
|
next_idx += 1;
|
||||||
join_set.spawn(async move {
|
join_set.spawn(async move {
|
||||||
let batch_timeout = if tcp_fallback {
|
let batch_timeout = if tcp_fallback {
|
||||||
@@ -303,18 +314,20 @@ async fn probe_stun_servers_parallel(
|
|||||||
STUN_BATCH_TIMEOUT
|
STUN_BATCH_TIMEOUT
|
||||||
};
|
};
|
||||||
let res = timeout(batch_timeout, async {
|
let res = timeout(batch_timeout, async {
|
||||||
let v4 = stun_probe_family_with_bind_and_tcp_fallback(
|
let v4 = stun_probe_family_with_bind_tcp_fallback_and_resolver(
|
||||||
&stun_addr,
|
&stun_addr,
|
||||||
IpFamily::V4,
|
IpFamily::V4,
|
||||||
bind_v4,
|
bind_v4,
|
||||||
tcp_fallback,
|
tcp_fallback,
|
||||||
|
Some(dns_resolver.as_ref()),
|
||||||
)
|
)
|
||||||
.await?;
|
.await?;
|
||||||
let v6 = stun_probe_family_with_bind_and_tcp_fallback(
|
let v6 = stun_probe_family_with_bind_tcp_fallback_and_resolver(
|
||||||
&stun_addr,
|
&stun_addr,
|
||||||
IpFamily::V6,
|
IpFamily::V6,
|
||||||
bind_v6,
|
bind_v6,
|
||||||
tcp_fallback,
|
tcp_fallback,
|
||||||
|
Some(dns_resolver.as_ref()),
|
||||||
)
|
)
|
||||||
.await?;
|
.await?;
|
||||||
Ok::<DualStunResult, crate::error::ProxyError>(DualStunResult { v4, v6 })
|
Ok::<DualStunResult, crate::error::ProxyError>(DualStunResult { v4, v6 })
|
||||||
|
|||||||
+33
-8
@@ -10,7 +10,7 @@ use tokio::time::{Duration, sleep, timeout};
|
|||||||
|
|
||||||
use crate::crypto::SecureRandom;
|
use crate::crypto::SecureRandom;
|
||||||
use crate::error::{ProxyError, Result};
|
use crate::error::{ProxyError, Result};
|
||||||
use crate::network::dns_overrides::{resolve, split_host_port};
|
use crate::network::dns_overrides::{GenerationDnsResolver, split_host_port};
|
||||||
|
|
||||||
fn stun_rng() -> &'static SecureRandom {
|
fn stun_rng() -> &'static SecureRandom {
|
||||||
static STUN_RNG: OnceLock<SecureRandom> = OnceLock::new();
|
static STUN_RNG: OnceLock<SecureRandom> = OnceLock::new();
|
||||||
@@ -80,13 +80,32 @@ pub async fn stun_probe_family_with_bind_and_tcp_fallback(
|
|||||||
family: IpFamily,
|
family: IpFamily,
|
||||||
bind_ip: Option<IpAddr>,
|
bind_ip: Option<IpAddr>,
|
||||||
tcp_fallback: bool,
|
tcp_fallback: bool,
|
||||||
|
) -> Result<Option<StunProbeResult>> {
|
||||||
|
stun_probe_family_with_bind_tcp_fallback_and_resolver(
|
||||||
|
stun_addr,
|
||||||
|
family,
|
||||||
|
bind_ip,
|
||||||
|
tcp_fallback,
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Probes one STUN family with an optional generation-owned DNS resolver.
|
||||||
|
pub async fn stun_probe_family_with_bind_tcp_fallback_and_resolver(
|
||||||
|
stun_addr: &str,
|
||||||
|
family: IpFamily,
|
||||||
|
bind_ip: Option<IpAddr>,
|
||||||
|
tcp_fallback: bool,
|
||||||
|
dns_resolver: Option<&GenerationDnsResolver>,
|
||||||
) -> Result<Option<StunProbeResult>> {
|
) -> Result<Option<StunProbeResult>> {
|
||||||
let udp_attempts = if tcp_fallback { 1 } else { 3 };
|
let udp_attempts = if tcp_fallback { 1 } else { 3 };
|
||||||
let udp_result = stun_probe_family_udp(stun_addr, family, bind_ip, udp_attempts).await?;
|
let udp_result =
|
||||||
|
stun_probe_family_udp(stun_addr, family, bind_ip, udp_attempts, dns_resolver).await?;
|
||||||
if udp_result.is_some() || !tcp_fallback {
|
if udp_result.is_some() || !tcp_fallback {
|
||||||
return Ok(udp_result);
|
return Ok(udp_result);
|
||||||
}
|
}
|
||||||
stun_probe_family_tcp(stun_addr, family, bind_ip).await
|
stun_probe_family_tcp(stun_addr, family, bind_ip, dns_resolver).await
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn stun_probe_family_udp(
|
async fn stun_probe_family_udp(
|
||||||
@@ -94,6 +113,7 @@ async fn stun_probe_family_udp(
|
|||||||
family: IpFamily,
|
family: IpFamily,
|
||||||
bind_ip: Option<IpAddr>,
|
bind_ip: Option<IpAddr>,
|
||||||
max_attempts: u8,
|
max_attempts: u8,
|
||||||
|
dns_resolver: Option<&GenerationDnsResolver>,
|
||||||
) -> Result<Option<StunProbeResult>> {
|
) -> Result<Option<StunProbeResult>> {
|
||||||
let bind_addr = match (family, bind_ip) {
|
let bind_addr = match (family, bind_ip) {
|
||||||
(IpFamily::V4, Some(IpAddr::V4(ip))) => SocketAddr::new(IpAddr::V4(ip), 0),
|
(IpFamily::V4, Some(IpAddr::V4(ip))) => SocketAddr::new(IpAddr::V4(ip), 0),
|
||||||
@@ -111,7 +131,7 @@ async fn stun_probe_family_udp(
|
|||||||
Err(e) => return Err(ProxyError::Proxy(format!("STUN bind failed: {e}"))),
|
Err(e) => return Err(ProxyError::Proxy(format!("STUN bind failed: {e}"))),
|
||||||
};
|
};
|
||||||
|
|
||||||
let target_addr = resolve_stun_addr(stun_addr, family).await?;
|
let target_addr = resolve_stun_addr(stun_addr, family, dns_resolver).await?;
|
||||||
if let Some(addr) = target_addr {
|
if let Some(addr) = target_addr {
|
||||||
match socket.connect(addr).await {
|
match socket.connect(addr).await {
|
||||||
Ok(()) => {}
|
Ok(()) => {}
|
||||||
@@ -182,8 +202,9 @@ async fn stun_probe_family_tcp(
|
|||||||
stun_addr: &str,
|
stun_addr: &str,
|
||||||
family: IpFamily,
|
family: IpFamily,
|
||||||
bind_ip: Option<IpAddr>,
|
bind_ip: Option<IpAddr>,
|
||||||
|
dns_resolver: Option<&GenerationDnsResolver>,
|
||||||
) -> Result<Option<StunProbeResult>> {
|
) -> Result<Option<StunProbeResult>> {
|
||||||
let target_addr = match resolve_stun_addr(stun_addr, family).await? {
|
let target_addr = match resolve_stun_addr(stun_addr, family, dns_resolver).await? {
|
||||||
Some(addr) => addr,
|
Some(addr) => addr,
|
||||||
None => return Ok(None),
|
None => return Ok(None),
|
||||||
};
|
};
|
||||||
@@ -360,7 +381,11 @@ fn parse_reflected_addr(buf: &[u8], txid: &[u8]) -> Option<SocketAddr> {
|
|||||||
None
|
None
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn resolve_stun_addr(stun_addr: &str, family: IpFamily) -> Result<Option<SocketAddr>> {
|
async fn resolve_stun_addr(
|
||||||
|
stun_addr: &str,
|
||||||
|
family: IpFamily,
|
||||||
|
dns_resolver: Option<&GenerationDnsResolver>,
|
||||||
|
) -> Result<Option<SocketAddr>> {
|
||||||
if let Ok(addr) = stun_addr.parse::<SocketAddr>() {
|
if let Ok(addr) = stun_addr.parse::<SocketAddr>() {
|
||||||
return Ok(match (addr.is_ipv4(), family) {
|
return Ok(match (addr.is_ipv4(), family) {
|
||||||
(true, IpFamily::V4) | (false, IpFamily::V6) => Some(addr),
|
(true, IpFamily::V4) | (false, IpFamily::V6) => Some(addr),
|
||||||
@@ -369,9 +394,9 @@ async fn resolve_stun_addr(stun_addr: &str, family: IpFamily) -> Result<Option<S
|
|||||||
}
|
}
|
||||||
|
|
||||||
if let Some((host, port)) = split_host_port(stun_addr)
|
if let Some((host, port)) = split_host_port(stun_addr)
|
||||||
&& let Some(ip) = resolve(&host, port)
|
&& let Some(addr) = dns_resolver
|
||||||
|
.and_then(|resolver| resolver.resolve_socket_addr(&host, port))
|
||||||
{
|
{
|
||||||
let addr = SocketAddr::new(ip, port);
|
|
||||||
return Ok(match (addr.is_ipv4(), family) {
|
return Ok(match (addr.is_ipv4(), family) {
|
||||||
(true, IpFamily::V4) | (false, IpFamily::V6) => Some(addr),
|
(true, IpFamily::V4) | (false, IpFamily::V6) => Some(addr),
|
||||||
_ => None,
|
_ => None,
|
||||||
|
|||||||
@@ -2,7 +2,7 @@
|
|||||||
|
|
||||||
use dashmap::DashMap;
|
use dashmap::DashMap;
|
||||||
use std::cmp::max;
|
use std::cmp::max;
|
||||||
use std::sync::OnceLock;
|
use std::sync::{Mutex, OnceLock};
|
||||||
use std::time::{Duration, Instant};
|
use std::time::{Duration, Instant};
|
||||||
|
|
||||||
const EMA_ALPHA: f64 = 0.2;
|
const EMA_ALPHA: f64 = 0.2;
|
||||||
@@ -294,6 +294,11 @@ fn profiles() -> &'static DashMap<String, UserAdaptiveProfile> {
|
|||||||
USER_PROFILES.get_or_init(DashMap::new)
|
USER_PROFILES.get_or_init(DashMap::new)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn profile_insert_guard() -> &'static Mutex<()> {
|
||||||
|
static PROFILE_INSERT_GUARD: OnceLock<Mutex<()>> = OnceLock::new();
|
||||||
|
PROFILE_INSERT_GUARD.get_or_init(|| Mutex::new(()))
|
||||||
|
}
|
||||||
|
|
||||||
/// Returns a fresh user's recent successful Direct tier, or `Base` when stale.
|
/// Returns a fresh user's recent successful Direct tier, or `Base` when stale.
|
||||||
#[allow(dead_code)]
|
#[allow(dead_code)]
|
||||||
pub fn seed_tier_for_user(user: &str) -> AdaptiveTier {
|
pub fn seed_tier_for_user(user: &str) -> AdaptiveTier {
|
||||||
@@ -320,29 +325,58 @@ pub fn record_user_tier(user: &str, tier: AdaptiveTier) {
|
|||||||
if user.len() > MAX_USER_KEY_BYTES {
|
if user.len() > MAX_USER_KEY_BYTES {
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
record_user_tier_with_cap(
|
||||||
|
profiles(),
|
||||||
|
profile_insert_guard(),
|
||||||
|
user,
|
||||||
|
tier,
|
||||||
|
MAX_USER_PROFILES_ENTRIES,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
fn record_user_tier_with_cap(
|
||||||
|
profiles: &DashMap<String, UserAdaptiveProfile>,
|
||||||
|
insert_guard: &Mutex<()>,
|
||||||
|
user: &str,
|
||||||
|
tier: AdaptiveTier,
|
||||||
|
max_entries: usize,
|
||||||
|
) {
|
||||||
let now = Instant::now();
|
let now = Instant::now();
|
||||||
let mut was_vacant = false;
|
if let Some(mut entry) = profiles.get_mut(user) {
|
||||||
match profiles().entry(user.to_string()) {
|
let effective = if now.saturating_duration_since(entry.seen_at) > PROFILE_TTL {
|
||||||
dashmap::mapref::entry::Entry::Occupied(mut entry) => {
|
tier
|
||||||
let existing = *entry.get();
|
} else {
|
||||||
let effective = if now.saturating_duration_since(existing.seen_at) > PROFILE_TTL {
|
max(entry.tier, tier)
|
||||||
tier
|
};
|
||||||
} else {
|
*entry = UserAdaptiveProfile {
|
||||||
max(existing.tier, tier)
|
tier: effective,
|
||||||
};
|
seen_at: now,
|
||||||
entry.insert(UserAdaptiveProfile {
|
};
|
||||||
tier: effective,
|
return;
|
||||||
seen_at: now,
|
|
||||||
});
|
|
||||||
}
|
|
||||||
dashmap::mapref::entry::Entry::Vacant(slot) => {
|
|
||||||
slot.insert(UserAdaptiveProfile { tier, seen_at: now });
|
|
||||||
was_vacant = true;
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
if was_vacant && profiles().len() > MAX_USER_PROFILES_ENTRIES {
|
|
||||||
profiles().retain(|_, v| now.saturating_duration_since(v.seen_at) <= PROFILE_TTL);
|
let _guard = insert_guard
|
||||||
|
.lock()
|
||||||
|
.unwrap_or_else(|poisoned| poisoned.into_inner());
|
||||||
|
if let Some(mut entry) = profiles.get_mut(user) {
|
||||||
|
let effective = if now.saturating_duration_since(entry.seen_at) > PROFILE_TTL {
|
||||||
|
tier
|
||||||
|
} else {
|
||||||
|
max(entry.tier, tier)
|
||||||
|
};
|
||||||
|
*entry = UserAdaptiveProfile {
|
||||||
|
tier: effective,
|
||||||
|
seen_at: now,
|
||||||
|
};
|
||||||
|
return;
|
||||||
}
|
}
|
||||||
|
if profiles.len() >= max_entries {
|
||||||
|
profiles.retain(|_, value| now.saturating_duration_since(value.seen_at) <= PROFILE_TTL);
|
||||||
|
}
|
||||||
|
if profiles.len() >= max_entries {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
profiles.insert(user.to_string(), UserAdaptiveProfile { tier, seen_at: now });
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
@@ -438,6 +472,17 @@ mod adaptive_direct_budget_policy_tests;
|
|||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn fresh_profile_cardinality_is_hard_bounded() {
|
||||||
|
let profiles = DashMap::new();
|
||||||
|
let insert_guard = Mutex::new(());
|
||||||
|
for index in 0..64 {
|
||||||
|
let user = format!("user-{index}");
|
||||||
|
record_user_tier_with_cap(&profiles, &insert_guard, &user, AdaptiveTier::Base, 16);
|
||||||
|
}
|
||||||
|
assert_eq!(profiles.len(), 16);
|
||||||
|
}
|
||||||
|
|
||||||
fn sample(
|
fn sample(
|
||||||
c2s_bytes: u64,
|
c2s_bytes: u64,
|
||||||
s2c_requested_bytes: u64,
|
s2c_requested_bytes: u64,
|
||||||
|
|||||||
+17
-2
@@ -50,7 +50,6 @@ use crate::proxy::handshake::{
|
|||||||
};
|
};
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
use crate::proxy::handshake::{handle_mtproto_handshake, handle_tls_handshake};
|
use crate::proxy::handshake::{handle_mtproto_handshake, handle_tls_handshake};
|
||||||
use crate::proxy::masking::handle_bad_client_with_shared;
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
use crate::proxy::route_mode::RelayRouteMode;
|
use crate::proxy::route_mode::RelayRouteMode;
|
||||||
use crate::proxy::route_mode::RouteRuntimeController;
|
use crate::proxy::route_mode::RouteRuntimeController;
|
||||||
@@ -247,6 +246,7 @@ fn masking_outcome<R, W>(
|
|||||||
peer: SocketAddr,
|
peer: SocketAddr,
|
||||||
local_addr: SocketAddr,
|
local_addr: SocketAddr,
|
||||||
config: Arc<ProxyConfig>,
|
config: Arc<ProxyConfig>,
|
||||||
|
upstream_manager: Arc<UpstreamManager>,
|
||||||
beobachten: Arc<BeobachtenStore>,
|
beobachten: Arc<BeobachtenStore>,
|
||||||
shared: Arc<ProxySharedState>,
|
shared: Arc<ProxySharedState>,
|
||||||
) -> HandshakeOutcome
|
) -> HandshakeOutcome
|
||||||
@@ -264,7 +264,7 @@ where
|
|||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
handle_bad_client_with_shared(
|
crate::proxy::masking::handle_bad_client_with_shared_resolver(
|
||||||
reader,
|
reader,
|
||||||
writer,
|
writer,
|
||||||
&initial_data,
|
&initial_data,
|
||||||
@@ -273,6 +273,7 @@ where
|
|||||||
&config,
|
&config,
|
||||||
&beobachten,
|
&beobachten,
|
||||||
shared.as_ref(),
|
shared.as_ref(),
|
||||||
|
Some(upstream_manager.as_ref()),
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
Ok(())
|
Ok(())
|
||||||
@@ -657,6 +658,7 @@ where
|
|||||||
real_peer,
|
real_peer,
|
||||||
local_addr,
|
local_addr,
|
||||||
config.clone(),
|
config.clone(),
|
||||||
|
upstream_manager.clone(),
|
||||||
beobachten.clone(),
|
beobachten.clone(),
|
||||||
shared.clone(),
|
shared.clone(),
|
||||||
));
|
));
|
||||||
@@ -679,6 +681,7 @@ where
|
|||||||
real_peer,
|
real_peer,
|
||||||
local_addr,
|
local_addr,
|
||||||
config.clone(),
|
config.clone(),
|
||||||
|
upstream_manager.clone(),
|
||||||
beobachten.clone(),
|
beobachten.clone(),
|
||||||
shared.clone(),
|
shared.clone(),
|
||||||
));
|
));
|
||||||
@@ -698,6 +701,7 @@ where
|
|||||||
real_peer,
|
real_peer,
|
||||||
local_addr,
|
local_addr,
|
||||||
config.clone(),
|
config.clone(),
|
||||||
|
upstream_manager.clone(),
|
||||||
beobachten.clone(),
|
beobachten.clone(),
|
||||||
shared.clone(),
|
shared.clone(),
|
||||||
));
|
));
|
||||||
@@ -729,6 +733,7 @@ where
|
|||||||
real_peer,
|
real_peer,
|
||||||
local_addr,
|
local_addr,
|
||||||
config.clone(),
|
config.clone(),
|
||||||
|
upstream_manager.clone(),
|
||||||
beobachten.clone(),
|
beobachten.clone(),
|
||||||
shared.clone(),
|
shared.clone(),
|
||||||
));
|
));
|
||||||
@@ -787,6 +792,7 @@ where
|
|||||||
real_peer,
|
real_peer,
|
||||||
local_addr,
|
local_addr,
|
||||||
config.clone(),
|
config.clone(),
|
||||||
|
upstream_manager.clone(),
|
||||||
beobachten.clone(),
|
beobachten.clone(),
|
||||||
shared.clone(),
|
shared.clone(),
|
||||||
));
|
));
|
||||||
@@ -817,6 +823,7 @@ where
|
|||||||
real_peer,
|
real_peer,
|
||||||
local_addr,
|
local_addr,
|
||||||
config.clone(),
|
config.clone(),
|
||||||
|
upstream_manager.clone(),
|
||||||
beobachten.clone(),
|
beobachten.clone(),
|
||||||
shared.clone(),
|
shared.clone(),
|
||||||
));
|
));
|
||||||
@@ -843,6 +850,7 @@ where
|
|||||||
real_peer,
|
real_peer,
|
||||||
local_addr,
|
local_addr,
|
||||||
config.clone(),
|
config.clone(),
|
||||||
|
upstream_manager.clone(),
|
||||||
beobachten.clone(),
|
beobachten.clone(),
|
||||||
shared.clone(),
|
shared.clone(),
|
||||||
));
|
));
|
||||||
@@ -1279,6 +1287,7 @@ impl RunningClientHandler {
|
|||||||
peer,
|
peer,
|
||||||
local_addr,
|
local_addr,
|
||||||
self.config.clone(),
|
self.config.clone(),
|
||||||
|
self.upstream_manager.clone(),
|
||||||
self.beobachten.clone(),
|
self.beobachten.clone(),
|
||||||
self.shared.clone(),
|
self.shared.clone(),
|
||||||
));
|
));
|
||||||
@@ -1301,6 +1310,7 @@ impl RunningClientHandler {
|
|||||||
peer,
|
peer,
|
||||||
local_addr,
|
local_addr,
|
||||||
self.config.clone(),
|
self.config.clone(),
|
||||||
|
self.upstream_manager.clone(),
|
||||||
self.beobachten.clone(),
|
self.beobachten.clone(),
|
||||||
self.shared.clone(),
|
self.shared.clone(),
|
||||||
));
|
));
|
||||||
@@ -1321,6 +1331,7 @@ impl RunningClientHandler {
|
|||||||
peer,
|
peer,
|
||||||
local_addr,
|
local_addr,
|
||||||
self.config.clone(),
|
self.config.clone(),
|
||||||
|
self.upstream_manager.clone(),
|
||||||
self.beobachten.clone(),
|
self.beobachten.clone(),
|
||||||
self.shared.clone(),
|
self.shared.clone(),
|
||||||
));
|
));
|
||||||
@@ -1377,6 +1388,7 @@ impl RunningClientHandler {
|
|||||||
peer,
|
peer,
|
||||||
local_addr,
|
local_addr,
|
||||||
config.clone(),
|
config.clone(),
|
||||||
|
self.upstream_manager.clone(),
|
||||||
self.beobachten.clone(),
|
self.beobachten.clone(),
|
||||||
self.shared.clone(),
|
self.shared.clone(),
|
||||||
));
|
));
|
||||||
@@ -1445,6 +1457,7 @@ impl RunningClientHandler {
|
|||||||
peer,
|
peer,
|
||||||
local_addr,
|
local_addr,
|
||||||
config.clone(),
|
config.clone(),
|
||||||
|
self.upstream_manager.clone(),
|
||||||
self.beobachten.clone(),
|
self.beobachten.clone(),
|
||||||
self.shared.clone(),
|
self.shared.clone(),
|
||||||
));
|
));
|
||||||
@@ -1493,6 +1506,7 @@ impl RunningClientHandler {
|
|||||||
peer,
|
peer,
|
||||||
local_addr,
|
local_addr,
|
||||||
self.config.clone(),
|
self.config.clone(),
|
||||||
|
self.upstream_manager.clone(),
|
||||||
self.beobachten.clone(),
|
self.beobachten.clone(),
|
||||||
self.shared.clone(),
|
self.shared.clone(),
|
||||||
));
|
));
|
||||||
@@ -1532,6 +1546,7 @@ impl RunningClientHandler {
|
|||||||
peer,
|
peer,
|
||||||
local_addr,
|
local_addr,
|
||||||
config.clone(),
|
config.clone(),
|
||||||
|
self.upstream_manager.clone(),
|
||||||
self.beobachten.clone(),
|
self.beobachten.clone(),
|
||||||
self.shared.clone(),
|
self.shared.clone(),
|
||||||
));
|
));
|
||||||
|
|||||||
@@ -98,7 +98,7 @@ pub(super) fn auth_probe_is_throttled_in(
|
|||||||
};
|
};
|
||||||
if auth_probe_state_expired(&entry, now) {
|
if auth_probe_state_expired(&entry, now) {
|
||||||
drop(entry);
|
drop(entry);
|
||||||
state.remove(&peer_ip);
|
state.remove_if(&peer_ip, |_, current| auth_probe_state_expired(current, now));
|
||||||
return false;
|
return false;
|
||||||
}
|
}
|
||||||
now < entry.blocked_until
|
now < entry.blocked_until
|
||||||
@@ -116,7 +116,7 @@ pub(super) fn auth_probe_saturation_grace_exhausted_in(
|
|||||||
};
|
};
|
||||||
if auth_probe_state_expired(&entry, now) {
|
if auth_probe_state_expired(&entry, now) {
|
||||||
drop(entry);
|
drop(entry);
|
||||||
state.remove(&peer_ip);
|
state.remove_if(&peer_ip, |_, current| auth_probe_state_expired(current, now));
|
||||||
return false;
|
return false;
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -264,11 +264,19 @@ pub(super) fn auth_probe_record_failure_with_state_in(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
let Some((evict_key, _, _)) = eviction_candidate else {
|
let Some((evict_key, evict_fail_streak, evict_last_seen)) = eviction_candidate else {
|
||||||
return;
|
return;
|
||||||
};
|
};
|
||||||
state.remove(&evict_key);
|
if state
|
||||||
break;
|
.remove_if(&evict_key, |_, current| {
|
||||||
|
current.fail_streak == evict_fail_streak
|
||||||
|
&& current.last_seen == evict_last_seen
|
||||||
|
})
|
||||||
|
.is_some()
|
||||||
|
{
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
continue;
|
||||||
}
|
}
|
||||||
|
|
||||||
let mut stale_keys = Vec::new();
|
let mut stale_keys = Vec::new();
|
||||||
@@ -334,18 +342,22 @@ pub(super) fn auth_probe_record_failure_with_state_in(
|
|||||||
}
|
}
|
||||||
|
|
||||||
for stale_key in stale_keys {
|
for stale_key in stale_keys {
|
||||||
state.remove(&stale_key);
|
state.remove_if(&stale_key, |_, current| {
|
||||||
|
auth_probe_state_expired(current, now)
|
||||||
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
if state.len() < AUTH_PROBE_TRACK_MAX_ENTRIES {
|
if state.len() < AUTH_PROBE_TRACK_MAX_ENTRIES {
|
||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
|
|
||||||
let Some((evict_key, _, _)) = eviction_candidate else {
|
let Some((evict_key, evict_fail_streak, evict_last_seen)) = eviction_candidate else {
|
||||||
auth_probe_note_saturation_in(shared, now);
|
auth_probe_note_saturation_in(shared, now);
|
||||||
return;
|
return;
|
||||||
};
|
};
|
||||||
state.remove(&evict_key);
|
state.remove_if(&evict_key, |_, current| {
|
||||||
|
current.fail_streak == evict_fail_streak && current.last_seen == evict_last_seen
|
||||||
|
});
|
||||||
auth_probe_note_saturation_in(shared, now);
|
auth_probe_note_saturation_in(shared, now);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -266,11 +266,11 @@ where
|
|||||||
return HandshakeResult::BadClient { reader, writer };
|
return HandshakeResult::BadClient { reader, writer };
|
||||||
}
|
}
|
||||||
|
|
||||||
|
let selected_tls_domain =
|
||||||
|
matched_tls_domain.unwrap_or(config.censorship.tls_domain.as_str());
|
||||||
let cached_entry = if config.censorship.tls_emulation {
|
let cached_entry = if config.censorship.tls_emulation {
|
||||||
if let Some(cache) = tls_cache.as_ref() {
|
if let Some(cache) = tls_cache.as_ref() {
|
||||||
let selected_domain =
|
let cached_entry = cache.get(selected_tls_domain).await;
|
||||||
matched_tls_domain.unwrap_or(config.censorship.tls_domain.as_str());
|
|
||||||
let cached_entry = cache.get(selected_domain).await;
|
|
||||||
Some(cached_entry)
|
Some(cached_entry)
|
||||||
} else {
|
} else {
|
||||||
None
|
None
|
||||||
@@ -322,6 +322,7 @@ where
|
|||||||
if let Some(cache) = tls_cache.as_ref() {
|
if let Some(cache) = tls_cache.as_ref() {
|
||||||
cache
|
cache
|
||||||
.take_full_cert_budget_for_ip(
|
.take_full_cert_budget_for_ip(
|
||||||
|
selected_tls_domain,
|
||||||
peer.ip(),
|
peer.ip(),
|
||||||
Duration::from_secs(config.censorship.tls_full_cert_ttl_secs),
|
Duration::from_secs(config.censorship.tls_full_cert_ttl_secs),
|
||||||
)
|
)
|
||||||
|
|||||||
+42
-7
@@ -1,7 +1,6 @@
|
|||||||
//! Masking - forward unrecognized traffic to mask host
|
//! Masking - forward unrecognized traffic to mask host
|
||||||
|
|
||||||
use crate::config::ProxyConfig;
|
use crate::config::ProxyConfig;
|
||||||
use crate::network::dns_overrides::resolve_socket_addr;
|
|
||||||
use crate::protocol::tls;
|
use crate::protocol::tls;
|
||||||
use crate::proxy::shared_state::ProxySharedState;
|
use crate::proxy::shared_state::ProxySharedState;
|
||||||
use crate::stats::beobachten::BeobachtenStore;
|
use crate::stats::beobachten::BeobachtenStore;
|
||||||
@@ -41,6 +40,7 @@ const MASK_RELAY_IDLE_TIMEOUT: Duration = Duration::from_millis(100);
|
|||||||
const MASK_BUFFER_SIZE: usize = 8192;
|
const MASK_BUFFER_SIZE: usize = 8192;
|
||||||
const MASK_BUFFER_GROW_AFTER_BYTES: usize = 256 * 1024;
|
const MASK_BUFFER_GROW_AFTER_BYTES: usize = 256 * 1024;
|
||||||
const MASK_BUFFER_MAX_SIZE: usize = 64 * 1024;
|
const MASK_BUFFER_MAX_SIZE: usize = 64 * 1024;
|
||||||
|
const MASK_DNS_RESULT_MAX_ADDRESSES: usize = 64;
|
||||||
#[cfg(unix)]
|
#[cfg(unix)]
|
||||||
#[cfg(not(test))]
|
#[cfg(not(test))]
|
||||||
const LOCAL_INTERFACE_CACHE_TTL: Duration = Duration::from_secs(300);
|
const LOCAL_INTERFACE_CACHE_TTL: Duration = Duration::from_secs(300);
|
||||||
@@ -532,19 +532,25 @@ fn parse_mask_host_ip_literal(host: &str) -> Option<IpAddr> {
|
|||||||
async fn resolve_mask_target_addrs(
|
async fn resolve_mask_target_addrs(
|
||||||
mask_host: &str,
|
mask_host: &str,
|
||||||
mask_port: u16,
|
mask_port: u16,
|
||||||
|
upstream_manager: Option<&crate::transport::UpstreamManager>,
|
||||||
) -> std::io::Result<Vec<SocketAddr>> {
|
) -> std::io::Result<Vec<SocketAddr>> {
|
||||||
if let Some(addr) = resolve_socket_addr(mask_host, mask_port) {
|
|
||||||
return Ok(vec![addr]);
|
|
||||||
}
|
|
||||||
|
|
||||||
if let Some(ip) = parse_mask_host_ip_literal(mask_host) {
|
if let Some(ip) = parse_mask_host_ip_literal(mask_host) {
|
||||||
return Ok(vec![SocketAddr::new(ip, mask_port)]);
|
return Ok(vec![SocketAddr::new(ip, mask_port)]);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if let Some(upstream_manager) = upstream_manager {
|
||||||
|
return upstream_manager
|
||||||
|
.resolve_all(mask_host, mask_port)
|
||||||
|
.await
|
||||||
|
.map_err(|error| IoError::new(ErrorKind::NotFound, error.to_string()));
|
||||||
|
}
|
||||||
|
|
||||||
let addrs = timeout(MASK_TIMEOUT, lookup_host((mask_host, mask_port)))
|
let addrs = timeout(MASK_TIMEOUT, lookup_host((mask_host, mask_port)))
|
||||||
.await
|
.await
|
||||||
.map_err(|_| IoError::new(ErrorKind::TimedOut, "mask target DNS lookup timed out"))??;
|
.map_err(|_| IoError::new(ErrorKind::TimedOut, "mask target DNS lookup timed out"))??;
|
||||||
let addrs = addrs.collect::<Vec<_>>();
|
let addrs = addrs
|
||||||
|
.take(MASK_DNS_RESULT_MAX_ADDRESSES)
|
||||||
|
.collect::<Vec<_>>();
|
||||||
if addrs.is_empty() {
|
if addrs.is_empty() {
|
||||||
return Err(IoError::new(
|
return Err(IoError::new(
|
||||||
ErrorKind::NotFound,
|
ErrorKind::NotFound,
|
||||||
@@ -999,6 +1005,34 @@ pub(crate) async fn handle_bad_client_with_shared<R, W>(
|
|||||||
) where
|
) where
|
||||||
R: AsyncRead + Unpin + Send + 'static,
|
R: AsyncRead + Unpin + Send + 'static,
|
||||||
W: AsyncWrite + Unpin + Send + 'static,
|
W: AsyncWrite + Unpin + Send + 'static,
|
||||||
|
{
|
||||||
|
handle_bad_client_with_shared_resolver(
|
||||||
|
reader,
|
||||||
|
writer,
|
||||||
|
initial_data,
|
||||||
|
peer,
|
||||||
|
local_addr,
|
||||||
|
config,
|
||||||
|
beobachten,
|
||||||
|
shared,
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) async fn handle_bad_client_with_shared_resolver<R, W>(
|
||||||
|
reader: R,
|
||||||
|
writer: W,
|
||||||
|
initial_data: &[u8],
|
||||||
|
peer: SocketAddr,
|
||||||
|
local_addr: SocketAddr,
|
||||||
|
config: &ProxyConfig,
|
||||||
|
beobachten: &BeobachtenStore,
|
||||||
|
shared: &ProxySharedState,
|
||||||
|
upstream_manager: Option<&crate::transport::UpstreamManager>,
|
||||||
|
) where
|
||||||
|
R: AsyncRead + Unpin + Send + 'static,
|
||||||
|
W: AsyncWrite + Unpin + Send + 'static,
|
||||||
{
|
{
|
||||||
let client_type = detect_client_type(initial_data);
|
let client_type = detect_client_type(initial_data);
|
||||||
if config.general.beobachten {
|
if config.general.beobachten {
|
||||||
@@ -1112,7 +1146,8 @@ pub(crate) async fn handle_bad_client_with_shared<R, W>(
|
|||||||
let mask_host = mask_target.host;
|
let mask_host = mask_target.host;
|
||||||
let mask_port = mask_target.port;
|
let mask_port = mask_target.port;
|
||||||
|
|
||||||
let resolved_mask_addrs = match resolve_mask_target_addrs(mask_host, mask_port).await {
|
let resolved_mask_addrs =
|
||||||
|
match resolve_mask_target_addrs(mask_host, mask_port, upstream_manager).await {
|
||||||
Ok(addrs) => addrs,
|
Ok(addrs) => addrs,
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
let outcome_started = Instant::now();
|
let outcome_started = Instant::now();
|
||||||
|
|||||||
@@ -34,7 +34,7 @@ use crate::stats::{
|
|||||||
MeD2cFlushReason, MeD2cQuotaRejectStage, MeD2cWriteMode, QuotaReserveError, Stats, UserStats,
|
MeD2cFlushReason, MeD2cQuotaRejectStage, MeD2cWriteMode, QuotaReserveError, Stats, UserStats,
|
||||||
};
|
};
|
||||||
use crate::stream::{BufferPool, CryptoReader, CryptoWriter, PooledBuffer};
|
use crate::stream::{BufferPool, CryptoReader, CryptoWriter, PooledBuffer};
|
||||||
use crate::transport::middle_proxy::{MePool, MeResponse, proto_flags_for_tag};
|
use crate::transport::middle_proxy::{ConnLease, MePool, MeResponse, proto_flags_for_tag};
|
||||||
|
|
||||||
mod c2me;
|
mod c2me;
|
||||||
mod d2c;
|
mod d2c;
|
||||||
|
|||||||
@@ -1,5 +1,42 @@
|
|||||||
use super::*;
|
use super::*;
|
||||||
|
|
||||||
|
struct RelayConnLease {
|
||||||
|
connection: Option<ConnLease>,
|
||||||
|
conn_id: u64,
|
||||||
|
shared: Arc<ProxySharedState>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl RelayConnLease {
|
||||||
|
fn new(connection: ConnLease, shared: Arc<ProxySharedState>) -> Self {
|
||||||
|
let conn_id = connection.conn_id();
|
||||||
|
Self {
|
||||||
|
connection: Some(connection),
|
||||||
|
conn_id,
|
||||||
|
shared,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn conn_id(&self) -> u64 {
|
||||||
|
self.conn_id
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn unregister(mut self) {
|
||||||
|
let Some(connection) = self.connection.take() else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
clear_relay_idle_candidate_in(self.shared.as_ref(), connection.conn_id());
|
||||||
|
connection.unregister().await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Drop for RelayConnLease {
|
||||||
|
fn drop(&mut self) {
|
||||||
|
if let Some(connection) = self.connection.as_ref() {
|
||||||
|
clear_relay_idle_candidate_in(self.shared.as_ref(), connection.conn_id());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// Runs Middle-End relay with explicit kernel-conntrack close publication policy.
|
/// Runs Middle-End relay with explicit kernel-conntrack close publication policy.
|
||||||
pub(crate) async fn handle_via_middle_proxy_with_conntrack<R, W>(
|
pub(crate) async fn handle_via_middle_proxy_with_conntrack<R, W>(
|
||||||
mut crypto_reader: CryptoReader<R>,
|
mut crypto_reader: CryptoReader<R>,
|
||||||
@@ -44,7 +81,11 @@ where
|
|||||||
"Routing via Middle-End"
|
"Routing via Middle-End"
|
||||||
);
|
);
|
||||||
|
|
||||||
let (conn_id, me_rx) = me_pool.registry().register().await;
|
let Some((connection, me_rx)) = me_pool.register_connection().await else {
|
||||||
|
return Err(ProxyError::MiddleConnectionLost);
|
||||||
|
};
|
||||||
|
let relay_connection = RelayConnLease::new(connection, Arc::clone(&shared));
|
||||||
|
let conn_id = relay_connection.conn_id();
|
||||||
let trace_id = session_id;
|
let trace_id = session_id;
|
||||||
let bytes_me2c = Arc::new(AtomicU64::new(0));
|
let bytes_me2c = Arc::new(AtomicU64::new(0));
|
||||||
let mut forensics = RelayForensicsState {
|
let mut forensics = RelayForensicsState {
|
||||||
@@ -76,7 +117,7 @@ where
|
|||||||
let _cutover_park_lease = stats.acquire_middle_cutover_park_lease();
|
let _cutover_park_lease = stats.acquire_middle_cutover_park_lease();
|
||||||
tokio::time::sleep(delay).await;
|
tokio::time::sleep(delay).await;
|
||||||
let _ = me_pool.send_close(conn_id).await;
|
let _ = me_pool.send_close(conn_id).await;
|
||||||
me_pool.registry().unregister(conn_id).await;
|
relay_connection.unregister().await;
|
||||||
return Err(ProxyError::RouteSwitched);
|
return Err(ProxyError::RouteSwitched);
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -825,8 +866,7 @@ where
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
clear_relay_idle_candidate_in(shared.as_ref(), conn_id);
|
relay_connection.unregister().await;
|
||||||
me_pool.registry().unregister(conn_id).await;
|
|
||||||
let pool_snapshot = buffer_pool.stats();
|
let pool_snapshot = buffer_pool.stats();
|
||||||
stats.set_buffer_pool_gauges(
|
stats.set_buffer_pool_gauges(
|
||||||
pool_snapshot.pooled,
|
pool_snapshot.pooled,
|
||||||
|
|||||||
@@ -229,7 +229,7 @@ async fn masking_fallback_down_mimics_timeout() {
|
|||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn masking_ssrf_resolve_internal_ranges_blocked() {
|
async fn masking_ssrf_resolve_internal_ranges_blocked() {
|
||||||
use crate::network::dns_overrides::resolve_socket_addr;
|
use crate::network::dns_overrides::DnsOverrides;
|
||||||
|
|
||||||
let blocked_ips = [
|
let blocked_ips = [
|
||||||
"127.0.0.1",
|
"127.0.0.1",
|
||||||
@@ -238,10 +238,11 @@ async fn masking_ssrf_resolve_internal_ranges_blocked() {
|
|||||||
"192.168.1.1",
|
"192.168.1.1",
|
||||||
"0.0.0.0",
|
"0.0.0.0",
|
||||||
];
|
];
|
||||||
|
let resolver = DnsOverrides::default();
|
||||||
|
|
||||||
for ip in blocked_ips {
|
for ip in blocked_ips {
|
||||||
assert!(
|
assert!(
|
||||||
resolve_socket_addr(ip, 80).is_none(),
|
resolver.resolve_socket_addr(ip, 80).is_none(),
|
||||||
"runtime DNS overrides must not resolve unconfigured literal host targets"
|
"runtime DNS overrides must not resolve unconfigured literal host targets"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,5 +1,5 @@
|
|||||||
use super::*;
|
use super::*;
|
||||||
use crate::network::dns_overrides::install_entries;
|
use std::sync::Arc;
|
||||||
use tokio::io::{AsyncReadExt, AsyncWriteExt, duplex};
|
use tokio::io::{AsyncReadExt, AsyncWriteExt, duplex};
|
||||||
use tokio::time::{Duration, Instant, timeout};
|
use tokio::time::{Duration, Instant, timeout};
|
||||||
|
|
||||||
@@ -8,6 +8,7 @@ async fn run_connect_failure_case(
|
|||||||
port: u16,
|
port: u16,
|
||||||
timing_normalization_enabled: bool,
|
timing_normalization_enabled: bool,
|
||||||
peer: SocketAddr,
|
peer: SocketAddr,
|
||||||
|
dns_overrides: Vec<String>,
|
||||||
) -> Duration {
|
) -> Duration {
|
||||||
let mut config = ProxyConfig::default();
|
let mut config = ProxyConfig::default();
|
||||||
config.general.beobachten = false;
|
config.general.beobachten = false;
|
||||||
@@ -21,6 +22,19 @@ async fn run_connect_failure_case(
|
|||||||
|
|
||||||
let local_addr: SocketAddr = "127.0.0.1:443".parse().unwrap();
|
let local_addr: SocketAddr = "127.0.0.1:443".parse().unwrap();
|
||||||
let beobachten = BeobachtenStore::new();
|
let beobachten = BeobachtenStore::new();
|
||||||
|
let upstream_manager = crate::transport::UpstreamManager::new(
|
||||||
|
Vec::new(),
|
||||||
|
1,
|
||||||
|
1,
|
||||||
|
100,
|
||||||
|
1,
|
||||||
|
1,
|
||||||
|
true,
|
||||||
|
Arc::new(crate::stats::Stats::new()),
|
||||||
|
)
|
||||||
|
.with_dns_overrides(&dns_overrides)
|
||||||
|
.unwrap();
|
||||||
|
let shared = ProxySharedState::new();
|
||||||
let probe = b"CONNECT example.org:443 HTTP/1.1\r\nHost: example.org\r\n\r\n";
|
let probe = b"CONNECT example.org:443 HTTP/1.1\r\nHost: example.org\r\n\r\n";
|
||||||
|
|
||||||
let (mut client_writer, client_reader) = duplex(1024);
|
let (mut client_writer, client_reader) = duplex(1024);
|
||||||
@@ -28,7 +42,7 @@ async fn run_connect_failure_case(
|
|||||||
|
|
||||||
let started = Instant::now();
|
let started = Instant::now();
|
||||||
let task = tokio::spawn(async move {
|
let task = tokio::spawn(async move {
|
||||||
handle_bad_client(
|
handle_bad_client_with_shared_resolver(
|
||||||
client_reader,
|
client_reader,
|
||||||
client_visible_writer,
|
client_visible_writer,
|
||||||
probe,
|
probe,
|
||||||
@@ -36,6 +50,8 @@ async fn run_connect_failure_case(
|
|||||||
local_addr,
|
local_addr,
|
||||||
&config,
|
&config,
|
||||||
&beobachten,
|
&beobachten,
|
||||||
|
shared.as_ref(),
|
||||||
|
Some(&upstream_manager),
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
});
|
});
|
||||||
@@ -71,8 +87,14 @@ async fn connect_failure_refusal_close_behavior_matrix() {
|
|||||||
.parse()
|
.parse()
|
||||||
.unwrap();
|
.unwrap();
|
||||||
let elapsed =
|
let elapsed =
|
||||||
run_connect_failure_case("127.0.0.1", unused_port, timing_normalization_enabled, peer)
|
run_connect_failure_case(
|
||||||
.await;
|
"127.0.0.1",
|
||||||
|
unused_port,
|
||||||
|
timing_normalization_enabled,
|
||||||
|
peer,
|
||||||
|
Vec::new(),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
|
||||||
if timing_normalization_enabled {
|
if timing_normalization_enabled {
|
||||||
assert!(
|
assert!(
|
||||||
@@ -94,9 +116,6 @@ async fn connect_failure_overridden_hostname_close_behavior_matrix() {
|
|||||||
let unused_port = temp_listener.local_addr().unwrap().port();
|
let unused_port = temp_listener.local_addr().unwrap().port();
|
||||||
drop(temp_listener);
|
drop(temp_listener);
|
||||||
|
|
||||||
// Make hostname resolution deterministic in tests so timing ceilings are meaningful.
|
|
||||||
install_entries(&[format!("mask.invalid:{}:127.0.0.1", unused_port)]).unwrap();
|
|
||||||
|
|
||||||
for (idx, timing_normalization_enabled) in [false, true].into_iter().enumerate() {
|
for (idx, timing_normalization_enabled) in [false, true].into_iter().enumerate() {
|
||||||
let peer: SocketAddr = format!("203.0.113.220:{}", 54200 + idx as u16)
|
let peer: SocketAddr = format!("203.0.113.220:{}", 54200 + idx as u16)
|
||||||
.parse()
|
.parse()
|
||||||
@@ -106,6 +125,7 @@ async fn connect_failure_overridden_hostname_close_behavior_matrix() {
|
|||||||
unused_port,
|
unused_port,
|
||||||
timing_normalization_enabled,
|
timing_normalization_enabled,
|
||||||
peer,
|
peer,
|
||||||
|
vec![format!("mask.invalid:{}:127.0.0.1", unused_port)],
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
@@ -121,6 +141,4 @@ async fn connect_failure_overridden_hostname_close_behavior_matrix() {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
install_entries(&[]).unwrap();
|
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -355,9 +355,10 @@ impl CidrBucket {
|
|||||||
}
|
}
|
||||||
|
|
||||||
fn acquire_user_share(&self, user: &str) -> Arc<CidrUserShare> {
|
fn acquire_user_share(&self, user: &str) -> Arc<CidrUserShare> {
|
||||||
let share = self.users.get_or_insert_with(user, CidrUserShare::new);
|
self.users
|
||||||
share.active_conns.fetch_add(1, Ordering::Relaxed);
|
.get_or_insert_with(user, CidrUserShare::new, |share| {
|
||||||
share
|
share.active_conns.fetch_add(1, Ordering::Relaxed);
|
||||||
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
fn release_user_share(&self, user: &str, share: &Arc<CidrUserShare>) {
|
fn release_user_share(&self, user: &str, share: &Arc<CidrUserShare>) {
|
||||||
@@ -487,15 +488,20 @@ impl<T> ShardedRegistry<T> {
|
|||||||
(hasher.finish() as usize) & self.mask
|
(hasher.finish() as usize) & self.mask
|
||||||
}
|
}
|
||||||
|
|
||||||
fn get_or_insert_with<F>(&self, key: &str, make: F) -> Arc<T>
|
fn get_or_insert_with<F, A>(&self, key: &str, make: F, activate: A) -> Arc<T>
|
||||||
where
|
where
|
||||||
F: FnOnce() -> T,
|
F: FnOnce() -> T,
|
||||||
|
A: FnOnce(&Arc<T>),
|
||||||
{
|
{
|
||||||
let shard = &self.shards[self.shard_index(key)];
|
let shard = &self.shards[self.shard_index(key)];
|
||||||
match shard.entry(key.to_string()) {
|
match shard.entry(key.to_string()) {
|
||||||
dashmap::mapref::entry::Entry::Occupied(entry) => Arc::clone(entry.get()),
|
dashmap::mapref::entry::Entry::Occupied(entry) => {
|
||||||
|
activate(entry.get());
|
||||||
|
Arc::clone(entry.get())
|
||||||
|
}
|
||||||
dashmap::mapref::entry::Entry::Vacant(slot) => {
|
dashmap::mapref::entry::Entry::Vacant(slot) => {
|
||||||
let value = Arc::new(make());
|
let value = Arc::new(make());
|
||||||
|
activate(&value);
|
||||||
slot.insert(Arc::clone(&value));
|
slot.insert(Arc::clone(&value));
|
||||||
value
|
value
|
||||||
}
|
}
|
||||||
@@ -516,14 +522,7 @@ impl<T> ShardedRegistry<T> {
|
|||||||
F: Fn(&Arc<T>) -> bool,
|
F: Fn(&Arc<T>) -> bool,
|
||||||
{
|
{
|
||||||
let shard = &self.shards[self.shard_index(key)];
|
let shard = &self.shards[self.shard_index(key)];
|
||||||
let should_remove = match shard.get(key) {
|
shard.remove_if(key, |_, value| predicate(value)).is_some()
|
||||||
Some(entry) => predicate(entry.value()),
|
|
||||||
None => false,
|
|
||||||
};
|
|
||||||
if !should_remove {
|
|
||||||
return false;
|
|
||||||
}
|
|
||||||
shard.remove(key).is_some()
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -743,9 +742,10 @@ impl TrafficLimiter {
|
|||||||
if let Some(limit) = policy.user_limits.get(user).copied() {
|
if let Some(limit) = policy.user_limits.get(user).copied() {
|
||||||
let bucket = self
|
let bucket = self
|
||||||
.user_buckets
|
.user_buckets
|
||||||
.get_or_insert_with(user, || UserBucket::new(limit));
|
.get_or_insert_with(user, || UserBucket::new(limit), |bucket| {
|
||||||
|
bucket.active_leases.fetch_add(1, Ordering::Relaxed);
|
||||||
|
});
|
||||||
bucket.set_rates(limit);
|
bucket.set_rates(limit);
|
||||||
bucket.active_leases.fetch_add(1, Ordering::Relaxed);
|
|
||||||
self.user_scope
|
self.user_scope
|
||||||
.active_leases
|
.active_leases
|
||||||
.fetch_add(1, Ordering::Relaxed);
|
.fetch_add(1, Ordering::Relaxed);
|
||||||
@@ -762,9 +762,10 @@ impl TrafficLimiter {
|
|||||||
};
|
};
|
||||||
let bucket = self
|
let bucket = self
|
||||||
.cidr_buckets
|
.cidr_buckets
|
||||||
.get_or_insert_with(key, || CidrBucket::new(limits));
|
.get_or_insert_with(key, || CidrBucket::new(limits), |bucket| {
|
||||||
|
bucket.active_leases.fetch_add(1, Ordering::Relaxed);
|
||||||
|
});
|
||||||
bucket.set_rates(limits);
|
bucket.set_rates(limits);
|
||||||
bucket.active_leases.fetch_add(1, Ordering::Relaxed);
|
|
||||||
self.cidr_scope
|
self.cidr_scope
|
||||||
.active_leases
|
.active_leases
|
||||||
.fetch_add(1, Ordering::Relaxed);
|
.fetch_add(1, Ordering::Relaxed);
|
||||||
|
|||||||
+376
-79
@@ -1,12 +1,18 @@
|
|||||||
use std::collections::BTreeMap;
|
use std::collections::{BTreeMap, BTreeSet};
|
||||||
use std::path::Path;
|
use std::io::Write;
|
||||||
|
use std::path::{Path, PathBuf};
|
||||||
|
use std::sync::Arc;
|
||||||
use std::time::{SystemTime, UNIX_EPOCH};
|
use std::time::{SystemTime, UNIX_EPOCH};
|
||||||
|
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
use tokio::io::AsyncWriteExt;
|
use tokio::io::AsyncReadExt;
|
||||||
|
use tokio::sync::Mutex;
|
||||||
use tracing::{info, warn};
|
use tracing::{info, warn};
|
||||||
|
|
||||||
use crate::stats::{Stats, UserQuotaSnapshot};
|
use crate::stats::{QuotaStore, UserQuotaSnapshot};
|
||||||
|
|
||||||
|
const QUOTA_STATE_MAX_BYTES: u64 = 16 * 1024 * 1024;
|
||||||
|
const QUOTA_STATE_MAX_USERS: usize = 65_536;
|
||||||
|
|
||||||
#[derive(Debug, Default, Serialize, Deserialize)]
|
#[derive(Debug, Default, Serialize, Deserialize)]
|
||||||
pub(crate) struct QuotaStateFile {
|
pub(crate) struct QuotaStateFile {
|
||||||
@@ -20,6 +26,148 @@ pub(crate) struct QuotaUserState {
|
|||||||
pub(crate) last_reset_epoch_secs: u64,
|
pub(crate) last_reset_epoch_secs: u64,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Serialized process owner for bounded durable quota checkpoints and resets.
|
||||||
|
pub(crate) struct QuotaStateOwner {
|
||||||
|
path: PathBuf,
|
||||||
|
store: Arc<QuotaStore>,
|
||||||
|
mutation: Arc<Mutex<()>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl QuotaStateOwner {
|
||||||
|
/// Creates a process-owned quota persistence coordinator.
|
||||||
|
pub(crate) fn new(path: PathBuf, store: Arc<QuotaStore>) -> Arc<Self> {
|
||||||
|
Arc::new(Self {
|
||||||
|
path,
|
||||||
|
store,
|
||||||
|
mutation: Arc::new(Mutex::new(())),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Returns the process-owned quota checkpoint path.
|
||||||
|
pub(crate) fn path(&self) -> &Path {
|
||||||
|
&self.path
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Loads only quota entries owned by currently configured users.
|
||||||
|
pub(crate) async fn load(&self, configured_users: &BTreeSet<String>) {
|
||||||
|
let _guard = self.mutation.lock().await;
|
||||||
|
let state = match read_state_file(&self.path).await {
|
||||||
|
Ok(Some(state)) => state,
|
||||||
|
Ok(None) => return,
|
||||||
|
Err(error) => {
|
||||||
|
warn!(
|
||||||
|
error = %error,
|
||||||
|
path = %self.path.display(),
|
||||||
|
"Failed to load quota state file"
|
||||||
|
);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let persisted_users = state.users.len();
|
||||||
|
let mut loaded_users = 0usize;
|
||||||
|
for (user, quota) in state.users {
|
||||||
|
if loaded_users >= QUOTA_STATE_MAX_USERS || !configured_users.contains(&user) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
self.store
|
||||||
|
.load(&user, quota.used_bytes, quota.last_reset_epoch_secs);
|
||||||
|
loaded_users = loaded_users.saturating_add(1);
|
||||||
|
}
|
||||||
|
info!(
|
||||||
|
path = %self.path.display(),
|
||||||
|
loaded_users,
|
||||||
|
skipped_users = persisted_users.saturating_sub(loaded_users),
|
||||||
|
"Loaded bounded per-user quota state"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Persists a checkpoint filtered to the active configured user set.
|
||||||
|
pub(crate) async fn save(
|
||||||
|
&self,
|
||||||
|
configured_users: &BTreeSet<String>,
|
||||||
|
) -> std::io::Result<()> {
|
||||||
|
let guard = Arc::clone(&self.mutation).lock_owned().await;
|
||||||
|
let state = self.state_for_users(configured_users, None);
|
||||||
|
let path = self.path.clone();
|
||||||
|
let task = tokio::task::spawn_blocking(move || {
|
||||||
|
let _guard = guard;
|
||||||
|
write_state_file_blocking(&path, &state)
|
||||||
|
});
|
||||||
|
wait_for_blocking_io(task).await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Durably commits a prospective reset before publishing it to live counters.
|
||||||
|
pub(crate) async fn reset_user(
|
||||||
|
&self,
|
||||||
|
configured_users: &BTreeSet<String>,
|
||||||
|
user: &str,
|
||||||
|
) -> std::io::Result<UserQuotaSnapshot> {
|
||||||
|
let guard = Arc::clone(&self.mutation).lock_owned().await;
|
||||||
|
let last_reset_epoch_secs = now_epoch_secs();
|
||||||
|
let prospective = UserQuotaSnapshot {
|
||||||
|
used_bytes: 0,
|
||||||
|
last_reset_epoch_secs,
|
||||||
|
};
|
||||||
|
let state = self.state_for_users(configured_users, Some((user, prospective.clone())));
|
||||||
|
let path = self.path.clone();
|
||||||
|
let store = Arc::clone(&self.store);
|
||||||
|
let user = user.to_string();
|
||||||
|
let task = tokio::task::spawn_blocking(move || {
|
||||||
|
let _guard = guard;
|
||||||
|
write_state_file_blocking(&path, &state)?;
|
||||||
|
Ok(store.reset(&user, last_reset_epoch_secs))
|
||||||
|
});
|
||||||
|
wait_for_blocking_io(task).await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Removes a deleted user's persisted and in-memory quota ownership.
|
||||||
|
pub(crate) async fn remove_user(
|
||||||
|
&self,
|
||||||
|
configured_users: &BTreeSet<String>,
|
||||||
|
user: &str,
|
||||||
|
) -> std::io::Result<()> {
|
||||||
|
let guard = Arc::clone(&self.mutation).lock_owned().await;
|
||||||
|
let state = self.state_for_users(configured_users, None);
|
||||||
|
let path = self.path.clone();
|
||||||
|
let store = Arc::clone(&self.store);
|
||||||
|
let user = user.to_string();
|
||||||
|
let task = tokio::task::spawn_blocking(move || {
|
||||||
|
let _guard = guard;
|
||||||
|
let persisted = write_state_file_blocking(&path, &state);
|
||||||
|
store.remove(&user);
|
||||||
|
persisted
|
||||||
|
});
|
||||||
|
wait_for_blocking_io(task).await
|
||||||
|
}
|
||||||
|
|
||||||
|
fn state_for_users(
|
||||||
|
&self,
|
||||||
|
configured_users: &BTreeSet<String>,
|
||||||
|
override_user: Option<(&str, UserQuotaSnapshot)>,
|
||||||
|
) -> QuotaStateFile {
|
||||||
|
let snapshot = self.store.snapshot();
|
||||||
|
let mut users = BTreeMap::new();
|
||||||
|
let mut last_reset_epoch_secs = 0;
|
||||||
|
for user in configured_users.iter().take(QUOTA_STATE_MAX_USERS) {
|
||||||
|
let quota = override_user
|
||||||
|
.as_ref()
|
||||||
|
.filter(|(override_name, _)| *override_name == user.as_str())
|
||||||
|
.map(|(_, quota)| quota.clone())
|
||||||
|
.or_else(|| snapshot.get(user).cloned());
|
||||||
|
let Some(quota) = quota else {
|
||||||
|
continue;
|
||||||
|
};
|
||||||
|
last_reset_epoch_secs = last_reset_epoch_secs.max(quota.last_reset_epoch_secs);
|
||||||
|
users.insert(user.clone(), quota_user_state(quota));
|
||||||
|
}
|
||||||
|
QuotaStateFile {
|
||||||
|
last_reset_epoch_secs,
|
||||||
|
users,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
fn now_epoch_secs() -> u64 {
|
fn now_epoch_secs() -> u64 {
|
||||||
SystemTime::now()
|
SystemTime::now()
|
||||||
.duration_since(UNIX_EPOCH)
|
.duration_since(UNIX_EPOCH)
|
||||||
@@ -27,83 +175,104 @@ fn now_epoch_secs() -> u64 {
|
|||||||
.as_secs()
|
.as_secs()
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) async fn load_quota_state(path: &Path, stats: &Stats) {
|
async fn read_state_file(path: &Path) -> std::io::Result<Option<QuotaStateFile>> {
|
||||||
let bytes = match tokio::fs::read(path).await {
|
let file = match tokio::fs::File::open(path).await {
|
||||||
Ok(bytes) => bytes,
|
Ok(file) => file,
|
||||||
Err(error) if error.kind() == std::io::ErrorKind::NotFound => return,
|
Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None),
|
||||||
Err(error) => {
|
Err(error) => return Err(error),
|
||||||
warn!(
|
};
|
||||||
error = %error,
|
if file.metadata().await?.len() > QUOTA_STATE_MAX_BYTES {
|
||||||
path = %path.display(),
|
return Err(std::io::Error::new(
|
||||||
"Failed to read quota state file"
|
std::io::ErrorKind::InvalidData,
|
||||||
);
|
"quota state file exceeds the 16 MiB limit",
|
||||||
return;
|
));
|
||||||
|
}
|
||||||
|
let mut payload = Vec::new();
|
||||||
|
file.take(QUOTA_STATE_MAX_BYTES.saturating_add(1))
|
||||||
|
.read_to_end(&mut payload)
|
||||||
|
.await?;
|
||||||
|
if payload.len() as u64 > QUOTA_STATE_MAX_BYTES {
|
||||||
|
return Err(std::io::Error::new(
|
||||||
|
std::io::ErrorKind::InvalidData,
|
||||||
|
"quota state file grew beyond the 16 MiB limit while reading",
|
||||||
|
));
|
||||||
|
}
|
||||||
|
let state = serde_json::from_slice(&payload).map_err(|error| {
|
||||||
|
std::io::Error::new(
|
||||||
|
std::io::ErrorKind::InvalidData,
|
||||||
|
format!("failed to parse quota state: {error}"),
|
||||||
|
)
|
||||||
|
})?;
|
||||||
|
Ok(Some(state))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn write_state_file(path: &Path, state: QuotaStateFile) -> std::io::Result<()> {
|
||||||
|
let path = path.to_path_buf();
|
||||||
|
let task = tokio::task::spawn_blocking(move || write_state_file_blocking(&path, &state));
|
||||||
|
wait_for_blocking_io(task).await
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn wait_for_blocking_io<T>(
|
||||||
|
task: tokio::task::JoinHandle<std::io::Result<T>>,
|
||||||
|
) -> std::io::Result<T> {
|
||||||
|
task.await
|
||||||
|
.map_err(|error| std::io::Error::other(format!("quota checkpoint join failed: {error}")))?
|
||||||
|
}
|
||||||
|
|
||||||
|
fn write_state_file_blocking(path: &Path, state: &QuotaStateFile) -> std::io::Result<()> {
|
||||||
|
let parent = path
|
||||||
|
.parent()
|
||||||
|
.filter(|parent| !parent.as_os_str().is_empty())
|
||||||
|
.unwrap_or_else(|| Path::new("."));
|
||||||
|
std::fs::create_dir_all(parent)?;
|
||||||
|
let mut payload = serde_json::to_vec_pretty(state)?;
|
||||||
|
payload.push(b'\n');
|
||||||
|
if payload.len() as u64 > QUOTA_STATE_MAX_BYTES {
|
||||||
|
return Err(std::io::Error::new(
|
||||||
|
std::io::ErrorKind::InvalidData,
|
||||||
|
"quota state payload exceeds the 16 MiB limit",
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut last_collision = None;
|
||||||
|
for _ in 0..8 {
|
||||||
|
let tmp_path = path.with_extension(format!(
|
||||||
|
"tmp.{}.{}",
|
||||||
|
std::process::id(),
|
||||||
|
rand::random::<u64>()
|
||||||
|
));
|
||||||
|
let mut file = match std::fs::OpenOptions::new()
|
||||||
|
.write(true)
|
||||||
|
.create_new(true)
|
||||||
|
.open(&tmp_path)
|
||||||
|
{
|
||||||
|
Ok(file) => file,
|
||||||
|
Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => {
|
||||||
|
last_collision = Some(error);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
Err(error) => return Err(error),
|
||||||
|
};
|
||||||
|
let result = (|| {
|
||||||
|
file.write_all(&payload)?;
|
||||||
|
file.sync_all()?;
|
||||||
|
drop(file);
|
||||||
|
std::fs::rename(&tmp_path, path)?;
|
||||||
|
#[cfg(unix)]
|
||||||
|
std::fs::File::open(parent)?.sync_all()?;
|
||||||
|
Ok(())
|
||||||
|
})();
|
||||||
|
if result.is_err() {
|
||||||
|
let _ = std::fs::remove_file(&tmp_path);
|
||||||
}
|
}
|
||||||
};
|
return result;
|
||||||
|
|
||||||
let state = match serde_json::from_slice::<QuotaStateFile>(&bytes) {
|
|
||||||
Ok(state) => state,
|
|
||||||
Err(error) => {
|
|
||||||
warn!(
|
|
||||||
error = %error,
|
|
||||||
path = %path.display(),
|
|
||||||
"Failed to parse quota state file"
|
|
||||||
);
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
let loaded_users = state.users.len();
|
|
||||||
for (user, quota) in state.users {
|
|
||||||
stats.load_user_quota_state(&user, quota.used_bytes, quota.last_reset_epoch_secs);
|
|
||||||
}
|
}
|
||||||
info!(
|
Err(last_collision.unwrap_or_else(|| {
|
||||||
path = %path.display(),
|
std::io::Error::new(
|
||||||
loaded_users,
|
std::io::ErrorKind::AlreadyExists,
|
||||||
"Loaded per-user quota state"
|
"failed to allocate a unique quota checkpoint temporary file",
|
||||||
);
|
)
|
||||||
}
|
}))
|
||||||
|
|
||||||
pub(crate) async fn save_quota_state(path: &Path, stats: &Stats) -> std::io::Result<()> {
|
|
||||||
let mut users = BTreeMap::new();
|
|
||||||
let mut last_reset_epoch_secs = 0;
|
|
||||||
for (user, quota) in stats.user_quota_snapshot() {
|
|
||||||
last_reset_epoch_secs = last_reset_epoch_secs.max(quota.last_reset_epoch_secs);
|
|
||||||
users.insert(user, quota_user_state(quota));
|
|
||||||
}
|
|
||||||
|
|
||||||
let state = QuotaStateFile {
|
|
||||||
last_reset_epoch_secs,
|
|
||||||
users,
|
|
||||||
};
|
|
||||||
write_state_file(path, &state).await
|
|
||||||
}
|
|
||||||
|
|
||||||
pub(crate) async fn reset_user_quota(
|
|
||||||
path: &Path,
|
|
||||||
stats: &Stats,
|
|
||||||
user: &str,
|
|
||||||
) -> std::io::Result<UserQuotaSnapshot> {
|
|
||||||
let snapshot = stats.reset_user_quota(user);
|
|
||||||
save_quota_state(path, stats).await?;
|
|
||||||
Ok(snapshot)
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn write_state_file(path: &Path, state: &QuotaStateFile) -> std::io::Result<()> {
|
|
||||||
if let Some(parent) = path.parent()
|
|
||||||
&& !parent.as_os_str().is_empty()
|
|
||||||
{
|
|
||||||
tokio::fs::create_dir_all(parent).await?;
|
|
||||||
}
|
|
||||||
|
|
||||||
let tmp_path = path.with_extension(format!("tmp.{}", now_epoch_secs()));
|
|
||||||
let payload = serde_json::to_vec_pretty(state)?;
|
|
||||||
let mut file = tokio::fs::File::create(&tmp_path).await?;
|
|
||||||
file.write_all(&payload).await?;
|
|
||||||
file.write_all(b"\n").await?;
|
|
||||||
file.sync_all().await?;
|
|
||||||
drop(file);
|
|
||||||
tokio::fs::rename(&tmp_path, path).await
|
|
||||||
}
|
}
|
||||||
|
|
||||||
fn quota_user_state(quota: UserQuotaSnapshot) -> QuotaUserState {
|
fn quota_user_state(quota: UserQuotaSnapshot) -> QuotaUserState {
|
||||||
@@ -112,3 +281,131 @@ fn quota_user_state(quota: UserQuotaSnapshot) -> QuotaUserState {
|
|||||||
last_reset_epoch_secs: quota.last_reset_epoch_secs,
|
last_reset_epoch_secs: quota.last_reset_epoch_secs,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
fn users(names: &[&str]) -> BTreeSet<String> {
|
||||||
|
names.iter().map(|name| (*name).to_string()).collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn failed_durable_reset_does_not_publish_live_reset() {
|
||||||
|
let directory = tempfile::tempdir().unwrap();
|
||||||
|
let store = Arc::new(QuotaStore::default());
|
||||||
|
store.user("alice").charge(512);
|
||||||
|
let owner = QuotaStateOwner::new(directory.path().to_path_buf(), store.clone());
|
||||||
|
|
||||||
|
assert!(owner.reset_user(&users(&["alice"]), "alice").await.is_err());
|
||||||
|
assert_eq!(store.used("alice"), 512);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn concurrent_resets_serialize_without_losing_committed_users() {
|
||||||
|
let directory = tempfile::tempdir().unwrap();
|
||||||
|
let path = directory.path().join("quota.json");
|
||||||
|
let store = Arc::new(QuotaStore::default());
|
||||||
|
store.user("alice").charge(512);
|
||||||
|
store.user("bob").charge(1024);
|
||||||
|
let owner = QuotaStateOwner::new(path.clone(), store.clone());
|
||||||
|
let configured = Arc::new(users(&["alice", "bob"]));
|
||||||
|
|
||||||
|
let alice_owner = owner.clone();
|
||||||
|
let alice_users = configured.clone();
|
||||||
|
let alice = tokio::spawn(async move {
|
||||||
|
alice_owner.reset_user(&alice_users, "alice").await
|
||||||
|
});
|
||||||
|
let bob_owner = owner.clone();
|
||||||
|
let bob_users = configured.clone();
|
||||||
|
let bob = tokio::spawn(async move { bob_owner.reset_user(&bob_users, "bob").await });
|
||||||
|
alice.await.unwrap().unwrap();
|
||||||
|
bob.await.unwrap().unwrap();
|
||||||
|
|
||||||
|
let state = read_state_file(&path).await.unwrap().unwrap();
|
||||||
|
assert_eq!(state.users["alice"].used_bytes, 0);
|
||||||
|
assert_eq!(state.users["bob"].used_bytes, 0);
|
||||||
|
assert_eq!(store.used("alice"), 0);
|
||||||
|
assert_eq!(store.used("bob"), 0);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||||
|
async fn cancelled_waiter_does_not_release_blocking_mutation_ownership() {
|
||||||
|
let mutation = Arc::new(Mutex::new(()));
|
||||||
|
let guard = Arc::clone(&mutation).lock_owned().await;
|
||||||
|
let (started_tx, started_rx) = tokio::sync::oneshot::channel();
|
||||||
|
let (release_tx, release_rx) = std::sync::mpsc::channel();
|
||||||
|
let operation = tokio::task::spawn_blocking(move || {
|
||||||
|
let _guard = guard;
|
||||||
|
let _ = started_tx.send(());
|
||||||
|
release_rx.recv().unwrap();
|
||||||
|
Ok::<_, std::io::Error>(())
|
||||||
|
});
|
||||||
|
let waiter = tokio::spawn(wait_for_blocking_io(operation));
|
||||||
|
started_rx.await.unwrap();
|
||||||
|
|
||||||
|
waiter.abort();
|
||||||
|
assert!(waiter.await.unwrap_err().is_cancelled());
|
||||||
|
assert!(mutation.try_lock().is_err());
|
||||||
|
|
||||||
|
release_tx.send(()).unwrap();
|
||||||
|
tokio::time::timeout(std::time::Duration::from_secs(1), async {
|
||||||
|
loop {
|
||||||
|
if mutation.try_lock().is_ok() {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
tokio::task::yield_now().await;
|
||||||
|
}
|
||||||
|
})
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn load_rejects_oversized_state_before_parsing() {
|
||||||
|
let directory = tempfile::tempdir().unwrap();
|
||||||
|
let path = directory.path().join("quota.json");
|
||||||
|
let file = std::fs::File::create(&path).unwrap();
|
||||||
|
file.set_len(QUOTA_STATE_MAX_BYTES + 1).unwrap();
|
||||||
|
let store = Arc::new(QuotaStore::default());
|
||||||
|
let owner = QuotaStateOwner::new(path, store.clone());
|
||||||
|
|
||||||
|
owner.load(&users(&["alice"])).await;
|
||||||
|
|
||||||
|
assert_eq!(store.used("alice"), 0);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn load_filters_entries_without_configured_ownership() {
|
||||||
|
let directory = tempfile::tempdir().unwrap();
|
||||||
|
let path = directory.path().join("quota.json");
|
||||||
|
let state = QuotaStateFile {
|
||||||
|
last_reset_epoch_secs: 1,
|
||||||
|
users: BTreeMap::from([
|
||||||
|
(
|
||||||
|
"alice".to_string(),
|
||||||
|
QuotaUserState {
|
||||||
|
used_bytes: 512,
|
||||||
|
last_reset_epoch_secs: 1,
|
||||||
|
},
|
||||||
|
),
|
||||||
|
(
|
||||||
|
"orphan".to_string(),
|
||||||
|
QuotaUserState {
|
||||||
|
used_bytes: 1024,
|
||||||
|
last_reset_epoch_secs: 1,
|
||||||
|
},
|
||||||
|
),
|
||||||
|
]),
|
||||||
|
};
|
||||||
|
write_state_file(&path, state).await.unwrap();
|
||||||
|
let store = Arc::new(QuotaStore::default());
|
||||||
|
let owner = QuotaStateOwner::new(path, store.clone());
|
||||||
|
|
||||||
|
owner.load(&users(&["alice"])).await;
|
||||||
|
|
||||||
|
assert_eq!(store.used("alice"), 512);
|
||||||
|
assert_eq!(store.used("orphan"), 0);
|
||||||
|
assert!(!store.snapshot().contains_key("orphan"));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
+79
-36
@@ -1,17 +1,21 @@
|
|||||||
//! Per-IP forensic buckets for scanner and handshake failure observation.
|
//! Per-IP forensic buckets for scanner and handshake failure observation.
|
||||||
|
|
||||||
|
use std::collections::hash_map::RandomState;
|
||||||
use std::collections::{BTreeMap, HashMap};
|
use std::collections::{BTreeMap, HashMap};
|
||||||
|
use std::hash::{BuildHasher, Hash, Hasher};
|
||||||
use std::net::IpAddr;
|
use std::net::IpAddr;
|
||||||
use std::time::{Duration, Instant};
|
use std::time::{Duration, Instant};
|
||||||
|
|
||||||
use parking_lot::Mutex;
|
use parking_lot::Mutex;
|
||||||
|
|
||||||
const CLEANUP_INTERVAL: Duration = Duration::from_secs(30);
|
const CLEANUP_INTERVAL: Duration = Duration::from_secs(30);
|
||||||
const MAX_BEOBACHTEN_ENTRIES: usize = 65_536;
|
const BEOBACHTEN_SHARDS: usize = 64;
|
||||||
|
const BEOBACHTEN_ENTRIES_PER_SHARD: usize = 1024;
|
||||||
|
|
||||||
#[derive(Default)]
|
#[derive(Default)]
|
||||||
struct BeobachtenInner {
|
struct BeobachtenShard {
|
||||||
entries: HashMap<(String, IpAddr), BeobachtenEntry>,
|
classes: HashMap<String, HashMap<IpAddr, BeobachtenEntry>>,
|
||||||
|
entries: usize,
|
||||||
last_cleanup: Option<Instant>,
|
last_cleanup: Option<Instant>,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -23,7 +27,8 @@ struct BeobachtenEntry {
|
|||||||
|
|
||||||
/// In-memory, TTL-scoped per-IP counters keyed by source class.
|
/// In-memory, TTL-scoped per-IP counters keyed by source class.
|
||||||
pub struct BeobachtenStore {
|
pub struct BeobachtenStore {
|
||||||
inner: Mutex<BeobachtenInner>,
|
shards: Vec<Mutex<BeobachtenShard>>,
|
||||||
|
hash_builder: RandomState,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Default for BeobachtenStore {
|
impl Default for BeobachtenStore {
|
||||||
@@ -35,7 +40,10 @@ impl Default for BeobachtenStore {
|
|||||||
impl BeobachtenStore {
|
impl BeobachtenStore {
|
||||||
pub fn new() -> Self {
|
pub fn new() -> Self {
|
||||||
Self {
|
Self {
|
||||||
inner: Mutex::new(BeobachtenInner::default()),
|
shards: (0..BEOBACHTEN_SHARDS)
|
||||||
|
.map(|_| Mutex::new(BeobachtenShard::default()))
|
||||||
|
.collect(),
|
||||||
|
hash_builder: RandomState::new(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -45,27 +53,31 @@ impl BeobachtenStore {
|
|||||||
}
|
}
|
||||||
|
|
||||||
let now = Instant::now();
|
let now = Instant::now();
|
||||||
let mut guard = self.inner.lock();
|
let shard_index = self.shard_index(class, ip);
|
||||||
Self::cleanup_if_needed(&mut guard, now, ttl);
|
let mut shard = self.shards[shard_index].lock();
|
||||||
|
Self::cleanup_if_needed(&mut shard, now, ttl);
|
||||||
|
|
||||||
let key = (class.to_string(), ip);
|
if let Some(entry) = shard
|
||||||
if let Some(entry) = guard.entries.get_mut(&key) {
|
.classes
|
||||||
|
.get_mut(class)
|
||||||
|
.and_then(|entries| entries.get_mut(&ip))
|
||||||
|
{
|
||||||
entry.tries = entry.tries.saturating_add(1);
|
entry.tries = entry.tries.saturating_add(1);
|
||||||
entry.last_seen = now;
|
entry.last_seen = now;
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
|
||||||
if guard.entries.len() >= MAX_BEOBACHTEN_ENTRIES {
|
if shard.entries >= BEOBACHTEN_ENTRIES_PER_SHARD {
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
shard.classes.entry(class.to_string()).or_default().insert(
|
||||||
guard.entries.insert(
|
ip,
|
||||||
key,
|
|
||||||
BeobachtenEntry {
|
BeobachtenEntry {
|
||||||
tries: 1,
|
tries: 1,
|
||||||
last_seen: now,
|
last_seen: now,
|
||||||
},
|
},
|
||||||
);
|
);
|
||||||
|
shard.entries = shard.entries.saturating_add(1);
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn snapshot_text(&self, ttl: Duration) -> String {
|
pub fn snapshot_text(&self, ttl: Duration) -> String {
|
||||||
@@ -74,21 +86,15 @@ impl BeobachtenStore {
|
|||||||
}
|
}
|
||||||
|
|
||||||
let now = Instant::now();
|
let now = Instant::now();
|
||||||
let entries = {
|
|
||||||
let mut guard = self.inner.lock();
|
|
||||||
Self::cleanup(&mut guard, now, ttl);
|
|
||||||
guard.last_cleanup = Some(now);
|
|
||||||
|
|
||||||
guard
|
|
||||||
.entries
|
|
||||||
.iter()
|
|
||||||
.map(|((class, ip), entry)| (class.clone(), *ip, entry.tries))
|
|
||||||
.collect::<Vec<_>>()
|
|
||||||
};
|
|
||||||
|
|
||||||
let mut grouped = BTreeMap::<String, Vec<(IpAddr, u64)>>::new();
|
let mut grouped = BTreeMap::<String, Vec<(IpAddr, u64)>>::new();
|
||||||
for (class, ip, tries) in entries {
|
for shard in &self.shards {
|
||||||
grouped.entry(class).or_default().push((ip, tries));
|
let mut shard = shard.lock();
|
||||||
|
Self::cleanup(&mut shard, now, ttl);
|
||||||
|
shard.last_cleanup = Some(now);
|
||||||
|
for (class, entries) in &shard.classes {
|
||||||
|
let output = grouped.entry(class.clone()).or_default();
|
||||||
|
output.extend(entries.iter().map(|(ip, entry)| (*ip, entry.tries)));
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
if grouped.is_empty() {
|
if grouped.is_empty() {
|
||||||
@@ -111,24 +117,61 @@ impl BeobachtenStore {
|
|||||||
out.push_str(&format!("{ip}-{tries}\n"));
|
out.push_str(&format!("{ip}-{tries}\n"));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
out
|
out
|
||||||
}
|
}
|
||||||
|
|
||||||
fn cleanup_if_needed(inner: &mut BeobachtenInner, now: Instant, ttl: Duration) {
|
fn shard_index(&self, class: &str, ip: IpAddr) -> usize {
|
||||||
let should_cleanup = match inner.last_cleanup {
|
let mut hasher = self.hash_builder.build_hasher();
|
||||||
|
class.hash(&mut hasher);
|
||||||
|
ip.hash(&mut hasher);
|
||||||
|
(hasher.finish() as usize) % BEOBACHTEN_SHARDS
|
||||||
|
}
|
||||||
|
|
||||||
|
fn cleanup_if_needed(shard: &mut BeobachtenShard, now: Instant, ttl: Duration) {
|
||||||
|
let should_cleanup = match shard.last_cleanup {
|
||||||
Some(last) => now.saturating_duration_since(last) >= CLEANUP_INTERVAL,
|
Some(last) => now.saturating_duration_since(last) >= CLEANUP_INTERVAL,
|
||||||
None => true,
|
None => true,
|
||||||
};
|
};
|
||||||
if should_cleanup {
|
if should_cleanup {
|
||||||
Self::cleanup(inner, now, ttl);
|
Self::cleanup(shard, now, ttl);
|
||||||
inner.last_cleanup = Some(now);
|
shard.last_cleanup = Some(now);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn cleanup(inner: &mut BeobachtenInner, now: Instant, ttl: Duration) {
|
fn cleanup(shard: &mut BeobachtenShard, now: Instant, ttl: Duration) {
|
||||||
inner
|
for entries in shard.classes.values_mut() {
|
||||||
.entries
|
entries.retain(|_, entry| now.saturating_duration_since(entry.last_seen) <= ttl);
|
||||||
.retain(|_, entry| now.saturating_duration_since(entry.last_seen) <= ttl);
|
}
|
||||||
|
shard.classes.retain(|_, entries| !entries.is_empty());
|
||||||
|
shard.entries = shard.classes.values().map(HashMap::len).sum();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn one_shard_never_exceeds_its_allocation_bound() {
|
||||||
|
let store = BeobachtenStore::new();
|
||||||
|
let ttl = Duration::from_secs(60);
|
||||||
|
let target_shard = 0;
|
||||||
|
let mut inserted = 0usize;
|
||||||
|
for suffix in 0..u32::MAX {
|
||||||
|
let ip = IpAddr::V4(std::net::Ipv4Addr::from(suffix));
|
||||||
|
if store.shard_index("scanner", ip) != target_shard {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
store.record("scanner", ip, ttl);
|
||||||
|
inserted = inserted.saturating_add(1);
|
||||||
|
if inserted > BEOBACHTEN_ENTRIES_PER_SHARD + 64 {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
store.shards[target_shard].lock().entries,
|
||||||
|
BEOBACHTEN_ENTRIES_PER_SHARD
|
||||||
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -77,11 +77,45 @@ impl Stats {
|
|||||||
if !self.telemetry_me_allows_normal() {
|
if !self.telemetry_me_allows_normal() {
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
let entry = self
|
if let Some(entry) = self.me_handshake_error_codes.get(&code) {
|
||||||
.me_handshake_error_codes
|
entry.fetch_add(1, Ordering::Relaxed);
|
||||||
.entry(code)
|
return;
|
||||||
.or_insert_with(|| AtomicU64::new(0));
|
}
|
||||||
entry.fetch_add(1, Ordering::Relaxed);
|
|
||||||
|
let mut slots = self
|
||||||
|
.me_handshake_error_code_slots
|
||||||
|
.load(Ordering::Acquire);
|
||||||
|
loop {
|
||||||
|
if slots >= ME_HANDSHAKE_ERROR_CODE_MAX {
|
||||||
|
if let Some(entry) = self.me_handshake_error_codes.get(&code) {
|
||||||
|
entry.fetch_add(1, Ordering::Relaxed);
|
||||||
|
} else {
|
||||||
|
self.me_handshake_error_code_overflow_total
|
||||||
|
.fetch_add(1, Ordering::Relaxed);
|
||||||
|
}
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
match self.me_handshake_error_code_slots.compare_exchange_weak(
|
||||||
|
slots,
|
||||||
|
slots + 1,
|
||||||
|
Ordering::AcqRel,
|
||||||
|
Ordering::Acquire,
|
||||||
|
) {
|
||||||
|
Ok(_) => break,
|
||||||
|
Err(observed) => slots = observed,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
match self.me_handshake_error_codes.entry(code) {
|
||||||
|
dashmap::mapref::entry::Entry::Occupied(entry) => {
|
||||||
|
self.me_handshake_error_code_slots
|
||||||
|
.fetch_sub(1, Ordering::AcqRel);
|
||||||
|
entry.get().fetch_add(1, Ordering::Relaxed);
|
||||||
|
}
|
||||||
|
dashmap::mapref::entry::Entry::Vacant(entry) => {
|
||||||
|
entry.insert(AtomicU64::new(1));
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
pub fn increment_me_reader_eof_total(&self) {
|
pub fn increment_me_reader_eof_total(&self) {
|
||||||
if self.telemetry_me_allows_normal() {
|
if self.telemetry_me_allows_normal() {
|
||||||
@@ -440,3 +474,44 @@ impl Stats {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn handshake_error_code_cardinality_is_bounded_under_concurrency() {
|
||||||
|
const WORKERS: usize = 8;
|
||||||
|
const CODES_PER_WORKER: usize = 128;
|
||||||
|
|
||||||
|
let stats = Arc::new(Stats::new());
|
||||||
|
std::thread::scope(|scope| {
|
||||||
|
for worker in 0..WORKERS {
|
||||||
|
let stats = stats.clone();
|
||||||
|
scope.spawn(move || {
|
||||||
|
for code in 0..CODES_PER_WORKER {
|
||||||
|
stats.increment_me_handshake_error_code(
|
||||||
|
(worker * CODES_PER_WORKER + code) as i32,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
let counts = stats.get_me_handshake_error_code_counts();
|
||||||
|
assert!(counts.len() <= ME_HANDSHAKE_ERROR_CODE_MAX);
|
||||||
|
assert_eq!(
|
||||||
|
counts.iter().map(|(_, total)| *total).sum::<u64>()
|
||||||
|
+ stats.get_me_handshake_error_code_overflow_total(),
|
||||||
|
(WORKERS * CODES_PER_WORKER) as u64
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
stats
|
||||||
|
.me_handshake_error_code_slots
|
||||||
|
.load(Ordering::Acquire),
|
||||||
|
counts.len()
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -10,6 +10,10 @@ impl Stats {
|
|||||||
out.sort_by_key(|(code, _)| *code);
|
out.sort_by_key(|(code, _)| *code);
|
||||||
out
|
out
|
||||||
}
|
}
|
||||||
|
pub fn get_me_handshake_error_code_overflow_total(&self) -> u64 {
|
||||||
|
self.me_handshake_error_code_overflow_total
|
||||||
|
.load(Ordering::Relaxed)
|
||||||
|
}
|
||||||
pub fn get_me_route_drop_no_conn(&self) -> u64 {
|
pub fn get_me_route_drop_no_conn(&self) -> u64 {
|
||||||
self.me_route_drop_no_conn.load(Ordering::Relaxed)
|
self.me_route_drop_no_conn.load(Ordering::Relaxed)
|
||||||
}
|
}
|
||||||
|
|||||||
+5
-1
@@ -18,7 +18,7 @@ mod writer_counters;
|
|||||||
use dashmap::DashMap;
|
use dashmap::DashMap;
|
||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU64, Ordering};
|
use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU64, AtomicUsize, Ordering};
|
||||||
use std::time::Instant;
|
use std::time::Instant;
|
||||||
|
|
||||||
pub(crate) use self::quota_store::QuotaStore;
|
pub(crate) use self::quota_store::QuotaStore;
|
||||||
@@ -28,6 +28,8 @@ use self::telemetry::TelemetryPolicy;
|
|||||||
pub use self::tls_fingerprints::TlsFingerprintSnapshotRow;
|
pub use self::tls_fingerprints::TlsFingerprintSnapshotRow;
|
||||||
use crate::config::MeWriterPickMode;
|
use crate::config::MeWriterPickMode;
|
||||||
|
|
||||||
|
const ME_HANDSHAKE_ERROR_CODE_MAX: usize = 64;
|
||||||
|
|
||||||
#[derive(Clone, Copy)]
|
#[derive(Clone, Copy)]
|
||||||
enum RouteConnectionGauge {
|
enum RouteConnectionGauge {
|
||||||
Direct,
|
Direct,
|
||||||
@@ -219,6 +221,8 @@ pub struct Stats {
|
|||||||
me_floor_swap_idle_total: AtomicU64,
|
me_floor_swap_idle_total: AtomicU64,
|
||||||
me_floor_swap_idle_failed_total: AtomicU64,
|
me_floor_swap_idle_failed_total: AtomicU64,
|
||||||
me_handshake_error_codes: DashMap<i32, AtomicU64>,
|
me_handshake_error_codes: DashMap<i32, AtomicU64>,
|
||||||
|
me_handshake_error_code_slots: AtomicUsize,
|
||||||
|
me_handshake_error_code_overflow_total: AtomicU64,
|
||||||
me_route_drop_no_conn: AtomicU64,
|
me_route_drop_no_conn: AtomicU64,
|
||||||
me_route_drop_channel_closed: AtomicU64,
|
me_route_drop_channel_closed: AtomicU64,
|
||||||
me_route_drop_queue_full: AtomicU64,
|
me_route_drop_queue_full: AtomicU64,
|
||||||
|
|||||||
@@ -56,6 +56,10 @@ impl QuotaStore {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub(crate) fn remove(&self, user: &str) {
|
||||||
|
self.users.remove(user);
|
||||||
|
}
|
||||||
|
|
||||||
pub(crate) fn snapshot(&self) -> HashMap<String, UserQuotaSnapshot> {
|
pub(crate) fn snapshot(&self) -> HashMap<String, UserQuotaSnapshot> {
|
||||||
let mut out = HashMap::new();
|
let mut out = HashMap::new();
|
||||||
for entry in self.users.iter() {
|
for entry in self.users.iter() {
|
||||||
|
|||||||
+486
-183
@@ -1,13 +1,14 @@
|
|||||||
use std::collections::HashMap;
|
use std::collections::{HashMap, HashSet};
|
||||||
use std::collections::hash_map::DefaultHasher;
|
use std::collections::hash_map::RandomState;
|
||||||
use std::hash::{Hash, Hasher};
|
use std::hash::{BuildHasher, Hash, Hasher};
|
||||||
use std::net::IpAddr;
|
use std::net::IpAddr;
|
||||||
use std::path::{Path, PathBuf};
|
use std::path::{Path, PathBuf};
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use std::sync::atomic::{AtomicU64, Ordering};
|
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
|
||||||
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
|
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
|
||||||
|
|
||||||
use tokio::sync::RwLock;
|
use tokio::sync::RwLock;
|
||||||
|
use tokio::io::AsyncReadExt;
|
||||||
use tokio::time::sleep;
|
use tokio::time::sleep;
|
||||||
use tracing::{debug, info, warn};
|
use tracing::{debug, info, warn};
|
||||||
|
|
||||||
@@ -16,21 +17,160 @@ use crate::tls_front::types::{
|
|||||||
TlsProfileSource,
|
TlsProfileSource,
|
||||||
};
|
};
|
||||||
|
|
||||||
const FULL_CERT_SENT_SWEEP_INTERVAL_SECS: u64 = 30;
|
const FULL_CERT_SENT_SWEEP_INTERVAL_SECS: u64 = 1;
|
||||||
const FULL_CERT_SENT_MAX_IPS: usize = 65_536;
|
const FULL_CERT_SENT_MAX_ENTRIES: usize = 65_536;
|
||||||
const FULL_CERT_SENT_SHARDS: usize = 64;
|
const FULL_CERT_SENT_SHARDS: usize = 64;
|
||||||
|
const FULL_CERT_SENT_MAX_ENTRIES_PER_SHARD: usize =
|
||||||
|
(FULL_CERT_SENT_MAX_ENTRIES / FULL_CERT_SENT_SHARDS) * 2;
|
||||||
|
const TLS_FRONT_DISK_ENTRY_MAX_BYTES: u64 = 1024 * 1024;
|
||||||
|
|
||||||
static FULL_CERT_SENT_IPS_GAUGE: AtomicU64 = AtomicU64::new(0);
|
#[derive(Debug, Clone, Hash, PartialEq, Eq)]
|
||||||
static FULL_CERT_SENT_CAP_DROPS: AtomicU64 = AtomicU64::new(0);
|
struct FullCertBudgetKey {
|
||||||
|
domain: Arc<str>,
|
||||||
/// Current number of IPs tracked by the TLS full-cert budget gate.
|
client_ip: IpAddr,
|
||||||
pub(crate) fn full_cert_sent_ips_for_metrics() -> u64 {
|
|
||||||
FULL_CERT_SENT_IPS_GAUGE.load(Ordering::Relaxed)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Number of new IPs denied a full-cert budget slot because the cap was reached.
|
#[derive(Debug)]
|
||||||
pub(crate) fn full_cert_sent_cap_drops_for_metrics() -> u64 {
|
struct FullCertBudgetEntry {
|
||||||
FULL_CERT_SENT_CAP_DROPS.load(Ordering::Relaxed)
|
expires_at: Option<Instant>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Process-owned, bounded TLS full-certificate admission history.
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub(crate) struct TlsFullCertBudget {
|
||||||
|
shards: Vec<RwLock<HashMap<FullCertBudgetKey, FullCertBudgetEntry>>>,
|
||||||
|
hash_builder: RandomState,
|
||||||
|
entries: AtomicUsize,
|
||||||
|
cap_drops: AtomicU64,
|
||||||
|
last_sweep_epoch_secs: AtomicU64,
|
||||||
|
sweep_cursor: AtomicUsize,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl TlsFullCertBudget {
|
||||||
|
/// Creates an empty process-owned full-certificate budget.
|
||||||
|
pub(crate) fn new() -> Self {
|
||||||
|
Self {
|
||||||
|
shards: (0..FULL_CERT_SENT_SHARDS)
|
||||||
|
.map(|_| RwLock::new(HashMap::new()))
|
||||||
|
.collect(),
|
||||||
|
hash_builder: RandomState::new(),
|
||||||
|
entries: AtomicUsize::new(0),
|
||||||
|
cap_drops: AtomicU64::new(0),
|
||||||
|
last_sweep_epoch_secs: AtomicU64::new(0),
|
||||||
|
sweep_cursor: AtomicUsize::new(0),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn shard_index(&self, key: &FullCertBudgetKey) -> usize {
|
||||||
|
let mut hasher = self.hash_builder.build_hasher();
|
||||||
|
key.hash(&mut hasher);
|
||||||
|
(hasher.finish() as usize) % FULL_CERT_SENT_SHARDS
|
||||||
|
}
|
||||||
|
|
||||||
|
fn decrement_entries(&self, amount: usize) {
|
||||||
|
if amount == 0 {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let _ = self
|
||||||
|
.entries
|
||||||
|
.fetch_update(Ordering::AcqRel, Ordering::Relaxed, |current| {
|
||||||
|
Some(current.saturating_sub(amount))
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
fn try_reserve_entry(&self) -> bool {
|
||||||
|
let mut current = self.entries.load(Ordering::Relaxed);
|
||||||
|
loop {
|
||||||
|
if current >= FULL_CERT_SENT_MAX_ENTRIES {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
match self.entries.compare_exchange_weak(
|
||||||
|
current,
|
||||||
|
current.saturating_add(1),
|
||||||
|
Ordering::AcqRel,
|
||||||
|
Ordering::Relaxed,
|
||||||
|
) {
|
||||||
|
Ok(_) => return true,
|
||||||
|
Err(actual) => current = actual,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn sweep_one_shard(&self, now: Instant) {
|
||||||
|
let shard_index = self.sweep_cursor.fetch_add(1, Ordering::Relaxed)
|
||||||
|
% FULL_CERT_SENT_SHARDS;
|
||||||
|
let mut guard = self.shards[shard_index].write().await;
|
||||||
|
let before = guard.len();
|
||||||
|
guard.retain(|_, entry| entry.expires_at.map_or(true, |expires_at| expires_at > now));
|
||||||
|
self.decrement_entries(before.saturating_sub(guard.len()));
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn maybe_sweep(&self, now: Instant) {
|
||||||
|
let now_epoch_secs = SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.unwrap_or_default()
|
||||||
|
.as_secs();
|
||||||
|
let should_sweep = self
|
||||||
|
.last_sweep_epoch_secs
|
||||||
|
.fetch_update(Ordering::AcqRel, Ordering::Relaxed, |last_sweep| {
|
||||||
|
if now_epoch_secs.saturating_sub(last_sweep)
|
||||||
|
>= FULL_CERT_SENT_SWEEP_INTERVAL_SECS
|
||||||
|
{
|
||||||
|
Some(now_epoch_secs)
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
}
|
||||||
|
})
|
||||||
|
.is_ok();
|
||||||
|
if should_sweep {
|
||||||
|
self.sweep_one_shard(now).await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn take(&self, domain: Arc<str>, client_ip: IpAddr, ttl: Duration) -> bool {
|
||||||
|
if ttl.is_zero() {
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
|
||||||
|
let now = Instant::now();
|
||||||
|
self.maybe_sweep(now).await;
|
||||||
|
let expires_at = now.checked_add(ttl);
|
||||||
|
let key = FullCertBudgetKey { domain, client_ip };
|
||||||
|
let shard_index = self.shard_index(&key);
|
||||||
|
let mut guard = self.shards[shard_index].write().await;
|
||||||
|
|
||||||
|
if let Some(entry) = guard.get_mut(&key) {
|
||||||
|
if entry.expires_at.is_some_and(|expires_at| expires_at <= now) {
|
||||||
|
entry.expires_at = expires_at;
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
|
||||||
|
if guard.len() >= FULL_CERT_SENT_MAX_ENTRIES_PER_SHARD {
|
||||||
|
let before = guard.len();
|
||||||
|
guard.retain(|_, entry| {
|
||||||
|
entry.expires_at.map_or(true, |expires_at| expires_at > now)
|
||||||
|
});
|
||||||
|
self.decrement_entries(before.saturating_sub(guard.len()));
|
||||||
|
}
|
||||||
|
if guard.len() >= FULL_CERT_SENT_MAX_ENTRIES_PER_SHARD || !self.try_reserve_entry() {
|
||||||
|
self.cap_drops.fetch_add(1, Ordering::Relaxed);
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
guard.insert(key, FullCertBudgetEntry { expires_at });
|
||||||
|
true
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Returns the current number of retained domain and client IP entries.
|
||||||
|
pub(crate) fn entries_for_metrics(&self) -> u64 {
|
||||||
|
self.entries.load(Ordering::Relaxed) as u64
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Returns the cumulative number of entries rejected by hard bounds.
|
||||||
|
pub(crate) fn cap_drops_for_metrics(&self) -> u64 {
|
||||||
|
self.cap_drops.load(Ordering::Relaxed)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Lightweight in-memory + optional on-disk cache for TLS fronting data.
|
/// Lightweight in-memory + optional on-disk cache for TLS fronting data.
|
||||||
@@ -38,8 +178,9 @@ pub(crate) fn full_cert_sent_cap_drops_for_metrics() -> u64 {
|
|||||||
pub struct TlsFrontCache {
|
pub struct TlsFrontCache {
|
||||||
memory: RwLock<HashMap<String, Arc<CachedTlsData>>>,
|
memory: RwLock<HashMap<String, Arc<CachedTlsData>>>,
|
||||||
default: Arc<CachedTlsData>,
|
default: Arc<CachedTlsData>,
|
||||||
full_cert_sent_shards: Vec<RwLock<HashMap<IpAddr, Instant>>>,
|
full_cert_budget: Arc<TlsFullCertBudget>,
|
||||||
full_cert_sent_last_sweep_epoch_secs: AtomicU64,
|
full_cert_domain_keys: HashMap<String, Arc<str>>,
|
||||||
|
disk_entry_names: HashSet<String>,
|
||||||
disk_path: PathBuf,
|
disk_path: PathBuf,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -91,6 +232,21 @@ fn key_share_group_label(group: Option<u16>) -> &'static str {
|
|||||||
#[allow(dead_code)]
|
#[allow(dead_code)]
|
||||||
impl TlsFrontCache {
|
impl TlsFrontCache {
|
||||||
pub fn new(domains: &[String], default_len: usize, disk_path: impl AsRef<Path>) -> Self {
|
pub fn new(domains: &[String], default_len: usize, disk_path: impl AsRef<Path>) -> Self {
|
||||||
|
Self::new_with_full_cert_budget(
|
||||||
|
domains,
|
||||||
|
default_len,
|
||||||
|
disk_path,
|
||||||
|
Arc::new(TlsFullCertBudget::new()),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Creates a generation-local cache backed by the process-owned full-cert budget.
|
||||||
|
pub(crate) fn new_with_full_cert_budget(
|
||||||
|
domains: &[String],
|
||||||
|
default_len: usize,
|
||||||
|
disk_path: impl AsRef<Path>,
|
||||||
|
full_cert_budget: Arc<TlsFullCertBudget>,
|
||||||
|
) -> Self {
|
||||||
let default_template = ParsedServerHello {
|
let default_template = ParsedServerHello {
|
||||||
version: [0x03, 0x03],
|
version: [0x03, 0x03],
|
||||||
random: [0u8; 32],
|
random: [0u8; 32],
|
||||||
@@ -112,17 +268,24 @@ impl TlsFrontCache {
|
|||||||
});
|
});
|
||||||
|
|
||||||
let mut map = HashMap::new();
|
let mut map = HashMap::new();
|
||||||
|
let mut full_cert_domain_keys = HashMap::new();
|
||||||
|
let mut disk_entry_names = HashSet::new();
|
||||||
for d in domains {
|
for d in domains {
|
||||||
map.insert(d.clone(), default.clone());
|
map.insert(d.clone(), default.clone());
|
||||||
|
disk_entry_names.insert(format!("{}.json", d.replace(['/', '\\'], "_")));
|
||||||
|
let canonical: Arc<str> = Arc::from(normalize_dns_name(d));
|
||||||
|
full_cert_domain_keys.insert(d.clone(), canonical.clone());
|
||||||
|
full_cert_domain_keys
|
||||||
|
.entry(canonical.to_string())
|
||||||
|
.or_insert(canonical);
|
||||||
}
|
}
|
||||||
|
|
||||||
Self {
|
Self {
|
||||||
memory: RwLock::new(map),
|
memory: RwLock::new(map),
|
||||||
default,
|
default,
|
||||||
full_cert_sent_shards: (0..FULL_CERT_SENT_SHARDS)
|
full_cert_budget,
|
||||||
.map(|_| RwLock::new(HashMap::new()))
|
full_cert_domain_keys,
|
||||||
.collect(),
|
disk_entry_names,
|
||||||
full_cert_sent_last_sweep_epoch_secs: AtomicU64::new(0),
|
|
||||||
disk_path: disk_path.as_ref().to_path_buf(),
|
disk_path: disk_path.as_ref().to_path_buf(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -202,114 +365,66 @@ impl TlsFrontCache {
|
|||||||
.collect()
|
.collect()
|
||||||
}
|
}
|
||||||
|
|
||||||
fn full_cert_sent_shard_index(client_ip: IpAddr) -> usize {
|
fn full_cert_domain_key(&self, domain: &str) -> Arc<str> {
|
||||||
let mut hasher = DefaultHasher::new();
|
self.full_cert_domain_keys
|
||||||
client_ip.hash(&mut hasher);
|
.get(domain)
|
||||||
(hasher.finish() as usize) % FULL_CERT_SENT_SHARDS
|
.cloned()
|
||||||
|
.unwrap_or_else(|| Arc::from(normalize_dns_name(domain)))
|
||||||
}
|
}
|
||||||
|
|
||||||
fn full_cert_sent_shard(&self, client_ip: IpAddr) -> &RwLock<HashMap<IpAddr, Instant>> {
|
/// Returns true when the selected domain and client IP may receive a full cert payload.
|
||||||
&self.full_cert_sent_shards[Self::full_cert_sent_shard_index(client_ip)]
|
pub async fn take_full_cert_budget_for_ip(
|
||||||
|
&self,
|
||||||
|
domain: &str,
|
||||||
|
client_ip: IpAddr,
|
||||||
|
ttl: Duration,
|
||||||
|
) -> bool {
|
||||||
|
self.full_cert_budget
|
||||||
|
.take(self.full_cert_domain_key(domain), client_ip, ttl)
|
||||||
|
.await
|
||||||
}
|
}
|
||||||
|
|
||||||
fn decrement_full_cert_sent_entries(amount: usize) {
|
/// Returns the current process-owned full-cert budget entry count.
|
||||||
if amount == 0 {
|
pub(crate) fn full_cert_budget_entries_for_metrics(&self) -> u64 {
|
||||||
return;
|
self.full_cert_budget.entries_for_metrics()
|
||||||
}
|
|
||||||
let amount = amount as u64;
|
|
||||||
let _ =
|
|
||||||
FULL_CERT_SENT_IPS_GAUGE.fetch_update(Ordering::AcqRel, Ordering::Relaxed, |current| {
|
|
||||||
Some(current.saturating_sub(amount))
|
|
||||||
});
|
|
||||||
}
|
}
|
||||||
|
|
||||||
fn try_reserve_full_cert_sent_entry() -> bool {
|
/// Returns the cumulative process-owned full-cert budget cap drops.
|
||||||
let mut current = FULL_CERT_SENT_IPS_GAUGE.load(Ordering::Relaxed);
|
pub(crate) fn full_cert_budget_cap_drops_for_metrics(&self) -> u64 {
|
||||||
loop {
|
self.full_cert_budget.cap_drops_for_metrics()
|
||||||
if current >= FULL_CERT_SENT_MAX_IPS as u64 {
|
|
||||||
return false;
|
|
||||||
}
|
|
||||||
match FULL_CERT_SENT_IPS_GAUGE.compare_exchange_weak(
|
|
||||||
current,
|
|
||||||
current.saturating_add(1),
|
|
||||||
Ordering::AcqRel,
|
|
||||||
Ordering::Relaxed,
|
|
||||||
) {
|
|
||||||
Ok(_) => return true,
|
|
||||||
Err(actual) => current = actual,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn sweep_full_cert_sent_shards(&self, now: Instant, ttl: Duration) {
|
|
||||||
for shard in &self.full_cert_sent_shards {
|
|
||||||
let mut guard = shard.write().await;
|
|
||||||
let before = guard.len();
|
|
||||||
guard.retain(|_, seen_at| now.duration_since(*seen_at) < ttl);
|
|
||||||
Self::decrement_full_cert_sent_entries(before.saturating_sub(guard.len()));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Returns true when full cert payload should be sent for client_ip
|
|
||||||
/// according to TTL policy.
|
|
||||||
pub async fn take_full_cert_budget_for_ip(&self, client_ip: IpAddr, ttl: Duration) -> bool {
|
|
||||||
if ttl.is_zero() {
|
|
||||||
return true;
|
|
||||||
}
|
|
||||||
|
|
||||||
let now = Instant::now();
|
|
||||||
let now_epoch_secs = SystemTime::now()
|
|
||||||
.duration_since(UNIX_EPOCH)
|
|
||||||
.unwrap_or_default()
|
|
||||||
.as_secs();
|
|
||||||
let should_sweep = self
|
|
||||||
.full_cert_sent_last_sweep_epoch_secs
|
|
||||||
.fetch_update(Ordering::AcqRel, Ordering::Relaxed, |last_sweep| {
|
|
||||||
if now_epoch_secs.saturating_sub(last_sweep) >= FULL_CERT_SENT_SWEEP_INTERVAL_SECS {
|
|
||||||
Some(now_epoch_secs)
|
|
||||||
} else {
|
|
||||||
None
|
|
||||||
}
|
|
||||||
})
|
|
||||||
.is_ok();
|
|
||||||
|
|
||||||
if should_sweep {
|
|
||||||
self.sweep_full_cert_sent_shards(now, ttl).await;
|
|
||||||
}
|
|
||||||
|
|
||||||
let mut guard = self.full_cert_sent_shard(client_ip).write().await;
|
|
||||||
let allowed = match guard.get_mut(&client_ip) {
|
|
||||||
Some(seen_at) => {
|
|
||||||
if now.duration_since(*seen_at) >= ttl {
|
|
||||||
*seen_at = now;
|
|
||||||
true
|
|
||||||
} else {
|
|
||||||
false
|
|
||||||
}
|
|
||||||
}
|
|
||||||
None => {
|
|
||||||
if !Self::try_reserve_full_cert_sent_entry() {
|
|
||||||
FULL_CERT_SENT_CAP_DROPS.fetch_add(1, Ordering::Relaxed);
|
|
||||||
return false;
|
|
||||||
}
|
|
||||||
guard.insert(client_ip, now);
|
|
||||||
true
|
|
||||||
}
|
|
||||||
};
|
|
||||||
allowed
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
async fn insert_full_cert_sent_for_tests(&self, client_ip: IpAddr, seen_at: Instant) {
|
async fn insert_full_cert_sent_for_tests(
|
||||||
let mut guard = self.full_cert_sent_shard(client_ip).write().await;
|
&self,
|
||||||
if guard.insert(client_ip, seen_at).is_none() {
|
domain: &str,
|
||||||
FULL_CERT_SENT_IPS_GAUGE.fetch_add(1, Ordering::Relaxed);
|
client_ip: IpAddr,
|
||||||
|
expires_at: Instant,
|
||||||
|
) {
|
||||||
|
let key = FullCertBudgetKey {
|
||||||
|
domain: self.full_cert_domain_key(domain),
|
||||||
|
client_ip,
|
||||||
|
};
|
||||||
|
let shard_index = self.full_cert_budget.shard_index(&key);
|
||||||
|
let mut guard = self.full_cert_budget.shards[shard_index].write().await;
|
||||||
|
if guard
|
||||||
|
.insert(
|
||||||
|
key,
|
||||||
|
FullCertBudgetEntry {
|
||||||
|
expires_at: Some(expires_at),
|
||||||
|
},
|
||||||
|
)
|
||||||
|
.is_none()
|
||||||
|
{
|
||||||
|
self.full_cert_budget
|
||||||
|
.entries
|
||||||
|
.fetch_add(1, Ordering::Relaxed);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
async fn full_cert_sent_is_empty_for_tests(&self) -> bool {
|
async fn full_cert_sent_is_empty_for_tests(&self) -> bool {
|
||||||
for shard in &self.full_cert_sent_shards {
|
for shard in &self.full_cert_budget.shards {
|
||||||
if !shard.read().await.is_empty() {
|
if !shard.read().await.is_empty() {
|
||||||
return false;
|
return false;
|
||||||
}
|
}
|
||||||
@@ -318,11 +433,16 @@ impl TlsFrontCache {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
async fn full_cert_sent_contains_for_tests(&self, client_ip: IpAddr) -> bool {
|
async fn full_cert_sent_contains_for_tests(&self, domain: &str, client_ip: IpAddr) -> bool {
|
||||||
self.full_cert_sent_shard(client_ip)
|
let key = FullCertBudgetKey {
|
||||||
|
domain: self.full_cert_domain_key(domain),
|
||||||
|
client_ip,
|
||||||
|
};
|
||||||
|
let shard_index = self.full_cert_budget.shard_index(&key);
|
||||||
|
self.full_cert_budget.shards[shard_index]
|
||||||
.read()
|
.read()
|
||||||
.await
|
.await
|
||||||
.contains_key(&client_ip)
|
.contains_key(&key)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn set(&self, domain: &str, data: CachedTlsData) {
|
pub async fn set(&self, domain: &str, data: CachedTlsData) {
|
||||||
@@ -336,54 +456,60 @@ impl TlsFrontCache {
|
|||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
let mut loaded = 0usize;
|
let mut loaded = 0usize;
|
||||||
if let Ok(mut dir) = tokio::fs::read_dir(&path).await {
|
for name in &self.disk_entry_names {
|
||||||
while let Ok(Some(entry)) = dir.next_entry().await {
|
let entry_path = path.join(name);
|
||||||
if let Ok(name) = entry.file_name().into_string() {
|
let Ok(metadata) = tokio::fs::symlink_metadata(&entry_path).await else {
|
||||||
if !name.ends_with(".json") {
|
continue;
|
||||||
continue;
|
};
|
||||||
}
|
if !metadata.file_type().is_file() {
|
||||||
if let Ok(data) = tokio::fs::read(entry.path()).await
|
continue;
|
||||||
&& let Ok(mut cached) = serde_json::from_slice::<CachedTlsData>(&data)
|
}
|
||||||
{
|
if let Ok(data) = read_disk_entry_bounded(&entry_path).await
|
||||||
if cached.domain.is_empty()
|
&& let Ok(mut cached) = serde_json::from_slice::<CachedTlsData>(&data)
|
||||||
|| cached.domain.len() > 255
|
{
|
||||||
|| !cached
|
if cached.domain.is_empty()
|
||||||
.domain
|
|| cached.domain.len() > 255
|
||||||
.chars()
|
|| !cached
|
||||||
.all(|c| c.is_ascii_alphanumeric() || c == '.' || c == '-')
|
.domain
|
||||||
{
|
.chars()
|
||||||
warn!(file = %name, "Skipping TLS cache entry with invalid domain");
|
.all(|c| c.is_ascii_alphanumeric() || c == '.' || c == '-')
|
||||||
continue;
|
{
|
||||||
}
|
warn!(file = %name, "Skipping TLS cache entry with invalid domain");
|
||||||
if !cert_info_matches_domain(&cached) {
|
continue;
|
||||||
warn!(
|
|
||||||
file = %name,
|
|
||||||
domain = %cached.domain,
|
|
||||||
"Skipping TLS cache entry with mismatched certificate metadata"
|
|
||||||
);
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
// fetched_at is skipped during deserialization; approximate with file mtime if available.
|
|
||||||
if let Ok(meta) = entry.metadata().await
|
|
||||||
&& let Ok(modified) = meta.modified()
|
|
||||||
{
|
|
||||||
cached.fetched_at = modified;
|
|
||||||
}
|
|
||||||
// Drop entries older than 72h
|
|
||||||
if let Ok(age) = cached.fetched_at.elapsed()
|
|
||||||
&& age > Duration::from_secs(72 * 3600)
|
|
||||||
{
|
|
||||||
warn!(domain = %cached.domain, "Skipping stale TLS cache entry (>72h)");
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
cached
|
|
||||||
.behavior_profile
|
|
||||||
.refresh_server_hello_summary(&cached.server_hello_template);
|
|
||||||
let domain = cached.domain.clone();
|
|
||||||
self.set(&domain, cached).await;
|
|
||||||
loaded += 1;
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
if !self.full_cert_domain_keys.contains_key(&cached.domain) {
|
||||||
|
warn!(
|
||||||
|
file = %name,
|
||||||
|
domain = %cached.domain,
|
||||||
|
"Skipping TLS cache entry outside configured domains"
|
||||||
|
);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
if !cert_info_matches_domain(&cached) {
|
||||||
|
warn!(
|
||||||
|
file = %name,
|
||||||
|
domain = %cached.domain,
|
||||||
|
"Skipping TLS cache entry with mismatched certificate metadata"
|
||||||
|
);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
// fetched_at is skipped during deserialization; approximate with file mtime if available.
|
||||||
|
if let Ok(modified) = metadata.modified() {
|
||||||
|
cached.fetched_at = modified;
|
||||||
|
}
|
||||||
|
// Drop entries older than 72h
|
||||||
|
if let Ok(age) = cached.fetched_at.elapsed()
|
||||||
|
&& age > Duration::from_secs(72 * 3600)
|
||||||
|
{
|
||||||
|
warn!(domain = %cached.domain, "Skipping stale TLS cache entry (>72h)");
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
cached
|
||||||
|
.behavior_profile
|
||||||
|
.refresh_server_hello_summary(&cached.server_hello_template);
|
||||||
|
let domain = cached.domain.clone();
|
||||||
|
self.set(&domain, cached).await;
|
||||||
|
loaded += 1;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
if loaded > 0 {
|
if loaded > 0 {
|
||||||
@@ -398,6 +524,14 @@ impl TlsFrontCache {
|
|||||||
let fname = format!("{}.json", domain.replace(['/', '\\'], "_"));
|
let fname = format!("{}.json", domain.replace(['/', '\\'], "_"));
|
||||||
let path = self.disk_path.join(fname);
|
let path = self.disk_path.join(fname);
|
||||||
if let Ok(json) = serde_json::to_vec_pretty(data) {
|
if let Ok(json) = serde_json::to_vec_pretty(data) {
|
||||||
|
if json.len() as u64 > TLS_FRONT_DISK_ENTRY_MAX_BYTES {
|
||||||
|
warn!(
|
||||||
|
domain,
|
||||||
|
bytes = json.len(),
|
||||||
|
"Skipping oversized TLS cache persistence"
|
||||||
|
);
|
||||||
|
return;
|
||||||
|
}
|
||||||
// best-effort write
|
// best-effort write
|
||||||
let _ = tokio::fs::write(path, json).await;
|
let _ = tokio::fs::write(path, json).await;
|
||||||
}
|
}
|
||||||
@@ -464,6 +598,27 @@ impl TlsFrontCache {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn read_disk_entry_bounded(path: &Path) -> std::io::Result<Vec<u8>> {
|
||||||
|
let file = tokio::fs::File::open(path).await?;
|
||||||
|
if file.metadata().await?.len() > TLS_FRONT_DISK_ENTRY_MAX_BYTES {
|
||||||
|
return Err(std::io::Error::new(
|
||||||
|
std::io::ErrorKind::InvalidData,
|
||||||
|
"TLS cache entry exceeds the 1 MiB limit",
|
||||||
|
));
|
||||||
|
}
|
||||||
|
let mut bytes = Vec::new();
|
||||||
|
file.take(TLS_FRONT_DISK_ENTRY_MAX_BYTES.saturating_add(1))
|
||||||
|
.read_to_end(&mut bytes)
|
||||||
|
.await?;
|
||||||
|
if bytes.len() as u64 > TLS_FRONT_DISK_ENTRY_MAX_BYTES {
|
||||||
|
return Err(std::io::Error::new(
|
||||||
|
std::io::ErrorKind::InvalidData,
|
||||||
|
"TLS cache entry grew beyond the 1 MiB limit while reading",
|
||||||
|
));
|
||||||
|
}
|
||||||
|
Ok(bytes)
|
||||||
|
}
|
||||||
|
|
||||||
fn cert_info_matches_domain(cached: &CachedTlsData) -> bool {
|
fn cert_info_matches_domain(cached: &CachedTlsData) -> bool {
|
||||||
let Some(cert_info) = cached.cert_info.as_ref() else {
|
let Some(cert_info) = cached.cert_info.as_ref() else {
|
||||||
return true;
|
return true;
|
||||||
@@ -581,12 +736,24 @@ mod tests {
|
|||||||
let ip: IpAddr = "127.0.0.1".parse().expect("ip");
|
let ip: IpAddr = "127.0.0.1".parse().expect("ip");
|
||||||
let ttl = Duration::from_millis(80);
|
let ttl = Duration::from_millis(80);
|
||||||
|
|
||||||
assert!(cache.take_full_cert_budget_for_ip(ip, ttl).await);
|
assert!(
|
||||||
assert!(!cache.take_full_cert_budget_for_ip(ip, ttl).await);
|
cache
|
||||||
|
.take_full_cert_budget_for_ip("example.com", ip, ttl)
|
||||||
|
.await
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
!cache
|
||||||
|
.take_full_cert_budget_for_ip("example.com", ip, ttl)
|
||||||
|
.await
|
||||||
|
);
|
||||||
|
|
||||||
tokio::time::sleep(Duration::from_millis(90)).await;
|
tokio::time::sleep(Duration::from_millis(90)).await;
|
||||||
|
|
||||||
assert!(cache.take_full_cert_budget_for_ip(ip, ttl).await);
|
assert!(
|
||||||
|
cache
|
||||||
|
.take_full_cert_budget_for_ip("example.com", ip, ttl)
|
||||||
|
.await
|
||||||
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
@@ -601,7 +768,11 @@ mod tests {
|
|||||||
((idx >> 8) & 0xff) as u8,
|
((idx >> 8) & 0xff) as u8,
|
||||||
(idx & 0xff) as u8,
|
(idx & 0xff) as u8,
|
||||||
));
|
));
|
||||||
assert!(cache.take_full_cert_budget_for_ip(ip, ttl).await);
|
assert!(
|
||||||
|
cache
|
||||||
|
.take_full_cert_budget_for_ip("example.com", ip, ttl)
|
||||||
|
.await
|
||||||
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
assert!(cache.full_cert_sent_is_empty_for_tests().await);
|
assert!(cache.full_cert_sent_is_empty_for_tests().await);
|
||||||
@@ -613,21 +784,42 @@ mod tests {
|
|||||||
let stale_ip: IpAddr = "127.0.0.1".parse().expect("ip");
|
let stale_ip: IpAddr = "127.0.0.1".parse().expect("ip");
|
||||||
let new_ip: IpAddr = "127.0.0.2".parse().expect("ip");
|
let new_ip: IpAddr = "127.0.0.2".parse().expect("ip");
|
||||||
let ttl = Duration::from_secs(1);
|
let ttl = Duration::from_secs(1);
|
||||||
let stale_seen_at = Instant::now()
|
let stale_expires_at = Instant::now()
|
||||||
.checked_sub(Duration::from_secs(10))
|
.checked_sub(Duration::from_secs(1))
|
||||||
.unwrap_or_else(Instant::now);
|
.unwrap_or_else(Instant::now);
|
||||||
|
|
||||||
cache
|
cache
|
||||||
.insert_full_cert_sent_for_tests(stale_ip, stale_seen_at)
|
.insert_full_cert_sent_for_tests("example.com", stale_ip, stale_expires_at)
|
||||||
.await;
|
.await;
|
||||||
|
let stale_key = FullCertBudgetKey {
|
||||||
|
domain: cache.full_cert_domain_key("example.com"),
|
||||||
|
client_ip: stale_ip,
|
||||||
|
};
|
||||||
|
cache.full_cert_budget.sweep_cursor.store(
|
||||||
|
cache.full_cert_budget.shard_index(&stale_key),
|
||||||
|
Ordering::Relaxed,
|
||||||
|
);
|
||||||
cache
|
cache
|
||||||
.full_cert_sent_last_sweep_epoch_secs
|
.full_cert_budget
|
||||||
|
.last_sweep_epoch_secs
|
||||||
.store(0, Ordering::Relaxed);
|
.store(0, Ordering::Relaxed);
|
||||||
|
|
||||||
assert!(cache.take_full_cert_budget_for_ip(new_ip, ttl).await);
|
assert!(
|
||||||
|
cache
|
||||||
|
.take_full_cert_budget_for_ip("example.com", new_ip, ttl)
|
||||||
|
.await
|
||||||
|
);
|
||||||
|
|
||||||
assert!(!cache.full_cert_sent_contains_for_tests(stale_ip).await);
|
assert!(
|
||||||
assert!(cache.full_cert_sent_contains_for_tests(new_ip).await);
|
!cache
|
||||||
|
.full_cert_sent_contains_for_tests("example.com", stale_ip)
|
||||||
|
.await
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
cache
|
||||||
|
.full_cert_sent_contains_for_tests("example.com", new_ip)
|
||||||
|
.await
|
||||||
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
@@ -636,8 +828,8 @@ mod tests {
|
|||||||
let stale_ip: IpAddr = "127.0.0.1".parse().expect("ip");
|
let stale_ip: IpAddr = "127.0.0.1".parse().expect("ip");
|
||||||
let new_ip: IpAddr = "127.0.0.2".parse().expect("ip");
|
let new_ip: IpAddr = "127.0.0.2".parse().expect("ip");
|
||||||
let ttl = Duration::from_secs(1);
|
let ttl = Duration::from_secs(1);
|
||||||
let stale_seen_at = Instant::now()
|
let stale_expires_at = Instant::now()
|
||||||
.checked_sub(Duration::from_secs(10))
|
.checked_sub(Duration::from_secs(1))
|
||||||
.unwrap_or_else(Instant::now);
|
.unwrap_or_else(Instant::now);
|
||||||
let now_epoch_secs = SystemTime::now()
|
let now_epoch_secs = SystemTime::now()
|
||||||
.duration_since(UNIX_EPOCH)
|
.duration_since(UNIX_EPOCH)
|
||||||
@@ -645,15 +837,126 @@ mod tests {
|
|||||||
.as_secs();
|
.as_secs();
|
||||||
|
|
||||||
cache
|
cache
|
||||||
.insert_full_cert_sent_for_tests(stale_ip, stale_seen_at)
|
.insert_full_cert_sent_for_tests("example.com", stale_ip, stale_expires_at)
|
||||||
.await;
|
.await;
|
||||||
cache
|
cache
|
||||||
.full_cert_sent_last_sweep_epoch_secs
|
.full_cert_budget
|
||||||
|
.last_sweep_epoch_secs
|
||||||
.store(now_epoch_secs, Ordering::Relaxed);
|
.store(now_epoch_secs, Ordering::Relaxed);
|
||||||
|
|
||||||
assert!(cache.take_full_cert_budget_for_ip(new_ip, ttl).await);
|
assert!(
|
||||||
|
cache
|
||||||
|
.take_full_cert_budget_for_ip("example.com", new_ip, ttl)
|
||||||
|
.await
|
||||||
|
);
|
||||||
|
|
||||||
assert!(cache.full_cert_sent_contains_for_tests(stale_ip).await);
|
assert!(
|
||||||
assert!(cache.full_cert_sent_contains_for_tests(new_ip).await);
|
cache
|
||||||
|
.full_cert_sent_contains_for_tests("example.com", stale_ip)
|
||||||
|
.await
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
cache
|
||||||
|
.full_cert_sent_contains_for_tests("example.com", new_ip)
|
||||||
|
.await
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn full_cert_budget_is_shared_across_cache_generations_and_scoped_by_domain() {
|
||||||
|
let budget = Arc::new(TlsFullCertBudget::new());
|
||||||
|
let domains = ["one.example".to_string(), "two.example".to_string()];
|
||||||
|
let first = TlsFrontCache::new_with_full_cert_budget(
|
||||||
|
&domains,
|
||||||
|
1024,
|
||||||
|
"tlsfront-test-cache",
|
||||||
|
budget.clone(),
|
||||||
|
);
|
||||||
|
let second = TlsFrontCache::new_with_full_cert_budget(
|
||||||
|
&domains,
|
||||||
|
1024,
|
||||||
|
"tlsfront-test-cache",
|
||||||
|
budget,
|
||||||
|
);
|
||||||
|
let ip: IpAddr = "127.0.0.1".parse().expect("ip");
|
||||||
|
let ttl = Duration::from_secs(60);
|
||||||
|
|
||||||
|
assert!(
|
||||||
|
first
|
||||||
|
.take_full_cert_budget_for_ip("one.example", ip, ttl)
|
||||||
|
.await
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
!second
|
||||||
|
.take_full_cert_budget_for_ip("one.example", ip, ttl)
|
||||||
|
.await
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
second
|
||||||
|
.take_full_cert_budget_for_ip("two.example", ip, ttl)
|
||||||
|
.await
|
||||||
|
);
|
||||||
|
assert_eq!(second.full_cert_budget_entries_for_metrics(), 2);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn existing_full_cert_entry_keeps_its_own_expiry_after_ttl_change() {
|
||||||
|
let cache = TlsFrontCache::new(&["example.com".to_string()], 1024, "tlsfront-test-cache");
|
||||||
|
let ip: IpAddr = "127.0.0.1".parse().expect("ip");
|
||||||
|
|
||||||
|
assert!(
|
||||||
|
cache
|
||||||
|
.take_full_cert_budget_for_ip("example.com", ip, Duration::from_millis(80))
|
||||||
|
.await
|
||||||
|
);
|
||||||
|
tokio::time::sleep(Duration::from_millis(20)).await;
|
||||||
|
assert!(
|
||||||
|
!cache
|
||||||
|
.take_full_cert_budget_for_ip("example.com", ip, Duration::from_millis(1))
|
||||||
|
.await
|
||||||
|
);
|
||||||
|
tokio::time::sleep(Duration::from_millis(70)).await;
|
||||||
|
assert!(
|
||||||
|
cache
|
||||||
|
.take_full_cert_budget_for_ip("example.com", ip, Duration::from_millis(1))
|
||||||
|
.await
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn disk_reader_rejects_an_entry_above_the_hard_limit() {
|
||||||
|
let directory = tempfile::tempdir().unwrap();
|
||||||
|
let path = directory.path().join("oversized.json");
|
||||||
|
tokio::fs::write(
|
||||||
|
&path,
|
||||||
|
vec![0u8; TLS_FRONT_DISK_ENTRY_MAX_BYTES as usize + 1],
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
let error = read_disk_entry_bounded(&path).await.unwrap_err();
|
||||||
|
|
||||||
|
assert_eq!(error.kind(), std::io::ErrorKind::InvalidData);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(unix)]
|
||||||
|
#[tokio::test]
|
||||||
|
async fn disk_loader_does_not_follow_a_configured_name_symlink() {
|
||||||
|
let directory = tempfile::tempdir().unwrap();
|
||||||
|
let target = directory.path().join("outside.json");
|
||||||
|
let cached = cached_with_cert_info("example.com", None, Vec::new());
|
||||||
|
tokio::fs::write(&target, serde_json::to_vec(&cached).unwrap())
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
std::os::unix::fs::symlink(&target, directory.path().join("example.com.json")).unwrap();
|
||||||
|
let cache = TlsFrontCache::new(
|
||||||
|
&["example.com".to_string()],
|
||||||
|
1024,
|
||||||
|
directory.path(),
|
||||||
|
);
|
||||||
|
|
||||||
|
cache.load_from_disk().await;
|
||||||
|
|
||||||
|
assert_eq!(cache.get("example.com").await.domain, "default");
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -30,7 +30,6 @@ use x509_parser::prelude::FromDer;
|
|||||||
|
|
||||||
use crate::config::TlsFetchProfile;
|
use crate::config::TlsFetchProfile;
|
||||||
use crate::crypto::{SecureRandom, sha256};
|
use crate::crypto::{SecureRandom, sha256};
|
||||||
use crate::network::dns_overrides::resolve_socket_addr;
|
|
||||||
use crate::protocol::constants::{
|
use crate::protocol::constants::{
|
||||||
TLS_RECORD_APPLICATION, TLS_RECORD_CHANGE_CIPHER, TLS_RECORD_HANDSHAKE,
|
TLS_RECORD_APPLICATION, TLS_RECORD_CHANGE_CIPHER, TLS_RECORD_HANDSHAKE,
|
||||||
};
|
};
|
||||||
@@ -901,9 +900,6 @@ async fn connect_with_dns_override(
|
|||||||
port: u16,
|
port: u16,
|
||||||
connect_timeout: Duration,
|
connect_timeout: Duration,
|
||||||
) -> Result<TcpStream> {
|
) -> Result<TcpStream> {
|
||||||
if let Some(addr) = resolve_socket_addr(host, port) {
|
|
||||||
return Ok(timeout(connect_timeout, TcpStream::connect(addr)).await??);
|
|
||||||
}
|
|
||||||
Ok(timeout(connect_timeout, TcpStream::connect((host, port))).await??)
|
Ok(timeout(connect_timeout, TcpStream::connect((host, port))).await??)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -14,7 +14,7 @@ use crate::error::Result;
|
|||||||
use crate::transport::UpstreamManager;
|
use crate::transport::UpstreamManager;
|
||||||
|
|
||||||
use super::MePool;
|
use super::MePool;
|
||||||
use super::http_fetch::https_get;
|
use super::http_fetch::{HTTPS_RESPONSE_BODY_MAX_BYTES, https_get};
|
||||||
use super::rotation::{MeReinitTrigger, enqueue_reinit_trigger};
|
use super::rotation::{MeReinitTrigger, enqueue_reinit_trigger};
|
||||||
use super::secret::download_proxy_secret_with_max_len_via_upstream;
|
use super::secret::download_proxy_secret_with_max_len_via_upstream;
|
||||||
use super::selftest::record_timeskew_sample;
|
use super::selftest::record_timeskew_sample;
|
||||||
@@ -106,7 +106,7 @@ pub async fn fetch_proxy_config_with_raw_via_upstream(
|
|||||||
url: &str,
|
url: &str,
|
||||||
upstream: Option<Arc<UpstreamManager>>,
|
upstream: Option<Arc<UpstreamManager>>,
|
||||||
) -> Result<(ProxyConfigData, String)> {
|
) -> Result<(ProxyConfigData, String)> {
|
||||||
let resp = https_get(url, upstream).await?;
|
let resp = https_get(url, upstream, HTTPS_RESPONSE_BODY_MAX_BYTES).await?;
|
||||||
let http_status = resp.status;
|
let http_status = resp.status;
|
||||||
|
|
||||||
if let Some(date_str) = resp.date_header.as_deref()
|
if let Some(date_str) = resp.date_header.as_deref()
|
||||||
|
|||||||
@@ -71,6 +71,40 @@ struct FamilyReconnectOutcome {
|
|||||||
endpoint_count: usize,
|
endpoint_count: usize,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
struct ScheduledReconnects<'a> {
|
||||||
|
inflight: &'a mut HashMap<(i32, IpFamily), usize>,
|
||||||
|
keys: Vec<(i32, IpFamily)>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ScheduledReconnects<'_> {
|
||||||
|
fn current(&self, key: &(i32, IpFamily)) -> usize {
|
||||||
|
self.inflight.get(key).copied().unwrap_or(0)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn reserve(&mut self, key: (i32, IpFamily)) {
|
||||||
|
self.keys.push(key);
|
||||||
|
*self.inflight.entry(key).or_insert(0) += 1;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Drop for ScheduledReconnects<'_> {
|
||||||
|
fn drop(&mut self) {
|
||||||
|
for key in self.keys.drain(..) {
|
||||||
|
let std::collections::hash_map::Entry::Occupied(mut entry) =
|
||||||
|
self.inflight.entry(key)
|
||||||
|
else {
|
||||||
|
continue;
|
||||||
|
};
|
||||||
|
let remaining = entry.get().saturating_sub(1);
|
||||||
|
if remaining == 0 {
|
||||||
|
entry.remove();
|
||||||
|
} else {
|
||||||
|
*entry.get_mut() = remaining;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
pub async fn me_health_monitor(pool: Arc<MePool>, rng: Arc<SecureRandom>, _min_connections: usize) {
|
pub async fn me_health_monitor(pool: Arc<MePool>, rng: Arc<SecureRandom>, _min_connections: usize) {
|
||||||
let mut backoff: HashMap<(i32, IpFamily), u64> = HashMap::new();
|
let mut backoff: HashMap<(i32, IpFamily), u64> = HashMap::new();
|
||||||
let mut next_attempt: HashMap<(i32, IpFamily), Instant> = HashMap::new();
|
let mut next_attempt: HashMap<(i32, IpFamily), Instant> = HashMap::new();
|
||||||
@@ -437,6 +471,10 @@ async fn check_family(
|
|||||||
let writer_idle_since = Arc::new(writer_idle_since);
|
let writer_idle_since = Arc::new(writer_idle_since);
|
||||||
let bound_clients_by_writer = Arc::new(bound_clients_by_writer);
|
let bound_clients_by_writer = Arc::new(bound_clients_by_writer);
|
||||||
let mut reconnect_set = JoinSet::<FamilyReconnectOutcome>::new();
|
let mut reconnect_set = JoinSet::<FamilyReconnectOutcome>::new();
|
||||||
|
let mut scheduled_reconnects = ScheduledReconnects {
|
||||||
|
inflight,
|
||||||
|
keys: Vec::new(),
|
||||||
|
};
|
||||||
|
|
||||||
for (dc, endpoints) in dc_endpoints {
|
for (dc, endpoints) in dc_endpoints {
|
||||||
if endpoints.is_empty() {
|
if endpoints.is_empty() {
|
||||||
@@ -562,7 +600,7 @@ async fn check_family(
|
|||||||
.reconnect_runtime
|
.reconnect_runtime
|
||||||
.me_reconnect_max_concurrent_per_dc
|
.me_reconnect_max_concurrent_per_dc
|
||||||
.max(1) as usize;
|
.max(1) as usize;
|
||||||
if *inflight.get(&key).unwrap_or(&0) >= max_concurrent {
|
if scheduled_reconnects.current(&key) >= max_concurrent {
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
if pool
|
if pool
|
||||||
@@ -579,7 +617,7 @@ async fn check_family(
|
|||||||
);
|
);
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
*inflight.entry(key).or_insert(0) += 1;
|
scheduled_reconnects.reserve(key);
|
||||||
let pool_for_reconnect = pool.clone();
|
let pool_for_reconnect = pool.clone();
|
||||||
let rng_for_reconnect = rng.clone();
|
let rng_for_reconnect = rng.clone();
|
||||||
let reconnect_sem_for_dc = reconnect_sem.clone();
|
let reconnect_sem_for_dc = reconnect_sem.clone();
|
||||||
@@ -740,9 +778,6 @@ async fn check_family(
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
if let Some(v) = inflight.get_mut(&outcome.key) {
|
|
||||||
*v = v.saturating_sub(1);
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
family_degraded
|
family_degraded
|
||||||
@@ -1701,15 +1736,35 @@ mod tests {
|
|||||||
use tokio::sync::mpsc;
|
use tokio::sync::mpsc;
|
||||||
use tokio_util::sync::CancellationToken;
|
use tokio_util::sync::CancellationToken;
|
||||||
|
|
||||||
use super::reap_draining_writers;
|
use super::{ScheduledReconnects, reap_draining_writers};
|
||||||
use crate::config::{GeneralConfig, MeRouteNoWriterMode, MeSocksKdfPolicy, MeWriterPickMode};
|
use crate::config::{GeneralConfig, MeRouteNoWriterMode, MeSocksKdfPolicy, MeWriterPickMode};
|
||||||
use crate::crypto::SecureRandom;
|
use crate::crypto::SecureRandom;
|
||||||
use crate::network::probe::NetworkDecision;
|
use crate::network::probe::NetworkDecision;
|
||||||
|
use crate::network::IpFamily;
|
||||||
use crate::stats::Stats;
|
use crate::stats::Stats;
|
||||||
use crate::transport::middle_proxy::codec::WriterCommand;
|
use crate::transport::middle_proxy::codec::WriterCommand;
|
||||||
use crate::transport::middle_proxy::pool::{MePool, MeWriter, WriterContour};
|
use crate::transport::middle_proxy::pool::{MePool, MeWriter, WriterContour};
|
||||||
use crate::transport::middle_proxy::registry::ConnMeta;
|
use crate::transport::middle_proxy::registry::ConnMeta;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn reconnect_batch_releases_every_reserved_key_after_join_failures() {
|
||||||
|
let retained = (1, IpFamily::V4);
|
||||||
|
let removed = (2, IpFamily::V6);
|
||||||
|
let mut inflight = HashMap::from([(retained, 1)]);
|
||||||
|
|
||||||
|
{
|
||||||
|
let mut scheduled = ScheduledReconnects {
|
||||||
|
inflight: &mut inflight,
|
||||||
|
keys: Vec::new(),
|
||||||
|
};
|
||||||
|
scheduled.reserve(retained);
|
||||||
|
scheduled.reserve(removed);
|
||||||
|
}
|
||||||
|
|
||||||
|
assert_eq!(inflight.get(&retained), Some(&1));
|
||||||
|
assert!(!inflight.contains_key(&removed));
|
||||||
|
}
|
||||||
|
|
||||||
async fn make_pool(me_pool_drain_threshold: u64) -> Arc<MePool> {
|
async fn make_pool(me_pool_drain_threshold: u64) -> Arc<MePool> {
|
||||||
let general = GeneralConfig {
|
let general = GeneralConfig {
|
||||||
me_pool_drain_threshold,
|
me_pool_drain_threshold,
|
||||||
@@ -1812,6 +1867,7 @@ mod tests {
|
|||||||
general.me_route_blocking_send_timeout_ms,
|
general.me_route_blocking_send_timeout_ms,
|
||||||
general.me_route_inline_recovery_attempts,
|
general.me_route_inline_recovery_attempts,
|
||||||
general.me_route_inline_recovery_wait_ms,
|
general.me_route_inline_recovery_wait_ms,
|
||||||
|
16_384,
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -1,7 +1,7 @@
|
|||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
|
|
||||||
use http_body_util::{BodyExt, Empty};
|
use http_body_util::{BodyExt, Empty, Limited};
|
||||||
use hyper::header::{CONNECTION, DATE, HOST, USER_AGENT};
|
use hyper::header::{CONNECTION, DATE, HOST, USER_AGENT};
|
||||||
use hyper::{Method, Request};
|
use hyper::{Method, Request};
|
||||||
use hyper_util::rt::TokioIo;
|
use hyper_util::rt::TokioIo;
|
||||||
@@ -12,11 +12,19 @@ use tokio_rustls::TlsConnector;
|
|||||||
use tracing::debug;
|
use tracing::debug;
|
||||||
|
|
||||||
use crate::error::{ProxyError, Result};
|
use crate::error::{ProxyError, Result};
|
||||||
use crate::network::dns_overrides::resolve_socket_addr;
|
|
||||||
use crate::transport::{UpstreamManager, UpstreamStream};
|
use crate::transport::{UpstreamManager, UpstreamStream};
|
||||||
|
|
||||||
const HTTP_CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
|
const HTTP_CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
|
||||||
const HTTP_REQUEST_TIMEOUT: Duration = Duration::from_secs(15);
|
const HTTP_REQUEST_TIMEOUT: Duration = Duration::from_secs(15);
|
||||||
|
pub(super) const HTTPS_RESPONSE_BODY_MAX_BYTES: usize = 1024 * 1024;
|
||||||
|
|
||||||
|
struct HttpsConnectionDriver(tokio::task::JoinHandle<()>);
|
||||||
|
|
||||||
|
impl Drop for HttpsConnectionDriver {
|
||||||
|
fn drop(&mut self) {
|
||||||
|
self.0.abort();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
pub(crate) struct HttpsGetResponse {
|
pub(crate) struct HttpsGetResponse {
|
||||||
pub(crate) status: u16,
|
pub(crate) status: u16,
|
||||||
@@ -81,14 +89,6 @@ async fn connect_https_transport(
|
|||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
if let Some(addr) = resolve_socket_addr(host, port) {
|
|
||||||
let stream = timeout(HTTP_CONNECT_TIMEOUT, TcpStream::connect(addr))
|
|
||||||
.await
|
|
||||||
.map_err(|_| ProxyError::Proxy(format!("connect timeout for {host}:{port}")))?
|
|
||||||
.map_err(|e| ProxyError::Proxy(format!("connect failed for {host}:{port}: {e}")))?;
|
|
||||||
return Ok(UpstreamStream::Tcp(stream));
|
|
||||||
}
|
|
||||||
|
|
||||||
let stream = timeout(HTTP_CONNECT_TIMEOUT, TcpStream::connect((host, port)))
|
let stream = timeout(HTTP_CONNECT_TIMEOUT, TcpStream::connect((host, port)))
|
||||||
.await
|
.await
|
||||||
.map_err(|_| ProxyError::Proxy(format!("connect timeout for {host}:{port}")))?
|
.map_err(|_| ProxyError::Proxy(format!("connect timeout for {host}:{port}")))?
|
||||||
@@ -99,7 +99,14 @@ async fn connect_https_transport(
|
|||||||
pub(crate) async fn https_get(
|
pub(crate) async fn https_get(
|
||||||
url: &str,
|
url: &str,
|
||||||
upstream: Option<Arc<UpstreamManager>>,
|
upstream: Option<Arc<UpstreamManager>>,
|
||||||
|
max_body_bytes: usize,
|
||||||
) -> Result<HttpsGetResponse> {
|
) -> Result<HttpsGetResponse> {
|
||||||
|
if max_body_bytes == 0 {
|
||||||
|
return Err(ProxyError::Proxy(
|
||||||
|
"HTTPS response body limit must be greater than zero".to_string(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
let max_body_bytes = max_body_bytes.min(HTTPS_RESPONSE_BODY_MAX_BYTES);
|
||||||
let (host, port, path_and_query) = extract_host_port_path(url)?;
|
let (host, port, path_and_query) = extract_host_port_path(url)?;
|
||||||
let stream = connect_https_transport(&host, port, upstream).await?;
|
let stream = connect_https_transport(&host, port, upstream).await?;
|
||||||
|
|
||||||
@@ -115,11 +122,11 @@ pub(crate) async fn https_get(
|
|||||||
.await
|
.await
|
||||||
.map_err(|e| ProxyError::Proxy(format!("HTTP handshake failed for {host}:{port}: {e}")))?;
|
.map_err(|e| ProxyError::Proxy(format!("HTTP handshake failed for {host}:{port}: {e}")))?;
|
||||||
|
|
||||||
tokio::spawn(async move {
|
let _connection_driver = HttpsConnectionDriver(tokio::spawn(async move {
|
||||||
if let Err(e) = connection.await {
|
if let Err(e) = connection.await {
|
||||||
debug!(error = %e, "HTTPS fetch connection task failed");
|
debug!(error = %e, "HTTPS fetch connection task failed");
|
||||||
}
|
}
|
||||||
});
|
}));
|
||||||
|
|
||||||
let host_header = if port == 443 {
|
let host_header = if port == 443 {
|
||||||
host.clone()
|
host.clone()
|
||||||
@@ -148,10 +155,17 @@ pub(crate) async fn https_get(
|
|||||||
.and_then(|value| value.to_str().ok())
|
.and_then(|value| value.to_str().ok())
|
||||||
.map(|value| value.to_string());
|
.map(|value| value.to_string());
|
||||||
|
|
||||||
let body = timeout(HTTP_REQUEST_TIMEOUT, response.into_body().collect())
|
let body = timeout(
|
||||||
|
HTTP_REQUEST_TIMEOUT,
|
||||||
|
Limited::new(response.into_body(), max_body_bytes).collect(),
|
||||||
|
)
|
||||||
.await
|
.await
|
||||||
.map_err(|_| ProxyError::Proxy(format!("HTTP body read timeout for {url}")))?
|
.map_err(|_| ProxyError::Proxy(format!("HTTP body read timeout for {url}")))?
|
||||||
.map_err(|e| ProxyError::Proxy(format!("HTTP body read failed for {url}: {e}")))?
|
.map_err(|e| {
|
||||||
|
ProxyError::Proxy(format!(
|
||||||
|
"HTTP body read failed or exceeded {max_body_bytes} bytes for {url}: {e}"
|
||||||
|
))
|
||||||
|
})?
|
||||||
.to_bytes()
|
.to_bytes()
|
||||||
.to_vec();
|
.to_vec();
|
||||||
|
|
||||||
@@ -161,3 +175,16 @@ pub(crate) async fn https_get(
|
|||||||
body,
|
body,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn limited_body_rejects_payload_above_caller_bound() {
|
||||||
|
let body = http_body_util::Full::new(bytes::Bytes::from_static(b"12345"));
|
||||||
|
let result = Limited::new(body, 4).collect().await;
|
||||||
|
|
||||||
|
assert!(result.is_err());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -22,6 +22,7 @@ mod ping;
|
|||||||
mod pool;
|
mod pool;
|
||||||
mod pool_config;
|
mod pool_config;
|
||||||
mod pool_init;
|
mod pool_init;
|
||||||
|
mod pool_lifecycle;
|
||||||
mod pool_nat;
|
mod pool_nat;
|
||||||
mod pool_refill;
|
mod pool_refill;
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
@@ -60,6 +61,7 @@ pub use ping::{
|
|||||||
MePingFamily, MePingReport, MePingSample, format_me_route, format_sample_line, run_me_ping,
|
MePingFamily, MePingReport, MePingSample, format_me_route, format_sample_line, run_me_ping,
|
||||||
};
|
};
|
||||||
pub use pool::MePool;
|
pub use pool::MePool;
|
||||||
|
pub(crate) use registry::ConnLease;
|
||||||
#[allow(unused_imports)]
|
#[allow(unused_imports)]
|
||||||
pub use pool_nat::{detect_public_ip, stun_probe};
|
pub use pool_nat::{detect_public_ip, stun_probe};
|
||||||
pub use registry::ConnRegistry;
|
pub use registry::ConnRegistry;
|
||||||
|
|||||||
@@ -10,6 +10,7 @@ use std::sync::atomic::{
|
|||||||
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
|
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
|
||||||
|
|
||||||
use arc_swap::ArcSwap;
|
use arc_swap::ArcSwap;
|
||||||
|
use parking_lot::Mutex as ParkingMutex;
|
||||||
use tokio::sync::{Mutex, RwLock, Semaphore, mpsc, watch};
|
use tokio::sync::{Mutex, RwLock, Semaphore, mpsc, watch};
|
||||||
use tokio_util::sync::CancellationToken;
|
use tokio_util::sync::CancellationToken;
|
||||||
|
|
||||||
@@ -23,6 +24,7 @@ use crate::transport::UpstreamManager;
|
|||||||
|
|
||||||
use super::ConnRegistry;
|
use super::ConnRegistry;
|
||||||
use super::codec::WriterCommand;
|
use super::codec::WriterCommand;
|
||||||
|
use super::pool_lifecycle::MePoolLifecycle;
|
||||||
|
|
||||||
const ME_FORCE_CLOSE_SAFETY_FALLBACK_SECS: u64 = 300;
|
const ME_FORCE_CLOSE_SAFETY_FALLBACK_SECS: u64 = 300;
|
||||||
|
|
||||||
@@ -32,12 +34,6 @@ pub(super) struct RefillDcKey {
|
|||||||
pub family: IpFamily,
|
pub family: IpFamily,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
|
||||||
pub(super) struct RefillEndpointKey {
|
|
||||||
pub dc: i32,
|
|
||||||
pub addr: SocketAddr,
|
|
||||||
}
|
|
||||||
|
|
||||||
#[derive(Clone)]
|
#[derive(Clone)]
|
||||||
pub struct MeWriter {
|
pub struct MeWriter {
|
||||||
pub id: u64,
|
pub id: u64,
|
||||||
@@ -148,6 +144,18 @@ pub(super) enum WriterContour {
|
|||||||
Draining = 2,
|
Draining = 2,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub(super) struct WriterOpenReservation<'a> {
|
||||||
|
counter: Option<&'a AtomicUsize>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Drop for WriterOpenReservation<'_> {
|
||||||
|
fn drop(&mut self) {
|
||||||
|
if let Some(counter) = self.counter {
|
||||||
|
counter.fetch_sub(1, Ordering::AcqRel);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
impl WriterContour {
|
impl WriterContour {
|
||||||
pub(super) fn as_u8(self) -> u8 {
|
pub(super) fn as_u8(self) -> u8 {
|
||||||
self as u8
|
self as u8
|
||||||
@@ -265,6 +273,10 @@ pub(super) struct ReinitCore {
|
|||||||
pub(super) pending_hardswap_generation: AtomicU64,
|
pub(super) pending_hardswap_generation: AtomicU64,
|
||||||
pub(super) pending_hardswap_started_at_epoch_secs: AtomicU64,
|
pub(super) pending_hardswap_started_at_epoch_secs: AtomicU64,
|
||||||
pub(super) pending_hardswap_map_hash: AtomicU64,
|
pub(super) pending_hardswap_map_hash: AtomicU64,
|
||||||
|
pub(super) scheduler_inflight: AtomicUsize,
|
||||||
|
pub(super) max_concurrency_effective: AtomicUsize,
|
||||||
|
pub(super) coordinator: ParkingMutex<ReinitCoordinatorState>,
|
||||||
|
pub(super) status: ArcSwap<ReinitStatusSnapshot>,
|
||||||
pub(super) hardswap: AtomicBool,
|
pub(super) hardswap: AtomicBool,
|
||||||
pub(super) me_hardswap_warmup_delay_min_ms: AtomicU64,
|
pub(super) me_hardswap_warmup_delay_min_ms: AtomicU64,
|
||||||
pub(super) me_hardswap_warmup_delay_max_ms: AtomicU64,
|
pub(super) me_hardswap_warmup_delay_max_ms: AtomicU64,
|
||||||
@@ -272,6 +284,39 @@ pub(super) struct ReinitCore {
|
|||||||
pub(super) me_hardswap_warmup_pass_backoff_base_ms: AtomicU64,
|
pub(super) me_hardswap_warmup_pass_backoff_base_ms: AtomicU64,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[derive(Clone, Debug)]
|
||||||
|
pub(super) struct ReinitStatusSnapshot {
|
||||||
|
pub(super) active_generation: u64,
|
||||||
|
pub(super) warm_generations: Vec<u64>,
|
||||||
|
pub(super) pending_hardswap_generation: u64,
|
||||||
|
pub(super) pending_hardswap_started_at_epoch_secs: u64,
|
||||||
|
pub(super) pending_hardswap_map_hash: u64,
|
||||||
|
pub(super) inflight: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Clone, Copy)]
|
||||||
|
pub(super) struct ReinitPendingState {
|
||||||
|
pub(super) generation: u64,
|
||||||
|
pub(super) started_at_epoch_secs: u64,
|
||||||
|
pub(super) map_hash: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Clone, Copy)]
|
||||||
|
pub(super) struct ReinitAttemptState {
|
||||||
|
pub(super) generation: u64,
|
||||||
|
pub(super) map_hash: u64,
|
||||||
|
pub(super) hardswap: bool,
|
||||||
|
pub(super) committed: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) struct ReinitCoordinatorState {
|
||||||
|
pub(super) next_attempt_id: u64,
|
||||||
|
pub(super) active_generation: u64,
|
||||||
|
pub(super) desired_map_hash: u64,
|
||||||
|
pub(super) pending: Option<ReinitPendingState>,
|
||||||
|
pub(super) attempts: HashMap<u64, ReinitAttemptState>,
|
||||||
|
}
|
||||||
|
|
||||||
pub(super) struct WriterLifecycleCore {
|
pub(super) struct WriterLifecycleCore {
|
||||||
pub(super) me_keepalive_enabled: bool,
|
pub(super) me_keepalive_enabled: bool,
|
||||||
pub(super) me_keepalive_interval: Duration,
|
pub(super) me_keepalive_interval: Duration,
|
||||||
@@ -421,6 +466,7 @@ pub struct MePool {
|
|||||||
pub(super) floor_runtime: Arc<FloorRuntimeCore>,
|
pub(super) floor_runtime: Arc<FloorRuntimeCore>,
|
||||||
pub(super) writer_selection_policy: Arc<WriterSelectionPolicyCore>,
|
pub(super) writer_selection_policy: Arc<WriterSelectionPolicyCore>,
|
||||||
pub(super) transport_policy: Arc<TransportPolicyCore>,
|
pub(super) transport_policy: Arc<TransportPolicyCore>,
|
||||||
|
pub(super) lifecycle: MePoolLifecycle,
|
||||||
pub(super) decision: NetworkDecision,
|
pub(super) decision: NetworkDecision,
|
||||||
pub(super) upstream: Option<Arc<UpstreamManager>>,
|
pub(super) upstream: Option<Arc<UpstreamManager>>,
|
||||||
pub(super) rng: Arc<SecureRandom>,
|
pub(super) rng: Arc<SecureRandom>,
|
||||||
@@ -431,9 +477,12 @@ pub struct MePool {
|
|||||||
pub(super) endpoint_dc_map: Arc<RwLock<HashMap<SocketAddr, Option<i32>>>>,
|
pub(super) endpoint_dc_map: Arc<RwLock<HashMap<SocketAddr, Option<i32>>>>,
|
||||||
pub(super) default_dc: AtomicI32,
|
pub(super) default_dc: AtomicI32,
|
||||||
pub(super) next_writer_id: AtomicU64,
|
pub(super) next_writer_id: AtomicU64,
|
||||||
|
pub(super) writer_connect_active_reserved: AtomicUsize,
|
||||||
|
pub(super) writer_connect_warm_reserved: AtomicUsize,
|
||||||
pub(super) rtt_stats: Arc<Mutex<HashMap<u64, (f64, f64)>>>,
|
pub(super) rtt_stats: Arc<Mutex<HashMap<u64, (f64, f64)>>>,
|
||||||
pub(super) refill_inflight: Arc<Mutex<HashSet<RefillEndpointKey>>>,
|
pub(super) refill_states: Arc<ParkingMutex<HashMap<RefillDcKey, Option<SocketAddr>>>>,
|
||||||
pub(super) refill_inflight_dc: Arc<Mutex<HashSet<RefillDcKey>>>,
|
pub(super) refill_running: AtomicUsize,
|
||||||
|
pub(super) refill_pending: AtomicUsize,
|
||||||
pub(super) conn_count: AtomicUsize,
|
pub(super) conn_count: AtomicUsize,
|
||||||
pub(super) draining_active_runtime: AtomicU64,
|
pub(super) draining_active_runtime: AtomicU64,
|
||||||
pub(super) stats: Arc<crate::stats::Stats>,
|
pub(super) stats: Arc<crate::stats::Stats>,
|
||||||
@@ -573,12 +622,14 @@ impl MePool {
|
|||||||
me_route_blocking_send_timeout_ms: u64,
|
me_route_blocking_send_timeout_ms: u64,
|
||||||
me_route_inline_recovery_attempts: u32,
|
me_route_inline_recovery_attempts: u32,
|
||||||
me_route_inline_recovery_wait_ms: u64,
|
me_route_inline_recovery_wait_ms: u64,
|
||||||
|
me_connection_cleanup_capacity: usize,
|
||||||
) -> Arc<Self> {
|
) -> Arc<Self> {
|
||||||
let endpoint_dc_map = Self::build_endpoint_dc_map_from_maps(&proxy_map_v4, &proxy_map_v6);
|
let endpoint_dc_map = Self::build_endpoint_dc_map_from_maps(&proxy_map_v4, &proxy_map_v6);
|
||||||
let preferred_endpoints_by_dc =
|
let preferred_endpoints_by_dc =
|
||||||
Self::build_preferred_endpoints_by_dc(&decision, &proxy_map_v4, &proxy_map_v6);
|
Self::build_preferred_endpoints_by_dc(&decision, &proxy_map_v4, &proxy_map_v6);
|
||||||
let registry = Arc::new(ConnRegistry::with_route_channel_capacity(
|
let registry = Arc::new(ConnRegistry::with_route_and_cleanup_capacity(
|
||||||
me_route_channel_capacity,
|
me_route_channel_capacity,
|
||||||
|
me_connection_cleanup_capacity,
|
||||||
));
|
));
|
||||||
registry.update_route_backpressure_policy(
|
registry.update_route_backpressure_policy(
|
||||||
me_route_backpressure_base_timeout_ms,
|
me_route_backpressure_base_timeout_ms,
|
||||||
@@ -587,6 +638,14 @@ impl MePool {
|
|||||||
);
|
);
|
||||||
let (writer_epoch, _) = watch::channel(0u64);
|
let (writer_epoch, _) = watch::channel(0u64);
|
||||||
let now_epoch_secs = Self::now_epoch_secs();
|
let now_epoch_secs = Self::now_epoch_secs();
|
||||||
|
let reinit_status = ReinitStatusSnapshot {
|
||||||
|
active_generation: 1,
|
||||||
|
warm_generations: Vec::new(),
|
||||||
|
pending_hardswap_generation: 0,
|
||||||
|
pending_hardswap_started_at_epoch_secs: 0,
|
||||||
|
pending_hardswap_map_hash: 0,
|
||||||
|
inflight: 0,
|
||||||
|
};
|
||||||
stats.set_me_writer_byte_budget_limit_bytes(me_writer_byte_budget_bytes);
|
stats.set_me_writer_byte_budget_limit_bytes(me_writer_byte_budget_bytes);
|
||||||
Arc::new(Self {
|
Arc::new(Self {
|
||||||
routing: Arc::new(RoutingCore {
|
routing: Arc::new(RoutingCore {
|
||||||
@@ -603,6 +662,16 @@ impl MePool {
|
|||||||
pending_hardswap_generation: AtomicU64::new(0),
|
pending_hardswap_generation: AtomicU64::new(0),
|
||||||
pending_hardswap_started_at_epoch_secs: AtomicU64::new(0),
|
pending_hardswap_started_at_epoch_secs: AtomicU64::new(0),
|
||||||
pending_hardswap_map_hash: AtomicU64::new(0),
|
pending_hardswap_map_hash: AtomicU64::new(0),
|
||||||
|
scheduler_inflight: AtomicUsize::new(0),
|
||||||
|
max_concurrency_effective: AtomicUsize::new(1),
|
||||||
|
coordinator: ParkingMutex::new(ReinitCoordinatorState {
|
||||||
|
next_attempt_id: 1,
|
||||||
|
active_generation: 1,
|
||||||
|
desired_map_hash: 0,
|
||||||
|
pending: None,
|
||||||
|
attempts: HashMap::new(),
|
||||||
|
}),
|
||||||
|
status: ArcSwap::from_pointee(reinit_status),
|
||||||
hardswap: AtomicBool::new(hardswap),
|
hardswap: AtomicBool::new(hardswap),
|
||||||
me_hardswap_warmup_delay_min_ms: AtomicU64::new(me_hardswap_warmup_delay_min_ms),
|
me_hardswap_warmup_delay_min_ms: AtomicU64::new(me_hardswap_warmup_delay_min_ms),
|
||||||
me_hardswap_warmup_delay_max_ms: AtomicU64::new(me_hardswap_warmup_delay_max_ms),
|
me_hardswap_warmup_delay_max_ms: AtomicU64::new(me_hardswap_warmup_delay_max_ms),
|
||||||
@@ -805,6 +874,7 @@ impl MePool {
|
|||||||
me_reader_route_data_wait_ms,
|
me_reader_route_data_wait_ms,
|
||||||
)),
|
)),
|
||||||
}),
|
}),
|
||||||
|
lifecycle: MePoolLifecycle::new(),
|
||||||
decision,
|
decision,
|
||||||
upstream,
|
upstream,
|
||||||
rng,
|
rng,
|
||||||
@@ -830,9 +900,12 @@ impl MePool {
|
|||||||
endpoint_dc_map: Arc::new(RwLock::new(endpoint_dc_map)),
|
endpoint_dc_map: Arc::new(RwLock::new(endpoint_dc_map)),
|
||||||
default_dc: AtomicI32::new(default_dc.unwrap_or(2)),
|
default_dc: AtomicI32::new(default_dc.unwrap_or(2)),
|
||||||
next_writer_id: AtomicU64::new(1),
|
next_writer_id: AtomicU64::new(1),
|
||||||
|
writer_connect_active_reserved: AtomicUsize::new(0),
|
||||||
|
writer_connect_warm_reserved: AtomicUsize::new(0),
|
||||||
rtt_stats: Arc::new(Mutex::new(HashMap::new())),
|
rtt_stats: Arc::new(Mutex::new(HashMap::new())),
|
||||||
refill_inflight: Arc::new(Mutex::new(HashSet::new())),
|
refill_states: Arc::new(ParkingMutex::new(HashMap::new())),
|
||||||
refill_inflight_dc: Arc::new(Mutex::new(HashSet::new())),
|
refill_running: AtomicUsize::new(0),
|
||||||
|
refill_pending: AtomicUsize::new(0),
|
||||||
conn_count: AtomicUsize::new(0),
|
conn_count: AtomicUsize::new(0),
|
||||||
draining_active_runtime: AtomicU64::new(0),
|
draining_active_runtime: AtomicU64::new(0),
|
||||||
endpoint_quarantine: Arc::new(Mutex::new(HashMap::new())),
|
endpoint_quarantine: Arc::new(Mutex::new(HashMap::new())),
|
||||||
@@ -1263,6 +1336,7 @@ impl MePool {
|
|||||||
self.translate_our_addr_with_reflection(addr, None)
|
self.translate_our_addr_with_reflection(addr, None)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[allow(dead_code)]
|
||||||
pub fn registry(&self) -> &Arc<ConnRegistry> {
|
pub fn registry(&self) -> &Arc<ConnRegistry> {
|
||||||
&self.registry
|
&self.registry
|
||||||
}
|
}
|
||||||
@@ -1733,6 +1807,69 @@ impl MePool {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub(super) async fn reserve_writer_open(
|
||||||
|
&self,
|
||||||
|
contour: WriterContour,
|
||||||
|
allow_coverage_override: bool,
|
||||||
|
writer_dc: i32,
|
||||||
|
) -> Option<WriterOpenReservation<'_>> {
|
||||||
|
let counter = match contour {
|
||||||
|
WriterContour::Active => &self.writer_connect_active_reserved,
|
||||||
|
WriterContour::Warm => &self.writer_connect_warm_reserved,
|
||||||
|
WriterContour::Draining => {
|
||||||
|
return Some(WriterOpenReservation { counter: None });
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
loop {
|
||||||
|
if !self
|
||||||
|
.can_open_writer_for_contour(contour, allow_coverage_override, writer_dc)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
let (active_writers, warm_writers, _) =
|
||||||
|
self.non_draining_writer_counts_by_contour().await;
|
||||||
|
let live = match contour {
|
||||||
|
WriterContour::Active => active_writers,
|
||||||
|
WriterContour::Warm => warm_writers,
|
||||||
|
WriterContour::Draining => 0,
|
||||||
|
};
|
||||||
|
let mut limit = match contour {
|
||||||
|
WriterContour::Active => self.adaptive_floor_active_cap_configured_total(),
|
||||||
|
WriterContour::Warm => self.adaptive_floor_warm_cap_configured_total(),
|
||||||
|
WriterContour::Draining => usize::MAX,
|
||||||
|
};
|
||||||
|
if contour == WriterContour::Active && allow_coverage_override {
|
||||||
|
limit = limit
|
||||||
|
.max(self.active_coverage_required_total().await)
|
||||||
|
.saturating_add(
|
||||||
|
self.reconnect_runtime
|
||||||
|
.me_reconnect_max_concurrent_per_dc
|
||||||
|
.max(1) as usize,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
let reserved = counter.load(Ordering::Acquire);
|
||||||
|
if live.saturating_add(reserved) >= limit {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
if counter
|
||||||
|
.compare_exchange_weak(
|
||||||
|
reserved,
|
||||||
|
reserved + 1,
|
||||||
|
Ordering::AcqRel,
|
||||||
|
Ordering::Acquire,
|
||||||
|
)
|
||||||
|
.is_ok()
|
||||||
|
{
|
||||||
|
return Some(WriterOpenReservation {
|
||||||
|
counter: Some(counter),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
pub(super) fn required_writers_for_dc_with_floor_mode(
|
pub(super) fn required_writers_for_dc_with_floor_mode(
|
||||||
&self,
|
&self,
|
||||||
endpoint_count: usize,
|
endpoint_count: usize,
|
||||||
|
|||||||
@@ -107,7 +107,7 @@ impl MePool {
|
|||||||
let pool = Arc::clone(self);
|
let pool = Arc::clone(self);
|
||||||
let rng_clone = Arc::clone(rng);
|
let rng_clone = Arc::clone(rng);
|
||||||
let dc_addrs_bg = dc_addrs.clone();
|
let dc_addrs_bg = dc_addrs.clone();
|
||||||
tokio::spawn(async move {
|
let saturation = async move {
|
||||||
let mut join_bg = tokio::task::JoinSet::new();
|
let mut join_bg = tokio::task::JoinSet::new();
|
||||||
for (dc, addrs) in dc_addrs_bg {
|
for (dc, addrs) in dc_addrs_bg {
|
||||||
if addrs.len() <= 1 {
|
if addrs.len() <= 1 {
|
||||||
@@ -135,7 +135,10 @@ impl MePool {
|
|||||||
current_pool_size = pool.connection_count(),
|
current_pool_size = pool.connection_count(),
|
||||||
"Background ME saturation warmup finished"
|
"Background ME saturation warmup finished"
|
||||||
);
|
);
|
||||||
});
|
};
|
||||||
|
if self.lifecycle.spawn_producer(saturation).is_err() {
|
||||||
|
debug!("Background ME saturation skipped: pool lifecycle closed");
|
||||||
|
}
|
||||||
|
|
||||||
if !self.decision.effective_multipath && self.connection_count() > 0 {
|
if !self.decision.effective_multipath && self.connection_count() > 0 {
|
||||||
break;
|
break;
|
||||||
|
|||||||
@@ -0,0 +1,326 @@
|
|||||||
|
use std::future::Future;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
|
||||||
|
use std::time::Duration;
|
||||||
|
|
||||||
|
use tokio::sync::Notify;
|
||||||
|
use tokio::sync::mpsc;
|
||||||
|
use tokio_util::sync::CancellationToken;
|
||||||
|
use tokio_util::task::TaskTracker;
|
||||||
|
|
||||||
|
use super::pool::MePool;
|
||||||
|
use super::registry::ConnLease;
|
||||||
|
|
||||||
|
const ME_TASK_ADMISSION_CLOSED: usize = 1 << (usize::BITS - 1);
|
||||||
|
const ME_TASK_REGISTRATION_COUNT: usize = ME_TASK_ADMISSION_CLOSED - 1;
|
||||||
|
|
||||||
|
struct MeTaskAdmission {
|
||||||
|
state: AtomicUsize,
|
||||||
|
registrations_drained: Notify,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// RAII ownership of one ME task-publication section.
|
||||||
|
pub(super) struct MeTaskRegistration<'a> {
|
||||||
|
admission: &'a MeTaskAdmission,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl MeTaskAdmission {
|
||||||
|
fn new() -> Self {
|
||||||
|
Self {
|
||||||
|
state: AtomicUsize::new(0),
|
||||||
|
registrations_drained: Notify::new(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn try_register(&self) -> Option<MeTaskRegistration<'_>> {
|
||||||
|
let mut state = self.state.load(Ordering::Acquire);
|
||||||
|
loop {
|
||||||
|
if state & ME_TASK_ADMISSION_CLOSED != 0
|
||||||
|
|| state & ME_TASK_REGISTRATION_COUNT == ME_TASK_REGISTRATION_COUNT
|
||||||
|
{
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
match self.state.compare_exchange_weak(
|
||||||
|
state,
|
||||||
|
state + 1,
|
||||||
|
Ordering::AcqRel,
|
||||||
|
Ordering::Acquire,
|
||||||
|
) {
|
||||||
|
Ok(_) => return Some(MeTaskRegistration { admission: self }),
|
||||||
|
Err(observed) => state = observed,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn close(&self) {
|
||||||
|
self.state
|
||||||
|
.fetch_or(ME_TASK_ADMISSION_CLOSED, Ordering::AcqRel);
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn wait_for_registrations(&self) {
|
||||||
|
loop {
|
||||||
|
let notified = self.registrations_drained.notified();
|
||||||
|
tokio::pin!(notified);
|
||||||
|
notified.as_mut().enable();
|
||||||
|
if self.state.load(Ordering::Acquire) & ME_TASK_REGISTRATION_COUNT == 0 {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
notified.await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Drop for MeTaskRegistration<'_> {
|
||||||
|
fn drop(&mut self) {
|
||||||
|
let previous = self.admission.state.fetch_sub(1, Ordering::AcqRel);
|
||||||
|
if previous & ME_TASK_REGISTRATION_COUNT == 1 {
|
||||||
|
self.admission.registrations_drained.notify_waiters();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Pool-owned admission, cancellation, and join authority for ME tasks.
|
||||||
|
pub(super) struct MePoolLifecycle {
|
||||||
|
admission: MeTaskAdmission,
|
||||||
|
producer_cancel: CancellationToken,
|
||||||
|
producer_tasks: TaskTracker,
|
||||||
|
writer_tasks: TaskTracker,
|
||||||
|
cleanup_cancel: CancellationToken,
|
||||||
|
cleanup_tasks: TaskTracker,
|
||||||
|
cleanup_started: AtomicBool,
|
||||||
|
shutdown_started: AtomicBool,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl MePoolLifecycle {
|
||||||
|
/// Creates an open ME task lifecycle.
|
||||||
|
pub(super) fn new() -> Self {
|
||||||
|
Self {
|
||||||
|
admission: MeTaskAdmission::new(),
|
||||||
|
producer_cancel: CancellationToken::new(),
|
||||||
|
producer_tasks: TaskTracker::new(),
|
||||||
|
writer_tasks: TaskTracker::new(),
|
||||||
|
cleanup_cancel: CancellationToken::new(),
|
||||||
|
cleanup_tasks: TaskTracker::new(),
|
||||||
|
cleanup_started: AtomicBool::new(false),
|
||||||
|
shutdown_started: AtomicBool::new(false),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Registers one task-publication section while lifecycle admission is open.
|
||||||
|
pub(super) fn try_register(&self) -> Option<MeTaskRegistration<'_>> {
|
||||||
|
self.admission.try_register()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Registers and spawns one cancellation-aware ME producer.
|
||||||
|
pub(super) fn spawn_producer<F>(&self, future: F) -> Result<(), F>
|
||||||
|
where
|
||||||
|
F: Future<Output = ()> + Send + 'static,
|
||||||
|
{
|
||||||
|
let Some(registration) = self.try_register() else {
|
||||||
|
return Err(future);
|
||||||
|
};
|
||||||
|
self.spawn_registered_producer(registration, future);
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Spawns a producer after its caller published cancellation cleanup ownership.
|
||||||
|
pub(super) fn spawn_registered_producer<F>(
|
||||||
|
&self,
|
||||||
|
registration: MeTaskRegistration<'_>,
|
||||||
|
future: F,
|
||||||
|
) where
|
||||||
|
F: Future<Output = ()> + Send + 'static,
|
||||||
|
{
|
||||||
|
let cancel = self.producer_cancel.clone();
|
||||||
|
self.producer_tasks.spawn(async move {
|
||||||
|
tokio::select! {
|
||||||
|
biased;
|
||||||
|
_ = cancel.cancelled() => {}
|
||||||
|
_ = future => {}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
drop(registration);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Spawns a writer after its caller completed task registration.
|
||||||
|
pub(super) fn spawn_registered_writer<F>(
|
||||||
|
&self,
|
||||||
|
registration: MeTaskRegistration<'_>,
|
||||||
|
future: F,
|
||||||
|
) where
|
||||||
|
F: Future<Output = ()> + Send + 'static,
|
||||||
|
{
|
||||||
|
self.writer_tasks.spawn(future);
|
||||||
|
drop(registration);
|
||||||
|
}
|
||||||
|
|
||||||
|
fn start_cleanup_worker(&self, pool: &Arc<MePool>) -> bool {
|
||||||
|
let Some(registration) = self.try_register() else {
|
||||||
|
return false;
|
||||||
|
};
|
||||||
|
if self
|
||||||
|
.cleanup_started
|
||||||
|
.compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
|
||||||
|
.is_err()
|
||||||
|
{
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
let Some(mut cleanup_rx) = pool.registry.take_cleanup_receiver() else {
|
||||||
|
self.cleanup_started.store(false, Ordering::Release);
|
||||||
|
return false;
|
||||||
|
};
|
||||||
|
let registry = Arc::clone(&pool.registry);
|
||||||
|
let cancel = self.cleanup_cancel.clone();
|
||||||
|
self.cleanup_tasks.spawn(async move {
|
||||||
|
loop {
|
||||||
|
tokio::select! {
|
||||||
|
biased;
|
||||||
|
cleanup = cleanup_rx.recv() => {
|
||||||
|
let Some(conn_id) = cleanup else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
registry.unregister(conn_id).await;
|
||||||
|
}
|
||||||
|
_ = cancel.cancelled() => {
|
||||||
|
while let Ok(conn_id) = cleanup_rx.try_recv() {
|
||||||
|
registry.unregister(conn_id).await;
|
||||||
|
}
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
drop(registration);
|
||||||
|
true
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Idempotently closes task admission and cancels ME producers.
|
||||||
|
pub(super) fn begin_shutdown(&self) {
|
||||||
|
if self.shutdown_started.swap(true, Ordering::AcqRel) {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
self.admission.close();
|
||||||
|
self.producer_cancel.cancel();
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn wait_until<F>(deadline: tokio::time::Instant, future: F) -> bool
|
||||||
|
where
|
||||||
|
F: Future<Output = ()>,
|
||||||
|
{
|
||||||
|
let now = tokio::time::Instant::now();
|
||||||
|
if now >= deadline {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
tokio::time::timeout_at(deadline, future).await.is_ok()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Joins producers, writers, and cleanup ownership under one deadline.
|
||||||
|
pub(super) async fn shutdown_pool(
|
||||||
|
&self,
|
||||||
|
pool: &Arc<MePool>,
|
||||||
|
timeout: Duration,
|
||||||
|
) -> bool {
|
||||||
|
let deadline = tokio::time::Instant::now() + timeout;
|
||||||
|
self.begin_shutdown();
|
||||||
|
pool.set_runtime_ready(false);
|
||||||
|
|
||||||
|
let registrations_stopped =
|
||||||
|
Self::wait_until(deadline, self.admission.wait_for_registrations()).await;
|
||||||
|
|
||||||
|
self.producer_tasks.close();
|
||||||
|
let producers_stopped = Self::wait_until(deadline, self.producer_tasks.wait()).await;
|
||||||
|
|
||||||
|
let close_signals_sent = if tokio::time::Instant::now() < deadline {
|
||||||
|
tokio::time::timeout_at(deadline, pool.shutdown_send_close_conn_all())
|
||||||
|
.await
|
||||||
|
.is_ok()
|
||||||
|
} else {
|
||||||
|
false
|
||||||
|
};
|
||||||
|
|
||||||
|
let writers = pool.writers.snapshot();
|
||||||
|
for writer in writers.iter() {
|
||||||
|
writer.cancel.cancel();
|
||||||
|
}
|
||||||
|
self.writer_tasks.close();
|
||||||
|
let writers_stopped = Self::wait_until(deadline, self.writer_tasks.wait()).await;
|
||||||
|
|
||||||
|
self.cleanup_cancel.cancel();
|
||||||
|
self.cleanup_tasks.close();
|
||||||
|
let cleanup_stopped = Self::wait_until(deadline, self.cleanup_tasks.wait()).await;
|
||||||
|
|
||||||
|
registrations_stopped
|
||||||
|
&& producers_stopped
|
||||||
|
&& close_signals_sent
|
||||||
|
&& writers_stopped
|
||||||
|
&& cleanup_stopped
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl MePool {
|
||||||
|
/// Registers one cancellation-safe client route in the bounded cleanup plane.
|
||||||
|
pub(crate) async fn register_connection(
|
||||||
|
self: &Arc<Self>,
|
||||||
|
) -> Option<(ConnLease, mpsc::Receiver<super::MeResponse>)> {
|
||||||
|
if !self.lifecycle.start_cleanup_worker(self) {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
let registration = self.lifecycle.try_register()?;
|
||||||
|
let registered = tokio::select! {
|
||||||
|
biased;
|
||||||
|
_ = self.lifecycle.producer_cancel.cancelled() => None,
|
||||||
|
registered = self.registry.register_leased() => registered,
|
||||||
|
};
|
||||||
|
drop(registration);
|
||||||
|
registered
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Closes ME task admission and joins pool-owned producers and writers.
|
||||||
|
pub(crate) async fn shutdown_until(self: &Arc<Self>, timeout: Duration) -> bool {
|
||||||
|
self.lifecycle.shutdown_pool(self, timeout).await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Terminally closes ME task admission without waiting for asynchronous teardown.
|
||||||
|
pub(crate) fn begin_shutdown(&self) {
|
||||||
|
self.lifecycle.begin_shutdown();
|
||||||
|
self.set_runtime_ready(false);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use std::sync::Arc;
|
||||||
|
use std::sync::atomic::{AtomicBool, Ordering};
|
||||||
|
use std::time::Duration;
|
||||||
|
|
||||||
|
use super::MePoolLifecycle;
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn shutdown_fence_rejects_late_producer_registration() {
|
||||||
|
struct DropSignal(Arc<AtomicBool>);
|
||||||
|
|
||||||
|
impl Drop for DropSignal {
|
||||||
|
fn drop(&mut self) {
|
||||||
|
self.0.store(true, Ordering::Release);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let lifecycle = Arc::new(MePoolLifecycle::new());
|
||||||
|
let dropped = Arc::new(AtomicBool::new(false));
|
||||||
|
let drop_signal = DropSignal(dropped.clone());
|
||||||
|
let spawned = lifecycle.spawn_producer(async move {
|
||||||
|
let _drop_signal = drop_signal;
|
||||||
|
std::future::pending::<()>().await;
|
||||||
|
});
|
||||||
|
assert!(spawned.is_ok());
|
||||||
|
|
||||||
|
lifecycle.begin_shutdown();
|
||||||
|
lifecycle.producer_tasks.close();
|
||||||
|
tokio::time::timeout(Duration::from_secs(1), lifecycle.producer_tasks.wait())
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
assert!(dropped.load(Ordering::Acquire));
|
||||||
|
assert!(lifecycle.spawn_producer(async {}).is_err());
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -9,7 +9,8 @@ use tracing::{debug, info};
|
|||||||
use crate::error::{ProxyError, Result};
|
use crate::error::{ProxyError, Result};
|
||||||
use crate::network::probe::{detect_public_ipv4_http, is_bogon};
|
use crate::network::probe::{detect_public_ipv4_http, is_bogon};
|
||||||
use crate::network::stun::{
|
use crate::network::stun::{
|
||||||
IpFamily, stun_probe_dual_with_tcp_fallback, stun_probe_family_with_bind_and_tcp_fallback,
|
IpFamily, stun_probe_dual_with_tcp_fallback,
|
||||||
|
stun_probe_family_with_bind_tcp_fallback_and_resolver,
|
||||||
};
|
};
|
||||||
|
|
||||||
use super::MePool;
|
use super::MePool;
|
||||||
@@ -66,10 +67,12 @@ impl MePool {
|
|||||||
let mut best_by_ip: HashMap<IpAddr, (usize, std::net::SocketAddr)> = HashMap::new();
|
let mut best_by_ip: HashMap<IpAddr, (usize, std::net::SocketAddr)> = HashMap::new();
|
||||||
let concurrency = self.nat_runtime.nat_probe_concurrency.max(1);
|
let concurrency = self.nat_runtime.nat_probe_concurrency.max(1);
|
||||||
let tcp_fallback = self.nat_runtime.stun_tcp_fallback;
|
let tcp_fallback = self.nat_runtime.stun_tcp_fallback;
|
||||||
|
let dns_resolver = self.upstream.as_ref().map(|manager| manager.dns_resolver());
|
||||||
|
|
||||||
while next_idx < servers.len() || !join_set.is_empty() {
|
while next_idx < servers.len() || !join_set.is_empty() {
|
||||||
while next_idx < servers.len() && join_set.len() < concurrency {
|
while next_idx < servers.len() && join_set.len() < concurrency {
|
||||||
let stun_addr = servers[next_idx].clone();
|
let stun_addr = servers[next_idx].clone();
|
||||||
|
let dns_resolver = dns_resolver.clone();
|
||||||
next_idx += 1;
|
next_idx += 1;
|
||||||
join_set.spawn(async move {
|
join_set.spawn(async move {
|
||||||
let batch_timeout = if tcp_fallback {
|
let batch_timeout = if tcp_fallback {
|
||||||
@@ -79,11 +82,12 @@ impl MePool {
|
|||||||
};
|
};
|
||||||
let res = timeout(
|
let res = timeout(
|
||||||
batch_timeout,
|
batch_timeout,
|
||||||
stun_probe_family_with_bind_and_tcp_fallback(
|
stun_probe_family_with_bind_tcp_fallback_and_resolver(
|
||||||
&stun_addr,
|
&stun_addr,
|
||||||
family,
|
family,
|
||||||
bind_ip,
|
bind_ip,
|
||||||
tcp_fallback,
|
tcp_fallback,
|
||||||
|
dns_resolver.as_deref(),
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
|
|||||||
@@ -9,13 +9,48 @@ use tracing::{debug, info, warn};
|
|||||||
use crate::crypto::SecureRandom;
|
use crate::crypto::SecureRandom;
|
||||||
use crate::network::IpFamily;
|
use crate::network::IpFamily;
|
||||||
|
|
||||||
use super::pool::{MePool, RefillDcKey, RefillEndpointKey, WriterContour};
|
use super::pool::{MePool, RefillDcKey, WriterContour};
|
||||||
|
|
||||||
const ME_FLAP_UPTIME_THRESHOLD_SECS: u64 = 20;
|
const ME_FLAP_UPTIME_THRESHOLD_SECS: u64 = 20;
|
||||||
const ME_FLAP_QUARANTINE_SECS: u64 = 25;
|
const ME_FLAP_QUARANTINE_SECS: u64 = 25;
|
||||||
const ME_FLAP_MIN_UPTIME_MILLIS: u64 = 500;
|
const ME_FLAP_MIN_UPTIME_MILLIS: u64 = 500;
|
||||||
const ME_REFILL_TOTAL_ATTEMPT_CAP: u32 = 20;
|
const ME_REFILL_TOTAL_ATTEMPT_CAP: u32 = 20;
|
||||||
|
|
||||||
|
struct RefillRunGuard {
|
||||||
|
pool: Arc<MePool>,
|
||||||
|
key: RefillDcKey,
|
||||||
|
active: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl RefillRunGuard {
|
||||||
|
fn next_or_finish(&mut self) -> Option<SocketAddr> {
|
||||||
|
let mut states = self.pool.refill_states.lock();
|
||||||
|
let next = states.get_mut(&self.key).and_then(Option::take);
|
||||||
|
if next.is_some() {
|
||||||
|
self.pool.refill_pending.fetch_sub(1, Ordering::AcqRel);
|
||||||
|
return next;
|
||||||
|
}
|
||||||
|
states.remove(&self.key);
|
||||||
|
self.pool.refill_running.fetch_sub(1, Ordering::AcqRel);
|
||||||
|
self.active = false;
|
||||||
|
None
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Drop for RefillRunGuard {
|
||||||
|
fn drop(&mut self) {
|
||||||
|
if !self.active {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
if let Some(pending) = self.pool.refill_states.lock().remove(&self.key)
|
||||||
|
&& pending.is_some()
|
||||||
|
{
|
||||||
|
self.pool.refill_pending.fetch_sub(1, Ordering::AcqRel);
|
||||||
|
}
|
||||||
|
self.pool.refill_running.fetch_sub(1, Ordering::AcqRel);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
impl MePool {
|
impl MePool {
|
||||||
pub(super) async fn sweep_endpoint_quarantine(&self) {
|
pub(super) async fn sweep_endpoint_quarantine(&self) {
|
||||||
let configured = self
|
let configured = self
|
||||||
@@ -131,8 +166,7 @@ impl MePool {
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub(super) async fn has_refill_inflight_for_dc_key(&self, key: RefillDcKey) -> bool {
|
pub(super) async fn has_refill_inflight_for_dc_key(&self, key: RefillDcKey) -> bool {
|
||||||
let guard = self.refill_inflight_dc.lock().await;
|
self.refill_states.lock().contains_key(&key)
|
||||||
guard.contains(&key)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(super) async fn connect_endpoints_round_robin(
|
pub(super) async fn connect_endpoints_round_robin(
|
||||||
@@ -327,62 +361,53 @@ impl MePool {
|
|||||||
addr: SocketAddr,
|
addr: SocketAddr,
|
||||||
writer_dc: i32,
|
writer_dc: i32,
|
||||||
) {
|
) {
|
||||||
let endpoint_key = RefillEndpointKey {
|
let Some(registration) = self.lifecycle.try_register() else {
|
||||||
dc: writer_dc,
|
return;
|
||||||
addr,
|
|
||||||
};
|
};
|
||||||
let pre_inserted = if let Ok(mut guard) = self.refill_inflight.try_lock() {
|
let dc_key = RefillDcKey {
|
||||||
if !guard.insert(endpoint_key) {
|
dc: writer_dc,
|
||||||
|
family: if addr.is_ipv4() {
|
||||||
|
IpFamily::V4
|
||||||
|
} else {
|
||||||
|
IpFamily::V6
|
||||||
|
},
|
||||||
|
};
|
||||||
|
{
|
||||||
|
let mut states = self.refill_states.lock();
|
||||||
|
if let Some(pending) = states.get_mut(&dc_key) {
|
||||||
|
if pending.is_none() {
|
||||||
|
self.refill_pending.fetch_add(1, Ordering::AcqRel);
|
||||||
|
}
|
||||||
|
*pending = Some(addr);
|
||||||
self.stats.increment_me_refill_skipped_inflight_total();
|
self.stats.increment_me_refill_skipped_inflight_total();
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
true
|
states.insert(dc_key, None);
|
||||||
} else {
|
self.refill_running.fetch_add(1, Ordering::AcqRel);
|
||||||
false
|
}
|
||||||
};
|
|
||||||
|
|
||||||
let pool = Arc::clone(self);
|
let pool = Arc::clone(self);
|
||||||
tokio::spawn(async move {
|
let mut run_guard = RefillRunGuard {
|
||||||
let dc_key = RefillDcKey {
|
pool: Arc::clone(&pool),
|
||||||
dc: writer_dc,
|
key: dc_key,
|
||||||
family: if addr.is_ipv4() {
|
active: true,
|
||||||
IpFamily::V4
|
};
|
||||||
} else {
|
self.lifecycle.spawn_registered_producer(registration, async move {
|
||||||
IpFamily::V6
|
let mut current_addr = addr;
|
||||||
},
|
loop {
|
||||||
};
|
pool.stats.increment_me_refill_triggered_total();
|
||||||
|
let restored = pool
|
||||||
if !pre_inserted {
|
.refill_writer_after_loss(current_addr, writer_dc)
|
||||||
let mut guard = pool.refill_inflight.lock().await;
|
.await;
|
||||||
if !guard.insert(endpoint_key) {
|
if !restored {
|
||||||
pool.stats.increment_me_refill_skipped_inflight_total();
|
warn!(%current_addr, dc = writer_dc, "ME immediate refill failed");
|
||||||
return;
|
|
||||||
}
|
}
|
||||||
}
|
|
||||||
|
|
||||||
{
|
let Some(next_addr) = run_guard.next_or_finish() else {
|
||||||
let mut dc_guard = pool.refill_inflight_dc.lock().await;
|
|
||||||
if dc_guard.contains(&dc_key) {
|
|
||||||
pool.stats.increment_me_refill_skipped_inflight_total();
|
|
||||||
drop(dc_guard);
|
|
||||||
let mut guard = pool.refill_inflight.lock().await;
|
|
||||||
guard.remove(&endpoint_key);
|
|
||||||
return;
|
return;
|
||||||
}
|
};
|
||||||
dc_guard.insert(dc_key);
|
current_addr = next_addr;
|
||||||
}
|
}
|
||||||
|
|
||||||
pool.stats.increment_me_refill_triggered_total();
|
|
||||||
let restored = pool.refill_writer_after_loss(addr, writer_dc).await;
|
|
||||||
if !restored {
|
|
||||||
warn!(%addr, dc = writer_dc, "ME immediate refill failed");
|
|
||||||
}
|
|
||||||
|
|
||||||
let mut guard = pool.refill_inflight.lock().await;
|
|
||||||
guard.remove(&endpoint_key);
|
|
||||||
drop(guard);
|
|
||||||
let mut dc_guard = pool.refill_inflight_dc.lock().await;
|
|
||||||
dc_guard.remove(&dc_key);
|
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -13,10 +13,106 @@ use tracing::{debug, info, warn};
|
|||||||
use crate::crypto::SecureRandom;
|
use crate::crypto::SecureRandom;
|
||||||
use crate::network::IpFamily;
|
use crate::network::IpFamily;
|
||||||
|
|
||||||
use super::pool::{MeDrainGateReason, MePool, WriterContour};
|
use super::pool::{
|
||||||
|
MeDrainGateReason, MePool, ReinitAttemptState, ReinitCoordinatorState, ReinitCore,
|
||||||
|
ReinitPendingState, ReinitStatusSnapshot, WriterContour,
|
||||||
|
};
|
||||||
|
|
||||||
const ME_HARDSWAP_PENDING_TTL_SECS: u64 = 1800;
|
const ME_HARDSWAP_PENDING_TTL_SECS: u64 = 1800;
|
||||||
|
|
||||||
|
struct ReinitAttemptGuard {
|
||||||
|
reinit: Arc<ReinitCore>,
|
||||||
|
attempt_id: u64,
|
||||||
|
generation: u64,
|
||||||
|
previous_generation: u64,
|
||||||
|
map_hash: u64,
|
||||||
|
hardswap: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Drop for ReinitAttemptGuard {
|
||||||
|
fn drop(&mut self) {
|
||||||
|
let mut state = self.reinit.coordinator.lock();
|
||||||
|
state.attempts.remove(&self.attempt_id);
|
||||||
|
publish_reinit_state(self.reinit.as_ref(), &state);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
struct ReinitReservation {
|
||||||
|
attempt: ReinitAttemptGuard,
|
||||||
|
pending_reused: bool,
|
||||||
|
pending_expired: bool,
|
||||||
|
pending_age_secs: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
fn publish_reinit_state(reinit: &ReinitCore, state: &ReinitCoordinatorState) {
|
||||||
|
let mut warm_generations = state
|
||||||
|
.attempts
|
||||||
|
.values()
|
||||||
|
.filter(|attempt| attempt.hardswap && !attempt.committed)
|
||||||
|
.map(|attempt| attempt.generation)
|
||||||
|
.collect::<Vec<_>>();
|
||||||
|
warm_generations.sort_unstable();
|
||||||
|
warm_generations.dedup();
|
||||||
|
let pending = state.pending;
|
||||||
|
let snapshot = ReinitStatusSnapshot {
|
||||||
|
active_generation: state.active_generation,
|
||||||
|
warm_generations,
|
||||||
|
pending_hardswap_generation: pending.map_or(0, |value| value.generation),
|
||||||
|
pending_hardswap_started_at_epoch_secs: pending
|
||||||
|
.map_or(0, |value| value.started_at_epoch_secs),
|
||||||
|
pending_hardswap_map_hash: pending.map_or(0, |value| value.map_hash),
|
||||||
|
inflight: state.attempts.len(),
|
||||||
|
};
|
||||||
|
reinit
|
||||||
|
.active_generation
|
||||||
|
.store(snapshot.active_generation, Ordering::Release);
|
||||||
|
reinit.warm_generation.store(
|
||||||
|
snapshot.warm_generations.last().copied().unwrap_or(0),
|
||||||
|
Ordering::Release,
|
||||||
|
);
|
||||||
|
reinit.pending_hardswap_generation.store(
|
||||||
|
snapshot.pending_hardswap_generation,
|
||||||
|
Ordering::Release,
|
||||||
|
);
|
||||||
|
reinit.pending_hardswap_started_at_epoch_secs.store(
|
||||||
|
snapshot.pending_hardswap_started_at_epoch_secs,
|
||||||
|
Ordering::Release,
|
||||||
|
);
|
||||||
|
reinit
|
||||||
|
.pending_hardswap_map_hash
|
||||||
|
.store(snapshot.pending_hardswap_map_hash, Ordering::Release);
|
||||||
|
reinit.status.store(Arc::new(snapshot));
|
||||||
|
}
|
||||||
|
|
||||||
|
fn commit_reinit_state(
|
||||||
|
state: &mut ReinitCoordinatorState,
|
||||||
|
attempt_id: u64,
|
||||||
|
generation: u64,
|
||||||
|
map_hash: u64,
|
||||||
|
hardswap: bool,
|
||||||
|
) -> bool {
|
||||||
|
let Some(record) = state.attempts.get(&attempt_id).copied() else {
|
||||||
|
return false;
|
||||||
|
};
|
||||||
|
if record.map_hash != state.desired_map_hash || record.map_hash != map_hash {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
if hardswap {
|
||||||
|
let pending_matches = state.pending.is_some_and(|pending| {
|
||||||
|
pending.generation == generation && pending.map_hash == map_hash
|
||||||
|
});
|
||||||
|
if !pending_matches || generation < state.active_generation {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
state.active_generation = generation;
|
||||||
|
state.pending = None;
|
||||||
|
}
|
||||||
|
if let Some(record) = state.attempts.get_mut(&attempt_id) {
|
||||||
|
record.committed = true;
|
||||||
|
}
|
||||||
|
true
|
||||||
|
}
|
||||||
|
|
||||||
impl MePool {
|
impl MePool {
|
||||||
fn desired_map_hash(desired_by_dc: &HashMap<i32, HashSet<SocketAddr>>) -> u64 {
|
fn desired_map_hash(desired_by_dc: &HashMap<i32, HashSet<SocketAddr>>) -> u64 {
|
||||||
let mut hasher = DefaultHasher::new();
|
let mut hasher = DefaultHasher::new();
|
||||||
@@ -36,36 +132,97 @@ impl MePool {
|
|||||||
hasher.finish()
|
hasher.finish()
|
||||||
}
|
}
|
||||||
|
|
||||||
fn clear_pending_hardswap_state(&self) {
|
fn reserve_reinit_attempt(
|
||||||
self.reinit
|
self: &Arc<Self>,
|
||||||
.pending_hardswap_generation
|
hardswap: bool,
|
||||||
.store(0, Ordering::Relaxed);
|
map_hash: u64,
|
||||||
self.reinit
|
now_epoch_secs: u64,
|
||||||
.pending_hardswap_started_at_epoch_secs
|
) -> ReinitReservation {
|
||||||
.store(0, Ordering::Relaxed);
|
let mut state = self.reinit.coordinator.lock();
|
||||||
self.reinit
|
state.desired_map_hash = map_hash;
|
||||||
.pending_hardswap_map_hash
|
let previous_generation = state.active_generation;
|
||||||
.store(0, Ordering::Relaxed);
|
let mut pending_reused = false;
|
||||||
self.reinit.warm_generation.store(0, Ordering::Relaxed);
|
let mut pending_expired = false;
|
||||||
|
let mut pending_age_secs = 0;
|
||||||
|
|
||||||
|
let generation = if hardswap {
|
||||||
|
let reusable = state.pending.filter(|pending| {
|
||||||
|
pending_age_secs = now_epoch_secs.saturating_sub(pending.started_at_epoch_secs);
|
||||||
|
pending_expired = pending.started_at_epoch_secs > 0
|
||||||
|
&& pending_age_secs > ME_HARDSWAP_PENDING_TTL_SECS;
|
||||||
|
pending.generation >= previous_generation
|
||||||
|
&& pending.map_hash == map_hash
|
||||||
|
&& !pending_expired
|
||||||
|
});
|
||||||
|
if let Some(pending) = reusable {
|
||||||
|
pending_reused = true;
|
||||||
|
pending.generation
|
||||||
|
} else {
|
||||||
|
let generation = self.reinit.generation.fetch_add(1, Ordering::AcqRel) + 1;
|
||||||
|
state.pending = Some(ReinitPendingState {
|
||||||
|
generation,
|
||||||
|
started_at_epoch_secs: now_epoch_secs,
|
||||||
|
map_hash,
|
||||||
|
});
|
||||||
|
generation
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
state.pending = None;
|
||||||
|
self.reinit.generation.fetch_add(1, Ordering::AcqRel) + 1
|
||||||
|
};
|
||||||
|
|
||||||
|
let attempt_id = state.next_attempt_id;
|
||||||
|
state.next_attempt_id = state.next_attempt_id.saturating_add(1);
|
||||||
|
state.attempts.insert(
|
||||||
|
attempt_id,
|
||||||
|
ReinitAttemptState {
|
||||||
|
generation,
|
||||||
|
map_hash,
|
||||||
|
hardswap,
|
||||||
|
committed: false,
|
||||||
|
},
|
||||||
|
);
|
||||||
|
publish_reinit_state(self.reinit.as_ref(), &state);
|
||||||
|
ReinitReservation {
|
||||||
|
attempt: ReinitAttemptGuard {
|
||||||
|
reinit: Arc::clone(&self.reinit),
|
||||||
|
attempt_id,
|
||||||
|
generation,
|
||||||
|
previous_generation,
|
||||||
|
map_hash,
|
||||||
|
hardswap,
|
||||||
|
},
|
||||||
|
pending_reused,
|
||||||
|
pending_expired,
|
||||||
|
pending_age_secs,
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn promote_warm_generation_to_active(&self, generation: u64) {
|
fn commit_reinit_attempt(&self, attempt: &ReinitAttemptGuard) -> bool {
|
||||||
self.reinit
|
let mut state = self.reinit.coordinator.lock();
|
||||||
.active_generation
|
if !commit_reinit_state(
|
||||||
.store(generation, Ordering::Relaxed);
|
&mut state,
|
||||||
self.reinit.warm_generation.store(0, Ordering::Relaxed);
|
attempt.attempt_id,
|
||||||
|
attempt.generation,
|
||||||
let ws = self.writers.read().await;
|
attempt.map_hash,
|
||||||
for writer in ws.iter() {
|
attempt.hardswap,
|
||||||
if writer.draining.load(Ordering::Relaxed) {
|
) {
|
||||||
continue;
|
return false;
|
||||||
}
|
}
|
||||||
if writer.generation == generation {
|
if attempt.hardswap {
|
||||||
writer
|
let writers = self.writers.snapshot();
|
||||||
.contour
|
for writer in writers.iter() {
|
||||||
.store(WriterContour::Active.as_u8(), Ordering::Relaxed);
|
if !writer.draining.load(Ordering::Relaxed)
|
||||||
|
&& writer.generation == attempt.generation
|
||||||
|
{
|
||||||
|
writer
|
||||||
|
.contour
|
||||||
|
.store(WriterContour::Active.as_u8(), Ordering::Release);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
publish_reinit_state(self.reinit.as_ref(), &state);
|
||||||
|
true
|
||||||
}
|
}
|
||||||
|
|
||||||
fn coverage_ratio(
|
fn coverage_ratio(
|
||||||
@@ -387,74 +544,32 @@ impl MePool {
|
|||||||
}
|
}
|
||||||
|
|
||||||
let desired_map_hash = Self::desired_map_hash(&desired_by_dc);
|
let desired_map_hash = Self::desired_map_hash(&desired_by_dc);
|
||||||
let previous_generation = self.current_generation();
|
|
||||||
let hardswap = self.reinit.hardswap.load(Ordering::Relaxed);
|
let hardswap = self.reinit.hardswap.load(Ordering::Relaxed);
|
||||||
let generation = if hardswap {
|
let reservation =
|
||||||
let pending_generation = self
|
self.reserve_reinit_attempt(hardswap, desired_map_hash, now_epoch_secs);
|
||||||
.reinit
|
let attempt = reservation.attempt;
|
||||||
.pending_hardswap_generation
|
let previous_generation = attempt.previous_generation;
|
||||||
.load(Ordering::Relaxed);
|
let generation = attempt.generation;
|
||||||
let pending_started_at = self
|
if reservation.pending_reused {
|
||||||
.reinit
|
self.stats.increment_me_hardswap_pending_reuse_total();
|
||||||
.pending_hardswap_started_at_epoch_secs
|
debug!(
|
||||||
.load(Ordering::Relaxed);
|
previous_generation,
|
||||||
let pending_map_hash = self
|
generation,
|
||||||
.reinit
|
pending_age_secs = reservation.pending_age_secs,
|
||||||
.pending_hardswap_map_hash
|
"ME hardswap continues with pending generation"
|
||||||
.load(Ordering::Relaxed);
|
);
|
||||||
let pending_age_secs = now_epoch_secs.saturating_sub(pending_started_at);
|
} else if reservation.pending_expired {
|
||||||
let pending_ttl_expired =
|
self.stats.increment_me_hardswap_pending_ttl_expired_total();
|
||||||
pending_started_at > 0 && pending_age_secs > ME_HARDSWAP_PENDING_TTL_SECS;
|
warn!(
|
||||||
let pending_matches_map = pending_map_hash != 0 && pending_map_hash == desired_map_hash;
|
previous_generation,
|
||||||
|
generation,
|
||||||
if pending_generation != 0
|
pending_age_secs = reservation.pending_age_secs,
|
||||||
&& pending_generation >= previous_generation
|
pending_ttl_secs = ME_HARDSWAP_PENDING_TTL_SECS,
|
||||||
&& pending_matches_map
|
"ME hardswap pending generation expired by TTL; starting fresh generation"
|
||||||
&& !pending_ttl_expired
|
);
|
||||||
{
|
}
|
||||||
self.stats.increment_me_hardswap_pending_reuse_total();
|
|
||||||
debug!(
|
|
||||||
previous_generation,
|
|
||||||
generation = pending_generation,
|
|
||||||
pending_age_secs,
|
|
||||||
"ME hardswap continues with pending generation"
|
|
||||||
);
|
|
||||||
pending_generation
|
|
||||||
} else {
|
|
||||||
if pending_generation != 0 && pending_ttl_expired {
|
|
||||||
self.stats.increment_me_hardswap_pending_ttl_expired_total();
|
|
||||||
warn!(
|
|
||||||
previous_generation,
|
|
||||||
generation = pending_generation,
|
|
||||||
pending_age_secs,
|
|
||||||
pending_ttl_secs = ME_HARDSWAP_PENDING_TTL_SECS,
|
|
||||||
"ME hardswap pending generation expired by TTL; starting fresh generation"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
let next_generation = self.reinit.generation.fetch_add(1, Ordering::Relaxed) + 1;
|
|
||||||
self.reinit
|
|
||||||
.pending_hardswap_generation
|
|
||||||
.store(next_generation, Ordering::Relaxed);
|
|
||||||
self.reinit
|
|
||||||
.pending_hardswap_started_at_epoch_secs
|
|
||||||
.store(now_epoch_secs, Ordering::Relaxed);
|
|
||||||
self.reinit
|
|
||||||
.pending_hardswap_map_hash
|
|
||||||
.store(desired_map_hash, Ordering::Relaxed);
|
|
||||||
self.reinit
|
|
||||||
.warm_generation
|
|
||||||
.store(next_generation, Ordering::Relaxed);
|
|
||||||
next_generation
|
|
||||||
}
|
|
||||||
} else {
|
|
||||||
self.clear_pending_hardswap_state();
|
|
||||||
self.reinit.generation.fetch_add(1, Ordering::Relaxed) + 1
|
|
||||||
};
|
|
||||||
|
|
||||||
if hardswap {
|
if hardswap {
|
||||||
self.reinit
|
|
||||||
.warm_generation
|
|
||||||
.store(generation, Ordering::Relaxed);
|
|
||||||
self.warmup_generation_for_all_dcs(rng, generation, &desired_by_dc)
|
self.warmup_generation_for_all_dcs(rng, generation, &desired_by_dc)
|
||||||
.await;
|
.await;
|
||||||
} else {
|
} else {
|
||||||
@@ -542,8 +657,13 @@ impl MePool {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
if hardswap {
|
if !self.commit_reinit_attempt(&attempt) {
|
||||||
self.promote_warm_generation_to_active(generation).await;
|
debug!(
|
||||||
|
previous_generation,
|
||||||
|
generation,
|
||||||
|
"ME reinit result discarded after a newer desired-map attempt"
|
||||||
|
);
|
||||||
|
return false;
|
||||||
}
|
}
|
||||||
|
|
||||||
let desired_addrs: HashSet<(i32, SocketAddr)> = desired_by_dc
|
let desired_addrs: HashSet<(i32, SocketAddr)> = desired_by_dc
|
||||||
@@ -566,9 +686,6 @@ impl MePool {
|
|||||||
drop(writers);
|
drop(writers);
|
||||||
|
|
||||||
if stale_writer_ids.is_empty() {
|
if stale_writer_ids.is_empty() {
|
||||||
if hardswap {
|
|
||||||
self.clear_pending_hardswap_state();
|
|
||||||
}
|
|
||||||
debug!("ME reinit cycle completed with no stale writers");
|
debug!("ME reinit cycle completed with no stale writers");
|
||||||
return true;
|
return true;
|
||||||
}
|
}
|
||||||
@@ -606,9 +723,6 @@ impl MePool {
|
|||||||
self.remove_writer_and_close_clients(writer_id).await;
|
self.remove_writer_and_close_clients(writer_id).await;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
if hardswap {
|
|
||||||
self.clear_pending_hardswap_state();
|
|
||||||
}
|
|
||||||
true
|
true
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -622,7 +736,10 @@ mod tests {
|
|||||||
use std::collections::{HashMap, HashSet};
|
use std::collections::{HashMap, HashSet};
|
||||||
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
|
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
|
||||||
|
|
||||||
use super::MePool;
|
use super::{MePool, commit_reinit_state};
|
||||||
|
use crate::transport::middle_proxy::pool::{
|
||||||
|
ReinitAttemptState, ReinitCoordinatorState, ReinitPendingState,
|
||||||
|
};
|
||||||
|
|
||||||
fn addr(octet: u8, port: u16) -> SocketAddr {
|
fn addr(octet: u8, port: u16) -> SocketAddr {
|
||||||
SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, octet)), port)
|
SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, octet)), port)
|
||||||
@@ -673,4 +790,44 @@ mod tests {
|
|||||||
assert_eq!(ratio, 0.0);
|
assert_eq!(ratio, 0.0);
|
||||||
assert_eq!(missing_dc, vec![1, 2]);
|
assert_eq!(missing_dc, vec![1, 2]);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn stale_concurrent_attempt_cannot_regress_active_generation() {
|
||||||
|
let mut state = ReinitCoordinatorState {
|
||||||
|
next_attempt_id: 3,
|
||||||
|
active_generation: 1,
|
||||||
|
desired_map_hash: 22,
|
||||||
|
pending: Some(ReinitPendingState {
|
||||||
|
generation: 3,
|
||||||
|
started_at_epoch_secs: 1,
|
||||||
|
map_hash: 22,
|
||||||
|
}),
|
||||||
|
attempts: HashMap::from([
|
||||||
|
(
|
||||||
|
1,
|
||||||
|
ReinitAttemptState {
|
||||||
|
generation: 2,
|
||||||
|
map_hash: 11,
|
||||||
|
hardswap: true,
|
||||||
|
committed: false,
|
||||||
|
},
|
||||||
|
),
|
||||||
|
(
|
||||||
|
2,
|
||||||
|
ReinitAttemptState {
|
||||||
|
generation: 3,
|
||||||
|
map_hash: 22,
|
||||||
|
hardswap: true,
|
||||||
|
committed: false,
|
||||||
|
},
|
||||||
|
),
|
||||||
|
]),
|
||||||
|
};
|
||||||
|
|
||||||
|
assert!(commit_reinit_state(&mut state, 2, 3, 22, true));
|
||||||
|
assert_eq!(state.active_generation, 3);
|
||||||
|
assert!(!commit_reinit_state(&mut state, 1, 2, 11, true));
|
||||||
|
assert_eq!(state.active_generation, 3);
|
||||||
|
assert!(state.pending.is_none());
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -15,6 +15,8 @@ pub(crate) struct MeApiRefillDcSnapshot {
|
|||||||
pub(crate) struct MeApiRefillSnapshot {
|
pub(crate) struct MeApiRefillSnapshot {
|
||||||
pub inflight_endpoints_total: usize,
|
pub inflight_endpoints_total: usize,
|
||||||
pub inflight_dc_total: usize,
|
pub inflight_dc_total: usize,
|
||||||
|
pub running_dc_total: usize,
|
||||||
|
pub pending_dc_total: usize,
|
||||||
pub by_dc: Vec<MeApiRefillDcSnapshot>,
|
pub by_dc: Vec<MeApiRefillDcSnapshot>,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -56,14 +58,21 @@ pub(crate) struct MeApiDrainGateSnapshot {
|
|||||||
|
|
||||||
impl MePool {
|
impl MePool {
|
||||||
pub(crate) async fn api_refill_snapshot(&self) -> MeApiRefillSnapshot {
|
pub(crate) async fn api_refill_snapshot(&self) -> MeApiRefillSnapshot {
|
||||||
let inflight_endpoints_total = self.refill_inflight.lock().await.len();
|
let refill_states = self.refill_states.lock();
|
||||||
let inflight_dc_keys = self
|
let inflight_endpoints_total = refill_states
|
||||||
.refill_inflight_dc
|
.values()
|
||||||
.lock()
|
.map(|pending| 1usize + usize::from(pending.is_some()))
|
||||||
.await
|
.sum();
|
||||||
.iter()
|
let running_dc_total = refill_states.len();
|
||||||
|
let pending_dc_total = refill_states
|
||||||
|
.values()
|
||||||
|
.filter(|pending| pending.is_some())
|
||||||
|
.count();
|
||||||
|
let inflight_dc_keys = refill_states
|
||||||
|
.keys()
|
||||||
.copied()
|
.copied()
|
||||||
.collect::<Vec<RefillDcKey>>();
|
.collect::<Vec<RefillDcKey>>();
|
||||||
|
drop(refill_states);
|
||||||
|
|
||||||
let mut by_dc_map = HashMap::<(i16, &'static str), usize>::new();
|
let mut by_dc_map = HashMap::<(i16, &'static str), usize>::new();
|
||||||
for key in inflight_dc_keys {
|
for key in inflight_dc_keys {
|
||||||
@@ -88,6 +97,8 @@ impl MePool {
|
|||||||
MeApiRefillSnapshot {
|
MeApiRefillSnapshot {
|
||||||
inflight_endpoints_total,
|
inflight_endpoints_total,
|
||||||
inflight_dc_total: by_dc.len(),
|
inflight_dc_total: by_dc.len(),
|
||||||
|
running_dc_total,
|
||||||
|
pending_dc_total,
|
||||||
by_dc,
|
by_dc,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,9 +1,10 @@
|
|||||||
use std::collections::{BTreeMap, BTreeSet, HashMap};
|
use std::collections::{BTreeMap, BTreeSet, HashMap};
|
||||||
use std::net::{IpAddr, SocketAddr};
|
use std::net::{IpAddr, SocketAddr};
|
||||||
|
use std::sync::Arc;
|
||||||
use std::sync::atomic::Ordering;
|
use std::sync::atomic::Ordering;
|
||||||
use std::time::Instant;
|
use std::time::Instant;
|
||||||
|
|
||||||
use super::pool::{MePool, WriterContour};
|
use super::pool::{MePool, ReinitStatusSnapshot, WriterContour};
|
||||||
use crate::config::{MeBindStaleMode, MeFloorMode, MeSocksKdfPolicy};
|
use crate::config::{MeBindStaleMode, MeFloorMode, MeSocksKdfPolicy};
|
||||||
use crate::transport::upstream::IpPreference;
|
use crate::transport::upstream::IpPreference;
|
||||||
|
|
||||||
@@ -87,8 +88,11 @@ pub(crate) struct MeApiDcPathSnapshot {
|
|||||||
pub(crate) struct MeApiRuntimeSnapshot {
|
pub(crate) struct MeApiRuntimeSnapshot {
|
||||||
pub active_generation: u64,
|
pub active_generation: u64,
|
||||||
pub warm_generation: u64,
|
pub warm_generation: u64,
|
||||||
|
pub warm_generations: Vec<u64>,
|
||||||
pub pending_hardswap_generation: u64,
|
pub pending_hardswap_generation: u64,
|
||||||
pub pending_hardswap_age_secs: Option<u64>,
|
pub pending_hardswap_age_secs: Option<u64>,
|
||||||
|
pub reinit_inflight: usize,
|
||||||
|
pub reinit_max_concurrency_effective: usize,
|
||||||
pub hardswap_enabled: bool,
|
pub hardswap_enabled: bool,
|
||||||
pub floor_mode: &'static str,
|
pub floor_mode: &'static str,
|
||||||
pub adaptive_floor_idle_secs: u64,
|
pub adaptive_floor_idle_secs: u64,
|
||||||
@@ -222,8 +226,16 @@ impl MePool {
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) async fn api_status_snapshot(&self) -> MeApiStatusSnapshot {
|
pub(crate) async fn api_status_snapshot(&self) -> MeApiStatusSnapshot {
|
||||||
|
let reinit = self.reinit.status.load_full();
|
||||||
|
self.api_status_snapshot_for_reinit(reinit.as_ref()).await
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn api_status_snapshot_for_reinit(
|
||||||
|
&self,
|
||||||
|
reinit: &ReinitStatusSnapshot,
|
||||||
|
) -> MeApiStatusSnapshot {
|
||||||
let now_epoch_secs = Self::now_epoch_secs();
|
let now_epoch_secs = Self::now_epoch_secs();
|
||||||
let active_generation = self.current_generation();
|
let active_generation = reinit.active_generation;
|
||||||
let drain_ttl_secs = self
|
let drain_ttl_secs = self
|
||||||
.drain_runtime
|
.drain_runtime
|
||||||
.me_pool_drain_ttl_secs
|
.me_pool_drain_ttl_secs
|
||||||
@@ -440,13 +452,19 @@ impl MePool {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[allow(dead_code)]
|
||||||
pub(crate) async fn api_runtime_snapshot(&self) -> MeApiRuntimeSnapshot {
|
pub(crate) async fn api_runtime_snapshot(&self) -> MeApiRuntimeSnapshot {
|
||||||
|
let reinit = self.reinit.status.load_full();
|
||||||
|
self.api_runtime_snapshot_for_reinit(reinit.as_ref()).await
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn api_runtime_snapshot_for_reinit(
|
||||||
|
&self,
|
||||||
|
reinit: &ReinitStatusSnapshot,
|
||||||
|
) -> MeApiRuntimeSnapshot {
|
||||||
let now = Instant::now();
|
let now = Instant::now();
|
||||||
let now_epoch_secs = Self::now_epoch_secs();
|
let now_epoch_secs = Self::now_epoch_secs();
|
||||||
let pending_started_at = self
|
let pending_started_at = reinit.pending_hardswap_started_at_epoch_secs;
|
||||||
.reinit
|
|
||||||
.pending_hardswap_started_at_epoch_secs
|
|
||||||
.load(Ordering::Relaxed);
|
|
||||||
let pending_hardswap_age_secs =
|
let pending_hardswap_age_secs =
|
||||||
(pending_started_at > 0).then_some(now_epoch_secs.saturating_sub(pending_started_at));
|
(pending_started_at > 0).then_some(now_epoch_secs.saturating_sub(pending_started_at));
|
||||||
|
|
||||||
@@ -486,13 +504,16 @@ impl MePool {
|
|||||||
}
|
}
|
||||||
|
|
||||||
MeApiRuntimeSnapshot {
|
MeApiRuntimeSnapshot {
|
||||||
active_generation: self.reinit.active_generation.load(Ordering::Relaxed),
|
active_generation: reinit.active_generation,
|
||||||
warm_generation: self.reinit.warm_generation.load(Ordering::Relaxed),
|
warm_generation: reinit.warm_generations.last().copied().unwrap_or(0),
|
||||||
pending_hardswap_generation: self
|
warm_generations: reinit.warm_generations.clone(),
|
||||||
.reinit
|
pending_hardswap_generation: reinit.pending_hardswap_generation,
|
||||||
.pending_hardswap_generation
|
|
||||||
.load(Ordering::Relaxed),
|
|
||||||
pending_hardswap_age_secs,
|
pending_hardswap_age_secs,
|
||||||
|
reinit_inflight: reinit.inflight,
|
||||||
|
reinit_max_concurrency_effective: self
|
||||||
|
.reinit
|
||||||
|
.max_concurrency_effective
|
||||||
|
.load(Ordering::Acquire),
|
||||||
hardswap_enabled: self.reinit.hardswap.load(Ordering::Relaxed),
|
hardswap_enabled: self.reinit.hardswap.load(Ordering::Relaxed),
|
||||||
floor_mode: floor_mode_label(self.floor_mode()),
|
floor_mode: floor_mode_label(self.floor_mode()),
|
||||||
adaptive_floor_idle_secs: self
|
adaptive_floor_idle_secs: self
|
||||||
@@ -662,6 +683,23 @@ impl MePool {
|
|||||||
network_path,
|
network_path,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub(crate) async fn api_coherent_snapshots(
|
||||||
|
&self,
|
||||||
|
) -> (MeApiStatusSnapshot, MeApiRuntimeSnapshot) {
|
||||||
|
let mut attempts = 0usize;
|
||||||
|
loop {
|
||||||
|
let reinit = self.reinit.status.load_full();
|
||||||
|
let status = self.api_status_snapshot_for_reinit(reinit.as_ref()).await;
|
||||||
|
let runtime = self
|
||||||
|
.api_runtime_snapshot_for_reinit(reinit.as_ref())
|
||||||
|
.await;
|
||||||
|
attempts += 1;
|
||||||
|
if Arc::ptr_eq(&reinit, &self.reinit.status.load_full()) || attempts >= 3 {
|
||||||
|
return (status, runtime);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn ratio_pct(part: usize, total: usize) -> f64 {
|
fn ratio_pct(part: usize, total: usize) -> f64 {
|
||||||
|
|||||||
@@ -269,7 +269,10 @@ async fn rpc_proxy_req_signal_loop(
|
|||||||
continue;
|
continue;
|
||||||
};
|
};
|
||||||
|
|
||||||
let (conn_id, mut service_rx) = pool.registry.register().await;
|
let Some((conn_lease, mut service_rx)) = pool.register_connection().await else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
let conn_id = conn_lease.conn_id();
|
||||||
// Service RPC_PROXY_REQ signal path is intentionally route-only:
|
// Service RPC_PROXY_REQ signal path is intentionally route-only:
|
||||||
// do not bind synthetic conn_id into regular writer/client accounting.
|
// do not bind synthetic conn_id into regular writer/client accounting.
|
||||||
|
|
||||||
@@ -286,7 +289,7 @@ async fn rpc_proxy_req_signal_loop(
|
|||||||
send_service_writer_command(&tx_signal, WriterCommand::DataAndFlush(payload)).await
|
send_service_writer_command(&tx_signal, WriterCommand::DataAndFlush(payload)).await
|
||||||
{
|
{
|
||||||
stats_signal.increment_me_rpc_proxy_req_signal_failed_total();
|
stats_signal.increment_me_rpc_proxy_req_signal_failed_total();
|
||||||
let _ = pool.registry.unregister(conn_id).await;
|
conn_lease.unregister().await;
|
||||||
match error {
|
match error {
|
||||||
ServiceWriterCommandSendError::Closed => return,
|
ServiceWriterCommandSendError::Closed => return,
|
||||||
ServiceWriterCommandSendError::TimedOut => continue,
|
ServiceWriterCommandSendError::TimedOut => continue,
|
||||||
@@ -313,7 +316,7 @@ async fn rpc_proxy_req_signal_loop(
|
|||||||
.await
|
.await
|
||||||
{
|
{
|
||||||
stats_signal.increment_me_rpc_proxy_req_signal_failed_total();
|
stats_signal.increment_me_rpc_proxy_req_signal_failed_total();
|
||||||
let _ = pool.registry.unregister(conn_id).await;
|
conn_lease.unregister().await;
|
||||||
match error {
|
match error {
|
||||||
ServiceWriterCommandSendError::Closed => return,
|
ServiceWriterCommandSendError::Closed => return,
|
||||||
ServiceWriterCommandSendError::TimedOut => continue,
|
ServiceWriterCommandSendError::TimedOut => continue,
|
||||||
@@ -321,7 +324,7 @@ async fn rpc_proxy_req_signal_loop(
|
|||||||
}
|
}
|
||||||
|
|
||||||
stats_signal.increment_me_rpc_proxy_req_signal_close_sent_total();
|
stats_signal.increment_me_rpc_proxy_req_signal_close_sent_total();
|
||||||
let _ = pool.registry.unregister(conn_id).await;
|
conn_lease.unregister().await;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -394,14 +397,14 @@ impl MePool {
|
|||||||
writer_dc: i32,
|
writer_dc: i32,
|
||||||
allow_coverage_override: bool,
|
allow_coverage_override: bool,
|
||||||
) -> Result<()> {
|
) -> Result<()> {
|
||||||
if !self
|
let Some(_writer_open_reservation) = self
|
||||||
.can_open_writer_for_contour(contour, allow_coverage_override, writer_dc)
|
.reserve_writer_open(contour, allow_coverage_override, writer_dc)
|
||||||
.await
|
.await
|
||||||
{
|
else {
|
||||||
return Err(ProxyError::Proxy(format!(
|
return Err(ProxyError::Proxy(format!(
|
||||||
"ME {contour:?} writer cap reached"
|
"ME {contour:?} writer cap reached"
|
||||||
)));
|
)));
|
||||||
}
|
};
|
||||||
|
|
||||||
let secret_len = self.proxy_secret.read().await.secret.len();
|
let secret_len = self.proxy_secret.read().await.secret.len();
|
||||||
if secret_len < 32 {
|
if secret_len < 32 {
|
||||||
@@ -415,6 +418,9 @@ impl MePool {
|
|||||||
let hs = self
|
let hs = self
|
||||||
.handshake_only(stream, addr, upstream_egress, rng)
|
.handshake_only(stream, addr, upstream_egress, rng)
|
||||||
.await?;
|
.await?;
|
||||||
|
let Some(task_registration) = self.lifecycle.try_register() else {
|
||||||
|
return Err(ProxyError::Proxy("ME pool lifecycle closed".into()));
|
||||||
|
};
|
||||||
|
|
||||||
let writer_id = self.next_writer_id.fetch_add(1, Ordering::Relaxed);
|
let writer_id = self.next_writer_id.fetch_add(1, Ordering::Relaxed);
|
||||||
let contour = Arc::new(AtomicU8::new(contour.as_u8()));
|
let contour = Arc::new(AtomicU8::new(contour.as_u8()));
|
||||||
@@ -499,7 +505,7 @@ impl MePool {
|
|||||||
let route_fairshare_enabled = self.transport_policy.me_route_fairshare_enabled.clone();
|
let route_fairshare_enabled = self.transport_policy.me_route_fairshare_enabled.clone();
|
||||||
let reader_route_data_wait_ms = self.transport_policy.me_reader_route_data_wait_ms.clone();
|
let reader_route_data_wait_ms = self.transport_policy.me_reader_route_data_wait_ms.clone();
|
||||||
|
|
||||||
tokio::spawn(async move {
|
self.lifecycle.spawn_registered_writer(task_registration, async move {
|
||||||
// Reader MUST be the first branch in biased select! to avoid read starvation.
|
// Reader MUST be the first branch in biased select! to avoid read starvation.
|
||||||
let exit = tokio::select! {
|
let exit = tokio::select! {
|
||||||
biased;
|
biased;
|
||||||
|
|||||||
@@ -1,6 +1,7 @@
|
|||||||
use std::collections::{HashMap, HashSet};
|
use std::collections::{HashMap, HashSet};
|
||||||
use std::net::SocketAddr;
|
use std::net::SocketAddr;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
|
use std::sync::Mutex as StdMutex;
|
||||||
use std::sync::atomic::{AtomicU8, AtomicU64, Ordering};
|
use std::sync::atomic::{AtomicU8, AtomicU64, Ordering};
|
||||||
use std::time::{SystemTime, UNIX_EPOCH};
|
use std::time::{SystemTime, UNIX_EPOCH};
|
||||||
|
|
||||||
@@ -118,6 +119,38 @@ pub struct ConnRegistry {
|
|||||||
route_backpressure_high_timeout_ms: AtomicU64,
|
route_backpressure_high_timeout_ms: AtomicU64,
|
||||||
route_backpressure_high_watermark_pct: AtomicU8,
|
route_backpressure_high_watermark_pct: AtomicU8,
|
||||||
route_byte_permits_per_conn: usize,
|
route_byte_permits_per_conn: usize,
|
||||||
|
cleanup_tx: mpsc::Sender<u64>,
|
||||||
|
cleanup_rx: StdMutex<Option<mpsc::Receiver<u64>>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Cancellation-safe ownership of one registered ME client route.
|
||||||
|
pub(crate) struct ConnLease {
|
||||||
|
registry: Arc<ConnRegistry>,
|
||||||
|
conn_id: u64,
|
||||||
|
cleanup_permit: Option<mpsc::OwnedPermit<u64>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ConnLease {
|
||||||
|
/// Returns the stable connection identifier owned by this lease.
|
||||||
|
pub(crate) fn conn_id(&self) -> u64 {
|
||||||
|
self.conn_id
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Completes asynchronous registry cleanup and disarms drop cleanup.
|
||||||
|
pub(crate) async fn unregister(mut self) {
|
||||||
|
self.registry.unregister(self.conn_id).await;
|
||||||
|
self.cleanup_permit.take();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Drop for ConnLease {
|
||||||
|
fn drop(&mut self) {
|
||||||
|
let Some(permit) = self.cleanup_permit.take() else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
self.registry.remove_route_now(self.conn_id);
|
||||||
|
permit.send(self.conn_id);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
impl ConnRegistry {
|
impl ConnRegistry {
|
||||||
@@ -133,15 +166,30 @@ impl ConnRegistry {
|
|||||||
Self::with_route_limits(
|
Self::with_route_limits(
|
||||||
route_channel_capacity,
|
route_channel_capacity,
|
||||||
Self::route_byte_permit_budget(route_channel_capacity),
|
Self::route_byte_permit_budget(route_channel_capacity),
|
||||||
|
16_384,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) fn with_route_and_cleanup_capacity(
|
||||||
|
route_channel_capacity: usize,
|
||||||
|
connection_cleanup_capacity: usize,
|
||||||
|
) -> Self {
|
||||||
|
let route_channel_capacity = route_channel_capacity.max(1);
|
||||||
|
Self::with_route_limits(
|
||||||
|
route_channel_capacity,
|
||||||
|
Self::route_byte_permit_budget(route_channel_capacity),
|
||||||
|
connection_cleanup_capacity,
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
fn with_route_limits(
|
fn with_route_limits(
|
||||||
route_channel_capacity: usize,
|
route_channel_capacity: usize,
|
||||||
route_byte_permits_per_conn: usize,
|
route_byte_permits_per_conn: usize,
|
||||||
|
connection_cleanup_capacity: usize,
|
||||||
) -> Self {
|
) -> Self {
|
||||||
let start = rand::random::<u64>() | 1;
|
let start = rand::random::<u64>() | 1;
|
||||||
let route_channel_capacity = route_channel_capacity.max(1);
|
let route_channel_capacity = route_channel_capacity.max(1);
|
||||||
|
let (cleanup_tx, cleanup_rx) = mpsc::channel(connection_cleanup_capacity.max(1));
|
||||||
Self {
|
Self {
|
||||||
routing: RoutingTable {
|
routing: RoutingTable {
|
||||||
map: DashMap::new(),
|
map: DashMap::new(),
|
||||||
@@ -168,6 +216,8 @@ impl ConnRegistry {
|
|||||||
ROUTE_BACKPRESSURE_HIGH_WATERMARK_PCT,
|
ROUTE_BACKPRESSURE_HIGH_WATERMARK_PCT,
|
||||||
),
|
),
|
||||||
route_byte_permits_per_conn: route_byte_permits_per_conn.max(1),
|
route_byte_permits_per_conn: route_byte_permits_per_conn.max(1),
|
||||||
|
cleanup_tx,
|
||||||
|
cleanup_rx: StdMutex::new(Some(cleanup_rx)),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -199,7 +249,7 @@ impl ConnRegistry {
|
|||||||
route_channel_capacity: usize,
|
route_channel_capacity: usize,
|
||||||
route_byte_permits_per_conn: usize,
|
route_byte_permits_per_conn: usize,
|
||||||
) -> Self {
|
) -> Self {
|
||||||
Self::with_route_limits(route_channel_capacity, route_byte_permits_per_conn)
|
Self::with_route_limits(route_channel_capacity, route_byte_permits_per_conn, 16_384)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn update_route_backpressure_policy(
|
pub fn update_route_backpressure_policy(
|
||||||
@@ -220,6 +270,29 @@ impl ConnRegistry {
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub async fn register(&self) -> (u64, mpsc::Receiver<MeResponse>) {
|
pub async fn register(&self) -> (u64, mpsc::Receiver<MeResponse>) {
|
||||||
|
self.register_route()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) async fn register_leased(
|
||||||
|
self: &Arc<Self>,
|
||||||
|
) -> Option<(ConnLease, mpsc::Receiver<MeResponse>)> {
|
||||||
|
let cleanup_permit = self.cleanup_tx.clone().reserve_owned().await.ok()?;
|
||||||
|
let (conn_id, rx) = self.register_route();
|
||||||
|
Some((
|
||||||
|
ConnLease {
|
||||||
|
registry: Arc::clone(self),
|
||||||
|
conn_id,
|
||||||
|
cleanup_permit: Some(cleanup_permit),
|
||||||
|
},
|
||||||
|
rx,
|
||||||
|
))
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn take_cleanup_receiver(&self) -> Option<mpsc::Receiver<u64>> {
|
||||||
|
self.cleanup_rx.lock().ok()?.take()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn register_route(&self) -> (u64, mpsc::Receiver<MeResponse>) {
|
||||||
let id = self.next_id.fetch_add(1, Ordering::Relaxed);
|
let id = self.next_id.fetch_add(1, Ordering::Relaxed);
|
||||||
let (tx, rx) = mpsc::channel(self.route_channel_capacity);
|
let (tx, rx) = mpsc::channel(self.route_channel_capacity);
|
||||||
self.routing.map.insert(id, tx);
|
self.routing.map.insert(id, tx);
|
||||||
@@ -229,6 +302,12 @@ impl ConnRegistry {
|
|||||||
);
|
);
|
||||||
(id, rx)
|
(id, rx)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn remove_route_now(&self, id: u64) {
|
||||||
|
self.routing.map.remove(&id);
|
||||||
|
self.routing.byte_budget.remove(&id);
|
||||||
|
self.hot_binding.map.remove(&id);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
|
|||||||
@@ -308,3 +308,21 @@ async fn non_empty_writer_ids_returns_only_writers_with_bound_clients() {
|
|||||||
assert!(!non_empty.contains(&20));
|
assert!(!non_empty.contains(&20));
|
||||||
assert!(!non_empty.contains(&30));
|
assert!(!non_empty.contains(&30));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn leased_registration_removes_hot_route_before_async_cleanup() {
|
||||||
|
let registry = Arc::new(ConnRegistry::with_route_and_cleanup_capacity(8, 1));
|
||||||
|
let mut cleanup_rx = registry.take_cleanup_receiver().unwrap();
|
||||||
|
let (lease, _rx) = registry.register_leased().await.unwrap();
|
||||||
|
let conn_id = lease.conn_id();
|
||||||
|
|
||||||
|
drop(lease);
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
registry.route_nowait(conn_id, MeResponse::Ack(1)).await,
|
||||||
|
RouteResult::NoConn
|
||||||
|
);
|
||||||
|
assert_eq!(cleanup_rx.recv().await, Some(conn_id));
|
||||||
|
registry.unregister(conn_id).await;
|
||||||
|
assert!(registry.active_conn_ids().await.is_empty());
|
||||||
|
}
|
||||||
|
|||||||
@@ -2,6 +2,7 @@ use std::sync::Arc;
|
|||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
|
|
||||||
use tokio::sync::{mpsc, watch};
|
use tokio::sync::{mpsc, watch};
|
||||||
|
use tokio::task::JoinSet;
|
||||||
use tracing::{debug, info, warn};
|
use tracing::{debug, info, warn};
|
||||||
|
|
||||||
use crate::config::ProxyConfig;
|
use crate::config::ProxyConfig;
|
||||||
@@ -42,6 +43,45 @@ pub fn enqueue_reinit_trigger(tx: &mpsc::Sender<MeReinitTrigger>, trigger: MeRei
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
const REINIT_TRIGGER_PERIODIC: u8 = 1;
|
||||||
|
const REINIT_TRIGGER_MAP_CHANGED: u8 = 2;
|
||||||
|
|
||||||
|
struct ReinitInflightGuard {
|
||||||
|
pool: Arc<MePool>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Drop for ReinitInflightGuard {
|
||||||
|
fn drop(&mut self) {
|
||||||
|
self.pool
|
||||||
|
.reinit
|
||||||
|
.scheduler_inflight
|
||||||
|
.fetch_sub(1, std::sync::atomic::Ordering::AcqRel);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn effective_reinit_concurrency(config: &ProxyConfig) -> usize {
|
||||||
|
if config.general.me_reinit_singleflight {
|
||||||
|
1
|
||||||
|
} else {
|
||||||
|
config.general.me_reinit_max_concurrency.clamp(1, 8)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn trigger_bit(trigger: MeReinitTrigger) -> u8 {
|
||||||
|
match trigger {
|
||||||
|
MeReinitTrigger::Periodic => REINIT_TRIGGER_PERIODIC,
|
||||||
|
MeReinitTrigger::MapChanged => REINIT_TRIGGER_MAP_CHANGED,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn trigger_reason(pending: u8) -> &'static str {
|
||||||
|
match pending {
|
||||||
|
REINIT_TRIGGER_PERIODIC => "periodic",
|
||||||
|
REINIT_TRIGGER_MAP_CHANGED => "map-change",
|
||||||
|
_ => "map-change+periodic",
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
pub async fn me_reinit_scheduler(
|
pub async fn me_reinit_scheduler(
|
||||||
pool: Arc<MePool>,
|
pool: Arc<MePool>,
|
||||||
rng: Arc<SecureRandom>,
|
rng: Arc<SecureRandom>,
|
||||||
@@ -50,68 +90,104 @@ pub async fn me_reinit_scheduler(
|
|||||||
me_ready_tx: watch::Sender<u64>,
|
me_ready_tx: watch::Sender<u64>,
|
||||||
) {
|
) {
|
||||||
info!("ME reinit scheduler started");
|
info!("ME reinit scheduler started");
|
||||||
|
let mut tasks = JoinSet::<bool>::new();
|
||||||
|
let mut pending = 0u8;
|
||||||
|
let mut pending_deadline = None;
|
||||||
|
let mut trigger_channel_open = true;
|
||||||
|
|
||||||
loop {
|
loop {
|
||||||
let Some(first_trigger) = trigger_rx.recv().await else {
|
|
||||||
warn!("ME reinit scheduler stopped: trigger channel closed");
|
|
||||||
break;
|
|
||||||
};
|
|
||||||
|
|
||||||
let mut map_change_seen = matches!(first_trigger, MeReinitTrigger::MapChanged);
|
|
||||||
let mut periodic_seen = matches!(first_trigger, MeReinitTrigger::Periodic);
|
|
||||||
let cfg = config_rx.borrow().clone();
|
let cfg = config_rx.borrow().clone();
|
||||||
let coalesce_window = Duration::from_millis(cfg.general.me_reinit_coalesce_window_ms);
|
let max_concurrency = effective_reinit_concurrency(&cfg);
|
||||||
if !coalesce_window.is_zero() {
|
pool.reinit
|
||||||
let deadline = tokio::time::Instant::now() + coalesce_window;
|
.max_concurrency_effective
|
||||||
loop {
|
.store(max_concurrency, std::sync::atomic::Ordering::Release);
|
||||||
let now = tokio::time::Instant::now();
|
|
||||||
if now >= deadline {
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
match tokio::time::timeout(deadline - now, trigger_rx.recv()).await {
|
|
||||||
Ok(Some(next)) => {
|
|
||||||
if next == MeReinitTrigger::MapChanged {
|
|
||||||
map_change_seen = true;
|
|
||||||
} else {
|
|
||||||
periodic_seen = true;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
Ok(None) => break,
|
|
||||||
Err(_) => break,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
let reason = if map_change_seen && periodic_seen {
|
let pending_ready = pending != 0
|
||||||
"map-change+periodic"
|
&& pending_deadline.is_none_or(|deadline| deadline <= tokio::time::Instant::now());
|
||||||
} else if map_change_seen {
|
if pending_ready && tasks.len() < max_concurrency {
|
||||||
"map-change"
|
let reason = trigger_reason(pending);
|
||||||
} else {
|
pending = 0;
|
||||||
"periodic"
|
pending_deadline = None;
|
||||||
};
|
debug!(reason, max_concurrency, "ME reinit scheduled");
|
||||||
|
|
||||||
if cfg.general.me_reinit_singleflight {
|
|
||||||
debug!(reason, "ME reinit scheduled (single-flight)");
|
|
||||||
if pool.zero_downtime_reinit_periodic(rng.as_ref()).await {
|
|
||||||
me_ready_tx.send_modify(|version| {
|
|
||||||
*version = version.saturating_add(1);
|
|
||||||
});
|
|
||||||
}
|
|
||||||
} else {
|
|
||||||
debug!(reason, "ME reinit scheduled (concurrent mode)");
|
|
||||||
let pool_clone = pool.clone();
|
let pool_clone = pool.clone();
|
||||||
let rng_clone = rng.clone();
|
let rng_clone = rng.clone();
|
||||||
let me_ready_tx_clone = me_ready_tx.clone();
|
pool.reinit
|
||||||
tokio::spawn(async move {
|
.scheduler_inflight
|
||||||
if pool_clone
|
.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
|
||||||
|
let inflight = ReinitInflightGuard {
|
||||||
|
pool: Arc::clone(&pool_clone),
|
||||||
|
};
|
||||||
|
tasks.spawn(async move {
|
||||||
|
let _inflight = inflight;
|
||||||
|
pool_clone
|
||||||
.zero_downtime_reinit_periodic(rng_clone.as_ref())
|
.zero_downtime_reinit_periodic(rng_clone.as_ref())
|
||||||
.await
|
.await
|
||||||
{
|
|
||||||
me_ready_tx_clone.send_modify(|version| {
|
|
||||||
*version = version.saturating_add(1);
|
|
||||||
});
|
|
||||||
}
|
|
||||||
});
|
});
|
||||||
|
continue;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if !trigger_channel_open && pending == 0 && tasks.is_empty() {
|
||||||
|
warn!("ME reinit scheduler stopped: trigger channel closed");
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
let timer_enabled = pending != 0 && tasks.len() < max_concurrency;
|
||||||
|
tokio::select! {
|
||||||
|
trigger = trigger_rx.recv(), if trigger_channel_open => {
|
||||||
|
match trigger {
|
||||||
|
Some(trigger) => {
|
||||||
|
if pending == 0 {
|
||||||
|
pending_deadline = Some(
|
||||||
|
tokio::time::Instant::now()
|
||||||
|
+ Duration::from_millis(
|
||||||
|
cfg.general.me_reinit_coalesce_window_ms,
|
||||||
|
),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
pending |= trigger_bit(trigger);
|
||||||
|
}
|
||||||
|
None => trigger_channel_open = false,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
joined = tasks.join_next(), if !tasks.is_empty() => {
|
||||||
|
match joined {
|
||||||
|
Some(Ok(true)) => {
|
||||||
|
me_ready_tx.send_modify(|version| {
|
||||||
|
*version = version.saturating_add(1);
|
||||||
|
});
|
||||||
|
}
|
||||||
|
Some(Ok(false)) => {}
|
||||||
|
Some(Err(error)) => {
|
||||||
|
warn!(error = %error, "ME reinit task failed");
|
||||||
|
}
|
||||||
|
None => {}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_ = async {
|
||||||
|
if let Some(deadline) = pending_deadline {
|
||||||
|
tokio::time::sleep_until(deadline).await;
|
||||||
|
}
|
||||||
|
}, if timer_enabled => {}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use crate::config::ProxyConfig;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn effective_concurrency_is_one_for_singleflight_and_bounded_otherwise() {
|
||||||
|
let mut config = ProxyConfig::default();
|
||||||
|
config.general.me_reinit_singleflight = true;
|
||||||
|
config.general.me_reinit_max_concurrency = 8;
|
||||||
|
assert_eq!(super::effective_reinit_concurrency(&config), 1);
|
||||||
|
|
||||||
|
config.general.me_reinit_singleflight = false;
|
||||||
|
config.general.me_reinit_max_concurrency = 2;
|
||||||
|
assert_eq!(super::effective_reinit_concurrency(&config), 2);
|
||||||
|
config.general.me_reinit_max_concurrency = usize::MAX;
|
||||||
|
assert_eq!(super::effective_reinit_concurrency(&config), 8);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -108,6 +108,7 @@ pub async fn download_proxy_secret_with_max_len_via_upstream(
|
|||||||
let resp = https_get(
|
let resp = https_get(
|
||||||
proxy_secret_url.unwrap_or("https://core.telegram.org/getProxySecret"),
|
proxy_secret_url.unwrap_or("https://core.telegram.org/getProxySecret"),
|
||||||
upstream,
|
upstream,
|
||||||
|
max_len,
|
||||||
)
|
)
|
||||||
.await?;
|
.await?;
|
||||||
|
|
||||||
|
|||||||
@@ -122,6 +122,7 @@ async fn make_pool(
|
|||||||
general.me_route_blocking_send_timeout_ms,
|
general.me_route_blocking_send_timeout_ms,
|
||||||
general.me_route_inline_recovery_attempts,
|
general.me_route_inline_recovery_attempts,
|
||||||
general.me_route_inline_recovery_wait_ms,
|
general.me_route_inline_recovery_wait_ms,
|
||||||
|
16_384,
|
||||||
);
|
);
|
||||||
|
|
||||||
(pool, rng)
|
(pool, rng)
|
||||||
|
|||||||
@@ -120,6 +120,7 @@ async fn make_pool(
|
|||||||
general.me_route_blocking_send_timeout_ms,
|
general.me_route_blocking_send_timeout_ms,
|
||||||
general.me_route_inline_recovery_attempts,
|
general.me_route_inline_recovery_attempts,
|
||||||
general.me_route_inline_recovery_wait_ms,
|
general.me_route_inline_recovery_wait_ms,
|
||||||
|
16_384,
|
||||||
);
|
);
|
||||||
(pool, rng)
|
(pool, rng)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -115,6 +115,7 @@ async fn make_pool(me_pool_drain_threshold: u64) -> Arc<MePool> {
|
|||||||
general.me_route_blocking_send_timeout_ms,
|
general.me_route_blocking_send_timeout_ms,
|
||||||
general.me_route_inline_recovery_attempts,
|
general.me_route_inline_recovery_attempts,
|
||||||
general.me_route_inline_recovery_wait_ms,
|
general.me_route_inline_recovery_wait_ms,
|
||||||
|
16_384,
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -1,6 +1,7 @@
|
|||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
|
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
|
use std::sync::atomic::Ordering;
|
||||||
use std::time::{Duration, Instant};
|
use std::time::{Duration, Instant};
|
||||||
|
|
||||||
use crate::config::{GeneralConfig, MeRouteNoWriterMode, MeSocksKdfPolicy, MeWriterPickMode};
|
use crate::config::{GeneralConfig, MeRouteNoWriterMode, MeSocksKdfPolicy, MeWriterPickMode};
|
||||||
@@ -104,6 +105,7 @@ async fn make_pool() -> Arc<MePool> {
|
|||||||
general.me_route_blocking_send_timeout_ms,
|
general.me_route_blocking_send_timeout_ms,
|
||||||
general.me_route_inline_recovery_attempts,
|
general.me_route_inline_recovery_attempts,
|
||||||
general.me_route_inline_recovery_wait_ms,
|
general.me_route_inline_recovery_wait_ms,
|
||||||
|
16_384,
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -159,3 +161,31 @@ async fn connectable_endpoints_releases_quarantine_lock_before_sleep() {
|
|||||||
.expect("task join failed");
|
.expect("task join failed");
|
||||||
assert_eq!(endpoints, vec![addr]);
|
assert_eq!(endpoints, vec![addr]);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test(flavor = "current_thread")]
|
||||||
|
async fn refill_coalesces_one_pending_endpoint_and_cleans_up_before_first_poll() {
|
||||||
|
let pool = make_pool().await;
|
||||||
|
let first = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 31, 0, 21)), 443);
|
||||||
|
let second = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 31, 0, 22)), 443);
|
||||||
|
let latest = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 31, 0, 23)), 443);
|
||||||
|
|
||||||
|
pool.trigger_immediate_refill_for_dc(first, 2);
|
||||||
|
pool.trigger_immediate_refill_for_dc(second, 2);
|
||||||
|
pool.trigger_immediate_refill_for_dc(latest, 2);
|
||||||
|
|
||||||
|
assert_eq!(pool.refill_states.lock().len(), 1);
|
||||||
|
assert_eq!(pool.refill_running.load(Ordering::Acquire), 1);
|
||||||
|
assert_eq!(pool.refill_pending.load(Ordering::Acquire), 1);
|
||||||
|
|
||||||
|
pool.begin_shutdown();
|
||||||
|
tokio::time::timeout(Duration::from_secs(1), async {
|
||||||
|
while pool.refill_running.load(Ordering::Acquire) != 0 {
|
||||||
|
tokio::task::yield_now().await;
|
||||||
|
}
|
||||||
|
})
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
assert!(pool.refill_states.lock().is_empty());
|
||||||
|
assert_eq!(pool.refill_pending.load(Ordering::Acquire), 0);
|
||||||
|
}
|
||||||
|
|||||||
@@ -109,6 +109,7 @@ async fn make_pool() -> Arc<MePool> {
|
|||||||
general.me_route_blocking_send_timeout_ms,
|
general.me_route_blocking_send_timeout_ms,
|
||||||
general.me_route_inline_recovery_attempts,
|
general.me_route_inline_recovery_attempts,
|
||||||
general.me_route_inline_recovery_wait_ms,
|
general.me_route_inline_recovery_wait_ms,
|
||||||
|
16_384,
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -119,6 +119,7 @@ async fn make_pool_with_decision(decision: NetworkDecision) -> (Arc<MePool>, Arc
|
|||||||
general.me_route_blocking_send_timeout_ms,
|
general.me_route_blocking_send_timeout_ms,
|
||||||
general.me_route_inline_recovery_attempts,
|
general.me_route_inline_recovery_attempts,
|
||||||
general.me_route_inline_recovery_wait_ms,
|
general.me_route_inline_recovery_wait_ms,
|
||||||
|
16_384,
|
||||||
);
|
);
|
||||||
|
|
||||||
(pool, rng)
|
(pool, rng)
|
||||||
|
|||||||
+30
-17
@@ -4,7 +4,6 @@
|
|||||||
|
|
||||||
#![allow(deprecated)]
|
#![allow(deprecated)]
|
||||||
|
|
||||||
use arc_swap::ArcSwap;
|
|
||||||
use rand::RngExt;
|
use rand::RngExt;
|
||||||
use std::collections::{BTreeSet, HashMap};
|
use std::collections::{BTreeSet, HashMap};
|
||||||
use std::net::{IpAddr, SocketAddr};
|
use std::net::{IpAddr, SocketAddr};
|
||||||
@@ -21,7 +20,7 @@ use tracing::{debug, info, trace, warn};
|
|||||||
|
|
||||||
use crate::config::{UpstreamConfig, UpstreamType};
|
use crate::config::{UpstreamConfig, UpstreamType};
|
||||||
use crate::error::{ProxyError, Result};
|
use crate::error::{ProxyError, Result};
|
||||||
use crate::network::dns_overrides::{DnsOverrides, split_host_port};
|
use crate::network::dns_overrides::{GenerationDnsResolver, split_host_port};
|
||||||
use crate::protocol::constants::{TG_DATACENTER_PORT, TG_DATACENTERS_V4, TG_DATACENTERS_V6};
|
use crate::protocol::constants::{TG_DATACENTER_PORT, TG_DATACENTERS_V4, TG_DATACENTERS_V6};
|
||||||
use crate::stats::Stats;
|
use crate::stats::Stats;
|
||||||
use crate::transport::shadowsocks::{
|
use crate::transport::shadowsocks::{
|
||||||
@@ -43,6 +42,7 @@ const HEALTH_CHECK_INTERVAL_SECS: u64 = 30;
|
|||||||
const HEALTH_CHECK_CONNECT_TIMEOUT_SECS: u64 = 10;
|
const HEALTH_CHECK_CONNECT_TIMEOUT_SECS: u64 = 10;
|
||||||
/// Upstream is considered healthy when at least this many DC groups are reachable.
|
/// Upstream is considered healthy when at least this many DC groups are reachable.
|
||||||
const MIN_HEALTHY_DC_GROUPS: usize = 3;
|
const MIN_HEALTHY_DC_GROUPS: usize = 3;
|
||||||
|
const DNS_RESULT_MAX_ADDRESSES: usize = 64;
|
||||||
|
|
||||||
// ============= RTT Tracking =============
|
// ============= RTT Tracking =============
|
||||||
|
|
||||||
@@ -334,7 +334,7 @@ pub struct UpstreamManager {
|
|||||||
no_upstreams_warn_epoch_ms: Arc<AtomicU64>,
|
no_upstreams_warn_epoch_ms: Arc<AtomicU64>,
|
||||||
no_healthy_warn_epoch_ms: Arc<AtomicU64>,
|
no_healthy_warn_epoch_ms: Arc<AtomicU64>,
|
||||||
stats: Arc<Stats>,
|
stats: Arc<Stats>,
|
||||||
dns_overrides: Arc<ArcSwap<DnsOverrides>>,
|
dns_resolver: Arc<GenerationDnsResolver>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl UpstreamManager {
|
impl UpstreamManager {
|
||||||
@@ -376,29 +376,42 @@ impl UpstreamManager {
|
|||||||
no_upstreams_warn_epoch_ms: Arc::new(AtomicU64::new(0)),
|
no_upstreams_warn_epoch_ms: Arc::new(AtomicU64::new(0)),
|
||||||
no_healthy_warn_epoch_ms: Arc::new(AtomicU64::new(0)),
|
no_healthy_warn_epoch_ms: Arc::new(AtomicU64::new(0)),
|
||||||
stats,
|
stats,
|
||||||
dns_overrides: Arc::new(ArcSwap::from_pointee(DnsOverrides::default())),
|
dns_resolver: Arc::new(GenerationDnsResolver::default()),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) fn with_dns_overrides(mut self, entries: &[String]) -> Result<Self> {
|
pub(crate) fn with_dns_overrides(self, entries: &[String]) -> Result<Self> {
|
||||||
self.dns_overrides = Arc::new(ArcSwap::from_pointee(DnsOverrides::from_entries(entries)?));
|
self.dns_resolver.apply_entries(entries)?;
|
||||||
Ok(self)
|
Ok(self)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) fn update_dns_overrides(&self, entries: &[String]) -> Result<()> {
|
pub(crate) fn update_dns_overrides(&self, entries: &[String]) -> Result<()> {
|
||||||
let snapshot = DnsOverrides::from_entries(entries)?;
|
self.dns_resolver.apply_entries(entries)
|
||||||
self.dns_overrides.store(Arc::new(snapshot));
|
}
|
||||||
Ok(())
|
|
||||||
|
pub(crate) fn dns_resolver(&self) -> Arc<GenerationDnsResolver> {
|
||||||
|
Arc::clone(&self.dns_resolver)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) async fn resolve_all(&self, host: &str, port: u16) -> Result<Vec<SocketAddr>> {
|
||||||
|
if let Some(addr) = self.dns_resolver.resolve_socket_addr(host, port) {
|
||||||
|
return Ok(vec![addr]);
|
||||||
|
}
|
||||||
|
let addrs = tokio::net::lookup_host((host, port))
|
||||||
|
.await
|
||||||
|
.map_err(ProxyError::Io)?
|
||||||
|
.take(DNS_RESULT_MAX_ADDRESSES)
|
||||||
|
.collect::<Vec<_>>();
|
||||||
|
if addrs.is_empty() {
|
||||||
|
return Err(ProxyError::Proxy(format!(
|
||||||
|
"DNS returned no addresses for {host}:{port}"
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
Ok(addrs)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) async fn resolve_hostname(&self, host: &str, port: u16) -> Result<SocketAddr> {
|
pub(crate) async fn resolve_hostname(&self, host: &str, port: u16) -> Result<SocketAddr> {
|
||||||
if let Some(addr) = self.dns_overrides.load().resolve_socket_addr(host, port) {
|
let addrs = self.resolve_all(host, port).await?;
|
||||||
return Ok(addr);
|
|
||||||
}
|
|
||||||
let addrs: Vec<SocketAddr> = tokio::net::lookup_host((host, port))
|
|
||||||
.await
|
|
||||||
.map_err(ProxyError::Io)?
|
|
||||||
.collect();
|
|
||||||
if let Some(addr) = addrs.iter().copied().find(SocketAddr::is_ipv4) {
|
if let Some(addr) = addrs.iter().copied().find(SocketAddr::is_ipv4) {
|
||||||
return Ok(addr);
|
return Ok(addr);
|
||||||
}
|
}
|
||||||
@@ -744,7 +757,7 @@ impl UpstreamManager {
|
|||||||
connect_timeout: Duration,
|
connect_timeout: Duration,
|
||||||
) -> Result<TcpStream> {
|
) -> Result<TcpStream> {
|
||||||
if let Some((host, port)) = split_host_port(address)
|
if let Some((host, port)) = split_host_port(address)
|
||||||
&& let Some(addr) = self.dns_overrides.load().resolve_socket_addr(&host, port)
|
&& let Some(addr) = self.dns_resolver.resolve_socket_addr(&host, port)
|
||||||
{
|
{
|
||||||
return match tokio::time::timeout(connect_timeout, TcpStream::connect(addr)).await {
|
return match tokio::time::timeout(connect_timeout, TcpStream::connect(addr)).await {
|
||||||
Ok(Ok(stream)) => Ok(stream),
|
Ok(Ok(stream)) => Ok(stream),
|
||||||
|
|||||||
@@ -53,6 +53,8 @@ pub(super) async fn reserve_data(
|
|||||||
}
|
}
|
||||||
let notify = runtime.budget_notify();
|
let notify = runtime.budget_notify();
|
||||||
let notified = notify.notified();
|
let notified = notify.notified();
|
||||||
|
tokio::pin!(notified);
|
||||||
|
notified.as_mut().enable();
|
||||||
if let Some(budget) = runtime.try_websocket_data_budget(owner, bytes.max(1)) {
|
if let Some(budget) = runtime.try_websocket_data_budget(owner, bytes.max(1)) {
|
||||||
return Ok(budget);
|
return Ok(budget);
|
||||||
}
|
}
|
||||||
@@ -104,6 +106,8 @@ pub(super) async fn process_lane(
|
|||||||
}
|
}
|
||||||
let notify = runtime.budget_notify();
|
let notify = runtime.budget_notify();
|
||||||
let notified = notify.notified();
|
let notified = notify.notified();
|
||||||
|
tokio::pin!(notified);
|
||||||
|
notified.as_mut().enable();
|
||||||
match session.process_websocket_lane(reservation, sequence, body) {
|
match session.process_websocket_lane(reservation, sequence, body) {
|
||||||
Ok(progressed) => return Ok(progressed),
|
Ok(progressed) => return Ok(progressed),
|
||||||
Err(ManagerError::Backpressure) => {}
|
Err(ManagerError::Backpressure) => {}
|
||||||
@@ -135,6 +139,8 @@ where
|
|||||||
}
|
}
|
||||||
let notify = runtime.budget_notify();
|
let notify = runtime.budget_notify();
|
||||||
let notified = notify.notified();
|
let notified = notify.notified();
|
||||||
|
tokio::pin!(notified);
|
||||||
|
notified.as_mut().enable();
|
||||||
match operation() {
|
match operation() {
|
||||||
Ok(value) => return Ok(value),
|
Ok(value) => return Ok(value),
|
||||||
Err(ManagerError::Backpressure) => {}
|
Err(ManagerError::Backpressure) => {}
|
||||||
|
|||||||
+32
-12
@@ -6,7 +6,7 @@ use std::time::Duration;
|
|||||||
|
|
||||||
use arc_swap::ArcSwap;
|
use arc_swap::ArcSwap;
|
||||||
use parking_lot::Mutex;
|
use parking_lot::Mutex;
|
||||||
use tokio::sync::{Notify, OwnedSemaphorePermit, Semaphore};
|
use tokio::sync::{Notify, OwnedSemaphorePermit, Semaphore, TryAcquireError};
|
||||||
use tokio_util::sync::CancellationToken;
|
use tokio_util::sync::CancellationToken;
|
||||||
use tokio_util::task::TaskTracker;
|
use tokio_util::task::TaskTracker;
|
||||||
|
|
||||||
@@ -71,6 +71,15 @@ pub(crate) type TokenHash = [u8; TOKEN_BYTES];
|
|||||||
/// Stable non-allocating key used for per-profile quotas.
|
/// Stable non-allocating key used for per-profile quotas.
|
||||||
pub(crate) type ProfileKey = [u8; TOKEN_BYTES];
|
pub(crate) type ProfileKey = [u8; TOKEN_BYTES];
|
||||||
|
|
||||||
|
/// Stable failure category for accepted-socket capacity admission.
|
||||||
|
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||||
|
pub(crate) enum HttpConnectionAdmissionError {
|
||||||
|
/// The bounded connection plane currently has no free permit.
|
||||||
|
AtCapacity,
|
||||||
|
/// Terminal runtime shutdown closed the connection plane.
|
||||||
|
Closed,
|
||||||
|
}
|
||||||
|
|
||||||
/// WEB manager operation failure category.
|
/// WEB manager operation failure category.
|
||||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||||
pub(crate) enum ManagerError {
|
pub(crate) enum ManagerError {
|
||||||
@@ -294,27 +303,38 @@ impl WebProcessRuntime {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Reserves one accepted HTTP connection.
|
/// Reserves one accepted HTTP connection.
|
||||||
pub(crate) fn try_http_connection(&self) -> Option<OwnedSemaphorePermit> {
|
pub(crate) fn try_http_connection(
|
||||||
let permit = Arc::clone(&self.http_connections).try_acquire_owned().ok();
|
&self,
|
||||||
if permit.is_none() {
|
) -> Result<OwnedSemaphorePermit, HttpConnectionAdmissionError> {
|
||||||
self.record_limit_hit();
|
match Arc::clone(&self.http_connections).try_acquire_owned() {
|
||||||
|
Ok(permit) => Ok(permit),
|
||||||
|
Err(TryAcquireError::NoPermits) => {
|
||||||
|
self.record_limit_hit();
|
||||||
|
Err(HttpConnectionAdmissionError::AtCapacity)
|
||||||
|
}
|
||||||
|
Err(TryAcquireError::Closed) => Err(HttpConnectionAdmissionError::Closed),
|
||||||
}
|
}
|
||||||
permit
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Waits for one accepted HTTP connection slot after bounded overload admission.
|
/// Waits for one accepted HTTP connection slot after bounded overload admission.
|
||||||
pub(crate) async fn acquire_http_connection(&self) -> Option<OwnedSemaphorePermit> {
|
pub(crate) async fn acquire_http_connection(
|
||||||
|
&self,
|
||||||
|
) -> Result<OwnedSemaphorePermit, HttpConnectionAdmissionError> {
|
||||||
Arc::clone(&self.http_connections)
|
Arc::clone(&self.http_connections)
|
||||||
.acquire_owned()
|
.acquire_owned()
|
||||||
.await
|
.await
|
||||||
.ok()
|
.map_err(|_| HttpConnectionAdmissionError::Closed)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Reserves one accepted socket outside ordinary HTTP connection capacity.
|
/// Reserves one accepted socket outside ordinary HTTP connection capacity.
|
||||||
pub(crate) fn try_http_overload_connection(&self) -> Option<OwnedSemaphorePermit> {
|
pub(crate) fn try_http_overload_connection(
|
||||||
Arc::clone(&self.http_overload_connections)
|
&self,
|
||||||
.try_acquire_owned()
|
) -> Result<OwnedSemaphorePermit, HttpConnectionAdmissionError> {
|
||||||
.ok()
|
match Arc::clone(&self.http_overload_connections).try_acquire_owned() {
|
||||||
|
Ok(permit) => Ok(permit),
|
||||||
|
Err(TryAcquireError::NoPermits) => Err(HttpConnectionAdmissionError::AtCapacity),
|
||||||
|
Err(TryAcquireError::Closed) => Err(HttpConnectionAdmissionError::Closed),
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Reserves one concurrently executing HTTP request handler.
|
/// Reserves one concurrently executing HTTP request handler.
|
||||||
|
|||||||
@@ -448,7 +448,10 @@ mod tests {
|
|||||||
tokio::task::yield_now().await;
|
tokio::task::yield_now().await;
|
||||||
let drain = runtime.begin_shutdown();
|
let drain = runtime.begin_shutdown();
|
||||||
|
|
||||||
assert!(runtime.try_http_connection().is_none());
|
assert_eq!(
|
||||||
|
runtime.try_http_connection().unwrap_err(),
|
||||||
|
super::super::HttpConnectionAdmissionError::Closed
|
||||||
|
);
|
||||||
assert!(runtime.try_http_handler().is_none());
|
assert!(runtime.try_http_handler().is_none());
|
||||||
assert!(runtime.try_lane_poll(false).is_none());
|
assert!(runtime.try_lane_poll(false).is_none());
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
|
|||||||
@@ -1,5 +1,5 @@
|
|||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
|
use std::sync::atomic::{AtomicU64, Ordering};
|
||||||
use std::time::{Duration, Instant};
|
use std::time::{Duration, Instant};
|
||||||
|
|
||||||
use arc_swap::ArcSwap;
|
use arc_swap::ArcSwap;
|
||||||
@@ -16,12 +16,14 @@ pub(crate) use status::{
|
|||||||
OperatorDrainOutcome, OperatorDrainState, OperatorDrainStatus, OperatorLifecycleState,
|
OperatorDrainOutcome, OperatorDrainState, OperatorDrainStatus, OperatorLifecycleState,
|
||||||
OperatorLifecycleStatus,
|
OperatorLifecycleStatus,
|
||||||
};
|
};
|
||||||
|
// Admission fencing and registration-drain synchronization.
|
||||||
|
mod admission;
|
||||||
|
use admission::{OperatorAdmission, OperatorAdmissionRejection};
|
||||||
|
pub(super) use admission::OperatorRegistration;
|
||||||
// Mutable and published lifecycle state storage.
|
// Mutable and published lifecycle state storage.
|
||||||
mod state;
|
mod state;
|
||||||
use state::{ActiveDrain, OperatorLifecycleInner, OperatorSnapshot, WorkCounts};
|
use state::{ActiveDrain, OperatorLifecycleInner, OperatorSnapshot, WorkCounts};
|
||||||
|
|
||||||
const OPERATOR_ADMISSION_CLOSED: usize = 1 << (usize::BITS - 1);
|
|
||||||
const OPERATOR_REGISTRATION_COUNT: usize = OPERATOR_ADMISSION_CLOSED - 1;
|
|
||||||
const DRAIN_REF_VERSION: &str = "wd1";
|
const DRAIN_REF_VERSION: &str = "wd1";
|
||||||
|
|
||||||
/// Stable operator-control rejection category.
|
/// Stable operator-control rejection category.
|
||||||
@@ -33,77 +35,7 @@ pub(crate) enum OperatorLifecycleError {
|
|||||||
OperationInProgress,
|
OperationInProgress,
|
||||||
}
|
}
|
||||||
|
|
||||||
struct OperatorAdmission {
|
/// Process-owned reversible lifecycle and admission authority.
|
||||||
state: AtomicUsize,
|
|
||||||
registrations_drained: Notify,
|
|
||||||
}
|
|
||||||
|
|
||||||
pub(super) struct OperatorRegistration<'a> {
|
|
||||||
admission: &'a OperatorAdmission,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl OperatorAdmission {
|
|
||||||
fn new() -> Self {
|
|
||||||
Self {
|
|
||||||
state: AtomicUsize::new(0),
|
|
||||||
registrations_drained: Notify::new(),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn try_register(&self) -> Option<OperatorRegistration<'_>> {
|
|
||||||
let mut state = self.state.load(Ordering::Acquire);
|
|
||||||
loop {
|
|
||||||
if state & OPERATOR_ADMISSION_CLOSED != 0
|
|
||||||
|| state & OPERATOR_REGISTRATION_COUNT == OPERATOR_REGISTRATION_COUNT
|
|
||||||
{
|
|
||||||
return None;
|
|
||||||
}
|
|
||||||
match self.state.compare_exchange_weak(
|
|
||||||
state,
|
|
||||||
state + 1,
|
|
||||||
Ordering::AcqRel,
|
|
||||||
Ordering::Acquire,
|
|
||||||
) {
|
|
||||||
Ok(_) => return Some(OperatorRegistration { admission: self }),
|
|
||||||
Err(observed) => state = observed,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn close(&self) {
|
|
||||||
self.state
|
|
||||||
.fetch_or(OPERATOR_ADMISSION_CLOSED, Ordering::AcqRel);
|
|
||||||
}
|
|
||||||
|
|
||||||
fn reopen(&self) {
|
|
||||||
self.state
|
|
||||||
.fetch_and(!OPERATOR_ADMISSION_CLOSED, Ordering::AcqRel);
|
|
||||||
}
|
|
||||||
|
|
||||||
fn is_closed(&self) -> bool {
|
|
||||||
self.state.load(Ordering::Acquire) & OPERATOR_ADMISSION_CLOSED != 0
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn wait_for_registrations(&self) {
|
|
||||||
loop {
|
|
||||||
let notified = self.registrations_drained.notified();
|
|
||||||
if self.state.load(Ordering::Acquire) & OPERATOR_REGISTRATION_COUNT == 0 {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
notified.await;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
impl Drop for OperatorRegistration<'_> {
|
|
||||||
fn drop(&mut self) {
|
|
||||||
let previous = self.admission.state.fetch_sub(1, Ordering::AcqRel);
|
|
||||||
if previous & OPERATOR_REGISTRATION_COUNT == 1 {
|
|
||||||
self.admission.registrations_drained.notify_waiters();
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
pub(super) struct OperatorLifecycle {
|
pub(super) struct OperatorLifecycle {
|
||||||
runtime_instance: Arc<str>,
|
runtime_instance: Arc<str>,
|
||||||
admission: OperatorAdmission,
|
admission: OperatorAdmission,
|
||||||
@@ -115,6 +47,7 @@ pub(super) struct OperatorLifecycle {
|
|||||||
}
|
}
|
||||||
|
|
||||||
impl OperatorLifecycle {
|
impl OperatorLifecycle {
|
||||||
|
/// Creates a running lifecycle for one immutable process instance.
|
||||||
pub(super) fn new(runtime_instance: Arc<str>) -> Self {
|
pub(super) fn new(runtime_instance: Arc<str>) -> Self {
|
||||||
let since = Instant::now();
|
let since = Instant::now();
|
||||||
let snapshot = OperatorSnapshot {
|
let snapshot = OperatorSnapshot {
|
||||||
@@ -142,16 +75,21 @@ impl OperatorLifecycle {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(super) fn try_register(&self) -> Option<OperatorRegistration<'_>> {
|
/// Registers one synchronous operator-fenced admission section.
|
||||||
|
pub(super) fn try_register(
|
||||||
|
&self,
|
||||||
|
) -> Result<OperatorRegistration<'_>, crate::web::telemetry::WebRejectionReason> {
|
||||||
self.admission.try_register()
|
self.admission.try_register()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Wakes an active drain after tracked work ownership changes.
|
||||||
pub(super) fn notify_work_changed(&self) {
|
pub(super) fn notify_work_changed(&self) {
|
||||||
if self.admission.is_closed() {
|
if self.admission.is_closed() {
|
||||||
self.work_changed.notify_waiters();
|
self.work_changed.notify_waiters();
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Returns the lock-free lifecycle snapshot with effective config admission.
|
||||||
pub(super) fn status(&self, config_enabled: bool) -> OperatorLifecycleStatus {
|
pub(super) fn status(&self, config_enabled: bool) -> OperatorLifecycleStatus {
|
||||||
let snapshot = self.published.load();
|
let snapshot = self.published.load();
|
||||||
let admission_open =
|
let admission_open =
|
||||||
@@ -170,30 +108,6 @@ impl OperatorLifecycle {
|
|||||||
self.published.load().terminal
|
self.published.load().terminal
|
||||||
}
|
}
|
||||||
|
|
||||||
fn rejection_reason(&self) -> crate::web::telemetry::WebRejectionReason {
|
|
||||||
let inner = self.inner.lock();
|
|
||||||
if inner.terminal {
|
|
||||||
return crate::web::telemetry::WebRejectionReason::RuntimeClosed;
|
|
||||||
}
|
|
||||||
match inner.state {
|
|
||||||
OperatorLifecycleState::Paused => {
|
|
||||||
crate::web::telemetry::WebRejectionReason::OperatorPaused
|
|
||||||
}
|
|
||||||
OperatorLifecycleState::Draining => {
|
|
||||||
crate::web::telemetry::WebRejectionReason::OperatorDraining
|
|
||||||
}
|
|
||||||
OperatorLifecycleState::ForceClosing => {
|
|
||||||
crate::web::telemetry::WebRejectionReason::OperatorForceClosing
|
|
||||||
}
|
|
||||||
OperatorLifecycleState::Drained => {
|
|
||||||
crate::web::telemetry::WebRejectionReason::OperatorDrained
|
|
||||||
}
|
|
||||||
OperatorLifecycleState::Running => {
|
|
||||||
crate::web::telemetry::WebRejectionReason::RuntimeClosed
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn publish_locked(&self, inner: &OperatorLifecycleInner) {
|
fn publish_locked(&self, inner: &OperatorLifecycleInner) {
|
||||||
self.published.store(Arc::new(OperatorSnapshot {
|
self.published.store(Arc::new(OperatorSnapshot {
|
||||||
state: inner.state,
|
state: inner.state,
|
||||||
@@ -246,6 +160,8 @@ impl OperatorLifecycle {
|
|||||||
return false;
|
return false;
|
||||||
}
|
}
|
||||||
self.transition_locked(&mut inner, OperatorLifecycleState::ForceClosing);
|
self.transition_locked(&mut inner, OperatorLifecycleState::ForceClosing);
|
||||||
|
self.admission
|
||||||
|
.close(OperatorAdmissionRejection::ForceClosing);
|
||||||
let Some(drain) = inner.drain.as_mut() else {
|
let Some(drain) = inner.drain.as_mut() else {
|
||||||
return false;
|
return false;
|
||||||
};
|
};
|
||||||
@@ -269,6 +185,7 @@ impl OperatorLifecycle {
|
|||||||
}
|
}
|
||||||
inner.active = None;
|
inner.active = None;
|
||||||
self.transition_locked(&mut inner, OperatorLifecycleState::Drained);
|
self.transition_locked(&mut inner, OperatorLifecycleState::Drained);
|
||||||
|
self.admission.close(OperatorAdmissionRejection::Drained);
|
||||||
if let Some(drain) = inner.drain.as_mut() {
|
if let Some(drain) = inner.drain.as_mut() {
|
||||||
drain.state = OperatorDrainState::Completed;
|
drain.state = OperatorDrainState::Completed;
|
||||||
drain.outcome = Some(if forced {
|
drain.outcome = Some(if forced {
|
||||||
@@ -286,7 +203,8 @@ impl OperatorLifecycle {
|
|||||||
|
|
||||||
fn close_terminal(&self) {
|
fn close_terminal(&self) {
|
||||||
let mut inner = self.inner.lock();
|
let mut inner = self.inner.lock();
|
||||||
self.admission.close();
|
self.admission
|
||||||
|
.close(OperatorAdmissionRejection::RuntimeClosed);
|
||||||
if inner.terminal {
|
if inner.terminal {
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
@@ -321,7 +239,6 @@ impl WebProcessRuntime {
|
|||||||
if inner.terminal || self.shutdown.is_cancelled() {
|
if inner.terminal || self.shutdown.is_cancelled() {
|
||||||
return Err(OperatorLifecycleError::Closed);
|
return Err(OperatorLifecycleError::Closed);
|
||||||
}
|
}
|
||||||
self.operator_lifecycle.admission.close();
|
|
||||||
if matches!(
|
if matches!(
|
||||||
inner.state,
|
inner.state,
|
||||||
OperatorLifecycleState::Running | OperatorLifecycleState::Drained
|
OperatorLifecycleState::Running | OperatorLifecycleState::Drained
|
||||||
@@ -329,6 +246,9 @@ impl WebProcessRuntime {
|
|||||||
self.operator_lifecycle
|
self.operator_lifecycle
|
||||||
.transition_locked(&mut inner, OperatorLifecycleState::Paused);
|
.transition_locked(&mut inner, OperatorLifecycleState::Paused);
|
||||||
}
|
}
|
||||||
|
self.operator_lifecycle.admission.close(
|
||||||
|
OperatorAdmissionRejection::for_state(inner.state),
|
||||||
|
);
|
||||||
self.operator_lifecycle.publish_locked(&inner);
|
self.operator_lifecycle.publish_locked(&inner);
|
||||||
}
|
}
|
||||||
self.operator_lifecycle
|
self.operator_lifecycle
|
||||||
@@ -355,7 +275,7 @@ impl WebProcessRuntime {
|
|||||||
let started = TokioInstant::now();
|
let started = TokioInstant::now();
|
||||||
let deadline = started + timeout;
|
let deadline = started + timeout;
|
||||||
let started_epoch_millis = crate::web::trace::store_epoch_millis();
|
let started_epoch_millis = crate::web::trace::store_epoch_millis();
|
||||||
{
|
let accepted = {
|
||||||
let mut inner = self.operator_lifecycle.inner.lock();
|
let mut inner = self.operator_lifecycle.inner.lock();
|
||||||
if inner.terminal || self.shutdown.is_cancelled() {
|
if inner.terminal || self.shutdown.is_cancelled() {
|
||||||
return Err(OperatorLifecycleError::Closed);
|
return Err(OperatorLifecycleError::Closed);
|
||||||
@@ -368,7 +288,6 @@ impl WebProcessRuntime {
|
|||||||
{
|
{
|
||||||
return Err(OperatorLifecycleError::OperationInProgress);
|
return Err(OperatorLifecycleError::OperationInProgress);
|
||||||
}
|
}
|
||||||
self.operator_lifecycle.admission.close();
|
|
||||||
inner.active = Some(ActiveDrain {
|
inner.active = Some(ActiveDrain {
|
||||||
sequence,
|
sequence,
|
||||||
cancellation: cancellation.clone(),
|
cancellation: cancellation.clone(),
|
||||||
@@ -392,26 +311,20 @@ impl WebProcessRuntime {
|
|||||||
});
|
});
|
||||||
self.operator_lifecycle
|
self.operator_lifecycle
|
||||||
.transition_locked(&mut inner, OperatorLifecycleState::Draining);
|
.transition_locked(&mut inner, OperatorLifecycleState::Draining);
|
||||||
|
self.operator_lifecycle
|
||||||
|
.admission
|
||||||
|
.close(OperatorAdmissionRejection::Draining);
|
||||||
|
let runtime = Arc::clone(self);
|
||||||
|
self.spawn_auxiliary(async move {
|
||||||
|
runtime
|
||||||
|
.run_operator_drain(sequence, deadline, cancellation)
|
||||||
|
.await;
|
||||||
|
});
|
||||||
self.operator_lifecycle.publish_locked(&inner);
|
self.operator_lifecycle.publish_locked(&inner);
|
||||||
}
|
let config_enabled = self.active_generation().config().web.enabled;
|
||||||
self.operator_lifecycle
|
self.operator_lifecycle.status(config_enabled)
|
||||||
.admission
|
};
|
||||||
.wait_for_registrations()
|
Ok(accepted)
|
||||||
.await;
|
|
||||||
if self.shutdown.is_cancelled() || self.operator_lifecycle.is_terminal() {
|
|
||||||
return Err(OperatorLifecycleError::Closed);
|
|
||||||
}
|
|
||||||
let counts = self.operator_work_counts();
|
|
||||||
if !self.operator_lifecycle.update_counts(sequence, counts) {
|
|
||||||
return Err(OperatorLifecycleError::Closed);
|
|
||||||
}
|
|
||||||
let runtime = Arc::clone(self);
|
|
||||||
self.spawn_auxiliary(async move {
|
|
||||||
runtime
|
|
||||||
.run_operator_drain(sequence, deadline, cancellation)
|
|
||||||
.await;
|
|
||||||
});
|
|
||||||
Ok(self.operator_lifecycle_status())
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Resumes operator admission and invalidates any active drain waiter.
|
/// Resumes operator admission and invalidates any active drain waiter.
|
||||||
@@ -444,23 +357,25 @@ impl WebProcessRuntime {
|
|||||||
Ok(self.operator_lifecycle_status())
|
Ok(self.operator_lifecycle_status())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Registers one WEB manager admission commit under the operator fence.
|
||||||
pub(super) fn try_operator_admission(
|
pub(super) fn try_operator_admission(
|
||||||
&self,
|
&self,
|
||||||
) -> Result<OperatorRegistration<'_>, super::ManagerError> {
|
) -> Result<OperatorRegistration<'_>, super::ManagerError> {
|
||||||
match self.operator_lifecycle.try_register() {
|
match self.operator_lifecycle.try_register() {
|
||||||
Some(registration) => Ok(registration),
|
Ok(registration) => Ok(registration),
|
||||||
None => {
|
Err(reason) => {
|
||||||
self.telemetry
|
self.telemetry.record_rejection(reason);
|
||||||
.record_rejection(self.operator_lifecycle.rejection_reason());
|
|
||||||
Err(super::ManagerError::AdmissionPaused)
|
Err(super::ManagerError::AdmissionPaused)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Notifies an active drain that tracked WEB work changed.
|
||||||
pub(super) fn notify_operator_work_changed(&self) {
|
pub(super) fn notify_operator_work_changed(&self) {
|
||||||
self.operator_lifecycle.notify_work_changed();
|
self.operator_lifecycle.notify_work_changed();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Terminally closes operator lifecycle during process shutdown.
|
||||||
pub(super) fn close_operator_lifecycle(&self) {
|
pub(super) fn close_operator_lifecycle(&self) {
|
||||||
self.operator_lifecycle.close_terminal();
|
self.operator_lifecycle.close_terminal();
|
||||||
}
|
}
|
||||||
@@ -482,9 +397,16 @@ impl WebProcessRuntime {
|
|||||||
deadline: TokioInstant,
|
deadline: TokioInstant,
|
||||||
cancellation: CancellationToken,
|
cancellation: CancellationToken,
|
||||||
) {
|
) {
|
||||||
|
tokio::select! {
|
||||||
|
biased;
|
||||||
|
_ = cancellation.cancelled() => return,
|
||||||
|
_ = self.operator_lifecycle.admission.wait_for_registrations() => {}
|
||||||
|
}
|
||||||
let mut forced = false;
|
let mut forced = false;
|
||||||
loop {
|
loop {
|
||||||
let notified = self.operator_lifecycle.work_changed.notified();
|
let notified = self.operator_lifecycle.work_changed.notified();
|
||||||
|
tokio::pin!(notified);
|
||||||
|
notified.as_mut().enable();
|
||||||
let counts = self.operator_work_counts();
|
let counts = self.operator_work_counts();
|
||||||
if !self.operator_lifecycle.update_counts(sequence, counts) {
|
if !self.operator_lifecycle.update_counts(sequence, counts) {
|
||||||
return;
|
return;
|
||||||
@@ -497,14 +419,14 @@ impl WebProcessRuntime {
|
|||||||
tokio::select! {
|
tokio::select! {
|
||||||
biased;
|
biased;
|
||||||
_ = cancellation.cancelled() => return,
|
_ = cancellation.cancelled() => return,
|
||||||
_ = notified => {}
|
_ = notified.as_mut() => {}
|
||||||
}
|
}
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
tokio::select! {
|
tokio::select! {
|
||||||
biased;
|
biased;
|
||||||
_ = cancellation.cancelled() => return,
|
_ = cancellation.cancelled() => return,
|
||||||
_ = notified => {},
|
_ = notified.as_mut() => {},
|
||||||
_ = tokio::time::sleep_until(deadline) => {
|
_ = tokio::time::sleep_until(deadline) => {
|
||||||
let counts = self.operator_work_counts();
|
let counts = self.operator_work_counts();
|
||||||
if counts.is_zero() {
|
if counts.is_zero() {
|
||||||
|
|||||||
@@ -0,0 +1,165 @@
|
|||||||
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||||
|
|
||||||
|
use tokio::sync::Notify;
|
||||||
|
|
||||||
|
use super::OperatorLifecycleState;
|
||||||
|
|
||||||
|
const OPERATOR_ADMISSION_CLOSED: usize = 1 << (usize::BITS - 1);
|
||||||
|
const OPERATOR_REJECTION_BITS: u32 = 3;
|
||||||
|
const OPERATOR_REJECTION_SHIFT: u32 = usize::BITS - 1 - OPERATOR_REJECTION_BITS;
|
||||||
|
const OPERATOR_REJECTION_MASK: usize =
|
||||||
|
((1usize << OPERATOR_REJECTION_BITS) - 1) << OPERATOR_REJECTION_SHIFT;
|
||||||
|
const OPERATOR_REGISTRATION_COUNT: usize = (1usize << OPERATOR_REJECTION_SHIFT) - 1;
|
||||||
|
|
||||||
|
#[derive(Clone, Copy)]
|
||||||
|
#[repr(usize)]
|
||||||
|
/// Stable rejection encoded into the closed admission-fence word.
|
||||||
|
pub(super) enum OperatorAdmissionRejection {
|
||||||
|
/// Operator pause closed admission.
|
||||||
|
Paused = 1,
|
||||||
|
/// Graceful drain closed admission.
|
||||||
|
Draining = 2,
|
||||||
|
/// Deadline-triggered forced closure remains in progress.
|
||||||
|
ForceClosing = 3,
|
||||||
|
/// The latest drain reached confirmed zero.
|
||||||
|
Drained = 4,
|
||||||
|
/// Terminal process shutdown closed admission.
|
||||||
|
RuntimeClosed = 5,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl OperatorAdmissionRejection {
|
||||||
|
/// Maps a closed lifecycle state to its admission rejection.
|
||||||
|
pub(super) fn for_state(state: OperatorLifecycleState) -> Self {
|
||||||
|
match state {
|
||||||
|
OperatorLifecycleState::Paused => Self::Paused,
|
||||||
|
OperatorLifecycleState::Draining => Self::Draining,
|
||||||
|
OperatorLifecycleState::ForceClosing => Self::ForceClosing,
|
||||||
|
OperatorLifecycleState::Drained => Self::Drained,
|
||||||
|
OperatorLifecycleState::Running => Self::RuntimeClosed,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn from_admission_state(state: usize) -> Self {
|
||||||
|
match (state & OPERATOR_REJECTION_MASK) >> OPERATOR_REJECTION_SHIFT {
|
||||||
|
value if value == Self::Paused as usize => Self::Paused,
|
||||||
|
value if value == Self::Draining as usize => Self::Draining,
|
||||||
|
value if value == Self::ForceClosing as usize => Self::ForceClosing,
|
||||||
|
value if value == Self::Drained as usize => Self::Drained,
|
||||||
|
_ => Self::RuntimeClosed,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn telemetry_reason(self) -> crate::web::telemetry::WebRejectionReason {
|
||||||
|
match self {
|
||||||
|
Self::Paused => crate::web::telemetry::WebRejectionReason::OperatorPaused,
|
||||||
|
Self::Draining => crate::web::telemetry::WebRejectionReason::OperatorDraining,
|
||||||
|
Self::ForceClosing => {
|
||||||
|
crate::web::telemetry::WebRejectionReason::OperatorForceClosing
|
||||||
|
}
|
||||||
|
Self::Drained => crate::web::telemetry::WebRejectionReason::OperatorDrained,
|
||||||
|
Self::RuntimeClosed => crate::web::telemetry::WebRejectionReason::RuntimeClosed,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Lock-free admission fence with bounded pre-cutover registration tracking.
|
||||||
|
pub(super) struct OperatorAdmission {
|
||||||
|
state: AtomicUsize,
|
||||||
|
registrations_drained: Notify,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// RAII ownership of one synchronous pre-cutover admission section.
|
||||||
|
pub(in crate::web::manager) struct OperatorRegistration<'a> {
|
||||||
|
admission: &'a OperatorAdmission,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl OperatorAdmission {
|
||||||
|
/// Creates an open admission fence.
|
||||||
|
pub(super) fn new() -> Self {
|
||||||
|
Self {
|
||||||
|
state: AtomicUsize::new(0),
|
||||||
|
registrations_drained: Notify::new(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Registers one pre-cutover section or returns its stable rejection reason.
|
||||||
|
pub(super) fn try_register(
|
||||||
|
&self,
|
||||||
|
) -> Result<OperatorRegistration<'_>, crate::web::telemetry::WebRejectionReason> {
|
||||||
|
let mut state = self.state.load(Ordering::Acquire);
|
||||||
|
loop {
|
||||||
|
if state & OPERATOR_ADMISSION_CLOSED != 0 {
|
||||||
|
return Err(
|
||||||
|
OperatorAdmissionRejection::from_admission_state(state).telemetry_reason(),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
if state & OPERATOR_REGISTRATION_COUNT == OPERATOR_REGISTRATION_COUNT {
|
||||||
|
return Err(crate::web::telemetry::WebRejectionReason::Concurrent);
|
||||||
|
}
|
||||||
|
match self.state.compare_exchange_weak(
|
||||||
|
state,
|
||||||
|
state + 1,
|
||||||
|
Ordering::AcqRel,
|
||||||
|
Ordering::Acquire,
|
||||||
|
) {
|
||||||
|
Ok(_) => return Ok(OperatorRegistration { admission: self }),
|
||||||
|
Err(observed) => state = observed,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Atomically closes admission while preserving active registration count.
|
||||||
|
pub(super) fn close(&self, reason: OperatorAdmissionRejection) {
|
||||||
|
let reason = (reason as usize) << OPERATOR_REJECTION_SHIFT;
|
||||||
|
let mut state = self.state.load(Ordering::Acquire);
|
||||||
|
loop {
|
||||||
|
let next = (state & OPERATOR_REGISTRATION_COUNT)
|
||||||
|
| OPERATOR_ADMISSION_CLOSED
|
||||||
|
| reason;
|
||||||
|
match self.state.compare_exchange_weak(
|
||||||
|
state,
|
||||||
|
next,
|
||||||
|
Ordering::AcqRel,
|
||||||
|
Ordering::Acquire,
|
||||||
|
) {
|
||||||
|
Ok(_) => return,
|
||||||
|
Err(observed) => state = observed,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Reopens admission without altering active registration ownership.
|
||||||
|
pub(super) fn reopen(&self) {
|
||||||
|
self.state.fetch_and(
|
||||||
|
!(OPERATOR_ADMISSION_CLOSED | OPERATOR_REJECTION_MASK),
|
||||||
|
Ordering::AcqRel,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Returns whether the admission fence is closed.
|
||||||
|
pub(super) fn is_closed(&self) -> bool {
|
||||||
|
self.state.load(Ordering::Acquire) & OPERATOR_ADMISSION_CLOSED != 0
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Waits until every pre-cutover registration has been released.
|
||||||
|
pub(super) async fn wait_for_registrations(&self) {
|
||||||
|
loop {
|
||||||
|
let notified = self.registrations_drained.notified();
|
||||||
|
tokio::pin!(notified);
|
||||||
|
notified.as_mut().enable();
|
||||||
|
if self.state.load(Ordering::Acquire) & OPERATOR_REGISTRATION_COUNT == 0 {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
notified.await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Drop for OperatorRegistration<'_> {
|
||||||
|
fn drop(&mut self) {
|
||||||
|
let previous = self.admission.state.fetch_sub(1, Ordering::AcqRel);
|
||||||
|
if previous & OPERATOR_REGISTRATION_COUNT == 1 {
|
||||||
|
self.admission.registrations_drained.notify_waiters();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -133,6 +133,43 @@ async fn empty_drain_completes_gracefully_and_stays_closed_until_resume() {
|
|||||||
stop_runtime(runtime, generation).await;
|
stop_runtime(runtime, generation).await;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn drain_request_returns_after_registering_its_worker() {
|
||||||
|
let (runtime, generation) = test_runtime();
|
||||||
|
let registration = runtime.try_operator_admission().unwrap();
|
||||||
|
let drain_runtime = Arc::clone(&runtime);
|
||||||
|
let drain = tokio::spawn(async move {
|
||||||
|
drain_runtime
|
||||||
|
.drain_operator(Duration::from_secs(30))
|
||||||
|
.await
|
||||||
|
});
|
||||||
|
|
||||||
|
tokio::time::timeout(Duration::from_secs(1), async {
|
||||||
|
while runtime.operator_lifecycle_status().state != OperatorLifecycleState::Draining {
|
||||||
|
tokio::task::yield_now().await;
|
||||||
|
}
|
||||||
|
})
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
let accepted = tokio::time::timeout(Duration::from_secs(1), drain)
|
||||||
|
.await
|
||||||
|
.unwrap()
|
||||||
|
.unwrap()
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(accepted.state, OperatorLifecycleState::Draining);
|
||||||
|
drop(registration);
|
||||||
|
|
||||||
|
tokio::time::timeout(Duration::from_secs(1), async {
|
||||||
|
while runtime.operator_lifecycle_status().state != OperatorLifecycleState::Drained {
|
||||||
|
tokio::task::yield_now().await;
|
||||||
|
}
|
||||||
|
})
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
stop_runtime(runtime, generation).await;
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn resume_after_force_commit_cancels_wait_but_preserves_force_evidence() {
|
async fn resume_after_force_commit_cancels_wait_but_preserves_force_evidence() {
|
||||||
let (runtime, generation) = test_runtime();
|
let (runtime, generation) = test_runtime();
|
||||||
@@ -140,7 +177,10 @@ async fn resume_after_force_commit_cancels_wait_but_preserves_force_evidence() {
|
|||||||
let cancellation = CancellationToken::new();
|
let cancellation = CancellationToken::new();
|
||||||
{
|
{
|
||||||
let mut inner = runtime.operator_lifecycle.inner.lock();
|
let mut inner = runtime.operator_lifecycle.inner.lock();
|
||||||
runtime.operator_lifecycle.admission.close();
|
runtime
|
||||||
|
.operator_lifecycle
|
||||||
|
.admission
|
||||||
|
.close(OperatorAdmissionRejection::Draining);
|
||||||
inner.active = Some(ActiveDrain {
|
inner.active = Some(ActiveDrain {
|
||||||
sequence,
|
sequence,
|
||||||
cancellation,
|
cancellation,
|
||||||
|
|||||||
@@ -64,6 +64,8 @@ impl WebSession {
|
|||||||
let poll = async {
|
let poll = async {
|
||||||
loop {
|
loop {
|
||||||
let notified = self.down_notify.notified();
|
let notified = self.down_notify.notified();
|
||||||
|
tokio::pin!(notified);
|
||||||
|
notified.as_mut().enable();
|
||||||
{
|
{
|
||||||
let mut state = self.state.lock();
|
let mut state = self.state.lock();
|
||||||
if state.down_epoch != epoch {
|
if state.down_epoch != epoch {
|
||||||
|
|||||||
@@ -151,6 +151,8 @@ impl WebSession {
|
|||||||
let poll = async {
|
let poll = async {
|
||||||
loop {
|
loop {
|
||||||
let notified = notify.notified();
|
let notified = notify.notified();
|
||||||
|
tokio::pin!(notified);
|
||||||
|
notified.as_mut().enable();
|
||||||
{
|
{
|
||||||
let mut state = self.state.lock();
|
let mut state = self.state.lock();
|
||||||
if state.closed {
|
if state.closed {
|
||||||
@@ -306,6 +308,8 @@ impl WebSession {
|
|||||||
let opened = tokio::time::timeout(deadline, async {
|
let opened = tokio::time::timeout(deadline, async {
|
||||||
loop {
|
loop {
|
||||||
let notified = self.lane_open_notify.notified();
|
let notified = self.lane_open_notify.notified();
|
||||||
|
tokio::pin!(notified);
|
||||||
|
notified.as_mut().enable();
|
||||||
{
|
{
|
||||||
let state = self.state.lock();
|
let state = self.state.lock();
|
||||||
if state.closed {
|
if state.closed {
|
||||||
|
|||||||
@@ -77,6 +77,8 @@ impl WebSession {
|
|||||||
pub(crate) async fn wait(&self) {
|
pub(crate) async fn wait(&self) {
|
||||||
loop {
|
loop {
|
||||||
let notified = self.tasks_done.notified();
|
let notified = self.tasks_done.notified();
|
||||||
|
tokio::pin!(notified);
|
||||||
|
notified.as_mut().enable();
|
||||||
if self.tasks_live.load(Ordering::Acquire) == 0 {
|
if self.tasks_live.load(Ordering::Acquire) == 0 {
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
|||||||
+23
-17
@@ -4,6 +4,10 @@ use std::time::Instant;
|
|||||||
|
|
||||||
use serde::Serialize;
|
use serde::Serialize;
|
||||||
|
|
||||||
|
const LAST_DECOY_OUTCOME_BITS: u32 = 4;
|
||||||
|
const LAST_DECOY_OUTCOME_MASK: u64 = (1 << LAST_DECOY_OUTCOME_BITS) - 1;
|
||||||
|
const LAST_DECOY_ELAPSED_MAX: u64 = u64::MAX >> LAST_DECOY_OUTCOME_BITS;
|
||||||
|
|
||||||
/// Stable operational rejection reason recorded at the decision point.
|
/// Stable operational rejection reason recorded at the decision point.
|
||||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||||
#[repr(usize)]
|
#[repr(usize)]
|
||||||
@@ -296,8 +300,7 @@ pub(crate) struct WebTelemetry {
|
|||||||
rejections: [AtomicU64; WebRejectionReason::ALL.len()],
|
rejections: [AtomicU64; WebRejectionReason::ALL.len()],
|
||||||
overload_outcomes: [AtomicU64; WebHttpConnectionOverloadOutcome::ALL.len()],
|
overload_outcomes: [AtomicU64; WebHttpConnectionOverloadOutcome::ALL.len()],
|
||||||
decoy_outcomes: [AtomicU64; WebDecoyUpstreamOutcome::ALL.len()],
|
decoy_outcomes: [AtomicU64; WebDecoyUpstreamOutcome::ALL.len()],
|
||||||
last_decoy_outcome: AtomicUsize,
|
last_decoy: AtomicU64,
|
||||||
last_decoy_elapsed_ms: AtomicU64,
|
|
||||||
sessions_created: AtomicU64,
|
sessions_created: AtomicU64,
|
||||||
sessions_closed: AtomicU64,
|
sessions_closed: AtomicU64,
|
||||||
streams_opened: AtomicU64,
|
streams_opened: AtomicU64,
|
||||||
@@ -318,8 +321,7 @@ impl WebTelemetry {
|
|||||||
rejections: std::array::from_fn(|_| AtomicU64::new(0)),
|
rejections: std::array::from_fn(|_| AtomicU64::new(0)),
|
||||||
overload_outcomes: std::array::from_fn(|_| AtomicU64::new(0)),
|
overload_outcomes: std::array::from_fn(|_| AtomicU64::new(0)),
|
||||||
decoy_outcomes: std::array::from_fn(|_| AtomicU64::new(0)),
|
decoy_outcomes: std::array::from_fn(|_| AtomicU64::new(0)),
|
||||||
last_decoy_outcome: AtomicUsize::new(usize::MAX),
|
last_decoy: AtomicU64::new(0),
|
||||||
last_decoy_elapsed_ms: AtomicU64::new(0),
|
|
||||||
sessions_created: AtomicU64::new(0),
|
sessions_created: AtomicU64::new(0),
|
||||||
sessions_closed: AtomicU64::new(0),
|
sessions_closed: AtomicU64::new(0),
|
||||||
streams_opened: AtomicU64::new(0),
|
streams_opened: AtomicU64::new(0),
|
||||||
@@ -408,11 +410,13 @@ impl WebTelemetry {
|
|||||||
/// Records one internal plain-HTTP decoy origin outcome.
|
/// Records one internal plain-HTTP decoy origin outcome.
|
||||||
pub(crate) fn record_decoy(&self, outcome: WebDecoyUpstreamOutcome) {
|
pub(crate) fn record_decoy(&self, outcome: WebDecoyUpstreamOutcome) {
|
||||||
self.decoy_outcomes[outcome as usize].fetch_add(1, Ordering::Relaxed);
|
self.decoy_outcomes[outcome as usize].fetch_add(1, Ordering::Relaxed);
|
||||||
let elapsed_ms = self.started.elapsed().as_millis().min(u128::from(u64::MAX)) as u64;
|
let elapsed_ms = self
|
||||||
self.last_decoy_elapsed_ms
|
.started
|
||||||
.store(elapsed_ms.saturating_add(1), Ordering::Relaxed);
|
.elapsed()
|
||||||
self.last_decoy_outcome
|
.as_millis()
|
||||||
.store(outcome as usize, Ordering::Release);
|
.min(u128::from(LAST_DECOY_ELAPSED_MAX)) as u64;
|
||||||
|
let packed = (elapsed_ms << LAST_DECOY_OUTCOME_BITS) | (outcome as u64 + 1);
|
||||||
|
self.last_decoy.store(packed, Ordering::Release);
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Returns one fixed internal decoy origin counter.
|
/// Returns one fixed internal decoy origin counter.
|
||||||
@@ -433,14 +437,16 @@ impl WebTelemetry {
|
|||||||
|
|
||||||
/// Returns the last decoy outcome and its monotonic age in milliseconds.
|
/// Returns the last decoy outcome and its monotonic age in milliseconds.
|
||||||
pub(crate) fn last_decoy(&self) -> Option<(&'static str, u64)> {
|
pub(crate) fn last_decoy(&self) -> Option<(&'static str, u64)> {
|
||||||
let raw = self.last_decoy_outcome.load(Ordering::Acquire);
|
let packed = self.last_decoy.load(Ordering::Acquire);
|
||||||
let outcome = WebDecoyUpstreamOutcome::ALL.get(raw).copied()?;
|
let outcome_index = (packed & LAST_DECOY_OUTCOME_MASK).checked_sub(1)? as usize;
|
||||||
let recorded = self.last_decoy_elapsed_ms.load(Ordering::Relaxed);
|
let outcome = WebDecoyUpstreamOutcome::ALL.get(outcome_index).copied()?;
|
||||||
let now = self.started.elapsed().as_millis().min(u128::from(u64::MAX)) as u64;
|
let recorded = packed >> LAST_DECOY_OUTCOME_BITS;
|
||||||
Some((
|
let now = self
|
||||||
outcome.as_str(),
|
.started
|
||||||
now.saturating_sub(recorded.saturating_sub(1)),
|
.elapsed()
|
||||||
))
|
.as_millis()
|
||||||
|
.min(u128::from(LAST_DECOY_ELAPSED_MAX)) as u64;
|
||||||
|
Some((outcome.as_str(), now.saturating_sub(recorded)))
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Records one created session incarnation.
|
/// Records one created session incarnation.
|
||||||
|
|||||||
Reference in New Issue
Block a user