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