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, Registration>>,
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        let previous = actors.insert(
139            id,
140            Registration {
141                actor,
142                identity: Rc::new(()),
143            },
144        );
145        drop(actors);
146        drop(previous);
147    }
148
149    pub fn get(&self, id: &Ustr) -> Option<Rc<UnsafeCell<dyn Actor>>> {
150        self.actors
151            .borrow()
152            .get(id)
153            .map(|entry| entry.actor.clone())
154    }
155
156    /// Returns the number of registered actors.
157    pub fn len(&self) -> usize {
158        self.actors.borrow().len()
159    }
160
161    /// Checks if the registry is empty.
162    pub fn is_empty(&self) -> bool {
163        self.actors.borrow().is_empty()
164    }
165
166    /// Removes an actor from the registry.
167    pub fn remove(&self, id: &Ustr) -> Option<Rc<UnsafeCell<dyn Actor>>> {
168        self.actors.borrow_mut().remove(id).map(|entry| entry.actor)
169    }
170
171    /// Checks if an actor with the `id` exists.
172    pub fn contains(&self, id: &Ustr) -> bool {
173        self.actors.borrow().contains_key(id)
174    }
175}
176
177pub fn with_actor_registry<R>(f: impl FnOnce(&ActorRegistry) -> R) -> R {
178    ACTOR_REGISTRY.with(f)
179}
180
181/// Registers an actor.
182pub fn register_actor<T>(actor: T) -> Rc<UnsafeCell<T>>
183where
184    T: Actor + 'static,
185{
186    let actor_id = actor.id();
187    let actor_ref = Rc::new(UnsafeCell::new(actor));
188
189    // Register as Actor (message handling only)
190    let actor_trait_ref: Rc<UnsafeCell<dyn Actor>> = actor_ref.clone();
191    with_actor_registry(|registry| registry.insert(actor_id, actor_trait_ref));
192
193    actor_ref
194}
195
196pub fn get_actor(id: &Ustr) -> Option<Rc<UnsafeCell<dyn Actor>>> {
197    with_actor_registry(|registry| registry.get(id))
198}
199
200/// Removes the actor with `id` from the registry.
201///
202/// Only the exact ID is removed, so unrelated actors sharing the thread-local registry are
203/// untouched.
204pub fn deregister_actor(id: &Ustr) {
205    with_actor_registry(|registry| registry.remove(id));
206}
207
208/// Returns a guard providing mutable access to the registered actor of type `T`.
209///
210/// The returned [`ActorRef`] holds an `Rc` to keep the actor alive, preventing
211/// use-after-free if the actor is removed from the registry.
212///
213/// # Panics
214///
215/// Panics if:
216/// - No actor with the specified `id` is found in the registry.
217/// - The stored actor is not of type `T`.
218#[must_use]
219pub fn get_actor_unchecked<T: Actor>(id: &Ustr) -> ActorRef<T> {
220    let actor_rc = with_actor_registry(|registry| registry.get(id))
221        .unwrap_or_else(|| panic!("Actor for {id} not found"));
222
223    match actor_ref_from_rc(actor_rc) {
224        Ok(actor_ref) => actor_ref,
225        Err(ActorRefError {
226            expected_type,
227            actual_type,
228        }) => {
229            panic!(
230                "Actor type mismatch for '{id}': expected {expected_type:?}, found {actual_type:?}"
231            )
232        }
233    }
234}
235
236/// Attempts to get a guard providing mutable access to the registered actor.
237///
238/// Returns `None` if the actor is not found or the type doesn't match.
239#[must_use]
240pub fn try_get_actor_unchecked<T: Actor>(id: &Ustr) -> Option<ActorRef<T>> {
241    let actor_rc = with_actor_registry(|registry| registry.get(id))?;
242    actor_ref_from_rc(actor_rc).ok()
243}
244
245#[derive(Debug)]
246struct ActorRefError {
247    expected_type: TypeId,
248    actual_type: TypeId,
249}
250
251fn actor_ref_from_rc<T: Actor>(
252    actor_rc: Rc<UnsafeCell<dyn Actor>>,
253) -> Result<ActorRef<T>, ActorRefError> {
254    // SAFETY: Get a reference to check the type before casting.
255    let actor_ref = unsafe { &*actor_rc.get() };
256    let actual_type = actor_ref.as_any().type_id();
257    let expected_type = TypeId::of::<T>();
258
259    if actual_type != expected_type {
260        return Err(ActorRefError {
261            expected_type,
262            actual_type,
263        });
264    }
265
266    Ok(ActorRef {
267        actor_rc,
268        _marker: PhantomData,
269    })
270}
271
272/// Checks if an actor with the `id` exists in the registry.
273pub fn actor_exists(id: &Ustr) -> bool {
274    with_actor_registry(|registry| registry.contains(id))
275}
276
277/// Returns the number of registered actors.
278pub fn actor_count() -> usize {
279    with_actor_registry(ActorRegistry::len)
280}
281
282#[derive(Clone)]
283struct Registration {
284    actor: Rc<UnsafeCell<dyn Actor>>,
285    identity: Rc<()>,
286}
287
288#[allow(
289    dead_code,
290    reason = "registration-bound admission remains inactive until runtime integration"
291)]
292pub(super) fn reserve_actor<T: Actor, E: 'static>(
293    id: Ustr,
294    heap_bytes: usize,
295) -> Option<ActorAdmission<T, E>> {
296    let registration = ACTOR_REGISTRY
297        .try_with(|registry| registry.actors.borrow().get(&id).cloned())
298        .ok()
299        .flatten()?;
300    let admission = super::dispatch::reserve(heap_bytes)?;
301    Some(ActorAdmission {
302        id,
303        registration,
304        admission,
305    })
306}
307
308#[allow(
309    dead_code,
310    reason = "registration-bound admission remains inactive until runtime integration"
311)]
312pub(super) struct ActorAdmission<T: Actor, E> {
313    id: Ustr,
314    registration: Registration,
315    admission: super::dispatch::Admission<ActorDelivery<T, E>>,
316}
317
318#[allow(
319    dead_code,
320    reason = "registration-bound admission remains inactive until runtime integration"
321)]
322impl<T: Actor, E: 'static> ActorAdmission<T, E> {
323    pub(super) fn commit(self, event: E, handler: fn(&mut T, &E)) {
324        self.admission.commit(
325            ActorDelivery {
326                id: self.id,
327                registration: self.registration,
328                event,
329                handler,
330            },
331            ActorDelivery::run,
332        );
333    }
334}
335
336struct ActorDelivery<T: Actor, E> {
337    id: Ustr,
338    registration: Registration,
339    event: E,
340    handler: fn(&mut T, &E),
341}
342
343impl<T: Actor, E> ActorDelivery<T, E> {
344    fn run(&mut self) -> bool {
345        let current =
346            ACTOR_REGISTRY
347                .try_with(|registry| {
348                    registry.actors.borrow().get(&self.id).is_some_and(|entry| {
349                        Rc::ptr_eq(&entry.identity, &self.registration.identity)
350                    })
351                })
352                .unwrap_or(false);
353
354        if !current {
355            return true;
356        }
357
358        match super::access::ActorGuard::acquire(self.registration.actor.clone()) {
359            Ok(mut actor) => {
360                (self.handler)(&mut actor, &self.event);
361                true
362            }
363            Err(super::access::ActorAccessError::Busy) => false,
364            Err(_) => {
365                super::dispatch::record_failure(super::dispatch::DispatchError::InvalidDestination);
366                true
367            }
368        }
369    }
370}
371
372#[cfg(test)]
373/// Clears the actor registry (for test isolation).
374pub fn clear_actor_registry() {
375    let actors = with_actor_registry(|registry| std::mem::take(&mut *registry.actors.borrow_mut()));
376    drop(actors);
377}
378
379#[cfg(test)]
380mod tests {
381    use std::any::Any;
382
383    use rstest::rstest;
384
385    use super::*;
386
387    #[derive(Debug)]
388    struct TestActor {
389        id: Ustr,
390        value: i32,
391    }
392
393    impl Actor for TestActor {
394        fn id(&self) -> Ustr {
395            self.id
396        }
397        fn handle(&mut self, _msg: &dyn Any) {}
398        fn as_any(&self) -> &dyn Any {
399            self
400        }
401    }
402
403    #[rstest]
404    fn owned_delivery_waits_for_access_and_cancels_same_allocation_registration() {
405        super::super::dispatch::clear().unwrap();
406        clear_actor_registry();
407        let id = Ustr::from("owned-registration");
408        let allocation = register_actor(TestActor { id, value: 11 });
409        let guard =
410            super::super::access::ActorGuard::<TestActor>::acquire(allocation.clone()).unwrap();
411        reserve_actor::<TestActor, _>(id, 0)
412            .unwrap()
413            .commit(23, |actor, value| actor.value = *value);
414        assert_eq!(super::super::dispatch::drain(1).unwrap().delivered, 0);
415        drop(guard);
416        assert_eq!(super::super::dispatch::drain(1).unwrap().delivered, 1);
417        assert_eq!(get_actor_unchecked::<TestActor>(&id).value, 23);
418        reserve_actor::<TestActor, _>(id, 0)
419            .unwrap()
420            .commit(37, |actor, value| actor.value = *value);
421        deregister_actor(&id);
422        with_actor_registry(|registry| registry.insert(id, allocation));
423        assert_eq!(super::super::dispatch::drain(1).unwrap().delivered, 1);
424        assert_eq!(get_actor_unchecked::<TestActor>(&id).value, 23);
425        super::super::dispatch::clear().unwrap();
426    }
427
428    #[rstest]
429    fn owned_delivery_retries_busy_allocation() {
430        super::super::dispatch::clear().unwrap();
431        clear_actor_registry();
432        let id = Ustr::from("busy-delivery");
433        let allocation = register_actor(TestActor { id, value: 11 });
434        let registration = with_actor_registry(|registry| registry.actors.borrow()[&id].clone());
435        let mut delivery = ActorDelivery {
436            id,
437            registration,
438            event: 23,
439            handler: |actor: &mut TestActor, value: &i32| actor.value = *value,
440        };
441        let guard = super::super::access::ActorGuard::<TestActor>::acquire(allocation).unwrap();
442        let busy = delivery.run();
443        assert!(!busy);
444        assert_eq!(guard.value, 11);
445        assert_eq!(super::super::dispatch::failure(), None);
446        drop(guard);
447        let delivered = delivery.run();
448        assert!(delivered);
449        assert_eq!(get_actor_unchecked::<TestActor>(&id).value, 23);
450        assert_eq!(super::super::dispatch::failure(), None);
451        clear_actor_registry();
452    }
453
454    #[rstest]
455    fn owned_delivery_rejects_wrong_actor_type() {
456        #[derive(Debug)]
457        struct OtherActor;
458
459        impl Actor for OtherActor {
460            fn id(&self) -> Ustr {
461                Ustr::from("wrong-delivery-type")
462            }
463            fn handle(&mut self, _msg: &dyn Any) {}
464            fn as_any(&self) -> &dyn Any {
465                self
466            }
467        }
468
469        super::super::dispatch::clear().unwrap();
470        clear_actor_registry();
471        let id = OtherActor.id();
472        register_actor(OtherActor);
473        reserve_actor::<TestActor, _>(id, 0)
474            .unwrap()
475            .commit(23, |_, _| panic!("wrong-type handler must not run"));
476        let result = super::super::dispatch::drain(1);
477        assert_eq!(
478            result,
479            Err(super::super::dispatch::DispatchError::InvalidDestination)
480        );
481        assert_eq!(
482            super::super::dispatch::failure(),
483            Some(super::super::dispatch::DispatchError::InvalidDestination)
484        );
485        super::super::dispatch::clear().unwrap();
486        clear_actor_registry();
487    }
488
489    #[rstest]
490    fn actor_reservation_rejects_after_registry_teardown() {
491        struct Probe;
492
493        impl Drop for Probe {
494            fn drop(&mut self) {
495                assert!(ACTOR_REGISTRY.try_with(|_| ()).is_err());
496                assert!(reserve_actor::<TestActor, ()>(Ustr::from("teardown"), 0).is_none());
497            }
498        }
499
500        thread_local! {
501            static PROBE: Probe = const { Probe };
502        }
503
504        std::thread::spawn(|| {
505            PROBE.with(|_| ());
506            register_actor(TestActor {
507                id: Ustr::from("teardown"),
508                value: 11,
509            });
510        })
511        .join()
512        .unwrap();
513    }
514
515    #[rstest]
516    fn test_register_and_get_actor() {
517        clear_actor_registry();
518
519        let id = Ustr::from("test-actor");
520        let actor = TestActor { id, value: 42 };
521        register_actor(actor);
522
523        let actor_ref = get_actor_unchecked::<TestActor>(&id);
524        assert_eq!(actor_ref.value, 42);
525    }
526
527    #[rstest]
528    fn test_mutation_through_reference() {
529        clear_actor_registry();
530
531        let id = Ustr::from("test-actor-mut");
532        let actor = TestActor { id, value: 0 };
533        register_actor(actor);
534
535        let mut actor_ref = get_actor_unchecked::<TestActor>(&id);
536        actor_ref.value = 999;
537        drop(actor_ref);
538
539        let actor_ref2 = get_actor_unchecked::<TestActor>(&id);
540        assert_eq!(actor_ref2.value, 999);
541    }
542
543    #[rstest]
544    fn test_try_get_returns_none_for_missing() {
545        clear_actor_registry();
546
547        let id = Ustr::from("nonexistent");
548        let result = try_get_actor_unchecked::<TestActor>(&id);
549        assert!(result.is_none());
550    }
551
552    #[rstest]
553    fn test_try_get_returns_none_for_wrong_type() {
554        #[derive(Debug)]
555        struct OtherActor {
556            id: Ustr,
557        }
558
559        impl Actor for OtherActor {
560            fn id(&self) -> Ustr {
561                self.id
562            }
563            fn handle(&mut self, _msg: &dyn Any) {}
564            fn as_any(&self) -> &dyn Any {
565                self
566            }
567        }
568
569        clear_actor_registry();
570
571        let id = Ustr::from("other-actor");
572        let actor = OtherActor { id };
573        register_actor(actor);
574
575        let result = try_get_actor_unchecked::<TestActor>(&id);
576        assert!(result.is_none());
577    }
578
579    #[rstest]
580    fn test_registry_is_thread_local() {
581        clear_actor_registry();
582
583        let id = Ustr::from("thread-local-actor");
584        let actor = TestActor { id, value: 42 };
585        register_actor(actor);
586
587        assert!(actor_exists(&id));
588        assert_eq!(actor_count(), 1);
589
590        let visible_on_other_thread = std::thread::spawn(move || {
591            // Each thread gets its own empty registry
592            (actor_exists(&id), actor_count())
593        })
594        .join()
595        .unwrap();
596
597        assert!(!visible_on_other_thread.0);
598        assert_eq!(visible_on_other_thread.1, 0);
599    }
600
601    #[rstest]
602    fn test_actor_ref_survives_registry_removal() {
603        clear_actor_registry();
604
605        let id = Ustr::from("removable-actor");
606        let actor = TestActor { id, value: 7 };
607        register_actor(actor);
608        assert_eq!(actor_count(), 1);
609
610        let mut guard = get_actor_unchecked::<TestActor>(&id);
611
612        with_actor_registry(|registry| {
613            registry.remove(&id);
614        });
615        assert!(!actor_exists(&id));
616        assert_eq!(actor_count(), 0);
617
618        assert_eq!(guard.value, 7);
619        guard.value = 99;
620        assert_eq!(guard.value, 99);
621    }
622
623    #[rstest]
624    fn test_deregister_actor_removes_only_requested_actor_and_retains_guard() {
625        clear_actor_registry();
626
627        let removed_id = Ustr::from("removed-actor");
628        let retained_id = Ustr::from("retained-actor");
629        register_actor(TestActor {
630            id: removed_id,
631            value: 7,
632        });
633        register_actor(TestActor {
634            id: retained_id,
635            value: 11,
636        });
637        let removed_guard = get_actor_unchecked::<TestActor>(&removed_id);
638
639        deregister_actor(&removed_id);
640
641        assert!(!actor_exists(&removed_id));
642        assert!(actor_exists(&retained_id));
643        assert_eq!(actor_count(), 1);
644        assert_eq!(removed_guard.value, 7);
645        assert_eq!(get_actor_unchecked::<TestActor>(&retained_id).value, 11);
646    }
647
648    #[rstest]
649    fn test_actor_ref_survives_same_id_replacement() {
650        clear_actor_registry();
651
652        let id = Ustr::from("replaceable-actor");
653        let actor_a = TestActor { id, value: 1 };
654        register_actor(actor_a);
655
656        let guard_a = get_actor_unchecked::<TestActor>(&id);
657        assert_eq!(guard_a.value, 1);
658
659        let actor_b = TestActor { id, value: 2 };
660        register_actor(actor_b);
661
662        // Old guard still sees actor A
663        assert_eq!(guard_a.value, 1);
664
665        // Fresh lookup sees actor B
666        let guard_b = get_actor_unchecked::<TestActor>(&id);
667        assert_eq!(guard_b.value, 2);
668        assert_eq!(actor_count(), 1);
669    }
670
671    #[should_panic(expected = "Actor type mismatch")]
672    #[rstest]
673    fn test_get_actor_unchecked_panics_on_type_mismatch() {
674        #[derive(Debug)]
675        struct OtherActor {
676            id: Ustr,
677        }
678
679        impl Actor for OtherActor {
680            fn id(&self) -> Ustr {
681                self.id
682            }
683            fn handle(&mut self, _msg: &dyn Any) {}
684            fn as_any(&self) -> &dyn Any {
685                self
686            }
687        }
688
689        clear_actor_registry();
690
691        let id = Ustr::from("typed-actor");
692        let actor = OtherActor { id };
693        register_actor(actor);
694
695        let _guard = get_actor_unchecked::<TestActor>(&id);
696    }
697}