endpoint_libs/libs/ws/
conn.rs1use 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}