zencodec_io/
drive.rs

1//! Drive a resumable decoder from a slice or a random-access source.
2
3use alloc::vec::Vec;
4
5use how_far::{Pulse, Total};
6use zencodec::decode::{
7    ByteRange, DecodeRowSink, Event, IncrementalDecode, Input, OutputInfo, StepError,
8};
9use zencodec::source::ReadAt;
10
11use crate::stepper::{DriveFailure, ReadAtCursor};
12use crate::window::check_need;
13
14/// Tuning for the drivers.
15#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16#[non_exhaustive]
17pub struct DriveOptions {
18    /// Requested ranges closer than this are merged into one read (HTTP:
19    /// 16–64 KiB).
20    pub max_gap: u64,
21    /// Bytes read per round for a sequential request, on top of the
22    /// decoder's minimum.
23    pub read_size: usize,
24    /// Admission for driver-owned input. Defaults to historical unlimited behavior.
25    pub limits: crate::IoLimits,
26    /// Scheduling quota for async cursors and compute batches, independent of output draining.
27    pub work_quota: crate::WorkQuota,
28}
29
30impl Default for DriveOptions {
31    fn default() -> Self {
32        Self {
33            max_gap: 0,
34            read_size: 64 * 1024,
35            limits: crate::IoLimits::unlimited(),
36            work_quota: crate::WorkQuota::unlimited(),
37        }
38    }
39}
40
41impl DriveOptions {
42    /// Options with the given coalescing gap and sequential read size.
43    pub const fn new(max_gap: u64, read_size: usize) -> Self {
44        Self {
45            max_gap,
46            read_size,
47            limits: crate::IoLimits::unlimited(),
48            work_quota: crate::WorkQuota::unlimited(),
49        }
50    }
51
52    /// Apply input admission to synchronous and asynchronous drivers alike.
53    pub const fn with_limits(mut self, limits: crate::IoLimits) -> Self {
54        self.limits = limits;
55        self
56    }
57
58    /// Bound uninterrupted async work without coupling it to destination reads.
59    pub const fn with_work_quota(mut self, quota: crate::WorkQuota) -> Self {
60        self.work_quota = quota;
61        self
62    }
63}
64
65/// Counters from one drive.
66#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
67pub struct DriveStats {
68    /// Steps taken.
69    pub steps: u64,
70    /// Source read calls issued.
71    pub reads: u64,
72    /// Bytes read from the source.
73    pub bytes_read: u64,
74    /// Rows the decoder reported writing.
75    pub rows: u64,
76}
77
78/// Failure of an image drive: a [`DriveFailure`] whose stage error is the
79/// decoder's [`StepError`].
80pub type DriveError<CE, SE> = DriveFailure<StepError<CE>, SE>;
81
82/// Merge overlapping or near-adjacent ranges (gap `<= max_gap`) into one.
83pub fn coalesce_byte_ranges(ranges: &[ByteRange], max_gap: u64) -> Vec<ByteRange> {
84    let mut out = Vec::new();
85    coalesce_into(ranges, max_gap, &mut out);
86    out
87}
88
89/// [`coalesce_byte_ranges`] writing into `out` (cleared first), so a loop
90/// reuses one allocation.
91pub fn coalesce_into(ranges: &[ByteRange], max_gap: u64, out: &mut Vec<ByteRange>) {
92    out.clear();
93    out.extend(ranges.iter().copied().filter(|r| !r.is_empty()));
94    if out.is_empty() {
95        return;
96    }
97    out.sort_by_key(|r| (r.offset, r.length));
98    let mut write = 0usize;
99    for i in 1..out.len() {
100        let r = out[i];
101        let cur = out[write];
102        if r.offset <= cur.end().saturating_add(max_gap) {
103            out[write].length = cur.end().max(r.end()) - cur.offset;
104        } else {
105            write += 1;
106            out[write] = r;
107        }
108    }
109    out.truncate(write + 1);
110}
111
112/// Report the row total once the header is known, and rows as they land.
113pub(crate) struct Progressor<'a> {
114    pulse: Option<&'a dyn Pulse>,
115    told_total: bool,
116}
117
118impl<'a> Progressor<'a> {
119    pub(crate) fn new(pulse: Option<&'a dyn Pulse>) -> Self {
120        Self {
121            pulse,
122            told_total: false,
123        }
124    }
125
126    #[inline]
127    pub(crate) fn observe<D: IncrementalDecode + ?Sized>(&mut self, dec: &D, event: &Event) {
128        let Some(p) = self.pulse else { return };
129        if !self.told_total
130            && let Some(info) = dec.info()
131        {
132            // A pulse that cannot take a total records the issue itself.
133            let _ = p.set_total(Total::Exact(info.height as u64));
134            self.told_total = true;
135        }
136        if let Event::Rows { rows, .. } = event {
137            p.advance(*rows as u64);
138        }
139    }
140}
141
142/// Run `decoder` over an in-memory image without copying it.
143///
144/// Every input is a sub-slice of `bytes` carrying the total length. A
145/// sequential request gets at least its minimum and up to `window` bytes
146/// (a small `window` exercises boundaries); a range gets exactly its span.
147///
148/// `pulse` is checked for cancellation each step (a plain `Stop` passes as
149/// `how_far::StopOnly::borrowed(&stop)`), told the row total once the
150/// header is known, and advanced by every strip written.
151pub fn drive_from_slice<D>(
152    decoder: &mut D,
153    bytes: &[u8],
154    sink: &mut dyn DecodeRowSink,
155    pulse: Option<&dyn Pulse>,
156    window: usize,
157) -> Result<(OutputInfo, DriveStats), DriveError<D::Error, core::convert::Infallible>>
158where
159    D: IncrementalDecode + ?Sized,
160{
161    let mut progress = Progressor::new(pulse);
162    let total = bytes.len() as u64;
163    let window = window.max(1);
164    let mut stats = DriveStats::default();
165    // The current input: `bytes[start..end]`. `start <= end <= len` holds:
166    // consumption never exceeds the offer and `check_need` keeps every served
167    // request inside `total`.
168    let (mut start, mut end) = (0usize, 0usize);
169    loop {
170        let input = Input::new(start as u64, &bytes[start..end]).with_total_len(total);
171        let p = decoder
172            .step(input, sink, pulse)
173            .map_err(DriveFailure::Step)?;
174        stats.steps += 1;
175        start += p.consumed.min(end - start);
176        progress.observe(decoder, &p.event);
177        match p.event {
178            Event::Done(info) => return Ok((info, stats)),
179            Event::Need => {}
180            Event::Rows { rows, .. } => {
181                stats.rows += rows as u64;
182                continue;
183            }
184            _ => continue,
185        }
186        let requests = decoder.requested();
187        check_need(requests, &input, p.consumed).map_err(DriveFailure::Protocol)?;
188        (start, end) = crate::stepper::slice_span(start, end, window, bytes.len(), requests);
189        stats.reads += 1;
190        stats.bytes_read += (end - start) as u64;
191    }
192}
193
194/// Run `decoder` over a random-access source.
195///
196/// Sequential requests read forward into a driver-owned window (short
197/// reads mark the end of the source; [`ReadAt::len`] is called only when the
198/// decoder asks for the length). Range requests are coalesced (`max_gap`)
199/// and fetched with one [`read_many_at`](ReadAt::read_many_at) per round, so
200/// a concurrent source overlaps their round trips; later ranges of the same
201/// round are served from the fetched buffers. A source that is already in
202/// memory ([`ReadAt::as_slice`]) is driven without copying. The same driver
203/// serves any [`ByteStepper`](crate::ByteStepper) through
204/// [`ReadAtCursor`].
205pub fn drive_from_read_at<D, S>(
206    decoder: &mut D,
207    source: &mut S,
208    sink: &mut dyn DecodeRowSink,
209    pulse: Option<&dyn Pulse>,
210    opts: DriveOptions,
211) -> Result<(OutputInfo, DriveStats), DriveError<D::Error, S::Error>>
212where
213    D: IncrementalDecode + ?Sized,
214    S: ReadAt + ?Sized,
215{
216    // A source already in memory is driven without copying, by the slice
217    // driver.
218    if let Some(bytes) = source.as_slice() {
219        return drive_from_slice(decoder, bytes, sink, pulse, opts.read_size)
220            .map_err(DriveFailure::widen);
221    }
222    let mut progress = Progressor::new(pulse);
223    let mut cursor = ReadAtCursor::new(opts);
224    let mut rows = 0u64;
225    loop {
226        let input = cursor.offer(source);
227        let p = decoder
228            .step(input, sink, pulse)
229            .map_err(DriveFailure::Step)?;
230        if p.event == Event::Need {
231            check_need(decoder.requested(), &input, p.consumed).map_err(DriveFailure::Protocol)?;
232        }
233        cursor.take(p.consumed);
234        progress.observe(decoder, &p.event);
235        match p.event {
236            Event::Done(info) => {
237                let mut stats = cursor.stats();
238                stats.rows = rows;
239                return Ok((info, stats));
240            }
241            Event::Need => cursor.serve(decoder.requested(), source)?,
242            Event::Rows { rows: n, .. } => rows += n as u64,
243            _ => {}
244        }
245    }
246}
247
248#[cfg(test)]
249mod tests {
250    use super::*;
251
252    #[test]
253    fn coalesce_merges_within_gap() {
254        let ranges = [
255            ByteRange::new(500, 100),
256            ByteRange::new(100, 50),
257            ByteRange::new(180, 20),
258        ];
259        assert_eq!(
260            coalesce_byte_ranges(&ranges, 64),
261            [ByteRange::new(100, 100), ByteRange::new(500, 100)]
262        );
263        assert_eq!(coalesce_byte_ranges(&ranges, 0).len(), 3);
264    }
265}
zencodec_io/drive.rs

zencodec_io/
drive.rs
