gateway_core/
governance.rs

1use std::{
2    collections::HashMap,
3    sync::Mutex,
4    time::{Duration, Instant},
5};
6
7#[derive(Debug, Clone, PartialEq, Eq, Hash)]
8pub struct GovernanceKey {
9    pub scope_id: String,
10    pub principal_id: String,
11    pub model: String,
12}
13
14impl GovernanceKey {
15    pub fn new(
16        scope_id: impl Into<String>,
17        principal_id: impl Into<String>,
18        model: impl Into<String>,
19    ) -> Self {
20        Self {
21            scope_id: scope_id.into(),
22            principal_id: principal_id.into(),
23            model: model.into(),
24        }
25    }
26}
27
28#[derive(Debug, Clone, Copy, PartialEq, Eq)]
29pub struct GovernanceLimits {
30    pub requests_per_window: u32,
31    pub tokens_per_window: u64,
32    pub window: Duration,
33}
34
35#[derive(Debug, Clone, Copy, PartialEq, Eq)]
36pub enum Admission {
37    Allowed,
38    RequestLimited { retry_after: Duration },
39    TokenLimited { retry_after: Duration },
40}
41
42#[derive(Debug, Clone, Copy)]
43struct Window {
44    started: Instant,
45    requests: u32,
46    tokens: u64,
47}
48
49pub struct Governance {
50    limits: GovernanceLimits,
51    windows: Mutex<HashMap<GovernanceKey, Window>>,
52}
53
54impl Governance {
55    pub fn new(limits: GovernanceLimits) -> Self {
56        Self {
57            limits,
58            windows: Mutex::new(HashMap::new()),
59        }
60    }
61
62    pub fn admit(&self, key: &GovernanceKey) -> Admission {
63        self.admit_at(key, Instant::now())
64    }
65
66    pub fn record_usage(&self, key: &GovernanceKey, tokens: u64) {
67        self.record_usage_at(key, tokens, Instant::now());
68    }
69
70    pub fn clear(&self, key: &GovernanceKey) {
71        self.lock().remove(key);
72    }
73
74    fn admit_at(&self, key: &GovernanceKey, now: Instant) -> Admission {
75        let mut windows = self.lock();
76        let window = window(&mut windows, key, now, self.limits.window);
77        let retry_after = self
78            .limits
79            .window
80            .saturating_sub(now.saturating_duration_since(window.started));
81        if window.tokens >= self.limits.tokens_per_window {
82            return Admission::TokenLimited { retry_after };
83        }
84        if window.requests >= self.limits.requests_per_window {
85            return Admission::RequestLimited { retry_after };
86        }
87        window.requests = window.requests.saturating_add(1);
88        Admission::Allowed
89    }
90
91    fn record_usage_at(&self, key: &GovernanceKey, tokens: u64, now: Instant) {
92        let mut windows = self.lock();
93        let window = window(&mut windows, key, now, self.limits.window);
94        window.tokens = window.tokens.saturating_add(tokens);
95    }
96
97    fn lock(&self) -> std::sync::MutexGuard<'_, HashMap<GovernanceKey, Window>> {
98        self.windows
99            .lock()
100            .unwrap_or_else(std::sync::PoisonError::into_inner)
101    }
102}
103
104fn window<'a>(
105    windows: &'a mut HashMap<GovernanceKey, Window>,
106    key: &GovernanceKey,
107    now: Instant,
108    duration: Duration,
109) -> &'a mut Window {
110    let current = windows.entry(key.clone()).or_insert(Window {
111        started: now,
112        requests: 0,
113        tokens: 0,
114    });
115    if now.saturating_duration_since(current.started) >= duration {
116        *current = Window {
117            started: now,
118            requests: 0,
119            tokens: 0,
120        };
121    }
122    current
123}
124
125#[cfg(test)]
126mod tests {
127    use super::*;
128
129    #[test]
130    fn enforces_requests_and_tokens_per_key() {
131        let governance = Governance::new(GovernanceLimits {
132            requests_per_window: 1,
133            tokens_per_window: 10,
134            window: Duration::from_secs(60),
135        });
136        let key = GovernanceKey::new("scope-a", "principal-a", "openai/model");
137        assert_eq!(governance.admit(&key), Admission::Allowed);
138        assert!(matches!(
139            governance.admit(&key),
140            Admission::RequestLimited { .. }
141        ));
142        governance.clear(&key);
143        governance.record_usage(&key, 10);
144        assert!(matches!(
145            governance.admit(&key),
146            Admission::TokenLimited { .. }
147        ));
148    }
149}
gateway_core/governance.rs

gateway_core/
governance.rs
