ab_direct_io_file/
lib.rs

1//! Cross-platform APIs for working with files using direct I/O.
2//!
3//! Depending on OS, this will use direct I/O, unbuffered, uncaches passthrough file reads/writes,
4//! bypassing as much of OS machinery as possible.
5//!
6//! NOTE: There are major alignment requirements described here:
7//! <https://learn.microsoft.com/en-us/windows/win32/fileio/file-buffering#alignment-and-file-access-requirements>
8//! <https://man7.org/linux/man-pages/man2/open.2.html>
9
10#![feature(core_io_borrowed_buf, maybe_uninit_as_bytes)]
11#![cfg_attr(unix, feature(read_buf_at))]
12#![cfg_attr(windows, feature(seek_read_exact_seek_write_all))]
13// TODO: Remove once https://github.com/rust-lang/rust-clippy/issues/17498 is resolved
14#![cfg_attr(
15    any(unix, windows),
16    expect(
17        clippy::items_after_statements,
18        reason = "https://github.com/rust-lang/rust-clippy/issues/17498"
19    )
20)]
21
22// TODO: Windows shims are incomplete under Miri: https://github.com/rust-lang/miri/issues/3482
23#[cfg(all(test, not(all(miri, windows))))]
24mod tests;
25
26use parking_lot::Mutex;
27use std::fs::{File, OpenOptions};
28use std::io::{BorrowedBuf, BorrowedCursor};
29use std::mem::MaybeUninit;
30use std::path::Path;
31use std::{io, mem};
32
33/// Restrict how much data to read from the disk in a single call to avoid very large memory usage
34const MAX_READ_SIZE: usize = 1024 * 1024;
35
36const _: () = {
37    assert!(MAX_READ_SIZE.is_multiple_of(AlignedPage::SIZE));
38};
39
40/// A wrapper data structure with 4096 bytes alignment, which is the most common alignment for
41/// direct I/O operations.
42#[derive(Debug, Copy, Clone)]
43#[repr(C, align(4096))]
44pub struct AlignedPage([u8; AlignedPage::SIZE]);
45
46const _: () = {
47    assert!(align_of::<AlignedPage>() == AlignedPage::SIZE);
48};
49
50impl Default for AlignedPage {
51    #[inline(always)]
52    fn default() -> Self {
53        Self([0; AlignedPage::SIZE])
54    }
55}
56
57impl AlignedPage {
58    /// 4096 is as a relatively safe size due to sector size on SSDs commonly being 512 or 4096
59    /// bytes
60    pub const SIZE: usize = 4096;
61
62    /// Convenient conversion from slice to underlying representation for efficiency purposes
63    #[inline(always)]
64    pub fn slice_to_repr(value: &[Self]) -> &[[u8; AlignedPage::SIZE]] {
65        // SAFETY: `AlignedPage` is `#[repr(C)]` and guaranteed to have the same memory layout
66        unsafe { mem::transmute(value) }
67    }
68
69    /// Convenient conversion from slice to underlying representation for efficiency purposes
70    #[inline(always)]
71    pub fn uninit_slice_to_repr(
72        value: &[MaybeUninit<Self>],
73    ) -> &[MaybeUninit<[u8; AlignedPage::SIZE]>] {
74        // SAFETY: `AlignedPage` is `#[repr(C)]` and guaranteed to have the same memory layout
75        unsafe { mem::transmute(value) }
76    }
77
78    /// Convenient conversion from a slice of underlying representation for efficiency purposes.
79    ///
80    /// Returns `None` if not correctly aligned.
81    #[inline]
82    pub fn try_slice_from_repr(value: &[[u8; AlignedPage::SIZE]]) -> Option<&[Self]> {
83        // SAFETY: All bit patterns are valid
84        let (before, slice, after) = unsafe { value.align_to::<Self>() };
85
86        if before.is_empty() && after.is_empty() {
87            Some(slice)
88        } else {
89            None
90        }
91    }
92
93    /// Convenient conversion from a slice of underlying representation for efficiency purposes.
94    ///
95    /// Returns `None` if not correctly aligned.
96    #[inline]
97    pub fn try_uninit_slice_from_repr(
98        value: &[MaybeUninit<[u8; AlignedPage::SIZE]>],
99    ) -> Option<&[MaybeUninit<Self>]> {
100        // SAFETY: All bit patterns are valid
101        let (before, slice, after) = unsafe { value.align_to::<MaybeUninit<Self>>() };
102
103        if before.is_empty() && after.is_empty() {
104            Some(slice)
105        } else {
106            None
107        }
108    }
109
110    /// Convenient conversion from mutable slice to underlying representation for efficiency
111    /// purposes
112    #[inline(always)]
113    pub fn slice_mut_to_repr(slice: &mut [Self]) -> &mut [[u8; AlignedPage::SIZE]] {
114        // SAFETY: `AlignedPage` is `#[repr(C)]` with a single `[u8; AlignedPage::SIZE]` field and
115        // alignment equal to its size, so it has no padding and the same size as the inner value,
116        // whose alignment is smaller. Any bytes written through the result are valid `AlignedPage`.
117        unsafe { mem::transmute(slice) }
118    }
119
120    /// Convenient conversion from mutable slice to underlying representation for efficiency
121    /// purposes
122    #[inline(always)]
123    pub fn uninit_slice_mut_to_repr(
124        slice: &mut [MaybeUninit<Self>],
125    ) -> &mut [MaybeUninit<[u8; AlignedPage::SIZE]>] {
126        // SAFETY: `AlignedPage` is `#[repr(C)]` with a single `[u8; AlignedPage::SIZE]` field and
127        // alignment equal to its size, so it has no padding and the same size as the inner value,
128        // whose alignment is smaller. `MaybeUninit` preserves the size and alignment of the wrapped
129        // type, and any bytes (initialized or not) written through the result are a valid
130        // `MaybeUninit<AlignedPage>`.
131        unsafe { mem::transmute(slice) }
132    }
133
134    /// Convenient conversion from a slice of underlying representation for efficiency purposes.
135    ///
136    /// Returns `None` if not correctly aligned.
137    #[inline]
138    pub fn try_slice_mut_from_repr(value: &mut [[u8; AlignedPage::SIZE]]) -> Option<&mut [Self]> {
139        // SAFETY: All bit patterns are valid
140        let (before, slice, after) = unsafe { value.align_to_mut::<Self>() };
141
142        if before.is_empty() && after.is_empty() {
143            Some(slice)
144        } else {
145            None
146        }
147    }
148
149    /// Convenient conversion from a slice of underlying representation for efficiency purposes.
150    ///
151    /// Returns `None` if not correctly aligned.
152    #[inline]
153    pub fn try_uninit_slice_mut_from_repr(
154        value: &mut [MaybeUninit<[u8; AlignedPage::SIZE]>],
155    ) -> Option<&mut [MaybeUninit<Self>]> {
156        // SAFETY: All bit patterns are valid
157        let (before, slice, after) = unsafe { value.align_to_mut::<MaybeUninit<Self>>() };
158
159        if before.is_empty() && after.is_empty() {
160            Some(slice)
161        } else {
162            None
163        }
164    }
165}
166
167/// Wrapper data structure for direct/unbuffered/uncached I/O.
168///
169/// Depending on OS, this will use direct I/O, unbuffered, uncaches passthrough file reads/writes,
170/// bypassing as much of OS machinery as possible.
171///
172/// NOTE: There are major alignment requirements described here:
173/// <https://learn.microsoft.com/en-us/windows/win32/fileio/file-buffering#alignment-and-file-access-requirements>
174/// <https://man7.org/linux/man-pages/man2/open.2.html>
175#[derive(Debug)]
176pub struct DirectIoFile {
177    file: File,
178    /// Scratch buffer of aligned memory for reads and writes
179    scratch_buffer: Mutex<Vec<AlignedPage>>,
180}
181
182impl DirectIoFile {
183    /// Open a file with basic open options at the specified path for direct/unbuffered I/O for
184    /// reads and writes.
185    ///
186    /// `options` allows configuring things like read/write/create/truncate, but custom options
187    /// will be overridden internally.
188    ///
189    /// This is especially important on Windows to prevent huge memory usage.
190    #[inline]
191    pub fn open<P>(
192        #[cfg(any(target_os = "linux", windows))] mut options: OpenOptions,
193        #[cfg(not(any(target_os = "linux", windows)))] options: OpenOptions,
194        path: P,
195    ) -> io::Result<Self>
196    where
197        P: AsRef<Path>,
198    {
199        // Direct I/O on Linux
200        #[cfg(target_os = "linux")]
201        // TODO: Unlock under Miri once supported: https://github.com/rust-lang/miri/issues/4462
202        if !cfg!(miri) {
203            use std::os::unix::fs::OpenOptionsExt;
204
205            options.custom_flags(libc::O_DIRECT);
206        }
207        // Unbuffered write-through on Windows
208        #[cfg(windows)]
209        // TODO: Unlock under Miri once supported: https://github.com/rust-lang/miri/issues/4462
210        if !cfg!(miri) {
211            use std::os::windows::fs::OpenOptionsExt;
212
213            options.custom_flags(
214                windows::Win32::Storage::FileSystem::FILE_FLAG_WRITE_THROUGH.0
215                    | windows::Win32::Storage::FileSystem::FILE_FLAG_NO_BUFFERING.0,
216            );
217        }
218        let file = options.open(path)?;
219
220        // Disable caching on macOS
221        #[cfg(target_os = "macos")]
222        // TODO: Unlock under Miri once supported: https://github.com/rust-lang/miri/issues/4462
223        if !cfg!(miri) {
224            use std::os::unix::io::AsRawFd;
225
226            // SAFETY: FFI call with correct file descriptor and arguments
227            if unsafe { libc::fcntl(file.as_raw_fd(), libc::F_NOCACHE, 1) } != 0 {
228                return Err(io::Error::last_os_error());
229            }
230        }
231
232        Ok(Self {
233            file,
234            // In many cases, we'll want to read this much at once, so pre-allocate it right away
235            scratch_buffer: Mutex::new(vec![
236                AlignedPage::default();
237                MAX_READ_SIZE / AlignedPage::SIZE
238            ]),
239        })
240    }
241
242    /// Get file size
243    #[inline]
244    pub fn len(&self) -> io::Result<u64> {
245        Ok(self.file.metadata()?.len())
246    }
247
248    /// Returns `Ok(true)` if the file is empty
249    #[inline]
250    pub fn is_empty(&self) -> io::Result<bool> {
251        Ok(self.len()? == 0)
252    }
253
254    /// Make sure the file has a specified number of bytes allocated on the disk.
255    ///
256    /// Later writes within `len` will not fail due to lack of disk space.
257    #[inline(always)]
258    pub fn allocate(&self, len: u64) -> io::Result<()> {
259        fs4::FileExt::allocate(&self.file, len)
260    }
261
262    /// Truncates or extends the underlying file, updating the size of this file to become `len`.
263    ///
264    /// Note if `len` is larger than the previous file size, it will result in a sparse file. If
265    /// you'd like to pre-allocate space on disk, use [`Self::allocate()`], which may be followed by
266    /// this method to truncate the file if the new file size is smaller than the previous
267    /// ([`Self::allocate()`] doesn't truncate the file).
268    #[inline(always)]
269    pub fn set_len(&self, len: u64) -> io::Result<()> {
270        self.file.set_len(len)
271    }
272
273    /// Read the exact number of bytes needed to fill `buf` at `offset`.
274    ///
275    /// NOTE: This uses locking and buffering internally, prefer [`Self::read_exact_at_raw()`] if
276    /// you can control data alignment.
277    pub fn read_exact_at(&self, buf: &mut [u8], mut offset: u64) -> io::Result<()> {
278        if buf.is_empty() {
279            return Ok(());
280        }
281
282        let mut scratch_buffer = self.scratch_buffer.lock();
283
284        // This is guaranteed by the constructor
285        debug_assert!(
286            AlignedPage::slice_to_repr(&scratch_buffer)
287                .as_flattened()
288                .len()
289                <= MAX_READ_SIZE
290        );
291
292        // First read up to `MAX_READ_SIZE - padding`
293        let padding = (offset % AlignedPage::SIZE as u64) as usize;
294        let first_unaligned_chunk_size = (MAX_READ_SIZE - padding).min(buf.len());
295        let (unaligned_start, buf) = buf.split_at_mut(first_unaligned_chunk_size);
296        {
297            let bytes_to_read = unaligned_start.len();
298            unaligned_start.copy_from_slice(self.read_exact_at_internal(
299                &mut scratch_buffer,
300                bytes_to_read,
301                offset,
302            )?);
303            offset += unaligned_start.len() as u64;
304        }
305
306        if buf.is_empty() {
307            return Ok(());
308        }
309
310        // Process the rest of the chunks, up to `MAX_READ_SIZE` at a time
311        for buf in buf.chunks_mut(MAX_READ_SIZE) {
312            let bytes_to_read = buf.len();
313            buf.copy_from_slice(self.read_exact_at_internal(
314                &mut scratch_buffer,
315                bytes_to_read,
316                offset,
317            )?);
318            offset += buf.len() as u64;
319        }
320
321        Ok(())
322    }
323
324    /// Write all bytes at `buf` at `offset`.
325    ///
326    /// NOTE: This uses locking and buffering internally, prefer [`Self::write_all_at_raw()`] if you
327    /// can control data alignment.
328    pub fn write_all_at(&self, buf: &[u8], mut offset: u64) -> io::Result<()> {
329        if buf.is_empty() {
330            return Ok(());
331        }
332
333        let mut scratch_buffer = self.scratch_buffer.lock();
334
335        // This is guaranteed by the constructor
336        debug_assert!(
337            AlignedPage::slice_to_repr(&scratch_buffer)
338                .as_flattened()
339                .len()
340                <= MAX_READ_SIZE
341        );
342
343        // First, write up to `MAX_READ_SIZE - padding`
344        let padding = (offset % AlignedPage::SIZE as u64) as usize;
345        let first_unaligned_chunk_size = (MAX_READ_SIZE - padding).min(buf.len());
346        let (unaligned_start, buf) = buf.split_at(first_unaligned_chunk_size);
347        {
348            self.write_all_at_internal(&mut scratch_buffer, unaligned_start, offset)?;
349            offset += unaligned_start.len() as u64;
350        }
351
352        if buf.is_empty() {
353            return Ok(());
354        }
355
356        // Process the rest of the chunks, up to `MAX_READ_SIZE` at a time
357        for buf in buf.chunks(MAX_READ_SIZE) {
358            self.write_all_at_internal(&mut scratch_buffer, buf, offset)?;
359            offset += buf.len() as u64;
360        }
361
362        Ok(())
363    }
364
365    /// Low-level reading into aligned memory.
366    ///
367    /// `offset` needs to be page-aligned as well or use [`Self::read_exact_at()`] if you're willing
368    /// to pay for the corresponding overhead.
369    ///
370    /// Successful result guarantees that the whole `buf` was filled.
371    #[inline]
372    pub fn read_exact_at_raw(
373        &self,
374        mut buf: BorrowedCursor<'_, AlignedPage>,
375        offset: u64,
376    ) -> io::Result<()> {
377        // SAFETY: Only bytes read from the file are written into pages
378        let pages = unsafe { buf.as_mut() };
379        let num_pages = pages.len();
380        let mut bytes = BorrowedBuf::from(pages.as_bytes_mut());
381
382        cfg_select! {
383            unix => {
384                use std::os::unix::fs::FileExt;
385
386                self.file.read_buf_exact_at(bytes.unfilled(), offset)?;
387            }
388            windows => {
389                use std::os::windows::fs::FileExt;
390
391                self.file.seek_read_buf_exact(bytes.unfilled(), offset)?;
392            }
393            _ => {
394                compile_error!("Unsupported platform (consider contributing)");
395            }
396        }
397
398        // SAFETY: All pages were just filled with bytes read from the file
399        unsafe {
400            buf.advance(num_pages);
401        }
402
403        Ok(())
404    }
405
406    /// Low-level writing from aligned memory.
407    ///
408    /// `offset` needs to be page-aligned as well or use [`Self::write_all_at()`] if you're willing
409    /// to pay for the corresponding overhead.
410    #[inline]
411    pub fn write_all_at_raw(&self, buf: &[AlignedPage], offset: u64) -> io::Result<()> {
412        let buf = AlignedPage::slice_to_repr(buf).as_flattened();
413
414        cfg_select! {
415            unix => {
416                use std::os::unix::fs::FileExt;
417
418                self.file.write_all_at(buf, offset)
419            }
420            windows => {
421                use std::os::windows::fs::FileExt;
422
423                self.file.seek_write_all(buf, offset)
424            }
425            _ => {
426                compile_error!("Unsupported platform (consider contributing)");
427            }
428        }
429    }
430
431    /// Access internal [`File`] instance
432    #[inline(always)]
433    pub fn file(&self) -> &File {
434        &self.file
435    }
436
437    fn read_exact_at_internal<'a>(
438        &self,
439        scratch_buffer: &'a mut [AlignedPage],
440        bytes_to_read: usize,
441        offset: u64,
442    ) -> io::Result<&'a [u8]> {
443        let page_aligned_offset = offset / AlignedPage::SIZE as u64 * AlignedPage::SIZE as u64;
444        let padding = (offset - page_aligned_offset) as usize;
445
446        // Make a scratch buffer of a size that is necessary to read aligned memory, accounting
447        // for extra bytes at the beginning and the end that will be thrown away
448        let pages_to_read = (padding + bytes_to_read).div_ceil(AlignedPage::SIZE);
449        let scratch_buffer = &mut scratch_buffer[..pages_to_read];
450
451        self.read_exact_at_raw(
452            BorrowedBuf::from(&mut *scratch_buffer).unfilled(),
453            page_aligned_offset,
454        )?;
455
456        Ok(&AlignedPage::slice_to_repr(scratch_buffer).as_flattened()[padding..][..bytes_to_read])
457    }
458
459    /// Panics on writes over `MAX_READ_SIZE` (including padding on both ends)
460    fn write_all_at_internal(
461        &self,
462        scratch_buffer: &mut [AlignedPage],
463        bytes_to_write: &[u8],
464        offset: u64,
465    ) -> io::Result<()> {
466        let page_aligned_offset = offset / AlignedPage::SIZE as u64 * AlignedPage::SIZE as u64;
467        let padding = (offset - page_aligned_offset) as usize;
468
469        // Calculate the size of the read including padding on both ends
470        let pages_to_read = (padding + bytes_to_write.len()).div_ceil(AlignedPage::SIZE);
471
472        if padding == 0 && bytes_to_write.len() == pages_to_read * AlignedPage::SIZE {
473            let scratch_buffer = &mut scratch_buffer[..pages_to_read];
474            AlignedPage::slice_mut_to_repr(scratch_buffer)
475                .as_flattened_mut()
476                .copy_from_slice(bytes_to_write);
477            self.write_all_at_raw(scratch_buffer, offset)?;
478        } else {
479            let scratch_buffer = &mut scratch_buffer[..pages_to_read];
480            // Read whole pages where `bytes_to_write` will be written
481            self.read_exact_at_raw(
482                BorrowedBuf::from(&mut *scratch_buffer).unfilled(),
483                page_aligned_offset,
484            )?;
485            // Update the contents of existing pages and write into the file
486            AlignedPage::slice_mut_to_repr(scratch_buffer).as_flattened_mut()[padding..]
487                [..bytes_to_write.len()]
488                .copy_from_slice(bytes_to_write);
489            self.write_all_at_raw(scratch_buffer, page_aligned_offset)?;
490        }
491
492        Ok(())
493    }
494}
ab_direct_io_file/lib.rs

ab_direct_io_file/
lib.rs
