molgfx_gpu/residency/
upload.rs

1//! Fixed-capacity upload staging with explicit backpressure and fence-safe
2//! retirement.
3//!
4//! Reservations advance monotonically through a byte ring and retire in FIFO
5//! order. Alignment and wrap padding remain occupied until the reservation's
6//! fence completes, preventing reuse while a backend may still read it.
7
8use super::upload_types::{
9    FenceValue, Retirement, UploadBackpressure, UploadMetrics, UploadReservation, UploadRingConfig,
10    UploadState, UploadTicket,
11};
12
13#[derive(Clone, Copy, Debug, PartialEq, Eq)]
14enum RecordState {
15    Vacant,
16    Live(UploadState),
17}
18
19#[derive(Clone, Copy, Debug)]
20struct UploadRecord {
21    ticket: UploadTicket,
22    payload_start: u64,
23    end: u64,
24    payload_len: usize,
25    state: RecordState,
26}
27
28impl UploadRecord {
29    const fn vacant() -> Self {
30        Self {
31            ticket: UploadTicket(0),
32            payload_start: 0,
33            end: 0,
34            payload_len: 0,
35            state: RecordState::Vacant,
36        }
37    }
38}
39
40/// A bounded host staging ring independent of any concrete GPU backend.
41#[derive(Debug)]
42pub struct UploadRing {
43    bytes: Vec<u8>,
44    records: Vec<UploadRecord>,
45    config: UploadRingConfig,
46    head: usize,
47    active: usize,
48    write_cursor: u64,
49    retired_cursor: u64,
50    next_ticket: u64,
51    metrics: UploadMetrics,
52}
53
54impl UploadRing {
55    /// Allocates byte and ticket storage once.
56    ///
57    /// # Errors
58    ///
59    /// Returns [`UploadBackpressure::InvalidConfiguration`] for empty
60    /// capacities, zero budget or non-power-of-two alignment.
61    pub fn new(config: UploadRingConfig) -> Result<Self, UploadBackpressure> {
62        config.validate()?;
63        Ok(Self {
64            bytes: vec![0; config.capacity_bytes],
65            records: vec![UploadRecord::vacant(); config.ticket_capacity],
66            config,
67            head: 0,
68            active: 0,
69            write_cursor: 0,
70            retired_cursor: 0,
71            next_ticket: 1,
72            metrics: UploadMetrics {
73                host_allocation_events: 2,
74                ..UploadMetrics::default()
75            },
76        })
77    }
78
79    /// Resets only the per-epoch payload counter.
80    pub fn begin_epoch(&mut self) {
81        self.metrics.epoch_bytes = 0;
82    }
83
84    /// Reserves one contiguous aligned staging range.
85    ///
86    /// # Errors
87    ///
88    /// Returns explicit backpressure without growing either storage vector.
89    pub fn reserve(&mut self, byte_len: usize) -> Result<UploadReservation, UploadBackpressure> {
90        if byte_len == 0 {
91            return Err(UploadBackpressure::EmptyUpload);
92        }
93        if self.metrics.epoch_bytes.saturating_add(byte_len as u64)
94            > self.config.epoch_budget_bytes as u64
95        {
96            return self.stall(byte_len, UploadBackpressure::EpochBudget);
97        }
98        if self.active == self.records.len() {
99            return self.stall(byte_len, UploadBackpressure::TicketCapacity);
100        }
101        if self.active == 0 {
102            self.write_cursor = 0;
103            self.retired_cursor = 0;
104        }
105        let (payload_start, end) = self.extent_for(byte_len)?;
106        if end - self.retired_cursor > self.bytes.len() as u64 {
107            return self.stall(byte_len, UploadBackpressure::RingFull);
108        }
109        let offset = self.byte_offset(payload_start)?;
110        let ticket = UploadTicket(self.next_ticket);
111        self.next_ticket = self.next_ticket.wrapping_add(1).max(1);
112        let index = (self.head + self.active) % self.records.len();
113        self.records[index] = UploadRecord {
114            ticket,
115            payload_start,
116            end,
117            payload_len: byte_len,
118            state: RecordState::Live(UploadState::Reserved),
119        };
120        self.active += 1;
121        self.write_cursor = end;
122        self.metrics.epoch_bytes = self.metrics.epoch_bytes.saturating_add(byte_len as u64);
123        self.refresh_gauges();
124        Ok(UploadReservation {
125            ticket,
126            offset,
127            len: byte_len,
128        })
129    }
130
131    /// Returns the writable payload range while a ticket is reserved.
132    ///
133    /// # Errors
134    ///
135    /// Returns [`UploadBackpressure::InvalidTicket`] for stale or already
136    /// committed reservations.
137    pub fn bytes_mut(
138        &mut self,
139        reservation: UploadReservation,
140    ) -> Result<&mut [u8], UploadBackpressure> {
141        let index = self.record_index(reservation.ticket, UploadState::Reserved)?;
142        let record = self.records[index];
143        let expected_offset = self.byte_offset(record.payload_start)?;
144        if record.payload_len != reservation.len || expected_offset != reservation.offset {
145            return Err(UploadBackpressure::InvalidTicket);
146        }
147        let start = reservation.offset;
148        let end = start + reservation.len;
149        Ok(&mut self.bytes[start..end])
150    }
151
152    /// Marks filled bytes ready for backend submission.
153    ///
154    /// # Errors
155    ///
156    /// The ticket must currently be reserved.
157    pub fn commit(&mut self, ticket: UploadTicket) -> Result<(), UploadBackpressure> {
158        let index = self.record_index(ticket, UploadState::Reserved)?;
159        self.records[index].state = RecordState::Live(UploadState::Ready);
160        self.metrics.bytes_staged = self
161            .metrics
162            .bytes_staged
163            .saturating_add(self.records[index].payload_len as u64);
164        Ok(())
165    }
166
167    /// Associates a ready upload with the fence of the submission reading it.
168    ///
169    /// # Errors
170    ///
171    /// The ticket must currently be ready.
172    pub fn submit(
173        &mut self,
174        ticket: UploadTicket,
175        fence: FenceValue,
176    ) -> Result<(), UploadBackpressure> {
177        let index = self.record_index(ticket, UploadState::Ready)?;
178        let payload_len = self.records[index].payload_len as u64;
179        if self.metrics.in_flight_bytes.saturating_add(payload_len)
180            > self.config.in_flight_budget_bytes as u64
181        {
182            return self.stall(
183                self.records[index].payload_len,
184                UploadBackpressure::InFlightBudget,
185            );
186        }
187        self.records[index].state = RecordState::Live(UploadState::InFlight(fence));
188        self.metrics.bytes_submitted = self.metrics.bytes_submitted.saturating_add(payload_len);
189        self.metrics.in_flight_bytes = self.metrics.in_flight_bytes.saturating_add(payload_len);
190        self.metrics.peak_in_flight_bytes = self
191            .metrics
192            .peak_in_flight_bytes
193            .max(self.metrics.in_flight_bytes);
194        Ok(())
195    }
196
197    /// Validates in-flight admission before a backend submission is issued.
198    ///
199    /// # Errors
200    ///
201    /// The ticket must be ready and fit the fixed in-flight byte budget.
202    pub fn ensure_submittable(&mut self, ticket: UploadTicket) -> Result<(), UploadBackpressure> {
203        let index = self.record_index(ticket, UploadState::Ready)?;
204        let payload_len = self.records[index].payload_len as u64;
205        if self.metrics.in_flight_bytes.saturating_add(payload_len)
206            > self.config.in_flight_budget_bytes as u64
207        {
208            return self.stall(
209                self.records[index].payload_len,
210                UploadBackpressure::InFlightBudget,
211            );
212        }
213        Ok(())
214    }
215
216    /// Cancels a reservation that has not been submitted.
217    ///
218    /// # Errors
219    ///
220    /// In-flight and stale tickets cannot be cancelled.
221    pub fn cancel(&mut self, ticket: UploadTicket) -> Result<(), UploadBackpressure> {
222        let Some(index) = self.find_ticket(ticket) else {
223            return Err(UploadBackpressure::InvalidTicket);
224        };
225        match self.records[index].state {
226            RecordState::Live(UploadState::Reserved | UploadState::Ready) => {
227                self.records[index].state = RecordState::Live(UploadState::Cancelled);
228                self.metrics.bytes_cancelled = self
229                    .metrics
230                    .bytes_cancelled
231                    .saturating_add(self.records[index].payload_len as u64);
232                self.reclaim_cancelled_front();
233                Ok(())
234            }
235            RecordState::Vacant
236            | RecordState::Live(UploadState::InFlight(_) | UploadState::Cancelled) => {
237                Err(UploadBackpressure::InvalidTicket)
238            }
239        }
240    }
241
242    /// Retires the contiguous FIFO prefix whose submission fences completed.
243    #[must_use]
244    pub fn retire(&mut self, completed: FenceValue) -> Retirement {
245        let mut retirement = Retirement::default();
246        while self.active > 0 {
247            let record = self.records[self.head];
248            let completed_record = matches!(
249                record.state,
250                RecordState::Live(UploadState::InFlight(fence)) if fence <= completed
251            );
252            if !completed_record {
253                break;
254            }
255            retirement.tickets = retirement.tickets.saturating_add(1);
256            retirement.bytes = retirement.bytes.saturating_add(record.payload_len as u64);
257            self.metrics.bytes_retired = self
258                .metrics
259                .bytes_retired
260                .saturating_add(record.payload_len as u64);
261            self.metrics.in_flight_bytes = self
262                .metrics
263                .in_flight_bytes
264                .saturating_sub(record.payload_len as u64);
265            self.reclaim_head(record.end);
266        }
267        retirement
268    }
269
270    /// Current lifecycle for a live ticket.
271    #[must_use]
272    pub fn state(&self, ticket: UploadTicket) -> Option<UploadState> {
273        let index = self.find_ticket(ticket)?;
274        match self.records[index].state {
275            RecordState::Live(state) => Some(state),
276            RecordState::Vacant => None,
277        }
278    }
279
280    /// Current cumulative counters.
281    #[must_use]
282    pub const fn metrics(&self) -> UploadMetrics {
283        self.metrics
284    }
285
286    /// Entire staging storage for backend copy commands.
287    #[must_use]
288    pub fn staging_bytes(&self) -> &[u8] {
289        &self.bytes
290    }
291
292    fn extent_for(&self, byte_len: usize) -> Result<(u64, u64), UploadBackpressure> {
293        let alignment = self.config.alignment as u64;
294        let mask = alignment - 1;
295        let begin = self.write_cursor;
296        let mut payload_start = begin
297            .checked_add(mask)
298            .map(|value| value & !mask)
299            .ok_or(UploadBackpressure::SizeOverflow)?;
300        let offset = payload_start % self.bytes.len() as u64;
301        if offset + byte_len as u64 > self.bytes.len() as u64 {
302            payload_start = payload_start
303                .checked_add(self.bytes.len() as u64 - offset)
304                .ok_or(UploadBackpressure::SizeOverflow)?;
305        }
306        let end = payload_start
307            .checked_add(byte_len as u64)
308            .ok_or(UploadBackpressure::SizeOverflow)?;
309        Ok((payload_start, end))
310    }
311
312    fn record_index(
313        &self,
314        ticket: UploadTicket,
315        expected: UploadState,
316    ) -> Result<usize, UploadBackpressure> {
317        let Some(index) = self.find_ticket(ticket) else {
318            return Err(UploadBackpressure::InvalidTicket);
319        };
320        if self.records[index].state != RecordState::Live(expected) {
321            return Err(UploadBackpressure::InvalidTicket);
322        }
323        Ok(index)
324    }
325
326    fn find_ticket(&self, ticket: UploadTicket) -> Option<usize> {
327        if ticket.0 == 0 || self.records.is_empty() {
328            return None;
329        }
330        let Ok(slot_count) = u64::try_from(self.records.len()) else {
331            return None;
332        };
333        let slot = (ticket.0 - 1) % slot_count;
334        let Ok(index) = usize::try_from(slot) else {
335            return None;
336        };
337        (self.records[index].ticket == ticket
338            && !matches!(self.records[index].state, RecordState::Vacant))
339        .then_some(index)
340    }
341
342    fn byte_offset(&self, cursor: u64) -> Result<usize, UploadBackpressure> {
343        let capacity =
344            u64::try_from(self.bytes.len()).map_err(|_| UploadBackpressure::SizeOverflow)?;
345        usize::try_from(cursor % capacity).map_err(|_| UploadBackpressure::SizeOverflow)
346    }
347
348    fn reclaim_cancelled_front(&mut self) {
349        while self.active > 0 {
350            let record = self.records[self.head];
351            if record.state != RecordState::Live(UploadState::Cancelled) {
352                break;
353            }
354            self.reclaim_head(record.end);
355        }
356    }
357
358    fn reclaim_head(&mut self, end: u64) {
359        self.records[self.head] = UploadRecord::vacant();
360        self.retired_cursor = end;
361        self.head = (self.head + 1) % self.records.len();
362        self.active -= 1;
363        self.refresh_gauges();
364    }
365
366    fn refresh_gauges(&mut self) {
367        self.metrics.active_tickets = self.active as u64;
368        self.metrics.occupied_bytes = self.write_cursor - self.retired_cursor;
369        self.metrics.peak_occupied_bytes = self
370            .metrics
371            .peak_occupied_bytes
372            .max(self.metrics.occupied_bytes);
373    }
374
375    fn stall<T>(
376        &mut self,
377        byte_len: usize,
378        error: UploadBackpressure,
379    ) -> Result<T, UploadBackpressure> {
380        self.metrics.stall_events = self.metrics.stall_events.saturating_add(1);
381        self.metrics.stalled_bytes = self.metrics.stalled_bytes.saturating_add(byte_len as u64);
382        Err(error)
383    }
384}
385
386#[cfg(test)]
387#[path = "upload_tests.rs"]
388mod tests;
molgfx_gpu/residency/upload.rs

molgfx_gpu/residency/
upload.rs
