From 1ea5f7a1d687bc09125880ba32f33c4eb7bc4c49 Mon Sep 17 00:00:00 2001 From: Alexey <247128645+axkurcom@users.noreply.github.com> Date: Sat, 12 Sep 2026 20:14:41 +0300 Subject: [PATCH] Atomic Lifecycle for Writer Refresh + Hardswap --- src/api/runtime_min.rs | 174 +------- src/api/runtime_min/me_pool.rs | 194 +++++++++ src/metrics.rs | 5 + src/metrics/render.rs | 4 + src/metrics/render/me_hardswap.rs | 156 ++++++++ src/metrics/tests.rs | 6 + src/transport/middle_proxy/health.rs | 8 +- src/transport/middle_proxy/health/drain.rs | 2 +- src/transport/middle_proxy/health/family.rs | 69 +++- .../middle_proxy/health/floor_plan.rs | 4 + .../middle_proxy/health/idle_refresh.rs | 234 +++++++---- src/transport/middle_proxy/health/recovery.rs | 32 +- src/transport/middle_proxy/health/tests.rs | 75 +++- .../middle_proxy/health/zombie_watchdog.rs | 2 +- src/transport/middle_proxy/mod.rs | 1 + src/transport/middle_proxy/pool.rs | 43 +- .../middle_proxy/pool/construction.rs | 1 + .../middle_proxy/pool/writer_admission.rs | 147 ++++++- src/transport/middle_proxy/pool_config.rs | 33 +- src/transport/middle_proxy/pool_init.rs | 14 +- src/transport/middle_proxy/pool_refill.rs | 243 +++++++---- src/transport/middle_proxy/pool_reinit.rs | 26 +- .../middle_proxy/pool_reinit/coordination.rs | 212 ++++++++-- .../middle_proxy/pool_reinit/reconcile.rs | 146 ++++--- .../middle_proxy/pool_reinit/tests.rs | 280 ++++++++++++- .../middle_proxy/pool_runtime_api.rs | 11 +- src/transport/middle_proxy/pool_status.rs | 17 + .../pool_status/hardswap_snapshot.rs | 102 +++++ .../pool_status/runtime_snapshot.rs | 8 + .../pool_status/status_snapshot.rs | 29 +- src/transport/middle_proxy/pool_writer.rs | 50 ++- .../middle_proxy/pool_writer/publication.rs | 180 +++++++++ .../middle_proxy/pool_writer/replacement.rs | 376 ++++++++++++++++++ .../middle_proxy/pool_writer/runtime.rs | 141 ++++--- src/transport/middle_proxy/registry.rs | 16 + .../middle_proxy/registry/publication.rs | 83 +++- .../middle_proxy/registry/replacement.rs | 142 +++++++ src/transport/middle_proxy/registry/tests.rs | 79 +++- src/transport/middle_proxy/registry/writer.rs | 32 +- src/transport/middle_proxy/send.rs | 44 +- .../tests/pool_refill_security_tests.rs | 32 +- .../tests/pool_writer_publication_tests.rs | 113 +++++- .../tests/pool_writer_security_tests.rs | 5 +- .../tests/send_adversarial_tests.rs | 30 ++ 44 files changed, 3020 insertions(+), 581 deletions(-) create mode 100644 src/api/runtime_min/me_pool.rs create mode 100644 src/metrics/render/me_hardswap.rs create mode 100644 src/transport/middle_proxy/pool_status/hardswap_snapshot.rs create mode 100644 src/transport/middle_proxy/pool_writer/publication.rs create mode 100644 src/transport/middle_proxy/pool_writer/replacement.rs create mode 100644 src/transport/middle_proxy/registry/replacement.rs diff --git a/src/api/runtime_min.rs b/src/api/runtime_min.rs index 5338baf..a89f5bc 100644 --- a/src/api/runtime_min.rs +++ b/src/api/runtime_min.rs @@ -1,4 +1,3 @@ -use std::collections::BTreeSet; use std::time::{SystemTime, UNIX_EPOCH}; use serde::Serialize; @@ -17,82 +16,6 @@ pub(super) struct SecurityWhitelistData { pub(super) entries: Vec, } -#[derive(Serialize)] -pub(super) struct RuntimeMePoolStateGenerationData { - pub(super) active_generation: u64, - pub(super) warm_generation: u64, - pub(super) warm_generations: Vec, - pub(super) pending_hardswap_generation: u64, - pub(super) pending_hardswap_age_secs: Option, - pub(super) reinit_inflight: usize, - pub(super) reinit_max_concurrency_effective: usize, - pub(super) draining_generations: Vec, -} - -#[derive(Serialize)] -pub(super) struct RuntimeMePoolStateHardswapData { - pub(super) enabled: bool, - pub(super) pending: bool, -} - -#[derive(Serialize)] -pub(super) struct RuntimeMePoolStateWriterContourData { - pub(super) warm: usize, - pub(super) active: usize, - pub(super) draining: usize, -} - -#[derive(Serialize)] -pub(super) struct RuntimeMePoolStateWriterHealthData { - pub(super) healthy: usize, - pub(super) degraded: usize, - pub(super) draining: usize, -} - -#[derive(Serialize)] -pub(super) struct RuntimeMePoolStateWriterData { - pub(super) total: usize, - pub(super) alive_non_draining: usize, - pub(super) draining: usize, - pub(super) degraded: usize, - pub(super) contour: RuntimeMePoolStateWriterContourData, - pub(super) health: RuntimeMePoolStateWriterHealthData, -} - -#[derive(Serialize)] -pub(super) struct RuntimeMePoolStateRefillDcData { - pub(super) dc: i16, - pub(super) family: &'static str, - pub(super) inflight: usize, -} - -#[derive(Serialize)] -pub(super) struct RuntimeMePoolStateRefillData { - pub(super) inflight_endpoints_total: usize, - pub(super) inflight_dc_total: usize, - pub(super) running_dc_total: usize, - pub(super) pending_dc_total: usize, - pub(super) by_dc: Vec, -} - -#[derive(Serialize)] -pub(super) struct RuntimeMePoolStatePayload { - pub(super) generations: RuntimeMePoolStateGenerationData, - pub(super) hardswap: RuntimeMePoolStateHardswapData, - pub(super) writers: RuntimeMePoolStateWriterData, - pub(super) refill: RuntimeMePoolStateRefillData, -} - -#[derive(Serialize)] -pub(super) struct RuntimeMePoolStateData { - pub(super) enabled: bool, - #[serde(skip_serializing_if = "Option::is_none")] - pub(super) reason: Option<&'static str>, - pub(super) generated_at_epoch_secs: u64, - #[serde(skip_serializing_if = "Option::is_none")] - pub(super) data: Option, -} - #[derive(Serialize)] pub(super) struct RuntimeMeQualityCountersData { pub(super) idle_close_by_peer_total: u64, @@ -285,100 +208,6 @@ pub(super) fn build_security_whitelist_data(cfg: &ProxyConfig) -> SecurityWhitel } } -pub(super) async fn build_runtime_me_pool_state_data(shared: &ApiShared) -> RuntimeMePoolStateData { - let now_epoch_secs = now_epoch_secs(); - let Some(pool) = shared.me_pool.read().await.clone() else { - return RuntimeMePoolStateData { - enabled: false, - reason: Some(SOURCE_UNAVAILABLE_REASON), - generated_at_epoch_secs: now_epoch_secs, - data: None, - }; - }; - - let (status, runtime) = pool.api_coherent_snapshots().await; - let refill = pool.api_refill_snapshot().await; - - let mut draining_generations = BTreeSet::::new(); - let mut contour_warm = 0usize; - let mut contour_active = 0usize; - let mut contour_draining = 0usize; - let mut draining = 0usize; - let mut degraded = 0usize; - let mut healthy = 0usize; - - for writer in &status.writers { - if writer.draining { - draining_generations.insert(writer.generation); - draining += 1; - } - if writer.degraded && !writer.draining { - degraded += 1; - } - if !writer.degraded && !writer.draining { - healthy += 1; - } - match writer.state { - "warm" => contour_warm += 1, - "active" => contour_active += 1, - _ => contour_draining += 1, - } - } - - RuntimeMePoolStateData { - enabled: true, - reason: None, - generated_at_epoch_secs: status.generated_at_epoch_secs, - data: Some(RuntimeMePoolStatePayload { - generations: RuntimeMePoolStateGenerationData { - active_generation: runtime.active_generation, - warm_generation: runtime.warm_generation, - warm_generations: runtime.warm_generations, - pending_hardswap_generation: runtime.pending_hardswap_generation, - pending_hardswap_age_secs: runtime.pending_hardswap_age_secs, - reinit_inflight: runtime.reinit_inflight, - reinit_max_concurrency_effective: runtime.reinit_max_concurrency_effective, - draining_generations: draining_generations.into_iter().collect(), - }, - hardswap: RuntimeMePoolStateHardswapData { - enabled: runtime.hardswap_enabled, - pending: runtime.pending_hardswap_generation != 0, - }, - writers: RuntimeMePoolStateWriterData { - total: status.writers.len(), - alive_non_draining: status.writers.len().saturating_sub(draining), - draining, - degraded, - contour: RuntimeMePoolStateWriterContourData { - warm: contour_warm, - active: contour_active, - draining: contour_draining, - }, - health: RuntimeMePoolStateWriterHealthData { - healthy, - degraded, - draining, - }, - }, - refill: RuntimeMePoolStateRefillData { - inflight_endpoints_total: refill.inflight_endpoints_total, - inflight_dc_total: refill.inflight_dc_total, - running_dc_total: refill.running_dc_total, - pending_dc_total: refill.pending_dc_total, - by_dc: refill - .by_dc - .into_iter() - .map(|entry| RuntimeMePoolStateRefillDcData { - dc: entry.dc, - family: entry.family, - inflight: entry.inflight, - }) - .collect(), - }, - }), - } -} - pub(super) async fn build_runtime_me_quality_data(shared: &ApiShared) -> RuntimeMeQualityData { let now_epoch_secs = now_epoch_secs(); let Some(pool) = shared.me_pool.read().await.clone() else { @@ -541,7 +370,10 @@ pub(super) async fn build_runtime_upstream_quality_data( } } +// ME pool runtime-state projection. +mod me_pool; // NAT/STUN runtime projection and timestamping. mod nat; +pub(super) use me_pool::build_runtime_me_pool_state_data; pub(super) use nat::build_runtime_nat_stun_data; use nat::now_epoch_secs; diff --git a/src/api/runtime_min/me_pool.rs b/src/api/runtime_min/me_pool.rs new file mode 100644 index 0000000..632449c --- /dev/null +++ b/src/api/runtime_min/me_pool.rs @@ -0,0 +1,194 @@ +//! ME pool runtime-state projection. + +use std::collections::BTreeSet; + +use serde::Serialize; + +use super::{now_epoch_secs, ApiShared, SOURCE_UNAVAILABLE_REASON}; + +#[derive(Serialize)] +struct RuntimeMePoolStateGenerationData { + active_generation: u64, + warm_generation: u64, + warm_generations: Vec, + pending_hardswap_generation: u64, + pending_hardswap_age_secs: Option, + reinit_inflight: usize, + reinit_max_concurrency_effective: usize, + draining_generations: Vec, +} + +#[derive(Serialize)] +struct RuntimeMePoolStateHardswapData { + enabled: bool, + pending: bool, + pending_writers_current: usize, + pending_writer_deficit: usize, + pending_missing_dc_groups: usize, + pending_map_current: Option, + orphan_warm_writers_current: usize, + replacement_preparing_current: usize, + replacement_retiring_current: usize, +} + +#[derive(Serialize)] +struct RuntimeMePoolStateWriterContourData { + warm: usize, + active: usize, + draining: usize, +} + +#[derive(Serialize)] +struct RuntimeMePoolStateWriterHealthData { + healthy: usize, + degraded: usize, + draining: usize, +} + +#[derive(Serialize)] +struct RuntimeMePoolStateWriterData { + total: usize, + alive_non_draining: usize, + draining: usize, + degraded: usize, + contour: RuntimeMePoolStateWriterContourData, + health: RuntimeMePoolStateWriterHealthData, +} + +#[derive(Serialize)] +struct RuntimeMePoolStateRefillDcData { + dc: i16, + family: &'static str, + inflight: usize, +} + +#[derive(Serialize)] +struct RuntimeMePoolStateRefillData { + inflight_endpoints_total: usize, + inflight_dc_total: usize, + running_dc_total: usize, + pending_dc_total: usize, + by_dc: Vec, +} + +#[derive(Serialize)] +struct RuntimeMePoolStatePayload { + generations: RuntimeMePoolStateGenerationData, + hardswap: RuntimeMePoolStateHardswapData, + writers: RuntimeMePoolStateWriterData, + refill: RuntimeMePoolStateRefillData, +} + +#[derive(Serialize)] +struct RuntimeMePoolStateData { + enabled: bool, + #[serde(skip_serializing_if = "Option::is_none")] + reason: Option<&'static str>, + generated_at_epoch_secs: u64, + #[serde(skip_serializing_if = "Option::is_none")] + data: Option, +} + +/// Builds the bounded runtime ME pool response projection. +pub(in crate::api) async fn build_runtime_me_pool_state_data( + shared: &ApiShared, +) -> impl Serialize { + let now_epoch_secs = now_epoch_secs(); + let Some(pool) = shared.me_pool.read().await.clone() else { + return RuntimeMePoolStateData { + enabled: false, + reason: Some(SOURCE_UNAVAILABLE_REASON), + generated_at_epoch_secs: now_epoch_secs, + data: None, + }; + }; + + let (status, runtime) = pool.api_coherent_snapshots().await; + let refill = pool.api_refill_snapshot().await; + + let mut draining_generations = BTreeSet::::new(); + let mut contour_warm = 0usize; + let mut contour_active = 0usize; + let mut contour_draining = 0usize; + let mut draining = 0usize; + let mut degraded = 0usize; + let mut healthy = 0usize; + + for writer in &status.writers { + if writer.draining { + draining_generations.insert(writer.generation); + draining += 1; + } + if writer.degraded && !writer.draining { + degraded += 1; + } + if !writer.degraded && !writer.draining { + healthy += 1; + } + match writer.state { + "warm" => contour_warm += 1, + "active" => contour_active += 1, + _ => contour_draining += 1, + } + } + + RuntimeMePoolStateData { + enabled: true, + reason: None, + generated_at_epoch_secs: status.generated_at_epoch_secs, + data: Some(RuntimeMePoolStatePayload { + generations: RuntimeMePoolStateGenerationData { + active_generation: runtime.active_generation, + warm_generation: runtime.warm_generation, + warm_generations: runtime.warm_generations, + pending_hardswap_generation: runtime.pending_hardswap_generation, + pending_hardswap_age_secs: runtime.pending_hardswap_age_secs, + reinit_inflight: runtime.reinit_inflight, + reinit_max_concurrency_effective: runtime.reinit_max_concurrency_effective, + draining_generations: draining_generations.into_iter().collect(), + }, + hardswap: RuntimeMePoolStateHardswapData { + enabled: runtime.hardswap_enabled, + pending: runtime.pending_hardswap_generation != 0, + pending_writers_current: runtime.pending_writers_current, + pending_writer_deficit: runtime.pending_writer_deficit, + pending_missing_dc_groups: runtime.pending_missing_dc_groups, + pending_map_current: runtime.pending_map_current, + orphan_warm_writers_current: runtime.orphan_warm_writers_current, + replacement_preparing_current: runtime.replacement_preparing_current, + replacement_retiring_current: runtime.replacement_retiring_current, + }, + writers: RuntimeMePoolStateWriterData { + total: status.writers.len(), + alive_non_draining: status.writers.len().saturating_sub(draining), + draining, + degraded, + contour: RuntimeMePoolStateWriterContourData { + warm: contour_warm, + active: contour_active, + draining: contour_draining, + }, + health: RuntimeMePoolStateWriterHealthData { + healthy, + degraded, + draining, + }, + }, + refill: RuntimeMePoolStateRefillData { + inflight_endpoints_total: refill.inflight_endpoints_total, + inflight_dc_total: refill.inflight_dc_total, + running_dc_total: refill.running_dc_total, + pending_dc_total: refill.pending_dc_total, + by_dc: refill + .by_dc + .into_iter() + .map(|entry| RuntimeMePoolStateRefillDcData { + dc: entry.dc, + family: entry.family, + inflight: entry.inflight, + }) + .collect(), + }, + }), + } +} diff --git a/src/metrics.rs b/src/metrics.rs index dfa69ba..74a7796 100644 --- a/src/metrics.rs +++ b/src/metrics.rs @@ -236,6 +236,10 @@ async fn handle( let config = runtime.config(); if req.uri().path() == "/metrics" { + let me_hardswap = match runtime.current_me_pool().await { + Some(pool) => Some(pool.api_hardswap_snapshot().await), + None => None, + }; let body = render_metrics( stats, shared_state, @@ -244,6 +248,7 @@ async fn handle( tls_cache, tls_full_cert_budget, web_publication, + me_hardswap.as_ref(), ) .await; let resp = Response::builder() diff --git a/src/metrics/render.rs b/src/metrics/render.rs index e7af27d..b900143 100644 --- a/src/metrics/render.rs +++ b/src/metrics/render.rs @@ -12,6 +12,8 @@ mod me_lifecycle; mod me_buffers; // ME writer selection, KDF, and hardswap metrics. mod me_policy; +// Live hardswap ownership and replacement progress metrics. +mod me_hardswap; // Adaptive-floor and writer-cap metrics. mod me_floor; // Desync, pool recovery, and refill metrics. @@ -27,6 +29,7 @@ pub(super) async fn render_metrics( tls_cache: Option<&TlsFrontCache>, tls_full_cert_budget: &TlsFullCertBudget, web_publication: &crate::web::control::WebRuntimePublication, + me_hardswap: Option<&crate::transport::middle_proxy::MeApiHardswapSnapshot>, ) -> String { let mut out = String::with_capacity(4096); let telemetry = stats.telemetry_policy(); @@ -62,6 +65,7 @@ pub(super) async fn render_metrics( me_allows_debug, ); me_policy::render(&mut out, stats, me_allows_normal, me_allows_debug); + me_hardswap::render(&mut out, me_hardswap, me_allows_normal); me_floor::render(&mut out, stats, config, me_allows_normal); me_recovery::render(&mut out, stats, me_allows_normal, me_allows_debug); users::render( diff --git a/src/metrics/render/me_hardswap.rs b/src/metrics/render/me_hardswap.rs new file mode 100644 index 0000000..e3efe21 --- /dev/null +++ b/src/metrics/render/me_hardswap.rs @@ -0,0 +1,156 @@ +use std::fmt::Write; + +use crate::transport::middle_proxy::MeApiHardswapSnapshot; + +pub(super) fn render( + out: &mut String, + snapshot: Option<&MeApiHardswapSnapshot>, + enabled: bool, +) { + let snapshot = enabled.then_some(snapshot).flatten(); + let pending = snapshot.is_some_and(|value| value.pending); + let pending_age_secs = snapshot + .and_then(|value| value.pending_age_secs) + .unwrap_or(0); + let pending_writers_current = snapshot + .map(|value| value.pending_writers_current) + .unwrap_or(0); + let pending_writer_deficit = snapshot + .map(|value| value.pending_writer_deficit) + .unwrap_or(0); + let pending_missing_dc_groups = snapshot + .map(|value| value.pending_missing_dc_groups) + .unwrap_or(0); + let pending_map_current = snapshot + .and_then(|value| value.pending_map_current) + .is_some_and(|value| value); + let orphan_warm_writers_current = snapshot + .map(|value| value.orphan_warm_writers_current) + .unwrap_or(0); + let replacement_preparing_current = snapshot + .map(|value| value.replacement_preparing_current) + .unwrap_or(0); + let replacement_retiring_current = snapshot + .map(|value| value.replacement_retiring_current) + .unwrap_or(0); + + let _ = writeln!( + out, + "# HELP telemt_me_hardswap_pending Whether an ME hardswap generation is pending" + ); + let _ = writeln!(out, "# TYPE telemt_me_hardswap_pending gauge"); + let _ = writeln!(out, "telemt_me_hardswap_pending {}", usize::from(pending)); + let _ = writeln!( + out, + "# HELP telemt_me_hardswap_pending_age_seconds Age of the pending ME hardswap generation" + ); + let _ = writeln!(out, "# TYPE telemt_me_hardswap_pending_age_seconds gauge"); + let _ = writeln!( + out, + "telemt_me_hardswap_pending_age_seconds {pending_age_secs}" + ); + let _ = writeln!( + out, + "# HELP telemt_me_hardswap_pending_writers_current Authoritative warm writers in the pending generation" + ); + let _ = writeln!( + out, + "# TYPE telemt_me_hardswap_pending_writers_current gauge" + ); + let _ = writeln!( + out, + "telemt_me_hardswap_pending_writers_current {pending_writers_current}" + ); + let _ = writeln!( + out, + "# HELP telemt_me_hardswap_pending_writer_deficit Writers missing from the pending generation floor" + ); + let _ = writeln!( + out, + "# TYPE telemt_me_hardswap_pending_writer_deficit gauge" + ); + let _ = writeln!( + out, + "telemt_me_hardswap_pending_writer_deficit {pending_writer_deficit}" + ); + let _ = writeln!( + out, + "# HELP telemt_me_hardswap_pending_missing_dc_groups Desired DC groups missing pending-generation coverage" + ); + let _ = writeln!( + out, + "# TYPE telemt_me_hardswap_pending_missing_dc_groups gauge" + ); + let _ = writeln!( + out, + "telemt_me_hardswap_pending_missing_dc_groups {pending_missing_dc_groups}" + ); + let _ = writeln!( + out, + "# HELP telemt_me_hardswap_pending_map_current Whether the pending generation targets the current endpoint map" + ); + let _ = writeln!(out, "# TYPE telemt_me_hardswap_pending_map_current gauge"); + let _ = writeln!( + out, + "telemt_me_hardswap_pending_map_current {}", + usize::from(pending_map_current) + ); + let _ = writeln!( + out, + "# HELP telemt_me_hardswap_orphan_warm_writers_current Warm writers not owned by the pending hardswap generation" + ); + let _ = writeln!( + out, + "# TYPE telemt_me_hardswap_orphan_warm_writers_current gauge" + ); + let _ = writeln!( + out, + "telemt_me_hardswap_orphan_warm_writers_current {orphan_warm_writers_current}" + ); + let _ = writeln!( + out, + "# HELP telemt_me_writer_replacement_current ME writer replacements by transaction phase" + ); + let _ = writeln!(out, "# TYPE telemt_me_writer_replacement_current gauge"); + let _ = writeln!( + out, + "telemt_me_writer_replacement_current{{state=\"preparing\"}} {replacement_preparing_current}" + ); + let _ = writeln!( + out, + "telemt_me_writer_replacement_current{{state=\"retiring\"}} {replacement_retiring_current}" + ); +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn renders_bounded_hardswap_and_replacement_state() { + let snapshot = MeApiHardswapSnapshot { + pending: true, + pending_age_secs: Some(42), + pending_writers_current: 3, + pending_writer_deficit: 4, + pending_missing_dc_groups: 2, + pending_map_current: Some(true), + orphan_warm_writers_current: 1, + replacement_preparing_current: 5, + replacement_retiring_current: 6, + }; + let mut out = String::new(); + + render(&mut out, Some(&snapshot), true); + + assert!(out.contains("telemt_me_hardswap_pending 1")); + assert!(out.contains("telemt_me_hardswap_pending_age_seconds 42")); + assert!(out.contains("telemt_me_hardswap_pending_writer_deficit 4")); + assert!(out.contains( + "telemt_me_writer_replacement_current{state=\"preparing\"} 5" + )); + assert!(out.contains( + "telemt_me_writer_replacement_current{state=\"retiring\"} 6" + )); + } +} diff --git a/src/metrics/tests.rs b/src/metrics/tests.rs index 04a0293..f157e3b 100644 --- a/src/metrics/tests.rs +++ b/src/metrics/tests.rs @@ -87,6 +87,7 @@ async fn test_render_metrics_format() { None, &TlsFullCertBudget::new(), &test_web_publication(), + None, ) .await; @@ -223,6 +224,7 @@ async fn test_render_tls_front_profile_health() { Some(&cache), &TlsFullCertBudget::new(), &test_web_publication(), + None, ) .await; @@ -293,6 +295,7 @@ async fn process_tls_budget_metrics_survive_a_generation_without_tls_cache() { None, budget.as_ref(), &test_web_publication(), + None, ) .await; @@ -313,6 +316,7 @@ async fn test_render_empty_stats() { None, &TlsFullCertBudget::new(), &test_web_publication(), + None, ) .await; assert!(output.contains("telemt_connections_total 0")); @@ -346,6 +350,7 @@ async fn test_render_uses_global_each_unique_ip_limit() { None, &TlsFullCertBudget::new(), &test_web_publication(), + None, ) .await; @@ -367,6 +372,7 @@ async fn test_render_has_type_annotations() { None, &TlsFullCertBudget::new(), &test_web_publication(), + None, ) .await; assert!(output.contains("# TYPE telemt_uptime_seconds gauge")); diff --git a/src/transport/middle_proxy/health.rs b/src/transport/middle_proxy/health.rs index f51b7f5..40c8743 100644 --- a/src/transport/middle_proxy/health.rs +++ b/src/transport/middle_proxy/health.rs @@ -16,9 +16,13 @@ use crate::crypto::SecureRandom; use crate::network::IpFamily; use super::MePool; -use super::pool::MeFamilyRuntimeState; +use super::pool::{ + MeFamilyRuntimeState, RefillTargetKey, WriterContour, WriterOpenIntent, WriterRole, +}; +use super::pool_writer::WriterReplacementPurpose; -const JITTER_FRAC_NUM: u64 = 2; // jitter up to 50% of backoff +// Bounds reconnect jitter to at most half of the current backoff. +const JITTER_FRAC_NUM: u64 = 2; #[allow(dead_code)] const MAX_CONCURRENT_PER_DC_DEFAULT: usize = 1; const SHADOW_ROTATE_RETRY_SECS: u64 = 30; diff --git a/src/transport/middle_proxy/health/drain.rs b/src/transport/middle_proxy/health/drain.rs index 336517b..5c9e629 100644 --- a/src/transport/middle_proxy/health/drain.rs +++ b/src/transport/middle_proxy/health/drain.rs @@ -31,7 +31,7 @@ pub(in crate::transport::middle_proxy) async fn reap_draining_writers( let mut force_close_writer_ids = Vec::::new(); let writers = pool.writers.read().await; for writer in writers.iter() { - if !writer.draining.load(std::sync::atomic::Ordering::Relaxed) { + if !writer.draining.load(std::sync::atomic::Ordering::Acquire) { continue; } if activity diff --git a/src/transport/middle_proxy/health/family.rs b/src/transport/middle_proxy/health/family.rs index c454675..b03677c 100644 --- a/src/transport/middle_proxy/health/family.rs +++ b/src/transport/middle_proxy/health/family.rs @@ -47,6 +47,7 @@ pub(super) async fn check_family( 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() @@ -62,6 +63,15 @@ pub(super) async fn check_family( ) { continue; } + if writer.generation != active_generation { + continue; + } + if !dc_endpoints + .get(&writer.writer_dc) + .is_some_and(|endpoints| endpoints.contains(&writer.addr)) + { + continue; + } let key = (writer.writer_dc, writer.addr); *live_addr_counts.entry(key).or_insert(0) += 1; live_writer_ids_by_addr @@ -93,6 +103,13 @@ pub(super) async fn check_family( 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); @@ -162,22 +179,23 @@ pub(super) async fn check_family( ); } + 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_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; maybe_rotate_single_endpoint_shadow( pool, rng, @@ -230,9 +248,11 @@ pub(super) async fn check_family( continue; } if pool - .has_refill_inflight_for_dc_key(crate::transport::middle_proxy::pool::RefillDcKey { + .has_refill_inflight_for_target(RefillTargetKey { dc, family, + generation: pool.current_generation(), + contour: WriterContour::Active, }) .await { @@ -254,6 +274,7 @@ pub(super) async fn check_family( 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; @@ -261,6 +282,8 @@ pub(super) async fn check_family( 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 { @@ -270,6 +293,8 @@ pub(super) async fn check_family( 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(), @@ -283,8 +308,6 @@ pub(super) async fn check_family( continue; } - let base_req = pool_for_reconnect - .required_writers_for_dc_with_floor_mode(endpoints_for_dc.len(), false); if alive + restored >= base_req { pool_for_reconnect .stats @@ -304,12 +327,20 @@ pub(super) async fn check_family( } } 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( + 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; diff --git a/src/transport/middle_proxy/health/floor_plan.rs b/src/transport/middle_proxy/health/floor_plan.rs index 0ad16eb..01b41b3 100644 --- a/src/transport/middle_proxy/health/floor_plan.rs +++ b/src/transport/middle_proxy/health/floor_plan.rs @@ -22,6 +22,7 @@ pub(super) async fn live_active_writers_for_dc_family( dc: i32, family: IpFamily, ) -> usize { + let active_generation = pool.current_generation(); let writers = pool.writers.read().await; writers .iter() @@ -32,6 +33,9 @@ pub(super) async fn live_active_writers_for_dc_family( if writer.writer_dc != dc { return false; } + if writer.generation != active_generation { + return false; + } if !matches!( crate::transport::middle_proxy::pool::WriterContour::from_u8( writer.contour.load(std::sync::atomic::Ordering::Relaxed), diff --git a/src/transport/middle_proxy/health/idle_refresh.rs b/src/transport/middle_proxy/health/idle_refresh.rs index 3cc26dc..e73b5be 100644 --- a/src/transport/middle_proxy/health/idle_refresh.rs +++ b/src/transport/middle_proxy/health/idle_refresh.rs @@ -6,16 +6,43 @@ pub(super) async fn maybe_swap_idle_writer_for_cap( dc: i32, family: IpFamily, endpoints: &[SocketAddr], + required: usize, + floor_targets_by_dc: &HashMap, live_writer_ids_by_addr: &HashMap<(i32, SocketAddr), Vec>, writer_idle_since: &HashMap, bound_clients_by_writer: &HashMap, ) -> bool { + let Some(replacement_endpoint) = endpoints + .iter() + .min_by_key(|endpoint| { + live_writer_ids_by_addr + .get(&(dc, **endpoint)) + .map_or(0, Vec::len) + }) + .copied() + else { + return false; + }; + + let mut alive_by_dc = HashMap::::new(); + for ((writer_dc, endpoint), writer_ids) in live_writer_ids_by_addr { + if endpoint.is_ipv4() == matches!(family, IpFamily::V4) { + *alive_by_dc.entry(*writer_dc).or_insert(0) += writer_ids.len(); + } + } + let now_epoch_secs = MePool::now_epoch_secs(); - let mut candidate: Option<(u64, SocketAddr, u64)> = None; - for endpoint in endpoints { - let Some(writer_ids) = live_writer_ids_by_addr.get(&(dc, *endpoint)) else { + let mut candidates = Vec::<(u64, SocketAddr, i32, u64, usize)>::new(); + for ((writer_dc, endpoint), writer_ids) in live_writer_ids_by_addr { + if endpoint.is_ipv4() != matches!(family, IpFamily::V4) { + continue; + } + let Some(donor_floor) = floor_targets_by_dc.get(writer_dc).copied() else { continue; }; + if alive_by_dc.get(writer_dc).copied().unwrap_or(0) <= donor_floor { + continue; + } for writer_id in writer_ids { if bound_clients_by_writer.get(writer_id).copied().unwrap_or(0) > 0 { continue; @@ -23,67 +50,96 @@ pub(super) async fn maybe_swap_idle_writer_for_cap( let Some(idle_since_epoch_secs) = writer_idle_since.get(writer_id).copied() else { continue; }; - let idle_age_secs = now_epoch_secs.saturating_sub(idle_since_epoch_secs); - if candidate - .as_ref() - .map(|(_, _, age)| idle_age_secs > *age) - .unwrap_or(true) - { - candidate = Some((*writer_id, *endpoint, idle_age_secs)); + candidates.push(( + *writer_id, + *endpoint, + *writer_dc, + now_epoch_secs.saturating_sub(idle_since_epoch_secs), + donor_floor, + )); + } + } + candidates.sort_unstable_by(|left, right| right.3.cmp(&left.3)); + + for (old_writer_id, donor_endpoint, donor_dc, idle_age_secs, donor_floor) in candidates { + let expected_role = { + let writers = pool.writers.read().await; + writers + .iter() + .find(|writer| writer.id == old_writer_id) + .map(WriterRole::from_writer) + }; + let Some(expected_role) = expected_role else { + continue; + }; + if expected_role.dc != donor_dc + || expected_role.family != family + || expected_role.contour != WriterContour::Active + { + continue; + } + let Some(mut reservation) = pool + .registry + .try_reserve_writer_replacement(old_writer_id) + .await + else { + continue; + }; + let replace = pool.replace_writer_with_generation_contour_for_dc( + replacement_endpoint, + rng.as_ref(), + pool.current_generation(), + WriterContour::Active, + dc, + expected_role, + WriterReplacementPurpose::FloorRebalance { + donor_floor, + receiver_floor: required, + }, + &mut reservation, + ); + match tokio::time::timeout(pool.reconnect_runtime.me_one_timeout, replace).await { + Ok(Ok(())) => { + info!( + dc = %dc, + ?family, + %replacement_endpoint, + donor_dc = %donor_dc, + %donor_endpoint, + old_writer_id, + idle_age_secs, + "Adaptive floor cap rebalance committed" + ); + return true; + } + Ok(Err(error)) => { + debug!( + dc = %dc, + ?family, + %replacement_endpoint, + donor_dc = %donor_dc, + old_writer_id, + idle_age_secs, + %error, + "Adaptive floor cap rebalance failed" + ); + return false; + } + Err(_) => { + debug!( + dc = %dc, + ?family, + %replacement_endpoint, + donor_dc = %donor_dc, + old_writer_id, + idle_age_secs, + "Adaptive floor cap rebalance timed out" + ); + return false; } } } - - let Some((old_writer_id, endpoint, idle_age_secs)) = candidate else { - return false; - }; - - let connected = match tokio::time::timeout( - pool.reconnect_runtime.me_one_timeout, - pool.connect_one_for_dc(endpoint, dc, rng.as_ref()), - ) - .await - { - Ok(Ok(())) => true, - Ok(Err(error)) => { - debug!( - dc = %dc, - ?family, - %endpoint, - old_writer_id, - idle_age_secs, - %error, - "Adaptive floor cap swap connect failed" - ); - false - } - Err(_) => { - debug!( - dc = %dc, - ?family, - %endpoint, - old_writer_id, - idle_age_secs, - "Adaptive floor cap swap connect timed out" - ); - false - } - }; - if !connected { - return false; - } - - pool.mark_writer_draining_with_timeout(old_writer_id, pool.force_close_timeout(), false) - .await; - info!( - dc = %dc, - ?family, - %endpoint, - old_writer_id, - idle_age_secs, - "Adaptive floor cap swap: idle writer rotated" - ); - true + false } pub(super) async fn maybe_refresh_idle_writer_for_dc( @@ -100,10 +156,6 @@ pub(super) async fn maybe_refresh_idle_writer_for_dc( bound_clients_by_writer: &HashMap, idle_refresh_next_attempt: &mut HashMap<(i32, IpFamily), Instant>, ) { - if alive < required { - return; - } - let now = Instant::now(); if let Some(next) = idle_refresh_next_attempt.get(&key) && now < *next @@ -143,13 +195,46 @@ pub(super) async fn maybe_refresh_idle_writer_for_dc( let Some((old_writer_id, endpoint, idle_age_secs, threshold_secs)) = candidate else { return; }; - - let rotate_ok = match tokio::time::timeout( - pool.reconnect_runtime.me_one_timeout, - pool.connect_one_for_dc(endpoint, dc, rng.as_ref()), - ) - .await + let expected_role = { + let writers = pool.writers.read().await; + writers + .iter() + .find(|writer| writer.id == old_writer_id) + .map(WriterRole::from_writer) + }; + let Some(expected_role) = expected_role else { + return; + }; + if expected_role.dc != dc + || expected_role.family != family + || expected_role.contour != WriterContour::Active { + return; + } + let Some(mut reservation) = pool + .registry + .try_reserve_writer_replacement(old_writer_id) + .await + else { + return; + }; + let generation = pool.current_generation(); + let purpose = if expected_role.generation == generation { + WriterReplacementPurpose::IdleRefresh + } else { + WriterReplacementPurpose::GenerationConvergence + }; + let replace = pool.replace_writer_with_generation_contour_for_dc( + endpoint, + rng.as_ref(), + generation, + WriterContour::Active, + dc, + expected_role, + purpose, + &mut reservation, + ); + let rotate_ok = match tokio::time::timeout(pool.reconnect_runtime.me_one_timeout, replace).await { Ok(Ok(())) => true, Ok(Err(error)) => { debug!( @@ -160,7 +245,7 @@ pub(super) async fn maybe_refresh_idle_writer_for_dc( idle_age_secs, threshold_secs, %error, - "Idle writer pre-refresh connect failed" + "Idle writer pre-refresh replacement failed" ); false } @@ -172,7 +257,7 @@ pub(super) async fn maybe_refresh_idle_writer_for_dc( old_writer_id, idle_age_secs, threshold_secs, - "Idle writer pre-refresh connect timed out" + "Idle writer pre-refresh replacement timed out" ); false } @@ -183,8 +268,6 @@ pub(super) async fn maybe_refresh_idle_writer_for_dc( return; } - pool.mark_writer_draining_with_timeout(old_writer_id, pool.force_close_timeout(), false) - .await; idle_refresh_next_attempt.insert( key, now + Duration::from_secs(IDLE_REFRESH_SUCCESS_GUARD_SECS), @@ -198,6 +281,7 @@ pub(super) async fn maybe_refresh_idle_writer_for_dc( threshold_secs, alive, required, + generation_convergence = expected_role.generation != generation, "Idle writer refreshed before upstream idle timeout" ); } diff --git a/src/transport/middle_proxy/health/recovery.rs b/src/transport/middle_proxy/health/recovery.rs index cfeee67..709ca0b 100644 --- a/src/transport/middle_proxy/health/recovery.rs +++ b/src/transport/middle_proxy/health/recovery.rs @@ -209,10 +209,38 @@ pub(super) async fn maybe_rotate_single_endpoint_shadow( ); return; }; + let expected_role = { + let writers = pool.writers.read().await; + writers + .iter() + .find(|writer| writer.id == old_writer_id) + .map(WriterRole::from_writer) + }; + let Some(expected_role) = expected_role else { + shadow_rotate_deadline.insert(key, now + Duration::from_secs(SHADOW_ROTATE_RETRY_SECS)); + return; + }; + let Some(mut reservation) = pool + .registry + .try_reserve_writer_replacement(old_writer_id) + .await + else { + shadow_rotate_deadline.insert(key, now + Duration::from_secs(SHADOW_ROTATE_RETRY_SECS)); + return; + }; let rotate_ok = match tokio::time::timeout( pool.reconnect_runtime.me_one_timeout, - pool.connect_one_for_dc(endpoint, dc, rng.as_ref()), + pool.replace_writer_with_generation_contour_for_dc( + endpoint, + rng.as_ref(), + pool.current_generation(), + WriterContour::Active, + dc, + expected_role, + WriterReplacementPurpose::ShadowRotation, + &mut reservation, + ), ) .await { @@ -246,8 +274,6 @@ pub(super) async fn maybe_rotate_single_endpoint_shadow( return; } - pool.mark_writer_draining_with_timeout(old_writer_id, pool.force_close_timeout(), false) - .await; pool.stats .increment_me_single_endpoint_shadow_rotate_total(); shadow_rotate_deadline.insert(key, now + interval); diff --git a/src/transport/middle_proxy/health/tests.rs b/src/transport/middle_proxy/health/tests.rs index 0e612ef..963d664 100644 --- a/src/transport/middle_proxy/health/tests.rs +++ b/src/transport/middle_proxy/health/tests.rs @@ -7,7 +7,7 @@ use std::time::{Duration, Instant}; use tokio::sync::mpsc; use tokio_util::sync::CancellationToken; -use super::{ScheduledReconnects, reap_draining_writers}; +use super::{ScheduledReconnects, maybe_refresh_idle_writer_for_dc, reap_draining_writers}; use crate::config::{GeneralConfig, MeRouteNoWriterMode, MeSocksKdfPolicy, MeWriterPickMode}; use crate::crypto::SecureRandom; use crate::network::IpFamily; @@ -226,6 +226,79 @@ async fn insert_live_writer(pool: &Arc, writer_id: u64, writer_dc: i32) pool.conn_count.fetch_add(1, Ordering::Relaxed); } +async fn insert_active_writer_at( + pool: &Arc, + writer_id: u64, + writer_dc: i32, + addr: SocketAddr, +) -> MeWriter { + let (tx, _writer_rx) = mpsc::channel::(8); + let byte_budget = pool.new_writer_byte_budget(); + let writer = MeWriter { + id: writer_id, + addr, + source_ip: addr.ip(), + writer_dc, + generation: pool.current_generation(), + contour: Arc::new(AtomicU8::new(WriterContour::Active.as_u8())), + created_at: Instant::now(), + tx: tx.clone(), + byte_budget: byte_budget.clone(), + cancel: CancellationToken::new(), + degraded: Arc::new(AtomicBool::new(false)), + rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)), + draining: Arc::new(AtomicBool::new(false)), + draining_started_at_epoch_secs: Arc::new(AtomicU64::new(0)), + drain_deadline_epoch_secs: Arc::new(AtomicU64::new(0)), + allow_drain_fallback: Arc::new(AtomicBool::new(false)), + }; + pool.writers.write().await.push(writer.clone()); + pool.registry + .register_writer(writer_id, tx, byte_budget) + .await; + pool.conn_count.fetch_add(1, Ordering::Relaxed); + writer +} + +#[tokio::test] +async fn under_floor_idle_writer_still_enters_transactional_refresh() { + let pool = make_pool(128).await; + let listener = std::net::TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).unwrap(); + let endpoint = listener.local_addr().unwrap(); + drop(listener); + let writer_id = 7001; + let writer = insert_active_writer_at(&pool, writer_id, 2, endpoint).await; + let key = (2, IpFamily::V4); + let live_writer_ids_by_addr = HashMap::from([((2, endpoint), vec![writer_id])]); + let writer_idle_since = HashMap::from([( + writer_id, + MePool::now_epoch_secs().saturating_sub(60), + )]); + let bound_clients_by_writer = HashMap::from([(writer_id, 0)]); + let mut next_attempt = HashMap::new(); + let rng = Arc::new(SecureRandom::new()); + + maybe_refresh_idle_writer_for_dc( + &pool, + &rng, + key, + 2, + IpFamily::V4, + &[endpoint], + 1, + 10, + &live_writer_ids_by_addr, + &writer_idle_since, + &bound_clients_by_writer, + &mut next_attempt, + ) + .await; + + assert!(next_attempt.contains_key(&key)); + assert!(!writer.draining.load(Ordering::Acquire)); + assert_eq!(pool.registry.writer_replacement_counts(), (0, 0)); +} + #[tokio::test] async fn reap_draining_writers_force_closes_oldest_over_threshold() { let pool = make_pool(2).await; diff --git a/src/transport/middle_proxy/health/zombie_watchdog.rs b/src/transport/middle_proxy/health/zombie_watchdog.rs index 074cad4..62ccd6c 100644 --- a/src/transport/middle_proxy/health/zombie_watchdog.rs +++ b/src/transport/middle_proxy/health/zombie_watchdog.rs @@ -35,7 +35,7 @@ pub async fn me_zombie_writer_watchdog(pool: Arc) { continue; }; ws.iter() - .filter(|w| w.draining.load(std::sync::atomic::Ordering::Relaxed)) + .filter(|w| w.draining.load(std::sync::atomic::Ordering::Acquire)) .filter_map(|w| { let deadline = w .drain_deadline_epoch_secs diff --git a/src/transport/middle_proxy/mod.rs b/src/transport/middle_proxy/mod.rs index 8adc45c..41ea1a2 100644 --- a/src/transport/middle_proxy/mod.rs +++ b/src/transport/middle_proxy/mod.rs @@ -64,6 +64,7 @@ pub use ping::{ MePingFamily, MePingReport, MePingSample, format_me_route, format_sample_line, run_me_ping, }; pub use pool::MePool; +pub(crate) use pool_status::MeApiHardswapSnapshot; #[allow(unused_imports)] pub use pool_nat::{detect_public_ip, stun_probe}; pub(crate) use registry::ConnLease; diff --git a/src/transport/middle_proxy/pool.rs b/src/transport/middle_proxy/pool.rs index cd4451a..093de8c 100644 --- a/src/transport/middle_proxy/pool.rs +++ b/src/transport/middle_proxy/pool.rs @@ -29,9 +29,25 @@ use super::pool_lifecycle::MePoolLifecycle; const ME_FORCE_CLOSE_SAFETY_FALLBACK_SECS: u64 = 300; #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] -pub(super) struct RefillDcKey { +/// Exact lifecycle role used to coalesce refill work without cross-generation drift. +pub(super) struct RefillTargetKey { + /// Telegram DC owning the writer. pub dc: i32, + /// Address family of the writer endpoint. pub family: IpFamily, + /// Generation that retains publication authority. + pub generation: u64, + /// Lifecycle contour that the replacement must preserve. + pub contour: WriterContour, +} + +#[derive(Debug, Clone, Copy, Default)] +/// Bounded queued-loss state for one exact refill target. +pub(super) struct RefillTargetState { + /// Additional lost writers waiting behind the active refill producer. + pub(super) pending_count: usize, + /// Most recently lost endpoint, used as the next same-endpoint preference. + pub(super) next_addr: Option, } #[derive(Clone)] @@ -121,6 +137,13 @@ impl DerefMut for WritersWriteGuard<'_> { } } +impl WritersWriteGuard<'_> { + /// Publishes the current vector while retaining exclusive mutation ownership. + pub(super) fn publish_current(&self) { + self.state.store_guarded(self.writers.clone()); + } +} + impl Drop for WritersWriteGuard<'_> { fn drop(&mut self) { let writers = std::mem::take(&mut self.writers); @@ -128,7 +151,7 @@ impl Drop for WritersWriteGuard<'_> { } } -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] #[repr(u8)] pub(super) enum WriterContour { Warm = 0, @@ -136,18 +159,6 @@ pub(super) enum WriterContour { Draining = 2, } -pub(super) struct WriterOpenReservation<'a> { - counter: Option<&'a AtomicUsize>, -} - -impl Drop for WriterOpenReservation<'_> { - fn drop(&mut self) { - if let Some(counter) = self.counter { - counter.fetch_sub(1, Ordering::AcqRel); - } - } -} - impl WriterContour { pub(super) fn as_u8(self) -> u8 { self as u8 @@ -471,8 +482,9 @@ pub struct MePool { pub(super) next_writer_id: AtomicU64, pub(super) writer_connect_active_reserved: AtomicUsize, pub(super) writer_connect_warm_reserved: AtomicUsize, + pub(super) writer_replacement_open_reserved: AtomicUsize, pub(super) rtt_stats: Arc>>, - pub(super) refill_states: Arc>>>, + pub(super) refill_states: Arc>>, pub(super) refill_running: AtomicUsize, pub(super) refill_pending: AtomicUsize, pub(super) conn_count: AtomicUsize, @@ -508,5 +520,6 @@ mod transport_policy; mod selection_policy; // Bounded writer-open admission and coverage accounting. mod writer_admission; +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 dd76e56..a14a086 100644 --- a/src/transport/middle_proxy/pool/construction.rs +++ b/src/transport/middle_proxy/pool/construction.rs @@ -398,6 +398,7 @@ impl MePool { 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), 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/writer_admission.rs b/src/transport/middle_proxy/pool/writer_admission.rs index 49b8c01..0041c3d 100644 --- a/src/transport/middle_proxy/pool/writer_admission.rs +++ b/src/transport/middle_proxy/pool/writer_admission.rs @@ -1,42 +1,110 @@ use super::*; +const WRITER_REPLACEMENT_OPEN_LIMIT_MAX: usize = 128; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +/// Immutable identity of a writer lifecycle role at replacement reservation time. +pub(in crate::transport::middle_proxy) struct WriterRole { + /// Telegram DC owning the writer. + pub(in crate::transport::middle_proxy) dc: i32, + /// Address family of the writer endpoint. + pub(in crate::transport::middle_proxy) family: IpFamily, + /// Pool generation owning the writer. + pub(in crate::transport::middle_proxy) generation: u64, + /// Lifecycle contour assigned to the writer. + pub(in crate::transport::middle_proxy) contour: WriterContour, +} + +impl WriterRole { + /// Captures the current role of an installed writer. + pub(in crate::transport::middle_proxy) fn from_writer(writer: &MeWriter) -> Self { + Self { + dc: writer.writer_dc, + family: if writer.addr.is_ipv4() { + IpFamily::V4 + } else { + IpFamily::V6 + }, + generation: writer.generation, + contour: WriterContour::from_u8(writer.contour.load(Ordering::Acquire)), + } + } + + /// Revalidates that an installed writer still has this exact role. + pub(in crate::transport::middle_proxy) fn matches(self, writer: &MeWriter) -> bool { + self == Self::from_writer(writer) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +/// Capacity policy applied while opening a writer. +pub(in crate::transport::middle_proxy) enum WriterOpenIntent { + /// Ordinary pool growth constrained by the configured contour cap. + Normal, + /// Required active coverage allowed to exceed an undersized configured cap temporarily. + Coverage, + /// Replacement-before-drain capacity owned by an existing victim reservation. + Replacement, +} + +/// RAII ownership of one bounded in-flight writer open. +pub(in crate::transport::middle_proxy) struct WriterOpenReservation<'a> { + counter: Option<&'a AtomicUsize>, +} + +impl Drop for WriterOpenReservation<'_> { + fn drop(&mut self) { + if let Some(counter) = self.counter { + counter.fetch_sub(1, Ordering::AcqRel); + } + } +} + impl MePool { + /// Computes the authoritative active-writer floor across enabled families and DCs. pub(in crate::transport::middle_proxy) async fn active_coverage_required_total(&self) -> usize { let now_epoch_secs = Self::now_epoch_secs(); - let mut endpoints_by_dc = HashMap::>::new(); + let mut required_total = 0usize; if self.family_enabled_for_drain_coverage(IpFamily::V4, now_epoch_secs) { let map = self.proxy_map_v4.read().await; - for (dc, addrs) in map.iter() { - let entry = endpoints_by_dc.entry(*dc).or_default(); + for addrs in map.values() { + let mut endpoints = HashSet::::new(); for (ip, port) in addrs.iter().copied() { - entry.insert(SocketAddr::new(ip, port)); + endpoints.insert(SocketAddr::new(ip, port)); } + required_total = required_total.saturating_add( + self.required_writers_for_dc_with_floor_mode(endpoints.len(), false), + ); } } if self.family_enabled_for_drain_coverage(IpFamily::V6, now_epoch_secs) { let map = self.proxy_map_v6.read().await; - for (dc, addrs) in map.iter() { - let entry = endpoints_by_dc.entry(*dc).or_default(); + for addrs in map.values() { + let mut endpoints = HashSet::::new(); for (ip, port) in addrs.iter().copied() { - entry.insert(SocketAddr::new(ip, port)); + endpoints.insert(SocketAddr::new(ip, port)); } + required_total = required_total.saturating_add( + self.required_writers_for_dc_with_floor_mode(endpoints.len(), false), + ); } } - endpoints_by_dc - .values() - .map(|endpoints| self.required_writers_for_dc_with_floor_mode(endpoints.len(), false)) - .sum() + required_total } + /// Reports whether one writer may be opened under the selected contour policy. pub(in crate::transport::middle_proxy) async fn can_open_writer_for_contour( &self, contour: WriterContour, - allow_coverage_override: bool, + intent: WriterOpenIntent, writer_dc: i32, ) -> bool { + if intent == WriterOpenIntent::Replacement { + return true; + } let (active_writers, warm_writers, _) = self.non_draining_writer_counts_by_contour().await; match contour { WriterContour::Active => { @@ -44,7 +112,7 @@ impl MePool { if active_writers < active_cap { return true; } - if !allow_coverage_override { + if intent != WriterOpenIntent::Coverage { return false; } @@ -64,12 +132,14 @@ impl MePool { if endpoints_len > 0 { let base_req = self.required_writers_for_dc_with_floor_mode(endpoints_len, false); + let active_generation = self.reinit.status.load().active_generation; let active_for_dc = { let ws = self.writers.read().await; ws.iter() .filter(|w| { !w.draining.load(std::sync::atomic::Ordering::Relaxed) && w.writer_dc == writer_dc + && w.generation == active_generation && matches!( WriterContour::from_u8( w.contour.load(std::sync::atomic::Ordering::Relaxed), @@ -92,10 +162,11 @@ impl MePool { } } + /// Reserves bounded transient capacity for a writer open attempt. pub(in crate::transport::middle_proxy) async fn reserve_writer_open( &self, contour: WriterContour, - allow_coverage_override: bool, + intent: WriterOpenIntent, writer_dc: i32, ) -> Option> { let counter = match contour { @@ -106,9 +177,52 @@ impl MePool { } }; + if intent == WriterOpenIntent::Replacement { + let configured_cap = match contour { + WriterContour::Active => self.adaptive_floor_active_cap_configured_total(), + WriterContour::Warm => self.adaptive_floor_warm_cap_configured_total(), + WriterContour::Draining => usize::MAX, + }; + let effective_cap = match contour { + WriterContour::Active => self + .floor_runtime + .me_adaptive_floor_active_cap_effective + .load(Ordering::Acquire) as usize, + WriterContour::Warm => self + .floor_runtime + .me_adaptive_floor_warm_cap_effective + .load(Ordering::Acquire) as usize, + WriterContour::Draining => usize::MAX, + }; + let replacement_limit = configured_cap + .max(effective_cap) + .max(1) + .min(WRITER_REPLACEMENT_OPEN_LIMIT_MAX); + loop { + let reserved = self.writer_replacement_open_reserved.load(Ordering::Acquire); + if reserved >= replacement_limit { + return None; + } + if self + .writer_replacement_open_reserved + .compare_exchange_weak( + reserved, + reserved + 1, + Ordering::AcqRel, + Ordering::Acquire, + ) + .is_ok() + { + return Some(WriterOpenReservation { + counter: Some(&self.writer_replacement_open_reserved), + }); + } + } + } + loop { if !self - .can_open_writer_for_contour(contour, allow_coverage_override, writer_dc) + .can_open_writer_for_contour(contour, intent, writer_dc) .await { return None; @@ -125,7 +239,7 @@ impl MePool { WriterContour::Warm => self.adaptive_floor_warm_cap_configured_total(), WriterContour::Draining => usize::MAX, }; - if contour == WriterContour::Active && allow_coverage_override { + if contour == WriterContour::Active && intent == WriterOpenIntent::Coverage { limit = limit .max(self.active_coverage_required_total().await) .saturating_add( @@ -150,6 +264,7 @@ impl MePool { } } + /// Resolves a DC writer floor for static or adaptive-idle operation. pub(in crate::transport::middle_proxy) fn required_writers_for_dc_with_floor_mode( &self, endpoint_count: usize, diff --git a/src/transport/middle_proxy/pool_config.rs b/src/transport/middle_proxy/pool_config.rs index 6e29918..73695b7 100644 --- a/src/transport/middle_proxy/pool_config.rs +++ b/src/transport/middle_proxy/pool_config.rs @@ -5,7 +5,8 @@ use std::time::Duration; use tracing::warn; -use super::pool::MePool; +use super::pool::{MePool, WriterContour, WriterRole}; +use super::pool_writer::WriterReplacementPurpose; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum SnapshotApplyOutcome { @@ -113,11 +114,33 @@ impl MePool { pub async fn reconnect_all(self: &Arc) { let ws = self.writers.read().await.clone(); for w in ws.iter() { - if let Ok(()) = self - .connect_one_for_dc(w.addr, w.writer_dc, self.rng.as_ref()) - .await + let role = WriterRole::from_writer(w); + if w.draining.load(std::sync::atomic::Ordering::Acquire) + || role.contour == WriterContour::Draining + { + continue; + } + let Some(mut reservation) = self + .registry + .try_reserve_writer_replacement_preserving_clients(w.id) + .await + else { + continue; + }; + if self + .replace_writer_with_generation_contour_for_dc( + w.addr, + self.rng.as_ref(), + role.generation, + role.contour, + role.dc, + role, + WriterReplacementPurpose::SecretRotation, + &mut reservation, + ) + .await + .is_ok() { - self.mark_writer_draining(w.id).await; tokio::time::sleep(Duration::from_secs(2)).await; } } diff --git a/src/transport/middle_proxy/pool_init.rs b/src/transport/middle_proxy/pool_init.rs index 0936246..eaa0476 100644 --- a/src/transport/middle_proxy/pool_init.rs +++ b/src/transport/middle_proxy/pool_init.rs @@ -9,7 +9,7 @@ use tracing::{debug, info, warn}; use crate::crypto::SecureRandom; use crate::error::{ProxyError, Result}; -use super::pool::MePool; +use super::pool::{MePool, WriterOpenIntent}; impl MePool { pub async fn init(self: &Arc, pool_size: usize, rng: &Arc) -> Result<()> { @@ -76,7 +76,7 @@ impl MePool { target_writers, rng_clone, connect_concurrency, - true, + WriterOpenIntent::Coverage, ) .await }); @@ -125,7 +125,7 @@ impl MePool { target_writers, rng_clone_local, connect_concurrency, - false, + WriterOpenIntent::Normal, ) .await }); @@ -162,7 +162,7 @@ impl MePool { target_writers: usize, rng: Arc, connect_concurrency: usize, - allow_coverage_override: bool, + intent: WriterOpenIntent, ) -> bool { if addrs.is_empty() { return false; @@ -204,7 +204,7 @@ impl MePool { rng_clone.as_ref(), generation, super::pool::WriterContour::Active, - allow_coverage_override, + intent, ) .await }); @@ -238,7 +238,9 @@ impl MePool { if !progress { let active_writers_current = self.active_contour_writer_count_total().await; let active_cap_configured = self.adaptive_floor_active_cap_configured_total(); - if !allow_coverage_override && active_writers_current >= active_cap_configured { + if intent != WriterOpenIntent::Coverage + && active_writers_current >= active_cap_configured + { info!( dc = %dc, alive = alive_after, diff --git a/src/transport/middle_proxy/pool_refill.rs b/src/transport/middle_proxy/pool_refill.rs index 30a7809..c6f1747 100644 --- a/src/transport/middle_proxy/pool_refill.rs +++ b/src/transport/middle_proxy/pool_refill.rs @@ -9,26 +9,42 @@ use tracing::{debug, info, warn}; use crate::crypto::SecureRandom; use crate::network::IpFamily; -use super::pool::{MePool, RefillDcKey, WriterContour}; +use super::pool::{ + MePool, RefillTargetKey, RefillTargetState, WriterContour, WriterOpenIntent, WriterRole, +}; const ME_FLAP_UPTIME_THRESHOLD_SECS: u64 = 20; const ME_FLAP_QUARANTINE_SECS: u64 = 25; const ME_FLAP_MIN_UPTIME_MILLIS: u64 = 500; const ME_REFILL_TOTAL_ATTEMPT_CAP: u32 = 20; +const ME_REFILL_PENDING_PER_TARGET_MAX: usize = 128; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum RefillOutcome { + Restored, + Failed, + Obsolete, +} struct RefillRunGuard { pool: Arc, - key: RefillDcKey, + key: RefillTargetKey, active: bool, } impl RefillRunGuard { fn next_or_finish(&mut self) -> Option { let mut states = self.pool.refill_states.lock(); - let next = states.get_mut(&self.key).and_then(Option::take); - if next.is_some() { + let next = states.get_mut(&self.key).and_then(|state| { + if state.pending_count == 0 { + return None; + } + state.pending_count -= 1; + state.next_addr + }); + if let Some(next) = next { self.pool.refill_pending.fetch_sub(1, Ordering::AcqRel); - return next; + return Some(next); } states.remove(&self.key); self.pool.refill_running.fetch_sub(1, Ordering::AcqRel); @@ -42,10 +58,10 @@ impl Drop for RefillRunGuard { if !self.active { return; } - if let Some(pending) = self.pool.refill_states.lock().remove(&self.key) - && pending.is_some() - { - self.pool.refill_pending.fetch_sub(1, Ordering::AcqRel); + if let Some(state) = self.pool.refill_states.lock().remove(&self.key) { + self.pool + .refill_pending + .fetch_sub(state.pending_count, Ordering::AcqRel); } self.pool.refill_running.fetch_sub(1, Ordering::AcqRel); } @@ -165,7 +181,8 @@ impl MePool { self.connectable_endpoints(endpoints).await } - pub(super) async fn has_refill_inflight_for_dc_key(&self, key: RefillDcKey) -> bool { + /// Reports whether the exact generation and contour already has a refill producer. + pub(super) async fn has_refill_inflight_for_target(&self, key: RefillTargetKey) -> bool { self.refill_states.lock().contains_key(&key) } @@ -181,7 +198,7 @@ impl MePool { rng, self.current_generation(), WriterContour::Active, - false, + WriterOpenIntent::Normal, ) .await } @@ -193,14 +210,14 @@ impl MePool { rng: &SecureRandom, generation: u64, contour: WriterContour, - allow_coverage_override: bool, + intent: WriterOpenIntent, ) -> bool { let mut candidates = self.connectable_endpoints(endpoints).await; if candidates.is_empty() { return false; } if candidates.len() > 1 { - let mut active_by_endpoint = HashMap::::new(); + let mut matching_by_endpoint = HashMap::::new(); let ws = self.writers.read().await; for writer in ws.iter() { if writer.draining.load(Ordering::Relaxed) { @@ -209,32 +226,31 @@ impl MePool { if writer.writer_dc != dc { continue; } - if !matches!( - super::pool::WriterContour::from_u8(writer.contour.load(Ordering::Relaxed),), - super::pool::WriterContour::Active - ) { + if writer.generation != generation + || WriterContour::from_u8(writer.contour.load(Ordering::Acquire)) != contour + { continue; } if candidates.contains(&writer.addr) { - *active_by_endpoint.entry(writer.addr).or_insert(0) += 1; + *matching_by_endpoint.entry(writer.addr).or_insert(0) += 1; } } drop(ws); candidates - .sort_by_key(|addr| (active_by_endpoint.get(addr).copied().unwrap_or(0), *addr)); + .sort_by_key(|addr| (matching_by_endpoint.get(addr).copied().unwrap_or(0), *addr)); } let start = (self.rr.fetch_add(1, Ordering::Relaxed) as usize) % candidates.len(); for offset in 0..candidates.len() { let idx = (start + offset) % candidates.len(); let addr = candidates[idx]; match self - .connect_one_with_generation_contour_for_dc_with_cap_policy( + .connect_one_with_generation_contour_for_dc_with_intent( addr, rng, generation, contour, dc, - allow_coverage_override, + intent, ) .await { @@ -245,43 +261,74 @@ impl MePool { false } - async fn endpoints_for_dc(&self, target_dc: i32) -> Vec { + async fn endpoints_for_refill_target(&self, target: RefillTargetKey) -> Vec { let now_epoch_secs = Self::now_epoch_secs(); - let mut endpoints = HashSet::::new(); - - if self.family_enabled_for_drain_coverage(IpFamily::V4, now_epoch_secs) { - let map = self.proxy_map_v4.read().await; - if let Some(addrs) = map.get(&target_dc) { - for (ip, port) in addrs { - endpoints.insert(SocketAddr::new(*ip, *port)); - } - } + if !self.family_enabled_for_drain_coverage(target.family, now_epoch_secs) { + return Vec::new(); } - - if self.family_enabled_for_drain_coverage(IpFamily::V6, now_epoch_secs) { - let map = self.proxy_map_v6.read().await; - if let Some(addrs) = map.get(&target_dc) { - for (ip, port) in addrs { - endpoints.insert(SocketAddr::new(*ip, *port)); - } - } - } - - let mut sorted: Vec = endpoints.into_iter().collect(); - sorted.sort_unstable(); - sorted + let map = match target.family { + IpFamily::V4 => self.proxy_map_v4.read().await, + IpFamily::V6 => self.proxy_map_v6.read().await, + }; + let mut endpoints = map + .get(&target.dc) + .into_iter() + .flatten() + .map(|(ip, port)| SocketAddr::new(*ip, *port)) + .collect::>(); + endpoints.sort_unstable(); + endpoints.dedup(); + endpoints } - async fn refill_writer_after_loss(self: &Arc, addr: SocketAddr, writer_dc: i32) -> bool { + fn refill_target_is_authoritative(&self, target: RefillTargetKey) -> bool { + if !self.family_enabled_for_drain_coverage(target.family, Self::now_epoch_secs()) { + return false; + } + let status = self.reinit.status.load(); + let role_is_authoritative = match target.contour { + WriterContour::Active => target.generation == status.active_generation, + WriterContour::Warm => status.pending_hardswap_generation != 0 + && target.generation == status.pending_hardswap_generation, + WriterContour::Draining => false, + }; + role_is_authoritative + && self + .preferred_endpoints_by_dc + .load() + .get(&target.dc) + .is_some_and(|endpoints| { + endpoints.iter().any(|endpoint| match target.family { + IpFamily::V4 => endpoint.is_ipv4(), + IpFamily::V6 => endpoint.is_ipv6(), + }) + }) + } + + async fn refill_writer_after_loss( + self: &Arc, + addr: SocketAddr, + target: RefillTargetKey, + ) -> RefillOutcome { + if !self.refill_target_is_authoritative(target) { + return RefillOutcome::Obsolete; + } + let open_intent = if target.contour == WriterContour::Active { + WriterOpenIntent::Coverage + } else { + WriterOpenIntent::Normal + }; let fast_retries = self.reconnect_runtime.me_reconnect_fast_retry_count.max(1); let mut total_attempts = 0u32; let same_endpoint_quarantined = self.is_endpoint_quarantined(addr).await; - let dc_endpoints = self.endpoints_for_dc(writer_dc).await; + let dc_endpoints = self.endpoints_for_refill_target(target).await; let single_endpoint_dc = dc_endpoints.len() == 1 && dc_endpoints[0] == addr; let bypass_quarantine_for_single_endpoint = single_endpoint_dc && self.single_endpoint_outage_disable_quarantine(); - if !same_endpoint_quarantined || bypass_quarantine_for_single_endpoint { + if dc_endpoints.contains(&addr) + && (!same_endpoint_quarantined || bypass_quarantine_for_single_endpoint) + { if same_endpoint_quarantined && bypass_quarantine_for_single_endpoint { debug!( %addr, @@ -289,13 +336,23 @@ impl MePool { ); } for attempt in 0..fast_retries { + if !self.refill_target_is_authoritative(target) { + return RefillOutcome::Obsolete; + } if total_attempts >= ME_REFILL_TOTAL_ATTEMPT_CAP { break; } total_attempts = total_attempts.saturating_add(1); self.stats.increment_me_reconnect_attempt(); match self - .connect_one_for_dc(addr, writer_dc, self.rng.as_ref()) + .connect_one_with_generation_contour_for_dc_with_intent( + addr, + self.rng.as_ref(), + target.generation, + target.contour, + target.dc, + open_intent, + ) .await { Ok(()) => { @@ -307,7 +364,7 @@ impl MePool { attempt = attempt + 1, "ME writer restored on the same endpoint" ); - return true; + return RefillOutcome::Restored; } Err(e) => { debug!( @@ -328,17 +385,27 @@ impl MePool { if dc_endpoints.is_empty() { self.stats.increment_me_refill_failed_total(); - return false; + return RefillOutcome::Failed; } for attempt in 0..fast_retries { + if !self.refill_target_is_authoritative(target) { + return RefillOutcome::Obsolete; + } if total_attempts >= ME_REFILL_TOTAL_ATTEMPT_CAP { break; } total_attempts = total_attempts.saturating_add(1); self.stats.increment_me_reconnect_attempt(); if self - .connect_endpoints_round_robin(writer_dc, &dc_endpoints, self.rng.as_ref()) + .connect_endpoints_round_robin_with_generation_contour( + target.dc, + &dc_endpoints, + self.rng.as_ref(), + target.generation, + target.contour, + open_intent, + ) .await { self.stats.increment_me_reconnect_success(); @@ -348,12 +415,12 @@ impl MePool { attempt = attempt + 1, "ME writer restored via DC fallback endpoint" ); - return true; + return RefillOutcome::Restored; } } self.stats.increment_me_refill_failed_total(); - false + RefillOutcome::Failed } pub(crate) fn trigger_immediate_refill_for_dc( @@ -361,35 +428,58 @@ impl MePool { addr: SocketAddr, writer_dc: i32, ) { + self.trigger_immediate_refill_for_role( + addr, + WriterRole { + dc: writer_dc, + family: if addr.is_ipv4() { + IpFamily::V4 + } else { + IpFamily::V6 + }, + generation: self.current_generation(), + contour: WriterContour::Active, + }, + ); + } + + /// Coalesces unexpected writer loss without changing its generation or contour role. + pub(super) fn trigger_immediate_refill_for_role( + self: &Arc, + addr: SocketAddr, + role: WriterRole, + ) { + let target = RefillTargetKey { + dc: role.dc, + family: role.family, + generation: role.generation, + contour: role.contour, + }; + if !self.refill_target_is_authoritative(target) { + return; + } let Some(registration) = self.lifecycle.try_register() else { return; }; - let dc_key = RefillDcKey { - dc: writer_dc, - family: if addr.is_ipv4() { - IpFamily::V4 - } else { - IpFamily::V6 - }, - }; { let mut states = self.refill_states.lock(); - if let Some(pending) = states.get_mut(&dc_key) { - if pending.is_none() { + if let Some(state) = states.get_mut(&target) { + if state.pending_count < ME_REFILL_PENDING_PER_TARGET_MAX { + state.pending_count += 1; + state.next_addr = Some(addr); self.refill_pending.fetch_add(1, Ordering::AcqRel); } - *pending = Some(addr); self.stats.increment_me_refill_skipped_inflight_total(); return; } - states.insert(dc_key, None); + states.insert(target, RefillTargetState::default()); self.refill_running.fetch_add(1, Ordering::AcqRel); } let pool = Arc::clone(self); let mut run_guard = RefillRunGuard { pool: Arc::clone(&pool), - key: dc_key, + key: target, active: true, }; self.lifecycle @@ -397,9 +487,24 @@ impl MePool { let mut current_addr = addr; loop { pool.stats.increment_me_refill_triggered_total(); - let restored = pool.refill_writer_after_loss(current_addr, writer_dc).await; - if !restored { - warn!(%current_addr, dc = writer_dc, "ME immediate refill failed"); + let outcome = pool.refill_writer_after_loss(current_addr, target).await; + if outcome == RefillOutcome::Failed { + warn!( + %current_addr, + dc = target.dc, + generation = target.generation, + contour = ?target.contour, + "ME immediate refill failed" + ); + } else if outcome == RefillOutcome::Obsolete { + debug!( + %current_addr, + dc = target.dc, + generation = target.generation, + contour = ?target.contour, + "ME immediate refill target is no longer authoritative" + ); + return; } let Some(next_addr) = run_guard.next_or_finish() else { diff --git a/src/transport/middle_proxy/pool_reinit.rs b/src/transport/middle_proxy/pool_reinit.rs index 46174cd..f0a3db5 100644 --- a/src/transport/middle_proxy/pool_reinit.rs +++ b/src/transport/middle_proxy/pool_reinit.rs @@ -11,11 +11,12 @@ use std::collections::hash_map::DefaultHasher; use tracing::{debug, info, warn}; use crate::crypto::SecureRandom; +use crate::config::MeBindStaleMode; use crate::network::IpFamily; use super::pool::{ MeDrainGateReason, MePool, ReinitAttemptState, ReinitCoordinatorState, ReinitCore, - ReinitPendingState, ReinitStatusSnapshot, WriterContour, + ReinitPendingState, ReinitStatusSnapshot, WriterContour, WriterOpenIntent, }; // Reinitialization admission, generation state, and coverage checks. @@ -51,6 +52,26 @@ struct ReinitReservation { pending_age_secs: u64, } +struct ReinitCommitOutcome { + coverage_ratio: f32, + missing_dc: Vec, + stale_writer_ids: Vec, + force_close_writer_ids: Vec, +} + +#[derive(Debug)] +enum ReinitCommitFailure { + Superseded, + Coverage { + coverage_ratio: f32, + missing_dc: Vec, + }, + Redundancy { + coverage_ratio: f32, + missing_dc: Vec, + }, +} + fn publish_reinit_state(reinit: &ReinitCore, state: &ReinitCoordinatorState) { let mut warm_generations = state .attempts @@ -58,6 +79,9 @@ fn publish_reinit_state(reinit: &ReinitCore, state: &ReinitCoordinatorState) { .filter(|attempt| attempt.hardswap && !attempt.committed) .map(|attempt| attempt.generation) .collect::>(); + if let Some(pending) = state.pending { + warm_generations.push(pending.generation); + } warm_generations.sort_unstable(); warm_generations.dedup(); let pending = state.pending; diff --git a/src/transport/middle_proxy/pool_reinit/coordination.rs b/src/transport/middle_proxy/pool_reinit/coordination.rs index 64e79d3..a81b672 100644 --- a/src/transport/middle_proxy/pool_reinit/coordination.rs +++ b/src/transport/middle_proxy/pool_reinit/coordination.rs @@ -1,7 +1,10 @@ use super::*; impl MePool { - pub(super) fn desired_map_hash(desired_by_dc: &HashMap>) -> u64 { + /// Hashes the sorted desired endpoint map for generation authority checks. + pub(in crate::transport::middle_proxy) fn desired_map_hash( + desired_by_dc: &HashMap>, + ) -> u64 { let mut hasher = DefaultHasher::new(); let mut dcs: Vec = desired_by_dc.keys().copied().collect(); dcs.sort_unstable(); @@ -19,6 +22,7 @@ impl MePool { hasher.finish() } + /// Reserves one generation attempt and publishes its pending ownership snapshot. pub(super) fn reserve_reinit_attempt( self: &Arc, hardswap: bool, @@ -85,8 +89,62 @@ impl MePool { } } - pub(super) fn commit_reinit_attempt(&self, attempt: &ReinitAttemptGuard) -> bool { + /// Revalidates coverage and commits generation ownership under the publication barrier. + pub(super) async fn commit_reinit_attempt( + &self, + attempt: &ReinitAttemptGuard, + desired_by_dc: &HashMap>, + min_ratio: f32, + ) -> std::result::Result { + let writers = self.writers.write().await; + let mut registry_registration = self.registry.prepare_writer_registration().await; let mut state = self.reinit.coordinator.lock(); + let Some(record) = state.attempts.get(&attempt.attempt_id).copied() else { + return Err(ReinitCommitFailure::Superseded); + }; + if record.generation != attempt.generation + || record.map_hash != state.desired_map_hash + || record.map_hash != attempt.map_hash + || (attempt.hardswap + && !state.pending.is_some_and(|pending| { + pending.generation == attempt.generation + && pending.map_hash == attempt.map_hash + })) + { + return Err(ReinitCommitFailure::Superseded); + } + + let authoritative_writer_addrs = writers + .iter() + .filter(|writer| !writer.draining.load(Ordering::Acquire)) + .filter(|writer| { + if attempt.hardswap { + writer.generation == attempt.generation + } else { + writer.generation == state.active_generation + && WriterContour::from_u8(writer.contour.load(Ordering::Acquire)) + == WriterContour::Active + } + }) + .map(|writer| (writer.writer_dc, writer.addr)) + .collect::>(); + let (coverage_ratio, missing_dc) = + Self::coverage_ratio(desired_by_dc, &authoritative_writer_addrs); + if coverage_ratio < min_ratio { + return Err(ReinitCommitFailure::Coverage { + coverage_ratio, + missing_dc, + }); + } + if attempt.hardswap + && !missing_dc.is_empty() + && self.bind_stale_mode() == MeBindStaleMode::Never + { + return Err(ReinitCommitFailure::Redundancy { + coverage_ratio, + missing_dc, + }); + } if !commit_reinit_state( &mut state, attempt.attempt_id, @@ -94,12 +152,12 @@ impl MePool { attempt.map_hash, attempt.hardswap, ) { - return false; + return Err(ReinitCommitFailure::Superseded); } + if attempt.hardswap { - let writers = self.writers.snapshot(); for writer in writers.iter() { - if !writer.draining.load(Ordering::Relaxed) + if !writer.draining.load(Ordering::Acquire) && writer.generation == attempt.generation { writer @@ -108,10 +166,59 @@ impl MePool { } } } + + let desired_addrs = desired_by_dc + .iter() + .flat_map(|(dc, endpoints)| endpoints.iter().copied().map(|addr| (*dc, addr))) + .collect::>(); + let missing_dc_set = missing_dc.iter().copied().collect::>(); + let mut stale_writer_ids = Vec::::new(); + let mut force_close_writer_ids = Vec::::new(); + for writer in writers.iter() { + if writer.draining.load(Ordering::Acquire) { + continue; + } + let stale = if attempt.hardswap { + writer.generation < attempt.generation + } else { + !desired_addrs.contains(&(writer.writer_dc, writer.addr)) + }; + if !stale { + continue; + } + + let preserve_fallback = attempt.hardswap + && missing_dc_set.contains(&writer.writer_dc); + if !preserve_fallback && attempt.hardswap { + registry_registration.retire(writer.id); + } + self.apply_writer_draining_state( + writer, + self.force_close_timeout(), + preserve_fallback || !attempt.hardswap, + ); + stale_writer_ids.push(writer.id); + if (attempt.hardswap && !preserve_fallback) + || (!attempt.hardswap && missing_dc.is_empty()) + { + force_close_writer_ids.push(writer.id); + } + } publish_reinit_state(self.reinit.as_ref(), &state); - true + drop(state); + drop(registry_registration); + drop(writers); + self.notify_writer_epoch(); + + Ok(ReinitCommitOutcome { + coverage_ratio, + missing_dc, + stale_writer_ids, + force_close_writer_ids, + }) } + /// Computes desired DC-group coverage and returns missing groups in stable order. pub(super) fn coverage_ratio( desired_by_dc: &HashMap>, active_writer_addrs: &HashSet<(i32, SocketAddr)>, @@ -146,6 +253,7 @@ impl MePool { (ratio, missing_dc) } + /// Restores at least one active writer for every enabled desired DC group. pub async fn reconcile_connections(self: &Arc, rng: &SecureRandom) { for family in self.family_order() { let map = self.proxy_map_for_family(family).await; @@ -175,7 +283,10 @@ impl MePool { } } - pub(super) async fn desired_dc_endpoints(&self) -> HashMap> { + /// Returns the currently authoritative endpoint set for drain and coverage decisions. + pub(in crate::transport::middle_proxy) async fn desired_dc_endpoints( + &self, + ) -> HashMap> { let now_epoch_secs = Self::now_epoch_secs(); let mut out: HashMap> = HashMap::new(); @@ -202,38 +313,63 @@ impl MePool { out } - pub(in crate::transport::middle_proxy) async fn has_non_draining_writer_per_desired_dc_group( - &self, - ) -> bool { - let desired_by_dc = self.desired_dc_endpoints().await; - let required_dcs: HashSet = desired_by_dc - .iter() - .filter_map(|(dc, endpoints)| { - if endpoints.is_empty() { - None - } else { - Some(*dc) - } - }) - .collect(); - if required_dcs.is_empty() { - return true; - } + /// Promotes authoritative warm writers and drains warm or active generation orphans. + pub(super) async fn reconcile_writer_generation_roles(&self) -> usize { + let writers = self.writers.write().await; + let mut registry_registration = self.registry.prepare_writer_registration().await; + let state = self.reinit.coordinator.lock(); + let active_generation = state.active_generation; + let pending_generation = state.pending.map(|pending| pending.generation); + let preferred = self.preferred_endpoints_by_dc.load(); + let now_epoch_secs = Self::now_epoch_secs(); + let mut changed = 0usize; - let ws = self.writers.read().await; - let mut covered_dcs = HashSet::::with_capacity(required_dcs.len()); - for writer in ws.iter() { - if writer.draining.load(Ordering::Relaxed) { + for writer in writers.iter() { + if writer.draining.load(Ordering::Acquire) { continue; } - if required_dcs.contains(&writer.writer_dc) { - covered_dcs.insert(writer.writer_dc); - if covered_dcs.len() == required_dcs.len() { - return true; - } + let contour = WriterContour::from_u8(writer.contour.load(Ordering::Acquire)); + let family = if writer.addr.is_ipv4() { + IpFamily::V4 + } else { + IpFamily::V6 + }; + let endpoint_is_current = self + .family_enabled_for_drain_coverage(family, now_epoch_secs) + && preferred + .get(&writer.writer_dc) + .is_some_and(|endpoints| endpoints.contains(&writer.addr)); + if contour == WriterContour::Warm + && writer.generation == active_generation + && endpoint_is_current + { + writer + .contour + .store(WriterContour::Active.as_u8(), Ordering::Release); + changed = changed.saturating_add(1); + continue; } + let authoritative_warm = contour == WriterContour::Warm + && pending_generation == Some(writer.generation) + && endpoint_is_current; + let stale_active = contour == WriterContour::Active + && writer.generation != active_generation; + if authoritative_warm || (contour == WriterContour::Active && !stale_active) { + continue; + } + + registry_registration.retire(writer.id); + self.apply_writer_draining_state(writer, self.force_close_timeout(), false); + changed = changed.saturating_add(1); } - false + drop(preferred); + drop(state); + drop(registry_registration); + drop(writers); + if changed > 0 { + self.notify_writer_epoch(); + } + changed } pub(super) fn hardswap_warmup_connect_delay_ms(&self) -> u64 { @@ -270,6 +406,7 @@ impl MePool { core.saturating_add(rand::rng().random_range(0..=jitter)) } + /// Counts non-draining writers owned by one generation and desired DC endpoint set. pub(super) async fn fresh_writer_count_for_dc_endpoints( &self, generation: u64, @@ -285,14 +422,21 @@ impl MePool { .count() } + /// Counts authoritative active writers for one desired DC endpoint set. pub(in crate::transport::middle_proxy) async fn active_writer_count_for_dc_endpoints( &self, dc: i32, endpoints: &HashSet, ) -> usize { + let generation = self.current_generation(); let ws = self.writers.read().await; ws.iter() .filter(|w| !w.draining.load(Ordering::Relaxed)) + .filter(|w| w.generation == generation) + .filter(|w| { + WriterContour::from_u8(w.contour.load(Ordering::Acquire)) + == WriterContour::Active + }) .filter(|w| w.writer_dc == dc) .filter(|w| endpoints.contains(&w.addr)) .count() diff --git a/src/transport/middle_proxy/pool_reinit/reconcile.rs b/src/transport/middle_proxy/pool_reinit/reconcile.rs index 5c900fe..b2b5ee4 100644 --- a/src/transport/middle_proxy/pool_reinit/reconcile.rs +++ b/src/transport/middle_proxy/pool_reinit/reconcile.rs @@ -56,7 +56,7 @@ impl MePool { rng, generation, WriterContour::Warm, - false, + WriterOpenIntent::Normal, ) .await; debug!( @@ -141,6 +141,15 @@ impl MePool { let attempt = reservation.attempt; let previous_generation = attempt.previous_generation; let generation = attempt.generation; + let reconciled_roles = self.reconcile_writer_generation_roles().await; + if reconciled_roles > 0 { + info!( + reconciled_roles, + active_generation = previous_generation, + pending_generation = generation, + "ME writer generation roles reconciled" + ); + } if reservation.pending_reused { self.stats.increment_me_hardswap_pending_reuse_total(); debug!( @@ -180,14 +189,10 @@ impl MePool { ); let (coverage_ratio, missing_dc) = Self::coverage_ratio(&desired_by_dc, &active_writer_addrs); - let mut route_quorum_ok = coverage_ratio >= min_ratio; - let mut redundancy_ok = missing_dc.is_empty(); - let mut redundancy_missing_dc = missing_dc.clone(); - let mut gate_coverage_ratio = coverage_ratio; if !hardswap && coverage_ratio < min_ratio { self.set_last_drain_gate( false, - redundancy_ok, + missing_dc.is_empty(), MeDrainGateReason::CoverageQuorum, now_epoch_secs, ); @@ -211,14 +216,10 @@ impl MePool { .collect(); let (fresh_coverage_ratio, fresh_missing_dc) = Self::coverage_ratio(&desired_by_dc, &fresh_writer_addrs); - route_quorum_ok = fresh_coverage_ratio >= min_ratio; - redundancy_ok = fresh_missing_dc.is_empty(); - redundancy_missing_dc = fresh_missing_dc.clone(); - gate_coverage_ratio = fresh_coverage_ratio; if fresh_coverage_ratio < min_ratio { self.set_last_drain_gate( false, - redundancy_ok, + fresh_missing_dc.is_empty(), MeDrainGateReason::CoverageQuorum, now_epoch_secs, ); @@ -233,49 +234,77 @@ impl MePool { } } + drop(writers); + let commit = self + .commit_reinit_attempt(&attempt, &desired_by_dc, min_ratio) + .await; + let outcome = match commit { + Ok(outcome) => outcome, + Err(ReinitCommitFailure::Superseded) => { + debug!( + previous_generation, + generation, + "ME reinit result discarded after a newer desired-map attempt" + ); + return false; + } + Err(ReinitCommitFailure::Coverage { + coverage_ratio, + missing_dc, + }) => { + self.set_last_drain_gate( + false, + missing_dc.is_empty(), + MeDrainGateReason::CoverageQuorum, + now_epoch_secs, + ); + warn!( + previous_generation, + generation, + coverage_ratio = format_args!("{coverage_ratio:.3}"), + min_ratio = format_args!("{min_ratio:.3}"), + missing_dc = ?missing_dc, + "ME reinit coverage changed before commit; keeping current generation" + ); + return false; + } + Err(ReinitCommitFailure::Redundancy { + coverage_ratio, + missing_dc, + }) => { + self.set_last_drain_gate( + true, + false, + MeDrainGateReason::Redundancy, + now_epoch_secs, + ); + warn!( + previous_generation, + generation, + coverage_ratio = format_args!("{coverage_ratio:.3}"), + min_ratio = format_args!("{min_ratio:.3}"), + missing_dc = ?missing_dc, + "ME hardswap weighted quorum requires stale-binding fallback" + ); + return false; + } + }; self.set_last_drain_gate( - route_quorum_ok, - redundancy_ok, + true, + outcome.missing_dc.is_empty(), MeDrainGateReason::Open, now_epoch_secs, ); - if !redundancy_ok { + if !outcome.missing_dc.is_empty() { warn!( - missing_dc = ?redundancy_missing_dc, - coverage_ratio = format_args!("{gate_coverage_ratio:.3}"), + missing_dc = ?outcome.missing_dc, + coverage_ratio = format_args!("{:.3}", outcome.coverage_ratio), min_ratio = format_args!("{min_ratio:.3}"), - "ME reinit proceeds with weighted quorum while some DC groups remain uncovered" + "ME reinit committed with bounded stale fallback for uncovered DC groups" ); } - if !self.commit_reinit_attempt(&attempt) { - debug!( - previous_generation, - generation, "ME reinit result discarded after a newer desired-map attempt" - ); - return false; - } - - let desired_addrs: HashSet<(i32, SocketAddr)> = desired_by_dc - .iter() - .flat_map(|(dc, set)| set.iter().copied().map(|addr| (*dc, addr))) - .collect(); - - let stale_writer_ids: Vec = writers - .iter() - .filter(|w| !w.draining.load(Ordering::Relaxed)) - .filter(|w| { - if hardswap { - w.generation < generation - } else { - !desired_addrs.contains(&(w.writer_dc, w.addr)) - } - }) - .map(|w| w.id) - .collect(); - drop(writers); - - if stale_writer_ids.is_empty() { + if outcome.stale_writer_ids.is_empty() { debug!("ME reinit cycle completed with no stale writers"); return true; } @@ -283,35 +312,20 @@ impl MePool { let drain_timeout = self.force_close_timeout(); let drain_timeout_secs = drain_timeout.map(|d| d.as_secs()).unwrap_or(0); info!( - stale_writers = stale_writer_ids.len(), + stale_writers = outcome.stale_writer_ids.len(), + force_close_writers = outcome.force_close_writer_ids.len(), previous_generation, generation, hardswap, - coverage_ratio = format_args!("{coverage_ratio:.3}"), + coverage_ratio = format_args!("{:.3}", outcome.coverage_ratio), min_ratio = format_args!("{min_ratio:.3}"), drain_timeout_secs, "ME reinit cycle covered; processing stale writers" ); self.stats.increment_pool_swap_total(); - let can_drop_with_replacement = self.has_non_draining_writer_per_desired_dc_group().await; - if can_drop_with_replacement { - info!( - stale_writers = stale_writer_ids.len(), - "ME reinit stale writers: replacement coverage ready, force-closing clients for fast rebind" - ); - } else { - warn!( - stale_writers = stale_writer_ids.len(), - "ME reinit stale writers: replacement coverage incomplete, keeping draining fallback" - ); - } - for writer_id in stale_writer_ids { - self.mark_writer_draining_with_timeout(writer_id, drain_timeout, !hardswap) - .await; - if can_drop_with_replacement { - self.stats.increment_pool_force_close_total(); - self.remove_writer_and_close_clients(writer_id).await; - } + for writer_id in outcome.force_close_writer_ids { + self.stats.increment_pool_force_close_total(); + self.remove_writer_and_close_clients(writer_id).await; } true } diff --git a/src/transport/middle_proxy/pool_reinit/tests.rs b/src/transport/middle_proxy/pool_reinit/tests.rs index cb85871..4f500e1 100644 --- a/src/transport/middle_proxy/pool_reinit/tests.rs +++ b/src/transport/middle_proxy/pool_reinit/tests.rs @@ -1,15 +1,68 @@ use std::collections::{HashMap, HashSet}; use std::net::{IpAddr, Ipv4Addr, SocketAddr}; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU32, AtomicU64, Ordering}; +use std::time::Instant; -use super::{MePool, commit_reinit_state}; +use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; + +use super::{MePool, ReinitCommitFailure, commit_reinit_state}; +use crate::config::MeBindStaleMode; +use crate::transport::middle_proxy::codec::WriterCommand; use crate::transport::middle_proxy::pool::{ - ReinitAttemptState, ReinitCoordinatorState, ReinitPendingState, + MeWriter, ReinitAttemptState, ReinitCoordinatorState, ReinitPendingState, WriterContour, }; +use crate::transport::middle_proxy::pool_writer_security_tests::make_pool; fn addr(octet: u8, port: u16) -> SocketAddr { SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, octet)), port) } +async fn insert_writer( + pool: &Arc, + writer_id: u64, + writer_dc: i32, + endpoint: SocketAddr, + generation: u64, + contour: WriterContour, +) -> MeWriter { + let (tx, _rx) = mpsc::channel::(8); + let byte_budget = pool.new_writer_byte_budget(); + let writer = MeWriter { + id: writer_id, + addr: endpoint, + source_ip: endpoint.ip(), + writer_dc, + generation, + contour: Arc::new(AtomicU8::new(contour.as_u8())), + created_at: Instant::now(), + tx: tx.clone(), + byte_budget: byte_budget.clone(), + cancel: CancellationToken::new(), + degraded: Arc::new(AtomicBool::new(false)), + rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)), + draining: Arc::new(AtomicBool::new(false)), + draining_started_at_epoch_secs: Arc::new(AtomicU64::new(0)), + drain_deadline_epoch_secs: Arc::new(AtomicU64::new(0)), + allow_drain_fallback: Arc::new(AtomicBool::new(false)), + }; + + pool.registry + .register_writer(writer_id, tx, byte_budget) + .await; + pool.writers.write().await.push(writer.clone()); + pool.conn_count.fetch_add(1, Ordering::Relaxed); + writer +} + +fn desired_two_dcs() -> HashMap> { + HashMap::from([ + (1, HashSet::from([addr(1, 2001)])), + (2, HashSet::from([addr(2, 2002)])), + ]) +} + #[test] fn coverage_ratio_counts_dc_coverage_not_floor() { let dc1 = addr(1, 2001); @@ -95,3 +148,226 @@ fn stale_concurrent_attempt_cannot_regress_active_generation() { assert_eq!(state.active_generation, 3); assert!(state.pending.is_none()); } + +#[tokio::test] +async fn partial_hardswap_is_rejected_when_stale_binding_is_disabled() { + let pool = make_pool().await; + let desired_by_dc = desired_two_dcs(); + let active_generation = pool.current_generation(); + let old_dc1 = insert_writer( + &pool, + 101, + 1, + addr(1, 2001), + active_generation, + WriterContour::Active, + ) + .await; + let old_dc2 = insert_writer( + &pool, + 102, + 2, + addr(2, 2002), + active_generation, + WriterContour::Active, + ) + .await; + let map_hash = MePool::desired_map_hash(&desired_by_dc); + let reservation = pool.reserve_reinit_attempt(true, map_hash, 100); + insert_writer( + &pool, + 201, + 1, + addr(1, 2001), + reservation.attempt.generation, + WriterContour::Warm, + ) + .await; + + let result = pool + .commit_reinit_attempt(&reservation.attempt, &desired_by_dc, 0.5) + .await; + + assert!(matches!(result, Err(ReinitCommitFailure::Redundancy { .. }))); + assert_eq!(pool.current_generation(), active_generation); + assert!(!old_dc1.draining.load(Ordering::Acquire)); + assert!(!old_dc2.draining.load(Ordering::Acquire)); +} + +#[tokio::test] +async fn partial_hardswap_preserves_fallback_only_for_missing_dc() { + let pool = make_pool().await; + pool.binding_policy + .me_bind_stale_mode + .store(MeBindStaleMode::Ttl.as_u8(), Ordering::Release); + let desired_by_dc = desired_two_dcs(); + let active_generation = pool.current_generation(); + let old_dc1 = insert_writer( + &pool, + 301, + 1, + addr(1, 2001), + active_generation, + WriterContour::Active, + ) + .await; + let old_dc2 = insert_writer( + &pool, + 302, + 2, + addr(9, 2999), + active_generation, + WriterContour::Active, + ) + .await; + let map_hash = MePool::desired_map_hash(&desired_by_dc); + let reservation = pool.reserve_reinit_attempt(true, map_hash, 100); + let fresh_dc1 = insert_writer( + &pool, + 401, + 1, + addr(1, 2001), + reservation.attempt.generation, + WriterContour::Warm, + ) + .await; + + let outcome = pool + .commit_reinit_attempt(&reservation.attempt, &desired_by_dc, 0.5) + .await + .expect("partial hardswap must commit when bounded stale fallback is enabled"); + + assert_eq!(pool.current_generation(), reservation.attempt.generation); + assert_eq!(outcome.missing_dc, vec![2]); + assert_eq!(outcome.force_close_writer_ids, vec![301]); + assert!(old_dc1.draining.load(Ordering::Acquire)); + assert!(!old_dc1.allow_drain_fallback.load(Ordering::Acquire)); + assert!(old_dc2.draining.load(Ordering::Acquire)); + assert!(old_dc2.allow_drain_fallback.load(Ordering::Acquire)); + assert_eq!( + WriterContour::from_u8(fresh_dc1.contour.load(Ordering::Acquire)), + WriterContour::Active + ); +} + +#[tokio::test] +async fn complete_hardswap_promotes_fresh_generation_and_retires_old_writers() { + let pool = make_pool().await; + let desired_by_dc = desired_two_dcs(); + let active_generation = pool.current_generation(); + let old_dc1 = insert_writer( + &pool, + 501, + 1, + addr(1, 2001), + active_generation, + WriterContour::Active, + ) + .await; + let old_dc2 = insert_writer( + &pool, + 502, + 2, + addr(2, 2002), + active_generation, + WriterContour::Active, + ) + .await; + let map_hash = MePool::desired_map_hash(&desired_by_dc); + let reservation = pool.reserve_reinit_attempt(true, map_hash, 100); + let fresh_dc1 = insert_writer( + &pool, + 601, + 1, + addr(1, 2001), + reservation.attempt.generation, + WriterContour::Warm, + ) + .await; + let fresh_dc2 = insert_writer( + &pool, + 602, + 2, + addr(2, 2002), + reservation.attempt.generation, + WriterContour::Warm, + ) + .await; + + let outcome = pool + .commit_reinit_attempt(&reservation.attempt, &desired_by_dc, 1.0) + .await + .expect("fully covered hardswap must commit"); + + assert_eq!(pool.current_generation(), reservation.attempt.generation); + assert!(outcome.missing_dc.is_empty()); + assert_eq!(outcome.force_close_writer_ids, vec![501, 502]); + assert!(old_dc1.draining.load(Ordering::Acquire)); + assert!(old_dc2.draining.load(Ordering::Acquire)); + assert_eq!( + WriterContour::from_u8(fresh_dc1.contour.load(Ordering::Acquire)), + WriterContour::Active + ); + assert_eq!( + WriterContour::from_u8(fresh_dc2.contour.load(Ordering::Acquire)), + WriterContour::Active + ); +} + +#[tokio::test] +async fn generation_role_reconciliation_promotes_active_warm_and_drains_orphans() { + let pool = make_pool().await; + let endpoint = addr(1, 2001); + pool.preferred_endpoints_by_dc + .store(Arc::new(HashMap::from([(1, vec![endpoint])]))); + let desired_by_dc = HashMap::from([(1, HashSet::from([endpoint]))]); + let map_hash = MePool::desired_map_hash(&desired_by_dc); + let reservation = pool.reserve_reinit_attempt(true, map_hash, 100); + let active_warm = insert_writer( + &pool, + 701, + 1, + endpoint, + pool.current_generation(), + WriterContour::Warm, + ) + .await; + let pending_warm = insert_writer( + &pool, + 702, + 1, + endpoint, + reservation.attempt.generation, + WriterContour::Warm, + ) + .await; + let orphan_warm = insert_writer( + &pool, + 703, + 1, + endpoint, + reservation.attempt.generation + 10, + WriterContour::Warm, + ) + .await; + + let changed = pool.reconcile_writer_generation_roles().await; + + assert_eq!(changed, 2); + assert_eq!( + WriterContour::from_u8(active_warm.contour.load(Ordering::Acquire)), + WriterContour::Active + ); + assert_eq!( + WriterContour::from_u8(pending_warm.contour.load(Ordering::Acquire)), + WriterContour::Warm + ); + assert!(orphan_warm.draining.load(Ordering::Acquire)); + assert!(!orphan_warm.allow_drain_fallback.load(Ordering::Acquire)); + assert_eq!( + pool.api_hardswap_snapshot() + .await + .orphan_warm_writers_current, + 0 + ); +} diff --git a/src/transport/middle_proxy/pool_runtime_api.rs b/src/transport/middle_proxy/pool_runtime_api.rs index 78c50e7..5ac9b83 100644 --- a/src/transport/middle_proxy/pool_runtime_api.rs +++ b/src/transport/middle_proxy/pool_runtime_api.rs @@ -1,7 +1,7 @@ use std::collections::HashMap; use std::time::Instant; -use super::pool::{MeDrainGateReason, MePool, RefillDcKey}; +use super::pool::{MeDrainGateReason, MePool, RefillTargetKey}; use crate::network::IpFamily; #[derive(Clone, Debug)] @@ -61,14 +61,17 @@ impl MePool { let refill_states = self.refill_states.lock(); let inflight_endpoints_total = refill_states .values() - .map(|pending| 1usize + usize::from(pending.is_some())) + .map(|state| 1usize.saturating_add(state.pending_count)) .sum(); let running_dc_total = refill_states.len(); let pending_dc_total = refill_states .values() - .filter(|pending| pending.is_some()) + .filter(|state| state.pending_count > 0) .count(); - let inflight_dc_keys = refill_states.keys().copied().collect::>(); + let inflight_dc_keys = refill_states + .keys() + .copied() + .collect::>(); drop(refill_states); let mut by_dc_map = HashMap::<(i16, &'static str), usize>::new(); diff --git a/src/transport/middle_proxy/pool_status.rs b/src/transport/middle_proxy/pool_status.rs index dcf7fc3..0e3cecc 100644 --- a/src/transport/middle_proxy/pool_status.rs +++ b/src/transport/middle_proxy/pool_status.rs @@ -12,6 +12,9 @@ use crate::transport::upstream::IpPreference; mod status_snapshot; // ME runtime policy and coherent snapshot assembly. mod runtime_snapshot; +// Hardswap ownership, coverage, and writer-replacement lifecycle state. +mod hardswap_snapshot; +pub(crate) use hardswap_snapshot::MeApiHardswapSnapshot; #[derive(Clone, Debug)] pub(crate) struct MeApiWriterStatusSnapshot { pub writer_id: u64, @@ -98,6 +101,20 @@ pub(crate) struct MeApiRuntimeSnapshot { pub reinit_inflight: usize, pub reinit_max_concurrency_effective: usize, pub hardswap_enabled: bool, + /// Number of authoritative warm writers in the pending generation. + pub pending_writers_current: usize, + /// Number of writers still required to reach the pending generation floor. + pub pending_writer_deficit: usize, + /// Number of desired DC groups without pending-generation coverage. + pub pending_missing_dc_groups: usize, + /// Whether the pending generation targets the current desired endpoint map. + pub pending_map_current: Option, + /// Number of warm writers not owned by the current pending generation. + pub orphan_warm_writers_current: usize, + /// Number of writer replacements in the preparatory phase. + pub replacement_preparing_current: usize, + /// Number of replacement victims closed to new client binds. + pub replacement_retiring_current: usize, pub floor_mode: &'static str, pub adaptive_floor_idle_secs: u64, pub adaptive_floor_min_writers_single_endpoint: u8, diff --git a/src/transport/middle_proxy/pool_status/hardswap_snapshot.rs b/src/transport/middle_proxy/pool_status/hardswap_snapshot.rs new file mode 100644 index 0000000..d332249 --- /dev/null +++ b/src/transport/middle_proxy/pool_status/hardswap_snapshot.rs @@ -0,0 +1,102 @@ +use super::*; + +#[derive(Clone, Debug)] +/// Bounded control-plane snapshot of hardswap and replacement progress. +pub(crate) struct MeApiHardswapSnapshot { + /// Whether a hardswap generation is pending. + pub pending: bool, + /// Age of the pending generation when one exists. + pub pending_age_secs: Option, + /// Number of authoritative warm writers in the pending generation. + pub pending_writers_current: usize, + /// Number of writers still required to reach the pending generation floor. + pub pending_writer_deficit: usize, + /// Number of desired DC groups without a pending-generation writer. + pub pending_missing_dc_groups: usize, + /// Whether the pending generation targets the current desired endpoint map. + pub pending_map_current: Option, + /// Number of warm writers not owned by the current pending generation. + pub orphan_warm_writers_current: usize, + /// Number of writer replacements still in the preparatory phase. + pub replacement_preparing_current: usize, + /// Number of replacement victims already closed to new bindings. + pub replacement_retiring_current: usize, +} + +impl MePool { + /// Returns bounded hardswap progress without exposing generation or endpoint labels. + pub(crate) async fn api_hardswap_snapshot(&self) -> MeApiHardswapSnapshot { + let reinit = self.reinit.status.load_full(); + self.api_hardswap_snapshot_for_reinit(reinit.as_ref()).await + } + + pub(super) async fn api_hardswap_snapshot_for_reinit( + &self, + reinit: &ReinitStatusSnapshot, + ) -> MeApiHardswapSnapshot { + let desired_by_dc = self.desired_dc_endpoints().await; + let desired_hash = Self::desired_map_hash(&desired_by_dc); + let writers = self.writers.read().await; + let pending_generation = reinit.pending_hardswap_generation; + let pending = pending_generation != 0; + let mut pending_writers_current = 0usize; + let mut pending_by_dc = HashMap::::new(); + let mut orphan_warm_writers_current = 0usize; + + for writer in writers.iter() { + if writer.draining.load(Ordering::Acquire) { + continue; + } + let contour = WriterContour::from_u8(writer.contour.load(Ordering::Acquire)); + if contour == WriterContour::Warm && writer.generation != pending_generation { + orphan_warm_writers_current = orphan_warm_writers_current.saturating_add(1); + } + if pending + && writer.generation == pending_generation + && contour == WriterContour::Warm + && desired_by_dc + .get(&writer.writer_dc) + .is_some_and(|endpoints| endpoints.contains(&writer.addr)) + { + pending_writers_current = pending_writers_current.saturating_add(1); + *pending_by_dc.entry(writer.writer_dc).or_insert(0) += 1; + } + } + + let mut pending_writer_deficit = 0usize; + let mut pending_missing_dc_groups = 0usize; + if pending { + for (dc, endpoints) in &desired_by_dc { + if endpoints.is_empty() { + continue; + } + let alive = pending_by_dc.get(dc).copied().unwrap_or(0); + let required = self.required_writers_for_dc(endpoints.len()); + pending_writer_deficit = pending_writer_deficit + .saturating_add(required.saturating_sub(alive)); + if alive == 0 { + pending_missing_dc_groups = pending_missing_dc_groups.saturating_add(1); + } + } + } + let (replacement_preparing_current, replacement_retiring_current) = + self.registry.writer_replacement_counts(); + let pending_age_secs = pending.then(|| { + Self::now_epoch_secs() + .saturating_sub(reinit.pending_hardswap_started_at_epoch_secs) + }); + + MeApiHardswapSnapshot { + pending, + pending_age_secs, + pending_writers_current, + pending_writer_deficit, + pending_missing_dc_groups, + pending_map_current: pending + .then_some(reinit.pending_hardswap_map_hash == desired_hash), + orphan_warm_writers_current, + replacement_preparing_current, + replacement_retiring_current, + } + } +} diff --git a/src/transport/middle_proxy/pool_status/runtime_snapshot.rs b/src/transport/middle_proxy/pool_status/runtime_snapshot.rs index 9291d85..f466977 100644 --- a/src/transport/middle_proxy/pool_status/runtime_snapshot.rs +++ b/src/transport/middle_proxy/pool_status/runtime_snapshot.rs @@ -16,6 +16,7 @@ impl MePool { let pending_started_at = reinit.pending_hardswap_started_at_epoch_secs; let pending_hardswap_age_secs = (pending_started_at > 0).then_some(now_epoch_secs.saturating_sub(pending_started_at)); + let hardswap = self.api_hardswap_snapshot_for_reinit(reinit).await; let mut quarantined_endpoints = Vec::::new(); { @@ -64,6 +65,13 @@ impl MePool { .max_concurrency_effective .load(Ordering::Acquire), hardswap_enabled: self.reinit.hardswap.load(Ordering::Relaxed), + pending_writers_current: hardswap.pending_writers_current, + pending_writer_deficit: hardswap.pending_writer_deficit, + pending_missing_dc_groups: hardswap.pending_missing_dc_groups, + pending_map_current: hardswap.pending_map_current, + orphan_warm_writers_current: hardswap.orphan_warm_writers_current, + replacement_preparing_current: hardswap.replacement_preparing_current, + replacement_retiring_current: hardswap.replacement_retiring_current, floor_mode: floor_mode_label(self.floor_mode()), adaptive_floor_idle_secs: self .floor_runtime diff --git a/src/transport/middle_proxy/pool_status/status_snapshot.rs b/src/transport/middle_proxy/pool_status/status_snapshot.rs index 660db57..4dc6d86 100644 --- a/src/transport/middle_proxy/pool_status/status_snapshot.rs +++ b/src/transport/middle_proxy/pool_status/status_snapshot.rs @@ -16,10 +16,15 @@ impl MePool { return false; } + let active_generation = self.reinit.status.load().active_generation; let writers = self.writers.read().await.clone(); let mut live_writers_by_dc = HashMap::::new(); for writer in writers.iter() { - if writer.draining.load(Ordering::Relaxed) { + if writer.draining.load(Ordering::Relaxed) + || writer.generation != active_generation + || WriterContour::from_u8(writer.contour.load(Ordering::Relaxed)) + != WriterContour::Active + { continue; } if let Ok(dc) = i16::try_from(writer.writer_dc) { @@ -53,10 +58,15 @@ impl MePool { return false; } + let active_generation = self.reinit.status.load().active_generation; let writers = self.writers.read().await.clone(); let mut live_writers_by_dc = HashMap::::new(); for writer in writers.iter() { - if writer.draining.load(Ordering::Relaxed) { + if writer.draining.load(Ordering::Relaxed) + || writer.generation != active_generation + || WriterContour::from_u8(writer.contour.load(Ordering::Relaxed)) + != WriterContour::Active + { continue; } if let Ok(dc) = i16::try_from(writer.writer_dc) { @@ -127,7 +137,7 @@ impl MePool { for writer in writers.iter() { let endpoint = writer.addr; let dc = i16::try_from(writer.writer_dc).ok(); - let draining = writer.draining.load(Ordering::Relaxed); + let draining = writer.draining.load(Ordering::Acquire); let degraded = writer.degraded.load(Ordering::Relaxed); let matches_active_generation = writer.generation == active_generation; let in_desired_map = dc @@ -156,13 +166,18 @@ impl MePool { && drain_ttl_secs > 0 && drain_started_at_epoch_secs .is_some_and(|started| now_epoch_secs.saturating_sub(started) > drain_ttl_secs); - let state = match WriterContour::from_u8(writer.contour.load(Ordering::Relaxed)) { + let contour = WriterContour::from_u8(writer.contour.load(Ordering::Relaxed)); + let state = match contour { WriterContour::Warm => "warm", WriterContour::Active => "active", WriterContour::Draining => "draining", }; - if !draining && let Some(dc_idx) = dc { + let authoritative_active = !draining + && contour == WriterContour::Active + && matches_active_generation + && in_desired_map; + if authoritative_active && let Some(dc_idx) = dc { *live_writers_by_dc_endpoint .entry((dc_idx, endpoint)) .or_insert(0) += 1; @@ -172,9 +187,7 @@ impl MePool { entry.0 += ema_ms; entry.1 += 1; } - if matches_active_generation && in_desired_map { - *fresh_writers_by_dc.entry(dc_idx).or_insert(0) += 1; - } + *fresh_writers_by_dc.entry(dc_idx).or_insert(0) += 1; } writer_rows.push(MeApiWriterStatusSnapshot { diff --git a/src/transport/middle_proxy/pool_writer.rs b/src/transport/middle_proxy/pool_writer.rs index f116225..a02c868 100644 --- a/src/transport/middle_proxy/pool_writer.rs +++ b/src/transport/middle_proxy/pool_writer.rs @@ -2,6 +2,7 @@ use std::collections::HashMap; use std::future::Future; use std::io::ErrorKind; use std::net::SocketAddr; +use std::pin::Pin; use std::sync::Arc; use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU32, AtomicU64, Ordering}; use std::time::{Duration, Instant}; @@ -19,14 +20,61 @@ use crate::error::{ProxyError, Result}; use crate::protocol::constants::{RPC_CLOSE_EXT_U32, RPC_PING_U32}; use super::codec::{RpcWriter, WriterCommand, build_control_payload}; -use super::pool::{MePool, MeWriter, WriterContour}; +use super::pool::{ + MePool, MeWriter, WriterContour, WriterOpenIntent, WriterOpenReservation, WriterRole, +}; use super::pool_lifecycle::MeTaskRegistration; use super::reader::reader_loop; +use super::registry::WriterReplacementReservation; use super::wire::build_proxy_req_payload; // Writer admission, teardown, and drain-state transitions. +mod publication; +mod replacement; mod runtime; +struct PreparedWriter<'a> { + writer: MeWriter, + tx: mpsc::Sender, + byte_budget: Arc, + task_registration: MeTaskRegistration<'a>, + writer_task: Pin + Send + 'static>>, + intent: WriterOpenIntent, + _open_reservation: WriterOpenReservation<'a>, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +/// Commit-time invariant applied to a transactional writer replacement. +pub(super) enum WriterReplacementPurpose { + /// Replaces an idle writer before the upstream idle lifetime expires. + IdleRefresh, + /// Moves one idle slot from a donor above floor to a receiver below floor. + FloorRebalance { + /// Minimum writer count that must remain in the donor group. + donor_floor: usize, + /// Target count above which the receiver no longer needs the slot. + receiver_floor: usize, + }, + /// Replaces a stale-generation writer with an active-generation writer. + GenerationConvergence, + /// Rotates a single-endpoint shadow while retaining coverage. + ShadowRotation, + /// Replaces a writer after cryptographic secret rotation. + SecretRotation, +} + +impl WriterReplacementPurpose { + fn as_str(self) -> &'static str { + match self { + Self::IdleRefresh => "idle_refresh", + Self::FloorRebalance { .. } => "floor_rebalance", + Self::GenerationConvergence => "generation_convergence", + Self::ShadowRotation => "shadow_rotation", + Self::SecretRotation => "secret_rotation", + } + } +} + const ME_ACTIVE_PING_SECS: u64 = 25; const ME_ACTIVE_PING_JITTER_SECS: i64 = 5; const ME_IDLE_KEEPALIVE_MAX_SECS: u64 = 5; diff --git a/src/transport/middle_proxy/pool_writer/publication.rs b/src/transport/middle_proxy/pool_writer/publication.rs new file mode 100644 index 0000000..2027e33 --- /dev/null +++ b/src/transport/middle_proxy/pool_writer/publication.rs @@ -0,0 +1,180 @@ +use super::*; + +impl MePool { + /// Publishes a connected writer only while its generation owns the requested role. + pub(super) async fn publish_connected_writer( + self: &Arc, + prepared: PreparedWriter<'_>, + ) -> Result<()> { + let PreparedWriter { + writer, + tx, + byte_budget, + task_registration, + writer_task, + intent, + _open_reservation, + } = prepared; + let writer_id = writer.id; + let mut writers = self.writers.write().await; + 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())?; + writer.contour.store(contour.as_u8(), Ordering::Release); + registry_registration.install(writer_id, tx, byte_budget); + writers.push(writer); + self.conn_count.fetch_add(1, Ordering::Relaxed); + self.lifecycle + .spawn_registered_writer(task_registration, writer_task); + drop(coordinator); + drop(registry_registration); + drop(writers); + self.notify_writer_epoch(); + Ok(()) + } + + /// Resolves the writer role against the linearized reinitialization authority. + pub(super) fn authorize_writer_publication( + &self, + writer: &MeWriter, + coordinator: &crate::transport::middle_proxy::pool::ReinitCoordinatorState, + ) -> Result { + let endpoint_is_current = self + .preferred_endpoints_by_dc + .load() + .get(&writer.writer_dc) + .is_some_and(|endpoints| endpoints.contains(&writer.addr)); + if !endpoint_is_current { + return Err(ProxyError::Proxy( + "ME writer target changed before publication".into(), + )); + } + + let requested = WriterContour::from_u8(writer.contour.load(Ordering::Acquire)); + if writer.generation == coordinator.active_generation + && matches!(requested, WriterContour::Active | WriterContour::Warm) + { + return Ok(WriterContour::Active); + } + if requested == WriterContour::Warm + && coordinator.pending.is_some_and(|pending| { + pending.generation == writer.generation + && pending.map_hash == coordinator.desired_map_hash + }) + { + return Ok(WriterContour::Warm); + } + Err(ProxyError::Proxy( + "ME writer generation lost publication authority".into(), + )) + } + + /// Revalidates role-local capacity at the serialized publication boundary. + pub(in crate::transport::middle_proxy) fn authorize_writer_publication_capacity( + &self, + writer: &MeWriter, + contour: WriterContour, + intent: WriterOpenIntent, + writers: &[MeWriter], + ) -> Result<()> { + let family = if writer.addr.is_ipv4() { + crate::network::IpFamily::V4 + } else { + crate::network::IpFamily::V6 + }; + let now_epoch_secs = Self::now_epoch_secs(); + if !self.family_enabled_for_drain_coverage(family, now_epoch_secs) { + return Err(ProxyError::Proxy( + "ME writer family lost publication authority".into(), + )); + } + if intent == WriterOpenIntent::Replacement || contour == WriterContour::Draining { + return Ok(()); + } + let preferred = self.preferred_endpoints_by_dc.load(); + let Some(endpoints) = preferred.get(&writer.writer_dc) else { + return Err(ProxyError::Proxy( + "ME writer target changed before publication".into(), + )); + }; + let required = match contour { + WriterContour::Active => self.required_writers_for_dc( + endpoints + .iter() + .filter(|endpoint| endpoint.is_ipv4() == writer.addr.is_ipv4()) + .count(), + ), + WriterContour::Warm => self.required_writers_for_dc( + endpoints + .iter() + .filter(|endpoint| { + let endpoint_family = if endpoint.is_ipv4() { + crate::network::IpFamily::V4 + } else { + crate::network::IpFamily::V6 + }; + self.family_enabled_for_drain_coverage( + endpoint_family, + now_epoch_secs, + ) + }) + .count(), + ), + WriterContour::Draining => 0, + }; + let current = writers + .iter() + .filter(|candidate| { + !candidate.draining.load(Ordering::Acquire) + && candidate.writer_dc == writer.writer_dc + && candidate.generation == writer.generation + && WriterContour::from_u8(candidate.contour.load(Ordering::Acquire)) + == contour + && (contour == WriterContour::Warm + || candidate.addr.is_ipv4() == writer.addr.is_ipv4()) + && endpoints.contains(&candidate.addr) + && (contour != WriterContour::Warm + || self.family_enabled_for_drain_coverage( + if candidate.addr.is_ipv4() { + crate::network::IpFamily::V4 + } else { + crate::network::IpFamily::V6 + }, + now_epoch_secs, + )) + }) + .count(); + if current >= required { + return Err(ProxyError::Proxy( + "ME writer floor was restored before publication".into(), + )); + } + Ok(()) + } + + /// Commits writer visibility and lifecycle ownership after all cancellation points. + #[allow(clippy::too_many_arguments)] + pub(in crate::transport::middle_proxy) async fn publish_prepared_writer( + self: &Arc, + writer: MeWriter, + tx: mpsc::Sender, + byte_budget: Arc, + task_registration: MeTaskRegistration<'_>, + writer_task: F, + ) where + F: Future + Send + 'static, + { + // Writer publication follows the global writers -> registry binding lock order. + let mut writers = self.writers.write().await; + let mut registry_registration = self.registry.prepare_writer_registration().await; + registry_registration.install(writer.id, tx, byte_budget); + writers.push(writer); + self.conn_count.fetch_add(1, Ordering::Relaxed); + self.lifecycle + .spawn_registered_writer(task_registration, writer_task); + drop(writers); + drop(registry_registration); + self.notify_writer_epoch(); + } +} diff --git a/src/transport/middle_proxy/pool_writer/replacement.rs b/src/transport/middle_proxy/pool_writer/replacement.rs new file mode 100644 index 0000000..85a9927 --- /dev/null +++ b/src/transport/middle_proxy/pool_writer/replacement.rs @@ -0,0 +1,376 @@ +use super::*; + +impl MePool { + /// Opens a replacement and atomically publishes it before retiring its reserved victim. + pub(in crate::transport::middle_proxy) async fn replace_writer_with_generation_contour_for_dc( + self: &Arc, + addr: SocketAddr, + rng: &SecureRandom, + generation: u64, + contour: WriterContour, + writer_dc: i32, + expected_victim_role: WriterRole, + purpose: WriterReplacementPurpose, + reservation: &mut WriterReplacementReservation<'_>, + ) -> Result<()> { + let prepared = self + .prepare_writer_with_intent( + addr, + rng, + generation, + contour, + writer_dc, + WriterOpenIntent::Replacement, + ) + .await?; + self.publish_prepared_replacement_writer( + prepared, + expected_victim_role, + purpose, + reservation, + ) + .await + } + + async fn publish_prepared_replacement_writer( + self: &Arc, + prepared: PreparedWriter<'_>, + expected_victim_role: WriterRole, + purpose: WriterReplacementPurpose, + reservation: &mut WriterReplacementReservation<'_>, + ) -> Result<()> { + let PreparedWriter { + writer, + tx, + byte_budget, + task_registration, + writer_task, + intent, + _open_reservation, + } = prepared; + let replacement_writer_id = writer.id; + let victim_writer_id = reservation.writer_id(); + + // Lock order is writers -> registry binding. No cancellation point follows acquisition of + // the registry guard, so publication and victim retirement commit as one state change. + let mut writers = self.writers.write().await; + 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())?; + writer.contour.store(contour.as_u8(), Ordering::Release); + let Some(victim_pos) = writers + .iter() + .position(|candidate| candidate.id == victim_writer_id) + else { + return Err(ProxyError::Proxy( + "ME replacement victim disappeared before commit".into(), + )); + }; + let victim = &writers[victim_pos]; + if victim.draining.load(Ordering::Acquire) || !expected_victim_role.matches(victim) { + return Err(ProxyError::Proxy( + "ME replacement victim changed role before commit".into(), + )); + } + if let WriterReplacementPurpose::FloorRebalance { + donor_floor, + receiver_floor, + } = purpose + { + if expected_victim_role.generation != coordinator.active_generation + || writer.generation != coordinator.active_generation + { + return Err(ProxyError::Proxy( + "ME floor rebalance lost active-generation authority".into(), + )); + } + let preferred = self.preferred_endpoints_by_dc.load(); + let donor_count = writers + .iter() + .filter(|candidate| { + !candidate.draining.load(Ordering::Acquire) + && candidate.writer_dc == expected_victim_role.dc + && candidate.generation == expected_victim_role.generation + && WriterContour::from_u8(candidate.contour.load(Ordering::Acquire)) + == WriterContour::Active + && preferred + .get(&candidate.writer_dc) + .is_some_and(|endpoints| endpoints.contains(&candidate.addr)) + && (candidate.addr.is_ipv4() + == matches!(expected_victim_role.family, crate::network::IpFamily::V4)) + }) + .count(); + let receiver_family = if writer.addr.is_ipv4() { + crate::network::IpFamily::V4 + } else { + crate::network::IpFamily::V6 + }; + let receiver_count = writers + .iter() + .filter(|candidate| { + !candidate.draining.load(Ordering::Acquire) + && candidate.writer_dc == writer.writer_dc + && candidate.generation == writer.generation + && WriterContour::from_u8(candidate.contour.load(Ordering::Acquire)) + == WriterContour::Active + && preferred + .get(&candidate.writer_dc) + .is_some_and(|endpoints| endpoints.contains(&candidate.addr)) + && (if candidate.addr.is_ipv4() { + crate::network::IpFamily::V4 + } else { + crate::network::IpFamily::V6 + }) == receiver_family + }) + .count(); + if donor_count <= donor_floor || receiver_count >= receiver_floor { + return Err(ProxyError::Proxy( + "ME floor rebalance became unnecessary before commit".into(), + )); + } + } + if !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); + writers.push(writer); + self.conn_count.fetch_add(1, Ordering::Relaxed); + writers.publish_current(); + self.apply_writer_draining_state( + &writers[victim_pos], + self.force_close_timeout(), + false, + ); + self.lifecycle + .spawn_registered_writer(task_registration, writer_task); + reservation.mark_committed(); + drop(coordinator); + drop(registry_registration); + drop(writers); + self.notify_writer_epoch(); + info!( + victim_writer_id, + replacement_writer_id, + purpose = purpose.as_str(), + "ME writer replacement committed" + ); + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::net::{IpAddr, Ipv4Addr}; + + use crate::transport::middle_proxy::pool_writer_security_tests::make_pool; + + use super::*; + + fn endpoint(octet: u8) -> SocketAddr { + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, octet)), 443) + } + + async fn install_writer( + pool: &Arc, + writer_id: u64, + writer_dc: i32, + addr: SocketAddr, + ) -> MeWriter { + let (tx, _rx) = mpsc::channel::(8); + let byte_budget = pool.new_writer_byte_budget(); + let writer = MeWriter { + id: writer_id, + addr, + source_ip: addr.ip(), + writer_dc, + generation: pool.current_generation(), + contour: Arc::new(AtomicU8::new(WriterContour::Active.as_u8())), + created_at: Instant::now(), + tx: tx.clone(), + byte_budget: byte_budget.clone(), + cancel: CancellationToken::new(), + degraded: Arc::new(AtomicBool::new(false)), + rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)), + draining: Arc::new(AtomicBool::new(false)), + draining_started_at_epoch_secs: Arc::new(AtomicU64::new(0)), + drain_deadline_epoch_secs: Arc::new(AtomicU64::new(0)), + allow_drain_fallback: Arc::new(AtomicBool::new(false)), + }; + let mut writers = pool.writers.write().await; + let mut registry_registration = pool.registry.prepare_writer_registration().await; + registry_registration.install(writer_id, tx, byte_budget); + writers.push(writer.clone()); + pool.conn_count.fetch_add(1, Ordering::Relaxed); + drop(registry_registration); + drop(writers); + writer + } + + async fn prepared_writer<'a>( + pool: &'a Arc, + writer_id: u64, + writer_dc: i32, + addr: SocketAddr, + ) -> PreparedWriter<'a> { + let open_reservation = pool + .reserve_writer_open( + WriterContour::Active, + WriterOpenIntent::Replacement, + writer_dc, + ) + .await + .expect("replacement open must be admitted"); + let task_registration = pool + .lifecycle + .try_register() + .expect("test pool lifecycle must be open"); + let (tx, rx) = mpsc::channel::(8); + let byte_budget = pool.new_writer_byte_budget(); + PreparedWriter { + writer: MeWriter { + id: writer_id, + addr, + source_ip: addr.ip(), + writer_dc, + generation: pool.current_generation(), + contour: Arc::new(AtomicU8::new(WriterContour::Active.as_u8())), + created_at: Instant::now(), + tx: tx.clone(), + byte_budget: byte_budget.clone(), + cancel: CancellationToken::new(), + degraded: Arc::new(AtomicBool::new(false)), + rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)), + draining: Arc::new(AtomicBool::new(false)), + draining_started_at_epoch_secs: Arc::new(AtomicU64::new(0)), + drain_deadline_epoch_secs: Arc::new(AtomicU64::new(0)), + allow_drain_fallback: Arc::new(AtomicBool::new(false)), + }, + tx, + byte_budget, + task_registration, + writer_task: Box::pin(async move { + drop(rx); + }), + intent: WriterOpenIntent::Replacement, + _open_reservation: open_reservation, + } + } + + #[tokio::test] + async fn replacement_commit_publishes_successor_before_draining_victim() { + let pool = make_pool().await; + let addr = endpoint(1); + pool.preferred_endpoints_by_dc + .store(Arc::new(HashMap::from([(2, vec![addr])]))); + let victim = install_writer(&pool, 1001, 2, addr).await; + let expected_role = WriterRole::from_writer(&victim); + let mut reservation = pool + .registry + .try_reserve_writer_replacement(victim.id) + .await + .expect("idle victim must be reservable"); + let prepared = prepared_writer(&pool, 1002, 2, addr).await; + + pool.publish_prepared_replacement_writer( + prepared, + expected_role, + WriterReplacementPurpose::IdleRefresh, + &mut reservation, + ) + .await + .expect("replacement commit must succeed"); + + let writers = pool.writers.read().await; + assert_eq!(writers.len(), 2); + assert_eq!( + writers + .iter() + .filter(|writer| !writer.draining.load(Ordering::Acquire)) + .count(), + 1 + ); + assert!(victim.draining.load(Ordering::Acquire)); + assert!(writers.iter().any(|writer| writer.id == 1002)); + drop(writers); + assert_eq!(pool.conn_count.load(Ordering::Acquire), 2); + assert_eq!(pool.registry.writer_replacement_counts(), (0, 1)); + } + + #[tokio::test] + async fn cancelled_replacement_waiting_for_publication_restores_all_reservations() { + let pool = make_pool().await; + let addr = endpoint(2); + pool.preferred_endpoints_by_dc + .store(Arc::new(HashMap::from([(2, vec![addr])]))); + let victim = install_writer(&pool, 2001, 2, addr).await; + let expected_role = WriterRole::from_writer(&victim); + let mut reservation = pool + .registry + .try_reserve_writer_replacement(victim.id) + .await + .expect("idle victim must be reservable"); + let prepared = prepared_writer(&pool, 2002, 2, addr).await; + let writers_guard = pool.writers.write().await; + + let result = tokio::time::timeout( + Duration::from_millis(10), + pool.publish_prepared_replacement_writer( + prepared, + expected_role, + WriterReplacementPurpose::IdleRefresh, + &mut reservation, + ), + ) + .await; + + assert!(result.is_err()); + drop(writers_guard); + drop(reservation); + assert_eq!(pool.writer_replacement_open_reserved.load(Ordering::Acquire), 0); + assert_eq!(pool.registry.writer_replacement_counts(), (0, 0)); + assert!(!victim.draining.load(Ordering::Acquire)); + assert!(!pool.writers.read().await.iter().any(|writer| writer.id == 2002)); + } + + #[tokio::test] + async fn floor_rebalance_commit_rejects_a_donor_without_surplus() { + let pool = make_pool().await; + let donor_addr = endpoint(3); + let receiver_addr = endpoint(4); + pool.preferred_endpoints_by_dc.store(Arc::new(HashMap::from([ + (1, vec![donor_addr]), + (2, vec![receiver_addr]), + ]))); + let victim = install_writer(&pool, 3001, 1, donor_addr).await; + let expected_role = WriterRole::from_writer(&victim); + let mut reservation = pool + .registry + .try_reserve_writer_replacement(victim.id) + .await + .expect("idle victim must be reservable"); + let prepared = prepared_writer(&pool, 3002, 2, receiver_addr).await; + + let result = pool + .publish_prepared_replacement_writer( + prepared, + expected_role, + WriterReplacementPurpose::FloorRebalance { + donor_floor: 1, + receiver_floor: 1, + }, + &mut reservation, + ) + .await; + + assert!(result.is_err()); + drop(reservation); + assert!(!victim.draining.load(Ordering::Acquire)); + assert!(!pool.writers.read().await.iter().any(|writer| writer.id == 3002)); + assert_eq!(pool.registry.writer_replacement_counts(), (0, 0)); + } +} diff --git a/src/transport/middle_proxy/pool_writer/runtime.rs b/src/transport/middle_proxy/pool_writer/runtime.rs index 427a197..382f844 100644 --- a/src/transport/middle_proxy/pool_writer/runtime.rs +++ b/src/transport/middle_proxy/pool_writer/runtime.rs @@ -54,23 +54,45 @@ impl MePool { contour: WriterContour, writer_dc: i32, ) -> Result<()> { - self.connect_one_with_generation_contour_for_dc_with_cap_policy( - addr, rng, generation, contour, writer_dc, false, + self.connect_one_with_generation_contour_for_dc_with_intent( + addr, + rng, + generation, + contour, + writer_dc, + WriterOpenIntent::Normal, ) .await } - pub(in crate::transport::middle_proxy) async fn connect_one_with_generation_contour_for_dc_with_cap_policy( + /// Connects and publishes one writer under an explicit capacity intent. + pub(in crate::transport::middle_proxy) async fn connect_one_with_generation_contour_for_dc_with_intent( self: &Arc, addr: SocketAddr, rng: &SecureRandom, generation: u64, contour: WriterContour, writer_dc: i32, - allow_coverage_override: bool, + intent: WriterOpenIntent, ) -> Result<()> { - let Some(_writer_open_reservation) = self - .reserve_writer_open(contour, allow_coverage_override, writer_dc) + let prepared = self + .prepare_writer_with_intent(addr, rng, generation, contour, writer_dc, intent) + .await?; + self.publish_connected_writer(prepared).await + } + + /// Completes all cancellable connection work without publishing pool visibility. + pub(super) async fn prepare_writer_with_intent<'a>( + self: &'a Arc, + addr: SocketAddr, + rng: &SecureRandom, + generation: u64, + contour: WriterContour, + writer_dc: i32, + intent: WriterOpenIntent, + ) -> Result> { + let Some(writer_open_reservation) = self + .reserve_writer_open(contour, intent, writer_dc) .await else { return Err(ProxyError::Proxy(format!( @@ -264,36 +286,15 @@ impl MePool { }) }; - self.publish_prepared_writer(writer, tx, byte_budget, task_registration, writer_task) - .await; - - Ok(()) - } - - /// Commits writer visibility and lifecycle ownership after all cancellation points. - #[allow(clippy::too_many_arguments)] - pub(in crate::transport::middle_proxy) async fn publish_prepared_writer( - self: &Arc, - writer: MeWriter, - tx: mpsc::Sender, - byte_budget: Arc, - task_registration: MeTaskRegistration<'_>, - writer_task: F, - ) where - F: Future + Send + 'static, - { - let (mut writers, mut registry_registration) = tokio::join!( - self.writers.write(), - self.registry.prepare_writer_registration() - ); - registry_registration.install(writer.id, tx, byte_budget); - writers.push(writer); - self.conn_count.fetch_add(1, Ordering::Relaxed); - self.lifecycle - .spawn_registered_writer(task_registration, writer_task); - drop(writers); - drop(registry_registration); - self.notify_writer_epoch(); + Ok(PreparedWriter { + writer, + tx, + byte_budget, + task_registration, + writer_task, + intent, + _open_reservation: writer_open_reservation, + }) } pub(crate) async fn remove_writer_and_close_clients(self: &Arc, writer_id: u64) { @@ -331,7 +332,7 @@ impl MePool { ) -> bool { let mut close_tx: Option> = None; let mut removed_addr: Option = None; - let mut removed_dc: Option = None; + let mut removed_role: Option = None; let mut removed_uptime: Option = None; let mut trigger_refill = false; let mut removed = false; @@ -352,7 +353,7 @@ impl MePool { self.stats.increment_me_writer_removed_total(); w.cancel.cancel(); removed_addr = Some(w.addr); - removed_dc = Some(w.writer_dc); + removed_role = Some(WriterRole::from_writer(&w)); removed_uptime = Some(w.created_at.elapsed()); trigger_refill = !was_draining; if trigger_refill { @@ -397,8 +398,8 @@ impl MePool { ); } } - if trigger_refill && let Some(writer_dc) = removed_dc { - self.trigger_immediate_refill_for_dc(addr, writer_dc); + if trigger_refill && let Some(role) = removed_role { + self.trigger_immediate_refill_for_role(addr, role); } } if removed { @@ -415,26 +416,13 @@ impl MePool { ) { let timeout = timeout.filter(|d| !d.is_zero()); let found = { - let mut ws = self.writers.write().await; - if let Some(w) = ws.iter_mut().find(|w| w.id == writer_id) { - let already_draining = w.draining.swap(true, Ordering::Relaxed); - w.allow_drain_fallback - .store(allow_drain_fallback, Ordering::Relaxed); - let now_epoch_secs = Self::now_epoch_secs(); - w.draining_started_at_epoch_secs - .store(now_epoch_secs, Ordering::Relaxed); - let drain_deadline_epoch_secs = timeout - .map(|duration| now_epoch_secs.saturating_add(duration.as_secs())) - .unwrap_or(0); - w.drain_deadline_epoch_secs - .store(drain_deadline_epoch_secs, Ordering::Relaxed); - if !already_draining { - self.stats.increment_pool_drain_active(); - self.increment_draining_active_runtime(); + let ws = self.writers.write().await; + let mut registry_registration = self.registry.prepare_writer_registration().await; + if let Some(w) = ws.iter().find(|w| w.id == writer_id) { + if !allow_drain_fallback { + registry_registration.retire(writer_id); } - w.contour - .store(WriterContour::Draining.as_u8(), Ordering::Relaxed); - w.draining.store(true, Ordering::Relaxed); + self.apply_writer_draining_state(w, timeout, allow_drain_fallback); true } else { false @@ -452,6 +440,39 @@ impl MePool { ); } + /// Publishes drain metadata before the release-store that makes draining observable. + pub(in crate::transport::middle_proxy) fn apply_writer_draining_state( + &self, + writer: &MeWriter, + timeout: Option, + allow_drain_fallback: bool, + ) -> bool { + let timeout = timeout.filter(|duration| !duration.is_zero()); + let already_draining = writer.draining.load(Ordering::Acquire); + writer + .allow_drain_fallback + .store(allow_drain_fallback, Ordering::Release); + let now_epoch_secs = Self::now_epoch_secs(); + writer + .draining_started_at_epoch_secs + .store(now_epoch_secs, Ordering::Release); + let drain_deadline_epoch_secs = timeout + .map(|duration| now_epoch_secs.saturating_add(duration.as_secs())) + .unwrap_or(0); + writer + .drain_deadline_epoch_secs + .store(drain_deadline_epoch_secs, Ordering::Release); + writer + .contour + .store(WriterContour::Draining.as_u8(), Ordering::Release); + writer.draining.store(true, Ordering::Release); + if !already_draining { + self.stats.increment_pool_drain_active(); + self.increment_draining_active_runtime(); + } + !already_draining + } + pub(crate) async fn mark_writer_draining(self: &Arc, writer_id: u64) { self.mark_writer_draining_with_timeout(writer_id, Some(Duration::from_secs(300)), false) .await; @@ -461,7 +482,7 @@ impl MePool { &self, writer: &MeWriter, ) -> bool { - if !writer.draining.load(Ordering::Relaxed) { + if !writer.draining.load(Ordering::Acquire) { return true; } if !writer.allow_drain_fallback.load(Ordering::Relaxed) { diff --git a/src/transport/middle_proxy/registry.rs b/src/transport/middle_proxy/registry.rs index 62ca021..29e4315 100644 --- a/src/transport/middle_proxy/registry.rs +++ b/src/transport/middle_proxy/registry.rs @@ -20,8 +20,14 @@ const ROUTE_QUEUED_MAX_FRAME_PERMITS: usize = 1024; // Transactional writer registry publication. mod publication; +// Cancellation-safe idle-writer replacement reservations. +mod replacement; mod writer; +pub(in crate::transport::middle_proxy) use replacement::{ + WriterBindOutcome, WriterReplacementReservation, +}; + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum RouteResult { Routed, @@ -74,6 +80,7 @@ struct WriterTable { struct WriterRoute { tx: mpsc::Sender, byte_budget: Arc, + replacement_state: Arc, } #[derive(Clone)] @@ -94,6 +101,15 @@ struct BindingState { last_meta_for_writer: DashMap, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[repr(u8)] +enum WriterReplacementState { + Open = 0, + Preparing, + Retiring, + Draining, +} + struct BindingInner { writer_for_conn: HashMap, conns_for_writer: HashMap>, diff --git a/src/transport/middle_proxy/registry/publication.rs b/src/transport/middle_proxy/registry/publication.rs index 8e29a55..9549c3a 100644 --- a/src/transport/middle_proxy/registry/publication.rs +++ b/src/transport/middle_proxy/registry/publication.rs @@ -1,9 +1,11 @@ use std::sync::Arc; +use std::sync::atomic::{AtomicU8, Ordering}; use tokio::sync::{MutexGuard, Semaphore, mpsc}; use super::super::codec::WriterCommand; -use super::{BindingInner, ConnRegistry, WriterRoute}; +use super::replacement::WriterReplacementReservation; +use super::{BindingInner, ConnRegistry, WriterReplacementState, WriterRoute}; /// Holds registry binding ownership until pool writer visibility is published. pub(in crate::transport::middle_proxy) struct WriterRegistrationGuard<'a> { @@ -24,6 +26,73 @@ impl ConnRegistry { } impl WriterRegistrationGuard<'_> { + /// Prevents new bindings while preserving all existing writer associations. + pub(in crate::transport::middle_proxy) fn retire(&mut self, writer_id: u64) -> bool { + let Some(state) = self + .registry + .writers + .map + .get(&writer_id) + .map(|route| Arc::clone(&route.replacement_state)) + else { + return false; + }; + loop { + let current = state.load(Ordering::Acquire); + if current == WriterReplacementState::Retiring as u8 + || current == WriterReplacementState::Draining as u8 + { + return true; + } + if state + .compare_exchange_weak( + current, + WriterReplacementState::Draining as u8, + Ordering::AcqRel, + Ordering::Acquire, + ) + .is_ok() + { + return true; + } + } + } + + /// Revalidates an idle replacement victim and prevents subsequent client bindings. + pub(in crate::transport::middle_proxy) fn prepare_replacement_commit( + &mut self, + reservation: &WriterReplacementReservation<'_>, + ) -> bool { + let Some(route_state) = self + .registry + .writers + .map + .get(&reservation.writer_id()) + .map(|route| Arc::clone(&route.replacement_state)) + else { + return false; + }; + if !std::ptr::eq(self.registry, reservation.registry()) + || !Arc::ptr_eq(&route_state, reservation.state()) + || reservation.requires_idle() + && self + .binding + .conns_for_writer + .get(&reservation.writer_id()) + .is_none_or(|conn_ids| !conn_ids.is_empty()) + { + return false; + } + route_state + .compare_exchange( + WriterReplacementState::Preparing as u8, + WriterReplacementState::Retiring as u8, + Ordering::AcqRel, + Ordering::Acquire, + ) + .is_ok() + } + /// Installs registry state while retaining binding ownership for pool publication. pub(in crate::transport::middle_proxy) fn install( &mut self, @@ -42,9 +111,13 @@ impl WriterRegistrationGuard<'_> { .writer_idle_since_epoch_secs .entry(writer_id) .or_insert_with(ConnRegistry::now_epoch_secs); - self.registry - .writers - .map - .insert(writer_id, WriterRoute { tx, byte_budget }); + self.registry.writers.map.insert( + writer_id, + WriterRoute { + tx, + byte_budget, + replacement_state: Arc::new(AtomicU8::new(WriterReplacementState::Open as u8)), + }, + ); } } diff --git a/src/transport/middle_proxy/registry/replacement.rs b/src/transport/middle_proxy/registry/replacement.rs new file mode 100644 index 0000000..ae7b5a6 --- /dev/null +++ b/src/transport/middle_proxy/registry/replacement.rs @@ -0,0 +1,142 @@ +use std::sync::Arc; +use std::sync::atomic::{AtomicU8, Ordering}; + +use super::{ConnRegistry, WriterReplacementState}; + +/// Result of atomically committing a client-to-writer binding. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(in crate::transport::middle_proxy) enum WriterBindOutcome { + /// The client route was bound to the selected writer. + Bound, + /// The client route disappeared before the binding commit. + RouteMissing, + /// The selected writer disappeared before the binding commit. + WriterMissing, + /// The selected writer is retiring and no longer accepts new clients. + WriterRetiring, +} + +/// Cancellation-safe ownership of one prospective writer replacement. +pub(in crate::transport::middle_proxy) struct WriterReplacementReservation<'a> { + registry: &'a ConnRegistry, + writer_id: u64, + state: Arc, + require_idle: bool, + committed: bool, +} + +impl WriterReplacementReservation<'_> { + /// Returns the stable identifier of the prospective victim. + pub(in crate::transport::middle_proxy) fn writer_id(&self) -> u64 { + self.writer_id + } + + /// Returns the registry whose binding lock linearizes this reservation. + pub(super) fn registry(&self) -> &ConnRegistry { + self.registry + } + + /// Returns the writer-local replacement state identity captured at reservation time. + pub(super) fn state(&self) -> &Arc { + &self.state + } + + /// Reports whether commit must revalidate that the victim remains unbound. + pub(super) fn requires_idle(&self) -> bool { + self.require_idle + } + + /// Transfers retirement ownership to the published replacement. + pub(in crate::transport::middle_proxy) fn mark_committed(&mut self) { + self.committed = true; + } +} + +impl Drop for WriterReplacementReservation<'_> { + fn drop(&mut self) { + if self.committed { + return; + } + let preparing = WriterReplacementState::Preparing as u8; + let open = WriterReplacementState::Open as u8; + let _ = self + .state + .compare_exchange(preparing, open, Ordering::AcqRel, Ordering::Acquire); + } +} + +impl ConnRegistry { + /// Claims an idle writer as a prospective replacement victim. + /// + /// Preparing prevents duplicate replacement work but intentionally permits new client binds. + /// Commit revalidates idleness under the existing binding lock before blocking future binds. + pub(in crate::transport::middle_proxy) async fn try_reserve_writer_replacement( + &self, + writer_id: u64, + ) -> Option> { + self.reserve_writer_replacement(writer_id, true).await + } + + /// Claims a writer while retaining its existing client associations through replacement. + pub(in crate::transport::middle_proxy) async fn try_reserve_writer_replacement_preserving_clients( + &self, + writer_id: u64, + ) -> Option> { + self.reserve_writer_replacement(writer_id, false).await + } + + async fn reserve_writer_replacement( + &self, + writer_id: u64, + require_idle: bool, + ) -> Option> { + let binding = self.binding.inner.lock().await; + let state = self + .writers + .map + .get(&writer_id) + .map(|route| Arc::clone(&route.replacement_state))?; + if require_idle + && binding + .conns_for_writer + .get(&writer_id) + .is_none_or(|conn_ids| !conn_ids.is_empty()) + { + return None; + } + state + .compare_exchange( + WriterReplacementState::Open as u8, + WriterReplacementState::Preparing as u8, + Ordering::AcqRel, + Ordering::Acquire, + ) + .ok()?; + + Some(WriterReplacementReservation { + registry: self, + writer_id, + state, + require_idle, + committed: false, + }) + } + + /// Returns fixed-cardinality gauges for preparing and retiring replacements. + pub(in crate::transport::middle_proxy) fn writer_replacement_counts(&self) -> (usize, usize) { + let mut preparing = 0usize; + let mut retiring = 0usize; + for route in &self.writers.map { + match route.replacement_state.load(Ordering::Acquire) { + state if state == WriterReplacementState::Preparing as u8 => { + preparing = preparing.saturating_add(1); + } + state if state == WriterReplacementState::Retiring as u8 => { + retiring = retiring.saturating_add(1); + } + _ => {} + } + } + (preparing, retiring) + } +} diff --git a/src/transport/middle_proxy/registry/tests.rs b/src/transport/middle_proxy/registry/tests.rs index c90daf0..99e442f 100644 --- a/src/transport/middle_proxy/registry/tests.rs +++ b/src/transport/middle_proxy/registry/tests.rs @@ -4,13 +4,90 @@ use std::sync::Arc; use bytes::Bytes; use tokio::sync::Semaphore; -use super::{ConnMeta, ConnRegistry, RouteResult}; +use super::{ConnMeta, ConnRegistry, RouteResult, WriterBindOutcome}; use crate::transport::middle_proxy::MeResponse; fn writer_byte_budget() -> Arc { Arc::new(Semaphore::new(2049)) } +fn conn_meta() -> ConnMeta { + let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443); + ConnMeta { + target_dc: 2, + client_addr: addr, + our_addr: addr, + proto_flags: 0, + } +} + +#[tokio::test] +async fn preparing_replacement_allows_bind_and_idle_commit_revalidates() { + let registry = ConnRegistry::new(); + let (writer_tx, _writer_rx) = tokio::sync::mpsc::channel(8); + registry + .register_writer(10, writer_tx, writer_byte_budget()) + .await; + let reservation = registry + .try_reserve_writer_replacement(10) + .await + .expect("idle writer must be reservable"); + let (conn_id, _rx) = registry.register().await; + + assert_eq!( + registry + .bind_writer_with_outcome(conn_id, 10, conn_meta()) + .await, + WriterBindOutcome::Bound + ); + let mut registration = registry.prepare_writer_registration().await; + assert!(!registration.prepare_replacement_commit(&reservation)); + drop(registration); +} + +#[tokio::test] +async fn retiring_replacement_rejects_new_bind() { + let registry = ConnRegistry::new(); + let (writer_tx, _writer_rx) = tokio::sync::mpsc::channel(8); + registry + .register_writer(10, writer_tx, writer_byte_budget()) + .await; + let mut reservation = registry + .try_reserve_writer_replacement(10) + .await + .expect("idle writer must be reservable"); + let mut registration = registry.prepare_writer_registration().await; + assert!(registration.prepare_replacement_commit(&reservation)); + drop(registration); + reservation.mark_committed(); + let (conn_id, _rx) = registry.register().await; + + assert_eq!( + registry + .bind_writer_with_outcome(conn_id, 10, conn_meta()) + .await, + WriterBindOutcome::WriterRetiring + ); +} + +#[tokio::test] +async fn dropped_preparing_reservation_restores_writer_admission() { + let registry = ConnRegistry::new(); + let (writer_tx, _writer_rx) = tokio::sync::mpsc::channel(8); + registry + .register_writer(10, writer_tx, writer_byte_budget()) + .await; + let reservation = registry + .try_reserve_writer_replacement(10) + .await + .expect("idle writer must be reservable"); + assert!(registry.try_reserve_writer_replacement(10).await.is_none()); + + drop(reservation); + + assert!(registry.try_reserve_writer_replacement(10).await.is_some()); +} + #[tokio::test] async fn writer_activity_snapshot_tracks_writer_and_dc_load() { let registry = ConnRegistry::new(); diff --git a/src/transport/middle_proxy/registry/writer.rs b/src/transport/middle_proxy/registry/writer.rs index ca0a2e7..efe18bc 100644 --- a/src/transport/middle_proxy/registry/writer.rs +++ b/src/transport/middle_proxy/registry/writer.rs @@ -12,6 +12,7 @@ use super::{ BoundConn, ConnMeta, ConnRegistry, ConnWriter, HotConnBinding, RouteResult, WriterActivitySnapshot, }; +use super::replacement::WriterBindOutcome; impl ConnRegistry { fn set_writer_bound_count(&self, writer_id: u64, count: usize) { @@ -286,17 +287,30 @@ impl ConnRegistry { } } - pub async fn bind_writer(&self, conn_id: u64, writer_id: u64, meta: ConnMeta) -> bool { + /// Atomically binds one client route while rejecting retiring writer generations. + pub(in crate::transport::middle_proxy) async fn bind_writer_with_outcome( + &self, + conn_id: u64, + writer_id: u64, + meta: ConnMeta, + ) -> WriterBindOutcome { let mut binding = self.binding.inner.lock().await; // ROUTING IS THE SOURCE OF TRUTH: // never keep/attach writer binding for a connection that is already // absent from the routing table. if !self.routing.map.contains_key(&conn_id) { - return false; + return WriterBindOutcome::RouteMissing; } - if !self.writers.map.contains_key(&writer_id) { - return false; + let Some(writer_route) = self.writers.map.get(&writer_id) else { + return WriterBindOutcome::WriterMissing; + }; + let writer_state = writer_route.replacement_state.load(Ordering::Acquire); + if writer_state == super::WriterReplacementState::Retiring as u8 + || writer_state == super::WriterReplacementState::Draining as u8 + { + return WriterBindOutcome::WriterRetiring; } + drop(writer_route); let previous_writer_id = binding.writer_for_conn.insert(conn_id, writer_id); if let Some(previous_writer_id) = previous_writer_id @@ -331,7 +345,15 @@ impl ConnRegistry { self.hot_binding .map .insert(conn_id, HotConnBinding { writer_id, meta }); - true + WriterBindOutcome::Bound + } + + pub async fn bind_writer(&self, conn_id: u64, writer_id: u64, meta: ConnMeta) -> bool { + matches!( + self.bind_writer_with_outcome(conn_id, writer_id, meta) + .await, + WriterBindOutcome::Bound + ) } pub async fn get_last_writer_meta(&self, writer_id: u64) -> Option { diff --git a/src/transport/middle_proxy/send.rs b/src/transport/middle_proxy/send.rs index c74aa74..88e73b0 100644 --- a/src/transport/middle_proxy/send.rs +++ b/src/transport/middle_proxy/send.rs @@ -12,7 +12,7 @@ use tracing::{debug, warn}; use super::MePool; use super::codec::{ProxyReqCommand, WriterBytePermit, WriterCommand}; -use super::registry::ConnMeta; +use super::registry::{ConnMeta, WriterBindOutcome}; use super::wire::{build_proxy_req_payload, proxy_req_payload_len}; use crate::config::defaults::ME_WRITER_BYTE_PERMIT_UNIT_BYTES; use crate::config::{MeRouteNoWriterMode, MeWriterPickMode}; @@ -648,13 +648,30 @@ impl MePool { // Keep the advertised proxy IP aligned with the selected ME writer source. let effective_our_addr = SocketAddr::new(w.source_ip, our_addr.port()); let (payload, meta) = build_routed_payload(effective_our_addr); - if !self.registry.bind_writer(conn_id, w.id, meta).await { + let bind_outcome = self + .registry + .bind_writer_with_outcome(conn_id, w.id, meta) + .await; + if bind_outcome != WriterBindOutcome::Bound { + drop(permit); + if bind_outcome == WriterBindOutcome::WriterRetiring { + debug!( + conn_id, + writer_id = w.id, + "ME writer entered replacement retirement before bind commit" + ); + continue; + } + if bind_outcome == WriterBindOutcome::RouteMissing { + return Err(ProxyError::Proxy( + "ME client route disappeared before writer bind".into(), + )); + } debug!( conn_id, writer_id = w.id, "ME writer disappeared before bind commit, pruning stale writer" ); - drop(permit); self.remove_writer_and_close_clients(w.id).await; continue; } @@ -742,13 +759,30 @@ impl MePool { // Keep the advertised proxy IP aligned with the selected ME writer source. let effective_our_addr = SocketAddr::new(w.source_ip, our_addr.port()); let (payload, meta) = build_routed_payload(effective_our_addr); - if !self.registry.bind_writer(conn_id, w.id, meta).await { + let bind_outcome = self + .registry + .bind_writer_with_outcome(conn_id, w.id, meta) + .await; + if bind_outcome != WriterBindOutcome::Bound { + drop(permit); + if bind_outcome == WriterBindOutcome::WriterRetiring { + debug!( + conn_id, + writer_id = w.id, + "ME writer entered replacement retirement before fallback bind commit" + ); + continue; + } + if bind_outcome == WriterBindOutcome::RouteMissing { + return Err(ProxyError::Proxy( + "ME client route disappeared before writer bind".into(), + )); + } debug!( conn_id, writer_id = w.id, "ME writer disappeared before fallback bind commit, pruning stale writer" ); - drop(permit); self.remove_writer_and_close_clients(w.id).await; continue; } diff --git a/src/transport/middle_proxy/tests/pool_refill_security_tests.rs b/src/transport/middle_proxy/tests/pool_refill_security_tests.rs index a513d1a..0cae5c5 100644 --- a/src/transport/middle_proxy/tests/pool_refill_security_tests.rs +++ b/src/transport/middle_proxy/tests/pool_refill_security_tests.rs @@ -30,7 +30,10 @@ async fn make_pool() -> Arc { HashMap::new(), HashMap::new(), None, - NetworkDecision::default(), + NetworkDecision { + ipv4_me: true, + ..NetworkDecision::default() + }, None, Arc::new(SecureRandom::new()), Arc::new(Stats::default()), @@ -162,12 +165,26 @@ async fn connectable_endpoints_releases_quarantine_lock_before_sleep() { assert_eq!(endpoints, vec![addr]); } +#[tokio::test] +async fn refill_does_not_queue_a_removed_dc_target() { + let pool = make_pool().await; + let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 31, 0, 20)), 443); + + pool.trigger_immediate_refill_for_dc(addr, 2); + + assert!(pool.refill_states.lock().is_empty()); + assert_eq!(pool.refill_running.load(Ordering::Acquire), 0); + assert_eq!(pool.refill_pending.load(Ordering::Acquire), 0); +} + #[tokio::test(flavor = "current_thread")] -async fn refill_coalesces_one_pending_endpoint_and_cleans_up_before_first_poll() { +async fn refill_preserves_bounded_pending_cardinality_and_cleans_up_before_first_poll() { let pool = make_pool().await; let first = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 31, 0, 21)), 443); let second = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 31, 0, 22)), 443); let latest = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 31, 0, 23)), 443); + pool.preferred_endpoints_by_dc + .store(Arc::new(HashMap::from([(2, vec![first, second, latest])]))); pool.trigger_immediate_refill_for_dc(first, 2); pool.trigger_immediate_refill_for_dc(second, 2); @@ -175,7 +192,16 @@ async fn refill_coalesces_one_pending_endpoint_and_cleans_up_before_first_poll() assert_eq!(pool.refill_states.lock().len(), 1); assert_eq!(pool.refill_running.load(Ordering::Acquire), 1); - assert_eq!(pool.refill_pending.load(Ordering::Acquire), 1); + assert_eq!(pool.refill_pending.load(Ordering::Acquire), 2); + let state = pool + .refill_states + .lock() + .values() + .copied() + .next() + .expect("refill state"); + assert_eq!(state.pending_count, 2); + assert_eq!(state.next_addr, Some(latest)); pool.begin_shutdown(); tokio::time::timeout(Duration::from_secs(1), async { diff --git a/src/transport/middle_proxy/tests/pool_writer_publication_tests.rs b/src/transport/middle_proxy/tests/pool_writer_publication_tests.rs index 0eb8f01..4135df1 100644 --- a/src/transport/middle_proxy/tests/pool_writer_publication_tests.rs +++ b/src/transport/middle_proxy/tests/pool_writer_publication_tests.rs @@ -7,10 +7,121 @@ use tokio::sync::mpsc; use tokio_util::sync::CancellationToken; use super::codec::WriterCommand; -use super::pool::{MeWriter, WriterContour}; +use super::pool::{MeWriter, WriterContour, WriterOpenIntent}; use super::pool_writer_security_tests::make_pool; use super::registry::ConnMeta; +fn unregistered_writer( + pool: &Arc, + writer_id: u64, + addr: SocketAddr, + generation: u64, + contour: WriterContour, +) -> MeWriter { + let (tx, _rx) = mpsc::channel::(8); + MeWriter { + id: writer_id, + addr, + source_ip: addr.ip(), + writer_dc: 2, + generation, + contour: Arc::new(AtomicU8::new(contour.as_u8())), + created_at: Instant::now(), + tx, + byte_budget: pool.new_writer_byte_budget(), + cancel: CancellationToken::new(), + degraded: Arc::new(AtomicBool::new(false)), + rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)), + draining: Arc::new(AtomicBool::new(false)), + draining_started_at_epoch_secs: Arc::new(AtomicU64::new(0)), + drain_deadline_epoch_secs: Arc::new(AtomicU64::new(0)), + allow_drain_fallback: Arc::new(AtomicBool::new(false)), + } +} + +#[tokio::test] +async fn normal_warm_publication_cannot_race_past_the_dc_floor() { + let pool = make_pool().await; + let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443); + pool.preferred_endpoints_by_dc + .store(Arc::new(std::collections::HashMap::from([(2, vec![addr])]))); + let generation = 2; + let writers = (1..=3) + .map(|writer_id| { + unregistered_writer(&pool, writer_id, addr, generation, WriterContour::Warm) + }) + .collect::>(); + let candidate = unregistered_writer(&pool, 4, addr, generation, WriterContour::Warm); + + assert!( + pool.authorize_writer_publication_capacity( + &candidate, + WriterContour::Warm, + WriterOpenIntent::Normal, + &writers[..2], + ) + .is_ok() + ); + assert!( + pool.authorize_writer_publication_capacity( + &candidate, + WriterContour::Warm, + WriterOpenIntent::Normal, + &writers, + ) + .is_err() + ); + assert!( + pool.authorize_writer_publication_capacity( + &candidate, + WriterContour::Warm, + WriterOpenIntent::Replacement, + &writers, + ) + .is_ok() + ); +} + +#[tokio::test] +async fn normal_active_publication_cannot_race_past_the_family_floor() { + let pool = make_pool().await; + let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443); + pool.preferred_endpoints_by_dc + .store(Arc::new(std::collections::HashMap::from([(2, vec![addr])]))); + let generation = pool.current_generation(); + let writers = (1..=3) + .map(|writer_id| { + unregistered_writer( + &pool, + writer_id, + addr, + generation, + WriterContour::Active, + ) + }) + .collect::>(); + let candidate = unregistered_writer(&pool, 4, addr, generation, WriterContour::Active); + + assert!( + pool.authorize_writer_publication_capacity( + &candidate, + WriterContour::Active, + WriterOpenIntent::Coverage, + &writers[..2], + ) + .is_ok() + ); + assert!( + pool.authorize_writer_publication_capacity( + &candidate, + WriterContour::Active, + WriterOpenIntent::Coverage, + &writers, + ) + .is_err() + ); +} + #[tokio::test] async fn successful_writer_publication_is_fully_visible_and_removable() { let pool = make_pool().await; diff --git a/src/transport/middle_proxy/tests/pool_writer_security_tests.rs b/src/transport/middle_proxy/tests/pool_writer_security_tests.rs index f468b47..e47524e 100644 --- a/src/transport/middle_proxy/tests/pool_writer_security_tests.rs +++ b/src/transport/middle_proxy/tests/pool_writer_security_tests.rs @@ -35,7 +35,10 @@ pub(super) async fn make_pool() -> Arc { HashMap::new(), HashMap::new(), None, - NetworkDecision::default(), + NetworkDecision { + ipv4_me: true, + ..NetworkDecision::default() + }, None, Arc::new(SecureRandom::new()), Arc::new(Stats::new()), diff --git a/src/transport/middle_proxy/tests/send_adversarial_tests.rs b/src/transport/middle_proxy/tests/send_adversarial_tests.rs index 667eb33..8f78562 100644 --- a/src/transport/middle_proxy/tests/send_adversarial_tests.rs +++ b/src/transport/middle_proxy/tests/send_adversarial_tests.rs @@ -324,6 +324,36 @@ async fn send_proxy_req_does_not_replay_when_first_bind_commit_fails() { assert_eq!(bound.expect("writer should be bound").writer_id, 11); } +#[tokio::test] +async fn missing_client_route_does_not_prune_a_healthy_writer() { + let (pool, _rng) = make_pool().await; + let writer_id = 12; + let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 12)), 443); + let _writer_rx = insert_writer(&pool, writer_id, 2, addr, true).await; + + let result = pool + .send_proxy_req( + 999_999, + 2, + SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 30005), + SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443), + b"cancelled-route", + 0, + None, + None, + ) + .await; + + assert!(result.is_err()); + assert!( + pool.writers + .read() + .await + .iter() + .any(|writer| writer.id == writer_id) + ); +} + #[tokio::test] async fn send_proxy_req_prunes_iterative_stale_bind_failures_without_data_replay() { let (pool, _rng) = make_pool().await;