zencodec_io/
admission.rs

1//! Admission shared by synchronous and asynchronous input drivers.
2
3use alloc::collections::TryReserveError;
4use alloc::vec::Vec;
5use core::fmt;
6use zencodec::decode::{ByteRange, InputRequest};
7
8/// Limits on storage and requests owned by one input driver.
9///
10/// This bounds retained payload capacities, stored range descriptors and the
11/// number of ranges handed to a source at once. It excludes caller-owned
12/// slices, allocator bookkeeping/transient reallocation, source caches, codec
13/// allocations and output. It is not a process-memory limit.
14///
15/// Zero is valid: no owned payload, range descriptors or reads, respectively.
16/// Borrowed resident input remains usable with all three limits set to zero.
17#[derive(Debug, Clone, Copy, PartialEq, Eq)]
18pub struct IoLimits {
19    input_bytes: usize,
20    ranges: usize,
21    in_flight: usize,
22}
23
24impl IoLimits {
25    /// Construct finite limits. Requests are checked before allocation or I/O.
26    pub const fn new(input_bytes: usize, ranges: usize, in_flight: usize) -> Self {
27        Self {
28            input_bytes,
29            ranges,
30            in_flight,
31        }
32    }
33
34    /// Preserve the historical unbounded driver configuration.
35    pub const fn unlimited() -> Self {
36        Self::new(usize::MAX, usize::MAX, usize::MAX)
37    }
38
39    /// Maximum retained driver-owned payload capacity.
40    pub const fn max_input_bytes(self) -> usize {
41        self.input_bytes
42    }
43    /// Maximum range descriptors in a fetch plan (excluding codec-owned requests).
44    pub const fn max_ranges(self) -> usize {
45        self.ranges
46    }
47    /// Maximum ranges passed to a source in one batch.
48    pub const fn max_in_flight(self) -> usize {
49        self.in_flight
50    }
51
52    pub(crate) fn bytes(self, needed: u64) -> Result<usize, InputError> {
53        if needed > self.input_bytes as u64 {
54            return Err(InputError::Limit {
55                needed,
56                limit: self.input_bytes,
57            });
58        }
59        usize::try_from(needed).map_err(|_| InputError::Limit {
60            needed,
61            limit: self.input_bytes,
62        })
63    }
64
65    pub(crate) fn reading(self) -> Result<(), InputError> {
66        if self.in_flight == 0 {
67            Err(InputError::ReadsDisabled)
68        } else {
69            Ok(())
70        }
71    }
72}
73
74impl Default for IoLimits {
75    fn default() -> Self {
76        Self::unlimited()
77    }
78}
79
80/// Driver admission or allocation failure, separate from codec/source errors.
81#[derive(Debug)]
82#[non_exhaustive]
83pub enum InputError {
84    /// A contiguous request or retained capacity exceeds the input allowance.
85    Limit {
86        /// Required payload capacity.
87        needed: u64,
88        /// Configured maximum capacity.
89        limit: usize,
90    },
91    /// The configuration allows no stored range descriptors.
92    RangesDisabled,
93    /// The configuration allows no source reads.
94    ReadsDisabled,
95    /// Fallible storage reservation failed. The original allocator error remains available.
96    Allocation {
97        /// Requested allocation size (saturating for descriptor overflow).
98        bytes: usize,
99        /// Original allocation failure.
100        source: TryReserveError,
101    },
102}
103
104impl fmt::Display for InputError {
105    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
106        match self {
107            Self::Limit { needed, limit } => {
108                write!(f, "input needs {needed} bytes; limit is {limit}")
109            }
110            Self::RangesDisabled => f.write_str("input range descriptors are disabled"),
111            Self::ReadsDisabled => f.write_str("input reads are disabled"),
112            Self::Allocation { bytes, .. } => {
113                write!(f, "could not reserve {bytes} input-storage bytes")
114            }
115        }
116    }
117}
118
119impl core::error::Error for InputError {
120    fn source(&self) -> Option<&(dyn core::error::Error + 'static)> {
121        match self {
122            Self::Allocation { source, .. } => Some(source),
123            _ => None,
124        }
125    }
126}
127
128pub(crate) fn reserve<T>(v: &mut Vec<T>, len: usize) -> Result<(), InputError> {
129    v.try_reserve_exact(len.saturating_sub(v.len()))
130        .map_err(|source| InputError::Allocation {
131            bytes: len.saturating_mul(core::mem::size_of::<T>()),
132            source,
133        })
134}
135
136/// Choose an affordable alternative, preserving the codec's order of preference.
137/// The borrowed request list need not be copied to make this choice.
138pub(crate) fn select(
139    requests: &[InputRequest],
140    limits: IoLimits,
141) -> Result<InputRequest, InputError> {
142    let mut first_error = None;
143    for &request in requests {
144        let result = match request {
145            InputRequest::AtLeast { min_len, .. } => {
146                limits.bytes(min_len.max(1) as u64).map(|_| ())
147            }
148            InputRequest::Range(r) => limits.bytes(r.length).and({
149                if limits.ranges == 0 {
150                    Err(InputError::RangesDisabled)
151                } else {
152                    Ok(())
153                }
154            }),
155            _ => Ok(()),
156        };
157        match result {
158            Ok(()) => return Ok(request),
159            Err(e) => {
160                first_error.get_or_insert(e);
161            }
162        }
163    }
164    Err(first_error.unwrap_or(InputError::RangesDisabled))
165}
166
167/// A bounded prefetch plan for alternatives, not an all-dependency plan.
168/// Always includes `first`. Optional ranges and gap bytes may be omitted to
169/// fit the budget. Overlaps are merged without charging their bytes twice.
170pub(crate) fn plan(
171    first: ByteRange,
172    requests: &[InputRequest],
173    gap: u64,
174    limits: IoLimits,
175    out: &mut Vec<ByteRange>,
176) -> Result<(), InputError> {
177    limits.reading()?;
178    limits.bytes(first.length)?;
179    if limits.ranges == 0 {
180        return Err(InputError::RangesDisabled);
181    }
182    out.clear();
183    reserve(out, requests.len().min(limits.ranges))?;
184    out.push(first);
185    let mut bytes = first.length;
186    for request in requests {
187        let InputRequest::Range(r) = *request else {
188            continue;
189        };
190        if r.is_empty() || r == first || out.len() == limits.ranges {
191            continue;
192        }
193        // Conservative admission before sorting. Overlaps may be fetched in a
194        // later round if counting them separately would exceed this allowance.
195        if let Some(sum) = bytes.checked_add(r.length)
196            && sum <= limits.input_bytes as u64
197        {
198            out.push(r);
199            bytes = sum;
200        }
201    }
202    out.sort_unstable_by_key(|r| (r.offset, r.length));
203    let mut write = 0;
204    for i in 1..out.len() {
205        let r = out[i];
206        let current = out[write];
207        let extra = r.offset.saturating_sub(current.end());
208        if r.offset <= current.end().saturating_add(gap)
209            && extra <= (limits.input_bytes as u64).saturating_sub(bytes)
210        {
211            let merged = current.end().max(r.end()) - current.offset;
212            bytes = bytes - current.length - r.length + merged;
213            out[write].length = merged;
214        } else {
215            write += 1;
216            out[write] = r;
217        }
218    }
219    out.truncate(write + 1);
220    Ok(())
221}
222
223pub(crate) fn capacity(fetched: &[(u64, Vec<u8>)]) -> usize {
224    fetched
225        .iter()
226        .fold(0usize, |sum, (_, b)| sum.saturating_add(b.capacity()))
227}
228
229/// Async image drivers cannot borrow a non-Sync decoder across an await. Copy
230/// only the selected request and a bounded number of affordable alternatives.
231#[cfg(feature = "tokio")]
232pub(crate) fn copy_requests(
233    requests: &[InputRequest],
234    limits: IoLimits,
235    out: &mut Vec<InputRequest>,
236) -> Result<(), InputError> {
237    let first = select(requests, limits)?;
238    out.clear();
239    let count = if matches!(first, InputRequest::Range(_)) {
240        requests.len().min(limits.ranges)
241    } else {
242        1
243    };
244    reserve(out, count)?;
245    out.push(first);
246    if matches!(first, InputRequest::Range(_)) {
247        for &r in requests {
248            if out.len() == count {
249                break;
250            }
251            if r != first
252                && matches!(r, InputRequest::Range(b) if b.length <= limits.input_bytes as u64)
253            {
254                out.push(r);
255            }
256        }
257    }
258    Ok(())
259}
260
261pub(crate) fn make_room(
262    window: &crate::window::Window,
263    fetched: &mut Vec<(u64, Vec<u8>)>,
264    target: usize,
265    limits: IoLimits,
266) {
267    if capacity(fetched).saturating_add(window.capacity().max(target)) > limits.input_bytes {
268        fetched.clear();
269    }
270}
271
272pub(crate) fn allocate_ranges(
273    ranges: &[ByteRange],
274    limits: IoLimits,
275    fetched: &mut Vec<(u64, Vec<u8>)>,
276) -> Result<(), InputError> {
277    // Reuse allocations when their capacities fit. Discard them before growth
278    // otherwise; old and new rounds must not silently double retained payload.
279    fetched.truncate(ranges.len());
280    let reusable = ranges.iter().enumerate().fold(0usize, |sum, (i, r)| {
281        sum.saturating_add(
282            fetched
283                .get(i)
284                .map_or(0, |(_, b)| b.capacity())
285                .max(r.length as usize),
286        )
287    });
288    if reusable > limits.input_bytes {
289        fetched.clear();
290    }
291    reserve(fetched, ranges.len())?;
292    fetched.resize_with(ranges.len(), Default::default);
293    let mut used = 0usize;
294    for (range, (offset, buf)) in ranges.iter().zip(fetched.iter_mut()) {
295        let len = limits.bytes(range.length)?;
296        *offset = range.offset;
297        reserve(buf, len)?;
298        used = used.saturating_add(buf.capacity());
299        limits.bytes(used as u64)?;
300        buf.resize(len, 0);
301    }
302    Ok(())
303}
304
305#[cfg(all(test, feature = "tokio"))]
306mod tests {
307    use super::*;
308
309    #[test]
310    fn async_request_copy_does_not_copy_an_unbounded_alternative_list() {
311        let requests: Vec<_> = (0..10000)
312            .map(|offset| InputRequest::range(offset * 8, 8))
313            .collect();
314        let mut copied = Vec::new();
315        copy_requests(&requests, IoLimits::new(16, 2, 1), &mut copied).unwrap();
316        assert_eq!(copied, requests[..2]);
317        assert!(copied.capacity() <= 2);
318        copy_requests(
319            &[InputRequest::range(0, 17), InputRequest::at_least(32, 8)],
320            IoLimits::new(16, 0, 1),
321            &mut copied,
322        )
323        .unwrap();
324        assert_eq!(copied, [InputRequest::at_least(32, 8)]);
325    }
326}
zencodec_io/admission.rs

zencodec_io/
admission.rs
