rust_elm/runtime/subscription/
mod.rs

1//! Subscription interpreter — starts/stops long-lived Tokio sources declared by [`Sub`].
2mod websocket;
3
4use std::collections::{HashMap, HashSet};
5use std::sync::atomic::{AtomicBool, Ordering};
6use std::sync::Arc;
7
8use tokio::runtime::Handle;
9use tokio::task::AbortHandle;
10use tokio::time::{interval, MissedTickBehavior};
11
12use crate::bus::BusSender;
13use crate::store::StoreBackend;
14use crate::sub::Sub;
15
16pub(crate) struct SubscriptionHandles {
17    pub handles: parking_lot::Mutex<HashMap<u64, AbortHandle>>,
18}
19
20impl SubscriptionHandles {
21    pub fn new() -> Self {
22        Self {
23            handles: parking_lot::Mutex::new(HashMap::new()),
24        }
25    }
26
27    pub fn abort_all(&self) {
28        self.handles
29            .lock()
30            .drain()
31            .for_each(|(_, handle)| handle.abort());
32    }
33}
34
35/// Reconcile running subscription tasks with the desired [`Sub`] tree.
36pub(crate) fn sync_subscriptions<S, M>(
37    sub: &Sub<M>,
38    registry: &SubscriptionHandles,
39    handle: Handle,
40    tx: BusSender<M>,
41    backend: StoreBackend<S, M>,
42    shutdown: Arc<AtomicBool>,
43) where
44    S: Send + 'static,
45    M: Send + 'static,
46{
47    let mut desired = HashSet::new();
48    collect_sub_ids(sub, &mut desired);
49
50    let mut active = registry.handles.lock();
51    active.retain(|id, abort| {
52        if desired.contains(id) {
53            true
54        } else {
55            abort.abort();
56            false
57        }
58    });
59
60    spawn_subscriptions(
61        sub,
62        &mut active,
63        handle,
64        tx,
65        backend,
66        shutdown,
67    );
68}
69
70fn spawn_subscriptions<S, M>(
71    sub: &Sub<M>,
72    active: &mut HashMap<u64, AbortHandle>,
73    handle: Handle,
74    tx: BusSender<M>,
75    backend: StoreBackend<S, M>,
76    shutdown: Arc<AtomicBool>,
77) where
78    S: Send + 'static,
79    M: Send + 'static,
80{
81    match sub {
82        Sub::None => {}
83        Sub::Tick { id, every, produce } => {
84            spawn_if_new(*id, active, || {
85                spawn_tick(*every, *produce, handle, tx, backend, shutdown)
86            });
87        }
88        Sub::Stream { id, name, every, produce } => {
89            spawn_if_new(*id, active, || {
90                spawn_stream(name, *every, *produce, handle, tx, backend, shutdown)
91            });
92        }
93        Sub::WebSocket { id, url, every, produce } => {
94            spawn_if_new(*id, active, || {
95                websocket::spawn_websocket(url, *every, *produce, handle, tx, backend, shutdown)
96            });
97        }
98        Sub::MapMsg { inner, map } => spawn_unit_subscriptions(
99            inner.as_ref(),
100            *map,
101            active,
102            handle,
103            tx,
104            backend,
105            shutdown,
106        ),
107        Sub::Batch(items) => {
108            for item in items {
109                spawn_subscriptions(
110                    item,
111                    active,
112                    handle.clone(),
113                    tx.clone(),
114                    backend.clone(),
115                    shutdown.clone(),
116                );
117            }
118        }
119    }
120}
121
122fn spawn_unit_subscriptions<S, M>(
123    sub: &Sub<()>,
124    map: fn(()) -> M,
125    active: &mut HashMap<u64, AbortHandle>,
126    handle: Handle,
127    tx: BusSender<M>,
128    backend: StoreBackend<S, M>,
129    shutdown: Arc<AtomicBool>,
130) where
131    S: Send + 'static,
132    M: Send + 'static,
133{
134    match sub {
135        Sub::None => {}
136        Sub::Tick { id, every, .. } => {
137            spawn_if_new(*id, active, || {
138                spawn_tick_mapped(*every, map, handle, tx, backend, shutdown)
139            });
140        }
141        Sub::Stream { id, name, every, .. } => {
142            spawn_if_new(*id, active, || {
143                spawn_stream_mapped(name, *every, map, handle, tx, backend, shutdown)
144            });
145        }
146        Sub::WebSocket { id, url, every, .. } => {
147            spawn_if_new(*id, active, || {
148                websocket::spawn_websocket_mapped(url, *every, map, handle, tx, backend, shutdown)
149            });
150        }
151        Sub::MapMsg { .. } => {}
152        Sub::Batch(items) => {
153            for item in items {
154                spawn_unit_subscriptions(
155                    item,
156                    map,
157                    active,
158                    handle.clone(),
159                    tx.clone(),
160                    backend.clone(),
161                    shutdown.clone(),
162                );
163            }
164        }
165    }
166}
167
168fn spawn_if_new(id: u64, active: &mut HashMap<u64, AbortHandle>, spawn: impl FnOnce() -> AbortHandle) {
169    if active.contains_key(&id) {
170        return;
171    }
172    active.insert(id, spawn());
173}
174
175fn spawn_tick<S, M>(
176    every: std::time::Duration,
177    produce: fn() -> M,
178    handle: Handle,
179    tx: BusSender<M>,
180    backend: StoreBackend<S, M>,
181    shutdown: Arc<AtomicBool>,
182) -> AbortHandle
183where
184    S: Send + 'static,
185    M: Send + 'static,
186{
187    handle
188        .spawn(async move {
189            let mut ticker = interval(every);
190            ticker.set_missed_tick_behavior(MissedTickBehavior::Skip);
191            ticker.tick().await;
192            loop {
193                if shutdown.load(Ordering::Relaxed) {
194                    break;
195                }
196                ticker.tick().await;
197                if shutdown.load(Ordering::Relaxed) {
198                    break;
199                }
200                crate::runtime::dispatch::dispatch_from_subscription(&backend, &tx, produce());
201            }
202        })
203        .abort_handle()
204}
205
206fn spawn_tick_mapped<S, M>(
207    every: std::time::Duration,
208    map: fn(()) -> M,
209    handle: Handle,
210    tx: BusSender<M>,
211    backend: StoreBackend<S, M>,
212    shutdown: Arc<AtomicBool>,
213) -> AbortHandle
214where
215    S: Send + 'static,
216    M: Send + 'static,
217{
218    handle
219        .spawn(async move {
220            let mut ticker = interval(every);
221            ticker.set_missed_tick_behavior(MissedTickBehavior::Skip);
222            ticker.tick().await;
223            loop {
224                if shutdown.load(Ordering::Relaxed) {
225                    break;
226                }
227                ticker.tick().await;
228                if shutdown.load(Ordering::Relaxed) {
229                    break;
230                }
231                crate::runtime::dispatch::dispatch_from_subscription(&backend, &tx, map(()));
232            }
233        })
234        .abort_handle()
235}
236
237fn spawn_stream<S, M>(
238    _name: &'static str,
239    every: std::time::Duration,
240    produce: fn() -> M,
241    handle: Handle,
242    tx: BusSender<M>,
243    backend: StoreBackend<S, M>,
244    shutdown: Arc<AtomicBool>,
245) -> AbortHandle
246where
247    S: Send + 'static,
248    M: Send + 'static,
249{
250    handle
251        .spawn(async move {
252            let mut ticker = interval(every);
253            ticker.set_missed_tick_behavior(MissedTickBehavior::Skip);
254            loop {
255                if shutdown.load(Ordering::Relaxed) {
256                    break;
257                }
258                ticker.tick().await;
259                if shutdown.load(Ordering::Relaxed) {
260                    break;
261                }
262                crate::runtime::dispatch::dispatch_from_subscription(&backend, &tx, produce());
263            }
264        })
265        .abort_handle()
266}
267
268fn spawn_stream_mapped<S, M>(
269    _name: &'static str,
270    every: std::time::Duration,
271    map: fn(()) -> M,
272    handle: Handle,
273    tx: BusSender<M>,
274    backend: StoreBackend<S, M>,
275    shutdown: Arc<AtomicBool>,
276) -> AbortHandle
277where
278    S: Send + 'static,
279    M: Send + 'static,
280{
281    handle
282        .spawn(async move {
283            let mut ticker = interval(every);
284            ticker.set_missed_tick_behavior(MissedTickBehavior::Skip);
285            loop {
286                if shutdown.load(Ordering::Relaxed) {
287                    break;
288                }
289                ticker.tick().await;
290                if shutdown.load(Ordering::Relaxed) {
291                    break;
292                }
293                crate::runtime::dispatch::dispatch_from_subscription(&backend, &tx, map(()));
294            }
295        })
296        .abort_handle()
297}
298
299pub(crate) fn collect_sub_ids<M>(sub: &Sub<M>, out: &mut HashSet<u64>) {
300    match sub {
301        Sub::None => {}
302        Sub::Tick { id, .. }
303        | Sub::Stream { id, .. }
304        | Sub::WebSocket { id, .. } => {
305            out.insert(*id);
306        }
307        Sub::MapMsg { inner, .. } => collect_sub_ids(inner, out),
308        Sub::Batch(items) => {
309            for item in items {
310                collect_sub_ids(item, out);
311            }
312        }
313    }
314}
315
316#[cfg(test)]
317mod tests {
318    use super::*;
319    use std::time::Duration;
320
321    fn unit() {}
322
323    fn mapped(_: ()) -> i32 {
324        7
325    }
326
327    #[test]
328    fn collect_ids_from_batch() {
329        fn zero() -> i32 {
330            0
331        }
332        let sub = Sub::batch([
333            Sub::tick(1, Duration::from_secs(1), zero),
334            Sub::stream(2, "events", Duration::from_secs(1), zero),
335        ]);
336        let mut ids = HashSet::new();
337        collect_sub_ids(&sub, &mut ids);
338        assert_eq!(ids, HashSet::from([1, 2]));
339    }
340
341    #[test]
342    fn collect_ids_from_map_msg() {
343        let sub = Sub::MapMsg {
344            inner: Box::new(Sub::tick(9, Duration::from_millis(1), unit)),
345            map: mapped,
346        };
347        let mut ids = HashSet::new();
348        collect_sub_ids(&sub, &mut ids);
349        assert_eq!(ids, HashSet::from([9]));
350    }
351}
rust_elm/runtime/subscription/mod.rs

rust_elm/runtime/subscription/
mod.rs
