use crate::ProcessLock;
use std::time::{Duration, Instant};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Mode {
Overlap,
Exclusive,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Step {
Recheck,
Prepare,
Drain,
Start,
Verify,
Validate,
Commit,
Retire,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Progress {
Done,
Pending,
Unchanged,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Outcome {
Unchanged,
Complete,
RetirementPending,
}
#[derive(Debug)]
pub struct Failure<E> {
pub step: Step,
pub committed: bool,
pub reason: Reason<E>,
}
#[derive(Debug)]
pub enum Reason<E> {
Deadline,
Adapter(E),
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> {}
#[derive(Clone, Copy, Debug)]
pub struct Budgets {
pub setup: Duration,
pub drain: Option<Duration>,
}
pub trait Driver {
type Error;
fn perform(&mut self, step: Step, deadline: Option<Instant>) -> Result<Progress, Self::Error>;
}
#[cfg(feature = "async")]
pub trait AsyncDriver {
type Error;
fn perform(
&mut self,
step: Step,
deadline: Option<tokio::time::Instant>,
) -> impl std::future::Future<Output = Result<Progress, Self::Error>> + Send;
}
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)
}
}
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)))?;
if step != Step::Commit && deadline.is_some_and(|limit| Instant::now() >= limit) {
return Err(machine.failure(Reason::Deadline));
}
if let Some(outcome) = machine.advance(progress)? {
return Ok(outcome);
}
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()))
});
std::thread::sleep(wait);
}
}
}
#[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;
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));
}
if let Some(outcome) = machine.advance(progress)? {
return Ok(outcome);
}
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;
tokio::time::sleep_until(deadline.map_or(wake, |limit| limit.min(wake))).await;
}
}
}