Skip to main content

o_sfu_core/engine/
packet_sink_registry.rs

1use std::{
2    collections::HashMap,
3    fmt,
4    sync::{
5        Arc, RwLock,
6        atomic::{AtomicBool, AtomicU64, Ordering},
7    },
8    time::Instant,
9};
10
11use super::{
12    RoomInstanceId,
13    media_transport::{TransportMediaId, TransportSessionKey},
14    metrics::RtpForwardDestinationKind,
15    sync::{read_unpoisoned, write_unpoisoned},
16};
17
18#[cfg(test)]
19#[expect(non_snake_case, reason = "test modules map to local TESTS directories")]
20#[path = "packet_sink_registry/TESTS/mod.rs"]
21mod TESTS;
22
23/// Observes origin-side RTP payloads routed to a room packet sink.
24///
25/// Packet gates do not filter this source-side stream. Relayed packets are not
26/// observed again.
27pub trait PacketSink: Send + Sync {
28    /// Records one source RTP payload.
29    ///
30    /// `session_key` and `transport_media_id` identify the source transport
31    /// media. The RTC engine invokes this synchronously on a packet worker.
32    /// Different workers may invoke the same sink concurrently, so
33    /// implementations must return promptly and must not perform blocking I/O.
34    fn record_packet(
35        &self,
36        session_key: &TransportSessionKey,
37        transport_media_id: TransportMediaId,
38        received_at: Instant,
39        payload: &[u8],
40    );
41}
42
43/// Forwarding labels available to source-side packet sinks.
44///
45/// Local RTC egress requires a successful RTC send and cannot be reported by a sink.
46#[derive(Clone, Copy, Debug)]
47pub enum PacketSinkKind {
48    Recording,
49    IntraNodeRelay,
50}
51
52#[derive(Clone)]
53pub struct RegisteredPacketSink {
54    sink: Arc<dyn PacketSink>,
55    kind: PacketSinkKind,
56}
57
58impl RegisteredPacketSink {
59    /// Registers a source-side sink without local RTC egress accounting.
60    pub fn new(sink: Arc<dyn PacketSink>, kind: PacketSinkKind) -> Self {
61        Self { sink, kind }
62    }
63
64    pub fn record_packet(
65        &self,
66        session_key: &TransportSessionKey,
67        transport_media_id: TransportMediaId,
68        received_at: Instant,
69        payload: &[u8],
70    ) {
71        self.sink
72            .record_packet(session_key, transport_media_id, received_at, payload);
73    }
74
75    #[must_use]
76    pub const fn forward_destination_kind(&self) -> RtpForwardDestinationKind {
77        match self.kind {
78            PacketSinkKind::Recording => RtpForwardDestinationKind::Recording,
79            PacketSinkKind::IntraNodeRelay => RtpForwardDestinationKind::IntraNodeRelay,
80        }
81    }
82}
83
84impl fmt::Debug for RegisteredPacketSink {
85    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
86        formatter
87            .debug_struct("RegisteredPacketSink")
88            .field("kind", &self.kind)
89            .finish_non_exhaustive()
90    }
91}
92
93pub struct RoomPacketSinkRegistry {
94    any_active: AtomicBool,
95    generation: AtomicU64,
96    active_rooms: RwLock<HashMap<RoomInstanceId, RegisteredPacketSink>>,
97}
98
99impl Default for RoomPacketSinkRegistry {
100    fn default() -> Self {
101        Self {
102            any_active: AtomicBool::new(false),
103            generation: AtomicU64::new(0),
104            active_rooms: RwLock::new(HashMap::new()),
105        }
106    }
107}
108
109#[derive(Default)]
110pub struct PacketSinkRouteCache {
111    generation: u64,
112    active_rooms: HashMap<RoomInstanceId, RegisteredPacketSink>,
113}
114
115impl PacketSinkRouteCache {
116    /// Refreshes the cached room routes to one registry generation.
117    ///
118    /// Later registry changes remain invisible through [`Self::sink_for_room`]
119    /// until `refresh_from` runs again.
120    pub fn refresh_from(&mut self, registry: &RoomPacketSinkRegistry) {
121        let generation = registry.generation();
122        if self.generation == generation {
123            return;
124        }
125        let snapshot = registry.snapshot();
126        self.generation = snapshot.generation;
127        self.active_rooms = snapshot.active_rooms;
128    }
129
130    #[inline]
131    pub fn sink_for_room(&self, room_instance_id: RoomInstanceId) -> Option<RegisteredPacketSink> {
132        self.active_rooms.get(&room_instance_id).cloned()
133    }
134}
135
136struct PacketSinkRegistrySnapshot {
137    generation: u64,
138    active_rooms: HashMap<RoomInstanceId, RegisteredPacketSink>,
139}
140
141impl RoomPacketSinkRegistry {
142    #[inline]
143    pub fn sink_for_room(&self, room_instance_id: RoomInstanceId) -> Option<RegisteredPacketSink> {
144        if !self.any_active.load(Ordering::Acquire) {
145            return None;
146        }
147        read_unpoisoned(&self.active_rooms)
148            .get(&room_instance_id)
149            .cloned()
150    }
151
152    fn generation(&self) -> u64 {
153        self.generation.load(Ordering::Acquire)
154    }
155
156    fn snapshot(&self) -> PacketSinkRegistrySnapshot {
157        // Both writers advance the generation before releasing this map's guard.
158        let active_rooms = read_unpoisoned(&self.active_rooms);
159        PacketSinkRegistrySnapshot {
160            generation: self.generation(),
161            active_rooms: active_rooms.clone(),
162        }
163    }
164
165    /// Registers `sink` for `room_instance_id`, replacing the current entry.
166    ///
167    /// Previously cloned [`RegisteredPacketSink`] handles are not revoked.
168    pub fn register_room(
169        &self,
170        room_instance_id: RoomInstanceId,
171        sink: Arc<dyn PacketSink>,
172        kind: PacketSinkKind,
173    ) {
174        let mut active_rooms = write_unpoisoned(&self.active_rooms);
175        active_rooms.insert(room_instance_id, RegisteredPacketSink::new(sink, kind));
176        self.any_active.store(true, Ordering::Release);
177        self.generation.fetch_add(1, Ordering::AcqRel);
178        drop(active_rooms);
179    }
180
181    /// Removes the current entry for `room_instance_id`.
182    ///
183    /// Cached or cloned sink handles may still receive packets after
184    /// `unregister_room` returns.
185    pub fn unregister_room(&self, room_instance_id: RoomInstanceId) {
186        let mut active_rooms = write_unpoisoned(&self.active_rooms);
187        active_rooms.remove(&room_instance_id);
188        self.any_active
189            .store(!active_rooms.is_empty(), Ordering::Release);
190        self.generation.fetch_add(1, Ordering::AcqRel);
191        drop(active_rooms);
192    }
193
194    fn active_room_count(&self) -> usize {
195        read_unpoisoned(&self.active_rooms).len()
196    }
197}
198
199impl fmt::Debug for RoomPacketSinkRegistry {
200    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
201        formatter
202            .debug_struct("RoomPacketSinkRegistry")
203            .field("any_active", &self.any_active.load(Ordering::Relaxed))
204            .field("active_room_count", &self.active_room_count())
205            .finish_non_exhaustive()
206    }
207}