Skip to main content

arete_server/
bus.rs

1use bytes::Bytes;
2use std::collections::HashMap;
3use std::sync::Arc;
4use tokio::sync::{broadcast, watch, RwLock};
5
6/// Message sent through the event bus
7#[derive(Debug, Clone)]
8pub struct BusMessage {
9    pub key: String,
10    pub entity: String,
11    pub payload: Arc<Bytes>,
12}
13
14/// The latest frame published to one state key, and how many frames the key
15/// has published in total.
16///
17/// A state bus keeps only the latest frame, so frames published faster than a
18/// subscriber reads them are overwritten. The count lets the subscriber tell:
19/// if it moved by more than one since its last read, frames were skipped.
20#[derive(Debug, Clone, Default)]
21pub struct StateUpdate {
22    pub published: u64,
23    pub payload: Arc<Bytes>,
24}
25
26#[derive(Clone)]
27#[allow(clippy::type_complexity)]
28pub struct BusManager {
29    state_buses: Arc<RwLock<HashMap<(String, String), watch::Sender<StateUpdate>>>>,
30    list_buses: Arc<RwLock<HashMap<String, broadcast::Sender<Arc<BusMessage>>>>>,
31    broadcast_capacity: usize,
32}
33
34impl BusManager {
35    pub fn new() -> Self {
36        Self::with_capacity(1000)
37    }
38
39    pub fn with_capacity(capacity: usize) -> Self {
40        Self {
41            state_buses: Arc::new(RwLock::new(HashMap::new())),
42            list_buses: Arc::new(RwLock::new(HashMap::new())),
43            broadcast_capacity: capacity,
44        }
45    }
46
47    /// Get or create a state bus (latest-value semantics)
48    /// Each (view_id, key) pair gets its own watch channel
49    pub async fn get_or_create_state_bus(
50        &self,
51        view_id: &str,
52        key: &str,
53    ) -> watch::Receiver<StateUpdate> {
54        let mut buses = self.state_buses.write().await;
55        let entry = (view_id.to_string(), key.to_string());
56
57        let tx = buses
58            .entry(entry)
59            .or_insert_with(|| watch::channel(StateUpdate::default()).0)
60            .clone();
61
62        tx.subscribe()
63    }
64
65    pub async fn get_or_create_list_bus(
66        &self,
67        view_id: &str,
68    ) -> broadcast::Receiver<Arc<BusMessage>> {
69        let mut buses = self.list_buses.write().await;
70
71        let tx = buses
72            .entry(view_id.to_string())
73            .or_insert_with(|| broadcast::channel(self.broadcast_capacity).0)
74            .clone();
75
76        tx.subscribe()
77    }
78
79    /// Publish to a state bus (latest-value)
80    pub async fn publish_state(&self, view_id: &str, key: &str, frame: Arc<Bytes>) {
81        let buses = self.state_buses.read().await;
82        if let Some(tx) = buses.get(&(view_id.to_string(), key.to_string())) {
83            tx.send_modify(|update| {
84                update.published += 1;
85                update.payload = frame;
86            });
87        }
88    }
89
90    pub async fn publish_list(&self, view_id: &str, message: Arc<BusMessage>) {
91        let buses = self.list_buses.read().await;
92        if let Some(tx) = buses.get(view_id) {
93            let _ = tx.send(message);
94        }
95    }
96
97    pub async fn cleanup_stale_state_buses(&self) -> usize {
98        let mut buses = self.state_buses.write().await;
99        let before = buses.len();
100
101        buses.retain(|_, tx| tx.receiver_count() > 0);
102
103        before - buses.len()
104    }
105
106    pub async fn cleanup_stale_list_buses(&self) -> usize {
107        let mut buses = self.list_buses.write().await;
108        let before = buses.len();
109
110        buses.retain(|_, tx| tx.receiver_count() > 0);
111
112        before - buses.len()
113    }
114
115    pub async fn bus_counts(&self) -> (usize, usize) {
116        let state_count = self.state_buses.read().await.len();
117        let list_count = self.list_buses.read().await.len();
118        (state_count, list_count)
119    }
120}
121
122impl Default for BusManager {
123    fn default() -> Self {
124        Self::new()
125    }
126}