Skip to main content

grate_limiter/quota/
mod.rs

1use serde::{Deserialize, Serialize};
2
3/// Quota dimension — what resource is being tracked.
4#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
5#[serde(rename_all = "snake_case")]
6pub enum Dimension {
7    /// Number of requests.
8    Requests,
9    /// Number of tokens (e.g., LLM tokens).
10    Tokens,
11    /// Concurrent active requests.
12    Concurrency,
13    /// Cost in USD.
14    CostUsd,
15    /// Bytes transferred.
16    Bytes,
17}
18
19/// Time window for quota reset.
20#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
21#[serde(rename_all = "snake_case")]
22pub enum Window {
23    /// Per-second window.
24    Second,
25    /// Per-minute window.
26    Minute,
27    /// Per-hour window.
28    Hour,
29    /// Per-day window.
30    Day,
31}
32
33impl Window {
34    /// Duration of this window in nanoseconds.
35    pub fn as_nanos(&self) -> u64 {
36        match self {
37            Window::Second => 1_000_000_000,
38            Window::Minute => 60_000_000_000,
39            Window::Hour => 3_600_000_000_000,
40            Window::Day => 86_400_000_000_000,
41        }
42    }
43
44    /// Duration of this window in seconds.
45    pub fn as_secs(&self) -> u64 {
46        match self {
47            Window::Second => 1,
48            Window::Minute => 60,
49            Window::Hour => 3_600,
50            Window::Day => 86_400,
51        }
52    }
53}
54
55/// Configuration for a single quota dimension on a provider.
56#[derive(Debug, Clone, Serialize, Deserialize)]
57pub struct QuotaConfig {
58    /// What resource is being limited.
59    pub dimension: Dimension,
60    /// Maximum allowed usage within the window.
61    pub limit: u64,
62    /// Time window for the limit. `None` for concurrency (no window).
63    pub window: Option<Window>,
64}
65
66mod concurrency;
67mod fixed_window;
68mod sliding_window;
69mod strategy;
70mod token_bucket;
71
72pub(crate) use concurrency::ConcurrencyLimiter;
73pub(crate) use strategy::QuotaTracker;
74pub(crate) use token_bucket::TokenBucket;
75
76use crate::clock::Timestamp;
77
78/// Create appropriate quota tracker for a given config.
79pub(crate) fn create_tracker(config: &QuotaConfig, now: Timestamp) -> Box<dyn QuotaTracker> {
80    match config.dimension {
81        Dimension::Concurrency => Box::new(ConcurrencyLimiter::new(config.limit)),
82        _ => match config.window {
83            Some(window) => Box::new(TokenBucket::new(config.limit, window, now)),
84            None => Box::new(TokenBucket::new(config.limit, Window::Minute, now)),
85        },
86    }
87}