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//!     Only available on glibc: `libc` does not expose `statx` on 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#[cfg(all(target_os = "linux", target_env = "gnu"))]
111fn statx_dioalign(file: &std::fs::File) -> StatxAlign {
112    let mut stx: libc::statx = unsafe { std::mem::zeroed() };
113    let rc = unsafe {
114        libc::statx(
115            file.as_raw_fd(),
116            c"".as_ptr(),
117            libc::AT_EMPTY_PATH,
118            libc::STATX_DIOALIGN,
119            &mut stx,
120        )
121    };
122    if rc != 0 || stx.stx_mask & libc::STATX_DIOALIGN == 0 {
123        return StatxAlign::Unknown;
124    }
125
126    let offset = stx.stx_dio_offset_align as usize;
127    let memory = stx.stx_dio_mem_align as usize;
128
129    if offset == 0 || memory == 0 {
130        StatxAlign::Unsupported
131    } else {
132        StatxAlign::Known(DioAlign { offset, memory })
133    }
134}
135
136/// `/sys/dev/block/<major>:<minor>/queue/logical_block_size`.
137///
138/// Keyed by device number, so this needs no device-name lookup. Used when the
139/// kernel predates `STATX_DIOALIGN`, or when `statx` is unavailable at all
140/// (musl).
141#[cfg(target_os = "linux")]
142fn sysfs_logical_block_size(file: &std::fs::File) -> Option<DioAlign> {
143    use std::os::unix::fs::MetadataExt;
144
145    let dev = file.metadata().ok()?.dev();
146    let (major, minor) = (libc::major(dev), libc::minor(dev));
147    let path = format!("/sys/dev/block/{major}:{minor}/queue/logical_block_size");
148    let v: usize = std::fs::read_to_string(path).ok()?.trim().parse().ok()?;
149
150    // The device sector size constrains offset and length. Buffer alignment is
151    // not reported here, so assume the same per historical behavior.
152    Some(DioAlign {
153        offset: v,
154        memory: v,
155    })
156}
157
158#[cfg(test)]
159mod tests {
160    use super::*;
161
162    #[test]
163    fn span_invariants() {
164        for align in [512usize, 4096] {
165            let a = DioAlign {
166                offset: align,
167                memory: align,
168            };
169            for (off, len) in [(206959u64, 206955usize), (4095, 2), (0, 1), (0, align)] {
170                let (lo, hi, pad) = a.span(off, len);
171                assert_eq!(lo as usize % align, 0, "lo unaligned");
172                assert_eq!((hi - lo) as usize % align, 0, "span unaligned");
173                assert!(hi >= off + len as u64, "span does not cover the range");
174                assert!(pad + len <= (hi - lo) as usize, "pad + len exceeds span");
175            }
176        }
177    }
178}
polars_io/utils/byte_source/dio_align.rs

polars_io/utils/byte_source/
dio_align.rs
