kunobi_daemon/

replacement.rs

//! Replacement ordering shared by blocking clients and async servers.
//!
//! The caller holds the persistent upgrade lock throughout the transaction.
//! Adapters perform one operation at a time; this module decides their order.
use crate::ProcessLock;
use std::time::{Duration, Instant};

/// Whether the application permits old and new processes to coexist.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[cfg_attr(kani, derive(kani::Arbitrary))]
pub enum Mode {
    /// Verify and select the candidate before retiring the incumbent.
    Overlap,
    /// Drain and release the incumbent before starting the candidate.
    Exclusive,
}

/// One operation supplied by an application adapter.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Step {
    /// Re-read selection under the transaction lock. Unchanged skips replacement.
    Recheck,
    /// Validate the artifact and capture application preparation state.
    Prepare,
    /// Close incumbent admission and release exclusive resources.
    Drain,
    /// Launch the candidate through its process owner.
    Start,
    /// Obtain a live identity/readiness proof. Pending retries this step only.
    Verify,
    /// Revalidate selection and application preparation after the live proof.
    Validate,
    /// Atomically select the verified candidate; never return an error after committing.
    Commit,
    /// Start retiring the incumbent. Pending keeps its obligations alive.
    Retire,
}

/// Result of one adapter operation.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[cfg_attr(kani, derive(kani::Arbitrary))]
pub enum Progress {
    /// The operation completed.
    Done,
    /// The proof/drain is not ready, or retirement remains in progress.
    Pending,
    /// Recheck found no replacement necessary or another updater already won.
    Unchanged,
}

/// A replacement result never implies replay of application work.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub enum Outcome {
    /// No replacement was performed.
    Unchanged,
    /// Selection committed and the incumbent has no remaining obligations.
    Complete,
    /// Selection committed; the incumbent must remain alive for its obligations.
    RetirementPending,
}

/// Failure retains the phase and whether selection has already committed.
#[derive(Debug)]
#[non_exhaustive]
pub struct Failure<E> {
    /// Failed operation.
    pub step: Step,
    /// Whether the new selection is authoritative despite this failure.
    pub committed: bool,
    /// Why the operation failed.
    pub reason: Reason<E>,
}

/// Cause of a replacement failure.
#[derive(Debug)]
#[non_exhaustive]
pub enum Reason<E> {
    /// Setup or the consumer's explicit drain budget expired.
    Deadline,
    /// Application or operating system operation failed.
    Adapter(E),
    /// The adapter returned a result that is not valid for the current step.
    InvalidProgress,
}

impl<E: std::fmt::Display> std::fmt::Display for Failure<E> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "replacement {:?} failed: ", self.step)?;
        match &self.reason {
            Reason::Deadline => f.write_str("deadline expired"),
            Reason::Adapter(error) => error.fmt(f),
            Reason::InvalidProgress => f.write_str("invalid adapter progress"),
        }
    }
}
impl<E: std::error::Error + 'static> std::error::Error for Failure<E> {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match &self.reason {
            Reason::Adapter(error) => Some(error),
            Reason::Deadline | Reason::InvalidProgress => None,
        }
    }
}

/// Short setup budget and independent application-owned drain policy.
#[derive(Clone, Copy, Debug)]
pub struct Budgets {
    /// Absolute budget shared by setup operations, restarted after exclusive drain.
    pub setup: Duration,
    /// None allows drain to take arbitrarily long. Expiry never kills work.
    pub drain: Option<Duration>,
}

/// Blocking application operations; no runtime is introduced into the client.
pub trait Driver {
    /// Application error.
    type Error;
    /// Perform exactly the requested operation, bounding I/O by the given deadline.
    fn perform(&mut self, step: Step, deadline: Option<Instant>) -> Result<Progress, Self::Error>;
    /// Client shims to prefault after selection commits. Empty by default.
    ///
    /// Typical entries are the relay or compiler wrapper next to this daemon.
    /// Prefault errors are discarded: publication has already succeeded.
    fn warmup_paths(&self) -> Vec<std::path::PathBuf> {
        Vec::new()
    }
}

/// Async application operations.
#[cfg(feature = "async")]
pub trait AsyncDriver {
    /// Application error.
    type Error;
    /// Perform exactly the requested operation. Cancellation must preserve ownership.
    fn perform(
        &mut self,
        step: Step,
        deadline: Option<tokio::time::Instant>,
    ) -> impl std::future::Future<Output = Result<Progress, Self::Error>> + Send;
    /// Client shims to prefault after selection commits. Empty by default.
    fn warmup_paths(&self) -> Vec<std::path::PathBuf> {
        Vec::new()
    }
}

