feat(bonfire): don't fan out UserUpdate (server) by default

2026-07-14 05:26:59 +00:00 · 2024-06-19 18:24:51 +01:00
parent ba411fef43
commit 36ecb48c7b
7 changed files with 237 additions and 94 deletions
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -138,11 +138,23 @@ version = "1.6.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "2114d64672151c0c5eaa5e131ec84a74f06e1e559830dabba01ca30605d66319"
 dependencies = [
- "concurrent-queue",
- "event-listener",
+ "concurrent-queue 1.2.2",
+ "event-listener 2.5.2",
 "futures-core",
 ]

+[[package]]
+name = "async-channel"
+version = "2.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "89b47800b0be77592da0afd425cc03468052844aff33b84e33cc696f64e77b6a"
+dependencies = [
+ "concurrent-queue 2.5.0",
+ "event-listener-strategy",
+ "futures-core",
+ "pin-project-lite 0.2.13",
+]
+
 [[package]]
 name = "async-executor"
 version = "1.4.1"
@@ -150,7 +162,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "871f9bb5e0a22eeb7e8cf16641feb87c9dc67032ccf8ff49e772eb9941d3a965"
 dependencies = [
 "async-task",
- "concurrent-queue",
+ "concurrent-queue 1.2.2",
 "fastrand",
 "futures-lite",
 "once_cell",
@@ -163,7 +175,7 @@ version = "2.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "fd8b508d585e01084059b60f06ade4cb7415cd2e4084b71dd1cb44e7d3fb9880"
 dependencies = [
- "async-channel",
+ "async-channel 1.6.1",
 "async-executor",
 "async-io",
 "async-lock",
@@ -180,7 +192,7 @@ version = "1.7.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e5e18f61464ae81cde0a23e713ae8fd299580c54d697a35820cfd0625b8b0e07"
 dependencies = [
- "concurrent-queue",
+ "concurrent-queue 1.2.2",
 "futures-lite",
 "libc",
 "log",
@@ -199,7 +211,7 @@ version = "2.8.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "287272293e9d8c41773cec55e365490fe034813a2f172f502d6ddcf75b2f582b"
 dependencies = [
- "event-listener",
+ "event-listener 2.5.2",
 ]

 [[package]]
@@ -211,7 +223,7 @@ dependencies = [
 "async-io",
 "blocking",
 "cfg-if",
- "event-listener",
+ "event-listener 2.5.2",
 "futures-lite",
 "libc",
 "once_cell",
@@ -237,7 +249,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "62565bb4402e926b29953c785397c6dc0391b7b446e45008b0049eb43cec6f5d"
 dependencies = [
 "async-attributes",
- "async-channel",
+ "async-channel 1.6.1",
 "async-global-executor",
 "async-io",
 "async-lock",
@@ -531,7 +543,7 @@ version = "1.2.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "c6ccb65d468978a086b69884437ded69a90faab3bbe6e67f242173ea728acccc"
 dependencies = [
- "async-channel",
+ "async-channel 1.6.1",
 "async-task",
 "atomic-waker",
 "fastrand",
@@ -729,6 +741,15 @@ dependencies = [
 "cache-padded",
 ]

+[[package]]
+name = "concurrent-queue"
+version = "2.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4ca0197aee26d1ae37445ee532fefce43251d24cc7c166799f4d46817f1d3973"
+dependencies = [
+ "crossbeam-utils",
+]
+
 [[package]]
 name = "config"
 version = "0.13.3"
@@ -833,13 +854,9 @@ dependencies = [

 [[package]]
 name = "crossbeam-utils"
-version = "0.8.8"
+version = "0.8.20"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0bf124c720b7686e3c2663cf54062ab0f68a88af2fb6a030e87e30bf721fcb38"
-dependencies = [
- "cfg-if",
- "lazy_static",
-]
+checksum = "22ec99545bb0ed0ea7bb9b8e1e9122ea386ff8a48c0922e43f36d45ab09e0e80"

 [[package]]
 name = "crypto-bigint"
@@ -1295,6 +1312,27 @@ version = "2.5.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "77f3309417938f28bf8228fcff79a4a37103981e3e186d2ccd19c74b38f4eb71"

+[[package]]
+name = "event-listener"
+version = "5.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6032be9bd27023a771701cc49f9f053c751055f71efb2e0ae5c15809093675ba"
+dependencies = [
+ "concurrent-queue 2.5.0",
+ "parking",
+ "pin-project-lite 0.2.13",
+]
+
+[[package]]
+name = "event-listener-strategy"
+version = "0.5.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0f214dc438f977e6d4e3500aaa277f5ad94ca83fbbd9b1a15713ce2344ccc5a1"
+dependencies = [
+ "event-listener 5.3.1",
+ "pin-project-lite 0.2.13",
+]
+
 [[package]]
 name = "fake-simd"
 version = "0.1.2"
@@ -1976,13 +2014,13 @@ version = "1.7.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "334e04b4d781f436dc315cb1e7515bd96826426345d498149e4bde36b67f8ee9"
 dependencies = [
- "async-channel",
+ "async-channel 1.6.1",
 "castaway",
 "crossbeam-utils",
 "curl",
 "curl-sys",
 "encoding_rs",
- "event-listener",
+ "event-listener 2.5.2",
 "futures-lite",
 "http",
 "log",
@@ -2266,6 +2304,12 @@ dependencies = [
 "linked-hash-map",
 ]

+[[package]]
+name = "lru_time_cache"
+version = "0.11.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9106e1d747ffd48e6be5bb2d97fa706ed25b144fbee4d5c02eae110cd8d6badd"
+
 [[package]]
 name = "maplit"
 version = "1.0.2"
@@ -3410,6 +3454,7 @@ dependencies = [
 name = "revolt-bonfire"
 version = "0.7.5"
 dependencies = [
+ "async-channel 2.3.1",
 "async-std",
 "async-tungstenite",
 "authifier",
@@ -3418,6 +3463,7 @@ dependencies = [
 "futures",
 "log",
 "lru 0.7.6",
+ "lru_time_cache",
 "once_cell",
 "querystring",
 "redis-kiss",
@@ -3500,7 +3546,7 @@ dependencies = [
 name = "revolt-delta"
 version = "0.7.5"
 dependencies = [
- "async-channel",
+ "async-channel 1.6.1",
 "async-std",
 "authifier",
 "bitfield",
@@ -4362,7 +4408,7 @@ version = "0.5.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "6d7400c0eff44aa2fcb5e31a5f24ba9716ed90138769e4977a2ba6014ae63eb5"
 dependencies = [
- "async-channel",
+ "async-channel 1.6.1",
 "futures-core",
 "futures-io",
 ]
--- a/crates/bonfire/Cargo.toml
+++ b/crates/bonfire/Cargo.toml
@@ -14,6 +14,8 @@ lru = "0.7.6"
 ulid = "0.5.0"
 once_cell = "1.9.0"
 redis-kiss = "0.1.4"
+lru_time_cache = "0.11.11"
+async-channel = "2.3.1"

 # parsing
 querystring = "1.1.0"
--- a/crates/bonfire/src/events/impl.rs
+++ b/crates/bonfire/src/events/impl.rs
@@ -92,6 +92,7 @@ impl State {
    /// Generate a Ready packet for the current user
    pub async fn generate_ready_payload(&mut self, db: &Database) -> Result<EventV1> {
        let user = self.clone_user();
+        self.cache.is_bot = user.bot.is_some();

        // Find all relationships to the user.
        let mut user_ids: HashSet<String> = user
@@ -191,6 +192,10 @@ impl State {

        for server in &servers {
            self.insert_subscription(server.id.clone()).await;
+
+            if self.cache.is_bot {
+                self.insert_subscription(format!("{}u", server.id)).await;
+            }
        }

        for channel in &channels {
@@ -397,6 +402,11 @@ impl State {
                emojis: _,
            } => {
                self.insert_subscription(id.clone()).await;
+
+                if self.cache.is_bot {
+                    self.insert_subscription(format!("{}u", id)).await;
+                }
+
                self.cache.servers.insert(id.clone(), server.clone().into());
                let member = Member {
                    id: MemberCompositeKey {
--- a/crates/bonfire/src/events/state.rs
+++ b/crates/bonfire/src/events/state.rs
@@ -1,10 +1,12 @@
 use std::{
    collections::{HashMap, HashSet},
    sync::Arc,
+    time::Duration,
 };

-use async_std::sync::RwLock;
+use async_std::sync::{Mutex, RwLock};
 use lru::LruCache;
+use lru_time_cache::{LruCache as LruTimeCache, TimedEntry};
 use revolt_database::{Channel, Member, Server, User};

 /// Enumeration representing some change in subscriptions
@@ -34,6 +36,7 @@ pub enum SubscriptionStateChange {
 #[derive(Debug)]
 pub struct Cache {
    pub user_id: String,
+    pub is_bot: bool,

    pub users: HashMap<String, User>,
    pub channels: HashMap<String, Channel>,
@@ -47,6 +50,7 @@ impl Default for Cache {
    fn default() -> Self {
        Cache {
            user_id: Default::default(),
+            is_bot: false,

            users: Default::default(),
            channels: Default::default(),
@@ -65,7 +69,9 @@ pub struct State {
    pub session_id: String,
    pub private_topic: String,
    pub state: SubscriptionStateChange,
+
    pub subscribed: Arc<RwLock<HashSet<String>>>,
+    pub active_servers: Arc<Mutex<LruTimeCache<String, ()>>>,
 }

 impl State {
@@ -86,6 +92,10 @@ impl State {
        State {
            cache,
            subscribed: Arc::new(RwLock::new(subscribed)),
+            active_servers: Arc::new(Mutex::new(LruTimeCache::with_expiry_duration_and_capacity(
+                Duration::from_secs(900),
+                5,
+            ))),
            session_id,
            private_topic,
            state: SubscriptionStateChange::Reset,
@@ -94,6 +104,39 @@ impl State {

    /// Apply currently queued state
    pub async fn apply_state(&mut self) -> SubscriptionStateChange {
+        // Check if we need to change subscriptions to member event topics
+        if !self.cache.is_bot {
+            enum Server {
+                Subscribe(String),
+                Unsubscribe(String),
+            }
+
+            let active_server_changes: Vec<Server> = {
+                let mut active_servers = self.active_servers.lock().await;
+                active_servers
+                    .notify_iter()
+                    .map(|e| match e {
+                        TimedEntry::Valid(k, _) => Server::Subscribe(format!("{}u", k)),
+                        TimedEntry::Expired(k, _) => Server::Unsubscribe(format!("{}u", k)),
+                    })
+                    .collect()
+                // It is bad practice to open more than one Mutex at once and could
+                // lead to a deadlock, so instead we choose to collect the changes.
+            };
+
+            for entry in active_server_changes {
+                match entry {
+                    Server::Subscribe(k) => {
+                        self.insert_subscription(k).await;
+                    }
+                    Server::Unsubscribe(k) => {
+                        self.remove_subscription(&k).await;
+                    }
+                }
+            }
+        }
+
+        // Flush changes to subscriptions
        let state = std::mem::replace(&mut self.state, SubscriptionStateChange::None);
        let mut subscribed = self.subscribed.write().await;
        if let SubscriptionStateChange::Change { add, remove } = &state {
--- a/crates/bonfire/src/websocket.rs
+++ b/crates/bonfire/src/websocket.rs
@@ -129,11 +129,23 @@ pub async fn client(db: &'static Database, stream: TcpStream, addr: SocketAddr)
    {
        let write = Mutex::new(write);
        let subscribed = state.subscribed.clone();
+        let active_servers = state.active_servers.clone();
+        let (signal_s, signal_r) = async_channel::unbounded();

        // Create a PubSub connection to poll on.
-        let listener = listener(db, &mut state, addr, &config, &write).fuse();
+        let listener = listener(db, &mut state, addr, &config, signal_r, &write).fuse();
        // Read from WebSocket stream.
-        let worker = worker(addr, subscribed, user_id.clone(), &config, read, &write).fuse();
+        let worker = worker(
+            addr,
+            subscribed,
+            active_servers,
+            user_id.clone(),
+            &config,
+            signal_s,
+            read,
+            &write,
+        )
+        .fuse();

        // Pin both tasks.
        pin_mut!(listener, worker);
@@ -158,6 +170,7 @@ async fn listener(
    state: &mut State,
    addr: SocketAddr,
    config: &ProtocolConfiguration,
+    signal_r: async_channel::Receiver<()>,
    write: &Mutex<WsWriter>,
 ) {
    let redis_config = RedisConfig::from_url(&REDIS_URI).unwrap();
@@ -227,94 +240,107 @@ async fn listener(
            SubscriptionStateChange::None => {}
        }

-        // Handle incoming events.
-        let message = match message_rx.recv().await {
-            Ok(message) => message,
-            Err(e) => {
-                error!("Error while consuming pub/sub messages: {e:?}");
-                sentry::capture_error(&e);
-                return;
-            }
-        };
+        let t1 = message_rx.recv().fuse();
+        let t2 = signal_r.recv().fuse();

-        let event = match *REDIS_PAYLOAD_TYPE {
-            PayloadType::Json => message
-                .value
-                .as_str()
-                .and_then(|s| serde_json::from_str::<EventV1>(s.as_ref()).ok()),
-            PayloadType::Msgpack => message
-                .value
-                .as_bytes()
-                .and_then(|b| rmp_serde::from_slice::<EventV1>(b).ok()),
-            PayloadType::Bincode => message
-                .value
-                .as_bytes()
-                .and_then(|b| bincode::deserialize::<EventV1>(b).ok()),
-        };
+        pin_mut!(t1, t2);

-        let Some(mut event) = event else {
-            let err = format!(
-                "Failed to deserialise an event for {}! Introspection: `{:?}`",
-                message.channel,
-                message
-                    .value
-                    .as_string()
-                    .map(|x| x.chars().take(32).collect::<String>())
-            );
+        select! {
+            _ = t2 => {},
+            message = t1 => {
+                // Handle incoming events.
+                let message = match message {
+                    Ok(message) => message,
+                    Err(e) => {
+                        error!("Error while consuming pub/sub messages: {e:?}");
+                        sentry::capture_error(&e);
+                        return;
+                    }
+                };

-            error!("{}", err);
-            sentry::capture_message(&err, Level::Error);
-            return;
-        };
+                let event = match *REDIS_PAYLOAD_TYPE {
+                    PayloadType::Json => message
+                        .value
+                        .as_str()
+                        .and_then(|s| serde_json::from_str::<EventV1>(s.as_ref()).ok()),
+                    PayloadType::Msgpack => message
+                        .value
+                        .as_bytes()
+                        .and_then(|b| rmp_serde::from_slice::<EventV1>(b).ok()),
+                    PayloadType::Bincode => message
+                        .value
+                        .as_bytes()
+                        .and_then(|b| bincode::deserialize::<EventV1>(b).ok()),
+                };

-        if let EventV1::Auth(auth) = &event {
-            if let AuthifierEvent::DeleteSession { session_id, .. } = auth {
-                if &state.session_id == session_id {
-                    event = EventV1::Logout;
-                }
-            } else if let AuthifierEvent::DeleteAllSessions {
-                exclude_session_id, ..
-            } = auth
-            {
-                if let Some(excluded) = exclude_session_id {
-                    if &state.session_id != excluded {
-                        event = EventV1::Logout;
+                let Some(mut event) = event else {
+                    let err = format!(
+                        "Failed to deserialise an event for {}! Introspection: `{:?}`",
+                        message.channel,
+                        message
+                            .value
+                            .as_string()
+                            .map(|x| x.chars().take(32).collect::<String>())
+                    );
+
+                    error!("{}", err);
+                    sentry::capture_message(&err, Level::Error);
+                    return;
+                };
+
+                if let EventV1::Auth(auth) = &event {
+                    if let AuthifierEvent::DeleteSession { session_id, .. } = auth {
+                        if &state.session_id == session_id {
+                            event = EventV1::Logout;
+                        }
+                    } else if let AuthifierEvent::DeleteAllSessions {
+                        exclude_session_id, ..
+                    } = auth
+                    {
+                        if let Some(excluded) = exclude_session_id {
+                            if &state.session_id != excluded {
+                                event = EventV1::Logout;
+                            }
+                        } else {
+                            event = EventV1::Logout;
+                        }
                    }
                } else {
-                    event = EventV1::Logout;
+                    let should_send = state.handle_incoming_event_v1(db, &mut event).await;
+                    if !should_send {
+                        continue;
+                    }
+                }
+
+                let result = write.lock().await.send(config.encode(&event)).await;
+                if let Err(e) = result {
+                    use async_tungstenite::tungstenite::Error;
+                    if !matches!(e, Error::AlreadyClosed | Error::ConnectionClosed) {
+                        let err = format!("Error while sending an event to {addr:?}: {e:?}");
+                        warn!("{}", err);
+                        sentry::capture_message(&err, Level::Warning);
+                    }
+
+                    return;
+                }
+
+                if let EventV1::Logout = event {
+                    info!("User {addr:?} received log out event!");
+                    return;
                }
            }
-        } else {
-            let should_send = state.handle_incoming_event_v1(db, &mut event).await;
-            if !should_send {
-                continue;
-            }
-        }
-
-        let result = write.lock().await.send(config.encode(&event)).await;
-        if let Err(e) = result {
-            use async_tungstenite::tungstenite::Error;
-            if !matches!(e, Error::AlreadyClosed | Error::ConnectionClosed) {
-                let err = format!("Error while sending an event to {addr:?}: {e:?}");
-                warn!("{}", err);
-                sentry::capture_message(&err, Level::Warning);
-            }
-
-            return;
-        }
-
-        if let EventV1::Logout = event {
-            info!("User {addr:?} received log out event!");
-            return;
        }
    }
 }

+#[allow(clippy::too_many_arguments)]
 async fn worker(
    addr: SocketAddr,
    subscribed: Arc<RwLock<HashSet<String>>>,
+    active_servers: Arc<Mutex<lru_time_cache::LruCache<String, ()>>>,
    user_id: String,
    config: &ProtocolConfiguration,
+    signal_s: async_channel::Sender<()>,
    mut read: WsReader,
    write: &Mutex<WsWriter>,
 ) {
@@ -368,6 +394,16 @@ async fn worker(
                .p(channel.clone())
                .await;
            }
+            ClientMessage::Subscribe { server_id } => {
+                let mut servers = active_servers.lock().await;
+                let has_item = servers.contains_key(&server_id);
+                servers.insert(server_id, ());
+
+                if !has_item {
+                    // Poke the listener to adjust subscriptions
+                    signal_s.send(()).await.ok();
+                }
+            }
            ClientMessage::Ping { data, responded } => {
                if responded.is_none() {
                    write
--- a/crates/core/database/src/events/client.rs
+++ b/crates/core/database/src/events/client.rs
@@ -249,7 +249,7 @@ impl EventV1 {
        // TODO: this should be captured by member list in the future and not immediately fanned out to users
        if let Ok(members) = db.fetch_all_memberships(&id).await {
            for member in members {
-                self.clone().p(member.id.server).await;
+                self.clone().server(member.id.server).await;
            }
        }
    }
@@ -259,6 +259,11 @@ impl EventV1 {
        self.p(format!("{id}!")).await;
    }

+    /// Publish server member event
+    pub async fn server(self, id: String) {
+        self.p(format!("{id}u")).await;
+    }
+
    /// Publish internal global event
    pub async fn global(self) {
        self.p("global".to_string()).await;
--- a/crates/core/database/src/events/server.rs
+++ b/crates/core/database/src/events/server.rs
@@ -8,5 +8,6 @@ pub enum ClientMessage {
    Authenticate { token: String },
    BeginTyping { channel: String },
    EndTyping { channel: String },
+    Subscribe { server_id: String },
    Ping { data: Ping, responded: Option<()> },
 }