jmug/stoatchat
1
0
Fork 1
Code Issues Packages Projects Releases Wiki Activity

feat: switch to RabbitMQ streams from Redis Pub/Sub

Browse Source 
refactor: redo the bonfire client structure
This commit is contained in:
izzy
2025-06-28 13:35:18 +01:00
parent aab1734615
commit 85989a2138
13 changed files with 1095 additions and 512 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
crates/bonfire/Cargo.toml
View File

@@ -9,6 +9,7 @@ edition = "2021"
[dependencies]
# util
log = "*"
rand = "*"
sentry = "0.31.5"
lru = "0.7.6"
ulid = "0.5.0"
@@ -46,3 +47,6 @@ revolt-presence = { path = "../core/presence", features = ["redis-is-patched"] }
# redis
fred = { version = "8.0.1", features = ["subscriber-client"] }
# rabbit
lapin = { version = "3.0.0" }

129
crates/bonfire/src/client/core.rs Normal file
View File

@@ -0,0 +1,129 @@
use async_std::{net::TcpStream, sync::Mutex};
use async_tungstenite::WebSocketStream;
use futures::{join, SinkExt, StreamExt, TryStreamExt};
use revolt_config::report_internal_error;
use revolt_database::{
events::{client::EventV1, server::ClientMessage},
iso8601_timestamp::Timestamp,
Database, User, UserHint,
};
use revolt_presence::{create_session, delete_session};
use revolt_result::create_error;
use crate::{
client::{
subscriber::client_subscriber,
worker::{client_worker, WorkerRef},
},
config::ProtocolConfiguration,
events::state::State,
};
/// Core event loop of gateway clients
pub async fn client_core(
db: &'static Database,
ws: WebSocketStream<TcpStream>,
mut config: ProtocolConfiguration,
) {
// Split the socket for simultaneously read and write.
let (mut write, mut read) = ws.split();
// If the user has not provided authentication, request information.
if config.get_session_token().is_none() {
while let Ok(Some(message)) = read.try_next().await {
if let Ok(ClientMessage::Authenticate { token }) = config.decode(&message) {
config.set_session_token(token);
break;
}
}
}
// Try to authenticate the user.
let Some(token) = config.get_session_token().as_ref() else {
write
.send(config.encode(&EventV1::Error {
data: create_error!(InvalidSession),
}))
.await
.ok();
return;
};
let (user, session_id) = match User::from_token(db, token, UserHint::Any).await {
Ok(user) => user,
Err(err) => {
write
.send(config.encode(&EventV1::Error { data: err }))
.await
.ok();
return;
}
};
info!(
"Authenticated user {}#{}",
user.username, user.discriminator
);
db.update_session_last_seen(&session_id, Timestamp::now_utc())
.await
.ok();
// Create local state.
let mut state = State::from(user, session_id);
let user_id = state.cache.user_id.clone();
// Notify socket we have authenticated.
if report_internal_error!(write.send(config.encode(&EventV1::Authenticated)).await).is_err() {
return;
}
// Download required data to local cache and send Ready payload.
let ready_payload = match report_internal_error!(
state
.generate_ready_payload(db, config.get_ready_payload_fields())
.await
) {
Ok(ready_payload) => ready_payload,
Err(_) => return,
};
if report_internal_error!(write.send(config.encode(&ready_payload)).await).is_err() {
return;
}
// Create presence session.
let (first_session, session_id) = create_session(&user_id, 0).await;
// If this was the first session, notify other users that we just went online.
if first_session {
state.broadcast_presence_change(true).await;
}
{
let worker_ref = WorkerRef::from(&state);
let write = Mutex::new(write);
let (reload, reloaded) = async_channel::bounded(1);
let (cancel_1, cancelled_1) = async_channel::bounded(1);
let (cancel_2, cancelled_2) = async_channel::bounded(1);
join!(
async {
client_subscriber(&write, cancelled_1, reloaded, &config, db, &mut state).await;
cancel_2.send(()).await.ok();
},
async {
client_worker(read, &write, cancelled_2, reload, &config, worker_ref).await;
cancel_1.send(()).await.ok();
}
);
}
// Clean up presence session.
let last_session = delete_session(&user_id, session_id).await;
// If this was the last session, notify other users that we just went offline.
if last_session {
state.broadcast_presence_change(false).await;
}
}

