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/// Removes the actor with `id` from the registry.
190///
191/// Only the exact ID is removed, so unrelated actors sharing the thread-local registry are
192/// untouched.
193pub fn deregister_actor(id: &Ustr) {
194    with_actor_registry(|registry| registry.remove(id));
195}
196
197/// Returns a guard providing mutable access to the registered actor of type `T`.
198///
199/// The returned [`ActorRef`] holds an `Rc` to keep the actor alive, preventing
200/// use-after-free if the actor is removed from the registry.
201///
202/// # Panics
203///
204/// - Panics if no actor with the specified `id` is found in the registry.
205/// - Panics if the stored actor is not of type `T`.
206#[must_use]
207pub fn get_actor_unchecked<T: Actor>(id: &Ustr) -> ActorRef<T> {
208    let actor_rc = with_actor_registry(|registry| registry.get(id))
209        .unwrap_or_else(|| panic!("Actor for {id} not found"));
210
211    match actor_ref_from_rc(actor_rc) {
212        Ok(actor_ref) => actor_ref,
213        Err(ActorRefError {
214            expected_type,
215            actual_type,
216        }) => {
217            panic!(
218                "Actor type mismatch for '{id}': expected {expected_type:?}, found {actual_type:?}"
219            )
220        }
221    }
222}
223
224/// Attempts to get a guard providing mutable access to the registered actor.
225///
226/// Returns `None` if the actor is not found or the type doesn't match.
227#[must_use]
228pub fn try_get_actor_unchecked<T: Actor>(id: &Ustr) -> Option<ActorRef<T>> {
229    let actor_rc = with_actor_registry(|registry| registry.get(id))?;
230    actor_ref_from_rc(actor_rc).ok()
231}
232
233#[derive(Debug)]
234struct ActorRefError {
235    expected_type: TypeId,
236    actual_type: TypeId,
237}
238
239fn actor_ref_from_rc<T: Actor>(
240    actor_rc: Rc<UnsafeCell<dyn Actor>>,
241) -> Result<ActorRef<T>, ActorRefError> {
242    // SAFETY: Get a reference to check the type before casting.
243    let actor_ref = unsafe { &*actor_rc.get() };
244    let actual_type = actor_ref.as_any().type_id();
245    let expected_type = TypeId::of::<T>();
246
247    if actual_type != expected_type {
248        return Err(ActorRefError {
249            expected_type,
250            actual_type,
251        });
252    }
253
254    Ok(ActorRef {
255        actor_rc,
256        _marker: PhantomData,
257    })
258}
259
260/// Checks if an actor with the `id` exists in the registry.
261pub fn actor_exists(id: &Ustr) -> bool {
262    with_actor_registry(|registry| registry.contains(id))
263}
264
265/// Returns the number of registered actors.
266pub fn actor_count() -> usize {
267    with_actor_registry(ActorRegistry::len)
268}
269
270#[cfg(test)]
271/// Clears the actor registry (for test isolation).
272pub fn clear_actor_registry() {
273    with_actor_registry(|registry| registry.actors.borrow_mut().clear());
274}
275
276#[cfg(test)]
277mod tests {
278    use std::any::Any;
279
280    use rstest::rstest;
281
282    use super::*;
283
284    #[derive(Debug)]
285    struct TestActor {
286        id: Ustr,
287        value: i32,
288    }
289
290    impl Actor for TestActor {
291        fn id(&self) -> Ustr {
292            self.id
293        }
294        fn handle(&mut self, _msg: &dyn Any) {}
295        fn as_any(&self) -> &dyn Any {
296            self
297        }
298    }
299
300    #[rstest]
301    fn test_register_and_get_actor() {
302        clear_actor_registry();
303
304        let id = Ustr::from("test-actor");
305        let actor = TestActor { id, value: 42 };
306        register_actor(actor);
307
308        let actor_ref = get_actor_unchecked::<TestActor>(&id);
309        assert_eq!(actor_ref.value, 42);
310    }
311
312    #[rstest]
313    fn test_mutation_through_reference() {
314        clear_actor_registry();
315
316        let id = Ustr::from("test-actor-mut");
317        let actor = TestActor { id, value: 0 };
318        register_actor(actor);
319
320        let mut actor_ref = get_actor_unchecked::<TestActor>(&id);
321        actor_ref.value = 999;
322        drop(actor_ref);
323
324        let actor_ref2 = get_actor_unchecked::<TestActor>(&id);
325        assert_eq!(actor_ref2.value, 999);
326    }
327
328    #[rstest]
329    fn test_try_get_returns_none_for_missing() {
330        clear_actor_registry();
331
332        let id = Ustr::from("nonexistent");
333        let result = try_get_actor_unchecked::<TestActor>(&id);
334        assert!(result.is_none());
335    }
336
337    #[rstest]
338    fn test_try_get_returns_none_for_wrong_type() {
339        #[derive(Debug)]
340        struct OtherActor {
341            id: Ustr,
342        }
343
344        impl Actor for OtherActor {
345            fn id(&self) -> Ustr {
346                self.id
347            }
348            fn handle(&mut self, _msg: &dyn Any) {}
349            fn as_any(&self) -> &dyn Any {
350                self
351            }
352        }
353
354        clear_actor_registry();
355
356        let id = Ustr::from("other-actor");
357        let actor = OtherActor { id };
358        register_actor(actor);
359
360        let result = try_get_actor_unchecked::<TestActor>(&id);
361        assert!(result.is_none());
362    }
363
364    #[rstest]
365    fn test_registry_is_thread_local() {
366        clear_actor_registry();
367
368        let id = Ustr::from("thread-local-actor");
369        let actor = TestActor { id, value: 42 };
370        register_actor(actor);
371
372        assert!(actor_exists(&id));
373        assert_eq!(actor_count(), 1);
374
375        let visible_on_other_thread = std::thread::spawn(move || {
376            // Each thread gets its own empty registry
377            (actor_exists(&id), actor_count())
378        })
379        .join()
380        .unwrap();
381
382        assert!(!visible_on_other_thread.0);
383        assert_eq!(visible_on_other_thread.1, 0);
384    }
385
386    #[rstest]
387    fn test_actor_ref_survives_registry_removal() {
388        clear_actor_registry();
389
390        let id = Ustr::from("removable-actor");
391        let actor = TestActor { id, value: 7 };
392        register_actor(actor);
393        assert_eq!(actor_count(), 1);
394
395        let mut guard = get_actor_unchecked::<TestActor>(&id);
396
397        with_actor_registry(|registry| {
398            registry.remove(&id);
399        });
400        assert!(!actor_exists(&id));
401        assert_eq!(actor_count(), 0);
402
403        assert_eq!(guard.value, 7);
404        guard.value = 99;
405        assert_eq!(guard.value, 99);
406    }
407
408    #[rstest]
409    fn test_deregister_actor_removes_only_requested_actor_and_retains_guard() {
410        clear_actor_registry();
411
412        let removed_id = Ustr::from("removed-actor");
413        let retained_id = Ustr::from("retained-actor");
414        register_actor(TestActor {
415            id: removed_id,
416            value: 7,
417        });
418        register_actor(TestActor {
419            id: retained_id,
420            value: 11,
421        });
422        let removed_guard = get_actor_unchecked::<TestActor>(&removed_id);
423
424        deregister_actor(&removed_id);
425
426        assert!(!actor_exists(&removed_id));
427        assert!(actor_exists(&retained_id));
428        assert_eq!(actor_count(), 1);
429        assert_eq!(removed_guard.value, 7);
430        assert_eq!(get_actor_unchecked::<TestActor>(&retained_id).value, 11);
431    }
432
433    #[rstest]
434    fn test_actor_ref_survives_same_id_replacement() {
435        clear_actor_registry();
436
437        let id = Ustr::from("replaceable-actor");
438        let actor_a = TestActor { id, value: 1 };
439        register_actor(actor_a);
440
441        let guard_a = get_actor_unchecked::<TestActor>(&id);
442        assert_eq!(guard_a.value, 1);
443
444        let actor_b = TestActor { id, value: 2 };
445        register_actor(actor_b);
446
447        // Old guard still sees actor A
448        assert_eq!(guard_a.value, 1);
449
450        // Fresh lookup sees actor B
451        let guard_b = get_actor_unchecked::<TestActor>(&id);
452        assert_eq!(guard_b.value, 2);
453        assert_eq!(actor_count(), 1);
454    }
455
456    #[should_panic(expected = "Actor type mismatch")]
457    #[rstest]
458    fn test_get_actor_unchecked_panics_on_type_mismatch() {
459        #[derive(Debug)]
460        struct OtherActor {
461            id: Ustr,
462        }
463
464        impl Actor for OtherActor {
465            fn id(&self) -> Ustr {
466                self.id
467            }
468            fn handle(&mut self, _msg: &dyn Any) {}
469            fn as_any(&self) -> &dyn Any {
470                self
471            }
472        }
473
474        clear_actor_registry();
475
476        let id = Ustr::from("typed-actor");
477        let actor = OtherActor { id };
478        register_actor(actor);
479
480        let _guard = get_actor_unchecked::<TestActor>(&id);
481    }
482}