sz_orm_websocket/
pusher.rs

1use crate::error::WsError;
2use serde::{Deserialize, Serialize};
3use std::collections::HashMap;
4use std::sync::Arc;
5use tokio::sync::RwLock;
6
7/// Shared in-memory message buffer that captures messages sent to each connection.
8/// Each connection_id maps to a list of messages that have been pushed to it.
9pub struct InMemorySender {
10    messages: Arc<RwLock<HashMap<String, Vec<Vec<u8>>>>>,
11}
12
13impl InMemorySender {
14    pub fn new() -> Self {
15        Self {
16            messages: Arc::new(RwLock::new(HashMap::new())),
17        }
18    }
19
20    pub async fn send(&self, connection_id: &str, message: Vec<u8>) -> Result<(), WsError> {
21        let mut messages = self.messages.write().await;
22        messages
23            .entry(connection_id.to_string())
24            .or_insert_with(Vec::new)
25            .push(message);
26        Ok(())
27    }
28
29    pub async fn close(&self, connection_id: &str) -> Result<(), WsError> {
30        let mut messages = self.messages.write().await;
31        messages.remove(connection_id);
32        Ok(())
33    }
34
35    pub async fn messages_for(&self, connection_id: &str) -> Vec<Vec<u8>> {
36        let messages = self.messages.read().await;
37        messages.get(connection_id).cloned().unwrap_or_default()
38    }
39
40    pub async fn message_count(&self, connection_id: &str) -> usize {
41        let messages = self.messages.read().await;
42        messages.get(connection_id).map(|v| v.len()).unwrap_or(0)
43    }
44
45    pub async fn total_message_count(&self) -> usize {
46        let messages = self.messages.read().await;
47        messages.values().map(|v| v.len()).sum()
48    }
49}
50
51impl Default for InMemorySender {
52    fn default() -> Self {
53        Self::new()
54    }
55}
56
57struct ConnectionState {
58    user_id: Option<i64>,
59    subscriptions: Vec<String>,
60}
61
62pub struct RealtimePusher {
63    connections: Arc<RwLock<HashMap<String, ConnectionState>>>,
64    rooms: Arc<RwLock<HashMap<String, Vec<String>>>>,
65    sender: InMemorySender,
66}
67
68impl RealtimePusher {
69    pub fn new() -> Self {
70        Self {
71            connections: Arc::new(RwLock::new(HashMap::new())),
72            rooms: Arc::new(RwLock::new(HashMap::new())),
73            sender: InMemorySender::new(),
74        }
75    }
76
77    pub async fn register_connection(&self, connection_id: impl Into<String>) {
78        let mut connections = self.connections.write().await;
79        connections.insert(
80            connection_id.into(),
81            ConnectionState {
82                user_id: None,
83                subscriptions: Vec::new(),
84            },
85        );
86    }
87
88    pub async fn register_connection_with_user(
89        &self,
90        connection_id: impl Into<String>,
91        user_id: i64,
92    ) {
93        let mut connections = self.connections.write().await;
94        connections.insert(
95            connection_id.into(),
96            ConnectionState {
97                user_id: Some(user_id),
98                subscriptions: Vec::new(),
99            },
100        );
101    }
102
103    pub async fn unregister_connection(&self, connection_id: &str) {
104        let mut connections = self.connections.write().await;
105        connections.remove(connection_id);
106        drop(connections);
107
108        let mut rooms = self.rooms.write().await;
109        for ids in rooms.values_mut() {
110            ids.retain(|id| id != connection_id);
111        }
112        drop(rooms);
113
114        let _ = self.sender.close(connection_id).await;
115    }
116
117    pub async fn subscribe(&self, connection_id: &str, room: &str) -> Result<(), WsError> {
118        let mut connections = self.connections.write().await;
119        let conn = connections
120            .get_mut(connection_id)
121            .ok_or_else(|| WsError::Connection("Connection not found".to_string()))?;
122
123        if !conn.subscriptions.contains(&room.to_string()) {
124            conn.subscriptions.push(room.to_string());
125        }
126
127        drop(connections);
128
129        let mut rooms = self.rooms.write().await;
130        let entry = rooms.entry(room.to_string()).or_insert_with(Vec::new);
131        if !entry.contains(&connection_id.to_string()) {
132            entry.push(connection_id.to_string());
133        }
134
135        Ok(())
136    }
137
138    pub async fn unsubscribe(&self, connection_id: &str, room: &str) {
139        let mut connections = self.connections.write().await;
140        if let Some(conn) = connections.get_mut(connection_id) {
141            conn.subscriptions.retain(|r| r != room);
142        }
143
144        drop(connections);
145
146        let mut rooms = self.rooms.write().await;
147        if let Some(ids) = rooms.get_mut(room) {
148            ids.retain(|id| id != connection_id);
149        }
150    }
151
152    pub async fn push_to_connection(
153        &self,
154        connection_id: &str,
155        message: Vec<u8>,
156    ) -> Result<(), WsError> {
157        let connections = self.connections.read().await;
158        let _conn = connections
159            .get(connection_id)
160            .ok_or_else(|| WsError::Connection("Connection not found".to_string()))?;
161        drop(connections);
162
163        self.sender.send(connection_id, message).await
164    }
165
166    pub async fn push_to_room(&self, room: &str, message: Vec<u8>) -> Result<usize, WsError> {
167        let rooms = self.rooms.read().await;
168        let connection_ids = rooms.get(room).cloned().unwrap_or_default();
169        drop(rooms);
170
171        let mut success_count = 0;
172
173        for conn_id in connection_ids {
174            if self.sender.send(&conn_id, message.clone()).await.is_ok() {
175                success_count += 1;
176            }
177        }
178
179        Ok(success_count)
180    }
181
182    pub async fn push_to_user(&self, user_id: i64, message: Vec<u8>) -> Result<usize, WsError> {
183        let connections = self.connections.read().await;
184        let matching_ids: Vec<String> = connections
185            .iter()
186            .filter(|(_, conn)| conn.user_id == Some(user_id))
187            .map(|(id, _)| id.clone())
188            .collect();
189        drop(connections);
190
191        let mut success_count = 0;
192
193        for conn_id in matching_ids {
194            if self.sender.send(&conn_id, message.clone()).await.is_ok() {
195                success_count += 1;
196            }
197        }
198
199        Ok(success_count)
200    }
201
202    pub async fn push_order_status(
203        &self,
204        user_id: i64,
205        order_id: i64,
206        status: &str,
207    ) -> Result<usize, WsError> {
208        let payload = serde_json::json!({
209            "type": "order_status",
210            "order_id": order_id,
211            "status": status,
212        });
213
214        let message = serde_json::to_vec(&payload)?;
215        self.push_to_user(user_id, message).await
216    }
217
218    pub async fn push_customer_message(
219        &self,
220        room_id: &str,
221        sender_id: i64,
222        content: &str,
223    ) -> Result<usize, WsError> {
224        let payload = serde_json::json!({
225            "type": "customer_message",
226            "sender_id": sender_id,
227            "content": content,
228        });
229
230        let message = serde_json::to_vec(&payload)?;
231        self.push_to_room(room_id, message).await
232    }
233
234    pub async fn broadcast(&self, message: Vec<u8>) -> Result<usize, WsError> {
235        let connections = self.connections.read().await;
236        let conn_ids: Vec<String> = connections.keys().cloned().collect();
237        drop(connections);
238
239        let mut success_count = 0;
240
241        for conn_id in conn_ids {
242            if self.sender.send(&conn_id, message.clone()).await.is_ok() {
243                success_count += 1;
244            }
245        }
246
247        Ok(success_count)
248    }
249
250    pub async fn connection_count(&self) -> usize {
251        let connections = self.connections.read().await;
252        connections.len()
253    }
254
255    pub async fn room_count(&self, room: &str) -> usize {
256        let rooms = self.rooms.read().await;
257        rooms.get(room).map(|v| v.len()).unwrap_or(0)
258    }
259
260    pub async fn room_list(&self) -> Vec<String> {
261        let rooms = self.rooms.read().await;
262        rooms.keys().cloned().collect()
263    }
264
265    /// Returns all messages sent to the given connection, in delivery order.
266    pub async fn messages_for(&self, connection_id: &str) -> Vec<Vec<u8>> {
267        self.sender.messages_for(connection_id).await
268    }
269
270    /// Returns the number of messages sent to the given connection.
271    pub async fn message_count(&self, connection_id: &str) -> usize {
272        self.sender.message_count(connection_id).await
273    }
274}
275
276impl Default for RealtimePusher {
277    fn default() -> Self {
278        Self::new()
279    }
280}
281
282#[derive(Debug, Clone, Serialize, Deserialize)]
283pub struct PushResult {
284    pub total: usize,
285    pub success: usize,
286    pub failed: usize,
287}
288
289impl PushResult {
290    pub fn new(total: usize) -> Self {
291        Self {
292            total,
293            success: 0,
294            failed: 0,
295        }
296    }
297
298    pub fn add_success(&mut self) {
299        self.success += 1;
300    }
301
302    pub fn add_failure(&mut self) {
303        self.failed += 1;
304    }
305
306    pub fn success_rate(&self) -> f64 {
307        if self.total == 0 {
308            return 0.0;
309        }
310        (self.success as f64 / self.total as f64) * 100.0
311    }
312}
313
314#[cfg(test)]
315mod tests {
316    use super::*;
317
318    #[tokio::test]
319    async fn test_in_memory_sender_send_and_messages_for() {
320        let sender = InMemorySender::new();
321        sender.send("conn1", b"hello".to_vec()).await.unwrap();
322        sender.send("conn1", b"world".to_vec()).await.unwrap();
323
324        let messages = sender.messages_for("conn1").await;
325        assert_eq!(messages.len(), 2);
326        assert_eq!(messages[0], b"hello");
327        assert_eq!(messages[1], b"world");
328    }
329
330    #[tokio::test]
331    async fn test_in_memory_sender_message_count() {
332        let sender = InMemorySender::new();
333        assert_eq!(sender.message_count("conn1").await, 0);
334
335        sender.send("conn1", b"a".to_vec()).await.unwrap();
336        sender.send("conn1", b"b".to_vec()).await.unwrap();
337        assert_eq!(sender.message_count("conn1").await, 2);
338    }
339
340    #[tokio::test]
341    async fn test_in_memory_sender_close() {
342        let sender = InMemorySender::new();
343        sender.send("conn1", b"data".to_vec()).await.unwrap();
344        assert_eq!(sender.message_count("conn1").await, 1);
345
346        sender.close("conn1").await.unwrap();
347        assert_eq!(sender.message_count("conn1").await, 0);
348    }
349
350    #[tokio::test]
351    async fn test_in_memory_sender_isolation() {
352        let sender = InMemorySender::new();
353        sender.send("conn1", b"a".to_vec()).await.unwrap();
354        sender.send("conn2", b"b".to_vec()).await.unwrap();
355
356        assert_eq!(sender.messages_for("conn1").await.len(), 1);
357        assert_eq!(sender.messages_for("conn2").await.len(), 1);
358        assert_eq!(sender.messages_for("conn3").await.len(), 0);
359    }
360
361    #[tokio::test]
362    async fn test_push_to_user_filters_by_user_id() {
363        let pusher = RealtimePusher::new();
364        pusher.register_connection_with_user("conn1", 100).await;
365        pusher.register_connection_with_user("conn2", 200).await;
366        pusher.register_connection_with_user("conn3", 100).await;
367        pusher.register_connection("conn4").await;
368
369        let count = pusher
370            .push_to_user(100, b"hi-user-100".to_vec())
371            .await
372            .unwrap();
373        assert_eq!(count, 2);
374
375        assert_eq!(pusher.message_count("conn1").await, 1);
376        assert_eq!(pusher.message_count("conn2").await, 0);
377        assert_eq!(pusher.message_count("conn3").await, 1);
378        assert_eq!(pusher.message_count("conn4").await, 0);
379
380        let messages = pusher.messages_for("conn1").await;
381        assert_eq!(messages[0], b"hi-user-100");
382    }
383
384    #[tokio::test]
385    async fn test_push_to_user_no_matching_connections() {
386        let pusher = RealtimePusher::new();
387        pusher.register_connection_with_user("conn1", 100).await;
388
389        let count = pusher.push_to_user(999, b"nope".to_vec()).await.unwrap();
390        assert_eq!(count, 0);
391        assert_eq!(pusher.message_count("conn1").await, 0);
392    }
393
394    #[tokio::test]
395    async fn test_push_to_user_with_unregistered_connections() {
396        let pusher = RealtimePusher::new();
397        pusher.register_connection("conn1").await;
398        pusher.register_connection("conn2").await;
399
400        let count = pusher
401            .push_to_user(100, b"no-users".to_vec())
402            .await
403            .unwrap();
404        assert_eq!(count, 0);
405    }
406
407    #[tokio::test]
408    async fn test_push_to_connection_delivers_message() {
409        let pusher = RealtimePusher::new();
410        pusher.register_connection("conn1").await;
411
412        pusher
413            .push_to_connection("conn1", b"direct-msg".to_vec())
414            .await
415            .unwrap();
416
417        let messages = pusher.messages_for("conn1").await;
418        assert_eq!(messages.len(), 1);
419        assert_eq!(messages[0], b"direct-msg");
420    }
421
422    #[tokio::test]
423    async fn test_push_to_connection_not_found() {
424        let pusher = RealtimePusher::new();
425        let result = pusher.push_to_connection("missing", b"data".to_vec()).await;
426        assert!(result.is_err());
427        assert!(matches!(result.unwrap_err(), WsError::Connection(_)));
428    }
429
430    #[tokio::test]
431    async fn test_push_to_room_delivers_to_all_members() {
432        let pusher = RealtimePusher::new();
433        pusher.register_connection("conn1").await;
434        pusher.register_connection("conn2").await;
435        pusher.register_connection("conn3").await;
436
437        pusher.subscribe("conn1", "room1").await.unwrap();
438        pusher.subscribe("conn2", "room1").await.unwrap();
439        pusher.subscribe("conn3", "room2").await.unwrap();
440
441        let count = pusher
442            .push_to_room("room1", b"room-msg".to_vec())
443            .await
444            .unwrap();
445        assert_eq!(count, 2);
446
447        assert_eq!(pusher.message_count("conn1").await, 1);
448        assert_eq!(pusher.message_count("conn2").await, 1);
449        assert_eq!(pusher.message_count("conn3").await, 0);
450    }
451
452    #[tokio::test]
453    async fn test_push_to_room_empty_room() {
454        let pusher = RealtimePusher::new();
455        let count = pusher
456            .push_to_room("nonexistent", b"data".to_vec())
457            .await
458            .unwrap();
459        assert_eq!(count, 0);
460    }
461
462    #[tokio::test]
463    async fn test_broadcast_delivers_to_all() {
464        let pusher = RealtimePusher::new();
465        pusher.register_connection("conn1").await;
466        pusher.register_connection("conn2").await;
467
468        let count = pusher.broadcast(b"broadcast".to_vec()).await.unwrap();
469        assert_eq!(count, 2);
470
471        assert_eq!(pusher.messages_for("conn1").await[0], b"broadcast");
472        assert_eq!(pusher.messages_for("conn2").await[0], b"broadcast");
473    }
474
475    #[tokio::test]
476    async fn test_unregister_clears_messages() {
477        let pusher = RealtimePusher::new();
478        pusher.register_connection("conn1").await;
479        pusher
480            .push_to_connection("conn1", b"data".to_vec())
481            .await
482            .unwrap();
483        assert_eq!(pusher.message_count("conn1").await, 1);
484
485        pusher.unregister_connection("conn1").await;
486        assert_eq!(pusher.message_count("conn1").await, 0);
487        assert_eq!(pusher.connection_count().await, 0);
488    }
489
490    #[tokio::test]
491    async fn test_push_order_status_delivers_to_matching_user() {
492        let pusher = RealtimePusher::new();
493        pusher.register_connection_with_user("conn1", 123).await;
494        pusher.register_connection_with_user("conn2", 456).await;
495
496        let count = pusher.push_order_status(123, 789, "shipped").await.unwrap();
497        assert_eq!(count, 1);
498
499        let messages = pusher.messages_for("conn1").await;
500        assert_eq!(messages.len(), 1);
501        let parsed: serde_json::Value = serde_json::from_slice(&messages[0]).unwrap();
502        assert_eq!(parsed["type"], "order_status");
503        assert_eq!(parsed["order_id"], 789);
504        assert_eq!(parsed["status"], "shipped");
505    }
506
507    #[tokio::test]
508    async fn test_push_customer_message_delivers_to_room() {
509        let pusher = RealtimePusher::new();
510        pusher.register_connection_with_user("conn1", 100).await;
511        pusher.register_connection_with_user("conn2", 200).await;
512        pusher.subscribe("conn1", "room1").await.unwrap();
513        pusher.subscribe("conn2", "room1").await.unwrap();
514
515        let count = pusher
516            .push_customer_message("room1", 100, "hello room")
517            .await
518            .unwrap();
519        assert_eq!(count, 2);
520
521        for conn_id in &["conn1", "conn2"] {
522            let messages = pusher.messages_for(conn_id).await;
523            assert_eq!(messages.len(), 1);
524            let parsed: serde_json::Value = serde_json::from_slice(&messages[0]).unwrap();
525            assert_eq!(parsed["type"], "customer_message");
526            assert_eq!(parsed["sender_id"], 100);
527            assert_eq!(parsed["content"], "hello room");
528        }
529    }
530
531    #[tokio::test]
532    async fn test_multiple_pushes_to_same_connection() {
533        let pusher = RealtimePusher::new();
534        pusher.register_connection("conn1").await;
535
536        pusher
537            .push_to_connection("conn1", b"msg1".to_vec())
538            .await
539            .unwrap();
540        pusher
541            .push_to_connection("conn1", b"msg2".to_vec())
542            .await
543            .unwrap();
544        pusher
545            .push_to_connection("conn1", b"msg3".to_vec())
546            .await
547            .unwrap();
548
549        let messages = pusher.messages_for("conn1").await;
550        assert_eq!(messages.len(), 3);
551        assert_eq!(messages[0], b"msg1");
552        assert_eq!(messages[1], b"msg2");
553        assert_eq!(messages[2], b"msg3");
554    }
555
556    #[tokio::test]
557    async fn test_subscribe_not_found() {
558        let pusher = RealtimePusher::new();
559        let result = pusher.subscribe("missing", "room1").await;
560        assert!(result.is_err());
561        assert!(matches!(result.unwrap_err(), WsError::Connection(_)));
562    }
563}
sz_orm_websocket/pusher.rs

sz_orm_websocket/
pusher.rs
