forest/utils/io/
mmap.rs

1// Copyright 2019-2025 ChainSafe Systems
2// SPDX-License-Identifier: Apache-2.0, MIT
3
4use std::{fs, io, path::Path};
5
6use memmap2::MmapAsRawDesc;
7use positioned_io::{RandomAccessFile, ReadAt, Size};
8
9use crate::utils::misc::env::is_env_truthy;
10
11/// Wrapper type of [`memmap2::Mmap`] that implements [`ReadAt`] and [`Size`]
12pub struct Mmap(memmap2::Mmap);
13
14impl Mmap {
15    pub fn map(file: impl MmapAsRawDesc) -> io::Result<Self> {
16        Ok(Self(unsafe { memmap2::Mmap::map(file)? }))
17    }
18
19    pub fn map_path(path: impl AsRef<Path>) -> io::Result<Self> {
20        Self::map(&fs::File::open(path.as_ref())?)
21    }
22}
23
24impl ReadAt for Mmap {
25    #[allow(clippy::indexing_slicing)]
26    fn read_at(&self, pos: u64, buf: &mut [u8]) -> io::Result<usize> {
27        let start = pos as usize;
28        if start >= self.0.len() {
29            // This matches the behaviour for seeking past the end of a file
30            return Ok(0);
31        }
32        let end = start + buf.len();
33        if end <= self.0.len() {
34            buf.copy_from_slice(&self.0[start..end]);
35            Ok(buf.len())
36        } else {
37            let len = self.0.len() - start;
38            buf[..len].copy_from_slice(&self.0[start..]);
39            Ok(len)
40        }
41    }
42}
43
44impl Size for Mmap {
45    fn size(&self) -> io::Result<Option<u64>> {
46        Ok(Some(self.0.len() as _))
47    }
48}
49
50pub enum EitherMmapOrRandomAccessFile {
51    Mmap(Mmap),
52    RandomAccessFile(RandomAccessFile),
53}
54
55impl EitherMmapOrRandomAccessFile {
56    pub fn open(path: impl AsRef<Path>) -> io::Result<Self> {
57        Ok(if should_use_file_io() {
58            Self::RandomAccessFile(RandomAccessFile::open(path)?)
59        } else {
60            Self::Mmap(Mmap::map_path(path)?)
61        })
62    }
63}
64
65impl ReadAt for EitherMmapOrRandomAccessFile {
66    fn read_at(&self, pos: u64, buf: &mut [u8]) -> io::Result<usize> {
67        use EitherMmapOrRandomAccessFile::*;
68        match self {
69            Mmap(mmap) => mmap.read_at(pos, buf),
70            RandomAccessFile(file) => file.read_at(pos, buf),
71        }
72    }
73}
74
75impl Size for EitherMmapOrRandomAccessFile {
76    fn size(&self) -> io::Result<Option<u64>> {
77        use EitherMmapOrRandomAccessFile::*;
78        match self {
79            Mmap(mmap) => mmap.size(),
80            RandomAccessFile(file) => file.size(),
81        }
82    }
83}
84
85fn should_use_file_io() -> bool {
86    // Use mmap by default, switch to file-io when `FOREST_CAR_LOADER_FILE_IO` is set to `1` or `true`
87    is_env_truthy("FOREST_CAR_LOADER_FILE_IO")
88}
89
90#[cfg(test)]
91mod tests {
92    use std::fs;
93
94    use super::*;
95    use quickcheck_macros::quickcheck;
96
97    #[quickcheck]
98    fn test_mmap_read_at_and_size(bytes: Vec<u8>) -> anyhow::Result<()> {
99        let tmp = tempfile::Builder::new().tempfile()?.into_temp_path();
100        fs::write(&tmp, &bytes)?;
101        let mmap = Mmap::map(&fs::File::open(&tmp)?)?;
102
103        assert_eq!(mmap.size()?.unwrap_or_default() as usize, bytes.len());
104
105        let mut buffer = [0; 128];
106        for pos in 0..bytes.len() {
107            let size = mmap.read_at(pos as _, &mut buffer)?;
108            assert_eq!(&bytes[pos..(pos + size)], &buffer[..size]);
109        }
110
111        Ok(())
112    }
113
114    #[test]
115    fn test_out_of_band_mmap_read() {
116        let temp_file = tempfile::Builder::new()
117            .tempfile()
118            .unwrap()
119            .into_temp_path();
120        let mmap = Mmap::map(&fs::File::open(&temp_file).unwrap()).unwrap();
121
122        let mut buffer = [];
123        // This matches the behaviour for seeking past the end of a file
124        assert_eq!(mmap.read_at(u64::MAX, &mut buffer).unwrap(), 0);
125    }
126}
forest/utils/io/mmap.rs

forest/utils/io/
mmap.rs
