Skip to main content

reddb_engine/storage/btree/
reclaim.rs

1use super::*;
2
3impl<P: PageStore> BTree<P> {
4    pub fn has_retirement_readers(&self) -> bool {
5        Arc::strong_count(&self.retirement_pin) != 1
6    }
7
8    /// Enumerate a retired generation before changing the allocator. Never
9    /// follow leaf sibling links as additional ownership edges: interior
10    /// children already own every leaf. Validate overflow chains as well.
11    pub fn retired_pages(&self) -> BTreeResult<Vec<u32>> {
12        let mut pending = vec![self.root_page_id()];
13        let mut seen = std::collections::HashSet::new();
14        let mut pages = Vec::new();
15        let page_count = self.pager.page_count()?;
16        while let Some(id) = pending.pop() {
17            if id == 0 {
18                continue;
19            }
20            if id >= page_count || !seen.insert(id) {
21                return Err(BTreeError::Corrupted(
22                    "invalid retired tree page ownership".into(),
23                ));
24            }
25            let page = self.pager.read_page(id)?;
26            match page.page_type()? {
27                PageType::BTreeInterior => {
28                    let (_, children) = read_interior_keys_children(&page)?;
29                    pending.extend(children);
30                }
31                PageType::BTreeLeaf => {
32                    for index in 0..usize::from(page.cell_count()) {
33                        let (_, stored) = read_leaf_cell(&page, index)?;
34                        if let Some(head) = value_layout::pointer_head(&stored) {
35                            pending.push(head);
36                        }
37                    }
38                }
39                PageType::Overflow => {
40                    let content = page.content();
41                    let next = u32::from_le_bytes(content[..4].try_into().expect("chain header"));
42                    let len = u32::from_le_bytes(content[4..8].try_into().expect("chain length"));
43                    if len as usize > crate::storage::overflow::OVERFLOW_PAYLOAD_PER_PAGE {
44                        return Err(BTreeError::Corrupted(
45                            "invalid retired overflow length".into(),
46                        ));
47                    }
48                    pending.push(next);
49                }
50                _ => {
51                    return Err(BTreeError::Corrupted(
52                        "unexpected retired tree page type".into(),
53                    ))
54                }
55            }
56            pages.push(id);
57        }
58        Ok(pages)
59    }
60
61    /// Stream the verified leaves of a damaged tree. Only checksum failures at
62    /// the proven leaf depth may be skipped. The caller must separately prove
63    /// that no allocated pages remain unowned before publishing any output.
64    pub fn visit_recoverable_leaves<C, R>(
65        &self,
66        mut claim: C,
67        mut record: R,
68    ) -> BTreeResult<Vec<u32>>
69    where
70        C: FnMut(u32) -> BTreeResult<()>,
71        R: FnMut(Vec<u8>, Vec<u8>) -> BTreeResult<()>,
72    {
73        struct Work {
74            id: u32,
75            depth: usize,
76            low: Option<Vec<u8>>,
77            high: Option<Vec<u8>>,
78        }
79        let mut pending = vec![Work {
80            id: self.root_page_id(),
81            depth: 0,
82            low: None,
83            high: None,
84        }];
85        let mut seen = std::collections::HashSet::new();
86        let count = self.pager.page_count()?;
87        let mut leaves: Vec<(u32, Option<u32>)> = Vec::new();
88        let mut skipped = Vec::new();
89        let mut leaf_depth = None;
90        let mut skipped_depths = Vec::new();
91        let mut previous: Option<Vec<u8>> = None;
92        while let Some(work) = pending.pop() {
93            if work.id == 0 || work.id >= count || work.depth > 64 || !seen.insert(work.id) {
94                return Err(BTreeError::Corrupted(
95                    "invalid recovery tree ownership".into(),
96                ));
97            }
98            claim(work.id)?;
99            let page = match self.pager.read_page(work.id) {
100                Ok(page) => page,
101                Err(PagerError::Page(super::super::page::PageError::ChecksumMismatch {
102                    ..
103                })) if work.depth > 0 => {
104                    skipped.push(work.id);
105                    skipped_depths.push(work.depth);
106                    leaves.push((work.id, None));
107                    continue;
108                }
109                Err(error) => return Err(error.into()),
110            };
111            match page.page_type()? {
112                PageType::BTreeInterior => {
113                    let (keys, children) = read_interior_keys_children(&page)?;
114                    if children.len() != keys.len() + 1
115                        || children.iter().any(|id| *id == 0)
116                        || keys.windows(2).any(|pair| pair[0] >= pair[1])
117                        || keys.iter().any(|key| {
118                            work.low.as_ref().is_some_and(|low| key < low)
119                                || work.high.as_ref().is_some_and(|high| key >= high)
120                        })
121                    {
122                        return Err(BTreeError::Corrupted(
123                            "invalid recovery interior bounds".into(),
124                        ));
125                    }
126                    for i in (0..children.len()).rev() {
127                        pending.push(Work {
128                            id: children[i],
129                            depth: work.depth + 1,
130                            low: if i == 0 {
131                                work.low.clone()
132                            } else {
133                                Some(keys[i - 1].clone())
134                            },
135                            high: if i == keys.len() {
136                                work.high.clone()
137                            } else {
138                                Some(keys[i].clone())
139                            },
140                        });
141                    }
142                }
143                PageType::BTreeLeaf => {
144                    if leaf_depth.is_some_and(|depth| depth != work.depth) {
145                        return Err(BTreeError::Corrupted(
146                            "inconsistent recovery leaf depth".into(),
147                        ));
148                    }
149                    leaf_depth = Some(work.depth);
150                    let bytes = page.content();
151                    let next = u32::from_le_bytes(bytes[4..8].try_into().expect("leaf header"));
152                    leaves.push((work.id, Some(next)));
153                    for index in 0..usize::from(page.cell_count()) {
154                        let (key, stored) = read_leaf_cell(&page, index)?;
155                        if work.low.as_ref().is_some_and(|low| &key < low)
156                            || work.high.as_ref().is_some_and(|high| &key >= high)
157                            || previous.as_ref().is_some_and(|old| old >= &key)
158                        {
159                            return Err(BTreeError::Corrupted(
160                                "invalid recovery leaf ordering".into(),
161                            ));
162                        }
163                        previous = Some(key.clone());
164                        if let Some(mut overflow) = value_layout::pointer_head(&stored) {
165                            while overflow != 0 {
166                                if overflow >= count || !seen.insert(overflow) {
167                                    return Err(BTreeError::Corrupted(
168                                        "invalid recovery overflow ownership".into(),
169                                    ));
170                                }
171                                claim(overflow)?;
172                                let page = self.pager.read_page(overflow)?;
173                                if page.page_type()? != PageType::Overflow {
174                                    return Err(BTreeError::Corrupted(
175                                        "invalid recovery overflow page".into(),
176                                    ));
177                                }
178                                let bytes = page.content();
179                                let length = u32::from_le_bytes(
180                                    bytes[4..8].try_into().expect("overflow header"),
181                                );
182                                if length as usize
183                                    > crate::storage::overflow::OVERFLOW_PAYLOAD_PER_PAGE
184                                {
185                                    return Err(BTreeError::Corrupted(
186                                        "invalid recovery overflow length".into(),
187                                    ));
188                                }
189                                overflow = u32::from_le_bytes(
190                                    bytes[..4].try_into().expect("overflow header"),
191                                );
192                            }
193                        }
194                        record(
195                            key,
196                            value_layout::decode_with_budget(
197                                self.pager.as_ref(),
198                                &stored,
199                                32 * 1024 * 1024,
200                            )?,
201                        )?;
202                    }
203                }
204                _ => {
205                    return Err(BTreeError::Corrupted(
206                        "invalid recovery tree page type".into(),
207                    ))
208                }
209            }
210        }
211        if leaf_depth.is_none()
212            || skipped_depths
213                .iter()
214                .any(|depth| Some(*depth) != leaf_depth)
215        {
216            return Err(BTreeError::Corrupted(
217                "damaged page is not at proven leaf depth".into(),
218            ));
219        }
220        for (i, (_, next)) in leaves.iter().enumerate() {
221            if next.is_some_and(|id| id != leaves.get(i + 1).map_or(0, |leaf| leaf.0)) {
222                return Err(BTreeError::Corrupted(
223                    "invalid recovery forward leaf chain".into(),
224                ));
225            }
226        }
227        Ok(skipped)
228    }
229}