Skip to main content

nautilus_common/actor/
registry.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16//! Thread-local actor registry with access guards.
17//!
18//! # Design
19//!
20//! The actor registry stores actors in thread-local storage and provides access via
21//! [`ActorRef<T>`] guards. This design addresses several constraints:
22//!
23//! - **Use-after-free prevention**: `ActorRef` holds an `Rc` clone, keeping the actor
24//!   alive even if removed from the registry while the guard exists.
25//! - **Scoped registry access**: Registry access stays tied to the thread-local storage
26//!   access callback.
27//! - **Thread-local only**: Guards must not be sent across threads.
28//!
29//! # Limitations
30//!
31//! - **Aliasing not prevented**: Two guards can exist for the same actor simultaneously,
32//!   allowing aliased mutable access. This is undefined behavior if both guards create
33//!   overlapping references to the same actor. The current actor dispatch model relies
34//!   on same-actor re-entrant lookups, so fixing this requires a broader dispatch and
35//!   ownership redesign.
36//!
37//! # Invariants
38//!
39//! These contracts must hold regardless of how the registry is implemented
40//! internally. The first three are verified by tests in this module. The
41//! fourth is a usage discipline enforced by convention.
42//!
43//! - **Thread isolation**: Each thread has its own registry instance. An actor
44//!   registered on one thread is never visible from another.
45//! - **Guard survival**: An [`ActorRef`] keeps its actor alive via reference
46//!   counting. Removing or replacing an actor in the registry does not invalidate
47//!   existing guards.
48//! - **Type safety**: [`get_actor_unchecked`] and [`try_get_actor_unchecked`]
49//!   verify the concrete type at runtime before casting. A type mismatch panics
50//!   or returns `None`, respectively.
51//! - **Short-lived guards**: Guards must be obtained, used, and dropped within a
52//!   single synchronous scope. Never store an [`ActorRef`] in a struct or hold
53//!   one across an `.await` point.
54
55use std::{
56    any::TypeId,
57    cell::{RefCell, UnsafeCell},
58    fmt::Debug,
59    marker::PhantomData,
60    ops::{Deref, DerefMut},
61    rc::Rc,
62};
63
64use ahash::AHashMap;
65use ustr::Ustr;
66
67use super::Actor;
68
69/// A guard providing mutable access to an actor.
70///
71/// This guard holds an `Rc` reference to keep the actor alive.
72pub struct ActorRef<T: Actor> {
73    actor_rc: Rc<UnsafeCell<dyn Actor>>,
74    _marker: PhantomData<T>,
75}
76
77impl<T: Actor> Debug for ActorRef<T> {
78    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
79        f.debug_struct(stringify!(ActorRef))
80            .field("actor_id", &self.deref().id())
81            .finish()
82    }
83}
84
85impl<T: Actor> Deref for ActorRef<T> {
86    type Target = T;
87
88    fn deref(&self) -> &Self::Target {
89        // SAFETY: Type was verified at construction time.
90        unsafe { &*(self.actor_rc.get() as *const T) }
91    }
92}
93
94impl<T: Actor> DerefMut for ActorRef<T> {
95    fn deref_mut(&mut self) -> &mut Self::Target {
96        // SAFETY: Type was verified at construction time.
97        unsafe { &mut *self.actor_rc.get().cast::<T>() }
98    }
99}
100
101thread_local! {
102    static ACTOR_REGISTRY: ActorRegistry = ActorRegistry::new();
103}
104
105/// Registry for storing actors.
106pub struct ActorRegistry {
107    actors: RefCell<AHashMap<Ustr, Rc<UnsafeCell<dyn Actor>>>>,
108}
109
110impl Debug for ActorRegistry {
111    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
112        let actors_ref = self.actors.borrow();
113        let keys: Vec<&Ustr> = actors_ref.keys().collect();
114        f.debug_struct(stringify!(ActorRegistry))
115            .field("actors", &keys)
116            .finish()
117    }
118}
119
120impl Default for ActorRegistry {
121    fn default() -> Self {
122        Self::new()
123    }
124}
125
126impl ActorRegistry {
127    pub fn new() -> Self {
128        Self {
129            actors: RefCell::new(AHashMap::new()),
130        }
131    }
132
133    pub fn insert(&self, id: Ustr, actor: Rc<UnsafeCell<dyn Actor>>) {
134        let mut actors = self.actors.borrow_mut();
135        if actors.contains_key(&id) {
136            log::warn!("Replacing existing actor with id: {id}");
137        }
138        actors.insert(id, actor);
139    }
140
141    pub fn get(&self, id: &Ustr) -> Option<Rc<UnsafeCell<dyn Actor>>> {
142        self.actors.borrow().get(id).cloned()
143    }
144
145    /// Returns the number of registered actors.
146    pub fn len(&self) -> usize {
147        self.actors.borrow().len()
148    }
149
150    /// Checks if the registry is empty.
151    pub fn is_empty(&self) -> bool {
152        self.actors.borrow().is_empty()
153    }
154
155    /// Removes an actor from the registry.
156    pub fn remove(&self, id: &Ustr) -> Option<Rc<UnsafeCell<dyn Actor>>> {
157        self.actors.borrow_mut().remove(id)
158    }
159
160    /// Checks if an actor with the `id` exists.
161    pub fn contains(&self, id: &Ustr) -> bool {
162        self.actors.borrow().contains_key(id)
163    }
164}
165
166pub fn with_actor_registry<R>(f: impl FnOnce(&ActorRegistry) -> R) -> R {
167    ACTOR_REGISTRY.with(f)
168}
169
170/// Registers an actor.
171pub fn register_actor<T>(actor: T) -> Rc<UnsafeCell<T>>
172where
173    T: Actor + 'static,
174{
175    let actor_id = actor.id();
176    let actor_ref = Rc::new(UnsafeCell::new(actor));
177
178    // Register as Actor (message handling only)
179    let actor_trait_ref: Rc<UnsafeCell<dyn Actor>> = actor_ref.clone();
180    with_actor_registry(|registry| registry.insert(actor_id, actor_trait_ref));
181
182    actor_ref
183}
184
185pub fn get_actor(id: &Ustr) -> Option<Rc<UnsafeCell<dyn Actor>>> {
186    with_actor_registry(|registry| registry.get(id))
187}
188
189/// Returns a guard providing mutable access to the registered actor of type `T`.
190///
191/// The returned [`ActorRef`] holds an `Rc` to keep the actor alive, preventing
192/// use-after-free if the actor is removed from the registry.
193///
194/// # Panics
195///
196/// - Panics if no actor with the specified `id` is found in the registry.
197/// - Panics if the stored actor is not of type `T`.
198#[must_use]
199pub fn get_actor_unchecked<T: Actor>(id: &Ustr) -> ActorRef<T> {
200    let actor_rc = with_actor_registry(|registry| registry.get(id))
201        .unwrap_or_else(|| panic!("Actor for {id} not found"));
202
203    match actor_ref_from_rc(actor_rc) {
204        Ok(actor_ref) => actor_ref,
205        Err(ActorRefError {
206            expected_type,
207            actual_type,
208        }) => {
209            panic!(
210                "Actor type mismatch for '{id}': expected {expected_type:?}, found {actual_type:?}"
211            )
212        }
213    }
214}
215
216/// Attempts to get a guard providing mutable access to the registered actor.
217///
218/// Returns `None` if the actor is not found or the type doesn't match.
219#[must_use]
220pub fn try_get_actor_unchecked<T: Actor>(id: &Ustr) -> Option<ActorRef<T>> {
221    let actor_rc = with_actor_registry(|registry| registry.get(id))?;
222    actor_ref_from_rc(actor_rc).ok()
223}
224
225#[derive(Debug)]
226struct ActorRefError {
227    expected_type: TypeId,
228    actual_type: TypeId,
229}
230
231fn actor_ref_from_rc<T: Actor>(
232    actor_rc: Rc<UnsafeCell<dyn Actor>>,
233) -> Result<ActorRef<T>, ActorRefError> {
234    // SAFETY: Get a reference to check the type before casting.
235    let actor_ref = unsafe { &*actor_rc.get() };
236    let actual_type = actor_ref.as_any().type_id();
237    let expected_type = TypeId::of::<T>();
238
239    if actual_type != expected_type {
240        return Err(ActorRefError {
241            expected_type,
242            actual_type,
243        });
244    }
245
246    Ok(ActorRef {
247        actor_rc,
248        _marker: PhantomData,
249    })
250}
251
252/// Checks if an actor with the `id` exists in the registry.
253pub fn actor_exists(id: &Ustr) -> bool {
254    with_actor_registry(|registry| registry.contains(id))
255}
256
257/// Returns the number of registered actors.
258pub fn actor_count() -> usize {
259    with_actor_registry(ActorRegistry::len)
260}
261
262#[cfg(test)]
263/// Clears the actor registry (for test isolation).
264pub fn clear_actor_registry() {
265    with_actor_registry(|registry| registry.actors.borrow_mut().clear());
266}
267
268#[cfg(test)]
269mod tests {
270    use std::any::Any;
271
272    use rstest::rstest;
273
274    use super::*;
275
276    #[derive(Debug)]
277    struct TestActor {
278        id: Ustr,
279        value: i32,
280    }
281
282    impl Actor for TestActor {
283        fn id(&self) -> Ustr {
284            self.id
285        }
286        fn handle(&mut self, _msg: &dyn Any) {}
287        fn as_any(&self) -> &dyn Any {
288            self
289        }
290    }
291
292    #[rstest]
293    fn test_register_and_get_actor() {
294        clear_actor_registry();
295
296        let id = Ustr::from("test-actor");
297        let actor = TestActor { id, value: 42 };
298        register_actor(actor);
299
300        let actor_ref = get_actor_unchecked::<TestActor>(&id);
301        assert_eq!(actor_ref.value, 42);
302    }
303
304    #[rstest]
305    fn test_mutation_through_reference() {
306        clear_actor_registry();
307
308        let id = Ustr::from("test-actor-mut");
309        let actor = TestActor { id, value: 0 };
310        register_actor(actor);
311
312        let mut actor_ref = get_actor_unchecked::<TestActor>(&id);
313        actor_ref.value = 999;
314        drop(actor_ref);
315
316        let actor_ref2 = get_actor_unchecked::<TestActor>(&id);
317        assert_eq!(actor_ref2.value, 999);
318    }
319
320    #[rstest]
321    fn test_try_get_returns_none_for_missing() {
322        clear_actor_registry();
323
324        let id = Ustr::from("nonexistent");
325        let result = try_get_actor_unchecked::<TestActor>(&id);
326        assert!(result.is_none());
327    }
328
329    #[rstest]
330    fn test_try_get_returns_none_for_wrong_type() {
331        #[derive(Debug)]
332        struct OtherActor {
333            id: Ustr,
334        }
335
336        impl Actor for OtherActor {
337            fn id(&self) -> Ustr {
338                self.id
339            }
340            fn handle(&mut self, _msg: &dyn Any) {}
341            fn as_any(&self) -> &dyn Any {
342                self
343            }
344        }
345
346        clear_actor_registry();
347
348        let id = Ustr::from("other-actor");
349        let actor = OtherActor { id };
350        register_actor(actor);
351
352        let result = try_get_actor_unchecked::<TestActor>(&id);
353        assert!(result.is_none());
354    }
355
356    #[rstest]
357    fn test_registry_is_thread_local() {
358        clear_actor_registry();
359
360        let id = Ustr::from("thread-local-actor");
361        let actor = TestActor { id, value: 42 };
362        register_actor(actor);
363
364        assert!(actor_exists(&id));
365        assert_eq!(actor_count(), 1);
366
367        let visible_on_other_thread = std::thread::spawn(move || {
368            // Each thread gets its own empty registry
369            (actor_exists(&id), actor_count())
370        })
371        .join()
372        .unwrap();
373
374        assert!(!visible_on_other_thread.0);
375        assert_eq!(visible_on_other_thread.1, 0);
376    }
377
378    #[rstest]
379    fn test_actor_ref_survives_registry_removal() {
380        clear_actor_registry();
381
382        let id = Ustr::from("removable-actor");
383        let actor = TestActor { id, value: 7 };
384        register_actor(actor);
385        assert_eq!(actor_count(), 1);
386
387        let mut guard = get_actor_unchecked::<TestActor>(&id);
388
389        with_actor_registry(|registry| {
390            registry.remove(&id);
391        });
392        assert!(!actor_exists(&id));
393        assert_eq!(actor_count(), 0);
394
395        assert_eq!(guard.value, 7);
396        guard.value = 99;
397        assert_eq!(guard.value, 99);
398    }
399
400    #[rstest]
401    fn test_actor_ref_survives_same_id_replacement() {
402        clear_actor_registry();
403
404        let id = Ustr::from("replaceable-actor");
405        let actor_a = TestActor { id, value: 1 };
406        register_actor(actor_a);
407
408        let guard_a = get_actor_unchecked::<TestActor>(&id);
409        assert_eq!(guard_a.value, 1);
410
411        let actor_b = TestActor { id, value: 2 };
412        register_actor(actor_b);
413
414        // Old guard still sees actor A
415        assert_eq!(guard_a.value, 1);
416
417        // Fresh lookup sees actor B
418        let guard_b = get_actor_unchecked::<TestActor>(&id);
419        assert_eq!(guard_b.value, 2);
420        assert_eq!(actor_count(), 1);
421    }
422
423    #[should_panic(expected = "Actor type mismatch")]
424    #[rstest]
425    fn test_get_actor_unchecked_panics_on_type_mismatch() {
426        #[derive(Debug)]
427        struct OtherActor {
428            id: Ustr,
429        }
430
431        impl Actor for OtherActor {
432            fn id(&self) -> Ustr {
433                self.id
434            }
435            fn handle(&mut self, _msg: &dyn Any) {}
436            fn as_any(&self) -> &dyn Any {
437                self
438            }
439        }
440
441        clear_actor_registry();
442
443        let id = Ustr::from("typed-actor");
444        let actor = OtherActor { id };
445        register_actor(actor);
446
447        let _guard = get_actor_unchecked::<TestActor>(&id);
448    }
449}