Races in admission + accounting + publication,+ PID fixed

src/transport/middle_proxy/mod.rs

@@ -33,6 +33,9 @@ mod pool_runtime_api;
 mod pool_status;
 mod pool_writer;
 #[cfg(test)]
+#[path = "tests/pool_writer_publication_tests.rs"]
+mod pool_writer_publication_tests;
+#[cfg(test)]
 #[path = "tests/pool_writer_security_tests.rs"]
 mod pool_writer_security_tests;
 mod reader;

src/transport/middle_proxy/pool.rs

@@ -88,14 +88,6 @@ impl WritersState {
         }
     }
 
-    pub(super) async fn update<F, R>(&self, f: F) -> R
-    where
-        F: FnOnce(&mut Vec<MeWriter>) -> R,
-    {
-        let mut guard = self.write().await;
-        f(&mut guard)
-    }
-
     fn debug_assert_store_guarded(&self) {
         debug_assert!(
             self.writers_write_guard.try_lock().is_err(),

src/transport/middle_proxy/pool_writer.rs

@@ -1,4 +1,5 @@
 use std::collections::HashMap;
+use std::future::Future;
 use std::io::ErrorKind;
 use std::net::SocketAddr;
 use std::sync::Arc;
@@ -19,6 +20,7 @@ use crate::protocol::constants::{RPC_CLOSE_EXT_U32, RPC_PING_U32};
 
 use super::codec::{RpcWriter, WriterCommand, build_control_payload};
 use super::pool::{MePool, MeWriter, WriterContour};
+use super::pool_lifecycle::MeTaskRegistration;
 use super::reader::reader_loop;
 use super::wire::build_proxy_req_payload;
 

src/transport/middle_proxy/pool_writer/runtime.rs

@@ -132,16 +132,6 @@ impl MePool {
             drain_deadline_epoch_secs: drain_deadline_epoch_secs.clone(),
             allow_drain_fallback: allow_drain_fallback.clone(),
         };
-        self.writers
-            .update(|writers| writers.push(writer.clone()))
-            .await;
-        self.registry
-            .register_writer(writer_id, tx.clone(), byte_budget)
-            .await;
-        self.registry.mark_writer_idle(writer_id).await;
-        self.conn_count.fetch_add(1, Ordering::Relaxed);
-        self.notify_writer_epoch();
-
         let reg = self.registry.clone();
         let writers_arc = self.writers_arc();
         let ping_tracker = Arc::new(tokio::sync::Mutex::new(HashMap::<i64, Instant>::new()));
@@ -177,8 +167,10 @@ impl MePool {
         let route_fairshare_enabled = self.transport_policy.me_route_fairshare_enabled.clone();
         let reader_route_data_wait_ms = self.transport_policy.me_reader_route_data_wait_ms.clone();
 
-        self.lifecycle
-            .spawn_registered_writer(task_registration, async move {
+        let writer_task = {
+            // Keep transport ownership behind a stable-size pointer while publication waits for
+            // locks.
+            Box::pin(async move {
                 // Reader MUST be the first branch in biased select! to avoid read starvation.
                 let exit = tokio::select! {
                     biased;
@@ -269,11 +261,41 @@ impl MePool {
 
                 let remaining = writers_arc.read().await.len();
                 debug!(writer_id, remaining, "ME writer lifecycle task finished");
-            });
+            })
+        };
+
+        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();
+    }
+
     pub(crate) async fn remove_writer_and_close_clients(self: &Arc<Self>, writer_id: u64) {
         // Full client cleanup now happens inside `registry.writer_lost` to keep
         // writer reap/remove paths strictly non-blocking per connection.
@@ -337,7 +359,12 @@ impl MePool {
                     self.stats.increment_me_writer_removed_unexpected_total();
                 }
                 close_tx = Some(w.tx.clone());
-                self.conn_count.fetch_sub(1, Ordering::Relaxed);
+                // Teardown remains idempotent if the advisory count was already reconciled.
+                let _ =
+                    self.conn_count
+                        .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |count| {
+                            count.checked_sub(1)
+                        });
                 removed = true;
             }
         }

src/transport/middle_proxy/registry.rs

@@ -18,6 +18,8 @@ const ROUTE_QUEUED_BYTE_PERMIT_UNIT: usize = 16 * 1024;
 const ROUTE_QUEUED_PERMITS_PER_SLOT: usize = 4;
 const ROUTE_QUEUED_MAX_FRAME_PERMITS: usize = 1024;
 
+// Transactional writer registry publication.
+mod publication;
 mod writer;
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]

src/transport/middle_proxy/registry/publication.rs

@@ -0,0 +1,50 @@
+use std::sync::Arc;
+
+use tokio::sync::{MutexGuard, Semaphore, mpsc};
+
+use super::super::codec::WriterCommand;
+use super::{BindingInner, ConnRegistry, WriterRoute};
+
+/// Holds registry binding ownership until pool writer visibility is published.
+pub(in crate::transport::middle_proxy) struct WriterRegistrationGuard<'a> {
+    registry: &'a ConnRegistry,
+    binding: MutexGuard<'a, BindingInner>,
+}
+
+impl ConnRegistry {
+    /// Acquires the only cancellation point required for writer registration.
+    pub(in crate::transport::middle_proxy) async fn prepare_writer_registration(
+        &self,
+    ) -> WriterRegistrationGuard<'_> {
+        WriterRegistrationGuard {
+            registry: self,
+            binding: self.binding.inner.lock().await,
+        }
+    }
+}
+
+impl WriterRegistrationGuard<'_> {
+    /// Installs registry state while retaining binding ownership for pool publication.
+    pub(in crate::transport::middle_proxy) fn install(
+        &mut self,
+        writer_id: u64,
+        tx: mpsc::Sender<WriterCommand>,
+        byte_budget: Arc<Semaphore>,
+    ) {
+        self.binding.conns_for_writer.entry(writer_id).or_default();
+        self.registry
+            .binding
+            .bound_clients_by_writer
+            .entry(writer_id)
+            .or_insert(0);
+        self.registry
+            .binding
+            .writer_idle_since_epoch_secs
+            .entry(writer_id)
+            .or_insert_with(ConnRegistry::now_epoch_secs);
+        self.registry
+            .writers
+            .map
+            .insert(writer_id, WriterRoute { tx, byte_budget });
+    }
+}

