Refactor health management: implement remove_writer_if_empty method for cleaner writer removal logic and update related functions to enhance efficiency in handling closed writers.

src/transport/middle_proxy/health.rs

@@ -239,7 +239,9 @@ pub(super) async fn reap_draining_writers(
         if !closed_writer_ids.insert(writer_id) {
             continue;
         }
-        pool.remove_writer_and_close_clients(writer_id).await;
+        if !pool.remove_writer_if_empty(writer_id).await {
+            continue;
+        }
         closed_total = closed_total.saturating_add(1);
     }
 

src/transport/middle_proxy/health_regression_tests.rs

@@ -592,3 +592,67 @@ async fn reap_draining_writers_mixed_backlog_converges_without_leaking_warn_stat
 fn general_config_default_drain_threshold_remains_enabled() {
     assert_eq!(GeneralConfig::default().me_pool_drain_threshold, 128);
 }
+
+#[tokio::test]
+async fn reap_draining_writers_does_not_close_writer_that_became_non_empty_after_snapshot() {
+    let pool = make_pool(128).await;
+    let now_epoch_secs = MePool::now_epoch_secs();
+
+    let empty_writer_id = 700u64;
+    insert_draining_writer(
+        &pool,
+        empty_writer_id,
+        now_epoch_secs.saturating_sub(60),
+        0,
+        0,
+    )
+    .await;
+
+    let stale_empty_snapshot = vec![empty_writer_id];
+    let (rebound_conn_id, _rx) = pool.registry.register().await;
+    assert!(
+        pool.registry
+            .bind_writer(
+                rebound_conn_id,
+                empty_writer_id,
+                ConnMeta {
+                    target_dc: 2,
+                    client_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 9050),
+                    our_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443),
+                    proto_flags: 0,
+                },
+            )
+            .await,
+        "writer should accept a new bind after stale empty snapshot"
+    );
+
+    for writer_id in stale_empty_snapshot {
+        assert!(
+            !pool.remove_writer_if_empty(writer_id).await,
+            "atomic empty cleanup must reject writers that gained bound clients"
+        );
+    }
+
+    assert!(
+        writer_exists(&pool, empty_writer_id).await,
+        "empty-path cleanup must not remove a writer that gained a bound client"
+    );
+    assert_eq!(
+        pool.registry.get_writer(rebound_conn_id).await.map(|w| w.writer_id),
+        Some(empty_writer_id)
+    );
+
+    let _ = pool.registry.unregister(rebound_conn_id).await;
+}
+
+#[tokio::test]
+async fn prune_closed_writers_closes_bound_clients_when_writer_is_non_empty() {
+    let pool = make_pool(128).await;
+    let now_epoch_secs = MePool::now_epoch_secs();
+    let conn_ids = insert_draining_writer(&pool, 910, now_epoch_secs.saturating_sub(60), 1, 0).await;
+
+    pool.prune_closed_writers().await;
+
+    assert!(!writer_exists(&pool, 910).await);
+    assert!(pool.registry.get_writer(conn_ids[0]).await.is_none());
+}

src/transport/middle_proxy/pool_writer.rs

@@ -42,11 +42,10 @@ impl MePool {
         }
 
         for writer_id in closed_writer_ids {
-            if self.registry.is_writer_empty(writer_id).await {
-                let _ = self.remove_writer_only(writer_id).await;
-            } else {
-                let _ = self.remove_writer_and_close_clients(writer_id).await;
+            if self.remove_writer_if_empty(writer_id).await {
+                continue;
             }
+            let _ = self.remove_writer_and_close_clients(writer_id).await;
         }
     }
 
@@ -501,6 +500,17 @@ impl MePool {
         }
     }
 
+    pub(crate) async fn remove_writer_if_empty(self: &Arc<Self>, writer_id: u64) -> bool {
+        if !self.registry.unregister_writer_if_empty(writer_id).await {
+            return false;
+        }
+
+        // The registry empty-check and unregister are atomic with respect to binds,
+        // so remove_writer_only cannot return active bound sessions here.
+        let _ = self.remove_writer_only(writer_id).await;
+        true
+    }
+
     async fn remove_writer_only(self: &Arc<Self>, writer_id: u64) -> Vec<BoundConn> {
         let mut close_tx: Option<mpsc::Sender<WriterCommand>> = None;
         let mut removed_addr: Option<SocketAddr> = None;

src/transport/middle_proxy/registry.rs

@@ -437,6 +437,23 @@ impl ConnRegistry {
             .unwrap_or(true)
     }
 
+    pub async fn unregister_writer_if_empty(&self, writer_id: u64) -> bool {
+        let mut inner = self.inner.write().await;
+        let Some(conn_ids) = inner.conns_for_writer.get(&writer_id) else {
+            // Writer is already absent from the registry.
+            return true;
+        };
+        if !conn_ids.is_empty() {
+            return false;
+        }
+
+        inner.writers.remove(&writer_id);
+        inner.last_meta_for_writer.remove(&writer_id);
+        inner.writer_idle_since_epoch_secs.remove(&writer_id);
+        inner.conns_for_writer.remove(&writer_id);
+        true
+    }
+
     pub(super) async fn non_empty_writer_ids(&self, writer_ids: &[u64]) -> HashSet<u64> {
         let inner = self.inner.read().await;
         let mut out = HashSet::<u64>::with_capacity(writer_ids.len());