diff --git a/cmux-tui/crates/cmux-tui-core/src/diagnostics.rs b/cmux-tui/crates/cmux-tui-core/src/diagnostics.rs index 9a3b62d6b1ab..d05fe44236cf 100644 --- a/cmux-tui/crates/cmux-tui-core/src/diagnostics.rs +++ b/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>, /// The socket these connections come in on, once a server attached the /// session to it. socket: std::sync::OnceLock, @@ -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) -> 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) -> 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) { + 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); } diff --git a/cmux-tui/crates/cmux-tui-core/src/diagnostics/connection_leases.rs b/cmux-tui/crates/cmux-tui-core/src/diagnostics/connection_leases.rs index b486ae6511c9..353e3572b4cd 100644 --- a/cmux-tui/crates/cmux-tui-core/src/diagnostics/connection_leases.rs +++ b/cmux-tui/crates/cmux-tui-core/src/diagnostics/connection_leases.rs @@ -32,8 +32,9 @@ pub struct ConnectionLease { hung_up: AtomicBool, peer_closed: Option, /// 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>, } impl ConnectionLease { @@ -43,7 +44,7 @@ impl ConnectionLease { pub fn new( stats: Arc, peer_closed: Option, - terminal: bool, + terminal: Option>, ) -> Arc { 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; } diff --git a/cmux-tui/crates/cmux-tui-core/src/server.rs b/cmux-tui/crates/cmux-tui-core/src/server.rs index 090a846144d8..ed2e6f8f281a 100644 --- a/cmux-tui/crates/cmux-tui-core/src/server.rs +++ b/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; diff --git a/cmux-tui/crates/cmux-tui-core/src/server/bounded_outbound.rs b/cmux-tui/crates/cmux-tui-core/src/server/bounded_outbound.rs index 196ab4f90204..7e064e6590d2 100644 --- a/cmux-tui/crates/cmux-tui-core/src/server/bounded_outbound.rs +++ b/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, ) -> Option { - 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, stream: &dyn transport::Stream, ) -> Option { - 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, stream: &dyn transport::Stream, + client: Option, ) -> Option { - 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, ) -> Option { - 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, + client: u32, +) -> Option { + claim_connection_watching(connections, None, Some(Some(client))) } fn watch_peer(stream: &dyn transport::Stream) -> Option { @@ -119,17 +131,22 @@ fn watch_peer(stream: &dyn transport::Stream) -> Option, peer_closed: Option, - terminal: bool, + terminal: Option>, ) -> Option { 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; diff --git a/cmux-tui/crates/cmux-tui-core/src/server/listen.rs b/cmux-tui/crates/cmux-tui-core/src/server/listen.rs index 70e74b010442..a5963587db66 100644 --- a/cmux-tui/crates/cmux-tui-core/src/server/listen.rs +++ b/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) }; diff --git a/cmux-tui/crates/cmux-tui-core/src/server/terminal_sessions.rs b/cmux-tui/crates/cmux-tui-core/src/server/terminal_sessions.rs index 923fc4ac5ebb..2714fdd959a7 100644 --- a/cmux-tui/crates/cmux-tui-core/src/server/terminal_sessions.rs +++ b/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 { + 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> { diff --git a/cmux-tui/crates/cmux-tui-core/src/server/tests/outbound_and_scheduler.rs b/cmux-tui/crates/cmux-tui-core/src/server/tests/outbound_and_scheduler.rs index 85731a2f730b..bdab683930b2 100644 --- a/cmux-tui/crates/cmux-tui-core/src/server/tests/outbound_and_scheduler.rs +++ b/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 = (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 = (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);