├── .gitignore ├── Cargo.toml ├── LICENSE ├── examples └── basic.rs ├── README.md ├── Cargo.lock └── src └── lib.rs /.gitignore: -------------------------------------------------------------------------------- 1 | /target 2 | Cargo.lock 3 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "axum-ws-rooms" 3 | version = "0.7.0" 4 | edition = "2021" 5 | license = "MIT" 6 | description = "Room Manager for axum websocket" 7 | repository = "https://github.com/mohammadjavad948/axum-ws-rooms" 8 | keywords = ["axum", "axum-ws", "websocket"] 9 | 10 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 11 | 12 | [dependencies] 13 | tokio = {version = "1.42.0", features = ["rt", "sync", "macros"]} 14 | 15 | [dev-dependencies] 16 | tokio = {version = "1.42.0", features = ["rt", "rt-multi-thread", "sync", "macros", "time"]} 17 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright 2022-present Mohammad Javad Kianpour 2 | 3 | Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: 4 | 5 | The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. 6 | 7 | THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 8 | -------------------------------------------------------------------------------- /examples/basic.rs: -------------------------------------------------------------------------------- 1 | use axum_ws_rooms::RoomsManager; 2 | 3 | type Manager = RoomsManager; 4 | 5 | #[tokio::main] 6 | async fn main() { 7 | // init manager 8 | let manager = std::sync::Arc::new(Manager::new()); 9 | 10 | // create two rooms 11 | manager.new_room(1, None).await; 12 | manager.new_room(2, None).await; 13 | 14 | // spawn a task that acts as a receiver 15 | let receiver = tokio::spawn({ 16 | let manager = manager.clone(); 17 | 18 | async move { 19 | // initialise a user and join room 1 and start receiving messages 20 | let mut receiver = manager.init_user(1, None).await; 21 | 22 | let _ = manager.join_room(1, 1).await; 23 | 24 | while let Ok(data) = receiver.recv().await { 25 | println!("received data `{}`", data); 26 | } 27 | } 28 | }); 29 | 30 | // spawn a task that acts as sender 31 | let sender = tokio::spawn({ 32 | let manager = manager.clone(); 33 | 34 | async move { 35 | tokio::time::sleep(std::time::Duration::from_secs(1)).await; 36 | 37 | let _ = manager 38 | .send_message_to_room(&1, "test message to room 1".into()) 39 | .await; 40 | 41 | let _ = manager 42 | .send_message_to_room(&2, "test message to room 2".into()) 43 | .await; 44 | } 45 | }); 46 | 47 | let _ = tokio::join!(sender, receiver); 48 | } 49 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Axum Websocket Rooms 2 | 3 | room manager for websocket connections 4 | 5 | ```rust 6 | use axum_ws_rooms::RoomsManager; 7 | 8 | type Manager = RoomsManager; 9 | 10 | #[tokio::main] 11 | async fn main() { 12 | // init manager 13 | let manager = std::sync::Arc::new(Manager::new()); 14 | 15 | // create two rooms 16 | manager.new_room(1, None).await; 17 | manager.new_room(2, None).await; 18 | 19 | // spawn a task that acts as a receiver 20 | let receiver = tokio::spawn({ 21 | let manager = manager.clone(); 22 | 23 | async move { 24 | // initialise a user and join room 1 and start receiving messages 25 | let mut receiver = manager.init_user(1, None).await; 26 | 27 | let _ = manager.join_room(1, 1).await; 28 | 29 | while let Ok(data) = receiver.recv().await { 30 | println!("received data `{}`", data); 31 | } 32 | } 33 | }); 34 | 35 | // spawn a task that acts as sender 36 | let sender = tokio::spawn({ 37 | let manager = manager.clone(); 38 | 39 | async move { 40 | tokio::time::sleep(std::time::Duration::from_secs(1)).await; 41 | 42 | let _ = manager 43 | .send_message_to_room(&1, "test message to room 1".into()) 44 | .await; 45 | 46 | let _ = manager 47 | .send_message_to_room(&2, "test message to room 2".into()) 48 | .await; 49 | } 50 | }); 51 | 52 | let _ = tokio::join!(sender, receiver); 53 | } 54 | ``` 55 | 56 | Why? 57 | ------- 58 | 59 | For a project i wanted to group users into rooms and send messages to them so i created this library. 60 | 61 | Example 62 | -------- 63 | examples are available in `/examples` directory 64 | -------------------------------------------------------------------------------- /Cargo.lock: -------------------------------------------------------------------------------- 1 | # This file is automatically @generated by Cargo. 2 | # It is not intended for manual editing. 3 | version = 4 4 | 5 | [[package]] 6 | name = "addr2line" 7 | version = "0.24.2" 8 | source = "registry+https://github.com/rust-lang/crates.io-index" 9 | checksum = "dfbe277e56a376000877090da837660b4427aad530e3028d44e0bffe4f89a1c1" 10 | dependencies = [ 11 | "gimli", 12 | ] 13 | 14 | [[package]] 15 | name = "adler2" 16 | version = "2.0.0" 17 | source = "registry+https://github.com/rust-lang/crates.io-index" 18 | checksum = "512761e0bb2578dd7380c6baaa0f4ce03e84f95e960231d1dec8bf4d7d6e2627" 19 | 20 | [[package]] 21 | name = "axum-ws-rooms" 22 | version = "0.7.0" 23 | dependencies = [ 24 | "tokio", 25 | ] 26 | 27 | [[package]] 28 | name = "backtrace" 29 | version = "0.3.74" 30 | source = "registry+https://github.com/rust-lang/crates.io-index" 31 | checksum = "8d82cb332cdfaed17ae235a638438ac4d4839913cc2af585c3c6746e8f8bee1a" 32 | dependencies = [ 33 | "addr2line", 34 | "cfg-if", 35 | "libc", 36 | "miniz_oxide", 37 | "object", 38 | "rustc-demangle", 39 | "windows-targets", 40 | ] 41 | 42 | [[package]] 43 | name = "cfg-if" 44 | version = "1.0.0" 45 | source = "registry+https://github.com/rust-lang/crates.io-index" 46 | checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" 47 | 48 | [[package]] 49 | name = "gimli" 50 | version = "0.31.1" 51 | source = "registry+https://github.com/rust-lang/crates.io-index" 52 | checksum = "07e28edb80900c19c28f1072f2e8aeca7fa06b23cd4169cefe1af5aa3260783f" 53 | 54 | [[package]] 55 | name = "libc" 56 | version = "0.2.169" 57 | source = "registry+https://github.com/rust-lang/crates.io-index" 58 | checksum = "b5aba8db14291edd000dfcc4d620c7ebfb122c613afb886ca8803fa4e128a20a" 59 | 60 | [[package]] 61 | name = "memchr" 62 | version = "2.7.4" 63 | source = "registry+https://github.com/rust-lang/crates.io-index" 64 | checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3" 65 | 66 | [[package]] 67 | name = "miniz_oxide" 68 | version = "0.8.2" 69 | source = "registry+https://github.com/rust-lang/crates.io-index" 70 | checksum = "4ffbe83022cedc1d264172192511ae958937694cd57ce297164951b8b3568394" 71 | dependencies = [ 72 | "adler2", 73 | ] 74 | 75 | [[package]] 76 | name = "object" 77 | version = "0.36.7" 78 | source = "registry+https://github.com/rust-lang/crates.io-index" 79 | checksum = "62948e14d923ea95ea2c7c86c71013138b66525b86bdc08d2dcc262bdb497b87" 80 | dependencies = [ 81 | "memchr", 82 | ] 83 | 84 | [[package]] 85 | name = "pin-project-lite" 86 | version = "0.2.15" 87 | source = "registry+https://github.com/rust-lang/crates.io-index" 88 | checksum = "915a1e146535de9163f3987b8944ed8cf49a18bb0056bcebcdcece385cece4ff" 89 | 90 | [[package]] 91 | name = "proc-macro2" 92 | version = "1.0.92" 93 | source = "registry+https://github.com/rust-lang/crates.io-index" 94 | checksum = "37d3544b3f2748c54e147655edb5025752e2303145b5aefb3c3ea2c78b973bb0" 95 | dependencies = [ 96 | "unicode-ident", 97 | ] 98 | 99 | [[package]] 100 | name = "quote" 101 | version = "1.0.38" 102 | source = "registry+https://github.com/rust-lang/crates.io-index" 103 | checksum = "0e4dccaaaf89514f546c693ddc140f729f958c247918a13380cccc6078391acc" 104 | dependencies = [ 105 | "proc-macro2", 106 | ] 107 | 108 | [[package]] 109 | name = "rustc-demangle" 110 | version = "0.1.24" 111 | source = "registry+https://github.com/rust-lang/crates.io-index" 112 | checksum = "719b953e2095829ee67db738b3bfa9fa368c94900df327b3f07fe6e794d2fe1f" 113 | 114 | [[package]] 115 | name = "syn" 116 | version = "2.0.91" 117 | source = "registry+https://github.com/rust-lang/crates.io-index" 118 | checksum = "d53cbcb5a243bd33b7858b1d7f4aca2153490815872d86d955d6ea29f743c035" 119 | dependencies = [ 120 | "proc-macro2", 121 | "quote", 122 | "unicode-ident", 123 | ] 124 | 125 | [[package]] 126 | name = "tokio" 127 | version = "1.42.0" 128 | source = "registry+https://github.com/rust-lang/crates.io-index" 129 | checksum = "5cec9b21b0450273377fc97bd4c33a8acffc8c996c987a7c5b319a0083707551" 130 | dependencies = [ 131 | "backtrace", 132 | "pin-project-lite", 133 | "tokio-macros", 134 | ] 135 | 136 | [[package]] 137 | name = "tokio-macros" 138 | version = "2.4.0" 139 | source = "registry+https://github.com/rust-lang/crates.io-index" 140 | checksum = "693d596312e88961bc67d7f1f97af8a70227d9f90c31bba5806eec004978d752" 141 | dependencies = [ 142 | "proc-macro2", 143 | "quote", 144 | "syn", 145 | ] 146 | 147 | [[package]] 148 | name = "unicode-ident" 149 | version = "1.0.14" 150 | source = "registry+https://github.com/rust-lang/crates.io-index" 151 | checksum = "adb9e6ca4f869e1180728b7950e35922a7fc6397f7b641499e8f3ef06e50dc83" 152 | 153 | [[package]] 154 | name = "windows-targets" 155 | version = "0.52.6" 156 | source = "registry+https://github.com/rust-lang/crates.io-index" 157 | checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" 158 | dependencies = [ 159 | "windows_aarch64_gnullvm", 160 | "windows_aarch64_msvc", 161 | "windows_i686_gnu", 162 | "windows_i686_gnullvm", 163 | "windows_i686_msvc", 164 | "windows_x86_64_gnu", 165 | "windows_x86_64_gnullvm", 166 | "windows_x86_64_msvc", 167 | ] 168 | 169 | [[package]] 170 | name = "windows_aarch64_gnullvm" 171 | version = "0.52.6" 172 | source = "registry+https://github.com/rust-lang/crates.io-index" 173 | checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" 174 | 175 | [[package]] 176 | name = "windows_aarch64_msvc" 177 | version = "0.52.6" 178 | source = "registry+https://github.com/rust-lang/crates.io-index" 179 | checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" 180 | 181 | [[package]] 182 | name = "windows_i686_gnu" 183 | version = "0.52.6" 184 | source = "registry+https://github.com/rust-lang/crates.io-index" 185 | checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" 186 | 187 | [[package]] 188 | name = "windows_i686_gnullvm" 189 | version = "0.52.6" 190 | source = "registry+https://github.com/rust-lang/crates.io-index" 191 | checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" 192 | 193 | [[package]] 194 | name = "windows_i686_msvc" 195 | version = "0.52.6" 196 | source = "registry+https://github.com/rust-lang/crates.io-index" 197 | checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" 198 | 199 | [[package]] 200 | name = "windows_x86_64_gnu" 201 | version = "0.52.6" 202 | source = "registry+https://github.com/rust-lang/crates.io-index" 203 | checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" 204 | 205 | [[package]] 206 | name = "windows_x86_64_gnullvm" 207 | version = "0.52.6" 208 | source = "registry+https://github.com/rust-lang/crates.io-index" 209 | checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" 210 | 211 | [[package]] 212 | name = "windows_x86_64_msvc" 213 | version = "0.52.6" 214 | source = "registry+https://github.com/rust-lang/crates.io-index" 215 | checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" 216 | -------------------------------------------------------------------------------- /src/lib.rs: -------------------------------------------------------------------------------- 1 | use std::{ 2 | collections::HashMap, 3 | error::Error, 4 | fmt, 5 | sync::atomic::{AtomicU32, Ordering}, 6 | }; 7 | 8 | use tokio::{ 9 | sync::{broadcast, RwLock}, 10 | task::JoinHandle, 11 | }; 12 | 13 | static CHANNEL_SIZE: usize = 100; 14 | 15 | /// each room has a name and it contains `broadcast::sender` which can be accessed 16 | /// by `get_sender` method and you can send message to a roome by calling `send` on room. 17 | /// each room counts how many user it has and there is a method to check if its empty 18 | /// each room track its joined users and stores spawned tasks handlers 19 | struct Room { 20 | name: K, 21 | tx: broadcast::Sender, 22 | inner_user: RwLock>, 23 | user_count: AtomicU32, 24 | } 25 | 26 | /// struct that contains task handler that forwards messages 27 | struct UserTask { 28 | task: JoinHandle<()>, 29 | } 30 | 31 | /// use in combination with `Arc` to share it between threads 32 | /// 33 | /// internally it uses `RwLock` so it can handle concurrent requests without a problem 34 | /// 35 | /// when a user connects to ws endpoint you have to call `init_user` and it gives you a guard that 36 | /// when dropped will remove user from all rooms 37 | /// 38 | /// # Generics 39 | /// `K` is type used to identify each room 40 | /// 41 | /// `U` is type used to identify each user 42 | /// 43 | /// `T` is message type that is sent between rooms and users 44 | /// # Examples 45 | /// 46 | /// ```rust 47 | /// use axum_ws_rooms::RoomsManager; 48 | /// 49 | /// type Manager = RoomsManager; 50 | /// 51 | /// #[tokio::main] 52 | /// async fn main() { 53 | /// // init manager 54 | /// let manager = std::sync::Arc::new(Manager::new()); 55 | /// 56 | /// // create two rooms 57 | /// manager.new_room(1, None).await; 58 | /// manager.new_room(2, None).await; 59 | /// 60 | /// // spawn a task that acts as a receiver 61 | /// let receiver = tokio::spawn({ 62 | /// let manager = manager.clone(); 63 | /// 64 | /// async move { 65 | /// // initialise a user and join room 1 and start receiving messages 66 | /// let mut receiver = manager.init_user(1, None).await; 67 | /// 68 | /// let _ = manager.join_room(1, 1).await; 69 | /// 70 | /// while let Ok(data) = receiver.recv().await { 71 | /// println!("received data `{}`", data); 72 | /// } 73 | /// } 74 | /// }); 75 | /// 76 | /// // spawn a task that acts as sender 77 | /// let sender = tokio::spawn({ 78 | /// let manager = manager.clone(); 79 | /// 80 | /// async move { 81 | /// tokio::time::sleep(std::time::Duration::from_secs(1)).await; 82 | /// 83 | /// let _ = manager 84 | /// .send_message_to_room(&1, "test message to room 1".into()) 85 | /// .await; 86 | /// 87 | /// let _ = manager 88 | /// .send_message_to_room(&2, "test message to room 2".into()) 89 | /// .await; 90 | /// } 91 | /// }); 92 | /// 93 | /// let _ = tokio::join!(sender, receiver); 94 | /// } 95 | /// ``` 96 | pub struct RoomsManager { 97 | inner: RwLock>>, 98 | user_reciever: RwLock>>, 99 | } 100 | 101 | #[derive(Debug)] 102 | pub enum RoomError { 103 | /// room does not exists 104 | RoomNotFound, 105 | /// can not send message to room 106 | MessageSendFail, 107 | /// you have not called init_user 108 | NotInitiated, 109 | } 110 | 111 | impl Error for RoomError {} 112 | 113 | impl fmt::Display for RoomError { 114 | fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 115 | match self { 116 | RoomError::RoomNotFound => { 117 | write!(f, "target room not found") 118 | } 119 | RoomError::NotInitiated => { 120 | write!(f, "user is not initiated") 121 | } 122 | RoomError::MessageSendFail => { 123 | write!(f, "failed to send message to the room") 124 | } 125 | } 126 | } 127 | } 128 | 129 | /// guard that cntains users receiver 130 | /// 131 | /// when this guard drops it also removes the associated user from all rooms 132 | pub struct UserReceiverGuard<'a, K, U, T> 133 | where 134 | T: Clone + Send + 'static, 135 | K: Eq + std::hash::Hash + Clone, 136 | U: Eq + std::hash::Hash + Clone, 137 | { 138 | receiver: broadcast::Receiver, 139 | user: U, 140 | manager: &'a RoomsManager, 141 | } 142 | 143 | impl Room 144 | where 145 | T: Clone + Send + 'static, 146 | K: Eq + std::hash::Hash, 147 | U: Eq + std::hash::Hash, 148 | { 149 | /// creates new room with a given name 150 | /// capacity is the underlying channel capacity and its default is 100 151 | fn new(name: K, capacity: Option) -> Room { 152 | let (tx, _rx) = broadcast::channel(capacity.unwrap_or(CHANNEL_SIZE)); 153 | 154 | Room { 155 | name, 156 | tx, 157 | inner_user: RwLock::new(HashMap::new()), 158 | user_count: AtomicU32::new(0), 159 | } 160 | } 161 | 162 | /// join the rooms with a unique user 163 | /// if user has joined before, it does nothing 164 | async fn join(&self, user: U, user_sender: broadcast::Sender) { 165 | let mut inner = self.inner_user.write().await; 166 | 167 | match inner.entry(user) { 168 | std::collections::hash_map::Entry::Occupied(_) => {} 169 | std::collections::hash_map::Entry::Vacant(data) => { 170 | let mut room_rec = self.get_sender().subscribe(); 171 | 172 | let task = tokio::spawn(async move { 173 | while let Ok(data) = room_rec.recv().await { 174 | let _ = user_sender.send(data); 175 | } 176 | }); 177 | 178 | data.insert(UserTask { task }); 179 | 180 | self.user_count.fetch_add(1, Ordering::SeqCst); 181 | } 182 | } 183 | } 184 | 185 | /// leave the room with user 186 | /// if user has left before it wont do anything 187 | async fn leave(&self, user: U) { 188 | let mut inner = self.inner_user.write().await; 189 | 190 | match inner.entry(user) { 191 | std::collections::hash_map::Entry::Vacant(_) => {} 192 | std::collections::hash_map::Entry::Occupied(data) => { 193 | let data = data.remove(); 194 | 195 | data.task.abort(); 196 | 197 | self.user_count.fetch_sub(1, Ordering::SeqCst); 198 | } 199 | } 200 | } 201 | 202 | fn blocking_leave(&self, user: U) { 203 | let mut inner = self.inner_user.blocking_write(); 204 | 205 | match inner.entry(user) { 206 | std::collections::hash_map::Entry::Vacant(_) => {} 207 | std::collections::hash_map::Entry::Occupied(data) => { 208 | let data = data.remove(); 209 | 210 | data.task.abort(); 211 | 212 | self.user_count.fetch_sub(1, Ordering::SeqCst); 213 | } 214 | } 215 | } 216 | 217 | async fn clear_tasks(&self) { 218 | let mut inner = self.inner_user.write().await; 219 | 220 | inner.values().for_each(|value| { 221 | value.task.abort(); 222 | }); 223 | 224 | inner.clear(); 225 | 226 | self.user_count.store(0, Ordering::SeqCst); 227 | } 228 | 229 | /// check if user is in the room 230 | async fn contains_user(&self, user: &U) -> bool { 231 | let inner = self.inner_user.read().await; 232 | 233 | inner.contains_key(user) 234 | } 235 | 236 | /// checks if room is empty 237 | fn is_empty(&self) -> bool { 238 | self.user_count.load(Ordering::SeqCst) == 0 239 | } 240 | 241 | /// get sender without joining room 242 | fn get_sender(&self) -> broadcast::Sender { 243 | self.tx.clone() 244 | } 245 | 246 | ///send message to room 247 | fn send(&self, data: T) -> Result> { 248 | self.tx.send(data) 249 | } 250 | 251 | /// get user count of room 252 | async fn user_count(&self) -> u32 { 253 | self.user_count.load(Ordering::SeqCst) 254 | } 255 | } 256 | 257 | impl RoomsManager 258 | where 259 | T: Clone + Send + 'static, 260 | K: Eq + std::hash::Hash + Clone, 261 | U: Eq + std::hash::Hash + Clone, 262 | { 263 | /// create a room manager 264 | /// you can use it with combination of `Arc` to share it between threads 265 | /// # Example 266 | /// ```rust 267 | /// let manager = std::sync::Arc::new(RoomsManager::::new()); 268 | /// ``` 269 | pub fn new() -> Self { 270 | RoomsManager { 271 | inner: RwLock::new(HashMap::new()), 272 | user_reciever: RwLock::new(HashMap::new()), 273 | } 274 | } 275 | 276 | /// creates a new room with given name 277 | /// capacity is the underlying channels capacity and is set to `100` by default 278 | pub async fn new_room(&self, name: K, capacity: Option) { 279 | let mut rooms = self.inner.write().await; 280 | 281 | rooms.insert(name.clone(), Room::new(name, capacity)); 282 | } 283 | 284 | /// checks if given room exists 285 | pub async fn room_exists(&self, name: &K) -> bool { 286 | let rooms = self.inner.read().await; 287 | 288 | rooms.get(name).is_some() 289 | } 290 | 291 | /// joins specified user into room if it exists or creates a new room and joins immediately 292 | pub async fn join_or_create(&self, user: U, room: K) -> Result<(), RoomError> { 293 | match self.room_exists(&room).await { 294 | true => self.join_room(room, user).await, 295 | false => { 296 | self.new_room(room.clone(), None).await; 297 | 298 | self.join_room(room, user).await 299 | } 300 | } 301 | } 302 | 303 | /// send a message to a room 304 | /// 305 | /// it will fail if there are no users in the room or 306 | /// if room does not exists 307 | pub async fn send_message_to_room(&self, name: &K, data: T) -> Result { 308 | let rooms = self.inner.read().await; 309 | 310 | rooms 311 | .get(name) 312 | .ok_or(RoomError::RoomNotFound)? 313 | .send(data) 314 | .map_err(|_| RoomError::MessageSendFail) 315 | } 316 | 317 | /// call this at first of your code to initialize user notifier 318 | pub async fn init_user( 319 | &self, 320 | user: U, 321 | capacity: Option, 322 | ) -> UserReceiverGuard<'_, K, U, T> { 323 | let mut user_reciever = self.user_reciever.write().await; 324 | 325 | match user_reciever.entry(user.clone()) { 326 | std::collections::hash_map::Entry::Occupied(channel) => UserReceiverGuard { 327 | user, 328 | receiver: channel.get().subscribe(), 329 | manager: self, 330 | }, 331 | std::collections::hash_map::Entry::Vacant(v) => { 332 | let (tx, rx) = broadcast::channel(capacity.unwrap_or(CHANNEL_SIZE)); 333 | v.insert(tx); 334 | 335 | UserReceiverGuard { 336 | user, 337 | receiver: rx, 338 | manager: self, 339 | } 340 | } 341 | } 342 | } 343 | 344 | /// call this at end of your code to remove user from all rooms 345 | pub fn end_user(&self, user: U) { 346 | let rooms = self.inner.blocking_write(); 347 | let mut user_reciever = self.user_reciever.blocking_write(); 348 | 349 | for (_key, room) in rooms.iter() { 350 | room.blocking_leave(user.clone()); 351 | } 352 | 353 | match user_reciever.entry(user.clone()) { 354 | std::collections::hash_map::Entry::Occupied(o) => { 355 | o.remove(); 356 | } 357 | std::collections::hash_map::Entry::Vacant(_) => {} 358 | } 359 | } 360 | 361 | /// join user to room 362 | pub async fn join_room(&self, name: K, user: U) -> Result<(), RoomError> { 363 | let rooms = self.inner.read().await; 364 | let user_reciever = self.user_reciever.read().await; 365 | 366 | let user_reciever = user_reciever 367 | .get(&user) 368 | .ok_or(RoomError::NotInitiated)? 369 | .clone(); 370 | 371 | rooms 372 | .get(&name) 373 | .ok_or(RoomError::RoomNotFound)? 374 | .join(user.clone(), user_reciever) 375 | .await; 376 | 377 | Ok(()) 378 | } 379 | 380 | pub async fn remove_room(&self, room: K) { 381 | let mut rooms = self.inner.write().await; 382 | 383 | match rooms.entry(room.clone()) { 384 | std::collections::hash_map::Entry::Vacant(_) => {} 385 | std::collections::hash_map::Entry::Occupied(el) => { 386 | let room = el.remove(); 387 | 388 | room.clear_tasks().await; 389 | } 390 | } 391 | } 392 | 393 | pub async fn leave_room(&self, name: K, user: U) -> Result<(), RoomError> { 394 | let rooms = self.inner.read().await; 395 | 396 | rooms 397 | .get(&name) 398 | .ok_or(RoomError::RoomNotFound)? 399 | .leave(user.clone()) 400 | .await; 401 | 402 | Ok(()) 403 | } 404 | 405 | pub async fn is_room_empty(&self, name: K) -> Result { 406 | let rooms = self.inner.read().await; 407 | 408 | Ok(rooms.get(&name).ok_or(RoomError::RoomNotFound)?.is_empty()) 409 | } 410 | 411 | pub async fn rooms_count(&self) -> usize { 412 | let rooms = self.inner.read().await; 413 | 414 | rooms.len() 415 | } 416 | } 417 | 418 | impl Default for RoomsManager 419 | where 420 | T: Clone + Send + 'static, 421 | K: Eq + std::hash::Hash + Clone, 422 | U: Eq + std::hash::Hash + Clone, 423 | { 424 | fn default() -> Self { 425 | Self::new() 426 | } 427 | } 428 | 429 | impl std::ops::Deref for UserReceiverGuard<'_, K, U, T> 430 | where 431 | T: Clone + Send + 'static, 432 | K: Eq + std::hash::Hash + Clone, 433 | U: Eq + std::hash::Hash + Clone, 434 | { 435 | type Target = broadcast::Receiver; 436 | 437 | fn deref(&self) -> &Self::Target { 438 | &self.receiver 439 | } 440 | } 441 | 442 | impl std::ops::DerefMut for UserReceiverGuard<'_, K, U, T> 443 | where 444 | T: Clone + Send + 'static, 445 | K: Eq + std::hash::Hash + Clone, 446 | U: Eq + std::hash::Hash + Clone, 447 | { 448 | fn deref_mut(&mut self) -> &mut Self::Target { 449 | &mut self.receiver 450 | } 451 | } 452 | 453 | impl Drop for UserReceiverGuard<'_, K, U, T> 454 | where 455 | T: Clone + Send + 'static, 456 | K: Eq + std::hash::Hash + Clone, 457 | U: Eq + std::hash::Hash + Clone, 458 | { 459 | fn drop(&mut self) { 460 | self.manager.end_user(self.user.clone()); 461 | } 462 | } 463 | --------------------------------------------------------------------------------