Atomic Lifecycle for Writer Refresh + Hardswap

src/transport/middle_proxy/pool_writer/publication.rs

@@ -0,0 +1,180 @@
+use super::*;
+
+impl MePool {
+    /// Publishes a connected writer only while its generation owns the requested role.
+    pub(super) async fn publish_connected_writer(
+        self: &Arc<Self>,
+        prepared: PreparedWriter<'_>,
+    ) -> Result<()> {
+        let PreparedWriter {
+            writer,
+            tx,
+            byte_budget,
+            task_registration,
+            writer_task,
+            intent,
+            _open_reservation,
+        } = prepared;
+        let writer_id = writer.id;
+        let mut writers = self.writers.write().await;
+        let mut registry_registration = self.registry.prepare_writer_registration().await;
+        let coordinator = self.reinit.coordinator.lock();
+        let contour = self.authorize_writer_publication(&writer, &coordinator)?;
+        self.authorize_writer_publication_capacity(&writer, contour, intent, writers.as_slice())?;
+        writer.contour.store(contour.as_u8(), Ordering::Release);
+        registry_registration.install(writer_id, tx, byte_budget);
+        writers.push(writer);
+        self.conn_count.fetch_add(1, Ordering::Relaxed);
+        self.lifecycle
+            .spawn_registered_writer(task_registration, writer_task);
+        drop(coordinator);
+        drop(registry_registration);
+        drop(writers);
+        self.notify_writer_epoch();
+        Ok(())
+    }
+
+    /// Resolves the writer role against the linearized reinitialization authority.
+    pub(super) fn authorize_writer_publication(
+        &self,
+        writer: &MeWriter,
+        coordinator: &crate::transport::middle_proxy::pool::ReinitCoordinatorState,
+    ) -> Result<WriterContour> {
+        let endpoint_is_current = self
+            .preferred_endpoints_by_dc
+            .load()
+            .get(&writer.writer_dc)
+            .is_some_and(|endpoints| endpoints.contains(&writer.addr));
+        if !endpoint_is_current {
+            return Err(ProxyError::Proxy(
+                "ME writer target changed before publication".into(),
+            ));
+        }
+
+        let requested = WriterContour::from_u8(writer.contour.load(Ordering::Acquire));
+        if writer.generation == coordinator.active_generation
+            && matches!(requested, WriterContour::Active | WriterContour::Warm)
+        {
+            return Ok(WriterContour::Active);
+        }
+        if requested == WriterContour::Warm
+            && coordinator.pending.is_some_and(|pending| {
+                pending.generation == writer.generation
+                    && pending.map_hash == coordinator.desired_map_hash
+            })
+        {
+            return Ok(WriterContour::Warm);
+        }
+        Err(ProxyError::Proxy(
+            "ME writer generation lost publication authority".into(),
+        ))
+    }
+
+    /// Revalidates role-local capacity at the serialized publication boundary.
+    pub(in crate::transport::middle_proxy) fn authorize_writer_publication_capacity(
+        &self,
+        writer: &MeWriter,
+        contour: WriterContour,
+        intent: WriterOpenIntent,
+        writers: &[MeWriter],
+    ) -> Result<()> {
+        let family = if writer.addr.is_ipv4() {
+            crate::network::IpFamily::V4
+        } else {
+            crate::network::IpFamily::V6
+        };
+        let now_epoch_secs = Self::now_epoch_secs();
+        if !self.family_enabled_for_drain_coverage(family, now_epoch_secs) {
+            return Err(ProxyError::Proxy(
+                "ME writer family lost publication authority".into(),
+            ));
+        }
+        if intent == WriterOpenIntent::Replacement || contour == WriterContour::Draining {
+            return Ok(());
+        }
+        let preferred = self.preferred_endpoints_by_dc.load();
+        let Some(endpoints) = preferred.get(&writer.writer_dc) else {
+            return Err(ProxyError::Proxy(
+                "ME writer target changed before publication".into(),
+            ));
+        };
+        let required = match contour {
+            WriterContour::Active => self.required_writers_for_dc(
+                endpoints
+                    .iter()
+                    .filter(|endpoint| endpoint.is_ipv4() == writer.addr.is_ipv4())
+                    .count(),
+            ),
+            WriterContour::Warm => self.required_writers_for_dc(
+                endpoints
+                    .iter()
+                    .filter(|endpoint| {
+                        let endpoint_family = if endpoint.is_ipv4() {
+                            crate::network::IpFamily::V4
+                        } else {
+                            crate::network::IpFamily::V6
+                        };
+                        self.family_enabled_for_drain_coverage(
+                            endpoint_family,
+                            now_epoch_secs,
+                        )
+                    })
+                    .count(),
+            ),
+            WriterContour::Draining => 0,
+        };
+        let current = writers
+            .iter()
+            .filter(|candidate| {
+                !candidate.draining.load(Ordering::Acquire)
+                    && candidate.writer_dc == writer.writer_dc
+                    && candidate.generation == writer.generation
+                    && WriterContour::from_u8(candidate.contour.load(Ordering::Acquire))
+                        == contour
+                    && (contour == WriterContour::Warm
+                        || candidate.addr.is_ipv4() == writer.addr.is_ipv4())
+                    && endpoints.contains(&candidate.addr)
+                    && (contour != WriterContour::Warm
+                        || self.family_enabled_for_drain_coverage(
+                            if candidate.addr.is_ipv4() {
+                                crate::network::IpFamily::V4
+                            } else {
+                                crate::network::IpFamily::V6
+                            },
+                            now_epoch_secs,
+                        ))
+            })
+            .count();
+        if current >= required {
+            return Err(ProxyError::Proxy(
+                "ME writer floor was restored before publication".into(),
+            ));
+        }
+        Ok(())
+    }
+
+    /// Commits writer visibility and lifecycle ownership after all cancellation points.
+    #[allow(clippy::too_many_arguments)]
+    pub(in crate::transport::middle_proxy) async fn publish_prepared_writer<F>(
+        self: &Arc<Self>,
+        writer: MeWriter,
+        tx: mpsc::Sender<WriterCommand>,
+        byte_budget: Arc<tokio::sync::Semaphore>,
+        task_registration: MeTaskRegistration<'_>,
+        writer_task: F,
+    ) where
+        F: Future<Output = ()> + Send + 'static,
+    {
+        // Writer publication follows the global writers -> registry binding lock order.
+        let mut writers = self.writers.write().await;
+        let mut registry_registration = self.registry.prepare_writer_registration().await;
+        registry_registration.install(writer.id, tx, byte_budget);
+        writers.push(writer);
+        self.conn_count.fetch_add(1, Ordering::Relaxed);
+        self.lifecycle
+            .spawn_registered_writer(task_registration, writer_task);
+        drop(writers);
+        drop(registry_registration);
+        self.notify_writer_epoch();
+    }
+}

