Runtime Ownership hardened

Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
This commit is contained in:
Alexey
2026-08-30 08:38:03 +03:00
parent 1bb6b0bdda
commit 281f63f940
91 changed files with 3972 additions and 1239 deletions
+2 -2
View File
@@ -14,7 +14,7 @@ use crate::error::Result;
use crate::transport::UpstreamManager;
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::secret::download_proxy_secret_with_max_len_via_upstream;
use super::selftest::record_timeskew_sample;
@@ -106,7 +106,7 @@ pub async fn fetch_proxy_config_with_raw_via_upstream(
url: &str,
upstream: Option<Arc<UpstreamManager>>,
) -> 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;
if let Some(date_str) = resp.date_header.as_deref()
+62 -6
View File
@@ -71,6 +71,40 @@ struct FamilyReconnectOutcome {
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) {
let mut backoff: HashMap<(i32, IpFamily), u64> = 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 bound_clients_by_writer = Arc::new(bound_clients_by_writer);
let mut reconnect_set = JoinSet::<FamilyReconnectOutcome>::new();
let mut scheduled_reconnects = ScheduledReconnects {
inflight,
keys: Vec::new(),
};
for (dc, endpoints) in dc_endpoints {
if endpoints.is_empty() {
@@ -562,7 +600,7 @@ async fn check_family(
.reconnect_runtime
.me_reconnect_max_concurrent_per_dc
.max(1) as usize;
if *inflight.get(&key).unwrap_or(&0) >= max_concurrent {
if scheduled_reconnects.current(&key) >= max_concurrent {
continue;
}
if pool
@@ -579,7 +617,7 @@ async fn check_family(
);
continue;
}
*inflight.entry(key).or_insert(0) += 1;
scheduled_reconnects.reserve(key);
let pool_for_reconnect = pool.clone();
let rng_for_reconnect = rng.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
@@ -1701,15 +1736,35 @@ mod tests {
use tokio::sync::mpsc;
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::crypto::SecureRandom;
use crate::network::probe::NetworkDecision;
use crate::network::IpFamily;
use crate::stats::Stats;
use crate::transport::middle_proxy::codec::WriterCommand;
use crate::transport::middle_proxy::pool::{MePool, MeWriter, WriterContour};
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> {
let general = GeneralConfig {
me_pool_drain_threshold,
@@ -1812,6 +1867,7 @@ mod tests {
general.me_route_blocking_send_timeout_ms,
general.me_route_inline_recovery_attempts,
general.me_route_inline_recovery_wait_ms,
16_384,
)
}
+41 -14
View File
@@ -1,7 +1,7 @@
use std::sync::Arc;
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::{Method, Request};
use hyper_util::rt::TokioIo;
@@ -12,11 +12,19 @@ use tokio_rustls::TlsConnector;
use tracing::debug;
use crate::error::{ProxyError, Result};
use crate::network::dns_overrides::resolve_socket_addr;
use crate::transport::{UpstreamManager, UpstreamStream};
const HTTP_CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
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) 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)))
.await
.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(
url: &str,
upstream: Option<Arc<UpstreamManager>>,
max_body_bytes: usize,
) -> 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 stream = connect_https_transport(&host, port, upstream).await?;
@@ -115,11 +122,11 @@ pub(crate) async fn https_get(
.await
.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 {
debug!(error = %e, "HTTPS fetch connection task failed");
}
});
}));
let host_header = if port == 443 {
host.clone()
@@ -148,10 +155,17 @@ pub(crate) async fn https_get(
.and_then(|value| value.to_str().ok())
.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
.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_vec();
@@ -161,3 +175,16 @@ pub(crate) async fn https_get(
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());
}
}
+2
View File
@@ -22,6 +22,7 @@ mod ping;
mod pool;
mod pool_config;
mod pool_init;
mod pool_lifecycle;
mod pool_nat;
mod pool_refill;
#[cfg(test)]
@@ -60,6 +61,7 @@ pub use ping::{
MePingFamily, MePingReport, MePingSample, format_me_route, format_sample_line, run_me_ping,
};
pub use pool::MePool;
pub(crate) use registry::ConnLease;
#[allow(unused_imports)]
pub use pool_nat::{detect_public_ip, stun_probe};
pub use registry::ConnRegistry;
+148 -11
View File
@@ -10,6 +10,7 @@ use std::sync::atomic::{
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use arc_swap::ArcSwap;
use parking_lot::Mutex as ParkingMutex;
use tokio::sync::{Mutex, RwLock, Semaphore, mpsc, watch};
use tokio_util::sync::CancellationToken;
@@ -23,6 +24,7 @@ use crate::transport::UpstreamManager;
use super::ConnRegistry;
use super::codec::WriterCommand;
use super::pool_lifecycle::MePoolLifecycle;
const ME_FORCE_CLOSE_SAFETY_FALLBACK_SECS: u64 = 300;
@@ -32,12 +34,6 @@ pub(super) struct RefillDcKey {
pub family: IpFamily,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub(super) struct RefillEndpointKey {
pub dc: i32,
pub addr: SocketAddr,
}
#[derive(Clone)]
pub struct MeWriter {
pub id: u64,
@@ -148,6 +144,18 @@ pub(super) enum WriterContour {
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 {
pub(super) fn as_u8(self) -> u8 {
self as u8
@@ -265,6 +273,10 @@ pub(super) struct ReinitCore {
pub(super) pending_hardswap_generation: AtomicU64,
pub(super) pending_hardswap_started_at_epoch_secs: 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) me_hardswap_warmup_delay_min_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,
}
#[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) me_keepalive_enabled: bool,
pub(super) me_keepalive_interval: Duration,
@@ -421,6 +466,7 @@ pub struct MePool {
pub(super) floor_runtime: Arc<FloorRuntimeCore>,
pub(super) writer_selection_policy: Arc<WriterSelectionPolicyCore>,
pub(super) transport_policy: Arc<TransportPolicyCore>,
pub(super) lifecycle: MePoolLifecycle,
pub(super) decision: NetworkDecision,
pub(super) upstream: Option<Arc<UpstreamManager>>,
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) default_dc: AtomicI32,
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) refill_inflight: Arc<Mutex<HashSet<RefillEndpointKey>>>,
pub(super) refill_inflight_dc: Arc<Mutex<HashSet<RefillDcKey>>>,
pub(super) refill_states: Arc<ParkingMutex<HashMap<RefillDcKey, Option<SocketAddr>>>>,
pub(super) refill_running: AtomicUsize,
pub(super) refill_pending: AtomicUsize,
pub(super) conn_count: AtomicUsize,
pub(super) draining_active_runtime: AtomicU64,
pub(super) stats: Arc<crate::stats::Stats>,
@@ -573,12 +622,14 @@ impl MePool {
me_route_blocking_send_timeout_ms: u64,
me_route_inline_recovery_attempts: u32,
me_route_inline_recovery_wait_ms: u64,
me_connection_cleanup_capacity: usize,
) -> Arc<Self> {
let endpoint_dc_map = Self::build_endpoint_dc_map_from_maps(&proxy_map_v4, &proxy_map_v6);
let preferred_endpoints_by_dc =
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_connection_cleanup_capacity,
));
registry.update_route_backpressure_policy(
me_route_backpressure_base_timeout_ms,
@@ -587,6 +638,14 @@ impl MePool {
);
let (writer_epoch, _) = watch::channel(0u64);
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);
Arc::new(Self {
routing: Arc::new(RoutingCore {
@@ -603,6 +662,16 @@ impl MePool {
pending_hardswap_generation: AtomicU64::new(0),
pending_hardswap_started_at_epoch_secs: 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),
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),
@@ -805,6 +874,7 @@ impl MePool {
me_reader_route_data_wait_ms,
)),
}),
lifecycle: MePoolLifecycle::new(),
decision,
upstream,
rng,
@@ -830,9 +900,12 @@ impl MePool {
endpoint_dc_map: Arc::new(RwLock::new(endpoint_dc_map)),
default_dc: AtomicI32::new(default_dc.unwrap_or(2)),
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())),
refill_inflight: Arc::new(Mutex::new(HashSet::new())),
refill_inflight_dc: Arc::new(Mutex::new(HashSet::new())),
refill_states: Arc::new(ParkingMutex::new(HashMap::new())),
refill_running: AtomicUsize::new(0),
refill_pending: AtomicUsize::new(0),
conn_count: AtomicUsize::new(0),
draining_active_runtime: AtomicU64::new(0),
endpoint_quarantine: Arc::new(Mutex::new(HashMap::new())),
@@ -1263,6 +1336,7 @@ impl MePool {
self.translate_our_addr_with_reflection(addr, None)
}
#[allow(dead_code)]
pub fn registry(&self) -> &Arc<ConnRegistry> {
&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(
&self,
endpoint_count: usize,
+5 -2
View File
@@ -107,7 +107,7 @@ impl MePool {
let pool = Arc::clone(self);
let rng_clone = Arc::clone(rng);
let dc_addrs_bg = dc_addrs.clone();
tokio::spawn(async move {
let saturation = async move {
let mut join_bg = tokio::task::JoinSet::new();
for (dc, addrs) in dc_addrs_bg {
if addrs.len() <= 1 {
@@ -135,7 +135,10 @@ impl MePool {
current_pool_size = pool.connection_count(),
"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 {
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());
}
}
+6 -2
View File
@@ -9,7 +9,8 @@ use tracing::{debug, info};
use crate::error::{ProxyError, Result};
use crate::network::probe::{detect_public_ipv4_http, is_bogon};
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;
@@ -66,10 +67,12 @@ impl MePool {
let mut best_by_ip: HashMap<IpAddr, (usize, std::net::SocketAddr)> = HashMap::new();
let concurrency = self.nat_runtime.nat_probe_concurrency.max(1);
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.len() < concurrency {
let stun_addr = servers[next_idx].clone();
let dns_resolver = dns_resolver.clone();
next_idx += 1;
join_set.spawn(async move {
let batch_timeout = if tcp_fallback {
@@ -79,11 +82,12 @@ impl MePool {
};
let res = timeout(
batch_timeout,
stun_probe_family_with_bind_and_tcp_fallback(
stun_probe_family_with_bind_tcp_fallback_and_resolver(
&stun_addr,
family,
bind_ip,
tcp_fallback,
dns_resolver.as_deref(),
),
)
.await;
+74 -49
View File
@@ -9,13 +9,48 @@ use tracing::{debug, info, warn};
use crate::crypto::SecureRandom;
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_QUARANTINE_SECS: u64 = 25;
const ME_FLAP_MIN_UPTIME_MILLIS: u64 = 500;
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 {
pub(super) async fn sweep_endpoint_quarantine(&self) {
let configured = self
@@ -131,8 +166,7 @@ impl MePool {
}
pub(super) async fn has_refill_inflight_for_dc_key(&self, key: RefillDcKey) -> bool {
let guard = self.refill_inflight_dc.lock().await;
guard.contains(&key)
self.refill_states.lock().contains_key(&key)
}
pub(super) async fn connect_endpoints_round_robin(
@@ -327,62 +361,53 @@ impl MePool {
addr: SocketAddr,
writer_dc: i32,
) {
let endpoint_key = RefillEndpointKey {
dc: writer_dc,
addr,
let Some(registration) = self.lifecycle.try_register() else {
return;
};
let pre_inserted = if let Ok(mut guard) = self.refill_inflight.try_lock() {
if !guard.insert(endpoint_key) {
let dc_key = RefillDcKey {
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();
return;
}
true
} else {
false
};
states.insert(dc_key, None);
self.refill_running.fetch_add(1, Ordering::AcqRel);
}
let pool = Arc::clone(self);
tokio::spawn(async move {
let dc_key = RefillDcKey {
dc: writer_dc,
family: if addr.is_ipv4() {
IpFamily::V4
} else {
IpFamily::V6
},
};
if !pre_inserted {
let mut guard = pool.refill_inflight.lock().await;
if !guard.insert(endpoint_key) {
pool.stats.increment_me_refill_skipped_inflight_total();
return;
let mut run_guard = RefillRunGuard {
pool: Arc::clone(&pool),
key: dc_key,
active: true,
};
self.lifecycle.spawn_registered_producer(registration, async move {
let mut current_addr = addr;
loop {
pool.stats.increment_me_refill_triggered_total();
let restored = pool
.refill_writer_after_loss(current_addr, writer_dc)
.await;
if !restored {
warn!(%current_addr, dc = writer_dc, "ME immediate refill failed");
}
}
{
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);
let Some(next_addr) = run_guard.next_or_finish() else {
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);
});
}
}
+258 -101
View File
@@ -13,10 +13,106 @@ use tracing::{debug, info, warn};
use crate::crypto::SecureRandom;
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;
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 {
fn desired_map_hash(desired_by_dc: &HashMap<i32, HashSet<SocketAddr>>) -> u64 {
let mut hasher = DefaultHasher::new();
@@ -36,36 +132,97 @@ impl MePool {
hasher.finish()
}
fn clear_pending_hardswap_state(&self) {
self.reinit
.pending_hardswap_generation
.store(0, Ordering::Relaxed);
self.reinit
.pending_hardswap_started_at_epoch_secs
.store(0, Ordering::Relaxed);
self.reinit
.pending_hardswap_map_hash
.store(0, Ordering::Relaxed);
self.reinit.warm_generation.store(0, Ordering::Relaxed);
fn reserve_reinit_attempt(
self: &Arc<Self>,
hardswap: bool,
map_hash: u64,
now_epoch_secs: u64,
) -> ReinitReservation {
let mut state = self.reinit.coordinator.lock();
state.desired_map_hash = map_hash;
let previous_generation = state.active_generation;
let mut pending_reused = false;
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) {
self.reinit
.active_generation
.store(generation, Ordering::Relaxed);
self.reinit.warm_generation.store(0, Ordering::Relaxed);
let ws = self.writers.read().await;
for writer in ws.iter() {
if writer.draining.load(Ordering::Relaxed) {
continue;
}
if writer.generation == generation {
writer
.contour
.store(WriterContour::Active.as_u8(), Ordering::Relaxed);
fn commit_reinit_attempt(&self, attempt: &ReinitAttemptGuard) -> bool {
let mut state = self.reinit.coordinator.lock();
if !commit_reinit_state(
&mut state,
attempt.attempt_id,
attempt.generation,
attempt.map_hash,
attempt.hardswap,
) {
return false;
}
if attempt.hardswap {
let writers = self.writers.snapshot();
for writer in writers.iter() {
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(
@@ -387,74 +544,32 @@ impl MePool {
}
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 generation = if hardswap {
let pending_generation = self
.reinit
.pending_hardswap_generation
.load(Ordering::Relaxed);
let pending_started_at = self
.reinit
.pending_hardswap_started_at_epoch_secs
.load(Ordering::Relaxed);
let pending_map_hash = self
.reinit
.pending_hardswap_map_hash
.load(Ordering::Relaxed);
let pending_age_secs = now_epoch_secs.saturating_sub(pending_started_at);
let pending_ttl_expired =
pending_started_at > 0 && pending_age_secs > ME_HARDSWAP_PENDING_TTL_SECS;
let pending_matches_map = pending_map_hash != 0 && pending_map_hash == desired_map_hash;
if pending_generation != 0
&& pending_generation >= previous_generation
&& pending_matches_map
&& !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
};
let reservation =
self.reserve_reinit_attempt(hardswap, desired_map_hash, now_epoch_secs);
let attempt = reservation.attempt;
let previous_generation = attempt.previous_generation;
let generation = attempt.generation;
if reservation.pending_reused {
self.stats.increment_me_hardswap_pending_reuse_total();
debug!(
previous_generation,
generation,
pending_age_secs = reservation.pending_age_secs,
"ME hardswap continues with pending generation"
);
} else if reservation.pending_expired {
self.stats.increment_me_hardswap_pending_ttl_expired_total();
warn!(
previous_generation,
generation,
pending_age_secs = reservation.pending_age_secs,
pending_ttl_secs = ME_HARDSWAP_PENDING_TTL_SECS,
"ME hardswap pending generation expired by TTL; starting fresh generation"
);
}
if hardswap {
self.reinit
.warm_generation
.store(generation, Ordering::Relaxed);
self.warmup_generation_for_all_dcs(rng, generation, &desired_by_dc)
.await;
} else {
@@ -542,8 +657,13 @@ impl MePool {
);
}
if hardswap {
self.promote_warm_generation_to_active(generation).await;
if !self.commit_reinit_attempt(&attempt) {
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
@@ -566,9 +686,6 @@ impl MePool {
drop(writers);
if stale_writer_ids.is_empty() {
if hardswap {
self.clear_pending_hardswap_state();
}
debug!("ME reinit cycle completed with no stale writers");
return true;
}
@@ -606,9 +723,6 @@ impl MePool {
self.remove_writer_and_close_clients(writer_id).await;
}
}
if hardswap {
self.clear_pending_hardswap_state();
}
true
}
@@ -622,7 +736,10 @@ mod tests {
use std::collections::{HashMap, HashSet};
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 {
SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, octet)), port)
@@ -673,4 +790,44 @@ mod tests {
assert_eq!(ratio, 0.0);
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());
}
}
+17 -6
View File
@@ -15,6 +15,8 @@ pub(crate) struct MeApiRefillDcSnapshot {
pub(crate) struct MeApiRefillSnapshot {
pub inflight_endpoints_total: usize,
pub inflight_dc_total: usize,
pub running_dc_total: usize,
pub pending_dc_total: usize,
pub by_dc: Vec<MeApiRefillDcSnapshot>,
}
@@ -56,14 +58,21 @@ pub(crate) struct MeApiDrainGateSnapshot {
impl MePool {
pub(crate) async fn api_refill_snapshot(&self) -> MeApiRefillSnapshot {
let inflight_endpoints_total = self.refill_inflight.lock().await.len();
let inflight_dc_keys = self
.refill_inflight_dc
.lock()
.await
.iter()
let refill_states = self.refill_states.lock();
let inflight_endpoints_total = refill_states
.values()
.map(|pending| 1usize + usize::from(pending.is_some()))
.sum();
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()
.collect::<Vec<RefillDcKey>>();
drop(refill_states);
let mut by_dc_map = HashMap::<(i16, &'static str), usize>::new();
for key in inflight_dc_keys {
@@ -88,6 +97,8 @@ impl MePool {
MeApiRefillSnapshot {
inflight_endpoints_total,
inflight_dc_total: by_dc.len(),
running_dc_total,
pending_dc_total,
by_dc,
}
}
+50 -12
View File
@@ -1,9 +1,10 @@
use std::collections::{BTreeMap, BTreeSet, HashMap};
use std::net::{IpAddr, SocketAddr};
use std::sync::Arc;
use std::sync::atomic::Ordering;
use std::time::Instant;
use super::pool::{MePool, WriterContour};
use super::pool::{MePool, ReinitStatusSnapshot, WriterContour};
use crate::config::{MeBindStaleMode, MeFloorMode, MeSocksKdfPolicy};
use crate::transport::upstream::IpPreference;
@@ -87,8 +88,11 @@ pub(crate) struct MeApiDcPathSnapshot {
pub(crate) struct MeApiRuntimeSnapshot {
pub active_generation: u64,
pub warm_generation: u64,
pub warm_generations: Vec<u64>,
pub pending_hardswap_generation: u64,
pub pending_hardswap_age_secs: Option<u64>,
pub reinit_inflight: usize,
pub reinit_max_concurrency_effective: usize,
pub hardswap_enabled: bool,
pub floor_mode: &'static str,
pub adaptive_floor_idle_secs: u64,
@@ -222,8 +226,16 @@ impl MePool {
}
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 active_generation = self.current_generation();
let active_generation = reinit.active_generation;
let drain_ttl_secs = self
.drain_runtime
.me_pool_drain_ttl_secs
@@ -440,13 +452,19 @@ impl MePool {
}
}
#[allow(dead_code)]
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_epoch_secs = Self::now_epoch_secs();
let pending_started_at = self
.reinit
.pending_hardswap_started_at_epoch_secs
.load(Ordering::Relaxed);
let pending_started_at = reinit.pending_hardswap_started_at_epoch_secs;
let pending_hardswap_age_secs =
(pending_started_at > 0).then_some(now_epoch_secs.saturating_sub(pending_started_at));
@@ -486,13 +504,16 @@ impl MePool {
}
MeApiRuntimeSnapshot {
active_generation: self.reinit.active_generation.load(Ordering::Relaxed),
warm_generation: self.reinit.warm_generation.load(Ordering::Relaxed),
pending_hardswap_generation: self
.reinit
.pending_hardswap_generation
.load(Ordering::Relaxed),
active_generation: reinit.active_generation,
warm_generation: reinit.warm_generations.last().copied().unwrap_or(0),
warm_generations: reinit.warm_generations.clone(),
pending_hardswap_generation: reinit.pending_hardswap_generation,
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),
floor_mode: floor_mode_label(self.floor_mode()),
adaptive_floor_idle_secs: self
@@ -662,6 +683,23 @@ impl MePool {
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 {
+15 -9
View File
@@ -269,7 +269,10 @@ async fn rpc_proxy_req_signal_loop(
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:
// 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
{
stats_signal.increment_me_rpc_proxy_req_signal_failed_total();
let _ = pool.registry.unregister(conn_id).await;
conn_lease.unregister().await;
match error {
ServiceWriterCommandSendError::Closed => return,
ServiceWriterCommandSendError::TimedOut => continue,
@@ -313,7 +316,7 @@ async fn rpc_proxy_req_signal_loop(
.await
{
stats_signal.increment_me_rpc_proxy_req_signal_failed_total();
let _ = pool.registry.unregister(conn_id).await;
conn_lease.unregister().await;
match error {
ServiceWriterCommandSendError::Closed => return,
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();
let _ = pool.registry.unregister(conn_id).await;
conn_lease.unregister().await;
}
}
@@ -394,14 +397,14 @@ impl MePool {
writer_dc: i32,
allow_coverage_override: bool,
) -> Result<()> {
if !self
.can_open_writer_for_contour(contour, allow_coverage_override, writer_dc)
let Some(_writer_open_reservation) = self
.reserve_writer_open(contour, allow_coverage_override, writer_dc)
.await
{
else {
return Err(ProxyError::Proxy(format!(
"ME {contour:?} writer cap reached"
)));
}
};
let secret_len = self.proxy_secret.read().await.secret.len();
if secret_len < 32 {
@@ -415,6 +418,9 @@ impl MePool {
let hs = self
.handshake_only(stream, addr, upstream_egress, rng)
.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 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 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.
let exit = tokio::select! {
biased;
+80 -1
View File
@@ -1,6 +1,7 @@
use std::collections::{HashMap, HashSet};
use std::net::SocketAddr;
use std::sync::Arc;
use std::sync::Mutex as StdMutex;
use std::sync::atomic::{AtomicU8, AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};
@@ -118,6 +119,38 @@ pub struct ConnRegistry {
route_backpressure_high_timeout_ms: AtomicU64,
route_backpressure_high_watermark_pct: AtomicU8,
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 {
@@ -133,15 +166,30 @@ impl ConnRegistry {
Self::with_route_limits(
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(
route_channel_capacity: usize,
route_byte_permits_per_conn: usize,
connection_cleanup_capacity: usize,
) -> Self {
let start = rand::random::<u64>() | 1;
let route_channel_capacity = route_channel_capacity.max(1);
let (cleanup_tx, cleanup_rx) = mpsc::channel(connection_cleanup_capacity.max(1));
Self {
routing: RoutingTable {
map: DashMap::new(),
@@ -168,6 +216,8 @@ impl ConnRegistry {
ROUTE_BACKPRESSURE_HIGH_WATERMARK_PCT,
),
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_byte_permits_per_conn: usize,
) -> 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(
@@ -220,6 +270,29 @@ impl ConnRegistry {
}
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 (tx, rx) = mpsc::channel(self.route_channel_capacity);
self.routing.map.insert(id, tx);
@@ -229,6 +302,12 @@ impl ConnRegistry {
);
(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)]
@@ -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(&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());
}
+129 -53
View File
@@ -2,6 +2,7 @@ use std::sync::Arc;
use std::time::Duration;
use tokio::sync::{mpsc, watch};
use tokio::task::JoinSet;
use tracing::{debug, info, warn};
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(
pool: Arc<MePool>,
rng: Arc<SecureRandom>,
@@ -50,68 +90,104 @@ pub async fn me_reinit_scheduler(
me_ready_tx: watch::Sender<u64>,
) {
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 {
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 coalesce_window = Duration::from_millis(cfg.general.me_reinit_coalesce_window_ms);
if !coalesce_window.is_zero() {
let deadline = tokio::time::Instant::now() + coalesce_window;
loop {
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 max_concurrency = effective_reinit_concurrency(&cfg);
pool.reinit
.max_concurrency_effective
.store(max_concurrency, std::sync::atomic::Ordering::Release);
let reason = if map_change_seen && periodic_seen {
"map-change+periodic"
} else if map_change_seen {
"map-change"
} else {
"periodic"
};
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 pending_ready = pending != 0
&& pending_deadline.is_none_or(|deadline| deadline <= tokio::time::Instant::now());
if pending_ready && tasks.len() < max_concurrency {
let reason = trigger_reason(pending);
pending = 0;
pending_deadline = None;
debug!(reason, max_concurrency, "ME reinit scheduled");
let pool_clone = pool.clone();
let rng_clone = rng.clone();
let me_ready_tx_clone = me_ready_tx.clone();
tokio::spawn(async move {
if pool_clone
pool.reinit
.scheduler_inflight
.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())
.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);
}
}
+1
View File
@@ -108,6 +108,7 @@ pub async fn download_proxy_secret_with_max_len_via_upstream(
let resp = https_get(
proxy_secret_url.unwrap_or("https://core.telegram.org/getProxySecret"),
upstream,
max_len,
)
.await?;
@@ -122,6 +122,7 @@ async fn make_pool(
general.me_route_blocking_send_timeout_ms,
general.me_route_inline_recovery_attempts,
general.me_route_inline_recovery_wait_ms,
16_384,
);
(pool, rng)
@@ -120,6 +120,7 @@ async fn make_pool(
general.me_route_blocking_send_timeout_ms,
general.me_route_inline_recovery_attempts,
general.me_route_inline_recovery_wait_ms,
16_384,
);
(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_inline_recovery_attempts,
general.me_route_inline_recovery_wait_ms,
16_384,
)
}
@@ -1,6 +1,7 @@
use std::collections::HashMap;
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use std::sync::Arc;
use std::sync::atomic::Ordering;
use std::time::{Duration, Instant};
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_inline_recovery_attempts,
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");
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_inline_recovery_attempts,
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_inline_recovery_attempts,
general.me_route_inline_recovery_wait_ms,
16_384,
);
(pool, rng)