src/transport/middle_proxy/registry/writer.rs

@@ -58,28 +58,16 @@ impl ConnRegistry {
     }
 
     /// Registers one writer command route and its matching memory budget atomically.
+    #[allow(dead_code)]
     pub async fn register_writer(
         &self,
         writer_id: u64,
         tx: mpsc::Sender<WriterCommand>,
         byte_budget: Arc<tokio::sync::Semaphore>,
     ) {
-        let mut binding = self.binding.inner.lock().await;
-        binding
-            .conns_for_writer
-            .entry(writer_id)
-            .or_insert_with(HashSet::new);
-        self.binding
-            .bound_clients_by_writer
-            .entry(writer_id)
-            .or_insert(0);
-        self.binding
-            .writer_idle_since_epoch_secs
-            .entry(writer_id)
-            .or_insert_with(Self::now_epoch_secs);
-        self.writers
-            .map
-            .insert(writer_id, super::WriterRoute { tx, byte_budget });
+        self.prepare_writer_registration()
+            .await
+            .install(writer_id, tx, byte_budget);
     }
 
     /// Unregister connection, returning associated writer_id if any.
@@ -346,20 +334,6 @@ impl ConnRegistry {
         true
     }
 
-    pub async fn mark_writer_idle(&self, writer_id: u64) {
-        let mut binding = self.binding.inner.lock().await;
-        binding
-            .conns_for_writer
-            .entry(writer_id)
-            .or_insert_with(HashSet::new);
-        let count = binding
-            .conns_for_writer
-            .get(&writer_id)
-            .map(|set| set.len())
-            .unwrap_or(0);
-        self.set_writer_bound_count(writer_id, count);
-    }
-
     pub async fn get_last_writer_meta(&self, writer_id: u64) -> Option<ConnMeta> {
         self.binding
             .last_meta_for_writer

src/transport/middle_proxy/tests/pool_writer_publication_tests.rs

@@ -0,0 +1,147 @@
+use std::net::{IpAddr, Ipv4Addr, SocketAddr};
+use std::sync::Arc;
+use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU32, AtomicU64, Ordering};
+use std::time::{Duration, Instant};
+
+use tokio::sync::mpsc;
+use tokio_util::sync::CancellationToken;
+
+use super::codec::WriterCommand;
+use super::pool::{MeWriter, WriterContour};
+use super::pool_writer_security_tests::make_pool;
+use super::registry::ConnMeta;
+
+#[tokio::test]
+async fn successful_writer_publication_is_fully_visible_and_removable() {
+    let pool = make_pool().await;
+    let writer_id = 76_002;
+    let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443);
+    let (tx, _rx) = mpsc::channel::<WriterCommand>(8);
+    let byte_budget = pool.new_writer_byte_budget();
+    let cancel = CancellationToken::new();
+    let writer = MeWriter {
+        id: writer_id,
+        addr,
+        source_ip: addr.ip(),
+        writer_dc: 2,
+        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: cancel.clone(),
+        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 task_started = Arc::new(AtomicBool::new(false));
+    let task_started_writer = Arc::clone(&task_started);
+    let writer_task = async move {
+        task_started_writer.store(true, Ordering::Release);
+        cancel.cancelled().await;
+    };
+    let task_registration = pool.lifecycle.try_register().unwrap();
+
+    pool.publish_prepared_writer(writer, tx, byte_budget, task_registration, writer_task)
+        .await;
+
+    assert_eq!(pool.writers.read().await.len(), 1);
+    assert_eq!(pool.conn_count.load(Ordering::Relaxed), 1);
+    tokio::time::timeout(Duration::from_secs(1), async {
+        while !task_started.load(Ordering::Acquire) {
+            tokio::task::yield_now().await;
+        }
+    })
+    .await
+    .expect("published writer task must start");
+
+    let (conn_id, _response_rx) = pool.registry.register().await;
+    assert!(
+        pool.registry
+            .bind_writer(
+                conn_id,
+                writer_id,
+                ConnMeta {
+                    target_dc: 2,
+                    client_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 7300),
+                    our_addr: addr,
+                    proto_flags: 0,
+                },
+            )
+            .await
+    );
+    assert_eq!(
+        pool.registry.get_writer(conn_id).await.unwrap().writer_id,
+        writer_id
+    );
+
+    pool.remove_writer_and_close_clients(writer_id).await;
+
+    assert!(pool.writers.read().await.is_empty());
+    assert_eq!(pool.conn_count.load(Ordering::Relaxed), 0);
+    assert!(pool.registry.get_writer(conn_id).await.is_none());
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+async fn published_writer_removal_race_preserves_count() {
+    let pool = make_pool().await;
+
+    for writer_id in 80_000..90_000 {
+        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443);
+        let (tx, _rx) = mpsc::channel::<WriterCommand>(1);
+        let byte_budget = pool.new_writer_byte_budget();
+        let cancel = CancellationToken::new();
+        let writer = MeWriter {
+            id: writer_id,
+            addr,
+            source_ip: addr.ip(),
+            writer_dc: 2,
+            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: cancel.clone(),
+            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 writer_task = async move {
+            cancel.cancelled().await;
+        };
+        let task_registration = pool.lifecycle.try_register().unwrap();
+        let remover_pool = Arc::clone(&pool);
+        let remover = tokio::spawn(async move {
+            loop {
+                if remover_pool
+                    .writers
+                    .snapshot()
+                    .iter()
+                    .any(|writer| writer.id == writer_id)
+                {
+                    remover_pool
+                        .remove_writer_and_close_clients(writer_id)
+                        .await;
+                    return;
+                }
+                tokio::task::yield_now().await;
+            }
+        });
+
+        pool.publish_prepared_writer(writer, tx, byte_budget, task_registration, writer_task)
+            .await;
+        tokio::time::timeout(Duration::from_secs(1), remover)
+            .await
+            .expect("published writer must become removable")
+            .expect("writer remover task must not panic");
+
+        assert!(pool.writers.read().await.is_empty());
+        assert_eq!(pool.conn_count.load(Ordering::Relaxed), 0);
+    }
+}

