mirror of
https://github.com/telemt/telemt.git
synced 2026-10-08 18:35:58 +03:00
Atomic Lifecycle for Writer Refresh + Hardswap
This commit is contained in:
+3
-171
@@ -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<String>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub(super) struct RuntimeMePoolStateGenerationData {
|
||||
pub(super) active_generation: u64,
|
||||
pub(super) warm_generation: u64,
|
||||
pub(super) warm_generations: Vec<u64>,
|
||||
pub(super) pending_hardswap_generation: u64,
|
||||
pub(super) pending_hardswap_age_secs: Option<u64>,
|
||||
pub(super) reinit_inflight: usize,
|
||||
pub(super) reinit_max_concurrency_effective: usize,
|
||||
pub(super) draining_generations: Vec<u64>,
|
||||
}
|
||||
|
||||
#[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<RuntimeMePoolStateRefillDcData>,
|
||||
}
|
||||
|
||||
#[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<RuntimeMePoolStatePayload>,
|
||||
}
|
||||
|
||||
#[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::<u64>::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;
|
||||
|
||||
@@ -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<u64>,
|
||||
pending_hardswap_generation: u64,
|
||||
pending_hardswap_age_secs: Option<u64>,
|
||||
reinit_inflight: usize,
|
||||
reinit_max_concurrency_effective: usize,
|
||||
draining_generations: Vec<u64>,
|
||||
}
|
||||
|
||||
#[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<bool>,
|
||||
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<RuntimeMePoolStateRefillDcData>,
|
||||
}
|
||||
|
||||
#[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<RuntimeMePoolStatePayload>,
|
||||
}
|
||||
|
||||
/// 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::<u64>::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(),
|
||||
},
|
||||
}),
|
||||
}
|
||||
}
|
||||
@@ -236,6 +236,10 @@ async fn handle<B>(
|
||||
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<B>(
|
||||
tls_cache,
|
||||
tls_full_cert_budget,
|
||||
web_publication,
|
||||
me_hardswap.as_ref(),
|
||||
)
|
||||
.await;
|
||||
let resp = Response::builder()
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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"
|
||||
));
|
||||
}
|
||||
}
|
||||
@@ -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"));
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -31,7 +31,7 @@ pub(in crate::transport::middle_proxy) async fn reap_draining_writers(
|
||||
let mut force_close_writer_ids = Vec::<u64>::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
|
||||
|
||||
@@ -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<u64>>::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::<HashMap<_, _>>(),
|
||||
);
|
||||
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;
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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<i32, usize>,
|
||||
live_writer_ids_by_addr: &HashMap<(i32, SocketAddr), Vec<u64>>,
|
||||
writer_idle_since: &HashMap<u64, u64>,
|
||||
bound_clients_by_writer: &HashMap<u64, usize>,
|
||||
) -> 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::<i32, usize>::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<u64, usize>,
|
||||
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"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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<MePool>, writer_id: u64, writer_dc: i32)
|
||||
pool.conn_count.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
async fn insert_active_writer_at(
|
||||
pool: &Arc<MePool>,
|
||||
writer_id: u64,
|
||||
writer_dc: i32,
|
||||
addr: SocketAddr,
|
||||
) -> MeWriter {
|
||||
let (tx, _writer_rx) = mpsc::channel::<WriterCommand>(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;
|
||||
|
||||
@@ -35,7 +35,7 @@ pub async fn me_zombie_writer_watchdog(pool: Arc<MePool>) {
|
||||
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
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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<SocketAddr>,
|
||||
}
|
||||
|
||||
#[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<Mutex<HashMap<u64, (f64, f64)>>>,
|
||||
pub(super) refill_states: Arc<ParkingMutex<HashMap<RefillDcKey, Option<SocketAddr>>>>,
|
||||
pub(super) refill_states: Arc<ParkingMutex<HashMap<RefillTargetKey, RefillTargetState>>>,
|
||||
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;
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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::<i32, HashSet<SocketAddr>>::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::<SocketAddr>::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::<SocketAddr>::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<WriterOpenReservation<'_>> {
|
||||
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,
|
||||
|
||||
@@ -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<Self>) {
|
||||
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;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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<Self>, pool_size: usize, rng: &Arc<SecureRandom>) -> 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<SecureRandom>,
|
||||
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,
|
||||
|
||||
@@ -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<MePool>,
|
||||
key: RefillDcKey,
|
||||
key: RefillTargetKey,
|
||||
active: bool,
|
||||
}
|
||||
|
||||
impl RefillRunGuard {
|
||||
fn next_or_finish(&mut self) -> Option<SocketAddr> {
|
||||
let mut states = self.pool.refill_states.lock();
|
||||
let next = states.get_mut(&self.key).and_then(Option::take);
|
||||
if next.is_some() {
|
||||
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::<SocketAddr, usize>::new();
|
||||
let mut matching_by_endpoint = HashMap::<SocketAddr, usize>::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<SocketAddr> {
|
||||
async fn endpoints_for_refill_target(&self, target: RefillTargetKey) -> Vec<SocketAddr> {
|
||||
let now_epoch_secs = Self::now_epoch_secs();
|
||||
let mut endpoints = HashSet::<SocketAddr>::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<SocketAddr> = 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::<Vec<_>>();
|
||||
endpoints.sort_unstable();
|
||||
endpoints.dedup();
|
||||
endpoints
|
||||
}
|
||||
|
||||
async fn refill_writer_after_loss(self: &Arc<Self>, 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<Self>,
|
||||
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<Self>,
|
||||
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 {
|
||||
|
||||
@@ -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<i32>,
|
||||
stale_writer_ids: Vec<u64>,
|
||||
force_close_writer_ids: Vec<u64>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
enum ReinitCommitFailure {
|
||||
Superseded,
|
||||
Coverage {
|
||||
coverage_ratio: f32,
|
||||
missing_dc: Vec<i32>,
|
||||
},
|
||||
Redundancy {
|
||||
coverage_ratio: f32,
|
||||
missing_dc: Vec<i32>,
|
||||
},
|
||||
}
|
||||
|
||||
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::<Vec<_>>();
|
||||
if let Some(pending) = state.pending {
|
||||
warm_generations.push(pending.generation);
|
||||
}
|
||||
warm_generations.sort_unstable();
|
||||
warm_generations.dedup();
|
||||
let pending = state.pending;
|
||||
|
||||
@@ -1,7 +1,10 @@
|
||||
use super::*;
|
||||
|
||||
impl MePool {
|
||||
pub(super) fn desired_map_hash(desired_by_dc: &HashMap<i32, HashSet<SocketAddr>>) -> 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<i32, HashSet<SocketAddr>>,
|
||||
) -> u64 {
|
||||
let mut hasher = DefaultHasher::new();
|
||||
let mut dcs: Vec<i32> = 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<Self>,
|
||||
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<i32, HashSet<SocketAddr>>,
|
||||
min_ratio: f32,
|
||||
) -> std::result::Result<ReinitCommitOutcome, ReinitCommitFailure> {
|
||||
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::<HashSet<_>>();
|
||||
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::<HashSet<_>>();
|
||||
let missing_dc_set = missing_dc.iter().copied().collect::<HashSet<_>>();
|
||||
let mut stale_writer_ids = Vec::<u64>::new();
|
||||
let mut force_close_writer_ids = Vec::<u64>::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<i32, HashSet<SocketAddr>>,
|
||||
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<Self>, 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<i32, HashSet<SocketAddr>> {
|
||||
/// Returns the currently authoritative endpoint set for drain and coverage decisions.
|
||||
pub(in crate::transport::middle_proxy) async fn desired_dc_endpoints(
|
||||
&self,
|
||||
) -> HashMap<i32, HashSet<SocketAddr>> {
|
||||
let now_epoch_secs = Self::now_epoch_secs();
|
||||
let mut out: HashMap<i32, HashSet<SocketAddr>> = 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<i32> = 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::<i32>::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<SocketAddr>,
|
||||
) -> 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()
|
||||
|
||||
@@ -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<u64> = 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
|
||||
}
|
||||
|
||||
@@ -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<MePool>,
|
||||
writer_id: u64,
|
||||
writer_dc: i32,
|
||||
endpoint: SocketAddr,
|
||||
generation: u64,
|
||||
contour: WriterContour,
|
||||
) -> MeWriter {
|
||||
let (tx, _rx) = mpsc::channel::<WriterCommand>(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<i32, HashSet<SocketAddr>> {
|
||||
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
|
||||
);
|
||||
}
|
||||
|
||||
@@ -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::<Vec<RefillDcKey>>();
|
||||
let inflight_dc_keys = refill_states
|
||||
.keys()
|
||||
.copied()
|
||||
.collect::<Vec<RefillTargetKey>>();
|
||||
drop(refill_states);
|
||||
|
||||
let mut by_dc_map = HashMap::<(i16, &'static str), usize>::new();
|
||||
|
||||
@@ -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<bool>,
|
||||
/// 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,
|
||||
|
||||
@@ -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<u64>,
|
||||
/// 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<bool>,
|
||||
/// 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::<i32, usize>::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,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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::<MeApiQuarantinedEndpointSnapshot>::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
|
||||
|
||||
@@ -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::<i16, usize>::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::<i16, usize>::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 {
|
||||
|
||||
@@ -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<WriterCommand>,
|
||||
byte_budget: Arc<tokio::sync::Semaphore>,
|
||||
task_registration: MeTaskRegistration<'a>,
|
||||
writer_task: Pin<Box<dyn Future<Output = ()> + 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;
|
||||
|
||||
@@ -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<Self>,
|
||||
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<WriterContour> {
|
||||
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<F>(
|
||||
self: &Arc<Self>,
|
||||
writer: MeWriter,
|
||||
tx: mpsc::Sender<WriterCommand>,
|
||||
byte_budget: Arc<tokio::sync::Semaphore>,
|
||||
task_registration: MeTaskRegistration<'_>,
|
||||
writer_task: F,
|
||||
) where
|
||||
F: Future<Output = ()> + 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();
|
||||
}
|
||||
}
|
||||
@@ -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<Self>,
|
||||
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<Self>,
|
||||
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<MePool>,
|
||||
writer_id: u64,
|
||||
writer_dc: i32,
|
||||
addr: SocketAddr,
|
||||
) -> MeWriter {
|
||||
let (tx, _rx) = mpsc::channel::<WriterCommand>(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<MePool>,
|
||||
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::<WriterCommand>(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));
|
||||
}
|
||||
}
|
||||
@@ -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<Self>,
|
||||
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<Self>,
|
||||
addr: SocketAddr,
|
||||
rng: &SecureRandom,
|
||||
generation: u64,
|
||||
contour: WriterContour,
|
||||
writer_dc: i32,
|
||||
intent: WriterOpenIntent,
|
||||
) -> Result<PreparedWriter<'a>> {
|
||||
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<F>(
|
||||
self: &Arc<Self>,
|
||||
writer: MeWriter,
|
||||
tx: mpsc::Sender<WriterCommand>,
|
||||
byte_budget: Arc<tokio::sync::Semaphore>,
|
||||
task_registration: MeTaskRegistration<'_>,
|
||||
writer_task: F,
|
||||
) where
|
||||
F: Future<Output = ()> + 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<Self>, writer_id: u64) {
|
||||
@@ -331,7 +332,7 @@ impl MePool {
|
||||
) -> bool {
|
||||
let mut close_tx: Option<mpsc::Sender<WriterCommand>> = None;
|
||||
let mut removed_addr: Option<SocketAddr> = None;
|
||||
let mut removed_dc: Option<i32> = None;
|
||||
let mut removed_role: Option<WriterRole> = None;
|
||||
let mut removed_uptime: Option<Duration> = 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<Duration>,
|
||||
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<Self>, 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) {
|
||||
|
||||
@@ -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<WriterCommand>,
|
||||
byte_budget: Arc<Semaphore>,
|
||||
replacement_state: Arc<AtomicU8>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
@@ -94,6 +101,15 @@ struct BindingState {
|
||||
last_meta_for_writer: DashMap<u64, ConnMeta>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
#[repr(u8)]
|
||||
enum WriterReplacementState {
|
||||
Open = 0,
|
||||
Preparing,
|
||||
Retiring,
|
||||
Draining,
|
||||
}
|
||||
|
||||
struct BindingInner {
|
||||
writer_for_conn: HashMap<u64, u64>,
|
||||
conns_for_writer: HashMap<u64, HashSet<u64>>,
|
||||
|
||||
@@ -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)),
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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<AtomicU8>,
|
||||
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<AtomicU8> {
|
||||
&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<WriterReplacementReservation<'_>> {
|
||||
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<WriterReplacementReservation<'_>> {
|
||||
self.reserve_writer_replacement(writer_id, false).await
|
||||
}
|
||||
|
||||
async fn reserve_writer_replacement(
|
||||
&self,
|
||||
writer_id: u64,
|
||||
require_idle: bool,
|
||||
) -> Option<WriterReplacementReservation<'_>> {
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -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<Semaphore> {
|
||||
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();
|
||||
|
||||
@@ -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<ConnMeta> {
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
@@ -30,7 +30,10 @@ async fn make_pool() -> Arc<MePool> {
|
||||
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 {
|
||||
|
||||
@@ -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<super::pool::MePool>,
|
||||
writer_id: u64,
|
||||
addr: SocketAddr,
|
||||
generation: u64,
|
||||
contour: WriterContour,
|
||||
) -> MeWriter {
|
||||
let (tx, _rx) = mpsc::channel::<WriterCommand>(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::<Vec<_>>();
|
||||
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::<Vec<_>>();
|
||||
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;
|
||||
|
||||
@@ -35,7 +35,10 @@ pub(super) async fn make_pool() -> Arc<MePool> {
|
||||
HashMap::new(),
|
||||
HashMap::new(),
|
||||
None,
|
||||
NetworkDecision::default(),
|
||||
NetworkDecision {
|
||||
ipv4_me: true,
|
||||
..NetworkDecision::default()
|
||||
},
|
||||
None,
|
||||
Arc::new(SecureRandom::new()),
|
||||
Arc::new(Stats::new()),
|
||||
|
||||
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user