use std::time::Duration;
use ozzy_proto::RequestId;
use super::prefix_valid;
use crate::{Digest, OpNumber, Prefix, Scope};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Probe {
pub scope: Scope,
pub request_id: RequestId,
pub tail: Prefix,
pub available: OpNumber,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ProbeTiming {
pub initial: Duration,
pub maximum: Duration,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
pub enum ProbeError {
#[error("invalid replica probe parameters")]
Invalid,
#[error("replica probe scope changed without invalidation")]
Scope,
#[error("replica probe clock regressed")]
Clock,
#[error("replica probe deadline exhausted")]
TimeExhausted,
#[error("replica probe request identities exhausted")]
IdsExhausted,
}
#[derive(Debug)]
pub struct ProbeScheduler {
timing: ProbeTiming,
next_id: u128,
now: Duration,
due: Duration,
delay: Duration,
pending: Option<Probe>,
advertised: OpNumber,
}
impl ProbeScheduler {
pub fn new(first: RequestId, timing: ProbeTiming, now: Duration) -> Result<Self, ProbeError> {
let next_id = u128::from_be_bytes(*first.as_bytes());
if next_id == 0 || timing.initial.is_zero() || timing.maximum < timing.initial {
return Err(ProbeError::Invalid);
}
now.checked_add(timing.maximum)
.ok_or(ProbeError::TimeExhausted)?;
Ok(Self {
timing,
next_id,
now,
due: now,
delay: timing.initial,
pending: None,
advertised: OpNumber(0),
})
}
pub const fn pending(&self) -> Option<Probe> {
self.pending
}
pub fn poll(
&mut self,
scope: Scope,
tail: Prefix,
available: OpNumber,
now: Duration,
) -> Result<Option<Probe>, ProbeError> {
self.check_time(now)?;
if scope.group_id.as_bytes() == &[0; 16]
|| scope.configuration_digest == Digest::ZERO
|| !prefix_valid(tail)
|| tail.op.0 == u64::MAX
|| available < tail.op
|| available.0 == u64::MAX
{
return Err(ProbeError::Invalid);
}
if self.pending.is_some_and(|request| request.scope != scope) {
return Err(ProbeError::Scope);
}
let new_work = self.pending.is_none() && available > self.advertised;
if now < self.due && !new_work {
self.now = now;
return Ok(None);
}
let (request, next_id) = if let Some(request) = self.pending {
(request, self.next_id)
} else {
let next_id = self
.next_id
.checked_add(1)
.ok_or(ProbeError::IdsExhausted)?;
(
Probe {
scope,
tail,
available,
request_id: RequestId::from_bytes(self.next_id.to_be_bytes()),
},
next_id,
)
};
let due = now
.checked_add(self.delay)
.ok_or(ProbeError::TimeExhausted)?;
self.now = now;
self.due = due;
self.delay = self.delay.saturating_mul(2).min(self.timing.maximum);
self.next_id = next_id;
self.pending = Some(request);
self.advertised = request.available;
Ok(Some(request))
}
pub fn complete(
&mut self,
scope: Scope,
id: RequestId,
now: Duration,
) -> Result<bool, ProbeError> {
let Some(due) = self.completion_deadline(scope, id, now)? else {
self.now = now;
return Ok(false);
};
self.now = now;
self.due = due;
self.delay = self.timing.initial;
self.pending = None;
Ok(true)
}
pub(super) fn completion_deadline(
&self,
scope: Scope,
id: RequestId,
now: Duration,
) -> Result<Option<Duration>, ProbeError> {
self.check_time(now)?;
if self
.pending
.is_none_or(|request| request.scope != scope || request.request_id != id)
{
return Ok(None);
}
now.checked_add(self.timing.initial)
.map(Some)
.ok_or(ProbeError::TimeExhausted)
}
pub fn invalidate(&mut self, now: Duration) -> Result<(), ProbeError> {
self.check_time(now)?;
self.now = now;
self.due = now;
self.delay = self.timing.initial;
self.pending = None;
Ok(())
}
pub(super) fn expedite(&mut self, now: Duration) -> Result<(), ProbeError> {
self.check_time(now)?;
self.now = now;
self.due = self.due.min(now);
self.delay = self.timing.initial;
Ok(())
}
fn check_time(&self, now: Duration) -> Result<(), ProbeError> {
if now < self.now {
Err(ProbeError::Clock)
} else {
Ok(())
}
}
}