bonsai/storage/
manager.rs

1//! Maps (relation, fork, block) to a file and an offset, and handles storing and retrieving blocks.
2//! Postgres: `smgr.c` and `md.c`.
3
4use super::{BLCKSZ, BlockCount, BlockNumber, page::Page};
5
6use anyhow::Context;
7
8use std::{
9    collections::HashMap,
10    fmt::Write,
11    fs::{self, File},
12    io,
13    os::unix::fs::FileExt,
14    path::PathBuf,
15    sync::{Arc, Mutex},
16};
17
18#[derive(Clone, Copy, Debug)]
19pub struct BlockIdentifier {
20    pub fork: ForkIdentifier,
21    pub block_number: BlockNumber,
22}
23
24#[derive(Clone, Copy, Debug, Hash, Eq, PartialEq)]
25pub struct ForkIdentifier {
26    pub rel: RelFileNumber,
27    pub kind: ForkKind,
28}
29
30#[derive(Clone, Copy, Debug, Hash, Eq, PartialEq)]
31pub struct RelFileNumber(pub u32);
32
33#[derive(Clone, Copy, Debug, Hash, Eq, PartialEq)]
34pub enum ForkKind {
35    Main = 0,
36    FreeSpaceMap = 1,
37    VisibilityMap = 2,
38    Init = 3,
39}
40
41pub struct StorageManager {
42    dir: PathBuf,
43    // Not necessarily equal to RELSEG_SIZE. This allows emulating smaller segment sizes for
44    // testing.
45    n_block_per_segment: BlockCount,
46    opened_forks: Mutex<HashMap<ForkIdentifier, Arc<OpenFork>>>,
47}
48
49#[derive(Debug)]
50pub struct BlockError {
51    pub block: BlockIdentifier,
52    pub inner: BlockErrorInner,
53}
54
55#[derive(Debug)]
56pub struct ForkError {
57    pub fork: ForkIdentifier,
58    pub inner: ForkErrorInner,
59}
60
61#[derive(Debug)]
62pub enum ForkErrorInner {
63    Io(io::Error),
64    SegmentFileInconsistentlyBig {
65        path: PathBuf,
66        file_size: u64,
67        n_blocks_per_segment: BlockCount,
68    },
69}
70
71#[derive(Debug)]
72pub enum BlockErrorInner {
73    BlockDoesNotExist { fork_size: BlockCount },
74    Io(io::Error),
75    ForkError(ForkError),
76}
77
78impl StorageManager {
79    /// Opens `dir`, the directory holding the relation files: `base/<database
80    /// OID>` in a Postgres cluster.
81    pub fn open(dir: impl Into<PathBuf>) -> anyhow::Result<Self> {
82        Self::with_segment_size(dir, RELSEG_SIZE)
83    }
84
85    /// Fetch the block from the disk and store its bits into `page`.
86    pub fn read(&self, block: BlockIdentifier, page: &mut Page) -> Result<(), BlockError> {
87        let fork = self.get_fork(block.fork).map_err(|err| BlockError {
88            block,
89            inner: BlockErrorInner::ForkError(err),
90        })?;
91        if !fork.nb_blocks.contains(block.block_number) {
92            return Err(BlockError {
93                block,
94                inner: BlockErrorInner::BlockDoesNotExist {
95                    fork_size: fork.nb_blocks,
96                },
97            });
98        }
99        let Position {
100            segno,
101            offset_bytes,
102        } = self.position(block.block_number);
103        fork.segments[segno as usize]
104            .read_exact_at(page, offset_bytes)
105            .map_err(|err| BlockError {
106                block,
107                inner: BlockErrorInner::Io(err),
108            })
109    }
110
111    fn with_segment_size(
112        dir: impl Into<PathBuf>,
113        n_block_per_segment: BlockCount,
114    ) -> anyhow::Result<Self> {
115        assert!(
116            n_block_per_segment.0 > 0,
117            "segments must hold at least one block"
118        );
119        let dir = dir.into();
120        fs::read_dir(&dir)
121            .with_context(|| format!("could not open directory {}", dir.display()))?;
122        Ok(Self {
123            dir,
124            n_block_per_segment,
125            opened_forks: Mutex::new(HashMap::new()),
126        })
127    }
128
129    fn get_fork(&self, fork: ForkIdentifier) -> Result<Arc<OpenFork>, ForkError> {
130        let mut forks = self.opened_forks.lock().unwrap();
131        if let Some(open) = forks.get(&fork) {
132            return Ok(Arc::clone(open));
133        }
134        let open = Arc::new(self.load_fork(fork)?);
135        forks.insert(fork, Arc::clone(&open));
136        Ok(open)
137    }
138
139    fn load_fork(&self, fork: ForkIdentifier) -> Result<OpenFork, ForkError> {
140        let empty = OpenFork {
141            segments: Vec::new(),
142            nb_blocks: BlockCount(0),
143        };
144        let mut previous_full = true;
145        (0..)
146            // Stops at the first missing segment, or after the first one that
147            // isn't full: only the last segment can be partly full.
148            .map_while(|segno| {
149                if !previous_full {
150                    return None;
151                }
152                let segment = self.open_segment(fork, segno).transpose()?;
153                previous_full = segment
154                    .as_ref()
155                    .is_ok_and(|segment| segment.nb_blocks == self.n_block_per_segment);
156                Some(segment)
157            })
158            .try_fold(empty, |mut open, segment| {
159                let segment = segment?;
160                open.segments.push(segment.file);
161                open.nb_blocks.0 += segment.nb_blocks.0;
162                Ok(open)
163            })
164    }
165
166    /// Opens segment `segno` of `fork` and counts its blocks. `None` if there
167    /// is no such segment, except for segment 0: a missing fork is an error.
168    fn open_segment(&self, fork: ForkIdentifier, segno: u32) -> Result<Option<Segment>, ForkError> {
169        let path = self.path(fork, segno);
170        let file = match File::open(&path) {
171            Ok(file) => file,
172            // The previous segment was full and nothing comes after it.
173            Err(err) if segno > 0 && err.kind() == io::ErrorKind::NotFound => return Ok(None),
174            Err(err) => {
175                return Err(ForkError {
176                    fork,
177                    inner: ForkErrorInner::Io(err),
178                });
179            }
180        };
181        // A partly written block at the end, left by a crash, doesn't
182        // count, as in `mdnblocks`.
183        let file_size = file
184            .metadata()
185            .map_err(|err| ForkError {
186                fork,
187                inner: err.into(),
188            })?
189            .len();
190        let blocks = file_size / (BLCKSZ as u64);
191        if blocks > u64::from(self.n_block_per_segment.0) {
192            return Err(ForkError {
193                fork,
194                inner: ForkErrorInner::SegmentFileInconsistentlyBig {
195                    path,
196                    file_size,
197                    n_blocks_per_segment: self.n_block_per_segment,
198                },
199            });
200        }
201        Ok(Some(Segment {
202            file,
203            nb_blocks: BlockCount(blocks as u32),
204        }))
205    }
206
207    /// Path to the file holding the segment `segno` of `fork`.
208    fn path(&self, fork: ForkIdentifier, segno: u32) -> PathBuf {
209        let mut name = fork.rel.0.to_string();
210        match fork.kind {
211            ForkKind::Main => {}
212            ForkKind::FreeSpaceMap => {
213                name.push_str("_fsm");
214            }
215            ForkKind::VisibilityMap => {
216                name.push_str("_vm");
217            }
218            ForkKind::Init => {
219                name.push_str("_init");
220            }
221        }
222        if segno > 0 {
223            write!(&mut name, ".{segno}").unwrap();
224        }
225        self.dir.join(name)
226    }
227
228    fn position(&self, block: BlockNumber) -> Position {
229        let segno = block.0 / self.n_block_per_segment.0;
230        let offset = u64::from(block.0 % self.n_block_per_segment.0) * BLCKSZ as u64;
231        Position {
232            segno,
233            offset_bytes: offset,
234        }
235    }
236}
237
238/// How many blocks fit in a 1GB segment file.
239pub const RELSEG_SIZE: BlockCount = BlockCount(1024 * 1024 * 1024 / (BLCKSZ as u32));
240
241/// A fork whose segment files are open.
242struct OpenFork {
243    segments: Vec<File>,
244    nb_blocks: BlockCount,
245}
246
247/// One open segment file of a fork.
248struct Segment {
249    file: File,
250    nb_blocks: BlockCount,
251}
252
253struct Position {
254    segno: u32,
255    offset_bytes: u64,
256}
257
258impl From<io::Error> for ForkErrorInner {
259    fn from(err: io::Error) -> Self {
260        ForkErrorInner::Io(err)
261    }
262}
263
264impl From<io::Error> for BlockErrorInner {
265    fn from(err: io::Error) -> Self {
266        BlockErrorInner::Io(err)
267    }
268}
269
270#[cfg(test)]
271mod tests {
272    use super::*;
273
274    use std::path::Path;
275
276    const MAIN: ForkIdentifier = ForkIdentifier {
277        rel: RelFileNumber(16384),
278        kind: ForkKind::Main,
279    };
280
281    /// Writes a file of `count` blocks into `dir`, block `i` filled with the
282    /// byte `first + i`.
283    fn write_blocks(dir: &Path, name: &str, first: u8, count: u8) {
284        let bytes: Vec<u8> = (first..first + count)
285            .flat_map(|byte| [byte; BLCKSZ])
286            .collect();
287        fs::write(dir.join(name), bytes).unwrap();
288    }
289
290    /// A storage manager whose segments hold two blocks, so that a few blocks
291    /// are enough to span several segments.
292    fn open_with_small_segments(dir: &Path) -> StorageManager {
293        StorageManager::with_segment_size(dir, BlockCount(2)).unwrap()
294    }
295
296    fn nb_blocks(smgr: &StorageManager) -> BlockCount {
297        smgr.get_fork(MAIN).unwrap().nb_blocks
298    }
299
300    fn read(smgr: &StorageManager, block_number: u32) -> Result<Page, BlockError> {
301        let block = BlockIdentifier {
302            fork: MAIN,
303            block_number: BlockNumber(block_number),
304        };
305        let mut page = Page { bytes: [0; BLCKSZ] };
306        smgr.read(block, &mut page).map(|()| page)
307    }
308
309    fn is_past_the_end(result: &Result<Page, BlockError>, fork_size: u32) -> bool {
310        matches!(
311            result,
312            Err(BlockError {
313                inner: BlockErrorInner::BlockDoesNotExist { fork_size: size },
314                ..
315            }) if *size == BlockCount(fork_size)
316        )
317    }
318
319    #[test]
320    fn blocks_are_found_across_segments() {
321        let dir = tempfile::tempdir().unwrap();
322        write_blocks(dir.path(), "16384", 0, 2);
323        write_blocks(dir.path(), "16384.1", 2, 2);
324        write_blocks(dir.path(), "16384.2", 4, 1);
325        let smgr = open_with_small_segments(dir.path());
326        assert_eq!(nb_blocks(&smgr), BlockCount(5));
327        for block_number in 0..5 {
328            let page = read(&smgr, block_number).unwrap();
329            assert!(
330                page.bytes == [block_number as u8; BLCKSZ],
331                "block {block_number}"
332            );
333        }
334        assert!(is_past_the_end(&read(&smgr, 5), 5));
335    }
336
337    #[test]
338    fn a_fork_can_end_with_a_full_segment() {
339        let dir = tempfile::tempdir().unwrap();
340        write_blocks(dir.path(), "16384", 0, 2);
341        write_blocks(dir.path(), "16384.1", 2, 2);
342        let smgr = open_with_small_segments(dir.path());
343        assert_eq!(nb_blocks(&smgr), BlockCount(4));
344        assert!(is_past_the_end(&read(&smgr, 4), 4));
345    }
346
347    #[test]
348    fn an_empty_fork_has_no_blocks() {
349        let dir = tempfile::tempdir().unwrap();
350        fs::write(dir.path().join("16384"), []).unwrap();
351        let smgr = open_with_small_segments(dir.path());
352        assert_eq!(nb_blocks(&smgr), BlockCount(0));
353        assert!(is_past_the_end(&read(&smgr, 0), 0));
354    }
355
356    /// As in `mdnblocks`, counting stops at the first segment that isn't
357    /// full, whatever files come after it.
358    #[test]
359    fn segments_after_a_partial_one_are_ignored() {
360        let dir = tempfile::tempdir().unwrap();
361        write_blocks(dir.path(), "16384", 0, 1);
362        write_blocks(dir.path(), "16384.1", 9, 2);
363        let smgr = open_with_small_segments(dir.path());
364        assert_eq!(nb_blocks(&smgr), BlockCount(1));
365    }
366
367    #[test]
368    fn a_partial_block_at_the_end_does_not_count() {
369        let dir = tempfile::tempdir().unwrap();
370        let mut bytes = vec![0; BLCKSZ];
371        bytes.extend([1; 100]);
372        fs::write(dir.path().join("16384"), bytes).unwrap();
373        let smgr = StorageManager::open(dir.path()).unwrap();
374        assert_eq!(nb_blocks(&smgr), BlockCount(1));
375        assert!(is_past_the_end(&read(&smgr, 1), 1));
376    }
377
378    #[test]
379    fn a_missing_fork_is_an_error() {
380        let dir = tempfile::tempdir().unwrap();
381        let smgr = open_with_small_segments(dir.path());
382        let result = read(&smgr, 0);
383        assert!(
384            matches!(
385                &result,
386                Err(BlockError {
387                    inner: BlockErrorInner::ForkError(ForkError {
388                        inner: ForkErrorInner::Io(err),
389                        ..
390                    }),
391                    ..
392                }) if err.kind() == io::ErrorKind::NotFound
393            ),
394            "{result:?}"
395        );
396    }
397
398    #[test]
399    fn a_segment_bigger_than_the_segment_size_is_an_error() {
400        let dir = tempfile::tempdir().unwrap();
401        write_blocks(dir.path(), "16384", 0, 3);
402        let smgr = open_with_small_segments(dir.path());
403        let result = smgr.get_fork(MAIN).map(|_| ());
404        assert!(
405            matches!(
406                &result,
407                Err(ForkError {
408                    inner: ForkErrorInner::SegmentFileInconsistentlyBig {
409                        file_size,
410                        n_blocks_per_segment: BlockCount(2),
411                        ..
412                    },
413                    ..
414                }) if *file_size == 3 * BLCKSZ as u64
415            ),
416            "{result:?}"
417        );
418    }
419}
bonsai/storage/manager.rs

bonsai/storage/
manager.rs
