moq_net/model/

group.rs

//! A group is a stream of frames, split into a [GroupProducer] and [GroupConsumer] handle.
//!
//! A [GroupProducer] writes an ordered stream of frames.
//! Frames can be written all at once, or in chunks.
//!
//! A [GroupConsumer] reads an ordered stream of frames.
//! The reader can be cloned, in which case each reader receives a copy of each frame. (fanout)
//!
//! The stream is closed with [Error] when all writers or readers are dropped.
use std::collections::VecDeque;
use std::task::{Poll, ready};

use bytes::Bytes;

use crate::{Error, MAX_FRAME_SIZE, Result};

use super::{Frame, FrameConsumer, FrameProducer};

/// Maximum total size of frames cached in a group before old frames are evicted.
///
/// Kept equal to `MAX_FRAME_SIZE` so a single maximum-size frame can fill a group's cache.
const MAX_GROUP_CACHE: u64 = 32 * 1024 * 1024; // 32 MB

/// Maximum number of frames cached in a group before old frames are evicted.
const MAX_GROUP_FRAMES: usize = 1024;

/// A group contains a sequence number because they can arrive out of order.
///
/// You can use [crate::TrackProducer::append_group] if you just want to +1 the sequence number.
#[derive(Clone, Debug, Hash, Eq, PartialEq, Ord, PartialOrd)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct Group {
	/// Per-track sequence number used to detect ordering and gaps. Higher numbers
	/// supersede lower ones; consumers may skip late arrivals.
	pub sequence: u64,
}

impl Group {
	/// Consume this [`Group`] to create a producer that owns its sequence number.
	pub fn produce(self) -> GroupProducer {
		GroupProducer::new(self)
	}
}

impl From<usize> for Group {
	fn from(sequence: usize) -> Self {
		Self {
			sequence: sequence as u64,
		}
	}
}

impl From<u64> for Group {
	fn from(sequence: u64) -> Self {
		Self { sequence }
	}
}

impl From<u32> for Group {
	fn from(sequence: u32) -> Self {
		Self {
			sequence: sequence as u64,
		}
	}
}

impl From<u16> for Group {
	fn from(sequence: u16) -> Self {
		Self {
			sequence: sequence as u64,
		}
	}
}

#[derive(Default)]
struct GroupState {
	// The frames currently cached in the group.
	// Evicted frames are popped from the front; `offset` tracks how many.
	frames: VecDeque<FrameProducer>,

	// The number of frames evicted from the front of the group.
	offset: usize,

	// The total size (in bytes) of all cached frames.
	cache: u64,

	// Whether the group has been finalized (no more frames).
	fin: bool,

	// The error that caused the group to be aborted, if any.
	abort: Option<Error>,
}

impl GroupState {
	fn poll_get_frame(&self, index: usize) -> Poll<Result<Option<FrameConsumer>>> {
		if index < self.offset {
			Poll::Ready(Err(Error::CacheFull))
		} else if let Some(frame) = self.frames.get(index - self.offset) {
			Poll::Ready(Ok(Some(frame.consume())))
		} else if self.fin {
			Poll::Ready(Ok(None))
		} else if let Some(err) = &self.abort {
			Poll::Ready(Err(err.clone()))
		} else {
			Poll::Pending
		}
	}

	/// Poll for the full payload of the frame at `index`, reading it in place.
	///
	/// Unlike [`Self::poll_get_frame`] this never mints a [`FrameConsumer`] (which
	/// would churn the frame's consumer count and wake its waiters every poll); it
	/// reads the cached [`FrameProducer`] directly. `waiter` is registered on the
	/// frame's state so the reader wakes when it finishes.
	fn poll_frame_read_all(&self, index: usize, waiter: &kio::Waiter) -> Poll<Result<Option<Bytes>>> {
		if index < self.offset {
			return Poll::Ready(Err(Error::CacheFull));
		}
		match self.frames.get(index - self.offset) {
			Some(frame) => Poll::Ready(Ok(Some(ready!(frame.poll_read_all(waiter))?))),
			None if self.fin => Poll::Ready(Ok(None)),
			None => match &self.abort {
				Some(err) => Poll::Ready(Err(err.clone())),
				None => Poll::Pending,
			},
		}
	}

