use core::time::Duration;
use crate::rate_limit::RateLimit;
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ActionUpdate<E> {
Running,
Success,
Failed(E),
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct PollContext {
observation: u32,
progress: u8,
rate_limit: Option<RateLimit>,
}
impl PollContext {
#[must_use]
pub const fn observation(self) -> u32 {
self.observation
}
#[must_use]
pub const fn progress(self) -> u8 {
self.progress
}
#[must_use]
pub const fn rate_limit(self) -> Option<RateLimit> {
self.rate_limit
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum PollDecision {
Delay(Duration),
Cancel,
Timeout,
}
pub trait PollPolicy {
type Error;
fn decide(&mut self, context: PollContext) -> Result<PollDecision, Self::Error>;
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ActionPollStep<E> {
Delay(Duration),
Complete,
Failed(E),
Cancelled,
TimedOut,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ActionPollError<E> {
InvalidProgress,
ProgressRegressed,
ZeroDelay,
ObservationOverflow,
Terminal,
Policy(E),
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct ActionPoller {
observations: u32,
last_progress: Option<u8>,
terminal: bool,
}
impl ActionPoller {
#[must_use]
pub const fn new() -> Self {
Self {
observations: 0,
last_progress: None,
terminal: false,
}
}
#[must_use]
pub const fn observations(self) -> u32 {
self.observations
}
#[must_use]
pub const fn is_terminal(self) -> bool {
self.terminal
}
pub fn observe<E, P>(
&mut self,
update: ActionUpdate<E>,
progress: u8,
rate_limit: Option<RateLimit>,
policy: &mut P,
) -> Result<ActionPollStep<E>, ActionPollError<P::Error>>
where
P: PollPolicy,
{
if self.terminal {
return Err(ActionPollError::Terminal);
}
let observations = self
.observations
.checked_add(1)
.ok_or(ActionPollError::ObservationOverflow)?;
let step = match update {
ActionUpdate::Success => ActionPollStep::Complete,
ActionUpdate::Failed(error) => ActionPollStep::Failed(error),
ActionUpdate::Running => {
if progress > 100 {
return Err(ActionPollError::InvalidProgress);
}
if self.last_progress.is_some_and(|last| progress < last) {
return Err(ActionPollError::ProgressRegressed);
}
let context = PollContext {
observation: observations,
progress,
rate_limit,
};
match policy.decide(context).map_err(ActionPollError::Policy)? {
PollDecision::Delay(delay) if delay.is_zero() => {
return Err(ActionPollError::ZeroDelay);
}
PollDecision::Delay(delay) => ActionPollStep::Delay(delay),
PollDecision::Cancel => ActionPollStep::Cancelled,
PollDecision::Timeout => ActionPollStep::TimedOut,
}
}
};
self.observations = observations;
self.last_progress = Some(progress);
self.terminal = !matches!(step, ActionPollStep::Delay(_));
Ok(step)
}
}
impl Default for ActionPoller {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests;