Skip to main content

o_sfu_core/engine/media_transport/rtc/
bitrate.rs

1//! Worker-local RTC payload bitrate observations.
2//!
3//! The packet loop records tracked ingress payload and local RTP payload queued
4//! into str0m. Egress counters do not reflect UDP send success or transport
5//! overhead. Shared atomic counters keep per-packet updates off the registry
6//! lock, which is used only to add or remove counters and collect snapshots.
7//!
8//! Snapshots publish the most recently completed window and expire it after one
9//! second without packets.
10
11use std::{
12    collections::BTreeMap,
13    sync::{
14        Arc,
15        atomic::{AtomicBool, AtomicU64, Ordering},
16    },
17    time::{Duration, Instant},
18};
19
20use super::state::PacketLoopState;
21use crate::{
22    Bitrate,
23    engine::media_transport::{TransportBitrateSnapshot, TransportMediaId, TransportSessionKey},
24};
25
26const BITRATE_WINDOW_NANOS: u64 = 1_000_000_000;
27
28#[derive(Debug)]
29pub(super) struct MediaBitrateCounter {
30    origin: Instant,
31    window_start_nanos: AtomicU64,
32    last_observed_nanos: AtomicU64,
33    bytes_in_window: AtomicU64,
34    completed_bps: AtomicU64,
35    observed: AtomicBool,
36}
37
38/// Packet-loop follow-up produced by one incoming bitrate observation.
39#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
40pub(super) enum IncomingBitrateObservation {
41    #[default]
42    /// No ingress edge or completed sample needs a policy wake.
43    Unchanged,
44    /// First packet after registration or one full idle window.
45    IngressStarted,
46    /// A completed bitrate window is available to policy readers.
47    SampleUpdated,
48}
49
50impl IncomingBitrateObservation {
51    pub(super) const fn policy_dirty(self) -> bool {
52        !matches!(self, Self::Unchanged)
53    }
54
55    pub(super) const fn ingress_started(self) -> bool {
56        matches!(self, Self::IngressStarted)
57    }
58}
59
60impl MediaBitrateCounter {
61    pub(super) fn new(now: Instant) -> Self {
62        Self {
63            origin: now,
64            window_start_nanos: AtomicU64::new(0),
65            last_observed_nanos: AtomicU64::new(0),
66            bytes_in_window: AtomicU64::new(0),
67            completed_bps: AtomicU64::new(0),
68            observed: AtomicBool::new(false),
69        }
70    }
71
72    /// The packet loop is the only writer. Cross-thread readers consume the
73    /// completed sample and freshness atomics, never the in-progress byte bucket.
74    ///
75    /// Exact saturating addition is not part of the observable contract because
76    /// one RTP bitrate window cannot approach `u64::MAX`.
77    pub(super) fn record(&self, now: Instant, payload_bytes: usize) -> IncomingBitrateObservation {
78        let now_nanos = self.nanos_since_origin(now);
79        let was_observed = self.observed.load(Ordering::Acquire);
80        let previous_observed = self.last_observed_nanos.load(Ordering::Acquire);
81        let payload_bytes = u64::try_from(payload_bytes).unwrap_or(u64::MAX);
82        let observation = if was_observed {
83            let ingress_started =
84                now_nanos.saturating_sub(previous_observed) >= BITRATE_WINDOW_NANOS;
85            if ingress_started {
86                self.window_start_nanos.store(now_nanos, Ordering::Release);
87                self.bytes_in_window.store(payload_bytes, Ordering::Release);
88                self.completed_bps.store(0, Ordering::Release);
89                IncomingBitrateObservation::IngressStarted
90            } else {
91                let window_start = self.window_start_nanos.load(Ordering::Acquire);
92                let elapsed_nanos = now_nanos.saturating_sub(window_start);
93                if elapsed_nanos >= BITRATE_WINDOW_NANOS {
94                    let completed_bytes =
95                        self.bytes_in_window.swap(payload_bytes, Ordering::AcqRel);
96                    let completed_bps = bitrate_per_second(completed_bytes, elapsed_nanos);
97                    self.completed_bps
98                        .store(completed_bps.as_bps(), Ordering::Release);
99                    self.window_start_nanos.store(now_nanos, Ordering::Release);
100                    IncomingBitrateObservation::SampleUpdated
101                } else {
102                    self.bytes_in_window
103                        .fetch_add(payload_bytes, Ordering::Release);
104                    IncomingBitrateObservation::Unchanged
105                }
106            }
107        } else {
108            self.window_start_nanos.store(now_nanos, Ordering::Release);
109            self.bytes_in_window.store(payload_bytes, Ordering::Release);
110            IncomingBitrateObservation::IngressStarted
111        };
112        self.last_observed_nanos.store(now_nanos, Ordering::Release);
113        if !was_observed {
114            self.observed.store(true, Ordering::Release);
115        }
116        observation
117    }
118
119    pub(super) fn last_observed_age(&self, now: Instant) -> Option<Duration> {
120        if !self.observed.load(Ordering::Acquire) {
121            return None;
122        }
123        Some(Duration::from_nanos(
124            self.nanos_since_origin(now)
125                .saturating_sub(self.last_observed_nanos.load(Ordering::Acquire)),
126        ))
127    }
128
129    fn snapshot(&self, now: Instant) -> Bitrate {
130        let now_nanos = self.nanos_since_origin(now);
131        if !self.observed.load(Ordering::Acquire)
132            || now_nanos.saturating_sub(self.last_observed_nanos.load(Ordering::Acquire))
133                >= BITRATE_WINDOW_NANOS
134        {
135            return Bitrate::zero();
136        }
137        Bitrate::from_bps(self.completed_bps.load(Ordering::Acquire))
138    }
139
140    fn nanos_since_origin(&self, now: Instant) -> u64 {
141        let elapsed = now
142            .checked_duration_since(self.origin)
143            .unwrap_or(Duration::ZERO);
144        u64::try_from(elapsed.as_nanos()).unwrap_or(u64::MAX)
145    }
146}
147
148fn bitrate_per_second(bytes: u64, elapsed_nanos: u64) -> Bitrate {
149    let bits_per_second = u128::from(bytes)
150        .saturating_mul(8)
151        .saturating_mul(u128::from(BITRATE_WINDOW_NANOS))
152        / u128::from(elapsed_nanos.max(1));
153    Bitrate::from_bps(u64::try_from(bits_per_second).unwrap_or(u64::MAX))
154}
155
156#[derive(Debug, Default)]
157pub(super) struct SessionIncomingBitrates {
158    per_media: BTreeMap<TransportMediaId, Arc<MediaBitrateCounter>>,
159}
160
161impl SessionIncomingBitrates {
162    fn register(
163        &mut self,
164        transport_media_id: TransportMediaId,
165        now: Instant,
166    ) -> Arc<MediaBitrateCounter> {
167        Arc::clone(
168            self.per_media
169                .entry(transport_media_id)
170                .or_insert_with(|| Arc::new(MediaBitrateCounter::new(now))),
171        )
172    }
173
174    fn remove(&mut self, transport_media_id: TransportMediaId) {
175        self.per_media.remove(&transport_media_id);
176    }
177
178    fn is_empty(&self) -> bool {
179        self.per_media.is_empty()
180    }
181
182    fn extend_snapshot(&self, now: Instant, snapshot: &mut TransportBitrateSnapshot) {
183        for (&media_id, counter) in &self.per_media {
184            let bitrate = counter.snapshot(now);
185            snapshot.total = snapshot.total.saturating_add(bitrate);
186            if bitrate > Bitrate::zero() {
187                snapshot.per_media.push((media_id, bitrate));
188            }
189        }
190    }
191}
192
193#[derive(Debug, Default)]
194pub struct BitrateRegistry {
195    pub(super) incoming_bitrates_by_session: BTreeMap<TransportSessionKey, SessionIncomingBitrates>,
196    pub(super) egress_bitrates_by_session: BTreeMap<TransportSessionKey, Arc<MediaBitrateCounter>>,
197}
198
199impl BitrateRegistry {
200    pub(super) fn register_incoming_media(
201        &mut self,
202        session_key: &TransportSessionKey,
203        transport_media_id: TransportMediaId,
204        now: Instant,
205    ) -> Arc<MediaBitrateCounter> {
206        self.incoming_bitrates_by_session
207            .entry(session_key.clone())
208            .or_default()
209            .register(transport_media_id, now)
210    }
211
212    pub(super) fn register_session_egress(
213        &mut self,
214        session_key: &TransportSessionKey,
215        counter: Arc<MediaBitrateCounter>,
216    ) {
217        self.egress_bitrates_by_session
218            .insert(session_key.clone(), counter);
219    }
220
221    pub(super) fn remove_incoming_media(
222        &mut self,
223        session_key: &TransportSessionKey,
224        transport_media_id: TransportMediaId,
225    ) {
226        let Some(session_bitrates) = self.incoming_bitrates_by_session.get_mut(session_key) else {
227            return;
228        };
229        session_bitrates.remove(transport_media_id);
230        if session_bitrates.is_empty() {
231            self.incoming_bitrates_by_session.remove(session_key);
232        }
233    }
234
235    pub(super) fn remove_session(&mut self, session_key: &TransportSessionKey) {
236        self.incoming_bitrates_by_session.remove(session_key);
237        self.egress_bitrates_by_session.remove(session_key);
238    }
239
240    pub fn transport_bitrate_snapshot_at(
241        &self,
242        session_keys: &[TransportSessionKey],
243        now: Instant,
244    ) -> TransportBitrateSnapshot {
245        let mut snapshot = TransportBitrateSnapshot::default();
246        for session_key in session_keys {
247            let Some(session_bitrates) = self.incoming_bitrates_by_session.get(session_key) else {
248                continue;
249            };
250            session_bitrates.extend_snapshot(now, &mut snapshot);
251        }
252        snapshot
253    }
254
255    #[cfg(any(test, feature = "internal-benchmarks"))]
256    pub fn egress_bitrate_snapshot_at(
257        &self,
258        session_keys: &[TransportSessionKey],
259        now: Instant,
260    ) -> Bitrate {
261        session_keys
262            .iter()
263            .filter_map(|session_key| self.egress_bitrates_by_session.get(session_key))
264            .fold(Bitrate::zero(), |total, bitrate| {
265                total.saturating_add(bitrate.snapshot(now))
266            })
267    }
268
269    pub fn total_egress_bitrate_snapshot_at(&self, now: Instant) -> Bitrate {
270        self.egress_bitrates_by_session
271            .values()
272            .fold(Bitrate::zero(), |total, bitrate| {
273                total.saturating_add(bitrate.snapshot(now))
274            })
275    }
276}
277
278impl PacketLoopState {
279    pub(super) fn register_incoming_bitrate_counter(
280        &mut self,
281        transport_media_id: TransportMediaId,
282        counter: Arc<MediaBitrateCounter>,
283    ) {
284        self.incoming_bitrate_counters
285            .insert(transport_media_id, counter);
286    }
287
288    pub(super) fn remove_incoming_bitrate_counter(&mut self, transport_media_id: TransportMediaId) {
289        self.incoming_bitrate_counters.remove(&transport_media_id);
290    }
291
292    pub(super) fn record_incoming_bitrate(
293        &self,
294        transport_media_id: TransportMediaId,
295        now: Instant,
296        payload_bytes: usize,
297    ) -> Option<IncomingBitrateObservation> {
298        self.incoming_bitrate_counters
299            .get(&transport_media_id)
300            .map(|bitrate| bitrate.record(now, payload_bytes))
301    }
302}
303
304#[cfg(test)]
305#[path = "TESTS/bitrate.rs"]
306mod tests;