3
crates/bonfire/src/client/mod.rs Normal file
View File

@@ -0,0 +1,3 @@
pub mod core;
pub mod subscriber;
pub mod worker;

203
crates/bonfire/src/client/subscriber.rs Normal file
View File

@@ -0,0 +1,203 @@
use std::collections::HashMap;
use async_channel::Receiver;
use async_std::{net::TcpStream, sync::Mutex};
use async_tungstenite::WebSocketStream;
use authifier::AuthifierEvent;
use futures::{pin_mut, select, stream::SplitSink, FutureExt, SinkExt, StreamExt};
use lapin::{
options::BasicAckOptions,
types::{AMQPValue, FieldArray, FieldTable, LongLongInt},
};
use rand::Rng;
use revolt_config::report_internal_error;
use revolt_database::{
events::client::{get_event_stream_channel, EventV1},
Database,
};
use sentry::Level;
use crate::{
config::ProtocolConfiguration,
events::state::{State, SubscriptionStateChange},
};
static QUEUE_NAME: &str = "revolt.events";
/// Event subscriber loop
pub async fn client_subscriber(
write: &Mutex<SplitSink<WebSocketStream<TcpStream>, async_tungstenite::tungstenite::Message>>,
cancelled: Receiver<()>,
reloaded: Receiver<()>,
protocol_config: &ProtocolConfiguration,
db: &'static Database,
state: &mut State,
) {
// * --- CHANNEL CREATE ---
let channel = get_event_stream_channel().await;
// * --- CHANNEL END ---
let mut offset: Option<LongLongInt> = None;
let mut cancel = false;
loop {
// Build arguments for consumer
let mut args: FieldTable = Default::default();
// Configure stream filter to select topics we are listening for
{
let mut filter: FieldArray = Default::default();
for topic in state.subscribed.read().await.iter() {
filter.push(AMQPValue::LongString(topic.as_str().into()));
}
args.insert("x-stream-filter".into(), AMQPValue::FieldArray(filter));
}
// Set stream offset if applicable
if let Some(offset) = offset {
args.insert("x-stream-offset".into(), AMQPValue::LongLongInt(offset));
}
// Create the consumer
let tag: String = rand::thread_rng()
.sample_iter::<char, _>(&rand::distributions::Standard)
.take(32)
.collect();
let mut consumer = channel
.basic_consume(QUEUE_NAME, &tag, Default::default(), args)
.await
.unwrap();
// Read incoming events
loop {
let reloaded = reloaded.recv().fuse();
let cancelled = cancelled.recv().fuse();
let delivery = consumer.next().fuse();
pin_mut!(delivery, reloaded, cancelled);
select! {
_ = reloaded => {
break;
}
_ = cancelled => {
cancel = true;
break;
}
delivery = delivery => {
let delivery = delivery.expect("error in consumer").expect("error in consumer");
// Acknowledgement is required
delivery.ack(BasicAckOptions::default()).await.expect("ack");
// Parse the delivery headers
let headers: HashMap<String, AMQPValue> = delivery
.properties
.headers()
.as_ref()
.map(|table| {
table
.into_iter()
.map(|(k, v)| (k.to_string(), v.clone()))
.collect()
})
.unwrap_or_default();
// Client-side topic filtering (broker uses Bloom filter so may have false-positives)
let filter_value = headers
.get("x-stream-filter-value")
.expect("`x-stream-filter-value` not present in message!");
if state
.subscribed
.read()
.await
.contains(
&filter_value
.as_long_string()
.expect("`string`")
.to_string()
) {
// Deserialise the data
let mut event: EventV1 = rmp_serde::from_slice(&delivery.data).expect("`data`");
// Handle the event
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 {
let should_send = state.handle_incoming_event_v1(db, &mut event).await;
if !should_send {
continue;
}
}
let result = write.lock().await.send(protocol_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: {e:?}");
warn!("{}", err);
sentry::capture_message(&err, Level::Warning);
}
cancel = true;
break;
}
if let EventV1::Logout = event {
info!("User {} received log out event!", state.user_id);
cancel = true;
break;
}
}
// Keep track of the current offset
let stream_offset = headers
.get("x-stream-offset")
.expect("`x-stream-offset` not present in message!");
offset = Some(stream_offset.as_long_long_int().unwrap() + 1);
// Reload consumer if subscriptions change
if !matches!(state.apply_state().await, SubscriptionStateChange::None) {
break;
}
}
}
}
// Close the consumer
if let Err(err) = channel.basic_cancel(&tag, Default::default()).await {
eprintln!("Failed to close consumer! {:?}", err);
} else {
// Read the consumer to the end
while let Some(delivery) = consumer.next().await {
let delivery = delivery.expect("error in consumer");
delivery.ack(BasicAckOptions::default()).await.expect("ack");
}
}
// Break out if cancelled
if cancel {
break;
}
}
report_internal_error!(channel.close(0, "closing channel").await).ok();
}

