1use std::sync::Arc;
25use std::sync::atomic::{AtomicU64, Ordering};
26use std::time::{Duration, Instant};
27
28pub const DEFAULT_RULE_TIMEOUT: Duration = Duration::from_secs(1);
30
31pub const DEFAULT_GLOBAL_TIMEOUT: Duration = Duration::from_secs(15);
33
34pub const DEFAULT_MEMORY_BYTES: usize = 64 * 1024 * 1024;
36
37#[derive(Debug, Clone, Copy, PartialEq, Eq)]
39pub struct Limits {
40 pub rule_timeout: Duration,
46
47 pub global_timeout: Duration,
52
53 pub memory_bytes: usize,
55}
56
57impl Default for Limits {
58 fn default() -> Self {
59 Self {
60 rule_timeout: DEFAULT_RULE_TIMEOUT,
61 global_timeout: DEFAULT_GLOBAL_TIMEOUT,
62 memory_bytes: DEFAULT_MEMORY_BYTES,
63 }
64 }
65}
66
67impl Limits {
68 #[must_use]
73 pub const fn with_rule_timeout(mut self, timeout: Duration) -> Self {
74 self.rule_timeout = timeout;
75 self
76 }
77
78 #[must_use]
80 pub const fn with_global_timeout(mut self, timeout: Duration) -> Self {
81 self.global_timeout = timeout;
82 self
83 }
84
85 #[must_use]
87 pub const fn with_memory_bytes(mut self, bytes: usize) -> Self {
88 self.memory_bytes = bytes;
89 self
90 }
91}
92
93#[derive(Debug)]
98pub struct RunClock {
99 start: Instant,
100 global_timeout: Duration,
101}
102
103impl RunClock {
104 #[must_use]
106 pub fn start(global_timeout: Duration) -> Arc<Self> {
107 Arc::new(Self {
108 start: Instant::now(),
109 global_timeout,
110 })
111 }
112
113 #[must_use]
115 pub fn elapsed(&self) -> Duration {
116 self.start.elapsed()
117 }
118
119 #[must_use]
121 pub const fn global_timeout(&self) -> Duration {
122 self.global_timeout
123 }
124
125 #[must_use]
127 pub fn is_expired(&self) -> bool {
128 self.elapsed() >= self.global_timeout
129 }
130
131 fn elapsed_nanos(&self) -> u64 {
132 u64::try_from(self.start.elapsed().as_nanos()).unwrap_or(u64::MAX)
133 }
134}
135
136#[derive(Debug, Clone, Copy, PartialEq, Eq)]
138pub enum Trip {
139 Rule,
141 Run,
143}
144
145const TRIP_NONE: u64 = 0;
146const TRIP_RULE: u64 = 1;
147const TRIP_RUN: u64 = 2;
148
149#[derive(Debug)]
159pub struct Budget {
160 clock: Arc<RunClock>,
161 global_nanos: u64,
162 invocation_deadline_nanos: AtomicU64,
165 tripped: AtomicU64,
166}
167
168impl Budget {
169 pub fn new(clock: Arc<RunClock>) -> Arc<Self> {
171 let global_nanos = u64::try_from(clock.global_timeout.as_nanos()).unwrap_or(u64::MAX);
172 Arc::new(Self {
173 clock,
174 global_nanos,
175 invocation_deadline_nanos: AtomicU64::new(0),
176 tripped: AtomicU64::new(TRIP_NONE),
177 })
178 }
179
180 pub fn arm(&self, rule_timeout: Duration) {
182 let now = self.clock.elapsed_nanos();
183 let budget = u64::try_from(rule_timeout.as_nanos()).unwrap_or(u64::MAX);
184 self.invocation_deadline_nanos
187 .store(now.saturating_add(budget).max(1), Ordering::Relaxed);
188 self.tripped.store(TRIP_NONE, Ordering::Relaxed);
189 }
190
191 pub fn disarm(&self) {
193 self.invocation_deadline_nanos.store(0, Ordering::Relaxed);
194 }
195
196 pub fn should_interrupt(&self) -> bool {
199 let elapsed = self.clock.elapsed_nanos();
200
201 if elapsed >= self.global_nanos {
202 self.tripped.store(TRIP_RUN, Ordering::Relaxed);
203 return true;
204 }
205
206 let deadline = self.invocation_deadline_nanos.load(Ordering::Relaxed);
207 if deadline != 0 && elapsed >= deadline {
208 self.tripped.store(TRIP_RULE, Ordering::Relaxed);
209 return true;
210 }
211
212 false
213 }
214
215 pub fn take_trip(&self) -> Option<Trip> {
217 match self.tripped.swap(TRIP_NONE, Ordering::Relaxed) {
218 TRIP_RULE => Some(Trip::Rule),
219 TRIP_RUN => Some(Trip::Run),
220 _ => None,
221 }
222 }
223
224 pub fn clock(&self) -> &RunClock {
226 &self.clock
227 }
228}
229
230#[cfg(test)]
231mod tests {
232 use super::*;
233
234 #[test]
235 fn defaults_match_the_documented_budgets() {
236 let limits = Limits::default();
237 assert_eq!(limits.rule_timeout, Duration::from_secs(1));
238 assert_eq!(limits.global_timeout, Duration::from_secs(15));
239 assert_eq!(limits.memory_bytes, 64 * 1024 * 1024);
240 }
241
242 #[test]
243 fn the_rule_budget_is_well_under_the_global_one() {
244 let limits = Limits::default();
248 assert!(
249 limits.rule_timeout * 5 < limits.global_timeout,
250 "the per-invocation budget must leave room for the global limit to be a backstop"
251 );
252 }
253
254 #[test]
255 fn a_rule_cannot_raise_the_global_budget() {
256 let limits = Limits::default().with_rule_timeout(Duration::from_mins(1));
257 assert_eq!(limits.rule_timeout, Duration::from_mins(1));
258 assert_eq!(
259 limits.global_timeout, DEFAULT_GLOBAL_TIMEOUT,
260 "raising a rule's own budget must not extend the run"
261 );
262 }
263
264 #[test]
265 fn an_unarmed_budget_never_interrupts() {
266 let budget = Budget::new(RunClock::start(Duration::from_hours(1)));
267 assert!(!budget.should_interrupt());
268 assert_eq!(budget.take_trip(), None);
269 }
270
271 #[test]
272 fn an_expired_invocation_budget_interrupts_and_records_why() {
273 let budget = Budget::new(RunClock::start(Duration::from_hours(1)));
274 budget.arm(Duration::ZERO);
275 assert!(budget.should_interrupt());
276 assert_eq!(budget.take_trip(), Some(Trip::Rule));
277 }
278
279 #[test]
280 fn an_expired_run_budget_interrupts_and_records_why() {
281 let budget = Budget::new(RunClock::start(Duration::ZERO));
282 budget.arm(Duration::from_hours(1));
283 assert!(budget.should_interrupt());
284 assert_eq!(budget.take_trip(), Some(Trip::Run));
285 }
286
287 #[test]
288 fn the_run_budget_wins_when_both_are_spent() {
289 let budget = Budget::new(RunClock::start(Duration::ZERO));
292 budget.arm(Duration::ZERO);
293 assert!(budget.should_interrupt());
294 assert_eq!(budget.take_trip(), Some(Trip::Run));
295 }
296
297 #[test]
298 fn disarming_stops_invocation_enforcement() {
299 let budget = Budget::new(RunClock::start(Duration::from_hours(1)));
300 budget.arm(Duration::ZERO);
301 budget.disarm();
302 assert!(!budget.should_interrupt(), "no invocation is in flight");
303 }
304
305 #[test]
306 fn taking_the_trip_clears_it() {
307 let budget = Budget::new(RunClock::start(Duration::from_hours(1)));
308 budget.arm(Duration::ZERO);
309 assert!(budget.should_interrupt());
310 assert_eq!(budget.take_trip(), Some(Trip::Rule));
311 assert_eq!(
312 budget.take_trip(),
313 None,
314 "a trip must not be reported twice"
315 );
316 }
317
318 #[test]
319 fn arming_clears_a_previous_trip() {
320 let budget = Budget::new(RunClock::start(Duration::from_hours(1)));
323 budget.arm(Duration::ZERO);
324 assert!(budget.should_interrupt());
325
326 budget.arm(Duration::from_hours(1));
327 assert!(!budget.should_interrupt());
328 assert_eq!(budget.take_trip(), None);
329 }
330
331 #[test]
332 fn an_overflowing_budget_does_not_wrap_into_disarmed() {
333 let budget = Budget::new(RunClock::start(Duration::from_hours(1)));
336 budget.arm(Duration::MAX);
337 assert_ne!(
338 budget.invocation_deadline_nanos.load(Ordering::Relaxed),
339 0,
340 "an overflowing budget must not read as disarmed"
341 );
342 }
343
344 #[test]
345 fn the_clock_measures_from_one_origin() {
346 let clock = RunClock::start(Duration::from_hours(1));
347 let a = Arc::clone(&clock);
348 let b = Arc::clone(&clock);
349 assert!(!a.is_expired());
350 assert!(!b.is_expired());
351 assert_eq!(a.global_timeout(), Duration::from_hours(1));
352 }
353
354 #[test]
355 fn a_zero_global_budget_is_immediately_expired() {
356 assert!(RunClock::start(Duration::ZERO).is_expired());
357 }
358}