feat(gateway): runtime introspection for status endpoint
This commit is contained in:
parent
5d0cf5b070
commit
aa989cbaf1
@ -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<AtomicUsize>,
|
||||
}
|
||||
|
||||
impl OutboundDispatcher {
|
||||
@ -30,12 +32,14 @@ impl OutboundDispatcher {
|
||||
channel_manager: ChannelManager,
|
||||
task_supervisor: TaskSupervisor,
|
||||
write_locks: ConversationWriteLocks,
|
||||
active_lanes: Arc<AtomicUsize>,
|
||||
) -> 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<AtomicUsize>,
|
||||
}
|
||||
|
||||
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 {
|
||||
|
||||
@ -106,6 +106,29 @@ impl MessageBus {
|
||||
pub async fn consume_control(&self) -> Option<ControlMessage> {
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
@ -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::time::Instant> = 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<AtomicUsize>,
|
||||
/// Active outbound dispatcher lane count (shared with the dispatcher).
|
||||
pub outbound_lanes: Arc<AtomicUsize>,
|
||||
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<String>,
|
||||
port: Option<u16>,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
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()?;
|
||||
|
||||
@ -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<String>,
|
||||
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::<WsOutbound>(100);
|
||||
|
||||
@ -131,6 +134,23 @@ async fn handle_socket(
|
||||
tracing::info!(session_id = %session_id, "CLI session ended");
|
||||
}
|
||||
|
||||
struct ConnectionGuard {
|
||||
counter: Arc<AtomicUsize>,
|
||||
}
|
||||
|
||||
impl ConnectionGuard {
|
||||
fn new(counter: Arc<AtomicUsize>) -> 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);
|
||||
}
|
||||
}
|
||||
|
||||
@ -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,
|
||||
|
||||
@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user