Skip to main content

zencodec_media/
staging.rs

1//! Explicit compressed-packet staging for finite encoding.
2//!
3//! Disk staging writes while the encoder is alive and may block. It is not
4//! permission to backpressure a live encoder on the final output transport.
5//! Packet metadata stays in memory; payloads and additions are preserved without
6//! container rewriting. Disk files must be private scratch files exclusively
7//! controlled by this operation: aliases/external mutation are not supported.
8//! The caller owns the file's path, permissions, retention and deletion policy.
9use crate::{
10    packet::{Packet, PacketPayload},
11    staged::Cause,
12};
13use alloc::{collections::VecDeque, vec::Vec};
14use core::fmt;
15
16use std::{
17    fs::File,
18    io::{Read, Seek, SeekFrom, Write},
19};
20use zencodec::{
21    At,
22    operation::{StagingFacts, StagingKind},
23    storage::StorageBudget,
24};
25
26/// Independent caps for disk payload and its in-memory packet index.
27#[derive(Clone, Copy, Debug)]
28pub struct DiskLimits {
29    /// Total packet and addition bytes appended after the scratch origin.
30    pub max_bytes: u64,
31    /// Maximum number of staged packets.
32    pub max_packets: usize,
33    /// Maximum side-data entries on each packet, including empty entries.
34    pub max_additions_per_packet: usize,
35}
36impl DiskLimits {
37    /// Zero limits are real limits. No unbounded or unlimited sentinel is used.
38    pub const fn new(max_bytes: u64, max_packets: usize, max_additions_per_packet: usize) -> Self {
39        Self {
40            max_bytes,
41            max_packets,
42            max_additions_per_packet,
43        }
44    }
45}
46/// Staging refused before exceeding its configured scope.
47#[derive(Clone, Copy, Debug, PartialEq, Eq)]
48#[non_exhaustive]
49pub enum Refusal {
50    /// Appending this packet would exceed the disk byte allowance.
51    Bytes,
52    /// Appending this packet would exceed the packet-index allowance.
53    Packets,
54    /// Packet side-data count exceeds its index allowance.
55    Additions,
56    /// Scratch file must start at its existing end, preserving all prior bytes.
57    NotAtEnd,
58    /// A size/offset calculation cannot be represented.
59    Overflow,
60}
61impl fmt::Display for Refusal {
62    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
63        write!(f, "packet staging refused: {self:?}")
64    }
65}
66impl core::error::Error for Refusal {}
67fn refuse(reason: Refusal) -> Cause {
68    alloc::boxed::Box::new(reason)
69}
70fn check(pulse: &dyn zencodec::Pulse) -> Result<(), Cause> {
71    pulse
72        .check()
73        .map_err(|r| At::wrap(zencodec::CodecError::new(None, r.into())).into())
74}
75struct Entry {
76    metadata: Packet,
77    offset: u64,
78    lengths: Vec<usize>,
79    bytes: usize,
80}
81struct Disk {
82    file: File,
83    limits: DiskLimits,
84    replay: StorageBudget,
85    origin: u64,
86    entries: VecDeque<Entry>,
87}
88enum Store {
89    Memory {
90        budget: StorageBudget,
91        packets: VecDeque<Packet>,
92    },
93    Disk(Disk),
94}
95/// Retained store chosen before accepting the first frame.
96///
97/// Byte limits do not bound native codec/muxer workspace, packet metadata or
98/// caller-owned pictures. Disk replay allocates admitted packet owners; a muxer
99/// that retains packets keeps their charges, so replay admission can later fail.
100/// No automatic disk fallback, path creation, deletion or retry occurs.
101pub struct Staging {
102    store: Store,
103    bytes: u64,
104}
105impl Staging {
106    /// Retain all compressed packet/addition payloads in admitted memory.
107    pub fn memory(budget: StorageBudget) -> Self {
108        Self {
109            store: Store::Memory {
110                budget,
111                packets: VecDeque::new(),
112            },
113            bytes: 0,
114        }
115    }
116    /// Append to an owned private scratch file at its current end. Existing
117    /// prefixes are preserved. This may seek, but writes nothing at construction.
118    /// Replay uses `replay` to admit owned packet bytes after the encoder drops.
119    /// A failed write may leave a partial spool; the final destination is separate.
120    pub fn disk(mut file: File, limits: DiskLimits, replay: StorageBudget) -> Result<Self, Cause> {
121        let origin = file.stream_position()?;
122        let end = file.seek(SeekFrom::End(0))?;
123        if origin != end {
124            return Err(refuse(Refusal::NotAtEnd));
125        }
126        Ok(Self {
127            store: Store::Disk(Disk {
128                file,
129                limits,
130                replay,
131                origin,
132                entries: VecDeque::new(),
133            }),
134            bytes: 0,
135        })
136    }
137    /// Bounded policy evidence for the actual store. Disk permits blocking spool
138    /// I/O during encoding; neither kind qualifies native workspace bounds.
139    pub fn facts(&self) -> StagingFacts {
140        let (kind, limit) = match &self.store {
141            Store::Memory { budget, .. } => (StagingKind::Memory, budget.max_bytes() as u64),
142            Store::Disk(disk) => (StagingKind::Disk, disk.limits.max_bytes),
143        };
144        StagingFacts::new(kind, self.bytes, limit)
145    }
146    pub(crate) fn admitted_memory(&self) -> usize {
147        match &self.store {
148            Store::Memory { budget, .. } => budget.reserved_bytes(),
149            Store::Disk(disk) => disk.replay.reserved_bytes(),
150        }
151    }
152    pub(crate) fn push(
153        &mut self,
154        packet: Packet,
155        pulse: &dyn zencodec::Pulse,
156    ) -> Result<(), Cause> {
157        check(pulse)?;
158        let bytes = packet
159            .additions()
160            .iter()
161            .try_fold(packet.len(), |sum, a| sum.checked_add(a.payload().len()))
162            .ok_or_else(|| refuse(Refusal::Overflow))?;
163        let total = self
164            .bytes
165            .checked_add(bytes as u64)
166            .ok_or_else(|| refuse(Refusal::Overflow))?;
167        match &mut self.store {
168            Store::Memory { budget, packets } => {
169                packets.try_reserve(1)?;
170                packets.push_back(packet.accounted(budget)?);
171            }
172            Store::Disk(disk) => {
173                if total > disk.limits.max_bytes {
174                    return Err(refuse(Refusal::Bytes));
175                }
176                if disk.entries.len() >= disk.limits.max_packets {
177                    return Err(refuse(Refusal::Packets));
178                }
179                if packet.additions().len() > disk.limits.max_additions_per_packet {
180                    return Err(refuse(Refusal::Additions));
181                }
182                let offset = disk
183                    .origin
184                    .checked_add(self.bytes)
185                    .ok_or_else(|| refuse(Refusal::Overflow))?;
186                disk.origin
187                    .checked_add(total)
188                    .ok_or_else(|| refuse(Refusal::Overflow))?;
189                disk.entries.try_reserve(1)?;
190                let mut lengths = Vec::new();
191                lengths.try_reserve(packet.additions().len() + 1)?;
192                disk.file.seek(SeekFrom::Start(offset))?;
193                let metadata = packet.map_payloads(|payload| {
194                    lengths.push(payload.len());
195                    for chunk in payload.as_bytes().chunks(64 * 1024) {
196                        check(pulse)?;
197                        disk.file.write_all(chunk)?;
198                    }
199                    Ok(PacketPayload::Owned(Vec::new()))
200                })?;
201                disk.entries.push_back(Entry {
202                    metadata,
203                    offset,
204                    lengths,
205                    bytes,
206                });
207            }
208        }
209        self.bytes = total;
210        check(pulse)
211    }
212    pub(crate) fn begin_replay(&mut self, pulse: &dyn zencodec::Pulse) -> Result<(), Cause> {
213        check(pulse)?;
214        if let Store::Disk(disk) = &mut self.store {
215            disk.file.flush()?;
216        }
217        check(pulse)
218    }
219    pub(crate) fn next(&mut self, pulse: &dyn zencodec::Pulse) -> Result<Option<Packet>, Cause> {
220        check(pulse)?;
221        match &mut self.store {
222            Store::Memory { packets, .. } => Ok(packets.pop_front()),
223            Store::Disk(disk) => {
224                let Some(entry) = disk.entries.pop_front() else {
225                    return Ok(None);
226                };
227                let mut payload = disk.replay.bytes(entry.bytes)?;
228                disk.file.seek(SeekFrom::Start(entry.offset))?;
229                for chunk in payload.as_mut_slice().chunks_mut(64 * 1024) {
230                    check(pulse)?;
231                    disk.file.read_exact(chunk)?;
232                }
233                check(pulse)?;
234                let payload = payload.freeze();
235                let mut lengths = entry.lengths.into_iter();
236                let mut offset = 0;
237                let packet = entry.metadata.map_payloads(|_| {
238                    let len = lengths
239                        .next()
240                        .expect("private metadata/payload count agrees");
241                    let next = offset + len;
242                    let view = payload.slice(offset..next)?;
243                    offset = next;
244                    Ok(PacketPayload::Accounted(view))
245                })?;
246                Ok(Some(packet))
247            }
248        }
249    }
250}
251
252#[cfg(test)]
253mod tests {
254    use super::*;
255    use crate::packet::{PacketAddition, RandomAccess};
256    use alloc::vec;
257    use std::{
258        fs::OpenOptions,
259        path::PathBuf,
260        sync::atomic::{AtomicUsize, Ordering},
261    };
262    use zencodec::time::{TimeBase, Timestamp};
263    static NEXT: AtomicUsize = AtomicUsize::new(0);
264    fn file() -> (File, PathBuf) {
265        let root = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("target/staging-tests");
266        std::fs::create_dir_all(&root).unwrap();
267        for _ in 0..1000 {
268            let path = root.join(alloc::format!(
269                "{}-{}.scratch",
270                std::process::id(),
271                NEXT.fetch_add(1, Ordering::Relaxed)
272            ));
273            match OpenOptions::new()
274                .read(true)
275                .write(true)
276                .create_new(true)
277                .open(&path)
278            {
279                Ok(mut file) => {
280                    file.write_all(b"prefix").unwrap();
281                    return (file, path);
282                }
283                Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => continue,
284                Err(e) => panic!("scratch creation failed: {e}"),
285            }
286        }
287        panic!("scratch namespace exhausted")
288    }
289    fn packet() -> Packet {
290        Packet::from_vec(vec![1, 2, 3])
291            .unwrap()
292            .with_track(7)
293            .with_pts(Timestamp::new(21, TimeBase::MILLIS))
294            .with_dts(Timestamp::new(19, TimeBase::MILLIS))
295            .with_duration_ticks(2)
296            .with_config_epoch(3)
297            .with_ordinal(11)
298            .with_input_offset(101)
299            .with_random_access(RandomAccess::Yes)
300            .with_addition(PacketAddition::new(9, PacketPayload::Owned(vec![4, 5])))
301    }
302    #[test]
303    fn staging_releases_oversized_backing_owners_before_replay() {
304        use alloc::sync::Arc;
305        for disk in [false, true] {
306            for accounted in [false, true] {
307                // Five encoded bytes must not keep an unrelated backing slab
308                // alive while the final destination waits. Cover additions too.
309                let slab: Arc<[u8]> = vec![7; 4096].into();
310                let weak = Arc::downgrade(&slab);
311                let source_budget = StorageBudget::new(4096, 1);
312                let packet = if accounted {
313                    let mut owner = source_budget.bytes(4096).unwrap();
314                    owner.as_mut_slice().fill(7);
315                    let owner = owner.freeze();
316                    packet()
317                        .map_payloads(|payload| {
318                            Ok(PacketPayload::Accounted(owner.slice(0..payload.len())?))
319                        })
320                        .unwrap()
321                } else {
322                    Packet::from_shared(slab.clone(), 100, 3)
323                        .unwrap()
324                        .with_addition(PacketAddition::new(
325                            9,
326                            PacketPayload::Shared {
327                                bytes: slab.clone(),
328                                start: 200,
329                                len: 2,
330                            },
331                        ))
332                };
333                drop(slab);
334                if accounted {
335                    assert_eq!(source_budget.reserved_bytes(), 4096);
336                    assert_eq!(source_budget.reserved_objects(), 1);
337                } else {
338                    assert!(weak.upgrade().is_some());
339                }
340                let budget = StorageBudget::new(5, 2);
341                let (mut staging, path) = if disk {
342                    let (file, path) = file();
343                    (
344                        Staging::disk(file, DiskLimits::new(5, 1, 1), budget.clone()).unwrap(),
345                        Some(path),
346                    )
347                } else {
348                    (Staging::memory(budget.clone()), None)
349                };
350                staging.push(packet, &zencodec::NoPulse).unwrap();
351                assert!(
352                    weak.upgrade().is_none(),
353                    "disk={disk}, accounted={accounted}"
354                );
355                assert_eq!(source_budget.reserved_bytes(), 0);
356                assert_eq!(source_budget.reserved_objects(), 0);
357                assert_eq!(budget.reserved_bytes(), if disk { 0 } else { 5 });
358                staging.begin_replay(&zencodec::NoPulse).unwrap();
359                let output = staging.next(&zencodec::NoPulse).unwrap().unwrap();
360                drop(staging);
361                assert_eq!(output.bytes(), &[7; 3]);
362                assert_eq!(output.additions()[0].payload().as_bytes(), &[7; 2]);
363                assert_eq!(budget.reserved_bytes(), 5);
364                drop(output);
365                assert_eq!(budget.reserved_bytes(), 0);
366                assert_eq!(budget.reserved_objects(), 0);
367                if let Some(path) = path {
368                    std::fs::remove_file(path).unwrap();
369                }
370            }
371        }
372    }
373
374    #[test]
375    fn disk_preserves_packet_metadata_additions_prefix_and_replay_owners() {
376        let (file, path) = file();
377        let budget = StorageBudget::new(5, 1);
378        let mut staging = Staging::disk(file, DiskLimits::new(5, 1, 1), budget.clone()).unwrap();
379        let original = packet();
380        staging
381            .push(original.clone(), &zencodec::how_far::NoPulse)
382            .unwrap();
383        assert_eq!(budget.reserved_bytes(), 0);
384        assert_eq!(staging.facts().payload_bytes, 5);
385        staging.begin_replay(&zencodec::how_far::NoPulse).unwrap();
386        let output = staging.next(&zencodec::how_far::NoPulse).unwrap().unwrap();
387        assert_eq!(output.bytes(), original.bytes());
388        assert_eq!(output.track(), original.track());
389        assert_eq!(output.pts(), original.pts());
390        assert_eq!(output.dts(), original.dts());
391        assert_eq!(output.duration_ticks(), original.duration_ticks());
392        assert_eq!(output.config_epoch(), original.config_epoch());
393        assert_eq!(output.ordinal(), original.ordinal());
394        assert_eq!(output.input_offset(), original.input_offset());
395        assert_eq!(output.random_access(), original.random_access());
396        assert_eq!(output.additions()[0].id(), 9);
397        assert_eq!(output.additions()[0].payload().as_bytes(), &[4, 5]);
398        assert_eq!(budget.reserved_bytes(), 5);
399        assert_eq!(budget.reserved_objects(), 1);
400        assert!(staging.next(&zencodec::how_far::NoPulse).unwrap().is_none());
401        drop(staging);
402        assert_eq!(output.bytes(), &[1, 2, 3]);
403        drop(output);
404        assert_eq!(budget.reserved_bytes(), 0);
405        assert_eq!(std::fs::read(path).unwrap(), b"prefix\x01\x02\x03\x04\x05");
406    }
407    #[test]
408    fn disk_limits_reject_whole_packet_before_touching_scratch() {
409        for (limits, expected) in [
410            (DiskLimits::new(4, 1, 1), Refusal::Bytes),
411            (DiskLimits::new(5, 0, 1), Refusal::Packets),
412            (DiskLimits::new(5, 1, 0), Refusal::Additions),
413        ] {
414            let (file, path) = file();
415            let mut staging = Staging::disk(file, limits, StorageBudget::new(5, 1)).unwrap();
416            let error = staging
417                .push(packet(), &zencodec::how_far::NoPulse)
418                .unwrap_err();
419            assert_eq!(error.downcast_ref::<Refusal>(), Some(&expected));
420            assert_eq!(std::fs::read(path).unwrap(), b"prefix");
421        }
422    }
423    #[test]
424    fn replay_truncation_keeps_io_cause_and_releases_admission() {
425        let (file, path) = file();
426        let budget = StorageBudget::new(5, 1);
427        let mut staging = Staging::disk(file, DiskLimits::new(5, 1, 1), budget.clone()).unwrap();
428        staging.push(packet(), &zencodec::how_far::NoPulse).unwrap();
429        // Deliberate external interference is fault injection, not a supported
430        // scratch-sharing mode or proof of same-length mutation detection.
431        OpenOptions::new()
432            .write(true)
433            .open(path)
434            .unwrap()
435            .set_len(7)
436            .unwrap();
437        let error = staging.next(&zencodec::how_far::NoPulse).unwrap_err();
438        assert_eq!(
439            error.downcast_ref::<std::io::Error>().unwrap().kind(),
440            std::io::ErrorKind::UnexpectedEof
441        );
442        assert_eq!(budget.reserved_bytes(), 0);
443    }
444    #[test]
445    fn disk_requires_end_position_without_overwriting_existing_content() {
446        let (mut file, path) = file();
447        file.seek(SeekFrom::Start(0)).unwrap();
448        let error = Staging::disk(file, DiskLimits::new(5, 1, 1), StorageBudget::new(5, 1))
449            .err()
450            .unwrap();
451        assert_eq!(error.downcast_ref::<Refusal>(), Some(&Refusal::NotAtEnd));
452        assert_eq!(std::fs::read(path).unwrap(), b"prefix");
453    }
454}