o_sfu_core/engine/media_transport/rtc/
bitrate.rs1use 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#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
40pub(super) enum IncomingBitrateObservation {
41 #[default]
42 Unchanged,
44 IngressStarted,
46 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 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;