vyre_driver_wgpu/runtime/
readback_ring.rs

1//! Async readback ring (Innovation I.5).
2//!
3//! Blocking readback submits a copy + device.poll(Wait) that stalls
4//! the submit queue. Under high dispatch rate this ruins latency and
5//! throughput  -  the GPU goes idle while the CPU waits.
6//!
7//! The readback ring threads N staging buffers. Dispatch \`i\` writes
8//! to \`ring[i % N]\`; the copy submits immediately and readback
9//! happens asynchronously via \`map_async\`. Dispatch \`i+1\` runs in
10//! parallel with readback \`i\`'s copy.
11
12use crossbeam_channel::Receiver;
13use dashmap::mapref::entry::Entry;
14use dashmap::DashMap;
15use rustc_hash::FxHasher;
16use std::hash::BuildHasherDefault;
17use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
18use std::sync::Arc;
19use vyre_driver::accounting::{atomic_max_u64, rebasing_atomic_next_u64};
20use vyre_driver::backend::BackendError;
21
22use crate::staging_reserve::reserve_backend_vec;
23
24const MIN_RING_SIZE: usize = 2;
25const MAX_RING_SIZE: usize = 256;
26const DEFAULT_RING_SLOTS: usize = 256;
27const RING_CAPACITY_GRANULARITY: u64 = 4096;
28const SLOT_FREE: u8 = 0;
29const SLOT_PENDING: u8 = 1;
30const SLOT_READY: u8 = 2;
31const SLOT_ERROR: u8 = 3;
32
33/// Result type produced by one `map_async` callback.
34pub type MapResult = Result<(), wgpu::BufferAsyncError>;
35
36/// Statistics collected by the ring at runtime.
37#[derive(Debug, Default)]
38pub struct RingStats {
39    /// Total dispatches queued.
40    pub dispatches: AtomicU64,
41    /// Readbacks that blocked waiting on map_async.
42    pub readback_stalls: AtomicU64,
43    /// Max outstanding (in-flight) copies.
44    pub peak_inflight: AtomicU64,
45}
46
47impl RingStats {
48    /// Record one dispatch; returns the monotonic dispatch index.
49    pub fn record_dispatch(&self) -> u64 {
50        rebasing_atomic_next_u64(
51            &self.dispatches,
52            0,
53            Ordering::Relaxed,
54            Ordering::Relaxed,
55            Ordering::Relaxed,
56            |_, _| {
57                tracing::error!(
58                    "readback ring dispatch counter reached u64::MAX and was rebased to zero. Fix: shard readback rings or scrape counters before wrap."
59                );
60            },
61        )
62    }
63
64    /// Record a stall.
65    pub fn record_stall(&self) {
66        rebasing_atomic_next_u64(
67            &self.readback_stalls,
68            0,
69            Ordering::Relaxed,
70            Ordering::Relaxed,
71            Ordering::Relaxed,
72            |_, _| {
73                tracing::error!(
74                    "readback ring stall counter reached u64::MAX and was rebased to zero. Fix: shard readback rings or scrape counters before wrap."
75                );
76            },
77        );
78    }
79
80    /// Update the peak-in-flight watermark.
81    pub fn update_peak(&self, current: u64) {
82        atomic_max_u64(&self.peak_inflight, current, Ordering::AcqRel);
83    }
84}
85
86/// Lifecycle state for one ring slot.
87#[derive(Debug, Clone, Copy, PartialEq, Eq)]
88pub enum SlotState {
89    /// Slot is available for new writes.
90    Free,
91    /// Copy has been submitted, data will be ready after fence.
92    Pending,
93    /// Map has completed and data is visible to the host.
94    Ready,
95    /// Mapping failed and the slot must be collected as an error.
96    Error,
97}
98
99/// GPU-aware ring slot.
100pub struct GpuSlot {
101    /// Underlying wgpu buffer.
102    pub buffer: wgpu::Buffer,
103    /// Atomic lifecycle state (0: Free, 1: Pending, 2: Ready).
104    pub state: Arc<std::sync::atomic::AtomicU8>,
105    byte_len: AtomicU64,
106    mapped_len: AtomicU64,
107    capacity: u64,
108}
109
110/// Submitted copy ticket for one readback-ring slot.
111pub struct ReadbackTicket {
112    idx: usize,
113    byte_len: u64,
114    mapped_len: u64,
115}
116
117/// Size-classed collection of readback rings for direct dispatch.
118pub struct ReadbackRingSet {
119    rings: DashMap<u64, Arc<ReadbackRing>, BuildHasherDefault<FxHasher>>,
120    slots_per_ring: usize,
121}
122
123impl Default for ReadbackRingSet {
124    fn default() -> Self {
125        Self::new()
126    }
127}
128
129impl ReadbackRingSet {
130    /// Construct an empty ring set using the default slot count.
131    #[must_use]
132    pub fn new() -> Self {
133        Self {
134            rings: DashMap::with_hasher(BuildHasherDefault::<FxHasher>::default()),
135            slots_per_ring: readback_ring_slots_from_env(),
136        }
137    }
138
139    /// Construct an empty ring set from a raw slot-count setting.
140    ///
141    /// Passing `None` uses the production default. This keeps test and embedded
142    /// callers off process-global environment mutation while preserving the same
143    /// parser and clamping semantics as [`Self::new`].
144    #[must_use]
145    pub fn with_requested_slots(raw_slots: Option<&str>) -> Self {
146        Self {
147            rings: DashMap::with_hasher(BuildHasherDefault::<FxHasher>::default()),
148            slots_per_ring: readback_ring_slots_from_raw(raw_slots),
149        }
150    }
151
152    /// Return the ring whose staging slots can hold `byte_len`.
153    ///
154    /// # Errors
155    ///
156    /// Returns a backend error if the requested byte length overflows wgpu copy
157    /// alignment.
158    pub fn ring_for(
159        &self,
160        device: &wgpu::Device,
161        byte_len: u64,
162    ) -> Result<Arc<ReadbackRing>, BackendError> {
163        let capacity = Self::capacity_class_for(byte_len)?;
164        self.ring_for_capacity(device, capacity)
165    }
166
167    /// Return a ring for an already-normalized capacity class.
168    #[inline]
169    pub(crate) fn ring_for_capacity(
170        &self,
171        device: &wgpu::Device,
172        capacity: u64,
173    ) -> Result<Arc<ReadbackRing>, BackendError> {
174        Ok(match self.rings.entry(capacity) {
175            Entry::Occupied(entry) => Arc::clone(entry.get()),
176            Entry::Vacant(entry) => {
177                let ring = Arc::new(ReadbackRing::new(device, self.slots_per_ring, capacity)?);
178                entry.insert(Arc::clone(&ring));
179                ring
180            }
181        })
182    }
183
184    /// Convert an arbitrary byte length to the ring capacity class used for
185    /// ring sizing.
186    #[inline]
187    pub(crate) fn capacity_class(byte_len: u64) -> Result<u64, BackendError> {
188        Self::capacity_class_for(byte_len)
189    }
190
191    /// Convert an arbitrary byte length to the ring capacity class used for
192    /// ring sizing.
193    #[inline]
194    pub(crate) fn capacity_class_for(byte_len: u64) -> Result<u64, BackendError> {
195        ring_capacity_class(byte_len)
196    }
197
198    /// Return an existing size-classed ring without taking exclusive access.
199    ///
200    /// # Errors
201    ///
202    /// Returns a backend error if the requested byte length overflows wgpu copy
203    /// alignment.
204    pub fn existing_ring_for(
205        &self,
206        byte_len: u64,
207    ) -> Result<Option<Arc<ReadbackRing>>, BackendError> {
208        let capacity = Self::capacity_class(byte_len)?;
209        Ok(self.existing_ring_for_capacity(capacity))
210    }
211
212    /// Return an existing size-classed ring without taking exclusive access.
213    #[inline]
214    pub(crate) fn existing_ring_for_capacity(&self, capacity: u64) -> Option<Arc<ReadbackRing>> {
215        self.rings
216            .get(&capacity)
217            .map(|ring| Arc::clone(ring.value()))
218    }
219
220    /// Number of slots configured for each runtime ring instance.
221    #[must_use]
222    pub fn slots_per_ring(&self) -> usize {
223        self.slots_per_ring
224    }
225}
226
227/// Async readback ring buffer with GPU-resident staging buffers.
228pub struct ReadbackRing {
229    slots: Vec<GpuSlot>,
230    stats: Arc<RingStats>,
231    next_idx: AtomicU64,
232}
233
234impl ReadbackRing {
235    /// Construct a ring with N staging buffers.
236    #[must_use]
237    pub fn new(device: &wgpu::Device, size: usize, buffer_size: u64) -> Result<Self, BackendError> {
238        let size = size.clamp(MIN_RING_SIZE, MAX_RING_SIZE);
239        let capacity = staging_capacity(buffer_size)?;
240        let mut slots = Vec::new();
241        reserve_backend_vec(&mut slots, size, "readback ring slot table")?;
242        for i in 0..size {
243            let buffer = device.create_buffer(&wgpu::BufferDescriptor {
244                label: Some(&format!("vyre readback ring slot {i}")),
245                size: capacity,
246                usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
247                mapped_at_creation: false,
248            });
249            slots.push(GpuSlot {
250                buffer,
251                state: Arc::new(std::sync::atomic::AtomicU8::new(SLOT_FREE)),
252                byte_len: AtomicU64::new(0),
253                mapped_len: AtomicU64::new(0),
254                capacity,
255            });
256        }
257        Ok(Self {
258            slots,
259            stats: Arc::new(RingStats::default()),
260            next_idx: AtomicU64::new(0),
261        })
262    }
263
264    /// Ensure slot `idx` is reusable for a fresh readback: either already
265    /// `SLOT_FREE`, or `SLOT_PENDING` that completes to `SLOT_FREE` after one
266    /// device poll. Any other terminal state is a caller contract violation and
267    /// is reported as a distinct, fail-closed error.
268    ///
269    /// VYRE-WGPU-002: earlier code conflated `SLOT_READY` / `SLOT_ERROR` with a
270    /// single misleading wrap-overflow message. We now name each
271    /// state (but we DO NOT silently recycle an uncollected `SLOT_READY` slot).
272    /// Recycling would unmap and discard a completed-but-uncollected readback
273    /// a silent recall loss (Law 10). The caller MUST collect every readback
274    /// before the ring wraps back to its slot; if it has not, we fail closed so
275    /// the data loss is impossible to miss.
276    fn ensure_slot_reusable(
277        &self,
278        idx: usize,
279        slot: &GpuSlot,
280        device: &wgpu::Device,
281    ) -> Result<(), BackendError> {
282        let mut state = slot.state.load(Ordering::Acquire);
283        if state == SLOT_PENDING {
284            self.stats.record_stall();
285            crate::runtime::device::poll_device_once(device)?;
286            state = slot.state.load(Ordering::Acquire);
287        }
288        match state {
289            SLOT_FREE => Ok(()),
290            SLOT_READY => Err(BackendError::new(format!(
291                "readback ring slot {idx} holds an uncollected completed readback (SLOT_READY). Fix: collect every ReadbackTicket via collect_slot_into before the ring wraps back to this slot (recycling it would silently drop the prior result (a recall loss))."
292            ))),
293            SLOT_ERROR => Err(BackendError::new(format!(
294                "readback ring slot {idx} is in SLOT_ERROR (prior map_async failed) and was not collected before reuse. Fix: collect error slots via collect_slot_into before submitting new readbacks to the same slot."
295            ))),
296            SLOT_PENDING => Err(BackendError::new(format!(
297                "readback ring slot {idx} is still SLOT_PENDING after a device poll, the prior readback has not completed. Fix: increase ring depth (more slots) or collect outstanding readbacks before submitting more."
298            ))),
299            other => Err(BackendError::new(format!(
300                "readback ring slot {idx} has unexpected state {other}. Fix: do not modify readback ring slot state outside the ring API."
301            ))),
302        }
303    }
304
305    /// Record a readback copy into the next available ring slot.
306    ///
307    /// The caller must submit the encoder and then arm the returned ticket with
308    /// [`Self::arm_ticket`]. This path lets the main dispatch encoder copy into
309    /// preallocated ring slots instead of allocating a fresh staging buffer per
310    /// output.
311    ///
312    /// # Errors
313    ///
314    /// Returns [`BackendError`] if the byte range cannot be represented, the
315    /// ring slot is not reusable (an uncollected `SLOT_READY`/`SLOT_ERROR` slot
316    /// or a still-pending slot after a device poll), or the requested readback
317    /// exceeds slot capacity.
318    pub fn record_copy(
319        &self,
320        device: &wgpu::Device,
321        encoder: &mut wgpu::CommandEncoder,
322        src_buffer: &wgpu::Buffer,
323        src_offset: u64,
324        byte_len: u64,
325    ) -> Result<ReadbackTicket, BackendError> {
326        let idx = self.next_slot_index()?;
327        let slot = &self.slots[idx];
328        let mapped_len = aligned_copy_len(byte_len)?;
329        if mapped_len > slot.capacity {
330            return Err(BackendError::new(format!(
331                "readback request of {byte_len} bytes ({} bytes after wgpu copy alignment) exceeds ring slot capacity {} bytes. Fix: construct ReadbackRing with a buffer_size at least as large as the largest readback.",
332                mapped_len, slot.capacity
333            )));
334        }
335
336        self.ensure_slot_reusable(idx, slot, device)?;
337
338        slot.byte_len.store(byte_len, Ordering::Release);
339        slot.mapped_len.store(mapped_len, Ordering::Release);
340        slot.state.store(SLOT_PENDING, Ordering::Release);
341        if mapped_len != 0 {
342            encoder.copy_buffer_to_buffer(src_buffer, src_offset, &slot.buffer, 0, mapped_len);
343        } else {
344            slot.state.store(SLOT_READY, Ordering::Release);
345        }
346        self.stats.record_dispatch();
347        Ok(ReadbackTicket {
348            idx,
349            byte_len,
350            mapped_len,
351        })
352    }
353
354    /// Arm a submitted ticket by registering its `map_async` callback.
355    ///
356    /// # Errors
357    ///
358    /// Returns [`BackendError`] when `ticket` does not reference a live slot.
359    pub fn arm_ticket(
360        &self,
361        ticket: &ReadbackTicket,
362    ) -> Result<(Receiver<MapResult>, Arc<AtomicBool>), BackendError> {
363        let Some(slot) = self.slots.get(ticket.idx) else {
364            return Err(BackendError::new(format!(
365                "readback ring ticket slot {} is out of bounds for {} slots. Fix: keep tickets paired with their originating ring.",
366                ticket.idx,
367                self.slots.len()
368            )));
369        };
370        let (sender, receiver) = crossbeam_channel::bounded(1);
371        let ready = Arc::new(AtomicBool::new(false));
372        if ticket.mapped_len == 0 {
373            if let Err(error) = sender.send(Ok(())) {
374                tracing::error!(
375                    ?error,
376                    "readback ring zero-length callback result was lost because the receiver dropped"
377                );
378            }
379            ready.store(true, Ordering::Release);
380            return Ok((receiver, ready));
381        }
382
383        let state = Arc::clone(&slot.state);
384        let ready_cb = Arc::clone(&ready);
385        slot.buffer
386            .slice(0..ticket.mapped_len)
387            .map_async(wgpu::MapMode::Read, move |result| {
388                match &result {
389                    Ok(()) => state.store(SLOT_READY, Ordering::Release),
390                    Err(error) => {
391                        tracing::error!(
392                            "readback ring map_async failed: {error:?}. Fix: inspect device health and readback buffer usage."
393                        );
394                        state.store(SLOT_ERROR, Ordering::Release);
395                    }
396                }
397                if let Err(error) = sender.send(result) {
398                    tracing::error!(
399                        ?error,
400                        "readback ring callback result was lost because the receiver dropped"
401                    );
402                }
403                ready_cb.store(true, Ordering::Release);
404            });
405        Ok((receiver, ready))
406    }
407
408    /// Expose a ready ticket's mapped bytes to `visitor`, then free the slot.
409    ///
410    /// # Errors
411    ///
412    /// Returns [`BackendError`] when the ticket is stale, the slot is not ready,
413    /// or mapped length metadata is inconsistent.
414    pub fn with_mapped_ticket<R>(
415        &self,
416        ticket: &ReadbackTicket,
417        visitor: impl FnOnce(&[u8]) -> Result<R, BackendError>,
418    ) -> Result<R, BackendError> {
419        let Some(slot) = self.slots.get(ticket.idx) else {
420            return Err(BackendError::new(format!(
421                "readback ring ticket slot {} is out of bounds for {} slots. Fix: keep tickets paired with their originating ring.",
422                ticket.idx,
423                self.slots.len()
424            )));
425        };
426        match slot.state.load(Ordering::Acquire) {
427            SLOT_READY => {}
428            SLOT_ERROR => {
429                slot.byte_len.store(0, Ordering::Release);
430                slot.mapped_len.store(0, Ordering::Release);
431                slot.state.store(SLOT_FREE, Ordering::Release);
432                return Err(BackendError::new(
433                    "readback ring map_async failed. Fix: inspect GPU device health and ensure the slot buffer has MAP_READ usage.",
434                ));
435            }
436            _ => {
437                return Err(BackendError::new(
438                    "readback ring ticket was collected before its map callback completed. Fix: poll the device or wait for the submitted GPU work before collection.",
439                ));
440            }
441        }
442
443        let len = usize::try_from(ticket.byte_len).map_err(|source| {
444            BackendError::new(format!(
445                "readback ring byte length {} cannot fit usize: {source}. Fix: split the readback before collecting it.",
446                ticket.byte_len
447            ))
448        })?;
449        if ticket.mapped_len == 0 {
450            slot.byte_len.store(0, Ordering::Release);
451            slot.mapped_len.store(0, Ordering::Release);
452            slot.state.store(SLOT_FREE, Ordering::Release);
453            return visitor(&[]);
454        }
455        let view = slot.buffer.slice(0..ticket.mapped_len).get_mapped_range();
456        if len > view.len() {
457            let mapped_len = view.len();
458            drop(view);
459            slot.buffer.unmap();
460            slot.byte_len.store(0, Ordering::Release);
461            slot.mapped_len.store(0, Ordering::Release);
462            slot.state.store(SLOT_FREE, Ordering::Release);
463            return Err(BackendError::new(format!(
464                "readback ring mapped length {mapped_len} is shorter than requested length {len}. Fix: keep ticket and slot byte lengths synchronized."
465            )));
466        }
467        let result = visitor(&view[..len]);
468        drop(view);
469        slot.buffer.unmap();
470        slot.byte_len.store(0, Ordering::Release);
471        slot.mapped_len.store(0, Ordering::Release);
472        slot.state.store(SLOT_FREE, Ordering::Release);
473        result
474    }
475
476    /// Submit a copy from `src_buffer` at `src_offset` and mark the slot pending.
477    ///
478    /// `src_offset` is the byte offset within `src_buffer` to copy from. Pass
479    /// `0` to read from the start of the buffer. This mirrors the `src_offset`
480    /// parameter accepted by `record_copy`; callers that need a sub-range of
481    /// the source buffer must supply a non-zero offset here rather than
482    /// wrapping a slice (the wgpu copy API requires aligned buffer offsets).
483    ///
484    /// # Errors
485    /// Returns `BackendError` if encoder or queue submission fails.
486    pub fn submit_readback(
487        &self,
488        device: &wgpu::Device,
489        queue: &wgpu::Queue,
490        src_buffer: &wgpu::Buffer,
491        src_offset: u64,
492        byte_len: u64,
493    ) -> Result<usize, BackendError> {
494        let idx = self.next_slot_index()?;
495        let slot = &self.slots[idx];
496        let mapped_len = aligned_copy_len(byte_len)?;
497        if mapped_len > slot.capacity {
498            return Err(BackendError::new(format!(
499                "readback request of {byte_len} bytes ({} bytes after wgpu copy alignment) exceeds ring slot capacity {} bytes. Fix: construct ReadbackRing with a buffer_size at least as large as the largest readback.",
500                mapped_len, slot.capacity
501            )));
502        }
503
504        self.ensure_slot_reusable(idx, slot, device)?;
505
506        let state_clone = Arc::clone(&slot.state);
507        slot.byte_len.store(byte_len, Ordering::Release);
508        slot.mapped_len.store(mapped_len, Ordering::Release);
509        state_clone.store(SLOT_PENDING, Ordering::Release);
510
511        if mapped_len == 0 {
512            state_clone.store(SLOT_READY, Ordering::Release);
513            self.stats.record_dispatch();
514            return Ok(idx);
515        }
516
517        let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
518            label: Some("vyre readback ring copy"),
519        });
520        encoder.copy_buffer_to_buffer(src_buffer, src_offset, &slot.buffer, 0, mapped_len);
521        queue.submit(std::iter::once(encoder.finish()));
522
523        slot.buffer
524            .slice(0..mapped_len)
525            .map_async(wgpu::MapMode::Read, move |result| {
526                match result {
527                    Ok(()) => state_clone.store(SLOT_READY, Ordering::Release),
528                    Err(error) => {
529                        tracing::error!(
530                            "readback ring map_async failed: {error:?}. Fix: inspect device health and readback buffer usage."
531                        );
532                        state_clone.store(SLOT_ERROR, Ordering::Release);
533                    }
534                }
535            });
536
537        self.stats.record_dispatch();
538
539        Ok(idx)
540    }
541
542    /// Try to collect data from a specific slot.
543    ///
544    /// # Errors
545    ///
546    /// Returns [`BackendError`] when `idx` is out of bounds or `map_async`
547    /// failed for the slot.
548    pub fn collect_slot(
549        &self,
550        device: &wgpu::Device,
551        idx: usize,
552    ) -> Result<Option<Vec<u8>>, BackendError> {
553        let mut data = Vec::new();
554        if self.collect_slot_into(device, idx, &mut data)?.is_some() {
555            Ok(Some(data))
556        } else {
557            Ok(None)
558        }
559    }
560
561    /// Try to collect data from a specific slot into a caller-owned buffer.
562    ///
563    /// Reusing `out` avoids an allocation on every ready readback. The buffer is
564    /// cleared before bytes are appended.
565    ///
566    /// # Errors
567    ///
568    /// Returns [`BackendError`] when `idx` is out of bounds or `map_async`
569    /// failed for the slot.
570    pub fn collect_slot_into(
571        &self,
572        device: &wgpu::Device,
573        idx: usize,
574        out: &mut Vec<u8>,
575    ) -> Result<Option<usize>, BackendError> {
576        let Some(slot) = self.slots.get(idx) else {
577            return Err(BackendError::new(format!(
578                "readback ring slot index {idx} is out of bounds for {} slots. Fix: collect only indices returned by submit_readback.",
579                self.slots.len()
580            )));
581        };
582        match slot.state.load(Ordering::Acquire) {
583            SLOT_READY => {
584                let len = self.copy_ready_slot_into(idx, out)?;
585                Ok(Some(len))
586            }
587            SLOT_ERROR => {
588                slot.byte_len.store(0, Ordering::Release);
589                slot.mapped_len.store(0, Ordering::Release);
590                slot.state.store(SLOT_FREE, Ordering::Release);
591                Err(BackendError::new(
592                    "readback ring map_async failed. Fix: inspect GPU device health and ensure the slot buffer has MAP_READ usage.",
593                ))
594            }
595            _ => {
596                crate::runtime::device::poll_device_once(device)?;
597                Ok(None)
598            }
599        }
600    }
601
602    fn copy_ready_slot_into(&self, idx: usize, out: &mut Vec<u8>) -> Result<usize, BackendError> {
603        let slot = &self.slots[idx];
604        let byte_len = slot.byte_len.load(Ordering::Acquire);
605        let mapped_len = slot.mapped_len.load(Ordering::Acquire);
606        let len = usize::try_from(byte_len).map_err(|source| {
607            BackendError::new(format!(
608                "readback ring byte length {byte_len} cannot fit usize: {source}. Fix: split the readback before collecting it."
609            ))
610        })?;
611        if mapped_len != 0 {
612            let view = slot.buffer.slice(0..mapped_len).get_mapped_range();
613            let bytes = &view[..len];
614            if out.len() == len {
615                out.copy_from_slice(bytes);
616            } else {
617                if len > out.capacity() {
618                    let additional = len - out.capacity();
619                    out.try_reserve_exact(additional).map_err(|source| {
620                        BackendError::new(format!(
621                            "readback ring collection could not reserve {len} output bytes exactly: {source}. Fix: lower max_output_bytes or collect readback in smaller shards."
622                        ))
623                    })?;
624                }
625                out.clear();
626                out.extend_from_slice(bytes);
627            }
628            drop(view);
629            slot.buffer.unmap();
630        } else {
631            out.clear();
632        }
633        slot.byte_len.store(0, Ordering::Release);
634        slot.mapped_len.store(0, Ordering::Release);
635        slot.state.store(SLOT_FREE, Ordering::Release);
636        Ok(len)
637    }
638
639    #[inline]
640    fn next_slot_index(&self) -> Result<usize, BackendError> {
641        let slot_len = u64::try_from(self.slots.len()).map_err(|source| {
642            BackendError::new(format!(
643                "readback ring slot count {} cannot fit u64: {source}. Fix: reduce readback ring slot count.",
644                self.slots.len()
645            ))
646        })?;
647        if slot_len == 0 {
648            return Err(BackendError::new(
649                "readback ring has zero slots. Fix: construct rings with at least two slots.",
650            ));
651        }
652        let next = rebasing_atomic_next_u64(
653            &self.next_idx,
654            0,
655            Ordering::Relaxed,
656            Ordering::Relaxed,
657            Ordering::Relaxed,
658            |_, _| {
659                tracing::error!(
660                    "readback ring slot counter reached u64::MAX and was rebased to zero. Fix: shard readback rings or scrape counters before wrap."
661                );
662            },
663        );
664        usize::try_from(next % slot_len).map_err(|source| {
665            BackendError::new(format!(
666                "readback ring slot index cannot fit usize: {source}. Fix: reduce readback ring slot count."
667            ))
668        })
669    }
670}
671
672#[inline]
673
674fn staging_capacity(byte_len: u64) -> Result<u64, BackendError> {
675    aligned_copy_len(byte_len).map_err(|error| {
676        tracing::warn!(
677            "readback ring staging capacity overflowed for {byte_len} bytes: {error}. Fix: shard the readback buffer before constructing the ring."
678        );
679        error
680    }).map(|len| len.max(4))
681}
682
683#[inline]
684fn ring_capacity_class(byte_len: u64) -> Result<u64, BackendError> {
685    let aligned = aligned_copy_len(byte_len)?.max(4);
686    aligned
687        .checked_add(RING_CAPACITY_GRANULARITY - 1)
688        .map(|len| len & !(RING_CAPACITY_GRANULARITY - 1))
689        .ok_or_else(|| {
690            BackendError::new(
691                "readback ring capacity class overflows u64. Fix: split the readback before submitting it to the ring.",
692            )
693        })
694}
695
696#[inline]
697fn aligned_copy_len(byte_len: u64) -> Result<u64, BackendError> {
698    crate::numeric::WGPU_NUMERIC.align_up_u64(byte_len, 4, 0, "readback byte length")
699}
700
701fn readback_ring_slots_from_env() -> usize {
702    let raw = std::env::var("VYRE_WGPU_READBACK_RING_SLOTS").ok();
703    readback_ring_slots_from_raw(raw.as_deref())
704}
705
706fn readback_ring_slots_from_raw(raw: Option<&str>) -> usize {
707    let Some(raw) = raw else {
708        return DEFAULT_RING_SLOTS;
709    };
710    let slots = match raw.parse::<usize>() {
711        Ok(0) => {
712            tracing::warn!(
713                "VYRE_WGPU_READBACK_RING_SLOTS=0 is invalid for GPU readback rings; defaulting to {MIN_RING_SIZE}. Fix: set it to a positive integer between {MIN_RING_SIZE} and {MAX_RING_SIZE}, or unset it."
714            );
715            MIN_RING_SIZE
716        }
717        Ok(value) if value > MAX_RING_SIZE => {
718            tracing::warn!(
719                "VYRE_WGPU_READBACK_RING_SLOTS={value} exceeds the safe cap of {MAX_RING_SIZE}; clamping.
720                Fix: set it to an integer between {MIN_RING_SIZE} and {MAX_RING_SIZE}, or unset it."
721            );
722            MAX_RING_SIZE
723        }
724        Ok(value) => value,
725        Err(error) => {
726            tracing::warn!(
727                "VYRE_WGPU_READBACK_RING_SLOTS={raw:?} is invalid ({error:?}); defaulting to {DEFAULT_RING_SLOTS}. Fix: set it to a positive integer between {MIN_RING_SIZE} and {MAX_RING_SIZE}, or unset it."
728            );
729            DEFAULT_RING_SLOTS
730        }
731    };
732    slots.clamp(MIN_RING_SIZE, MAX_RING_SIZE)
733}
734
735#[cfg(test)]
736mod tests {
737    use super::*;
738
739    #[test]
740    fn capacity_class_classifies_by_alignment_and_granularity() {
741        assert_eq!(
742            ReadbackRingSet::capacity_class_for(16).unwrap(),
743            4096,
744            "16-byte requests must promote to 4096-byte slot class"
745        );
746        assert_eq!(
747            ReadbackRingSet::capacity_class_for(1).unwrap(),
748            4096,
749            "1-byte requests must promote to minimum aligned 4096-byte class"
750        );
751        assert_eq!(
752            ReadbackRingSet::capacity_class_for(4097).unwrap(),
753            8192,
754            "4KB boundary crossings must promote to the next class"
755        );
756    }
757
758    #[test]
759    fn existing_ring_for_and_capacity_variant_agree_on_lookup_key() {
760        let ring_set = ReadbackRingSet::new();
761        let from_raw = ring_set
762            .existing_ring_for(16)
763            .expect("Fix: lookup with raw byte length should not fail");
764        let from_class = ring_set.existing_ring_for_capacity(4096);
765        assert!(
766            from_raw.is_none() && from_class.is_none(),
767            "raw and capacity-based lookups should agree on an empty set"
768        );
769    }
770
771    #[test]
772    fn production_ring_construction_uses_fallible_slot_reservation() {
773        let production = include_str!("readback_ring.rs")
774            .split("\n#[cfg(test)]\nmod tests")
775            .next()
776            .expect("Fix: readback ring production section should precede tests");
777
778        assert!(
779            !production.contains("Vec::with_capacity(size)"),
780            "Fix: readback ring construction must not allocate slot tables infallibly."
781        );
782        assert!(
783            production.contains("reserve_backend_vec(&mut slots, size, \"readback ring slot table\")?"),
784            "Fix: readback ring construction should reserve slot tables through the shared WGPU staging helper."
785        );
786    }
787
788    /// VYRE-WGPU-002: the pre-use slot check must distinguish SLOT_READY and
789    /// SLOT_ERROR from a ring overflow with state-specific diagnostics, and it
790    /// must FAIL CLOSED on an uncollected SLOT_READY slot rather than silently
791    /// recycle it. Silently recycling unmaps and discards a completed-but-
792    /// uncollected readback, which is a silent recall loss (Law 10).
793    ///
794    /// Both `record_copy` and `submit_readback` route their pre-use check
795    /// through the single `ensure_slot_reusable` helper, so the contract is
796    /// expressed once. This source-text canary asserts that structural shape
797    /// without a live GPU; the behavioral round-trip lives in the GPU-gated
798    /// `readback_ring_liveness_contracts` integration test.
799    #[test]
800    fn slot_reuse_check_fails_closed_on_uncollected_ready_with_distinct_diagnostics() {
801        let src = include_str!("readback_ring.rs");
802        // Locate production code only (before the first test module).
803        let production = src
804            .split("\n#[cfg(test)]\nmod tests")
805            .next()
806            .expect("Fix: readback_ring.rs should have a test module");
807
808        // Both methods must funnel through the single dedup'd helper.
809        assert!(
810            production.contains("fn ensure_slot_reusable("),
811            "Fix: the slot reuse check must live in one ensure_slot_reusable helper, not be duplicated across record_copy / submit_readback"
812        );
813        assert_eq!(
814            production.matches("self.ensure_slot_reusable(idx, slot, device)?").count(),
815            2,
816            "Fix: both record_copy and submit_readback must call ensure_slot_reusable (one call site each)"
817        );
818
819        // Distinct, named arms for each terminal state.
820        assert!(
821            production.contains("SLOT_READY =>"),
822            "Fix: ensure_slot_reusable must have an explicit SLOT_READY arm"
823        );
824        assert!(
825            production.contains("SLOT_ERROR =>"),
826            "Fix: ensure_slot_reusable must have an explicit SLOT_ERROR arm with a distinct diagnostic"
827        );
828
829        // FAIL CLOSED, never silently recycle: an uncollected SLOT_READY slot
830        // must surface as an Err naming the loss, and the old recycle path
831        // (the "was SLOT_READY ... reused" tracing::warn that unmapped and
832        // reset the slot to FREE) must be gone entirely.
833        assert!(
834            production.contains("holds an uncollected completed readback (SLOT_READY)"),
835            "Fix: the SLOT_READY arm must fail closed with an error naming the uncollected readback, not recycle the slot"
836        );
837        assert!(
838            !production.contains("was SLOT_READY"),
839            "Fix: the silent recycle-on-reuse path (tracing::warn \"was SLOT_READY\" then unmap + store(SLOT_FREE)) is a Law-10 recall loss and must be removed, fail closed instead"
840        );
841        // The slot reuse check must not silently reset a non-FREE slot back to
842        // FREE; the only SLOT_FREE store is the post-submit transition. Count
843        // them: exactly the two PENDING transitions and zero recycle resets.
844        assert!(
845            !production.contains("slot.buffer.unmap();\n                slot.byte_len.store(0"),
846            "Fix: no recycle-and-continue (unmap + zero len + store(SLOT_FREE)) may remain in the reuse check"
847        );
848
849        // The old conflated message must be gone (it described every non-FREE
850        // state, READY and ERROR included, as a wrap).
851        assert!(
852            !production.contains("wrapped before collection"),
853            "Fix: the misleading 'wrapped before collection' message must be replaced by state-specific diagnostics"
854        );
855    }
856
857    /// VYRE-WGPU-003: `submit_readback` must accept a `src_offset` parameter
858    /// matching the `record_copy` signature.  Before the fix, `src_offset` was
859    /// hardcoded to 0, making sub-range reads silently return wrong data.
860    ///
861    /// This test verifies both the signature change and that the offset is
862    /// forwarded to the wgpu copy call (not discarded (without a live GPU)).
863    #[test]
864    fn submit_readback_has_src_offset_parameter_matching_record_copy() {
865        let src = include_str!("readback_ring.rs");
866        let production = src
867            .split("\n#[cfg(test)]\nmod tests")
868            .next()
869            .expect("Fix: readback_ring.rs should have a test module");
870
871        // The function signature must include src_offset.
872        assert!(
873            production.contains("pub fn submit_readback(")
874                && production.contains("src_offset: u64"),
875            "Fix: submit_readback must declare src_offset: u64 to match record_copy's signature"
876        );
877
878        // The copy call in submit_readback must forward src_offset, not
879        // hardcode 0.  We verify the copy_buffer_to_buffer call inside
880        // submit_readback uses `src_offset` as the second argument.
881        //
882        // Locate the submit_readback function body and check it does not
883        // contain `copy_buffer_to_buffer(src_buffer, 0,` (the old hardcoded
884        // form that silently read from offset 0).
885        let submit_body_start = production
886            .find("pub fn submit_readback(")
887            .expect("submit_readback must exist");
888        let submit_body = &production[submit_body_start..];
889        // Find the copy call within that body.
890        let copy_call_in_body = submit_body
891            .find("copy_buffer_to_buffer(src_buffer,")
892            .expect("submit_readback must contain a copy_buffer_to_buffer call");
893        let copy_call_text = &submit_body[copy_call_in_body..copy_call_in_body + 80];
894        assert!(
895            !copy_call_text.contains("copy_buffer_to_buffer(src_buffer, 0,"),
896            "Fix: submit_readback must forward src_offset to copy_buffer_to_buffer, not hardcode 0. Found: {copy_call_text:?}"
897        );
898        assert!(
899            copy_call_text.contains("copy_buffer_to_buffer(src_buffer, src_offset,"),
900            "Fix: submit_readback must pass src_offset as the second argument to copy_buffer_to_buffer. Found: {copy_call_text:?}"
901        );
902    }
903}
vyre_driver_wgpu/runtime/readback_ring.rs

vyre_driver_wgpu/runtime/
readback_ring.rs
