use std::fmt;
use std::future::Future;
use std::time::Duration;
use rustev_contract::time::DurationMs;
use tokio::time::Instant;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Budget {
deadline: Instant,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Expired;
impl Budget {
pub fn start(remaining: DurationMs) -> Self {
Self::from_ms(remaining.0)
}
pub fn from_ms(ms: u64) -> Self {
Self::from_start(Instant::now(), DurationMs(ms))
}
pub(crate) fn from_start(start: Instant, remaining: DurationMs) -> Self {
Budget {
deadline: start + Duration::from_millis(remaining.0),
}
}
pub fn deadline(&self) -> Instant {
self.deadline
}
pub fn min(self, other: Budget) -> Budget {
Budget {
deadline: self.deadline.min(other.deadline),
}
}
pub fn remaining(&self) -> Duration {
self.deadline.saturating_duration_since(Instant::now())
}
pub fn remaining_ms(&self) -> u64 {
u64::try_from(self.remaining().as_millis()).unwrap_or(u64::MAX)
}
pub fn expired(&self) -> bool {
Instant::now() >= self.deadline
}
pub fn told_ms(&self, transit_margin_ms: u64) -> u64 {
self.remaining_ms().saturating_sub(transit_margin_ms)
}
pub fn phase_deadline(&self, limit: Option<Duration>) -> Instant {
match limit {
Some(l) => self.deadline.min(Instant::now() + l),
None => self.deadline,
}
}
pub async fn run<F: Future>(
&self,
limit: Option<Duration>,
fut: F,
) -> Result<F::Output, Expired> {
tokio::time::timeout_at(self.phase_deadline(limit), fut)
.await
.map_err(|_| Expired)
}
pub async fn expiry(&self) {
tokio::time::sleep_until(self.deadline).await
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct TimeoutBeyondBudget {
pub transport_timeout_ms: u64,
pub budget_ms: u64,
}
impl fmt::Display for TimeoutBeyondBudget {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(
f,
"a transport timeout of {} ms is longer than the {} ms budget it serves",
self.transport_timeout_ms, self.budget_ms
)
}
}
impl std::error::Error for TimeoutBeyondBudget {}
pub fn check_transport_timeout(
transport_timeout_ms: Option<u64>,
budget_ms: u64,
) -> Result<(), TimeoutBeyondBudget> {
match transport_timeout_ms {
Some(t) if t > budget_ms => Err(TimeoutBeyondBudget {
transport_timeout_ms: t,
budget_ms,
}),
_ => Ok(()),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test(flavor = "current_thread", start_paused = true)]
async fn batch_preparation_does_not_extend_later_member_deadlines() {
let start = Instant::now();
let first = Budget::from_start(start, DurationMs(100));
tokio::time::advance(Duration::from_millis(40)).await;
let second = Budget::from_start(start, DurationMs(100));
let longer = Budget::from_start(start, DurationMs(150));
assert_eq!(second.remaining_ms(), 60);
first.expiry().await;
assert!(first.expired());
assert!(second.expired());
assert!(!longer.expired());
assert_eq!(longer.remaining_ms(), 50);
}
#[tokio::test(flavor = "current_thread", start_paused = true)]
async fn a_phase_limit_never_outlives_the_budget() {
let b = Budget::from_ms(100);
let t0 = Instant::now();
assert_eq!(
b.phase_deadline(Some(Duration::from_secs(60))),
b.deadline()
);
assert_eq!(
b.phase_deadline(Some(Duration::from_millis(10))),
t0 + Duration::from_millis(10)
);
let r = b
.run(Some(Duration::from_secs(60)), std::future::pending::<()>())
.await;
assert_eq!(r, Err(Expired));
assert_eq!(Instant::now(), t0 + Duration::from_millis(100));
assert!(b.expired());
assert_eq!(b.told_ms(5), 0);
}
#[tokio::test(flavor = "current_thread", start_paused = true)]
async fn the_remote_side_is_told_the_budget_minus_the_margin() {
let b = Budget::from_ms(1_000);
assert_eq!(b.told_ms(50), 950);
assert_eq!(b.told_ms(5_000), 0);
}
#[test]
fn a_transport_timeout_longer_than_the_budget_is_refused() {
assert!(check_transport_timeout(None, 10).is_ok());
assert!(check_transport_timeout(Some(10), 10).is_ok());
assert_eq!(
check_transport_timeout(Some(11), 10),
Err(TimeoutBeyondBudget {
transport_timeout_ms: 11,
budget_ms: 10
})
);
}
}