Skip to main content

radiate_engines/events/
stream.rs

1use crate::{
2    EventHandler,
3    events::{Event, Handler, Subscriber, Subscription, SubscriptionId},
4};
5use radiate_core::{Executor, error::RadiateResult};
6use std::{
7    any::{Any, TypeId},
8    collections::HashMap,
9    fmt::Debug,
10    sync::{Arc, RwLock},
11};
12
13type Payload = Arc<dyn Any + Send + Sync>;
14type Forward = Arc<dyn Fn(Payload) + Send + Sync>;
15
16#[derive(Clone)]
17struct Registration {
18    forward: Forward,
19    subscription: Subscription,
20}
21
22type SubscriberList = Arc<Vec<Registration>>;
23type SubscriberMap = HashMap<TypeId, SubscriberList>;
24
25#[derive(Clone, Default)]
26pub struct EventStream {
27    executor: Arc<Executor>,
28    subscribers: Arc<RwLock<SubscriberMap>>,
29}
30
31impl EventStream {
32    pub fn new(executor: Arc<Executor>) -> Self {
33        EventStream {
34            executor,
35            subscribers: Arc::default(),
36        }
37    }
38
39    pub fn set_executor(&mut self, executor: Arc<Executor>) {
40        self.executor = executor;
41    }
42
43    pub fn attatch<H: EventHandler>(&self, handler: H) -> RadiateResult<Subscriber<H>> {
44        let subscriber = Subscriber::new(handler, Arc::clone(&self.executor), self.clone());
45        subscriber.start()?;
46        Ok(subscriber)
47    }
48
49    pub fn subscribe<E: Event>(&self, handler: impl Handler<E>) -> Subscription {
50        let subscriber = Subscriber::new(handler, Arc::clone(&self.executor), self.clone());
51        self.subscribe_existing::<E, _>(&subscriber)
52    }
53
54    #[inline]
55    pub fn publish<E: Event>(&self, event: E) {
56        let type_id = TypeId::of::<E>();
57        let Some(group) = self.subscribers.read().unwrap().get(&type_id).cloned() else {
58            return;
59        };
60
61        self.dispatch(&group, Arc::new(event), false);
62    }
63
64    pub fn lazy_publish<E: Event>(&self, f: impl FnOnce() -> E) -> RadiateResult<()> {
65        let type_id = TypeId::of::<E>();
66        let Some(group) = self.subscribers.read().unwrap().get(&type_id).cloned() else {
67            return Ok(());
68        };
69
70        let any_due = group
71            .iter()
72            .any(|registration| registration.subscription.reserve().is_ok_and(|val| val));
73
74        if !any_due {
75            return Ok(());
76        }
77
78        self.dispatch(&group, Arc::new(f()), true);
79        Ok(())
80    }
81
82    pub fn unsubscribe(&self, id: SubscriptionId) {
83        let mut subscribers = self.subscribers.write().unwrap();
84        for group in subscribers.values_mut() {
85            Arc::make_mut(group).retain(|registration| registration.subscription.id() != id);
86        }
87    }
88
89    pub(super) fn subscribe_existing<E, H>(&self, subscriber: &Subscriber<H>) -> Subscription
90    where
91        E: Event,
92        H: Handler<E>,
93    {
94        let target = subscriber.clone();
95        let forward: Forward = Arc::new(move |payload: Payload| {
96            if let Ok(event) = payload.downcast::<E>() {
97                target.send_shared(event);
98            }
99        });
100
101        self.register::<E>(forward)
102    }
103
104    fn register<E: Event>(&self, forward: Forward) -> Subscription {
105        let subscription = Subscription::new();
106        let registration = Registration {
107            forward,
108            subscription: subscription.clone(),
109        };
110
111        let mut subscribers = self.subscribers.write().unwrap();
112        let type_id = TypeId::of::<E>();
113        let list = subscribers
114            .entry(type_id)
115            .or_insert_with(|| Arc::new(Vec::new()));
116        let list = Arc::make_mut(list);
117
118        list.retain(|registration| registration.subscription.is_alive());
119        list.push(registration);
120
121        subscription
122    }
123
124    #[inline]
125    fn dispatch(&self, group: &SubscriberList, payload: Payload, scheduled: bool) {
126        for registration in group.iter() {
127            if !registration.subscription.is_alive() {
128                continue;
129            }
130
131            if scheduled && !registration.subscription.take_permit() {
132                continue;
133            }
134
135            (registration.forward)(Arc::clone(&payload));
136        }
137    }
138}
139
140impl Debug for EventStream {
141    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
142        write!(
143            f,
144            "EventStream(subscribers={}, executor={:?})",
145            self.subscribers.read().unwrap().len(),
146            self.executor,
147        )
148    }
149}
150
151impl From<Executor> for EventStream {
152    fn from(executor: Executor) -> Self {
153        EventStream::new(Arc::new(executor))
154    }
155}