src/transport/middle_proxy/pool_writer/replacement.rs

@@ -0,0 +1,376 @@
+use super::*;
+
+impl MePool {
+    /// Opens a replacement and atomically publishes it before retiring its reserved victim.
+    pub(in crate::transport::middle_proxy) async fn replace_writer_with_generation_contour_for_dc(
+        self: &Arc<Self>,
+        addr: SocketAddr,
+        rng: &SecureRandom,
+        generation: u64,
+        contour: WriterContour,
+        writer_dc: i32,
+        expected_victim_role: WriterRole,
+        purpose: WriterReplacementPurpose,
+        reservation: &mut WriterReplacementReservation<'_>,
+    ) -> Result<()> {
+        let prepared = self
+            .prepare_writer_with_intent(
+                addr,
+                rng,
+                generation,
+                contour,
+                writer_dc,
+                WriterOpenIntent::Replacement,
+            )
+            .await?;
+        self.publish_prepared_replacement_writer(
+            prepared,
+            expected_victim_role,
+            purpose,
+            reservation,
+        )
+        .await
+    }
+
+    async fn publish_prepared_replacement_writer(
+        self: &Arc<Self>,
+        prepared: PreparedWriter<'_>,
+        expected_victim_role: WriterRole,
+        purpose: WriterReplacementPurpose,
+        reservation: &mut WriterReplacementReservation<'_>,
+    ) -> Result<()> {
+        let PreparedWriter {
+            writer,
+            tx,
+            byte_budget,
+            task_registration,
+            writer_task,
+            intent,
+            _open_reservation,
+        } = prepared;
+        let replacement_writer_id = writer.id;
+        let victim_writer_id = reservation.writer_id();
+
+        // Lock order is writers -> registry binding. No cancellation point follows acquisition of
+        // the registry guard, so publication and victim retirement commit as one state change.
+        let mut writers = self.writers.write().await;
+        let mut registry_registration = self.registry.prepare_writer_registration().await;
+        let coordinator = self.reinit.coordinator.lock();
+        let contour = self.authorize_writer_publication(&writer, &coordinator)?;
+        self.authorize_writer_publication_capacity(&writer, contour, intent, writers.as_slice())?;
+        writer.contour.store(contour.as_u8(), Ordering::Release);
+        let Some(victim_pos) = writers
+            .iter()
+            .position(|candidate| candidate.id == victim_writer_id)
+        else {
+            return Err(ProxyError::Proxy(
+                "ME replacement victim disappeared before commit".into(),
+            ));
+        };
+        let victim = &writers[victim_pos];
+        if victim.draining.load(Ordering::Acquire) || !expected_victim_role.matches(victim) {
+            return Err(ProxyError::Proxy(
+                "ME replacement victim changed role before commit".into(),
+            ));
+        }
+        if let WriterReplacementPurpose::FloorRebalance {
+            donor_floor,
+            receiver_floor,
+        } = purpose
+        {
+            if expected_victim_role.generation != coordinator.active_generation
+                || writer.generation != coordinator.active_generation
+            {
+                return Err(ProxyError::Proxy(
+                    "ME floor rebalance lost active-generation authority".into(),
+                ));
+            }
+            let preferred = self.preferred_endpoints_by_dc.load();
+            let donor_count = writers
+                .iter()
+                .filter(|candidate| {
+                    !candidate.draining.load(Ordering::Acquire)
+                        && candidate.writer_dc == expected_victim_role.dc
+                        && candidate.generation == expected_victim_role.generation
+                        && WriterContour::from_u8(candidate.contour.load(Ordering::Acquire))
+                            == WriterContour::Active
+                        && preferred
+                            .get(&candidate.writer_dc)
+                            .is_some_and(|endpoints| endpoints.contains(&candidate.addr))
+                        && (candidate.addr.is_ipv4()
+                            == matches!(expected_victim_role.family, crate::network::IpFamily::V4))
+                })
+                .count();
+            let receiver_family = if writer.addr.is_ipv4() {
+                crate::network::IpFamily::V4
+            } else {
+                crate::network::IpFamily::V6
+            };
+            let receiver_count = writers
+                .iter()
+                .filter(|candidate| {
+                    !candidate.draining.load(Ordering::Acquire)
+                        && candidate.writer_dc == writer.writer_dc
+                        && candidate.generation == writer.generation
+                        && WriterContour::from_u8(candidate.contour.load(Ordering::Acquire))
+                            == WriterContour::Active
+                        && preferred
+                            .get(&candidate.writer_dc)
+                            .is_some_and(|endpoints| endpoints.contains(&candidate.addr))
+                        && (if candidate.addr.is_ipv4() {
+                            crate::network::IpFamily::V4
+                        } else {
+                            crate::network::IpFamily::V6
+                        }) == receiver_family
+                })
+                .count();
+            if donor_count <= donor_floor || receiver_count >= receiver_floor {
+                return Err(ProxyError::Proxy(
+                    "ME floor rebalance became unnecessary before commit".into(),
+                ));
+            }
+        }
+        if !registry_registration.prepare_replacement_commit(reservation) {
+            return Err(ProxyError::Proxy(
+                "ME replacement victim became active before commit".into(),
+            ));
+        }
+
+        registry_registration.install(replacement_writer_id, tx, byte_budget);
+        writers.push(writer);
+        self.conn_count.fetch_add(1, Ordering::Relaxed);
+        writers.publish_current();
+        self.apply_writer_draining_state(
+            &writers[victim_pos],
+            self.force_close_timeout(),
+            false,
+        );
+        self.lifecycle
+            .spawn_registered_writer(task_registration, writer_task);
+        reservation.mark_committed();
+        drop(coordinator);
+        drop(registry_registration);
+        drop(writers);
+        self.notify_writer_epoch();
+        info!(
+            victim_writer_id,
+            replacement_writer_id,
+            purpose = purpose.as_str(),
+            "ME writer replacement committed"
+        );
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::collections::HashMap;
+    use std::net::{IpAddr, Ipv4Addr};
+
+    use crate::transport::middle_proxy::pool_writer_security_tests::make_pool;
+
+    use super::*;
+
+    fn endpoint(octet: u8) -> SocketAddr {
+        SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, octet)), 443)
+    }
+
+    async fn install_writer(
+        pool: &Arc<MePool>,
+        writer_id: u64,
+        writer_dc: i32,
+        addr: SocketAddr,
+    ) -> MeWriter {
+        let (tx, _rx) = mpsc::channel::<WriterCommand>(8);
+        let byte_budget = pool.new_writer_byte_budget();
+        let writer = MeWriter {
+            id: writer_id,
+            addr,
+            source_ip: addr.ip(),
+            writer_dc,
+            generation: pool.current_generation(),
+            contour: Arc::new(AtomicU8::new(WriterContour::Active.as_u8())),
+            created_at: Instant::now(),
+            tx: tx.clone(),
+            byte_budget: byte_budget.clone(),
+            cancel: CancellationToken::new(),
+            degraded: Arc::new(AtomicBool::new(false)),
+            rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)),
+            draining: Arc::new(AtomicBool::new(false)),
+            draining_started_at_epoch_secs: Arc::new(AtomicU64::new(0)),
+            drain_deadline_epoch_secs: Arc::new(AtomicU64::new(0)),
+            allow_drain_fallback: Arc::new(AtomicBool::new(false)),
+        };
+        let mut writers = pool.writers.write().await;
+        let mut registry_registration = pool.registry.prepare_writer_registration().await;
+        registry_registration.install(writer_id, tx, byte_budget);
+        writers.push(writer.clone());
+        pool.conn_count.fetch_add(1, Ordering::Relaxed);
+        drop(registry_registration);
+        drop(writers);
+        writer
+    }
+
+    async fn prepared_writer<'a>(
+        pool: &'a Arc<MePool>,
+        writer_id: u64,
+        writer_dc: i32,
+        addr: SocketAddr,
+    ) -> PreparedWriter<'a> {
+        let open_reservation = pool
+            .reserve_writer_open(
+                WriterContour::Active,
+                WriterOpenIntent::Replacement,
+                writer_dc,
+            )
+            .await
+            .expect("replacement open must be admitted");
+        let task_registration = pool
+            .lifecycle
+            .try_register()
+            .expect("test pool lifecycle must be open");
+        let (tx, rx) = mpsc::channel::<WriterCommand>(8);
+        let byte_budget = pool.new_writer_byte_budget();
+        PreparedWriter {
+            writer: MeWriter {
+                id: writer_id,
+                addr,
+                source_ip: addr.ip(),
+                writer_dc,
+                generation: pool.current_generation(),
+                contour: Arc::new(AtomicU8::new(WriterContour::Active.as_u8())),
+                created_at: Instant::now(),
+                tx: tx.clone(),
+                byte_budget: byte_budget.clone(),
+                cancel: CancellationToken::new(),
+                degraded: Arc::new(AtomicBool::new(false)),
+                rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)),
+                draining: Arc::new(AtomicBool::new(false)),
+                draining_started_at_epoch_secs: Arc::new(AtomicU64::new(0)),
+                drain_deadline_epoch_secs: Arc::new(AtomicU64::new(0)),
+                allow_drain_fallback: Arc::new(AtomicBool::new(false)),
+            },
+            tx,
+            byte_budget,
+            task_registration,
+            writer_task: Box::pin(async move {
+                drop(rx);
+            }),
+            intent: WriterOpenIntent::Replacement,
+            _open_reservation: open_reservation,
+        }
+    }
+
+    #[tokio::test]
+    async fn replacement_commit_publishes_successor_before_draining_victim() {
+        let pool = make_pool().await;
+        let addr = endpoint(1);
+        pool.preferred_endpoints_by_dc
+            .store(Arc::new(HashMap::from([(2, vec![addr])])));
+        let victim = install_writer(&pool, 1001, 2, addr).await;
+        let expected_role = WriterRole::from_writer(&victim);
+        let mut reservation = pool
+            .registry
+            .try_reserve_writer_replacement(victim.id)
+            .await
+            .expect("idle victim must be reservable");
+        let prepared = prepared_writer(&pool, 1002, 2, addr).await;
+
+        pool.publish_prepared_replacement_writer(
+            prepared,
+            expected_role,
+            WriterReplacementPurpose::IdleRefresh,
+            &mut reservation,
+        )
+        .await
+        .expect("replacement commit must succeed");
+
+        let writers = pool.writers.read().await;
+        assert_eq!(writers.len(), 2);
+        assert_eq!(
+            writers
+                .iter()
+                .filter(|writer| !writer.draining.load(Ordering::Acquire))
+                .count(),
+            1
+        );
+        assert!(victim.draining.load(Ordering::Acquire));
+        assert!(writers.iter().any(|writer| writer.id == 1002));
+        drop(writers);
+        assert_eq!(pool.conn_count.load(Ordering::Acquire), 2);
+        assert_eq!(pool.registry.writer_replacement_counts(), (0, 1));
+    }
+
+    #[tokio::test]
+    async fn cancelled_replacement_waiting_for_publication_restores_all_reservations() {
+        let pool = make_pool().await;
+        let addr = endpoint(2);
+        pool.preferred_endpoints_by_dc
+            .store(Arc::new(HashMap::from([(2, vec![addr])])));
+        let victim = install_writer(&pool, 2001, 2, addr).await;
+        let expected_role = WriterRole::from_writer(&victim);
+        let mut reservation = pool
+            .registry
+            .try_reserve_writer_replacement(victim.id)
+            .await
+            .expect("idle victim must be reservable");
+        let prepared = prepared_writer(&pool, 2002, 2, addr).await;
+        let writers_guard = pool.writers.write().await;
+
+        let result = tokio::time::timeout(
+            Duration::from_millis(10),
+            pool.publish_prepared_replacement_writer(
+                prepared,
+                expected_role,
+                WriterReplacementPurpose::IdleRefresh,
+                &mut reservation,
+            ),
+        )
+        .await;
+
+        assert!(result.is_err());
+        drop(writers_guard);
+        drop(reservation);
+        assert_eq!(pool.writer_replacement_open_reserved.load(Ordering::Acquire), 0);
+        assert_eq!(pool.registry.writer_replacement_counts(), (0, 0));
+        assert!(!victim.draining.load(Ordering::Acquire));
+        assert!(!pool.writers.read().await.iter().any(|writer| writer.id == 2002));
+    }
+
+    #[tokio::test]
+    async fn floor_rebalance_commit_rejects_a_donor_without_surplus() {
+        let pool = make_pool().await;
+        let donor_addr = endpoint(3);
+        let receiver_addr = endpoint(4);
+        pool.preferred_endpoints_by_dc.store(Arc::new(HashMap::from([
+            (1, vec![donor_addr]),
+            (2, vec![receiver_addr]),
+        ])));
+        let victim = install_writer(&pool, 3001, 1, donor_addr).await;
+        let expected_role = WriterRole::from_writer(&victim);
+        let mut reservation = pool
+            .registry
+            .try_reserve_writer_replacement(victim.id)
+            .await
+            .expect("idle victim must be reservable");
+        let prepared = prepared_writer(&pool, 3002, 2, receiver_addr).await;
+
+        let result = pool
+            .publish_prepared_replacement_writer(
+                prepared,
+                expected_role,
+                WriterReplacementPurpose::FloorRebalance {
+                    donor_floor: 1,
+                    receiver_floor: 1,
+                },
+                &mut reservation,
+            )
+            .await;
+
+        assert!(result.is_err());
+        drop(reservation);
+        assert!(!victim.draining.load(Ordering::Acquire));
+        assert!(!pool.writers.read().await.iter().any(|writer| writer.id == 3002));
+        assert_eq!(pool.registry.writer_replacement_counts(), (0, 0));
+    }
+}

