endpoint_libs/libs/ws/
conn.rs

1use dashmap::DashMap;
2use std::sync::Arc;
3
4use tokio_tungstenite::tungstenite::Message;
5
6use super::{ConnectionId, WsConnection};
7
8#[derive(Default)]
9pub struct WebsocketStates {
10    states: Arc<DashMap<ConnectionId, Arc<WsStreamState>>>,
11}
12
13impl WebsocketStates {
14    pub fn new() -> Self {
15        WebsocketStates::default()
16    }
17    pub fn remove(&self, connection_id: u32) {
18        self.states.remove(&connection_id);
19    }
20
21    pub fn get_state(&self, connection_id: u32) -> Option<Arc<WsStreamState>> {
22        self.states.get(&connection_id).map(|x| x.value().clone())
23    }
24    pub fn clone_states(&self) -> Arc<DashMap<u32, Arc<WsStreamState>>> {
25        Arc::clone(&self.states)
26    }
27    pub fn insert(
28        &self,
29        connection_id: u32,
30        message_queue: tokio::sync::mpsc::Sender<Message>,
31        conn: Arc<WsConnection>,
32    ) {
33        self.states.insert(
34            connection_id,
35            Arc::new(WsStreamState {
36                conn,
37                message_queue,
38            }),
39        );
40    }
41}
42
43pub struct WsStreamState {
44    pub conn: Arc<WsConnection>,
45    pub message_queue: tokio::sync::mpsc::Sender<Message>,
46}