1use bytes::Bytes;
2use std::collections::HashMap;
3use std::sync::Arc;
4use tokio::sync::{broadcast, watch, RwLock};
5
6#[derive(Debug, Clone)]
8pub struct BusMessage {
9 pub key: String,
10 pub entity: String,
11 pub payload: Arc<Bytes>,
12}
13
14#[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 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 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}