From 36ecb48c7b054a1c2a7c449fed9143324c7d3d52 Mon Sep 17 00:00:00 2001 From: Paul Makles Date: Wed, 19 Jun 2024 18:24:51 +0100 Subject: [PATCH] feat(bonfire): don't fan out UserUpdate (server) by default --- Cargo.lock | 84 +++++++--- crates/bonfire/Cargo.toml | 2 + crates/bonfire/src/events/impl.rs | 10 ++ crates/bonfire/src/events/state.rs | 45 +++++- crates/bonfire/src/websocket.rs | 182 +++++++++++++--------- crates/core/database/src/events/client.rs | 7 +- crates/core/database/src/events/server.rs | 1 + 7 files changed, 237 insertions(+), 94 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index ff5ccdb2..947600e4 100644 --- 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", ] diff --git a/crates/bonfire/Cargo.toml b/crates/bonfire/Cargo.toml index bc06f367..41d6ec28 100644 --- 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" diff --git a/crates/bonfire/src/events/impl.rs b/crates/bonfire/src/events/impl.rs index 45d9a262..46c854b8 100644 --- 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 { let user = self.clone_user(); + self.cache.is_bot = user.bot.is_some(); // Find all relationships to the user. let mut user_ids: HashSet = 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 { diff --git a/crates/bonfire/src/events/state.rs b/crates/bonfire/src/events/state.rs index 988d5143..eae39801 100644 --- 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, pub channels: HashMap, @@ -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>>, + pub active_servers: Arc>>, } 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 = { + 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 { diff --git a/crates/bonfire/src/websocket.rs b/crates/bonfire/src/websocket.rs index 5727ad75..5e9f0aac 100644 --- 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, ) { 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::(s.as_ref()).ok()), - PayloadType::Msgpack => message - .value - .as_bytes() - .and_then(|b| rmp_serde::from_slice::(b).ok()), - PayloadType::Bincode => message - .value - .as_bytes() - .and_then(|b| bincode::deserialize::(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::()) - ); + 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::(s.as_ref()).ok()), + PayloadType::Msgpack => message + .value + .as_bytes() + .and_then(|b| rmp_serde::from_slice::(b).ok()), + PayloadType::Bincode => message + .value + .as_bytes() + .and_then(|b| bincode::deserialize::(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::()) + ); + + 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>>, + active_servers: Arc>>, user_id: String, config: &ProtocolConfiguration, + signal_s: async_channel::Sender<()>, mut read: WsReader, write: &Mutex, ) { @@ -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 diff --git a/crates/core/database/src/events/client.rs b/crates/core/database/src/events/client.rs index 68802af7..dd20dd4f 100644 --- 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; diff --git a/crates/core/database/src/events/server.rs b/crates/core/database/src/events/server.rs index 883b92c3..cdaf6418 100644 --- 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<()> }, }