struct Machine {
    mode: Mode,
    step: Step,
    committed: bool,
}
impl Machine {
    fn new(mode: Mode) -> Self {
        Self {
            mode,
            step: Step::Recheck,
            committed: false,
        }
    }
    fn failure<E>(&self, reason: Reason<E>) -> Failure<E> {
        Failure {
            step: self.step,
            committed: self.committed,
            reason,
        }
    }
    fn advance<E>(&mut self, progress: Progress) -> Result<Option<Outcome>, Failure<E>> {
        use Step::*;
        match (self.step, progress) {
            (Recheck, Progress::Unchanged) => return Ok(Some(Outcome::Unchanged)),
            (Recheck | Verify | Drain, Progress::Pending) => return Ok(None),
            (Retire, Progress::Pending) => return Ok(Some(Outcome::RetirementPending)),
            (_, Progress::Done) => {}
            _ => return Err(self.failure(Reason::InvalidProgress)),
        }
        self.step = match self.step {
            Recheck => Prepare,
            Prepare if self.mode == Mode::Exclusive => Drain,
            Prepare | Drain => Start,
            Start => Verify,
            Verify => Validate,
            Validate => Commit,
            Commit => {
                self.committed = true;
                if self.mode == Mode::Exclusive {
                    return Ok(Some(Outcome::Complete));
                }
                Retire
            }
            Retire => return Ok(Some(Outcome::Complete)),
        };
        Ok(None)
    }
}

/// Execute a serialized replacement. The lock is retained by the caller on return.
///
/// The adapter must preserve a selected candidate when reporting post-commit
/// retirement failures. No callback automatically replays requests or kills a daemon.
pub fn run<D: Driver>(
    _ownership: &ProcessLock,
    mode: Mode,
    budgets: Budgets,
    driver: &mut D,
) -> Result<Outcome, Failure<D::Error>> {
    let mut machine = Machine::new(mode);
    let mut deadline = Some(Instant::now() + budgets.setup);
    loop {
        if deadline.is_some_and(|limit| Instant::now() >= limit) {
            return Err(machine.failure(Reason::Deadline));
        }
        let step = machine.step;
        let progress = driver
            .perform(step, deadline)
            .map_err(|e| machine.failure(Reason::Adapter(e)))?;
        // A committed result must be recorded even if a blocking commit ran late.
        if step != Step::Commit && deadline.is_some_and(|limit| Instant::now() >= limit) {
            return Err(machine.failure(Reason::Deadline));
        }
        let committed_before = machine.committed;
        if let Some(outcome) = machine.advance(progress)? {
            if machine.committed && !committed_before {
                crate::warmup::prefault_all(driver.warmup_paths());
            }
            return Ok(outcome);
        }
        if machine.committed && !committed_before {
            crate::warmup::prefault_all(driver.warmup_paths());
        }
        if machine.step != step {
            if machine.step == Step::Drain || machine.step == Step::Retire {
                deadline = budgets.drain.map(|budget| Instant::now() + budget);
            } else if step == Step::Drain {
                deadline = Some(Instant::now() + budgets.setup);
            }
        } else {
            let wait = deadline.map_or(crate::readiness::POLL_INTERVAL, |limit| {
                crate::readiness::POLL_INTERVAL.min(limit.saturating_duration_since(Instant::now()))
            });
            #[expect(
                clippy::disallowed_methods,
                reason = "The replacement driver returned Pending without a wake source; retry within the phase budget."
            )]
            std::thread::sleep(wait);
        }
    }
}

/// Execute the same transition machine on an async runtime.
#[cfg(feature = "async")]
pub async fn run_async<D: AsyncDriver>(
    _ownership: &ProcessLock,
    mode: Mode,
    budgets: Budgets,
    driver: &mut D,
) -> Result<Outcome, Failure<D::Error>> {
    let mut machine = Machine::new(mode);
    let mut deadline = Some(tokio::time::Instant::now() + budgets.setup);
    loop {
        if deadline.is_some_and(|limit| tokio::time::Instant::now() >= limit) {
            return Err(machine.failure(Reason::Deadline));
        }
        let step = machine.step;
        // Commit is an atomic publication operation. Interrupting it could leave its
        // result unknown, so its adapter must not await after publication.
        let result = if let Some(limit) = deadline.filter(|_| step != Step::Commit) {
            tokio::time::timeout_at(limit, driver.perform(step, deadline))
                .await
                .map_err(|_| machine.failure(Reason::Deadline))?
        } else {
            driver.perform(step, deadline).await
        };
        let progress = result.map_err(|e| machine.failure(Reason::Adapter(e)))?;
        if step != Step::Commit
            && deadline.is_some_and(|limit| tokio::time::Instant::now() >= limit)
        {
            return Err(machine.failure(Reason::Deadline));
        }
        let committed_before = machine.committed;
        if let Some(outcome) = machine.advance(progress)? {
            if machine.committed && !committed_before {
                prefault_async(driver.warmup_paths()).await;
            }
            return Ok(outcome);
        }
        if machine.committed && !committed_before {
            prefault_async(driver.warmup_paths()).await;
        }
        if machine.step != step {
            if machine.step == Step::Drain || machine.step == Step::Retire {
                deadline = budgets
                    .drain
                    .map(|budget| tokio::time::Instant::now() + budget);
            } else if step == Step::Drain {
                deadline = Some(tokio::time::Instant::now() + budgets.setup);
            }
        } else {
            let wake = tokio::time::Instant::now() + crate::readiness::POLL_INTERVAL;
            #[expect(
                clippy::disallowed_methods,
                reason = "The replacement driver returned Pending without a wake source; retry within the phase budget."
            )]
            tokio::time::sleep_until(deadline.map_or(wake, |limit| limit.min(wake))).await;
        }
    }
}

