aspero_engine/
budget.rs

1//! Providers' budgets as their last responses reported them, so a front
2//! end can show what's left without spending any.
3
4use std::sync::Mutex;
5
6use aspero_core::{Error, ProviderId, Usage};
7use jiff::Timestamp;
8
9/// What a provider last said was left of its budget.
10#[derive(Debug, Clone, Copy, PartialEq)]
11#[non_exhaustive]
12pub struct Budget {
13    pub remaining: Option<f64>,
14    pub resets_at: Option<Timestamp>,
15    /// When the provider said so.
16    pub seen_at: Timestamp,
17}
18
19#[derive(Debug, Default)]
20pub(crate) struct Budgets(Mutex<Vec<(ProviderId, Budget)>>);
21
22impl Budgets {
23    /// Remembers what a response (or a spent-budget error) said.
24    pub(crate) fn note(
25        &self,
26        provider: &ProviderId,
27        usage: Option<&Usage>,
28        error: Option<&Error>,
29        now: Timestamp,
30    ) {
31        let (remaining, resets_at) = match (usage, error) {
32            (Some(u), _) => (u.remaining, u.resets_at),
33            (None, Some(Error::QuotaExhausted { resets_at })) => (Some(0.0), *resets_at),
34            _ => return,
35        };
36        if remaining.is_none() && resets_at.is_none() {
37            return;
38        }
39        let budget = Budget {
40            remaining,
41            resets_at,
42            seen_at: now,
43        };
44        let mut all = self.0.lock().unwrap();
45        match all.iter_mut().find(|(p, _)| p == provider) {
46            Some((_, known)) => *known = budget,
47            None => all.push((provider.clone(), budget)),
48        }
49    }
50
51    pub(crate) fn all(&self) -> Vec<(ProviderId, Budget)> {
52        self.0.lock().unwrap().clone()
53    }
54}
aspero_engine/budget.rs

aspero_engine/
budget.rs
