Skip to main content

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}