lmjtfy.git / packages / budget / src / lib.rs
lib.rsannotatedlib.rssource290 lines · 9.7 KB · raw

Daily budgets shared by every visitor: the arithmetic, and the messages the Worker and its Durable Object exchange. No I/O and no clock; the Durable Object supplies the day and keeps the [Meter]s.

Spending is held before a call and settled after it, as jev-http's ledger does: the hold is the most the call can cost, so two visitors arriving together cannot both spend the last of the day.

8#![forbid(unsafe_code)]

10use serde::{Deserialize, Serialize};

Days since the Unix epoch, UTC: the day Workers AI's allocation resets on.

13pub fn day(unix_ms: f64) -> u64 {
14    (unix_ms / 86_400_000.0).floor().max(0.0) as u64
15}

A UTC day as YYYY-MM-DD, the form Cloudflare's analytics filter takes. (Days-to-civil, after Howard Hinnant's civil_from_days.)

19pub fn date(day: u64) -> String {
20    let z = day as i64 + 719_468;
21    let era = z.div_euclid(146_097);
22    let doe = z.rem_euclid(146_097);
23    let yoe = (doe - doe / 1_460 + doe / 36_524 - doe / 146_096) / 365;
24    let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
25    let mp = (5 * doy + 2) / 153;
26    let d = doy - (153 * mp + 2) / 5 + 1;
27    let m = if mp < 10 { mp + 3 } else { mp - 9 };
28    let y = yoe + era * 400 + i64::from(m <= 2);
29    format!("{y:04}-{m:02}-{d:02}")
30}

What has been spent, or is held, on one day.

33#[derive(Clone, Copy, Debug, Default, PartialEq, Serialize, Deserialize)]
34pub struct Meter {
35    day: u64,
36    used: f64,
37}

An amount set aside for a call that has not finished.

40#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
41pub struct Hold {
42    day: u64,
43    amount: f64,
44}

The day's budget cannot cover the call.

47#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
48pub struct Exhausted {
49    pub remaining: f64,
50}

52impl Meter {
53    fn roll(&mut self, today: u64) {
54        if self.day != today {
55            *self = Meter { day: today, used: 0.0 };
56        }
57    }

What is spent or held today.

60    pub fn used(&self, today: u64) -> f64 {
61        if self.day == today { self.used } else { 0.0 }
62    }

Sets amount aside, if today's per_day still covers it.

65    pub fn hold(&mut self, today: u64, amount: f64, per_day: f64) -> Result<Hold, Exhausted> {
66        self.roll(today);
67        let remaining = (per_day - self.used).max(0.0);
68        if amount > remaining {
69            return Err(Exhausted { remaining });
70        }
71        self.used += amount;
72        Ok(Hold { day: today, amount })
73    }

Takes in a total reported from outside: the account's own figure, which counts spending this meter never saw (other programs on the same allocation). The meter is raised to it and never lowered by it, because the figure lags: spending held or settled here since it was measured is already on top.

80    pub fn observe(&mut self, today: u64, total: f64) {
81        self.roll(today);
82        self.used = self.used.max(total);
83    }

Replaces a hold with what the call really cost. A hold from a day that has since ended was never counted against today, so only the cost is.

87    pub fn settle(&mut self, today: u64, hold: Hold, actual: f64) {
88        self.roll(today);
89        if hold.day == today {
90            self.used -= hold.amount;
91        }
92        self.used = (self.used + actual).max(0.0);
93    }
94}

Which budget.

97#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
98#[serde(rename_all = "snake_case")]
99pub enum Which {

Workers AI neurons, for the LLM.

101    Neurons,

Dollars, for Jev.

103    JevDollars,
104}

106impl Which {

The key the meter is stored under.

108    pub fn key(self) -> &'static str {
109        match self {
110            Which::Neurons => "neurons",
111            Which::JevDollars => "jev_dollars",
112        }
113    }
114}

How long a visitor's count runs before it starts again.

117pub const VISIT_WINDOW_MS: f64 = 60_000.0;

How many requests each visitor has made in their current minute, so one person cannot spend the day's budgets for everyone. Kept in memory only: who a visitor is (their address) is never written anywhere, and a count that is lost when the object restarts only lets someone start again.

123#[derive(Debug, Default)]
124pub struct Visits(std::collections::HashMap<(String, String), (f64, u32)>);

126impl Visits {

Counts one request of kind what by who, and says whether it is within per_minute. A request over the limit is not counted, so hammering does not push the end of the wait further off.

130    pub fn admit(&mut self, what: &str, who: &str, now_ms: f64, per_minute: u32) -> bool {
131        // Whoever's minute is over is forgotten.
132        self.0.retain(|_, (started, _)| now_ms - *started < VISIT_WINDOW_MS);
133        let (_, count) = self.0.entry((what.to_owned(), who.to_owned())).or_insert((now_ms, 0));
134        if *count >= per_minute {
135            return false;
136        }
137        *count += 1;
138        true
139    }

How many visitors are being counted right now.

142    pub fn len(&self) -> usize {
143        self.0.len()
144    }

146    pub fn is_empty(&self) -> bool {
147        self.0.is_empty()
148    }
149}

The Worker to the Durable Object.

