Add proper shutdown, rate limits, config validation, cors config, fix dangling cursors, cache regex, merge created texts

sync-server/src/app_state/websocket/broadcasts.rs

@@ -1,35 +1,52 @@
 use std::{collections::HashMap, sync::Arc};
 
-use anyhow::Context;
 use log::{debug, warn};
 use tokio::sync::{Mutex, broadcast};
 
 use super::models::WebSocketServerMessageWithOrigin;
-use crate::{
-    app_state::database::models::VaultId, config::server_config::ServerConfig, errors::server_error,
-};
+use crate::{app_state::database::models::VaultId, config::server_config::ServerConfig};
 
 #[derive(Debug, Clone)]
 pub struct Broadcasts {
-    max_clients_per_vault: usize,
+    broadcast_channel_capacity: usize,
     tx: Arc<Mutex<HashMap<VaultId, broadcast::Sender<WebSocketServerMessageWithOrigin>>>>,
 }
 
+type TxMap = HashMap<VaultId, broadcast::Sender<WebSocketServerMessageWithOrigin>>;
+
 impl Broadcasts {
     pub fn new(server_config: &ServerConfig) -> Self {
         Self {
-            max_clients_per_vault: server_config.max_clients_per_vault,
+            broadcast_channel_capacity: server_config.broadcast_channel_capacity,
             tx: Arc::new(Mutex::new(HashMap::new())),
         }
     }
 
+    /// Remove senders for vaults with no active receivers
+    fn prune_inactive_vaults(tx_map: &mut TxMap) {
+        tx_map.retain(|_, sender| sender.receiver_count() > 0);
+    }
+
     pub async fn get_receiver(
         &self,
         vault: VaultId,
-    ) -> broadcast::Receiver<WebSocketServerMessageWithOrigin> {
-        let tx = self.get_or_create(vault).await;
+        max_clients: usize,
+    ) -> Result<broadcast::Receiver<WebSocketServerMessageWithOrigin>, crate::errors::SyncServerError>
+    {
+        let mut tx_map = self.tx.lock().await;
+        Self::prune_inactive_vaults(&mut tx_map);
 
-        tx.subscribe()
+        let sender = tx_map
+            .entry(vault)
+            .or_insert_with(|| broadcast::channel(self.broadcast_channel_capacity).0);
+
+        if sender.receiver_count() >= max_clients {
+            return Err(crate::errors::client_error(anyhow::anyhow!(
+                "Vault has reached the maximum number of clients ({max_clients})"
+            )));
+        }
+
+        Ok(sender.subscribe())
     }
 
     /// Notify all clients (who are subscribed to the vault) about an update.
@@ -39,31 +56,20 @@ impl Broadcasts {
         vault: VaultId,
         document: WebSocketServerMessageWithOrigin,
     ) {
-        let tx = self.get_or_create(vault.clone()).await;
+        let mut tx_map = self.tx.lock().await;
+        Self::prune_inactive_vaults(&mut tx_map);
 
-        if tx.receiver_count() == 0 {
+        let sender = tx_map
+            .entry(vault.clone())
+            .or_insert_with(|| broadcast::channel(self.broadcast_channel_capacity).0);
+
+        if sender.receiver_count() == 0 {
             debug!("Skipping broadcast, no clients connected for vault `{vault}`");
             return;
         }
 
-        let result = tx
-            .send(document)
-            .context("Cannot broadcast server message to websocket listeners")
-            .map_err(server_error);
-
-        if result.is_err() {
-            warn!("Failed to send message: {result:?}");
+        if let Err(e) = sender.send(document) {
+            warn!("Failed to broadcast to vault `{vault}`: {e}");
         }
     }
-
-    async fn get_or_create(
-        &self,
-        vault: VaultId,
-    ) -> broadcast::Sender<WebSocketServerMessageWithOrigin> {
-        let mut tx = self.tx.lock().await;
-
-        tx.entry(vault)
-            .or_insert_with(|| broadcast::channel(self.max_clients_per_vault).0.clone())
-            .clone()
-    }
 }