lean_ctx/core/archive/
register.rs

1// SPDX-License-Identifier: Apache-2.0
2//! One persistence form for the small content-free registers next to the
3//! archive (`delivered`, `presented`, and the proxy's register of
4//! provider-sealed items): 256 JSON shards under
5//! `<data dir>/archives/<name>/<2 hex>.json`, owner-only, replaced atomically
6//! under a bounded file lock, with no in-memory state — a second process reads
7//! exactly what the first one wrote.
8//!
9//! Keys are 64 lowercase hex characters. Values are whatever the register
10//! stores; this module never sees content and logs nothing.
11
12use std::collections::BTreeMap;
13use std::path::PathBuf;
14use std::time::{Duration, Instant};
15
16use serde::{Deserialize, Serialize, de::DeserializeOwned};
17
18/// Upper bound of one shard file; a larger file is not read.
19pub(super) const MAX_SHARD_BYTES: u64 = 4 * 1024 * 1024;
20/// How long an update waits for the register lock before it gives up.
21const LOCK_WAIT: Duration = Duration::from_millis(250);
22
23const SHARD_VERSION: u8 = 1;
24
25/// The bounds of one register.
26#[derive(Debug, Clone, Copy)]
27pub struct Register {
28    /// Directory name under the archive store.
29    pub name: &'static str,
30    /// Keys per shard; the oldest record is evicted first.
31    pub shard_entries: usize,
32    /// A record older than this is neither returned nor kept.
33    pub max_age_seconds: i64,
34}
35
36#[derive(Debug, Serialize, Deserialize)]
37#[serde(deny_unknown_fields)]
38struct Shard<V> {
39    version: u8,
40    entries: BTreeMap<String, Stamped<V>>,
41}
42
43#[derive(Debug, Serialize, Deserialize)]
44#[serde(deny_unknown_fields)]
45struct Stamped<V> {
46    recorded_at: i64,
47    value: V,
48}
49
50pub fn now() -> i64 {
51    chrono::Utc::now().timestamp()
52}
53
54pub(super) fn valid_key(key: &str) -> bool {
55    key.len() == 64
56        && key
57            .bytes()
58            .all(|byte| byte.is_ascii_hexdigit() && !byte.is_ascii_uppercase())
59}
60
61impl Register {
62    fn dir(&self) -> PathBuf {
63        super::archive_base_dir().join(self.name)
64    }
65
66    pub(super) fn shard_path(&self, key: &str) -> PathBuf {
67        self.dir().join(format!("{}.json", &key[..2]))
68    }
69
70    fn fresh<V>(&self, entry: &Stamped<V>, now: i64) -> bool {
71        now.checked_sub(entry.recorded_at)
72            .is_some_and(|age| (0..=self.max_age_seconds).contains(&age))
73    }
74
75    /// A missing shard is empty. An oversized, foreign or undecodable one is
76    /// `None`: nothing is read from it.
77    fn load<V: DeserializeOwned>(&self, key: &str) -> Option<Shard<V>> {
78        let path = self.shard_path(key);
79        match std::fs::symlink_metadata(&path) {
80            Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
81                return Some(Shard {
82                    version: SHARD_VERSION,
83                    entries: BTreeMap::new(),
84                });
85            }
86            Err(_) => return None,
87            Ok(metadata) if metadata.len() > MAX_SHARD_BYTES => return None,
88            Ok(_) => {}
89        }
90        let raw =
91            crate::core::atomic_fs::read_store_file(&super::archive_base_dir(), &path).ok()?;
92        if raw.len() as u64 > MAX_SHARD_BYTES {
93            return None;
94        }
95        serde_json::from_str::<Shard<V>>(&raw)
96            .ok()
97            .filter(|shard| shard.version == SHARD_VERSION)
98    }
99
100    /// The value recorded under `key`, unless it is missing or expired.
101    /// Readers take no lock: a shard is only ever replaced atomically.
102    pub(super) fn read<V: DeserializeOwned>(&self, key: &str, now: i64) -> Option<V> {
103        if !valid_key(key) {
104            return None;
105        }
106        self.load::<V>(key)?
107            .entries
108            .remove(key)
109            .filter(|entry| self.fresh(entry, now))
110            .map(|entry| entry.value)
111    }
112
113    /// Replace the value under `key` with `merge(current)`. `merge` returning
114    /// `None` records nothing and leaves the register as it is. `false`: not
115    /// recorded (malformed key, a bound, the lock, the file system).
116    ///
117    /// An unreadable shard is started afresh: unknown is the safe state of
118    /// everything it held.
119    pub(super) fn update<V: Serialize + DeserializeOwned>(
120        &self,
121        key: &str,
122        now: i64,
123        merge: impl FnOnce(Option<V>) -> Option<V>,
124    ) -> bool {
125        if !valid_key(key) {
126            return false;
127        }
128        let path = self.shard_path(key);
129        self.locked(|| {
130            let mut shard = self.load::<V>(key).unwrap_or(Shard {
131                version: SHARD_VERSION,
132                entries: BTreeMap::new(),
133            });
134            shard.entries.retain(|_, entry| self.fresh(entry, now));
135            let current = shard.entries.remove(key).map(|entry| entry.value);
136            let Some(value) = merge(current) else {
137                return false;
138            };
139            shard.entries.insert(
140                key.to_owned(),
141                Stamped {
142                    recorded_at: now,
143                    value,
144                },
145            );
146            while shard.entries.len() > self.shard_entries {
147                let Some(oldest) = shard
148                    .entries
149                    .iter()
150                    .filter(|(candidate, _)| candidate.as_str() != key)
151                    .min_by_key(|(candidate, entry)| (entry.recorded_at, (*candidate).clone()))
152                    .map(|(candidate, _)| candidate.clone())
153                else {
154                    break;
155                };
156                shard.entries.remove(&oldest);
157            }
158            let Ok(encoded) = serde_json::to_vec(&shard) else {
159                return false;
160            };
161            encoded.len() as u64 <= MAX_SHARD_BYTES && write_private(&path, &encoded)
162        })
163        .unwrap_or(false)
164    }
165
166    /// Record many values at once under one lock, one write per touched
167    /// shard: `merge(key, current, new)` gives the value kept under a key
168    /// that already holds one. Keys that are malformed are skipped. `false`: the
169    /// lock was not obtained or a shard could not be written (its keys are
170    /// then not recorded; unknown is the safe state).
171    pub fn update_many<V: Serialize + DeserializeOwned>(
172        &self,
173        now: i64,
174        items: Vec<(String, V)>,
175        merge: impl Fn(&str, V, V) -> V,
176    ) -> bool {
177        let mut by_shard: BTreeMap<String, Vec<(String, V)>> = BTreeMap::new();
178        for (key, value) in items {
179            if valid_key(&key) {
180                by_shard
181                    .entry(key[..2].to_owned())
182                    .or_default()
183                    .push((key, value));
184            }
185        }
186        if by_shard.is_empty() {
187            return true;
188        }
189        self.locked(|| {
190            let mut written = true;
191            for values in by_shard.into_values() {
192                let Some((first, _)) = values.first() else {
193                    continue;
194                };
195                let path = self.shard_path(first);
196                let mut shard = self.load::<V>(first).unwrap_or(Shard {
197                    version: SHARD_VERSION,
198                    entries: BTreeMap::new(),
199                });
200                shard.entries.retain(|_, entry| self.fresh(entry, now));
201                for (key, value) in values {
202                    let value = match shard.entries.remove(&key) {
203                        Some(current) => merge(&key, current.value, value),
204                        None => value,
205                    };
206                    shard.entries.insert(
207                        key,
208                        Stamped {
209                            recorded_at: now,
210                            value,
211                        },
212                    );
213                }
214                // Bounded in count and in bytes: the oldest record goes first
215                // until the shard fits, so one large shard never stalls every
216                // later write to it until its records age out.
217                let encoded = loop {
218                    if shard.entries.len() <= self.shard_entries
219                        && let Ok(encoded) = serde_json::to_vec(&shard)
220                        && encoded.len() as u64 <= MAX_SHARD_BYTES
221                    {
222                        break Some(encoded);
223                    }
224                    let Some(oldest) = shard
225                        .entries
226                        .iter()
227                        .min_by_key(|(candidate, entry)| (entry.recorded_at, (*candidate).clone()))
228                        .map(|(candidate, _)| candidate.clone())
229                    else {
230                        break None;
231                    };
232                    shard.entries.remove(&oldest);
233                };
234                written &= encoded.is_some_and(|encoded| write_private(&path, &encoded));
235            }
236            written
237        })
238        .unwrap_or(false)
239    }
240
241    /// Every fresh record of every shard with the time it was recorded, for
242    /// a holder that keeps the register in memory and loads it once. A shard
243    /// that is oversized, foreign or undecodable contributes nothing.
244    pub fn read_all<V: DeserializeOwned>(&self, now: i64) -> Vec<(String, i64, V)> {
245        let mut out = Vec::new();
246        for shard in 0..=u8::MAX {
247            let probe = format!("{shard:02x}{}", "0".repeat(62));
248            let Some(loaded) = self.load::<V>(&probe) else {
249                continue;
250            };
251            for (key, entry) in loaded.entries {
252                if valid_key(&key) && key[..2] == probe[..2] && self.fresh(&entry, now) {
253                    out.push((key, entry.recorded_at, entry.value));
254                }
255            }
256        }
257        out
258    }
259
260    /// Run `operation` under the register's exclusive file lock, waiting at
261    /// most [`LOCK_WAIT`]. `None`: directory or lock unavailable, or the
262    /// budget ran out.
263    fn locked<T>(&self, operation: impl FnOnce() -> T) -> Option<T> {
264        use fs2::FileExt;
265
266        let directory = self.dir();
267        std::fs::create_dir_all(&directory).ok()?;
268        let mut options = std::fs::OpenOptions::new();
269        options.read(true).write(true).create(true).truncate(false);
270        #[cfg(unix)]
271        {
272            use std::os::unix::fs::OpenOptionsExt;
273            options.mode(0o600).custom_flags(libc::O_NOFOLLOW);
274        }
275        let lock = options.open(directory.join(".lock")).ok()?;
276        let deadline = Instant::now() + LOCK_WAIT;
277        loop {
278            if FileExt::try_lock_exclusive(&lock).is_ok() {
279                break;
280            }
281            if Instant::now() >= deadline {
282                return None;
283            }
284            std::thread::sleep(Duration::from_millis(5));
285        }
286        let result = operation();
287        let _ = FileExt::unlock(&lock);
288        Some(result)
289    }
290}
291
292fn write_private(path: &std::path::Path, bytes: &[u8]) -> bool {
293    #[cfg(unix)]
294    let permissions = {
295        use std::os::unix::fs::PermissionsExt;
296        Some(std::fs::Permissions::from_mode(0o600))
297    };
298    #[cfg(not(unix))]
299    let permissions: Option<std::fs::Permissions> = None;
300    crate::core::atomic_fs::try_atomic_write(path, bytes, permissions.as_ref()).is_ok()
301}
302
303#[cfg(test)]
304mod tests {
305    use super::*;
306
307    /// A shard over the byte bound evicts its oldest records until it fits,
308    /// instead of refusing every later write until its records age out.
309    #[test]
310    fn a_shard_over_the_byte_bound_evicts_oldest_until_it_fits() {
311        let _data = crate::core::data_dir::isolated_data_dir();
312        let store = Register {
313            name: "bound-test",
314            shard_entries: 64,
315            max_age_seconds: 3600,
316        };
317        let large = "x".repeat(1024 * 1024);
318        let key = |n: u8| format!("aa{n:062}");
319        let now = now();
320        let old: Vec<(String, String)> = (0..3).map(|n| (key(n), large.clone())).collect();
321        assert!(store.update_many(now - 10, old, |_, _, new| new));
322        let new: Vec<(String, String)> = (3..5).map(|n| (key(n), large.clone())).collect();
323        assert!(
324            store.update_many(now, new, |_, _, new| new),
325            "the write fits after eviction"
326        );
327        let kept: Vec<String> = store
328            .read_all::<String>(now)
329            .into_iter()
330            .map(|(key, _, _)| key)
331            .collect();
332        assert!(kept.contains(&key(3)) && kept.contains(&key(4)), "{kept:?}");
333        assert!(!kept.contains(&key(0)), "the oldest record went first");
334    }
335}
lean_ctx/core/archive/register.rs

lean_ctx/core/archive/
register.rs
