metrale_storage/backend/
io_uring.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The Linux file backend: `io_uring` fixed-buffer reads (SQPOLL,
4//! registered pinned buffers) into `qd` bounce buffers, each copied to the
5//! device on the caller's stream.
6//!
7//! Owner: metrale-storage high-speed swap.
8//! Invariants:
9//! - A bounce buffer is refilled only after the event recorded behind its last
10//!   host-to-device copy has completed (`wait_buffer_free`).
11
12use anyhow::{Context, Result, bail};
13use io_uring::{IoUring, opcode, types};
14use std::ffi::c_void;
15use std::os::fd::RawFd;
16
17use super::{ReadRequest, StorageBackend};
18use crate::cuda_min::{CudaEvent, PinnedBuffer, copy_h_to_d_async, stream_sync};
19use crate::group::{GroupKey, GroupLayout};
20use crate::layout::Layout;
21
22pub struct IoUringBackend {
23    layout: Layout,
24    ring: IoUring,
25    buffers: Vec<PinnedBuffer>,
26    // 2026-09-25: The event behind each buffer's last copy; `None` when no copy
27    // is pending.
28    events: Vec<Option<CudaEvent>>,
29    qd: usize,
30}
31
32impl IoUringBackend {
33    pub fn new(layout: Layout, qd: usize) -> Result<Self> {
34        if qd == 0 {
35            bail!("queue depth must be ≥ 1");
36        }
37        // 2026-09-25: The SQPOLL kernel thread idles 2000 ms before it sleeps.
38        let ring = IoUring::builder()
39            .setup_sqpoll(2_000)
40            .build(qd as u32)
41            .context("io_uring build")?;
42
43        let group_bytes = layout.group_bytes() as usize;
44        let mut buffers = Vec::with_capacity(qd);
45        for _ in 0..qd {
46            buffers.push(PinnedBuffer::new(group_bytes)?);
47        }
48        // 2026-09-25: After registration, `ReadFixed` with index `i` lands in
49        // `buffers[i]`.
50        let iovecs: Vec<libc::iovec> = buffers
51            .iter()
52            .map(|b| libc::iovec {
53                iov_base: b.ptr,
54                iov_len: b.bytes,
55            })
56            .collect();
57        unsafe {
58            ring.submitter()
59                .register_buffers(&iovecs)
60                .context("register_buffers")?;
61        }
62        let events: Vec<Option<CudaEvent>> = (0..qd).map(|_| None).collect();
63        Ok(Self {
64            layout,
65            ring,
66            buffers,
67            events,
68            qd,
69        })
70    }
71
72    pub fn layout(&self) -> &Layout {
73        &self.layout
74    }
75
76    /// 2026-09-25: Advise the kernel to drop every layer file's page cache
77    /// (`POSIX_FADV_DONTNEED`), so later reads go to the device.
78    pub fn drop_pagecache(&self) {
79        for layer in 0..self.layout.spec.num_layers {
80            let fd = self.layout.fd(layer);
81            unsafe { libc::posix_fadvise(fd, 0, 0, libc::POSIX_FADV_DONTNEED) };
82        }
83    }
84
85    /// 2026-09-25: Wait for the last copy out of `buf_idx`, if one is pending,
86    /// so the buffer can take a new read.
87    fn wait_buffer_free(&mut self, buf_idx: usize) -> Result<()> {
88        if let Some(ev) = self.events[buf_idx].take() {
89            ev.sync()?;
90        }
91        Ok(())
92    }
93
94    /// 2026-09-25: Queue one fixed-buffer read of `bytes` at `offset` into
95    /// `buf_idx`, tagged `user_data`. Fails when the submission queue is full.
96    fn submit_read(
97        &mut self,
98        fd: RawFd,
99        offset: u64,
100        bytes: u32,
101        buf_idx: u16,
102        user_data: u64,
103    ) -> Result<()> {
104        let buf_ptr = self.buffers[buf_idx as usize].ptr as *mut u8;
105        let read_e = opcode::ReadFixed::new(types::Fd(fd), buf_ptr, bytes, buf_idx)
106            .offset(offset)
107            .build()
108            .user_data(user_data);
109        unsafe {
110            self.ring
111                .submission()
112                .push(&read_e)
113                .map_err(|_| anyhow::anyhow!("io_uring SQ full"))?;
114        }
115        Ok(())
116    }
117}
118
119impl StorageBackend for IoUringBackend {
120    fn read(&mut self, requests: &[ReadRequest], stream: u64) -> Result<()> {
121        let bytes = self.layout.group_bytes() as u32;
122        // 2026-09-25: user_data is `req_index << 16 | buf_index`.
123        if requests.len() > u16::MAX as usize {
124            bail!("io_uring batch too large: {}", requests.len());
125        }
126
127        let mut next_submit = 0;
128        let mut completed = 0;
129        // 2026-09-25: Free buffers form a stack; a buffer leaves it at submit
130        // and returns when its completion is reaped.
131        let mut free_bufs: Vec<u16> = (0..self.qd as u16).rev().collect();
132
133        while completed < requests.len() {
134            while next_submit < requests.len() {
135                let Some(&buf_idx) = free_bufs.last() else {
136                    break;
137                };
138                self.wait_buffer_free(buf_idx as usize)?;
139                free_bufs.pop();
140                let req = &requests[next_submit];
141                let fd = self.layout.fd(req.group.layer);
142                let off = self.layout.offset(req.group);
143                let user = ((next_submit as u64) << 16) | (buf_idx as u64);
144                self.submit_read(fd, off, bytes, buf_idx, user)?;
145                next_submit += 1;
146            }
147            self.ring
148                .submit_and_wait(1)
149                .context("io_uring submit_and_wait")?;
150            let cq = self.ring.completion();
151            for cqe in cq {
152                let user = cqe.user_data();
153                let buf_idx = (user & 0xFFFF) as u16;
154                let req_idx = (user >> 16) as usize;
155                let result = cqe.result();
156                if result < 0 {
157                    bail!("io_uring read failed for req {req_idx}: errno {}", -result);
158                }
159                if result as u32 != bytes {
160                    bail!("io_uring short read: req {req_idx} got {result}, expected {bytes}");
161                }
162                let req = &requests[req_idx];
163                let buf = &self.buffers[buf_idx as usize];
164                copy_h_to_d_async(
165                    req.dst_dev_ptr,
166                    buf.ptr as *const c_void,
167                    bytes as usize,
168                    stream,
169                )?;
170                let ev = CudaEvent::new()?;
171                ev.record(stream)?;
172                self.events[buf_idx as usize] = Some(ev);
173                free_bufs.push(buf_idx);
174                completed += 1;
175            }
176        }
177        // 2026-09-25: The `read` contract: the data is visible to later work
178        // on `stream`.
179        stream_sync(stream)?;
180        // 2026-09-25: The sync above completed every recorded event.
181        for slot in self.events.iter_mut() {
182            *slot = None;
183        }
184        Ok(())
185    }
186
187    fn write_from_host(&mut self, key: GroupKey, src: &[u8]) -> Result<()> {
188        let bytes = self.layout.group_bytes() as usize;
189        if src.len() != bytes {
190            bail!(
191                "write_from_host: src len {} != group bytes {bytes}",
192                src.len()
193            );
194        }
195        // 2026-09-25: The layer files are O_DIRECT (layout.rs), so `src` is
196        // staged through pinned buffer 0.
197        self.wait_buffer_free(0)?;
198        unsafe {
199            std::ptr::copy_nonoverlapping(src.as_ptr(), self.buffers[0].ptr as *mut u8, bytes);
200        }
201        let fd = self.layout.fd(key.layer);
202        let off = self.layout.offset(key) as i64;
203        let n = unsafe { libc::pwrite(fd, self.buffers[0].ptr, bytes, off) };
204        if n != bytes as isize {
205            bail!(
206                "pwrite {bytes}@{off} returned {n}, errno {}",
207                std::io::Error::last_os_error()
208            );
209        }
210        Ok(())
211    }
212
213    fn group_layout(&self) -> GroupLayout {
214        self.layout.spec
215    }
216}
217
218#[cfg(test)]
219mod tests {
220    use super::*;
221    use crate::cuda_min::{CudaCtx, DeviceBuffer, copy_d_to_h_async};
222    use crate::group::{GroupKey, GroupLayout, KvKind};
223    use std::path::PathBuf;
224
225    fn tempdir(name: &str) -> PathBuf {
226        let p =
227            std::env::temp_dir().join(format!("metrale-iouring-{}-{}", name, std::process::id()));
228        let _ = std::fs::remove_dir_all(&p);
229        std::fs::create_dir_all(&p).unwrap();
230        p
231    }
232
233    #[test]
234    #[ignore = "requires GPU"]
235    fn write_then_read_round_trip() {
236        let _ctx = CudaCtx::new(0).expect("cuda init");
237        let dir = tempdir("rt");
238        let spec = GroupLayout::new(1, 4, 1, 16, 128, 2, 4096);
239        let layout = Layout::create(&dir, spec).unwrap();
240        let mut backend = IoUringBackend::new(layout, 4).unwrap();
241        let bytes = backend.layout().group_bytes() as usize;
242        // 2026-09-25: Four blocks with distinct patterns, as many as the queue
243        // depth.
244        let patterns: Vec<(GroupKey, Vec<u8>)> = (0..4u32)
245            .map(|b| {
246                let k = GroupKey::new(0, b, 0, KvKind::K);
247                let pat: Vec<u8> = (0..bytes)
248                    .map(|i| ((i + b as usize) & 0xFF) as u8)
249                    .collect();
250                (k, pat)
251            })
252            .collect();
253        for (k, p) in &patterns {
254            backend.write_from_host(*k, p).unwrap();
255        }
256        backend.drop_pagecache();
257        let dev: Vec<DeviceBuffer> = patterns
258            .iter()
259            .map(|_| DeviceBuffer::new(bytes).unwrap())
260            .collect();
261        let reqs: Vec<ReadRequest> = patterns
262            .iter()
263            .zip(&dev)
264            .map(|((k, _), d)| ReadRequest {
265                group: *k,
266                dst_dev_ptr: d.ptr,
267            })
268            .collect();
269        backend.read(&reqs, _ctx.stream).unwrap();
270        for ((_, expected), d) in patterns.iter().zip(&dev) {
271            let mut got = vec![0_u8; bytes];
272            copy_d_to_h_async(got.as_mut_ptr() as *mut c_void, d.ptr, bytes, _ctx.stream).unwrap();
273            stream_sync(_ctx.stream).unwrap();
274            assert_eq!(&got, expected);
275        }
276        std::fs::remove_dir_all(&dir).ok();
277    }
278}
metrale_storage/backend/io_uring.rs

metrale_storage/backend/
io_uring.rs
