HMC/handmade-revolt-backend
4
0
Fork 0
forked from jmug/stoatchat
Code Pull Requests Actions Packages Releases 2 Activity

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

Browse Source
This commit is contained in:
Paul Makles
2024-06-19 18:24:51 +01:00
parent ba411fef43
commit 36ecb48c7b
7 changed files with 237 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
crates/bonfire/Cargo.toml
View File

@@ -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"

10
crates/bonfire/src/events/impl.rs
View File

@@ -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 {

45
crates/bonfire/src/events/state.rs
View File

@@ -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 {

182
crates/bonfire/src/websocket.rs
View File

@@ -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