src/transport/middle_proxy/pool_writer/runtime.rs

@@ -54,23 +54,45 @@ impl MePool {
         contour: WriterContour,
         writer_dc: i32,
     ) -> Result<()> {
-        self.connect_one_with_generation_contour_for_dc_with_cap_policy(
-            addr, rng, generation, contour, writer_dc, false,
+        self.connect_one_with_generation_contour_for_dc_with_intent(
+            addr,
+            rng,
+            generation,
+            contour,
+            writer_dc,
+            WriterOpenIntent::Normal,
         )
         .await
     }
 
-    pub(in crate::transport::middle_proxy) async fn connect_one_with_generation_contour_for_dc_with_cap_policy(
+    /// Connects and publishes one writer under an explicit capacity intent.
+    pub(in crate::transport::middle_proxy) async fn connect_one_with_generation_contour_for_dc_with_intent(
         self: &Arc<Self>,
         addr: SocketAddr,
         rng: &SecureRandom,
         generation: u64,
         contour: WriterContour,
         writer_dc: i32,
-        allow_coverage_override: bool,
+        intent: WriterOpenIntent,
     ) -> Result<()> {
-        let Some(_writer_open_reservation) = self
-            .reserve_writer_open(contour, allow_coverage_override, writer_dc)
+        let prepared = self
+            .prepare_writer_with_intent(addr, rng, generation, contour, writer_dc, intent)
+            .await?;
+        self.publish_connected_writer(prepared).await
+    }
+
+    /// Completes all cancellable connection work without publishing pool visibility.
+    pub(super) async fn prepare_writer_with_intent<'a>(
+        self: &'a Arc<Self>,
+        addr: SocketAddr,
+        rng: &SecureRandom,
+        generation: u64,
+        contour: WriterContour,
+        writer_dc: i32,
+        intent: WriterOpenIntent,
+    ) -> Result<PreparedWriter<'a>> {
+        let Some(writer_open_reservation) = self
+            .reserve_writer_open(contour, intent, writer_dc)
             .await
         else {
             return Err(ProxyError::Proxy(format!(
@@ -264,36 +286,15 @@ impl MePool {
             })
         };
 
-        self.publish_prepared_writer(writer, tx, byte_budget, task_registration, writer_task)
-            .await;
-
-        Ok(())
-    }
-
-    /// Commits writer visibility and lifecycle ownership after all cancellation points.
-    #[allow(clippy::too_many_arguments)]
-    pub(in crate::transport::middle_proxy) async fn publish_prepared_writer<F>(
-        self: &Arc<Self>,
-        writer: MeWriter,
-        tx: mpsc::Sender<WriterCommand>,
-        byte_budget: Arc<tokio::sync::Semaphore>,
-        task_registration: MeTaskRegistration<'_>,
-        writer_task: F,
-    ) where
-        F: Future<Output = ()> + Send + 'static,
-    {
-        let (mut writers, mut registry_registration) = tokio::join!(
-            self.writers.write(),
-            self.registry.prepare_writer_registration()
-        );
-        registry_registration.install(writer.id, tx, byte_budget);
-        writers.push(writer);
-        self.conn_count.fetch_add(1, Ordering::Relaxed);
-        self.lifecycle
-            .spawn_registered_writer(task_registration, writer_task);
-        drop(writers);
-        drop(registry_registration);
-        self.notify_writer_epoch();
+        Ok(PreparedWriter {
+            writer,
+            tx,
+            byte_budget,
+            task_registration,
+            writer_task,
+            intent,
+            _open_reservation: writer_open_reservation,
+        })
     }
 
     pub(crate) async fn remove_writer_and_close_clients(self: &Arc<Self>, writer_id: u64) {
@@ -331,7 +332,7 @@ impl MePool {
     ) -> bool {
         let mut close_tx: Option<mpsc::Sender<WriterCommand>> = None;
         let mut removed_addr: Option<SocketAddr> = None;
-        let mut removed_dc: Option<i32> = None;
+        let mut removed_role: Option<WriterRole> = None;
         let mut removed_uptime: Option<Duration> = None;
         let mut trigger_refill = false;
         let mut removed = false;
@@ -352,7 +353,7 @@ impl MePool {
                 self.stats.increment_me_writer_removed_total();
                 w.cancel.cancel();
                 removed_addr = Some(w.addr);
-                removed_dc = Some(w.writer_dc);
+                removed_role = Some(WriterRole::from_writer(&w));
                 removed_uptime = Some(w.created_at.elapsed());
                 trigger_refill = !was_draining;
                 if trigger_refill {
@@ -397,8 +398,8 @@ impl MePool {
                     );
                 }
             }
-            if trigger_refill && let Some(writer_dc) = removed_dc {
-                self.trigger_immediate_refill_for_dc(addr, writer_dc);
+            if trigger_refill && let Some(role) = removed_role {
+                self.trigger_immediate_refill_for_role(addr, role);
             }
         }
         if removed {
@@ -415,26 +416,13 @@ impl MePool {
     ) {
         let timeout = timeout.filter(|d| !d.is_zero());
         let found = {
-            let mut ws = self.writers.write().await;
-            if let Some(w) = ws.iter_mut().find(|w| w.id == writer_id) {
-                let already_draining = w.draining.swap(true, Ordering::Relaxed);
-                w.allow_drain_fallback
-                    .store(allow_drain_fallback, Ordering::Relaxed);
-                let now_epoch_secs = Self::now_epoch_secs();
-                w.draining_started_at_epoch_secs
-                    .store(now_epoch_secs, Ordering::Relaxed);
-                let drain_deadline_epoch_secs = timeout
-                    .map(|duration| now_epoch_secs.saturating_add(duration.as_secs()))
-                    .unwrap_or(0);
-                w.drain_deadline_epoch_secs
-                    .store(drain_deadline_epoch_secs, Ordering::Relaxed);
-                if !already_draining {
-                    self.stats.increment_pool_drain_active();
-                    self.increment_draining_active_runtime();
+            let ws = self.writers.write().await;
+            let mut registry_registration = self.registry.prepare_writer_registration().await;
+            if let Some(w) = ws.iter().find(|w| w.id == writer_id) {
+                if !allow_drain_fallback {
+                    registry_registration.retire(writer_id);
                 }
-                w.contour
-                    .store(WriterContour::Draining.as_u8(), Ordering::Relaxed);
-                w.draining.store(true, Ordering::Relaxed);
+                self.apply_writer_draining_state(w, timeout, allow_drain_fallback);
                 true
             } else {
                 false
@@ -452,6 +440,39 @@ impl MePool {
         );
     }
 
+    /// Publishes drain metadata before the release-store that makes draining observable.
+    pub(in crate::transport::middle_proxy) fn apply_writer_draining_state(
+        &self,
+        writer: &MeWriter,
+        timeout: Option<Duration>,
+        allow_drain_fallback: bool,
+    ) -> bool {
+        let timeout = timeout.filter(|duration| !duration.is_zero());
+        let already_draining = writer.draining.load(Ordering::Acquire);
+        writer
+            .allow_drain_fallback
+            .store(allow_drain_fallback, Ordering::Release);
+        let now_epoch_secs = Self::now_epoch_secs();
+        writer
+            .draining_started_at_epoch_secs
+            .store(now_epoch_secs, Ordering::Release);
+        let drain_deadline_epoch_secs = timeout
+            .map(|duration| now_epoch_secs.saturating_add(duration.as_secs()))
+            .unwrap_or(0);
+        writer
+            .drain_deadline_epoch_secs
+            .store(drain_deadline_epoch_secs, Ordering::Release);
+        writer
+            .contour
+            .store(WriterContour::Draining.as_u8(), Ordering::Release);
+        writer.draining.store(true, Ordering::Release);
+        if !already_draining {
+            self.stats.increment_pool_drain_active();
+            self.increment_draining_active_runtime();
+        }
+        !already_draining
+    }
+
     pub(crate) async fn mark_writer_draining(self: &Arc<Self>, writer_id: u64) {
         self.mark_writer_draining_with_timeout(writer_id, Some(Duration::from_secs(300)), false)
             .await;
@@ -461,7 +482,7 @@ impl MePool {
         &self,
         writer: &MeWriter,
     ) -> bool {
-        if !writer.draining.load(Ordering::Relaxed) {
+        if !writer.draining.load(Ordering::Acquire) {
             return true;
         }
         if !writer.allow_drain_fallback.load(Ordering::Relaxed) {