headless_lms_credit_registration/runtime/suotar/
rate_limit.rs

1//! How much one worker process may ask of each rate-limited Suotar endpoint.
2//!
3//! One token bucket per endpoint, in memory like the circuit breaker: every endpoint is called from
4//! one process only. After its own failures, or when the breaker trips or closes again, an endpoint
5//! drops to a tenth of its rate and doubles back every five healthy minutes.
6//!
7//! Keyed by scope like the breaker, so a test's scoped ticks never spend what the live workers or
8//! another test may send.
9
10use std::time::{Duration, Instant};
11
12use headless_lms_models::suotar_api_calls::SuotarEndpoint;
13
14use crate::runtime::process_local::{LastReported, ProcessLocalMap, ScopeKey};
15
16/// What an endpoint may take at full rate.
17#[derive(Debug, Clone, Copy, PartialEq)]
18pub(super) struct EndpointRate {
19    /// Items per minute, or requests per minute for `list_by_course`.
20    pub per_minute: f64,
21    /// The most that may go out in one burst.
22    pub capacity: f64,
23}
24
25/// The rate an endpoint is limited to, or `None` for one that is not limited.
26pub(super) fn endpoint_rate(endpoint: SuotarEndpoint) -> Option<EndpointRate> {
27    match endpoint {
28        SuotarEndpoint::ResolveEnrolments | SuotarEndpoint::ResolvePersons => Some(EndpointRate {
29            per_minute: 500.0,
30            capacity: 1000.0,
31        }),
32        SuotarEndpoint::ListByCourse => Some(EndpointRate {
33            per_minute: 2.0,
34            capacity: 2.0,
35        }),
36        SuotarEndpoint::ImportAttainments
37        | SuotarEndpoint::VerifyAttainments
38        | SuotarEndpoint::ValidateCourseCodes => None,
39    }
40}
41
42/// The endpoints [`endpoint_rate`] limits.
43pub(super) const LIMITED_ENDPOINTS: [SuotarEndpoint; 3] = [
44    SuotarEndpoint::ResolvePersons,
45    SuotarEndpoint::ResolveEnrolments,
46    SuotarEndpoint::ListByCourse,
47];
48
49/// The share of the full rate an endpoint drops to.
50const FLOOR_SHARE: f64 = 0.1;
51/// The reduced share doubles after each of these without a failure.
52const RAMP_DOUBLING_INTERVAL: Duration = Duration::from_secs(5 * 60);
53
54#[derive(Debug, Clone)]
55struct Bucket {
56    tokens: f64,
57    refilled_at: Instant,
58    /// When the reduced rate last restarted from the floor; `None` at full rate.
59    floor_started_at: Option<Instant>,
60}
61
62impl Bucket {
63    fn new(rate: EndpointRate, now: Instant) -> Self {
64        Self {
65            tokens: rate.capacity,
66            refilled_at: now,
67            floor_started_at: None,
68        }
69    }
70
71    fn share(&mut self, now: Instant) -> f64 {
72        let Some(started) = self.floor_started_at else {
73            return 1.0;
74        };
75        let doublings = now.duration_since(started).as_secs() / RAMP_DOUBLING_INTERVAL.as_secs();
76        let share = FLOOR_SHARE * 2f64.powi(i32::try_from(doublings).unwrap_or(i32::MAX));
77        if share >= 1.0 {
78            self.floor_started_at = None;
79            return 1.0;
80        }
81        share
82    }
83
84    /// Refills the bucket, returning whether this call is the one that found the ramp back at full
85    /// rate.
86    fn refill(&mut self, rate: EndpointRate, now: Instant) -> bool {
87        let was_reduced = self.floor_started_at.is_some();
88        let share = self.share(now);
89        let elapsed_minutes = now.duration_since(self.refilled_at).as_secs_f64() / 60.0;
90        self.tokens =
91            (self.tokens + elapsed_minutes * rate.per_minute * share).min(burst_limit(rate, share));
92        self.refilled_at = now;
93        was_reduced && self.floor_started_at.is_none()
94    }
95
96    /// Whether the bucket is back where a new one starts, so dropping it changes nothing.
97    fn is_fresh(&mut self, rate: EndpointRate, now: Instant) -> bool {
98        self.refill(rate, now);
99        self.floor_started_at.is_none() && self.tokens >= rate.capacity
100    }
101}
102
103static BUCKETS: ProcessLocalMap<(ScopeKey, SuotarEndpoint), Bucket> = ProcessLocalMap::new();
104
105/// The global limiters as this process last wrote them for the dashboard.
106pub(super) static REPORTED: LastReported<SuotarEndpoint, LimiterSnapshot> = LastReported::new();
107
108fn with_bucket<T>(
109    key: &ScopeKey,
110    endpoint: SuotarEndpoint,
111    use_bucket: impl FnOnce(&mut Bucket, EndpointRate) -> T,
112) -> Option<T> {
113    let rate = endpoint_rate(endpoint)?;
114    let now = Instant::now();
115    let mut buckets = BUCKETS.lock();
116    let bucket_key = (key.clone(), endpoint);
117    if !buckets.contains_key(&bucket_key) {
118        // A scope never run again would otherwise stay in the map for the life of the process.
119        buckets.retain(|(_, endpoint), bucket| {
120            endpoint_rate(*endpoint).is_some_and(|rate| !bucket.is_fresh(rate, now))
121        });
122    }
123    let bucket = buckets
124        .entry(bucket_key)
125        .or_insert_with(|| Bucket::new(rate, now));
126    let recovered = bucket.refill(rate, now);
127    trace!(
128        ?endpoint,
129        tokens = bucket.tokens,
130        "Refilled Suotar rate limit bucket"
131    );
132    if recovered {
133        info!(
134            ?endpoint,
135            rate_per_minute = rate.per_minute,
136            "Suotar endpoint rate limit recovered to full rate"
137        );
138    }
139    Some(use_bucket(bucket, rate))
140}
141
142/// How many items, or requests, the endpoint may take right now. Unlimited endpoints answer
143/// `usize::MAX`.
144pub(super) fn available(key: &ScopeKey, endpoint: SuotarEndpoint) -> usize {
145    with_bucket(key, endpoint, |bucket, _| {
146        bucket.tokens.floor().max(0.0) as usize
147    })
148    .unwrap_or(usize::MAX)
149}
150
151/// Spends `count` of what [`available`] allowed.
152pub(super) fn take(key: &ScopeKey, endpoint: SuotarEndpoint, count: usize) {
153    with_bucket(key, endpoint, |bucket, _| {
154        bucket.tokens = (bucket.tokens - count as f64).max(0.0);
155    });
156}
157
158/// Spends `count` even past what [`available`] allows, into debt that later refills pay off first,
159/// capped at one capacity so recovery takes at most one refill. Only for a split batch's resent
160/// halves, which must not wait for a later iteration and re-claim the same batch.
161pub(super) fn overdraw(key: &ScopeKey, endpoint: SuotarEndpoint, count: usize) {
162    with_bucket(key, endpoint, |bucket, rate| {
163        bucket.tokens = (bucket.tokens - count as f64).max(-rate.capacity);
164    });
165}
166
167/// Drops `endpoints` to [`FLOOR_SHARE`] of their rate, restarting the climb back. Unlimited ones
168/// are skipped.
169pub(super) fn drop_to_floor(key: &ScopeKey, endpoints: &[SuotarEndpoint]) {
170    for &endpoint in endpoints {
171        with_bucket(key, endpoint, |bucket, rate| {
172            let was_at_floor = bucket.floor_started_at.is_some();
173            bucket.floor_started_at = Some(Instant::now());
174            bucket.tokens = bucket.tokens.min(burst_limit(rate, FLOOR_SHARE));
175            if !was_at_floor {
176                info!(
177                    ?endpoint,
178                    floor_share = FLOOR_SHARE,
179                    "Suotar endpoint rate limit dropped to floor"
180                );
181            }
182        });
183    }
184}
185
186/// At least one, so an endpoint whose capacity is a couple of requests still gets one through at
187/// the floor.
188fn burst_limit(rate: EndpointRate, share: f64) -> f64 {
189    (rate.capacity * share).max(1.0)
190}
191
192/// One endpoint's limiter as it stands, for the dashboard.
193#[derive(Debug, Clone, Copy, PartialEq)]
194pub(super) struct LimiterSnapshot {
195    pub share: f64,
196    pub available: usize,
197    pub rate: EndpointRate,
198}
199
200/// Forgets every bucket of `scope`, so its endpoints are back at full rate with a full burst.
201pub(super) fn reset(scope: &ScopeKey) {
202    BUCKETS.lock().retain(|(key, _), _| key != scope);
203}
204
205/// `None` for an unlimited endpoint.
206pub(super) fn snapshot(key: &ScopeKey, endpoint: SuotarEndpoint) -> Option<LimiterSnapshot> {
207    with_bucket(key, endpoint, |bucket, rate| LimiterSnapshot {
208        share: bucket.share(Instant::now()),
209        available: bucket.tokens.floor().max(0.0) as usize,
210        rate,
211    })
212}
213
214#[cfg(test)]
215mod tests {
216    use super::*;
217
218    #[test]
219    fn the_limited_endpoints_are_exactly_the_rated_ones() {
220        let every = [
221            SuotarEndpoint::ResolvePersons,
222            SuotarEndpoint::ResolveEnrolments,
223            SuotarEndpoint::ImportAttainments,
224            SuotarEndpoint::VerifyAttainments,
225            SuotarEndpoint::ListByCourse,
226            SuotarEndpoint::ValidateCourseCodes,
227        ];
228        let rated: Vec<_> = every
229            .into_iter()
230            .filter(|&endpoint| endpoint_rate(endpoint).is_some())
231            .collect();
232        assert_eq!(rated, LIMITED_ENDPOINTS);
233    }
234}
headless_lms_credit_registration/runtime/suotar/rate_limit.rs

headless_lms_credit_registration/runtime/suotar/
rate_limit.rs
