From e9940be9834cb08706851f79d9802ba744346fc2 Mon Sep 17 00:00:00 2001 From: damocles Date: Sat, 7 Nov 2026 08:59:13 +0100 Subject: [PATCH] hive-agent: cap concurrent connections on the in-agent todo socket --- hive-agent/src/todo_server.rs | 75 ++++++++++++++++++++++++++++++----- 1 file changed, 66 insertions(+), 9 deletions(-) diff --git a/hive-agent/src/todo_server.rs b/hive-agent/src/todo_server.rs index 6d755dad..a98807cf 100644 --- a/hive-agent/src/todo_server.rs +++ b/hive-agent/src/todo_server.rs @@ -47,6 +47,12 @@ const COMPACT_MIN_USAGE_FRACTION: f64 = 0.66; /// does not turn the accept loop into a busy spin. const ACCEPT_RETRY_DELAY: std::time::Duration = std::time::Duration::from_secs(1); +/// Most connections served at once. The socket is mode 0666, so any process +/// in the container can dial it; past this the accept loop stops accepting +/// (the kernel queues the rest) instead of spawning a task and holding an fd +/// per stalled client. +const MAX_CONNECTIONS: usize = 64; + /// Resolve the in-agent socket path from `HIVE_AGENT_SOCKET`. `None` when /// unset/empty (e.g. a standalone dev run) — the server then stays off. fn socket_path() -> Option { @@ -96,19 +102,39 @@ pub async fn run( }; let listener = bind(&path)?; tracing::info!(socket = %path.display(), "in-agent todo socket listening"); + accept_loop(listener, MAX_CONNECTIONS, move |stream| { + let (store, wake, reminders, bus) = + (store.clone(), wake.clone(), reminders.clone(), bus.clone()); + async move { + if let Err(e) = handle_conn(stream, &store, &wake, reminders.as_deref(), &bus).await { + tracing::warn!(error = ?e, "in-agent todo connection failed"); + } + } + }) + .await +} + +/// Accept connections for ever, serving at most `max` at a time with +/// `handle`. A permit is taken before `accept`, so at the limit the loop +/// waits there and the kernel's listen queue holds the rest. +async fn accept_loop(listener: UnixListener, max: usize, handle: F) -> Result<()> +where + F: Fn(UnixStream) -> Fut, + Fut: std::future::Future + Send + 'static, +{ + let permits = Arc::new(tokio::sync::Semaphore::new(max)); loop { + let permit = permits + .clone() + .acquire_owned() + .await + .context("connection limit closed")?; match listener.accept().await { Ok((stream, _)) => { - let store = store.clone(); - let wake = wake.clone(); - let reminders = reminders.clone(); - let bus = bus.clone(); + let served = handle(stream); tokio::spawn(async move { - if let Err(e) = - handle_conn(stream, &store, &wake, reminders.as_deref(), &bus).await - { - tracing::warn!(error = ?e, "in-agent todo connection failed"); - } + served.await; + drop(permit); }); } Err(e) => { @@ -502,6 +528,37 @@ mod tests { use hive_sock_client::REQUEST_LINE_MAX_BYTES; use tokio::io::AsyncBufReadExt; + /// At the limit the loop stops accepting; a freed slot lets the next + /// connection in. + #[tokio::test] + async fn the_accept_loop_serves_no_more_than_its_limit_at_once() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + let dir = tempfile::tempdir().expect("tempdir"); + let path = dir.path().join("s.sock"); + let listener = UnixListener::bind(&path).expect("bind"); + let started = Arc::new(AtomicUsize::new(0)); + let release = Arc::new(Notify::new()); + let (s, r) = (started.clone(), release.clone()); + let server = tokio::spawn(accept_loop(listener, 1, move |_stream| { + let (s, r) = (s.clone(), r.clone()); + async move { + s.fetch_add(1, Ordering::SeqCst); + r.notified().await; + } + })); + + let _a = UnixStream::connect(&path).await.expect("first"); + let _b = UnixStream::connect(&path).await.expect("second"); + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + assert_eq!(started.load(Ordering::SeqCst), 1, "second waited its turn"); + + release.notify_one(); + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + assert_eq!(started.load(Ordering::SeqCst), 2, "freed slot admitted it"); + server.abort(); + } + /// `saturating_age`'s own arms are exact (`hive-sh4re::inbox`); what /// is only checkable here is *which* stored timestamp each projection /// feeds it, and both read the clock internally — so these assert a