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