use std::time::Duration;
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[non_exhaustive]
pub enum Denial {
QuotaExceeded {
limit: u64,
retry_after: Duration,
},
StorageCapacity {
retry_after: Option<Duration>,
},
}
impl Denial {
pub const fn retry_after(&self) -> Option<Duration> {
match self {
Self::QuotaExceeded { retry_after, .. } => Some(*retry_after),
Self::StorageCapacity { retry_after } => *retry_after,
}
}
pub const fn retry_after_seconds(&self) -> Option<u64> {
match self.retry_after() {
Some(duration) => Some(ceil_seconds(duration)),
None => None,
}
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
struct Allowance {
limit: u64,
remaining: u64,
reset_after: Duration,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum Outcome {
Allowed(Allowance),
Denied(Denial),
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct Decision {
outcome: Outcome,
}
impl Decision {
pub const fn allowed(limit: u64, remaining: u64, reset_after: Duration) -> Self {
Self {
outcome: Outcome::Allowed(Allowance {
limit,
remaining,
reset_after,
}),
}
}
pub const fn denied(denial: Denial) -> Self {
Self {
outcome: Outcome::Denied(denial),
}
}
pub const fn is_allowed(&self) -> bool {
matches!(self.outcome, Outcome::Allowed(_))
}
pub const fn is_denied(&self) -> bool {
!self.is_allowed()
}
pub const fn limit(&self) -> Option<u64> {
match self.outcome {
Outcome::Allowed(allowance) => Some(allowance.limit),
Outcome::Denied(Denial::QuotaExceeded { limit, .. }) => Some(limit),
Outcome::Denied(Denial::StorageCapacity { .. }) => None,
}
}
pub const fn remaining(&self) -> Option<u64> {
match self.outcome {
Outcome::Allowed(allowance) => Some(allowance.remaining),
Outcome::Denied(_) => None,
}
}
pub const fn reset_after(&self) -> Option<Duration> {
match self.outcome {
Outcome::Allowed(allowance) => Some(allowance.reset_after),
Outcome::Denied(_) => None,
}
}
pub const fn retry_after(&self) -> Option<Duration> {
match self.outcome {
Outcome::Allowed(_) => None,
Outcome::Denied(denial) => denial.retry_after(),
}
}
pub const fn retry_after_seconds(&self) -> Option<u64> {
match self.outcome {
Outcome::Allowed(_) => None,
Outcome::Denied(denial) => denial.retry_after_seconds(),
}
}
pub const fn denial(&self) -> Option<&Denial> {
match &self.outcome {
Outcome::Allowed(_) => None,
Outcome::Denied(denial) => Some(denial),
}
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum BatchDecision {
Allowed(Vec<Decision>),
Denied {
index: usize,
denial: Denial,
},
}
impl BatchDecision {
pub fn try_into_single_decision(self) -> Result<Decision, Self> {
match self {
Self::Allowed(decisions) if matches!(decisions.as_slice(), [decision] if decision.is_allowed()) => {
Ok(decisions[0])
}
Self::Denied { index: 0, denial } => Ok(Decision::denied(denial)),
batch => Err(batch),
}
}
}
const fn ceil_seconds(duration: Duration) -> u64 {
let seconds = duration.as_secs();
if duration.subsec_nanos() == 0 {
seconds
} else {
seconds.saturating_add(1)
}
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use super::{BatchDecision, Decision, Denial};
#[test]
fn allowed_decision_exposes_remaining_and_reset() {
let decision = Decision::allowed(8, 7, Duration::from_millis(59_999));
assert!(decision.is_allowed());
assert!(!decision.is_denied());
assert_eq!(decision.limit(), Some(8));
assert_eq!(decision.remaining(), Some(7));
assert_eq!(decision.reset_after(), Some(Duration::from_millis(59_999)));
assert_eq!(decision.retry_after(), None);
assert_eq!(decision.denial(), None);
}
#[test]
fn quota_denial_exposes_exact_and_ceiling_retry_duration() {
let denial = Denial::QuotaExceeded {
limit: 8,
retry_after: Duration::from_millis(1_001),
};
let decision = Decision::denied(denial);
assert!(decision.is_denied());
assert_eq!(decision.limit(), Some(8));
assert_eq!(decision.remaining(), None);
assert_eq!(decision.reset_after(), None);
assert_eq!(decision.retry_after(), Some(Duration::from_millis(1_001)));
assert_eq!(decision.retry_after_seconds(), Some(2));
assert_eq!(decision.denial(), Some(&denial));
assert!(matches!(
denial,
Denial::QuotaExceeded {
limit: 8,
retry_after
} if retry_after == Duration::from_millis(1_001)
));
}
#[test]
fn retry_after_seconds_preserves_exact_seconds() {
let denial = Denial::QuotaExceeded {
limit: 1,
retry_after: Duration::from_secs(3),
};
assert_eq!(denial.retry_after_seconds(), Some(3));
}
#[test]
fn retry_after_seconds_saturates_without_losing_exact_duration() {
let duration = Duration::new(u64::MAX, 1);
let denial = Denial::QuotaExceeded {
limit: 1,
retry_after: duration,
};
assert_eq!(denial.retry_after(), Some(duration));
assert_eq!(denial.retry_after_seconds(), Some(u64::MAX));
}
#[test]
fn storage_capacity_retry_can_be_unknown() {
let denial = Denial::StorageCapacity { retry_after: None };
let decision = Decision::denied(denial);
assert!(matches!(
denial,
Denial::StorageCapacity { retry_after: None }
));
assert_eq!(denial.retry_after(), None);
assert_eq!(denial.retry_after_seconds(), None);
assert_eq!(decision.limit(), None);
assert_eq!(decision.retry_after(), None);
}
#[test]
fn batch_of_one_converts_to_a_single_decision() {
let allowed = Decision::allowed(8, 7, Duration::from_secs(60));
let denied = Denial::QuotaExceeded {
limit: 8,
retry_after: Duration::from_secs(60),
};
assert_eq!(
BatchDecision::Allowed(vec![allowed]).try_into_single_decision(),
Ok(allowed)
);
assert_eq!(
BatchDecision::Denied {
index: 0,
denial: denied
}
.try_into_single_decision(),
Ok(Decision::denied(denied))
);
}
#[test]
fn malformed_batch_of_one_is_rejected() {
let decision = Decision::allowed(8, 7, Duration::from_secs(60));
let denial = Denial::QuotaExceeded {
limit: 8,
retry_after: Duration::from_secs(60),
};
assert!(
BatchDecision::Allowed(Vec::new())
.try_into_single_decision()
.is_err()
);
assert!(
BatchDecision::Allowed(vec![decision, decision])
.try_into_single_decision()
.is_err()
);
assert!(
BatchDecision::Allowed(vec![Decision::denied(denial)])
.try_into_single_decision()
.is_err()
);
assert!(
BatchDecision::Denied { index: 1, denial }
.try_into_single_decision()
.is_err()
);
}
}