Skip to main content

boson_backend_sql_common/
enqueue_rate.rs

1//! Process-local sliding-window limiter for enqueues per task per second.
2//!
3//! Uses [`std::sync::Mutex`] intentionally: [`EnqueueRateLimiter::try_record`] holds the lock only
4//! for a short, non-`.await` critical section (prune window + push timestamp). Async enqueue paths
5//! call it between awaits, so a Tokio mutex would add overhead without preventing runtime stalls.
6//! Poisoned locks recover via [`std::sync::PoisonError::into_inner`].
7
8use std::collections::{HashMap, VecDeque};
9use std::sync::Mutex;
10use std::time::Instant;
11
12/// Tracks recent enqueue timestamps per task name (1-second sliding window).
13#[derive(Debug, Default)]
14pub struct EnqueueRateLimiter {
15    inner: Mutex<HashMap<String, VecDeque<Instant>>>,
16}
17
18impl EnqueueRateLimiter {
19    /// New empty limiter.
20    #[must_use]
21    pub fn new() -> Self {
22        Self::default()
23    }
24
25    /// Returns `true` if this enqueue is allowed under `max_per_second` for `task_name`.
26    ///
27    /// Records the enqueue timestamp when allowed. `max_per_second == 0` means unlimited.
28    pub fn try_record(&self, task_name: &str, max_per_second: u32) -> bool {
29        if max_per_second == 0 {
30            return true;
31        }
32        let mut guard = self
33            .inner
34            .lock()
35            .unwrap_or_else(std::sync::PoisonError::into_inner);
36        let now = Instant::now();
37        let allowed = {
38            let window = guard.entry(task_name.to_string()).or_default();
39            while let Some(front) = window.front().copied() {
40                if now.duration_since(front).as_secs() >= 1 {
41                    window.pop_front();
42                } else {
43                    break;
44                }
45            }
46            if window.len() >= max_per_second as usize {
47                false
48            } else {
49                window.push_back(now);
50                true
51            }
52        };
53        drop(guard);
54        allowed
55    }
56}