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
34            .insert(connection_id, Arc::new(WsStreamState { conn, message_queue }));
35    }
36}
37
38pub struct WsStreamState {
39    pub conn: Arc<WsConnection>,
40    pub message_queue: tokio::sync::mpsc::Sender<Message>,
41}