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
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}