diff --git a/Cargo.lock b/Cargo.lock index 7b175658..50160f66 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6573,6 +6573,7 @@ dependencies = [ "querystring", "rand 0.6.5", "redis-kiss", + "revolt-broker", "revolt-config", "revolt-database", "revolt-models", @@ -6586,6 +6587,19 @@ dependencies = [ "ulid 0.5.0", ] +[[package]] +name = "revolt-broker" +version = "0.8.8" +dependencies = [ + "async-std", + "lapin", + "log", + "rand 0.9.1", + "revolt-config", + "rmp-serde", + "serde", +] + [[package]] name = "revolt-config" version = "0.8.8" @@ -6645,6 +6659,7 @@ dependencies = [ "rand 0.8.5", "redis-kiss", "regex", + "revolt-broker", "revolt-config", "revolt-models", "revolt-parser", @@ -6694,6 +6709,7 @@ dependencies = [ "redis-kiss", "regex", "reqwest 0.11.27", + "revolt-broker", "revolt-config", "revolt-database", "revolt-models", diff --git a/crates/bonfire/Cargo.toml b/crates/bonfire/Cargo.toml index 02e1934f..ce73ae64 100644 --- a/crates/bonfire/Cargo.toml +++ b/crates/bonfire/Cargo.toml @@ -39,6 +39,7 @@ async-std = { version = "1.8.0", features = [ # core authifier = { version = "1.0.15" } revolt-result = { path = "../core/result" } +revolt-broker = { path = "../core/broker" } revolt-models = { path = "../core/models" } revolt-config = { path = "../core/config" } revolt-database = { path = "../core/database" } diff --git a/crates/bonfire/src/client/subscriber.rs b/crates/bonfire/src/client/subscriber.rs index a1c295d4..c0f13e3e 100644 --- a/crates/bonfire/src/client/subscriber.rs +++ b/crates/bonfire/src/client/subscriber.rs @@ -1,20 +1,10 @@ -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 futures::{pin_mut, select, stream::SplitSink, FutureExt, SinkExt}; +use revolt_broker::event_stream; +use revolt_database::{events::client::EventV1, Database}; use sentry::Level; use crate::{ @@ -22,8 +12,6 @@ use crate::{ events::state::{State, SubscriptionStateChange}, }; -static QUEUE_NAME: &str = "revolt.events"; - /// Event subscriber loop pub async fn client_subscriber( write: &Mutex, async_tungstenite::tungstenite::Message>>, @@ -33,43 +21,17 @@ pub async fn client_subscriber( db: &'static Database, state: &mut State, ) { - // * --- CHANNEL CREATE --- - let channel = get_event_stream_channel().await; - // * --- CHANNEL END --- + let mut consumer = event_stream::Consumer::new().await; + consumer.set_topics(state.subscribed.read().await.clone()); - let mut offset: Option = 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)); + // Reload consumer if subscriptions change + if !matches!(state.apply_state().await, SubscriptionStateChange::None) { + consumer.set_topics(state.subscribed.read().await.clone()); } - // 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::(&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(); @@ -85,43 +47,8 @@ pub async fn client_subscriber( 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 = 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`"); - + event = delivery => { + if let Some(mut event) = event { // Handle the event if let EventV1::Auth(auth) = &event { if let AuthifierEvent::DeleteSession { session_id, .. } = auth { @@ -165,39 +92,21 @@ pub async fn client_subscriber( 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; + } else { + cancel = true; 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(); + consumer.dispose().await; } diff --git a/crates/core/broker/Cargo.toml b/crates/core/broker/Cargo.toml new file mode 100644 index 00000000..57fc8c01 --- /dev/null +++ b/crates/core/broker/Cargo.toml @@ -0,0 +1,25 @@ +[package] +name = "revolt-broker" +version = "0.8.8" +edition = "2024" +license = "AGPL-3.0-or-later" +authors = ["Paul Makles "] +description = "Revolt Backend: Event Broker" + +[dependencies] +# Utility +log = "0.4" +rand = "0.9.1" + +# RabbitMQ/AMQP client +lapin = "3.0.0" + +# Async runtime +async-std = { version = "1.8.0" } + +# Serialisation +serde = "1" +rmp-serde = "1.3.0" + +# Core +revolt-config = { version = "0.8.8", path = "../config" } diff --git a/crates/core/broker/src/event_stream/consumer.rs b/crates/core/broker/src/event_stream/consumer.rs new file mode 100644 index 00000000..50126658 --- /dev/null +++ b/crates/core/broker/src/event_stream/consumer.rs @@ -0,0 +1,160 @@ +use std::{ + collections::{HashMap, HashSet}, + sync::Arc, +}; + +use async_std::stream::StreamExt; +use lapin::{ + Channel, + options::BasicAckOptions, + types::{AMQPValue, FieldArray, FieldTable, LongLongInt}, +}; +use rand::Rng; +use revolt_config::config; +use serde::de::DeserializeOwned; + +use crate::event_stream::get_channel; + +pub struct Consumer { + channel: Arc, + tag: String, + topics: HashSet, + topics_changed: bool, + consumer: Option, + offset: Option, +} + +impl Consumer { + /// Create a new event stream consumer + pub async fn new() -> Consumer { + Consumer { + channel: get_channel().await, + tag: rand::rng() + .sample_iter::(&rand::distr::StandardUniform) + .take(32) + .collect(), + topics: HashSet::new(), + topics_changed: false, + consumer: None, + offset: None, + } + } + + /// Update the set of topics + pub fn set_topics(&mut self, topics: HashSet) { + self.topics = topics; + self.topics_changed = true; + } + + /// Get the current consumer + pub async fn ensure_consumer(&mut self) { + if self.topics_changed { + self.consumer = None; + } + + if self.consumer.is_none() { + let config = config().await; + + // 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 &self.topics { + 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) = self.offset { + args.insert("x-stream-offset".into(), AMQPValue::LongLongInt(offset)); + } + + // Create the consumer + self.consumer = Some( + self.channel + .basic_consume( + &config.rabbit.event_stream.queue, + &self.tag, + Default::default(), + args, + ) + .await + .unwrap(), + ); + } + } + + /// Close the active consumer + pub async fn dispose(&mut self) { + // Close the channel -- don't do this actually + // capture_internal_error!(self.channel.close(0, "closing channel").await); + + // Close the consumer + if let Err(err) = self + .channel + .basic_cancel(&self.tag, Default::default()) + .await + { + eprintln!("Failed to close consumer! {:?}", err); + } + // is this necessary? + // 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"); + // } + // } + } + + /// Get the next item + pub async fn next(&mut self) -> Option { + self.ensure_consumer().await; + + let consumer = self.consumer.as_mut().unwrap(); + + while let Some(Ok(delivery)) = consumer.next().await { + // Acknowledgement is required + delivery.ack(BasicAckOptions::default()).await.expect("ack"); + + // Parse the delivery headers + let headers: HashMap = delivery + .properties + .headers() + .as_ref() + .map(|table| { + table + .into_iter() + .map(|(k, v)| (k.to_string(), v.clone())) + .collect() + }) + .unwrap_or_default(); + + // Keep track of the current offset + let stream_offset = headers + .get("x-stream-offset") + .expect("`x-stream-offset` not present in message!"); + + self.offset = Some(stream_offset.as_long_long_int().unwrap() + 1); + + // 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!") + .as_long_string() + .expect("`string`") + .to_string(); + + if self.topics.contains(&filter_value) { + // Deserialise the data + return Some(rmp_serde::from_slice(&delivery.data).expect("`data`")); + } + } + + None + } +} diff --git a/crates/core/broker/src/event_stream/mod.rs b/crates/core/broker/src/event_stream/mod.rs new file mode 100644 index 00000000..7b9d4b08 --- /dev/null +++ b/crates/core/broker/src/event_stream/mod.rs @@ -0,0 +1,7 @@ +mod consumer; +mod pool; +mod publish; + +pub use consumer::Consumer; +pub use pool::get_channel; +pub use publish::publish_event; diff --git a/crates/core/broker/src/event_stream/pool.rs b/crates/core/broker/src/event_stream/pool.rs new file mode 100644 index 00000000..ff706c19 --- /dev/null +++ b/crates/core/broker/src/event_stream/pool.rs @@ -0,0 +1,100 @@ +use async_std::sync::Mutex; +use lapin::{ + Channel, + options::QueueDeclareOptions, + types::{AMQPValue, FieldTable}, +}; +use log::{debug, warn}; +use revolt_config::{RabbitEventStream, config}; +use std::sync::Arc; + +use crate::create_client; + +/// Get a handle to the event stream +pub async fn get_channel() -> Arc { + let config = config().await; + + static CONNECTIONS: Mutex>> = Mutex::new(Vec::new()); + + let mut connections = CONNECTIONS.lock().await; + connections.retain(|item| { + if item.status().connected() { + true + } else { + warn!( + "Dropping connection with status {:?}", + item.status().state() + ); + + false + } + }); + + debug!( + "Connections: {}, Clients: {:?}", + connections.len(), + connections + .iter() + .map(Arc::strong_count) + .collect::>() + ); + + for channel in connections.iter() { + if Arc::strong_count(channel) < config.rabbit.event_stream.channels_per_conn { + return channel.clone(); + } + } + + let conn = create_client().await; + let channel = Arc::new(create_channel(conn, config.rabbit.event_stream).await); + + connections.push(channel.clone()); + channel +} + +/// Create a channel +async fn create_channel( + conn: lapin::Connection, + event_stream: RabbitEventStream, +) -> lapin::Channel { + 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(event_stream.stream_max_length_bytes), + ); + + args.insert( + // size of the Bloom filter + "x-stream-filter-size-bytes".into(), + AMQPValue::LongLongInt(event_stream.filter_size_bytes), + ); + + channel + .queue_declare( + &event_stream.queue, + QueueDeclareOptions { + durable: true, + ..Default::default() + }, + args, + ) + .await + .unwrap(); + + channel + .basic_qos(event_stream.qos_prefetch, Default::default()) + .await + .unwrap(); + + channel +} diff --git a/crates/core/broker/src/event_stream/publish.rs b/crates/core/broker/src/event_stream/publish.rs new file mode 100644 index 00000000..4be20efa --- /dev/null +++ b/crates/core/broker/src/event_stream/publish.rs @@ -0,0 +1,35 @@ +use lapin::{ + Error, + protocol::basic::AMQPProperties, + publisher_confirm::PublisherConfirm, + types::{AMQPValue, FieldTable}, +}; +use revolt_config::config; +use serde::Serialize; + +use crate::event_stream::get_channel; + +/// Publish an event to the message broker +pub async fn publish_event( + channel: &str, + data: &T, +) -> Result { + let config = config().await; + + let mut headers: FieldTable = Default::default(); + headers.insert( + "x-stream-filter-value".into(), + AMQPValue::LongString(channel.into()), + ); + + get_channel() + .await + .basic_publish( + &config.rabbit.event_stream.exchange, + &config.rabbit.event_stream.queue, + Default::default(), + &rmp_serde::to_vec_named(data).unwrap(), + AMQPProperties::default().with_headers(headers), + ) + .await +} diff --git a/crates/core/broker/src/lib.rs b/crates/core/broker/src/lib.rs new file mode 100644 index 00000000..fb2b5c4a --- /dev/null +++ b/crates/core/broker/src/lib.rs @@ -0,0 +1,18 @@ +use revolt_config::config; + +pub mod event_stream; + +/// Create a lapin client +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() +} diff --git a/crates/core/config/Revolt.toml b/crates/core/config/Revolt.toml index e3a04f2a..c1e868bc 100644 --- a/crates/core/config/Revolt.toml +++ b/crates/core/config/Revolt.toml @@ -28,6 +28,20 @@ port = 5672 username = "rabbituser" password = "rabbitpass" +[rabbit.event_stream] +# Configuration for event brokerage +# Using default/direct exchange +exchange = "" +queue = "revolt.events" +# Number of channels that can be opened per single TCP connection +channels_per_conn = 128 +# Maximum size of the stream +stream_max_length_bytes = 5_000_000_000 +# Size of the Bloom filter +filter_size_bytes = 26 +# Number of messages to prefetch +qos_prefetch = 100 + [api] [api.registration] diff --git a/crates/core/config/src/lib.rs b/crates/core/config/src/lib.rs index 7c050499..f315afbc 100644 --- a/crates/core/config/src/lib.rs +++ b/crates/core/config/src/lib.rs @@ -108,12 +108,24 @@ pub struct Database { pub redis: String, } +#[derive(Deserialize, Debug, Clone)] +pub struct RabbitEventStream { + pub exchange: String, + pub queue: String, + pub channels_per_conn: usize, + pub stream_max_length_bytes: i64, + pub filter_size_bytes: i64, + pub qos_prefetch: u16, +} + #[derive(Deserialize, Debug, Clone)] pub struct Rabbit { pub host: String, pub port: u16, pub username: String, pub password: String, + + pub event_stream: RabbitEventStream, } #[derive(Deserialize, Debug, Clone)] diff --git a/crates/core/database/Cargo.toml b/crates/core/database/Cargo.toml index e2ea93e5..854ce60e 100644 --- a/crates/core/database/Cargo.toml +++ b/crates/core/database/Cargo.toml @@ -37,6 +37,7 @@ revolt-permissions = { version = "0.8.8", path = "../permissions", features = [ "bson", ] } revolt-parser = { version = "0.8.8", path = "../parser" } +revolt-broker = { version = "0.8.8", path = "../broker" } # Utility log = "0.4" diff --git a/crates/core/database/src/events/client.rs b/crates/core/database/src/events/client.rs index a7415025..b3bfb4e5 100644 --- a/crates/core/database/src/events/client.rs +++ b/crates/core/database/src/events/client.rs @@ -1,13 +1,5 @@ -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_broker::event_stream; use revolt_result::Error; use serde::{Deserialize, Serialize}; @@ -259,129 +251,13 @@ 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 { - static CLIENTS_PER_CONN: usize = 128; - static CONNECTIONS: Mutex>> = 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::>() - ); - - 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(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; + let result = event_stream::publish_event(&channel, &self).await; #[cfg(not(debug_assertions))] result.ok(); diff --git a/crates/delta/Cargo.toml b/crates/delta/Cargo.toml index 192f9f77..2bf59e47 100644 --- a/crates/delta/Cargo.toml +++ b/crates/delta/Cargo.toml @@ -71,6 +71,7 @@ lapin = { version = "3.0.0" } # core authifier = "1.0.15" revolt-config = { path = "../core/config" } +revolt-broker = { path = "../core/broker" } revolt-database = { path = "../core/database", features = [ "rocket-impl", "redis-is-patched", diff --git a/crates/delta/src/util/test.rs b/crates/delta/src/util/test.rs index fce54d71..de47ecb6 100644 --- a/crates/delta/src/util/test.rs +++ b/crates/delta/src/util/test.rs @@ -7,9 +7,9 @@ use authifier::{ use futures::StreamExt; use lapin::types::AMQPValue; use rand::Rng; +use revolt_broker::event_stream; use revolt_database::{ - events::client::{get_event_stream_channel, EventV1}, - Channel, Database, Member, Message, Server, User, AMQP, + events::client::EventV1, Channel, Database, Member, Message, Server, User, AMQP, }; use revolt_database::{util::idempotency::IdempotencyKey, Role}; use revolt_models::v0; @@ -40,7 +40,7 @@ impl TestHarness { .collect(); static QUEUE_NAME: &str = "revolt.events"; - let consumer = get_event_stream_channel() + let consumer = event_stream::get_channel() .await .basic_consume(QUEUE_NAME, &tag, Default::default(), Default::default()) .await