diff --git a/README.md b/README.md index b882a69..5631e2e 100644 --- a/README.md +++ b/README.md @@ -23,12 +23,15 @@ Losing the owner connection stops the daemon and its terminals unless the owner first sends `{"token":"...","request":{"op":"prepare_handoff"}}` on that same socket. The response is `{"type":"handoff","ticket":"...","expires_at":123}`, where `expires_at` is Unix milliseconds, 120 seconds from preparation. Repeated -preparation during that window returns the same ticket and deadline. After the -old owner disconnects, a successor claims the same instance with the ticket in -its `own` request. A live owner cannot be displaced, and successful adoption -consumes the ticket. Expiry stops an unowned daemon; if the old owner is still -connected, expiry simply cancels the handoff. `shutdown` always stops the daemon, -including during handoff. No ownership or handoff state is persisted. +preparation during that window returns the same ticket and deadline. A successor +claims the same instance with the ticket in its `own` request, even while the old +owner is still connected. Successful adoption consumes the ticket atomically; +only one successor can claim it. The superseded connection can no longer prepare +handoffs or shut down the daemon, and its disconnect does not affect the new owner. +Expiry stops an unowned daemon; if the old owner is still connected, expiry simply +cancels the handoff. An ordinary authenticated `shutdown` request, or one from the +current owner, stops the daemon even during handoff. No ownership or handoff state +is persisted. ## Architecture diff --git a/src/daemon.rs b/src/daemon.rs index 219c6ac..2a45e15 100644 --- a/src/daemon.rs +++ b/src/daemon.rs @@ -233,19 +233,23 @@ mod unix { .map_err(|_| anyhow!("ownership lock poisoned"))? .claim(ticket.as_deref(), Instant::now()) }; - if let Err(error) = claim { - return write_frame( - &mut stream, - &Response::Error { - message: error.to_string(), - }, - ); - } - let result = owner_connection(&mut stream, registration, shutdown, ownership); + let generation = match claim { + Ok(generation) => generation, + Err(error) => { + return write_frame( + &mut stream, + &Response::Error { + message: error.to_string(), + }, + ); + } + }; + let result = + owner_connection(&mut stream, registration, shutdown, ownership, generation); ownership .lock() .map_err(|_| anyhow!("ownership lock poisoned"))? - .disconnect(Instant::now()); + .disconnect(generation, Instant::now()); return result; } if let Request::Subscribe { @@ -287,10 +291,23 @@ mod unix { registration: &Registration, shutdown: &AtomicBool, ownership: &Mutex, + generation: u64, ) -> Result<()> { write_frame(&mut *stream, &Response::Owned)?; while !shutdown.load(Ordering::Acquire) { let envelope: Envelope = read_frame(&mut *stream)?; + if !ownership + .lock() + .map_err(|_| anyhow!("ownership lock poisoned"))? + .is_owner(generation) + { + return write_frame( + &mut *stream, + &Response::Error { + message: "owner connection has been superseded".to_string(), + }, + ); + } let stopping = envelope.token == registration.token && matches!(envelope.request, Request::Shutdown); let response = if envelope.token != registration.token { @@ -304,6 +321,7 @@ mod unix { .lock() .map_err(|_| anyhow!("ownership lock poisoned"))? .prepare( + generation, Instant::now(), SystemTime::now().duration_since(UNIX_EPOCH)?.as_millis() as u64, )?; @@ -321,7 +339,13 @@ mod unix { }; let result = write_frame(&mut *stream, &response); if stopping { - shutdown.store(true, Ordering::Release); + // A takeover may have occurred while writing the response. + let owner = ownership + .lock() + .map_err(|_| anyhow!("ownership lock poisoned"))?; + if owner.is_owner(generation) { + shutdown.store(true, Ordering::Release); + } return result; } result?; diff --git a/src/ownership.rs b/src/ownership.rs index 72ac656..ef42a3c 100644 --- a/src/ownership.rs +++ b/src/ownership.rs @@ -12,7 +12,12 @@ pub(crate) struct Handoff { deadline: Instant, } -pub(crate) enum Ownership { +pub(crate) struct Ownership { + state: State, + generation: u64, +} + +enum State { Starting(Instant), Owned(Option), Waiting(Handoff), @@ -21,35 +26,49 @@ pub(crate) enum Ownership { impl Ownership { pub fn new(now: Instant) -> Self { - Self::Starting(now + ACQUIRE_TIMEOUT) + Self { + state: State::Starting(now + ACQUIRE_TIMEOUT), + generation: 0, + } } pub fn tick(&mut self, now: Instant) -> bool { - match self { - Self::Starting(deadline) if now >= *deadline => *self = Self::Stopped, - Self::Waiting(handoff) if now >= handoff.deadline => *self = Self::Stopped, - Self::Owned(Some(handoff)) if now >= handoff.deadline => *self = Self::Owned(None), + match &self.state { + State::Starting(deadline) if now >= *deadline => self.state = State::Stopped, + State::Waiting(handoff) if now >= handoff.deadline => self.state = State::Stopped, + State::Owned(Some(handoff)) if now >= handoff.deadline => { + self.state = State::Owned(None) + } _ => {} } - matches!(self, Self::Stopped) + matches!(self.state, State::Stopped) } - pub fn claim(&mut self, ticket: Option<&str>, now: Instant) -> Result<()> { + pub fn claim(&mut self, ticket: Option<&str>, now: Instant) -> Result { self.tick(now); - match &*self { - Self::Starting(_) if ticket.is_none() => {} - Self::Waiting(handoff) if ticket == Some(handoff.ticket.as_str()) => {} - Self::Owned(_) => bail!("daemon already has a live owner"), - Self::Stopped => bail!("daemon ownership deadline expired or daemon is stopping"), + match &self.state { + State::Starting(_) if ticket.is_none() => {} + State::Owned(Some(handoff)) | State::Waiting(handoff) + if ticket == Some(handoff.ticket.as_str()) => {} + State::Owned(_) => bail!("daemon already has a live owner"), + State::Stopped => bail!("daemon ownership deadline expired or daemon is stopping"), _ => bail!("invalid handoff ticket"), } - *self = Self::Owned(None); - Ok(()) + self.generation += 1; + self.state = State::Owned(None); + Ok(self.generation) } - pub fn prepare(&mut self, now: Instant, unix_ms: u64) -> Result { + pub fn is_owner(&self, generation: u64) -> bool { + self.generation == generation && matches!(self.state, State::Owned(_)) + } + + pub fn prepare(&mut self, generation: u64, now: Instant, unix_ms: u64) -> Result { self.tick(now); - let Self::Owned(handoff) = self else { + if self.generation != generation { + bail!("handoff requires the owner connection"); + } + let State::Owned(handoff) = &mut self.state else { bail!("handoff requires the owner connection"); }; Ok(handoff @@ -61,11 +80,15 @@ impl Ownership { .clone()) } - pub fn disconnect(&mut self, now: Instant) { + pub fn disconnect(&mut self, generation: u64, now: Instant) { + // A ticket may transfer ownership before the previous connection closes. + if !self.is_owner(generation) { + return; + } self.tick(now); - *self = match self { - Self::Owned(Some(handoff)) => Self::Waiting(handoff.clone()), - _ => Self::Stopped, + self.state = match &self.state { + State::Owned(Some(handoff)) => State::Waiting(handoff.clone()), + _ => State::Stopped, }; } } @@ -79,9 +102,9 @@ mod tests { let now = Instant::now(); let mut owner = Ownership::new(now); assert!(owner.claim(Some("unexpected"), now).is_err()); - owner.claim(None, now).unwrap(); + let generation = owner.claim(None, now).unwrap(); assert!(owner.claim(None, now).is_err()); - owner.disconnect(now); + owner.disconnect(generation, now); assert!(owner.tick(now)); assert!(owner.claim(None, now).is_err()); @@ -95,22 +118,51 @@ mod tests { fn handoff_is_nonrenewing_and_consumed_by_claim() { let now = Instant::now(); let mut owner = Ownership::new(now); - owner.claim(None, now).unwrap(); - let handoff = owner.prepare(now, 1000).unwrap(); + let generation = owner.claim(None, now).unwrap(); + let handoff = owner.prepare(generation, now, 1000).unwrap(); assert_eq!(handoff.expires_at, 121_000); assert_eq!( owner - .prepare(now + Duration::from_secs(60), 61_000) + .prepare(generation, now + Duration::from_secs(60), 61_000) .unwrap(), handoff ); + owner.disconnect(generation, now); + assert!(owner.claim(None, now).is_err()); + assert!(owner.claim(Some("wrong"), now).is_err()); + let successor = owner.claim(Some(&handoff.ticket), now).unwrap(); assert!(owner.claim(Some(&handoff.ticket), now).is_err()); - owner.disconnect(now); + owner.disconnect(successor, now); + assert!(owner.tick(now)); + } + + #[test] + fn valid_ticket_replaces_live_owner_and_fences_old_connection() { + let now = Instant::now(); + let mut owner = Ownership::new(now); + let generation = owner.claim(None, now).unwrap(); + let handoff = owner.prepare(generation, now, 1000).unwrap(); assert!(owner.claim(None, now).is_err()); assert!(owner.claim(Some("wrong"), now).is_err()); - owner.claim(Some(&handoff.ticket), now).unwrap(); + assert!(owner.is_owner(generation)); + + let successor = owner.claim(Some(&handoff.ticket), now).unwrap(); + assert!(!owner.is_owner(generation)); + assert!(owner.is_owner(successor)); assert!(owner.claim(Some(&handoff.ticket), now).is_err()); - owner.disconnect(now); + assert!(owner.prepare(generation, now, 1000).is_err()); + owner.disconnect(generation, now); + assert!(owner.is_owner(successor)); + assert!(!owner.tick(now)); + + let next = owner.prepare(successor, now, 1000).unwrap(); + owner.disconnect(generation, now); + assert!(owner.is_owner(successor)); + owner.disconnect(successor, now); + // Stale cleanup must not discard the successor's pending handoff either. + owner.disconnect(generation, now); + let last = owner.claim(Some(&next.ticket), now).unwrap(); + owner.disconnect(last, now); assert!(owner.tick(now)); } @@ -118,17 +170,23 @@ mod tests { fn expiry_only_stops_a_disconnected_owner() { let now = Instant::now(); let mut owner = Ownership::new(now); - owner.claim(None, now).unwrap(); - let handoff = owner.prepare(now, 1000).unwrap(); + let generation = owner.claim(None, now).unwrap(); + let handoff = owner.prepare(generation, now, 1000).unwrap(); assert!(!owner.tick(handoff.deadline)); - assert!(matches!(owner, Ownership::Owned(None))); - owner.disconnect(handoff.deadline); + assert!(matches!(owner.state, State::Owned(None))); + assert!( + owner + .claim(Some(&handoff.ticket), handoff.deadline) + .is_err() + ); + assert!(owner.is_owner(generation)); + owner.disconnect(generation, handoff.deadline); assert!(owner.tick(handoff.deadline)); let mut owner = Ownership::new(now); - owner.claim(None, now).unwrap(); - let handoff = owner.prepare(now, 1000).unwrap(); - owner.disconnect(now); + let generation = owner.claim(None, now).unwrap(); + let handoff = owner.prepare(generation, now, 1000).unwrap(); + owner.disconnect(generation, now); assert!(!owner.tick(handoff.deadline - Duration::from_millis(1))); assert!( owner diff --git a/tests/ownership.rs b/tests/ownership.rs index e574202..34e1204 100644 --- a/tests/ownership.rs +++ b/tests/ownership.rs @@ -1,9 +1,10 @@ #![cfg(unix)] -use std::io::Write; +use std::io::{Read, Write}; use std::os::unix::net::UnixStream; use std::path::PathBuf; use std::process::{Child, Command, Stdio}; +use std::sync::Barrier; use std::thread; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; @@ -222,23 +223,16 @@ fn handoff_preserves_daemon_and_terminal_and_shutdown_overrides_it() { } if repeated == ticket && deadline == expires_at) ); assert!(matches!( - daemon.own(Some(ticket.clone())).1, + daemon.own(Some("wrong-ticket".into())).1, Response::Error { .. } )); + // The old owner stays connected until the successor has acquired the daemon. + let (mut successor, response) = daemon.own(Some(ticket.clone())); + assert!(matches!(response, Response::Owned)); + assert!(matches!(daemon.own(Some(ticket)).1, Response::Error { .. })); drop(owner); let deadline = Instant::now() + Duration::from_secs(3); - let mut successor = loop { - let (stream, response) = daemon.own(Some(ticket.clone())); - if matches!(response, Response::Owned) { - break stream; - } - assert!( - Instant::now() < deadline, - "successor could not acquire: {response:?}" - ); - thread::sleep(Duration::from_millis(10)); - }; assert!( matches!(daemon.request(Request::Ping), Response::Pong { instance_id, pid, .. } if instance_id == daemon.registration.instance_id && pid == daemon.child.id()) @@ -276,6 +270,78 @@ fn handoff_preserves_daemon_and_terminal_and_shutdown_overrides_it() { assert_terminal_stopped(&terminal); } +#[test] +fn superseded_owner_cannot_prepare_handoff_or_shutdown() { + for request in [Request::PrepareHandoff, Request::Shutdown] { + let mut daemon = Daemon::start(); + let (mut owner, response) = daemon.own(None); + assert!(matches!(response, Response::Owned)); + let Response::Handoff { ticket, .. } = daemon.send(&mut owner, Request::PrepareHandoff) + else { + panic!("handoff response"); + }; + let (mut successor, response) = daemon.own(Some(ticket)); + assert!(matches!(response, Response::Owned)); + assert!(matches!( + daemon.send(&mut owner, request), + Response::Error { message } if message == "owner connection has been superseded" + )); + // EOF confirms the stale handler has run its disconnect cleanup. + assert_eq!(owner.read(&mut [0]).unwrap(), 0); + assert!(matches!( + daemon.request(Request::Ping), + Response::Pong { .. } + )); + assert!(matches!( + daemon.send(&mut successor, Request::PrepareHandoff), + Response::Handoff { .. } + )); + assert!(matches!( + daemon.send(&mut successor, Request::Shutdown), + Response::Ok + )); + daemon.wait(); + } +} + +#[test] +fn handoff_ticket_has_only_one_winner() { + let mut daemon = Daemon::start(); + let (mut owner, response) = daemon.own(None); + assert!(matches!(response, Response::Owned)); + let Response::Handoff { ticket, .. } = daemon.send(&mut owner, Request::PrepareHandoff) else { + panic!("handoff response"); + }; + let barrier = Barrier::new(2); + let results = thread::scope(|scope| { + let claim = || { + barrier.wait(); + daemon.own(Some(ticket.clone())) + }; + let first = scope.spawn(claim); + let second = scope.spawn(claim); + [first.join().unwrap(), second.join().unwrap()] + }); + assert_eq!( + results + .iter() + .filter(|(_, response)| matches!(response, Response::Owned)) + .count(), + 1 + ); + assert_eq!( + results + .iter() + .filter(|(_, response)| matches!(response, Response::Error { .. })) + .count(), + 1 + ); + // Losing the new owner must stop the daemon even if the old socket is still open. + drop(results); + daemon.wait(); + assert_eq!(owner.read(&mut [0]).unwrap(), 0); +} + #[test] fn unclaimed_daemon_times_out() { let mut daemon = Daemon::start();