feat: env var to control rabbitmq event ingress

Signed-off-by: Zomatree <me@zomatree.live>

Cargo.lock

@@ -6838,8 +6838,10 @@ dependencies = [
 name = "revolt-bonfire"
 version = "0.11.5"
 dependencies = [
+ "amqprs",
  "async-channel 2.5.0",
  "async-std",
+ "async-trait",
  "async-tungstenite",
  "authifier",
  "bincode",

crates/bonfire/Cargo.toml

@@ -36,6 +36,7 @@ async-std = { version = "1.8.0", features = [
     "tokio02",
     "attributes",
 ] }
+async-trait = "0.1.89"
 
 # core
 authifier = { version = "1.0.16" }
@@ -48,3 +49,6 @@ revolt-presence = { path = "../core/presence", features = ["redis-is-patched"] }
 
 # redis
 fred = { version = "8.0.1", features = ["subscriber-client"] }
+
+# Redis
+amqprs = { version = "1.7.0" }

crates/bonfire/src/main.rs

@@ -1,6 +1,18 @@
 use std::env;
 
+use amqprs::{
+    channel::{
+        BasicConsumeArguments, Channel, ExchangeDeclareArguments, QueueBindArguments,
+        QueueDeclareArguments,
+    },
+    connection::{Connection, OpenConnectionArguments},
+    consumer::AsyncConsumer,
+    BasicProperties, Deliver,
+};
 use async_std::net::TcpListener;
+use async_trait::async_trait;
+use redis_kiss::AsyncCommands;
+use revolt_database::util::rabbit::set_rabbitmq_connection;
 use revolt_presence::clear_region;
 
 #[macro_use]
@@ -31,6 +43,55 @@ async fn main() {
     let try_socket = TcpListener::bind(bind).await;
     let listener = try_socket.expect("Failed to bind");
 
+    let config = revolt_config::config().await;
+
+    let rmq_conn = Connection::open(&OpenConnectionArguments::new(
+        &config.rabbit.host,
+        config.rabbit.port,
+        &config.rabbit.username,
+        &config.rabbit.password,
+    ))
+    .await
+    .expect("Failed to connect to RabbitMQ");
+
+    set_rabbitmq_connection(rmq_conn.clone());
+
+    if std::env::var("ENABLE_RABBITMQ_INGRESS").as_deref().is_ok_and(|v| v == "1") {
+        let channel = rmq_conn
+            .open_channel(None)
+            .await
+            .expect("Failed to open RabbitMQ channel.");
+
+        channel
+            .exchange_declare(
+                ExchangeDeclareArguments::new("events", "fanout")
+                    .durable(true)
+                    .finish(),
+            )
+            .await
+            .expect("Failed to declare exchange");
+
+        channel
+            .queue_declare(QueueDeclareArguments::new("events").durable(true).finish())
+            .await
+            .expect("Failed to declare queue");
+
+        channel
+            .queue_bind(QueueBindArguments::new("events", "events", "events"))
+            .await
+            .expect("Failed to bind queue");
+
+        channel
+            .basic_consume(
+                RabbitToRedisConsumer,
+                BasicConsumeArguments::new("events", "")
+                    .manual_ack(false)
+                    .finish(),
+            )
+            .await
+            .expect("Failed to consume channel");
+    }
+
     // Start accepting new connections and spawn a client for each connection.
     while let Ok((stream, addr)) = listener.accept().await {
         async_std::task::spawn(async move {
@@ -40,3 +101,32 @@ async fn main() {
         });
     }
 }
+
+struct RabbitToRedisConsumer;
+
+#[async_trait]
+impl AsyncConsumer for RabbitToRedisConsumer {
+    async fn consume(
+        &mut self,
+        _channel: &Channel,
+        _deliver: Deliver,
+        basic_properties: BasicProperties,
+        content: Vec<u8>,
+    ) {
+        let mut redis_conn = redis_kiss::get_connection()
+            .await
+            .expect("Failed to connect to Redis.");
+
+        let pubsub_channel = basic_properties
+            .headers()
+            .expect("No headers")
+            .get(&"c".try_into().unwrap())
+            .expect("No channel header")
+            .to_string();
+
+        redis_conn
+            .publish(pubsub_channel, content)
+            .await
+            .expect("failed to publish")
+    }
+}

crates/bonfire/src/websocket.rs

@@ -464,30 +464,30 @@ async fn worker(
 
                 match payload {
                     // disabled events
-                    // ClientMessage::BeginTyping { channel } => {
-                    //     if !subscribed.read().await.contains(&channel) {
-                    //         continue;
-                    //     }
+                    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::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;
-                    // }
+                        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);

crates/core/config/src/lib.rs

@@ -449,7 +449,7 @@ pub async fn config() -> Settings {
     let mut config = read().await.try_deserialize::<Settings>().unwrap();
 
     // inject REDIS_URI for redis-kiss library
-    if std::env::var("REDIS_URL").is_err() {
+    if std::env::var("REDIS_URI").is_err() {
         std::env::set_var("REDIS_URI", config.database.redis.clone());
     }
 

crates/core/database/src/amqp/amqp.rs

@@ -11,6 +11,8 @@ use revolt_presence::filter_online;
 
 use serde_json::to_string;
 
+// TODO: move away from storing Connection and Channel to using thread local singletons - #659
+
 #[derive(Clone)]
 pub struct AMQP {
     #[allow(unused)]

crates/core/database/src/events/client.rs

@@ -1,12 +1,20 @@
+use amqprs::{
+    channel::{BasicPublishArguments},
+    BasicProperties, FieldTable,
+};
 use authifier::AuthifierEvent;
 use revolt_result::Error;
 use serde::{Deserialize, Serialize};
 
 use revolt_models::v0::{
-    AppendMessage, Channel, ChannelUnread, ChannelVoiceState, Emoji, FieldsChannel, FieldsMember, FieldsMessage, FieldsRole, FieldsServer, FieldsUser, FieldsWebhook, Member, MemberCompositeKey, Message, PartialChannel, PartialMember, PartialMessage, PartialRole, PartialServer, PartialUser, PartialUserVoiceState, PartialWebhook, PolicyChange, RemovalIntention, Report, Server, User, UserSettings, UserVoiceState, Webhook
+    AppendMessage, Channel, ChannelUnread, ChannelVoiceState, Emoji, FieldsChannel, FieldsMember,
+    FieldsMessage, FieldsRole, FieldsServer, FieldsUser, FieldsWebhook, Member, MemberCompositeKey,
+    Message, PartialChannel, PartialMember, PartialMessage, PartialRole, PartialServer,
+    PartialUser, PartialUserVoiceState, PartialWebhook, PolicyChange, RemovalIntention, Report,
+    Server, User, UserSettings, UserVoiceState, Webhook,
 };
 
-use crate::Database;
+use crate::{util::rabbit::get_channel, Database};
 
 /// Ping Packet
 #[derive(Serialize, Deserialize, Debug, Clone)]
@@ -304,14 +312,23 @@ pub enum EventV1 {
 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:?}");
 
-        #[cfg(debug_assertions)]
-        redis_kiss::publish(channel, self).await.unwrap();
+        let rmq = get_channel().await;
+
+        let mut headers = FieldTable::new();
+        headers.insert("c".try_into().unwrap(), channel.clone().into());
+        let mut properties = BasicProperties::default();
+        properties.with_headers(headers);
+
+        rmq.basic_publish(
+            properties,
+            serde_json::to_string(&self).unwrap().into_bytes(),
+            BasicPublishArguments::new("events", "events"),
+        )
+        .await
+        .unwrap();
     }
 
     /// Publish user event

crates/core/database/src/util/mod.rs

@@ -5,5 +5,6 @@ pub mod idempotency;
 pub mod permissions;
 pub mod reference;
 pub mod test_fixtures;
+pub mod rabbit;
 
 pub use funcs::*;

crates/core/database/src/util/rabbit.rs

@@ -0,0 +1,48 @@
+use amqprs::{channel::Channel, connection::Connection};
+use once_cell::sync::OnceCell;
+use std::{
+    collections::HashMap, future::ready, sync::{LazyLock, RwLock}, thread::{ThreadId, current}
+};
+
+static RABBIT_CONNECTION: OnceCell<Connection> = OnceCell::new();
+static RABBIT_CHANNELS: LazyLock<RwLock<HashMap<ThreadId, Channel>>> =
+    LazyLock::new(|| RwLock::new(HashMap::new()));
+
+pub async fn get_channel_with_init<F: AsyncFnOnce(Channel) -> Channel>(init: F) -> Channel {
+    let conn = RABBIT_CONNECTION
+        .get()
+        .expect("Rabbit connection is not initialised.");
+
+    let thread_id = current().id();
+
+    let channel = RABBIT_CHANNELS
+        .read()
+        .expect("Channels poisioned")
+        .get(&thread_id)
+        .cloned();
+
+    if let Some(channel) = channel {
+        channel
+    } else {
+        let mut channel =
+            conn.open_channel(None)
+                .await
+                .expect("Failed to open rabbitmq channel");
+
+        channel = init(channel).await;
+
+        RABBIT_CHANNELS
+            .write()
+            .expect("Channels poisioned")
+            .insert(thread_id, channel.clone());
+        channel
+    }
+}
+
+pub async fn get_channel() -> Channel {
+    get_channel_with_init(ready).await
+}
+
+pub fn set_rabbitmq_connection(connection: Connection) -> bool {
+    RABBIT_CONNECTION.set(connection).is_ok()
+}

crates/delta/src/main.rs

@@ -9,8 +9,7 @@ pub mod routes;
 pub mod util;
 
 use revolt_config::config;
-use revolt_database::events::client::EventV1;
-use revolt_database::AMQP;
+use revolt_database::{AMQP, util::rabbit::{get_channel, get_channel_with_init, set_rabbitmq_connection}};
 use revolt_ratelimits::rocket as ratelimiter;
 use rocket::{Build, Rocket};
 use rocket_cors::{AllowedOrigins, CorsOptions};
@@ -19,11 +18,9 @@ use std::net::Ipv4Addr;
 use std::str::FromStr;
 
 use amqprs::{
-    channel::ExchangeDeclareArguments,
+    channel::{Channel, ExchangeDeclareArguments},
     connection::{Connection, OpenConnectionArguments},
 };
-use async_std::channel::unbounded;
-use authifier::AuthifierEvent;
 use rocket::data::ToByteUnit;
 use revolt_database::voice::VoiceClient;
 
@@ -39,28 +36,9 @@ pub async fn web() -> Rocket<Build> {
     log::info!("database_here {db:?}");
     db.migrate_database().await.unwrap();
 
-    // Setup Authifier event channel
-    let (_, receiver) = unbounded();
-
     // Setup Authifier
     let authifier = db.clone().to_authifier().await;
 
-    // Launch a listener for Authifier events
-    async_std::task::spawn(async move {
-        while let Ok(event) = receiver.recv().await {
-            match &event {
-                AuthifierEvent::CreateSession { .. } | AuthifierEvent::CreateAccount { .. } => {
-                    EventV1::Auth(event).global().await
-                }
-                AuthifierEvent::DeleteSession { user_id, .. }
-                | AuthifierEvent::DeleteAllSessions { user_id, .. } => {
-                    let id = user_id.to_string();
-                    EventV1::Auth(event).private(id).await
-                }
-            }
-        }
-    });
-
     // Configure CORS
     let cors = CorsOptions {
         allowed_origins: AllowedOrigins::All,
@@ -121,19 +99,19 @@ pub async fn web() -> Rocket<Build> {
     .await
     .expect("Failed to connect to RabbitMQ");
 
-    let channel = connection
-        .open_channel(None)
-        .await
-        .expect("Failed to open RabbitMQ channel");
+    set_rabbitmq_connection(connection.clone());
+    let channel = get_channel_with_init(|channel: Channel| async {
+        channel
+            .exchange_declare(
+                ExchangeDeclareArguments::new(&config.pushd.exchange, "direct")
+                    .durable(true)
+                    .finish(),
+            )
+            .await
+            .expect("Failed to declare exchange");
 
-    channel
-        .exchange_declare(
-            ExchangeDeclareArguments::new(&config.pushd.exchange, "direct")
-                .durable(true)
-                .finish(),
-        )
-        .await
-        .expect("Failed to declare exchange");
+        channel
+    }).await;
 
     let amqp = AMQP::new(connection, channel);
 

crates/delta/src/util/test.rs

@@ -6,7 +6,9 @@ use futures::StreamExt;
 use rand::Rng;
 use redis_kiss::redis::aio::PubSub;
 use revolt_database::{
-    events::client::EventV1, Channel, Database, Member, Message, PartialRole, Server, User, AMQP,
+    events::client::EventV1,
+    util::rabbit::{get_channel, set_rabbitmq_connection},
+    Channel, Database, Member, Message, PartialRole, Server, User, AMQP,
 };
 use revolt_database::{util::idempotency::IdempotencyKey, Role};
 use revolt_models::v0;
@@ -59,10 +61,16 @@ impl TestHarness {
         )
         .await
         .unwrap();
-        let channel = connection.open_channel(None).await.unwrap();
+
+        set_rabbitmq_connection(connection.clone());
+        let channel = get_channel().await;
 
         let amqp = AMQP::new(connection, channel);
 
+        async_std::task::spawn(async {
+
+        });
+
         TestHarness {
             client,
             authifier,