diff --git a/crates/maestro-daemon/src/broker.rs b/crates/maestro-daemon/src/broker.rs index 2cd880c..05792a9 100644 --- a/crates/maestro-daemon/src/broker.rs +++ b/crates/maestro-daemon/src/broker.rs @@ -6,13 +6,33 @@ //! send and the reply correlation live here so the handler stays thin. use std::collections::HashMap; -use std::sync::Arc; +use std::hash::Hash; +use std::sync::{Arc, Mutex}; use std::time::Duration; use maestro_protocol::{AttachCandidates, DaemonOutbound, ReplyTo, RpcResult}; -use tokio::sync::{mpsc, oneshot, Mutex}; +use tokio::sync::{mpsc, oneshot}; use uuid::Uuid; +/// RAII cleanup for a pending-RPC entry. Removing the entry on drop guarantees +/// it is gone on *every* exit path of the awaiting call — success, error, +/// timeout, and crucially future cancellation (the axum handler is dropped when +/// the client hangs up) — so the map can't grow without bound (NAP-259). +/// +/// The map is a `std::sync::Mutex` (never held across an await) precisely so the +/// sync `Drop` can lock it. Removal is idempotent: on the happy path the reply +/// handler has already taken the entry and this is a harmless no-op. +struct PendingGuard<'a, K: Eq + Hash, V> { + map: &'a Mutex>, + key: K, +} + +impl Drop for PendingGuard<'_, K, V> { + fn drop(&mut self) { + self.map.lock().unwrap().remove(&self.key); + } +} + /// How long to wait for a `ResolveAttachResult` after the request goes out. /// Matches the registry's RPC timeout — a healthy registry replies in ms. pub const RESOLVE_TIMEOUT: Duration = Duration::from_secs(10); @@ -55,7 +75,13 @@ impl BrokerHandle { pub async fn resolve_attach(&self, session_id: Uuid) -> Result { let reply_to = Uuid::new_v4(); let (tx, rx) = oneshot::channel(); - self.pending.lock().await.insert(reply_to, tx); + self.pending.lock().unwrap().insert(reply_to, tx); + // From here on, every exit path (including cancellation of this future) + // removes the pending entry when `_guard` drops (NAP-259). + let _guard = PendingGuard { + map: &*self.pending, + key: reply_to, + }; if self .out_tx @@ -66,7 +92,6 @@ impl BrokerHandle { .await .is_err() { - self.pending.lock().await.remove(&reply_to); return Err(BrokerError::NoControlLink); } @@ -76,10 +101,7 @@ impl BrokerHandle { Err(BrokerError::Registry(error.code, error.message)) } Ok(Err(_)) => Err(BrokerError::ReplyClosed), - Err(_) => { - self.pending.lock().await.remove(&reply_to); - Err(BrokerError::Timeout(RESOLVE_TIMEOUT)) - } + Err(_) => Err(BrokerError::Timeout(RESOLVE_TIMEOUT)), } } @@ -91,8 +113,8 @@ impl BrokerHandle { reply_to: ReplyTo, result: RpcResult, ) -> bool { - let mut pending = self.pending.lock().await; - match pending.remove(&reply_to) { + let entry = self.pending.lock().unwrap().remove(&reply_to); + match entry { Some(tx) => tx.send(result).is_ok(), None => { tracing::debug!(?reply_to, "no waiter for ResolveAttachResult"); @@ -101,3 +123,39 @@ impl BrokerHandle { } } } + +#[cfg(test)] +mod tests { + use super::*; + + /// Cancelling (dropping) an in-flight `resolve_attach` must not leak its + /// pending-reply entry — otherwise a browser hanging up mid-attach grows + /// the map without bound (NAP-259). + #[tokio::test] + async fn cancelled_resolve_attach_removes_pending_entry() { + // Keep the receiver alive so `out_tx.send()` succeeds and the call + // parks on the reply oneshot rather than failing fast. + let (handle, _out_rx) = BrokerHandle::new(); + + { + let fut = handle.resolve_attach(Uuid::new_v4()); + tokio::pin!(fut); + tokio::select! { + biased; + _ = &mut fut => panic!("resolve must not complete without a reply"), + _ = tokio::time::sleep(Duration::from_millis(20)) => {} + } + assert_eq!( + handle.pending.lock().unwrap().len(), + 1, + "entry should be in-flight before cancellation" + ); + } + + assert_eq!( + handle.pending.lock().unwrap().len(), + 0, + "pending entry must be removed on cancellation (NAP-259)" + ); + } +} diff --git a/crates/maestro-registry/src/state.rs b/crates/maestro-registry/src/state.rs index cb53b06..25f25af 100644 --- a/crates/maestro-registry/src/state.rs +++ b/crates/maestro-registry/src/state.rs @@ -10,7 +10,8 @@ //! handful of daemons and these critical sections are tiny. use std::collections::{BTreeMap, HashMap}; -use std::sync::Arc; +use std::hash::Hash; +use std::sync::{Arc, Mutex as StdMutex}; use std::time::Duration; use maestro_protocol::{ @@ -22,6 +23,26 @@ use uuid::Uuid; use crate::jwt::issue_pair_attach_jwt; +/// RAII cleanup for a pending-RPC entry. Removing the entry on drop guarantees +/// it is gone on *every* exit path of the awaiting call — success, error, +/// timeout, and crucially future cancellation (an axum handler is dropped when +/// the HTTP client hangs up) — so the map can't grow without bound (NAP-259). +/// +/// The pending maps live behind a `std::sync::Mutex` (never held across an +/// await) precisely so this sync `Drop` can lock them. Removal is idempotent: +/// on the happy path the reply handler has already taken the entry, making the +/// drop a harmless no-op. +struct PendingGuard<'a, K: Eq + Hash, V> { + map: &'a StdMutex>, + key: K, +} + +impl Drop for PendingGuard<'_, K, V> { + fn drop(&mut self) { + self.map.lock().unwrap().remove(&self.key); + } +} + /// Capacity of the registry-wide broadcast channel that feeds /// `WS /api/events` subscribers. Subscribers that fall this far behind get a /// `Lagged` event and are expected to resync via `GET /api/sessions`. @@ -80,6 +101,11 @@ impl From for DaemonInfo { pub struct Registry { inner: Arc>, events: broadcast::Sender, + /// Outstanding R->D RPC waiters. Held behind a `std::sync::Mutex` (separate + /// from `State`'s async lock) so a [`PendingGuard`] can clean them up on + /// drop — including when the awaiting HTTP handler is cancelled (NAP-259). + pending_create: Arc>>>>, + pending_kill: Arc>>>>, } #[derive(Debug, Default)] @@ -87,8 +113,6 @@ struct State { daemons: HashMap, /// `(daemon_id, session_id) -> Session`. sessions: BTreeMap<(String, Uuid), Session>, - pending_create: HashMap>>, - pending_kill: HashMap>>, } #[derive(Debug)] @@ -126,6 +150,8 @@ impl Registry { Self { inner: Arc::new(Mutex::new(State::default())), events, + pending_create: Arc::new(StdMutex::new(HashMap::new())), + pending_kill: Arc::new(StdMutex::new(HashMap::new())), } } @@ -420,8 +446,8 @@ impl Registry { /// Resolve a CreateSession RPC reply. Returns true if a waiter was found. pub async fn complete_create(&self, reply_to: ReplyTo, result: RpcResult) -> bool { - let mut state = self.inner.lock().await; - match state.pending_create.remove(&reply_to) { + let entry = self.pending_create.lock().unwrap().remove(&reply_to); + match entry { Some(tx) => tx.send(result).is_ok(), None => false, } @@ -429,8 +455,8 @@ impl Registry { /// Resolve a KillSession RPC reply. Returns true if a waiter was found. pub async fn complete_kill(&self, reply_to: ReplyTo, result: RpcResult<()>) -> bool { - let mut state = self.inner.lock().await; - match state.pending_kill.remove(&reply_to) { + let entry = self.pending_kill.lock().unwrap().remove(&reply_to); + match entry { Some(tx) => tx.send(result).is_ok(), None => false, } @@ -450,8 +476,7 @@ impl Registry { let reply_to = Uuid::new_v4(); let (tx, rx) = oneshot::channel(); let outbound = { - let mut state = self.inner.lock().await; - state.pending_create.insert(reply_to, tx); + let state = self.inner.lock().await; state .daemons .get(daemon_id) @@ -459,16 +484,22 @@ impl Registry { .outbound .clone() }; - let send_result = outbound + self.pending_create.lock().unwrap().insert(reply_to, tx); + // Every exit path past here — send failure, timeout, a successful reply, + // or cancellation of this future — removes the pending entry when + // `_guard` drops (NAP-259). + let _guard = PendingGuard { + map: &*self.pending_create, + key: reply_to, + }; + if outbound .send(DaemonInbound::CreateSession { reply_to, request }) - .await; - if send_result.is_err() { - self.inner.lock().await.pending_create.remove(&reply_to); + .await + .is_err() + { return Err(RegistryError::DaemonGone); } - let rpc = await_rpc(rx, reply_to, |s| &mut s.pending_create, &self.inner) - .await - .map(rpc_into_result)?; + let rpc = await_rpc(rx).await.map(rpc_into_result)?; // Eagerly mirror the session into the catalog on success — the daemon // also emits `SessionAdded` asynchronously, but the HTTP caller would // otherwise see an empty `GET /api/sessions` immediately after a @@ -491,8 +522,7 @@ impl Registry { let reply_to = Uuid::new_v4(); let (tx, rx) = oneshot::channel(); let outbound = { - let mut state = self.inner.lock().await; - state.pending_kill.insert(reply_to, tx); + let state = self.inner.lock().await; state .daemons .get(daemon_id) @@ -500,16 +530,21 @@ impl Registry { .outbound .clone() }; - let send_result = outbound + self.pending_kill.lock().unwrap().insert(reply_to, tx); + // See `create_session`: `_guard` removes the pending entry on every exit + // path, including cancellation of this future (NAP-259). + let _guard = PendingGuard { + map: &*self.pending_kill, + key: reply_to, + }; + if outbound .send(DaemonInbound::KillSession { reply_to, id }) - .await; - if send_result.is_err() { - self.inner.lock().await.pending_kill.remove(&reply_to); + .await + .is_err() + { return Err(RegistryError::DaemonGone); } - let rpc = await_rpc(rx, reply_to, |s| &mut s.pending_kill, &self.inner) - .await - .map(rpc_into_result)?; + let rpc = await_rpc(rx).await.map(rpc_into_result)?; // Mirror of the create-path eager update: drop the session locally so // subsequent reads see consistent state. The daemon also emits // `SessionRemoved`; `remove_session` is idempotent. @@ -567,22 +602,14 @@ impl State { } } -/// Await `rx` with [`RPC_TIMEOUT`]; on timeout / channel close, also clean up -/// the pending map so a late reply doesn't leak. -async fn await_rpc( - rx: oneshot::Receiver>, - reply_to: ReplyTo, - pick: impl FnOnce(&mut State) -> &mut HashMap>>, - inner: &Arc>, -) -> Result, RegistryError> { +/// Await `rx` with [`RPC_TIMEOUT`]. Cleanup of the pending map on every exit +/// path (timeout, channel close, cancellation) is handled by the caller's +/// [`PendingGuard`], so this only has to classify the outcome. +async fn await_rpc(rx: oneshot::Receiver>) -> Result, RegistryError> { match tokio::time::timeout(RPC_TIMEOUT, rx).await { Ok(Ok(result)) => Ok(result), Ok(Err(_)) => Err(RegistryError::ReplyChannelClosed), - Err(_) => { - let mut state = inner.lock().await; - pick(&mut state).remove(&reply_to); - Err(RegistryError::Timeout(RPC_TIMEOUT)) - } + Err(_) => Err(RegistryError::Timeout(RPC_TIMEOUT)), } } @@ -675,3 +702,49 @@ mod tests { assert!(reg.daemons().await.is_empty()); } } + +#[cfg(test)] +mod nap259_tests { + use super::*; + + /// Cancelling an in-flight `kill_session` must not leak its pending-reply + /// entry — a client that hangs up mid-request would otherwise grow the map + /// without bound (NAP-259). + #[tokio::test] + async fn cancelled_kill_session_removes_pending_entry() { + let registry = Registry::new(); + // Register a daemon with a live outbound receiver so the R->D send + // succeeds and the call parks on the reply oneshot. + let (tx, _rx) = mpsc::channel::(8); + registry + .register( + "d1".into(), + "host".into(), + "0.1.0".into(), + "ws://127.0.0.1:1".into(), + tx, + ) + .await; + + { + let fut = registry.kill_session("d1", Uuid::new_v4()); + tokio::pin!(fut); + tokio::select! { + biased; + _ = &mut fut => panic!("kill must not complete without a reply"), + _ = tokio::time::sleep(Duration::from_millis(20)) => {} + } + assert_eq!( + registry.pending_kill.lock().unwrap().len(), + 1, + "entry should be in-flight before cancellation" + ); + } + + assert_eq!( + registry.pending_kill.lock().unwrap().len(), + 0, + "pending entry must be removed on cancellation (NAP-259)" + ); + } +}