coreutils_rs/common/
io.rs

1use std::fs::{self, File};
2use std::io::{self, Read};
3use std::ops::Deref;
4use std::path::Path;
5
6#[cfg(target_os = "linux")]
7use std::sync::atomic::{AtomicBool, Ordering};
8
9use memmap2::{Mmap, MmapOptions};
10
11/// Holds file data — either zero-copy mmap or an owned Vec.
12/// Dereferences to `&[u8]` for transparent use.
13pub enum FileData {
14    Mmap(Mmap),
15    Owned(Vec<u8>),
16}
17
18impl Deref for FileData {
19    type Target = [u8];
20
21    fn deref(&self) -> &[u8] {
22        match self {
23            FileData::Mmap(m) => m,
24            FileData::Owned(v) => v,
25        }
26    }
27}
28
29/// Threshold below which we use read() instead of mmap.
30/// For files under 1MB, read() is faster since mmap has setup/teardown overhead
31/// (page table creation for up to 256 pages, TLB flush on munmap) that exceeds
32/// the zero-copy benefit.
33const MMAP_THRESHOLD: u64 = 1024 * 1024;
34
35/// Track whether O_NOATIME is supported to avoid repeated failed open() attempts.
36/// After the first EPERM, we never try O_NOATIME again (saves one syscall per file).
37#[cfg(target_os = "linux")]
38static NOATIME_SUPPORTED: AtomicBool = AtomicBool::new(true);
39
40/// Open a file with O_NOATIME on Linux to avoid atime inode writes.
41/// Caches whether O_NOATIME works to avoid double-open on every file.
42#[cfg(target_os = "linux")]
43fn open_noatime(path: &Path) -> io::Result<File> {
44    use std::os::unix::fs::OpenOptionsExt;
45    if NOATIME_SUPPORTED.load(Ordering::Relaxed) {
46        match fs::OpenOptions::new()
47            .read(true)
48            .custom_flags(libc::O_NOATIME)
49            .open(path)
50        {
51            Ok(f) => return Ok(f),
52            Err(ref e) if e.raw_os_error() == Some(libc::EPERM) => {
53                // O_NOATIME requires file ownership or CAP_FOWNER — disable globally
54                NOATIME_SUPPORTED.store(false, Ordering::Relaxed);
55            }
56            Err(e) => return Err(e), // Real error, propagate
57        }
58    }
59    File::open(path)
60}
61
62#[cfg(not(target_os = "linux"))]
63fn open_noatime(path: &Path) -> io::Result<File> {
64    File::open(path)
65}
66
67/// Read a file with zero-copy mmap for large files or read() for small files.
68/// Opens once with O_NOATIME, uses fstat for metadata to save a syscall.
69pub fn read_file(path: &Path) -> io::Result<FileData> {
70    let file = open_noatime(path)?;
71    let metadata = file.metadata()?;
72    let len = metadata.len();
73
74    if len > 0 && metadata.file_type().is_file() {
75        // Small files: exact-size read from already-open fd.
76        // Uses read_full into pre-sized buffer instead of read_to_end,
77        // which avoids the grow-and-probe pattern (saves 1-2 extra read() syscalls).
78        if len < MMAP_THRESHOLD {
79            let mut buf = vec![0u8; len as usize];
80            let n = read_full(&mut &file, &mut buf)?;
81            buf.truncate(n);
82            return Ok(FileData::Owned(buf));
83        }
84
85        // SAFETY: Read-only mapping. MADV_SEQUENTIAL lets the kernel
86        // prefetch ahead of our sequential access pattern.
87        match unsafe { MmapOptions::new().populate().map(&file) } {
88            Ok(mmap) => {
89                #[cfg(target_os = "linux")]
90                {
91                    let _ = mmap.advise(memmap2::Advice::Sequential);
92                    // HUGEPAGE reduces TLB misses for large files (2MB+ = 1+ huge page).
93                    // With 4KB pages, a 100MB file needs 25,600 TLB entries; with 2MB
94                    // huge pages it needs only 50, reducing TLB miss overhead by ~500x.
95                    if len >= 2 * 1024 * 1024 {
96                        let _ = mmap.advise(memmap2::Advice::HugePage);
97                    }
98                }
99                Ok(FileData::Mmap(mmap))
100            }
101            Err(_) => {
102                // mmap failed — fall back to read
103                let mut buf = Vec::with_capacity(len as usize);
104                let mut reader = file;
105                reader.read_to_end(&mut buf)?;
106                Ok(FileData::Owned(buf))
107            }
108        }
109    } else if len > 0 {
110        // Non-regular file (special files) — read from open fd
111        let mut buf = Vec::new();
112        let mut reader = file;
113        reader.read_to_end(&mut buf)?;
114        Ok(FileData::Owned(buf))
115    } else {
116        Ok(FileData::Owned(Vec::new()))
117    }
118}
119
120/// Get file size without reading it (for byte-count-only optimization).
121pub fn file_size(path: &Path) -> io::Result<u64> {
122    Ok(fs::metadata(path)?.len())
123}
124
125/// Read all bytes from stdin into a Vec.
126pub fn read_stdin() -> io::Result<Vec<u8>> {
127    let mut buf = Vec::new();
128    io::stdin().lock().read_to_end(&mut buf)?;
129    Ok(buf)
130}
131
132/// Read as many bytes as possible into buf, retrying on partial reads.
133/// Ensures the full buffer is filled (or EOF reached), avoiding the
134/// probe-read overhead of read_to_end.
135#[inline]
136fn read_full(reader: &mut impl Read, buf: &mut [u8]) -> io::Result<usize> {
137    let mut total = 0;
138    while total < buf.len() {
139        match reader.read(&mut buf[total..]) {
140            Ok(0) => break,
141            Ok(n) => total += n,
142            Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
143            Err(e) => return Err(e),
144        }
145    }
146    Ok(total)
147}
coreutils_rs/common/io.rs

coreutils_rs/common/
io.rs
