use std::{
cell::Cell,
fs,
path::PathBuf,
thread,
time::{Duration, SystemTime, UNIX_EPOCH},
};
pub const MIN_INTERVAL: Duration = Duration::from_secs(5);
const LAST_REQUEST_FILE: &str = "last-request";
pub trait Timer {
fn now(&self) -> SystemTime;
fn wait(&self, duration: Duration);
}
#[derive(Debug, Default, Clone, Copy)]
pub struct SystemTimer;
impl Timer for SystemTimer {
fn now(&self) -> SystemTime {
SystemTime::now()
}
fn wait(&self, duration: Duration) {
thread::sleep(duration);
}
}
#[derive(Debug)]
pub struct Throttle<T = SystemTimer> {
path: PathBuf,
timer: T,
interval: Duration,
last: Cell<Option<SystemTime>>,
}
impl Throttle {
#[must_use]
pub fn new(state_dir: impl Into<PathBuf>) -> Self {
Self::with_timer(state_dir, SystemTimer, MIN_INTERVAL)
}
}
impl<T: Timer> Throttle<T> {
#[must_use]
pub fn with_timer(state_dir: impl Into<PathBuf>, timer: T, interval: Duration) -> Self {
Self {
path: state_dir.into().join(LAST_REQUEST_FILE),
timer,
interval,
last: Cell::new(None),
}
}
pub fn acquire(&self) {
let now = self.timer.now();
let wait = self.remaining(now);
if !wait.is_zero() {
self.timer.wait(wait);
}
self.record(now.checked_add(wait).unwrap_or(now));
}
fn remaining(&self, now: SystemTime) -> Duration {
self.last_request()
.and_then(|last| last.checked_add(self.interval))
.and_then(|earliest| earliest.duration_since(now).ok())
.map_or(Duration::ZERO, |wait| wait.min(self.interval))
}
fn last_request(&self) -> Option<SystemTime> {
let persisted = fs::read_to_string(&self.path)
.ok()
.and_then(|text| text.trim().parse::<u64>().ok())
.and_then(|millis| UNIX_EPOCH.checked_add(Duration::from_millis(millis)));
match (self.last.get(), persisted) {
(Some(memory), Some(disk)) => Some(memory.max(disk)),
(memory, disk) => memory.or(disk),
}
}
fn record(&self, at: SystemTime) {
self.last.set(Some(at));
let Ok(since_epoch) = at.duration_since(UNIX_EPOCH) else {
return;
};
let Ok(millis) = u64::try_from(since_epoch.as_millis()) else {
return;
};
let Some(parent) = self.path.parent() else {
return;
};
if fs::create_dir_all(parent).is_ok() {
let _ = fs::write(&self.path, millis.to_string());
}
}
}
#[cfg(test)]
pub(crate) mod fake {
use super::{Duration, SystemTime, Timer};
use std::cell::{Cell, RefCell};
#[derive(Debug)]
pub(crate) struct FakeTimer {
now: Cell<SystemTime>,
waits: RefCell<Vec<Duration>>,
}
impl FakeTimer {
pub(crate) fn new(now: SystemTime) -> Self {
Self {
now: Cell::new(now),
waits: RefCell::new(Vec::new()),
}
}
pub(crate) fn advance(&self, by: Duration) {
self.now.set(self.now.get() + by);
}
pub(crate) fn waits(&self) -> Vec<Duration> {
self.waits.borrow().clone()
}
}
impl Timer for FakeTimer {
fn now(&self) -> SystemTime {
self.now.get()
}
fn wait(&self, duration: Duration) {
self.waits.borrow_mut().push(duration);
self.advance(duration);
}
}
}
#[cfg(test)]
mod tests {
use super::{fake::FakeTimer, *};
const INTERVAL: Duration = Duration::from_secs(5);
fn at(seconds: u64) -> SystemTime {
UNIX_EPOCH + Duration::from_secs(seconds)
}
fn throttle(state_dir: impl Into<PathBuf>, now: u64) -> Throttle<FakeTimer> {
Throttle::with_timer(state_dir, FakeTimer::new(at(now)), INTERVAL)
}
#[test]
fn the_first_request_is_not_delayed() {
let dir = tempfile::tempdir().expect("temp dir");
let throttle = throttle(dir.path(), 1000);
throttle.acquire();
assert!(throttle.timer.waits().is_empty());
}
#[test]
fn a_burst_is_spaced_out() {
let dir = tempfile::tempdir().expect("temp dir");
let throttle = throttle(dir.path(), 1000);
throttle.acquire();
throttle.acquire();
throttle.acquire();
assert_eq!(throttle.timer.waits(), [INTERVAL, INTERVAL]);
}
#[test]
fn waiting_out_the_interval_costs_nothing() {
let dir = tempfile::tempdir().expect("temp dir");
let throttle = throttle(dir.path(), 1000);
throttle.acquire();
throttle.timer.advance(INTERVAL);
throttle.acquire();
assert!(throttle.timer.waits().is_empty());
}
#[test]
fn the_gap_is_honoured_across_invocations() {
let dir = tempfile::tempdir().expect("temp dir");
throttle(dir.path(), 1000).acquire();
let next = throttle(dir.path(), 1002);
next.acquire();
assert_eq!(next.timer.waits(), [Duration::from_secs(3)]);
}
#[test]
fn an_unwritable_state_directory_still_throttles_in_process() {
let throttle = throttle("/proc/definitely-not-writable", 1000);
throttle.acquire();
throttle.acquire();
assert_eq!(throttle.timer.waits(), [INTERVAL]);
}
#[test]
fn a_timestamp_from_the_future_waits_at_most_one_interval() {
let dir = tempfile::tempdir().expect("temp dir");
fs::write(dir.path().join(LAST_REQUEST_FILE), "99000000000000").expect("seed timestamp");
let throttle = throttle(dir.path(), 1000);
throttle.acquire();
assert_eq!(throttle.timer.waits(), [INTERVAL]);
}
#[test]
fn an_unreadable_record_does_not_block_the_first_request() {
let dir = tempfile::tempdir().expect("temp dir");
fs::write(dir.path().join(LAST_REQUEST_FILE), "not a timestamp").expect("seed timestamp");
let throttle = throttle(dir.path(), 1000);
throttle.acquire();
assert!(throttle.timer.waits().is_empty());
}
}