Skip to main content

rightkit_daemon/
reload.rs

1//! Crash-safe daemon reload: replace a running daemon with a new build without dropping the
2//! endpoint or losing recovery state (rightkit#666).
3//!
4//! One serialized saga, one coordinator for every caller (a CLI, a TUI, the desktop shell):
5//!
6//! 1. write the durable marker `daemon.reload.json` (`request_id`, `hash`, `phase`, `pid`, time,
7//!    `endpoint`, `detail`) in phase `starting`; the request id and hash bind the replacement to
8//!    the update that asked, so a stale marker is never taken for readiness;
9//! 2. host hook [`ReloadHost::persist_recovery`] (snapshot the sessions and state);
10//! 3. host hook [`ReloadHost::fence_intake`] (refuse new mutations);
11//! 4. marker `draining`, host hook [`ReloadHost::drain`] (finish or stop in-flight work);
12//! 5. marker `handoff`; when a [`StagedBinary`] was given, swap it onto its target path under an
13//!    [`ActivationFence`] (journalled, so a crash in the middle is undone by
14//!    [`ActivationFence::recover`]);
15//! 6. host hook [`ReloadHost::launch_replacement`] (spawn the new process with the same endpoint
16//!    and the exact launch arguments; the host stamps [`ENV_RELOAD_REQUEST`] and
17//!    [`ENV_RELOAD_HASH`] so the child knows it is a reload replacement);
18//! 7. host hook [`ReloadHost::wait_ready`]: poll [`read_marker`] and the live endpoint until the
19//!    replacement calls [`publish_socket_ready`] (marker `socket_ready`) for exactly that request
20//!    id and hash; a timeout is a failure;
21//! 8. host hook [`ReloadHost::resume`], then the activation is committed (backup removed).
22//!
23//! Rollback: any stage failure writes marker phase `failed` with `"stage: error"`, rolls the
24//! activation fence back (restores the previous binary), then calls [`ReloadHost::rollback`] (which
25//! terminates the never-ready replacement by pid, best effort, and restores the old runtime).
26//! Errors compose: the stage error, the rollback error and a failed marker write are all reported
27//! ([`ReloadError::StageAndRollback`]). The caller never silently continues after an incomplete
28//! reload. A second request while one runs is [`ReloadError::InFlight`], never a second saga.
29
30use crate::record::write_atomic;
31use serde::{Deserialize, Serialize};
32use std::io;
33use std::path::{Path, PathBuf};
34use std::sync::Mutex;
35use std::time::{Duration, SystemTime, UNIX_EPOCH};
36
37/// The marker's file name, in the directory the host chooses.
38pub const MARKER_FILE: &str = "daemon.reload.json";
39/// Environment variable the replacement receives: the reload request id.
40pub const ENV_RELOAD_REQUEST: &str = "RIGHTKIT_RELOAD_REQUEST";
41/// Environment variable the replacement receives: the hash of the update.
42pub const ENV_RELOAD_HASH: &str = "RIGHTKIT_RELOAD_HASH";
43
44/// Where a reload is, as persisted.
45#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
46#[serde(rename_all = "snake_case")]
47pub enum ReloadPhase {
48    Starting,
49    Draining,
50    Handoff,
51    /// The replacement published its endpoint for this request.
52    SocketReady,
53    Failed,
54}
55
56/// The durable record of a reload.
57#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
58pub struct ReloadMarker {
59    pub request_id: String,
60    pub hash: String,
61    pub phase: ReloadPhase,
62    /// The pid of the process that wrote this phase (the old daemon, then the replacement).
63    pub pid: u32,
64    pub timestamp_ms: u64,
65    pub endpoint: String,
66    #[serde(default)]
67    pub detail: String,
68}
69
70fn now_ms() -> u64 {
71    SystemTime::now()
72        .duration_since(UNIX_EPOCH)
73        .map(|d| d.as_millis() as u64)
74        .unwrap_or(0)
75}
76
77/// Read the marker in `dir`. `Ok(None)` when there is none.
78pub fn read_marker(dir: &Path) -> io::Result<Option<ReloadMarker>> {
79    match std::fs::read(dir.join(MARKER_FILE)) {
80        Ok(bytes) => serde_json::from_slice(&bytes)
81            .map(Some)
82            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e)),
83        Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(None),
84        Err(e) => Err(e),
85    }
86}
87
88fn write_marker(dir: &Path, marker: &ReloadMarker) -> io::Result<()> {
89    let body = serde_json::to_vec_pretty(marker)
90        .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
91    write_atomic(&dir.join(MARKER_FILE), &body)
92}
93
94/// Called by the REPLACEMENT once its endpoint is live: marker `socket_ready` for exactly
95/// `request_id` and `hash`. Refused when the marker is missing or names another request or hash,
96/// so a stale marker, or a process that was not launched by this reload, can never claim
97/// readiness.
98pub fn publish_socket_ready(
99    dir: &Path,
100    request_id: &str,
101    hash: &str,
102    endpoint: &str,
103) -> io::Result<()> {
104    let Some(mut marker) = read_marker(dir)? else {
105        return Err(io::Error::new(
106            io::ErrorKind::NotFound,
107            "no reload is in progress",
108        ));
109    };
110    if marker.request_id != request_id || marker.hash != hash || marker.phase == ReloadPhase::Failed
111    {
112        return Err(io::Error::new(
113            io::ErrorKind::PermissionDenied,
114            "the reload marker is for another request, or the reload failed",
115        ));
116    }
117    marker.phase = ReloadPhase::SocketReady;
118    marker.pid = std::process::id();
119    marker.timestamp_ms = now_ms();
120    marker.endpoint = endpoint.to_string();
121    write_marker(dir, &marker)
122}
123
124/// What the replacement is told, and what the hooks receive.
125#[derive(Debug, Clone, PartialEq, Eq)]
126pub struct ReloadRequest {
127    pub request_id: String,
128    pub hash: String,
129    /// The endpoint the replacement must serve (the same one).
130    pub endpoint: String,
131}
132
133/// A new binary staged next to the one that runs.
134#[derive(Debug, Clone, PartialEq, Eq)]
135pub struct StagedBinary {
136    pub staged: PathBuf,
137    pub target: PathBuf,
138}
139
140/// The host's side of the saga. Every hook returns `Err(text)` to fail its stage.
141pub trait ReloadHost {
142    fn persist_recovery(&self, request: &ReloadRequest) -> Result<(), String>;
143    fn fence_intake(&self, request: &ReloadRequest) -> Result<(), String>;
144    fn drain(&self, request: &ReloadRequest) -> Result<(), String>;
145    /// Spawn the replacement, detached, with the same endpoint and launch arguments, and with
146    /// [`ENV_RELOAD_REQUEST`] and [`ENV_RELOAD_HASH`] set. Returns its pid.
147    fn launch_replacement(&self, request: &ReloadRequest) -> Result<u32, String>;
148    /// Wait until the replacement is ready: normally `poll_ready(dir, request)` and a live
149    /// endpoint probe, until `timeout`. `Err` is a failure (a timeout is one).
150    fn wait_ready(
151        &self,
152        request: &ReloadRequest,
153        replacement_pid: u32,
154        timeout: Duration,
155    ) -> Result<(), String>;
156    /// Reopen for work (the old daemon was replaced; this is the resume of the host's state).
157    fn resume(&self, request: &ReloadRequest) -> Result<(), String>;
158    /// Undo a failed reload: terminate the never-ready replacement (`replacement_pid`, when one
159    /// was launched) and restore the old runtime.
160    fn rollback(&self, request: &ReloadRequest, replacement_pid: Option<u32>)
161        -> Result<(), String>;
162}
163
164/// True once the marker in `dir` is `socket_ready` for exactly this request (the check a
165/// [`ReloadHost::wait_ready`] polls).
166pub fn poll_ready(dir: &Path, request: &ReloadRequest) -> bool {
167    matches!(
168        read_marker(dir),
169        Ok(Some(m)) if m.phase == ReloadPhase::SocketReady
170            && m.request_id == request.request_id
171            && m.hash == request.hash
172    )
173}
174
175/// Why a reload did not complete.
176#[derive(Debug)]
177#[non_exhaustive]
178pub enum ReloadError {
179    /// Another reload is running in this process.
180    InFlight,
181    /// The marker could not be written at the start: nothing was changed.
182    Marker(io::Error),
183    /// A stage failed and the rollback succeeded.
184    Stage { stage: &'static str, error: String },
185    /// A stage failed and the rollback, or the failed-marker write, also failed. Every error is
186    /// kept: the host must not continue as if the old daemon were intact.
187    StageAndRollback {
188        stage: &'static str,
189        error: String,
190        rollback: Option<String>,
191        marker: Option<String>,
192    },
193}
194
195impl std::fmt::Display for ReloadError {
196    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
197        match self {
198            Self::InFlight => write!(f, "a reload is already running"),
199            Self::Marker(e) => write!(f, "the reload marker could not be written: {e}"),
200            Self::Stage { stage, error } => write!(f, "reload failed at {stage}: {error}"),
201            Self::StageAndRollback {
202                stage,
203                error,
204                rollback,
205                marker,
206            } => {
207                write!(f, "reload failed at {stage}: {error}")?;
208                if let Some(r) = rollback {
209                    write!(f, "; rollback failed: {r}")?;
210                }
211                if let Some(m) = marker {
212                    write!(f, "; the failed marker was not written: {m}")?;
213                }
214                Ok(())
215            }
216        }
217    }
218}
219
220impl std::error::Error for ReloadError {}
221
222/// The journalled swap of a staged binary onto its target path. The journal
223/// (`<target>.reload-journal`) names the target, the backup of the previous binary and the staged
224/// file, so [`recover`](Self::recover) can undo a swap that a crash cut in the middle.
225#[derive(Debug)]
226pub struct ActivationFence {
227    target: PathBuf,
228    backup: PathBuf,
229    journal: PathBuf,
230    staged: PathBuf,
231}
232
233#[derive(Serialize, Deserialize)]
234struct Journal {
235    target: PathBuf,
236    backup: PathBuf,
237    staged: PathBuf,
238    swapped: bool,
239}
240
241fn side(path: &Path, suffix: &str) -> PathBuf {
242    let mut name = path
243        .file_name()
244        .map(|n| n.to_os_string())
245        .unwrap_or_default();
246    name.push(suffix);
247    path.with_file_name(name)
248}
249
250impl ActivationFence {
251    pub fn new(binary: &StagedBinary) -> Self {
252        Self {
253            target: binary.target.clone(),
254            backup: side(&binary.target, ".reload-previous"),
255            journal: side(&binary.target, ".reload-journal"),
256            staged: binary.staged.clone(),
257        }
258    }
259
260    fn write_journal(&self, swapped: bool) -> io::Result<()> {
261        let body = serde_json::to_vec(&Journal {
262            target: self.target.clone(),
263            backup: self.backup.clone(),
264            staged: self.staged.clone(),
265            swapped,
266        })
267        .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
268        write_atomic(&self.journal, &body)
269    }
270
271    /// Swap: journal, target to backup, staged onto target, journal marks it done.
272    pub fn activate(&self) -> io::Result<()> {
273        if !self.staged.is_file() {
274            return Err(io::Error::new(
275                io::ErrorKind::NotFound,
276                "the staged binary is missing",
277            ));
278        }
279        let _ = std::fs::remove_file(&self.backup);
280        self.write_journal(false)?;
281        std::fs::rename(&self.target, &self.backup)?;
282        if let Err(e) = std::fs::rename(&self.staged, &self.target) {
283            let _ = std::fs::rename(&self.backup, &self.target);
284            let _ = std::fs::remove_file(&self.journal);
285            return Err(e);
286        }
287        self.write_journal(true)
288    }
289
290    /// Put the previous binary back and drop the journal.
291    pub fn rollback(&self) -> io::Result<()> {
292        if self.backup.is_file() {
293            // The new binary goes back to its staged name, the previous one to the target.
294            let _ = std::fs::rename(&self.target, &self.staged);
295            std::fs::rename(&self.backup, &self.target)?;
296        }
297        match std::fs::remove_file(&self.journal) {
298            Ok(()) => Ok(()),
299            Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(()),
300            Err(e) => Err(e),
301        }
302    }
303
304    /// The reload succeeded: drop the backup and the journal.
305    pub fn commit(&self) -> io::Result<()> {
306        let _ = std::fs::remove_file(&self.backup);
307        match std::fs::remove_file(&self.journal) {
308            Ok(()) => Ok(()),
309            Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(()),
310            Err(e) => Err(e),
311        }
312    }
313
314    /// At startup: a journal that was never committed means a reload was cut by a crash. When the
315    /// marker in `dir` is not `socket_ready` for the journalled swap, the previous binary is
316    /// restored. Returns whether a rollback ran.
317    pub fn recover(binary: &StagedBinary, dir: &Path) -> io::Result<bool> {
318        let fence = Self::new(binary);
319        let Ok(bytes) = std::fs::read(&fence.journal) else {
320            return Ok(false);
321        };
322        let Ok(journal) = serde_json::from_slice::<Journal>(&bytes) else {
323            return Err(io::Error::new(
324                io::ErrorKind::InvalidData,
325                "unreadable reload journal",
326            ));
327        };
328        let ready = matches!(read_marker(dir), Ok(Some(m)) if m.phase == ReloadPhase::SocketReady);
329        if journal.swapped && ready {
330            fence.commit()?;
331            return Ok(false);
332        }
333        fence.rollback()?;
334        Ok(true)
335    }
336}
337
338/// Runs reloads, one at a time.
339pub struct ReloadCoordinator {
340    dir: PathBuf,
341    ready_timeout: Duration,
342    running: Mutex<()>,
343}
344
345impl ReloadCoordinator {
346    /// `dir` holds the marker (the host's runtime directory).
347    pub fn new(dir: impl Into<PathBuf>) -> Self {
348        Self {
349            dir: dir.into(),
350            ready_timeout: Duration::from_secs(30),
351            running: Mutex::new(()),
352        }
353    }
354
355    /// How long the replacement gets to become ready. Default 30 s.
356    #[must_use]
357    pub fn ready_timeout(mut self, timeout: Duration) -> Self {
358        self.ready_timeout = timeout;
359        self
360    }
361
362    pub fn dir(&self) -> &Path {
363        &self.dir
364    }
365
366    /// Run the saga. `Ok` only when the replacement answered ready and the host resumed.
367    pub fn run(
368        &self,
369        host: &dyn ReloadHost,
370        request: &ReloadRequest,
371        staged: Option<&StagedBinary>,
372    ) -> Result<(), ReloadError> {
373        let Ok(_one) = self.running.try_lock() else {
374            return Err(ReloadError::InFlight);
375        };
376        let mut marker = ReloadMarker {
377            request_id: request.request_id.clone(),
378            hash: request.hash.clone(),
379            phase: ReloadPhase::Starting,
380            pid: std::process::id(),
381            timestamp_ms: now_ms(),
382            endpoint: request.endpoint.clone(),
383            detail: String::new(),
384        };
385        write_marker(&self.dir, &marker).map_err(ReloadError::Marker)?;
386        let fence = staged.map(ActivationFence::new);
387        let mut replacement = None;
388        let outcome = self.stages(host, request, &mut marker, fence.as_ref(), &mut replacement);
389        match outcome {
390            Ok(()) => {
391                if let Some(fence) = &fence {
392                    let _ = fence.commit();
393                }
394                Ok(())
395            }
396            Err((stage, error)) => {
397                marker.phase = ReloadPhase::Failed;
398                marker.detail = format!("{stage}: {error}");
399                marker.timestamp_ms = now_ms();
400                let marker_error = write_marker(&self.dir, &marker)
401                    .err()
402                    .map(|e| e.to_string());
403                let mut rollback_error = None;
404                if let Some(fence) = &fence {
405                    if let Err(e) = fence.rollback() {
406                        rollback_error = Some(format!("activation: {e}"));
407                    }
408                }
409                if let Err(e) = host.rollback(request, replacement) {
410                    rollback_error = Some(match rollback_error {
411                        Some(prior) => format!("{prior}; host: {e}"),
412                        None => format!("host: {e}"),
413                    });
414                }
415                if rollback_error.is_none() && marker_error.is_none() {
416                    Err(ReloadError::Stage { stage, error })
417                } else {
418                    Err(ReloadError::StageAndRollback {
419                        stage,
420                        error,
421                        rollback: rollback_error,
422                        marker: marker_error,
423                    })
424                }
425            }
426        }
427    }
428
429    fn stages(
430        &self,
431        host: &dyn ReloadHost,
432        request: &ReloadRequest,
433        marker: &mut ReloadMarker,
434        fence: Option<&ActivationFence>,
435        replacement: &mut Option<u32>,
436    ) -> Result<(), (&'static str, String)> {
437        host.persist_recovery(request)
438            .map_err(|e| ("persist_recovery", e))?;
439        host.fence_intake(request)
440            .map_err(|e| ("fence_intake", e))?;
441        self.advance(marker, ReloadPhase::Draining)?;
442        host.drain(request).map_err(|e| ("drain", e))?;
443        self.advance(marker, ReloadPhase::Handoff)?;
444        if let Some(fence) = fence {
445            fence.activate().map_err(|e| ("activate", e.to_string()))?;
446        }
447        let pid = host
448            .launch_replacement(request)
449            .map_err(|e| ("launch_replacement", e))?;
450        *replacement = Some(pid);
451        host.wait_ready(request, pid, self.ready_timeout)
452            .map_err(|e| ("wait_ready", e))?;
453        if !poll_ready(&self.dir, request) {
454            return Err((
455                "wait_ready",
456                "the replacement did not publish readiness for this request".to_string(),
457            ));
458        }
459        host.resume(request).map_err(|e| ("resume", e))
460    }
461
462    fn advance(
463        &self,
464        marker: &mut ReloadMarker,
465        phase: ReloadPhase,
466    ) -> Result<(), (&'static str, String)> {
467        marker.phase = phase;
468        marker.timestamp_ms = now_ms();
469        write_marker(&self.dir, marker).map_err(|e| ("marker", e.to_string()))
470    }
471}