src/transport/middle_proxy/tests/pool_writer_security_tests.rs

@@ -15,7 +15,8 @@ use crate::crypto::SecureRandom;
 use crate::network::probe::NetworkDecision;
 use crate::stats::Stats;
 
-async fn make_pool() -> Arc<MePool> {
+/// Builds an isolated ME pool for writer-state tests.
+pub(super) async fn make_pool() -> Arc<MePool> {
     let general = GeneralConfig::default();
 
     MePool::new(
@@ -154,6 +155,70 @@ async fn insert_writer(
     pool.conn_count.fetch_add(1, Ordering::Relaxed);
 }
 
+#[tokio::test]
+async fn cancelled_writer_publication_leaves_no_ghost_state() {
+    let pool = make_pool().await;
+    let writer_id = 76_001;
+    let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443);
+    let (tx, mut 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: 2,
+        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 task_started = Arc::new(AtomicBool::new(false));
+    let task_started_writer = Arc::clone(&task_started);
+    let writer_task = async move {
+        task_started_writer.store(true, Ordering::Release);
+        let _ = rx.recv().await;
+    };
+    let task_registration = pool.lifecycle.try_register().unwrap();
+    let held_registration = pool.registry.prepare_writer_registration().await;
+
+    let publication =
+        pool.publish_prepared_writer(writer, tx, byte_budget, task_registration, writer_task);
+    assert!(
+        tokio::time::timeout(Duration::from_millis(10), publication)
+            .await
+            .is_err()
+    );
+    drop(held_registration);
+
+    assert!(pool.writers.read().await.is_empty());
+    assert_eq!(pool.conn_count.load(Ordering::Relaxed), 0);
+    assert!(!task_started.load(Ordering::Acquire));
+    let (conn_id, _response_rx) = pool.registry.register().await;
+    assert!(
+        !pool
+            .registry
+            .bind_writer(
+                conn_id,
+                writer_id,
+                ConnMeta {
+                    target_dc: 2,
+                    client_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 7300),
+                    our_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443),
+                    proto_flags: 0,
+                },
+            )
+            .await
+    );
+}
+
 async fn current_writer_ids(pool: &Arc<MePool>) -> HashSet<u64> {
     pool.writers
         .read()