Adaptive Buffers + Session Eviction Method

src/transport/middle_proxy/registry.rs

@@ -394,6 +394,56 @@ impl ConnRegistry {
         inner.writer_for_conn.keys().copied().collect()
     }
 
+    pub(super) async fn bound_conn_ids_for_writer_limited(
+        &self,
+        writer_id: u64,
+        limit: usize,
+    ) -> Vec<u64> {
+        if limit == 0 {
+            return Vec::new();
+        }
+        let inner = self.inner.read().await;
+        let Some(conn_ids) = inner.conns_for_writer.get(&writer_id) else {
+            return Vec::new();
+        };
+        let mut out = conn_ids.iter().copied().collect::<Vec<_>>();
+        out.sort_unstable();
+        out.truncate(limit);
+        out
+    }
+
+    pub(super) async fn evict_bound_conn_if_writer(&self, conn_id: u64, writer_id: u64) -> bool {
+        let maybe_client_tx = {
+            let mut inner = self.inner.write().await;
+            if inner.writer_for_conn.get(&conn_id).copied() != Some(writer_id) {
+                return false;
+            }
+
+            let client_tx = inner.map.get(&conn_id).cloned();
+            inner.map.remove(&conn_id);
+            inner.meta.remove(&conn_id);
+            inner.writer_for_conn.remove(&conn_id);
+
+            let became_empty = if let Some(set) = inner.conns_for_writer.get_mut(&writer_id) {
+                set.remove(&conn_id);
+                set.is_empty()
+            } else {
+                false
+            };
+            if became_empty {
+                inner
+                    .writer_idle_since_epoch_secs
+                    .insert(writer_id, Self::now_epoch_secs());
+            }
+            client_tx
+        };
+
+        if let Some(client_tx) = maybe_client_tx {
+            let _ = client_tx.try_send(MeResponse::Close);
+        }
+        true
+    }
+
     pub async fn writer_lost(&self, writer_id: u64) -> Vec<BoundConn> {
         let mut inner = self.inner.write().await;
         inner.writers.remove(&writer_id);
@@ -444,6 +494,7 @@ mod tests {
 
     use super::ConnMeta;
     use super::ConnRegistry;
+    use super::MeResponse;
 
     #[tokio::test]
     async fn writer_activity_snapshot_tracks_writer_and_dc_load() {
@@ -634,4 +685,86 @@ mod tests {
         );
         assert!(registry.get_writer(conn_id).await.is_none());
     }
+
+    #[tokio::test]
+    async fn bound_conn_ids_for_writer_limited_is_sorted_and_bounded() {
+        let registry = ConnRegistry::new();
+        let (writer_tx, _writer_rx) = tokio::sync::mpsc::channel(8);
+        registry.register_writer(10, writer_tx).await;
+        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443);
+        let mut conn_ids = Vec::new();
+        for _ in 0..5 {
+            let (conn_id, _rx) = registry.register().await;
+            assert!(
+                registry
+                    .bind_writer(
+                        conn_id,
+                        10,
+                        ConnMeta {
+                            target_dc: 2,
+                            client_addr: addr,
+                            our_addr: addr,
+                            proto_flags: 0,
+                        },
+                    )
+                    .await
+            );
+            conn_ids.push(conn_id);
+        }
+        conn_ids.sort_unstable();
+
+        let limited = registry.bound_conn_ids_for_writer_limited(10, 3).await;
+        assert_eq!(limited.len(), 3);
+        assert_eq!(limited, conn_ids.into_iter().take(3).collect::<Vec<_>>());
+    }
+
+    #[tokio::test]
+    async fn evict_bound_conn_if_writer_does_not_touch_rebound_conn() {
+        let registry = ConnRegistry::new();
+        let (conn_id, mut rx) = registry.register().await;
+        let (writer_tx_a, _writer_rx_a) = tokio::sync::mpsc::channel(8);
+        let (writer_tx_b, _writer_rx_b) = tokio::sync::mpsc::channel(8);
+        registry.register_writer(10, writer_tx_a).await;
+        registry.register_writer(20, writer_tx_b).await;
+        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443);
+
+        assert!(
+            registry
+                .bind_writer(
+                    conn_id,
+                    10,
+                    ConnMeta {
+                        target_dc: 2,
+                        client_addr: addr,
+                        our_addr: addr,
+                        proto_flags: 0,
+                    },
+                )
+                .await
+        );
+        assert!(
+            registry
+                .bind_writer(
+                    conn_id,
+                    20,
+                    ConnMeta {
+                        target_dc: 2,
+                        client_addr: addr,
+                        our_addr: addr,
+                        proto_flags: 1,
+                    },
+                )
+                .await
+        );
+
+        let evicted = registry.evict_bound_conn_if_writer(conn_id, 10).await;
+        assert!(!evicted);
+        assert_eq!(registry.get_writer(conn_id).await.expect("writer").writer_id, 20);
+        assert!(rx.try_recv().is_err());
+
+        let evicted = registry.evict_bound_conn_if_writer(conn_id, 20).await;
+        assert!(evicted);
+        assert!(registry.get_writer(conn_id).await.is_none());
+        assert!(matches!(rx.try_recv(), Ok(MeResponse::Close)));
+    }
 }