refactor(process-daemon): route stop and ownership through Control (#59700)

packages/process-daemon/src/daemon.rs

@@ -69,6 +69,8 @@ pub fn read_registration(directory: &Path) -> Result<Registration> {
     }
 }
 
+#[cfg(any(unix, windows))]
+mod control;
 #[cfg(unix)]
 #[path = "daemon/unix.rs"]
 mod platform;

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<Ownership>,
+}
+
+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<T>(&self, f: impl FnOnce(&mut Ownership) -> T) -> Result<T> {
+        let mut ownership = self.lock()?;
+        Ok(f(&mut ownership))
+    }
+
+    pub fn is_owner(&self, generation: u64) -> Result<bool> {
+        Ok(self.lock()?.is_owner(generation))
+    }
+
+    /// Main loop only. Returns whether the daemon should stop.
+    pub fn tick(&self, now: Instant) -> Result<bool> {
+        let mut ownership = self.lock()?;
+        Ok(self.stop.load(Ordering::Acquire) || ownership.tick(now))
+    }
+
+    fn lock(&self) -> Result<MutexGuard<'_, Ownership>> {
+        self.ownership
+            .lock()
+            .map_err(|_| anyhow!("ownership lock poisoned"))
+    }
+}

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<AtomicBool>,
+    thread: thread::JoinHandle<()>,
+}
+
+/// Marks a handler finished once its connection has been dropped, even on panic.
+struct Done(Arc<AtomicBool>);
+
+impl Drop for Done {
+    fn drop(&mut self) {
+        self.0.store(true, Ordering::Release);
+    }
+}
+
+fn serve(runtime: platform::Runtime, mut listener: Listener, control: Arc<Control>) -> Result<()> {
     let registration = runtime.registration.clone();
     let registration = &registration;
-    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::<Handler>::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,
-                            &registration,
-                            &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, &registration, &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<Ownership>,
+    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<Ownership>,
+    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<Control>,
+        exited: Receiver<(Result<()>, Instant)>,
+        thread: Option<thread::JoinHandle<()>>,
+    }
+
+    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<String>) -> 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));
+    }
+}

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::<u128>()),
-                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));
+    }
 }