Proxy Shared User Drafts

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
@@ -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);
@@ -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;
@@ -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,