hive-agent: cap concurrent connections on the in-agent todo socket
This commit is contained in:
parent
533a39a00b
commit
e9940be983
1 changed files with 66 additions and 9 deletions
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in a new issue