Atomic Lifecycle for Writer Refresh + Hardswap

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));
+    }
+}