use crate::ProcessLock;
use std::{fmt, future::Future, io, path::Path, time::Duration};
use tokio::time::Instant;
#[derive(Debug)]
#[non_exhaustive]
pub enum UpgradeError<E> {
Probe(E),
Replace(E),
Lock(io::Error),
Deadline,
}
impl<E: fmt::Display> fmt::Display for UpgradeError<E> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Probe(e) => write!(f, "daemon health probe failed: {e}"),
Self::Replace(e) => write!(f, "daemon replacement failed: {e}"),
Self::Lock(e) => write!(f, "daemon upgrade lock failed: {e}"),
Self::Deadline => f.write_str("daemon did not become current before the deadline"),
}
}
}
impl<E: std::error::Error + 'static> std::error::Error for UpgradeError<E> {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::Probe(e) | Self::Replace(e) => Some(e),
Self::Lock(e) => Some(e),
Self::Deadline => None,
}
}
}
pub async fn ensure_current<T, E, P, PF, A, R, RF>(
upgrade_lock: &Path,
deadline: Instant,
mut probe: P,
accepts: A,
replace: R,
) -> Result<T, UpgradeError<E>>
where
P: FnMut() -> PF,
PF: Future<Output = Result<Option<T>, E>>,
A: Fn(&T) -> bool,
R: FnOnce() -> RF,
RF: Future<Output = Result<(), E>>,
{
if Instant::now() >= deadline {
return Err(UpgradeError::Deadline);
}
tokio::time::timeout_at(deadline, async {
if let Some(current) = probe().await.map_err(UpgradeError::Probe)?
&& accepts(¤t)
{
return Ok(current);
}
let _upgrade = ProcessLock::acquire_async(upgrade_lock)
.await
.map_err(UpgradeError::Lock)?;
if let Some(current) = probe().await.map_err(UpgradeError::Probe)?
&& accepts(¤t)
{
return Ok(current);
}
replace().await.map_err(UpgradeError::Replace)?;
loop {
if let Some(current) = probe().await.map_err(UpgradeError::Probe)?
&& accepts(¤t)
{
return Ok(current);
}
#[expect(clippy::disallowed_methods, reason = "The application probe exposes no readiness event; the outer timeout bounds retries.")]
tokio::time::sleep(Duration::from_millis(25)).await;
}
})
.await
.map_err(|_| UpgradeError::Deadline)?
}