hatchery_node/
local_infra.rs

1//! Discovery of local infrastructure the node can reuse instead of starting
2//! its own copy: a mail4agent courier, another gate4agent node.
3//!
4//! Rule: an existing process is never killed or relaunched just because it
5//! exists. The node probes the loopback address it would bind, reads the
6//! service's `GET /health` (service name, version, optional capabilities)
7//! and then owns, attaches to, or steps around it. Upgrading an older
8//! instance is a graceful handoff through that service's control endpoint
9//! (see the `local-infra-reuse` spec); this module only reports the need.
10
11use std::io::{Read, Write};
12use std::net::{SocketAddr, TcpListener, TcpStream};
13use std::time::Duration;
14
15/// What answers on a loopback address.
16#[derive(Debug, Clone, PartialEq, Eq)]
17pub enum Probe {
18    /// Nothing holds the port (or it is ephemeral): we may bind it.
19    Free,
20    /// The expected service answers.
21    Ours(ServiceInfo),
22    /// Something else holds the port.
23    Foreign(String),
24}
25
26/// What a reused service reported about itself.
27#[derive(Debug, Clone, PartialEq, Eq, Default)]
28pub struct ServiceInfo {
29    pub service: String,
30    pub version: Option<String>,
31    pub capabilities: Vec<String>,
32    pub pid: Option<u64>,
33    /// Owner-only control token file the service published in `/health`.
34    pub control_token_file: Option<String>,
35}
36
37impl ServiceInfo {
38    /// True when the reported version is known and older than `minimum`.
39    pub fn older_than(&self, minimum: &str) -> bool {
40        match (self.version.as_deref().and_then(parse_version), parse_version(minimum)) {
41            (Some(found), Some(min)) => found < min,
42            _ => false,
43        }
44    }
45
46    /// Graceful control verbs the service advertises (`control.*` capabilities).
47    pub fn supports(&self, capability: &str) -> bool {
48        self.capabilities.iter().any(|c| c == capability)
49    }
50}
51
52/// `MAJOR.MINOR.PATCH` (pre-release / build suffix ignored).
53pub fn parse_version(raw: &str) -> Option<(u64, u64, u64)> {
54    let core = raw.trim().trim_start_matches('v').split(['-', '+']).next()?;
55    let mut parts = core.split('.').map(|p| p.parse::<u64>().ok());
56    Some((parts.next()??, parts.next().flatten().unwrap_or(0), parts.next().flatten().unwrap_or(0)))
57}
58
59/// Probe `addr` for `service` (the `service` field its `/health` returns).
60pub fn probe(addr: SocketAddr, service: &str, timeout: Duration) -> Probe {
61    if addr.port() == 0 {
62        return Probe::Free;
63    }
64    match TcpListener::bind(addr) {
65        Ok(listener) => {
66            drop(listener);
67            return Probe::Free;
68        }
69        Err(error) if error.kind() != std::io::ErrorKind::AddrInUse => {
70            return Probe::Foreign(format!("cannot bind: {error}"));
71        }
72        Err(_) => {}
73    }
74    match read_health(addr, timeout) {
75        Ok(json) => {
76            let name = json.get("service").and_then(|v| v.as_str()).unwrap_or("");
77            if name != service {
78                return Probe::Foreign(format!("health names {name:?}, not {service:?}"));
79            }
80            Probe::Ours(ServiceInfo {
81                service: name.to_owned(),
82                version: json.get("version").and_then(|v| v.as_str()).map(str::to_owned),
83                capabilities: json
84                    .get("capabilities")
85                    .and_then(|v| v.as_array())
86                    .map(|a| a.iter().filter_map(|c| c.as_str().map(str::to_owned)).collect())
87                    .unwrap_or_default(),
88                pid: json.get("pid").and_then(|v| v.as_u64()),
89                control_token_file: json.get("control_token_file").and_then(|v| v.as_str()).map(str::to_owned),
90            })
91        }
92        Err(detail) => Probe::Foreign(detail),
93    }
94}
95
96/// Minimal HTTP/1.1 `GET /health` (unauthenticated in both services).
97fn read_health(addr: SocketAddr, timeout: Duration) -> Result<serde_json::Value, String> {
98    let mut stream = TcpStream::connect_timeout(&addr, timeout).map_err(|e| format!("connect: {e}"))?;
99    let _ = stream.set_read_timeout(Some(timeout));
100    let _ = stream.set_write_timeout(Some(timeout));
101    stream
102        .write_all(format!("GET /health HTTP/1.1\r\nHost: {addr}\r\nConnection: close\r\n\r\n").as_bytes())
103        .map_err(|e| format!("write: {e}"))?;
104    let mut raw = Vec::new();
105    let _ = (&mut stream).take(64 * 1024).read_to_end(&mut raw);
106    let text = String::from_utf8_lossy(&raw);
107    let (head, body) = text.split_once("\r\n\r\n").ok_or_else(|| "no HTTP response".to_owned())?;
108    let status = head.lines().next().unwrap_or("");
109    if !(status.starts_with("HTTP/1.1 200") || status.starts_with("HTTP/1.0 200")) {
110        return Err(format!("health status: {status}"));
111    }
112    let start = body.find('{').ok_or_else(|| "health body is not JSON".to_owned())?;
113    let end = body.rfind('}').ok_or_else(|| "health body is not JSON".to_owned())?;
114    serde_json::from_str(&body[start..=end]).map_err(|e| format!("health JSON: {e}"))
115}
116
117/// Control verbs a reusable service advertises in `capabilities`.
118pub mod control {
119    pub const DRAIN: &str = "control.drain";
120    pub const RELOAD: &str = "control.reload";
121    pub const HANDOFF: &str = "control.handoff";
122}
123
124/// Upgrade decision for a reused service.
125#[derive(Debug, Clone, PartialEq, Eq)]
126pub enum Upgrade {
127    NotNeeded,
128    /// Older (or takeover forced) and it offers graceful handoff.
129    Handoff,
130    /// Older, no control endpoint: keep using it, tell the operator.
131    ManualOnly,
132}
133
134/// `force` = take over even when not older (redeploy of the same version).
135pub fn upgrade_plan(info: &ServiceInfo, minimum: &str, force: bool) -> Upgrade {
136    if !force && !info.older_than(minimum) {
137        Upgrade::NotNeeded
138    } else if info.supports(control::HANDOFF) && info.control_token_file.is_some() {
139        Upgrade::Handoff
140    } else if force || info.older_than(minimum) {
141        Upgrade::ManualOnly
142    } else {
143        Upgrade::NotNeeded
144    }
145}
146
147/// `HATCHERY_LOCAL_UPGRADE=off` disables automatic handoff (default: auto).
148pub fn auto_upgrade_allowed() -> bool {
149    std::env::var("HATCHERY_LOCAL_UPGRADE").map_or(true, |v| !v.eq_ignore_ascii_case("off"))
150}
151
152/// Ask the service on `addr` to hand over, then wait until the port is free.
153/// The service exits by itself; nothing is killed.
154pub fn handoff(addr: SocketAddr, info: &ServiceInfo, release_after_ms: u64, wait: Duration) -> Result<(), String> {
155    let file = info.control_token_file.as_deref().ok_or_else(|| "service published no control token file".to_owned())?;
156    let token = std::fs::read_to_string(file).map_err(|e| format!("control token file: {}", e.kind()))?;
157    let token = token.trim();
158    let body = format!("{{\"release_after_ms\":{release_after_ms},\"successor\":\"pid:{}\"}}", std::process::id());
159    let timeout = Duration::from_secs(3);
160    let mut stream = TcpStream::connect_timeout(&addr, timeout).map_err(|e| format!("connect: {e}"))?;
161    let _ = stream.set_read_timeout(Some(timeout));
162    let request = format!(
163        "POST /control/handoff?release_after_ms={release_after_ms} HTTP/1.1\r\nHost: {addr}\r\nAuthorization: Bearer {token}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
164        body.len()
165    );
166    stream.write_all(request.as_bytes()).map_err(|e| format!("write: {e}"))?;
167    let mut raw = Vec::new();
168    let _ = (&mut stream).take(16 * 1024).read_to_end(&mut raw);
169    let status = String::from_utf8_lossy(&raw).lines().next().unwrap_or("").to_owned();
170    if !status.contains(" 202") {
171        return Err(format!("handoff refused: {status}"));
172    }
173    wait_port_free(addr, wait)
174}
175
176/// Poll until `addr` can be bound (the old owner released it).
177pub fn wait_port_free(addr: SocketAddr, wait: Duration) -> Result<(), String> {
178    let deadline = std::time::Instant::now() + wait;
179    loop {
180        if TcpListener::bind(addr).is_ok() {
181            return Ok(());
182        }
183        if std::time::Instant::now() >= deadline {
184            return Err(format!("{addr} still held after {}ms", wait.as_millis()));
185        }
186        std::thread::sleep(Duration::from_millis(100));
187    }
188}
189
190#[cfg(test)]
191mod tests {
192    use super::*;
193
194    pub(crate) fn serve_once(body: &'static str, status: &'static str) -> SocketAddr {
195        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
196        let addr = listener.local_addr().unwrap();
197        std::thread::spawn(move || {
198            if let Ok((mut s, _)) = listener.accept() {
199                let mut buf = [0u8; 1024];
200                let _ = s.read(&mut buf);
201                let _ = write!(s, "HTTP/1.1 {status}\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{body}", body.len());
202            }
203        });
204        addr
205    }
206
207    const T: Duration = Duration::from_secs(2);
208
209    #[test]
210    fn free_and_zero_ports() {
211        let l = TcpListener::bind("127.0.0.1:0").unwrap();
212        let addr = l.local_addr().unwrap();
213        drop(l);
214        assert_eq!(probe(addr, "svc", T), Probe::Free);
215        assert_eq!(probe("127.0.0.1:0".parse().unwrap(), "svc", T), Probe::Free);
216    }
217
218    #[test]
219    fn ours_reports_version_and_capabilities() {
220        let addr = serve_once(r#"{"ok":true,"service":"svc","version":"1.2.3","pid":7,"capabilities":["control.drain"]}"#, "200 OK");
221        let Probe::Ours(info) = probe(addr, "svc", T) else { panic!("not ours") };
222        assert_eq!(info.version.as_deref(), Some("1.2.3"));
223        assert_eq!(info.pid, Some(7));
224        assert!(info.supports(control::DRAIN));
225        assert!(!info.supports(control::HANDOFF));
226    }
227
228    #[test]
229    fn other_service_or_status_is_foreign() {
230        let addr = serve_once(r#"{"service":"other"}"#, "200 OK");
231        assert!(matches!(probe(addr, "svc", T), Probe::Foreign(_)));
232        let addr = serve_once(r#"{"service":"svc"}"#, "500 Internal Server Error");
233        assert!(matches!(probe(addr, "svc", T), Probe::Foreign(_)));
234    }
235
236    #[test]
237    fn silent_socket_is_foreign() {
238        let l = TcpListener::bind("127.0.0.1:0").unwrap();
239        let addr = l.local_addr().unwrap();
240        assert!(matches!(probe(addr, "svc", Duration::from_millis(300)), Probe::Foreign(_)));
241    }
242
243    #[test]
244    fn version_and_upgrade_plan() {
245        assert_eq!(parse_version("0.4.14"), Some((0, 4, 14)));
246        assert_eq!(parse_version("v1.2-rc1"), Some((1, 2, 0)));
247        let old = ServiceInfo { version: Some("0.4.9".into()), ..Default::default() };
248        assert_eq!(upgrade_plan(&old, "0.4.14", false), Upgrade::ManualOnly);
249        let old_ctl = ServiceInfo { version: Some("0.4.9".into()), capabilities: vec![control::HANDOFF.into()], control_token_file: Some("t".into()), ..Default::default() };
250        assert_eq!(upgrade_plan(&old_ctl, "0.4.14", false), Upgrade::Handoff);
251        let new = ServiceInfo { version: Some("0.5.0".into()), ..Default::default() };
252        assert_eq!(upgrade_plan(&new, "0.4.14", false), Upgrade::NotNeeded);
253        let unknown = ServiceInfo::default();
254        assert_eq!(upgrade_plan(&unknown, "0.4.14", false), Upgrade::NotNeeded);
255    }
256
257    #[test]
258    fn forced_takeover_of_same_version() {
259        let same = ServiceInfo { version: Some("1.0.0".into()), capabilities: vec![control::HANDOFF.into()], control_token_file: Some("t".into()), ..Default::default() };
260        assert_eq!(upgrade_plan(&same, "1.0.0", false), Upgrade::NotNeeded);
261        assert_eq!(upgrade_plan(&same, "1.0.0", true), Upgrade::Handoff);
262        let bare = ServiceInfo { version: Some("1.0.0".into()), ..Default::default() };
263        assert_eq!(upgrade_plan(&bare, "1.0.0", true), Upgrade::ManualOnly);
264    }
265
266    /// A fake service: answers /health, accepts one authorized handoff, then
267    /// releases its listener the way a real one does.
268    #[test]
269    fn handoff_releases_the_port() {
270        let dir = std::env::temp_dir().join(format!("li-handoff-{}", std::process::id()));
271        std::fs::create_dir_all(&dir).unwrap();
272        let token_file = dir.join("control.token");
273        std::fs::write(&token_file, "secret-token-for-test").unwrap();
274        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
275        let addr = listener.local_addr().unwrap();
276        let tf = token_file.display().to_string().replace('\\', "\\\\");
277        std::thread::spawn(move || {
278            for stream in listener.incoming() {
279                let mut s = stream.unwrap();
280                let mut buf = [0u8; 4096];
281                let n = s.read(&mut buf).unwrap_or(0);
282                let req = String::from_utf8_lossy(&buf[..n]).to_string();
283                if req.starts_with("GET /health") {
284                    let body = format!(r#"{{"service":"svc","version":"0.1.0","capabilities":["control.handoff"],"control_token_file":"{tf}"}}"#);
285                    let _ = write!(s, "HTTP/1.1 200 OK\r\nContent-Length: {}\r\n\r\n{body}", body.len());
286                } else if req.starts_with("POST /control/handoff") && req.contains("Bearer secret-token-for-test") {
287                    let _ = write!(s, "HTTP/1.1 202 Accepted\r\nContent-Length: 2\r\n\r\n{{}}");
288                    return; // drop listener = release port
289                } else {
290                    let _ = write!(s, "HTTP/1.1 401 Unauthorized\r\nContent-Length: 0\r\n\r\n");
291                }
292            }
293        });
294        let Probe::Ours(info) = probe(addr, "svc", T) else { panic!("not ours") };
295        assert_eq!(upgrade_plan(&info, "0.4.0", false), Upgrade::Handoff);
296        handoff(addr, &info, 50, Duration::from_secs(5)).unwrap();
297        assert_eq!(probe(addr, "svc", T), Probe::Free);
298        let _ = std::fs::remove_dir_all(dir);
299    }
300}
hatchery_node/local_infra.rs

hatchery_node/
local_infra.rs
