mirror of
https://github.com/telemt/telemt.git
synced 2026-09-29 05:55:58 +03:00
ME Authority + Quota Resets + WEB Replacement Rollback fixes
This commit is contained in:
+2
-3
@@ -104,14 +104,13 @@ impl QuotaStateOwner {
|
||||
used_bytes: 0,
|
||||
last_reset_epoch_secs,
|
||||
};
|
||||
let reset_target = self.store.current_or_legacy_handle(user);
|
||||
let state = self.state_for_users(configured_users, Some((user, prospective.clone())));
|
||||
let path = self.path.clone();
|
||||
let store = Arc::clone(&self.store);
|
||||
let user = user.to_string();
|
||||
let task = tokio::task::spawn_blocking(move || {
|
||||
let _guard = guard;
|
||||
write_state_file_blocking(&path, &state)?;
|
||||
Ok(store.reset(&user, last_reset_epoch_secs))
|
||||
Ok(reset_target.reset(last_reset_epoch_secs))
|
||||
});
|
||||
wait_for_blocking_io(task).await
|
||||
}
|
||||
|
||||
@@ -213,12 +213,7 @@ impl QuotaStore {
|
||||
}
|
||||
|
||||
pub(crate) fn reset(&self, user: &str, now_epoch_secs: u64) -> UserQuotaSnapshot {
|
||||
let state = self.current_or_legacy_handle(user);
|
||||
state.counters.replace(0, now_epoch_secs);
|
||||
UserQuotaSnapshot {
|
||||
used_bytes: 0,
|
||||
last_reset_epoch_secs: now_epoch_secs,
|
||||
}
|
||||
self.current_or_legacy_handle(user).reset(now_epoch_secs)
|
||||
}
|
||||
|
||||
pub(crate) fn remove(&self, user: &str) {
|
||||
@@ -359,6 +354,15 @@ impl UserQuotaHandle {
|
||||
) -> Result<QuotaReservation, QuotaReserveError> {
|
||||
self.counters.try_reserve(bytes, limit)
|
||||
}
|
||||
|
||||
/// Resets only the quota incarnation captured by this handle.
|
||||
pub(crate) fn reset(&self, now_epoch_secs: u64) -> UserQuotaSnapshot {
|
||||
self.counters.replace(0, now_epoch_secs);
|
||||
UserQuotaSnapshot {
|
||||
used_bytes: 0,
|
||||
last_reset_epoch_secs: now_epoch_secs,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl QuotaReservation {
|
||||
@@ -497,6 +501,22 @@ mod tests {
|
||||
assert_eq!(current.used(), 40);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn captured_reset_handle_cannot_reset_a_new_incarnation() {
|
||||
let store = QuotaStore::default();
|
||||
store.activate_fresh("alice", 1);
|
||||
let reset_target = store.handle_exact("alice", 1).unwrap();
|
||||
reset_target.charge(40);
|
||||
store.advance_preserving_usage("alice", 2);
|
||||
let current = store.handle_exact("alice", 2).unwrap();
|
||||
current.charge(20);
|
||||
|
||||
reset_target.reset(7);
|
||||
|
||||
assert_eq!(reset_target.used(), 0);
|
||||
assert_eq!(current.used(), 60);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stale_retirement_cannot_remove_newer_quota_owner() {
|
||||
let store = QuotaStore::default();
|
||||
|
||||
@@ -26,6 +26,7 @@ pub(super) async fn check_family(
|
||||
|
||||
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,
|
||||
@@ -253,6 +254,7 @@ pub(super) async fn check_family(
|
||||
dc,
|
||||
family,
|
||||
generation: pool.current_generation(),
|
||||
endpoint_revision,
|
||||
contour: WriterContour::Active,
|
||||
})
|
||||
.await
|
||||
|
||||
@@ -37,6 +37,8 @@ pub(super) struct RefillTargetKey {
|
||||
pub family: IpFamily,
|
||||
/// Generation that retains publication authority.
|
||||
pub generation: u64,
|
||||
/// Endpoint snapshot revision targeted by this refill producer.
|
||||
pub endpoint_revision: u64,
|
||||
/// Lifecycle contour that the replacement must preserve.
|
||||
pub contour: WriterContour,
|
||||
}
|
||||
@@ -310,6 +312,7 @@ pub(super) struct ReinitStatusSnapshot {
|
||||
pub(super) pending_hardswap_generation: u64,
|
||||
pub(super) pending_hardswap_started_at_epoch_secs: u64,
|
||||
pub(super) pending_hardswap_map_hash: u64,
|
||||
pub(super) pending_hardswap_endpoint_revision: u64,
|
||||
pub(super) inflight: usize,
|
||||
}
|
||||
|
||||
|
||||
@@ -143,6 +143,7 @@ impl MePool {
|
||||
pending_hardswap_generation: 0,
|
||||
pending_hardswap_started_at_epoch_secs: 0,
|
||||
pending_hardswap_map_hash: 0,
|
||||
pending_hardswap_endpoint_revision: 0,
|
||||
inflight: 0,
|
||||
};
|
||||
stats.set_me_writer_byte_budget_limit_bytes(me_writer_byte_budget_bytes);
|
||||
|
||||
@@ -1,5 +1,37 @@
|
||||
use super::*;
|
||||
|
||||
impl EndpointSnapshot {
|
||||
/// Returns authoritative configured endpoints for one exact DC and address family.
|
||||
pub(in crate::transport::middle_proxy) fn endpoints_for_dc_family(
|
||||
&self,
|
||||
dc: i32,
|
||||
family: IpFamily,
|
||||
) -> &[(IpAddr, u16)] {
|
||||
match family {
|
||||
IpFamily::V4 => self.map_v4.get(&dc),
|
||||
IpFamily::V6 => self.map_v6.get(&dc),
|
||||
}
|
||||
.map(Vec::as_slice)
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
/// Checks control-plane membership without applying data-plane family preference.
|
||||
pub(in crate::transport::middle_proxy) fn contains_dc_endpoint(
|
||||
&self,
|
||||
dc: i32,
|
||||
endpoint: SocketAddr,
|
||||
) -> bool {
|
||||
let family = if endpoint.is_ipv4() {
|
||||
IpFamily::V4
|
||||
} else {
|
||||
IpFamily::V6
|
||||
};
|
||||
self.endpoints_for_dc_family(dc, family)
|
||||
.iter()
|
||||
.any(|(ip, port)| *ip == endpoint.ip() && *port == endpoint.port())
|
||||
}
|
||||
}
|
||||
|
||||
impl MePool {
|
||||
pub(in crate::transport::middle_proxy) fn single_endpoint_outage_mode_enabled(&self) -> bool {
|
||||
self.single_endpoint_runtime
|
||||
|
||||
@@ -65,7 +65,7 @@ impl MePool {
|
||||
pub(in crate::transport::middle_proxy) async fn active_coverage_required_total(&self) -> usize {
|
||||
let now_epoch_secs = Self::now_epoch_secs();
|
||||
let mut required_total = 0usize;
|
||||
let endpoint_snapshot = self.endpoint_snapshot.load();
|
||||
let endpoint_snapshot = self.endpoint_snapshot.load_full();
|
||||
|
||||
if self.family_enabled_for_drain_coverage(IpFamily::V4, now_epoch_secs) {
|
||||
for addrs in endpoint_snapshot.map_v4.values() {
|
||||
@@ -100,11 +100,21 @@ impl MePool {
|
||||
contour: WriterContour,
|
||||
intent: WriterOpenIntent,
|
||||
writer_dc: i32,
|
||||
family: IpFamily,
|
||||
target_addr: SocketAddr,
|
||||
) -> bool {
|
||||
if intent == WriterOpenIntent::Replacement {
|
||||
return true;
|
||||
}
|
||||
let family = if target_addr.is_ipv4() {
|
||||
IpFamily::V4
|
||||
} else {
|
||||
IpFamily::V6
|
||||
};
|
||||
let endpoint_snapshot = self.endpoint_snapshot.load_full();
|
||||
let endpoints = endpoint_snapshot.endpoints_for_dc_family(writer_dc, family);
|
||||
if !endpoint_snapshot.contains_dc_endpoint(writer_dc, target_addr) {
|
||||
return false;
|
||||
}
|
||||
let (active_writers, warm_writers, _) = self.non_draining_writer_counts_by_contour().await;
|
||||
let live = match contour {
|
||||
WriterContour::Active => active_writers,
|
||||
@@ -123,13 +133,7 @@ impl MePool {
|
||||
return false;
|
||||
}
|
||||
|
||||
let endpoint_snapshot = self.endpoint_snapshot.load();
|
||||
let endpoint_count = match family {
|
||||
IpFamily::V4 => endpoint_snapshot.map_v4.get(&writer_dc),
|
||||
IpFamily::V6 => endpoint_snapshot.map_v6.get(&writer_dc),
|
||||
}
|
||||
.map(Vec::len)
|
||||
.unwrap_or(0);
|
||||
let endpoint_count = endpoints.len();
|
||||
if endpoint_count == 0 {
|
||||
return false;
|
||||
}
|
||||
@@ -150,6 +154,7 @@ impl MePool {
|
||||
&& writer.generation == generation
|
||||
&& WriterContour::from_u8(writer.contour.load(Ordering::Relaxed)) == contour
|
||||
&& writer.addr.is_ipv4() == (family == IpFamily::V4)
|
||||
&& endpoint_snapshot.contains_dc_endpoint(writer_dc, writer.addr)
|
||||
})
|
||||
.count()
|
||||
};
|
||||
@@ -157,7 +162,7 @@ impl MePool {
|
||||
return true;
|
||||
}
|
||||
|
||||
live < self.active_coverage_required_total().await
|
||||
false
|
||||
}
|
||||
|
||||
/// Reserves bounded transient capacity for a writer open attempt.
|
||||
@@ -166,7 +171,7 @@ impl MePool {
|
||||
contour: WriterContour,
|
||||
intent: WriterOpenIntent,
|
||||
writer_dc: i32,
|
||||
family: IpFamily,
|
||||
target_addr: SocketAddr,
|
||||
) -> Option<WriterOpenReservation<'_>> {
|
||||
let counter = match contour {
|
||||
WriterContour::Active => &self.writer_connect_active_reserved,
|
||||
@@ -221,7 +226,7 @@ impl MePool {
|
||||
|
||||
loop {
|
||||
if !self
|
||||
.can_open_writer_for_contour(contour, intent, writer_dc, family)
|
||||
.can_open_writer_for_contour(contour, intent, writer_dc, target_addr)
|
||||
.await
|
||||
{
|
||||
return None;
|
||||
|
||||
@@ -26,6 +26,13 @@ enum RefillOutcome {
|
||||
Obsolete,
|
||||
}
|
||||
|
||||
fn refill_open_intent(contour: WriterContour) -> WriterOpenIntent {
|
||||
match contour {
|
||||
WriterContour::Active | WriterContour::Warm => WriterOpenIntent::Coverage,
|
||||
WriterContour::Draining => WriterOpenIntent::Normal,
|
||||
}
|
||||
}
|
||||
|
||||
struct RefillRunGuard {
|
||||
pool: Arc<MePool>,
|
||||
key: RefillTargetKey,
|
||||
@@ -293,18 +300,15 @@ impl MePool {
|
||||
&& target.generation == status.pending_hardswap_generation,
|
||||
WriterContour::Draining => false,
|
||||
};
|
||||
let pending_revision_matches = target.contour != WriterContour::Warm
|
||||
|| target.endpoint_revision == status.pending_hardswap_endpoint_revision;
|
||||
let endpoint_snapshot = self.endpoint_snapshot.load();
|
||||
role_is_authoritative
|
||||
&& self
|
||||
.endpoint_snapshot
|
||||
.load()
|
||||
.preferred_endpoints_by_dc
|
||||
.get(&target.dc)
|
||||
.is_some_and(|endpoints| {
|
||||
endpoints.iter().any(|endpoint| match target.family {
|
||||
IpFamily::V4 => endpoint.is_ipv4(),
|
||||
IpFamily::V6 => endpoint.is_ipv6(),
|
||||
})
|
||||
})
|
||||
&& pending_revision_matches
|
||||
&& target.endpoint_revision == endpoint_snapshot.revision
|
||||
&& !endpoint_snapshot
|
||||
.endpoints_for_dc_family(target.dc, target.family)
|
||||
.is_empty()
|
||||
}
|
||||
|
||||
async fn refill_writer_after_loss(
|
||||
@@ -315,11 +319,7 @@ impl MePool {
|
||||
if !self.refill_target_is_authoritative(target) {
|
||||
return RefillOutcome::Obsolete;
|
||||
}
|
||||
let open_intent = if target.contour == WriterContour::Active {
|
||||
WriterOpenIntent::Coverage
|
||||
} else {
|
||||
WriterOpenIntent::Normal
|
||||
};
|
||||
let open_intent = refill_open_intent(target.contour);
|
||||
let fast_retries = self.reconnect_runtime.me_reconnect_fast_retry_count.max(1);
|
||||
let mut total_attempts = 0u32;
|
||||
let same_endpoint_quarantined = self.is_endpoint_quarantined(addr).await;
|
||||
@@ -455,6 +455,7 @@ impl MePool {
|
||||
dc: role.dc,
|
||||
family: role.family,
|
||||
generation: role.generation,
|
||||
endpoint_revision: self.endpoint_snapshot.load().revision,
|
||||
contour: role.contour,
|
||||
};
|
||||
if !self.refill_target_is_authoritative(target) {
|
||||
@@ -517,3 +518,16 @@ impl MePool {
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn warm_refill_preserves_required_floor_coverage() {
|
||||
assert_eq!(
|
||||
refill_open_intent(WriterContour::Warm),
|
||||
WriterOpenIntent::Coverage
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -24,6 +24,8 @@ mod coordination;
|
||||
// Generation warmup and stale-writer reconciliation.
|
||||
mod reconcile;
|
||||
|
||||
#[cfg(test)]
|
||||
mod dual_family_tests;
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
const ME_HARDSWAP_PENDING_TTL_SECS: u64 = 1800;
|
||||
@@ -115,6 +117,8 @@ fn publish_reinit_state(reinit: &ReinitCore, state: &ReinitCoordinatorState) {
|
||||
pending_hardswap_started_at_epoch_secs: pending
|
||||
.map_or(0, |value| value.started_at_epoch_secs),
|
||||
pending_hardswap_map_hash: pending.map_or(0, |value| value.map_hash),
|
||||
pending_hardswap_endpoint_revision: pending
|
||||
.map_or(0, |value| value.endpoint_revision),
|
||||
inflight: state.attempts.len(),
|
||||
};
|
||||
reinit
|
||||
|
||||
@@ -388,7 +388,7 @@ impl MePool {
|
||||
}
|
||||
|
||||
/// Projects desired per-DC endpoint sets from one immutable endpoint revision.
|
||||
pub(super) fn desired_dc_endpoints_from_snapshot(
|
||||
pub(in crate::transport::middle_proxy) fn desired_dc_endpoints_from_snapshot(
|
||||
&self,
|
||||
endpoint_snapshot: &EndpointSnapshot,
|
||||
) -> HashMap<i32, HashSet<SocketAddr>> {
|
||||
@@ -424,7 +424,6 @@ impl MePool {
|
||||
let active_generation = state.active_generation;
|
||||
let pending_generation = state.pending.map(|pending| pending.generation);
|
||||
let endpoint_snapshot = self.endpoint_snapshot.load();
|
||||
let preferred = &endpoint_snapshot.preferred_endpoints_by_dc;
|
||||
let now_epoch_secs = Self::now_epoch_secs();
|
||||
let mut changed = 0usize;
|
||||
|
||||
@@ -440,9 +439,7 @@ impl MePool {
|
||||
};
|
||||
let endpoint_is_current = self
|
||||
.family_enabled_for_drain_coverage(family, now_epoch_secs)
|
||||
&& preferred
|
||||
.get(&writer.writer_dc)
|
||||
.is_some_and(|endpoints| endpoints.contains(&writer.addr));
|
||||
&& endpoint_snapshot.contains_dc_endpoint(writer.writer_dc, writer.addr);
|
||||
if contour == WriterContour::Warm
|
||||
&& writer.generation == active_generation
|
||||
&& endpoint_is_current
|
||||
|
||||
@@ -0,0 +1,129 @@
|
||||
use std::collections::HashMap;
|
||||
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
|
||||
use std::sync::atomic::Ordering;
|
||||
|
||||
use super::tests::{insert_writer, insert_writer_floor};
|
||||
use crate::network::probe::NetworkDecision;
|
||||
use crate::transport::middle_proxy::pool::WriterContour;
|
||||
use crate::transport::middle_proxy::pool_writer_security_tests::make_pool_with_decision;
|
||||
|
||||
#[tokio::test]
|
||||
async fn nonpreferred_family_warm_generation_retains_authority_and_commits() {
|
||||
let pool = make_pool_with_decision(NetworkDecision {
|
||||
ipv4_me: true,
|
||||
ipv6_me: true,
|
||||
effective_prefer: 4,
|
||||
effective_multipath: false,
|
||||
..NetworkDecision::default()
|
||||
})
|
||||
.await;
|
||||
let v4 = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443);
|
||||
let v6 = SocketAddr::new(IpAddr::V6(Ipv6Addr::LOCALHOST), 443);
|
||||
pool.update_proxy_maps(
|
||||
HashMap::from([(2, vec![(v4.ip(), v4.port())])]),
|
||||
Some(HashMap::from([(2, vec![(v6.ip(), v6.port())])])),
|
||||
)
|
||||
.await;
|
||||
let desired = pool.desired_dc_endpoints().await;
|
||||
let map_hash = super::MePool::desired_map_hash(&desired);
|
||||
let endpoint_revision = pool.endpoint_snapshot.load().revision;
|
||||
let reservation = pool
|
||||
.reserve_reinit_attempt(true, map_hash, endpoint_revision, 100)
|
||||
.expect("current endpoint revision must admit the hardswap");
|
||||
assert!(pool.hardswap_warmup_is_authoritative(
|
||||
reservation.attempt.generation,
|
||||
map_hash,
|
||||
endpoint_revision,
|
||||
));
|
||||
let generation = reservation.attempt.generation;
|
||||
let fresh_v4 = insert_writer_floor(&pool, 10, 2, v4, generation, WriterContour::Warm).await;
|
||||
let fresh_v6 = insert_writer_floor(&pool, 20, 2, v6, generation, WriterContour::Warm).await;
|
||||
let fresh_v4_media =
|
||||
insert_writer_floor(&pool, 30, -2, v4, generation, WriterContour::Warm).await;
|
||||
let fresh_v6_media =
|
||||
insert_writer_floor(&pool, 40, -2, v6, generation, WriterContour::Warm).await;
|
||||
|
||||
assert_eq!(pool.reconcile_writer_generation_roles().await, 0);
|
||||
for writer in fresh_v4
|
||||
.iter()
|
||||
.chain(&fresh_v6)
|
||||
.chain(&fresh_v4_media)
|
||||
.chain(&fresh_v6_media)
|
||||
{
|
||||
assert!(!writer.draining.load(Ordering::Acquire));
|
||||
}
|
||||
|
||||
let outcome = pool
|
||||
.commit_reinit_attempt(&reservation.attempt, &desired, 1.0)
|
||||
.await
|
||||
.expect("full dual-family floor must commit without multipath selection");
|
||||
|
||||
assert!(outcome.missing_groups.is_empty());
|
||||
assert_eq!(pool.current_generation(), generation);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn endpoint_revision_fences_pending_hardswap_after_map_aba() {
|
||||
let pool = make_pool_with_decision(NetworkDecision {
|
||||
ipv4_me: true,
|
||||
effective_prefer: 4,
|
||||
..NetworkDecision::default()
|
||||
})
|
||||
.await;
|
||||
let endpoint_a = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 10)), 443);
|
||||
let endpoint_b = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 11)), 443);
|
||||
pool.update_proxy_maps(
|
||||
HashMap::from([(2, vec![(endpoint_a.ip(), endpoint_a.port())])]),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
let desired_a = pool.desired_dc_endpoints().await;
|
||||
let map_hash = super::MePool::desired_map_hash(&desired_a);
|
||||
let endpoint_revision = pool.endpoint_snapshot.load().revision;
|
||||
let reservation = pool
|
||||
.reserve_reinit_attempt(true, map_hash, endpoint_revision, 100)
|
||||
.expect("current endpoint revision must admit the hardswap");
|
||||
assert!(pool.hardswap_warmup_is_authoritative(
|
||||
reservation.attempt.generation,
|
||||
map_hash,
|
||||
endpoint_revision,
|
||||
));
|
||||
let warm = insert_writer(
|
||||
&pool,
|
||||
100,
|
||||
2,
|
||||
endpoint_a,
|
||||
reservation.attempt.generation,
|
||||
WriterContour::Warm,
|
||||
)
|
||||
.await;
|
||||
|
||||
pool.update_proxy_maps(
|
||||
HashMap::from([(2, vec![(endpoint_b.ip(), endpoint_b.port())])]),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
assert!(!pool.hardswap_warmup_is_authoritative(
|
||||
reservation.attempt.generation,
|
||||
map_hash,
|
||||
endpoint_revision,
|
||||
));
|
||||
pool.update_proxy_maps(
|
||||
HashMap::from([(2, vec![(endpoint_a.ip(), endpoint_a.port())])]),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
assert!(!pool.hardswap_warmup_is_authoritative(
|
||||
reservation.attempt.generation,
|
||||
map_hash,
|
||||
endpoint_revision,
|
||||
));
|
||||
|
||||
let snapshot = pool.api_hardswap_snapshot().await;
|
||||
|
||||
assert!(snapshot.pending);
|
||||
assert_eq!(snapshot.pending_map_current, Some(false));
|
||||
assert_eq!(snapshot.pending_writers_current, 0);
|
||||
assert_eq!(snapshot.orphan_warm_writers_current, 1);
|
||||
assert!(!warm.draining.load(Ordering::Acquire));
|
||||
}
|
||||
@@ -1,10 +1,29 @@
|
||||
use super::*;
|
||||
|
||||
impl MePool {
|
||||
/// Checks the exact pending-generation tuple before starting more warmup work.
|
||||
pub(super) fn hardswap_warmup_is_authoritative(
|
||||
&self,
|
||||
generation: u64,
|
||||
map_hash: u64,
|
||||
endpoint_revision: u64,
|
||||
) -> bool {
|
||||
let state = self.reinit.coordinator.lock();
|
||||
state.desired_map_hash == map_hash
|
||||
&& state.endpoint_revision == endpoint_revision
|
||||
&& state.pending.is_some_and(|pending| {
|
||||
pending.generation == generation
|
||||
&& pending.map_hash == map_hash
|
||||
&& pending.endpoint_revision == endpoint_revision
|
||||
})
|
||||
}
|
||||
|
||||
async fn warmup_generation_for_all_dcs(
|
||||
self: &Arc<Self>,
|
||||
rng: &SecureRandom,
|
||||
generation: u64,
|
||||
map_hash: u64,
|
||||
endpoint_revision: u64,
|
||||
desired_by_dc: &HashMap<i32, HashSet<SocketAddr>>,
|
||||
) {
|
||||
let extra_passes = self
|
||||
@@ -15,6 +34,13 @@ impl MePool {
|
||||
let total_passes = 1 + extra_passes;
|
||||
|
||||
for (dc, endpoints) in desired_by_dc {
|
||||
if !self.hardswap_warmup_is_authoritative(
|
||||
generation,
|
||||
map_hash,
|
||||
endpoint_revision,
|
||||
) {
|
||||
return;
|
||||
}
|
||||
for family in [IpFamily::V4, IpFamily::V6] {
|
||||
let family_endpoints = endpoints
|
||||
.iter()
|
||||
@@ -53,8 +79,22 @@ impl MePool {
|
||||
);
|
||||
|
||||
for attempt_idx in 0..missing {
|
||||
if !self.hardswap_warmup_is_authoritative(
|
||||
generation,
|
||||
map_hash,
|
||||
endpoint_revision,
|
||||
) {
|
||||
return;
|
||||
}
|
||||
let delay_ms = self.hardswap_warmup_connect_delay_ms();
|
||||
tokio::time::sleep(Duration::from_millis(delay_ms)).await;
|
||||
if !self.hardswap_warmup_is_authoritative(
|
||||
generation,
|
||||
map_hash,
|
||||
endpoint_revision,
|
||||
) {
|
||||
return;
|
||||
}
|
||||
|
||||
let connected = self
|
||||
.connect_endpoints_round_robin_with_generation_contour(
|
||||
@@ -100,6 +140,13 @@ impl MePool {
|
||||
}
|
||||
|
||||
if pass_idx + 1 < total_passes {
|
||||
if !self.hardswap_warmup_is_authoritative(
|
||||
generation,
|
||||
map_hash,
|
||||
endpoint_revision,
|
||||
) {
|
||||
return;
|
||||
}
|
||||
let backoff_ms = self.hardswap_warmup_backoff_ms(pass_idx);
|
||||
debug!(
|
||||
dc = *dc,
|
||||
@@ -198,8 +245,14 @@ impl MePool {
|
||||
}
|
||||
|
||||
if hardswap {
|
||||
self.warmup_generation_for_all_dcs(rng, generation, &desired_by_dc)
|
||||
.await;
|
||||
self.warmup_generation_for_all_dcs(
|
||||
rng,
|
||||
generation,
|
||||
desired_map_hash,
|
||||
endpoint_snapshot.revision,
|
||||
&desired_by_dc,
|
||||
)
|
||||
.await;
|
||||
} else {
|
||||
self.reconcile_connections(rng).await;
|
||||
}
|
||||
|
||||
@@ -26,7 +26,7 @@ fn addr_v6(segment: u16, port: u16) -> SocketAddr {
|
||||
)
|
||||
}
|
||||
|
||||
async fn insert_writer(
|
||||
pub(super) async fn insert_writer(
|
||||
pool: &Arc<MePool>,
|
||||
writer_id: u64,
|
||||
writer_dc: i32,
|
||||
@@ -63,7 +63,7 @@ async fn insert_writer(
|
||||
writer
|
||||
}
|
||||
|
||||
async fn insert_writer_floor(
|
||||
pub(super) async fn insert_writer_floor(
|
||||
pool: &Arc<MePool>,
|
||||
first_writer_id: u64,
|
||||
writer_dc: i32,
|
||||
@@ -489,7 +489,7 @@ async fn generation_role_reconciliation_promotes_active_warm_and_drains_orphans(
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
let desired_by_dc = HashMap::from([(1, HashSet::from([endpoint]))]);
|
||||
let desired_by_dc = pool.desired_dc_endpoints().await;
|
||||
let map_hash = MePool::desired_map_hash(&desired_by_dc);
|
||||
let endpoint_revision = pool.endpoint_snapshot.load().revision;
|
||||
let reservation = pool
|
||||
|
||||
@@ -15,6 +15,9 @@ mod runtime_snapshot;
|
||||
// Hardswap ownership, coverage, and writer-replacement lifecycle state.
|
||||
mod hardswap_snapshot;
|
||||
pub(crate) use hardswap_snapshot::MeApiHardswapSnapshot;
|
||||
#[cfg(test)]
|
||||
#[path = "pool_status/tests.rs"]
|
||||
mod status_invariant_tests;
|
||||
#[derive(Clone, Debug)]
|
||||
pub(crate) struct MeApiWriterStatusSnapshot {
|
||||
pub writer_id: u64,
|
||||
|
||||
@@ -35,11 +35,15 @@ impl MePool {
|
||||
&self,
|
||||
reinit: &ReinitStatusSnapshot,
|
||||
) -> MeApiHardswapSnapshot {
|
||||
let desired_by_dc = self.desired_dc_endpoints().await;
|
||||
let endpoint_snapshot = self.endpoint_snapshot.load_full();
|
||||
let desired_by_dc = self.desired_dc_endpoints_from_snapshot(&endpoint_snapshot);
|
||||
let desired_hash = Self::desired_map_hash(&desired_by_dc);
|
||||
let writers = self.writers.read().await;
|
||||
let pending_generation = reinit.pending_hardswap_generation;
|
||||
let pending = pending_generation != 0;
|
||||
let pending_map_current = pending
|
||||
&& reinit.pending_hardswap_map_hash == desired_hash
|
||||
&& reinit.pending_hardswap_endpoint_revision == endpoint_snapshot.revision;
|
||||
let mut pending_writers_current = 0usize;
|
||||
let mut pending_writer_addrs = Vec::<(i32, SocketAddr)>::new();
|
||||
let mut orphan_warm_writers_current = 0usize;
|
||||
@@ -49,10 +53,12 @@ impl MePool {
|
||||
continue;
|
||||
}
|
||||
let contour = WriterContour::from_u8(writer.contour.load(Ordering::Acquire));
|
||||
if contour == WriterContour::Warm && writer.generation != pending_generation {
|
||||
if contour == WriterContour::Warm
|
||||
&& (!pending_map_current || writer.generation != pending_generation)
|
||||
{
|
||||
orphan_warm_writers_current = orphan_warm_writers_current.saturating_add(1);
|
||||
}
|
||||
if pending
|
||||
if pending_map_current
|
||||
&& writer.generation == pending_generation
|
||||
&& contour == WriterContour::Warm
|
||||
&& desired_by_dc
|
||||
@@ -64,7 +70,7 @@ impl MePool {
|
||||
}
|
||||
}
|
||||
|
||||
let pending_coverage = pending
|
||||
let pending_coverage = pending_map_current
|
||||
.then(|| self.hardswap_coverage(&desired_by_dc, &pending_writer_addrs));
|
||||
let pending_writer_deficit = pending_coverage
|
||||
.as_ref()
|
||||
@@ -85,8 +91,7 @@ impl MePool {
|
||||
pending_writers_current,
|
||||
pending_writer_deficit,
|
||||
pending_missing_dc_groups,
|
||||
pending_map_current: pending
|
||||
.then_some(reinit.pending_hardswap_map_hash == desired_hash),
|
||||
pending_map_current: pending.then_some(pending_map_current),
|
||||
orphan_warm_writers_current,
|
||||
replacement_preparing_current,
|
||||
replacement_retiring_current,
|
||||
|
||||
@@ -62,7 +62,7 @@ impl MePool {
|
||||
|
||||
let active_generation = self.reinit.status.load().active_generation;
|
||||
let writers = self.writers.read().await.clone();
|
||||
let mut live_writers_by_dc = HashMap::<i16, usize>::new();
|
||||
let mut live_writers_by_group = HashMap::<(i16, bool), usize>::new();
|
||||
for writer in writers.iter() {
|
||||
if writer.draining.load(Ordering::Relaxed)
|
||||
|| writer.generation != active_generation
|
||||
@@ -72,19 +72,26 @@ impl MePool {
|
||||
continue;
|
||||
}
|
||||
if let Ok(dc) = i16::try_from(writer.writer_dc) {
|
||||
*live_writers_by_dc.entry(dc).or_insert(0) += 1;
|
||||
*live_writers_by_group
|
||||
.entry((dc, writer.addr.is_ipv4()))
|
||||
.or_insert(0) += 1;
|
||||
}
|
||||
}
|
||||
|
||||
for (dc, endpoints) in endpoints_by_dc {
|
||||
let endpoint_count = endpoints.len();
|
||||
if endpoint_count == 0 {
|
||||
return false;
|
||||
}
|
||||
let required = self.required_writers_for_dc_with_floor_mode(endpoint_count, false);
|
||||
let alive = live_writers_by_dc.get(&dc).copied().unwrap_or(0);
|
||||
if alive < required {
|
||||
return false;
|
||||
for (ipv4, endpoint_count) in endpoint_family_counts(&endpoints) {
|
||||
if endpoint_count == 0 {
|
||||
continue;
|
||||
}
|
||||
let required =
|
||||
self.required_writers_for_dc_with_floor_mode(endpoint_count, false);
|
||||
let alive = live_writers_by_group
|
||||
.get(&(dc, ipv4))
|
||||
.copied()
|
||||
.unwrap_or(0);
|
||||
if alive < required {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -123,7 +130,14 @@ impl MePool {
|
||||
|
||||
let required_writers = endpoints_by_dc
|
||||
.values()
|
||||
.map(|endpoints| self.required_writers_for_dc_with_floor_mode(endpoints.len(), false))
|
||||
.map(|endpoints| {
|
||||
endpoint_family_counts(endpoints)
|
||||
.into_iter()
|
||||
.map(|(_, count)| {
|
||||
self.required_writers_for_dc_with_floor_mode(count, false)
|
||||
})
|
||||
.sum::<usize>()
|
||||
})
|
||||
.sum();
|
||||
|
||||
let idle_since = self.registry.writer_idle_since_snapshot().await;
|
||||
@@ -227,39 +241,58 @@ impl MePool {
|
||||
.max(1);
|
||||
for (dc, endpoints) in endpoints_by_dc {
|
||||
let endpoint_count = endpoints.len();
|
||||
let family_counts = endpoint_family_counts(&endpoints);
|
||||
let dc_available_endpoints = endpoints
|
||||
.iter()
|
||||
.filter(|endpoint| live_writers_by_dc_endpoint.contains_key(&(dc, **endpoint)))
|
||||
.count();
|
||||
let base_required = self.required_writers_for_dc(endpoint_count);
|
||||
let dc_required_writers =
|
||||
self.required_writers_for_dc_with_floor_mode(endpoint_count, false);
|
||||
let floor_min = if endpoint_count <= 1 {
|
||||
(self
|
||||
.floor_runtime
|
||||
.me_adaptive_floor_min_writers_single_endpoint
|
||||
.load(Ordering::Relaxed) as usize)
|
||||
.max(1)
|
||||
.min(base_required.max(1))
|
||||
} else {
|
||||
(self
|
||||
.floor_runtime
|
||||
.me_adaptive_floor_min_writers_multi_endpoint
|
||||
.load(Ordering::Relaxed) as usize)
|
||||
.max(1)
|
||||
.min(base_required.max(1))
|
||||
};
|
||||
let extra_per_core = if endpoint_count <= 1 {
|
||||
self.floor_runtime
|
||||
.me_adaptive_floor_max_extra_writers_single_per_core
|
||||
.load(Ordering::Relaxed) as usize
|
||||
} else {
|
||||
self.floor_runtime
|
||||
.me_adaptive_floor_max_extra_writers_multi_per_core
|
||||
.load(Ordering::Relaxed) as usize
|
||||
};
|
||||
let floor_max =
|
||||
base_required.saturating_add(adaptive_cpu_cores.saturating_mul(extra_per_core));
|
||||
let base_required = family_counts
|
||||
.iter()
|
||||
.filter(|(_, count)| *count > 0)
|
||||
.map(|(_, count)| self.required_writers_for_dc(*count))
|
||||
.sum::<usize>();
|
||||
let dc_required_writers = family_counts
|
||||
.iter()
|
||||
.filter(|(_, count)| *count > 0)
|
||||
.map(|(_, count)| {
|
||||
self.required_writers_for_dc_with_floor_mode(*count, false)
|
||||
})
|
||||
.sum::<usize>();
|
||||
let floor_min = family_counts
|
||||
.iter()
|
||||
.filter(|(_, count)| *count > 0)
|
||||
.map(|(_, count)| {
|
||||
let family_base = self.required_writers_for_dc(*count);
|
||||
let configured = if *count <= 1 {
|
||||
self.floor_runtime
|
||||
.me_adaptive_floor_min_writers_single_endpoint
|
||||
.load(Ordering::Relaxed) as usize
|
||||
} else {
|
||||
self.floor_runtime
|
||||
.me_adaptive_floor_min_writers_multi_endpoint
|
||||
.load(Ordering::Relaxed) as usize
|
||||
};
|
||||
configured.max(1).min(family_base.max(1))
|
||||
})
|
||||
.sum::<usize>();
|
||||
let floor_max = family_counts
|
||||
.iter()
|
||||
.filter(|(_, count)| *count > 0)
|
||||
.map(|(_, count)| {
|
||||
let family_base = self.required_writers_for_dc(*count);
|
||||
let extra_per_core = if *count <= 1 {
|
||||
self.floor_runtime
|
||||
.me_adaptive_floor_max_extra_writers_single_per_core
|
||||
.load(Ordering::Relaxed) as usize
|
||||
} else {
|
||||
self.floor_runtime
|
||||
.me_adaptive_floor_max_extra_writers_multi_per_core
|
||||
.load(Ordering::Relaxed) as usize
|
||||
};
|
||||
family_base
|
||||
.saturating_add(adaptive_cpu_cores.saturating_mul(extra_per_core))
|
||||
})
|
||||
.sum::<usize>();
|
||||
let floor_capped =
|
||||
matches!(floor_mode, MeFloorMode::Adaptive) && dc_required_writers < base_required;
|
||||
let dc_alive_writers = live_writers_by_dc.get(&dc).copied().unwrap_or(0);
|
||||
@@ -322,3 +355,8 @@ impl MePool {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn endpoint_family_counts(endpoints: &BTreeSet<SocketAddr>) -> [(bool, usize); 2] {
|
||||
let ipv4 = endpoints.iter().filter(|endpoint| endpoint.is_ipv4()).count();
|
||||
[(true, ipv4), (false, endpoints.len().saturating_sub(ipv4))]
|
||||
}
|
||||
|
||||
@@ -0,0 +1,84 @@
|
||||
use std::collections::HashMap;
|
||||
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU32, AtomicU64};
|
||||
use std::time::Instant;
|
||||
|
||||
use tokio::sync::mpsc;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
use crate::network::probe::NetworkDecision;
|
||||
use crate::transport::middle_proxy::codec::WriterCommand;
|
||||
use crate::transport::middle_proxy::pool::{MePool, MeWriter, WriterContour};
|
||||
use crate::transport::middle_proxy::pool_writer_security_tests::make_pool_with_decision;
|
||||
|
||||
fn writer(
|
||||
pool: &Arc<MePool>,
|
||||
id: u64,
|
||||
dc: i32,
|
||||
addr: SocketAddr,
|
||||
) -> MeWriter {
|
||||
let (tx, _rx) = mpsc::channel::<WriterCommand>(8);
|
||||
MeWriter {
|
||||
id,
|
||||
addr,
|
||||
source_ip: addr.ip(),
|
||||
writer_dc: dc,
|
||||
generation: pool.current_generation(),
|
||||
contour: Arc::new(AtomicU8::new(WriterContour::Active.as_u8())),
|
||||
created_at: Instant::now(),
|
||||
tx,
|
||||
byte_budget: pool.new_writer_byte_budget(),
|
||||
cancel: CancellationToken::new(),
|
||||
degraded: Arc::new(AtomicBool::new(false)),
|
||||
rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)),
|
||||
draining: Arc::new(AtomicBool::new(false)),
|
||||
draining_started_at_epoch_secs: Arc::new(AtomicU64::new(0)),
|
||||
drain_deadline_epoch_secs: Arc::new(AtomicU64::new(0)),
|
||||
allow_drain_fallback: Arc::new(AtomicBool::new(false)),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dual_family_status_reports_each_family_floor() {
|
||||
let pool = make_pool_with_decision(NetworkDecision {
|
||||
ipv4_me: true,
|
||||
ipv6_me: true,
|
||||
effective_prefer: 4,
|
||||
effective_multipath: false,
|
||||
..NetworkDecision::default()
|
||||
})
|
||||
.await;
|
||||
let v4 = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443);
|
||||
let v6 = SocketAddr::new(IpAddr::V6(Ipv6Addr::LOCALHOST), 443);
|
||||
pool.update_proxy_maps(
|
||||
HashMap::from([(2, vec![(v4.ip(), v4.port())])]),
|
||||
Some(HashMap::from([(2, vec![(v6.ip(), v6.port())])])),
|
||||
)
|
||||
.await;
|
||||
let required_per_family = pool.required_writers_for_dc(1);
|
||||
let mut writers = pool.writers.write().await;
|
||||
for (group, dc) in [2, -2].into_iter().enumerate() {
|
||||
for offset in 0..required_per_family {
|
||||
writers.push(writer(
|
||||
&pool,
|
||||
(group as u64 * 100) + offset as u64,
|
||||
dc,
|
||||
v4,
|
||||
));
|
||||
}
|
||||
}
|
||||
drop(writers);
|
||||
|
||||
let snapshot = pool.api_status_snapshot().await;
|
||||
|
||||
assert_eq!(snapshot.required_writers, required_per_family * 4);
|
||||
assert_eq!(snapshot.alive_writers, required_per_family * 2);
|
||||
assert_eq!(snapshot.coverage_pct, 50.0);
|
||||
for dc in snapshot.dcs {
|
||||
assert_eq!(dc.required_writers, required_per_family * 2);
|
||||
assert_eq!(dc.alive_writers, required_per_family);
|
||||
assert_eq!(dc.coverage_pct, 50.0);
|
||||
}
|
||||
assert!(!pool.admission_ready_full_floor().await);
|
||||
}
|
||||
@@ -43,9 +43,7 @@ impl MePool {
|
||||
let endpoint_is_current = self
|
||||
.endpoint_snapshot
|
||||
.load()
|
||||
.preferred_endpoints_by_dc
|
||||
.get(&writer.writer_dc)
|
||||
.is_some_and(|endpoints| endpoints.contains(&writer.addr));
|
||||
.contains_dc_endpoint(writer.writer_dc, writer.addr);
|
||||
if !endpoint_is_current {
|
||||
return Err(ProxyError::Proxy(
|
||||
"ME writer target changed before publication".into(),
|
||||
@@ -95,21 +93,29 @@ impl MePool {
|
||||
return Ok(());
|
||||
}
|
||||
let endpoint_snapshot = self.endpoint_snapshot.load();
|
||||
let preferred = &endpoint_snapshot.preferred_endpoints_by_dc;
|
||||
let Some(endpoints) = preferred.get(&writer.writer_dc) else {
|
||||
if !endpoint_snapshot.contains_dc_endpoint(writer.writer_dc, writer.addr) {
|
||||
return Err(ProxyError::Proxy(
|
||||
"ME writer target changed before publication".into(),
|
||||
));
|
||||
};
|
||||
let required = match contour {
|
||||
WriterContour::Active | WriterContour::Warm => self.required_writers_for_dc(
|
||||
endpoints
|
||||
.iter()
|
||||
.filter(|endpoint| endpoint.is_ipv4() == writer.addr.is_ipv4())
|
||||
.count(),
|
||||
),
|
||||
WriterContour::Draining => 0,
|
||||
};
|
||||
}
|
||||
if contour == WriterContour::Active && intent == WriterOpenIntent::Normal {
|
||||
let current = writers
|
||||
.iter()
|
||||
.filter(|candidate| {
|
||||
!candidate.draining.load(Ordering::Acquire)
|
||||
&& WriterContour::from_u8(candidate.contour.load(Ordering::Acquire))
|
||||
== WriterContour::Active
|
||||
})
|
||||
.count();
|
||||
if current >= self.adaptive_floor_active_cap_configured_total() {
|
||||
return Err(ProxyError::Proxy(
|
||||
"ME active writer cap was reached before publication".into(),
|
||||
));
|
||||
}
|
||||
return Ok(());
|
||||
}
|
||||
let endpoints = endpoint_snapshot.endpoints_for_dc_family(writer.writer_dc, family);
|
||||
let required = self.required_writers_for_dc(endpoints.len());
|
||||
let current = writers
|
||||
.iter()
|
||||
.filter(|candidate| {
|
||||
@@ -119,7 +125,8 @@ impl MePool {
|
||||
&& WriterContour::from_u8(candidate.contour.load(Ordering::Acquire))
|
||||
== contour
|
||||
&& candidate.addr.is_ipv4() == writer.addr.is_ipv4()
|
||||
&& endpoints.contains(&candidate.addr)
|
||||
&& endpoint_snapshot
|
||||
.contains_dc_endpoint(candidate.writer_dc, candidate.addr)
|
||||
})
|
||||
.count();
|
||||
if current >= required {
|
||||
|
||||
@@ -86,7 +86,6 @@ impl MePool {
|
||||
));
|
||||
}
|
||||
let endpoint_snapshot = self.endpoint_snapshot.load();
|
||||
let preferred = &endpoint_snapshot.preferred_endpoints_by_dc;
|
||||
let donor_count = writers
|
||||
.iter()
|
||||
.filter(|candidate| {
|
||||
@@ -95,9 +94,8 @@ impl MePool {
|
||||
&& candidate.generation == expected_victim_role.generation
|
||||
&& WriterContour::from_u8(candidate.contour.load(Ordering::Acquire))
|
||||
== WriterContour::Active
|
||||
&& preferred
|
||||
.get(&candidate.writer_dc)
|
||||
.is_some_and(|endpoints| endpoints.contains(&candidate.addr))
|
||||
&& endpoint_snapshot
|
||||
.contains_dc_endpoint(candidate.writer_dc, candidate.addr)
|
||||
&& (candidate.addr.is_ipv4()
|
||||
== matches!(expected_victim_role.family, crate::network::IpFamily::V4))
|
||||
})
|
||||
@@ -115,9 +113,8 @@ impl MePool {
|
||||
&& candidate.generation == writer.generation
|
||||
&& WriterContour::from_u8(candidate.contour.load(Ordering::Acquire))
|
||||
== WriterContour::Active
|
||||
&& preferred
|
||||
.get(&candidate.writer_dc)
|
||||
.is_some_and(|endpoints| endpoints.contains(&candidate.addr))
|
||||
&& endpoint_snapshot
|
||||
.contains_dc_endpoint(candidate.writer_dc, candidate.addr)
|
||||
&& (if candidate.addr.is_ipv4() {
|
||||
crate::network::IpFamily::V4
|
||||
} else {
|
||||
@@ -223,11 +220,7 @@ mod tests {
|
||||
WriterContour::Active,
|
||||
WriterOpenIntent::Replacement,
|
||||
writer_dc,
|
||||
if addr.is_ipv4() {
|
||||
crate::network::IpFamily::V4
|
||||
} else {
|
||||
crate::network::IpFamily::V6
|
||||
},
|
||||
addr,
|
||||
)
|
||||
.await
|
||||
.expect("replacement open must be admitted");
|
||||
|
||||
@@ -96,11 +96,7 @@ impl MePool {
|
||||
contour,
|
||||
intent,
|
||||
writer_dc,
|
||||
if addr.is_ipv4() {
|
||||
crate::network::IpFamily::V4
|
||||
} else {
|
||||
crate::network::IpFamily::V6
|
||||
},
|
||||
addr,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
|
||||
@@ -1,26 +1,29 @@
|
||||
#![allow(clippy::too_many_arguments)]
|
||||
|
||||
use std::cmp::Reverse;
|
||||
use std::net::SocketAddr;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::Ordering;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use tokio::sync::mpsc::error::TrySendError;
|
||||
use tokio::sync::{OwnedSemaphorePermit, Semaphore, TryAcquireError, mpsc};
|
||||
use tokio::sync::{OwnedSemaphorePermit, TryAcquireError};
|
||||
use tracing::{debug, warn};
|
||||
|
||||
use super::MePool;
|
||||
use super::codec::{ProxyReqCommand, WriterBytePermit, WriterCommand};
|
||||
use super::codec::WriterCommand;
|
||||
use super::registry::{ConnMeta, WriterBindOutcome};
|
||||
use super::wire::{build_proxy_req_payload, proxy_req_payload_len};
|
||||
use crate::config::defaults::ME_WRITER_BYTE_PERMIT_UNIT_BYTES;
|
||||
use crate::config::{MeRouteNoWriterMode, MeWriterPickMode};
|
||||
use super::wire::build_proxy_req_payload;
|
||||
use crate::config::MeRouteNoWriterMode;
|
||||
use crate::error::{ProxyError, Result};
|
||||
use crate::stats::Stats;
|
||||
use crate::stream::PooledBuffer;
|
||||
use rand::seq::SliceRandom;
|
||||
|
||||
use self::bound::BoundWriterSendOutcome;
|
||||
use self::reservation::{
|
||||
LEGACY_PROXY_REQ_SOURCE_CAPACITY_OVERHEAD_BYTES, WriterByteReserveError,
|
||||
WriterCommandReserveError, proxy_req_resident_permits, reserve_writer_bytes,
|
||||
reserve_writer_command_slot, try_reserve_writer_bytes, writer_send_deadline,
|
||||
};
|
||||
|
||||
const IDLE_WRITER_PENALTY_MID_SECS: u64 = 45;
|
||||
const IDLE_WRITER_PENALTY_HIGH_SECS: u64 = 55;
|
||||
const HYBRID_GLOBAL_BURST_PERIOD_ROUNDS: u32 = 4;
|
||||
@@ -31,142 +34,15 @@ const PICK_PENALTY_WARM: u64 = 200;
|
||||
const PICK_PENALTY_DRAINING: u64 = 600;
|
||||
const PICK_PENALTY_STALE: u64 = 300;
|
||||
const PICK_PENALTY_DEGRADED: u64 = 250;
|
||||
const RPC_WRITER_FRAME_CAPACITY_OVERHEAD_BYTES: usize = 27;
|
||||
const LEGACY_PROXY_REQ_SOURCE_CAPACITY_OVERHEAD_BYTES: usize = 128;
|
||||
|
||||
// Send-path submodules isolate delivery, close handling, recovery, reservations, and selection.
|
||||
mod bound;
|
||||
mod close;
|
||||
mod pooled;
|
||||
mod recovery;
|
||||
mod reservation;
|
||||
mod selection;
|
||||
|
||||
enum WriterCommandReserveError {
|
||||
Closed,
|
||||
TimedOut,
|
||||
}
|
||||
|
||||
enum WriterByteReserveError {
|
||||
Closed,
|
||||
TimedOut,
|
||||
}
|
||||
|
||||
fn proxy_tag_array(tag: Option<&[u8]>) -> Option<[u8; 16]> {
|
||||
tag.and_then(|tag| <[u8; 16]>::try_from(tag).ok())
|
||||
}
|
||||
|
||||
fn proxy_req_payload_from_command(
|
||||
cmd: WriterCommand,
|
||||
) -> Option<(PooledBuffer, OwnedSemaphorePermit)> {
|
||||
match cmd {
|
||||
WriterCommand::ProxyReq(command) => Some((command.payload, command._permit)),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn payload_permit_from_data_command(cmd: WriterCommand) -> Option<OwnedSemaphorePermit> {
|
||||
match cmd {
|
||||
WriterCommand::Data { _permit, .. } => _permit,
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
async fn reserve_writer_command_slot(
|
||||
tx: &mpsc::Sender<WriterCommand>,
|
||||
deadline: Option<Instant>,
|
||||
) -> std::result::Result<mpsc::OwnedPermit<WriterCommand>, WriterCommandReserveError> {
|
||||
let reserve = tx.clone().reserve_owned();
|
||||
match deadline {
|
||||
Some(deadline) => {
|
||||
match tokio::time::timeout(deadline.saturating_duration_since(Instant::now()), reserve)
|
||||
.await
|
||||
{
|
||||
Ok(Ok(permit)) => Ok(permit),
|
||||
Ok(Err(_)) => Err(WriterCommandReserveError::Closed),
|
||||
Err(_) => Err(WriterCommandReserveError::TimedOut),
|
||||
}
|
||||
}
|
||||
None => reserve.await.map_err(|_| WriterCommandReserveError::Closed),
|
||||
}
|
||||
}
|
||||
|
||||
fn writer_send_deadline(wait: Option<Duration>) -> Option<Instant> {
|
||||
wait.map(|wait| Instant::now() + wait)
|
||||
}
|
||||
|
||||
fn writer_resident_permits(
|
||||
source_capacity: usize,
|
||||
encoded_payload_len: usize,
|
||||
) -> Option<(u32, usize)> {
|
||||
let resident_bytes = source_capacity
|
||||
.checked_add(encoded_payload_len)?
|
||||
.checked_add(RPC_WRITER_FRAME_CAPACITY_OVERHEAD_BYTES)?;
|
||||
let permits = resident_bytes.div_ceil(ME_WRITER_BYTE_PERMIT_UNIT_BYTES);
|
||||
let permits = u32::try_from(permits).ok()?;
|
||||
let reserved_bytes = (permits as usize).checked_mul(ME_WRITER_BYTE_PERMIT_UNIT_BYTES)?;
|
||||
Some((
|
||||
permits.max(1),
|
||||
reserved_bytes.max(ME_WRITER_BYTE_PERMIT_UNIT_BYTES),
|
||||
))
|
||||
}
|
||||
|
||||
fn proxy_req_resident_permits(
|
||||
source_capacity: usize,
|
||||
data_len: usize,
|
||||
proxy_tag: Option<&[u8]>,
|
||||
proto_flags: u32,
|
||||
) -> Option<(u32, usize)> {
|
||||
writer_resident_permits(
|
||||
source_capacity,
|
||||
proxy_req_payload_len(data_len, proxy_tag, proto_flags),
|
||||
)
|
||||
}
|
||||
|
||||
fn try_reserve_writer_bytes(
|
||||
byte_budget: &Arc<Semaphore>,
|
||||
permits: u32,
|
||||
reserved_bytes: usize,
|
||||
stats: &Arc<Stats>,
|
||||
) -> std::result::Result<WriterBytePermit, TryAcquireError> {
|
||||
byte_budget
|
||||
.clone()
|
||||
.try_acquire_many_owned(permits)
|
||||
.map(|permit| WriterBytePermit::new(permit, reserved_bytes, stats.clone()))
|
||||
}
|
||||
|
||||
async fn reserve_writer_bytes(
|
||||
byte_budget: &Arc<Semaphore>,
|
||||
permits: u32,
|
||||
reserved_bytes: usize,
|
||||
deadline: Option<Instant>,
|
||||
stats: &Arc<Stats>,
|
||||
) -> std::result::Result<WriterBytePermit, WriterByteReserveError> {
|
||||
match try_reserve_writer_bytes(byte_budget, permits, reserved_bytes, stats) {
|
||||
Ok(permit) => return Ok(permit),
|
||||
Err(TryAcquireError::Closed) => return Err(WriterByteReserveError::Closed),
|
||||
Err(TryAcquireError::NoPermits) => {
|
||||
stats.increment_me_writer_byte_budget_wait_total();
|
||||
}
|
||||
}
|
||||
|
||||
let acquire = byte_budget.clone().acquire_many_owned(permits);
|
||||
match deadline {
|
||||
Some(deadline) => {
|
||||
match tokio::time::timeout(deadline.saturating_duration_since(Instant::now()), acquire)
|
||||
.await
|
||||
{
|
||||
Ok(Ok(permit)) => Ok(WriterBytePermit::new(permit, reserved_bytes, stats.clone())),
|
||||
Ok(Err(_)) => Err(WriterByteReserveError::Closed),
|
||||
Err(_) => {
|
||||
stats.increment_me_writer_byte_budget_timeout_total();
|
||||
Err(WriterByteReserveError::TimedOut)
|
||||
}
|
||||
}
|
||||
}
|
||||
None => acquire
|
||||
.await
|
||||
.map(|permit| WriterBytePermit::new(permit, reserved_bytes, stats.clone()))
|
||||
.map_err(|_| WriterByteReserveError::Closed),
|
||||
}
|
||||
}
|
||||
|
||||
impl MePool {
|
||||
/// Send RPC_PROXY_REQ. `tag_override`: per-user ad_tag (from access.user_ad_tags); if None, uses pool default.
|
||||
/// `payload_permit` keeps optional client byte accounting alive until the writer consumes the command.
|
||||
@@ -243,78 +119,22 @@ impl MePool {
|
||||
let mut hybrid_wait_current = hybrid_wait_step;
|
||||
|
||||
loop {
|
||||
if let Some((current, current_meta)) = self.registry.get_writer_with_meta(conn_id).await
|
||||
{
|
||||
let deadline =
|
||||
writer_send_deadline(self.route_runtime.me_route_blocking_send_timeout);
|
||||
let writer_permit = match reserve_writer_bytes(
|
||||
¤t.byte_budget,
|
||||
match self
|
||||
.try_send_bound_writer(
|
||||
conn_id,
|
||||
client_addr,
|
||||
data,
|
||||
proto_flags,
|
||||
tag,
|
||||
writer_byte_permits,
|
||||
writer_reserved_bytes,
|
||||
deadline,
|
||||
&self.stats,
|
||||
payload_permit,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(permit) => permit,
|
||||
Err(WriterByteReserveError::TimedOut) => {
|
||||
self.stats
|
||||
.increment_me_writer_pick_full_total(self.writer_pick_mode());
|
||||
return Err(ProxyError::Proxy(
|
||||
"ME writer byte budget full within blocking send timeout".into(),
|
||||
));
|
||||
}
|
||||
Err(WriterByteReserveError::Closed) => {
|
||||
warn!(
|
||||
writer_id = current.writer_id,
|
||||
"ME writer byte budget closed"
|
||||
);
|
||||
self.remove_writer_and_close_clients(current.writer_id)
|
||||
.await;
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let (current_payload, _) = build_routed_payload(current_meta.our_addr);
|
||||
let command = WriterCommand::Data {
|
||||
payload: current_payload,
|
||||
_permit: payload_permit.take(),
|
||||
writer_permit,
|
||||
};
|
||||
match current.tx.try_send(command) {
|
||||
Ok(()) => {
|
||||
self.note_hybrid_route_success();
|
||||
return Ok(());
|
||||
}
|
||||
Err(TrySendError::Full(cmd)) => {
|
||||
match reserve_writer_command_slot(¤t.tx, deadline).await {
|
||||
Ok(permit) => {
|
||||
permit.send(cmd);
|
||||
self.note_hybrid_route_success();
|
||||
return Ok(());
|
||||
}
|
||||
Err(WriterCommandReserveError::TimedOut) => {
|
||||
self.stats
|
||||
.increment_me_writer_pick_full_total(self.writer_pick_mode());
|
||||
return Err(ProxyError::Proxy(
|
||||
"ME writer channel full within blocking send timeout".into(),
|
||||
));
|
||||
}
|
||||
Err(WriterCommandReserveError::Closed) => {
|
||||
payload_permit = payload_permit_from_data_command(cmd);
|
||||
}
|
||||
}
|
||||
warn!(writer_id = current.writer_id, "ME writer channel closed");
|
||||
self.remove_writer_and_close_clients(current.writer_id)
|
||||
.await;
|
||||
continue;
|
||||
}
|
||||
Err(TrySendError::Closed(cmd)) => {
|
||||
payload_permit = payload_permit_from_data_command(cmd);
|
||||
warn!(writer_id = current.writer_id, "ME writer channel closed");
|
||||
self.remove_writer_and_close_clients(current.writer_id)
|
||||
.await;
|
||||
continue;
|
||||
}
|
||||
.await?
|
||||
{
|
||||
BoundWriterSendOutcome::Sent => return Ok(()),
|
||||
BoundWriterSendOutcome::Retry(retry_permit) => {
|
||||
payload_permit = retry_permit;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -537,88 +357,9 @@ impl MePool {
|
||||
}
|
||||
hybrid_wait_current = hybrid_wait_step;
|
||||
let pick_mode = self.writer_pick_mode();
|
||||
let pick_sample_size = self.writer_pick_sample_size();
|
||||
let writer_ids: Vec<u64> = candidate_indices
|
||||
.iter()
|
||||
.map(|idx| writers_snapshot[*idx].id)
|
||||
.collect();
|
||||
let writer_idle_since = self
|
||||
.registry
|
||||
.writer_idle_since_for_writer_ids(&writer_ids)
|
||||
let ordered_candidate_indices = self
|
||||
.ordered_candidate_indices(candidate_indices, &writers_snapshot, pick_mode)
|
||||
.await;
|
||||
let now_epoch_secs = Self::now_epoch_secs();
|
||||
let start = self.rr.fetch_add(1, Ordering::Relaxed) as usize % candidate_indices.len();
|
||||
let ordered_candidate_indices = if pick_mode == MeWriterPickMode::P2c {
|
||||
self.p2c_ordered_candidate_indices(
|
||||
&candidate_indices,
|
||||
&writers_snapshot,
|
||||
&writer_idle_since,
|
||||
now_epoch_secs,
|
||||
start,
|
||||
pick_sample_size,
|
||||
)
|
||||
} else {
|
||||
if self
|
||||
.writer_selection_policy
|
||||
.me_deterministic_writer_sort
|
||||
.load(Ordering::Relaxed)
|
||||
{
|
||||
candidate_indices.sort_by(|lhs, rhs| {
|
||||
let left = &writers_snapshot[*lhs];
|
||||
let right = &writers_snapshot[*rhs];
|
||||
let left_key = (
|
||||
self.writer_contour_rank_for_selection(left),
|
||||
(left.generation < self.current_generation()) as usize,
|
||||
left.degraded.load(Ordering::Relaxed) as usize,
|
||||
self.writer_idle_rank_for_selection(
|
||||
left,
|
||||
&writer_idle_since,
|
||||
now_epoch_secs,
|
||||
),
|
||||
Reverse(left.tx.capacity()),
|
||||
left.addr,
|
||||
left.id,
|
||||
);
|
||||
let right_key = (
|
||||
self.writer_contour_rank_for_selection(right),
|
||||
(right.generation < self.current_generation()) as usize,
|
||||
right.degraded.load(Ordering::Relaxed) as usize,
|
||||
self.writer_idle_rank_for_selection(
|
||||
right,
|
||||
&writer_idle_since,
|
||||
now_epoch_secs,
|
||||
),
|
||||
Reverse(right.tx.capacity()),
|
||||
right.addr,
|
||||
right.id,
|
||||
);
|
||||
left_key.cmp(&right_key)
|
||||
});
|
||||
} else {
|
||||
candidate_indices.sort_by_key(|idx| {
|
||||
let w = &writers_snapshot[*idx];
|
||||
let degraded = w.degraded.load(Ordering::Relaxed);
|
||||
let stale = (w.generation < self.current_generation()) as usize;
|
||||
(
|
||||
self.writer_contour_rank_for_selection(w),
|
||||
stale,
|
||||
degraded as usize,
|
||||
self.writer_idle_rank_for_selection(
|
||||
w,
|
||||
&writer_idle_since,
|
||||
now_epoch_secs,
|
||||
),
|
||||
Reverse(w.tx.capacity()),
|
||||
)
|
||||
});
|
||||
}
|
||||
|
||||
let mut ordered = Vec::<usize>::with_capacity(candidate_indices.len());
|
||||
for offset in 0..candidate_indices.len() {
|
||||
ordered.push(candidate_indices[(start + offset) % candidate_indices.len()]);
|
||||
}
|
||||
ordered
|
||||
};
|
||||
let mut fallback_blocking_idx: Option<usize> = None;
|
||||
|
||||
for idx in ordered_candidate_indices {
|
||||
@@ -804,167 +545,4 @@ impl MePool {
|
||||
}
|
||||
}
|
||||
|
||||
/// Send RPC_PROXY_REQ while keeping the first bound-writer path allocation-light.
|
||||
/// The client byte permit follows the payload until writer completion or command drop.
|
||||
pub async fn send_proxy_req_pooled(
|
||||
self: &Arc<Self>,
|
||||
conn_id: u64,
|
||||
target_dc: i16,
|
||||
client_addr: SocketAddr,
|
||||
our_addr: SocketAddr,
|
||||
payload: PooledBuffer,
|
||||
_permit: OwnedSemaphorePermit,
|
||||
proto_flags: u32,
|
||||
tag_override: Option<[u8; 16]>,
|
||||
) -> Result<()> {
|
||||
let tag = tag_override.or_else(|| proxy_tag_array(self.proxy_tag.as_deref()));
|
||||
let Some((writer_byte_permits, writer_reserved_bytes)) = proxy_req_resident_permits(
|
||||
payload.capacity(),
|
||||
payload.len(),
|
||||
tag.as_ref().map(|tag| tag.as_slice()),
|
||||
proto_flags,
|
||||
) else {
|
||||
self.stats.increment_me_writer_byte_budget_oversize_total();
|
||||
return Err(ProxyError::Proxy(
|
||||
"ME writer payload residency calculation overflow".into(),
|
||||
));
|
||||
};
|
||||
if writer_byte_permits as usize > self.writer_lifecycle.writer_byte_budget_permits {
|
||||
self.stats.increment_me_writer_byte_budget_oversize_total();
|
||||
return Err(ProxyError::Proxy(
|
||||
"ME writer payload exceeds configured byte budget".into(),
|
||||
));
|
||||
}
|
||||
|
||||
if let Some((current, current_meta)) = self.registry.get_writer_with_meta(conn_id).await {
|
||||
let deadline = writer_send_deadline(self.route_runtime.me_route_blocking_send_timeout);
|
||||
let writer_permit = match reserve_writer_bytes(
|
||||
¤t.byte_budget,
|
||||
writer_byte_permits,
|
||||
writer_reserved_bytes,
|
||||
deadline,
|
||||
&self.stats,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(permit) => permit,
|
||||
Err(WriterByteReserveError::TimedOut) => {
|
||||
self.stats
|
||||
.increment_me_writer_pick_full_total(self.writer_pick_mode());
|
||||
return Err(ProxyError::Proxy(
|
||||
"ME writer byte budget full within blocking send timeout".into(),
|
||||
));
|
||||
}
|
||||
Err(WriterByteReserveError::Closed) => {
|
||||
warn!(
|
||||
writer_id = current.writer_id,
|
||||
"ME writer byte budget closed"
|
||||
);
|
||||
self.remove_writer_and_close_clients(current.writer_id)
|
||||
.await;
|
||||
return self
|
||||
.send_proxy_req(
|
||||
conn_id,
|
||||
target_dc,
|
||||
client_addr,
|
||||
our_addr,
|
||||
payload.as_ref(),
|
||||
proto_flags,
|
||||
tag.as_ref().map(|tag| tag.as_slice()),
|
||||
Some(_permit),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
};
|
||||
let command = WriterCommand::ProxyReq(ProxyReqCommand {
|
||||
conn_id,
|
||||
client_addr,
|
||||
our_addr: current_meta.our_addr,
|
||||
proto_flags,
|
||||
proxy_tag: tag,
|
||||
payload,
|
||||
_permit,
|
||||
writer_permit,
|
||||
});
|
||||
match current.tx.try_send(command) {
|
||||
Ok(()) => {
|
||||
self.note_hybrid_route_success();
|
||||
return Ok(());
|
||||
}
|
||||
Err(TrySendError::Full(cmd)) => {
|
||||
match reserve_writer_command_slot(¤t.tx, deadline).await {
|
||||
Ok(permit) => {
|
||||
permit.send(cmd);
|
||||
self.note_hybrid_route_success();
|
||||
return Ok(());
|
||||
}
|
||||
Err(WriterCommandReserveError::TimedOut) => {
|
||||
self.stats
|
||||
.increment_me_writer_pick_full_total(self.writer_pick_mode());
|
||||
return Err(ProxyError::Proxy(
|
||||
"ME writer channel full within blocking send timeout".into(),
|
||||
));
|
||||
}
|
||||
Err(WriterCommandReserveError::Closed) => {
|
||||
let Some((payload, _permit)) = proxy_req_payload_from_command(cmd)
|
||||
else {
|
||||
return Err(ProxyError::Proxy(
|
||||
"ME writer rejected unexpected command type".into(),
|
||||
));
|
||||
};
|
||||
warn!(writer_id = current.writer_id, "ME writer channel closed");
|
||||
self.remove_writer_and_close_clients(current.writer_id)
|
||||
.await;
|
||||
return self
|
||||
.send_proxy_req(
|
||||
conn_id,
|
||||
target_dc,
|
||||
client_addr,
|
||||
our_addr,
|
||||
payload.as_ref(),
|
||||
proto_flags,
|
||||
tag.as_ref().map(|tag| tag.as_slice()),
|
||||
Some(_permit),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(TrySendError::Closed(cmd)) => {
|
||||
let Some((payload, _permit)) = proxy_req_payload_from_command(cmd) else {
|
||||
return Err(ProxyError::Proxy(
|
||||
"ME writer rejected unexpected command type".into(),
|
||||
));
|
||||
};
|
||||
warn!(writer_id = current.writer_id, "ME writer channel closed");
|
||||
self.remove_writer_and_close_clients(current.writer_id)
|
||||
.await;
|
||||
return self
|
||||
.send_proxy_req(
|
||||
conn_id,
|
||||
target_dc,
|
||||
client_addr,
|
||||
our_addr,
|
||||
payload.as_ref(),
|
||||
proto_flags,
|
||||
tag.as_ref().map(|tag| tag.as_slice()),
|
||||
Some(_permit),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
self.send_proxy_req(
|
||||
conn_id,
|
||||
target_dc,
|
||||
client_addr,
|
||||
our_addr,
|
||||
payload.as_ref(),
|
||||
proto_flags,
|
||||
tag.as_ref().map(|tag| tag.as_slice()),
|
||||
Some(_permit),
|
||||
)
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,114 @@
|
||||
use std::net::SocketAddr;
|
||||
use std::sync::Arc;
|
||||
|
||||
use tokio::sync::OwnedSemaphorePermit;
|
||||
use tokio::sync::mpsc::error::TrySendError;
|
||||
use tracing::warn;
|
||||
|
||||
use super::super::MePool;
|
||||
use super::super::codec::WriterCommand;
|
||||
use super::super::wire::build_proxy_req_payload;
|
||||
use super::reservation::{
|
||||
WriterByteReserveError, WriterCommandReserveError, payload_permit_from_data_command,
|
||||
reserve_writer_bytes, reserve_writer_command_slot, writer_send_deadline,
|
||||
};
|
||||
use crate::error::{ProxyError, Result};
|
||||
|
||||
pub(super) enum BoundWriterSendOutcome {
|
||||
Sent,
|
||||
Retry(Option<OwnedSemaphorePermit>),
|
||||
}
|
||||
|
||||
impl MePool {
|
||||
pub(super) async fn try_send_bound_writer(
|
||||
self: &Arc<Self>,
|
||||
conn_id: u64,
|
||||
client_addr: SocketAddr,
|
||||
data: &[u8],
|
||||
proto_flags: u32,
|
||||
tag: Option<&[u8]>,
|
||||
writer_byte_permits: u32,
|
||||
writer_reserved_bytes: usize,
|
||||
payload_permit: Option<OwnedSemaphorePermit>,
|
||||
) -> Result<BoundWriterSendOutcome> {
|
||||
let Some((current, current_meta)) = self.registry.get_writer_with_meta(conn_id).await else {
|
||||
return Ok(BoundWriterSendOutcome::Retry(payload_permit));
|
||||
};
|
||||
let deadline = writer_send_deadline(self.route_runtime.me_route_blocking_send_timeout);
|
||||
let writer_permit = match reserve_writer_bytes(
|
||||
¤t.byte_budget,
|
||||
writer_byte_permits,
|
||||
writer_reserved_bytes,
|
||||
deadline,
|
||||
&self.stats,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(permit) => permit,
|
||||
Err(WriterByteReserveError::TimedOut) => {
|
||||
self.stats
|
||||
.increment_me_writer_pick_full_total(self.writer_pick_mode());
|
||||
return Err(ProxyError::Proxy(
|
||||
"ME writer byte budget full within blocking send timeout".into(),
|
||||
));
|
||||
}
|
||||
Err(WriterByteReserveError::Closed) => {
|
||||
warn!(
|
||||
writer_id = current.writer_id,
|
||||
"ME writer byte budget closed"
|
||||
);
|
||||
self.remove_writer_and_close_clients(current.writer_id)
|
||||
.await;
|
||||
return Ok(BoundWriterSendOutcome::Retry(payload_permit));
|
||||
}
|
||||
};
|
||||
let payload = build_proxy_req_payload(
|
||||
conn_id,
|
||||
client_addr,
|
||||
current_meta.our_addr,
|
||||
data,
|
||||
tag,
|
||||
proto_flags,
|
||||
);
|
||||
let command = WriterCommand::Data {
|
||||
payload,
|
||||
_permit: payload_permit,
|
||||
writer_permit,
|
||||
};
|
||||
match current.tx.try_send(command) {
|
||||
Ok(()) => {
|
||||
self.note_hybrid_route_success();
|
||||
Ok(BoundWriterSendOutcome::Sent)
|
||||
}
|
||||
Err(TrySendError::Full(cmd)) => {
|
||||
match reserve_writer_command_slot(¤t.tx, deadline).await {
|
||||
Ok(permit) => {
|
||||
permit.send(cmd);
|
||||
self.note_hybrid_route_success();
|
||||
return Ok(BoundWriterSendOutcome::Sent);
|
||||
}
|
||||
Err(WriterCommandReserveError::TimedOut) => {
|
||||
self.stats
|
||||
.increment_me_writer_pick_full_total(self.writer_pick_mode());
|
||||
return Err(ProxyError::Proxy(
|
||||
"ME writer channel full within blocking send timeout".into(),
|
||||
));
|
||||
}
|
||||
Err(WriterCommandReserveError::Closed) => {}
|
||||
}
|
||||
let payload_permit = payload_permit_from_data_command(cmd);
|
||||
warn!(writer_id = current.writer_id, "ME writer channel closed");
|
||||
self.remove_writer_and_close_clients(current.writer_id)
|
||||
.await;
|
||||
Ok(BoundWriterSendOutcome::Retry(payload_permit))
|
||||
}
|
||||
Err(TrySendError::Closed(cmd)) => {
|
||||
let payload_permit = payload_permit_from_data_command(cmd);
|
||||
warn!(writer_id = current.writer_id, "ME writer channel closed");
|
||||
self.remove_writer_and_close_clients(current.writer_id)
|
||||
.await;
|
||||
Ok(BoundWriterSendOutcome::Retry(payload_permit))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,182 @@
|
||||
use std::net::SocketAddr;
|
||||
use std::sync::Arc;
|
||||
|
||||
use tokio::sync::OwnedSemaphorePermit;
|
||||
use tokio::sync::mpsc::error::TrySendError;
|
||||
use tracing::warn;
|
||||
|
||||
use super::super::MePool;
|
||||
use super::super::codec::{ProxyReqCommand, WriterCommand};
|
||||
use super::reservation::{
|
||||
WriterByteReserveError, WriterCommandReserveError, proxy_req_payload_from_command,
|
||||
proxy_req_resident_permits, proxy_tag_array, reserve_writer_bytes,
|
||||
reserve_writer_command_slot, writer_send_deadline,
|
||||
};
|
||||
use crate::error::{ProxyError, Result};
|
||||
use crate::stream::PooledBuffer;
|
||||
|
||||
impl MePool {
|
||||
/// Send RPC_PROXY_REQ while keeping the first bound-writer path allocation-light.
|
||||
/// The client byte permit follows the payload until writer completion or command drop.
|
||||
pub async fn send_proxy_req_pooled(
|
||||
self: &Arc<Self>,
|
||||
conn_id: u64,
|
||||
target_dc: i16,
|
||||
client_addr: SocketAddr,
|
||||
our_addr: SocketAddr,
|
||||
payload: PooledBuffer,
|
||||
_permit: OwnedSemaphorePermit,
|
||||
proto_flags: u32,
|
||||
tag_override: Option<[u8; 16]>,
|
||||
) -> Result<()> {
|
||||
let tag = tag_override.or_else(|| proxy_tag_array(self.proxy_tag.as_deref()));
|
||||
let Some((writer_byte_permits, writer_reserved_bytes)) = proxy_req_resident_permits(
|
||||
payload.capacity(),
|
||||
payload.len(),
|
||||
tag.as_ref().map(|tag| tag.as_slice()),
|
||||
proto_flags,
|
||||
) else {
|
||||
self.stats.increment_me_writer_byte_budget_oversize_total();
|
||||
return Err(ProxyError::Proxy(
|
||||
"ME writer payload residency calculation overflow".into(),
|
||||
));
|
||||
};
|
||||
if writer_byte_permits as usize > self.writer_lifecycle.writer_byte_budget_permits {
|
||||
self.stats.increment_me_writer_byte_budget_oversize_total();
|
||||
return Err(ProxyError::Proxy(
|
||||
"ME writer payload exceeds configured byte budget".into(),
|
||||
));
|
||||
}
|
||||
|
||||
if let Some((current, current_meta)) = self.registry.get_writer_with_meta(conn_id).await {
|
||||
let deadline = writer_send_deadline(self.route_runtime.me_route_blocking_send_timeout);
|
||||
let writer_permit = match reserve_writer_bytes(
|
||||
¤t.byte_budget,
|
||||
writer_byte_permits,
|
||||
writer_reserved_bytes,
|
||||
deadline,
|
||||
&self.stats,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(permit) => permit,
|
||||
Err(WriterByteReserveError::TimedOut) => {
|
||||
self.stats
|
||||
.increment_me_writer_pick_full_total(self.writer_pick_mode());
|
||||
return Err(ProxyError::Proxy(
|
||||
"ME writer byte budget full within blocking send timeout".into(),
|
||||
));
|
||||
}
|
||||
Err(WriterByteReserveError::Closed) => {
|
||||
warn!(
|
||||
writer_id = current.writer_id,
|
||||
"ME writer byte budget closed"
|
||||
);
|
||||
self.remove_writer_and_close_clients(current.writer_id)
|
||||
.await;
|
||||
return self
|
||||
.send_proxy_req(
|
||||
conn_id,
|
||||
target_dc,
|
||||
client_addr,
|
||||
our_addr,
|
||||
payload.as_ref(),
|
||||
proto_flags,
|
||||
tag.as_ref().map(|tag| tag.as_slice()),
|
||||
Some(_permit),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
};
|
||||
let command = WriterCommand::ProxyReq(ProxyReqCommand {
|
||||
conn_id,
|
||||
client_addr,
|
||||
our_addr: current_meta.our_addr,
|
||||
proto_flags,
|
||||
proxy_tag: tag,
|
||||
payload,
|
||||
_permit,
|
||||
writer_permit,
|
||||
});
|
||||
match current.tx.try_send(command) {
|
||||
Ok(()) => {
|
||||
self.note_hybrid_route_success();
|
||||
return Ok(());
|
||||
}
|
||||
Err(TrySendError::Full(cmd)) => {
|
||||
match reserve_writer_command_slot(¤t.tx, deadline).await {
|
||||
Ok(permit) => {
|
||||
permit.send(cmd);
|
||||
self.note_hybrid_route_success();
|
||||
return Ok(());
|
||||
}
|
||||
Err(WriterCommandReserveError::TimedOut) => {
|
||||
self.stats
|
||||
.increment_me_writer_pick_full_total(self.writer_pick_mode());
|
||||
return Err(ProxyError::Proxy(
|
||||
"ME writer channel full within blocking send timeout".into(),
|
||||
));
|
||||
}
|
||||
Err(WriterCommandReserveError::Closed) => {
|
||||
let Some((payload, _permit)) = proxy_req_payload_from_command(cmd)
|
||||
else {
|
||||
return Err(ProxyError::Proxy(
|
||||
"ME writer rejected unexpected command type".into(),
|
||||
));
|
||||
};
|
||||
warn!(writer_id = current.writer_id, "ME writer channel closed");
|
||||
self.remove_writer_and_close_clients(current.writer_id)
|
||||
.await;
|
||||
return self
|
||||
.send_proxy_req(
|
||||
conn_id,
|
||||
target_dc,
|
||||
client_addr,
|
||||
our_addr,
|
||||
payload.as_ref(),
|
||||
proto_flags,
|
||||
tag.as_ref().map(|tag| tag.as_slice()),
|
||||
Some(_permit),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(TrySendError::Closed(cmd)) => {
|
||||
let Some((payload, _permit)) = proxy_req_payload_from_command(cmd) else {
|
||||
return Err(ProxyError::Proxy(
|
||||
"ME writer rejected unexpected command type".into(),
|
||||
));
|
||||
};
|
||||
warn!(writer_id = current.writer_id, "ME writer channel closed");
|
||||
self.remove_writer_and_close_clients(current.writer_id)
|
||||
.await;
|
||||
return self
|
||||
.send_proxy_req(
|
||||
conn_id,
|
||||
target_dc,
|
||||
client_addr,
|
||||
our_addr,
|
||||
payload.as_ref(),
|
||||
proto_flags,
|
||||
tag.as_ref().map(|tag| tag.as_slice()),
|
||||
Some(_permit),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
self.send_proxy_req(
|
||||
conn_id,
|
||||
target_dc,
|
||||
client_addr,
|
||||
our_addr,
|
||||
payload.as_ref(),
|
||||
proto_flags,
|
||||
tag.as_ref().map(|tag| tag.as_slice()),
|
||||
Some(_permit),
|
||||
)
|
||||
.await
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,143 @@
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use tokio::sync::{OwnedSemaphorePermit, Semaphore, TryAcquireError, mpsc};
|
||||
|
||||
use super::super::codec::{WriterBytePermit, WriterCommand};
|
||||
use crate::config::defaults::ME_WRITER_BYTE_PERMIT_UNIT_BYTES;
|
||||
use crate::stats::Stats;
|
||||
use crate::stream::PooledBuffer;
|
||||
|
||||
const RPC_WRITER_FRAME_CAPACITY_OVERHEAD_BYTES: usize = 27;
|
||||
pub(super) const LEGACY_PROXY_REQ_SOURCE_CAPACITY_OVERHEAD_BYTES: usize = 128;
|
||||
|
||||
pub(super) enum WriterCommandReserveError {
|
||||
Closed,
|
||||
TimedOut,
|
||||
}
|
||||
|
||||
pub(super) enum WriterByteReserveError {
|
||||
Closed,
|
||||
TimedOut,
|
||||
}
|
||||
|
||||
pub(super) fn proxy_tag_array(tag: Option<&[u8]>) -> Option<[u8; 16]> {
|
||||
tag.and_then(|tag| <[u8; 16]>::try_from(tag).ok())
|
||||
}
|
||||
|
||||
pub(super) fn proxy_req_payload_from_command(
|
||||
cmd: WriterCommand,
|
||||
) -> Option<(PooledBuffer, OwnedSemaphorePermit)> {
|
||||
match cmd {
|
||||
WriterCommand::ProxyReq(command) => Some((command.payload, command._permit)),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn payload_permit_from_data_command(
|
||||
cmd: WriterCommand,
|
||||
) -> Option<OwnedSemaphorePermit> {
|
||||
match cmd {
|
||||
WriterCommand::Data { _permit, .. } => _permit,
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn reserve_writer_command_slot(
|
||||
tx: &mpsc::Sender<WriterCommand>,
|
||||
deadline: Option<Instant>,
|
||||
) -> std::result::Result<mpsc::OwnedPermit<WriterCommand>, WriterCommandReserveError> {
|
||||
let reserve = tx.clone().reserve_owned();
|
||||
match deadline {
|
||||
Some(deadline) => {
|
||||
match tokio::time::timeout(deadline.saturating_duration_since(Instant::now()), reserve)
|
||||
.await
|
||||
{
|
||||
Ok(Ok(permit)) => Ok(permit),
|
||||
Ok(Err(_)) => Err(WriterCommandReserveError::Closed),
|
||||
Err(_) => Err(WriterCommandReserveError::TimedOut),
|
||||
}
|
||||
}
|
||||
None => reserve.await.map_err(|_| WriterCommandReserveError::Closed),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn writer_send_deadline(wait: Option<Duration>) -> Option<Instant> {
|
||||
wait.map(|wait| Instant::now() + wait)
|
||||
}
|
||||
|
||||
fn writer_resident_permits(
|
||||
source_capacity: usize,
|
||||
encoded_payload_len: usize,
|
||||
) -> Option<(u32, usize)> {
|
||||
let resident_bytes = source_capacity
|
||||
.checked_add(encoded_payload_len)?
|
||||
.checked_add(RPC_WRITER_FRAME_CAPACITY_OVERHEAD_BYTES)?;
|
||||
let permits = resident_bytes.div_ceil(ME_WRITER_BYTE_PERMIT_UNIT_BYTES);
|
||||
let permits = u32::try_from(permits).ok()?;
|
||||
let reserved_bytes = (permits as usize).checked_mul(ME_WRITER_BYTE_PERMIT_UNIT_BYTES)?;
|
||||
Some((
|
||||
permits.max(1),
|
||||
reserved_bytes.max(ME_WRITER_BYTE_PERMIT_UNIT_BYTES),
|
||||
))
|
||||
}
|
||||
|
||||
pub(super) fn proxy_req_resident_permits(
|
||||
source_capacity: usize,
|
||||
data_len: usize,
|
||||
proxy_tag: Option<&[u8]>,
|
||||
proto_flags: u32,
|
||||
) -> Option<(u32, usize)> {
|
||||
writer_resident_permits(
|
||||
source_capacity,
|
||||
super::super::wire::proxy_req_payload_len(data_len, proxy_tag, proto_flags),
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) fn try_reserve_writer_bytes(
|
||||
byte_budget: &Arc<Semaphore>,
|
||||
permits: u32,
|
||||
reserved_bytes: usize,
|
||||
stats: &Arc<Stats>,
|
||||
) -> std::result::Result<WriterBytePermit, TryAcquireError> {
|
||||
byte_budget
|
||||
.clone()
|
||||
.try_acquire_many_owned(permits)
|
||||
.map(|permit| WriterBytePermit::new(permit, reserved_bytes, stats.clone()))
|
||||
}
|
||||
|
||||
pub(super) async fn reserve_writer_bytes(
|
||||
byte_budget: &Arc<Semaphore>,
|
||||
permits: u32,
|
||||
reserved_bytes: usize,
|
||||
deadline: Option<Instant>,
|
||||
stats: &Arc<Stats>,
|
||||
) -> std::result::Result<WriterBytePermit, WriterByteReserveError> {
|
||||
match try_reserve_writer_bytes(byte_budget, permits, reserved_bytes, stats) {
|
||||
Ok(permit) => return Ok(permit),
|
||||
Err(TryAcquireError::Closed) => return Err(WriterByteReserveError::Closed),
|
||||
Err(TryAcquireError::NoPermits) => {
|
||||
stats.increment_me_writer_byte_budget_wait_total();
|
||||
}
|
||||
}
|
||||
|
||||
let acquire = byte_budget.clone().acquire_many_owned(permits);
|
||||
match deadline {
|
||||
Some(deadline) => {
|
||||
match tokio::time::timeout(deadline.saturating_duration_since(Instant::now()), acquire)
|
||||
.await
|
||||
{
|
||||
Ok(Ok(permit)) => Ok(WriterBytePermit::new(permit, reserved_bytes, stats.clone())),
|
||||
Ok(Err(_)) => Err(WriterByteReserveError::Closed),
|
||||
Err(_) => {
|
||||
stats.increment_me_writer_byte_budget_timeout_total();
|
||||
Err(WriterByteReserveError::TimedOut)
|
||||
}
|
||||
}
|
||||
}
|
||||
None => acquire
|
||||
.await
|
||||
.map(|permit| WriterBytePermit::new(permit, reserved_bytes, stats.clone()))
|
||||
.map_err(|_| WriterByteReserveError::Closed),
|
||||
}
|
||||
}
|
||||
@@ -1,3 +1,4 @@
|
||||
use std::cmp::Reverse;
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::sync::atomic::Ordering;
|
||||
|
||||
@@ -7,6 +8,7 @@ use super::{
|
||||
IDLE_WRITER_PENALTY_HIGH_SECS, IDLE_WRITER_PENALTY_MID_SECS, PICK_PENALTY_DEGRADED,
|
||||
PICK_PENALTY_DRAINING, PICK_PENALTY_STALE, PICK_PENALTY_WARM,
|
||||
};
|
||||
use crate::config::MeWriterPickMode;
|
||||
|
||||
impl MePool {
|
||||
pub(super) async fn candidate_indices_for_dc(
|
||||
@@ -184,4 +186,94 @@ impl MePool {
|
||||
}
|
||||
ordered
|
||||
}
|
||||
|
||||
pub(super) async fn ordered_candidate_indices(
|
||||
&self,
|
||||
mut candidate_indices: Vec<usize>,
|
||||
writers_snapshot: &[super::super::pool::MeWriter],
|
||||
pick_mode: MeWriterPickMode,
|
||||
) -> Vec<usize> {
|
||||
let pick_sample_size = self.writer_pick_sample_size();
|
||||
let writer_ids: Vec<u64> = candidate_indices
|
||||
.iter()
|
||||
.map(|idx| writers_snapshot[*idx].id)
|
||||
.collect();
|
||||
let writer_idle_since = self
|
||||
.registry
|
||||
.writer_idle_since_for_writer_ids(&writer_ids)
|
||||
.await;
|
||||
let now_epoch_secs = Self::now_epoch_secs();
|
||||
let start = self.rr.fetch_add(1, Ordering::Relaxed) as usize % candidate_indices.len();
|
||||
if pick_mode == MeWriterPickMode::P2c {
|
||||
return self.p2c_ordered_candidate_indices(
|
||||
&candidate_indices,
|
||||
writers_snapshot,
|
||||
&writer_idle_since,
|
||||
now_epoch_secs,
|
||||
start,
|
||||
pick_sample_size,
|
||||
);
|
||||
}
|
||||
|
||||
if self
|
||||
.writer_selection_policy
|
||||
.me_deterministic_writer_sort
|
||||
.load(Ordering::Relaxed)
|
||||
{
|
||||
candidate_indices.sort_by(|lhs, rhs| {
|
||||
let left = &writers_snapshot[*lhs];
|
||||
let right = &writers_snapshot[*rhs];
|
||||
let left_key = (
|
||||
self.writer_contour_rank_for_selection(left),
|
||||
(left.generation < self.current_generation()) as usize,
|
||||
left.degraded.load(Ordering::Relaxed) as usize,
|
||||
self.writer_idle_rank_for_selection(
|
||||
left,
|
||||
&writer_idle_since,
|
||||
now_epoch_secs,
|
||||
),
|
||||
Reverse(left.tx.capacity()),
|
||||
left.addr,
|
||||
left.id,
|
||||
);
|
||||
let right_key = (
|
||||
self.writer_contour_rank_for_selection(right),
|
||||
(right.generation < self.current_generation()) as usize,
|
||||
right.degraded.load(Ordering::Relaxed) as usize,
|
||||
self.writer_idle_rank_for_selection(
|
||||
right,
|
||||
&writer_idle_since,
|
||||
now_epoch_secs,
|
||||
),
|
||||
Reverse(right.tx.capacity()),
|
||||
right.addr,
|
||||
right.id,
|
||||
);
|
||||
left_key.cmp(&right_key)
|
||||
});
|
||||
} else {
|
||||
candidate_indices.sort_by_key(|idx| {
|
||||
let writer = &writers_snapshot[*idx];
|
||||
let degraded = writer.degraded.load(Ordering::Relaxed);
|
||||
let stale = (writer.generation < self.current_generation()) as usize;
|
||||
(
|
||||
self.writer_contour_rank_for_selection(writer),
|
||||
stale,
|
||||
degraded as usize,
|
||||
self.writer_idle_rank_for_selection(
|
||||
writer,
|
||||
&writer_idle_since,
|
||||
now_epoch_secs,
|
||||
),
|
||||
Reverse(writer.tx.capacity()),
|
||||
)
|
||||
});
|
||||
}
|
||||
|
||||
let mut ordered = Vec::<usize>::with_capacity(candidate_indices.len());
|
||||
for offset in 0..candidate_indices.len() {
|
||||
ordered.push(candidate_indices[(start + offset) % candidate_indices.len()]);
|
||||
}
|
||||
ordered
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,15 +1,19 @@
|
||||
use std::collections::HashMap;
|
||||
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
|
||||
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::Ordering;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use crate::config::{GeneralConfig, MeRouteNoWriterMode, MeSocksKdfPolicy, MeWriterPickMode};
|
||||
use crate::crypto::SecureRandom;
|
||||
use crate::network::IpFamily;
|
||||
use crate::network::probe::NetworkDecision;
|
||||
use crate::stats::Stats;
|
||||
|
||||
use super::pool::MePool;
|
||||
use super::pool::{
|
||||
MePool, ReinitStatusSnapshot, WriterContour, WriterRole,
|
||||
};
|
||||
use super::pool_writer_security_tests::make_pool_with_decision;
|
||||
|
||||
async fn make_pool() -> Arc<MePool> {
|
||||
let general = GeneralConfig::default();
|
||||
@@ -177,6 +181,91 @@ async fn refill_does_not_queue_a_removed_dc_target() {
|
||||
assert_eq!(pool.refill_pending.load(Ordering::Acquire), 0);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
async fn refill_accepts_enabled_nonpreferred_family_without_multipath() {
|
||||
let pool = make_pool_with_decision(NetworkDecision {
|
||||
ipv4_me: true,
|
||||
ipv6_me: true,
|
||||
effective_prefer: 4,
|
||||
effective_multipath: false,
|
||||
..NetworkDecision::default()
|
||||
})
|
||||
.await;
|
||||
let v4_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443);
|
||||
let v6_addr = SocketAddr::new(IpAddr::V6(Ipv6Addr::LOCALHOST), 443);
|
||||
pool.update_proxy_maps(
|
||||
HashMap::from([(2, vec![(v4_addr.ip(), v4_addr.port())])]),
|
||||
Some(HashMap::from([(
|
||||
2,
|
||||
vec![(v6_addr.ip(), v6_addr.port())],
|
||||
)])),
|
||||
)
|
||||
.await;
|
||||
|
||||
pool.trigger_immediate_refill_for_dc(v6_addr, 2);
|
||||
|
||||
assert_eq!(pool.refill_states.lock().len(), 1);
|
||||
assert_eq!(pool.refill_running.load(Ordering::Acquire), 1);
|
||||
pool.begin_shutdown();
|
||||
tokio::time::timeout(Duration::from_secs(1), async {
|
||||
while pool.refill_running.load(Ordering::Acquire) != 0 {
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(pool.refill_states.lock().is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
async fn stale_endpoint_revision_cancels_queued_warm_refill_before_connect() {
|
||||
let pool = make_pool().await;
|
||||
let old_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 31, 0, 30)), 443);
|
||||
let new_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 31, 0, 31)), 443);
|
||||
pool.update_proxy_maps(
|
||||
HashMap::from([(2, vec![(old_addr.ip(), old_addr.port())])]),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
let old_revision = pool.endpoint_snapshot.load().revision;
|
||||
let pending_generation = pool.current_generation().saturating_add(1);
|
||||
pool.reinit.status.store(Arc::new(ReinitStatusSnapshot {
|
||||
active_generation: pool.current_generation(),
|
||||
warm_generations: vec![pending_generation],
|
||||
pending_hardswap_generation: pending_generation,
|
||||
pending_hardswap_started_at_epoch_secs: 1,
|
||||
pending_hardswap_map_hash: 1,
|
||||
pending_hardswap_endpoint_revision: old_revision,
|
||||
inflight: 1,
|
||||
}));
|
||||
pool.trigger_immediate_refill_for_role(
|
||||
old_addr,
|
||||
WriterRole {
|
||||
dc: 2,
|
||||
family: IpFamily::V4,
|
||||
generation: pending_generation,
|
||||
contour: WriterContour::Warm,
|
||||
},
|
||||
);
|
||||
assert_eq!(pool.refill_states.lock().len(), 1);
|
||||
|
||||
pool.update_proxy_maps(
|
||||
HashMap::from([(2, vec![(new_addr.ip(), new_addr.port())])]),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
tokio::time::timeout(Duration::from_secs(1), async {
|
||||
while pool.refill_running.load(Ordering::Acquire) != 0 {
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(pool.refill_states.lock().is_empty());
|
||||
assert_eq!(pool.stats.get_me_reconnect_attempts(), 0);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
async fn refill_preserves_bounded_pending_cardinality_and_cleans_up_before_first_poll() {
|
||||
let pool = make_pool().await;
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
|
||||
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU32, AtomicU64, Ordering};
|
||||
use std::time::{Duration, Instant};
|
||||
@@ -8,8 +8,9 @@ use tokio_util::sync::CancellationToken;
|
||||
|
||||
use super::codec::WriterCommand;
|
||||
use super::pool::{MeWriter, WriterContour, WriterOpenIntent};
|
||||
use super::pool_writer_security_tests::make_pool;
|
||||
use super::pool_writer_security_tests::{make_pool, make_pool_with_decision};
|
||||
use super::registry::ConnMeta;
|
||||
use crate::network::probe::NetworkDecision;
|
||||
|
||||
fn unregistered_writer(
|
||||
pool: &Arc<super::pool::MePool>,
|
||||
@@ -128,6 +129,225 @@ async fn normal_active_publication_cannot_race_past_the_family_floor() {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn stale_same_family_writers_do_not_satisfy_current_endpoint_coverage() {
|
||||
let pool = make_pool().await;
|
||||
let current_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443);
|
||||
let stale_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 2)), 443);
|
||||
pool.update_proxy_maps(
|
||||
std::collections::HashMap::from([(
|
||||
2,
|
||||
vec![(current_addr.ip(), current_addr.port())],
|
||||
)]),
|
||||
None,
|
||||
)
|
||||
.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);
|
||||
let generation = pool.current_generation();
|
||||
let required = pool.required_writers_for_dc_with_floor_mode(1, false);
|
||||
pool.writers
|
||||
.write()
|
||||
.await
|
||||
.extend((1..=required.saturating_mul(2)).map(|writer_id| {
|
||||
unregistered_writer(
|
||||
&pool,
|
||||
writer_id as u64,
|
||||
stale_addr,
|
||||
generation,
|
||||
WriterContour::Active,
|
||||
)
|
||||
}));
|
||||
|
||||
assert!(
|
||||
pool.can_open_writer_for_contour(
|
||||
WriterContour::Active,
|
||||
WriterOpenIntent::Coverage,
|
||||
2,
|
||||
current_addr,
|
||||
)
|
||||
.await
|
||||
);
|
||||
assert!(
|
||||
!pool
|
||||
.can_open_writer_for_contour(
|
||||
WriterContour::Active,
|
||||
WriterOpenIntent::Coverage,
|
||||
2,
|
||||
stale_addr,
|
||||
)
|
||||
.await
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn covered_group_does_not_consume_another_group_coverage_slot() {
|
||||
let pool = make_pool().await;
|
||||
let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443);
|
||||
pool.update_proxy_maps(
|
||||
std::collections::HashMap::from([(2, vec![(addr.ip(), addr.port())])]),
|
||||
None,
|
||||
)
|
||||
.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);
|
||||
let generation = pool.current_generation();
|
||||
let required = pool.required_writers_for_dc_with_floor_mode(1, false);
|
||||
pool.writers
|
||||
.write()
|
||||
.await
|
||||
.extend((1..=required).map(|writer_id| {
|
||||
unregistered_writer(
|
||||
&pool,
|
||||
writer_id as u64,
|
||||
addr,
|
||||
generation,
|
||||
WriterContour::Active,
|
||||
)
|
||||
}));
|
||||
|
||||
assert!(
|
||||
!pool
|
||||
.can_open_writer_for_contour(
|
||||
WriterContour::Active,
|
||||
WriterOpenIntent::Coverage,
|
||||
2,
|
||||
addr,
|
||||
)
|
||||
.await
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn normal_active_publication_allows_adaptive_growth_above_family_floor() {
|
||||
let pool = make_pool().await;
|
||||
let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443);
|
||||
pool.update_proxy_maps(
|
||||
std::collections::HashMap::from([(2, vec![(addr.ip(), addr.port())])]),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
let generation = pool.current_generation();
|
||||
let writers = (1..=3)
|
||||
.map(|writer_id| {
|
||||
unregistered_writer(
|
||||
&pool,
|
||||
writer_id,
|
||||
addr,
|
||||
generation,
|
||||
WriterContour::Active,
|
||||
)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let candidate = unregistered_writer(&pool, 4, addr, generation, WriterContour::Active);
|
||||
|
||||
assert!(
|
||||
pool.authorize_writer_publication_capacity(
|
||||
&candidate,
|
||||
WriterContour::Active,
|
||||
WriterOpenIntent::Normal,
|
||||
&writers,
|
||||
)
|
||||
.is_ok()
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn normal_active_publication_rejects_configured_contour_cap() {
|
||||
let pool = make_pool().await;
|
||||
let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443);
|
||||
pool.update_proxy_maps(
|
||||
std::collections::HashMap::from([(2, vec![(addr.ip(), addr.port())])]),
|
||||
None,
|
||||
)
|
||||
.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(3, Ordering::Relaxed);
|
||||
pool.floor_runtime
|
||||
.me_adaptive_floor_max_active_writers_global
|
||||
.store(3, Ordering::Relaxed);
|
||||
let generation = pool.current_generation();
|
||||
let writers = (1..=3)
|
||||
.map(|writer_id| {
|
||||
unregistered_writer(
|
||||
&pool,
|
||||
writer_id,
|
||||
addr,
|
||||
generation,
|
||||
WriterContour::Active,
|
||||
)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let candidate = unregistered_writer(&pool, 4, addr, generation, WriterContour::Active);
|
||||
|
||||
assert!(
|
||||
pool.authorize_writer_publication_capacity(
|
||||
&candidate,
|
||||
WriterContour::Active,
|
||||
WriterOpenIntent::Normal,
|
||||
&writers,
|
||||
)
|
||||
.is_err()
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn nonpreferred_enabled_family_retains_writer_publication_authority() {
|
||||
let pool = make_pool_with_decision(NetworkDecision {
|
||||
ipv4_me: true,
|
||||
ipv6_me: true,
|
||||
effective_prefer: 4,
|
||||
effective_multipath: false,
|
||||
..NetworkDecision::default()
|
||||
})
|
||||
.await;
|
||||
let v4_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443);
|
||||
let v6_addr = SocketAddr::new(IpAddr::V6(Ipv6Addr::LOCALHOST), 443);
|
||||
pool.update_proxy_maps(
|
||||
std::collections::HashMap::from([(2, vec![(v4_addr.ip(), v4_addr.port())])]),
|
||||
Some(std::collections::HashMap::from([(
|
||||
2,
|
||||
vec![(v6_addr.ip(), v6_addr.port())],
|
||||
)])),
|
||||
)
|
||||
.await;
|
||||
let candidate = unregistered_writer(
|
||||
&pool,
|
||||
1,
|
||||
v6_addr,
|
||||
pool.current_generation(),
|
||||
WriterContour::Active,
|
||||
);
|
||||
|
||||
assert!(
|
||||
pool.authorize_writer_publication_capacity(
|
||||
&candidate,
|
||||
WriterContour::Active,
|
||||
WriterOpenIntent::Coverage,
|
||||
&[],
|
||||
)
|
||||
.is_ok()
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn successful_writer_publication_is_fully_visible_and_removable() {
|
||||
let pool = make_pool().await;
|
||||
|
||||
@@ -17,6 +17,15 @@ use crate::stats::Stats;
|
||||
|
||||
/// Builds an isolated ME pool for writer-state tests.
|
||||
pub(super) async fn make_pool() -> Arc<MePool> {
|
||||
make_pool_with_decision(NetworkDecision {
|
||||
ipv4_me: true,
|
||||
..NetworkDecision::default()
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
/// Builds an isolated ME pool with an explicit network-family policy.
|
||||
pub(super) async fn make_pool_with_decision(decision: NetworkDecision) -> Arc<MePool> {
|
||||
let general = GeneralConfig::default();
|
||||
|
||||
MePool::new(
|
||||
@@ -35,10 +44,7 @@ pub(super) async fn make_pool() -> Arc<MePool> {
|
||||
HashMap::new(),
|
||||
HashMap::new(),
|
||||
None,
|
||||
NetworkDecision {
|
||||
ipv4_me: true,
|
||||
..NetworkDecision::default()
|
||||
},
|
||||
decision,
|
||||
None,
|
||||
Arc::new(SecureRandom::new()),
|
||||
Arc::new(Stats::new()),
|
||||
|
||||
@@ -14,7 +14,6 @@ use crate::config::{WebCarrier, WebLimitsConfig};
|
||||
use crate::maestro::generation::RuntimeGeneration;
|
||||
use crate::web::telemetry::{WebRejectionReason, WebTelemetry};
|
||||
use crate::web::trace::WebTraceStore;
|
||||
|
||||
// Credential maps, quotas, and token-bucket helpers remain private to the manager.
|
||||
mod state;
|
||||
// Carrier attempt metadata remains explicit and independent from HTTP parsing.
|
||||
|
||||
@@ -57,6 +57,8 @@ impl WebProcessRuntime {
|
||||
return Err(ManagerError::Closed);
|
||||
};
|
||||
let Some(user_registration) = user_publication.take_registration() else {
|
||||
// Rollback callbacks can retire active user owners and must not run under authority.
|
||||
drop(user_publication);
|
||||
drop(state);
|
||||
self.cancel_replacement(bootstrap_hash, &replacement.old_session);
|
||||
return Err(ManagerError::Closed);
|
||||
@@ -65,6 +67,8 @@ impl WebProcessRuntime {
|
||||
self.record_limit_hit();
|
||||
self.telemetry
|
||||
.record_rejection(crate::web::telemetry::WebRejectionReason::SessionCapacity);
|
||||
drop(user_registration);
|
||||
drop(user_publication);
|
||||
drop(state);
|
||||
self.cancel_replacement(bootstrap_hash, &replacement.old_session);
|
||||
return Err(ManagerError::Limit);
|
||||
@@ -99,6 +103,7 @@ impl WebProcessRuntime {
|
||||
Some(user_registration),
|
||||
);
|
||||
let Some(supersede) = replacement.old_session.prepare_carrier_supersede() else {
|
||||
drop(user_publication);
|
||||
drop(state);
|
||||
self.cancel_replacement(bootstrap_hash, &replacement.old_session);
|
||||
session.close(crate::web::session::SessionCloseReason::Protocol);
|
||||
|
||||
Reference in New Issue
Block a user