Watch
0
0
Fork
You've already forked hyperhive
0

hive-agent: cap concurrent connections on the in-agent todo socket

This commit is contained in:
damocles 2026-11-07 08:59:13 +01:00 • committed by mara
commit e9940be983

View file

@ -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<PathBuf> {
@ -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<F, Fut>(listener: UnixListener, max: usize, handle: F) -> Result<()>
where
F: Fn(UnixStream) -> Fut,
Fut: std::future::Future<Output = ()> + 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