cmux-tui: each terminal gets a fair share of the terminal connection lane (on top of #27719's 48-of-64 cap)

cmux-tui/crates/cmux-tui-core/src/diagnostics.rs

@@ -10,8 +10,8 @@
 
 use std::collections::HashMap;
 use std::panic::Location;
-use std::sync::Mutex;
 use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
+use std::sync::{Mutex, PoisonError};
 use std::time::{Duration, Instant};
 
 use serde::Serialize;
@@ -657,6 +657,10 @@ pub struct ConnectionStats {
     /// terminal connection gives this slot back with its main one
     /// (`connection_leases`).
     terminal_active: AtomicU64,
+    /// Of `terminal_active`, the connections each terminal holds (by its
+    /// process session), for the terminal lane's fair share
+    /// (`try_claim_terminal`). A terminal with none is not listed.
+    terminal_clients: Mutex<HashMap<u32, u64>>,
     /// The socket these connections come in on, once a server attached the
     /// session to it.
     socket: std::sync::OnceLock<std::path::PathBuf>,
@@ -669,15 +673,18 @@ impl ConnectionStats {
     /// the slot back.
     /// At `limit` it first takes back the slots of clients that hung up
     /// (`reclaim_hung_up`, pool bounded by `pool_limit`), as `try_claim` does.
-    pub fn try_claim_terminal(&self, limit: u64, pool_limit: u64) -> bool {
-        let claim = |count: u64| (count < limit).then_some(count + 1);
-        let mut claimed =
-            self.terminal_active.fetch_update(Ordering::AcqRel, Ordering::Acquire, claim).is_ok();
+    ///
+    /// `client`: the terminal the connection comes from (its process session).
+    /// A terminal's first connection gets in whenever the lane has room; a
+    /// further one only while that terminal then holds no more than the lane
+    /// has left free. One terminal (an agent's hooks in a loop, a script
+    /// running `cmux`) stops at half the lane, so a second terminal still finds
+    /// room; several terminals in a loop together can still fill the lane. A
+    /// connection whose terminal is unknown is bounded by the lane alone.
+    pub fn try_claim_terminal(&self, limit: u64, pool_limit: u64, client: Option<u32>) -> bool {
+        let mut claimed = self.claim_terminal_slot(limit, client);
         if !claimed && self.reclaim_hung_up(pool_limit) > 0 {
-            claimed = self
-                .terminal_active
-                .fetch_update(Ordering::AcqRel, Ordering::Acquire, claim)
-                .is_ok();
+            claimed = self.claim_terminal_slot(limit, client);
         }
         if !claimed {
             self.accepted.fetch_sub(1, Ordering::Relaxed);
@@ -686,7 +693,38 @@ impl ConnectionStats {
         claimed
     }
 
-    pub fn release_terminal(&self) {
+    /// One terminal slot under `limit`, within `client`'s fair share. The
+    /// share and the count change together under `terminal_clients`, so two
+    /// claims never both take the last place a share allows.
+    fn claim_terminal_slot(&self, limit: u64, client: Option<u32>) -> bool {
+        let mut clients = self.terminal_clients.lock().unwrap_or_else(PoisonError::into_inner);
+        let active = self.terminal_active.load(Ordering::Acquire);
+        if active >= limit {
+            return false;
+        }
+        if let Some(client) = client {
+            let held = clients.get(&client).copied().unwrap_or(0);
+            let free_after = limit.saturating_sub(active.saturating_add(1));
+            if held > 0 && held.saturating_add(1) > free_after {
+                return false;
+            }
+            let held = clients.entry(client).or_insert(0);
+            *held = held.saturating_add(1);
+        }
+        self.terminal_active.fetch_add(1, Ordering::AcqRel);
+        true
+    }
+
+    pub fn release_terminal(&self, client: Option<u32>) {
+        if let Some(client) = client {
+            let mut clients = self.terminal_clients.lock().unwrap_or_else(PoisonError::into_inner);
+            if let Some(held) = clients.get_mut(&client) {
+                *held = held.saturating_sub(1);
+                if *held == 0 {
+                    clients.remove(&client);
+                }
+            }
+        }
         self.terminal_active.fetch_sub(1, Ordering::AcqRel);
     }
 

cmux-tui/crates/cmux-tui-core/src/diagnostics/connection_leases.rs

@@ -32,8 +32,9 @@ pub struct ConnectionLease {
     hung_up: AtomicBool,
     peer_closed: Option<PeerClosed>,
     /// Also holds a terminal slot (`ConnectionStats::try_claim_terminal`,
-    /// #27698), given back with the main slot.
-    terminal: bool,
+    /// #27698), given back with the main slot; inside, the terminal it came
+    /// from (nil when unknown).
+    terminal: Option<Option<u32>>,
 }
 
 impl ConnectionLease {
@@ -43,7 +44,7 @@ impl ConnectionLease {
     pub fn new(
         stats: Arc<ConnectionStats>,
         peer_closed: Option<PeerClosed>,
-        terminal: bool,
+        terminal: Option<Option<u32>>,
     ) -> Arc<Self> {
         let watched = peer_closed.is_some();
         let lease =
@@ -68,8 +69,8 @@ impl Drop for ConnectionLease {
             self.stats.hung_up.fetch_sub(1, Ordering::AcqRel);
         } else {
             self.stats.release();
-            if self.terminal {
-                self.stats.release_terminal();
+            if let Some(client) = self.terminal {
+                self.stats.release_terminal(client);
             }
         }
     }
@@ -107,8 +108,8 @@ impl ConnectionStats {
                 continue;
             }
             self.release();
-            if lease.terminal {
-                self.release_terminal();
+            if let Some(client) = lease.terminal {
+                self.release_terminal(client);
             }
             freed += 1;
         }

cmux-tui/crates/cmux-tui-core/src/server.rs

@@ -2644,7 +2644,7 @@ use bounded_outbound::{
 #[cfg(test)]
 use bounded_outbound::{
     ControlOutbound, MAX_TERMINAL_CONNECTIONS, claim_terminal_connection,
-    websocket_server_frame_header,
+    claim_terminal_connection_from, websocket_server_frame_header,
 };
 
 mod client_registry;

cmux-tui/crates/cmux-tui-core/src/server/bounded_outbound.rs

@@ -82,7 +82,7 @@ pub(super) struct ConnectionPermit {
 pub(super) fn claim_connection(
     connections: &Arc<crate::diagnostics::ConnectionStats>,
 ) -> Option<ConnectionPermit> {
-    claim_connection_watching(connections, None, false)
+    claim_connection_watching(connections, None, None)
 }
 
 /// A slot for a local client, taken back if the client hangs up while its
@@ -91,25 +91,37 @@ pub(super) fn claim_watched_connection(
     connections: &Arc<crate::diagnostics::ConnectionStats>,
     stream: &dyn transport::Stream,
 ) -> Option<ConnectionPermit> {
-    claim_connection_watching(connections, watch_peer(stream), false)
+    claim_connection_watching(connections, watch_peer(stream), None)
 }
 
 /// A slot for a process inside one of the session's terminals: it also takes
 /// a terminal slot, and is refused (its slot given back) when terminals
 /// already hold `MAX_TERMINAL_CONNECTIONS` (#27698). Watched like
 /// `claim_watched_connection`; a hang-up gives back both slots.
+/// `client`: the terminal the peer runs in (its process session), nil when
+/// unknown.
 pub(super) fn claim_watched_terminal_connection(
     connections: &Arc<crate::diagnostics::ConnectionStats>,
     stream: &dyn transport::Stream,
+    client: Option<u32>,
 ) -> Option<ConnectionPermit> {
-    claim_connection_watching(connections, watch_peer(stream), true)
+    claim_connection_watching(connections, watch_peer(stream), Some(client))
 }
 
 #[cfg(test)]
 pub(super) fn claim_terminal_connection(
     connections: &Arc<crate::diagnostics::ConnectionStats>,
 ) -> Option<ConnectionPermit> {
-    claim_connection_watching(connections, None, true)
+    claim_connection_watching(connections, None, Some(None))
+}
+
+/// A terminal slot for a connection from terminal `client` (its process session).
+#[cfg(test)]
+pub(super) fn claim_terminal_connection_from(
+    connections: &Arc<crate::diagnostics::ConnectionStats>,
+    client: u32,
+) -> Option<ConnectionPermit> {
+    claim_connection_watching(connections, None, Some(Some(client)))
 }
 
 fn watch_peer(stream: &dyn transport::Stream) -> Option<crate::diagnostics::PeerClosed> {
@@ -119,17 +131,22 @@ fn watch_peer(stream: &dyn transport::Stream) -> Option<crate::diagnostics::Peer
     }
 }
 
+/// `terminal`: nil for a connection from outside the session's terminals,
+/// else the terminal it comes from (nil inside when unknown).
 fn claim_connection_watching(
     connections: &Arc<crate::diagnostics::ConnectionStats>,
     peer_closed: Option<crate::diagnostics::PeerClosed>,
-    terminal: bool,
+    terminal: Option<Option<u32>>,
 ) -> Option<ConnectionPermit> {
     if !connections.try_claim(MAX_SERVER_CONNECTIONS as u64) {
         return None;
     }
-    if terminal
-        && !connections
-            .try_claim_terminal(MAX_TERMINAL_CONNECTIONS as u64, MAX_SERVER_CONNECTIONS as u64)
+    if let Some(client) = terminal
+        && !connections.try_claim_terminal(
+            MAX_TERMINAL_CONNECTIONS as u64,
+            MAX_SERVER_CONNECTIONS as u64,
+            client,
+        )
     {
         connections.release();
         return None;

cmux-tui/crates/cmux-tui-core/src/server/listen.rs

@@ -416,8 +416,10 @@ impl ServedSession {
         // state lock (`TerminalSessions`).
         let key = stream.peer_process_key();
         let connections = self.mux.connection_stats();
-        let permit = if self.terminal_sessions.contains_peer(&self.mux, key.as_deref()) {
-            claim_watched_terminal_connection(connections, &*stream)
+        let permit = if let Some(terminal) =
+            self.terminal_sessions.terminal_of_peer(&self.mux, key.as_deref())
+        {
+            claim_watched_terminal_connection(connections, &*stream, Some(terminal))
         } else {
             claim_watched_connection(connections, &*stream)
         };

cmux-tui/crates/cmux-tui-core/src/server/terminal_sessions.rs

@@ -28,20 +28,25 @@ pub(super) struct TerminalSessions {
 
 impl TerminalSessions {
     /// Whether the peer a peer key names runs in one of the session's
-    /// terminals.
+    /// terminals (tests; the accept loop asks ``terminal_of_peer``).
+    #[cfg(test)]
     pub(super) fn contains_peer(&self, mux: &Mux, peer_key: Option<&str>) -> bool {
-        let Some(session) = peer_key
+        self.terminal_of_peer(mux, peer_key).is_some()
+    }
+
+    /// The process session of the terminal the peer runs in (one client of the
+    /// terminal lane: its shell, the agent in it, the agent's hooks and a `cmux`
+    /// run in it), or nil when the peer runs in none.
+    pub(super) fn terminal_of_peer(&self, mux: &Mux, peer_key: Option<&str>) -> Option<u32> {
+        let session = peer_key
             .and_then(super::origin_gate::peer_pid)
-            .and_then(super::origin_gate::peer_session_id)
-        else {
-            return false;
-        };
+            .and_then(super::origin_gate::peer_session_id)?;
         // A known terminal answers at once: a burst of agents costs no read.
         // A miss (the app, a new terminal) reads the terminals again.
         if self.last.lock().recover_poison().contains(&session) {
-            return true;
+            return Some(session);
         }
-        self.current(mux).contains(&session)
+        self.current(mux).contains(&session).then_some(session)
     }
 
     fn current(&self, mux: &Mux) -> Arc<HashSet<u32>> {

cmux-tui/crates/cmux-tui-core/src/server/tests/outbound_and_scheduler.rs

@@ -305,6 +305,38 @@ fn terminal_connections_leave_slots_for_the_app() {
     assert_eq!(connections.active(), 0);
 }
 
+/// One terminal opening connections in a loop (an agent's hooks, a script
+/// running `cmux` in a pane) cannot take the terminal lane from the others: a
+/// terminal's first connection gets in whenever the lane has room, a further
+/// one only while that terminal then holds no more than the lane has left free.
+#[test]
+fn one_terminal_cannot_take_the_terminal_lane_from_the_others() {
+    let connections = Arc::new(crate::diagnostics::ConnectionStats::default());
+    let first: Vec<ConnectionPermit> = (0..MAX_TERMINAL_CONNECTIONS)
+        .map_while(|_| claim_terminal_connection_from(&connections, 101))
+        .collect();
+    assert!(
+        first.len() < MAX_TERMINAL_CONNECTIONS / 2 + 1,
+        "terminal 101 took {} of the {MAX_TERMINAL_CONNECTIONS} terminal connections",
+        first.len()
+    );
+    let second = claim_terminal_connection_from(&connections, 102)
+        .expect("another terminal's first connection finds room");
+    let third =
+        claim_terminal_connection_from(&connections, 103).expect("and so does a third terminal's");
+    assert!(
+        claim_connection(&connections).is_some(),
+        "the app's slots are untouched by the terminal lane"
+    );
+    let taken = first.len();
+    drop((first, second, third));
+    assert_eq!(connections.active(), 0);
+    let again: Vec<ConnectionPermit> = (0..MAX_TERMINAL_CONNECTIONS)
+        .map_while(|_| claim_terminal_connection_from(&connections, 101))
+        .collect();
+    assert_eq!(again.len(), taken, "a terminal's released connections count no more against it");
+}
+
 /// #27669: connection threads parked on the registry lock outlived their CLI
 /// clients and held every slot, so the app's reconnects were refused
 /// `server.busy`. A client that hung up gives its slot to the next one; a live
@@ -358,15 +390,15 @@ fn hung_up_terminal_clients_give_back_their_terminal_slots() {
         .map(|_| {
             let (daemon_side, client) = UnixStream::pair().unwrap();
             clients.push(client);
-            claim_watched_terminal_connection(&connections, &daemon_side)
+            claim_watched_terminal_connection(&connections, &daemon_side, None)
                 .expect("terminal slot below its share")
         })
         .collect();
     let (live_daemon_side, _live_client) = UnixStream::pair().unwrap();
-    assert!(claim_watched_terminal_connection(&connections, &live_daemon_side).is_none());
+    assert!(claim_watched_terminal_connection(&connections, &live_daemon_side, None).is_none());
 
     clients.pop();
-    let next = claim_watched_terminal_connection(&connections, &live_daemon_side)
+    let next = claim_watched_terminal_connection(&connections, &live_daemon_side, None)
         .expect("the terminal slot of a client that hung up serves the next one");
     let snapshot = connections.snapshot(MAX_SERVER_CONNECTIONS as u64);
     assert_eq!(snapshot.hung_up, 1);