	fn poll_finished(&self) -> Poll<Result<u64>> {
		if self.fin {
			Poll::Ready(Ok((self.offset + self.frames.len()) as u64))
		} else if let Some(err) = &self.abort {
			Poll::Ready(Err(err.clone()))
		} else {
			Poll::Pending
		}
	}

	/// Evict frames from the front of the group until within both limits.
	fn evict(&mut self) {
		while self.cache > MAX_GROUP_CACHE || self.frames.len() > MAX_GROUP_FRAMES {
			let Some(frame) = self.frames.pop_front() else {
				break;
			};
			self.cache -= frame.size;
			self.offset += 1;
		}
	}
}

fn modify(state: &kio::Producer<GroupState>) -> Result<kio::Mut<'_, GroupState>> {
	state.write().map_err(|r| r.abort.clone().unwrap_or(Error::Dropped))
}

/// Writes frames to a group in order.
///
/// Each group is delivered independently over a QUIC stream.
/// Use [Self::write_frame] for simple single-buffer frames,
/// or [Self::create_frame] for multi-chunk streaming writes.
pub struct GroupProducer {
	// Mutable stream state.
	state: kio::Producer<GroupState>,

	// The group header containing the sequence number.
	info: Group,
}

impl std::ops::Deref for GroupProducer {
	type Target = Group;

	fn deref(&self) -> &Self::Target {
		&self.info
	}
}

impl GroupProducer {
	/// Create a new group producer.
	pub fn new(info: Group) -> Self {
		Self {
			info,
			state: kio::Producer::default(),
		}
	}

	/// A helper method to write a frame from a single byte buffer.
	///
	/// If you want to write multiple chunks, use [Self::create_frame] to get a frame producer.
	/// But an upfront size is required.
	pub fn write_frame<B: Into<Bytes>>(&mut self, frame: B) -> Result<()> {
		let data = frame.into();
		let frame = Frame {
			size: data.len() as u64,
		};
		let mut frame = self.create_frame(frame)?;
		frame.write(data)?;
		frame.finish()?;
		Ok(())
	}

	/// Create a frame with an upfront size
	pub fn create_frame(&mut self, info: Frame) -> Result<FrameProducer> {
		// Reject before `produce()`: `FrameProducer::new` preallocates `size` bytes, so an oversized
		// frame must be caught here or it triggers the very allocation the limit exists to prevent.
		if info.size > MAX_FRAME_SIZE {
			return Err(Error::FrameTooLarge);
		}
		let frame = info.produce();
		self.append_frame(frame.clone())?;
		Ok(frame)
	}

	/// Append a frame producer to the group.
	pub fn append_frame(&mut self, frame: FrameProducer) -> Result<()> {
		// Backstop for direct callers (the buffer is already allocated by the time we hold a
		// FrameProducer); `create_frame` is the path that rejects before allocating.
		if frame.size > MAX_FRAME_SIZE {
			return Err(Error::FrameTooLarge);
		}
		let mut state = modify(&self.state)?;
		if state.fin {
			return Err(Error::Closed);
		}
		state.cache += frame.size;
		state.frames.push_back(frame);
		state.evict();
		Ok(())
	}

	/// Return the number of frames written so far.
	pub fn frame_count(&self) -> usize {
		let state = self.state.read();
		state.offset + state.frames.len()
	}

	/// Mark the group as complete; no more frames will be written.
	pub fn finish(&mut self) -> Result<()> {
		let mut state = modify(&self.state)?;
		state.fin = true;
		Ok(())
	}

