diff --git a/src/bus/dispatcher.rs b/src/bus/dispatcher.rs index edb830e..0ae9db2 100644 --- a/src/bus/dispatcher.rs +++ b/src/bus/dispatcher.rs @@ -1,5 +1,6 @@ use std::collections::HashMap; use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; use tokio::sync::mpsc; @@ -22,6 +23,7 @@ pub struct OutboundDispatcher { channel_manager: ChannelManager, task_supervisor: TaskSupervisor, write_locks: ConversationWriteLocks, + active_lanes: Arc, } impl OutboundDispatcher { @@ -30,12 +32,14 @@ impl OutboundDispatcher { channel_manager: ChannelManager, task_supervisor: TaskSupervisor, write_locks: ConversationWriteLocks, + active_lanes: Arc, ) -> Self { Self { bus, channel_manager, task_supervisor, write_locks, + active_lanes, } } @@ -125,9 +129,14 @@ impl OutboundDispatcher { chat_id: String, ) -> bool { let target_lock = self.write_locks.for_target(&channel_name, &chat_id); + let active_lanes = self.active_lanes.clone(); self.task_supervisor.spawn( format!("outbound-lane:{channel_name}:{chat_id}"), async move { + active_lanes.fetch_add(1, Ordering::Relaxed); + let _guard = LaneGuard { + counter: active_lanes, + }; loop { let msg = match tokio::time::timeout(LANE_IDLE_TIMEOUT, receiver.recv()).await { Ok(Some(msg)) => msg, @@ -180,6 +189,18 @@ impl OutboundDispatcher { } } +/// Decrements the active-lane counter exactly once when a lane task ends, +/// whether it exits normally, is cancelled, or is aborted. +struct LaneGuard { + counter: Arc, +} + +impl Drop for LaneGuard { + fn drop(&mut self) { + self.counter.fetch_sub(1, Ordering::Relaxed); + } +} + #[cfg(test)] mod tests { use super::*; @@ -275,6 +296,7 @@ mod tests { manager, supervisor.clone(), ConversationWriteLocks::default(), + Arc::new(AtomicUsize::new(0)), ); let task = tokio::spawn(async move { dispatcher.run().await }); bus.publish_outbound(outbound("slow", "slow-1")) @@ -318,6 +340,7 @@ mod tests { manager, supervisor.clone(), ConversationWriteLocks::default(), + Arc::new(AtomicUsize::new(0)), ); let task = tokio::spawn(async move { dispatcher.run().await }); @@ -348,6 +371,7 @@ mod tests { manager, supervisor.clone(), ConversationWriteLocks::default(), + Arc::new(AtomicUsize::new(0)), ); let task = tokio::spawn(async move { dispatcher.run().await }); @@ -360,6 +384,50 @@ mod tests { supervisor.shutdown(Duration::from_secs(1)).await; } + #[tokio::test] + async fn active_lane_count_tracks_task_lifetime() { + let bus = MessageBus::new(8); + let manager = ChannelManager::with_bus( + Arc::new(crate::channels::CliChatChannel::new()), + bus.clone(), + ); + manager + .register_channel( + "recording", + Arc::new(RecordingChannel { + sent: Mutex::new(Vec::new()), + notify: Notify::new(), + }), + ) + .await; + let supervisor = TaskSupervisor::new(); + let active_lanes = Arc::new(AtomicUsize::new(0)); + let dispatcher = OutboundDispatcher::new( + bus.clone(), + manager, + supervisor.clone(), + ConversationWriteLocks::default(), + active_lanes.clone(), + ); + let dispatcher_task = tokio::spawn(async move { dispatcher.run().await }); + + bus.publish_outbound(outbound("counted", "message")) + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(1), async { + while active_lanes.load(Ordering::Relaxed) != 1 { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + + supervisor.shutdown(Duration::from_secs(1)).await; + assert_eq!(active_lanes.load(Ordering::Relaxed), 0); + dispatcher_task.abort(); + let _ = dispatcher_task.await; + } + #[tokio::test] async fn permanent_send_failure_is_not_retried() { let channel = PermanentFailureChannel { diff --git a/src/bus/mod.rs b/src/bus/mod.rs index 0a4b1d2..9259288 100644 --- a/src/bus/mod.rs +++ b/src/bus/mod.rs @@ -106,6 +106,29 @@ impl MessageBus { pub async fn consume_control(&self) -> Option { self.control_rx.lock().await.recv().await } + + /// Snapshot of the current depth and capacity of each bus queue. + pub fn queue_depths(&self) -> QueueDepths { + QueueDepths { + inbound_depth: (self.inbound_tx.max_capacity() - self.inbound_tx.capacity()) as u64, + inbound_cap: self.inbound_tx.max_capacity() as u64, + outbound_depth: (self.outbound_tx.max_capacity() - self.outbound_tx.capacity()) as u64, + outbound_cap: self.outbound_tx.max_capacity() as u64, + control_depth: (self.control_tx.max_capacity() - self.control_tx.capacity()) as u64, + control_cap: self.control_tx.max_capacity() as u64, + } + } +} + +/// Read-only snapshot of MessageBus queue utilization. +#[derive(serde::Serialize)] +pub struct QueueDepths { + pub inbound_depth: u64, + pub inbound_cap: u64, + pub outbound_depth: u64, + pub outbound_cap: u64, + pub control_depth: u64, + pub control_cap: u64, } // ============================================================================ @@ -130,3 +153,38 @@ impl std::fmt::Display for BusError { } impl std::error::Error for BusError {} + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::HashMap; + + #[tokio::test] + async fn queue_depths_report_retained_sender_usage() { + let bus = MessageBus::new(3); + + let empty = bus.queue_depths(); + assert_eq!(empty.inbound_depth, 0); + assert_eq!(empty.inbound_cap, 3); + assert_eq!(empty.outbound_depth, 0); + assert_eq!(empty.outbound_cap, 3); + assert_eq!(empty.control_depth, 0); + assert_eq!(empty.control_cap, 3); + + bus.publish_outbound(OutboundMessage { + channel: "test".to_string(), + chat_id: "chat".to_string(), + content: "queued".to_string(), + reply_to: None, + media: vec![], + metadata: HashMap::new(), + delivery: None, + }) + .await + .unwrap(); + + assert_eq!(bus.queue_depths().outbound_depth, 1); + bus.consume_outbound().await.unwrap(); + assert_eq!(bus.queue_depths().outbound_depth, 0); + } +} diff --git a/src/gateway/mod.rs b/src/gateway/mod.rs index d0937e7..3a85eb2 100644 --- a/src/gateway/mod.rs +++ b/src/gateway/mod.rs @@ -8,6 +8,7 @@ pub mod ws; use axum::{Router, middleware, routing}; use std::net::SocketAddr; use std::sync::Arc; +use std::sync::atomic::AtomicUsize; use tokio::net::TcpListener; use crate::bus::{MessageBus, OutboundDispatcher}; @@ -21,6 +22,18 @@ use crate::scheduler::Scheduler; use crate::session::{SessionManager, SessionManagerServices}; use crate::task_supervisor::TaskSupervisor; +/// Process boot clock. A process-level static so uptime survives config reload, +/// which swaps GatewayState generations without restarting the process. +static STARTED: std::sync::OnceLock = std::sync::OnceLock::new(); + +/// Seconds elapsed since the gateway process started. +pub fn process_uptime_secs() -> u64 { + STARTED + .get_or_init(std::time::Instant::now) + .elapsed() + .as_secs() +} + pub struct GatewayState { pub config: Config, pub config_path: std::path::PathBuf, @@ -33,6 +46,10 @@ pub struct GatewayState { pub connection_shutdown: tokio_util::sync::CancellationToken, pub auth: auth::AuthManager, pub uploads: uploads::UploadRegistry, + /// Live WebSocket connection count. + pub ws_connections: Arc, + /// Active outbound dispatcher lane count (shared with the dispatcher). + pub outbound_lanes: Arc, pub(crate) reload: reload::ReloadHandle, pub(crate) admission: reload::RuntimeAdmission, } @@ -231,6 +248,8 @@ impl GatewayState { connection_shutdown, auth, uploads, + ws_connections: Arc::new(AtomicUsize::new(0)), + outbound_lanes: Arc::new(AtomicUsize::new(0)), reload, admission, }) @@ -312,6 +331,7 @@ impl GatewayState { self.channel_manager.clone(), self.task_supervisor.clone(), self.delivery_coordinator.write_locks(), + self.outbound_lanes.clone(), ); self.task_supervisor @@ -341,6 +361,7 @@ pub async fn run( host: Option, port: Option, ) -> Result<(), Box> { + STARTED.get_or_init(std::time::Instant::now); let config_path = crate::config::resolve_default_config_path(); let startup_process_env = Config::startup_process_env(); let startup_cwd = std::env::current_dir()?; diff --git a/src/gateway/ws.rs b/src/gateway/ws.rs index 064e38d..85eb6ea 100644 --- a/src/gateway/ws.rs +++ b/src/gateway/ws.rs @@ -7,6 +7,7 @@ use axum::response::Response; use futures_util::{SinkExt, StreamExt}; use serde::Deserialize; use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; use tokio::sync::mpsc; use tokio::time::{Duration, timeout}; @@ -42,6 +43,8 @@ async fn handle_socket( client_id: Option, identity: super::auth::AuthIdentity, ) { + let _connection_guard = ConnectionGuard::new(state.ws_connections.clone()); + // Create channel for sending outbound messages to this client let (sender, mut receiver) = mpsc::channel::(100); @@ -131,6 +134,23 @@ async fn handle_socket( tracing::info!(session_id = %session_id, "CLI session ended"); } +struct ConnectionGuard { + counter: Arc, +} + +impl ConnectionGuard { + fn new(counter: Arc) -> Self { + counter.fetch_add(1, Ordering::Relaxed); + Self { counter } + } +} + +impl Drop for ConnectionGuard { + fn drop(&mut self) { + self.counter.fetch_sub(1, Ordering::Relaxed); + } +} + #[cfg(test)] mod tests { use super::*; @@ -145,4 +165,14 @@ mod tests { assert!(valid_client_id(Some("x".repeat(65))).is_none()); assert!(valid_client_id(Some(String::new())).is_none()); } + + #[test] + fn connection_guard_tracks_its_scope() { + let connections = Arc::new(AtomicUsize::new(0)); + { + let _guard = ConnectionGuard::new(connections.clone()); + assert_eq!(connections.load(Ordering::Relaxed), 1); + } + assert_eq!(connections.load(Ordering::Relaxed), 0); + } } diff --git a/src/session/session.rs b/src/session/session.rs index bcd08e9..0d8c340 100644 --- a/src/session/session.rs +++ b/src/session/session.rs @@ -432,6 +432,7 @@ mod cancelled_partial_tests { channels, supervisor.clone(), ConversationWriteLocks::default(), + Arc::new(std::sync::atomic::AtomicUsize::new(0)), ); let dispatcher_task = tokio::spawn(async move { dispatcher.run().await }); @@ -2157,6 +2158,27 @@ impl SessionManager { } } + /// Number of live sessions currently tracked by the manager. + pub async fn session_count(&self) -> usize { + self.inner.lock().await.sessions.len() + } + + /// Number of sessions with an actively executing Turn. + pub async fn active_turn_count(&self) -> usize { + let sessions: Vec<_> = { + let inner = self.inner.lock().await; + inner.sessions.values().cloned().collect() + }; + let mut count = 0; + for session in sessions { + let session = session.lock().await; + if session.current_cancel.is_some() { + count += 1; + } + } + count + } + pub async fn create_session( &self, channel: &str, diff --git a/src/task_supervisor.rs b/src/task_supervisor.rs index c07c614..9d48f11 100644 --- a/src/task_supervisor.rs +++ b/src/task_supervisor.rs @@ -114,6 +114,13 @@ impl TaskSupervisor { true } + /// Number of currently running managed tasks, pruning finished handles. + pub fn running_count(&self) -> usize { + let mut state = self.inner.state.lock().unwrap_or_else(|e| e.into_inner()); + state.tasks.retain(|task| !task.handle.is_finished()); + state.tasks.len() + } + pub fn cancel(&self) { let mut state = self.inner.state.lock().unwrap_or_else(|e| e.into_inner()); state.stopping = true; @@ -199,4 +206,18 @@ mod tests { assert!(cleaned_up.load(std::sync::atomic::Ordering::SeqCst)); } + + #[tokio::test] + async fn running_count_prunes_finished_tasks() { + let supervisor = TaskSupervisor::new(); + assert!(supervisor.spawn("pending", std::future::pending())); + assert_eq!(supervisor.running_count(), 1); + + assert!(supervisor.spawn("finished", async {})); + tokio::task::yield_now().await; + + assert_eq!(supervisor.running_count(), 1); + supervisor.shutdown(Duration::from_secs(1)).await; + assert_eq!(supervisor.running_count(), 0); + } }