1use std::collections::HashMap;
2use std::sync::Arc;
3
4use tokio::sync::RwLock;
5
6use crate::channel::{
7 Channel, ChannelConfig, EventHandler, EventRegistration, LeaveHandler, OutgoingHandler,
8 OutgoingRegistration,
9};
10use crate::errors::Result;
11use crate::pubsub::PubSub;
12use crate::types::Options;
13
14pub struct Lobby {
15 path: String,
16 endpoint_path: String,
17 options: Options,
18 pubsub: Option<Arc<dyn PubSub>>,
19 channels: RwLock<HashMap<String, Arc<Channel>>>,
20 event_handlers: RwLock<Vec<EventRegistration>>,
21 outgoing_handlers: RwLock<Vec<OutgoingRegistration>>,
22 leave_handler: RwLock<Option<Arc<dyn LeaveHandler>>>,
23}
24
25impl Lobby {
26 pub fn new(
27 path: impl Into<String>,
28 endpoint_path: impl Into<String>,
29 options: Options,
30 pubsub: Option<Arc<dyn PubSub>>,
31 ) -> Arc<Self> {
32 Arc::new(Self {
33 path: path.into(),
34 endpoint_path: endpoint_path.into(),
35 options,
36 pubsub,
37 channels: RwLock::new(HashMap::new()),
38 event_handlers: RwLock::new(Vec::new()),
39 outgoing_handlers: RwLock::new(Vec::new()),
40 leave_handler: RwLock::new(None),
41 })
42 }
43
44 pub fn path(&self) -> &str {
45 &self.path
46 }
47
48 pub async fn on_message<H>(&self, pattern: impl Into<String>, handler: H)
49 where
50 H: EventHandler + 'static,
51 {
52 self.event_handlers.write().await.push(EventRegistration {
53 pattern: pattern.into(),
54 handler: Arc::new(handler),
55 });
56 }
57
58 pub async fn on_outgoing<H>(&self, pattern: impl Into<String>, handler: H)
59 where
60 H: OutgoingHandler + 'static,
61 {
62 self.outgoing_handlers
63 .write()
64 .await
65 .push(OutgoingRegistration {
66 pattern: pattern.into(),
67 handler: Arc::new(handler),
68 });
69 }
70
71 pub async fn on_leave<H>(&self, handler: H)
72 where
73 H: LeaveHandler + 'static,
74 {
75 *self.leave_handler.write().await = Some(Arc::new(handler));
76 }
77
78 pub async fn get_channel(&self, name: &str) -> Option<Arc<Channel>> {
79 self.channels.read().await.get(name).cloned()
80 }
81
82 pub async fn get_or_create_channel(self: &Arc<Self>, name: &str) -> Result<Arc<Channel>> {
83 let mut channels = self.channels.write().await;
84 self.get_or_create_locked(&mut channels, name).await
85 }
86
87 pub async fn acquire_for_join(self: &Arc<Self>, name: &str) -> Result<Arc<Channel>> {
88 let mut channels = self.channels.write().await;
89 let channel = self.get_or_create_locked(&mut channels, name).await?;
90 channel.begin_join();
91 Ok(channel)
92 }
93
94 pub async fn finish_join(self: &Arc<Self>, channel: Arc<Channel>) {
95 channel.end_join();
96 self.remove_channel_if_idle(&channel.name).await;
97 }
98
99 pub async fn remove_channel_if_idle(self: &Arc<Self>, name: &str) {
100 let mut channels = self.channels.write().await;
101 let Some(channel) = channels.get(name).cloned() else {
102 return;
103 };
104 if !channel.is_idle().await {
105 return;
106 }
107 channels.remove(name);
108 drop(channels);
109 let _ = channel.close().await;
110 }
111
112 async fn get_or_create_locked(
113 self: &Arc<Self>,
114 channels: &mut HashMap<String, Arc<Channel>>,
115 name: &str,
116 ) -> Result<Arc<Channel>> {
117 if let Some(channel) = channels.get(name) {
118 return Ok(channel.clone());
119 }
120 let channel = Channel::new(ChannelConfig {
121 name: name.to_owned(),
122 endpoint_path: self.endpoint_path.clone(),
123 options: self.options.clone(),
124 pubsub: self.pubsub.clone(),
125 event_handlers: self.event_handlers.read().await.clone(),
126 outgoing_handlers: self.outgoing_handlers.read().await.clone(),
127 leave_handler: self.leave_handler.read().await.clone(),
128 lobby: Arc::downgrade(self),
129 });
130 channel.start().await?;
131 channels.insert(name.to_owned(), channel.clone());
132 Ok(channel)
133 }
134
135 pub async fn list_channels(&self) -> Vec<Arc<Channel>> {
136 self.channels.read().await.values().cloned().collect()
137 }
138
139 pub async fn close(&self) -> Result<()> {
140 let channels = self.list_channels().await;
141 for channel in channels {
142 let _ = channel.close().await;
143 }
144 self.channels.write().await.clear();
145 Ok(())
146 }
147}
148
149#[cfg(test)]
150mod tests {
151 use super::*;
152 use crate::transport::MemoryTransport;
153 use pondsocket_common::PondAssigns;
154
155 #[tokio::test]
156 async fn adv_channel_is_removed_after_last_user_leaves() {
157 let lobby = Lobby::new("/room", "/", Options::default(), None);
158 let channel = lobby.get_or_create_channel("/room/1").await.unwrap();
159 let transport = Arc::new(MemoryTransport::new("u1", PondAssigns::new()));
160 channel.add_user(transport).await.unwrap();
161 assert_eq!(channel.user_count().await, 1);
162
163 channel.remove_user("u1", "explicit_leave").await.unwrap();
164 assert_eq!(channel.user_count().await, 0);
165
166 assert!(lobby.list_channels().await.is_empty());
167 assert!(lobby.get_channel("/room/1").await.is_none());
168 }
169
170 #[tokio::test]
171 async fn adv_declined_join_leaves_no_running_channel() {
172 let lobby = Lobby::new("/room", "/", Options::default(), None);
173 let channel = lobby.acquire_for_join("/room/1").await.unwrap();
174 lobby.finish_join(channel).await;
175
176 assert!(lobby.get_channel("/room/1").await.is_none());
177 }
178
179 #[tokio::test]
180 async fn adv_concurrent_join_never_rejected_during_teardown() {
181 let lobby = Lobby::new("/room", "/", Options::default(), None);
182 let mut handles = Vec::new();
183 for i in 0..256 {
184 let lobby = lobby.clone();
185 handles.push(tokio::spawn(async move {
186 let user_id = format!("u{i}");
187 let channel = lobby.acquire_for_join("/room/1").await.unwrap();
188 let transport = Arc::new(MemoryTransport::new(user_id.clone(), PondAssigns::new()));
189 let added = channel.add_user(transport).await;
190 channel
191 .remove_user(&user_id, "explicit_leave")
192 .await
193 .unwrap();
194 lobby.finish_join(channel).await;
195 added
196 }));
197 }
198 for handle in handles {
199 handle
200 .await
201 .unwrap()
202 .expect("a legitimate join was rejected while a channel was torn down");
203 }
204 assert!(lobby.get_channel("/room/1").await.is_none());
205 }
206
207 #[derive(Default)]
208 struct CountingPubSub {
209 counts: std::sync::Mutex<HashMap<String, usize>>,
210 next_id: std::sync::atomic::AtomicU64,
211 }
212
213 impl CountingPubSub {
214 fn count(&self, pattern: &str) -> usize {
215 *self.counts.lock().unwrap().get(pattern).unwrap_or(&0)
216 }
217 }
218
219 #[async_trait::async_trait]
220 impl crate::pubsub::PubSub for CountingPubSub {
221 async fn subscribe(
222 &self,
223 pattern: &str,
224 _handler: crate::pubsub::PubSubHandler,
225 ) -> Result<crate::pubsub::SubscriptionId> {
226 *self
227 .counts
228 .lock()
229 .unwrap()
230 .entry(pattern.to_owned())
231 .or_insert(0) += 1;
232 Ok(self
233 .next_id
234 .fetch_add(1, std::sync::atomic::Ordering::SeqCst))
235 }
236
237 async fn unsubscribe(
238 &self,
239 pattern: &str,
240 _id: crate::pubsub::SubscriptionId,
241 ) -> Result<()> {
242 let mut counts = self.counts.lock().unwrap();
243 if let Some(count) = counts.get_mut(pattern) {
244 *count -= 1;
245 if *count == 0 {
246 counts.remove(pattern);
247 }
248 }
249 Ok(())
250 }
251
252 async fn publish(&self, _topic: &str, _data: Vec<u8>) -> Result<()> {
253 Ok(())
254 }
255
256 async fn close(&self) -> Result<()> {
257 Ok(())
258 }
259 }
260
261 #[tokio::test]
262 async fn adv_channel_teardown_keeps_other_heartbeat_subscription() {
263 let pubsub = Arc::new(CountingPubSub::default());
264 let dyn_pubsub: Arc<dyn PubSub> = pubsub.clone();
265 let lobby = Lobby::new("/room", "/api", Options::default(), Some(dyn_pubsub));
266 let _a = lobby.get_or_create_channel("/room/a").await.unwrap();
267 let _b = lobby.get_or_create_channel("/room/b").await.unwrap();
268
269 let heartbeat = crate::pubsub::format_heartbeat_topic("default");
270 assert_eq!(pubsub.count(&heartbeat), 2);
271
272 lobby.remove_channel_if_idle("/room/a").await;
273
274 assert!(lobby.get_channel("/room/a").await.is_none());
275 assert!(lobby.get_channel("/room/b").await.is_some());
276 assert_eq!(pubsub.count(&heartbeat), 1);
277 }
278
279 #[tokio::test]
280 async fn adv_sequential_recreate_after_teardown_resubscribes() {
281 let pubsub = Arc::new(CountingPubSub::default());
282 let dyn_pubsub: Arc<dyn PubSub> = pubsub.clone();
283 let lobby = Lobby::new("/room", "/api", Options::default(), Some(dyn_pubsub));
284
285 let _first = lobby.get_or_create_channel("/room/x").await.unwrap();
286 let heartbeat = crate::pubsub::format_heartbeat_topic("default");
287 assert_eq!(pubsub.count(&heartbeat), 1);
288
289 lobby.remove_channel_if_idle("/room/x").await;
290 assert_eq!(pubsub.count(&heartbeat), 0);
291
292 let _second = lobby.get_or_create_channel("/room/x").await.unwrap();
293 assert_eq!(pubsub.count(&heartbeat), 1);
294 }
295
296 #[tokio::test]
297 async fn adv_eviction_of_last_user_delivers_evicted_then_tears_down() {
298 let lobby = Lobby::new("/room", "/", Options::default(), None);
299 let channel = lobby.get_or_create_channel("/room/1").await.unwrap();
300 let transport = Arc::new(MemoryTransport::new("u1", PondAssigns::new()));
301 channel.add_user(transport.clone()).await.unwrap();
302
303 channel.evict_user("u1", "kicked").await.unwrap();
304
305 let event = transport
306 .recv()
307 .await
308 .expect("no frame delivered on eviction");
309 assert_eq!(event.event, "EVICTED");
310 assert!(lobby.get_channel("/room/1").await.is_none());
311 }
312
313 #[tokio::test]
314 async fn adv_teardown_stress_never_deadlocks() {
315 let lobby = Lobby::new("/room", "/", Options::default(), None);
316 let run = async {
317 let mut handles = Vec::new();
318 for i in 0..128 {
319 let lobby = lobby.clone();
320 handles.push(tokio::spawn(async move {
321 let user_id = format!("u{i}");
322 let channel = lobby.acquire_for_join("/room/1").await.unwrap();
323 let transport =
324 Arc::new(MemoryTransport::new(user_id.clone(), PondAssigns::new()));
325 let _ = channel.add_user(transport).await;
326 let _ = channel.broadcast("ping", serde_json::json!({})).await;
327 let _ = lobby.get_channel("/room/1").await;
328 let _ = lobby.list_channels().await;
329 channel
330 .remove_user(&user_id, "explicit_leave")
331 .await
332 .unwrap();
333 lobby.finish_join(channel).await;
334 }));
335 }
336 for handle in handles {
337 handle.await.unwrap();
338 }
339 };
340 tokio::time::timeout(std::time::Duration::from_secs(20), run)
341 .await
342 .expect("teardown stress deadlocked or hung");
343 assert!(lobby.get_channel("/room/1").await.is_none());
344 }
345}