Skip to main content

moq_net/model/
group.rs

1//! A group is a stream of frames, split into a [Producer] and [Consumer] handle.
2//!
3//! A [Producer] writes an ordered stream of frames.
4//! Frames can be written all at once ([Producer::write_frame]), or in chunks
5//! ([Producer::create_frame]).
6//!
7//! A [Consumer] reads an ordered stream of frames.
8//! The reader can be cloned, in which case each reader receives a copy of each frame. (fanout)
9//!
10//! The stream is closed with [Error] when all writers or readers are dropped.
11use crate::cache;
12use crate::frame::{self, Frame, FrameBuf};
13use crate::{Timescale, stats, track};
14use std::collections::VecDeque;
15use std::mem::MaybeUninit;
16use std::sync::Arc;
17use std::task::{Poll, ready};
18
19use crate::{Error, IntoBytes, Result, Timestamp};
20
21/// Maximum total size of frames cached in a group before old frames are evicted.
22///
23/// Doubles as the per-frame size cap: a single frame can be at most this large (a
24/// larger declared size is refused before allocating), so one maximum-size frame can
25/// fill a group's cache.
26const MAX_GROUP_CACHE: u64 = 32 * 1024 * 1024; // 32 MB
27
28/// A group contains a sequence number because they can arrive out of order.
29///
30/// You can use [track::Producer::append_group] if you just want to +1 the sequence number.
31#[derive(Clone, Copy, Debug, Hash, Eq, PartialEq, Ord, PartialOrd)]
32pub struct Info {
33	/// Per-track sequence number used to detect ordering and gaps. Higher numbers
34	/// supersede lower ones; consumers may skip late arrivals.
35	pub sequence: u64,
36}
37
38impl Info {
39	/// Create an untimed producer for this group.
40	///
41	/// Test-only: real groups are created via [`track::Producer`], which
42	/// supplies the parent track's [`track::Info`]. This helper exists for in-crate
43	/// tests that don't exercise timestamps.
44	#[cfg(test)]
45	pub(crate) fn produce(self) -> Producer {
46		Producer::new(self, track::Info::default(), Default::default())
47	}
48}
49
50impl From<usize> for Info {
51	fn from(sequence: usize) -> Self {
52		Self {
53			sequence: sequence as u64,
54		}
55	}
56}
57
58impl From<u64> for Info {
59	fn from(sequence: u64) -> Self {
60		Self { sequence }
61	}
62}
63
64impl From<u32> for Info {
65	fn from(sequence: u32) -> Self {
66		Self {
67			sequence: sequence as u64,
68		}
69	}
70}
71
72impl From<u16> for Info {
73	fn from(sequence: u16) -> Self {
74		Self {
75			sequence: sequence as u64,
76		}
77	}
78}
79
80/// The in-flight (tail) frame being written. At most one exists at a time, since a
81/// group is a single ordered stream.
82pub(crate) struct Partial {
83	timestamp: Timestamp,
84	buf: FrameBuf,
85}
86
87/// Shared group state. `pub(crate)` so [`frame`] handles can observe the abort flag
88/// while streaming a partial frame.
89#[derive(Default)]
90pub(crate) struct GroupState {
91	// Completed frames, each a contiguous payload. Evicted frames are popped from the
92	// front; `offset` tracks how many.
93	pub(crate) frames: VecDeque<Frame>,
94
95	// The single in-flight frame, if one is open.
96	pub(crate) partial: Option<Partial>,
97
98	// The number of frames evicted from the front of the group.
99	pub(crate) offset: usize,
100
101	// The total size (in bytes) of all cached frames plus any in-flight frame.
102	pub(crate) cache: u64,
103
104	// Mirrors `cache` into the track's shared cache pool, so the group's bytes count
105	// against the byte budget tracks evict toward.
106	charge: cache::Charge,
107
108	// Once finalized, the total number of frames the group will ever contain. Recorded
109	// at finish so the count outlives an abort that clears the cache.
110	pub(crate) fin: Option<usize>,
111
112	// The error that caused the group to be aborted, if any.
113	pub(crate) abort: Option<Error>,
114}
115
116impl GroupState {
117	/// Resolve the source for the frame at `index`: a completed frame (whole) or the
118	/// in-flight tail (streamed). Used by [`Consumer::poll_next_frame`].
119	fn poll_frame_source(&self, index: usize) -> Poll<Result<Option<(frame::Info, frame::Source)>>> {
120		if index < self.offset {
121			return Poll::Ready(Err(Error::Lagged));
122		}
123		let local = index - self.offset;
124		if let Some(f) = self.frames.get(local) {
125			let info = frame::Info {
126				size: f.payload.len() as u64,
127				timestamp: f.timestamp,
128			};
129			return Poll::Ready(Ok(Some((info, frame::Source::Complete(f.payload.clone())))));
130		}
131		if local == self.frames.len()
132			&& let Some(p) = &self.partial
133		{
134			let info = frame::Info {
135				size: p.buf.capacity() as u64,
136				timestamp: p.timestamp,
137			};
138			return Poll::Ready(Ok(Some((info, frame::Source::Partial(p.buf.clone())))));
139		}
140		ready!(self.poll_terminal(index))?;
141		Poll::Ready(Ok(None))
142	}
143
144	/// Resolve the group's terminal state for a reader positioned at `index`.
145	///
146	/// A finished group is still aborted once its frames are released to free memory
147	/// (aged out of the track's latency window, or evicted by the cache pool). A reader
148	/// that already consumed every frame is missing nothing, so it gets the clean end of
149	/// group; one that fell short sees the abort rather than a silently truncated stream.
150	fn poll_terminal(&self, index: usize) -> Poll<Result<()>> {
151		match (self.fin, &self.abort) {
152			(Some(total), Some(err)) if index < total => Poll::Ready(Err(err.clone())),
153			(Some(_), _) => Poll::Ready(Ok(())),
154			(None, Some(err)) => Poll::Ready(Err(err.clone())),
155			(None, None) => Poll::Pending,
156		}
157	}
158
159	fn poll_finished(&self) -> Poll<Result<u64>> {
160		// The count is recorded at finish, so a later abort that cleared the cache
161		// doesn't turn a complete group into an error.
162		if let Some(total) = self.fin {
163			Poll::Ready(Ok(total as u64))
164		} else if let Some(err) = &self.abort {
165			Poll::Ready(Err(err.clone()))
166		} else {
167			Poll::Pending
168		}
169	}
170
171	/// Evict completed frames from the front until within the byte budget.
172	fn evict(&mut self) {
173		while self.cache > MAX_GROUP_CACHE {
174			let Some(frame) = self.frames.pop_front() else {
175				break;
176			};
177			let size = frame.payload.len() as u64;
178			self.cache -= size;
179			self.charge.sub(size);
180			self.offset += 1;
181		}
182	}
183
184	/// Drop the cached frames (and any in-flight tail) and release their pool charge.
185	fn release(&mut self) {
186		self.frames.clear();
187		self.partial = None;
188		self.cache = 0;
189		self.charge.clear();
190	}
191}
192
193fn modify(state: &kio::Producer<GroupState>) -> Result<kio::Mut<'_, GroupState>> {
194	state.write().map_err(|r| r.abort.clone().unwrap_or(Error::Dropped))
195}
196
197/// Writes frames to a group in order.
198///
199/// Each group is delivered independently over a QUIC stream.
200/// Use [Self::write_frame] for simple single-buffer frames,
201/// or [Self::create_frame] for multi-chunk streaming writes.
202pub struct Producer {
203	// Mutable stream state.
204	state: kio::Producer<GroupState>,
205
206	// The group header containing the sequence number. A small `Copy` value,
207	// inherited by each frame (see [`Self::create_frame`]).
208	info: Info,
209
210	// The parent track's properties, inherited rather than passed piecemeal. Its
211	// `timescale` is used by [`Self::create_frame`] to normalize every frame's
212	// timestamp into the track scale before it enters the stream. Threaded down by
213	// value from [`track::Producer::create_group`] / `append_group`.
214	track: track::Info,
215
216	// The parent track's account against the shared cache pool. Held here as well as
217	// in the group's `cache::Charge` so a frame write can settle the track's eviction
218	// debt with the group lock released.
219	cache: Arc<cache::Track>,
220
221	// Ingress payload meter, set by a tagged [`track::Producer`] via
222	// [`Self::with_meter`]. Empty (no-op) for an untagged group.
223	stats: stats::Meter,
224
225	// Shared by every clone: its `Drop` is the abrupt-teardown, running exactly once
226	// when the last of them goes.
227	alive: Arc<Alive>,
228}
229
230/// Ends the group when the last [`Producer`] clone drops, including the clone the
231/// parent track holds in its cache.
232///
233/// A refcount rather than a "am I the last one?" check inside `Drop`: that answer is
234/// a snapshot, and acting on it is exactly what can invalidate it. Holding a producer
235/// of its own also keeps the state writable until the teardown has run, whatever order
236/// the last owner's fields drop in.
237struct Alive {
238	info: Info,
239	state: kio::Producer<GroupState>,
240}
241
242impl Drop for Alive {
243	fn drop(&mut self) {
244		// See track::Alive: the last producer dropping without a clean finish releases
245		// the cached frames so a stale consumer can't pin their buffers forever. A
246		// finished group keeps its cache so consumers can drain.
247		if let Ok(mut state) = modify(&self.state)
248			&& state.fin.is_none()
249		{
250			// Dropped without finish() or abort(), so consumers will see
251			// Error::Dropped mid-group. Deliberate ends go through finish()/abort().
252			tracing::warn!(
253				sequence = self.info.sequence,
254				"group::Producer dropped without finish() or abort()"
255			);
256			state.release();
257		}
258	}
259}
260
261impl std::ops::Deref for Producer {
262	type Target = Info;
263
264	fn deref(&self) -> &Self::Target {
265		&self.info
266	}
267}
268
269impl Producer {
270	/// Create a group producer bound to its parent track's [`track::Info`] and cache
271	/// account.
272	///
273	/// Crate-private: groups are only constructed via [`track::Producer`], which
274	/// threads both down so properties like the timescale are inherited rather than
275	/// passed in. Every frame added to this group is normalized to the track's
276	/// timescale by [`Self::create_frame`].
277	///
278	/// Charges the group into `cache`, so its cached bytes count against the budget the
279	/// track evicts toward under memory pressure.
280	pub(crate) fn new(info: Info, track: track::Info, cache: Arc<cache::Track>) -> Self {
281		let state = kio::Producer::<GroupState>::default();
282		state.write().ok().expect("a new group is open").charge = cache.charge();
283		let alive = Arc::new(Alive {
284			info,
285			state: state.clone(),
286		});
287		Self {
288			info,
289			state,
290			track,
291			cache,
292			stats: stats::Meter::default(),
293			alive,
294		}
295	}
296
297	/// Attach an ingress payload meter, counting this as one delivered group.
298	/// Called by a tagged [`track::Producer`] when it creates the group.
299	pub(crate) fn with_meter(mut self, meter: stats::Meter) -> Self {
300		meter.group();
301		self.stats = meter;
302		self
303	}
304
305	/// The group header.
306	pub(crate) fn info(&self) -> Info {
307		self.info
308	}
309
310	/// The parent track's timescale.
311	pub fn timescale(&self) -> Timescale {
312		self.track.timescale
313	}
314
315	/// A helper method to write a frame from a single byte buffer.
316	///
317	/// If you want to write multiple chunks, use [Self::create_frame] to get a frame producer.
318	/// But an upfront size is required.
319	///
320	/// `timestamp` is converted into the parent track's timescale. For data without
321	/// a presentation time, pass [`Timestamp::now`] explicitly.
322	pub fn write_frame<B: IntoBytes>(&mut self, timestamp: Timestamp, data: B) -> Result<()> {
323		let timestamp = timestamp
324			.convert(self.track.timescale)
325			.map_err(|_| Error::TimestampMismatch)?;
326		let payload = data.into_bytes();
327		if payload.len() as u64 > MAX_GROUP_CACHE {
328			return Err(Error::FrameTooLarge);
329		}
330
331		let mut state = modify(&self.state)?;
332		if state.fin.is_some() {
333			return Err(Error::Closed);
334		}
335		debug_assert!(state.partial.is_none(), "a frame is already open");
336		let size = payload.len() as u64;
337		state.cache += size;
338		state.charge.add(size);
339		state.frames.push_back(Frame { timestamp, payload });
340		state.evict();
341		drop(state);
342
343		// With the group lock released (lock order is track then group), settle
344		// eviction debt if enough has been written since the track last paid.
345		self.cache.settle();
346
347		// Ingress payload: one whole frame written.
348		self.stats.frames(1);
349		self.stats.bytes(size);
350		Ok(())
351	}
352
353	/// Create a frame with an upfront size and presentation timestamp, streamed in
354	/// chunks. Borrows the group exclusively until the returned [`frame::Producer`]
355	/// is finished or dropped, so only one frame is open at a time.
356	///
357	/// The `timestamp` is converted into the parent track's timescale, so the scale you
358	/// build it with doesn't have to match the track. Returns [`Error::FrameTooLarge`]
359	/// if the declared size exceeds the group's byte budget (refused before allocating)
360	/// or [`Error::TimestampMismatch`] if the timestamp can't be converted (overflow).
361	pub fn create_frame(&mut self, frame: frame::Info) -> Result<frame::Producer<'_>> {
362		let timestamp = frame
363			.timestamp
364			.convert(self.track.timescale)
365			.map_err(|_| Error::TimestampMismatch)?;
366		if frame.size > MAX_GROUP_CACHE {
367			return Err(Error::FrameTooLarge);
368		}
369		let buf = FrameBuf::new(frame.size as usize);
370
371		let mut state = modify(&self.state)?;
372		if state.fin.is_some() {
373			return Err(Error::Closed);
374		}
375		debug_assert!(state.partial.is_none(), "a frame is already open");
376		state.cache += frame.size;
377		state.charge.add(frame.size);
378		state.partial = Some(Partial {
379			timestamp,
380			buf: buf.clone(),
381		});
382		state.evict();
383		drop(state);
384
385		// With the group lock released (lock order is track then group), settle
386		// eviction debt if enough has been written since the track last paid.
387		self.cache.settle();
388
389		// Ingress payload: one frame opened; its bytes are counted per chunk as the
390		// frame::Producer writes them.
391		self.stats.frames(1);
392		let meter = self.stats.clone();
393
394		let info = frame::Info {
395			size: frame.size,
396			timestamp,
397		};
398		Ok(frame::Producer::new(self, buf, info).with_meter(meter))
399	}
400
401	/// Wake consumers parked on the group channel (called after a partial write).
402	pub(crate) fn frame_notify(&self) {
403		// Taking the write lock and dropping it triggers kio's notify.
404		let _ = self.state.write();
405	}
406
407	/// Commit the in-flight frame as a completed frame (called by [`frame::Producer::finish`]).
408	pub(crate) fn frame_commit(&mut self, frame: Frame) -> Result<()> {
409		let mut state = modify(&self.state)?;
410		// Bytes were already counted against the cache (and the pool charge) when the
411		// frame was created; committing just moves the tail into the completed set.
412		state.partial = None;
413		state.frames.push_back(frame);
414		Ok(())
415	}
416
417	/// Fail the group because an in-flight frame couldn't complete (called by
418	/// [`frame::Producer::abort`] / its drop).
419	pub(crate) fn frame_abort(&mut self, err: Error) {
420		let _ = self.clone().abort(err);
421	}
422
423	/// Return the number of frames written so far (completed plus any in-flight).
424	pub fn frame_count(&self) -> usize {
425		let state = self.state.read();
426		state.offset + state.frames.len() + state.partial.is_some() as usize
427	}
428
429	/// Mark the group as complete; no more frames will be written.
430	///
431	/// Borrows rather than consumes, so a later failure can still be reported through
432	/// [`abort`](Self::abort). The handle also keeps the cached frames readable.
433	pub fn finish(&mut self) -> Result<()> {
434		let mut state = modify(&self.state)?;
435		state.fin = Some(state.offset + state.frames.len());
436		Ok(())
437	}
438
439	/// Abort the group with the given error.
440	///
441	/// Consumes the handle. Drops the cached frames so a stale [`Consumer`] can't pin
442	/// their buffers in memory forever; consumers that haven't drained yet surface the
443	/// abort error instead of the leftover cache.
444	pub fn abort(self, err: Error) -> Result<()> {
445		let mut guard = modify(&self.state)?;
446		guard.abort = Some(err);
447		guard.release();
448		guard.close();
449		Ok(())
450	}
451
452	/// Whether the group has been aborted (including pool eviction). The track's
453	/// read paths treat an aborted cached group as absent.
454	pub(crate) fn is_aborted(&self) -> bool {
455		self.state.read().abort.is_some()
456	}
457
458	/// The group's full cached footprint (payload plus fixed overhead), used by the
459	/// track to size this group as an eviction victim.
460	pub(crate) fn cache_size(&self) -> u64 {
461		self.state.read().charge.size()
462	}
463
464	/// Tick of the group's last cache access, driving eviction protection and age
465	/// expiry (see [`cache::Pool::average`]).
466	pub(crate) fn cache_accessed(&self) -> u64 {
467		self.state.read().charge.accessed()
468	}
469
470	/// Enter the group into the evictable population: demoted from the live edge,
471	/// or inserted behind it. Idempotent; a no-op once the group is closed.
472	pub(crate) fn cache_demote(&self) {
473		if let Ok(mut state) = self.state.write() {
474			state.charge.demote();
475		}
476	}
477
478	/// Record a cache access (a FETCH hit, or a fetched backfill's birth),
479	/// protecting the group from eviction and restarting its expiry clock. A no-op
480	/// once the group is closed.
481	pub(crate) fn cache_refresh(&self) {
482		if let Ok(mut state) = self.state.write() {
483			state.charge.refresh();
484		}
485	}
486
487	/// Create a new consumer for the group.
488	pub fn consume(&self) -> Consumer {
489		Consumer {
490			info: self.info,
491			state: self.state.consume(),
492			track: self.track.clone(),
493			index: 0,
494			prefetch: Prefetch::default(),
495			// Untagged: a tagged track attaches the egress meter via `with_meter`
496			// when it hands the consumer to a subscriber/fetch.
497			stats: stats::Meter::default(),
498		}
499	}
500
501	/// Block until the group is closed or aborted.
502	pub async fn closed(&self) -> Error {
503		kio::wait(|waiter| self.poll_closed(waiter)).await
504	}
505
506	/// Poll until the group is closed or aborted; ready with the cause.
507	pub fn poll_closed(&self, waiter: &kio::Waiter) -> Poll<Error> {
508		self.state.poll_closed(waiter).map(|()| self.abort_reason())
509	}
510
511	/// Block until there are no active consumers.
512	pub async fn unused(&self) -> Result<()> {
513		self.state.unused().await.map_err(|_| self.abort_reason())
514	}
515
516	/// The recorded abort reason, or [`Error::Dropped`] if the group closed without one.
517	fn abort_reason(&self) -> Error {
518		self.state.read().abort.clone().unwrap_or(Error::Dropped)
519	}
520}
521
522impl Clone for Producer {
523	fn clone(&self) -> Self {
524		Self {
525			info: self.info,
526			state: self.state.clone(),
527			track: self.track.clone(),
528			cache: self.cache.clone(),
529			stats: self.stats.clone(),
530			alive: self.alive.clone(),
531		}
532	}
533}
534
535/// A small inline batch of completed frames, drained from the shared group state
536/// under one lock and then handed out without re-locking.
537///
538/// Each [`Consumer::read_frame`] otherwise takes the group mutex and allocates a
539/// waker just to clone one `Bytes`; draining a batch amortizes both across `CAP`
540/// frames. Storage is inline and uninitialized (no heap), so a consumer that never
541/// reads whole frames, or drains through a higher-level buffer, pays nothing.
542struct Prefetch {
543	// Initialized, not-yet-taken frames are `frames[pos..len]`; the rest are uninitialized.
544	frames: [MaybeUninit<Frame>; Self::CAP],
545	pos: usize,
546	len: usize,
547}
548
549impl Prefetch {
550	const CAP: usize = 8;
551
552	/// Take the next buffered frame, or `None` if the batch is drained.
553	fn pop(&mut self) -> Option<Frame> {
554		if self.pos == self.len {
555			return None;
556		}
557		// SAFETY: `pos < len`, so this slot was written by `fill` and not yet taken.
558		let frame = unsafe { self.frames[self.pos].assume_init_read() };
559		self.pos += 1;
560		Some(frame)
561	}
562
563	/// Refill with up to `CAP` frames. Must be drained first (`pop` returned `None`).
564	fn fill(&mut self, frames: impl Iterator<Item = Frame>) {
565		debug_assert_eq!(self.pos, self.len, "fill on a non-empty batch would leak frames");
566		self.pos = 0;
567		self.len = 0;
568		for frame in frames.take(Self::CAP) {
569			self.frames[self.len].write(frame);
570			self.len += 1;
571		}
572	}
573
574	/// `(frame count, total payload bytes)` of the buffered, not-yet-taken frames.
575	/// Read once per fill to bump the egress payload counters for the whole batch.
576	fn buffered(&self) -> (u64, u64) {
577		let mut bytes = 0u64;
578		for slot in &self.frames[self.pos..self.len] {
579			// SAFETY: slots in `pos..len` are initialized (written by `fill`, not yet popped).
580			bytes += unsafe { slot.assume_init_ref() }.payload.len() as u64;
581		}
582		((self.len - self.pos) as u64, bytes)
583	}
584}
585
586impl Default for Prefetch {
587	fn default() -> Self {
588		Self {
589			frames: [const { MaybeUninit::uninit() }; Self::CAP],
590			pos: 0,
591			len: 0,
592		}
593	}
594}
595
596impl Drop for Prefetch {
597	fn drop(&mut self) {
598		for slot in &mut self.frames[self.pos..self.len] {
599			// SAFETY: slots in `pos..len` are initialized and were never taken.
600			unsafe { slot.assume_init_drop() };
601		}
602	}
603}
604
605/// Consume a group, frame-by-frame.
606pub struct Consumer {
607	// Shared state with the producer.
608	state: kio::Consumer<GroupState>,
609
610	// Immutable stream state.
611	info: Info,
612
613	// The parent track's info, inherited from the producer. Its `timescale` lets the
614	// wire publisher emit per-frame timestamps at the right scale for a fetched group.
615	track: track::Info,
616
617	// The number of frames we've read.
618	// NOTE: Cloned readers inherit this offset, but then run in parallel.
619	index: usize,
620
621	// A batch of completed frames drained ahead under one lock (whole-frame reads only).
622	prefetch: Prefetch,
623
624	// Egress payload meter, set by a tagged track via [`Self::with_meter`]. Empty
625	// (no-op) for an untagged group.
626	stats: stats::Meter,
627}
628
629impl Clone for Consumer {
630	fn clone(&self) -> Self {
631		// A clone shares the channel and inherits `index`, but starts with an empty
632		// prefetch: it re-reads its batch from the shared state, in parallel.
633		Self {
634			state: self.state.clone(),
635			info: self.info,
636			track: self.track.clone(),
637			index: self.index,
638			prefetch: Prefetch::default(),
639			// Inherit the meter without re-counting the group: the original already
640			// counted it when the track handed it out.
641			stats: self.stats.clone(),
642		}
643	}
644}
645
646impl std::ops::Deref for Consumer {
647	type Target = Info;
648
649	fn deref(&self) -> &Self::Target {
650		&self.info
651	}
652}
653
654impl Consumer {
655	/// Attach an egress payload meter, counting this as one delivered group.
656	/// Called by a tagged track when it hands the consumer to a subscriber or fetch.
657	pub(crate) fn with_meter(mut self, meter: stats::Meter) -> Self {
658		meter.group();
659		self.stats = meter;
660		self
661	}
662
663	/// The parent track's timescale.
664	pub fn timescale(&self) -> Timescale {
665		self.track.timescale
666	}
667
668	// A helper to automatically apply Dropped if the state is closed without an error.
669	fn poll<F, R>(&self, waiter: &kio::Waiter, f: F) -> Poll<Result<R>>
670	where
671		F: Fn(&kio::Ref<'_, GroupState>) -> Poll<Result<R>>,
672	{
673		Poll::Ready(match ready!(self.state.poll(waiter, f)) {
674			Ok(res) => res,
675			// We try to clone abort just in case the function forgot to check for terminal state.
676			Err(state) => Err(state.abort.clone().unwrap_or(Error::Dropped)),
677		})
678	}
679
680	/// Return a consumer for the next frame for chunked reading.
681	pub async fn next_frame(&mut self) -> Result<Option<frame::Consumer>> {
682		kio::wait(|waiter| self.poll_next_frame(waiter)).await
683	}
684
685	/// Poll for the next frame, without blocking.
686	///
687	/// Returns None if the group is finished and the index is out of range.
688	pub fn poll_next_frame(&mut self, waiter: &kio::Waiter) -> Poll<Result<Option<frame::Consumer>>> {
689		// Hand out any frames a prior read_frame prefetched before touching the tail.
690		// Their bytes were already counted at the batch fill, so the frame::Consumer
691		// carries no meter.
692		if let Some(frame) = self.prefetch.pop() {
693			self.index += 1;
694			let info = frame::Info {
695				size: frame.payload.len() as u64,
696				timestamp: frame.timestamp,
697			};
698			let source = frame::Source::Complete(frame.payload);
699			return Poll::Ready(Ok(Some(frame::Consumer::new(self.state.clone(), info, source))));
700		}
701
702		let index = self.index;
703		let Some((info, source)) = ready!(self.poll(waiter, |state| state.poll_frame_source(index))?) else {
704			return Poll::Ready(Ok(None));
705		};
706
707		self.index += 1;
708		// A direct read (not prefetched): count the frame here; the frame::Consumer
709		// counts its bytes per chunk as they're read out.
710		self.stats.frames(1);
711		Poll::Ready(Ok(Some(
712			frame::Consumer::new(self.state.clone(), info, source).with_meter(self.stats.clone()),
713		)))
714	}
715
716	/// Read the next frame (timestamp and payload) all at once, without blocking.
717	pub fn poll_read_frame(&mut self, waiter: &kio::Waiter) -> Poll<Result<Option<frame::Frame>>> {
718		// Fast path: serve from the prefetched batch without locking or allocating a waker.
719		if let Some(frame) = self.prefetch.pop() {
720			self.index += 1;
721			return Poll::Ready(Ok(Some(frame)));
722		}
723
724		// The batch is drained: refill it under a single lock, registering the waiter if
725		// nothing is ready. Borrow the two fields disjointly so the closure can fill.
726		let index = self.index;
727		let prefetch = &mut self.prefetch;
728		let res = self.state.poll(waiter, |state| {
729			if index < state.offset {
730				return Poll::Ready(Err(Error::Lagged));
731			}
732			// `local` can run past the buffered count when frames were cleared or evicted out
733			// from under us (abort, unfinished drop, an eviction gap); clamp so `range` never
734			// panics on an out-of-bounds start. `fill` always resets the batch, so an empty
735			// range leaves `len == 0` and the terminal checks below resolve abort/fin/pending.
736			let local = (index - state.offset).min(state.frames.len());
737			prefetch.fill(state.frames.range(local..).cloned());
738			if prefetch.len > 0 {
739				return Poll::Ready(Ok(()));
740			}
741			// Nothing completed at `index`: an in-flight tail waits, otherwise resolve
742			// the terminal state (whole-frame reads never stream the partial).
743			state.poll_terminal(index)
744		});
745
746		match ready!(res) {
747			Ok(Ok(())) => {}
748			Ok(Err(err)) => return Poll::Ready(Err(err)),
749			Err(state) => return Poll::Ready(Err(state.abort.clone().unwrap_or(Error::Dropped))),
750		}
751
752		// A fresh batch was just filled (empty only on a clean end). Count the whole
753		// batch once here, under no lock, so the drained pops that follow stay free.
754		let (frames, bytes) = self.prefetch.buffered();
755		self.stats.frames(frames);
756		self.stats.bytes(bytes);
757
758		Poll::Ready(Ok(self.prefetch.pop().inspect(|_| {
759			self.index += 1;
760		})))
761	}
762
763	/// Read the next frame (timestamp and payload) all at once.
764	pub async fn read_frame(&mut self) -> Result<Option<frame::Frame>> {
765		// Serve from the prefetched batch without building a future or allocating a waker.
766		if let Some(frame) = self.prefetch.pop() {
767			self.index += 1;
768			return Ok(Some(frame));
769		}
770		kio::wait(|waiter| self.poll_read_frame(waiter)).await
771	}
772
773	/// Poll for the final number of frames in the group.
774	pub fn poll_finished(&mut self, waiter: &kio::Waiter) -> Poll<Result<u64>> {
775		self.poll(waiter, |state| state.poll_finished())
776	}
777
778	/// Block until the group is finished, returning the number of frames in the group.
779	pub async fn finished(&mut self) -> Result<u64> {
780		kio::wait(|waiter| self.poll_finished(waiter)).await
781	}
782}
783
784/// Options for a one-shot [`track::Consumer::fetch_group`] of a past group.
785#[derive(Clone, Debug, Default)]
786#[non_exhaustive]
787pub struct Fetch {
788	/// Delivery priority for the fetched group's stream. Defaults to 0.
789	pub priority: u8,
790}
791
792impl Fetch {
793	/// Set the delivery priority, returning `self` for chaining.
794	pub fn with_priority(mut self, priority: u8) -> Self {
795		self.priority = priority;
796		self
797	}
798}
799
800#[cfg(test)]
801mod test {
802	use super::*;
803	use bytes::Bytes;
804	use futures::FutureExt;
805
806	#[test]
807	fn basic_frame_reading() {
808		let mut producer = Info { sequence: 0 }.produce();
809		producer
810			.write_frame(Timestamp::ZERO, Bytes::from_static(b"frame0"))
811			.unwrap();
812		producer
813			.write_frame(Timestamp::ZERO, Bytes::from_static(b"frame1"))
814			.unwrap();
815		producer.finish().unwrap();
816
817		let mut consumer = producer.consume();
818		let f0 = consumer.next_frame().now_or_never().unwrap().unwrap().unwrap();
819		assert_eq!(f0.size, 6);
820		let f1 = consumer.next_frame().now_or_never().unwrap().unwrap().unwrap();
821		assert_eq!(f1.size, 6);
822		let end = consumer.next_frame().now_or_never().unwrap().unwrap();
823		assert!(end.is_none());
824	}
825
826	#[test]
827	fn read_frame_all_at_once() {
828		let mut producer = Info { sequence: 0 }.produce();
829		producer
830			.write_frame(Timestamp::ZERO, Bytes::from_static(b"hello"))
831			.unwrap();
832		producer.finish().unwrap();
833
834		let mut consumer = producer.consume();
835		let frame = consumer.read_frame().now_or_never().unwrap().unwrap().unwrap();
836		assert_eq!(frame.payload, Bytes::from_static(b"hello"));
837	}
838
839	#[test]
840	fn read_frame_preserves_timestamp() {
841		let mut producer = Info { sequence: 0 }.produce();
842		let timestamp = Timestamp::from_micros(20_000).unwrap();
843		producer.write_frame(timestamp, Bytes::from_static(b"hello")).unwrap();
844		producer.finish().unwrap();
845
846		let mut consumer = producer.consume();
847		let frame = consumer.read_frame().now_or_never().unwrap().unwrap().unwrap();
848		assert_eq!(frame.timestamp.as_micros(), 20_000);
849		assert_eq!(frame.payload, Bytes::from_static(b"hello"));
850	}
851
852	#[test]
853	fn chunked_frame_reads_whole() {
854		let mut producer = Info { sequence: 0 }.produce();
855		{
856			let mut frame = producer
857				.create_frame(frame::Info {
858					size: 10,
859					timestamp: Timestamp::ZERO,
860				})
861				.unwrap();
862			frame.write(Bytes::from_static(b"hello")).unwrap();
863			frame.write(Bytes::from_static(b"world")).unwrap();
864			frame.finish().unwrap();
865		}
866		producer.finish().unwrap();
867
868		// Frame data is held in a single per-frame buffer; a whole-frame read returns
869		// the full contents in one slice.
870		let mut consumer = producer.consume();
871		let frame = consumer.read_frame().now_or_never().unwrap().unwrap().unwrap();
872		assert_eq!(frame.payload, Bytes::from_static(b"helloworld"));
873	}
874
875	#[test]
876	fn chunked_frame_streams_partial() {
877		let mut producer = Info { sequence: 0 }.produce();
878		let mut consumer = producer.consume();
879
880		let mut frame = producer
881			.create_frame(frame::Info {
882				size: 6,
883				timestamp: Timestamp::ZERO,
884			})
885			.unwrap();
886		frame.write(Bytes::from_static(b"foo")).unwrap();
887
888		// A consumer can stream the in-flight tail before it's finished.
889		let mut f = consumer.next_frame().now_or_never().unwrap().unwrap().unwrap();
890		let c1 = f.read_chunk().now_or_never().unwrap().unwrap();
891		assert_eq!(c1, Some(Bytes::from_static(b"foo")));
892		assert!(f.read_chunk().now_or_never().is_none());
893
894		frame.write(Bytes::from_static(b"bar")).unwrap();
895		frame.finish().unwrap();
896
897		let c2 = f.read_chunk().now_or_never().unwrap().unwrap();
898		assert_eq!(c2, Some(Bytes::from_static(b"bar")));
899		let c3 = f.read_chunk().now_or_never().unwrap().unwrap();
900		assert_eq!(c3, None);
901	}
902
903	#[test]
904	fn group_finish_returns_none() {
905		let mut producer = Info { sequence: 0 }.produce();
906		producer.finish().unwrap();
907
908		let mut consumer = producer.consume();
909		let end = consumer.next_frame().now_or_never().unwrap().unwrap();
910		assert!(end.is_none());
911	}
912
913	#[test]
914	fn abort_propagates() {
915		let producer = Info { sequence: 0 }.produce();
916		let mut consumer = producer.consume();
917		producer.abort(crate::Error::Cancel).unwrap();
918
919		let result = consumer.next_frame().now_or_never().unwrap();
920		assert!(matches!(result, Err(crate::Error::Cancel)));
921	}
922
923	#[test]
924	fn abort_clears_cached_frames() {
925		let mut producer = Info { sequence: 0 }.produce();
926		producer
927			.write_frame(Timestamp::ZERO, Bytes::from_static(b"data"))
928			.unwrap();
929
930		// A stale consumer that never reads must not pin the cached frames.
931		let _consumer = producer.consume();
932		assert_eq!(producer.state.read().frames.len(), 1);
933
934		producer.clone().abort(crate::Error::Cancel).unwrap();
935
936		let state = producer.state.read();
937		assert!(state.frames.is_empty(), "cached frames should be dropped on abort");
938		assert_eq!(state.cache, 0);
939	}
940
941	#[test]
942	fn drop_unfinished_clears_cached_frames() {
943		let producer = Info { sequence: 0 }.produce();
944		let mut writer = producer.clone();
945		writer
946			.write_frame(Timestamp::ZERO, Bytes::from_static(b"data"))
947			.unwrap();
948
949		// A stale consumer keeps the channel (and thus the cache) alive.
950		let mut consumer = producer.consume();
951		assert_eq!(producer.state.read().frames.len(), 1);
952
953		// Drop every producer without finishing: the cache is released.
954		drop(writer);
955		drop(producer);
956
957		let result = consumer.next_frame().now_or_never().unwrap();
958		assert!(matches!(result, Err(crate::Error::Dropped)));
959	}
960
961	#[test]
962	fn drop_finished_keeps_cached_frames() {
963		let mut producer = Info { sequence: 0 }.produce();
964		producer
965			.write_frame(Timestamp::ZERO, Bytes::from_static(b"data"))
966			.unwrap();
967		producer.finish().unwrap();
968
969		let mut consumer = producer.consume();
970		drop(producer);
971
972		// A cleanly finished group keeps its cache so the consumer can still drain.
973		let frame = consumer.read_frame().now_or_never().unwrap().unwrap().unwrap();
974		assert_eq!(frame.payload, Bytes::from_static(b"data"));
975	}
976
977	#[tokio::test]
978	async fn pending_then_ready() {
979		let mut producer = Info { sequence: 0 }.produce();
980		let mut consumer = producer.consume();
981
982		// Consumer blocks because no frames yet.
983		assert!(consumer.next_frame().now_or_never().is_none());
984
985		producer
986			.write_frame(Timestamp::ZERO, Bytes::from_static(b"data"))
987			.unwrap();
988		producer.finish().unwrap();
989
990		let frame = consumer.next_frame().now_or_never().unwrap().unwrap().unwrap();
991		assert_eq!(frame.size, 4);
992	}
993
994	#[test]
995	fn eviction_drops_old_frames() {
996		let mut producer = Info { sequence: 0 }.produce();
997
998		// Write frames that total more than MAX_GROUP_CACHE.
999		let big = Bytes::from(vec![0u8; MAX_GROUP_CACHE as usize]);
1000		producer.write_frame(Timestamp::ZERO, big.clone()).unwrap();
1001		producer.write_frame(Timestamp::ZERO, big).unwrap();
1002
1003		// The first frame should have been evicted (tombstoned via offset).
1004		let state = producer.state.read();
1005		assert_eq!(state.offset, 1);
1006		assert_eq!(state.frames.len(), 1);
1007		assert_eq!(state.frames[0].payload.len(), MAX_GROUP_CACHE as usize);
1008	}
1009
1010	#[test]
1011	fn next_frame_returns_cache_full_on_tombstone() {
1012		let mut producer = Info { sequence: 0 }.produce();
1013
1014		let big = Bytes::from(vec![0u8; MAX_GROUP_CACHE as usize]);
1015		producer.write_frame(Timestamp::ZERO, big.clone()).unwrap();
1016		producer.write_frame(Timestamp::ZERO, big).unwrap();
1017
1018		let mut consumer = producer.consume();
1019		// First frame was evicted, next_frame should return Lagged.
1020		let result = consumer.next_frame().now_or_never().unwrap();
1021		assert!(matches!(result, Err(crate::Error::Lagged)));
1022	}
1023
1024	#[test]
1025	fn no_eviction_under_budget() {
1026		let mut producer = Info { sequence: 0 }.produce();
1027		// Many small frames stay cached: there is no frame-count cap, only a byte budget.
1028		for _ in 0..100_000 {
1029			producer.write_frame(Timestamp::ZERO, Bytes::from_static(b"x")).unwrap();
1030		}
1031		producer.finish().unwrap();
1032
1033		let state = producer.state.read();
1034		assert_eq!(state.offset, 0);
1035		assert_eq!(state.frames.len(), 100_000);
1036	}
1037
1038	#[test]
1039	fn clone_consumer_independent() {
1040		let mut producer = Info { sequence: 0 }.produce();
1041		producer.write_frame(Timestamp::ZERO, Bytes::from_static(b"a")).unwrap();
1042
1043		let mut c1 = producer.consume();
1044		// Read one frame from c1
1045		let _ = c1.next_frame().now_or_never().unwrap().unwrap().unwrap();
1046
1047		// Clone c1, inheriting its index (past first frame).
1048		let mut c2 = c1.clone();
1049
1050		producer.write_frame(Timestamp::ZERO, Bytes::from_static(b"b")).unwrap();
1051		producer.finish().unwrap();
1052
1053		// c2 should get the second frame (inherited index)
1054		let f = c2.next_frame().now_or_never().unwrap().unwrap().unwrap();
1055		assert_eq!(f.size, 1); // "b"
1056
1057		let end = c2.next_frame().now_or_never().unwrap().unwrap();
1058		assert!(end.is_none());
1059	}
1060
1061	/// Reading more than one prefetch batch drains every frame in order across the
1062	/// batch boundary (the refill starts exactly where the previous batch ended).
1063	#[test]
1064	fn read_frame_crosses_prefetch_batches() {
1065		let n = Prefetch::CAP * 3 + 5;
1066		let mut producer = Info { sequence: 0 }.produce();
1067		for i in 0..n {
1068			producer
1069				.write_frame(Timestamp::ZERO, Bytes::from(vec![i as u8; 4]))
1070				.unwrap();
1071		}
1072		producer.finish().unwrap();
1073
1074		let mut consumer = producer.consume();
1075		for i in 0..n {
1076			let frame = consumer.read_frame().now_or_never().unwrap().unwrap().unwrap();
1077			assert_eq!(frame.payload, Bytes::from(vec![i as u8; 4]));
1078		}
1079		assert!(consumer.read_frame().now_or_never().unwrap().unwrap().is_none());
1080	}
1081
1082	/// A finished group is still aborted once its frames are released to free memory (the
1083	/// track's latency window, or the cache pool). A reader that already drained every frame
1084	/// is missing nothing, so it must see the clean end of group rather than the abort.
1085	#[test]
1086	fn abort_after_finish_keeps_the_clean_end_for_a_drained_reader() {
1087		let mut producer = Info { sequence: 0 }.produce();
1088		producer
1089			.write_frame(Timestamp::ZERO, Bytes::from_static(b"hello"))
1090			.unwrap();
1091		producer.finish().unwrap();
1092
1093		let mut drained = producer.consume();
1094		let mut behind = producer.consume();
1095		let frame = drained.read_frame().now_or_never().unwrap().unwrap().unwrap();
1096		assert_eq!(frame.payload, Bytes::from_static(b"hello"));
1097
1098		producer.abort(Error::Old).unwrap();
1099
1100		// Drained everything before the abort: nothing is missing.
1101		assert!(drained.read_frame().now_or_never().unwrap().unwrap().is_none());
1102		assert!(drained.next_frame().now_or_never().unwrap().unwrap().is_none());
1103
1104		// Never read the frame, and its bytes are gone: a truncated stream, not a clean end.
1105		assert!(matches!(behind.read_frame().now_or_never().unwrap(), Err(Error::Old)));
1106	}
1107
1108	/// The frame count is fixed at finish, so an abort that clears the cache can't turn a
1109	/// complete group into an error.
1110	#[test]
1111	fn finished_survives_a_later_abort() {
1112		let mut producer = Info { sequence: 0 }.produce();
1113		producer.write_frame(Timestamp::ZERO, Bytes::from_static(b"a")).unwrap();
1114		producer.write_frame(Timestamp::ZERO, Bytes::from_static(b"b")).unwrap();
1115		producer.finish().unwrap();
1116
1117		let mut consumer = producer.consume();
1118		producer.abort(Error::Old).unwrap();
1119
1120		assert_eq!(consumer.finished().now_or_never().unwrap().unwrap(), 2);
1121	}
1122
1123	/// `next_frame` drains frames a prior `read_frame` prefetched, preserving order.
1124	#[test]
1125	fn interleave_read_and_next_frame() {
1126		let mut producer = Info { sequence: 0 }.produce();
1127		for i in 0..5u8 {
1128			producer.write_frame(Timestamp::ZERO, Bytes::from(vec![i; 1])).unwrap();
1129		}
1130		producer.finish().unwrap();
1131
1132		let mut consumer = producer.consume();
1133		// The first whole-frame read prefetches all five frames into the batch.
1134		let f0 = consumer.read_frame().now_or_never().unwrap().unwrap().unwrap();
1135		assert_eq!(f0.payload, Bytes::from(vec![0u8; 1]));
1136
1137		// next_frame must continue from the batch, not skip ahead or repeat.
1138		for i in 1..5u8 {
1139			let mut f = consumer.next_frame().now_or_never().unwrap().unwrap().unwrap();
1140			let data = f.read_all().now_or_never().unwrap().unwrap();
1141			assert_eq!(data, Bytes::from(vec![i; 1]));
1142		}
1143		assert!(consumer.next_frame().now_or_never().unwrap().unwrap().is_none());
1144	}
1145
1146	/// A `read_frame` whose index sits past the buffered frames (cleared by an abort, or an
1147	/// eviction gap) must surface the error, not panic on an out-of-range `range(local..)`.
1148	#[test]
1149	fn read_frame_past_cleared_frames_does_not_panic() {
1150		let mut producer = Info { sequence: 0 }.produce();
1151		producer.write_frame(Timestamp::ZERO, Bytes::from_static(b"a")).unwrap();
1152		producer.write_frame(Timestamp::ZERO, Bytes::from_static(b"b")).unwrap();
1153
1154		let mut consumer = producer.consume();
1155		consumer.read_frame().now_or_never().unwrap().unwrap().unwrap();
1156		consumer.read_frame().now_or_never().unwrap().unwrap().unwrap();
1157
1158		// Abort clears the cached frames but leaves the consumer's index (2) past them, so the
1159		// refill's `local` (2) exceeds `frames.len()` (0).
1160		producer.abort(Error::Cancel).unwrap();
1161
1162		let result = consumer.read_frame().now_or_never().unwrap();
1163		assert!(matches!(result, Err(Error::Cancel)), "expected Cancel, got {result:?}");
1164	}
1165
1166	/// Dropping a consumer mid-batch must drop the buffered-but-untaken frames
1167	/// (exercises the `MaybeUninit` Drop path; run under miri to catch leaks/UB).
1168	#[test]
1169	fn drop_with_partial_batch() {
1170		let mut producer = Info { sequence: 0 }.produce();
1171		for _ in 0..Prefetch::CAP {
1172			producer.write_frame(Timestamp::ZERO, Bytes::from_static(b"x")).unwrap();
1173		}
1174		producer.finish().unwrap();
1175
1176		let mut consumer = producer.consume();
1177		// Take one frame so the batch is filled but only partially drained.
1178		let _ = consumer.read_frame().now_or_never().unwrap().unwrap().unwrap();
1179		drop(consumer);
1180	}
1181
1182	/// A frame whose timestamp is at a different scale is converted to the group's
1183	/// scale by `create_frame`.
1184	#[test]
1185	fn create_frame_converts_mismatched_scale() {
1186		use crate::{Timescale, Timestamp};
1187
1188		let mut producer = Producer::new(
1189			Info { sequence: 0 },
1190			track::Info::default().with_timescale(Timescale::MICRO),
1191			Default::default(),
1192		);
1193		let frame = frame::Info {
1194			size: 3,
1195			timestamp: Timestamp::from_millis(1).unwrap(), // 1ms -> 1000µs
1196		};
1197		let writer = producer.create_frame(frame).unwrap();
1198		assert_eq!(writer.timestamp.scale(), Timescale::MICRO);
1199		assert_eq!(writer.timestamp.value(), 1000);
1200	}
1201
1202	/// An explicit current timestamp is converted to the group's scale.
1203	#[tokio::test]
1204	async fn create_frame_converts_current_timestamp() {
1205		use crate::Timescale;
1206
1207		let mut producer = Producer::new(
1208			Info { sequence: 0 },
1209			track::Info::default().with_timescale(Timescale::MICRO),
1210			Default::default(),
1211		);
1212		let writer = producer
1213			.create_frame(frame::Info {
1214				size: 3,
1215				timestamp: Timestamp::now(),
1216			})
1217			.unwrap();
1218		assert_eq!(writer.timestamp.scale(), Timescale::MICRO);
1219		assert!(!writer.timestamp.is_zero(), "local clock should be non-zero");
1220	}
1221
1222	/// The per-frame size cap (the group byte budget) is enforced before allocating.
1223	#[test]
1224	fn create_frame_rejects_oversized() {
1225		let mut producer = Info { sequence: 0 }.produce();
1226		let result = producer.create_frame(frame::Info {
1227			size: MAX_GROUP_CACHE + 1,
1228			timestamp: Timestamp::ZERO,
1229		});
1230		assert!(matches!(result, Err(Error::FrameTooLarge)));
1231	}
1232}