diff --git a/crates/maestro-daemon/src/registry_client.rs b/crates/maestro-daemon/src/registry_client.rs index 66cd187..f3eb98f 100644 --- a/crates/maestro-daemon/src/registry_client.rs +++ b/crates/maestro-daemon/src/registry_client.rs @@ -7,9 +7,12 @@ //! 2. Send `SessionsSnapshot { sessions: snapshot() }`. //! 3. Drive a `select!` loop that pumps: //! - inbound `CreateSession` / `KillSession` → `SessionManager` → reply -//! - outbound `SessionAdded/Removed/Updated` from the local event bus +//! - outbound `SessionAdded/Removed/Updated` from the local event bus, +//! except the event that an in-flight RPC reply already conveys, which is +//! swallowed so the registry writes each change once (see [`RpcCovered`]). //! - a 10s `Heartbeat` ticker. +use std::collections::HashSet; use std::time::Duration; use futures_util::{SinkExt, StreamExt}; @@ -18,6 +21,7 @@ use maestro_protocol::{ }; use tokio::sync::broadcast::error::RecvError; use tokio_tungstenite::tungstenite::Message; +use uuid::Uuid; use tokio::sync::mpsc; @@ -32,6 +36,31 @@ pub const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(10); const BACKOFF_MIN: Duration = Duration::from_millis(250); const BACKOFF_MAX: Duration = Duration::from_secs(30); +/// Session ids whose next lifecycle event the registry already learns about via +/// an RPC reply, so forwarding the matching async `SessionAdded` / +/// `SessionRemoved` would be a redundant *second* write to the registry catalog +/// (NAP-268). +/// +/// A registry-driven `CreateSession` / `KillSession` is reported back to the +/// registry in two ways: the `CreateSessionResult` / `KillSessionResult` reply +/// (which the registry applies eagerly), *and* the async lifecycle event this +/// daemon emits on its own `SessionManager` stream. Both used to land, so the +/// registry's catalog write had to be idempotent to absorb the duplicate. We +/// drop the duplicate at the source instead: when an RPC creates or kills a +/// session, its id is parked here and the matching async event is swallowed +/// (consumed exactly once) before it reaches the registry. Spontaneous changes +/// — sessions created via the daemon's local HTTP API, or reaped on their own — +/// are not parked and still forward normally, since the registry has no other +/// way to learn about them. +/// +/// Tracked per-connection: a reconnect starts fresh (the post-handshake +/// `SessionsSnapshot` reconciles the registry's full set anyway). +#[derive(Debug, Default)] +struct RpcCovered { + added: HashSet, + removed: HashSet, +} + #[derive(Debug, thiserror::Error)] pub enum ClientError { #[error("ws error: {0}")] @@ -225,6 +254,11 @@ impl RegistryClient { let mut hb = tokio::time::interval(HEARTBEAT_INTERVAL); hb.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + // Session ids whose async lifecycle event the in-flight RPC reply + // already conveys to the registry — swallowed here so the registry + // writes each change exactly once (NAP-268). + let mut covered = RpcCovered::default(); + loop { tokio::select! { msg = stream.next() => { @@ -233,13 +267,13 @@ impl RegistryClient { Some(Err(e)) => return Err(e.into()), Some(Ok(Message::Close(_))) => break, Some(Ok(Message::Text(text))) => { - self.handle_inbound(&text, &mut sink).await?; + self.handle_inbound(&text, &mut sink, &mut covered).await?; } Some(Ok(_)) => {} } } ev = events.recv() => { - if !forward_session_event(&self.manager, ev, &mut sink).await? { + if !forward_session_event(&self.manager, ev, &mut sink, &mut covered).await? { break; } } @@ -260,7 +294,12 @@ impl RegistryClient { Ok(()) } - async fn handle_inbound(&self, text: &str, sink: &mut S) -> Result<(), ClientError> + async fn handle_inbound( + &self, + text: &str, + sink: &mut S, + covered: &mut RpcCovered, + ) -> Result<(), ClientError> where S: SinkExt + Unpin, { @@ -280,7 +319,13 @@ impl RegistryClient { } DaemonInbound::CreateSession { reply_to, request } => { let result = match self.manager.create(request).await { - Ok(s) => RpcResult::Ok { value: s }, + Ok(s) => { + // The reply carries the session, so the registry applies + // it on receipt; swallow the async `SessionAdded` this + // create just queued on our event stream (NAP-268). + covered.added.insert(s.id); + RpcResult::Ok { value: s } + } Err(e) => RpcResult::Err { error: e.to_error_info(), }, @@ -292,7 +337,13 @@ impl RegistryClient { } DaemonInbound::KillSession { reply_to, id } => { let result = match self.manager.kill(id).await { - Ok(()) => RpcResult::Ok { value: () }, + Ok(()) => { + // The reply confirms the removal; swallow the async + // `SessionRemoved` it queued so the registry drops the + // session exactly once (NAP-268). + covered.removed.insert(id); + RpcResult::Ok { value: () } + } Err(e) => RpcResult::Err { error: e.to_error_info(), }, @@ -351,29 +402,48 @@ impl RegistryClient { /// `sink`. Returns `Ok(false)` once the stream has closed and the caller /// should stop the loop. /// +/// An `Added` / `Removed` whose id is parked in `covered` was driven by an RPC +/// whose reply already reports the change to the registry; it is consumed and +/// *not* forwarded, so the registry writes each change exactly once (NAP-268). +/// /// On `Lagged` the daemon's own view of its sessions skipped frames — which /// may include a `Removed`. Emitting nothing (the previous behaviour) left the /// registry advertising a now-dead session until the next reconnect, since the /// registry only reconciles a daemon's full session set on `apply_snapshot`. /// So we re-send a full `SessionsSnapshot`, which the registry applies by -/// replacing this daemon's set — healing the missed removal immediately. +/// replacing this daemon's set — healing the missed removal immediately. The +/// snapshot is the full truth, so any still-parked `covered` ids are stale and +/// cleared. async fn forward_session_event( manager: &SessionManager, ev: Result, sink: &mut S, + covered: &mut RpcCovered, ) -> Result where S: SinkExt + Unpin, { let frame = match ev { - Ok(SessionEvent::Added(session)) => DaemonOutbound::SessionAdded { session }, - Ok(SessionEvent::Removed(id)) => DaemonOutbound::SessionRemoved { id }, + Ok(SessionEvent::Added(session)) => { + if covered.added.remove(&session.id) { + return Ok(true); + } + DaemonOutbound::SessionAdded { session } + } + Ok(SessionEvent::Removed(id)) => { + if covered.removed.remove(&id) { + return Ok(true); + } + DaemonOutbound::SessionRemoved { id } + } Ok(SessionEvent::Updated(session)) => DaemonOutbound::SessionUpdated { session }, Err(RecvError::Lagged(skipped)) => { tracing::warn!( skipped, "session-event stream lagged — resyncing registry with a full snapshot" ); + covered.added.clear(); + covered.removed.clear(); DaemonOutbound::SessionsSnapshot { sessions: manager.snapshot().await, } @@ -440,6 +510,19 @@ mod tests { .collect() } + fn sample_session(id: Uuid) -> maestro_protocol::Session { + maestro_protocol::Session { + id, + daemon_id: "test-daemon".into(), + cwd: std::path::PathBuf::from("/tmp"), + command: "bash".into(), + env: Default::default(), + metadata: Default::default(), + status: maestro_protocol::SessionStatus::Running, + created_at: chrono::Utc::now(), + } + } + // Regression test for NAP-291 (defect B): a lagged event stream may have // dropped a `Removed`, so the daemon must resync the registry's full set // rather than silently skip — otherwise a dead session lingers in the @@ -448,9 +531,14 @@ mod tests { async fn lagged_stream_resyncs_with_full_snapshot() { let (mgr, _dir) = test_manager(); let mut sink = VecSink::default(); - let cont = forward_session_event(&mgr, Err(RecvError::Lagged(9)), &mut sink) - .await - .unwrap(); + let cont = forward_session_event( + &mgr, + Err(RecvError::Lagged(9)), + &mut sink, + &mut RpcCovered::default(), + ) + .await + .unwrap(); assert!(cont, "lag is recoverable; the loop must continue"); let frames = decode(&sink); assert!( @@ -464,9 +552,14 @@ mod tests { let (mgr, _dir) = test_manager(); let mut sink = VecSink::default(); let id = Uuid::new_v4(); - let cont = forward_session_event(&mgr, Ok(SessionEvent::Removed(id)), &mut sink) - .await - .unwrap(); + let cont = forward_session_event( + &mgr, + Ok(SessionEvent::Removed(id)), + &mut sink, + &mut RpcCovered::default(), + ) + .await + .unwrap(); assert!(cont); let frames = decode(&sink); assert!( @@ -479,10 +572,159 @@ mod tests { async fn closed_stream_signals_break_without_sending() { let (mgr, _dir) = test_manager(); let mut sink = VecSink::default(); - let cont = forward_session_event(&mgr, Err(RecvError::Closed), &mut sink) - .await - .unwrap(); + let cont = forward_session_event( + &mgr, + Err(RecvError::Closed), + &mut sink, + &mut RpcCovered::default(), + ) + .await + .unwrap(); assert!(!cont, "a closed stream must signal the loop to break"); assert!(sink.0.is_empty(), "nothing should be sent on close"); } + + // NAP-268: an RPC-driven create/kill reports the change to the registry via + // its reply, so the matching async lifecycle event must be swallowed — never + // forwarded — to keep the registry's catalog write single (not idempotent). + #[tokio::test] + async fn rpc_covered_added_is_swallowed_once() { + let (mgr, _dir) = test_manager(); + let mut sink = VecSink::default(); + let id = Uuid::new_v4(); + let mut covered = RpcCovered::default(); + covered.added.insert(id); + + let cont = forward_session_event( + &mgr, + Ok(SessionEvent::Added(sample_session(id))), + &mut sink, + &mut covered, + ) + .await + .unwrap(); + + assert!(cont, "swallowing an event must not stop the loop"); + assert!( + sink.0.is_empty(), + "an RPC-covered SessionAdded must not reach the registry; got {:?}", + decode(&sink), + ); + assert!( + !covered.added.contains(&id), + "the parked id must be consumed exactly once", + ); + } + + #[tokio::test] + async fn rpc_covered_removed_is_swallowed_once() { + let (mgr, _dir) = test_manager(); + let mut sink = VecSink::default(); + let id = Uuid::new_v4(); + let mut covered = RpcCovered::default(); + covered.removed.insert(id); + + let cont = + forward_session_event(&mgr, Ok(SessionEvent::Removed(id)), &mut sink, &mut covered) + .await + .unwrap(); + + assert!(cont); + assert!( + sink.0.is_empty(), + "an RPC-covered SessionRemoved must not reach the registry; got {:?}", + decode(&sink), + ); + assert!(!covered.removed.contains(&id)); + } + + // A spontaneous change (e.g. a session created via the daemon's local HTTP + // API, or reaped on its own) is the registry's only signal, so it must still + // forward even while *other* ids are parked. + #[tokio::test] + async fn uncovered_added_still_forwards() { + let (mgr, _dir) = test_manager(); + let mut sink = VecSink::default(); + let parked = Uuid::new_v4(); + let spontaneous = Uuid::new_v4(); + let mut covered = RpcCovered::default(); + covered.added.insert(parked); + + let cont = forward_session_event( + &mgr, + Ok(SessionEvent::Added(sample_session(spontaneous))), + &mut sink, + &mut covered, + ) + .await + .unwrap(); + + assert!(cont); + let frames = decode(&sink); + assert!( + matches!( + frames.as_slice(), + [DaemonOutbound::SessionAdded { session }] if session.id == spontaneous + ), + "an unparked create must forward; got {frames:?}", + ); + assert!( + covered.added.contains(&parked), + "an unrelated parked id must be left intact", + ); + } + + // A second event for the same id (e.g. the RPC was swallowed, then the + // session is reaped) is no longer covered and forwards normally. + #[tokio::test] + async fn covered_consumed_then_subsequent_event_forwards() { + let (mgr, _dir) = test_manager(); + let mut sink = VecSink::default(); + let id = Uuid::new_v4(); + let mut covered = RpcCovered::default(); + covered.added.insert(id); + + // First event: the RPC-covered create — swallowed. + forward_session_event( + &mgr, + Ok(SessionEvent::Added(sample_session(id))), + &mut sink, + &mut covered, + ) + .await + .unwrap(); + assert!(sink.0.is_empty()); + + // Second event: a spontaneous removal of that same session — forwarded, + // since only the create was parked. + forward_session_event(&mgr, Ok(SessionEvent::Removed(id)), &mut sink, &mut covered) + .await + .unwrap(); + let frames = decode(&sink); + assert!( + matches!(frames.as_slice(), [DaemonOutbound::SessionRemoved { id: got }] if *got == id), + "got {frames:?}", + ); + } + + // A lag resync is the registry's full truth, so any still-parked ids are + // stale and must be cleared (else a later unrelated event could be wrongly + // swallowed). + #[tokio::test] + async fn lag_resync_clears_parked_ids() { + let (mgr, _dir) = test_manager(); + let mut sink = VecSink::default(); + let mut covered = RpcCovered::default(); + covered.added.insert(Uuid::new_v4()); + covered.removed.insert(Uuid::new_v4()); + + forward_session_event(&mgr, Err(RecvError::Lagged(3)), &mut sink, &mut covered) + .await + .unwrap(); + + assert!( + covered.added.is_empty() && covered.removed.is_empty(), + "a full-snapshot resync must clear stale parked ids", + ); + } } diff --git a/crates/maestro-registry/src/state.rs b/crates/maestro-registry/src/state.rs index 3fc89fc..55983c4 100644 --- a/crates/maestro-registry/src/state.rs +++ b/crates/maestro-registry/src/state.rs @@ -599,12 +599,12 @@ impl Registry { return Err(RegistryError::DaemonGone); } 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 - // successful POST (NAP-246 made the race wider by adding extra - // outbound traffic on the daemon control link). `add_session` is - // idempotent so a later `SessionAdded` for the same id is a no-op. + // Apply the session to the catalog from the reply — the authoritative, + // single write for an RPC-initiated create. The daemon swallows the + // matching async `SessionAdded` precisely because we apply it here + // (NAP-268), so this is not racing a duplicate and `add_session` need + // not be idempotent for correctness. Without it the HTTP caller would + // see an empty `GET /api/sessions` immediately after a successful POST. if let Ok(ref session) = rpc { self.add_session(daemon_id, session.clone()).await; } @@ -644,9 +644,10 @@ impl Registry { return Err(RegistryError::DaemonGone); } 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. + // Authoritative single write for an RPC-initiated kill (see + // `create_session`): drop the session from the catalog so subsequent + // reads are consistent. The daemon swallows the matching async + // `SessionRemoved` (NAP-268), so this is not racing a duplicate. if rpc.is_ok() { self.remove_session(daemon_id, id).await; }