vyre_runtime/uring/
stream.rs

1//! `AsyncUringStream`  -  drives io_uring reads into GPU-visible memory
2//! and advances the megakernel tail pointer on each completion.
3//!
4//! The critical safety contract: every byte read lands in a
5//! [`GpuMappedBuffer`]. Compatibility ingest uses registered
6//! host-visible GPU mappings; canonical native ingest uses BAR1 peer memory
7//! via [`GpuMappedBuffer::from_bar1_peer_with_owner`] plus NVMe passthrough.
8//! The io_uring writer never targets an ordinary userspace bounce buffer.
9
10use super::ring::IoUringState;
11use crate::PipelineError;
12use core::marker::PhantomData;
13use core::sync::atomic::{AtomicU32, Ordering};
14
15/// Minimal `iovec` struct matching the Linux ABI for `readv`.
16#[repr(C)]
17#[derive(Debug, Clone, Copy)]
18pub struct Iovec {
19    /// Target buffer address for this chunk of the read.
20    pub iov_base: *mut core::ffi::c_void,
21    /// Byte length of the target buffer.
22    pub iov_len: usize,
23}
24
25/// `IORING_OP_READV`  -  scatter-read into an array of iovecs.
26pub const IORING_OP_READV: u8 = 1;
27/// `IORING_OP_READ_FIXED`  -  read into a pre-registered buffer.
28pub const IORING_OP_READ_FIXED: u8 = 22;
29/// `IORING_OP_URING_CMD`  -  vendor-specific passthrough (NVMe). Kernel 6.0+.
30pub const IORING_OP_URING_CMD: u8 = 46;
31
32/// GPU-visible memory region that io_uring is allowed to DMA into.
33///
34/// Compatibility constructors cover host-visible shared mappings. The BAR1
35/// constructor covers the native GPUDirect path where NVMe DMA lands directly
36/// in GPU-owned memory.
37pub struct GpuMappedBuffer<'a> {
38    ptr: *mut u8,
39    len: usize,
40    _owner: PhantomData<&'a mut [u8]>,
41}
42
43// SAFETY: Send + Sync because (a) the constructor's safety contract
44// requires the caller to commit the lifetime invariant, and (b) the
45// raw pointer is only dereferenced by the kernel via io_uring  -
46// vyre-runtime never reads through it directly.
47unsafe impl Send for GpuMappedBuffer<'_> {}
48unsafe impl Sync for GpuMappedBuffer<'_> {}
49
50macro_rules! define_mapped_owner_constructor {
51    ($name:ident, $ptr:ident, $doc:expr) => {
52        #[doc = $doc]
53        pub unsafe fn $name<O: ?Sized>(_owner: &'a mut O, $ptr: *mut u8, len: usize) -> Self {
54            Self {
55                ptr: $ptr,
56                len,
57                _owner: PhantomData,
58            }
59        }
60    };
61}
62
63impl<'a> GpuMappedBuffer<'a> {
64    /// Construct from a borrowed host-visible byte slice.
65    ///
66    /// # Safety
67    ///
68    /// The caller asserts:
69    /// - `slice` aliases a device allocation created with host-visible
70    ///   host-shared usage bits by the concrete backend.
71    /// - No other code reads or writes through `slice` while the
72    ///   returned handle is alive.
73    pub unsafe fn from_host_visible_slice(slice: &'a mut [u8]) -> Self {
74        Self {
75            ptr: slice.as_mut_ptr(),
76            len: slice.len(),
77            _owner: PhantomData,
78        }
79    }
80
81    define_mapped_owner_constructor!(
82        from_host_visible_owner,
83        ptr,
84        concat!(
85            "Construct from a raw pointer plus an explicit owner anchor.\n\n",
86            "The borrow on `owner` forces the mapped region to outlive every derived ",
87            "[`AsyncUringStream`].\n\n",
88            "# Safety\n\n",
89            "The caller must ensure that `ptr` names a `len`-byte host-visible GPU ",
90            "allocation owned by `owner`, and that no other code accesses the region ",
91            "while the returned handle is alive."
92        )
93    );
94
95    /// Duplicate the mapped-buffer handle for the same underlying region.
96    ///
97    /// # Safety
98    ///
99    /// The caller must uphold the same aliasing and lifetime guarantees as
100    /// [`GpuMappedBuffer::from_host_visible_slice`]. This does not clone memory;
101    /// it creates another handle to the same mapped bytes.
102    pub unsafe fn duplicate(&self) -> Self {
103        Self {
104            ptr: self.ptr,
105            len: self.len,
106            _owner: PhantomData,
107        }
108    }
109
110    /// Carve out a sub-region of this mapped buffer.
111    ///
112    /// This preserves the original constructor contract: the returned
113    /// handle aliases the same host-visible GPU allocation and carries
114    /// no ownership of its own.
115    ///
116    /// # Errors
117    ///
118    /// Returns [`PipelineError::QueueFull`] when `offset + len`
119    /// exceeds the mapped buffer bounds.
120    pub fn sub_region(&self, offset: usize, len: usize) -> Result<Self, crate::PipelineError> {
121        let _end = vyre_driver::accounting::checked_usize_byte_range_end_lazy(
122            offset,
123            len,
124            self.len,
125            || {
126                crate::PipelineError::QueueFull {
127                queue: "submission",
128                fix: "GpuMappedBuffer::sub_region offset + len overflows usize; reduce slot size or enlarge the staging buffer",
129            }
130            },
131            |_| {
132                crate::PipelineError::QueueFull {
133                queue: "submission",
134                fix: "GpuMappedBuffer::sub_region exceeds the mapped allocation; reduce slot size or enlarge the staging buffer",
135            }
136            },
137        )?;
138        Ok(Self {
139            ptr: self.ptr.wrapping_add(offset),
140            len,
141            _owner: PhantomData,
142        })
143    }
144
145    /// Byte length of the mapped region.
146    #[must_use]
147    pub fn len(&self) -> usize {
148        self.len
149    }
150
151    /// Whether the region is empty.
152    #[must_use]
153    pub fn is_empty(&self) -> bool {
154        self.len == 0
155    }
156
157    /// Raw pointer for io_uring submission. Crate-private.
158    pub(crate) fn as_ptr(&self) -> *mut u8 {
159        self.ptr
160    }
161
162    /// Borrow the mapped bytes as a mutable slice.
163    ///
164    /// # Safety
165    ///
166    /// The caller must ensure exclusive mutable access to the region for the
167    /// lifetime of the returned slice.
168    pub unsafe fn as_mut_slice(&mut self) -> &mut [u8] {
169        // SAFETY: Safe FFI / low-level operation verified and audited for Release compliance.
170        unsafe { core::slice::from_raw_parts_mut(self.ptr, self.len) }
171    }
172
173    define_mapped_owner_constructor!(
174        from_bar1_peer_with_owner,
175        peer_ptr,
176        concat!(
177            "Construct from a PCIe peer-memory pointer for direct storage DMA.\n\n",
178            "# Safety\n\n",
179            "The caller must ensure that `peer_ptr` names a GPU allocation suitable ",
180            "for peer DMA, that the allocation outlives the handle, and that the ",
181            "io_uring kernel and storage driver both support DMA mapping."
182        )
183    );
184}
185
186/// Streaming reader that pushes chunked reads into an io_uring SQ and
187/// advances an atomic tail pointer the megakernel observes.
188pub struct AsyncUringStream<'a> {
189    pub(crate) ring_state: IoUringState,
190    pub(crate) gpu_buffer: GpuMappedBuffer<'a>,
191    pub(crate) megakernel_tail: &'a AtomicU32,
192    pub(crate) inflight: u32,
193    pub(crate) pending_submissions: u32,
194}
195
196// SAFETY: raw pointer fields covered by GpuMappedBuffer's contract +
197// the constructor's safety commitment on megakernel_tail_ptr.
198unsafe impl Send for AsyncUringStream<'_> {}
199unsafe impl Sync for AsyncUringStream<'_> {}
200
201impl<'a> AsyncUringStream<'a> {
202    /// Create a stream bound to the given ring state, GPU-mapped
203    /// buffer, and megakernel tail pointer.
204    pub fn new(
205        ring_state: IoUringState,
206        gpu_buffer: GpuMappedBuffer<'a>,
207        megakernel_tail: &'a AtomicU32,
208    ) -> Self {
209        Self {
210            ring_state,
211            gpu_buffer,
212            megakernel_tail,
213            inflight: 0,
214            pending_submissions: 0,
215        }
216    }
217
218    /// Rebind the target mapped buffer for future submissions.
219    pub fn replace_buffer(&mut self, gpu_buffer: GpuMappedBuffer<'a>) {
220        self.gpu_buffer = gpu_buffer;
221    }
222
223    /// Submit a scattered read of `len` bytes at file offset `offset`
224    /// into the slot at `chunk_idx * len` within the GPU buffer.
225    ///
226    /// # Errors
227    ///
228    /// - [`PipelineError::QueueFull`] if the SQ is full OR the
229    ///   destination slot exceeds buffer bounds.
230    /// - Range errors surface later as [`PipelineError::IoUringSyscall`]
231    ///   on `poll` if the kernel rejects the SQE.
232    ///
233    /// # Safety
234    ///
235    /// `iovs_storage` must live until this SQE's completion is reaped;
236    /// the kernel dereferences `iov_base` at I/O time, not submit time.
237    pub unsafe fn submit_read_to_gpu(
238        &mut self,
239        fd: i32,
240        offset: u64,
241        len: u32,
242        chunk_idx: usize,
243        iovs_storage: &mut [Iovec],
244    ) -> Result<(), PipelineError> {
245        if iovs_storage.is_empty() {
246            return Err(PipelineError::QueueFull {
247                queue: "submission",
248                fix: "caller supplied empty iovs_storage; pass at least one slot",
249            });
250        }
251        let target_offset = checked_chunk_target_offset(chunk_idx, len)?;
252        // SAFETY: Safe FFI / low-level operation verified and audited for Release compliance.
253        unsafe { self.submit_read_to_gpu_at(fd, offset, len, target_offset, iovs_storage) }
254    }
255
256    /// Submit a read directly into a byte offset inside the mapped buffer.
257    ///
258    /// Unlike [`AsyncUringStream::submit_read_to_gpu`], this method does not
259    /// derive the destination from a fixed chunk index. Wrappers that stream
260    /// variable-sized shards can place each read contiguously in a staging
261    /// buffer without being forced into `chunk_idx * len` layout.
262    ///
263    /// # Errors
264    ///
265    /// - [`PipelineError::QueueFull`] if the SQ is full OR the target range
266    ///   exceeds the mapped buffer bounds.
267    ///
268    /// # Safety
269    ///
270    /// `iovs_storage` must live until this SQE's completion is reaped.
271    pub unsafe fn submit_read_to_gpu_at(
272        &mut self,
273        fd: i32,
274        offset: u64,
275        len: u32,
276        target_offset: u64,
277        iovs_storage: &mut [Iovec],
278    ) -> Result<(), PipelineError> {
279        // SAFETY: registered fixed buffers + file index are valid for the lifetime
280        // of the ring; the SQE is built on the ring's own SQ slot.
281        unsafe {
282            self.submit_read_to_gpu_at_with_user_data(
283                fd,
284                offset,
285                len,
286                target_offset,
287                target_offset,
288                iovs_storage,
289            )
290        }
291    }
292
293    /// Submit a read into an arbitrary byte offset and preserve caller-defined
294    /// `user_data` for completion correlation.
295    ///
296    /// # Errors
297    ///
298    /// Returns [`PipelineError::QueueFull`] when the SQ is full, the iovec
299    /// storage is empty, or the target range exceeds the mapped GPU buffer.
300    ///
301    /// # Safety
302    ///
303    /// `iovs_storage` must live until this SQE's completion is reaped.
304    pub unsafe fn submit_read_to_gpu_at_with_user_data(
305        &mut self,
306        fd: i32,
307        offset: u64,
308        len: u32,
309        target_offset: u64,
310        user_data: u64,
311        iovs_storage: &mut [Iovec],
312    ) -> Result<(), PipelineError> {
313        if iovs_storage.is_empty() {
314            return Err(PipelineError::QueueFull {
315                queue: "submission",
316                fix: "caller supplied empty iovs_storage; pass at least one slot",
317            });
318        }
319        let end = checked_target_end(target_offset, len)?;
320        let gpu_len = usize_to_u64(self.gpu_buffer.len(), "mapped GPU buffer length")?;
321        if end > gpu_len {
322            return Err(PipelineError::QueueFull {
323                queue: "submission",
324                fix: "target_offset + len exceeds GpuMappedBuffer length; enlarge the buffer or reduce the read size",
325            });
326        }
327
328        let Some(sqe) = self.ring_state.get_sqe() else {
329            return Err(PipelineError::QueueFull {
330                queue: "submission",
331                fix: "SQ full; call AsyncUringStream::poll to drain completions then retry",
332            });
333        };
334
335        // SAFETY: bounds-checked above; writing to a sub-region of
336        // the host-visible GpuMappedBuffer the caller committed.
337        let target_addr = unsafe {
338            self.gpu_buffer
339                .as_ptr()
340                .add(u64_to_usize(target_offset, "target offset")?)
341        };
342
343        iovs_storage[0] = Iovec {
344            iov_base: target_addr.cast::<core::ffi::c_void>(),
345            iov_len: u32_to_usize(len, "read length")?,
346        };
347
348        sqe.opcode = IORING_OP_READV;
349        sqe.fd = fd;
350        sqe.user_data_or_off = offset;
351        sqe.addr = pointer_addr_u64(iovs_storage.as_ptr(), "readv iovec pointer")?;
352        sqe.len = 1;
353        sqe.user_data = user_data;
354
355        self.ring_state.commit_sqe();
356        increment_queue_counter(&mut self.inflight, "inflight SQE count")?;
357        increment_queue_counter(&mut self.pending_submissions, "pending submission count")?;
358
359        Ok(())
360    }
361
362    /// Submit any queued SQEs to the kernel.
363    ///
364    /// SQPOLL can pick up tail updates on its own, but wrappers that rely on
365    /// bounded latency should not depend on the polling thread waking
366    /// promptly. Flushing pending submissions makes progress explicit.
367    pub fn flush_submissions(&mut self) -> Result<(), PipelineError> {
368        if self.pending_submissions == 0 {
369            return Ok(());
370        }
371        if self.ring_state.uses_sqpoll() {
372            if self.ring_state.sq_needs_wakeup() {
373                self.ring_state.wake_sqpoll()?;
374            }
375        } else {
376            self.ring_state.enter(self.pending_submissions, 0, 0)?;
377        }
378        self.pending_submissions = 0;
379        Ok(())
380    }
381
382    /// Reap available completions, advancing the megakernel tail
383    /// pointer once per success. Returns completions reaped.
384    ///
385    /// # Errors
386    ///
387    /// Returns [`PipelineError::IoUringSyscall`] on the first CQE
388    /// reporting `res < 0`. Remaining CQEs are still drained so the
389    /// ring does not overflow, but only the first failure is
390    /// returned  -  caller re-polls to pick up subsequent errors or
391    /// successes.
392    pub fn poll(&mut self) -> Result<u32, PipelineError> {
393        self.flush_submissions()?;
394        let mut completed: u32 = 0;
395        let mut first_error: Option<PipelineError> = None;
396
397        while let Some(cqe) = self.ring_state.peek_cqe() {
398            let res = cqe.res;
399            self.ring_state.advance_cq();
400            decrement_queue_counter(&mut self.inflight, "inflight SQE count")?;
401
402            if res < 0 {
403                if first_error.is_none() {
404                    first_error = Some(PipelineError::IoUringSyscall {
405                        syscall: "io_uring_cqe",
406                        errno: -res,
407                        fix: "inspect user_data to identify the failed SQE; common causes: EIO on disk, EFAULT on bad iovec, EINVAL on misaligned offset",
408                    });
409                }
410                continue;
411            }
412
413            // Successful DMA: bytes landed in GPU-visible memory. Tail
414            // publication is batched after CQ drain so one poll with N
415            // completions performs one release atomic instead of N.
416            completed = vyre_driver::accounting::checked_add_u32_value(
417                completed,
418                1,
419                PipelineError::QueueFull {
420                    queue: "completion",
421                    fix: "io_uring completion count overflowed u32; drain completions more frequently",
422                },
423            )?;
424        }
425
426        if completed != 0 {
427            self.megakernel_tail.fetch_add(completed, Ordering::Release);
428        }
429
430        match first_error {
431            Some(err) => Err(err),
432            None => Ok(completed),
433        }
434    }
435
436    /// Flush pending submissions + wait for at least one completion.
437    ///
438    /// # Errors
439    ///
440    /// Returns [`PipelineError::IoUringSyscall`] if `io_uring_enter`
441    /// fails.
442    pub fn wait_for_completion(&mut self) -> Result<(), PipelineError> {
443        if self.inflight > 0 {
444            self.flush_submissions()?;
445            self.ring_state.enter(0, 1, 1)?;
446            self.poll()?;
447        }
448        Ok(())
449    }
450
451    /// Number of submissions still awaiting completion.
452    #[must_use]
453    pub fn inflight(&self) -> u32 {
454        self.inflight
455    }
456
457    /// Submit an NVMe passthrough command via `IORING_OP_URING_CMD`.
458    /// Requires the `uring-cmd-nvme` feature and Linux kernel 6.0+.
459    ///
460    /// The NVMe SQE layout is encoded by the caller in `nvme_sqe_bytes`
461    /// (64 bytes)  -  the SQE is memcpy'd into the `addr3`+`addr` slots
462    /// the kernel forwards to the NVMe driver. `user_data` is returned
463    /// on the matching CQE so the caller can correlate completions.
464    ///
465    /// # Errors
466    ///
467    /// - [`PipelineError::NvmePassthroughDisabled`] if the
468    ///   `uring-cmd-nvme` feature is not enabled at compile time. This
469    ///   variant is unreachable in this cfg-gated method but remains
470    ///   part of the public error contract shared with the feature-gated
471    ///   implementation.
472    /// - [`PipelineError::QueueFull`] if the SQ is full or the NVMe
473    ///   command buffer is malformed (must be exactly 64 bytes).
474    ///
475    /// # Safety
476    ///
477    /// - `fd` must be an open character device the caller has
478    ///   `IORING_SETUP_CQE32`-compatible access to (e.g. `/dev/ng0n1`).
479    /// - `nvme_sqe_bytes` must encode a valid NVMe command  -  kernel
480    ///   rejection returns an errno on the CQE, but a forged payload
481    ///   can still trigger device-level misbehavior.
482    #[cfg(feature = "uring-cmd-nvme")]
483    pub unsafe fn submit_nvme_passthrough(
484        &mut self,
485        fd: i32,
486        user_data: u64,
487        nvme_sqe_bytes: &[u8],
488    ) -> Result<(), PipelineError> {
489        if nvme_sqe_bytes.len() != 64 {
490            return Err(PipelineError::QueueFull {
491                queue: "submission",
492                fix: "NVMe passthrough SQE must be exactly 64 bytes; see linux/nvme_ioctl.h",
493            });
494        }
495
496        let Some(sqe) = self.ring_state.get_sqe() else {
497            return Err(PipelineError::QueueFull {
498                queue: "submission",
499                fix: "SQ full; call AsyncUringStream::poll to drain completions then retry",
500            });
501        };
502
503        // SAFETY: caller-provided slice is 64 bytes as validated
504        // above; we copy into the 64-byte NVMe passthrough region
505        // the kernel expects (addr + addr3 cover the first 40 bytes;
506        // the remaining 24 live in the SQE's inline payload).
507        let nvme_ptr = nvme_sqe_bytes.as_ptr();
508        sqe.opcode = IORING_OP_URING_CMD;
509        sqe.fd = fd;
510        sqe.user_data_or_off = 0;
511        // `cmd_op` in the first 4 bytes of addr3 (kernel reads it as u32).
512        sqe.addr = pointer_addr_u64(nvme_ptr, "NVMe command pointer")?;
513        sqe.len = 64;
514        sqe.user_data = user_data;
515        // The kernel reads the remaining payload bytes out of addr3
516        // directly; downstream NVMe drivers dereference this pointer.
517        sqe.addr3 = pointer_addr_u64(nvme_ptr, "NVMe command addr3 pointer")?;
518
519        self.ring_state.commit_sqe();
520        increment_queue_counter(&mut self.inflight, "inflight SQE count")?;
521        increment_queue_counter(&mut self.pending_submissions, "pending submission count")?;
522
523        Ok(())
524    }
525
526    /// Submit an `IORING_OP_READ_FIXED` into a pre-registered buffer.
527    ///
528    /// Requires the caller to have previously called
529    /// [`super::ring::IoUringState::register_buffers`] with an iovec
530    /// slice whose entry `buf_index` covers the target range. Because
531    /// the kernel skips per-SQE iovec validation, this path is 20-40%
532    /// lower latency than `submit_read_to_gpu` on hot loops.
533    ///
534    /// # Errors
535    ///
536    /// - [`PipelineError::QueueFull`] if the SQ is full or the
537    ///   destination range exceeds the GPU buffer bounds.
538    ///
539    /// # Safety
540    ///
541    /// The `buf_index` must reference a still-registered iovec whose
542    /// region overlaps `chunk_idx * len .. (chunk_idx + 1) * len`
543    /// inside the [`GpuMappedBuffer`]. Mis-indexing produces a kernel
544    /// DMA into the wrong region  -  silent data corruption.
545    pub unsafe fn submit_read_fixed(
546        &mut self,
547        fd: i32,
548        offset: u64,
549        len: u32,
550        chunk_idx: usize,
551        buf_index: u16,
552    ) -> Result<(), PipelineError> {
553        let target_offset = checked_chunk_target_offset(chunk_idx, len)?;
554        // SAFETY: Safe FFI / low-level operation verified and audited for Release compliance.
555        unsafe {
556            self.submit_read_fixed_at(
557                fd,
558                offset,
559                len,
560                target_offset,
561                buf_index,
562                usize_to_u64(chunk_idx, "chunk index")?,
563            )
564        }
565    }
566
567    /// Submit an `IORING_OP_READ_FIXED` into a registered buffer at an
568    /// explicit destination offset inside the mapped buffer.
569    ///
570    /// Unlike [`AsyncUringStream::submit_read_fixed`], this variant decouples
571    /// the CQE `user_data` from the destination layout so higher-level
572    /// drivers can publish their own slot ids while still using a fixed slot
573    /// stride.
574    ///
575    /// # Errors
576    ///
577    /// - [`PipelineError::QueueFull`] if the SQ is full or the
578    ///   destination range exceeds the GPU buffer bounds.
579    ///
580    /// # Safety
581    ///
582    /// `buf_index` must reference a still-registered iovec covering the
583    /// target range, and `user_data` must remain meaningful to the caller
584    /// until the CQE is reaped.
585    pub unsafe fn submit_read_fixed_at(
586        &mut self,
587        fd: i32,
588        offset: u64,
589        len: u32,
590        target_offset: u64,
591        buf_index: u16,
592        user_data: u64,
593    ) -> Result<(), PipelineError> {
594        let end = checked_target_end(target_offset, len)?;
595        let gpu_len = usize_to_u64(self.gpu_buffer.len(), "mapped GPU buffer length")?;
596        if end > gpu_len {
597            return Err(PipelineError::QueueFull {
598                queue: "submission",
599                fix: "chunk_idx * len exceeds GpuMappedBuffer length",
600            });
601        }
602
603        let Some(sqe) = self.ring_state.get_sqe() else {
604            return Err(PipelineError::QueueFull {
605                queue: "submission",
606                fix: "SQ full; call AsyncUringStream::poll to drain completions then retry",
607            });
608        };
609
610        // SAFETY: bounds-checked target address inside the host-visible
611        // GpuMappedBuffer the caller committed at construction.
612        let target_addr = unsafe {
613            self.gpu_buffer
614                .as_ptr()
615                .add(u64_to_usize(target_offset, "target offset")?)
616        };
617
618        sqe.opcode = IORING_OP_READ_FIXED;
619        sqe.fd = fd;
620        sqe.user_data_or_off = offset;
621        sqe.addr = pointer_addr_u64(target_addr, "fixed-read target pointer")?;
622        sqe.len = len;
623        sqe.buf_index = buf_index;
624        sqe.user_data = user_data;
625
626        self.ring_state.commit_sqe();
627        increment_queue_counter(&mut self.inflight, "inflight SQE count")?;
628        increment_queue_counter(&mut self.pending_submissions, "pending submission count")?;
629
630        Ok(())
631    }
632
633    /// Submit a read using a registered-file-table index instead of a
634    /// live fd. Use with
635    /// [`super::ring::IoUringState::register_files`]  -  avoids the
636    /// per-SQE file refcount bump.
637    ///
638    /// # Errors
639    ///
640    /// Same surface as [`AsyncUringStream::submit_read_to_gpu`].
641    ///
642    /// # Safety
643    ///
644    /// `file_index` must name a still-registered fd.
645    /// `iovs_storage` must outlive the completion. All other
646    /// conditions match `submit_read_to_gpu`.
647    pub unsafe fn submit_read_to_gpu_fixed_file(
648        &mut self,
649        file_index: i32,
650        offset: u64,
651        len: u32,
652        chunk_idx: usize,
653        iovs_storage: &mut [Iovec],
654    ) -> Result<(), PipelineError> {
655        if iovs_storage.is_empty() {
656            return Err(PipelineError::QueueFull {
657                queue: "submission",
658                fix: "caller supplied empty iovs_storage; pass at least one slot",
659            });
660        }
661        let target_offset = checked_chunk_target_offset(chunk_idx, len)?;
662        let end = checked_target_end(target_offset, len)?;
663        let gpu_len = usize_to_u64(self.gpu_buffer.len(), "mapped GPU buffer length")?;
664        if end > gpu_len {
665            return Err(PipelineError::QueueFull {
666                queue: "submission",
667                fix: "chunk_idx * len exceeds GpuMappedBuffer length",
668            });
669        }
670
671        let Some(sqe) = self.ring_state.get_sqe() else {
672            return Err(PipelineError::QueueFull {
673                queue: "submission",
674                fix: "SQ full; call AsyncUringStream::poll to drain completions then retry",
675            });
676        };
677
678        // SAFETY: same invariants as submit_read_to_gpu, plus the
679        // caller committed that file_index is a registered fd.
680        let target_addr = unsafe {
681            self.gpu_buffer
682                .as_ptr()
683                .add(u64_to_usize(target_offset, "target offset")?)
684        };
685        iovs_storage[0] = Iovec {
686            iov_base: target_addr.cast::<core::ffi::c_void>(),
687            iov_len: u32_to_usize(len, "read length")?,
688        };
689
690        sqe.opcode = IORING_OP_READV;
691        sqe.flags = super::ring::IOSQE_FIXED_FILE;
692        sqe.fd = file_index;
693        sqe.user_data_or_off = offset;
694        sqe.addr = pointer_addr_u64(iovs_storage.as_ptr(), "fixed-file readv iovec pointer")?;
695        sqe.len = 1;
696        sqe.user_data = usize_to_u64(chunk_idx, "chunk index")?;
697
698        self.ring_state.commit_sqe();
699        increment_queue_counter(&mut self.inflight, "inflight SQE count")?;
700        increment_queue_counter(&mut self.pending_submissions, "pending submission count")?;
701
702        Ok(())
703    }
704
705    /// Disabled-feature implementation for NVMe passthrough. Always returns
706    /// [`PipelineError::NvmePassthroughDisabled`] so callers get a
707    /// structured error rather than a link failure.
708    #[cfg(not(feature = "uring-cmd-nvme"))]
709    #[allow(clippy::unused_self, clippy::missing_safety_doc)]
710    pub unsafe fn submit_nvme_passthrough(
711        &mut self,
712        _fd: i32,
713        _user_data: u64,
714        _nvme_sqe_bytes: &[u8],
715    ) -> Result<(), PipelineError> {
716        Err(PipelineError::NvmePassthroughDisabled)
717    }
718}
719
720fn checked_chunk_target_offset(chunk_idx: usize, len: u32) -> Result<u64, PipelineError> {
721    let chunk_idx = usize_to_u64(chunk_idx, "chunk index")?;
722    vyre_driver::accounting::checked_mul_u64_lazy(chunk_idx, u64::from(len), || {
723        PipelineError::QueueFull {
724            queue: "submission",
725            fix: "chunk_idx * len overflows u64; split the IO batch before submission",
726        }
727    })
728}
729
730fn checked_target_end(target_offset: u64, len: u32) -> Result<u64, PipelineError> {
731    vyre_driver::accounting::checked_add_u64_lazy(target_offset, u64::from(len), || {
732        PipelineError::QueueFull {
733            queue: "submission",
734            fix: "target_offset + len overflows u64; split the IO batch before submission",
735        }
736    })
737}
738
739fn increment_queue_counter(counter: &mut u32, label: &'static str) -> Result<(), PipelineError> {
740    *counter = vyre_driver::accounting::checked_add_u32_value(
741        *counter,
742        1,
743        PipelineError::QueueFull {
744            queue: "submission",
745            fix: match label {
746                "inflight SQE count" => {
747                    "inflight SQE count overflowed u32; poll completions before submitting more work"
748                }
749                "pending submission count" => {
750                    "pending submission count overflowed u32; flush submissions before queuing more work"
751                }
752                _ => {
753                    "io_uring queue counter overflowed u32; drain the queue before submitting more work"
754                }
755            },
756        },
757    )?;
758    Ok(())
759}
760
761fn decrement_queue_counter(counter: &mut u32, label: &'static str) -> Result<(), PipelineError> {
762    *counter = counter.checked_sub(1).ok_or(PipelineError::QueueFull {
763        queue: "completion",
764        fix: match label {
765            "inflight SQE count" => {
766                "io_uring completion arrived with no inflight SQE; rebuild the stream state"
767            }
768            _ => "io_uring queue counter underflowed; rebuild the stream state",
769        },
770    })?;
771    Ok(())
772}
773
774fn usize_to_u64(value: usize, label: &'static str) -> Result<u64, PipelineError> {
775    u64::try_from(value).map_err(|_| PipelineError::QueueFull {
776        queue: "submission",
777        fix: match label {
778            "chunk index" => "chunk index cannot fit u64; split the IO batch before submission",
779            "mapped GPU buffer length" => {
780                "mapped GPU buffer length cannot fit u64; split the staging allocation"
781            }
782            _ => "host usize value cannot fit u64; split the IO batch before submission",
783        },
784    })
785}
786
787fn pointer_addr_u64<T>(ptr: *const T, label: &'static str) -> Result<u64, PipelineError> {
788    usize_to_u64(ptr.addr(), label)
789}
790
791fn u64_to_usize(value: u64, label: &'static str) -> Result<usize, PipelineError> {
792    usize::try_from(value).map_err(|_| PipelineError::QueueFull {
793        queue: "submission",
794        fix: match label {
795            "target offset" => {
796                "target offset cannot fit usize; split the IO batch before submission"
797            }
798            _ => "u64 value cannot fit usize; split the IO batch before submission",
799        },
800    })
801}
802
803fn u32_to_usize(value: u32, label: &'static str) -> Result<usize, PipelineError> {
804    usize::try_from(value).map_err(|_| PipelineError::QueueFull {
805        queue: "submission",
806        fix: match label {
807            "read length" => "read length cannot fit usize; split the IO request before submission",
808            _ => "u32 value cannot fit usize; split the IO request before submission",
809        },
810    })
811}
812
813#[cfg(test)]
814mod tests {
815    use super::*;
816
817    #[test]
818    fn mapped_slice_roundtrip_is_miri_clean() {
819        let mut backing = [1_u8, 2, 3, 4];
820        // SAFETY: `backing` stays live and uniquely borrowed for the mapped buffer lifetime.
821        let mut mapped = unsafe { GpuMappedBuffer::from_host_visible_slice(&mut backing) };
822        // SAFETY: the mapped buffer was built from `backing` and remains uniquely borrowed.
823        let slice = unsafe { mapped.as_mut_slice() };
824        slice[0] = 9;
825        slice[3] = 7;
826        assert_eq!(backing, [9, 2, 3, 7]);
827    }
828}
vyre_runtime/uring/stream.rs

vyre_runtime/uring/
stream.rs
