1use 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#[derive(Clone, Copy, Debug)]
28pub struct DiskLimits {
29 pub max_bytes: u64,
31 pub max_packets: usize,
33 pub max_additions_per_packet: usize,
35}
36impl DiskLimits {
37 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#[derive(Clone, Copy, Debug, PartialEq, Eq)]
48#[non_exhaustive]
49pub enum Refusal {
50 Bytes,
52 Packets,
54 Additions,
56 NotAtEnd,
58 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}
95pub struct Staging {
102 store: Store,
103 bytes: u64,
104}
105impl Staging {
106 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 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 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 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 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}