diff --git a/src/db/collection_id.rs b/src/db/collection_id.rs new file mode 100644 --- /dev/null +++ b/src/db/collection_id.rs @@ -0,0 +1,308 @@ +//! collection id encoding for composite keys. +//! +//! collections are interned to integer ids and written into composite keys as +//! LEB128 varints in place of the NSID text, so `app.bsky.feed.like` costs 1 +//! byte instead of 18 in every record key. +//! +//! # why a varint +//! +//! a varint decouples id headroom from key size: one byte covers the handful of +//! collections that carry almost every record, five bytes reach 2^35, and only +//! rare collections pay the wide encoding. a fixed width would force a choice +//! between headroom and size. +//! +//! # properties this module guarantees +//! +//! - **prefix-free.** the terminating byte has its high bit clear, so no +//! encoding is a prefix of another. the segment is self-delimiting, so +//! composite keys need no separator after it and a prefix scan over +//! `{did}|{id}` still isolates exactly one collection. +//! - **separator-safe.** the segment is never split on `keys::SEP`, so a varint +//! byte that happens to be `0x7C` is harmless. +//! - **terminating byte below `0x80`.** incrementing the last byte of a key +//! prefix to form an exclusive upper bound can never overflow. +//! - **never ASCII-alphabetic in the first byte.** an NSID always starts with +//! `[a-zA-Z]`, so [`is_interned`] tells an interned segment from legacy NSID +//! text with no ambiguity. the migration relies on this to stay idempotent: +//! a multi-byte id sorts *after* the text it replaced, so a chunked rewrite +//! does re-scan its own output. +//! +//! the encoding is deliberately **not** order-preserving: `129` (`[0x81, 0x01]`) +//! sorts after `256` (`[0x80, 0x02]`). id order was never NSID order, and a +//! `(did, collection)` group stays contiguous because it is an exact key prefix +//! rather than a range. if ordered ids are ever needed, a length-prefixed +//! varint is order-preserving at the same byte cost and drops straight in. + +use miette::Result; + +/// widest LEB128 encoding of a [`CollectionId`]. +pub const MAX_ENCODED_LEN: usize = 5; + +/// an interned collection. +/// +/// construct only through [`CollectionId::first`] and [`CollectionId::next`] so +/// that reserved ids are never handed out. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub struct CollectionId(u32); + +/// true if `id`'s single-byte encoding would collide with NSID text. +/// +/// only single-byte encodings can collide: a multi-byte encoding's first byte +/// has its high bit set, so it is at least `0x80` and never ASCII. so this only +/// removes `A..=Z` and `a..=z` from the single-byte range, leaving 76 ids that +/// still encode in one byte. +fn is_reserved(id: u32) -> bool { + id < 0x80 && (id as u8).is_ascii_alphabetic() +} + +impl CollectionId { + /// the id handed out from an empty dictionary. + pub const fn first() -> Self { + Self(0) + } + + /// the next allocatable id, skipping reserved values. + /// + /// `None` only once the id space is exhausted, which takes 2^32 distinct + /// collections. + pub fn next(self) -> Option { + let mut next = self.0.checked_add(1)?; + while is_reserved(next) { + next = next.checked_add(1)?; + } + Some(Self(next)) + } + + /// reconstruct an id read back from the dictionary. + /// + /// rejects reserved values, so a corrupt or hand-written dictionary entry + /// cannot produce an id whose encoding is indistinguishable from NSID text. + pub fn from_raw(id: u32) -> Result { + if is_reserved(id) { + miette::bail!("collection id {id} is reserved (would encode as ASCII text)"); + } + Ok(Self(id)) + } + + pub fn get(self) -> u32 { + self.0 + } + + /// bytes this id occupies in a composite key. + pub fn encoded_len(self) -> usize { + let mut len = 1; + let mut rest = self.0 >> 7; + while rest != 0 { + len += 1; + rest >>= 7; + } + len + } + + /// append the LEB128 encoding to `buf`. + /// + /// the allocation-free counterpart of `jacquard_common`'s `encode_uvarint`, + /// which returns a fresh `Vec` per call and so cannot be used on a path that + /// builds a key per record. + pub fn write_to_vec(self, buf: &mut Vec) { + let mut x = self.0; + while x >= 0x80 { + buf.push((x as u8) | 0x80); + x >>= 7; + } + buf.push(x as u8); + } + + /// decode an id from the front of `bytes`, returning it with the remainder. + pub fn read_from(bytes: &[u8]) -> Result<(Self, &[u8])> { + let mut id: u32 = 0; + for (i, &byte) in bytes.iter().take(MAX_ENCODED_LEN).enumerate() { + let payload = u32::from(byte & 0x7F); + // the 5th byte of a u32 varint carries only 4 usable bits + id |= payload + .checked_shl(7 * i as u32) + .filter(|shifted| shifted >> (7 * i as u32) == payload) + .ok_or_else(|| miette::miette!("collection id overflows u32"))?; + if byte & 0x80 == 0 { + return Ok((Self::from_raw(id)?, &bytes[i + 1..])); + } + } + if bytes.is_empty() { + miette::bail!("collection id is missing"); + } + miette::bail!("collection id is truncated or overlong"); + } +} + +/// whether `segment` starts an interned collection id rather than NSID text. +/// +/// an NSID matches `^[a-zA-Z]` per the atproto spec, and no allocatable id +/// encodes to a leading ASCII letter, so the two forms never overlap. an empty +/// segment is neither, and is reported as not interned so the caller surfaces it +/// as malformed legacy text rather than decoding garbage. +pub fn is_interned(segment: &[u8]) -> bool { + segment + .first() + .is_some_and(|byte| !byte.is_ascii_alphabetic()) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// every id an allocator can hand out, up to `limit` allocations. + fn allocatable(limit: usize) -> Vec { + let mut ids = vec![CollectionId::first()]; + while ids.len() < limit { + ids.push(ids.last().expect("non-empty").next().expect("id space")); + } + ids + } + + fn encode(id: CollectionId) -> Vec { + let mut buf = Vec::new(); + id.write_to_vec(&mut buf); + buf + } + + #[test] + fn roundtrips_through_the_encoding() -> Result<()> { + for id in allocatable(5_000) { + let mut buf = encode(id); + buf.extend_from_slice(b"trailing"); + let (decoded, rest) = CollectionId::read_from(&buf)?; + assert_eq!(decoded, id, "id {} did not roundtrip", id.get()); + assert_eq!(rest, b"trailing"); + } + Ok(()) + } + + #[test] + fn encoded_len_matches_what_is_written() { + for id in allocatable(5_000) { + assert_eq!(encode(id).len(), id.encoded_len(), "id {}", id.get()); + assert!(id.encoded_len() <= MAX_ENCODED_LEN); + } + } + + /// the property that lets composite keys drop the separator after the id. + #[test] + fn no_encoding_is_a_prefix_of_another() { + let encodings: Vec> = allocatable(2_000).into_iter().map(encode).collect(); + for (i, a) in encodings.iter().enumerate() { + for (j, b) in encodings.iter().enumerate() { + if i == j { + continue; + } + assert!(!b.starts_with(a), "{a:?} is a prefix of {b:?}"); + } + } + } + + /// the property that keeps the migration idempotent, and that lets a legacy + /// key be told from a migrated one. + #[test] + fn no_allocatable_id_encodes_to_leading_ascii_text() { + for id in allocatable(10_000) { + let encoded = encode(id); + let first = encoded[0]; + assert!( + !first.is_ascii_alphabetic(), + "id {} encodes to leading {first:#x}, which reads as NSID text", + id.get() + ); + assert!(is_interned(&encoded), "id {} not detected", id.get()); + } + } + + /// the property that keeps prefix-to-exclusive-upper-bound from overflowing. + #[test] + fn the_terminating_byte_is_always_below_0x80() { + for id in allocatable(5_000) { + let encoded = encode(id); + assert!( + *encoded.last().expect("non-empty") < 0x80, + "id {}", + id.get() + ); + } + } + + #[test] + fn exactly_76_ids_still_fit_in_one_byte() { + let single: Vec = allocatable(200) + .into_iter() + .filter(|id| id.encoded_len() == 1) + .collect(); + assert_eq!(single.len(), 76); + // the ranges left over after removing A-Z and a-z + assert_eq!(single.first().copied().map(CollectionId::get), Some(0)); + assert_eq!(single.last().copied().map(CollectionId::get), Some(127)); + assert!( + single + .iter() + .all(|id| !encode(*id)[0].is_ascii_alphabetic()) + ); + } + + #[test] + fn allocation_skips_both_reserved_letter_ranges() { + let ids: Vec = allocatable(200) + .into_iter() + .map(CollectionId::get) + .collect(); + // 'A' is 0x41 and 'Z' is 0x5A, 'a' is 0x61 and 'z' is 0x7A + assert!(ids.contains(&0x40)); + assert!(ids.contains(&0x5B)); + assert!(ids.contains(&0x60)); + assert!(ids.contains(&0x7B)); + for reserved in (0x41..=0x5Au32).chain(0x61..=0x7A) { + assert!(!ids.contains(&reserved), "{reserved:#x} was allocated"); + } + } + + #[test] + fn is_interned_rejects_every_legal_nsid_first_character() { + for first in (b'a'..=b'z').chain(b'A'..=b'Z') { + let segment = [first, b'p', b'p']; + assert!(!is_interned(&segment), "{} read as interned", first as char); + } + } + + #[test] + fn is_interned_says_no_for_an_empty_segment() { + assert!(!is_interned(&[])); + } + + #[test] + fn from_raw_rejects_reserved_ids() { + assert!(CollectionId::from_raw(0x41).is_err()); + assert!(CollectionId::from_raw(0x7A).is_err()); + assert!(CollectionId::from_raw(0x40).is_ok()); + assert!(CollectionId::from_raw(0x80).is_ok()); + assert!(CollectionId::from_raw(u32::MAX).is_ok()); + } + + #[test] + fn read_from_rejects_malformed_input() { + assert!(CollectionId::read_from(&[]).is_err()); + // continuation bit set with nothing following + assert!(CollectionId::read_from(&[0x81]).is_err()); + // six continuation bytes cannot terminate within a u32 + assert!(CollectionId::read_from(&[0x80, 0x80, 0x80, 0x80, 0x80, 0x01]).is_err()); + // decodes to 'a', which is reserved and must not appear on disk + assert!(CollectionId::read_from(&[0x61]).is_err()); + } + + #[test] + fn read_from_rejects_ids_past_u32() { + // 2^32 - 1 is the largest value that may decode + let max = CollectionId::from_raw(u32::MAX).expect("u32::MAX allocatable"); + let encoded = encode(max); + assert_eq!(encoded.len(), MAX_ENCODED_LEN); + assert_eq!(CollectionId::read_from(&encoded).expect("decodes").0, max); + + // 2^32, one past the top, must not silently wrap + assert!(CollectionId::read_from(&[0x80, 0x80, 0x80, 0x80, 0x10]).is_err()); + } +} diff --git a/src/db/keyspaces.rs b/src/db/keyspaces.rs --- a/src/db/keyspaces.rs +++ b/src/db/keyspaces.rs @@ -105,7 +105,6 @@ self.blocks.range(range) } - #[cfg(test)] #[allow(dead_code)] pub(crate) fn stage_block(&self, batch: &mut fjall::OwnedWriteBatch, key: K, value: V) diff --git a/src/db/mod.rs b/src/db/mod.rs --- a/src/db/mod.rs +++ b/src/db/mod.rs @@ -11,6 +11,10 @@ use std::sync::{Arc, Mutex}; use url::Url; +// consumed once composite keys carry ids instead of NSID text; until then the +// encoding invariants the key format depends on are held up by its own tests +#[allow(dead_code)] +pub mod collection_id; pub mod compaction; pub mod counts; #[cfg(feature = "indexer")]