diff --git a/sync-server/src/app_state/websocket/broadcasts.rs b/sync-server/src/app_state/websocket/broadcasts.rs index 1b157b13..4980ca10 100644 --- a/sync-server/src/app_state/websocket/broadcasts.rs +++ b/sync-server/src/app_state/websocket/broadcasts.rs @@ -12,23 +12,42 @@ pub enum Notification { #[derive(Debug, Clone)] pub struct Broadcasts { capacity: usize, + admission: Arc>>>, senders: Arc>>>, } impl Broadcasts { pub fn new(config: &ServerConfig) -> Self { Self { - capacity: config.max_clients_per_vault.max(1), + capacity: config.max_clients_per_vault, + admission: Arc::default(), senders: Arc::default(), } } + pub fn try_admit(&self, vault: &VaultId) -> Option { + let mut limits = self.admission.lock().unwrap(); + + limits.retain(|_, limit| limit.strong_count() > 0); + + let limit = limits + .get(vault) + .and_then(std::sync::Weak::upgrade) + .unwrap_or_else(|| { + let limit = Arc::new(tokio::sync::Semaphore::new(self.capacity)); + limits.insert(vault.clone(), Arc::downgrade(&limit)); + limit + }); + + limit.try_acquire_owned().ok() + } + async fn sender(&self, vault: VaultId) -> broadcast::Sender { self.senders .lock() .await .entry(vault) - .or_insert_with(|| broadcast::channel(self.capacity).0) + .or_insert_with(|| broadcast::channel(self.capacity.max(1)).0) .clone() } diff --git a/sync-server/src/server/endpoints/websocket.rs b/sync-server/src/server/endpoints/websocket.rs index ad1fe8c0..0000c53a 100644 --- a/sync-server/src/server/endpoints/websocket.rs +++ b/sync-server/src/server/endpoints/websocket.rs @@ -32,7 +32,16 @@ pub async fn websocket_handler( ) -> Result { debug!("Upgrading WebSocket connection for vault `{vault_id}`"); + let Some(permit) = state.broadcasts.try_admit(&vault_id) else { + return Ok(( + StatusCode::TOO_MANY_REQUESTS, + "Vault connection limit reached", + ) + .into_response()); + }; + Ok(ws.on_upgrade(move |socket| async move { + let _permit = permit; if let Err(error) = websocket(state, socket, vault_id).await { debug!("WebSocket disconnected: {error}"); }