From 556affc63add668db950c0794e33098712d00bb1 Mon Sep 17 00:00:00 2001 From: Alexey <247128645+axkurcom@users.noreply.github.com> Date: Mon, 5 Oct 2026 06:02:08 +0300 Subject: [PATCH] ME: Autonomous&Lifecycle-bounded Pool Convergence Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com> --- src/transport/middle_proxy/handshake.rs | 32 +- src/transport/middle_proxy/health.rs | 4 + src/transport/middle_proxy/health/coverage.rs | 109 ++++ src/transport/middle_proxy/health/family.rs | 482 ++---------------- src/transport/middle_proxy/health/monitor.rs | 340 ++++++++++-- .../middle_proxy/health/scheduler.rs | 311 +++++++++++ .../middle_proxy/health/scheduler/tests.rs | 67 +++ src/transport/middle_proxy/pool.rs | 14 +- .../middle_proxy/pool/construction.rs | 7 +- .../middle_proxy/pool/floor_authority.rs | 59 +++ .../pool/floor_authority/tests.rs | 96 ++++ .../middle_proxy/pool/runtime_policy.rs | 13 + .../middle_proxy/pool/selection_policy.rs | 11 + .../middle_proxy/pool/writer_admission.rs | 16 +- src/transport/middle_proxy/pool_lifecycle.rs | 25 + src/transport/middle_proxy/pool_nat.rs | 78 +-- src/transport/middle_proxy/pool_nat/http.rs | 108 ++++ src/transport/middle_proxy/pool_nat/tests.rs | 7 +- .../middle_proxy/pool_nat/tests/freshness.rs | 75 +++ .../middle_proxy/pool_nat/tests/recovery.rs | 139 +++-- .../middle_proxy/pool_reinit/coordination.rs | 6 +- .../middle_proxy/pool_reinit/reconcile.rs | 11 + .../middle_proxy/pool_reinit/tests.rs | 1 + src/transport/middle_proxy/pool_writer.rs | 14 +- .../middle_proxy/pool_writer/overlap.rs | 87 ++++ .../middle_proxy/pool_writer/overlap/tests.rs | 238 +++++++++ .../middle_proxy/pool_writer/publication.rs | 11 +- .../pool_writer/recovery_tests.rs | 14 +- .../middle_proxy/pool_writer/replacement.rs | 175 ++++++- .../replacement/tests/convergence.rs | 198 +++++++ .../middle_proxy/pool_writer/runtime.rs | 3 +- src/transport/middle_proxy/registry.rs | 1 + .../middle_proxy/registry/publication.rs | 29 ++ src/transport/middle_proxy/registry/writer.rs | 30 +- src/transport/middle_proxy/rotation.rs | 177 +++++-- .../middle_proxy/rotation/objective.rs | 55 ++ .../middle_proxy/rotation/retry_tests.rs | 188 +++++++ 37 files changed, 2546 insertions(+), 685 deletions(-) create mode 100644 src/transport/middle_proxy/health/coverage.rs create mode 100644 src/transport/middle_proxy/health/scheduler.rs create mode 100644 src/transport/middle_proxy/health/scheduler/tests.rs create mode 100644 src/transport/middle_proxy/pool/floor_authority.rs create mode 100644 src/transport/middle_proxy/pool/floor_authority/tests.rs create mode 100644 src/transport/middle_proxy/pool_nat/http.rs create mode 100644 src/transport/middle_proxy/pool_nat/tests/freshness.rs create mode 100644 src/transport/middle_proxy/pool_writer/overlap.rs create mode 100644 src/transport/middle_proxy/pool_writer/overlap/tests.rs create mode 100644 src/transport/middle_proxy/pool_writer/replacement/tests/convergence.rs create mode 100644 src/transport/middle_proxy/rotation/objective.rs create mode 100644 src/transport/middle_proxy/rotation/retry_tests.rs diff --git a/src/transport/middle_proxy/handshake.rs b/src/transport/middle_proxy/handshake.rs index b43d240..fe007c3 100644 --- a/src/transport/middle_proxy/handshake.rs +++ b/src/transport/middle_proxy/handshake.rs @@ -265,7 +265,6 @@ impl MePool { let transport_peer_addr = stream.peer_addr().map_err(ProxyError::Io)?; let peer_addr = addr; - let _ = self.maybe_detect_nat_ip(local_addr.ip()).await; let family = if local_addr.ip().is_ipv4() { IpFamily::V4 } else { @@ -324,9 +323,34 @@ impl MePool { None }; - let local_addr_nat = self.translate_our_addr_with_reflection(local_addr, reflected); - let peer_addr_nat = - SocketAddr::new(self.translate_ip_for_nat(peer_addr.ip()), peer_addr.port()); + // Authoritative sources do not wait behind unrelated HTTP discovery. Capture once + // before deriving either side of the tuple so a refresh cannot mix observations. + let detected = if socks_bound_kdf_addr.is_some() + || self.nat_runtime.nat_ip_cfg.is_some() + || reflected.is_some() + { + self.cached_http_nat_ip() + } else { + self.maybe_detect_nat_ip(local_addr.ip()).await + }; + if self.nat_runtime.nat_probe + && local_addr.is_ipv4() + && is_bogon(local_addr.ip()) + && socks_bound_kdf_addr.is_none() + && self.nat_runtime.nat_ip_cfg.is_none() + && reflected.is_none() + && detected.is_none() + { + return Err(ProxyError::Io(std::io::Error::new( + std::io::ErrorKind::WouldBlock, + "ME NAT source observation is unavailable or expired", + ))); + } + let local_addr_nat = self.translate_our_addr_with_nat(local_addr, reflected, detected); + let peer_addr_nat = SocketAddr::new( + Self::translate_ip_with_nat(peer_addr.ip(), self.nat_runtime.nat_ip_cfg.or(detected)), + peer_addr.port(), + ); let client_addr_for_kdf = socks_bound_kdf_addr.unwrap_or(local_addr_nat); if let Some(upstream_info) = upstream_egress { record_upstream_bnd_status( diff --git a/src/transport/middle_proxy/health.rs b/src/transport/middle_proxy/health.rs index 40c8743..f71b9cf 100644 --- a/src/transport/middle_proxy/health.rs +++ b/src/transport/middle_proxy/health.rs @@ -20,6 +20,8 @@ use super::pool::{ MeFamilyRuntimeState, RefillTargetKey, WriterContour, WriterOpenIntent, WriterRole, }; use super::pool_writer::WriterReplacementPurpose; +// Bounded cross-family coverage transfers preserve the donor's current base floor. +mod coverage; // Bounds reconnect jitter to at most half of the current backoff. const JITTER_FRAC_NUM: u64 = 2; @@ -110,6 +112,8 @@ impl Drop for ScheduledReconnects<'_> { // Periodic family health monitor orchestration. mod monitor; +// Shared-family fair scheduling and independently tracked writer-operation quanta. +mod scheduler; // Draining-writer deadline enforcement. mod drain; // Per-family reconnect scheduling and health state. diff --git a/src/transport/middle_proxy/health/coverage.rs b/src/transport/middle_proxy/health/coverage.rs new file mode 100644 index 0000000..c7e327c --- /dev/null +++ b/src/transport/middle_proxy/health/coverage.rs @@ -0,0 +1,109 @@ +use super::*; +use std::sync::atomic::Ordering; + +/// Selects globally because the Active cap is shared across DCs and address families. +pub(super) async fn transfer_coverage( + pool: &Arc, + rng: &Arc, + dc: i32, + endpoint: SocketAddr, + required: usize, +) -> Option { + let capacity = pool + .reserve_writer_open( + WriterContour::Active, + WriterOpenIntent::Replacement, + dc, + endpoint, + ) + .await?; + drop(capacity); + let authority = pool.floor_authority(); + let endpoints = pool.endpoint_snapshot.load_full(); + let generation = pool.current_generation(); + let writers = pool.writers.read().await; + let activity = pool.registry.writer_activity_snapshot().await; + let mut counts = HashMap::new(); + for writer in writers.iter() { + if !writer.draining.load(Ordering::Acquire) + && writer.generation == generation + && WriterContour::from_u8(writer.contour.load(Ordering::Acquire)) + == WriterContour::Active + && endpoints.contains_dc_endpoint(writer.writer_dc, writer.addr) + { + *counts + .entry((writer.writer_dc, writer.addr.is_ipv4())) + .or_insert(0usize) += 1; + } + } + let mut candidates = Vec::new(); + for writer in writers.iter() { + if writer.draining.load(Ordering::Acquire) + || WriterContour::from_u8(writer.contour.load(Ordering::Acquire)) + != WriterContour::Active + { + continue; + } + let role = WriterRole::from_writer(writer); + let obsolete = writer.generation != generation + || !endpoints.contains_dc_endpoint(writer.writer_dc, writer.addr); + let protected = pool.required_writers_for_dc_with_floor_mode( + endpoints + .endpoints_for_dc_family(writer.writer_dc, role.family) + .len(), + false, + ); + if !obsolete + && counts + .get(&(writer.writer_dc, writer.addr.is_ipv4())) + .copied() + .unwrap_or(0) + <= protected + { + continue; + } + let busy = activity + .bound_clients_by_writer + .get(&writer.id) + .copied() + .unwrap_or(0) + > 0; + candidates.push((busy, !obsolete, writer.created_at, writer.id, role)); + } + drop(writers); + candidates + .sort_unstable_by_key(|candidate| (candidate.0, candidate.1, candidate.2, candidate.3)); + for (busy, _, _, id, role) in candidates { + let reservation = if busy { + pool.registry + .try_reserve_writer_replacement_preserving_clients(id) + .await + } else { + pool.registry.try_reserve_writer_replacement(id).await + }; + let Some(mut reservation) = reservation else { + continue; + }; + return Some(matches!( + tokio::time::timeout( + pool.reconnect_runtime.me_one_timeout, + pool.replace_writer_with_generation_contour_for_dc( + endpoint, + rng.as_ref(), + generation, + WriterContour::Active, + dc, + role, + WriterReplacementPurpose::CoverageTransfer { + authority, + receiver_floor: required + }, + &mut reservation, + ), + ) + .await, + Ok(Ok(())) + )); + } + None +} diff --git a/src/transport/middle_proxy/health/family.rs b/src/transport/middle_proxy/health/family.rs index a23a02d..af5a51d 100644 --- a/src/transport/middle_proxy/health/family.rs +++ b/src/transport/middle_proxy/health/family.rs @@ -1,449 +1,67 @@ +use super::scheduler::HealthObservation; use super::*; -pub(super) async fn check_family( - family: IpFamily, - pool: &Arc, - rng: &Arc, - backoff: &mut HashMap<(i32, IpFamily), u64>, - next_attempt: &mut HashMap<(i32, IpFamily), Instant>, - inflight: &mut HashMap<(i32, IpFamily), usize>, - outage_backoff: &mut HashMap<(i32, IpFamily), u64>, - outage_next_attempt: &mut HashMap<(i32, IpFamily), Instant>, - single_endpoint_outage: &mut HashSet<(i32, IpFamily)>, - shadow_rotate_deadline: &mut HashMap<(i32, IpFamily), Instant>, - idle_refresh_next_attempt: &mut HashMap<(i32, IpFamily), Instant>, - floor_warn_next_allowed: &mut HashMap<(i32, IpFamily), Instant>, -) -> bool { - let enabled = match family { - IpFamily::V4 => pool.decision.ipv4_me, - IpFamily::V6 => pool.decision.ipv6_me, +/// Takes a control-plane observation without awaiting any network operation. +pub(super) async fn check_family(family: IpFamily, pool: &Arc) -> HealthObservation { + let authority = pool.floor_authority(); + let generation = pool.current_generation(); + let snapshot = pool.endpoint_snapshot.load_full(); + let enabled = if family == IpFamily::V4 { + pool.decision.ipv4_me + } else { + pool.decision.ipv6_me }; - if !enabled { - return false; - } - - let mut family_degraded = false; - - let mut dc_endpoints = HashMap::>::new(); - let endpoint_snapshot = pool.endpoint_snapshot.load(); - let endpoint_revision = endpoint_snapshot.revision; - let map_guard = match family { - IpFamily::V4 => &endpoint_snapshot.map_v4, - IpFamily::V6 => &endpoint_snapshot.map_v6, + let source = if family == IpFamily::V4 { + &snapshot.map_v4 + } else { + &snapshot.map_v6 }; - for (dc, addrs) in map_guard.iter() { - let entry = dc_endpoints.entry(*dc).or_default(); - for (ip, port) in addrs.iter().copied() { - entry.push(SocketAddr::new(ip, port)); + let mut endpoints = HashMap::new(); + if enabled { + for (dc, addresses) in source { + let mut addresses: Vec<_> = addresses + .iter() + .map(|(ip, port)| SocketAddr::new(*ip, *port)) + .collect(); + addresses.sort_unstable(); + addresses.dedup(); + endpoints.insert(*dc, addresses); } } - drop(endpoint_snapshot); - for endpoints in dc_endpoints.values_mut() { - endpoints.sort_unstable(); - endpoints.dedup(); - } - let reconnect_budget = health_reconnect_budget(pool, dc_endpoints.len()); - let reconnect_sem = Arc::new(Semaphore::new(reconnect_budget)); - - if pool.floor_mode() == MeFloorMode::Static {} - - let mut live_addr_counts = HashMap::<(i32, SocketAddr), usize>::new(); - let mut live_writer_ids_by_addr = HashMap::<(i32, SocketAddr), Vec>::new(); - let active_generation = pool.current_generation(); - for writer in pool - .writers - .read() - .await - .iter() - .filter(|w| !w.draining.load(std::sync::atomic::Ordering::Relaxed)) - { - if !matches!( - crate::transport::middle_proxy::pool::WriterContour::from_u8( - writer.contour.load(std::sync::atomic::Ordering::Relaxed), - ), - crate::transport::middle_proxy::pool::WriterContour::Active - ) { - continue; - } - if writer.generation != active_generation { - continue; - } - if !dc_endpoints - .get(&writer.writer_dc) - .is_some_and(|endpoints| endpoints.contains(&writer.addr)) + let mut counts = HashMap::new(); + let mut ids = HashMap::<_, Vec<_>>::new(); + for writer in pool.writers.read().await.iter() { + if !writer.draining.load(std::sync::atomic::Ordering::Acquire) + && writer.generation == generation + && WriterContour::from_u8(writer.contour.load(std::sync::atomic::Ordering::Acquire)) + == WriterContour::Active + && endpoints + .get(&writer.writer_dc) + .is_some_and(|addresses| addresses.contains(&writer.addr)) { - continue; + *counts + .entry((writer.writer_dc, writer.addr)) + .or_insert(0usize) += 1; + ids.entry((writer.writer_dc, writer.addr)) + .or_default() + .push(writer.id); } - let key = (writer.writer_dc, writer.addr); - *live_addr_counts.entry(key).or_insert(0) += 1; - live_writer_ids_by_addr - .entry(key) - .or_default() - .push(writer.id); } - let writer_idle_since = pool.registry.writer_idle_since_snapshot().await; - let bound_clients_by_writer = pool + let idle = pool.registry.writer_idle_since_snapshot().await; + let bound = pool .registry .writer_activity_snapshot() .await .bound_clients_by_writer; - let floor_plan = build_family_floor_plan( - pool, - family, - &dc_endpoints, - &live_addr_counts, - &live_writer_ids_by_addr, - &bound_clients_by_writer, - ) - .await; - pool.set_adaptive_floor_runtime_caps( - floor_plan.active_cap_configured_total, - floor_plan.active_cap_effective_total, - floor_plan.warm_cap_configured_total, - floor_plan.warm_cap_effective_total, - floor_plan.target_writers_total, - floor_plan.active_writers_current, - floor_plan.warm_writers_current, - ); - let floor_targets_by_dc = Arc::new( - floor_plan - .by_dc - .iter() - .map(|(dc, entry)| (*dc, entry.target_required)) - .collect::>(), - ); - let live_writer_ids_by_addr = Arc::new(live_writer_ids_by_addr); - 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::::new(); - let mut scheduled_reconnects = ScheduledReconnects { - inflight, - keys: Vec::new(), - }; - - for (dc, endpoints) in dc_endpoints { - if endpoints.is_empty() { - continue; - } - let key = (dc, family); - let required = floor_plan - .by_dc - .get(&dc) - .map(|entry| entry.target_required) - .unwrap_or_else(|| { - pool.required_writers_for_dc_with_floor_mode(endpoints.len(), false) - }); - let alive = endpoints - .iter() - .map(|addr| *live_addr_counts.get(&(dc, *addr)).unwrap_or(&0)) - .sum::(); - - if endpoints.len() == 1 && pool.single_endpoint_outage_mode_enabled() && alive == 0 { - family_degraded = true; - if single_endpoint_outage.insert(key) { - pool.stats.increment_me_single_endpoint_outage_enter_total(); - warn!( - dc = %dc, - ?family, - required, - endpoint_count = endpoints.len(), - "Single-endpoint DC outage detected" - ); - } - - recover_single_endpoint_outage( - pool, - rng, - key, - endpoints[0], - required, - outage_backoff, - outage_next_attempt, - &reconnect_sem, - ) - .await; - continue; - } - - if single_endpoint_outage.remove(&key) { - pool.stats.increment_me_single_endpoint_outage_exit_total(); - outage_backoff.remove(&key); - outage_next_attempt.remove(&key); - shadow_rotate_deadline.remove(&key); - idle_refresh_next_attempt.remove(&key); - info!( - dc = %dc, - ?family, - alive, - required, - endpoint_count = endpoints.len(), - "Single-endpoint DC outage recovered" - ); - } - - maybe_refresh_idle_writer_for_dc( - pool, - rng, - key, - dc, - family, - &endpoints, - alive, - required, - live_writer_ids_by_addr.as_ref(), - writer_idle_since.as_ref(), - bound_clients_by_writer.as_ref(), - idle_refresh_next_attempt, - ) - .await; - - if alive >= required { - maybe_rotate_single_endpoint_shadow( - pool, - rng, - key, - dc, - family, - &endpoints, - alive, - required, - live_writer_ids_by_addr.as_ref(), - bound_clients_by_writer.as_ref(), - shadow_rotate_deadline, - ) - .await; - continue; - } - let missing = required - alive; - family_degraded = true; - - let now = Instant::now(); - if reconnect_sem.available_permits() == 0 { - let base_ms = pool.reconnect_runtime.me_reconnect_backoff_base.as_millis() as u64; - let next_ms = (*backoff.get(&key).unwrap_or(&base_ms)).max(base_ms); - let jitter = next_ms / JITTER_FRAC_NUM; - let wait = Duration::from_millis(next_ms) - + Duration::from_millis(rand::rng().random_range(0..=jitter.max(1))); - next_attempt.insert(key, now + wait); - debug!( - dc = %dc, - ?family, - alive, - required, - endpoint_count = endpoints.len(), - reconnect_budget, - "Skipping reconnect due to per-tick health reconnect budget" - ); - continue; - } - if let Some(ts) = next_attempt.get(&key) - && now < *ts - { - continue; - } - - let max_concurrent = pool - .reconnect_runtime - .me_reconnect_max_concurrent_per_dc - .max(1) as usize; - if scheduled_reconnects.current(&key) >= max_concurrent { - continue; - } - if pool - .has_refill_inflight_for_target(RefillTargetKey { - dc, - family, - generation: pool.current_generation(), - endpoint_revision, - contour: WriterContour::Active, - }) - .await - { - debug!( - dc = %dc, - ?family, - alive, - required, - endpoint_count = endpoints.len(), - "Skipping health reconnect: immediate refill is already in flight for this DC group" - ); - continue; - } - scheduled_reconnects.reserve(key); - let pool_for_reconnect = pool.clone(); - let rng_for_reconnect = rng.clone(); - let reconnect_sem_for_dc = reconnect_sem.clone(); - let endpoints_for_dc = endpoints.clone(); - let live_writer_ids_by_addr_for_dc = live_writer_ids_by_addr.clone(); - let writer_idle_since_for_dc = writer_idle_since.clone(); - let bound_clients_by_writer_for_dc = bound_clients_by_writer.clone(); - let floor_targets_by_dc_for_reconnect = floor_targets_by_dc.clone(); - let active_cap_effective_total = floor_plan.active_cap_effective_total; - reconnect_set.spawn(async move { - let mut restored = 0usize; - for _ in 0..missing { - let Ok(reconnect_permit) = reconnect_sem_for_dc.clone().try_acquire_owned() else { - break; - }; - let base_req = pool_for_reconnect - .required_writers_for_dc_with_floor_mode(endpoints_for_dc.len(), false); - if pool_for_reconnect.active_contour_writer_count_total().await - >= active_cap_effective_total - { - let swapped = maybe_swap_idle_writer_for_cap( - &pool_for_reconnect, - &rng_for_reconnect, - dc, - family, - &endpoints_for_dc, - required, - floor_targets_by_dc_for_reconnect.as_ref(), - live_writer_ids_by_addr_for_dc.as_ref(), - writer_idle_since_for_dc.as_ref(), - bound_clients_by_writer_for_dc.as_ref(), - ) - .await; - if swapped { - pool_for_reconnect - .stats - .increment_me_floor_swap_idle_total(); - restored += 1; - continue; - } - - if alive + restored >= base_req { - pool_for_reconnect - .stats - .increment_me_floor_cap_block_total(); - pool_for_reconnect - .stats - .increment_me_floor_swap_idle_failed_total(); - debug!( - dc = %dc, - ?family, - alive, - required, - active_cap_effective_total, - "Adaptive floor cap reached, reconnect attempt blocked" - ); - break; - } - } - pool_for_reconnect.stats.increment_me_reconnect_attempt(); - let intent = if alive + restored < base_req { - WriterOpenIntent::Coverage - } else { - WriterOpenIntent::Normal - }; - let res = tokio::time::timeout( - pool_for_reconnect.reconnect_runtime.me_one_timeout, - pool_for_reconnect.connect_endpoints_round_robin_with_generation_contour( - dc, - &endpoints_for_dc, - rng_for_reconnect.as_ref(), - pool_for_reconnect.current_generation(), - WriterContour::Active, - intent, - ), - ) - .await; - match res { - Ok(true) => { - restored += 1; - pool_for_reconnect.stats.increment_me_reconnect_success(); - } - Ok(false) => { - debug!(dc = %dc, ?family, "ME round-robin reconnect failed") - } - Err(_) => { - debug!(dc = %dc, ?family, "ME reconnect timed out"); - } - } - drop(reconnect_permit); - } - - FamilyReconnectOutcome { - key, - dc, - family, - required, - endpoint_count: endpoints_for_dc.len(), - } - }); + let floor = build_family_floor_plan(pool, family, &endpoints, &counts, &ids, &bound).await; + HealthObservation { + authority, + generation, + floor, + ids, + idle, + bound, } - - while let Some(joined) = reconnect_set.join_next().await { - let outcome = match joined { - Ok(outcome) => outcome, - Err(join_error) => { - debug!(error = %join_error, "Health reconnect task failed"); - continue; - } - }; - let now = Instant::now(); - let now_alive = live_active_writers_for_dc_family(pool, outcome.dc, outcome.family).await; - if now_alive >= outcome.required { - info!( - dc = %outcome.dc, - family = ?outcome.family, - alive = now_alive, - required = outcome.required, - endpoint_count = outcome.endpoint_count, - "ME writer floor restored for DC" - ); - backoff.insert( - outcome.key, - pool.reconnect_runtime.me_reconnect_backoff_base.as_millis() as u64, - ); - let jitter = pool.reconnect_runtime.me_reconnect_backoff_base.as_millis() as u64 - / JITTER_FRAC_NUM; - let wait = pool.reconnect_runtime.me_reconnect_backoff_base - + Duration::from_millis(rand::rng().random_range(0..=jitter.max(1))); - next_attempt.insert(outcome.key, now + wait); - } else { - let curr = *backoff - .get(&outcome.key) - .unwrap_or(&(pool.reconnect_runtime.me_reconnect_backoff_base.as_millis() as u64)); - let next_ms = (curr.saturating_mul(2)) - .min(pool.reconnect_runtime.me_reconnect_backoff_cap.as_millis() as u64); - backoff.insert(outcome.key, next_ms); - let jitter = next_ms / JITTER_FRAC_NUM; - let wait = Duration::from_millis(next_ms) - + Duration::from_millis(rand::rng().random_range(0..=jitter.max(1))); - next_attempt.insert(outcome.key, now + wait); - if pool.is_runtime_ready() { - let warn_cooldown = pool.warn_rate_limit_duration(); - if should_emit_rate_limited_warn( - floor_warn_next_allowed, - outcome.key, - now, - warn_cooldown, - ) { - warn!( - dc = %outcome.dc, - family = ?outcome.family, - alive = now_alive, - required = outcome.required, - endpoint_count = outcome.endpoint_count, - backoff_ms = next_ms, - "DC writer floor is below required level, scheduled reconnect" - ); - } - } else { - info!( - dc = %outcome.dc, - family = ?outcome.family, - alive = now_alive, - required = outcome.required, - endpoint_count = outcome.endpoint_count, - backoff_ms = next_ms, - "DC writer floor is below required level during startup, scheduled reconnect" - ); - } - } - } - - family_degraded } pub(super) fn health_reconnect_budget(pool: &Arc, dc_groups: usize) -> usize { diff --git a/src/transport/middle_proxy/health/monitor.rs b/src/transport/middle_proxy/health/monitor.rs index 1388613..df0cfae 100644 --- a/src/transport/middle_proxy/health/monitor.rs +++ b/src/transport/middle_proxy/health/monitor.rs @@ -1,59 +1,293 @@ +use super::scheduler::{ + HealthJob, HealthObservation, HealthScheduler, JobKind, JobOutcome, idle_refresh_ready, run_job, +}; use super::*; +/// Observes family health independently of bounded, lifecycle-owned network operations. pub async fn me_health_monitor(pool: Arc, rng: Arc, _min_connections: usize) { - let mut backoff: HashMap<(i32, IpFamily), u64> = HashMap::new(); - let mut next_attempt: HashMap<(i32, IpFamily), Instant> = HashMap::new(); - let mut inflight: HashMap<(i32, IpFamily), usize> = HashMap::new(); - let mut outage_backoff: HashMap<(i32, IpFamily), u64> = HashMap::new(); - let mut outage_next_attempt: HashMap<(i32, IpFamily), Instant> = HashMap::new(); - let mut single_endpoint_outage: HashSet<(i32, IpFamily)> = HashSet::new(); - let mut shadow_rotate_deadline: HashMap<(i32, IpFamily), Instant> = HashMap::new(); - let mut idle_refresh_next_attempt: HashMap<(i32, IpFamily), Instant> = HashMap::new(); - let mut floor_warn_next_allowed: HashMap<(i32, IpFamily), Instant> = HashMap::new(); - let mut drain_warn_next_allowed: HashMap = HashMap::new(); - let mut degraded_interval = true; + let mut scheduler = HealthScheduler::default(); + let mut tasks = JoinSet::new(); + let mut observations = HashMap::>::new(); + let mut drain_warn_next_allowed = HashMap::new(); + let mut observe = true; + let mut health_tick = true; + let mut next_tick = tokio::time::Instant::now(); + loop { - let interval = if degraded_interval { - pool.health_interval_unhealthy() - } else { - pool.health_interval_healthy() - }; - tokio::time::sleep(interval).await; - pool.prune_closed_writers().await; - pool.sweep_endpoint_quarantine().await; - reap_draining_writers(&pool, &mut drain_warn_next_allowed).await; - let v4_degraded = check_family( - IpFamily::V4, - &pool, - &rng, - &mut backoff, - &mut next_attempt, - &mut inflight, - &mut outage_backoff, - &mut outage_next_attempt, - &mut single_endpoint_outage, - &mut shadow_rotate_deadline, - &mut idle_refresh_next_attempt, - &mut floor_warn_next_allowed, - ) - .await; - let v6_degraded = check_family( - IpFamily::V6, - &pool, - &rng, - &mut backoff, - &mut next_attempt, - &mut inflight, - &mut outage_backoff, - &mut outage_next_attempt, - &mut single_endpoint_outage, - &mut shadow_rotate_deadline, - &mut idle_refresh_next_attempt, - &mut floor_warn_next_allowed, - ) - .await; - update_family_runtime_state(&pool, IpFamily::V4, v4_degraded); - update_family_runtime_state(&pool, IpFamily::V6, v6_degraded); - degraded_interval = v4_degraded || v6_degraded; + if observe { + if health_tick { + pool.prune_closed_writers().await; + pool.sweep_endpoint_quarantine().await; + reap_draining_writers(&pool, &mut drain_warn_next_allowed).await; + } + let mut target_total = 0usize; + let mut effective_cap = 0usize; + let mut cap_observation = None; + let mut degraded = false; + for family in [IpFamily::V4, IpFamily::V6] { + let observation = Arc::new(check_family(family, &pool).await); + let family_degraded = observation + .floor + .by_dc + .values() + .any(|entry| entry.alive < entry.target_required); + if health_tick { + update_family_runtime_state(&pool, family, family_degraded); + } + degraded |= family_degraded; + target_total += observation + .floor + .by_dc + .values() + .map(|entry| entry.target_required) + .sum::(); + effective_cap = effective_cap.max(observation.floor.active_cap_effective_total); + if cap_observation.is_none() { + cap_observation = Some(observation.clone()); + } + observations.insert(family, observation); + } + if let Some(observation) = cap_observation + && observations + .values() + .all(|other| other.authority == observation.authority) + { + let floor = &observation.floor; + pool.set_adaptive_floor_runtime_caps( + observation.authority, + floor.active_cap_configured_total, + effective_cap.max(target_total), + floor.warm_cap_configured_total, + floor.warm_cap_effective_total, + target_total, + floor.active_writers_current, + floor.warm_writers_current, + ); + } + scheduler.observe(&observations); + for (key, state) in &mut scheduler.groups { + if health_tick { + state.deferred = false; + } + let Some(entry) = observations + .get(&key.1) + .and_then(|o| o.floor.by_dc.get(&key.0)) + else { + continue; + }; + let outage = entry.endpoints.len() == 1 + && entry.alive == 0 + && pool.single_endpoint_outage_mode_enabled(); + if outage != state.outage { + if outage { + pool.stats.increment_me_single_endpoint_outage_enter_total(); + warn!(dc = key.0, family = ?key.1, required = entry.target_required, + endpoint_count = entry.endpoints.len(), "Single-endpoint DC outage detected"); + } else { + pool.stats.increment_me_single_endpoint_outage_exit_total(); + info!(dc = key.0, family = ?key.1, alive = entry.alive, required = entry.target_required, + endpoint_count = entry.endpoints.len(), "Single-endpoint DC outage recovered"); + } + state.outage = outage; + // Outage retries and ordinary missing-floor rounds use different bounds. + // A successful first writer must not inherit an outage failure delay. + state.backoff_ms = 0; + state.due = None; + state.round_left = 0; + state.round_active = false; + } + if state.round_active + && (state.round_left == 0 || entry.alive >= entry.target_required) + { + let (base, cap) = if state.outage { + pool.single_endpoint_outage_backoff_bounds_ms() + } else { + ( + pool.reconnect_runtime.me_reconnect_backoff_base.as_millis() as u64, + pool.reconnect_runtime.me_reconnect_backoff_cap.as_millis() as u64, + ) + }; + let restored = entry.alive >= entry.target_required; + state.backoff_ms = if restored { + base + } else { + state.backoff_ms.max(base).saturating_mul(2).min(cap) + }; + let wait = Duration::from_millis(state.backoff_ms) + + Duration::from_millis( + rand::rng() + .random_range(0..=(state.backoff_ms / JITTER_FRAC_NUM).max(1)), + ); + state.due = Some(Instant::now() + wait); + state.round_active = false; + if restored { + info!(dc = key.0, family = ?key.1, alive = entry.alive, required = entry.target_required, + endpoint_count = entry.endpoints.len(), "ME writer floor restored for DC"); + } + if !restored && state.warn_due.is_none_or(|due| due <= Instant::now()) { + state.warn_due = Some(Instant::now() + pool.warn_rate_limit_duration()); + if state.outage { + warn!(dc = key.0, family = ?key.1, endpoint = %entry.endpoints[0], required = entry.target_required, + backoff_ms = state.backoff_ms, "Single-endpoint outage reconnect scheduled"); + } else if pool.is_runtime_ready() { + warn!(dc = key.0, family = ?key.1, alive = entry.alive, required = entry.target_required, + endpoint_count = entry.endpoints.len(), backoff_ms = state.backoff_ms, + "DC writer floor is below required level, scheduled reconnect"); + } else { + info!(dc = key.0, family = ?key.1, alive = entry.alive, required = entry.target_required, + endpoint_count = entry.endpoints.len(), backoff_ms = state.backoff_ms, + "DC writer floor is below required level during startup, scheduled reconnect"); + } + } + } + } + if health_tick { + next_tick = tokio::time::Instant::now() + + if degraded { + pool.health_interval_unhealthy() + } else { + pool.health_interval_healthy() + }; + } + observe = false; + health_tick = false; + } + + let budget = family::health_reconnect_budget(&pool, scheduler.groups.len()); + for _ in 0..scheduler.queue.len() { + if tasks.len() >= budget { + break; + } + let Some(key) = scheduler.queue.pop_front() else { + break; + }; + scheduler.queue.push_back(key); + if scheduler.busy(key) { + continue; + } + let Some(observation) = observations.get(&key.1).cloned() else { + continue; + }; + if pool.floor_authority() != observation.authority + || pool.current_generation() != observation.generation + { + continue; + } + let Some(entry) = observation.floor.by_dc.get(&key.0).cloned() else { + continue; + }; + if entry.endpoints.is_empty() { + continue; + } + let Some(state) = scheduler.groups.get_mut(&key) else { + continue; + }; + if state.deferred { + continue; + } + let now = Instant::now(); + let recover = + entry.alive < entry.target_required && state.due.is_none_or(|due| due <= now); + let refresh = state.refresh_due.is_none_or(|due| due <= now) + && idle_refresh_ready(&observation, &entry); + let mut kind = if refresh && (!recover || state.prefer_refresh) { + JobKind::Refresh + } else if recover { + JobKind::Recover + } else if entry.alive >= entry.target_required + && entry.endpoints.len() == 1 + && pool.single_endpoint_shadow_rotate_interval().is_some() + && state.shadow_due.is_none_or(|due| due <= now) + { + JobKind::Shadow + } else { + continue; + }; + if kind == JobKind::Recover { + if pool + .has_refill_inflight_for_target(RefillTargetKey { + dc: key.0, + family: key.1, + generation: observation.generation, + endpoint_revision: observation.authority.0, + contour: WriterContour::Active, + }) + .await + { + continue; + } + // Preflight distinguishes capacity pressure from unreachable endpoints without + // retaining an opening token while a group waits in the fair queue. + let reservation = pool + .reserve_writer_open( + WriterContour::Active, + WriterOpenIntent::Coverage, + key.0, + entry.endpoints[0], + ) + .await; + if reservation.is_none() { + kind = JobKind::Transfer; + } + drop(reservation); + if !state.round_active { + state.round = state.round.wrapping_add(1); + state.round_left = if state.outage { + 1 + } else { + entry.target_required.saturating_sub(entry.alive) + }; + state.round_active = true; + } + state.prefer_refresh = true; + } else if kind == JobKind::Refresh { + state.prefer_refresh = false; + } + let job = HealthJob { + key, + kind, + observation, + entry, + round: state.round, + }; + let future = run_job(pool.clone(), rng.clone(), job.clone()); + let Ok(tracked) = pool.lifecycle.track_producer(future) else { + return; + }; + let id = tasks.spawn(tracked).id(); + scheduler.owners.insert(id, job); + } + + let now = Instant::now(); + let deadline = scheduler + .groups + .values() + .flat_map(|state| [state.due, state.refresh_due, state.shadow_due]) + .flatten() + .filter(|due| *due > now) + .min() + .map(tokio::time::Instant::from_std) + .map_or(next_tick, |due| due.min(next_tick)); + tokio::select! { + joined = tasks.join_next_with_id(), if !tasks.is_empty() => { + let Some(joined) = joined else { continue; }; + let (id, output) = match joined { + Ok((id, output)) => (id, output), + Err(error) => { + debug!(error = %error, "Health writer operation failed"); + (error.id(), Some(JobOutcome::Completed(None))) + } + }; + if let Some(job) = scheduler.owners.remove(&id) + && let Some(deadline) = output { + scheduler.completed(&pool, job, deadline); + } + observe = true; + } + _ = tokio::time::sleep_until(deadline) => { + observe = true; + health_tick = tokio::time::Instant::now() >= next_tick; + } + } } } diff --git a/src/transport/middle_proxy/health/scheduler.rs b/src/transport/middle_proxy/health/scheduler.rs new file mode 100644 index 0000000..25d4a73 --- /dev/null +++ b/src/transport/middle_proxy/health/scheduler.rs @@ -0,0 +1,311 @@ +use super::*; +use std::collections::VecDeque; +use tokio::task::Id; +#[cfg(test)] +mod tests; + +/// Exact fairness and single-owner domain; DC sign remains part of the identity. +pub(super) type Group = (i32, IpFamily); + +/// Immutable observation is fenced again before dispatch and when accepting completion. +pub(super) struct HealthObservation { + /// Endpoint and floor-policy revisions used to derive this plan. + pub(super) authority: (u64, u64), + /// Active generation whose writers contribute coverage. + pub(super) generation: u64, + /// Family-local targets and global capacity projection. + pub(super) floor: FamilyFloorPlan, + /// Current non-draining writer IDs grouped by DC and endpoint. + pub(super) ids: HashMap<(i32, SocketAddr), Vec>, + /// Registry idle timestamps used only for refresh selection. + pub(super) idle: HashMap, + /// Captured bindings; commit revalidates victim idleness when required. + pub(super) bound: HashMap, +} + +/// Mutually exclusive network work selected for one group quantum. +#[derive(Clone, Copy, PartialEq, Eq)] +pub(super) enum JobKind { + /// Restores one missing writer through ordinary bounded admission. + Recover, + /// Restores coverage by replacing an eligible donor. + Transfer, + /// Replaces an idle writer before upstream expiry. + Refresh, + /// Applies the configured single-endpoint rotation cadence. + Shadow, +} + +/// Resource deferral consumes no network retry budget and preserves the eligibility deadline. +pub(super) enum JobOutcome { + /// Work consumed its quantum and may return a refresh or shadow deadline. + Completed(Option), + /// Capacity or authority prevented an operation; do not charge network backoff. + Deferred, +} + +/// Captured authority accompanies both successful results and JoinError ownership. +#[derive(Clone)] +pub(super) struct HealthJob { + /// Group holding the unique inflight slot. + pub(super) key: Group, + /// Determines which deadline or round may change on completion. + pub(super) kind: JobKind, + /// Immutable selection evidence shared by jobs from one observation. + pub(super) observation: Arc, + /// The selected group's target and endpoint candidates. + pub(super) entry: DcFloorPlanEntry, + /// Recovery round identity; refresh and shadow operations do not consume it. + pub(super) round: u64, +} + +/// Retry and cooldown state owned exclusively by the monitor. +#[derive(Default)] +pub(super) struct GroupState { + /// Endpoint, policy and generation fence for every mutable field below. + pub(super) authority: Option<((u64, u64), u64)>, + /// Earliest recovery eligibility; resource deferral leaves it unchanged. + pub(super) due: Option, + /// Idle-refresh cooldown independent of missing-floor recovery. + pub(super) refresh_due: Option, + /// Single-endpoint rotation cooldown. + pub(super) shadow_due: Option, + /// Backoff advances once per completed recovery round. + pub(super) backoff_ms: u64, + /// Remaining writer operations in the current round. + pub(super) round_left: usize, + /// Distinguishes a finalized round from one that has no operations left yet. + pub(super) round_active: bool, + /// Monotonic identity prevents late outcomes from changing a newer round. + pub(super) round: u64, + /// Alternates independently eligible refresh and recovery work. + pub(super) prefer_refresh: bool, + /// Selects the single-endpoint outage retry namespace. + pub(super) outage: bool, + /// Rate limit for missing-floor warnings. + pub(super) warn_due: Option, + /// Resource-blocked groups wait until the next health tick. + pub(super) deferred: bool, +} + +/// One queue and one owner table bound work across both address families. +#[derive(Default)] +pub(super) struct HealthScheduler { + /// Persistent cursor order survives observations and worker completions. + pub(super) queue: VecDeque, + /// Coalesced work and deadlines for configured groups only. + pub(super) groups: HashMap, + /// Task IDs release group ownership even after panic or abort. + pub(super) owners: HashMap, +} + +impl HealthScheduler { + /// Refreshes desired groups without resetting the round-robin cursor. + pub(super) fn observe(&mut self, observations: &HashMap>) { + let mut current = HashSet::new(); + for (family, observation) in observations { + for dc in observation.floor.by_dc.keys() { + let key = (*dc, *family); + current.insert(key); + let state = self.groups.entry(key).or_default(); + let authority = (observation.authority, observation.generation); + if state.authority != Some(authority) { + *state = GroupState { + authority: Some(authority), + ..GroupState::default() + }; + } + if !self.queue.contains(&key) { + self.queue.push_back(key); + } + } + } + self.queue.retain(|key| current.contains(key)); + self.groups.retain(|key, _| current.contains(key)); + } + + /// Includes obsolete tasks until their actual completion releases ownership. + pub(super) fn busy(&self, key: Group) -> bool { + self.owners.values().any(|job| job.key == key) + } + + /// Accepts only an authoritative outcome from a still-active recovery round. + pub(super) fn completed(&mut self, pool: &Arc, job: HealthJob, outcome: JobOutcome) { + let Some(state) = self.groups.get_mut(&job.key) else { + return; + }; + if state.authority != Some((job.observation.authority, job.observation.generation)) + || pool.floor_authority() != job.observation.authority + || pool.current_generation() != job.observation.generation + { + return; + } + if matches!(job.kind, JobKind::Recover | JobKind::Transfer) + && (!state.round_active || state.round != job.round) + { + return; + } + let JobOutcome::Completed(deadline) = outcome else { + state.deferred = true; + return; + }; + match job.kind { + JobKind::Refresh => { + state.refresh_due = Some(deadline.unwrap_or_else(|| { + Instant::now() + Duration::from_secs(IDLE_REFRESH_RETRY_SECS) + })) + } + JobKind::Shadow => { + state.shadow_due = Some(deadline.unwrap_or_else(|| { + Instant::now() + Duration::from_secs(SHADOW_ROTATE_RETRY_SECS) + })) + } + JobKind::Recover | JobKind::Transfer => { + state.round_left = state.round_left.saturating_sub(1); + // Round finalization uses the next fresh observation, never the job's old count. + state.due = None; + } + } + } +} + +/// Finds refresh work without waiting for another group's network operation. +pub(super) fn idle_refresh_ready( + observation: &HealthObservation, + entry: &DcFloorPlanEntry, +) -> bool { + let now = MePool::now_epoch_secs(); + entry.endpoints.iter().any(|endpoint| { + observation + .ids + .get(&(entry.dc, *endpoint)) + .is_some_and(|ids| { + ids.iter().any(|id| { + observation.bound.get(id).copied().unwrap_or(0) == 0 + && observation.idle.get(id).is_some_and(|since| { + now.saturating_sub(*since) + >= IDLE_REFRESH_TRIGGER_BASE_SECS + + id % (IDLE_REFRESH_TRIGGER_JITTER_SECS + 1) + }) + }) + }) + }) +} + +/// Executes at most one bounded writer operation, never a whole missing-writer batch. +pub(super) async fn run_job( + pool: Arc, + rng: Arc, + job: HealthJob, +) -> JobOutcome { + if pool.floor_authority() != job.observation.authority + || pool.current_generation() != job.observation.generation + { + return JobOutcome::Deferred; + } + let key = job.key; + let entry = &job.entry; + let observation = &job.observation; + let mut deadlines = HashMap::new(); + match job.kind { + JobKind::Refresh => { + maybe_refresh_idle_writer_for_dc( + &pool, + &rng, + key, + key.0, + key.1, + &entry.endpoints, + entry.alive, + entry.target_required, + &observation.ids, + &observation.idle, + &observation.bound, + &mut deadlines, + ) + .await + } + JobKind::Shadow => { + maybe_rotate_single_endpoint_shadow( + &pool, + &rng, + key, + key.0, + key.1, + &entry.endpoints, + entry.alive, + entry.target_required, + &observation.ids, + &observation.bound, + &mut deadlines, + ) + .await + } + JobKind::Transfer => { + if coverage::transfer_coverage( + &pool, + &rng, + key.0, + entry.endpoints[0], + entry.target_required, + ) + .await + .is_none() + { + return JobOutcome::Deferred; + } + } + JobKind::Recover => { + let outage = entry.endpoints.len() == 1 + && entry.alive == 0 + && pool.single_endpoint_outage_mode_enabled(); + pool.stats.increment_me_reconnect_attempt(); + if outage { + pool.stats + .increment_me_single_endpoint_outage_reconnect_attempt_total(); + } + let base = pool.required_writers_for_dc_with_floor_mode(entry.endpoints.len(), false); + let intent = if entry.alive < base { + WriterOpenIntent::Coverage + } else { + WriterOpenIntent::Normal + }; + let connected = tokio::time::timeout(pool.reconnect_runtime.me_one_timeout, async { + if outage && pool.single_endpoint_outage_disable_quarantine() { + pool.stats + .increment_me_single_endpoint_quarantine_bypass_total(); + pool.connect_one_with_generation_contour_for_dc_with_intent( + entry.endpoints[0], + rng.as_ref(), + observation.generation, + WriterContour::Active, + key.0, + intent, + ) + .await + .is_ok() + } else { + pool.connect_endpoints_round_robin_with_generation_contour( + key.0, + &entry.endpoints, + rng.as_ref(), + observation.generation, + WriterContour::Active, + intent, + ) + .await + } + }) + .await + .unwrap_or(false); + if connected { + pool.stats.increment_me_reconnect_success(); + if outage { + pool.stats + .increment_me_single_endpoint_outage_reconnect_success_total(); + } + } + } + } + JobOutcome::Completed(deadlines.remove(&key)) +} diff --git a/src/transport/middle_proxy/health/scheduler/tests.rs b/src/transport/middle_proxy/health/scheduler/tests.rs new file mode 100644 index 0000000..4993a71 --- /dev/null +++ b/src/transport/middle_proxy/health/scheduler/tests.rs @@ -0,0 +1,67 @@ +use super::*; +use crate::transport::middle_proxy::pool_writer_security_tests::make_pool; + +#[tokio::test] +async fn deferred_and_stale_completions_preserve_round_budget_and_deadlines() { + let pool = make_pool().await; + let addr: SocketAddr = "127.0.0.1:443".parse().unwrap(); + pool.update_proxy_maps(HashMap::from([(2, vec![(addr.ip(), addr.port())])]), None) + .await; + let observation = Arc::new(check_family(IpFamily::V4, &pool).await); + let mut scheduler = HealthScheduler::default(); + scheduler.observe(&HashMap::from([(IpFamily::V4, observation.clone())])); + let key = (2, IpFamily::V4); + let entry = observation.floor.by_dc[&2].clone(); + let job = HealthJob { + key, + kind: JobKind::Transfer, + observation, + entry, + round: 4, + }; + let due = Instant::now() - Duration::from_secs(1); + let state = scheduler.groups.get_mut(&key).unwrap(); + state.round = 4; + state.round_left = 3; + state.round_active = true; + state.due = Some(due); + scheduler.completed(&pool, job.clone(), JobOutcome::Deferred); + let state = &scheduler.groups[&key]; + assert!(state.deferred); + assert_eq!(state.round_left, 3); + assert_eq!(state.due, Some(due)); + let state = scheduler.groups.get_mut(&key).unwrap(); + state.deferred = false; + state.round_active = false; + scheduler.completed(&pool, job.clone(), JobOutcome::Completed(None)); + scheduler.completed(&pool, job, JobOutcome::Deferred); + let state = &scheduler.groups[&key]; + assert!(!state.deferred); + assert_eq!(state.round_left, 3); + assert_eq!(state.due, Some(due)); +} + +#[tokio::test] +async fn queue_survives_repeated_observations_without_starving_later_groups() { + let pool = make_pool().await; + pool.update_proxy_maps( + (1..=12) + .map(|dc| (dc, vec![("127.0.0.1".parse().unwrap(), 443)])) + .collect(), + None, + ) + .await; + let observation = Arc::new(check_family(IpFamily::V4, &pool).await); + let observations = HashMap::from([(IpFamily::V4, observation)]); + let mut scheduler = HealthScheduler::default(); + scheduler.observe(&observations); + let expected = scheduler.queue.iter().copied().collect::>(); + let mut visited = HashSet::new(); + for _ in 0..expected.len() { + let key = scheduler.queue.pop_front().unwrap(); + scheduler.queue.push_back(key); + visited.insert(key); + scheduler.observe(&observations); + } + assert_eq!(visited, expected); +} diff --git a/src/transport/middle_proxy/pool.rs b/src/transport/middle_proxy/pool.rs index 95ba2c2..9543779 100644 --- a/src/transport/middle_proxy/pool.rs +++ b/src/transport/middle_proxy/pool.rs @@ -342,6 +342,8 @@ pub(super) struct ReinitCoordinatorState { /// Latest endpoint authority revision accepted by the coordinator. pub(super) endpoint_revision: u64, pub(super) pending: Option, + /// Revision of the coherent floor and capacity policy. + pub(super) floor_policy_revision: u64, pub(super) attempts: HashMap, } @@ -409,7 +411,7 @@ pub(super) struct BindingPolicyCore { pub(super) struct NatRuntimeCore { pub(super) discovery: super::pool_nat::NatDiscovery, pub(super) nat_ip_cfg: Option, - pub(super) nat_ip_detected: Arc>>, + pub(super) nat_ip_detected: Arc>>, pub(super) nat_probe: bool, pub(super) nat_stun: Option, pub(super) nat_stun_servers: Vec, @@ -503,10 +505,10 @@ pub struct MePool { pub(super) proxy_secret: Arc>, 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, - /// Replacement connections opened but not yet committed to writer visibility. - pub(super) writer_replacement_open_reserved: AtomicUsize, + pub(super) writer_connect_active_reserved: Arc, + pub(super) writer_connect_warm_reserved: Arc, + /// Physical replacement overlap: preparation plus victims awaiting transport destruction. + pub(super) writer_replacement_open_reserved: Arc, pub(super) rtt_stats: Arc>>, /// Coalesced refill state keyed by exact generation and contour ownership. pub(super) refill_states: Arc>>, @@ -545,6 +547,8 @@ mod transport_policy; mod selection_policy; // Bounded writer-open admission and coverage accounting. mod writer_admission; +// Coherent policy revisions fence derived health plans and replacement decisions. +mod floor_authority; pub(super) use writer_admission::{WriterOpenIntent, WriterOpenReservation, WriterRole}; // Endpoint-to-DC routing and health timing policy. mod routing; diff --git a/src/transport/middle_proxy/pool/construction.rs b/src/transport/middle_proxy/pool/construction.rs index 83bd429..d93805d 100644 --- a/src/transport/middle_proxy/pool/construction.rs +++ b/src/transport/middle_proxy/pool/construction.rs @@ -165,6 +165,7 @@ impl MePool { active_generation: 1, desired_map_hash: 0, endpoint_revision: 1, + floor_policy_revision: 1, pending: None, attempts: HashMap::new(), }), @@ -395,9 +396,9 @@ impl MePool { pool_size: 2, 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), - writer_replacement_open_reserved: AtomicUsize::new(0), + writer_connect_active_reserved: Arc::new(AtomicUsize::new(0)), + writer_connect_warm_reserved: Arc::new(AtomicUsize::new(0)), + writer_replacement_open_reserved: Arc::new(AtomicUsize::new(0)), rtt_stats: Arc::new(Mutex::new(HashMap::new())), refill_states: Arc::new(ParkingMutex::new(HashMap::new())), refill_running: AtomicUsize::new(0), diff --git a/src/transport/middle_proxy/pool/floor_authority.rs b/src/transport/middle_proxy/pool/floor_authority.rs new file mode 100644 index 0000000..8bca96a --- /dev/null +++ b/src/transport/middle_proxy/pool/floor_authority.rs @@ -0,0 +1,59 @@ +use super::*; +#[cfg(test)] +mod tests; + +impl MePool { + /// Captures the two independent authorities used by a health plan. + pub(in crate::transport::middle_proxy) fn floor_authority(&self) -> (u64, u64) { + let coordinator = self.reinit.coordinator.lock(); + ( + coordinator.endpoint_revision, + coordinator.floor_policy_revision, + ) + } + + /// Captures floor-affecting values under the coordinator without recursively acquiring it. + pub(super) fn floor_policy_values(&self) -> [u64; 14] { + let floor = &self.floor_runtime; + [ + self.single_endpoint_runtime + .me_single_endpoint_shadow_writers + .load(Ordering::Relaxed) as u64, + floor.me_floor_mode.load(Ordering::Relaxed) as u64, + floor.me_adaptive_floor_idle_secs.load(Ordering::Relaxed), + floor + .me_adaptive_floor_min_writers_single_endpoint + .load(Ordering::Relaxed) as u64, + floor + .me_adaptive_floor_min_writers_multi_endpoint + .load(Ordering::Relaxed) as u64, + floor + .me_adaptive_floor_recover_grace_secs + .load(Ordering::Relaxed), + floor + .me_adaptive_floor_writers_per_core_total + .load(Ordering::Relaxed) as u64, + floor + .me_adaptive_floor_cpu_cores_override + .load(Ordering::Relaxed) as u64, + floor + .me_adaptive_floor_max_extra_writers_single_per_core + .load(Ordering::Relaxed) as u64, + floor + .me_adaptive_floor_max_extra_writers_multi_per_core + .load(Ordering::Relaxed) as u64, + floor + .me_adaptive_floor_max_active_writers_per_core + .load(Ordering::Relaxed) as u64, + floor + .me_adaptive_floor_max_warm_writers_per_core + .load(Ordering::Relaxed) as u64, + floor + .me_adaptive_floor_max_active_writers_global + .load(Ordering::Relaxed) as u64, + floor + .me_adaptive_floor_max_warm_writers_global + .load(Ordering::Relaxed) as u64, + ] + } +} diff --git a/src/transport/middle_proxy/pool/floor_authority/tests.rs b/src/transport/middle_proxy/pool/floor_authority/tests.rs new file mode 100644 index 0000000..49ccee0 --- /dev/null +++ b/src/transport/middle_proxy/pool/floor_authority/tests.rs @@ -0,0 +1,96 @@ +use super::*; +use crate::config::GeneralConfig; +use crate::transport::middle_proxy::pool_writer_security_tests::make_pool; + +fn reload(pool: &MePool, cfg: &GeneralConfig) { + pool.update_runtime_reinit_policy( + cfg.hardswap, + cfg.me_pool_drain_ttl_secs, + cfg.me_instadrain, + cfg.me_pool_drain_threshold, + cfg.me_pool_drain_soft_evict_enabled, + cfg.me_pool_drain_soft_evict_grace_secs, + cfg.me_pool_drain_soft_evict_per_writer, + cfg.me_pool_drain_soft_evict_budget_per_core, + cfg.me_pool_drain_soft_evict_cooldown_ms, + cfg.effective_me_pool_force_close_secs(), + cfg.me_pool_min_fresh_ratio, + cfg.me_hardswap_warmup_delay_min_ms, + cfg.me_hardswap_warmup_delay_max_ms, + cfg.me_hardswap_warmup_extra_passes, + cfg.me_hardswap_warmup_pass_backoff_base_ms, + cfg.me_bind_stale_mode, + cfg.me_bind_stale_ttl_secs, + cfg.me_secret_atomic_snapshot, + cfg.me_deterministic_writer_sort, + cfg.me_writer_pick_mode, + cfg.me_writer_pick_sample_size, + cfg.me_single_endpoint_shadow_writers, + cfg.me_single_endpoint_outage_mode_enabled, + cfg.me_single_endpoint_outage_disable_quarantine, + cfg.me_single_endpoint_outage_backoff_min_ms, + cfg.me_single_endpoint_outage_backoff_max_ms, + cfg.me_single_endpoint_shadow_rotate_every_secs, + cfg.me_floor_mode, + cfg.me_adaptive_floor_idle_secs, + cfg.me_adaptive_floor_min_writers_single_endpoint, + cfg.me_adaptive_floor_min_writers_multi_endpoint, + cfg.me_adaptive_floor_recover_grace_secs, + cfg.me_adaptive_floor_writers_per_core_total, + cfg.me_adaptive_floor_cpu_cores_override, + cfg.me_adaptive_floor_max_extra_writers_single_per_core, + cfg.me_adaptive_floor_max_extra_writers_multi_per_core, + cfg.me_adaptive_floor_max_active_writers_per_core, + cfg.me_adaptive_floor_max_warm_writers_per_core, + cfg.me_adaptive_floor_max_active_writers_global, + cfg.me_adaptive_floor_max_warm_writers_global, + cfg.me_health_interval_ms_unhealthy, + cfg.me_health_interval_ms_healthy, + cfg.me_warn_rate_limit_ms, + ); +} + +#[tokio::test] +async fn policy_reload_invalidates_old_caps_and_rejects_stale_plan_publication() { + let pool = make_pool().await; + let mut cfg = GeneralConfig::default(); + reload(&pool, &cfg); + let authority = pool.floor_authority(); + pool.set_adaptive_floor_runtime_caps(authority, 10, 12, 10, 12, 6, 6, 0); + assert_eq!( + pool.floor_runtime + .me_adaptive_floor_active_cap_effective + .load(Ordering::Acquire), + 12 + ); + cfg.me_single_endpoint_shadow_writers += 1; + reload(&pool, &cfg); + assert_ne!(pool.floor_authority(), authority); + assert_eq!( + pool.floor_runtime + .me_adaptive_floor_active_cap_effective + .load(Ordering::Acquire), + 0 + ); + pool.set_adaptive_floor_runtime_caps(authority, 10, 128, 10, 128, 100, 6, 0); + assert_eq!( + pool.floor_runtime + .me_adaptive_floor_active_cap_effective + .load(Ordering::Acquire), + 0 + ); + let new_authority = pool.floor_authority(); + reload(&pool, &cfg); + assert_eq!( + pool.floor_authority(), + new_authority, + "unchanged reload must not supersede work" + ); + pool.set_adaptive_floor_runtime_caps(new_authority, 10, 11, 10, 11, 6, 6, 0); + assert_eq!( + pool.floor_runtime + .me_adaptive_floor_active_cap_effective + .load(Ordering::Acquire), + 11 + ); +} diff --git a/src/transport/middle_proxy/pool/runtime_policy.rs b/src/transport/middle_proxy/pool/runtime_policy.rs index 3c11bf8..3b92735 100644 --- a/src/transport/middle_proxy/pool/runtime_policy.rs +++ b/src/transport/middle_proxy/pool/runtime_policy.rs @@ -305,6 +305,8 @@ impl MePool { if previous_writer_pick_mode != writer_pick_mode { self.stats.increment_me_writer_pick_mode_switch_total(); } + let mut coordinator = self.reinit.coordinator.lock(); + let previous_floor_policy = self.floor_policy_values(); self.single_endpoint_runtime .me_single_endpoint_shadow_writers .store(single_endpoint_shadow_writers, Ordering::Relaxed); @@ -381,6 +383,17 @@ impl MePool { self.floor_runtime .me_adaptive_floor_max_warm_writers_global .store(adaptive_floor_max_warm_writers_global, Ordering::Relaxed); + if previous_floor_policy != self.floor_policy_values() { + coordinator.floor_policy_revision = coordinator.floor_policy_revision.wrapping_add(1); + // Cached caps belong to the old policy until a new health plan is accepted. + self.floor_runtime + .me_adaptive_floor_active_cap_effective + .store(0, Ordering::Release); + self.floor_runtime + .me_adaptive_floor_warm_cap_effective + .store(0, Ordering::Release); + } + drop(coordinator); self.health_runtime .me_health_interval_ms_unhealthy .store(me_health_interval_ms_unhealthy.max(1), Ordering::Relaxed); diff --git a/src/transport/middle_proxy/pool/selection_policy.rs b/src/transport/middle_proxy/pool/selection_policy.rs index b0d0b93..2f92ce0 100644 --- a/src/transport/middle_proxy/pool/selection_policy.rs +++ b/src/transport/middle_proxy/pool/selection_policy.rs @@ -230,8 +230,10 @@ impl MePool { configured } + /// Publishes derived caps only while both endpoint and floor-policy authority remain current. pub(in crate::transport::middle_proxy) fn set_adaptive_floor_runtime_caps( &self, + authority: (u64, u64), active_cap_configured: usize, active_cap_effective: usize, warm_cap_configured: usize, @@ -240,6 +242,15 @@ impl MePool { active_writers_current: usize, warm_writers_current: usize, ) { + let coordinator = self.reinit.coordinator.lock(); + if authority + != ( + coordinator.endpoint_revision, + coordinator.floor_policy_revision, + ) + { + return; + } self.floor_runtime .me_adaptive_floor_global_cap_raw .store(active_cap_configured as u64, Ordering::Relaxed); diff --git a/src/transport/middle_proxy/pool/writer_admission.rs b/src/transport/middle_proxy/pool/writer_admission.rs index a8afd76..b64e3ed 100644 --- a/src/transport/middle_proxy/pool/writer_admission.rs +++ b/src/transport/middle_proxy/pool/writer_admission.rs @@ -48,13 +48,13 @@ pub(in crate::transport::middle_proxy) enum WriterOpenIntent { } /// RAII ownership of one bounded in-flight writer open. -pub(in crate::transport::middle_proxy) struct WriterOpenReservation<'a> { - counter: Option<&'a AtomicUsize>, +pub(in crate::transport::middle_proxy) struct WriterOpenReservation { + counter: Option>, } -impl Drop for WriterOpenReservation<'_> { +impl Drop for WriterOpenReservation { fn drop(&mut self) { - if let Some(counter) = self.counter { + if let Some(counter) = &self.counter { counter.fetch_sub(1, Ordering::AcqRel); } } @@ -172,7 +172,7 @@ impl MePool { intent: WriterOpenIntent, writer_dc: i32, target_addr: SocketAddr, - ) -> Option> { + ) -> Option { let counter = match contour { WriterContour::Active => &self.writer_connect_active_reserved, WriterContour::Warm => &self.writer_connect_warm_reserved, @@ -182,6 +182,8 @@ impl MePool { }; if intent == WriterOpenIntent::Replacement { + // Policy reload and replacement admission share the same linearization point. + let _coordinator = self.reinit.coordinator.lock(); let configured_cap = match contour { WriterContour::Active => self.adaptive_floor_active_cap_configured_total(), WriterContour::Warm => self.adaptive_floor_warm_cap_configured_total(), @@ -220,7 +222,7 @@ impl MePool { .is_ok() { return Some(WriterOpenReservation { - counter: Some(&self.writer_replacement_open_reserved), + counter: Some(self.writer_replacement_open_reserved.clone()), }); } } @@ -266,7 +268,7 @@ impl MePool { .is_ok() { return Some(WriterOpenReservation { - counter: Some(counter), + counter: Some(counter.clone()), }); } } diff --git a/src/transport/middle_proxy/pool_lifecycle.rs b/src/transport/middle_proxy/pool_lifecycle.rs index dcfe93c..7734bff 100644 --- a/src/transport/middle_proxy/pool_lifecycle.rs +++ b/src/transport/middle_proxy/pool_lifecycle.rs @@ -111,6 +111,31 @@ impl MePoolLifecycle { self.admission.try_register() } + /// Registers a cancellation-aware producer for a caller-owned JoinSet without a second spawn. + pub(super) fn track_producer( + &self, + future: F, + ) -> Result> + Send + 'static + use, F> + where + F: Future + Send + 'static, + F::Output: Send + 'static, + { + let Some(registration) = self.try_register() else { + return Err(future); + }; + let cancel = self.producer_cancel.clone(); + let tracked = self.producer_tasks.track_future(async move { + tokio::select! { + biased; + _ = cancel.cancelled() => None, + output = future => Some(output), + } + }); + // The tracker owns destruction before shutdown can observe no registrations. + drop(registration); + Ok(tracked) + } + /// Registers and spawns one cancellation-aware ME producer. pub(super) fn spawn_producer(&self, future: F) -> Result<(), F> where diff --git a/src/transport/middle_proxy/pool_nat.rs b/src/transport/middle_proxy/pool_nat.rs index 23ff765..d50396d 100644 --- a/src/transport/middle_proxy/pool_nat.rs +++ b/src/transport/middle_proxy/pool_nat.rs @@ -26,6 +26,9 @@ const HTTP_FAILURE_TTL: Duration = Duration::from_secs(60); // Discovery ownership and source-bound cache admission for this pool generation. mod state; pub(super) use state::NatDiscovery; +// Successful HTTP observations expire independently of failure backoff. +mod http; +pub(super) use http::HttpNatObservation; // Local network regressions exercise cancellation before discovery publication. #[cfg(test)] @@ -167,14 +170,16 @@ impl MePool { /// Applies the configured or HTTP-detected NAT IP to non-public addresses. pub(super) fn translate_ip_for_nat(&self, ip: IpAddr) -> IpAddr { - let nat_ip = self.nat_runtime.nat_ip_cfg.or_else(|| { + Self::translate_ip_with_nat( + ip, self.nat_runtime - .nat_ip_detected - .try_read() - .ok() - .and_then(|g| *g) - }); + .nat_ip_cfg + .or_else(|| self.cached_http_nat_ip()), + ) + } + /// Applies a captured NAT observation without rereading mutable discovery state. + pub(super) fn translate_ip_with_nat(ip: IpAddr, nat_ip: Option) -> IpAddr { let Some(nat_ip) = nat_ip else { return ip; }; @@ -197,6 +202,16 @@ impl MePool { &self, addr: std::net::SocketAddr, reflected: Option, + ) -> std::net::SocketAddr { + self.translate_our_addr_with_nat(addr, reflected, self.cached_http_nat_ip()) + } + + /// Retains the actual TCP port while using one handshake-local NAT snapshot. + pub(super) fn translate_our_addr_with_nat( + &self, + addr: std::net::SocketAddr, + reflected: Option, + detected: Option, ) -> std::net::SocketAddr { let ip = if let Some(nat_ip) = self.nat_runtime.nat_ip_cfg { match (addr.ip(), nat_ip) { @@ -209,63 +224,16 @@ impl MePool { if is_bogon(addr.ip()) || addr.ip().is_loopback() || addr.ip().is_unspecified() { r.ip() } else { - self.translate_ip_for_nat(addr.ip()) + Self::translate_ip_with_nat(addr.ip(), detected) } } else { - self.translate_ip_for_nat(addr.ip()) + Self::translate_ip_with_nat(addr.ip(), detected) }; // Keep the kernel-assigned TCP source port; STUN port can differ. std::net::SocketAddr::new(ip, addr.port()) } - /// Waits for pool-owned HTTP discovery without owning its cancellation. - pub(super) async fn maybe_detect_nat_ip(self: &Arc, local_ip: IpAddr) -> Option { - if self.nat_runtime.nat_ip_cfg.is_some() { - return self.nat_runtime.nat_ip_cfg; - } - - if !self.nat_runtime.nat_probe { - return None; - } - - if !(is_bogon(local_ip) || local_ip.is_loopback() || local_ip.is_unspecified()) { - return None; - } - - if let Some(ip) = *self.nat_runtime.nat_ip_detected.read().await { - return Some(ip); - } - - let mut flight = self.nat_runtime.discovery.http.clone().lock_owned().await; - if let Some(ip) = *self.nat_runtime.nat_ip_detected.read().await { - return Some(ip); - } - if flight.is_some_and(|until| Instant::now() < until) { - return None; - } - - let pool = self.clone(); - let (send, receive) = oneshot::channel(); - // Transfer singleflight ownership before the producer can be cancelled or polled. - self.lifecycle - .spawn_producer(async move { - let detected = detect_public_ipv4_http(&pool.nat_runtime.http_ip_detect_urls) - .await - .map(IpAddr::V4); - if let Some(ip) = detected { - *pool.nat_runtime.nat_ip_detected.write().await = Some(ip); - *flight = None; - info!(public_ip = %ip, "Auto-detected public IP for NAT translation"); - } else { - *flight = Some(Instant::now() + HTTP_FAILURE_TTL); - } - let _ = send.send(detected); - }) - .ok()?; - receive.await.ok().flatten() - } - /// Waits for a pool-owned STUN refresh in the exact source-address context. pub(super) async fn maybe_reflect_public_addr( self: &Arc, diff --git a/src/transport/middle_proxy/pool_nat/http.rs b/src/transport/middle_proxy/pool_nat/http.rs new file mode 100644 index 0000000..540ea73 --- /dev/null +++ b/src/transport/middle_proxy/pool_nat/http.rs @@ -0,0 +1,108 @@ +use super::*; + +const HTTP_FRESH_TTL: Duration = Duration::from_secs(600); +const HTTP_MAX_AGE: Duration = Duration::from_secs(1200); + +/// A successful HTTP observation has an immutable monotonic timestamp. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(in crate::transport::middle_proxy) struct HttpNatObservation { + /// Successful detector result, never updated by a failed refresh. + pub(in crate::transport::middle_proxy) ip: IpAddr, + /// Original observation time used for both freshness and bounded stale fallback. + pub(super) observed_at: tokio::time::Instant, +} + +impl HttpNatObservation { + fn usable(self) -> Option { + (self.observed_at.elapsed() < HTTP_MAX_AGE).then_some(self.ip) + } + + fn fresh(self) -> bool { + self.observed_at.elapsed() < HTTP_FRESH_TTL + } +} + +impl MePool { + /// Reads a bounded cache snapshot without discovery or data-plane waiting. + pub(in crate::transport::middle_proxy) fn cached_http_nat_ip(&self) -> Option { + self.nat_runtime + .nat_ip_detected + .try_read() + .ok() + .and_then(|cache| cache.and_then(HttpNatObservation::usable)) + } + + /// Refreshes once per pool, retaining stale success only within its original grace period. + pub(in crate::transport::middle_proxy) async fn maybe_detect_nat_ip( + self: &Arc, + local_ip: IpAddr, + ) -> Option { + if self.nat_runtime.nat_ip_cfg.is_some() { + return self.nat_runtime.nat_ip_cfg; + } + if !self.nat_runtime.nat_probe + || !local_ip.is_ipv4() + || !(is_bogon(local_ip) || local_ip.is_loopback() || local_ip.is_unspecified()) + { + return None; + } + if let Some(cached) = *self.nat_runtime.nat_ip_detected.read().await + && cached.fresh() + { + return Some(cached.ip); + } + let mut flight = self.nat_runtime.discovery.http.clone().lock_owned().await; + let cached = *self.nat_runtime.nat_ip_detected.read().await; + if let Some(cached) = cached + && cached.fresh() + { + return Some(cached.ip); + } + if flight.is_some_and(|until| Instant::now() < until) { + return cached.and_then(HttpNatObservation::usable); + } + let pool = self.clone(); + let (send, receive) = oneshot::channel(); + // Discovery owns the flight; cancellation of a handshake cannot cancel publication. + self.lifecycle + .spawn_producer(async move { + let detected = detect_public_ipv4_http(&pool.nat_runtime.http_ip_detect_urls) + .await + .map(IpAddr::V4); + if let Some(ip) = detected { + *pool.nat_runtime.nat_ip_detected.write().await = Some(HttpNatObservation { + ip, + observed_at: tokio::time::Instant::now(), + }); + *flight = None; + info!(public_ip = %ip, "Auto-detected public IP for NAT translation"); + } else { + *flight = Some(Instant::now() + HTTP_FAILURE_TTL); + } + let _ = send.send(()); + }) + .ok()?; + let _ = receive.await; + self.cached_http_nat_ip() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test(start_paused = true)] + async fn observation_age_is_not_renewed_by_reads_or_failure() { + let cached = HttpNatObservation { + ip: "91.108.56.1".parse().unwrap(), + observed_at: tokio::time::Instant::now(), + }; + tokio::time::advance(HTTP_FRESH_TTL).await; + assert!(!cached.fresh()); + for _ in 0..100 { + assert_eq!(cached.usable(), Some(cached.ip)); + } + tokio::time::advance(HTTP_MAX_AGE - HTTP_FRESH_TTL).await; + assert_eq!(cached.usable(), None); + } +} diff --git a/src/transport/middle_proxy/pool_nat/tests.rs b/src/transport/middle_proxy/pool_nat/tests.rs index 0a1c68c..243bc14 100644 --- a/src/transport/middle_proxy/pool_nat/tests.rs +++ b/src/transport/middle_proxy/pool_nat/tests.rs @@ -18,6 +18,7 @@ pub(super) const PUBLIC_IP: Ipv4Addr = Ipv4Addr::new(91, 108, 56, 1); pub(super) const TEST_DEADLINE: Duration = Duration::from_secs(2); // Concurrency and end-to-end regressions share gated local discovery servers. +mod freshness; pub(super) mod lifecycle; mod recovery; @@ -243,7 +244,11 @@ async fn http_discovery_survives_reconnect_timeout() { .await .expect("pool-owned HTTP discovery must publish after caller cancellation"); assert_eq!( - *pool.nat_runtime.nat_ip_detected.read().await, + pool.nat_runtime + .nat_ip_detected + .read() + .await + .map(|observation| observation.ip), Some(PUBLIC_IP.into()) ); assert_eq!(http.requests.load(Ordering::Relaxed), 1); diff --git a/src/transport/middle_proxy/pool_nat/tests/freshness.rs b/src/transport/middle_proxy/pool_nat/tests/freshness.rs new file mode 100644 index 0000000..d2796f7 --- /dev/null +++ b/src/transport/middle_proxy/pool_nat/tests/freshness.rs @@ -0,0 +1,75 @@ +use super::*; + +#[tokio::test] +async fn failed_refresh_retains_original_age_but_never_expired_material() { + let http = http_server("invalid address").await; + http.release.add_permits(8); + let pool = nat_pool(vec![], vec![format!("http://{}/", http.addr)]).await; + let observed_at = tokio::time::Instant::now() - Duration::from_secs(601); + *pool.nat_runtime.nat_ip_detected.write().await = Some(HttpNatObservation { + ip: PUBLIC_IP.into(), + observed_at, + }); + assert_eq!( + pool.maybe_detect_nat_ip(Ipv4Addr::LOCALHOST.into()).await, + Some(PUBLIC_IP.into()) + ); + assert_eq!( + pool.nat_runtime + .nat_ip_detected + .read() + .await + .unwrap() + .observed_at, + observed_at + ); + assert_eq!(http.requests.load(Ordering::Acquire), 1); + assert_eq!( + pool.maybe_detect_nat_ip(Ipv4Addr::LOCALHOST.into()).await, + Some(PUBLIC_IP.into()) + ); + assert_eq!( + http.requests.load(Ordering::Acquire), + 1, + "failure backoff remains singleflight" + ); + pool.nat_runtime + .nat_ip_detected + .write() + .await + .as_mut() + .unwrap() + .observed_at = tokio::time::Instant::now() - Duration::from_secs(1201); + assert_eq!( + pool.maybe_detect_nat_ip(Ipv4Addr::LOCALHOST.into()).await, + None + ); + assert_eq!( + pool.translate_ip_for_nat(Ipv4Addr::LOCALHOST.into()), + IpAddr::V4(Ipv4Addr::LOCALHOST) + ); + assert!(pool.shutdown_until(TEST_DEADLINE).await); +} + +#[tokio::test] +async fn expired_http_observation_defers_handshake_before_sending_kdf_material() { + let pool = nat_pool(vec![], vec![]).await; + *pool.nat_runtime.nat_ip_detected.write().await = Some(HttpNatObservation { + ip: PUBLIC_IP.into(), + observed_at: tokio::time::Instant::now() - Duration::from_secs(1201), + }); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let stream = tokio::net::TcpStream::connect(addr).await.unwrap(); + let (mut peer, _) = listener.accept().await.unwrap(); + let error = pool + .handshake_only(stream, addr, None, pool.rng.as_ref()) + .await + .err() + .unwrap(); + assert!( + matches!(error, ProxyError::Io(ref error) if error.kind() == std::io::ErrorKind::WouldBlock) + ); + assert_eq!(peer.read(&mut [0u8; 1]).await.unwrap(), 0); + assert!(pool.shutdown_until(TEST_DEADLINE).await); +} diff --git a/src/transport/middle_proxy/pool_nat/tests/recovery.rs b/src/transport/middle_proxy/pool_nat/tests/recovery.rs index 911fb32..0c1a445 100644 --- a/src/transport/middle_proxy/pool_nat/tests/recovery.rs +++ b/src/transport/middle_proxy/pool_nat/tests/recovery.rs @@ -24,7 +24,11 @@ impl Drop for MiddlePeer { impl MiddlePeer { async fn new(server_ip: Ipv4Addr) -> Self { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + Self::with_bind("127.0.0.1:0", server_ip.into(), PUBLIC_IP.into()).await + } + + async fn with_bind(bind: &str, server_ip: IpAddr, client_ip: IpAddr) -> Self { + let listener = TcpListener::bind(bind).await.unwrap(); let addr = listener.local_addr().unwrap(); let handshakes = Arc::new(AtomicUsize::new(0)); let cancelled = Arc::new(AtomicUsize::new(0)); @@ -37,7 +41,7 @@ impl MiddlePeer { tokio::select! { accepted = listener.accept() => { let (stream, _) = accepted.unwrap(); - pending.spawn(accept_handshake(stream, server_ip)); + pending.spawn(accept_handshake(stream, server_ip, client_ip)); } completed = pending.join_next(), if !pending.is_empty() => { match completed.unwrap().expect("fake ME handshake task must not panic") { @@ -60,7 +64,11 @@ impl MiddlePeer { } } -async fn accept_handshake(mut stream: TcpStream, server_ip: Ipv4Addr) -> Option { +async fn accept_handshake( + mut stream: TcpStream, + server_ip: IpAddr, + client_ip: IpAddr, +) -> Option { let client_port = stream.peer_addr().unwrap().port(); let server_port = stream.local_addr().unwrap().port(); let (seq, payload) = match read_rpc_frame_plaintext(&mut stream).await { @@ -79,35 +87,43 @@ async fn accept_handshake(mut stream: TcpStream, server_ip: Ipv4Addr) -> Option< .write_all(&build_rpc_frame(-2, &nonce, RpcChecksumMode::Crc32)) .await .unwrap(); - let mut client_ip = PUBLIC_IP.octets(); - client_ip.reverse(); - let mut server_ip = server_ip.octets(); - server_ip.reverse(); + let split = |ip: IpAddr| match ip { + IpAddr::V4(ip) => { + let mut octets = ip.octets(); + octets.reverse(); + (Some(octets), None) + } + IpAddr::V6(ip) => (None, Some(ip.octets())), + }; + let (client_v4, client_v6) = split(client_ip); + let (server_v4, server_v6) = split(server_ip); + let client_ip = client_v4.unwrap_or([0; 4]); + let server_ip = server_v4.unwrap_or([0; 4]); let (read_key, read_iv) = derive_middleproxy_keys( &server_nonce, &client_nonce, ×tamp.to_le_bytes(), - Some(&server_ip), + server_v4.as_ref().map(|ip| ip.as_slice()), &client_port.to_le_bytes(), b"CLIENT", - Some(&client_ip), + client_v4.as_ref().map(|ip| ip.as_slice()), &server_port.to_le_bytes(), &[1u8; 32], - None, - None, + client_v6.as_ref(), + server_v6.as_ref(), ); let (write_key, write_iv) = derive_middleproxy_keys( &server_nonce, &client_nonce, ×tamp.to_le_bytes(), - Some(&server_ip), + server_v4.as_ref().map(|ip| ip.as_slice()), &client_port.to_le_bytes(), b"SERVER", - Some(&client_ip), + client_v4.as_ref().map(|ip| ip.as_slice()), &server_port.to_le_bytes(), &[1u8; 32], - None, - None, + client_v6.as_ref(), + server_v6.as_ref(), ); let mut encrypted = [0u8; 48]; stream.read_exact(&mut encrypted).await.unwrap(); @@ -151,12 +167,8 @@ async fn suppressed_pool_recovers_after_delayed_discovery(http_slow: bool) { let http = http_server_with_delay("91.108.56.1", delay).await; stun.release.add_permits(64); http.release.add_permits(64); - let peer = MiddlePeer::new(if http_slow { - PUBLIC_IP - } else { - Ipv4Addr::LOCALHOST - }) - .await; + // A valid STUN source bypasses HTTP; the peer remains the configured loopback endpoint. + let peer = MiddlePeer::new(Ipv4Addr::LOCALHOST).await; let urls = if http_slow { vec![format!("http://{}/", http.addr)] } else { @@ -178,6 +190,9 @@ async fn suppressed_pool_recovers_after_delayed_discovery(http_slow: bool) { pool.single_endpoint_runtime .me_single_endpoint_shadow_writers .store(0, Ordering::Relaxed); + pool.single_endpoint_runtime + .me_single_endpoint_shadow_rotate_every_secs + .store(0, Ordering::Relaxed); pool.update_proxy_maps( HashMap::from([(2, vec![(peer.addr.ip(), peer.addr.port())])]), None, @@ -206,20 +221,18 @@ async fn suppressed_pool_recovers_after_delayed_discovery(http_slow: bool) { Arc::new(SecureRandom::new()), 0, )); - if http_slow { - http.received().await; - } else { + if !http_slow { stun.received().await; + timeout(Duration::from_secs(5), async { + while peer.cancelled.load(Ordering::Acquire) == 0 { + assert!(!peer.task.is_finished()); + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("the actual 1200ms reconnect must cancel a pre-KDF TCP connection"); + assert_eq!(peer.handshakes.load(Ordering::Acquire), 0); } - timeout(Duration::from_secs(5), async { - while peer.cancelled.load(Ordering::Acquire) == 0 { - assert!(!peer.task.is_finished()); - tokio::time::sleep(Duration::from_millis(10)).await; - } - }) - .await - .expect("the actual 1200ms reconnect must cancel at least one pre-KDF TCP connection"); - assert_eq!(peer.handshakes.load(Ordering::Acquire), 0); let expected = 2 * pool.required_writers_for_dc(1); timeout(Duration::from_secs(10), async { while !pool.admission_ready_full_floor().await @@ -244,10 +257,7 @@ async fn suppressed_pool_recovers_after_delayed_discovery(http_slow: bool) { ); assert_eq!(pool.writers.read().await.len(), expected); assert!(peer.handshakes.load(Ordering::Acquire) >= expected); - assert_eq!( - http.requests.load(Ordering::Relaxed), - usize::from(http_slow) - ); + assert_eq!(http.requests.load(Ordering::Relaxed), 0); assert_eq!( stun.requests.load(Ordering::Relaxed), usize::from(!http_slow) @@ -264,3 +274,58 @@ async fn suppressed_pool_recovers_full_floor_after_expired_stun_refresh() { async fn suppressed_pool_recovers_full_floor_after_slow_http_with_fresh_stun() { suppressed_pool_recovers_after_delayed_discovery(true).await; } + +#[tokio::test] +async fn stalled_ipv4_does_not_delay_real_ipv6_handshakes() { + let stalled = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let stalled_addr = stalled.local_addr().unwrap(); + let healthy = + MiddlePeer::with_bind("[::1]:0", "::1".parse().unwrap(), "::1".parse().unwrap()).await; + let mut pool = make_pool_with_decision(NetworkDecision { + ipv4_me: true, + ipv6_me: true, + ..NetworkDecision::default() + }) + .await; + Arc::get_mut(&mut Arc::get_mut(&mut pool).unwrap().reconnect_runtime) + .unwrap() + .me_one_timeout = Duration::from_secs(5); + pool.health_runtime + .me_health_interval_ms_unhealthy + .store(10, Ordering::Relaxed); + pool.single_endpoint_runtime + .me_single_endpoint_shadow_rotate_every_secs + .store(0, Ordering::Relaxed); + pool.update_proxy_maps( + HashMap::from([(2, vec![(stalled_addr.ip(), stalled_addr.port())])]), + Some(HashMap::from([( + 2, + vec![(healthy.addr.ip(), healthy.addr.port())], + )])), + ) + .await; + let monitor = tokio::spawn(me_health_monitor(pool.clone(), pool.rng.clone(), 0)); + let (blocked_stream, _) = timeout(TEST_DEADLINE, stalled.accept()) + .await + .unwrap() + .unwrap(); + timeout(TEST_DEADLINE, async { + while healthy.handshakes.load(Ordering::Acquire) < pool.required_writers_for_dc(1) { + assert!(!healthy.task.is_finished()); + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("IPv6 must restore coverage while IPv4 still owns a stalled handshake"); + monitor.abort(); + let _ = monitor.await; + assert!( + pool.shutdown_until(TEST_DEADLINE).await, + "tracked child transports must join on shutdown" + ); + drop(blocked_stream); + assert_eq!( + pool.writer_connect_active_reserved.load(Ordering::Acquire), + 0 + ); +} diff --git a/src/transport/middle_proxy/pool_reinit/coordination.rs b/src/transport/middle_proxy/pool_reinit/coordination.rs index d928fde..876445c 100644 --- a/src/transport/middle_proxy/pool_reinit/coordination.rs +++ b/src/transport/middle_proxy/pool_reinit/coordination.rs @@ -491,7 +491,11 @@ impl MePool { rand::rng().random_range(min_ms..=max_ms) } - pub(super) fn hardswap_warmup_backoff_ms(&self, pass_idx: usize) -> u64 { + /// Returns capped exponential retry delay with bounded positive jitter. + pub(in crate::transport::middle_proxy) fn hardswap_warmup_backoff_ms( + &self, + pass_idx: usize, + ) -> u64 { let base_ms = self .reinit .me_hardswap_warmup_pass_backoff_base_ms diff --git a/src/transport/middle_proxy/pool_reinit/reconcile.rs b/src/transport/middle_proxy/pool_reinit/reconcile.rs index be133da..f22c05a 100644 --- a/src/transport/middle_proxy/pool_reinit/reconcile.rs +++ b/src/transport/middle_proxy/pool_reinit/reconcile.rs @@ -169,11 +169,21 @@ impl MePool { } } + /// Reconciles one immutable endpoint snapshot without draining uncovered groups. pub async fn zero_downtime_reinit_after_map_change( self: &Arc, rng: &SecureRandom, ) -> bool { let endpoint_snapshot = self.endpoint_snapshot.load_full(); + self.reinit_endpoint_snapshot(rng, endpoint_snapshot).await + } + + /// Reconciles exactly the revision owned by the scheduler's retry objective. + pub(in crate::transport::middle_proxy) async fn reinit_endpoint_snapshot( + self: &Arc, + rng: &SecureRandom, + endpoint_snapshot: Arc, + ) -> bool { let desired_by_dc = self.desired_dc_endpoints_from_snapshot(&endpoint_snapshot); let now_epoch_secs = Self::now_epoch_secs(); let v4_suppressed = self.is_family_temporarily_suppressed(IpFamily::V4, now_epoch_secs); @@ -407,6 +417,7 @@ impl MePool { true } + /// Applies periodic rotation using the same coverage and generation fences as map changes. pub async fn zero_downtime_reinit_periodic(self: &Arc, rng: &SecureRandom) -> bool { self.zero_downtime_reinit_after_map_change(rng).await } diff --git a/src/transport/middle_proxy/pool_reinit/tests.rs b/src/transport/middle_proxy/pool_reinit/tests.rs index bb274be..1f1b492 100644 --- a/src/transport/middle_proxy/pool_reinit/tests.rs +++ b/src/transport/middle_proxy/pool_reinit/tests.rs @@ -149,6 +149,7 @@ fn stale_concurrent_attempt_cannot_regress_active_generation() { active_generation: 1, desired_map_hash: 22, endpoint_revision: 7, + floor_policy_revision: 1, pending: Some(ReinitPendingState { generation: 3, started_at_epoch_secs: 1, diff --git a/src/transport/middle_proxy/pool_writer.rs b/src/transport/middle_proxy/pool_writer.rs index 7bf97c0..b9298d8 100644 --- a/src/transport/middle_proxy/pool_writer.rs +++ b/src/transport/middle_proxy/pool_writer.rs @@ -32,6 +32,9 @@ use super::wire::build_proxy_req_payload; mod publication; mod replacement; mod runtime; +// Physical replacement overlap survives logical writer removal until transport destruction. +pub(super) mod overlap; +use overlap::WriterTransport; // Recovery regressions exercise the serialized production publication boundary. #[cfg(test)] @@ -42,9 +45,8 @@ struct PreparedWriter<'a> { tx: mpsc::Sender, byte_budget: Arc, task_registration: MeTaskRegistration<'a>, - writer_task: Pin + Send + 'static>>, + writer_task: WriterTransport, intent: WriterOpenIntent, - _open_reservation: WriterOpenReservation<'a>, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -59,6 +61,13 @@ pub(super) enum WriterReplacementPurpose { /// Target count above which the receiver no longer needs the slot. receiver_floor: usize, }, + /// Transfers one slot without taking a current donor below its protected base floor. + CoverageTransfer { + /// Endpoint and floor-policy authority captured before selecting the donor. + authority: (u64, u64), + /// Planned receiver target, revalidated against current policy at commit. + receiver_floor: usize, + }, /// Replaces a stale-generation writer with an active-generation writer. GenerationConvergence, /// Rotates a single-endpoint shadow while retaining coverage. @@ -72,6 +81,7 @@ impl WriterReplacementPurpose { match self { Self::IdleRefresh => "idle_refresh", Self::FloorRebalance { .. } => "floor_rebalance", + Self::CoverageTransfer { .. } => "coverage_transfer", Self::GenerationConvergence => "generation_convergence", Self::ShadowRotation => "shadow_rotation", Self::SecretRotation => "secret_rotation", diff --git a/src/transport/middle_proxy/pool_writer/overlap.rs b/src/transport/middle_proxy/pool_writer/overlap.rs new file mode 100644 index 0000000..8e494b9 --- /dev/null +++ b/src/transport/middle_proxy/pool_writer/overlap.rs @@ -0,0 +1,87 @@ +use super::*; +use std::task::{Context, Poll}; + +/// Writer-local physical lifetime, independent of registry visibility. +pub(in crate::transport::middle_proxy) struct ReplacementHandoff { + state: parking_lot::Mutex, +} + +enum HandoffState { + Open(Option), + Closed, +} + +impl ReplacementHandoff { + /// Creates a handoff that may accept one physical-overlap token. + pub(in crate::transport::middle_proxy) fn new() -> Arc { + Arc::new(Self { + state: parking_lot::Mutex::new(HandoffState::Open(None)), + }) + } + + fn close(&self) { + *self.state.lock() = HandoffState::Closed; + } + + /// Keeps attachment reversible until the final registry validation has succeeded. + pub(super) fn commit( + &self, + permit: &mut Option, + validate: impl FnOnce() -> bool, + ) -> bool { + let mut state = self.state.lock(); + match &mut *state { + HandoffState::Open(slot) if slot.is_none() && permit.is_some() && validate() => { + *slot = permit.take(); + true + } + _ => false, + } + } +} + +/// Destroys transport futures before releasing either opening or retiring capacity. +pub(super) struct WriterTransport { + future: Option + Send + 'static>>>, + /// Shared with the registry but closed only by transport destruction. + pub(super) lifetime: Arc, + /// Moves to the victim on replacement or is released after normal publication. + pub(super) opening: Option, +} + +impl WriterTransport { + /// Owns prepared transports even if their task is never polled or published. + pub(super) fn new( + future: Pin + Send + 'static>>, + opening: WriterOpenReservation, + ) -> Self { + Self { + future: Some(future), + lifetime: ReplacementHandoff::new(), + opening: Some(opening), + } + } +} + +impl Future for WriterTransport { + type Output = (); + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { + match self.future.as_mut() { + Some(future) => future.as_mut().poll(cx), + None => Poll::Ready(()), + } + } +} + +impl Drop for WriterTransport { + fn drop(&mut self) { + // Explicit order also applies to unpolled futures, panic unwinding and abort. + drop(self.future.take()); + self.lifetime.close(); + drop(self.opening.take()); + } +} + +#[cfg(test)] +mod tests; diff --git a/src/transport/middle_proxy/pool_writer/overlap/tests.rs b/src/transport/middle_proxy/pool_writer/overlap/tests.rs new file mode 100644 index 0000000..5a3ec7f --- /dev/null +++ b/src/transport/middle_proxy/pool_writer/overlap/tests.rs @@ -0,0 +1,238 @@ +use super::*; +use crate::transport::middle_proxy::pool_writer_security_tests::make_pool; +use std::sync::atomic::AtomicUsize; + +struct DestructionProbe { + counter: Arc, + destroyed: Arc, +} + +impl Drop for DestructionProbe { + fn drop(&mut self) { + assert!( + self.counter.load(Ordering::Acquire) > 0, + "transport must be destroyed before its overlap token" + ); + self.destroyed.store(true, Ordering::Release); + } +} + +async fn transport(pool: &Arc, destroyed: Arc) -> WriterTransport { + let permit = pool + .reserve_writer_open( + WriterContour::Active, + WriterOpenIntent::Replacement, + 2, + "127.0.0.1:443".parse().unwrap(), + ) + .await + .unwrap(); + let probe = DestructionProbe { + counter: pool.writer_replacement_open_reserved.clone(), + destroyed, + }; + WriterTransport::new( + Box::pin(async move { + let _probe = probe; + std::future::pending::<()>().await; + }), + permit, + ) +} + +#[tokio::test] +async fn unpolled_transport_and_aborted_task_release_after_destruction() { + let pool = make_pool().await; + for abort in [false, true] { + let destroyed = Arc::new(AtomicBool::new(false)); + let task = transport(&pool, destroyed.clone()).await; + if abort { + let task = tokio::spawn(task); + task.abort(); + let _ = task.await; + } else { + drop(task); + } + assert!(destroyed.load(Ordering::Acquire)); + assert_eq!( + pool.writer_replacement_open_reserved + .load(Ordering::Acquire), + 0 + ); + } +} + +#[tokio::test] +async fn retiring_transport_owns_capacity_after_logical_removal() { + let pool = make_pool().await; + let destroyed = Arc::new(AtomicBool::new(false)); + let mut victim = transport(&pool, destroyed.clone()).await; + drop(victim.opening.take()); + let lifetime = victim.lifetime.clone(); + let mut successor = transport(&pool, Arc::new(AtomicBool::new(false))).await; + successor.future = Some(Box::pin(std::future::pending())); + assert!(lifetime.commit(&mut successor.opening, || true)); + assert!(successor.opening.is_none()); + // Registry removal only drops its handle; the real transport retains the token. + drop(lifetime); + assert_eq!( + pool.writer_replacement_open_reserved + .load(Ordering::Acquire), + 1 + ); + drop(victim); + assert!(destroyed.load(Ordering::Acquire)); + assert_eq!( + pool.writer_replacement_open_reserved + .load(Ordering::Acquire), + 0 + ); + // Successor is no longer overlap after its victim has physically gone. + drop(successor); +} + +#[tokio::test] +async fn failed_or_closed_handoff_keeps_the_prepared_token() { + let pool = make_pool().await; + let mut prepared = transport(&pool, Arc::new(AtomicBool::new(false))).await; + let lifetime = ReplacementHandoff::new(); + assert!(!lifetime.commit(&mut prepared.opening, || false)); + assert!(prepared.opening.is_some()); + lifetime.close(); + assert!(!lifetime.commit(&mut prepared.opening, || true)); + assert!(prepared.opening.is_some()); + drop(prepared); + assert_eq!( + pool.writer_replacement_open_reserved + .load(Ordering::Acquire), + 0 + ); +} + +#[tokio::test] +async fn overlap_limit_survives_commit_and_cap_shrink_until_teardown() { + let pool = make_pool().await; + pool.floor_runtime + .me_adaptive_floor_cpu_cores_override + .store(1, Ordering::Relaxed); + pool.floor_runtime + .me_adaptive_floor_max_active_writers_per_core + .store(1, Ordering::Relaxed); + pool.floor_runtime + .me_adaptive_floor_max_active_writers_global + .store(1, Ordering::Relaxed); + pool.set_adaptive_floor_runtime_caps(pool.floor_authority(), 1, 2, 1, 1, 1, 0, 0); + let mut victim = transport(&pool, Arc::new(AtomicBool::new(false))).await; + drop(victim.opening.take()); + let mut successor = transport(&pool, Arc::new(AtomicBool::new(false))).await; + successor.future = Some(Box::pin(std::future::pending())); + assert!(victim.lifetime.commit(&mut successor.opening, || true)); + let second = transport(&pool, Arc::new(AtomicBool::new(false))).await; + pool.set_adaptive_floor_runtime_caps(pool.floor_authority(), 1, 1, 1, 1, 1, 1, 0); + let addr = "127.0.0.1:443".parse().unwrap(); + assert!( + pool.reserve_writer_open( + WriterContour::Active, + WriterOpenIntent::Replacement, + 2, + addr + ) + .await + .is_none() + ); + drop(second); + assert!( + pool.reserve_writer_open( + WriterContour::Active, + WriterOpenIntent::Replacement, + 2, + addr + ) + .await + .is_none() + ); + drop(victim); + assert!( + pool.reserve_writer_open( + WriterContour::Active, + WriterOpenIntent::Replacement, + 2, + addr + ) + .await + .is_some() + ); + drop(successor); + assert_eq!( + pool.writer_replacement_open_reserved + .load(Ordering::Acquire), + 0 + ); +} + +#[tokio::test] +async fn panicking_transport_releases_capacity_after_unwinding_its_resources() { + let pool = make_pool().await; + let permit = pool + .reserve_writer_open( + WriterContour::Active, + WriterOpenIntent::Replacement, + 2, + "127.0.0.1:443".parse().unwrap(), + ) + .await + .unwrap(); + let destroyed = Arc::new(AtomicBool::new(false)); + let probe = DestructionProbe { + counter: pool.writer_replacement_open_reserved.clone(), + destroyed: destroyed.clone(), + }; + let task = WriterTransport::new( + Box::pin(async move { + let _probe = probe; + panic!("injected transport panic"); + }), + permit, + ); + assert!(tokio::spawn(task).await.unwrap_err().is_panic()); + assert!(destroyed.load(Ordering::Acquire)); + assert_eq!( + pool.writer_replacement_open_reserved + .load(Ordering::Acquire), + 0 + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn racing_commit_and_close_release_each_token_exactly_once() { + let pool = make_pool().await; + for _ in 0..128 { + let permit = pool + .reserve_writer_open( + WriterContour::Active, + WriterOpenIntent::Replacement, + 2, + "127.0.0.1:443".parse().unwrap(), + ) + .await + .unwrap(); + let lifetime = ReplacementHandoff::new(); + let closing = lifetime.clone(); + let close = tokio::spawn(async move { + tokio::task::yield_now().await; + closing.close(); + }); + let commit = tokio::spawn(async move { + let mut permit = Some(permit); + lifetime.commit(&mut permit, || true); + drop(permit); + }); + close.await.unwrap(); + commit.await.unwrap(); + assert_eq!( + pool.writer_replacement_open_reserved + .load(Ordering::Acquire), + 0 + ); + } +} diff --git a/src/transport/middle_proxy/pool_writer/publication.rs b/src/transport/middle_proxy/pool_writer/publication.rs index 212ac45..d148ca9 100644 --- a/src/transport/middle_proxy/pool_writer/publication.rs +++ b/src/transport/middle_proxy/pool_writer/publication.rs @@ -11,9 +11,8 @@ impl MePool { tx, byte_budget, task_registration, - writer_task, + mut writer_task, intent, - _open_reservation, } = prepared; let writer_id = writer.id; let mut writers = self.writers.write().await; @@ -22,9 +21,15 @@ impl MePool { let contour = self.authorize_writer_publication(&writer, &coordinator)?; self.authorize_writer_publication_capacity(&writer, contour, intent, writers.as_slice())?; writer.contour.store(contour.as_u8(), Ordering::Release); - registry_registration.install(writer_id, tx, byte_budget); + registry_registration.install_with_lifetime( + writer_id, + tx, + byte_budget, + writer_task.lifetime.clone(), + ); writers.push(writer); self.conn_count.fetch_add(1, Ordering::Relaxed); + drop(writer_task.opening.take()); self.lifecycle .spawn_registered_writer(task_registration, writer_task); drop(coordinator); diff --git a/src/transport/middle_proxy/pool_writer/recovery_tests.rs b/src/transport/middle_proxy/pool_writer/recovery_tests.rs index ee7859f..448d5f9 100644 --- a/src/transport/middle_proxy/pool_writer/recovery_tests.rs +++ b/src/transport/middle_proxy/pool_writer/recovery_tests.rs @@ -36,7 +36,7 @@ fn prepared<'a>( pool: &'a Arc, writer: MeWriter, intent: WriterOpenIntent, - reservation: WriterOpenReservation<'a>, + reservation: WriterOpenReservation, task_started: Arc, ) -> PreparedWriter<'a> { let cancel = writer.cancel.clone(); @@ -45,12 +45,14 @@ fn prepared<'a>( byte_budget: writer.byte_budget.clone(), writer, task_registration: pool.lifecycle.try_register().unwrap(), - writer_task: Box::pin(async move { - task_started.store(true, Ordering::Release); - cancel.cancelled().await; - }), + writer_task: WriterTransport::new( + Box::pin(async move { + task_started.store(true, Ordering::Release); + cancel.cancelled().await; + }), + reservation, + ), intent, - _open_reservation: reservation, } } diff --git a/src/transport/middle_proxy/pool_writer/replacement.rs b/src/transport/middle_proxy/pool_writer/replacement.rs index b195096..5437dbe 100644 --- a/src/transport/middle_proxy/pool_writer/replacement.rs +++ b/src/transport/middle_proxy/pool_writer/replacement.rs @@ -44,9 +44,8 @@ impl MePool { tx, byte_budget, task_registration, - writer_task, - intent, - _open_reservation, + mut writer_task, + intent: _, } = prepared; let replacement_writer_id = writer.id; let victim_writer_id = reservation.writer_id(); @@ -57,7 +56,25 @@ impl MePool { let mut registry_registration = self.registry.prepare_writer_registration().await; let coordinator = self.reinit.coordinator.lock(); let contour = self.authorize_writer_publication(&writer, &coordinator)?; - self.authorize_writer_publication_capacity(&writer, contour, intent, writers.as_slice())?; + // Only this victim-backed transaction may replace Active coverage during suppression. + let family = if writer.addr.is_ipv4() { + crate::network::IpFamily::V4 + } else { + crate::network::IpFamily::V6 + }; + let configured = if writer.addr.is_ipv4() { + self.decision.ipv4_me + } else { + self.decision.ipv6_me + }; + if !configured + || (contour != WriterContour::Active + && self.is_family_temporarily_suppressed(family, Self::now_epoch_secs())) + { + return Err(ProxyError::Proxy( + "ME replacement family lost publication authority".into(), + )); + } writer.contour.store(contour.as_u8(), Ordering::Release); let Some(victim_pos) = writers .iter() @@ -73,6 +90,17 @@ impl MePool { "ME replacement victim changed role before commit".into(), )); } + if !matches!( + purpose, + WriterReplacementPurpose::FloorRebalance { .. } + | WriterReplacementPurpose::CoverageTransfer { .. } + ) && (expected_victim_role.dc != writer.writer_dc + || expected_victim_role.family != family) + { + return Err(ProxyError::Proxy( + "ME replacement purpose cannot move coverage between groups".into(), + )); + } if let WriterReplacementPurpose::FloorRebalance { donor_floor, receiver_floor, @@ -122,19 +150,103 @@ impl MePool { }) == receiver_family }) .count(); - if donor_count <= donor_floor || receiver_count >= receiver_floor { + let protected_floor = self.required_writers_for_dc_with_floor_mode( + endpoint_snapshot + .endpoints_for_dc_family(expected_victim_role.dc, expected_victim_role.family) + .len(), + false, + ); + let current_receiver_floor = self.required_writers_for_dc_with_floor_mode( + endpoint_snapshot + .endpoints_for_dc_family(writer.writer_dc, receiver_family) + .len(), + false, + ); + if donor_count <= donor_floor.max(protected_floor) + || receiver_count >= receiver_floor.min(current_receiver_floor) + { return Err(ProxyError::Proxy( "ME floor rebalance became unnecessary before commit".into(), )); } } - if !registry_registration.prepare_replacement_commit(reservation) { + if let WriterReplacementPurpose::CoverageTransfer { + authority, + receiver_floor, + } = purpose + { + if authority + != ( + coordinator.endpoint_revision, + coordinator.floor_policy_revision, + ) + || contour != WriterContour::Active + || expected_victim_role.contour != WriterContour::Active + { + return Err(ProxyError::Proxy( + "ME coverage transfer lost policy authority".into(), + )); + } + let endpoints = self.endpoint_snapshot.load(); + let count = |dc, family| { + writers + .iter() + .filter(|candidate| { + !candidate.draining.load(Ordering::Acquire) + && candidate.generation == coordinator.active_generation + && candidate.writer_dc == dc + && WriterContour::from_u8(candidate.contour.load(Ordering::Acquire)) + == WriterContour::Active + && candidate.addr.is_ipv4() + == matches!(family, crate::network::IpFamily::V4) + && endpoints.contains_dc_endpoint(dc, candidate.addr) + }) + .count() + }; + let donor_is_current = victim.generation == coordinator.active_generation + && endpoints.contains_dc_endpoint(victim.writer_dc, victim.addr); + let donor_floor = self.required_writers_for_dc_with_floor_mode( + endpoints + .endpoints_for_dc_family(victim.writer_dc, expected_victim_role.family) + .len(), + false, + ); + if (donor_is_current + && count(victim.writer_dc, expected_victim_role.family) <= donor_floor) + || count(writer.writer_dc, family) + >= receiver_floor.min( + self.required_writers_for_dc_with_floor_mode( + endpoints + .endpoints_for_dc_family(writer.writer_dc, family) + .len(), + false, + ), + ) + { + return Err(ProxyError::Proxy( + "ME coverage transfer no longer preserves floors".into(), + )); + } + } + let lifetime = registry_registration + .writer_lifetime(victim_writer_id) + .ok_or_else(|| { + ProxyError::Proxy("ME replacement victim transport disappeared".into()) + })?; + if !lifetime.commit(&mut writer_task.opening, || { + registry_registration.prepare_replacement_commit(reservation) + }) { return Err(ProxyError::Proxy( "ME replacement victim became active before commit".into(), )); } - registry_registration.install(replacement_writer_id, tx, byte_budget); + registry_registration.install_with_lifetime( + replacement_writer_id, + tx, + byte_budget, + writer_task.lifetime.clone(), + ); writers.push(writer); self.conn_count.fetch_add(1, Ordering::Relaxed); writers.publish_current(); @@ -158,6 +270,7 @@ impl MePool { #[cfg(test)] mod tests { + mod convergence; use std::collections::HashMap; use std::net::{IpAddr, Ipv4Addr}; @@ -248,14 +361,54 @@ mod tests { tx, byte_budget, task_registration, - writer_task: Box::pin(async move { - drop(rx); - }), + writer_task: WriterTransport::new( + Box::pin(async move { + drop(rx); + }), + open_reservation, + ), intent: WriterOpenIntent::Replacement, - _open_reservation: open_reservation, } } + #[tokio::test] + async fn suppressed_victim_backed_replacement_keeps_coverage() { + let pool = make_pool().await; + let addr = endpoint(8); + pool.update_proxy_maps(HashMap::from([(2, vec![(addr.ip(), addr.port())])]), None) + .await; + let victim = install_writer(&pool, 8001, 2, addr).await; + pool.set_family_runtime_state( + crate::network::IpFamily::V4, + crate::transport::middle_proxy::pool::MeFamilyRuntimeState::Suppressed, + MePool::now_epoch_secs(), + MePool::now_epoch_secs() + 3600, + 5, + 0, + ); + let mut reservation = pool + .registry + .try_reserve_writer_replacement(victim.id) + .await + .unwrap(); + let prepared = prepared_writer(&pool, 8002, 2, addr).await; + pool.publish_prepared_replacement_writer( + prepared, + WriterRole::from_writer(&victim), + WriterReplacementPurpose::IdleRefresh, + &mut reservation, + ) + .await + .expect("a reserved active victim must remain replaceable during suppression"); + assert!(victim.draining.load(Ordering::Acquire)); + assert_eq!( + pool.writer_replacement_open_reserved + .load(Ordering::Acquire), + 1, + "overlap is owned until the victim transport is destroyed" + ); + } + #[tokio::test] async fn replacement_commit_publishes_successor_before_draining_victim() { let pool = make_pool().await; diff --git a/src/transport/middle_proxy/pool_writer/replacement/tests/convergence.rs b/src/transport/middle_proxy/pool_writer/replacement/tests/convergence.rs new file mode 100644 index 0000000..c3f7761 --- /dev/null +++ b/src/transport/middle_proxy/pool_writer/replacement/tests/convergence.rs @@ -0,0 +1,198 @@ +use super::*; +use crate::transport::middle_proxy::registry::ConnMeta; + +#[tokio::test] +async fn busy_obsolete_donor_preserves_clients_and_rechecks_policy() { + for stale in [false, true] { + let pool = make_pool().await; + let obsolete = endpoint(21); + let receiver = endpoint(22); + pool.update_proxy_maps( + HashMap::from([(2, vec![(receiver.ip(), receiver.port())])]), + None, + ) + .await; + let victim = install_writer(&pool, 9001, 1, obsolete).await; + let (client, _messages) = pool.registry.register().await; + assert!( + pool.registry + .bind_writer( + client, + victim.id, + ConnMeta { + target_dc: 1, + client_addr: obsolete, + our_addr: obsolete, + proto_flags: 0, + } + ) + .await + ); + let authority = pool.floor_authority(); + let mut reservation = pool + .registry + .try_reserve_writer_replacement_preserving_clients(victim.id) + .await + .unwrap(); + let prepared = prepared_writer(&pool, 9002, 2, receiver).await; + if stale { + pool.reinit.coordinator.lock().floor_policy_revision += 1; + } + let result = pool + .publish_prepared_replacement_writer( + prepared, + WriterRole::from_writer(&victim), + WriterReplacementPurpose::CoverageTransfer { + authority, + receiver_floor: pool.required_writers_for_dc(1), + }, + &mut reservation, + ) + .await; + assert_eq!(result.is_ok(), !stale); + assert_eq!(victim.draining.load(Ordering::Acquire), !stale); + assert_eq!( + pool.registry.get_writer(client).await.unwrap().writer_id, + victim.id + ); + assert_eq!( + pool.writer_replacement_open_reserved + .load(Ordering::Acquire), + usize::from(!stale) + ); + } +} + +#[tokio::test] +async fn transfer_commit_matches_independent_donor_floor_model() { + // The model counts only current endpoints; obsolete sockets contribute no coverage. + for obsolete in [false, true] { + for donor_delta in [0usize, 1] { + for receiver_full in [false, true] { + let pool = make_pool().await; + let donor = endpoint(31); + let receiver = endpoint(32); + pool.update_proxy_maps( + HashMap::from([ + (1, vec![(donor.ip(), donor.port())]), + (2, vec![(receiver.ip(), receiver.port())]), + ]), + None, + ) + .await; + let floor = pool.required_writers_for_dc(1); + let mut victim = None; + for id in 0..(floor + donor_delta) { + let writer = install_writer(&pool, 10000 + id as u64, 1, donor).await; + if id == 0 { + victim = Some(writer); + } + } + let mut victim = victim.unwrap(); + if obsolete { + let mut writers = pool.writers.write().await; + writers + .iter_mut() + .find(|writer| writer.id == victim.id) + .unwrap() + .addr = endpoint(33); + victim.addr = endpoint(33); + } + if receiver_full { + for id in 0..floor { + install_writer(&pool, 11000 + id as u64, 2, receiver).await; + } + } + let mut reservation = pool + .registry + .try_reserve_writer_replacement(victim.id) + .await + .unwrap(); + let prepared = prepared_writer(&pool, 12000, 2, receiver).await; + let result = pool + .publish_prepared_replacement_writer( + prepared, + WriterRole::from_writer(&victim), + WriterReplacementPurpose::CoverageTransfer { + authority: pool.floor_authority(), + receiver_floor: floor, + }, + &mut reservation, + ) + .await; + assert_eq!( + result.is_ok(), + !receiver_full && (obsolete || donor_delta > 0), + "obsolete={obsolete}, surplus={donor_delta}, receiver_full={receiver_full}" + ); + } + } + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn concurrent_transfers_cannot_overfill_receiver_or_steal_donor_floor() { + let pool = make_pool().await; + let donor = endpoint(41); + let receiver = endpoint(42); + pool.update_proxy_maps( + HashMap::from([ + (1, vec![(donor.ip(), donor.port())]), + (2, vec![(receiver.ip(), receiver.port())]), + ]), + None, + ) + .await; + pool.set_adaptive_floor_runtime_caps(pool.floor_authority(), 128, 128, 128, 128, 6, 0, 0); + let floor = pool.required_writers_for_dc(1); + let count = floor + 8; + let mut victims = Vec::new(); + for id in 0..count { + victims.push(install_writer(&pool, 20000 + id as u64, 1, donor).await); + } + let mut tasks = tokio::task::JoinSet::new(); + for (index, victim) in victims.into_iter().enumerate() { + let pool = pool.clone(); + tasks.spawn(async move { + let authority = pool.floor_authority(); + let mut reservation = pool + .registry + .try_reserve_writer_replacement(victim.id) + .await + .unwrap(); + let prepared = prepared_writer(&pool, 21000 + index as u64, 2, receiver).await; + tokio::task::yield_now().await; + pool.publish_prepared_replacement_writer( + prepared, + WriterRole::from_writer(&victim), + WriterReplacementPurpose::CoverageTransfer { + authority, + receiver_floor: floor, + }, + &mut reservation, + ) + .await + .is_ok() + }); + } + let mut committed = 0; + while let Some(result) = tasks.join_next().await { + committed += usize::from(result.unwrap()); + } + assert_eq!(committed, floor); + let writers = pool.writers.read().await; + let live = |dc| { + writers + .iter() + .filter(|writer| writer.writer_dc == dc && !writer.draining.load(Ordering::Acquire)) + .count() + }; + assert_eq!(live(2), floor); + assert_eq!(live(1), count - floor); + assert!(live(1) >= floor); + assert_eq!( + pool.writer_replacement_open_reserved + .load(Ordering::Acquire), + committed + ); +} diff --git a/src/transport/middle_proxy/pool_writer/runtime.rs b/src/transport/middle_proxy/pool_writer/runtime.rs index 8ec9cdf..d15824a 100644 --- a/src/transport/middle_proxy/pool_writer/runtime.rs +++ b/src/transport/middle_proxy/pool_writer/runtime.rs @@ -291,9 +291,8 @@ impl MePool { tx, byte_budget, task_registration, - writer_task, + writer_task: WriterTransport::new(writer_task, writer_open_reservation), intent, - _open_reservation: writer_open_reservation, }) } diff --git a/src/transport/middle_proxy/registry.rs b/src/transport/middle_proxy/registry.rs index 29e4315..1655c4b 100644 --- a/src/transport/middle_proxy/registry.rs +++ b/src/transport/middle_proxy/registry.rs @@ -81,6 +81,7 @@ struct WriterRoute { tx: mpsc::Sender, byte_budget: Arc, replacement_state: Arc, + lifetime: Arc, } #[derive(Clone)] diff --git a/src/transport/middle_proxy/registry/publication.rs b/src/transport/middle_proxy/registry/publication.rs index d5870cf..97f8892 100644 --- a/src/transport/middle_proxy/registry/publication.rs +++ b/src/transport/middle_proxy/registry/publication.rs @@ -99,6 +99,34 @@ impl WriterRegistrationGuard<'_> { writer_id: u64, tx: mpsc::Sender, byte_budget: Arc, + ) { + self.install_with_lifetime( + writer_id, + tx, + byte_budget, + super::super::pool_writer::overlap::ReplacementHandoff::new(), + ); + } + + /// Retrieves the physical lifetime while binding ownership fences victim replacement. + pub(in crate::transport::middle_proxy) fn writer_lifetime( + &self, + writer_id: u64, + ) -> Option> { + self.registry + .writers + .map + .get(&writer_id) + .map(|route| route.lifetime.clone()) + } + + /// Publishes a route associated with its transport destruction authority. + pub(in crate::transport::middle_proxy) fn install_with_lifetime( + &mut self, + writer_id: u64, + tx: mpsc::Sender, + byte_budget: Arc, + lifetime: Arc, ) { self.binding.conns_for_writer.entry(writer_id).or_default(); self.registry @@ -117,6 +145,7 @@ impl WriterRegistrationGuard<'_> { tx, byte_budget, replacement_state: Arc::new(AtomicU8::new(WriterReplacementState::Open as u8)), + lifetime, }, ); } diff --git a/src/transport/middle_proxy/registry/writer.rs b/src/transport/middle_proxy/registry/writer.rs index f242a68..efe3c34 100644 --- a/src/transport/middle_proxy/registry/writer.rs +++ b/src/transport/middle_proxy/registry/writer.rs @@ -415,15 +415,10 @@ impl ConnRegistry { .map .get(&conn_id) .map(|entry| entry.writer_id)?; - let writer = self - .writers - .map - .get(&writer_id) - .map(|entry| entry.value().clone())?; - Some(ConnWriter { + self.writers.map.get(&writer_id).map(|entry| ConnWriter { writer_id, - tx: writer.tx, - byte_budget: writer.byte_budget, + tx: entry.tx.clone(), + byte_budget: entry.byte_budget.clone(), }) } @@ -436,19 +431,12 @@ impl ConnRegistry { let hot = self.hot_binding.map.get(&conn_id)?; let writer_id = hot.writer_id; let meta = hot.meta.clone(); - let writer = self - .writers - .map - .get(&writer_id) - .map(|entry| entry.value().clone())?; - Some(( - ConnWriter { - writer_id, - tx: writer.tx, - byte_budget: writer.byte_budget, - }, - meta, - )) + let writer = self.writers.map.get(&writer_id).map(|entry| ConnWriter { + writer_id, + tx: entry.tx.clone(), + byte_budget: entry.byte_budget.clone(), + })?; + Some((writer, meta)) } pub async fn active_conn_ids(&self) -> Vec { diff --git a/src/transport/middle_proxy/rotation.rs b/src/transport/middle_proxy/rotation.rs index 90a08e1..1320374 100644 --- a/src/transport/middle_proxy/rotation.rs +++ b/src/transport/middle_proxy/rotation.rs @@ -9,10 +9,17 @@ use crate::config::ProxyConfig; use crate::crypto::SecureRandom; use super::MePool; +// Durable scheduler intent is independent of pending generation creation. +mod objective; +#[cfg(test)] +mod retry_tests; +/// External causes that request generation reconciliation. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum MeReinitTrigger { + /// The configured rotation interval elapsed. Periodic, + /// Endpoint authority changed. MapChanged, } @@ -25,6 +32,7 @@ impl MeReinitTrigger { } } +/// Coalesces an external request; endpoint epochs independently preserve map-change intent. pub fn enqueue_reinit_trigger(tx: &mpsc::Sender, trigger: MeReinitTrigger) { match tx.try_send(trigger) { Ok(()) => {} @@ -82,6 +90,7 @@ fn trigger_reason(pending: u8) -> &'static str { } } +/// Retries durable reinitialization objectives until their exact endpoint revision converges. pub async fn me_reinit_scheduler( pool: Arc, rng: Arc, @@ -90,9 +99,14 @@ pub async fn me_reinit_scheduler( me_ready_tx: watch::Sender, ) { info!("ME reinit scheduler started"); - let mut tasks = JoinSet::::new(); - let mut pending = 0u8; - let mut pending_deadline = None; + let mut tasks = JoinSet::>::new(); + let mut owners = std::collections::HashMap::new(); + let mut objective: Option = None; + let mut next_id = 1u64; + let mut epochs = pool.writer_epoch.subscribe(); + let mut epoch_open = true; + let mut observed_revision = pool.endpoint_snapshot.load().revision; + let mut completed_revision = None; let mut trigger_channel_open = true; loop { @@ -101,73 +115,152 @@ pub async fn me_reinit_scheduler( pool.reinit .max_concurrency_effective .store(max_concurrency, std::sync::atomic::Ordering::Release); - - 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 snapshot = pool.endpoint_snapshot.load_full(); + if snapshot.revision != observed_revision { + observed_revision = snapshot.revision; + objective = Some(objective::RetryObjective::new( + next_id, + observed_revision, + REINIT_TRIGGER_MAP_CHANGED, + Duration::from_millis(cfg.general.me_reinit_coalesce_window_ms), + )); + next_id = next_id.wrapping_add(1); + } + let owned = objective + .as_ref() + .is_some_and(|current| owners.values().any(|identity| current.matches(*identity))); + if let Some(current) = objective.as_mut() + && !owned + && current.attempt > 0 + { + let status = pool.reinit.status.load_full(); + let authority = pool.floor_authority(); + let readiness = pool.api_hardswap_snapshot().await; + // An API projection is usable only while its exact publication and policy survive + // the await. A changed snapshot is unknown, not a not-ready transition. + if Arc::ptr_eq(&status, &pool.reinit.status.load_full()) + && authority == pool.floor_authority() + { + if status.pending_hardswap_endpoint_revision == current.revision + && readiness.pending_map_current == Some(true) + && readiness.pending_writer_deficit == 0 + { + current.ready(( + status.pending_hardswap_generation, + status.pending_hardswap_map_hash, + status.pending_hardswap_endpoint_revision, + )); + } else { + current.not_ready(); + } + } + } + let can_launch = !owned + && tasks.len() < max_concurrency + && objective + .as_ref() + .is_some_and(|current| current.attempt == 0 || tasks.is_empty()); + if can_launch + && objective.as_ref().is_some_and(|current| { + current + .deadline + .is_some_and(|deadline| deadline <= tokio::time::Instant::now()) + }) + { + if pool.endpoint_snapshot.load().revision != snapshot.revision { + continue; + } + let Some(current) = objective.as_mut() else { + continue; + }; + let identity = (current.id, current.revision); + current.deadline = None; + debug!( + reason = trigger_reason(current.triggers), + attempt = current.attempt, + endpoint_revision = current.revision, + "ME reinit scheduled" + ); let pool_clone = pool.clone(); let rng_clone = rng.clone(); pool.reinit .scheduler_inflight .fetch_add(1, std::sync::atomic::Ordering::AcqRel); let inflight = ReinitInflightGuard { - pool: Arc::clone(&pool_clone), + pool: pool_clone.clone(), }; - tasks.spawn(async move { + let future = async move { let _inflight = inflight; pool_clone - .zero_downtime_reinit_periodic(rng_clone.as_ref()) + .reinit_endpoint_snapshot(rng_clone.as_ref(), snapshot) .await - }); + }; + // Pool shutdown joins destruction of attempts as well as their writer transports. + let Ok(tracked) = pool.lifecycle.track_producer(future) else { + return; + }; + let task = tasks.spawn(tracked); + owners.insert(task.id(), identity); continue; } - - if !trigger_channel_open && pending == 0 && tasks.is_empty() { - warn!("ME reinit scheduler stopped: trigger channel closed"); + if !trigger_channel_open && objective.is_none() && tasks.is_empty() { return; } - - let timer_enabled = pending != 0 && tasks.len() < max_concurrency; + let deadline = objective.as_ref().and_then(|current| current.deadline); 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, - ), - ); + let revision = pool.endpoint_snapshot.load().revision; + if trigger == MeReinitTrigger::MapChanged && completed_revision == Some(revision) { + continue; } - pending |= trigger_bit(trigger); + if let Some(current) = objective.as_mut() + && current.revision == revision + && (trigger == MeReinitTrigger::MapChanged + || !owners.values().any(|identity| current.matches(*identity))) { + // Endpoint epochs and updater notifications describe the same objective. + current.triggers |= trigger_bit(trigger); + } else { + objective = Some(objective::RetryObjective::new(next_id, revision, + trigger_bit(trigger), Duration::from_millis(cfg.general.me_reinit_coalesce_window_ms))); + next_id = next_id.wrapping_add(1); + } + observed_revision = revision; } 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)) => { + changed = epochs.changed(), if epoch_open => { + if changed.is_err() { epoch_open = false; } + } + joined = tasks.join_next_with_id(), if !tasks.is_empty() => { + let Some(joined) = joined else { continue; }; + let (id, succeeded) = match joined { + Ok((id, Some(succeeded))) => (id, succeeded), + Ok((_, None)) => return, + Err(error) => { warn!(error = %error, "ME reinit task failed"); + (error.id(), false) + } + }; + let Some(identity) = owners.remove(&id) else { continue; }; + if succeeded { + me_ready_tx.send_modify(|version| *version = version.saturating_add(1)); + } + if let Some(current) = objective.as_mut() + && current.matches(identity) { + if succeeded { + completed_revision = Some(current.revision); + objective = None; + } else { + current.failed(Duration::from_millis(pool.hardswap_warmup_backoff_ms(current.attempt))); } - None => {} } } _ = async { - if let Some(deadline) = pending_deadline { - tokio::time::sleep_until(deadline).await; - } - }, if timer_enabled => {} + if let Some(deadline) = deadline { tokio::time::sleep_until(deadline).await; } + }, if can_launch && deadline.is_some() => {} } } } diff --git a/src/transport/middle_proxy/rotation/objective.rs b/src/transport/middle_proxy/rotation/objective.rs new file mode 100644 index 0000000..a6bf3d0 --- /dev/null +++ b/src/transport/middle_proxy/rotation/objective.rs @@ -0,0 +1,55 @@ +use std::time::Duration; +use tokio::time::Instant; + +/// Durable intent exists even if a reinit returns before creating a pending generation. +pub(super) struct RetryObjective { + /// Distinguishes separate external requests for the same endpoint revision. + pub(super) id: u64, + /// Exact endpoint snapshot to reconcile. + pub(super) revision: u64, + /// Coalesced external causes used for diagnostic logging. + pub(super) triggers: u8, + /// Completed failures; the count saturates but never terminates recovery. + pub(super) attempt: usize, + /// A missing deadline means the objective currently owns a running task. + pub(super) deadline: Option, + ready_wake: Option<(u64, u64, u64)>, +} + +impl RetryObjective { + /// Creates retryable intent before any asynchronous attempt can fail. + pub(super) fn new(id: u64, revision: u64, triggers: u8, delay: Duration) -> Self { + Self { + id, + revision, + triggers, + attempt: 0, + deadline: Some(Instant::now() + delay), + ready_wake: None, + } + } + + /// Rejects completions from superseded objectives, including the same endpoint revision. + pub(super) fn matches(&self, identity: (u64, u64)) -> bool { + (self.id, self.revision) == identity + } + + /// Measures backoff from completion and never caps the number of retry attempts. + pub(super) fn failed(&mut self, delay: Duration) { + self.attempt = self.attempt.saturating_add(1); + self.deadline = Some(Instant::now() + delay); + } + + /// Accelerates once per observed ready tuple without allowing an epoch storm to spin. + pub(super) fn ready(&mut self, key: (u64, u64, u64)) { + if self.ready_wake != Some(key) { + self.ready_wake = Some(key); + self.deadline = Some(Instant::now()); + } + } + + /// Re-arms acceleration only after a coherent observation of lost readiness. + pub(super) fn not_ready(&mut self) { + self.ready_wake = None; + } +} diff --git a/src/transport/middle_proxy/rotation/retry_tests.rs b/src/transport/middle_proxy/rotation/retry_tests.rs new file mode 100644 index 0000000..df11185 --- /dev/null +++ b/src/transport/middle_proxy/rotation/retry_tests.rs @@ -0,0 +1,188 @@ +use super::objective::RetryObjective; +use super::*; + +#[tokio::test(start_paused = true)] +async fn retry_objective_survives_unlimited_failures_and_ready_storms() { + let mut objective = RetryObjective::new(1, 2, REINIT_TRIGGER_PERIODIC, Duration::ZERO); + assert!(objective.deadline.is_some()); + for _ in 0..100 { + objective.deadline = None; + tokio::time::advance(Duration::from_secs(70)).await; + let completed = tokio::time::Instant::now(); + objective.failed(Duration::from_secs(30)); + assert_eq!( + objective.deadline, + Some(completed + Duration::from_secs(30)) + ); + } + objective.ready((3, 4, 2)); + objective.failed(Duration::from_secs(30)); + let deadline = objective.deadline; + for _ in 0..100 { + objective.ready((3, 4, 2)); + } + assert_eq!(objective.deadline, deadline); + objective.not_ready(); + objective.ready((3, 4, 2)); + assert!( + objective.deadline < deadline, + "a genuine loss and recovery of the same floor re-arms wakeup" + ); + assert!(!objective.matches((1, 1))); + assert!(!objective.matches((2, 2))); +} + +#[tokio::test] +async fn pre_pending_failure_remains_live_without_a_periodic_trigger() { + let pool = crate::transport::middle_proxy::pool_writer_security_tests::make_pool().await; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + pool.update_proxy_maps( + std::collections::HashMap::from([(2, vec![(addr.ip(), addr.port())])]), + None, + ) + .await; + pool.set_family_runtime_state( + crate::network::IpFamily::V4, + crate::transport::middle_proxy::pool::MeFamilyRuntimeState::Suppressed, + MePool::now_epoch_secs(), + MePool::now_epoch_secs() + 3600, + 5, + 0, + ); + pool.reinit + .me_hardswap_warmup_pass_backoff_base_ms + .store(10, std::sync::atomic::Ordering::Relaxed); + pool.reinit + .me_hardswap_warmup_delay_min_ms + .store(0, std::sync::atomic::Ordering::Relaxed); + pool.reinit + .me_hardswap_warmup_delay_max_ms + .store(0, std::sync::atomic::Ordering::Relaxed); + let mut config = ProxyConfig::default(); + config.general.me_reinit_coalesce_window_ms = 0; + let (_config_tx, config_rx) = watch::channel(Arc::new(config)); + let (tx, rx) = mpsc::channel(1); + let (ready_tx, _) = watch::channel(0); + tx.send(MeReinitTrigger::Periodic).await.unwrap(); + drop(tx); + let scheduler = tokio::spawn(me_reinit_scheduler( + pool.clone(), + pool.rng.clone(), + config_rx, + rx, + ready_tx, + )); + tokio::time::sleep(Duration::from_millis(50)).await; + assert!( + !scheduler.is_finished(), + "false before pending creation must not discard recovery intent" + ); + assert!(pool.reinit.coordinator.lock().pending.is_none()); + // Family policy changes do not emit endpoint or writer epochs. + pool.set_family_runtime_state( + crate::network::IpFamily::V4, + crate::transport::middle_proxy::pool::MeFamilyRuntimeState::Healthy, + MePool::now_epoch_secs(), + 0, + 0, + 0, + ); + let (stream, _) = tokio::time::timeout(Duration::from_secs(2), listener.accept()) + .await + .expect("autonomous retry must actually reach TCP without another trigger") + .unwrap(); + assert!(pool.shutdown_until(Duration::from_secs(1)).await); + tokio::time::timeout(Duration::from_secs(1), scheduler) + .await + .unwrap() + .unwrap(); + assert_eq!( + pool.reinit + .scheduler_inflight + .load(std::sync::atomic::Ordering::Acquire), + 0 + ); + drop(stream); +} + +#[tokio::test] +async fn panic_retains_the_task_identity_for_retry() { + let mut tasks = JoinSet::::new(); + let id = tasks.spawn(async { panic!("injected reinit panic") }).id(); + let error = tasks.join_next_with_id().await.unwrap().unwrap_err(); + assert_eq!(error.id(), id); +} + +async fn check_inflight_trigger(trigger: MeReinitTrigger, starts_successor: bool) { + let pool = crate::transport::middle_proxy::pool_writer_security_tests::make_pool().await; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + pool.update_proxy_maps( + std::collections::HashMap::from([(2, vec![(addr.ip(), addr.port())])]), + None, + ) + .await; + pool.reinit + .me_hardswap_warmup_delay_min_ms + .store(0, std::sync::atomic::Ordering::Relaxed); + pool.reinit + .me_hardswap_warmup_delay_max_ms + .store(0, std::sync::atomic::Ordering::Relaxed); + let mut config = ProxyConfig::default(); + config.general.me_reinit_coalesce_window_ms = 0; + config.general.me_reinit_singleflight = false; + config.general.me_reinit_max_concurrency = 2; + let (_config_tx, config_rx) = watch::channel(Arc::new(config)); + let (tx, rx) = mpsc::channel(1); + let (ready_tx, _) = watch::channel(0); + tx.send(MeReinitTrigger::MapChanged).await.unwrap(); + let scheduler = tokio::spawn(me_reinit_scheduler( + pool.clone(), + pool.rng.clone(), + config_rx, + rx, + ready_tx, + )); + let (first, _) = tokio::time::timeout(Duration::from_secs(1), listener.accept()) + .await + .unwrap() + .unwrap(); + assert_eq!( + pool.reinit + .scheduler_inflight + .load(std::sync::atomic::Ordering::Acquire), + 1 + ); + // Keep the first attempt blocked before handshake completion while delivering another trigger. + tx.send(trigger).await.unwrap(); + let successor = tokio::time::timeout(Duration::from_millis(100), listener.accept()).await; + let inflight = pool + .reinit + .scheduler_inflight + .load(std::sync::atomic::Ordering::Acquire); + assert!(pool.shutdown_until(Duration::from_secs(1)).await); + tokio::time::timeout(Duration::from_secs(1), scheduler) + .await + .unwrap() + .unwrap(); + assert_eq!(successor.is_ok(), starts_successor); + assert_eq!(inflight, if starts_successor { 2 } else { 1 }); + assert_eq!( + pool.reinit + .scheduler_inflight + .load(std::sync::atomic::Ordering::Acquire), + 0 + ); + drop(first); +} + +#[tokio::test] +async fn duplicate_map_change_coalesces_with_its_inflight_revision() { + check_inflight_trigger(MeReinitTrigger::MapChanged, false).await; +} + +#[tokio::test] +async fn periodic_trigger_retains_successor_intent_while_map_change_runs() { + check_inflight_trigger(MeReinitTrigger::Periodic, true).await; +}