polars_io/utils/byte_source/
dio_align.rs

1//! O_DIRECT alignment discovery.
2//!
3//! `O_DIRECT` requires the file offset, the transfer length, and the buffer
4//! address to be aligned.
5//!
6//! Sources, in order of accuracy:
7//!  1. `statx(STATX_DIOALIGN)` (Linux 6.1+) - authoritative, per-file, and
8//!     distinguishes memory alignment from offset/length alignment. It can
9//!     also say *definitively* that a file does not support direct I/O.
10//!  2. `/sys/dev/block/<major>:<minor>/queue/logical_block_size` - the device
11//!     sector size. Right for most filesystems, but misses cases where the
12//!     filesystem imposes something stricter.
13//!
14//! When neither can answer, the alignment stays unknown and the caller reads
15//! through the page cache.
16
17#[cfg(target_os = "linux")]
18use std::os::fd::AsRawFd;
19
20/// Substituted for a reported value that cannot be an alignment.
21#[cfg(target_os = "linux")]
22const FALLBACK_ALIGN: usize = 4096;
23
24/// Floor for the buffer alignment, applied even when the kernel reports less.
25///
26/// `statx` can report a memory alignment below the sector size (ext4 reports 4),
27/// and the other two probes infer the alignment rather than being told it.
28#[cfg(target_os = "linux")]
29const MIN_MEM_ALIGN: usize = 512;
30
31#[derive(Debug, Clone, Copy)]
32pub struct DioAlign {
33    /// Alignment required for the file offset and the transfer length.
34    pub offset: usize,
35    /// Alignment required for the buffer address.
36    /// Invariant: no less than `MIN_MEM_ALIGN`.
37    pub memory: usize,
38}
39
40/// What `statx` was able to tell us.
41#[cfg(target_os = "linux")]
42#[derive(Debug, Clone)]
43enum StatxAlign {
44    /// The kernel reported concrete alignments.
45    Known(DioAlign),
46    /// The kernel reported zero: this file does not support direct I/O.
47    /// Distinct from `Unknown` - here we must not guess.
48    Unsupported,
49    /// The syscall failed or the mask is unavailable (kernel < 6.1).
50    Unknown,
51}
52
53impl DioAlign {
54    /// Round `offset` down and `end` up to the offset alignment.
55    ///
56    /// Returns `(lo, hi, pad)` where `pad` is how far into the aligned span the
57    /// caller's data begins. Note `hi` is deliberately *not* clamped to the
58    /// file size: the length must stay aligned, so a read at EOF is expected to
59    /// come up short and the caller must handle the tail.
60    pub fn span(&self, offset: u64, len: usize) -> (u64, u64, usize) {
61        debug_assert!(self.offset > 0);
62
63        let a = self.offset as u64;
64        let lo = offset & !(a - 1);
65        let hi = (offset + len as u64).next_multiple_of(a);
66        (lo, hi, (offset - lo) as usize)
67    }
68
69    /// Query the alignment `O_DIRECT` requires for this file.
70    ///
71    /// `None` means direct I/O is not supported here and the caller should use
72    /// buffered reads.
73    #[cfg(target_os = "linux")]
74    pub fn probe(file: &std::fs::File) -> Option<Self> {
75        match statx_dioalign(file) {
76            StatxAlign::Unsupported => None,
77            StatxAlign::Known(a) => Some(Self::new(a.offset, a.memory)),
78            StatxAlign::Unknown => {
79                sysfs_logical_block_size(file).map(|a| Self::new(a.offset, a.memory))
80            },
81        }
82    }
83
84    #[cfg(not(target_os = "linux"))]
85    pub fn probe(_file: &std::fs::File) -> Option<Self> {
86        None
87    }
88
89    /// Satisfy type invariants.
90    #[cfg(target_os = "linux")]
91    fn new(offset: usize, memory: usize) -> Self {
92        Self {
93            offset: normalize(offset),
94            memory: normalize(memory).max(MIN_MEM_ALIGN),
95        }
96    }
97}
98
99#[cfg(target_os = "linux")]
100fn normalize(v: usize) -> usize {
101    if v == 0 || !v.is_power_of_two() {
102        FALLBACK_ALIGN
103    } else {
104        v
105    }
106}
107
108#[cfg(target_os = "linux")]
109fn statx_dioalign(file: &std::fs::File) -> StatxAlign {
110    let mut stx: libc::statx = unsafe { std::mem::zeroed() };
111    let rc = unsafe {
112        libc::statx(
113            file.as_raw_fd(),
114            c"".as_ptr(),
115            libc::AT_EMPTY_PATH,
116            libc::STATX_DIOALIGN,
117            &mut stx,
118        )
119    };
120    if rc != 0 || stx.stx_mask & libc::STATX_DIOALIGN == 0 {
121        return StatxAlign::Unknown;
122    }
123
124    let offset = stx.stx_dio_offset_align as usize;
125    let memory = stx.stx_dio_mem_align as usize;
126
127    if offset == 0 || memory == 0 {
128        StatxAlign::Unsupported
129    } else {
130        StatxAlign::Known(DioAlign { offset, memory })
131    }
132}
133
134/// `/sys/dev/block/<major>:<minor>/queue/logical_block_size`.
135///
136/// Keyed by device number, so this needs no device-name lookup. Used when the
137/// kernel predates `STATX_DIOALIGN`.
138#[cfg(target_os = "linux")]
139fn sysfs_logical_block_size(file: &std::fs::File) -> Option<DioAlign> {
140    use std::os::unix::fs::MetadataExt;
141
142    let dev = file.metadata().ok()?.dev();
143    let (major, minor) = (libc::major(dev), libc::minor(dev));
144    let path = format!("/sys/dev/block/{major}:{minor}/queue/logical_block_size");
145    let v: usize = std::fs::read_to_string(path).ok()?.trim().parse().ok()?;
146
147    // The device sector size constrains offset and length. Buffer alignment is
148    // not reported here, so assume the same per historical behavior.
149    Some(DioAlign {
150        offset: v,
151        memory: v,
152    })
153}
154
155#[cfg(test)]
156mod tests {
157    use super::*;
158
159    #[test]
160    fn span_invariants() {
161        for align in [512usize, 4096] {
162            let a = DioAlign {
163                offset: align,
164                memory: align,
165            };
166            for (off, len) in [(206959u64, 206955usize), (4095, 2), (0, 1), (0, align)] {
167                let (lo, hi, pad) = a.span(off, len);
168                assert_eq!(lo as usize % align, 0, "lo unaligned");
169                assert_eq!((hi - lo) as usize % align, 0, "span unaligned");
170                assert!(hi >= off + len as u64, "span does not cover the range");
171                assert!(pad + len <= (hi - lo) as usize, "pad + len exceeds span");
172            }
173        }
174    }
175}
polars_io/utils/byte_source/dio_align.rs

polars_io/utils/byte_source/
dio_align.rs
