twobitreader 0.2.2

// Asynchronous hints for prefetching ranges of a memory mapped file.
//
// Reading from a cold 2bit file is almost entirely IO-bound. The utilities here support
// the TwobitReader::prefetch function to alleviate this IO bottleneck. The goal is to do so
// using a single thread, rather than forcing the user to rely on multiple threads as a
// workaround to initiate the page faults asynchronously.

// Standard library
use std::fs::File;
use std::ops::Range;

// Dependencies
use memmap2::Mmap;

// Sends prefetch hints to the operating system, possibly batched.
//
// On operating systems (Mac, Linux) where individual hints are efficient to send,
// no actual batching takes place. On Windows, we must collect and send a batch.
pub(crate) struct PrefetchBatcher<'a> {
    #[allow(dead_code)] // Unused on windows
    file: &'a File,
    mmap: &'a Mmap,
    #[cfg(windows)]
    ranges: Vec<Range<usize>>, // Used on Windows only
}

impl<'a> PrefetchBatcher<'a> {
    pub(crate) fn new(file: &'a File, mmap: &'a Mmap) -> Self {
        Self {
            file,
            mmap,
            #[cfg(windows)]
            ranges: Vec::new(),
        }
    }

    #[cfg(target_vendor = "apple")]
    pub(crate) fn push(&mut self, range: Range<usize>) {
        use std::os::unix::io::AsRawFd;

        // If range clamped to mapping extents was non-empty, continue.
        let Some(range) = clamp_range(range, self.mmap.len()) else {
            return;
        };

        // radvisory counts bytes using a c_int, so extremely long ranges are hinted in 1 GB chunks.
        const MAX_HINT_LEN: usize = 1 << 30;

        let fd = self.file.as_raw_fd();
        for offset in range.clone().step_by(MAX_HINT_LEN) {
            let len = (range.end - offset).min(MAX_HINT_LEN);
            let radvisory = libc::radvisory { ra_offset: offset as libc::off_t, ra_count: len as libc::c_int };

            // SAFETY: fd is valid for as long as `file` is borrowed, and fcntl(F_RDADVISE) only
            // reads the radvisory struct, which outlives the call. The call populates the
            // kernel's read cache and cannot modify any memory owned by this process.
            _ = unsafe { libc::fcntl(fd, libc::F_RDADVISE, &radvisory) };
        }
    }

    #[cfg(all(unix, not(target_vendor = "apple")))]
    pub(crate) fn push(&mut self, range: Range<usize>) {
        // If range clamped to mapping extents was non-empty, continue.
        let Some(range) = clamp_range(range, self.mmap.len()) else {
            return;
        };

        // madvise(MADV_WILLNEED) starts asynchronous readahead for the range, and memmap2
        // already wraps it in a safe interface, so no unsafe code is needed here.
        let _ = self.mmap.advise_range(memmap2::Advice::WillNeed, range.start, range.end - range.start);
    }

    #[cfg(windows)]
    pub(crate) fn push(&mut self, range: Range<usize>) {
        // If range clamped to mapping extents was non-empty, collect it in the batch.
        let Some(range) = clamp_range(range, self.mmap.len()) else {
            return;
        };
        self.ranges.push(range);
    }

    #[cfg(not(any(unix, windows)))]
    pub(crate) fn push(&mut self, _range: Range<usize>) {}

    #[cfg(windows)]
    pub(crate) fn flush(&mut self) {
        // PrefetchVirtualMemory takes the whole set of ranges in one call and returns before
        // the reads complete, which is exactly the primitive wanted here. It requires Windows 8.
        use std::ffi::c_void;
        use windows_sys::Win32::System::Memory::{PrefetchVirtualMemory, WIN32_MEMORY_RANGE_ENTRY};
        use windows_sys::Win32::System::Threading::GetCurrentProcess;

        let base = self.mmap.as_ptr();
        let entries = self
            .ranges
            .iter()
            .map(|range| WIN32_MEMORY_RANGE_ENTRY {
                // SAFETY: ranges were clamped to 0..mmap.len(), so this stays within the mapping.
                VirtualAddress: unsafe { base.add(range.start) } as *mut c_void,
                NumberOfBytes: range.end - range.start,
            })
            .collect::<Vec<_>>();
        self.ranges.clear();

        // SAFETY: entries points to entries.len() initialized structs, each describing a range of a
        // mapping owned by this process. PrefetchVirtualMemory only adds pages to the working set.
        unsafe { PrefetchVirtualMemory(GetCurrentProcess(), entries.len(), entries.as_ptr(), 0) };
    }

    #[cfg(not(windows))]
    pub(crate) fn flush(&mut self) {}
}

fn clamp_range(range: Range<usize>, mmap_len: usize) -> Option<Range<usize>> {
    // Fast path: what we expect in any correctly-formed input ranges.
    if range.start < range.end && range.end <= mmap_len {
        return Some(range);
    }

    // Otherwise sanitize the range for safety.
    let clamped = range.start.min(mmap_len)..range.end.min(mmap_len);
    if clamped.start < clamped.end { Some(clamped) } else { None }
}

impl<'a> Drop for PrefetchBatcher<'a> {
    // Flush just in case the user forgot to explicitly do so after pushing hints.
    fn drop(&mut self) {
        self.flush();
    }
}