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