	/// Abort the group with the given error.
	///
	/// No updates can be made after this point. Drops the cached frames so a stale
	/// [`GroupConsumer`] can't pin their buffers in memory forever; consumers that
	/// haven't drained yet surface the abort error instead of the leftover cache.
	/// Child frames are independent: a consumer that already pulled a
	/// [`FrameConsumer`] keeps its own handle and can finish reading it.
	pub fn abort(&mut self, err: Error) -> Result<()> {
		let mut guard = modify(&self.state)?;
		guard.abort = Some(err);
		guard.frames.clear();
		guard.cache = 0;
		guard.close();
		Ok(())
	}

	/// Create a new consumer for the group.
	pub fn consume(&self) -> GroupConsumer {
		GroupConsumer {
			info: self.info.clone(),
			state: self.state.consume(),
			index: 0,
		}
	}

	/// Block until the group is closed or aborted.
	pub async fn closed(&self) -> Error {
		self.state.closed().await;
		self.state.read().abort.clone().unwrap_or(Error::Dropped)
	}

	/// Block until there are no active consumers.
	pub async fn unused(&self) -> Result<()> {
		self.state
			.unused()
			.await
			.map_err(|r| r.abort.clone().unwrap_or(Error::Dropped))
	}
}

impl Clone for GroupProducer {
	fn clone(&self) -> Self {
		Self {
			info: self.info.clone(),
			state: self.state.clone(),
		}
	}
}

impl Drop for GroupProducer {
	fn drop(&mut self) {
		// See TrackProducer::drop: the last producer dropping without a clean
		// finish releases the cached frames so a stale consumer can't pin their
		// buffers forever. A finished group keeps its cache so consumers can drain.
		if !self.state.is_last() {
			return;
		}
		if let Ok(mut state) = modify(&self.state)
			&& !state.fin
		{
			state.frames.clear();
			state.cache = 0;
		}
	}
}

impl From<Group> for GroupProducer {
	fn from(info: Group) -> Self {
		GroupProducer::new(info)
	}
}

/// Consume a group, frame-by-frame.
#[derive(Clone)]
pub struct GroupConsumer {
	// Shared state with the producer.
	state: kio::Consumer<GroupState>,

	// Immutable stream state.
	info: Group,

	// The number of frames we've read.
	// NOTE: Cloned readers inherit this offset, but then run in parallel.
	index: usize,
}

impl std::ops::Deref for GroupConsumer {
	type Target = Group;

	fn deref(&self) -> &Self::Target {
		&self.info
	}
}

