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

aspero_engine/
budget.rs