152#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
153#[serde(rename_all = "snake_case")]
154pub enum Ask {
155    Hold { which: Which, amount: f64, per_day: f64 },
156    Settle { which: Which, hold: Hold, actual: f64 },

What is used so far today, for the page.

158    Status,

One request of kind what by the visitor who: is it within per_minute?

161    Visit { what: String, who: String, per_minute: u32 },
162}

Whose count of the neurons this is.

165#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
166#[serde(rename_all = "snake_case")]
167pub enum Counted {

Cloudflare's figure for the whole account, with this Worker's own spending since it was read on top.

170    Account,

Only what this Worker spent: the account's figure could not be read.

172    Own,
173}

Today's use of both budgets.

176#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
177pub struct Status {
178    pub neurons: f64,
179    pub counted: Counted,
180    pub jev_dollars: f64,
181}

The Durable Object to the Worker.

184#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
185#[serde(rename_all = "snake_case")]
186pub enum Told {
187    Held(Hold),
188    Exhausted(Exhausted),
189    Settled,
190    Status(Status),

Whether the visit is within its limit.

192    Visit(bool),
193}

195#[cfg(test)]
196mod tests {
197    use super::*;
198
199    #[test]
200    fn a_visitor_is_stopped_at_the_limit_and_starts_again_after_a_minute() {
201        let mut visits = Visits::default();
202        for n in 0..10 {
203            assert!(visits.admit("ask", "a", f64::from(n) * 1000.0, 10), "{n}");
204        }
205        assert!(!visits.admit("ask", "a", 10_000.0, 10));
206        assert!(!visits.admit("ask", "a", 59_999.0, 10));
207        // A minute after the first one, not after the last refusal.
208        assert!(visits.admit("ask", "a", 60_000.0, 10));
209    }
210
211    #[test]
212    fn each_visitor_and_each_kind_of_request_is_counted_apart() {
213        let mut visits = Visits::default();
214        assert!(visits.admit("ask", "a", 0.0, 1));
215        assert!(!visits.admit("ask", "a", 1.0, 1));
216        assert!(visits.admit("ask", "b", 1.0, 1));
217        assert!(visits.admit("glance", "a", 1.0, 1));
218    }
219
220    #[test]
221    fn nobody_is_remembered_past_their_minute() {
222        let mut visits = Visits::default();
223        visits.admit("ask", "a", 0.0, 10);
224        visits.admit("ask", "b", 30_000.0, 10);
225        assert_eq!(visits.len(), 2);
226        visits.admit("ask", "c", 61_000.0, 10);
227        assert_eq!(visits.len(), 2);
228        visits.admit("ask", "d", 95_000.0, 10);
229        assert_eq!(visits.len(), 2);
230    }
231
232    #[test]
233    fn a_day_starts_at_midnight_utc() {
234        assert_eq!(day(0.0), 0);
235        assert_eq!(day(86_399_999.0), 0);
236        assert_eq!(day(86_400_000.0), 1);
237    }
238
239    #[test]
240    fn a_day_is_written_as_a_date() {
241        assert_eq!(date(0), "1970-01-01");
242        assert_eq!(date(day(1_790_967_600_000.0)), "2026-10-02");
243        assert_eq!(date(19_782), "2024-02-29");
244        assert_eq!(date(19_783), "2024-03-01");
245    }
246
247    #[test]
248    fn an_observed_total_raises_the_meter_and_never_lowers_it() {
249        let mut meter = Meter::default();
250        // The account has spent 900 that this meter never saw.
251        meter.observe(7, 900.0);
252        let hold = meter.hold(7, 50.0, 1000.0).unwrap();
253        meter.settle(7, hold, 16.0);
254        assert_eq!(meter.used(7), 916.0);
255        // The account's figure has not caught up with the 16 yet.
256        meter.observe(7, 905.0);
257        assert_eq!(meter.used(7), 916.0);
258        // Now it has, and something else spent 30 more.
259        meter.observe(7, 946.0);
260        assert_eq!(meter.used(7), 946.0);
261        assert_eq!(meter.hold(7, 60.0, 1000.0), Err(Exhausted { remaining: 54.0 }));
262    }
263
264    #[test]
265    fn a_hold_is_refused_once_the_day_cannot_cover_it() {
266        let mut meter = Meter::default();
267        let first = meter.hold(7, 60.0, 100.0).unwrap();
268        assert_eq!(meter.hold(7, 60.0, 100.0), Err(Exhausted { remaining: 40.0 }));
269        meter.settle(7, first, 10.0);
270        assert_eq!(meter.used(7), 10.0);
271        assert!(meter.hold(7, 60.0, 100.0).is_ok());
272    }
273
274    #[test]
275    fn a_new_day_starts_empty() {
276        let mut meter = Meter::default();
277        meter.hold(7, 100.0, 100.0).unwrap();
278        assert_eq!(meter.used(8), 0.0);
279        assert!(meter.hold(8, 100.0, 100.0).is_ok());
280    }
281
282    #[test]
283    fn a_hold_settled_after_midnight_costs_only_what_was_spent() {
284        let mut meter = Meter::default();
285        let late = meter.hold(7, 50.0, 100.0).unwrap();
286        meter.hold(8, 20.0, 100.0).unwrap();
287        meter.settle(8, late, 5.0);
288        assert_eq!(meter.used(8), 25.0);
289    }
290}