astelm/telemt
1
0
Fork 0
mirror of https://github.com/telemt/telemt.git synced 2026-11-01 05:12:02 +03:00
Code Issues Packages Projects Releases Wiki Activity

Proxy Shared User Drafts

Browse Source
This commit is contained in:
Alexey
2026-09-14 20:19:48 +03:00
parent b37f1ebdeb
commit 0ac236955a
59 changed files with 1485 additions and 401 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/transport/middle_proxy/health/family.rs
+4 -3
View File
@@ -25,9 +25,10 @@ pub(super) async fn check_family(
let mut family_degraded = false;
let mut dc_endpoints = HashMap::<i32, Vec<SocketAddr>>::new();
let endpoint_snapshot = pool.endpoint_snapshot.load();
let map_guard = match family {
IpFamily::V4 => pool.proxy_map_v4.read().await,
IpFamily::V6 => pool.proxy_map_v6.read().await,
IpFamily::V4 => &endpoint_snapshot.map_v4,
IpFamily::V6 => &endpoint_snapshot.map_v6,
};
for (dc, addrs) in map_guard.iter() {
let entry = dc_endpoints.entry(*dc).or_default();
@@ -35,7 +36,7 @@ pub(super) async fn check_family(
entry.push(SocketAddr::new(ip, port));
}
}
drop(map_guard);
drop(endpoint_snapshot);
for endpoints in dc_endpoints.values_mut() {
endpoints.sort_unstable();
endpoints.dedup();
src/transport/middle_proxy/ping.rs
+3 -2
View File
@@ -329,13 +329,14 @@ mod tests {
pub async fn run_me_ping(pool: &Arc<MePool>, rng: &SecureRandom) -> Vec<MePingReport> {
let mut reports = Vec::new();
let endpoint_snapshot = pool.endpoint_snapshot.load_full();
let v4_map = if pool.decision.ipv4_me {
pool.proxy_map_v4.read().await.clone()
endpoint_snapshot.map_v4.clone()
} else {
HashMap::new()
};
let v6_map = if pool.decision.ipv6_me {
pool.proxy_map_v6.read().await.clone()
endpoint_snapshot.map_v6.clone()
} else {
HashMap::new()
};
src/transport/middle_proxy/pool.rs
+14 -4
View File
@@ -266,7 +266,17 @@ pub struct RoutingCore {
pub(super) writers: Arc<WritersState>,
pub(super) rr: AtomicU64,
pub(super) writer_epoch: watch::Sender<u64>,
pub(super) preferred_endpoints_by_dc: ArcSwap<HashMap<i32, Vec<SocketAddr>>>,
pub(super) endpoint_snapshot: ArcSwap<EndpointSnapshot>,
}
/// Immutable endpoint routing authority published as one coherent revision.
#[derive(Clone, Debug)]
pub(super) struct EndpointSnapshot {
pub(super) revision: u64,
pub(super) map_v4: HashMap<i32, Vec<(IpAddr, u16)>>,
pub(super) map_v6: HashMap<i32, Vec<(IpAddr, u16)>>,
pub(super) endpoint_dc_map: HashMap<SocketAddr, Option<i32>>,
pub(super) preferred_endpoints_by_dc: HashMap<i32, Vec<SocketAddr>>,
}
pub(super) struct ReinitCore {
@@ -302,12 +312,14 @@ pub(super) struct ReinitPendingState {
pub(super) generation: u64,
pub(super) started_at_epoch_secs: u64,
pub(super) map_hash: u64,
pub(super) endpoint_revision: u64,
}
#[derive(Clone, Copy)]
pub(super) struct ReinitAttemptState {
pub(super) generation: u64,
pub(super) map_hash: u64,
pub(super) endpoint_revision: u64,
pub(super) hardswap: bool,
pub(super) committed: bool,
}
@@ -316,6 +328,7 @@ pub(super) struct ReinitCoordinatorState {
pub(super) next_attempt_id: u64,
pub(super) active_generation: u64,
pub(super) desired_map_hash: u64,
pub(super) endpoint_revision: u64,
pub(super) pending: Option<ReinitPendingState>,
pub(super) attempts: HashMap<u64, ReinitAttemptState>,
}
@@ -475,9 +488,6 @@ pub struct MePool {
pub(super) rng: Arc<SecureRandom>,
pub(super) proxy_tag: Option<Vec<u8>>,
pub(super) proxy_secret: Arc<RwLock<SecretSnapshot>>,
pub(super) proxy_map_v4: Arc<RwLock<HashMap<i32, Vec<(IpAddr, u16)>>>>,
pub(super) proxy_map_v6: Arc<RwLock<HashMap<i32, Vec<(IpAddr, u16)>>>>,
pub(super) endpoint_dc_map: Arc<RwLock<HashMap<SocketAddr, Option<i32>>>>,
pub(super) default_dc: AtomicI32,
pub(super) next_writer_id: AtomicU64,
pub(super) writer_connect_active_reserved: AtomicUsize,
src/transport/middle_proxy/pool/construction.rs
+8 -7
View File
@@ -120,9 +120,12 @@ impl MePool {
me_route_inline_recovery_wait_ms: u64,
me_connection_cleanup_capacity: usize,
) -> Arc<Self> {
let endpoint_dc_map = Self::build_endpoint_dc_map_from_maps(&proxy_map_v4, &proxy_map_v6);
let preferred_endpoints_by_dc =
Self::build_preferred_endpoints_by_dc(&decision, &proxy_map_v4, &proxy_map_v6);
let endpoint_snapshot = Self::build_endpoint_snapshot(
&decision,
proxy_map_v4,
proxy_map_v6,
1,
);
let registry = Arc::new(ConnRegistry::with_route_and_cleanup_capacity(
me_route_channel_capacity,
me_connection_cleanup_capacity,
@@ -149,7 +152,7 @@ impl MePool {
writers: Arc::new(WritersState::new()),
rr: AtomicU64::new(0),
writer_epoch,
preferred_endpoints_by_dc: ArcSwap::from_pointee(preferred_endpoints_by_dc),
endpoint_snapshot: ArcSwap::from_pointee(endpoint_snapshot),
}),
reinit: Arc::new(ReinitCore {
generation: AtomicU64::new(1),
@@ -164,6 +167,7 @@ impl MePool {
next_attempt_id: 1,
active_generation: 1,
desired_map_hash: 0,
endpoint_revision: 1,
pending: None,
attempts: HashMap::new(),
}),
@@ -391,9 +395,6 @@ impl MePool {
})),
stats,
pool_size: 2,
proxy_map_v4: Arc::new(RwLock::new(proxy_map_v4)),
proxy_map_v6: Arc::new(RwLock::new(proxy_map_v6)),
endpoint_dc_map: Arc::new(RwLock::new(endpoint_dc_map)),
default_dc: AtomicI32::new(default_dc.unwrap_or(2)),
next_writer_id: AtomicU64::new(1),
writer_connect_active_reserved: AtomicUsize::new(0),
src/transport/middle_proxy/pool/routing.rs
+64 -18
View File
@@ -76,16 +76,23 @@ impl MePool {
&self,
dc: i32,
) -> bool {
let snapshot = self.endpoint_snapshot.load();
if self.decision.ipv4_me {
let map = self.proxy_map_v4.read().await;
if map.get(&dc).is_some_and(|endpoints| !endpoints.is_empty()) {
if snapshot
.map_v4
.get(&dc)
.is_some_and(|endpoints| !endpoints.is_empty())
{
return true;
}
}
if self.decision.ipv6_me {
let map = self.proxy_map_v6.read().await;
if map.get(&dc).is_some_and(|endpoints| !endpoints.is_empty()) {
if snapshot
.map_v6
.get(&dc)
.is_some_and(|endpoints| !endpoints.is_empty())
{
return true;
}
}
@@ -112,7 +119,12 @@ impl MePool {
&self,
addr: SocketAddr,
) -> i32 {
if let Some(cached) = self.endpoint_dc_map.read().await.get(&addr).copied()
if let Some(cached) = self
.endpoint_snapshot
.load()
.endpoint_dc_map
.get(&addr)
.copied()
&& let Some(dc) = cached
{
return dc;
@@ -125,9 +137,45 @@ impl MePool {
&self,
family: IpFamily,
) -> HashMap<i32, Vec<(IpAddr, u16)>> {
let snapshot = self.endpoint_snapshot.load();
match family {
IpFamily::V4 => self.proxy_map_v4.read().await.clone(),
IpFamily::V6 => self.proxy_map_v6.read().await.clone(),
IpFamily::V4 => snapshot.map_v4.clone(),
IpFamily::V6 => snapshot.map_v6.clone(),
}
}
pub(in crate::transport::middle_proxy) fn build_endpoint_snapshot(
decision: &NetworkDecision,
mut map_v4: HashMap<i32, Vec<(IpAddr, u16)>>,
mut map_v6: HashMap<i32, Vec<(IpAddr, u16)>>,
revision: u64,
) -> EndpointSnapshot {
Self::mirror_negative_dcs(&mut map_v4);
Self::mirror_negative_dcs(&mut map_v6);
let endpoint_dc_map = Self::build_endpoint_dc_map_from_maps(&map_v4, &map_v6);
let preferred_endpoints_by_dc =
Self::build_preferred_endpoints_by_dc(decision, &map_v4, &map_v6);
EndpointSnapshot {
revision,
map_v4,
map_v6,
endpoint_dc_map,
preferred_endpoints_by_dc,
}
}
fn mirror_negative_dcs(map: &mut HashMap<i32, Vec<(IpAddr, u16)>>) {
let positive_dcs = map
.keys()
.copied()
.filter(|dc| *dc > 0)
.collect::<Vec<_>>();
for dc in positive_dcs {
if !map.contains_key(&-dc)
&& let Some(endpoints) = map.get(&dc).cloned()
{
map.insert(-dc, endpoints);
}
}
}
@@ -224,17 +272,11 @@ impl MePool {
endpoint_dc_map
}
pub(in crate::transport::middle_proxy) async fn rebuild_endpoint_dc_map(&self) {
let map_v4 = self.proxy_map_v4.read().await.clone();
let map_v6 = self.proxy_map_v6.read().await.clone();
let rebuilt = Self::build_endpoint_dc_map_from_maps(&map_v4, &map_v6);
let preferred = Self::build_preferred_endpoints_by_dc(&self.decision, &map_v4, &map_v6);
*self.endpoint_dc_map.write().await = rebuilt;
self.preferred_endpoints_by_dc.store(Arc::new(preferred));
pub(in crate::transport::middle_proxy) async fn prune_endpoint_runtime_state(&self) {
let configured_endpoints = self
.endpoint_snapshot
.load()
.endpoint_dc_map
.read()
.await
.keys()
.copied()
.collect::<HashSet<SocketAddr>>();
@@ -253,8 +295,12 @@ impl MePool {
&self,
dc: i32,
) -> Vec<SocketAddr> {
let guard = self.preferred_endpoints_by_dc.load();
guard.get(&dc).cloned().unwrap_or_default()
self.endpoint_snapshot
.load()
.preferred_endpoints_by_dc
.get(&dc)
.cloned()
.unwrap_or_default()
}
pub(in crate::transport::middle_proxy) fn health_interval_unhealthy(&self) -> Duration {
src/transport/middle_proxy/pool/writer_admission.rs
+6 -6
View File
@@ -65,10 +65,10 @@ 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();
if self.family_enabled_for_drain_coverage(IpFamily::V4, now_epoch_secs) {
let map = self.proxy_map_v4.read().await;
for addrs in map.values() {
for addrs in endpoint_snapshot.map_v4.values() {
let mut endpoints = HashSet::<SocketAddr>::new();
for (ip, port) in addrs.iter().copied() {
endpoints.insert(SocketAddr::new(ip, port));
@@ -80,8 +80,7 @@ impl MePool {
}
if self.family_enabled_for_drain_coverage(IpFamily::V6, now_epoch_secs) {
let map = self.proxy_map_v6.read().await;
for addrs in map.values() {
for addrs in endpoint_snapshot.map_v6.values() {
let mut endpoints = HashSet::<SocketAddr>::new();
for (ip, port) in addrs.iter().copied() {
endpoints.insert(SocketAddr::new(ip, port));
@@ -118,13 +117,14 @@ impl MePool {
let mut endpoints_len = 0;
let now_epoch = Self::now_epoch_secs();
let endpoint_snapshot = self.endpoint_snapshot.load();
if self.family_enabled_for_drain_coverage(IpFamily::V4, now_epoch) {
if let Some(addrs) = self.proxy_map_v4.read().await.get(&writer_dc) {
if let Some(addrs) = endpoint_snapshot.map_v4.get(&writer_dc) {
endpoints_len += addrs.len();
}
}
if self.family_enabled_for_drain_coverage(IpFamily::V6, now_epoch) {
if let Some(addrs) = self.proxy_map_v6.read().await.get(&writer_dc) {
if let Some(addrs) = endpoint_snapshot.map_v6.get(&writer_dc) {
endpoints_len += addrs.len();
}
}
src/transport/middle_proxy/pool_config.rs
+27 -40
View File
@@ -31,48 +31,35 @@ impl MePool {
return SnapshotApplyOutcome::RejectedEmpty;
}
let mut changed = false;
{
let mut guard = self.proxy_map_v4.write().await;
if !new_v4.is_empty() && *guard != new_v4 {
*guard = new_v4;
changed = true;
let changed = {
// Endpoint publication and reinit commit share this barrier.
let mut coordinator = self.reinit.coordinator.lock();
let current = self.endpoint_snapshot.load_full();
let map_v4 = if new_v4.is_empty() {
current.map_v4.clone()
} else {
new_v4
};
let map_v6 = match new_v6 {
Some(map) if !map.is_empty() => map,
_ => current.map_v6.clone(),
};
let candidate = Self::build_endpoint_snapshot(
&self.decision,
map_v4,
map_v6,
current.revision.saturating_add(1),
);
if candidate.map_v4 == current.map_v4 && candidate.map_v6 == current.map_v6 {
false
} else {
coordinator.endpoint_revision = candidate.revision;
self.endpoint_snapshot.store(Arc::new(candidate));
true
}
}
if let Some(v6) = new_v6 {
let mut guard = self.proxy_map_v6.write().await;
if !v6.is_empty() && *guard != v6 {
*guard = v6;
changed = true;
}
}
// Ensure negative DC entries mirror positives when absent (Telegram convention).
{
let mut guard = self.proxy_map_v4.write().await;
let keys: Vec<i32> = guard.keys().cloned().collect();
for k in keys.iter().cloned().filter(|k| *k > 0) {
if !guard.contains_key(&-k)
&& let Some(addrs) = guard.get(&k).cloned()
{
guard.insert(-k, addrs);
changed = true;
}
}
}
{
let mut guard = self.proxy_map_v6.write().await;
let keys: Vec<i32> = guard.keys().cloned().collect();
for k in keys.iter().cloned().filter(|k| *k > 0) {
if !guard.contains_key(&-k)
&& let Some(addrs) = guard.get(&k).cloned()
{
guard.insert(-k, addrs);
changed = true;
}
}
}
};
if changed {
self.rebuild_endpoint_dc_map().await;
self.prune_endpoint_runtime_state().await;
self.notify_writer_epoch();
}
if changed {
src/transport/middle_proxy/pool_init.rs
+1 -1
View File
@@ -19,7 +19,7 @@ impl MePool {
.me_reconnect_max_concurrent_per_dc
.max(1) as usize;
let ks = self.key_selector().await;
let me_servers = self.proxy_map_v4.read().await.len();
let me_servers = self.endpoint_snapshot.load().map_v4.len();
let secret_len = self.proxy_secret.read().await.secret.len();
info!(
me_servers,
src/transport/middle_proxy/pool_refill.rs
+7 -5
View File
@@ -70,9 +70,9 @@ impl Drop for RefillRunGuard {
impl MePool {
pub(super) async fn sweep_endpoint_quarantine(&self) {
let configured = self
.endpoint_snapshot
.load()
.endpoint_dc_map
.read()
.await
.keys()
.copied()
.collect::<HashSet<SocketAddr>>();
@@ -266,9 +266,10 @@ impl MePool {
if !self.family_enabled_for_drain_coverage(target.family, now_epoch_secs) {
return Vec::new();
}
let snapshot = self.endpoint_snapshot.load();
let map = match target.family {
IpFamily::V4 => self.proxy_map_v4.read().await,
IpFamily::V6 => self.proxy_map_v6.read().await,
IpFamily::V4 => &snapshot.map_v4,
IpFamily::V6 => &snapshot.map_v6,
};
let mut endpoints = map
.get(&target.dc)
@@ -294,8 +295,9 @@ impl MePool {
};
role_is_authoritative
&& self
.preferred_endpoints_by_dc
.endpoint_snapshot
.load()
.preferred_endpoints_by_dc
.get(&target.dc)
.is_some_and(|endpoints| {
endpoints.iter().any(|endpoint| match target.family {
src/transport/middle_proxy/pool_reinit.rs
+12 -4
View File
@@ -15,8 +15,8 @@ use crate::config::MeBindStaleMode;
use crate::network::IpFamily;
use super::pool::{
MeDrainGateReason, MePool, ReinitAttemptState, ReinitCoordinatorState, ReinitCore,
ReinitPendingState, ReinitStatusSnapshot, WriterContour, WriterOpenIntent,
EndpointSnapshot, MeDrainGateReason, MePool, ReinitAttemptState, ReinitCoordinatorState,
ReinitCore, ReinitPendingState, ReinitStatusSnapshot, WriterContour, WriterOpenIntent,
};
// Reinitialization admission, generation state, and coverage checks.
@@ -34,6 +34,7 @@ struct ReinitAttemptGuard {
generation: u64,
previous_generation: u64,
map_hash: u64,
endpoint_revision: u64,
hardswap: bool,
}
@@ -119,17 +120,24 @@ fn commit_reinit_state(
attempt_id: u64,
generation: u64,
map_hash: u64,
endpoint_revision: u64,
hardswap: bool,
) -> bool {
let Some(record) = state.attempts.get(&attempt_id).copied() else {
return false;
};
if record.map_hash != state.desired_map_hash || record.map_hash != map_hash {
if record.map_hash != state.desired_map_hash
|| record.map_hash != map_hash
|| record.endpoint_revision != endpoint_revision
|| record.endpoint_revision != state.endpoint_revision
{
return false;
}
if hardswap {
let pending_matches = state.pending.is_some_and(|pending| {
pending.generation == generation && pending.map_hash == map_hash
pending.generation == generation
&& pending.map_hash == map_hash
&& pending.endpoint_revision == endpoint_revision
});
if !pending_matches || generation < state.active_generation {
return false;
src/transport/middle_proxy/pool_reinit/coordination.rs
+38 -15
View File
@@ -27,9 +27,13 @@ impl MePool {
self: &Arc<Self>,
hardswap: bool,
map_hash: u64,
endpoint_revision: u64,
now_epoch_secs: u64,
) -> ReinitReservation {
) -> Option<ReinitReservation> {
let mut state = self.reinit.coordinator.lock();
if state.endpoint_revision != endpoint_revision {
return None;
}
state.desired_map_hash = map_hash;
let previous_generation = state.active_generation;
let mut pending_reused = false;
@@ -43,6 +47,7 @@ impl MePool {
&& pending_age_secs > ME_HARDSWAP_PENDING_TTL_SECS;
pending.generation >= previous_generation
&& pending.map_hash == map_hash
&& pending.endpoint_revision == endpoint_revision
&& !pending_expired
});
if let Some(pending) = reusable {
@@ -54,6 +59,7 @@ impl MePool {
generation,
started_at_epoch_secs: now_epoch_secs,
map_hash,
endpoint_revision,
});
generation
}
@@ -69,24 +75,26 @@ impl MePool {
ReinitAttemptState {
generation,
map_hash,
endpoint_revision,
hardswap,
committed: false,
},
);
publish_reinit_state(self.reinit.as_ref(), &state);
ReinitReservation {
Some(ReinitReservation {
attempt: ReinitAttemptGuard {
reinit: Arc::clone(&self.reinit),
attempt_id,
generation,
previous_generation,
map_hash,
endpoint_revision,
hardswap,
},
pending_reused,
pending_expired,
pending_age_secs,
}
})
}
/// Revalidates coverage and commits generation ownership under the publication barrier.
@@ -105,10 +113,13 @@ impl MePool {
if record.generation != attempt.generation
|| record.map_hash != state.desired_map_hash
|| record.map_hash != attempt.map_hash
|| record.endpoint_revision != attempt.endpoint_revision
|| record.endpoint_revision != state.endpoint_revision
|| (attempt.hardswap
&& !state.pending.is_some_and(|pending| {
pending.generation == attempt.generation
&& pending.map_hash == attempt.map_hash
&& pending.endpoint_revision == attempt.endpoint_revision
}))
{
return Err(ReinitCommitFailure::Superseded);
@@ -150,6 +161,7 @@ impl MePool {
attempt.attempt_id,
attempt.generation,
attempt.map_hash,
attempt.endpoint_revision,
attempt.hardswap,
) {
return Err(ReinitCommitFailure::Superseded);
@@ -255,9 +267,13 @@ impl MePool {
/// Restores at least one active writer for every enabled desired DC group.
pub async fn reconcile_connections(self: &Arc<Self>, rng: &SecureRandom) {
let endpoint_snapshot = self.endpoint_snapshot.load_full();
for family in self.family_order() {
let map = self.proxy_map_for_family(family).await;
for (dc, addrs) in &map {
let map = match family {
IpFamily::V4 => &endpoint_snapshot.map_v4,
IpFamily::V6 => &endpoint_snapshot.map_v6,
};
for (dc, addrs) in map {
let dc_addrs: Vec<SocketAddr> = addrs
.iter()
.map(|(ip, port)| SocketAddr::new(*ip, *port))
@@ -286,26 +302,32 @@ impl MePool {
/// Returns the currently authoritative endpoint set for drain and coverage decisions.
pub(in crate::transport::middle_proxy) async fn desired_dc_endpoints(
&self,
) -> HashMap<i32, HashSet<SocketAddr>> {
let endpoint_snapshot = self.endpoint_snapshot.load_full();
self.desired_dc_endpoints_from_snapshot(&endpoint_snapshot)
}
pub(super) fn desired_dc_endpoints_from_snapshot(
&self,
endpoint_snapshot: &EndpointSnapshot,
) -> HashMap<i32, HashSet<SocketAddr>> {
let now_epoch_secs = Self::now_epoch_secs();
let mut out: HashMap<i32, HashSet<SocketAddr>> = HashMap::new();
if self.family_enabled_for_drain_coverage(IpFamily::V4, now_epoch_secs) {
let map_v4 = self.proxy_map_v4.read().await.clone();
for (dc, addrs) in map_v4 {
let entry = out.entry(dc).or_default();
for (dc, addrs) in &endpoint_snapshot.map_v4 {
let entry = out.entry(*dc).or_default();
for (ip, port) in addrs {
entry.insert(SocketAddr::new(ip, port));
entry.insert(SocketAddr::new(*ip, *port));
}
}
}
if self.family_enabled_for_drain_coverage(IpFamily::V6, now_epoch_secs) {
let map_v6 = self.proxy_map_v6.read().await.clone();
for (dc, addrs) in map_v6 {
let entry = out.entry(dc).or_default();
for (dc, addrs) in &endpoint_snapshot.map_v6 {
let entry = out.entry(*dc).or_default();
for (ip, port) in addrs {
entry.insert(SocketAddr::new(ip, port));
entry.insert(SocketAddr::new(*ip, *port));
}
}
}
@@ -320,7 +342,8 @@ impl MePool {
let state = self.reinit.coordinator.lock();
let active_generation = state.active_generation;
let pending_generation = state.pending.map(|pending| pending.generation);
let preferred = self.preferred_endpoints_by_dc.load();
let 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;
@@ -362,7 +385,7 @@ impl MePool {
self.apply_writer_draining_state(writer, self.force_close_timeout(), false);
changed = changed.saturating_add(1);
}
drop(preferred);
drop(endpoint_snapshot);
drop(state);
drop(registry_registration);
drop(writers);
src/transport/middle_proxy/pool_reinit/reconcile.rs
+14 -2
View File
@@ -118,7 +118,8 @@ impl MePool {
self: &Arc<Self>,
rng: &SecureRandom,
) -> bool {
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 now_epoch_secs = Self::now_epoch_secs();
let v4_suppressed = self.is_family_temporarily_suppressed(IpFamily::V4, now_epoch_secs);
let v6_suppressed = self.is_family_temporarily_suppressed(IpFamily::V6, now_epoch_secs);
@@ -137,7 +138,18 @@ impl MePool {
let desired_map_hash = Self::desired_map_hash(&desired_by_dc);
let hardswap = self.reinit.hardswap.load(Ordering::Relaxed);
let reservation = self.reserve_reinit_attempt(hardswap, desired_map_hash, now_epoch_secs);
let Some(reservation) = self.reserve_reinit_attempt(
hardswap,
desired_map_hash,
endpoint_snapshot.revision,
now_epoch_secs,
) else {
debug!(
endpoint_revision = endpoint_snapshot.revision,
"ME reinit snapshot superseded before reservation"
);
return false;
};
let attempt = reservation.attempt;
let previous_generation = attempt.previous_generation;
let generation = attempt.generation;
src/transport/middle_proxy/pool_reinit/tests.rs
+86 -8
View File
@@ -115,10 +115,12 @@ fn stale_concurrent_attempt_cannot_regress_active_generation() {
next_attempt_id: 3,
active_generation: 1,
desired_map_hash: 22,
endpoint_revision: 7,
pending: Some(ReinitPendingState {
generation: 3,
started_at_epoch_secs: 1,
map_hash: 22,
endpoint_revision: 7,
}),
attempts: HashMap::from([
(
@@ -126,6 +128,7 @@ fn stale_concurrent_attempt_cannot_regress_active_generation() {
ReinitAttemptState {
generation: 2,
map_hash: 11,
endpoint_revision: 6,
hardswap: true,
committed: false,
},
@@ -135,6 +138,7 @@ fn stale_concurrent_attempt_cannot_regress_active_generation() {
ReinitAttemptState {
generation: 3,
map_hash: 22,
endpoint_revision: 7,
hardswap: true,
committed: false,
},
@@ -142,9 +146,9 @@ fn stale_concurrent_attempt_cannot_regress_active_generation() {
]),
};
assert!(commit_reinit_state(&mut state, 2, 3, 22, true));
assert!(commit_reinit_state(&mut state, 2, 3, 22, 7, true));
assert_eq!(state.active_generation, 3);
assert!(!commit_reinit_state(&mut state, 1, 2, 11, true));
assert!(!commit_reinit_state(&mut state, 1, 2, 11, 6, true));
assert_eq!(state.active_generation, 3);
assert!(state.pending.is_none());
}
@@ -173,7 +177,10 @@ async fn partial_hardswap_is_rejected_when_stale_binding_is_disabled() {
)
.await;
let map_hash = MePool::desired_map_hash(&desired_by_dc);
let reservation = pool.reserve_reinit_attempt(true, map_hash, 100);
let endpoint_revision = pool.endpoint_snapshot.load().revision;
let reservation = pool
.reserve_reinit_attempt(true, map_hash, endpoint_revision, 100)
.expect("endpoint revision must remain current");
insert_writer(
&pool,
201,
@@ -221,7 +228,10 @@ async fn partial_hardswap_preserves_fallback_only_for_missing_dc() {
)
.await;
let map_hash = MePool::desired_map_hash(&desired_by_dc);
let reservation = pool.reserve_reinit_attempt(true, map_hash, 100);
let endpoint_revision = pool.endpoint_snapshot.load().revision;
let reservation = pool
.reserve_reinit_attempt(true, map_hash, endpoint_revision, 100)
.expect("endpoint revision must remain current");
let fresh_dc1 = insert_writer(
&pool,
401,
@@ -274,7 +284,10 @@ async fn complete_hardswap_promotes_fresh_generation_and_retires_old_writers() {
)
.await;
let map_hash = MePool::desired_map_hash(&desired_by_dc);
let reservation = pool.reserve_reinit_attempt(true, map_hash, 100);
let endpoint_revision = pool.endpoint_snapshot.load().revision;
let reservation = pool
.reserve_reinit_attempt(true, map_hash, endpoint_revision, 100)
.expect("endpoint revision must remain current");
let fresh_dc1 = insert_writer(
&pool,
601,
@@ -314,15 +327,80 @@ async fn complete_hardswap_promotes_fresh_generation_and_retires_old_writers() {
);
}
#[tokio::test]
async fn endpoint_revision_change_supersedes_hardswap_before_writer_drain() {
let pool = make_pool().await;
let old_endpoint = addr(1, 2001);
pool.update_proxy_maps(
HashMap::from([(1, vec![(old_endpoint.ip(), old_endpoint.port())])]),
None,
)
.await;
let active_generation = pool.current_generation();
let old_writer = insert_writer(
&pool,
651,
1,
old_endpoint,
active_generation,
WriterContour::Active,
)
.await;
let desired_by_dc = HashMap::from([(1, HashSet::from([old_endpoint]))]);
let map_hash = MePool::desired_map_hash(&desired_by_dc);
let endpoint_revision = pool.endpoint_snapshot.load().revision;
let reservation = pool
.reserve_reinit_attempt(true, map_hash, endpoint_revision, 100)
.expect("endpoint revision must remain current");
let fresh_writer = insert_writer(
&pool,
652,
1,
old_endpoint,
reservation.attempt.generation,
WriterContour::Warm,
)
.await;
let replacement_endpoint = addr(3, 2003);
pool.update_proxy_maps(
HashMap::from([(
1,
vec![(replacement_endpoint.ip(), replacement_endpoint.port())],
)]),
None,
)
.await;
let result = pool
.commit_reinit_attempt(&reservation.attempt, &desired_by_dc, 1.0)
.await;
assert!(matches!(result, Err(ReinitCommitFailure::Superseded)));
assert_eq!(pool.current_generation(), active_generation);
assert!(!old_writer.draining.load(Ordering::Acquire));
assert!(!fresh_writer.draining.load(Ordering::Acquire));
assert_eq!(
WriterContour::from_u8(fresh_writer.contour.load(Ordering::Acquire)),
WriterContour::Warm
);
}
#[tokio::test]
async fn generation_role_reconciliation_promotes_active_warm_and_drains_orphans() {
let pool = make_pool().await;
let endpoint = addr(1, 2001);
pool.preferred_endpoints_by_dc
.store(Arc::new(HashMap::from([(1, vec![endpoint])])));
pool.update_proxy_maps(
HashMap::from([(1, vec![(endpoint.ip(), endpoint.port())])]),
None,
)
.await;
let desired_by_dc = HashMap::from([(1, HashSet::from([endpoint]))]);
let map_hash = MePool::desired_map_hash(&desired_by_dc);
let reservation = pool.reserve_reinit_attempt(true, map_hash, 100);
let endpoint_revision = pool.endpoint_snapshot.load().revision;
let reservation = pool
.reserve_reinit_attempt(true, map_hash, endpoint_revision, 100)
.expect("endpoint revision must remain current");
let active_warm = insert_writer(
&pool,
701,
src/transport/middle_proxy/pool_status/status_snapshot.rs
+9 -6
View File
@@ -3,12 +3,13 @@ use super::*;
impl MePool {
pub(crate) async fn admission_ready_conditional_cast(&self) -> bool {
let mut endpoints_by_dc = BTreeMap::<i16, BTreeSet<SocketAddr>>::new();
let endpoint_snapshot = self.endpoint_snapshot.load_full();
if self.decision.ipv4_me {
let map = self.proxy_map_v4.read().await.clone();
let map = endpoint_snapshot.map_v4.clone();
extend_signed_endpoints(&mut endpoints_by_dc, map);
}
if self.decision.ipv6_me {
let map = self.proxy_map_v6.read().await.clone();
let map = endpoint_snapshot.map_v6.clone();
extend_signed_endpoints(&mut endpoints_by_dc, map);
}
@@ -45,12 +46,13 @@ impl MePool {
#[allow(dead_code)]
pub(crate) async fn admission_ready_full_floor(&self) -> bool {
let mut endpoints_by_dc = BTreeMap::<i16, BTreeSet<SocketAddr>>::new();
let endpoint_snapshot = self.endpoint_snapshot.load_full();
if self.decision.ipv4_me {
let map = self.proxy_map_v4.read().await.clone();
let map = endpoint_snapshot.map_v4.clone();
extend_signed_endpoints(&mut endpoints_by_dc, map);
}
if self.decision.ipv6_me {
let map = self.proxy_map_v6.read().await.clone();
let map = endpoint_snapshot.map_v6.clone();
extend_signed_endpoints(&mut endpoints_by_dc, map);
}
@@ -106,12 +108,13 @@ impl MePool {
.load(Ordering::Relaxed);
let mut endpoints_by_dc = BTreeMap::<i16, BTreeSet<SocketAddr>>::new();
let endpoint_snapshot = self.endpoint_snapshot.load_full();
if self.decision.ipv4_me {
let map = self.proxy_map_v4.read().await.clone();
let map = endpoint_snapshot.map_v4.clone();
extend_signed_endpoints(&mut endpoints_by_dc, map);
}
if self.decision.ipv6_me {
let map = self.proxy_map_v6.read().await.clone();
let map = endpoint_snapshot.map_v6.clone();
extend_signed_endpoints(&mut endpoints_by_dc, map);
}
src/transport/middle_proxy/pool_writer/publication.rs
+5 -2
View File
@@ -41,8 +41,9 @@ impl MePool {
coordinator: &crate::transport::middle_proxy::pool::ReinitCoordinatorState,
) -> Result<WriterContour> {
let endpoint_is_current = self
.preferred_endpoints_by_dc
.endpoint_snapshot
.load()
.preferred_endpoints_by_dc
.get(&writer.writer_dc)
.is_some_and(|endpoints| endpoints.contains(&writer.addr));
if !endpoint_is_current {
@@ -61,6 +62,7 @@ impl MePool {
&& coordinator.pending.is_some_and(|pending| {
pending.generation == writer.generation
&& pending.map_hash == coordinator.desired_map_hash
&& pending.endpoint_revision == coordinator.endpoint_revision
})
{
return Ok(WriterContour::Warm);
@@ -92,7 +94,8 @@ impl MePool {
if intent == WriterOpenIntent::Replacement || contour == WriterContour::Draining {
return Ok(());
}
let preferred = self.preferred_endpoints_by_dc.load();
let endpoint_snapshot = self.endpoint_snapshot.load();
let preferred = &endpoint_snapshot.preferred_endpoints_by_dc;
let Some(endpoints) = preferred.get(&writer.writer_dc) else {
return Err(ProxyError::Proxy(
"ME writer target changed before publication".into(),
src/transport/middle_proxy/pool_writer/replacement.rs
+20 -9
View File
@@ -85,7 +85,8 @@ impl MePool {
"ME floor rebalance lost active-generation authority".into(),
));
}
let preferred = self.preferred_endpoints_by_dc.load();
let endpoint_snapshot = self.endpoint_snapshot.load();
let preferred = &endpoint_snapshot.preferred_endpoints_by_dc;
let donor_count = writers
.iter()
.filter(|candidate| {
@@ -265,8 +266,11 @@ mod tests {
async fn replacement_commit_publishes_successor_before_draining_victim() {
let pool = make_pool().await;
let addr = endpoint(1);
pool.preferred_endpoints_by_dc
.store(Arc::new(HashMap::from([(2, vec![addr])])));
pool.update_proxy_maps(
HashMap::from([(2, vec![(addr.ip(), addr.port())])]),
None,
)
.await;
let victim = install_writer(&pool, 1001, 2, addr).await;
let expected_role = WriterRole::from_writer(&victim);
let mut reservation = pool
@@ -305,8 +309,11 @@ mod tests {
async fn cancelled_replacement_waiting_for_publication_restores_all_reservations() {
let pool = make_pool().await;
let addr = endpoint(2);
pool.preferred_endpoints_by_dc
.store(Arc::new(HashMap::from([(2, vec![addr])])));
pool.update_proxy_maps(
HashMap::from([(2, vec![(addr.ip(), addr.port())])]),
None,
)
.await;
let victim = install_writer(&pool, 2001, 2, addr).await;
let expected_role = WriterRole::from_writer(&victim);
let mut reservation = pool
@@ -342,10 +349,14 @@ mod tests {
let pool = make_pool().await;
let donor_addr = endpoint(3);
let receiver_addr = endpoint(4);
pool.preferred_endpoints_by_dc.store(Arc::new(HashMap::from([
(1, vec![donor_addr]),
(2, vec![receiver_addr]),
])));
pool.update_proxy_maps(
HashMap::from([
(1, vec![(donor_addr.ip(), donor_addr.port())]),
(2, vec![(receiver_addr.ip(), receiver_addr.port())]),
]),
None,
)
.await;
let victim = install_writer(&pool, 3001, 1, donor_addr).await;
let expected_role = WriterRole::from_writer(&victim);
let mut reservation = pool
src/transport/middle_proxy/send.rs
+6 -3
View File
@@ -348,8 +348,10 @@ impl MePool {
for _ in
0..self.route_runtime.me_route_inline_recovery_attempts.max(1)
{
let preferred = self.preferred_endpoints_by_dc.load_full();
for (dc, addrs) in preferred.iter() {
let endpoint_snapshot = self.endpoint_snapshot.load_full();
for (dc, addrs) in
&endpoint_snapshot.preferred_endpoints_by_dc
{
for addr in addrs {
let _ = self
.connect_one_for_dc(*addr, *dc, self.rng.as_ref())
@@ -470,8 +472,9 @@ impl MePool {
}
emergency_attempts += 1;
let mut endpoints = self
.preferred_endpoints_by_dc
.endpoint_snapshot
.load()
.preferred_endpoints_by_dc
.get(&routed_dc)
.cloned()
.unwrap_or_default();
src/transport/middle_proxy/send/recovery.rs
+2 -1
View File
@@ -88,7 +88,8 @@ impl MePool {
pub(super) async fn trigger_async_recovery_global(self: &Arc<Self>) {
self.stats.increment_me_async_recovery_trigger_total();
let preferred = self.preferred_endpoints_by_dc.load();
let endpoint_snapshot = self.endpoint_snapshot.load();
let preferred = &endpoint_snapshot.preferred_endpoints_by_dc;
let mut triggered = 0usize;
for (dc, addrs) in preferred.iter() {
for addr in addrs {
src/transport/middle_proxy/send/selection.rs
+2 -1
View File
@@ -15,7 +15,8 @@ impl MePool {
routed_dc: i32,
include_warm: bool,
) -> Vec<usize> {
let preferred_snapshot = self.preferred_endpoints_by_dc.load();
let endpoint_snapshot = self.endpoint_snapshot.load();
let preferred_snapshot = &endpoint_snapshot.preferred_endpoints_by_dc;
let mut out = Vec::new();
if let Some(preferred) = preferred_snapshot
.get(&routed_dc)
src/transport/middle_proxy/tests/pool_refill_security_tests.rs
+12 -2
View File
@@ -183,8 +183,18 @@ async fn refill_preserves_bounded_pending_cardinality_and_cleans_up_before_first
let first = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 31, 0, 21)), 443);
let second = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 31, 0, 22)), 443);
let latest = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 31, 0, 23)), 443);
pool.preferred_endpoints_by_dc
.store(Arc::new(HashMap::from([(2, vec![first, second, latest])])));
pool.update_proxy_maps(
HashMap::from([(
2,
vec![
(first.ip(), first.port()),
(second.ip(), second.port()),
(latest.ip(), latest.port()),
],
)]),
None,
)
.await;
pool.trigger_immediate_refill_for_dc(first, 2);
pool.trigger_immediate_refill_for_dc(second, 2);
src/transport/middle_proxy/tests/pool_writer_publication_tests.rs
+10 -4
View File
@@ -43,8 +43,11 @@ fn unregistered_writer(
async fn normal_warm_publication_cannot_race_past_the_dc_floor() {
let pool = make_pool().await;
let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443);
pool.preferred_endpoints_by_dc
.store(Arc::new(std::collections::HashMap::from([(2, vec![addr])])));
pool.update_proxy_maps(
std::collections::HashMap::from([(2, vec![(addr.ip(), addr.port())])]),
None,
)
.await;
let generation = 2;
let writers = (1..=3)
.map(|writer_id| {
@@ -86,8 +89,11 @@ async fn normal_warm_publication_cannot_race_past_the_dc_floor() {
async fn normal_active_publication_cannot_race_past_the_family_floor() {
let pool = make_pool().await;
let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443);
pool.preferred_endpoints_by_dc
.store(Arc::new(std::collections::HashMap::from([(2, vec![addr])])));
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| {
src/transport/middle_proxy/tests/send_adversarial_tests.rs
+6 -12
View File
@@ -154,13 +154,11 @@ async fn insert_writer(
};
pool.writers.write().await.push(writer);
{
let mut map = pool.proxy_map_v4.write().await;
map.entry(writer_dc)
.or_insert_with(Vec::new)
.push((addr.ip(), addr.port()));
}
pool.rebuild_endpoint_dc_map().await;
let mut map = pool.endpoint_snapshot.load().map_v4.clone();
map.entry(writer_dc)
.or_insert_with(Vec::new)
.push((addr.ip(), addr.port()));
pool.update_proxy_maps(map, None).await;
if register_in_registry {
pool.registry
.register_writer(writer_id, tx, byte_budget)
@@ -241,11 +239,7 @@ async fn send_proxy_req_uses_live_same_dc_writer_while_preferred_endpoint_refill
assert!(pool.admission_ready_conditional_cast().await);
assert_eq!(
pool.preferred_endpoints_by_dc
.load()
.get(&2)
.cloned()
.unwrap_or_default(),
pool.preferred_endpoints_for_dc(2).await,
vec![new_addr]
);