o_sfu_core/engine/media_transport/rtc/slots.rs
1//! worker-local slots for packet-loop identities
2//!
3//! slots solve the mismatch between stable transport ids and packet-loop work
4//!
5//! ```text
6//! room commands, diagnostics and teardown
7//! use stable public keys such as TransportSessionKey or TransportMediaId
8//! those keys are meaningful outside one media worker
9//!
10//! packet-loop queues, timeout heaps and route destinations
11//! use tiny copy handles
12//! those handles are only meaningful inside the store that created them
13//! ```
14//!
15//! the worker must accept that some work is already queued when a session,
16//! media entry or consumer route is removed
17//! a bare index would let that delayed work reach a replacement occupant after
18//! the index is recycled
19//! [`SlotHandle`] prevents that by pairing the index with the generation that was
20//! current when the handle was created
21//! removal advances the generation before reuse, so stale dirty-session marks,
22//! timeout heap entries and route destinations become ordinary no-ops instead
23//! of touching new state
24//!
25//! the tag parameter is part of the contract
26//! it keeps session handles, media handles and consumer stream handles in
27//! separate type namespaces even though every handle is represented by the same
28//! compact index plus generation pair
29
30use std::{collections::BTreeMap, marker::PhantomData};
31
32use super::{media_registry::RegisteredMediaHandle, state::RtcSessionState};
33use crate::engine::media_transport::{TransportMediaId, TransportSessionKey};
34
35/// handle namespace for live `RtcSessionState` entries
36///
37/// session handles are allowed to live in scheduler queues after the public
38/// session key has been removed
39/// the store validates the generation before the packet loop polls a session
40#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
41pub(super) struct SessionSlot;
42
43/// handle namespace for registered media entries
44///
45/// media slots keep transport media lookup state worker-local while room and
46/// diagnostics paths continue to name media by [`TransportMediaId`]
47#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
48pub(super) struct MediaSlot;
49
50/// handle namespace for downstream RTP rewrite state
51///
52/// route destinations carry these handles so per-packet local forwarding can
53/// reach receiver-local rewrite state without rebuilding a lookup key
54#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
55pub(super) struct ConsumerStreamSlot;
56
57/// generation-checked session handle used by packet-loop scheduler queues
58pub(super) type SessionHandle = SlotHandle<SessionSlot>;
59
60/// generation-checked handle stored on route destinations for local RTP rewrite state
61pub(super) type ConsumerStreamHandle = SlotHandle<ConsumerStreamSlot>;
62
63/// session table keyed by the public session identity at the worker boundary
64///
65/// commands enter through [`TransportSessionKey`]
66/// the packet loop converts that key into [`SessionHandle`] only for queued work
67pub(super) type SessionStore = KeyedSlotStore<TransportSessionKey, RtcSessionState, SessionSlot>;
68
69/// media table keyed by the stable transport media id exposed outside the worker
70pub(super) type MediaStore = KeyedSlotStore<TransportMediaId, RegisteredMediaHandle, MediaSlot>;
71
72/// copy identity for one reusable worker-local slot
73///
74/// callers may queue, copy and compare handles freely because the value is only
75/// an access token
76/// every read, write or removal must still go through the owning store so stale
77/// generations are rejected at the point where state would be touched
78#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
79pub(super) struct SlotHandle<Tag> {
80 index: usize,
81 generation: u64,
82 _tag: PhantomData<fn() -> Tag>,
83}
84
85impl<Tag> Clone for SlotHandle<Tag> {
86 fn clone(&self) -> Self {
87 *self
88 }
89}
90
91impl<Tag> Copy for SlotHandle<Tag> {}
92
93impl<Tag> Default for SlotHandle<Tag> {
94 /// create an invalid handle that cannot match a live generation
95 fn default() -> Self {
96 Self {
97 index: usize::MAX,
98 generation: 0,
99 _tag: PhantomData,
100 }
101 }
102}
103
104/// dense reusable storage for one worker-local identity class
105///
106/// `SlotStore` provides O(1) generational array access for packet-loop hot paths.
107/// Callers hold lightweight `Copy` access tokens ([`SlotHandle`]) while the store
108/// manages entry allocation, free-list recycling, and ABA safety.
109///
110/// # Generational Indexing & ABA Safety
111///
112/// In an asynchronous packet loop, scheduler queues, timeout heaps, and route tables
113/// may hold references to an occupant after it has been removed. A bare array index
114/// would allow delayed work for a removed occupant to mistakenly corrupt a new occupant
115/// that recycled the same index (the classic ABA problem).
116///
117/// `SlotStore` uses these checks during one slot's generation cycle:
118/// 1. Each slot tracks an independent `generation: u64` counter alongside `value: Option<T>`.
119/// 2. `insert` assigns the slot's current generation to the returned `SlotHandle`.
120/// 3. `get` / `get_mut` validate `handle.generation == entry.generation` before granting access.
121/// 4. `remove` takes the value, advances the slot's generation, and returns the index to the LIFO `free` list.
122/// 5. Any delayed work attempting access with an older `SlotHandle` fails the generation check
123/// and evaluates to `None`, safely turning stale queue items into no-ops.
124///
125/// A generation value may be shared by separate slot indices. For one slot index, a
126/// generation can repeat only after its `u64` counter wraps.
127///
128/// # Type-Safe Tagging
129///
130/// The `Tag` marker parameter ([`SessionSlot`], [`MediaSlot`], [`ConsumerStreamSlot`]) ensures
131/// distinct identity-class handle types at compile time with zero memory overhead
132/// (`PhantomData<fn() -> Tag>`). It does not identify a particular store instance, so callers
133/// must use a handle only with the store that created it.
134///
135/// ```text
136/// 1. Initial State (2 active sessions, 1 free slot):
137/// entries:
138/// [0] generation: 1 | value: Some(Session A) <--- SlotHandle { idx: 0, gen: 1 } (active)
139/// [1] generation: 3 | value: Some(Session B) <--- SlotHandle { idx: 1, gen: 3 } (active)
140/// [2] generation: 2 | value: None
141/// free stack: [ 2 ]
142///
143/// 2. Session A is Removed (`remove(handle)`):
144/// - value taken -> None
145/// - generation advanced -> 2
146/// - index 0 pushed to free stack
147/// entries:
148/// [0] generation: 2 | value: None
149/// [1] generation: 3 | value: Some(Session B)
150/// [2] generation: 2 | value: None
151/// free stack: [ 2, 0 ]
152///
153/// 3. Session C is Inserted (`insert(Session C)`):
154/// - index 0 popped from free stack
155/// - value set -> Some(Session C)
156/// - generation retained -> 2
157/// entries:
158/// [0] generation: 2 | value: Some(Session C) <--- SlotHandle { idx: 0, gen: 2 } (new handle)
159/// [1] generation: 3 | value: Some(Session B)
160/// [2] generation: 2 | value: None
161/// free stack: [ 2 ]
162///
163/// 4. Delayed Stale Queue Item Executes:
164/// - queue item holds stale SlotHandle { idx: 0, gen: 1 } (from Session A)
165/// - check: handle.gen (1) == entry.gen (2) ?
166/// |
167/// v (mismatch)
168/// - access returns `None` --> stale work safely discarded with 0 side-effects
169/// ```
170pub(super) struct SlotStore<T, Tag> {
171 entries: Vec<SlotEntry<T>>,
172 free: Vec<usize>,
173 _tag: PhantomData<fn() -> Tag>,
174}
175
176impl<T, Tag> Default for SlotStore<T, Tag> {
177 fn default() -> Self {
178 Self {
179 entries: Vec::new(),
180 free: Vec::new(),
181 _tag: PhantomData,
182 }
183 }
184}
185
186struct SlotEntry<T> {
187 generation: u64,
188 value: Option<T>,
189}
190
191impl<T, Tag> SlotStore<T, Tag> {
192 /// allocate or reuse a slot and return the generation that names this value
193 pub(super) fn insert(&mut self, value: T) -> SlotHandle<Tag> {
194 let index = self.free.pop().unwrap_or_else(|| {
195 let index = self.entries.len();
196 self.entries.push(SlotEntry {
197 generation: 1,
198 value: None,
199 });
200 index
201 });
202 let Some(entry) = self.entries.get_mut(index) else {
203 return SlotHandle::default();
204 };
205 entry.value = Some(value);
206 SlotHandle {
207 index,
208 generation: entry.generation,
209 _tag: PhantomData,
210 }
211 }
212
213 /// return the value only when the handle still names the current generation
214 pub(super) fn get(&self, handle: SlotHandle<Tag>) -> Option<&T> {
215 let entry = self.entries.get(handle.index)?;
216 (entry.generation == handle.generation)
217 .then_some(entry.value.as_ref())
218 .flatten()
219 }
220
221 /// return the mutable value only when the handle still names the current generation
222 pub(super) fn get_mut(&mut self, handle: SlotHandle<Tag>) -> Option<&mut T> {
223 let entry = self.entries.get_mut(handle.index)?;
224 (entry.generation == handle.generation)
225 .then_some(entry.value.as_mut())
226 .flatten()
227 }
228
229 pub(super) fn values_mut(&mut self) -> impl Iterator<Item = &mut T> {
230 self.entries
231 .iter_mut()
232 .filter_map(|entry| entry.value.as_mut())
233 }
234
235 /// remove the value only when the handle still names the current generation
236 ///
237 /// successful removal invalidates every copied handle for the old occupant
238 pub(super) fn remove(&mut self, handle: SlotHandle<Tag>) -> Option<T> {
239 let entry = self.entries.get_mut(handle.index)?;
240 if entry.generation != handle.generation {
241 return None;
242 }
243 let value = entry.value.take()?;
244 entry.generation = next_generation(entry.generation);
245 self.free.push(handle.index);
246 Some(value)
247 }
248}
249
250/// public-key index backed by generation slots
251///
252/// `KeyedSlotStore` bridges public identity at worker boundaries ([`TransportSessionKey`],
253/// [`TransportMediaId`]) with worker-internal generational slots ([`SlotHandle`]).
254///
255/// ```text
256/// Public Boundary (O(log N) tree):
257/// TransportSessionKey("user-1", conn_id) ---> BTreeMap lookup ---> SlotHandle { idx: 0, gen: 2 }
258/// |
259/// Packet-Loop Hot Path (O(1) direct slot): |
260/// SlotStore entries[0] --------------------------------------------------+
261/// - KeyedSlot { key: TransportSessionKey(...), value: RtcSessionState }
262/// ```
263///
264/// - **Worker Boundary**: Callers query by public key (`get(&key)`). `handle_for_key`
265/// resolves that key to its current slot handle through the `BTreeMap`.
266/// - **Packet Loop**: Queues and heaps store lightweight `SlotHandle` copies, avoiding
267/// per-packet map lookups.
268/// - **Reverse Translation**: `key_for_handle` translates a live handle back to its public
269/// key in O(1) time by reading `KeyedSlot.key`.
270pub(super) struct KeyedSlotStore<K, V, Tag> {
271 by_key: BTreeMap<K, SlotHandle<Tag>>,
272 slots: SlotStore<KeyedSlot<K, V>, Tag>,
273}
274
275struct KeyedSlot<K, V> {
276 key: K,
277 value: V,
278}
279
280impl<K, V, Tag> Default for KeyedSlotStore<K, V, Tag> {
281 fn default() -> Self {
282 Self {
283 by_key: BTreeMap::new(),
284 slots: SlotStore::default(),
285 }
286 }
287}
288
289impl<K: Ord + Clone, V, Tag> KeyedSlotStore<K, V, Tag> {
290 pub(super) fn contains_key(&self, key: &K) -> bool {
291 self.by_key.contains_key(key)
292 }
293
294 /// read by public key after validating the current slot generation
295 pub(super) fn get(&self, key: &K) -> Option<&V> {
296 self.get_by_handle(self.handle_for_key(key)?)
297 }
298
299 /// mutably read by public key after validating the current slot generation
300 pub(super) fn get_mut(&mut self, key: &K) -> Option<&mut V> {
301 self.get_mut_by_handle(self.handle_for_key(key)?)
302 }
303
304 /// replace the value for a public key and invalidate any old handle
305 pub(super) fn insert(&mut self, key: K, value: V) -> Option<V> {
306 let previous = self
307 .by_key
308 .remove(&key)
309 .and_then(|handle| self.slots.remove(handle))
310 .map(|entry| entry.value);
311 let handle = self.slots.insert(KeyedSlot {
312 key: key.clone(),
313 value,
314 });
315 self.by_key.insert(key, handle);
316 previous
317 }
318
319 /// remove the value for a public key and invalidate its handle
320 pub(super) fn remove(&mut self, key: &K) -> Option<V> {
321 let handle = self.by_key.remove(key)?;
322 self.slots.remove(handle).map(|entry| entry.value)
323 }
324
325 pub(super) fn len(&self) -> usize {
326 self.by_key.len()
327 }
328
329 pub(super) fn keys(&self) -> impl Iterator<Item = &K> {
330 self.by_key.keys()
331 }
332
333 pub(super) fn iter(&self) -> impl Iterator<Item = (&K, &V)> {
334 self.by_key
335 .iter()
336 .filter_map(|(key, handle)| self.slots.get(*handle).map(|entry| (key, &entry.value)))
337 }
338
339 /// translate a public key to the handle used by packet-loop queues
340 pub(super) fn handle_for_key(&self, key: &K) -> Option<SlotHandle<Tag>> {
341 self.by_key.get(key).copied()
342 }
343
344 /// translate a live handle back to its public key
345 ///
346 /// `None` means the handle is stale or invalid for this store
347 pub(super) fn key_for_handle(&self, handle: SlotHandle<Tag>) -> Option<&K> {
348 self.slots.get(handle).map(|entry| &entry.key)
349 }
350
351 /// read by handle after validating the current slot generation
352 pub(super) fn get_by_handle(&self, handle: SlotHandle<Tag>) -> Option<&V> {
353 self.slots.get(handle).map(|entry| &entry.value)
354 }
355
356 /// mutably read by handle after validating the current slot generation
357 pub(super) fn get_mut_by_handle(&mut self, handle: SlotHandle<Tag>) -> Option<&mut V> {
358 self.slots.get_mut(handle).map(|entry| &mut entry.value)
359 }
360
361 /// mutably read by handle while borrowing the matching public key
362 pub(super) fn get_key_value_mut_by_handle(
363 &mut self,
364 handle: SlotHandle<Tag>,
365 ) -> Option<(&K, &mut V)> {
366 self.slots
367 .get_mut(handle)
368 .map(|entry| (&entry.key, &mut entry.value))
369 }
370}
371
372fn next_generation(generation: u64) -> u64 {
373 // keep generation 0 reserved for invalid handles after wraparound
374 generation.wrapping_add(1).max(1)
375}