citadeldb_buffer/
allocator.rs

1//! Page allocator for the CoW B+ tree.
2//!
3//! Uses a two-phase pending-free model:
4//! - New pages are allocated from `ready_to_use` (reclaimed) or high water mark
5//! - Freed pages go into `freed_this_txn` (not reusable until committed + no older readers)
6//! - On-disk pending-free chain persistence happens during commit
7
8use citadel_core::types::PageId;
9
10/// In-memory page allocator state.
11pub struct PageAllocator {
12    /// Next page ID to allocate from (high water mark).
13    next_page_id: u32,
14    /// Pages reclaimed from pending-free chain (safe to reuse).
15    ready_to_use: Vec<PageId>,
16    /// Pages freed in the current write transaction.
17    freed_this_txn: Vec<PageId>,
18}
19
20impl PageAllocator {
21    /// Create a new allocator starting from the given high water mark.
22    pub fn new(high_water_mark: u32) -> Self {
23        Self {
24            next_page_id: high_water_mark,
25            ready_to_use: Vec::new(),
26            freed_this_txn: Vec::new(),
27        }
28    }
29
30    /// Allocate a new page ID.
31    /// Prefers reusing reclaimed pages; falls back to incrementing the high water mark.
32    pub fn allocate(&mut self) -> PageId {
33        if let Some(id) = self.ready_to_use.pop() {
34            id
35        } else {
36            let id = PageId(self.next_page_id);
37            self.next_page_id += 1;
38            id
39        }
40    }
41
42    /// Mark a page as freed in the current transaction.
43    /// The page is NOT immediately reusable — it goes into the pending-free list.
44    pub fn free(&mut self, page_id: PageId) {
45        self.freed_this_txn.push(page_id);
46    }
47
48    /// Current high water mark (next page ID that would be allocated from disk).
49    pub fn high_water_mark(&self) -> u32 {
50        self.next_page_id
51    }
52
53    /// Pages freed in the current transaction.
54    pub fn freed_this_txn(&self) -> &[PageId] {
55        &self.freed_this_txn
56    }
57
58    /// Add pages that are safe to reuse (reclaimed from pending-free chain).
59    pub fn add_ready_to_use(&mut self, pages: Vec<PageId>) {
60        self.ready_to_use.extend(pages);
61    }
62
63    /// Commit: take ownership of freed pages (for writing to pending-free chain).
64    /// Returns the list of pages freed in this transaction.
65    pub fn commit(&mut self) -> Vec<PageId> {
66        std::mem::take(&mut self.freed_this_txn)
67    }
68
69    /// Rollback: discard freed pages (transaction aborted).
70    pub fn rollback(&mut self) {
71        self.freed_this_txn.clear();
72    }
73
74    /// Number of pages available for immediate reuse.
75    pub fn ready_count(&self) -> usize {
76        self.ready_to_use.len()
77    }
78
79    /// Number of pages freed in the current transaction.
80    pub fn freed_count(&self) -> usize {
81        self.freed_this_txn.len()
82    }
83}
84
85#[cfg(test)]
86mod tests {
87    use super::*;
88
89    #[test]
90    fn allocate_from_hwm() {
91        let mut alloc = PageAllocator::new(0);
92        assert_eq!(alloc.allocate(), PageId(0));
93        assert_eq!(alloc.allocate(), PageId(1));
94        assert_eq!(alloc.allocate(), PageId(2));
95        assert_eq!(alloc.high_water_mark(), 3);
96    }
97
98    #[test]
99    fn allocate_from_ready_to_use() {
100        let mut alloc = PageAllocator::new(10);
101        alloc.add_ready_to_use(vec![PageId(3), PageId(7)]);
102        // Should use ready_to_use first (LIFO)
103        assert_eq!(alloc.allocate(), PageId(7));
104        assert_eq!(alloc.allocate(), PageId(3));
105        // Now falls back to HWM
106        assert_eq!(alloc.allocate(), PageId(10));
107    }
108
109    #[test]
110    fn free_and_commit() {
111        let mut alloc = PageAllocator::new(5);
112        alloc.free(PageId(1));
113        alloc.free(PageId(3));
114        assert_eq!(alloc.freed_count(), 2);
115
116        let freed = alloc.commit();
117        assert_eq!(freed.len(), 2);
118        assert_eq!(alloc.freed_count(), 0);
119    }
120
121    #[test]
122    fn rollback_clears_freed() {
123        let mut alloc = PageAllocator::new(5);
124        alloc.free(PageId(1));
125        alloc.free(PageId(3));
126        alloc.rollback();
127        assert_eq!(alloc.freed_count(), 0);
128    }
129}
citadeldb_buffer/allocator.rs

citadeldb_buffer/
allocator.rs