// Filesystem reads must not occupy a Tokio worker after a successful commit.
// The paths are owned and bounded by the caller; failed warming is non-fatal.
#[cfg(feature = "async")]
async fn prefault_async(paths: Vec<std::path::PathBuf>) {
    if !paths.is_empty() {
        let _ = tokio::task::spawn_blocking(move || crate::warmup::prefault_all(paths)).await;
    }
}

/// Properties of the transition machine for every adapter answer, not only the
/// sequences a test author thought of. Each one restates a promise from this
/// module's documentation.
#[cfg(kani)]
mod proofs {
    use super::*;

    /// Long enough to reach `Complete` in either mode with a `Pending` at every
    /// step that allows one: Recheck, Prepare, Drain, Start, Verify, Validate,
    /// Commit and Retire, each preceded by a retry.
    const ANSWERS: usize = 17;

    #[kani::proof]
    #[kani::unwind(18)]
    fn replacement_never_claims_what_it_has_not_done() {
        let mode: Mode = kani::any();
        let mut machine = Machine::new(mode);
        let mut drained = false;
        let mut verified = false;
        let mut validated = false;

        for _ in 0..ANSWERS {
            let step = machine.step;
            let was_committed = machine.committed;
            let progress: Progress = kani::any();
            let result = machine.advance::<()>(progress);

            // Selection is never withdrawn once it has been made.
            assert!(!was_committed || machine.committed);

            if progress == Progress::Done {
                drained |= step == Step::Drain;
                verified |= step == Step::Verify;
                validated |= step == Step::Validate;
            }

            // Overlap keeps the incumbent serving: it never drains it.
            if mode == Mode::Overlap {
                assert!(machine.step != Step::Drain);
            }
            // Exclusive releases the incumbent before starting the candidate.
            if mode == Mode::Exclusive && machine.step == Step::Start {
                assert!(drained);
            }
            // A candidate is selected only after a live proof and revalidation.
            if machine.committed {
                assert!(verified && validated);
            }
            // The incumbent is retired only after the new selection is committed.
            if machine.step == Step::Retire {
                assert!(machine.committed);
            }

            match result {
                Ok(Some(Outcome::Unchanged)) => {
                    // "Unchanged" is only ever the answer to Recheck, before
                    // anything was replaced.
                    assert!(step == Step::Recheck && !machine.committed);
                    return;
                }
                Ok(Some(Outcome::Complete | Outcome::RetirementPending)) => {
                    // Success is never reported for a candidate that was not
                    // selected.
                    assert!(machine.committed);
                    return;
                }
                Ok(None) => {}
                Err(failure) => {
                    // A failure after commit says so, so the caller does not
                    // treat an authoritative selection as rolled back.
                    assert!(failure.committed == machine.committed);
                    assert!(failure.step == step);
                    return;
                }
            }
        }
    }
}

#[cfg(test)]
mod error_tests {
    use super::*;

    #[test]
    fn an_adapter_failure_preserves_its_source() {
        let failure = Failure {
            step: Step::Start,
            committed: false,
            reason: Reason::Adapter(std::io::Error::other("failed to spawn")),
        };
        assert_eq!(
            std::error::Error::source(&failure).unwrap().to_string(),
            "failed to spawn"
        );
        for reason in [Reason::Deadline, Reason::InvalidProgress] {
            let failure: Failure<std::io::Error> = Failure {
                step: Step::Start,
                committed: false,
                reason,
            };
            assert!(std::error::Error::source(&failure).is_none());
        }
    }
}