125
crates/bonfire/src/client/worker.rs Normal file
View File

@@ -0,0 +1,125 @@
use std::{collections::HashSet, sync::Arc};
use async_channel::{Receiver, Sender};
use async_std::{
net::TcpStream,
sync::{Mutex, RwLock},
};
use async_tungstenite::WebSocketStream;
use futures::{
pin_mut, select,
stream::{SplitSink, SplitStream},
FutureExt, SinkExt, TryStreamExt,
};
use revolt_database::events::{client::EventV1, server::ClientMessage};
use sentry::Level;
use crate::{config::ProtocolConfiguration, events::state::State};
pub struct WorkerRef {
user_id: String,
active_servers: Arc<Mutex<lru_time_cache::LruCache<String, ()>>>,
subscribed: Arc<RwLock<HashSet<String>>>,
}
impl WorkerRef {
pub fn from(state: &State) -> WorkerRef {
WorkerRef {
user_id: state.user_id.clone(),
active_servers: state.active_servers.clone(),
subscribed: state.subscribed.clone(),
}
}
}
/// Incoming message handling
pub async fn client_worker(
mut read: SplitStream<WebSocketStream<TcpStream>>,
write: &Mutex<SplitSink<WebSocketStream<TcpStream>, async_tungstenite::tungstenite::Message>>,
cancelled: Receiver<()>,
reload: Sender<()>,
config: &ProtocolConfiguration,
state: WorkerRef,
) {
loop {
let read = read.try_next().fuse();
let cancelled = cancelled.recv().fuse();
pin_mut!(read, cancelled);
select! {
_ = cancelled => { return; },
msg = read => {
let msg = match msg {
Ok(Some(msg)) => msg,
Ok(None) => {
warn!("Received a None message!");
sentry::capture_message("Received a None message!", Level::Warning);
return;
}
Err(e) => {
use async_tungstenite::tungstenite::Error;
if !matches!(e, Error::AlreadyClosed | Error::ConnectionClosed) {
let err = format!("Error while reading an event: {e:?}");
warn!("{}", err);
sentry::capture_message(&err, Level::Warning);
}
return;
}
};
let Ok(payload) = config.decode(&msg) else {
continue;
};
match payload {
ClientMessage::BeginTyping { channel } => {
if !state.subscribed.read().await.contains(&channel) {
continue;
}
EventV1::ChannelStartTyping {
id: channel.clone(),
user: state.user_id.clone(),
}
.p(channel.clone())
.await;
}
ClientMessage::EndTyping { channel } => {
if !state.subscribed.read().await.contains(&channel) {
continue;
}
EventV1::ChannelStopTyping {
id: channel.clone(),
user: state.user_id.clone(),
}
.p(channel.clone())
.await;
}
ClientMessage::Subscribe { server_id } => {
let mut servers = state.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
reload.send(()).await.ok();
}
}
ClientMessage::Ping { data, responded } => {
if responded.is_none() {
write
.lock()
.await
.send(config.encode(&EventV1::Pong { data }))
.await
.ok();
}
}
_ => {}
}
}
}
}
}

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

