use parking_lot::Mutex;
use std::time::{SystemTime, UNIX_EPOCH};
/// per-UTC-day counter for newly admitted PDS sources.
///
/// the in-memory state is initialised from the database on startup (see
/// [`crate::db::load_pds_daily_adds`]). a reservation remains held until the
/// matching source write succeeds, then gets refunded if that write fails.
pub(crate) struct PdsDailyLimit {
limit: Option<u64>,
state: Mutex<DailyCount>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct DailyCount {
day: u64,
count: u64,
}
/// a slot reserved until its source record is durably written.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum PdsDailyReservation {
Unlimited,
Limited {
previous: DailyCount,
persisted: DailyCount,
},
}
impl PdsDailyReservation {
pub(crate) fn persisted(self) -> Option<(u64, u64)> {
match self {
Self::Unlimited => None,
Self::Limited { persisted, .. } => Some((persisted.day, persisted.count)),
}
}
}
impl PdsDailyLimit {
/// construct from the previously-persisted `(day, count)` pair. a stale day
/// starts a fresh counter.
pub(crate) fn new(limit: Option<u64>, stored: Option<(u64, u64)>) -> Self {
Self::for_day(limit, stored, utc_day())
}
fn for_day(limit: Option<u64>, stored: Option<(u64, u64)>, today: u64) -> Self {
let count = stored
.filter(|(day, _)| *day == today)
.map(|(_, count)| count)
.unwrap_or(0);
Self {
limit,
state: Mutex::new(DailyCount { day: today, count }),
}
}
/// reserve one slot. the caller must persist [`PdsDailyReservation::persisted`]
/// before reporting admission, or call [`Self::refund`] on failure.
pub(crate) fn reserve(&self) -> Option<PdsDailyReservation> {
self.reserve_for_day(utc_day())
}
fn reserve_for_day(&self, today: u64) -> Option<PdsDailyReservation> {
let Some(limit) = self.limit else {
return Some(PdsDailyReservation::Unlimited);
};
let mut state = self.state.lock();
if state.day != today {
*state = DailyCount {
day: today,
count: 0,
};
}
if state.count >= limit {
return None;
}
let previous = *state;
state.count += 1;
Some(PdsDailyReservation::Limited {
previous,
persisted: *state,
})
}
/// release an unpersisted reservation.
///
/// admission serializes reserve/refund, so equality proves no later
/// reservation has superseded this one.
pub(crate) fn refund(&self, reservation: PdsDailyReservation) {
let PdsDailyReservation::Limited {
previous,
persisted,
} = reservation
else {
return;
};
let mut state = self.state.lock();
if *state == persisted {
*state = previous;
}
}
}
fn utc_day() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs()
/ 86400
}
#[cfg(test)]
mod tests {
use super::*;
fn persisted(reservation: Option<PdsDailyReservation>) -> Option<(u64, u64)> {
reservation.and_then(PdsDailyReservation::persisted)
}
#[test]
fn resets_stale_persisted_usage_at_rollover() {
let limit = PdsDailyLimit::for_day(Some(2), Some((7, 2)), 8);
assert_eq!(persisted(limit.reserve_for_day(8)), Some((8, 1)));
assert_eq!(persisted(limit.reserve_for_day(8)), Some((8, 2)));
assert!(limit.reserve_for_day(8).is_none());
}
#[test]
fn refunds_an_unpersisted_reservation() {
let limit = PdsDailyLimit::for_day(Some(1), None, 7);
let reservation = limit
.reserve_for_day(7)
.expect("first reservation must succeed");
assert!(limit.reserve_for_day(7).is_none());
limit.refund(reservation);
assert_eq!(persisted(limit.reserve_for_day(7)), Some((7, 1)));
}
}