Skip to main content

moq_binary/stream/
consumer.rs

1//! Consuming an ordered log of binary payloads from a track.
2
3use std::task::Poll;
4
5use bytes::Bytes;
6
7use crate::Result;
8
9pub use super::Config;
10
11/// Consumes an ordered log of binary payloads from a track, yielding every one.
12///
13/// The log is a single group. That is what makes the mode lossless: rolling to a second group
14/// means the records that would have completed the first are gone, so a publisher that cannot
15/// write ends the track instead. A second group is therefore a broken publisher, and reading it
16/// would present a gap as a continuous log, so it fails with
17/// [`Error::Rolled`](crate::Error::Rolled) rather than yielding the remainder.
18///
19/// The failure does not wait for the first group to end: whatever has already arrived in it is
20/// yielded, and the read then fails rather than blocking on a group a broken publisher may never
21/// finish.
22pub struct Consumer {
23	track: moq_net::track::Subscriber,
24	group: Option<moq_net::group::Consumer>,
25	/// Whether the log's one group has been taken, so a second is a rolled log rather than the first.
26	taken: bool,
27	/// Sticky once a second group is seen: the payloads it displaced are gone, so every later read
28	/// fails too rather than reporting the rest of the log as a whole one.
29	rolled: bool,
30	/// The DEFLATE decoder for the group, `Some` while decompressing.
31	flate: Option<moq_flate::Decoder>,
32	compression: bool,
33}
34
35impl Consumer {
36	/// Create a consumer reading from the given track subscriber.
37	///
38	/// Set [`Config::compression`] to match the producer that wrote the track.
39	pub fn new(track: moq_net::track::Subscriber, config: Config) -> Self {
40		Self {
41			track,
42			group: None,
43			taken: false,
44			rolled: false,
45			flate: None,
46			compression: config.compression.is_deflate(),
47		}
48	}
49
50	/// Get the next payload, or `None` once the track ends.
51	pub async fn next(&mut self) -> Result<Option<Bytes>> {
52		kio::wait(|waiter| self.poll_next(waiter)).await
53	}
54
55	/// Poll for the next payload, without blocking.
56	pub fn poll_next(&mut self, waiter: &kio::Waiter) -> Poll<Result<Option<Bytes>>> {
57		loop {
58			if self.rolled {
59				return Poll::Ready(Err(crate::Error::Rolled));
60			}
61
62			let Some(group) = &mut self.group else {
63				// Arrival order rather than sequence order, because there is only ever one group to
64				// take and a second one has to be seen whatever its sequence. The monotonic
65				// `poll_next_group` would drop a late lower sequence, which is the very loss this
66				// has to report.
67				match self.track.poll_recv_group(waiter)? {
68					Poll::Ready(Some(_)) if self.taken => self.rolled = true,
69					Poll::Ready(Some(group)) => {
70						self.taken = true;
71						self.flate = self.compression.then(moq_flate::Decoder::new);
72						self.group = Some(group);
73					}
74					Poll::Ready(None) => return Poll::Ready(Ok(None)),
75					Poll::Pending => return Poll::Pending,
76				}
77				continue;
78			};
79
80			match group.poll_read_frame(waiter)? {
81				Poll::Ready(Some(frame)) => return Poll::Ready(self.decode(&frame.payload).map(Some)),
82				Poll::Ready(None) => {
83					// The log's one group is exhausted. Keep polling the track so a clean end still
84					// reports the log as complete, and so a second group is caught as `Rolled`.
85					self.group = None;
86				}
87				// Nothing more in the group yet, so ask the track before parking on it. A publisher
88				// that opens a second group and leaves the first open would otherwise hold this read
89				// open forever, on a log that already lost the payloads the second one displaced.
90				// Both polls register, so either source wakes this read.
91				Poll::Pending => match self.track.poll_recv_group(waiter)? {
92					Poll::Ready(Some(_)) => self.rolled = true,
93					// A finished track does not truncate the group in hand; its frames may still arrive.
94					Poll::Ready(None) | Poll::Pending => return Poll::Pending,
95				},
96			}
97		}
98	}
99
100	/// Decompress one frame, if the track is compressed.
101	fn decode(&mut self, payload: &Bytes) -> Result<Bytes> {
102		Ok(match self.flate.as_mut() {
103			Some(flate) => flate.frame(payload)?,
104			None => payload.clone(),
105		})
106	}
107}