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

polars_io/utils/byte_source/
dio_align.rs
