plausible_cli/rate_limit/
mod.rs

1use governor::{
2    clock::DefaultClock,
3    state::{direct::NotKeyed, InMemoryState},
4    Jitter, Quota, RateLimiter as GovernorRateLimiter,
5};
6use serde::{Deserialize, Deserializer, Serialize, Serializer};
7use std::fs;
8use std::num::NonZeroU32;
9use std::path::PathBuf;
10use std::sync::Arc;
11use std::time::Duration as StdDuration;
12use time::{
13    format_description::FormatItem, macros::format_description, Date, Duration as TimeDuration,
14    OffsetDateTime,
15};
16use tokio::sync::Mutex;
17
18use crate::config::ConfigPaths;
19
20const DEFAULT_JITTER_UP_TO_MS: u64 = 250;
21
22/// Rate limiting configuration combining hourly quota, optional daily ceiling, and jitter.
23#[derive(Debug, Clone)]
24pub struct RateLimitConfig {
25    pub hourly_quota: NonZeroU32,
26    pub daily_quota: Option<NonZeroU32>,
27    pub jitter_max: StdDuration,
28    quota: Quota,
29}
30
31impl RateLimitConfig {
32    pub fn new(hourly_quota: NonZeroU32) -> Self {
33        let quota = Quota::per_hour(hourly_quota);
34        Self {
35            hourly_quota,
36            daily_quota: Some(hourly_quota),
37            jitter_max: StdDuration::from_millis(DEFAULT_JITTER_UP_TO_MS),
38            quota,
39        }
40    }
41
42    pub fn with_daily_quota(mut self, daily: Option<NonZeroU32>) -> Self {
43        self.daily_quota = daily;
44        self
45    }
46
47    #[cfg(test)]
48    pub fn with_quota(mut self, quota: Quota) -> Self {
49        self.quota = quota;
50        self
51    }
52
53    fn jitter(&self) -> Jitter {
54        if self.jitter_max.is_zero() {
55            Jitter::up_to(StdDuration::from_millis(0))
56        } else {
57            Jitter::up_to(self.jitter_max)
58        }
59    }
60
61    fn quota(&self) -> Quota {
62        self.quota
63    }
64}
65
66impl Default for RateLimitConfig {
67    fn default() -> Self {
68        RateLimitConfig::new(NonZeroU32::new(600).expect("non zero"))
69    }
70}
71
72/// Aggregate rate limiter coordinating in-memory throttling with persisted usage counters.
73#[derive(Debug, Clone)]
74pub struct RateLimiter {
75    limiter: Arc<GovernorRateLimiter<NotKeyed, InMemoryState, DefaultClock>>,
76    usage: Arc<Mutex<UsageLedger>>,
77    config: RateLimitConfig,
78    jitter: Jitter,
79}
80
81impl RateLimiter {
82    pub async fn new(
83        paths: ConfigPaths,
84        account: &str,
85        config: RateLimitConfig,
86    ) -> Result<Self, RateLimitError> {
87        paths.ensure_exists()?;
88        let usage = UsageLedger::load(paths, account.to_string())?;
89        let limiter = GovernorRateLimiter::direct(config.quota());
90        Ok(Self {
91            limiter: Arc::new(limiter),
92            usage: Arc::new(Mutex::new(usage)),
93            jitter: config.jitter(),
94            config,
95        })
96    }
97
98    /// Wait until the configured quota permits executing `weight` cost units.
99    pub async fn acquire(&self, weight: NonZeroU32) {
100        let _ = self
101            .limiter
102            .until_n_ready_with_jitter(weight, self.jitter)
103            .await;
104    }
105
106    /// Record a successful request, updating persisted counters.
107    pub async fn record_success(
108        &self,
109        weight: u32,
110        now: OffsetDateTime,
111    ) -> Result<RateStatus, RateLimitError> {
112        let mut usage = self.usage.lock().await;
113        usage.record(weight, now)?;
114        usage.persist()?;
115        Ok(usage.status(&self.config, now))
116    }
117
118    /// Retrieve the current counter status without mutating state.
119    pub async fn status(&self, now: OffsetDateTime) -> Result<RateStatus, RateLimitError> {
120        let usage = self.usage.lock().await;
121        Ok(usage.status(&self.config, now))
122    }
123}
124
125#[derive(Debug, Clone)]
126struct UsageLedger {
127    path: PathBuf,
128    snapshot: UsageSnapshot,
129}
130
131impl UsageLedger {
132    fn load(paths: ConfigPaths, account: String) -> Result<Self, RateLimitError> {
133        let path = paths.usage_dir().join(format!("{account}.json"));
134        let snapshot = match fs::read(&path) {
135            Ok(bytes) => {
136                if bytes.is_empty() {
137                    UsageSnapshot::new(account.clone())
138                } else {
139                    serde_json::from_slice(&bytes).map_err(RateLimitError::Deserialize)?
140                }
141            }
142            Err(err) if err.kind() == std::io::ErrorKind::NotFound => {
143                UsageSnapshot::new(account.clone())
144            }
145            Err(source) => return Err(RateLimitError::Io { path, source }),
146        };
147        Ok(Self { path, snapshot })
148    }
149
150    fn persist(&self) -> Result<(), RateLimitError> {
151        let contents =
152            serde_json::to_vec_pretty(&self.snapshot).map_err(RateLimitError::Serialize)?;
153        if let Some(parent) = self.path.parent() {
154            fs::create_dir_all(parent).map_err(|source| RateLimitError::Io {
155                path: parent.to_path_buf(),
156                source,
157            })?;
158        }
159        fs::write(&self.path, contents).map_err(|source| RateLimitError::Io {
160            path: self.path.clone(),
161            source,
162        })?;
163        #[cfg(unix)]
164        {
165            use std::os::unix::fs::PermissionsExt;
166            let mut perms = fs::metadata(&self.path)
167                .map_err(|source| RateLimitError::Io {
168                    path: self.path.clone(),
169                    source,
170                })?
171                .permissions();
172            perms.set_mode(0o600);
173            fs::set_permissions(&self.path, perms).map_err(|source| RateLimitError::Io {
174                path: self.path.clone(),
175                source,
176            })?;
177        }
178        Ok(())
179    }
180
181    fn record(&mut self, amount: u32, now: OffsetDateTime) -> Result<(), RateLimitError> {
182        let hour = hour_bucket(now);
183        let entry = self.snapshot.hourly.get_or_insert(HourWindow {
184            start: hour,
185            count: 0,
186        });
187        if entry.start != hour {
188            entry.start = hour;
189            entry.count = 0;
190        }
191        entry.count = entry.count.saturating_add(amount);
192
193        let day = now.date();
194        let daily = self.snapshot.daily.get_or_insert(DayWindow {
195            date: day,
196            count: 0,
197        });
198        if daily.date != day {
199            daily.date = day;
200            daily.count = 0;
201        }
202        daily.count = daily.count.saturating_add(amount);
203        Ok(())
204    }
205
206    fn status(&self, config: &RateLimitConfig, now: OffsetDateTime) -> RateStatus {
207        let hour = hour_bucket(now);
208        let day = now.date();
209        let hourly_used = self
210            .snapshot
211            .hourly
212            .as_ref()
213            .filter(|window| window.start == hour)
214            .map(|window| window.count)
215            .unwrap_or(0);
216        let hourly_remaining = config.hourly_quota.get().saturating_sub(hourly_used);
217        let hourly_reset_at = hour.checked_add(TimeDuration::hours(1)).unwrap_or(now);
218
219        let (daily_used, daily_remaining, daily_reset_at) = match config.daily_quota {
220            Some(limit) => {
221                let used = self
222                    .snapshot
223                    .daily
224                    .as_ref()
225                    .filter(|window| window.date == day)
226                    .map(|window| window.count)
227                    .unwrap_or(0);
228                let remaining = limit.get().saturating_sub(used);
229                let reset = day
230                    .next_day()
231                    .and_then(|d| d.with_hms(0, 0, 0).ok())
232                    .map(|dt| dt.assume_utc());
233                (Some(used), Some(remaining), reset)
234            }
235            None => (None, None, None),
236        };
237
238        RateStatus {
239            hourly_remaining,
240            hourly_limit: config.hourly_quota.get(),
241            hourly_reset_at,
242            hourly_used,
243            daily_limit: config.daily_quota.map(|n| n.get()),
244            daily_remaining,
245            daily_used,
246            daily_reset_at,
247        }
248    }
249}
250
251fn hour_bucket(now: OffsetDateTime) -> OffsetDateTime {
252    now.replace_minute(0)
253        .unwrap()
254        .replace_second(0)
255        .unwrap()
256        .replace_nanosecond(0)
257        .unwrap()
258}
259
260#[derive(Debug, Clone, Serialize, Deserialize)]
261struct UsageSnapshot {
262    account: String,
263    #[serde(default)]
264    hourly: Option<HourWindow>,
265    #[serde(default)]
266    daily: Option<DayWindow>,
267}
268
269impl UsageSnapshot {
270    fn new(account: String) -> Self {
271        Self {
272            account,
273            hourly: None,
274            daily: None,
275        }
276    }
277}
278
279#[derive(Debug, Clone, Serialize, Deserialize)]
280struct HourWindow {
281    #[serde(with = "time::serde::rfc3339")]
282    start: OffsetDateTime,
283    count: u32,
284}
285
286#[derive(Debug, Clone, Serialize, Deserialize)]
287struct DayWindow {
288    #[serde(with = "serde_date")]
289    date: Date,
290    count: u32,
291}
292
293mod serde_date {
294    use super::*;
295
296    const FORMAT: &[FormatItem<'static>] = format_description!("[year]-[month]-[day]");
297
298    pub fn serialize<S>(date: &Date, serializer: S) -> Result<S::Ok, S::Error>
299    where
300        S: Serializer,
301    {
302        let formatted = date.format(FORMAT).map_err(serde::ser::Error::custom)?;
303        serializer.serialize_str(&formatted)
304    }
305
306    pub fn deserialize<'de, D>(deserializer: D) -> Result<Date, D::Error>
307    where
308        D: Deserializer<'de>,
309    {
310        let value = String::deserialize(deserializer)?;
311        Date::parse(&value, FORMAT).map_err(serde::de::Error::custom)
312    }
313}
314
315/// Snapshot of remaining budget for presentation.
316#[derive(Debug, Clone, PartialEq)]
317pub struct RateStatus {
318    pub hourly_remaining: u32,
319    pub hourly_limit: u32,
320    pub hourly_used: u32,
321    pub hourly_reset_at: OffsetDateTime,
322    pub daily_limit: Option<u32>,
323    pub daily_used: Option<u32>,
324    pub daily_remaining: Option<u32>,
325    pub daily_reset_at: Option<OffsetDateTime>,
326}
327
328#[derive(thiserror::Error, Debug)]
329pub enum RateLimitError {
330    #[error("I/O error at {path:?}")]
331    Io {
332        path: PathBuf,
333        #[source]
334        source: std::io::Error,
335    },
336    #[error("failed to deserialize usage file: {0}")]
337    Deserialize(#[source] serde_json::Error),
338    #[error("failed to serialize usage file: {0}")]
339    Serialize(#[source] serde_json::Error),
340    #[error(transparent)]
341    Config(#[from] crate::config::ConfigError),
342}
343
344#[cfg(test)]
345mod tests {
346    use super::*;
347    use crate::config::ConfigPaths;
348    use std::num::NonZeroU32;
349    use tempfile::tempdir;
350    use tokio::time::Instant;
351
352    fn temp_paths() -> (ConfigPaths, tempfile::TempDir) {
353        let tmp = tempdir().expect("tempdir");
354        let paths = ConfigPaths::from_base_dir(tmp.path());
355        (paths, tmp)
356    }
357
358    #[tokio::test]
359    async fn records_and_persists_usage() {
360        let (paths, _guard) = temp_paths();
361        let config = RateLimitConfig::new(NonZeroU32::new(10).unwrap());
362        let limiter = RateLimiter::new(paths.clone(), "test", config)
363            .await
364            .expect("limiter");
365        let now = OffsetDateTime::now_utc();
366        let status = limiter.record_success(5, now).await.expect("record");
367        assert_eq!(status.hourly_used, 5);
368        assert_eq!(status.hourly_remaining, 5);
369
370        // Reload from disk and ensure counts persisted.
371        let limiter2 = RateLimiter::new(
372            paths,
373            "test",
374            RateLimitConfig::new(NonZeroU32::new(10).unwrap()),
375        )
376        .await
377        .expect("limiter2");
378        let status2 = limiter2.status(now).await.expect("status");
379        assert_eq!(status2.hourly_used, 5);
380    }
381
382    #[tokio::test]
383    async fn hourly_window_resets_after_boundary() {
384        let (paths, _guard) = temp_paths();
385        let config = RateLimitConfig::new(NonZeroU32::new(10).unwrap());
386        let limiter = RateLimiter::new(paths, "test", config)
387            .await
388            .expect("limiter");
389        let now = OffsetDateTime::now_utc();
390        limiter.record_success(10, now).await.expect("record");
391        let next_hour = hour_bucket(now)
392            .checked_add(TimeDuration::hours(1))
393            .unwrap();
394        let status = limiter.status(next_hour).await.expect("status");
395        assert_eq!(status.hourly_used, 0);
396        assert_eq!(status.hourly_remaining, 10);
397    }
398
399    #[tokio::test]
400    async fn rate_limiter_waits_when_quota_exceeded() {
401        let (paths, _guard) = temp_paths();
402        let quota = Quota::per_second(NonZeroU32::new(1).unwrap());
403        let config = RateLimitConfig::new(NonZeroU32::new(1).unwrap()).with_quota(quota);
404        let limiter = RateLimiter::new(paths, "test", config)
405            .await
406            .expect("limiter");
407
408        limiter.acquire(NonZeroU32::new(1).unwrap()).await;
409        let start = Instant::now();
410        limiter.acquire(NonZeroU32::new(1).unwrap()).await;
411        let elapsed = start.elapsed();
412        assert!(elapsed >= StdDuration::from_millis(900));
413    }
414}
plausible_cli/rate_limit/mod.rs

plausible_cli/rate_limit/
mod.rs
