didbot_policy_count/
counts.rs

1//! The counts, and the ring each one is kept in.
2
3use std::collections::HashMap;
4use std::sync::Mutex;
5
6use crate::statement::Window;
7
8/// Hours in the ring, which is how long a day window is.
9pub const BUCKETS: usize = 24;
10
11/// Seconds in a bucket.
12const SECONDS_PER_BUCKET: i64 = 3_600;
13
14/// How many counts one [`Counts`] holds before it drops the ones that have
15/// aged out. A deployment holds far fewer accounts than this; the number
16/// is here so a deployment whose accounts come and go cannot grow the map
17/// without bound between deletions.
18pub const HELD: usize = 65_536;
19
20/// The hour `now` falls in, in whole hours since the epoch.
21#[must_use]
22pub fn hour_of(now: time::OffsetDateTime) -> i64 {
23    now.unix_timestamp().div_euclid(SECONDS_PER_BUCKET)
24}
25
26/// A ring of hourly counters, newest at `head`.
27#[derive(Debug, Clone, Copy)]
28struct Ring {
29    head: i64,
30    counts: [u32; BUCKETS],
31}
32
33impl Ring {
34    fn new(at: i64) -> Self {
35        Self {
36            head: at,
37            counts: [0; BUCKETS],
38        }
39    }
40
41    /// Moves the head to `at`, clearing the buckets passed on the way. An
42    /// instant at or before the head leaves the head where it is: `now` is
43    /// supplied per write, and a write presented with an earlier instant
44    /// than one already counted lands in the newest bucket rather than
45    /// reopening a bucket that has closed.
46    fn advance(&mut self, at: i64) {
47        if at <= self.head {
48            return;
49        }
50        let steps = (at - self.head).min(BUCKETS as i64);
51        for step in 1..=steps {
52            self.counts[index(self.head + step)] = 0;
53        }
54        self.head = at;
55    }
56
57    fn read(&self, window: Window) -> u64 {
58        match window {
59            Window::Hour => u64::from(self.counts[index(self.head)]),
60            Window::Day => self.counts.iter().copied().map(u64::from).sum(),
61        }
62    }
63
64    /// Counts one more in the newest bucket. Saturates: something that has
65    /// reached `u32::MAX` in an hour is past every limit there is, and
66    /// wrapping would put it back under one.
67    fn add(&mut self) {
68        let slot = &mut self.counts[index(self.head)];
69        *slot = slot.saturating_add(1);
70    }
71
72    fn aged_out(&self, at: i64) -> bool {
73        self.head + (BUCKETS as i64) <= at
74    }
75}
76
77fn index(hour: i64) -> usize {
78    (hour.rem_euclid(BUCKETS as i64)) as usize
79}
80
81/// What one count is kept under: the statement, and the account when the
82/// statement counts each account on its own.
83type Key = (String, Option<String>);
84
85/// Every count this process holds.
86///
87/// Memory only. A restart is a fresh window for every statement, every
88/// account and the deployment, which is `plan/policy.md`'s
89/// "the reset-on-restart behaviour must be stated rather than discovered".
90///
91/// [`BUCKETS`] counters and a head per account per statement: about 150
92/// bytes with a DID for a key. A statement counting each account on its
93/// own is that much per account; one counting the deployment is that much
94/// once.
95#[derive(Default)]
96pub struct Counts {
97    rings: Mutex<HashMap<Key, Ring>>,
98}
99
100impl Counts {
101    /// Nothing counted yet.
102    #[must_use]
103    pub fn new() -> Self {
104        Self::default()
105    }
106
107    /// What `key` holds in `window` at `at`.
108    #[must_use]
109    pub fn read(&self, counter: &str, account: Option<&str>, window: Window, at: i64) -> u64 {
110        let key = (counter.to_owned(), account.map(str::to_owned));
111        let mut rings = self.lock();
112        match rings.get_mut(&key) {
113            Some(ring) => {
114                ring.advance(at);
115                ring.read(window)
116            }
117            None => 0,
118        }
119    }
120
121    /// Counts one more against `key` at `at`.
122    pub fn add(&self, counter: &str, account: Option<&str>, at: i64) {
123        let key = (counter.to_owned(), account.map(str::to_owned));
124        let mut rings = self.lock();
125        if rings.len() >= HELD {
126            rings.retain(|_, ring| !ring.aged_out(at));
127        }
128        let ring = rings.entry(key).or_insert_with(|| Ring::new(at));
129        ring.advance(at);
130        ring.add();
131    }
132
133    /// Drops everything counted against `account`, whatever statement
134    /// counted it.
135    pub fn forget(&self, account: &str) {
136        self.lock()
137            .retain(|(_, held), _| held.as_deref() != Some(account));
138    }
139
140    /// Drops every count, which is what a restart does. Not reachable from
141    /// any request: nothing on the write path calls it.
142    pub fn forget_everything(&self) {
143        self.lock().clear();
144    }
145
146    /// How many counts are held.
147    #[must_use]
148    pub fn held(&self) -> usize {
149        self.lock().len()
150    }
151
152    fn lock(&self) -> std::sync::MutexGuard<'_, HashMap<Key, Ring>> {
153        self.rings
154            .lock()
155            .unwrap_or_else(std::sync::PoisonError::into_inner)
156    }
157}
158
159#[cfg(test)]
160mod tests {
161    use super::*;
162
163    const COUNTER: &str = "a statement";
164
165    #[test]
166    fn the_hour_is_the_newest_bucket_and_the_day_is_the_ring() {
167        let counts = Counts::new();
168        for hour in 0..5 {
169            counts.add(COUNTER, Some("did:web:one"), hour);
170        }
171        assert_eq!(
172            counts.read(COUNTER, Some("did:web:one"), Window::Hour, 4),
173            1
174        );
175        assert_eq!(counts.read(COUNTER, Some("did:web:one"), Window::Day, 4), 5);
176        // A whole ring later, the day window holds none of it.
177        let later = 4 + BUCKETS as i64;
178        assert_eq!(
179            counts.read(COUNTER, Some("did:web:one"), Window::Day, later),
180            0
181        );
182    }
183
184    #[test]
185    fn accounts_and_statements_each_count_for_themselves() {
186        let counts = Counts::new();
187        counts.add(COUNTER, Some("did:web:one"), 10);
188        counts.add(COUNTER, Some("did:web:two"), 10);
189        counts.add("another statement", Some("did:web:one"), 10);
190        counts.add(COUNTER, None, 10);
191        assert_eq!(
192            counts.read(COUNTER, Some("did:web:one"), Window::Hour, 10),
193            1
194        );
195        assert_eq!(counts.read(COUNTER, None, Window::Hour, 10), 1);
196        assert_eq!(counts.held(), 4);
197        counts.forget("did:web:one");
198        assert_eq!(
199            counts.read(COUNTER, Some("did:web:one"), Window::Hour, 10),
200            0
201        );
202        assert_eq!(
203            counts.read(COUNTER, Some("did:web:two"), Window::Hour, 10),
204            1
205        );
206        assert_eq!(
207            counts.read(COUNTER, None, Window::Hour, 10),
208            1,
209            "a deleted account takes nothing of the deployment's with it"
210        );
211    }
212}
didbot_policy_count/counts.rs

didbot_policy_count/
counts.rs