impl GroupConsumer {
	// A helper to automatically apply Dropped if the state is closed without an error.
	fn poll<F, R>(&self, waiter: &kio::Waiter, f: F) -> Poll<Result<R>>
	where
		F: Fn(&kio::Ref<'_, GroupState>) -> Poll<Result<R>>,
	{
		Poll::Ready(match ready!(self.state.poll(waiter, f)) {
			Ok(res) => res,
			// We try to clone abort just in case the function forgot to check for terminal state.
			Err(state) => Err(state.abort.clone().unwrap_or(Error::Dropped)),
		})
	}

	/// Block until the frame at the given index is available.
	///
	/// Returns None if the group is finished and the index is out of range.
	pub async fn get_frame(&self, index: usize) -> Result<Option<FrameConsumer>> {
		kio::wait(|waiter| self.poll_get_frame(waiter, index)).await
	}

	/// Poll for the frame at the given index, without blocking.
	///
	/// Returns None if the group is finished and the index is out of range.
	pub fn poll_get_frame(&self, waiter: &kio::Waiter, index: usize) -> Poll<Result<Option<FrameConsumer>>> {
		self.poll(waiter, |state| state.poll_get_frame(index))
	}

	/// Return a consumer for the next frame for chunked reading.
	pub async fn next_frame(&mut self) -> Result<Option<FrameConsumer>> {
		kio::wait(|waiter| self.poll_next_frame(waiter)).await
	}

	/// Poll for the next frame, without blocking.
	///
	/// Returns None if the group is finished and the index is out of range.
	pub fn poll_next_frame(&mut self, waiter: &kio::Waiter) -> Poll<Result<Option<FrameConsumer>>> {
		let Some(frame) = ready!(self.poll(waiter, |state| state.poll_get_frame(self.index))?) else {
			return Poll::Ready(Ok(None));
		};

		self.index += 1;
		Poll::Ready(Ok(Some(frame)))
	}

	/// Read the next frame's data all at once, without blocking.
	pub fn poll_read_frame(&mut self, waiter: &kio::Waiter) -> Poll<Result<Option<Bytes>>> {
		let index = self.index;
		let Some(data) = ready!(self.poll(waiter, |state| state.poll_frame_read_all(index, waiter))?) else {
			return Poll::Ready(Ok(None));
		};

		self.index += 1;
		Poll::Ready(Ok(Some(data)))
	}

	/// Read the next frame's data all at once.
	pub async fn read_frame(&mut self) -> Result<Option<Bytes>> {
		kio::wait(|waiter| self.poll_read_frame(waiter)).await
	}

	/// Read all of the chunks of the next frame, without blocking.
	pub fn poll_read_frame_chunks(&mut self, waiter: &kio::Waiter) -> Poll<Result<Option<Vec<Bytes>>>> {
		let index = self.index;
		let Some(data) = ready!(self.poll(waiter, |state| state.poll_frame_read_all(index, waiter))?) else {
			return Poll::Ready(Ok(None));
		};

		self.index += 1;
		// In-place reads return the whole frame as one slice; keep the chunked API
		// shape (empty payload -> no chunks).
		Poll::Ready(Ok(Some(if data.is_empty() { Vec::new() } else { vec![data] })))
	}

	/// Read all of the chunks of the next frame.
	pub async fn read_frame_chunks(&mut self) -> Result<Option<Vec<Bytes>>> {
		kio::wait(|waiter| self.poll_read_frame_chunks(waiter)).await
	}

	/// Poll for the final number of frames in the group.
	pub fn poll_finished(&mut self, waiter: &kio::Waiter) -> Poll<Result<u64>> {
		self.poll(waiter, |state| state.poll_finished())
	}

	/// Block until the group is finished, returning the number of frames in the group.
	pub async fn finished(&mut self) -> Result<u64> {
		kio::wait(|waiter| self.poll_finished(waiter)).await
	}
}

#[cfg(test)]
mod test {
	use super::*;
	use futures::FutureExt;

	#[test]
	fn basic_frame_reading() {
		let mut producer = Group { sequence: 0 }.produce();
		producer.write_frame(Bytes::from_static(b"frame0")).unwrap();
		producer.write_frame(Bytes::from_static(b"frame1")).unwrap();
		producer.finish().unwrap();

		let mut consumer = producer.consume();
		let f0 = consumer.next_frame().now_or_never().unwrap().unwrap().unwrap();
		assert_eq!(f0.size, 6);
		let f1 = consumer.next_frame().now_or_never().unwrap().unwrap().unwrap();
		assert_eq!(f1.size, 6);
		let end = consumer.next_frame().now_or_never().unwrap().unwrap();
		assert!(end.is_none());
	}

	#[test]
	fn read_frame_all_at_once() {
		let mut producer = Group { sequence: 0 }.produce();
		producer.write_frame(Bytes::from_static(b"hello")).unwrap();
		producer.finish().unwrap();

		let mut consumer = producer.consume();
		let data = consumer.read_frame().now_or_never().unwrap().unwrap().unwrap();
		assert_eq!(data, Bytes::from_static(b"hello"));
	}

	#[test]
	fn read_frame_chunks() {
		let mut producer = Group { sequence: 0 }.produce();
		let mut frame = producer.create_frame(Frame { size: 10 }).unwrap();
		frame.write(Bytes::from_static(b"hello")).unwrap();
		frame.write(Bytes::from_static(b"world")).unwrap();
		frame.finish().unwrap();
		producer.finish().unwrap();

		// Frame data is held in a single per-frame buffer; consumers see the full
		// contents in one chunk rather than the individual write boundaries.
		let mut consumer = producer.consume();
		let chunks = consumer.read_frame_chunks().now_or_never().unwrap().unwrap().unwrap();
		assert_eq!(chunks.len(), 1);
		assert_eq!(chunks[0], Bytes::from_static(b"helloworld"));
	}

