diff --git a/packages/process-daemon/src/daemon.rs b/packages/process-daemon/src/daemon.rs index c3d38a30bf42..af03b3630110 100644 --- a/packages/process-daemon/src/daemon.rs +++ b/packages/process-daemon/src/daemon.rs @@ -69,6 +69,8 @@ pub fn read_registration(directory: &Path) -> Result { } } +#[cfg(any(unix, windows))] +mod control; #[cfg(unix)] #[path = "daemon/unix.rs"] mod platform; diff --git a/packages/process-daemon/src/daemon/control.rs b/packages/process-daemon/src/daemon/control.rs new file mode 100644 index 000000000000..d5b2c7d52ebc --- /dev/null +++ b/packages/process-daemon/src/daemon/control.rs @@ -0,0 +1,62 @@ +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Mutex, MutexGuard, PoisonError}; +use std::time::Instant; + +use anyhow::{Result, anyhow}; + +use crate::ownership::Ownership; + +/// Daemon stop flag and ownership state, shared by the main loop and handlers. +/// +/// Every write to either goes through a method here; the fields stay private to +/// this module so no caller can change them behind the main loop's back. Both +/// writes happen under the ownership lock. +pub(super) struct Control { + stop: AtomicBool, + ownership: Mutex, +} + +impl Control { + pub fn new(ownership: Ownership) -> Self { + Self { + stop: AtomicBool::new(false), + ownership: Mutex::new(ownership), + } + } + + /// Never fails: a poisoned lock is recovered, because stopping must always work. + /// Must not be called from an `update_ownership` closure. + pub fn request_stop(&self) { + let _ownership = self + .ownership + .lock() + .unwrap_or_else(PoisonError::into_inner); + self.stop.store(true, Ordering::Release); + } + + pub fn stopping(&self) -> bool { + self.stop.load(Ordering::Acquire) + } + + /// `f` must not call back into `Control`; the lock is not reentrant. + pub fn update_ownership(&self, f: impl FnOnce(&mut Ownership) -> T) -> Result { + let mut ownership = self.lock()?; + Ok(f(&mut ownership)) + } + + pub fn is_owner(&self, generation: u64) -> Result { + Ok(self.lock()?.is_owner(generation)) + } + + /// Main loop only. Returns whether the daemon should stop. + pub fn tick(&self, now: Instant) -> Result { + let mut ownership = self.lock()?; + Ok(self.stop.load(Ordering::Acquire) || ownership.tick(now)) + } + + fn lock(&self) -> Result> { + self.ownership + .lock() + .map_err(|_| anyhow!("ownership lock poisoned")) + } +} diff --git a/packages/process-daemon/src/daemon/server.rs b/packages/process-daemon/src/daemon/server.rs index 48a20a84ad89..ccd1ed3a85d3 100644 --- a/packages/process-daemon/src/daemon/server.rs +++ b/packages/process-daemon/src/daemon/server.rs @@ -1,59 +1,80 @@ use std::path::Path; +use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::{Arc, Mutex}; use std::thread; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; use anyhow::{Context, Result, anyhow}; use base64::Engine; +use super::control::Control; use super::{Registration, platform}; use crate::ownership::Ownership; use crate::protocol::{Envelope, Request, Response, read_frame, write_frame, write_output_frame}; use crate::service::{CreateTerminal, StreamEvent, TerminalService}; -use crate::transport::{Cancellation, Connection}; +use crate::transport::{Cancellation, Connection, Listener}; pub fn run(directory: &Path) -> Result<()> { - let (runtime, mut listener) = platform::Runtime::bind(directory)?; + let (runtime, listener) = platform::Runtime::bind(directory)?; + // Created after bind so bind I/O does not consume the acquire window. + let control = Arc::new(Control::new(Ownership::new(Instant::now()))); + serve(runtime, listener, control) +} + +struct Handler { + cancellation: Cancellation, + done: Arc, + thread: thread::JoinHandle<()>, +} + +/// Marks a handler finished once its connection has been dropped, even on panic. +struct Done(Arc); + +impl Drop for Done { + fn drop(&mut self) { + self.0.store(true, Ordering::Release); + } +} + +fn serve(runtime: platform::Runtime, mut listener: Listener, control: Arc) -> Result<()> { let registration = runtime.registration.clone(); let registration = ®istration; - let ownership = Arc::new(Mutex::new(Ownership::new(Instant::now()))); let service = Arc::new(TerminalService::default()); - let shutdown = Arc::new(AtomicBool::new(false)); - let mut handlers = Vec::<(Cancellation, thread::JoinHandle<()>)>::new(); + let mut handlers = Vec::::new(); let result = (|| -> Result<()> { - while !shutdown.load(Ordering::Acquire) { - if ownership - .lock() - .map_err(|_| anyhow!("ownership lock poisoned"))? - .tick(Instant::now()) - { - shutdown.store(true, Ordering::Release); + loop { + if control.tick(Instant::now())? { break; } - for (_, handler) in handlers.extract_if(.., |(_, handler)| handler.is_finished()) { - let _ = handler.join(); + for handler in handlers.extract_if(.., |handler| { + handler.done.load(Ordering::Acquire) || handler.thread.is_finished() + }) { + let _ = handler.thread.join(); } match listener.accept() { Ok(Some(stream)) => { - let control = stream.cancellation()?; - let service = Arc::clone(&service); - let shutdown = Arc::clone(&shutdown); - let ownership = Arc::clone(&ownership); - let registration = registration.clone(); - let handle = thread::spawn(move || { - if let Err(error) = handle_connection( - stream, - &service, - ®istration, - &shutdown, - &ownership, - ) { - eprintln!("opencode-pty request failed: {error:#}"); - } + let cancellation = stream.cancellation()?; + let done = Arc::new(AtomicBool::new(false)); + let thread = { + let done = Arc::clone(&done); + let service = Arc::clone(&service); + let control = Arc::clone(&control); + let registration = registration.clone(); + thread::spawn(move || { + let _done = Done(done); + if let Err(error) = + handle_connection(stream, &service, ®istration, &control) + { + eprintln!("opencode-pty request failed: {error:#}"); + } + }) + }; + handlers.push(Handler { + cancellation, + done, + thread, }); - handlers.push((control, handle)); } Ok(None) => { thread::sleep(Duration::from_millis(10)); @@ -64,12 +85,12 @@ pub fn run(directory: &Path) -> Result<()> { Ok(()) })(); - shutdown.store(true, Ordering::Release); + control.request_stop(); listener.stop(); // Unblock partial requests, owner reads, and backpressured subscriptions // before joining. PTY workers still use their existing termination path. - for (cancellation, _) in &handlers { - cancellation.cancel(); + for handler in &handlers { + handler.cancellation.cancel(); } let (cleanup_tx, cleanup_rx) = crossbeam_channel::bounded::<()>(1); let cleanup = runtime.cleanup.clone(); @@ -81,8 +102,8 @@ pub fn run(directory: &Path) -> Result<()> { } }); service.shutdown(); - for (_, handler) in handlers { - let _ = handler.join(); + for handler in handlers { + let _ = handler.thread.join(); } drop(service); let cleanup = runtime.finish(); @@ -96,8 +117,7 @@ fn handle_connection( mut stream: Connection, service: &TerminalService, registration: &Registration, - shutdown: &AtomicBool, - ownership: &Mutex, + control: &Control, ) -> Result<()> { let envelope: Envelope = read_frame(&mut stream)?; if envelope.token != registration.token { @@ -115,13 +135,11 @@ fn handle_connection( { let claim = if instance_id != registration.instance_id { Err(anyhow!("daemon instance_id mismatch")) - } else if shutdown.load(Ordering::Acquire) { + } else if control.stopping() { Err(anyhow!("daemon is stopping")) } else { - ownership - .lock() - .map_err(|_| anyhow!("ownership lock poisoned"))? - .claim(ticket.as_deref(), Instant::now()) + control + .update_ownership(|ownership| ownership.claim(ticket.as_deref(), Instant::now()))? }; let generation = match claim { Ok(generation) => generation, @@ -134,11 +152,8 @@ fn handle_connection( ); } }; - let result = owner_connection(&mut stream, registration, shutdown, ownership, generation); - ownership - .lock() - .map_err(|_| anyhow!("ownership lock poisoned"))? - .disconnect(generation, Instant::now()); + let result = owner_connection(&mut stream, registration, control, generation); + control.update_ownership(|ownership| ownership.disconnect(generation, Instant::now()))?; return result; } if let Request::Subscribe { @@ -152,7 +167,7 @@ fn handle_connection( return stream_subscription( &mut stream, service, - shutdown, + control, SubscriptionRequest { id, offset, @@ -169,7 +184,7 @@ fn handle_connection( }); let result = send_response(&mut stream, &response); if stopping { - shutdown.store(true, Ordering::Release); + control.request_stop(); } result } @@ -177,18 +192,13 @@ fn handle_connection( fn owner_connection( stream: &mut Connection, registration: &Registration, - shutdown: &AtomicBool, - ownership: &Mutex, + control: &Control, generation: u64, ) -> Result<()> { write_frame(&mut *stream, &Response::Owned)?; - while !shutdown.load(Ordering::Acquire) { + while !control.stopping() { let envelope: Envelope = read_frame(&mut *stream)?; - if !ownership - .lock() - .map_err(|_| anyhow!("ownership lock poisoned"))? - .is_owner(generation) - { + if !control.is_owner(generation)? { return send_response( stream, &Response::Error { @@ -205,14 +215,10 @@ fn owner_connection( } else { match envelope.request { Request::PrepareHandoff => { - let handoff = ownership - .lock() - .map_err(|_| anyhow!("ownership lock poisoned"))? - .prepare( - generation, - Instant::now(), - SystemTime::now().duration_since(UNIX_EPOCH)?.as_millis() as u64, - )?; + let unix_ms = SystemTime::now().duration_since(UNIX_EPOCH)?.as_millis() as u64; + let handoff = control.update_ownership(|ownership| { + ownership.prepare(generation, Instant::now(), unix_ms) + })??; Response::Handoff { ticket: handoff.ticket, expires_at: handoff.expires_at, @@ -228,12 +234,8 @@ fn owner_connection( if stopping { let result = send_response(stream, &response); // A takeover may have occurred while writing the response. - if ownership - .lock() - .map_err(|_| anyhow!("ownership lock poisoned"))? - .is_owner(generation) - { - shutdown.store(true, Ordering::Release); + if control.is_owner(generation)? { + control.request_stop(); } return result; } @@ -253,7 +255,7 @@ struct SubscriptionRequest { fn stream_subscription( stream: &mut Connection, service: &TerminalService, - shutdown: &AtomicBool, + control: &Control, request: SubscriptionRequest, ) -> Result<()> { let attachment = service.attach( @@ -279,7 +281,7 @@ fn stream_subscription( )?; let monitor = stream.monitor_disconnect()?; let result = (|| loop { - if shutdown.load(Ordering::Acquire) { + if control.stopping() { break Ok(()); } let event = crossbeam_channel::select! { @@ -458,3 +460,190 @@ fn dispatch( }, }) } + +#[cfg(all(test, unix))] +mod tests { + use std::path::PathBuf; + use std::sync::atomic::AtomicUsize; + + use crossbeam_channel::Receiver; + + use super::*; + + const TEST_HANDOFF: Duration = Duration::from_secs(60); + + struct TestDaemon { + root: PathBuf, + registration: Registration, + start: Instant, + control: Arc, + exited: Receiver<(Result<()>, Instant)>, + thread: Option>, + } + + impl TestDaemon { + fn start(acquire: Duration, handoff: Duration) -> Self { + static NEXT: AtomicUsize = AtomicUsize::new(0); + let root = std::env::temp_dir().join(format!( + "opencode-pty-server-{}-{}", + std::process::id(), + NEXT.fetch_add(1, Ordering::Relaxed) + )); + // The sweeper scans the runtime directory's siblings; keep it alone in `root`. + let (runtime, listener) = platform::Runtime::bind(&root.join("runtime")).unwrap(); + let registration = runtime.registration.clone(); + let start = Instant::now(); + let control = Arc::new(Control::new(Ownership::with_timeouts( + start, acquire, handoff, + ))); + let (sender, exited) = crossbeam_channel::bounded(1); + let thread = { + let control = Arc::clone(&control); + thread::spawn(move || { + let result = serve(runtime, listener, control); + let _ = sender.send((result, Instant::now())); + }) + }; + Self { + root, + registration, + start, + control, + exited, + thread: Some(thread), + } + } + + fn connect(&self) -> Connection { + let stream = Connection::connect(&self.registration.socket).unwrap(); + stream + .set_read_timeout(Some(Duration::from_secs(3))) + .unwrap(); + stream + .set_write_timeout(Some(Duration::from_secs(3))) + .unwrap(); + stream + } + + fn send(&self, stream: &mut Connection, request: Request) -> Response { + write_frame( + &mut *stream, + &Envelope { + token: self.registration.token.clone(), + request, + }, + ) + .unwrap(); + read_frame(stream).unwrap() + } + + fn own(&self, ticket: Option) -> Connection { + let mut stream = self.connect(); + let request = Request::Own { + instance_id: self.registration.instance_id.clone(), + ticket, + }; + match self.send(&mut stream, request) { + Response::Owned => stream, + response => panic!("own: {response:?}"), + } + } + + fn claim(&self) -> Connection { + self.own(None) + } + + fn prepare(&self, owner: &mut Connection) -> String { + match self.send(owner, Request::PrepareHandoff) { + Response::Handoff { ticket, .. } => ticket, + response => panic!("prepare: {response:?}"), + } + } + + fn exit_within(&self, timeout: Duration) -> Option<(Result<()>, Instant)> { + self.exited.recv_timeout(timeout).ok() + } + + fn stopped_within(&self, timeout: Duration) -> Instant { + let (result, returned_at) = self + .exit_within(timeout) + .expect("daemon did not stop in time"); + result.unwrap(); + returned_at + } + } + + impl Drop for TestDaemon { + fn drop(&mut self) { + self.control.request_stop(); + if let Some(thread) = self.thread.take() { + let _ = thread.join(); + } + let _ = std::fs::remove_dir_all(&self.root); + } + } + + #[test] + fn shutdown_request_stops_immediately() { + let daemon = TestDaemon::start(Duration::from_secs(60), TEST_HANDOFF); + let _owner = daemon.claim(); + assert!(matches!( + daemon.send(&mut daemon.connect(), Request::Shutdown), + Response::Ok + )); + daemon.stopped_within(Duration::from_secs(1)); + } + + #[test] + fn owner_shutdown_stops_immediately() { + let daemon = TestDaemon::start(Duration::from_secs(60), TEST_HANDOFF); + let mut owner = daemon.claim(); + assert!(matches!( + daemon.send(&mut owner, Request::Shutdown), + Response::Ok + )); + daemon.stopped_within(Duration::from_secs(1)); + } + + #[test] + fn owner_disconnect_without_handoff_stops_immediately() { + let daemon = TestDaemon::start(Duration::from_secs(60), TEST_HANDOFF); + drop(daemon.claim()); + daemon.stopped_within(Duration::from_secs(1)); + } + + #[test] + fn starting_deadline_fires_on_time() { + let daemon = TestDaemon::start(Duration::from_millis(300), TEST_HANDOFF); + let elapsed = daemon.stopped_within(Duration::from_secs(3)) - daemon.start; + assert!( + (Duration::from_millis(300)..=Duration::from_millis(1500)).contains(&elapsed), + "{elapsed:?}" + ); + } + + #[test] + fn waiting_deadline_fires_on_time() { + let daemon = TestDaemon::start(Duration::from_secs(60), Duration::from_millis(500)); + let mut owner = daemon.claim(); + let before = Instant::now(); + daemon.prepare(&mut owner); + let after = Instant::now(); + drop(owner); + let returned_at = daemon.stopped_within(Duration::from_secs(3)); + assert!(returned_at >= before + Duration::from_millis(500)); + assert!(returned_at <= after + Duration::from_millis(1500)); + } + + #[test] + fn handoff_claim_transfers_without_spurious_stop() { + let daemon = TestDaemon::start(Duration::from_secs(60), TEST_HANDOFF); + let mut owner = daemon.claim(); + let ticket = daemon.prepare(&mut owner); + let successor = daemon.own(Some(ticket)); + drop(owner); + assert!(daemon.exit_within(Duration::from_millis(300)).is_none()); + drop(successor); + daemon.stopped_within(Duration::from_secs(1)); + } +} diff --git a/packages/process-daemon/src/ownership.rs b/packages/process-daemon/src/ownership.rs index ef42a3c6a4f4..eb127dd4e51f 100644 --- a/packages/process-daemon/src/ownership.rs +++ b/packages/process-daemon/src/ownership.rs @@ -15,6 +15,7 @@ pub(crate) struct Handoff { pub(crate) struct Ownership { state: State, generation: u64, + handoff_timeout: Duration, } enum State { @@ -26,9 +27,14 @@ enum State { impl Ownership { pub fn new(now: Instant) -> Self { + Self::with_timeouts(now, ACQUIRE_TIMEOUT, HANDOFF_TIMEOUT) + } + + pub fn with_timeouts(now: Instant, acquire: Duration, handoff: Duration) -> Self { Self { - state: State::Starting(now + ACQUIRE_TIMEOUT), + state: State::Starting(now + acquire), generation: 0, + handoff_timeout: handoff, } } @@ -74,8 +80,8 @@ impl Ownership { Ok(handoff .get_or_insert_with(|| Handoff { ticket: format!("{:032x}", rand::random::()), - expires_at: unix_ms + HANDOFF_TIMEOUT.as_millis() as u64, - deadline: now + HANDOFF_TIMEOUT, + expires_at: unix_ms + self.handoff_timeout.as_millis() as u64, + deadline: now + self.handoff_timeout, }) .clone()) } @@ -195,4 +201,23 @@ mod tests { ); assert!(owner.tick(handoff.deadline)); } + + #[test] + fn custom_timeouts_bound_acquire_and_handoff() { + let now = Instant::now(); + let acquire = Duration::from_millis(300); + let handoff = Duration::from_millis(500); + let mut starting = Ownership::with_timeouts(now, acquire, handoff); + assert!(!starting.tick(now + acquire - Duration::from_millis(1))); + assert!(starting.tick(now + acquire)); + + let mut owner = Ownership::with_timeouts(now, acquire, handoff); + let generation = owner.claim(None, now).unwrap(); + let ticket = owner.prepare(generation, now, 1000).unwrap(); + assert_eq!(ticket.expires_at, 1500); + assert_eq!(ticket.deadline, now + handoff); + owner.disconnect(generation, now); + assert!(!owner.tick(now + handoff - Duration::from_millis(1))); + assert!(owner.tick(now + handoff)); + } }