Skip to main content

rgx/
pagination.rs

1//! Server-side pagination store: the daemon keeps the small cursor blob (query + keyset position) for
2//! a couple of minutes and hands the client a short opaque token in its place, so the printed
3//! `--cursor` is tiny instead of a base64 blob. The blob is exactly what the self-contained cursor
4//! used to carry — paging still re-runs the search, so memory is a few dozen bytes per live page, not
5//! the result set. Tokens are stamped with a per-process session so a restarted daemon's old tokens
6//! miss cleanly (the client just re-runs the search). `take` is single-use: following a page deletes
7//! the token it came from.
8
9use std::collections::HashMap;
10use std::time::{Duration, Instant};
11
12/// How long a minted token stays resolvable. Long enough for an agent to read a page and ask for the
13/// next, short enough that abandoned paginations cost nothing.
14pub const DEFAULT_TTL: Duration = Duration::from_secs(120);
15
16pub struct PaginationStore {
17    /// Per-process stamp (hex) prefixed onto every token; a token from another process can't resolve.
18    session: String,
19    counter: u64,
20    ttl: Duration,
21    entries: HashMap<u64, (Vec<u8>, Instant)>,
22}
23
24impl PaginationStore {
25    pub fn new(session: u64, ttl: Duration) -> Self {
26        Self {
27            // 32 bits of the seed is plenty to make a previous daemon's tokens miss; keeps the prefix
28            // (and thus the printed token) short.
29            session: format!("{:08x}", session as u32),
30            counter: 0,
31            ttl,
32            entries: HashMap::new(),
33        }
34    }
35
36    /// Store `blob` and return its token. `now` is injected so the store stays testable.
37    pub fn store(&mut self, blob: Vec<u8>, now: Instant) -> String {
38        self.evict(now);
39        self.counter += 1;
40        let id = self.counter;
41        self.entries.insert(id, (blob, now));
42        format!("{}{id:x}", self.session)
43    }
44
45    /// Resolve and consume `token`, returning its blob if present and unexpired. A token minted by a
46    /// different session (e.g. a previous daemon) or already taken returns `None`.
47    pub fn take(&mut self, token: &str, now: Instant) -> Option<Vec<u8>> {
48        self.evict(now);
49        let rest = token.strip_prefix(&self.session)?;
50        let id: u64 = u64::from_str_radix(rest, 16).ok()?;
51        let (blob, minted) = self.entries.remove(&id)?;
52        (now.duration_since(minted) < self.ttl).then_some(blob)
53    }
54
55    fn evict(&mut self, now: Instant) {
56        let ttl = self.ttl;
57        self.entries
58            .retain(|_, (_, minted)| now.duration_since(*minted) < ttl);
59    }
60
61    #[cfg(test)]
62    fn len(&self) -> usize {
63        self.entries.len()
64    }
65}
66
67#[cfg(test)]
68mod tests {
69    use super::*;
70
71    #[test]
72    fn store_then_take_roundtrips_and_is_single_use() {
73        let mut s = PaginationStore::new(0xABCD, DEFAULT_TTL);
74        let t0 = Instant::now();
75        let tok = s.store(b"hello".to_vec(), t0);
76        assert!(tok.starts_with("0000abcd"));
77        assert_eq!(s.take(&tok, t0).as_deref(), Some(&b"hello"[..]));
78        // single use: the token is gone after the first take.
79        assert_eq!(s.take(&tok, t0), None);
80        assert_eq!(s.len(), 0);
81    }
82
83    #[test]
84    fn expired_token_and_foreign_session_miss() {
85        let mut s = PaginationStore::new(1, Duration::from_secs(120));
86        let t0 = Instant::now();
87        let tok = s.store(b"x".to_vec(), t0);
88        assert_eq!(s.take(&tok, t0 + Duration::from_secs(121)), None);
89
90        // A token shaped for a different session never resolves.
91        let other = PaginationStore::new(2, DEFAULT_TTL).store(b"y".to_vec(), t0);
92        let mut s = PaginationStore::new(1, DEFAULT_TTL);
93        assert_eq!(s.take(&other, t0), None);
94        assert_eq!(s.take("not-a-token", t0), None);
95    }
96}