	#[test]
	fn append_rejects_oversized_frame() {
		let mut producer = Group { sequence: 0 }.produce();
		let err = producer.create_frame(Frame {
			size: MAX_FRAME_SIZE + 1,
		});
		assert!(
			matches!(err, Err(Error::FrameTooLarge)),
			"a frame over the limit is rejected"
		);
		// A frame at the limit is still accepted.
		assert!(producer.create_frame(Frame { size: MAX_FRAME_SIZE }).is_ok());
	}

	#[test]
	fn get_frame_by_index() {
		let mut producer = Group { sequence: 0 }.produce();
		producer.write_frame(Bytes::from_static(b"a")).unwrap();
		producer.write_frame(Bytes::from_static(b"bb")).unwrap();
		producer.finish().unwrap();

		let consumer = producer.consume();
		let f0 = consumer.get_frame(0).now_or_never().unwrap().unwrap().unwrap();
		assert_eq!(f0.size, 1);
		let f1 = consumer.get_frame(1).now_or_never().unwrap().unwrap().unwrap();
		assert_eq!(f1.size, 2);
		let f2 = consumer.get_frame(2).now_or_never().unwrap().unwrap();
		assert!(f2.is_none());
	}

	#[test]
	fn group_finish_returns_none() {
		let mut producer = Group { sequence: 0 }.produce();
		producer.finish().unwrap();

		let mut consumer = producer.consume();
		let end = consumer.next_frame().now_or_never().unwrap().unwrap();
		assert!(end.is_none());
	}

	#[test]
	fn abort_propagates() {
		let mut producer = Group { sequence: 0 }.produce();
		let mut consumer = producer.consume();
		producer.abort(crate::Error::Cancel).unwrap();

		let result = consumer.next_frame().now_or_never().unwrap();
		assert!(matches!(result, Err(crate::Error::Cancel)));
	}

	#[test]
	fn abort_clears_cached_frames() {
		let mut producer = Group { sequence: 0 }.produce();
		producer.write_frame(Bytes::from_static(b"data")).unwrap();

		// A stale consumer that never reads must not pin the cached frames.
		let _consumer = producer.consume();
		assert_eq!(producer.state.read().frames.len(), 1);

		producer.abort(crate::Error::Cancel).unwrap();

		let state = producer.state.read();
		assert!(state.frames.is_empty(), "cached frames should be dropped on abort");
		assert_eq!(state.cache, 0);
	}

	#[test]
	fn drop_unfinished_clears_cached_frames() {
		let producer = Group { sequence: 0 }.produce();
		let mut writer = producer.clone();
		writer.write_frame(Bytes::from_static(b"data")).unwrap();

		// A stale consumer keeps the channel (and thus the cache) alive.
		let mut consumer = producer.consume();
		assert_eq!(producer.state.read().frames.len(), 1);

		// Drop every producer without finishing: the cache is released.
		drop(writer);
		drop(producer);

		let result = consumer.next_frame().now_or_never().unwrap();
		assert!(matches!(result, Err(crate::Error::Dropped)));
	}

	#[test]
	fn drop_finished_keeps_cached_frames() {
		let mut producer = Group { sequence: 0 }.produce();
		producer.write_frame(Bytes::from_static(b"data")).unwrap();
		producer.finish().unwrap();

		let mut consumer = producer.consume();
		drop(producer);

		// A cleanly finished group keeps its cache so the consumer can still drain.
		let frame = consumer.read_frame().now_or_never().unwrap().unwrap().unwrap();
		assert_eq!(frame, Bytes::from_static(b"data"));
	}

