hyperdb_mcp/daemon/
health.rs

1// Copyright (c) 2026, Salesforce, Inc. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0 OR MIT
3
4//! TCP health listener for the daemon.
5//!
6//! The health listener serves two purposes:
7//! 1. **Single-instance lock** — binding the port guarantees at most one daemon per user.
8//! 2. **Liveness probe + heartbeat** — clients connect and send simple text commands.
9//!
10//! Protocol (line-based, newline-terminated):
11//! - `PING\n` → `PONG\n` (liveness check)
12//! - `HEARTBEAT\n` → `OK\n` (resets idle timer)
13//! - `STOP\n` → `STOPPING\n` (triggers graceful shutdown)
14//! - `STATUS\n` → JSON line with daemon info (reports the *current* hyperd
15//!   endpoint, which can change after a restart).
16//! - `REPORT_HYPERD_ERROR\n` → `OK\n` (sets the restart-requested flag —
17//!   the monitor task picks it up on its next tick).
18
19use std::io::{BufRead, BufReader, Write};
20use std::net::{TcpListener, TcpStream};
21use std::sync::atomic::{AtomicBool, Ordering};
22use std::sync::{Arc, Mutex};
23use std::time::{Duration, Instant};
24
25use tracing::{debug, warn};
26
27use super::discovery::DaemonInfo;
28
29/// Handle to the health listener, used to check binding success and manage lifecycle.
30#[derive(Debug)]
31pub struct HealthListener {
32    listener: TcpListener,
33    pub port: u16,
34}
35
36/// Shared state between the health listener and the daemon main loop.
37#[derive(Debug)]
38pub struct DaemonState {
39    /// Last time any client sent a heartbeat or query.
40    pub last_activity: Mutex<Instant>,
41    /// Signal to shut down the daemon.
42    pub shutdown: AtomicBool,
43    /// Set by clients reporting that hyperd looks dead from over there;
44    /// consumed by the daemon's restart monitor.
45    pub restart_requested: AtomicBool,
46}
47
48impl Default for DaemonState {
49    fn default() -> Self {
50        Self::new()
51    }
52}
53
54impl DaemonState {
55    pub fn new() -> Self {
56        Self {
57            last_activity: Mutex::new(Instant::now()),
58            shutdown: AtomicBool::new(false),
59            restart_requested: AtomicBool::new(false),
60        }
61    }
62
63    /// Record activity (resets idle timer).
64    ///
65    /// # Panics
66    /// Panics if the internal mutex is poisoned.
67    pub fn touch(&self) {
68        *self.last_activity.lock().expect("mutex poisoned") = Instant::now();
69    }
70
71    /// Duration since the last activity.
72    ///
73    /// # Panics
74    /// Panics if the internal mutex is poisoned.
75    pub fn idle_duration(&self) -> Duration {
76        self.last_activity.lock().expect("mutex poisoned").elapsed()
77    }
78
79    pub fn request_shutdown(&self) {
80        self.shutdown.store(true, Ordering::Release);
81    }
82
83    pub fn should_shutdown(&self) -> bool {
84        self.shutdown.load(Ordering::Acquire)
85    }
86
87    /// Signal that hyperd appears to have died and a restart is needed.
88    pub fn request_restart(&self) {
89        self.restart_requested.store(true, Ordering::Release);
90    }
91
92    /// Atomically read-and-clear the restart-request flag.
93    /// Returns true if a restart was requested since the last call.
94    pub fn consume_restart_request(&self) -> bool {
95        self.restart_requested.swap(false, Ordering::AcqRel)
96    }
97}
98
99impl HealthListener {
100    /// Try to bind the health port.
101    ///
102    /// # Errors
103    /// Returns `Err` if the port is already in use (another daemon is running)
104    /// or the bind fails for another reason.
105    pub fn bind(port: u16) -> std::io::Result<Self> {
106        let addr = std::net::SocketAddr::from(([127, 0, 0, 1], port));
107        let listener = TcpListener::bind(addr)?;
108        listener.set_nonblocking(true)?;
109        let port = listener.local_addr()?.port();
110        Ok(Self { listener, port })
111    }
112
113    /// Run the health listener loop. Spawns per-connection threads until shutdown.
114    /// Consumes `self` because this is intended to be called from a dedicated thread.
115    ///
116    /// `info` is shared (`Arc<Mutex<DaemonInfo>>`) so the listener reports the
117    /// *current* hyperd endpoint after a restart — the monitor task updates the
118    /// same Arc once a new hyperd is running.
119    #[expect(
120        clippy::needless_pass_by_value,
121        reason = "Arcs are cloned into per-connection threads"
122    )]
123    pub fn run(self, state: Arc<DaemonState>, info: Arc<Mutex<DaemonInfo>>) {
124        loop {
125            if state.should_shutdown() {
126                break;
127            }
128
129            match self.listener.accept() {
130                Ok((stream, _addr)) => {
131                    let state = Arc::clone(&state);
132                    let info = Arc::clone(&info);
133                    std::thread::spawn(move || {
134                        handle_client(stream, &state, &info);
135                    });
136                }
137                Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => {
138                    std::thread::sleep(Duration::from_millis(100));
139                }
140                Err(e) => {
141                    warn!(error = %e, "health listener accept error");
142                    std::thread::sleep(Duration::from_millis(500));
143                }
144            }
145        }
146        debug!("health listener shut down");
147    }
148}
149
150#[expect(
151    clippy::needless_pass_by_value,
152    reason = "TcpStream must be owned for BufReader"
153)]
154fn handle_client(stream: TcpStream, state: &DaemonState, info: &Mutex<DaemonInfo>) {
155    let _ = stream.set_read_timeout(Some(Duration::from_secs(5)));
156    let mut reader = BufReader::new(&stream);
157    let mut writer = &stream;
158    let mut line = String::new();
159
160    loop {
161        line.clear();
162        match reader.read_line(&mut line) {
163            Ok(0) => break,
164            Ok(_) => {
165                let cmd = line.trim();
166                let response = match cmd {
167                    "PING" => "PONG\n".to_string(),
168                    "HEARTBEAT" => {
169                        state.touch();
170                        "OK\n".to_string()
171                    }
172                    "STOP" => {
173                        state.request_shutdown();
174                        "STOPPING\n".to_string()
175                    }
176                    "STATUS" => {
177                        // Brief lock — only to clone the current snapshot.
178                        let snapshot = info.lock().expect("DaemonInfo mutex poisoned").clone();
179                        let json = serde_json::to_string(&snapshot).unwrap_or_default();
180                        format!("{json}\n")
181                    }
182                    "REPORT_HYPERD_ERROR" => {
183                        state.request_restart();
184                        "OK\n".to_string()
185                    }
186                    _ => "ERR unknown command\n".to_string(),
187                };
188                if writer.write_all(response.as_bytes()).is_err() {
189                    break;
190                }
191            }
192            Err(_) => break,
193        }
194    }
195}
196
197/// Send a command to the daemon's health port and return the response.
198///
199/// Uses generous timeouts (2s connect, 5s read) suitable for `STOP`/`STATUS`
200/// where the caller is willing to wait. Use [`send_command_with_timeout`] for
201/// best-effort fire-and-forget calls (e.g. heartbeat, error reporting).
202///
203/// # Errors
204/// Returns an error if the connection fails or the response cannot be read.
205pub fn send_command(port: u16, command: &str) -> std::io::Result<String> {
206    send_command_with_timeout(
207        port,
208        command,
209        Duration::from_secs(2),
210        Duration::from_secs(5),
211    )
212}
213
214/// Best-effort fire-and-forget: tell the running daemon that hyperd appears to
215/// be dead from this client's perspective. Uses short timeouts (200ms each) so
216/// the calling tool handler isn't stalled if the daemon itself is slow.
217/// Errors are logged at debug level and otherwise ignored.
218pub fn report_hyperd_error_to_daemon() {
219    let port = super::discovery::resolve_port();
220    let timeout = Duration::from_millis(200);
221    match send_command_with_timeout(port, "REPORT_HYPERD_ERROR", timeout, timeout) {
222        Ok(response) => {
223            debug!(response = %response.trim(), "reported hyperd error to daemon");
224        }
225        Err(e) => {
226            debug!(error = %e, "could not report hyperd error to daemon (best-effort)");
227        }
228    }
229}
230
231/// Send a command with caller-specified connect/read timeouts.
232///
233/// # Errors
234/// Returns an error if the connection fails or the response cannot be read
235/// within the supplied timeouts.
236pub fn send_command_with_timeout(
237    port: u16,
238    command: &str,
239    connect_timeout: Duration,
240    read_timeout: Duration,
241) -> std::io::Result<String> {
242    let addr = std::net::SocketAddr::from(([127, 0, 0, 1], port));
243    let mut stream = TcpStream::connect_timeout(&addr, connect_timeout)?;
244    stream.set_read_timeout(Some(read_timeout))?;
245
246    let msg = format!("{command}\n");
247    stream.write_all(msg.as_bytes())?;
248    stream.flush()?;
249
250    let mut reader = BufReader::new(&stream);
251    let mut response = String::new();
252    reader.read_line(&mut response)?;
253    Ok(response)
254}
hyperdb_mcp/daemon/health.rs

hyperdb_mcp/daemon/
health.rs
