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