1//! Daily budgets shared by every visitor: the arithmetic, and the messages
2//! the Worker and its Durable Object exchange. No I/O and no clock; the
3//! Durable Object supplies the day and keeps the [`Meter`]s.
4//!
5//! Spending is held before a call and settled after it, as jev-http's ledger
6//! does: the hold is the most the call can cost, so two visitors arriving
7//! together cannot both spend the last of the day.
8#![forbid(unsafe_code)]
9
10use serde::{Deserialize, Serialize};
11
12/// 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}
16
17/// A UTC day as `YYYY-MM-DD`, the form Cloudflare's analytics filter takes.
18/// (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}
31
32/// 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}
38
39/// 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}
45
46/// The day's budget cannot cover the call.
47#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
48pub struct Exhausted {
49    pub remaining: f64,
50}
51
52impl Meter {
53    fn roll(&mut self, today: u64) {
54        if self.day != today {
55            *self = Meter { day: today, used: 0.0 };
56        }
57    }
58
59    /// 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    }
63
64    /// 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    }
74
75    /// Takes in a total reported from outside: the account's own figure,
76    /// which counts spending this meter never saw (other programs on the
77    /// same allocation). The meter is raised to it and never lowered by it,
78    /// because the figure lags: spending held or settled here since it was
79    /// 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    }
84
85    /// Replaces a hold with what the call really cost. A hold from a day that
86    /// 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}
95
96/// Which budget.
97#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
98#[serde(rename_all = "snake_case")]
99pub enum Which {
100    /// Workers AI neurons, for the LLM.
101    Neurons,
102    /// Dollars, for Jev.
103    JevDollars,
104}
105
106impl Which {
107    /// 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}
115
116/// How long a visitor's count runs before it starts again.
117pub const VISIT_WINDOW_MS: f64 = 60_000.0;
118
119/// How many requests each visitor has made in their current minute, so one
120/// person cannot spend the day's budgets for everyone. Kept in memory only:
121/// who a visitor is (their address) is never written anywhere, and a count
122/// 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)>);
125
126impl Visits {
127    /// Counts one request of kind `what` by `who`, and says whether it is
128    /// within `per_minute`. A request over the limit is not counted, so
129    /// 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    }
140
141    /// How many visitors are being counted right now.
142    pub fn len(&self) -> usize {
143        self.0.len()
144    }
145
146    pub fn is_empty(&self) -> bool {
147        self.0.is_empty()
148    }
149}
150
151/// 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 },
157    /// What is used so far today, for the page.
158    Status,
159    /// One request of kind `what` by the visitor `who`: is it within
160    /// `per_minute`?
161    Visit { what: String, who: String, per_minute: u32 },
162}
163
164/// 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 {
168    /// Cloudflare's figure for the whole account, with this Worker's own
169    /// spending since it was read on top.
170    Account,
171    /// Only what this Worker spent: the account's figure could not be read.
172    Own,
173}
174
175/// 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}
182
183/// 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),
191    /// Whether the visit is within its limit.
192    Visit(bool),
193}
194
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}