1#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
9pub struct Limits {
10 pub max_units: Option<u64>,
12 pub max_spend_micros: Option<u64>,
14}
15
16#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
18pub struct Usage {
19 pub units: u64,
21 pub spend_micros: u64,
23}
24
25#[derive(Debug, Clone, Copy, PartialEq, Eq)]
27pub enum LimitDecision {
28 Allowed(Usage),
30 Rejected(LimitReason),
32}
33
34impl LimitDecision {
35 #[must_use]
37 pub const fn is_allowed(self) -> bool {
38 matches!(self, Self::Allowed(_))
39 }
40}
41
42#[derive(Debug, Clone, Copy, PartialEq, Eq)]
44pub enum LimitReason {
45 QuotaExceeded,
47 SpendCapExceeded,
49 ArithmeticOverflow,
51}
52
53#[must_use]
58pub fn assess_limits(
59 current: Usage,
60 requested_units: u64,
61 unit_price_micros: u64,
62 limits: Limits,
63) -> LimitDecision {
64 let Some(units) = current.units.checked_add(requested_units) else {
65 return LimitDecision::Rejected(LimitReason::ArithmeticOverflow);
66 };
67 if limits.max_units.is_some_and(|limit| units > limit) {
68 return LimitDecision::Rejected(LimitReason::QuotaExceeded);
69 }
70
71 let Some(incremental_spend) = requested_units.checked_mul(unit_price_micros) else {
72 return LimitDecision::Rejected(LimitReason::ArithmeticOverflow);
73 };
74 let Some(spend_micros) = current.spend_micros.checked_add(incremental_spend) else {
75 return LimitDecision::Rejected(LimitReason::ArithmeticOverflow);
76 };
77 if limits
78 .max_spend_micros
79 .is_some_and(|limit| spend_micros > limit)
80 {
81 return LimitDecision::Rejected(LimitReason::SpendCapExceeded);
82 }
83
84 LimitDecision::Allowed(Usage {
85 units,
86 spend_micros,
87 })
88}
89
90#[cfg(test)]
91mod tests {
92 use super::*;
93
94 #[test]
95 fn exact_boundaries_are_allowed() {
96 let decision = assess_limits(
97 Usage {
98 units: 8,
99 spend_micros: 800,
100 },
101 2,
102 100,
103 Limits {
104 max_units: Some(10),
105 max_spend_micros: Some(1_000),
106 },
107 );
108 assert_eq!(
109 decision,
110 LimitDecision::Allowed(Usage {
111 units: 10,
112 spend_micros: 1_000
113 })
114 );
115 }
116
117 #[test]
118 fn quota_is_checked_before_spend() {
119 let decision = assess_limits(
120 Usage::default(),
121 11,
122 1_000,
123 Limits {
124 max_units: Some(10),
125 max_spend_micros: Some(1),
126 },
127 );
128 assert_eq!(
129 decision,
130 LimitDecision::Rejected(LimitReason::QuotaExceeded)
131 );
132 }
133
134 #[test]
135 fn spend_cap_rejects_the_first_unit_over() {
136 let decision = assess_limits(
137 Usage {
138 units: 2,
139 spend_micros: 200,
140 },
141 1,
142 101,
143 Limits {
144 max_units: None,
145 max_spend_micros: Some(300),
146 },
147 );
148 assert_eq!(
149 decision,
150 LimitDecision::Rejected(LimitReason::SpendCapExceeded)
151 );
152 }
153
154 #[test]
155 fn arithmetic_overflow_never_wraps_into_an_allowance() {
156 assert_eq!(
157 assess_limits(Usage::default(), u64::MAX, 2, Limits::default()),
158 LimitDecision::Rejected(LimitReason::ArithmeticOverflow)
159 );
160 assert_eq!(
161 assess_limits(
162 Usage {
163 units: u64::MAX,
164 spend_micros: 0
165 },
166 1,
167 0,
168 Limits::default()
169 ),
170 LimitDecision::Rejected(LimitReason::ArithmeticOverflow)
171 );
172 }
173}