@@ -66,6 +66,7 @@ impl Default for Cache {
pub struct State {
pub cache: Cache,
pub user_id: String,
pub session_id: String,
pub private_topic: String,
pub state: SubscriptionStateChange,
@@ -87,6 +88,7 @@ impl State {
..Default::default()
};
let user_id = user.id.clone();
cache.users.insert(user.id.clone(), user);
State {
@@ -96,6 +98,7 @@ impl State {
Duration::from_secs(900),
5,
))),
user_id,
session_id,
private_topic,
state: SubscriptionStateChange::Reset,

5
crates/bonfire/src/main.rs
View File

@@ -6,10 +6,11 @@ use revolt_presence::clear_region;
#[macro_use]
extern crate log;
pub mod config;
pub mod events;
pub mod client;
mod config;
mod database;
mod events;
mod websocket;
#[async_std::main]

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

@@ -1,40 +1,11 @@
use std::{collections::HashSet, net::SocketAddr, sync::Arc};
use std::net::SocketAddr;
use async_tungstenite::WebSocketStream;
use authifier::AuthifierEvent;
use fred::{
error::RedisErrorKind,
interfaces::{ClientLike, EventInterface, PubsubInterface},
types::RedisConfig,
};
use futures::{
channel::oneshot,
join, pin_mut, select,
stream::{SplitSink, SplitStream},
FutureExt, SinkExt, StreamExt, TryStreamExt,
};
use redis_kiss::{PayloadType, REDIS_PAYLOAD_TYPE, REDIS_URI};
use revolt_config::report_internal_error;
use revolt_database::{
events::{client::EventV1, server::ClientMessage},
iso8601_timestamp::Timestamp,
Database, User, UserHint,
};
use revolt_presence::{create_session, delete_session};
use async_std::net::TcpStream;
use futures::channel::oneshot;
use revolt_database::Database;
use async_std::{
net::TcpStream,
sync::{Mutex, RwLock},
task::spawn,
};
use revolt_result::create_error;
use sentry::Level;
use crate::config::{ProtocolConfiguration, WebsocketHandshakeCallback};
use crate::events::state::{State, SubscriptionStateChange};
type WsReader = SplitStream<WebSocketStream<TcpStream>>;
type WsWriter = SplitSink<WebSocketStream<TcpStream>, async_tungstenite::tungstenite::Message>;
use crate::client::core::client_core;
use crate::config::WebsocketHandshakeCallback;
/// Start a new WebSocket client worker given access to the database,
/// the relevant TCP stream and the remote address of the client.
@@ -54,7 +25,7 @@ pub async fn client(db: &'static Database, stream: TcpStream, addr: SocketAddr)
};
// Verify we've received a valid config, otherwise we should just drop the connection.
let Ok(mut config) = receiver.await else {
let Ok(config) = receiver.await else {
return;
};
@@ -64,446 +35,5 @@ pub async fn client(db: &'static Database, stream: TcpStream, addr: SocketAddr)
config.get_protocol_format()
);
// Split the socket for simultaneously read and write.
let (mut write, mut read) = ws.split();
// If the user has not provided authentication, request information.
if config.get_session_token().is_none() {
while let Ok(Some(message)) = read.try_next().await {
if let Ok(ClientMessage::Authenticate { token }) = config.decode(&message) {
config.set_session_token(token);
break;
}
}
}
// Try to authenticate the user.
let Some(token) = config.get_session_token().as_ref() else {
write
.send(config.encode(&EventV1::Error {
data: create_error!(InvalidSession),
}))
.await
.ok();
return;
};
let (user, session_id) = match User::from_token(db, token, UserHint::Any).await {
Ok(user) => user,
Err(err) => {
write
.send(config.encode(&EventV1::Error { data: err }))
.await
.ok();
return;
}
};
info!("User {addr:?} authenticated as @{}", user.username);
db.update_session_last_seen(&session_id, Timestamp::now_utc())
.await
.ok();
// Create local state.
let mut state = State::from(user, session_id);
let user_id = state.cache.user_id.clone();
// Notify socket we have authenticated.
if report_internal_error!(write.send(config.encode(&EventV1::Authenticated)).await).is_err() {
return;
}
// Download required data to local cache and send Ready payload.
let ready_payload = match report_internal_error!(
state
.generate_ready_payload(db, config.get_ready_payload_fields())
.await
) {
Ok(ready_payload) => ready_payload,
Err(_) => return,
};
if report_internal_error!(write.send(config.encode(&ready_payload)).await).is_err() {
return;
}
// Create presence session.
let (first_session, session_id) = create_session(&user_id, 0).await;
// If this was the first session, notify other users that we just went online.
if first_session {
state.broadcast_presence_change(true).await;
}
{
// Setup channels and mutexes
let write = Mutex::new(write);
let subscribed = state.subscribed.clone();
let active_servers = state.active_servers.clone();
let (topic_signal_s, topic_signal_r) = async_channel::unbounded();
// TODO: this needs to be rewritten
// Create channels through which the tasks can signal to each other they need to clean up
let (kill_signal_1_s, kill_signal_1_r) = async_channel::bounded(1);
let (kill_signal_2_s, kill_signal_2_r) = async_channel::bounded(1);
// Create a PubSub connection to poll on.
let listener = listener_with_kill_signal(
db,
&mut state,
addr,
&config,
topic_signal_r,
kill_signal_1_r,
&write,
kill_signal_2_s,
);
// Read from WebSocket stream.
let worker = worker_with_kill_signal(
addr,
subscribed,
active_servers,
user_id.clone(),
&config,
topic_signal_s,
kill_signal_2_r,
read,
&write,
kill_signal_1_s,
);
join!(listener, worker);
}
// Clean up presence session.
let last_session = delete_session(&user_id, session_id).await;
// If this was the last session, notify other users that we just went offline.
if last_session {
state.broadcast_presence_change(false).await;
}
}
#[allow(clippy::too_many_arguments)]
async fn listener_with_kill_signal(
db: &'static Database,
state: &mut State,
addr: SocketAddr,
config: &ProtocolConfiguration,
topic_signal_r: async_channel::Receiver<()>,
kill_signal_r: async_channel::Receiver<()>,
write: &Mutex<WsWriter>,
kill_signal_s: async_channel::Sender<()>,
) {
listener(
db,
state,
addr,
config,
topic_signal_r,
kill_signal_r,
write,
)
.await;
kill_signal_s.send(()).await.ok();
}
async fn listener(
db: &'static Database,
state: &mut State,
addr: SocketAddr,
config: &ProtocolConfiguration,
topic_signal_r: async_channel::Receiver<()>,
kill_signal_r: async_channel::Receiver<()>,
write: &Mutex<WsWriter>,
) {
let redis_config = RedisConfig::from_url(&REDIS_URI).unwrap();
let subscriber = match report_internal_error!(
fred::types::Builder::from_config(redis_config).build_subscriber_client()
) {
Ok(subscriber) => subscriber,
Err(_) => return,
};
if report_internal_error!(subscriber.init().await).is_err() {
return;
}
// Handle Redis connection dropping
let (clean_up_s, clean_up_r) = async_channel::bounded(1);
let clean_up_s = Arc::new(Mutex::new(clean_up_s));
subscriber.on_error(move |err| {
if let RedisErrorKind::Canceled = err.kind() {
let clean_up_s = clean_up_s.clone();
spawn(async move {
clean_up_s.lock().await.send(()).await.ok();
});
}
Ok(())
});
let mut message_rx = subscriber.message_rx();
'out: loop {
// Check for state changes for subscriptions.
match state.apply_state().await {
SubscriptionStateChange::Reset => {
if report_internal_error!(subscriber.unsubscribe_all().await).is_err() {
break 'out;
}
let subscribed = state.subscribed.read().await;
for id in subscribed.iter() {
if report_internal_error!(subscriber.subscribe(id).await).is_err() {
break 'out;
}
}
#[cfg(debug_assertions)]
info!("{addr:?} has reset their subscriptions");
}
SubscriptionStateChange::Change { add, remove } => {
for id in remove {
#[cfg(debug_assertions)]
info!("{addr:?} unsubscribing from {id}");
if report_internal_error!(subscriber.unsubscribe(id).await).is_err() {
break 'out;
}
}
for id in add {
#[cfg(debug_assertions)]
info!("{addr:?} subscribing to {id}");
if report_internal_error!(subscriber.subscribe(id).await).is_err() {
break 'out;
}
}
}
SubscriptionStateChange::None => {}
}
let t1 = message_rx.recv().fuse();
let t2 = topic_signal_r.recv().fuse();
let t3 = kill_signal_r.recv().fuse();
let t4 = clean_up_r.recv().fuse();
pin_mut!(t1, t2, t3, t4);
select! {
_ = t4 => {
break 'out;
},
_ = t3 => {
break 'out;
},
_ = t2 => {},
message = t1 => {
// Handle incoming events.
let message = match report_internal_error!(message) {
Ok(message) => message,
Err(_) => break 'out
};
let event = match *REDIS_PAYLOAD_TYPE {
PayloadType::Json => message
.value
.as_str()
.and_then(|s| report_internal_error!(serde_json::from_str::<EventV1>(s.as_ref())).ok()),
PayloadType::Msgpack => message
.value
.as_bytes()
.and_then(|b| report_internal_error!(rmp_serde::from_slice::<EventV1>(b)).ok()),
PayloadType::Bincode => message
.value
.as_bytes()
.and_then(|b| report_internal_error!(bincode::deserialize::<EventV1>(b)).ok()),
};
let Some(mut event) = event else {
let err = format!(
"Failed to deserialise event for {}: `{:?}`",
message.channel,
message
.value
);
error!("{}", err);
sentry::capture_message(&err, Level::Error);
break 'out;
};
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 {
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);
}
break 'out;
}
if let EventV1::Logout = event {
info!("User {addr:?} received log out event!");
break 'out;
}
}
}
}
report_internal_error!(subscriber.quit().await).ok();
}
#[allow(clippy::too_many_arguments)]
async fn worker_with_kill_signal(
addr: SocketAddr,
subscribed: Arc<RwLock<HashSet<String>>>,
active_servers: Arc<Mutex<lru_time_cache::LruCache<String, ()>>>,
user_id: String,
config: &ProtocolConfiguration,
topic_signal_s: async_channel::Sender<()>,
kill_signal_r: async_channel::Receiver<()>,
read: WsReader,
write: &Mutex<WsWriter>,
kill_signal_s: async_channel::Sender<()>,
) {
worker(
addr,
subscribed,
active_servers,
user_id,
config,
topic_signal_s,
kill_signal_r,
read,
write,
)
.await;
kill_signal_s.send(()).await.ok();
}
#[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,
topic_signal_s: async_channel::Sender<()>,
kill_signal_r: async_channel::Receiver<()>,
mut read: WsReader,
write: &Mutex<WsWriter>,
) {
loop {
let t1 = read.try_next().fuse();
let t2 = kill_signal_r.recv().fuse();
pin_mut!(t1, t2);
select! {
_ = t2 => {
return;
},
result = t1 => {
let msg = match result {
Ok(Some(msg)) => msg,
Ok(None) => {
warn!("Received a None message!");
sentry::capture_message("Received a None message!", Level::Warning);
return;
}
Err(e) => {
use async_tungstenite::tungstenite::Error;
if !matches!(e, Error::AlreadyClosed | Error::ConnectionClosed) {
let err = format!("Error while reading an event from {addr:?}: {e:?}");
warn!("{}", err);
sentry::capture_message(&err, Level::Warning);
}
return;
}
};
let Ok(payload) = config.decode(&msg) else {
continue;
};
match payload {
ClientMessage::BeginTyping { channel } => {
if !subscribed.read().await.contains(&channel) {
continue;
}
EventV1::ChannelStartTyping {
id: channel.clone(),
user: user_id.clone(),
}
.p(channel.clone())
.await;
}
ClientMessage::EndTyping { channel } => {
if !subscribed.read().await.contains(&channel) {
continue;
}
EventV1::ChannelStopTyping {
id: channel.clone(),
user: user_id.clone(),
}
.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
topic_signal_s.send(()).await.ok();
}
}
ClientMessage::Ping { data, responded } => {
if responded.is_none() {
write
.lock()
.await
.send(config.encode(&EventV1::Pong { data }))
.await
.ok();
}
}
_ => {}
}
}
}
}
client_core(db, ws, config).await;
}

