Skip to main content

kunobi_daemon/
upgrade.rs

1use crate::ProcessLock;
2use std::{fmt, future::Future, io, path::Path, time::Duration};
3use tokio::time::Instant;
4
5/// Failure to obtain a live daemon satisfying the caller's version policy.
6#[derive(Debug)]
7#[non_exhaustive]
8pub enum UpgradeError<E> {
9    /// A health probe failed, rather than reporting an absent daemon.
10    Probe(E),
11    /// The caller's replacement operation failed.
12    Replace(E),
13    /// The upgrade lock could not be opened or acquired.
14    Lock(io::Error),
15    /// The shared deadline expired during a probe, lock wait, or replacement.
16    Deadline,
17}
18
19impl<E: fmt::Display> fmt::Display for UpgradeError<E> {
20    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
21        match self {
22            Self::Probe(e) => write!(f, "daemon health probe failed: {e}"),
23            Self::Replace(e) => write!(f, "daemon replacement failed: {e}"),
24            Self::Lock(e) => write!(f, "daemon upgrade lock failed: {e}"),
25            Self::Deadline => f.write_str("daemon did not become current before the deadline"),
26        }
27    }
28}
29
30impl<E: std::error::Error + 'static> std::error::Error for UpgradeError<E> {
31    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
32        match self {
33            Self::Probe(e) | Self::Replace(e) => Some(e),
34            Self::Lock(e) => Some(e),
35            Self::Deadline => None,
36        }
37    }
38}
39
40/// Return a live response accepted by the caller, replacing the daemon if needed.
41///
42/// `probe` must perform a live, cheap health request; `None` means no daemon is
43/// available yet. Errors (including permission errors) are propagated. `accepts`
44/// defines build/protocol compatibility; the crate does not order build strings.
45///
46/// A replacement holds the persistent `upgrade_lock` and probes again after
47/// acquiring it, so concurrent upgraders can reuse the winner's replacement.
48/// `replace` must use the existing service manager when it owns the daemon.
49/// Success always comes from a fresh accepted probe, never an exit code or a
50/// retiring daemon's response.
51///
52/// All awaits share `deadline`. Callbacks must yield, bound their own blocking
53/// I/O, and tolerate cancellation. Dropping this future releases its upgrade
54/// lock; it does not kill a process spawned by `replace`.
55pub async fn ensure_current<T, E, P, PF, A, R, RF>(
56    upgrade_lock: &Path,
57    deadline: Instant,
58    mut probe: P,
59    accepts: A,
60    replace: R,
61) -> Result<T, UpgradeError<E>>
62where
63    P: FnMut() -> PF,
64    PF: Future<Output = Result<Option<T>, E>>,
65    A: Fn(&T) -> bool,
66    R: FnOnce() -> RF,
67    RF: Future<Output = Result<(), E>>,
68{
69    if Instant::now() >= deadline {
70        return Err(UpgradeError::Deadline);
71    }
72    tokio::time::timeout_at(deadline, async {
73        if let Some(current) = probe().await.map_err(UpgradeError::Probe)?
74            && accepts(&current)
75        {
76            return Ok(current);
77        }
78        let _upgrade = ProcessLock::acquire_async(upgrade_lock)
79            .await
80            .map_err(UpgradeError::Lock)?;
81        if let Some(current) = probe().await.map_err(UpgradeError::Probe)?
82            && accepts(&current)
83        {
84            return Ok(current);
85        }
86        replace().await.map_err(UpgradeError::Replace)?;
87        loop {
88            if let Some(current) = probe().await.map_err(UpgradeError::Probe)?
89                && accepts(&current)
90            {
91                return Ok(current);
92            }
93            #[expect(clippy::disallowed_methods, reason = "The application probe exposes no readiness event; the outer timeout bounds retries.")]
94            tokio::time::sleep(Duration::from_millis(25)).await;
95        }
96    })
97    .await
98    .map_err(|_| UpgradeError::Deadline)?
99}