Skip to main content

citadel_sdk/media/
receiver.rs

1use super::eos::EosTracker;
2use super::transport::{BoxedSource, MediaTransportKind};
3use citadel_io::time::{sleep_until, Duration, Instant};
4use citadel_io::ErrorCode;
5use citadel_media::{
6    ControlMessage, JitterBuffer, MediaConfig, MediaFrame, MediaInstant, MediaStats,
7    MediaTrackDescriptor, PopResult, PushResult, ReassembleOutcome, Reassembler, TrackId,
8    FIRST_SEQUENCE,
9};
10use citadel_proto::prelude::{NetworkError, SecBuffer};
11use futures::stream::{select_all, SelectAll};
12use futures::StreamExt;
13use std::collections::VecDeque;
14
15/// What [`MediaReceiver::next_event`] yields.
16#[derive(Debug, Clone, PartialEq, Eq)]
17pub enum MediaEvent {
18    /// The peer announced (or accepted) these tracks.
19    Tracks(Vec<MediaTrackDescriptor>),
20    /// A frame in sequence order for its track.
21    Frame(MediaFrame),
22    /// Frames `missing_from..=missing_to` on `track` were skipped; the next
23    /// event is the frame that follows the gap.
24    Gap {
25        track: TrackId,
26        missing_from: u32,
27        missing_to: u32,
28    },
29    /// The peer finished `track`.
30    EndOfStream(TrackId),
31    /// A transport stream ended; no further events will follow.
32    Closed,
33}
34
35enum Input {
36    Datagram(Option<SecBuffer>),
37    Deadline,
38}
39
40/// Inbound half of a media endpoint.
41///
42/// Owns the transport receive halves. In unreliable mode that includes the UDP
43/// receive half, so dropping this receiver signals `DisconnectUDP` to the peer.
44pub struct MediaReceiver {
45    kind: MediaTransportKind,
46    /// Media and control lanes merged; ends only once every lane has ended.
47    lanes: SelectAll<BoxedSource>,
48    reassembler: Reassembler,
49    jitter: JitterBuffer,
50    /// Events decoded ahead of delivery (frame following a gap, control).
51    ready: VecDeque<MediaEvent>,
52    /// End-of-stream announcements waiting for their track to fully drain.
53    eos: EosTracker,
54    start: Instant,
55    closed: bool,
56    stats: MediaStats,
57}
58
59impl std::fmt::Debug for MediaReceiver {
60    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
61        f.debug_struct("MediaReceiver")
62            .field("kind", &self.kind)
63            .field("closed", &self.closed)
64            .field("stats", &self.stats)
65            .finish()
66    }
67}
68
69impl MediaReceiver {
70    pub(crate) fn new(
71        kind: MediaTransportKind,
72        media: BoxedSource,
73        control: Option<BoxedSource>,
74        config: MediaConfig,
75        start: Instant,
76    ) -> Result<Self, NetworkError> {
77        Ok(Self {
78            kind,
79            lanes: select_all(std::iter::once(media).chain(control)),
80            reassembler: Reassembler::new(config).map_err(super::error::from_media_error)?,
81            jitter: JitterBuffer::new(config).map_err(super::error::from_media_error)?,
82            ready: VecDeque::new(),
83            eos: EosTracker::new(config.jitter_depth_micros),
84            start,
85            closed: false,
86            stats: MediaStats::new(),
87        })
88    }
89
90    pub fn kind(&self) -> MediaTransportKind {
91        self.kind
92    }
93
94    pub fn stats(&self) -> MediaStats {
95        self.stats
96    }
97
98    /// The only place wall time is read for the inbound media path.
99    fn now(&self) -> MediaInstant {
100        MediaInstant::from_micros(self.start.elapsed().as_micros() as u64)
101    }
102
103    fn to_instant(&self, at: MediaInstant) -> Instant {
104        self.start + Duration::from_micros(at.as_micros())
105    }
106
107    /// Waits for the next event. After [`MediaEvent::Closed`] every call
108    /// returns `Closed` again.
109    pub async fn next_event(&mut self) -> MediaEvent {
110        loop {
111            if let Some(event) = self.ready.pop_front() {
112                return event;
113            }
114            if let Some(event) = self.pop_jitter() {
115                return event;
116            }
117            if self.resolve_eos(false) {
118                continue;
119            }
120            if self.closed {
121                // Flush frames still parked behind a gap before reporting closure.
122                match self.jitter.next_deadline() {
123                    Some(at) if self.jitter.buffered_len() > 0 => {
124                        sleep_until(self.to_instant(at)).await;
125                        continue;
126                    }
127                    _ => {}
128                }
129                // No more input can arrive: expire pending end-of-stream
130                // records now instead of waiting out their deadlines.
131                if self.resolve_eos(true) {
132                    continue;
133                }
134                return MediaEvent::Closed;
135            }
136            match self.wait_input().await {
137                Input::Datagram(Some(buf)) => self.ingest(buf.as_ref()),
138                Input::Datagram(None) => self.closed = true,
139                Input::Deadline => {
140                    let now = self.now();
141                    self.stats.frames_evicted_incomplete +=
142                        self.reassembler.evict_stale(now) as u64;
143                }
144            }
145        }
146    }
147
148    /// Bridges [`EosTracker::resolve`] to this receiver's state.
149    fn resolve_eos(&mut self, force: bool) -> bool {
150        let now = self.now();
151        let Self {
152            eos,
153            jitter,
154            ready,
155            stats,
156            ..
157        } = self;
158        eos.resolve(now, force, jitter, ready, stats)
159    }
160
161    fn pop_jitter(&mut self) -> Option<MediaEvent> {
162        let now = self.now();
163        let pop = self.jitter.pop_ready(now);
164        if queue_pop(pop, &mut self.ready, &mut self.stats) {
165            self.ready.pop_front()
166        } else {
167            None
168        }
169    }
170
171    async fn wait_input(&mut self) -> Input {
172        let deadline = self
173            .jitter
174            .next_deadline()
175            .into_iter()
176            .chain(self.eos.next_deadline())
177            .min()
178            .map(|at| self.to_instant(at));
179        let timer = async move {
180            match deadline {
181                Some(at) => sleep_until(at).await,
182                None => futures::future::pending().await,
183            }
184        };
185        citadel_io::tokio::select! {
186            item = self.lanes.next() => Input::Datagram(item),
187            _ = timer => Input::Deadline,
188        }
189    }
190
191    fn ingest(&mut self, datagram: &[u8]) {
192        let now = self.now();
193        self.stats.fragments_received += 1;
194        self.stats.bytes_received += datagram.len() as u64;
195        match self.reassembler.push(datagram, now) {
196            ReassembleOutcome::Complete(frame) => {
197                self.stats.frames_completed += 1;
198                // The peer's `MediaSender` numbers every track from a fresh packetizer.
199                self.jitter.expect_from(frame.header.track, FIRST_SEQUENCE);
200                match self.jitter.push(frame, now) {
201                    PushResult::Buffered => {}
202                    PushResult::Late => self.stats.frames_late += 1,
203                    PushResult::TooOld => self.stats.frames_too_old += 1,
204                    PushResult::Duplicate => self.stats.fragments_duplicate += 1,
205                }
206            }
207            ReassembleOutcome::Partial { .. } => {}
208            ReassembleOutcome::Duplicate => self.stats.fragments_duplicate += 1,
209            ReassembleOutcome::Control(body) => match Self::decode_control(&body) {
210                Ok(ControlMessage::AnnounceTracks(list) | ControlMessage::AcceptTracks(list)) => {
211                    self.ready.push_back(MediaEvent::Tracks(list))
212                }
213                Ok(ControlMessage::EndOfStream { track, frames_sent }) => {
214                    self.eos.record(track, frames_sent, now)
215                }
216                Ok(ControlMessage::TransportOffer { .. }) => {
217                    self.stats.fragments_rejected += 1;
218                    log::warn!(target: "citadel", "media: transport offer after negotiation ignored");
219                }
220                Err(err) => {
221                    self.stats.fragments_rejected += 1;
222                    log::warn!(target: "citadel", "{err}");
223                }
224            },
225            ReassembleOutcome::Rejected(err) => {
226                self.stats.fragments_rejected += 1;
227                log::debug!(target: "citadel", "media fragment rejected: {err}");
228            }
229        }
230    }
231
232    fn decode_control(body: &[u8]) -> Result<ControlMessage, NetworkError> {
233        ControlMessage::decode(body)
234            .map_err(|err| citadel_io::error!(ErrorCode::MediaControlDecode, err))
235    }
236}
237
238/// Queues the events for one jitter pop, with its accounting. Returns whether
239/// anything was queued.
240pub(super) fn queue_pop(
241    pop: PopResult,
242    ready: &mut VecDeque<MediaEvent>,
243    stats: &mut MediaStats,
244) -> bool {
245    match pop {
246        PopResult::Frame(frame) => {
247            stats.frames_delivered += 1;
248            ready.push_back(MediaEvent::Frame(frame));
249        }
250        PopResult::Gap {
251            track,
252            missing_from,
253            missing_to,
254            next,
255        } => {
256            stats.gaps_skipped += 1;
257            stats.frames_missing +=
258                u64::from(missing_to.wrapping_sub(missing_from).wrapping_add(1));
259            stats.frames_delivered += 1;
260            ready.push_back(MediaEvent::Gap {
261                track,
262                missing_from,
263                missing_to,
264            });
265            ready.push_back(MediaEvent::Frame(next));
266        }
267        PopResult::NotReady => return false,
268    }
269    true
270}