2
crates/core/database/Cargo.toml
View File

@@ -56,6 +56,7 @@ isahc = { optional = true, version = "1.7", features = ["json"] }
# Serialisation
serde_json = "1"
rmp-serde = "1.0.0"
revolt_optional_struct = "0.2.0"
serde = { version = "1", features = ["derive"] }
iso8601-timestamp = { version = "0.2.10", features = ["serde", "bson"] }
@@ -101,3 +102,4 @@ authifier = { version = "1.0.15", features = ["rocket_impl"] }
# RabbitMQ
amqprs = { version = "1.7.0" }
lapin = { version = "3.0.0" }

135
crates/core/database/src/events/client.rs
View File

@@ -1,4 +1,13 @@
use std::sync::Arc;
use async_std::sync::Mutex;
use authifier::AuthifierEvent;
use lapin::{
options::QueueDeclareOptions,
protocol::basic::AMQPProperties,
types::{AMQPValue, FieldTable},
};
use revolt_config::config;
use revolt_result::Error;
use serde::{Deserialize, Serialize};
@@ -250,17 +259,135 @@ pub enum EventV1 {
Auth(AuthifierEvent),
}
// * not sure where to put this yet...
pub async fn create_client() -> lapin::Connection {
let config = config().await;
lapin::Connection::connect(
&format!(
"amqp://{}:{}@{}:{}/%2f",
config.rabbit.username, config.rabbit.password, config.rabbit.host, config.rabbit.port
),
Default::default(),
)
.await
.unwrap()
}
pub async fn create_event_stream_channel(conn: lapin::Connection) -> lapin::Channel {
static STREAM_NAME: &str = "revolt.events";
let channel = conn.create_channel().await.unwrap();
let mut args: FieldTable = Default::default();
args.insert(
// set queue type to stream
"x-queue-type".into(),
AMQPValue::LongString("stream".into()),
);
args.insert(
// max. size of the stream
"x-max-length-bytes".into(),
AMQPValue::LongLongInt(5_000_000_000), // 5 GB
);
args.insert(
// size of the Bloom filter
"x-stream-filter-size-bytes".into(),
AMQPValue::LongLongInt(26), // 26 B
);
channel
.queue_declare(
STREAM_NAME,
QueueDeclareOptions {
durable: true,
..Default::default()
},
args,
)
.await
.unwrap();
channel.basic_qos(100, Default::default()).await.unwrap();
channel
}
/// simple connection pooling algorithm
pub async fn get_event_stream_channel() -> Arc<lapin::Channel> {
static CLIENTS_PER_CONN: usize = 128;
static CONNECTIONS: Mutex<Vec<Arc<lapin::Channel>>> = Mutex::new(Vec::new());
let mut connections = CONNECTIONS.lock().await;
connections.retain(|item| {
if item.status().connected() {
true
} else {
info!(
"Dropping connection with status {:?}",
item.status().state()
);
false
}
});
info!(
"Connections: {}, Clients: {:?}",
connections.len(),
connections
.iter()
.map(Arc::strong_count)
.collect::<Vec<usize>>()
);
for channel in connections.iter() {
if Arc::strong_count(channel) < CLIENTS_PER_CONN {
return channel.clone();
}
}
let conn = create_client().await;
let channel = Arc::new(create_event_stream_channel(conn).await);
connections.push(channel.clone());
channel
}
// * ---
impl EventV1 {
/// Publish helper wrapper
pub async fn p(self, channel: String) {
#[cfg(not(debug_assertions))]
redis_kiss::p(channel, self).await;
#[cfg(debug_assertions)]
info!("Publishing event to {channel}: {self:?}");
static QUEUE_NAME: &str = "revolt.events";
let mut headers: FieldTable = Default::default();
headers.insert(
"x-stream-filter-value".into(),
AMQPValue::LongString(channel.into()),
);
let properties: AMQPProperties = Default::default();
let result = get_event_stream_channel()
.await
.basic_publish(
"",
QUEUE_NAME,
Default::default(),
&rmp_serde::to_vec_named(&self).unwrap(),
properties.with_headers(headers),
)
.await;
#[cfg(not(debug_assertions))]
result.ok();
#[cfg(debug_assertions)]
redis_kiss::publish(channel, self).await.unwrap();
result.unwrap();
}
/// Publish user event

2
crates/delta/Cargo.toml
View File

@@ -32,6 +32,7 @@ ulid = "0.4.1"
nanoid = "0.4.0"
# serde
rmp-serde = "1.0.0"
serde_json = "1.0.57"
serde = { version = "1.0.115", features = ["derive"] }
validator = { version = "0.16", features = ["derive"] }
@@ -65,6 +66,7 @@ revolt_rocket_okapi = { version = "0.10.0", features = ["swagger"] }
# rabbit
amqprs = { version = "1.7.0" }
lapin = { version = "3.0.0" }
# core
authifier = "1.0.15"

54
crates/delta/src/util/test.rs
View File

@@ -1,12 +1,15 @@
use std::collections::HashMap;
use authifier::{
models::{Account, EmailVerification, Session},
Authifier,
};
use futures::StreamExt;
use lapin::types::AMQPValue;
use rand::Rng;
use redis_kiss::redis::aio::PubSub;
use revolt_database::{
events::client::EventV1, Channel, Database, Member, Message, Server, User, AMQP,
events::client::{get_event_stream_channel, EventV1},
Channel, Database, Member, Message, Server, User, AMQP,
};
use revolt_database::{util::idempotency::IdempotencyKey, Role};
use revolt_models::v0;
@@ -19,7 +22,7 @@ pub struct TestHarness {
authifier: Authifier,
pub db: Database,
pub amqp: AMQP,
sub: PubSub,
consumer: lapin::Consumer,
event_buffer: Vec<(String, EventV1)>,
}
@@ -31,11 +34,17 @@ impl TestHarness {
.await
.expect("valid rocket instance");
let mut sub = redis_kiss::open_pubsub_connection()
.await
.expect("`PubSub`");
let tag: String = rand::thread_rng()
.sample_iter::<char, _>(&rand::distributions::Standard)
.take(32)
.collect();
sub.psubscribe("*").await.unwrap();
static QUEUE_NAME: &str = "revolt.events";
let consumer = get_event_stream_channel()
.await
.basic_consume(QUEUE_NAME, &tag, Default::default(), Default::default())
.await
.unwrap();
let db = client
.rocket()
@@ -68,7 +77,7 @@ impl TestHarness {
authifier,
db,
amqp,
sub,
consumer,
event_buffer: vec![],
}
}
@@ -227,16 +236,33 @@ impl TestHarness {
}
}
let mut stream = self.sub.on_message();
while let Some(item) = stream.next().await {
let msg_topic = item.get_channel_name();
let payload: EventV1 = redis_kiss::decode_payload(&item).unwrap();
while let Some(Ok(delivery)) = self.consumer.next().await {
let headers: HashMap<String, AMQPValue> = delivery
.properties
.headers()
.as_ref()
.map(|table| {
table
.into_iter()
.map(|(k, v)| (k.to_string(), v.clone()))
.collect()
})
.unwrap_or_default();
if topic == msg_topic && predicate(&payload) {
let filter_value = headers
.get("x-stream-filter-value")
.expect("`x-stream-filter-value` not present in message!")
.as_long_string()
.unwrap()
.to_string();
let payload: EventV1 = rmp_serde::from_slice(&delivery.data).unwrap();
if topic == filter_value && predicate(&payload) {
return payload;
}
self.event_buffer.push((msg_topic.to_string(), payload));
self.event_buffer.push((filter_value, payload));
}
// WARNING: if predicate is never satisfied, this will never return