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

ab_direct_io_file/
lib.rs
