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