	#[tokio::test]
	async fn pending_then_ready() {
		let mut producer = Group { sequence: 0 }.produce();
		let mut consumer = producer.consume();

		// Consumer blocks because no frames yet.
		assert!(consumer.next_frame().now_or_never().is_none());

		producer.write_frame(Bytes::from_static(b"data")).unwrap();
		producer.finish().unwrap();

		let frame = consumer.next_frame().now_or_never().unwrap().unwrap().unwrap();
		assert_eq!(frame.size, 4);
	}

	#[test]
	fn eviction_drops_old_frames() {
		let mut producer = Group { sequence: 0 }.produce();

		// Write frames that total more than MAX_GROUP_CACHE.
		let big = Bytes::from(vec![0u8; MAX_GROUP_CACHE as usize]);
		producer.write_frame(big.clone()).unwrap();
		producer.write_frame(big).unwrap();

		// The first frame should have been evicted.
		let consumer = producer.consume();
		let result = consumer.get_frame(0).now_or_never().unwrap();
		assert!(matches!(result, Err(crate::Error::CacheFull)));

		// The second frame should still be available.
		let f1 = consumer.get_frame(1).now_or_never().unwrap().unwrap().unwrap();
		assert_eq!(f1.size, MAX_GROUP_CACHE);
	}

	#[test]
	fn no_eviction_under_limit() {
		let mut producer = Group { sequence: 0 }.produce();
		producer.write_frame(Bytes::from_static(b"small")).unwrap();
		producer.write_frame(Bytes::from_static(b"frames")).unwrap();
		producer.finish().unwrap();

		let consumer = producer.consume();
		let f0 = consumer.get_frame(0).now_or_never().unwrap().unwrap().unwrap();
		assert_eq!(f0.size, 5);
		let f1 = consumer.get_frame(1).now_or_never().unwrap().unwrap().unwrap();
		assert_eq!(f1.size, 6);
	}

	#[test]
	fn eviction_by_frame_count() {
		let mut producer = Group { sequence: 0 }.produce();

		// Write more than MAX_GROUP_FRAMES frames.
		for _ in 0..=MAX_GROUP_FRAMES {
			producer.write_frame(Bytes::from_static(b"x")).unwrap();
		}

		// The first frame should have been evicted.
		let consumer = producer.consume();
		let result = consumer.get_frame(0).now_or_never().unwrap();
		assert!(matches!(result, Err(crate::Error::CacheFull)));

		// The last frame should still be available.
		let f = consumer
			.get_frame(MAX_GROUP_FRAMES)
			.now_or_never()
			.unwrap()
			.unwrap()
			.unwrap();
		assert_eq!(f.size, 1);
	}

	#[test]
	fn next_frame_returns_cache_full_on_tombstone() {
		let mut producer = Group { sequence: 0 }.produce();

		let big = Bytes::from(vec![0u8; MAX_GROUP_CACHE as usize]);
		producer.write_frame(big.clone()).unwrap();
		producer.write_frame(big).unwrap();

		let mut consumer = producer.consume();
		// First frame was evicted, next_frame should return CacheFull.
		let result = consumer.next_frame().now_or_never().unwrap();
		assert!(matches!(result, Err(crate::Error::CacheFull)));
	}

	#[test]
	fn clone_consumer_independent() {
		let mut producer = Group { sequence: 0 }.produce();
		producer.write_frame(Bytes::from_static(b"a")).unwrap();

		let mut c1 = producer.consume();
		// Read one frame from c1
		let _ = c1.next_frame().now_or_never().unwrap().unwrap().unwrap();

		// Clone c1 — inherits index (past first frame)
		let mut c2 = c1.clone();

		producer.write_frame(Bytes::from_static(b"b")).unwrap();
		producer.finish().unwrap();

		// c2 should get the second frame (inherited index)
		let f = c2.next_frame().now_or_never().unwrap().unwrap().unwrap();
		assert_eq!(f.size, 1); // "b"

		let end = c2.next_frame().now_or_never().unwrap().unwrap();
		assert!(end.is_none());
	}
}