rushdb_storage/
id_alloc.rs

1//! Store-wide allocation of monotone ids (record and relationship ordinals,
2//! property families and definitions).
3//!
4//! An id is reserved when a statement is planned, like a PostgreSQL sequence:
5//! two concurrent transactions never receive the same id, and an id is never
6//! handed out twice, even when the transaction that reserved it aborts (that
7//! leaves a gap). The durable counter row is a watermark that commits advance
8//! to `max(current, id + 1)` ([`crate::Txn::set_watermark`]); it is neither a
9//! read dependency nor a write conflict of a transaction, so transactions that
10//! create disjoint records do not conflict on it.
11//!
12//! The reservations live in memory. After a restart the next id comes from
13//! the durable watermark: every id of a committed write is below it, and an id
14//! reserved by a transaction that never committed is free again, since that
15//! transaction did not survive the restart.
16
17use std::collections::HashMap;
18use std::sync::Mutex;
19
20use crate::StorageError;
21
22/// Ids reserved so far, per counter key: the next id each counter hands out
23/// when its durable watermark is not higher.
24#[derive(Debug, Default)]
25pub struct IdAllocator {
26    next: Mutex<HashMap<Vec<u8>, u64>>,
27}
28
29impl IdAllocator {
30    /// Reserve the next id of the counter at `key`, given the counter's
31    /// watermark as the caller's view sees it.
32    ///
33    /// # Errors
34    /// Returns [`StorageError::Builder`] if the reservation table's lock is
35    /// poisoned, or [`StorageError::ResourceExhausted`] if the counter would
36    /// pass `u64::MAX`.
37    pub fn reserve(&self, key: &[u8], watermark: u64) -> Result<u64, StorageError> {
38        let mut next = self
39            .next
40            .lock()
41            .map_err(|_| StorageError::Builder("id allocator lock poisoned".into()))?;
42        let slot = next.entry(key.to_vec()).or_insert(0);
43        let id = (*slot).max(watermark);
44        *slot = id
45            .checked_add(1)
46            .ok_or_else(|| StorageError::ResourceExhausted("id counter exhausted".into()))?;
47        Ok(id)
48    }
49}
50
51#[cfg(test)]
52mod tests {
53    use super::IdAllocator;
54
55    #[test]
56    fn reservations_are_unique_and_follow_the_watermark() {
57        let ids = IdAllocator::default();
58        // Two planners that read the same stale watermark get distinct ids.
59        assert_eq!(ids.reserve(b"c", 5).unwrap(), 5);
60        assert_eq!(ids.reserve(b"c", 5).unwrap(), 6);
61        // A watermark advanced by a commit moves the next id past it.
62        assert_eq!(ids.reserve(b"c", 10).unwrap(), 10);
63        assert_eq!(ids.reserve(b"c", 0).unwrap(), 11);
64        // Counters are independent.
65        assert_eq!(ids.reserve(b"d", 0).unwrap(), 0);
66    }
67}
rushdb_storage/id_alloc.rs

rushdb_storage/
id_alloc.rs
