1use crate::ProcessLock;
2use std::{fmt, future::Future, io, path::Path, time::Duration};
3use tokio::time::Instant;
4
5#[derive(Debug)]
7#[non_exhaustive]
8pub enum UpgradeError<E> {
9 Probe(E),
11 Replace(E),
13 Lock(io::Error),
15 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
40pub 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(¤t)
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(¤t)
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(¤t)
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}