chore: update es code to lapin

Signed-off-by: Zomatree <me@zomatree.live>
This commit is contained in:
Zomatree
2026-06-25 06:20:24 +01:00
parent 6a421dcd37
commit 553437e1df
30 changed files with 344 additions and 293 deletions

View File

@@ -31,7 +31,7 @@ default = ["mongodb", "tokio-runtime", "tasks"]
[dependencies]
# Core
revolt-config = { workspace = true, features = ["report-macros"] }
revolt-config = { workspace = true, features = ["report-macros", "anyhow"] }
revolt-result = { workspace = true }
revolt-models = { workspace = true, features = ["validator"] }
revolt-presence = { workspace = true }
@@ -56,6 +56,7 @@ validator = { workspace = true, features = ["derive"] }
isahc = { workspace = true, features = ["json"], optional = true }
base32 = { workspace = true }
sha1 = { workspace = true }
anyhow = { workspace = true }
# Serialisation
serde_json = { workspace = true }

View File

@@ -1,7 +1,5 @@
use std::{
future::{ready, Future},
pin::Pin,
sync::Arc,
future::{Future, ready}, marker::PhantomData, pin::Pin, sync::Arc
};
use anyhow::Result;
@@ -12,14 +10,15 @@ use lapin::{
BasicProperties, Channel, Connection, ConsumerDelegate, Error as AMQPError,
};
use log::debug;
use revolt_database::Database;
use crate::Database;
#[async_trait]
pub trait Consumer: Clone + Send + Sync + 'static {
pub trait Consumer<T: Clone = ()>: Clone + Send + Sync + 'static {
async fn create(
db: Database,
connection: Arc<Connection>,
channel: Arc<Channel>,
data: T,
) -> Self;
fn channel(&self) -> &Arc<Channel>;
async fn consume(&self, delivery: Delivery) -> Result<()>;
@@ -65,9 +64,15 @@ pub trait Consumer: Clone + Send + Sync + 'static {
}
}
pub struct Delegate<C: Consumer>(pub C);
pub struct Delegate<C: Consumer<D>, D: Clone>(C, PhantomData<D>);
impl<C: Consumer> ConsumerDelegate for Delegate<C> {
impl<C: Consumer<D>, D: Clone> Delegate<C, D> {
pub fn new(consumer: C) -> Self {
Self(consumer, PhantomData)
}
}
impl<C: Consumer<D>, D: Clone + Send + Sync> ConsumerDelegate for Delegate<C, D> {
fn on_new_delivery(
&self,
delivery: DeliveryResult,

View File

@@ -1,2 +1,3 @@
#[allow(clippy::module_inception)]
pub mod amqp;
pub mod consumer;

View File

@@ -109,7 +109,7 @@ pub mod events;
#[cfg(feature = "tasks")]
pub mod tasks;
mod amqp;
pub mod amqp;
pub use amqp::amqp::AMQP;
#[cfg(feature = "voice")]

View File

@@ -882,13 +882,13 @@ impl User {
/// - deletes owned bots, servers and messages
/// - removes user from all groups
/// - clears relationships
pub async fn delete(&mut self, db: &Database) -> Result<()> {
pub async fn delete(&mut self, db: &Database, amqp: Option<&AMQP>) -> Result<()> {
for bot in db.fetch_bots_by_user(&self.id).await? {
bot.delete(db).await?;
}
for server in db.fetch_owned_servers(&self.id).await? {
server.delete(db).await?;
server.delete(db, amqp).await?;
}
self.remove_from_all_groups(db).await?;

View File

@@ -12,7 +12,6 @@ pub mod permissions;
pub mod reference;
pub mod shield;
pub mod test_fixtures;
pub mod chunked;
pub use chunked::ChunkedDatabaseGenerator;
pub use funcs::*;
pub use chunked::ChunkedDatabaseGenerator;

View File

@@ -1,19 +1,19 @@
[package]
name = "revolt-search"
version = "0.12.0"
version = "0.13.7"
edition = "2024"
license = "MIT"
publish = false
[dependencies]
revolt-config = { version = "0.12.0", path = "../config" }
revolt-database = { version = "0.12.0", path = "../database" }
revolt-models = { version = "0.12.0", path = "../models" }
elasticsearch = "9.1.0-alpha.1"
serde_json = "1"
ulid = "1.2.1"
iso8601-timestamp = { version = "0.2.10", features = ["serde", "bson"] }
futures = "0.3.32"
elasticsearch-dsl = "0.4"
serde = "1"
linkify = "0.10.0"
revolt-config = { workspace = true }
revolt-database = { workspace = true }
revolt-models = { workspace = true }
elasticsearch = { workspace = true }
serde_json = { workspace = true }
ulid = { workspace = true }
iso8601-timestamp = { workspace = true, features = ["serde", "bson"] }
futures = { workspace = true }
elasticsearch-dsl = { workspace = true }
serde = { workspace = true }
linkify = { workspace = true }

View File

@@ -1,10 +1,10 @@
use std::time::Duration;
use revolt_database::Database;
use revolt_database::{Database, AMQP};
use revolt_result::Result;
use tokio::time::sleep;
pub async fn task(db: Database, _: revolt_database::AMQP) -> Result<()> {
pub async fn task(db: Database, amqp: AMQP) -> Result<()> {
loop {
let accounts = db.fetch_accounts_due_for_deletion().await?;
let count = accounts.len();
@@ -12,7 +12,7 @@ pub async fn task(db: Database, _: revolt_database::AMQP) -> Result<()> {
for mut account in accounts {
let mut user = db.fetch_user(&account.id).await?;
user.delete(&db).await?;
user.delete(&db, Some(&amqp)).await?;
account.mark_deleted(&db).await?;
}

View File

@@ -1,10 +1,9 @@
use std::sync::Arc;
use crate::utils::Consumer;
use anyhow::Result;
use async_trait::async_trait;
use lapin::{message::Delivery, Channel, Connection};
use revolt_database::{events::rabbit::*, Database};
use revolt_database::{amqp::consumer::Consumer, events::rabbit::*, Database};
#[derive(Clone)]
#[allow(unused)]
@@ -16,11 +15,7 @@ pub struct AckConsumer {
#[async_trait]
impl Consumer for AckConsumer {
async fn create(
db: Database,
connection: Arc<Connection>,
channel: Arc<Channel>,
) -> Self {
async fn create(db: Database, connection: Arc<Connection>, channel: Arc<Channel>, _: ()) -> Self {
Self {
db,
connection,

View File

@@ -1,11 +1,10 @@
use std::{collections::HashMap, sync::Arc};
use crate::utils::Consumer;
use anyhow::Result;
use async_trait::async_trait;
use lapin::{message::Delivery, Channel, Connection};
use log::debug;
use revolt_database::{events::rabbit::*, Database};
use revolt_database::{amqp::consumer::Consumer, events::rabbit::*, Database};
#[derive(Clone)]
#[allow(unused)]
@@ -17,11 +16,7 @@ pub struct DmCallConsumer {
#[async_trait]
impl Consumer for DmCallConsumer {
async fn create(
db: Database,
connection: Arc<Connection>,
channel: Arc<Channel>,
) -> Self {
async fn create(db: Database, connection: Arc<Connection>, channel: Arc<Channel>, _: ()) -> Self {
Self {
db,
connection,

View File

@@ -1,11 +1,10 @@
use std::{collections::HashMap, sync::Arc};
use crate::utils::Consumer;
use anyhow::Result;
use async_trait::async_trait;
use lapin::{message::Delivery, Channel, Connection};
use log::debug;
use revolt_database::{events::rabbit::*, Database};
use revolt_database::{amqp::consumer::Consumer, events::rabbit::*, Database};
#[derive(Clone)]
#[allow(unused)]
@@ -17,11 +16,7 @@ pub struct FRAcceptedConsumer {
#[async_trait]
impl Consumer for FRAcceptedConsumer {
async fn create(
db: Database,
connection: Arc<Connection>,
channel: Arc<Channel>,
) -> Self {
async fn create(db: Database, connection: Arc<Connection>, channel: Arc<Channel>, _: ()) -> Self {
Self {
db,
connection,

View File

@@ -1,11 +1,10 @@
use std::{collections::HashMap, sync::Arc};
use crate::utils::Consumer;
use anyhow::Result;
use async_trait::async_trait;
use lapin::{message::Delivery, Channel, Connection};
use log::debug;
use revolt_database::{events::rabbit::*, Database};
use revolt_database::{amqp::consumer::Consumer, events::rabbit::*, Database};
#[derive(Clone)]
#[allow(unused)]
@@ -17,11 +16,7 @@ pub struct FRReceivedConsumer {
#[async_trait]
impl Consumer for FRReceivedConsumer {
async fn create(
db: Database,
connection: Arc<Connection>,
channel: Arc<Channel>,
) -> Self {
async fn create(db: Database, connection: Arc<Connection>, channel: Arc<Channel>, _: ()) -> Self {
Self {
db,
connection,

View File

@@ -1,11 +1,10 @@
use std::{collections::HashMap, sync::Arc};
use crate::utils::Consumer;
use anyhow::Result;
use async_trait::async_trait;
use lapin::{message::Delivery, Channel, Connection};
use log::debug;
use revolt_database::{events::rabbit::*, Database};
use revolt_database::{amqp::consumer::Consumer, events::rabbit::*, Database};
#[derive(Clone)]
#[allow(unused)]
@@ -17,11 +16,7 @@ pub struct GenericConsumer {
#[async_trait]
impl Consumer for GenericConsumer {
async fn create(
db: Database,
connection: Arc<Connection>,
channel: Arc<Channel>,
) -> Self {
async fn create(db: Database, connection: Arc<Connection>, channel: Arc<Channel>, _: ()) -> Self {
Self {
db,
connection,

View File

@@ -4,13 +4,13 @@ use std::{
sync::Arc,
};
use crate::utils::{render_notification_content, Consumer};
use crate::utils::render_notification_content;
use anyhow::Result;
use async_trait::async_trait;
use lapin::{message::Delivery, Channel, Connection};
use revolt_database::{
events::rabbit::*, util::bulk_permissions::BulkDatabasePermissionQuery, Database, Member,
MessageFlagsValue,
amqp::consumer::Consumer, events::rabbit::*,
util::bulk_permissions::BulkDatabasePermissionQuery, Database, Member, MessageFlagsValue,
};
use revolt_models::v0::{MessageFlags, PushNotification};
use revolt_result::ToRevoltError;
@@ -29,11 +29,7 @@ impl MassMessageConsumer {
push: &PushNotification,
users: &[String],
) -> Result<()> {
if let Ok(sessions) = self
.db
.fetch_sessions_with_subscription(users)
.await
{
if let Ok(sessions) = self.db.fetch_sessions_with_subscription(users).await {
let config = revolt_config::config().await;
for session in sessions {
if let Some(sub) = session.subscription {
@@ -72,11 +68,7 @@ impl MassMessageConsumer {
#[async_trait]
impl Consumer for MassMessageConsumer {
async fn create(
db: Database,
connection: Arc<Connection>,
channel: Arc<Channel>,
) -> Self {
async fn create(db: Database, connection: Arc<Connection>, channel: Arc<Channel>, _: ()) -> Self {
Self {
db,
connection,

View File

@@ -1,11 +1,11 @@
use std::{collections::HashMap, sync::Arc};
use crate::utils::{render_notification_content, Consumer};
use crate::utils::render_notification_content;
use anyhow::Result;
use async_trait::async_trait;
use lapin::{message::Delivery, Channel, Connection};
use log::debug;
use revolt_database::{events::rabbit::*, Database};
use revolt_database::{amqp::consumer::Consumer, events::rabbit::*, Database};
#[derive(Clone)]
#[allow(unused)]
@@ -17,11 +17,7 @@ pub struct MessageConsumer {
#[async_trait]
impl Consumer for MessageConsumer {
async fn create(
db: Database,
connection: Arc<Connection>,
channel: Arc<Channel>,
) -> Self {
async fn create(db: Database, connection: Arc<Connection>, channel: Arc<Channel>, _: ()) -> Self {
Self {
db,
connection,

View File

@@ -1,6 +1,5 @@
use std::{borrow::Cow, collections::BTreeMap, io::Cursor, sync::Arc};
use crate::utils::Consumer;
use anyhow::Result;
use async_trait::async_trait;
use base64::{
@@ -15,7 +14,7 @@ use revolt_a2::{
},
Client, ClientConfig, Endpoint, Error, ErrorBody, ErrorReason, Priority, PushType, Response,
};
use revolt_database::{events::rabbit::*, Database};
use revolt_database::{amqp::consumer::Consumer, events::rabbit::*, Database};
use revolt_models::v0::{Channel, Message, PushNotification};
use serde::Serialize;
@@ -123,11 +122,7 @@ impl ApnsOutboundConsumer {
#[async_trait]
impl Consumer for ApnsOutboundConsumer {
async fn create(
db: Database,
connection: Arc<Connection>,
channel: Arc<AMQPChannel>,
) -> Self {
async fn create(db: Database, connection: Arc<Connection>, channel: Arc<AMQPChannel>, _: ()) -> Self {
let config = revolt_config::config().await;
if config.pushd.apn.pkcs8.is_empty()

View File

@@ -1,6 +1,5 @@
use std::{collections::HashMap, sync::Arc, time::Duration};
use crate::utils::Consumer;
use anyhow::{bail, Result};
use async_trait::async_trait;
use fcm_v1::{
@@ -10,7 +9,7 @@ use fcm_v1::{
};
use lapin::{message::Delivery, Channel as AMQPChannel, Connection};
use revolt_config::config;
use revolt_database::{events::rabbit::*, Database};
use revolt_database::{amqp::consumer::Consumer, events::rabbit::*, Database};
use serde_json::Value;
/// Custom notification data
@@ -126,11 +125,7 @@ pub struct FcmOutboundConsumer {
#[async_trait]
impl Consumer for FcmOutboundConsumer {
async fn create(
db: Database,
connection: Arc<Connection>,
channel: Arc<AMQPChannel>,
) -> Self {
async fn create(db: Database, connection: Arc<Connection>, channel: Arc<AMQPChannel>, _: ()) -> Self {
let config = revolt_config::config().await;
Self {

View File

@@ -1,7 +1,5 @@
use std::{collections::HashMap, sync::Arc};
use crate::utils::Consumer;
use anyhow::{anyhow, bail, Result};
use async_trait::async_trait;
use base64::{
@@ -9,7 +7,9 @@ use base64::{
Engine as _,
};
use lapin::{message::Delivery, Channel as AMQPChannel, Connection};
use revolt_database::{events::rabbit::*, util::format_display_name, Database};
use revolt_database::{
amqp::consumer::Consumer, events::rabbit::*, util::format_display_name, Database,
};
use web_push::{
ContentEncoding, IsahcWebPushClient, SubscriptionInfo, SubscriptionKeys, VapidSignatureBuilder,
WebPushClient, WebPushError, WebPushMessageBuilder,
@@ -27,11 +27,7 @@ pub struct VapidOutboundConsumer {
#[async_trait]
impl Consumer for VapidOutboundConsumer {
async fn create(
db: Database,
connection: Arc<Connection>,
channel: Arc<AMQPChannel>,
) -> Self {
async fn create(db: Database, connection: Arc<Connection>, channel: Arc<AMQPChannel>, _: ()) -> Self {
let config = revolt_config::config().await;
if config.pushd.vapid.private_key.is_empty() || config.pushd.vapid.public_key.is_empty() {

View File

@@ -9,7 +9,7 @@ use lapin::{
Channel, Connection, ConnectionProperties,
};
use revolt_config::{config, Settings};
use revolt_database::Database;
use revolt_database::{Database, amqp::consumer::{Consumer, Delegate}};
use tokio::signal::ctrl_c;
mod consumers;
@@ -23,8 +23,6 @@ use consumers::{
outbound::{apn::ApnsOutboundConsumer, fcm::FcmOutboundConsumer, vapid::VapidOutboundConsumer},
};
use crate::utils::{Consumer, Delegate};
#[tokio::main(flavor = "multi_thread", worker_threads = 2)]
async fn main() {
// Configure logging and environment
@@ -275,11 +273,12 @@ where
consumer.tag()
);
let delegate = Delegate(
let delegate = Delegate::new(
F::create(
db.clone(),
connection.clone(),
channel.clone(),
()
)
.await,
);

View File

@@ -1,5 +1,3 @@
mod renderer;
mod consumer;
pub use renderer::render_notification_content;
pub use consumer::{Consumer, Delegate};
pub use renderer::render_notification_content;

View File

@@ -1,16 +1,17 @@
[package]
name = "revolt-searchd"
version = "0.12.0"
version = "0.13.7"
edition = "2024"
license = "AGPL-3.0-or-later"
publish = false
[dependencies]
tokio = { version = "1.50.0", features = ["macros"] }
revolt-database = { version = "0.12.0", path = "../../core/database" }
revolt-config = { version = "0.12.0", path = "../../core/config" }
revolt-search = { version = "0.12.0", path = "../../core/search" }
amqprs = "1.7.0"
async-trait = "0.1.89"
log = "0.4.29"
serde_json = "1"
tokio = { workspace = true, features = ["macros"] }
revolt-database = { workspace = true }
revolt-config = { workspace = true }
revolt-search = { workspace = true }
lapin = { workspace = true }
async-trait = { workspace = true }
log = { workspace = true }
serde_json = { workspace = true }
anyhow = { workspace = true }

View File

@@ -1,34 +1,37 @@
use amqprs::{
BasicProperties, Deliver,
channel::{BasicAckArguments, BasicRejectArguments, Channel},
consumer::AsyncConsumer,
};
use std::sync::Arc;
use lapin::{Channel, Connection, message::Delivery, options::{BasicAckOptions, BasicRejectOptions}};
use async_trait::async_trait;
use revolt_database::{Database, events::rabbit::ChannelDeletePayload};
use revolt_database::{Database, events::rabbit::ChannelDeletePayload, amqp::consumer::Consumer};
use revolt_search::ElasticsearchClient;
use anyhow::Result;
#[allow(unused)]
#[derive(Clone)]
pub struct ChannelDeleteConsumer {
client: ElasticsearchClient,
database: Database,
}
impl ChannelDeleteConsumer {
pub fn new(client: ElasticsearchClient, database: Database) -> Self {
Self { client, database }
}
connection: Arc<Connection>,
channel: Arc<Channel>,
}
#[async_trait]
impl AsyncConsumer for ChannelDeleteConsumer {
async fn consume(
&mut self,
channel: &Channel,
deliver: Deliver,
_basic_properties: BasicProperties,
content: Vec<u8>,
) {
let payload = serde_json::from_slice::<ChannelDeletePayload>(&content)
impl Consumer<ElasticsearchClient> for ChannelDeleteConsumer {
async fn create(database: Database, connection: Arc<Connection>, channel: Arc<Channel>, client: ElasticsearchClient) -> Self {
Self {
client,
database,
connection,
channel,
}
}
fn channel(&self) -> &Arc<Channel> {
&self.channel
}
async fn consume(&self, delivery: Delivery) -> Result<()> {
let payload = serde_json::from_slice::<ChannelDeletePayload>(&delivery.data)
.expect("Failed to decode message");
log::debug!("Received channel delete {payload:?}");
@@ -38,15 +41,19 @@ impl AsyncConsumer for ChannelDeleteConsumer {
.await
.is_ok()
{
channel
.basic_ack(BasicAckArguments::new(deliver.delivery_tag(), false))
self.channel
.basic_ack(delivery.delivery_tag, BasicAckOptions::default())
.await
.expect("Failed to ack");
} else {
channel
.basic_reject(BasicRejectArguments::new(deliver.delivery_tag(), true))
self.channel
.basic_reject(delivery.delivery_tag, BasicRejectOptions {
requeue: true,
})
.await
.expect("Failed to reject");
}
};
Ok(())
}
}

View File

@@ -1,33 +1,38 @@
use amqprs::{
BasicProperties, Deliver,
channel::{BasicAckArguments, BasicRejectArguments, Channel},
consumer::AsyncConsumer,
};
use async_trait::async_trait;
use revolt_database::{Database, events::rabbit::MessageCreatePayload};
use revolt_search::ElasticsearchClient;
use std::sync::Arc;
use lapin::{Channel, Connection, message::Delivery, options::{BasicAckOptions, BasicRejectOptions}};
use async_trait::async_trait;
use revolt_database::{Database, events::rabbit::MessageCreatePayload, amqp::consumer::Consumer};
use revolt_search::ElasticsearchClient;
use anyhow::Result;
#[allow(unused)]
#[derive(Clone)]
pub struct MessageConsumer {
client: ElasticsearchClient,
database: Database,
}
impl MessageConsumer {
pub fn new(client: ElasticsearchClient, database: Database) -> Self {
Self { client, database }
}
connection: Arc<Connection>,
channel: Arc<Channel>,
}
#[async_trait]
impl AsyncConsumer for MessageConsumer {
async fn consume(
&mut self,
channel: &Channel,
deliver: Deliver,
_basic_properties: BasicProperties,
content: Vec<u8>,
) {
let payload = serde_json::from_slice::<MessageCreatePayload>(&content)
impl Consumer<ElasticsearchClient> for MessageConsumer {
async fn create(database: Database, connection: Arc<Connection>, channel: Arc<Channel>, client: ElasticsearchClient) -> Self {
Self {
client,
database,
connection,
channel,
}
}
fn channel(&self) -> &Arc<Channel> {
&self.channel
}
async fn consume(&self, delivery: Delivery) -> Result<()> {
let payload = serde_json::from_slice::<MessageCreatePayload>(&delivery.data)
.expect("Failed to decode message");
log::debug!("Received message {payload:?}");
@@ -37,15 +42,19 @@ impl AsyncConsumer for MessageConsumer {
.await
.is_ok()
{
channel
.basic_ack(BasicAckArguments::new(deliver.delivery_tag(), false))
self.channel
.basic_ack(delivery.delivery_tag, BasicAckOptions::default())
.await
.expect("Failed to ack");
} else {
channel
.basic_reject(BasicRejectArguments::new(deliver.delivery_tag(), true))
self.channel
.basic_reject(delivery.delivery_tag, BasicRejectOptions {
requeue: true,
})
.await
.expect("Failed to reject");
}
};
Ok(())
}
}

View File

@@ -1,34 +1,39 @@
use amqprs::{
BasicProperties, Deliver,
channel::{BasicAckArguments, BasicRejectArguments, Channel},
consumer::AsyncConsumer,
};
use std::sync::Arc;
use lapin::{Channel, Connection, message::Delivery, options::{BasicAckOptions, BasicRejectOptions}};
use async_trait::async_trait;
use revolt_database::{Database, events::rabbit::MessageDeletePayload};
use revolt_database::{Database, events::rabbit::MessageDeletePayload, amqp::consumer::Consumer};
use revolt_search::ElasticsearchClient;
use anyhow::Result;
#[allow(unused)]
#[derive(Clone)]
pub struct MessageDeleteConsumer {
client: ElasticsearchClient,
database: Database,
connection: Arc<Connection>,
channel: Arc<Channel>,
}
impl MessageDeleteConsumer {
pub fn new(client: ElasticsearchClient, database: Database) -> Self {
Self { client, database }
}
}
#[async_trait]
impl AsyncConsumer for MessageDeleteConsumer {
async fn consume(
&mut self,
channel: &Channel,
deliver: Deliver,
_basic_properties: BasicProperties,
content: Vec<u8>,
) {
let payload = serde_json::from_slice::<MessageDeletePayload>(&content)
impl Consumer<ElasticsearchClient> for MessageDeleteConsumer {
async fn create(database: Database, connection: Arc<Connection>, channel: Arc<Channel>, client: ElasticsearchClient) -> Self {
Self {
client,
database,
connection,
channel,
}
}
fn channel(&self) -> &Arc<Channel> {
&self.channel
}
async fn consume(&self, delivery: Delivery) -> Result<()> {
let payload = serde_json::from_slice::<MessageDeletePayload>(&delivery.data)
.expect("Failed to decode message");
log::debug!("Received message delete {payload:?}");
@@ -38,15 +43,19 @@ impl AsyncConsumer for MessageDeleteConsumer {
.await
.is_ok()
{
channel
.basic_ack(BasicAckArguments::new(deliver.delivery_tag(), false))
self.channel
.basic_ack(delivery.delivery_tag, BasicAckOptions::default())
.await
.expect("Failed to ack");
} else {
channel
.basic_reject(BasicRejectArguments::new(deliver.delivery_tag(), true))
self.channel
.basic_reject(delivery.delivery_tag, BasicRejectOptions {
requeue: true,
})
.await
.expect("Failed to reject");
}
};
Ok(())
}
}

View File

@@ -1,33 +1,38 @@
use amqprs::{
BasicProperties, Deliver,
channel::{BasicAckArguments, BasicRejectArguments, Channel},
consumer::AsyncConsumer,
};
use async_trait::async_trait;
use revolt_database::{Database, events::rabbit::MessageEditPayload};
use revolt_search::ElasticsearchClient;
use std::sync::Arc;
use lapin::{Channel, Connection, message::Delivery, options::{BasicAckOptions, BasicRejectOptions}};
use async_trait::async_trait;
use revolt_database::{Database, events::rabbit::MessageEditPayload, amqp::consumer::Consumer};
use revolt_search::ElasticsearchClient;
use anyhow::Result;
#[derive(Clone)]
pub struct MessageEditConsumer {
client: ElasticsearchClient,
database: Database,
}
impl MessageEditConsumer {
pub fn new(client: ElasticsearchClient, database: Database) -> Self {
Self { client, database }
}
connection: Arc<Connection>,
channel: Arc<Channel>,
}
#[async_trait]
impl AsyncConsumer for MessageEditConsumer {
async fn consume(
&mut self,
channel: &Channel,
deliver: Deliver,
_basic_properties: BasicProperties,
content: Vec<u8>,
) {
let payload = serde_json::from_slice::<MessageEditPayload>(&content)
impl Consumer<ElasticsearchClient> for MessageEditConsumer {
async fn create(database: Database, connection: Arc<Connection>, channel: Arc<Channel>, client: ElasticsearchClient) -> Self {
Self {
client,
database,
connection,
channel,
}
}
fn channel(&self) -> &Arc<Channel> {
&self.channel
}
async fn consume(&self, delivery: Delivery) -> Result<()> {
let payload = serde_json::from_slice::<MessageEditPayload>(&delivery.data)
.expect("Failed to decode message");
log::debug!("Received edit message {payload:?}");
@@ -37,15 +42,19 @@ impl AsyncConsumer for MessageEditConsumer {
.await
.is_ok()
{
channel
.basic_ack(BasicAckArguments::new(deliver.delivery_tag(), false))
self.channel
.basic_ack(delivery.delivery_tag, BasicAckOptions::default())
.await
.expect("Failed to ack");
} else {
channel
.basic_reject(BasicRejectArguments::new(deliver.delivery_tag(), true))
self.channel
.basic_reject(delivery.delivery_tag, BasicRejectOptions {
requeue: true,
})
.await
.expect("Failed to reject");
}
};
Ok(())
}
}

View File

@@ -1,16 +1,14 @@
use amqprs::{
FieldTable,
channel::{
BasicConsumeArguments, Channel, ExchangeDeclareArguments, QueueBindArguments,
QueueDeclareArguments,
},
connection::{Connection, OpenConnectionArguments},
consumer::AsyncConsumer,
use lapin::{
options::{BasicConsumeOptions, ExchangeDeclareOptions, QueueBindOptions, QueueDeclareOptions},
types::{AMQPValue, FieldTable},
Channel, Connection, ConnectionProperties,
};
use log::info;
use revolt_config::{Settings, config, configure};
use revolt_database::DatabaseInfo;
use revolt_database::{Database, DatabaseInfo, amqp::consumer::{Consumer, Delegate}};
use revolt_search::ElasticsearchClient;
use tokio::{signal::ctrl_c, spawn};
use std::{marker::PhantomData, sync::Arc};
mod consumers;
mod index;
@@ -38,53 +36,63 @@ async fn _main() {
client.setup_indexes().await.unwrap();
};
let connection = Connection::open(&OpenConnectionArguments::new(
&config.rabbit.host,
config.rabbit.port,
&config.rabbit.username,
&config.rabbit.password,
))
.await
.expect("Failed to connect to RabbitMQ");
let connection = Arc::new(
Connection::connect(
&format!(
"amqp://{}:{}@{}:{}",
&config.rabbit.username,
&config.rabbit.password,
&config.rabbit.host,
&config.rabbit.port,
),
ConnectionProperties::default(),
)
.await
.expect("Failed to connect to RabbitMQ"),
);
let mut channels = Vec::new();
channels.push(
make_queue_and_consume(
&config,
make_queue_and_consume::<consumers::MessageConsumer>(
&client,
&db,
&connection,
&config,
&config.elasticsearch.message_queue,
consumers::MessageConsumer::new(client.clone(), db.clone()),
)
.await,
);
channels.push(
make_queue_and_consume(
&config,
make_queue_and_consume::<consumers::MessageEditConsumer>(
&client,
&db,
&connection,
&config,
&config.elasticsearch.message_edit_queue,
consumers::MessageEditConsumer::new(client.clone(), db.clone()),
)
.await,
);
channels.push(
make_queue_and_consume(
&config,
make_queue_and_consume::<consumers::MessageDeleteConsumer>(
&client,
&db,
&connection,
&config,
&config.elasticsearch.message_delete_queue,
consumers::MessageDeleteConsumer::new(client.clone(), db.clone()),
)
.await,
);
channels.push(
make_queue_and_consume(
&config,
make_queue_and_consume::<consumers::ChannelDeleteConsumer>(
&client,
&db,
&connection,
&config,
&config.elasticsearch.channel_delete_queue,
consumers::ChannelDeleteConsumer::new(client.clone(), db.clone()),
)
.await,
);
@@ -98,10 +106,10 @@ async fn _main() {
task = Some(spawn(index::index_existing_messages(db, client.clone())));
}
ctrl_c().await.expect("Failed to wait for ctrl-c");
ctrl_c().await.unwrap();
for channel in channels {
channel.close().await.unwrap();
let _ = channel.close(0, "close".into()).await;
}
if let Some(task) = task {
@@ -109,52 +117,85 @@ async fn _main() {
}
}
async fn make_queue_and_consume<F: AsyncConsumer + Send + 'static>(
async fn make_queue_and_consume<F: Consumer<ElasticsearchClient>>(
client: &ElasticsearchClient,
db: &Database,
connection: &Arc<Connection>,
config: &Settings,
connection: &Connection,
queue: &str,
consumer: F,
) -> Channel {
let channel = connection.open_channel(None).await.unwrap();
queue_name: &str,
) -> Arc<Channel> {
let channel = Arc::new(connection.create_channel().await.unwrap());
channel
.exchange_declare(
ExchangeDeclareArguments::new(&config.elasticsearch.exchange, "direct")
.durable(true)
.finish(),
config.elasticsearch.exchange.clone().into(),
lapin::ExchangeKind::Direct,
ExchangeDeclareOptions {
durable: true,
..Default::default()
},
FieldTable::default(),
)
.await
.expect("Failed to declare pushd exchange");
.expect("Failed to declare exchange");
let mut table = FieldTable::new();
table.insert("x-queue-type".try_into().unwrap(), "quorum".into());
let mut table = FieldTable::default();
table.insert("x-queue-type".try_into().unwrap(), AMQPValue::LongString("quorum".into()));
_ = channel
.queue_declare(
QueueDeclareArguments::new(queue)
.durable(true)
.arguments(table)
.finish(),
)
let args = QueueDeclareOptions {
durable: true,
..Default::default()
};
channel
.queue_declare(queue_name.into(), args, table)
.await
.unwrap()
.unwrap();
channel
.queue_bind(QueueBindArguments::new(
queue,
&config.elasticsearch.exchange,
queue,
))
.queue_bind(
queue_name.into(),
config.elasticsearch.exchange.clone().into(),
queue_name.into(),
QueueBindOptions::default(),
FieldTable::default(),
)
.await
.expect("This probably means the revolt.messages exchange does not exist in rabbitmq!");
.expect(
"This probably means the revolt.messages exchange does not exist in rabbitmq!",
);
let args = BasicConsumeArguments::new(queue, "")
.manual_ack(true)
.finish();
let tag = channel.basic_consume(consumer, args).await.unwrap();
log::info!("Consuming queue {queue}, tag {tag}");
let consumer = channel
.basic_consume(
queue_name.into(),
"".into(),
BasicConsumeOptions {
no_ack: true,
..Default::default()
},
FieldTable::default(),
)
.await
.unwrap();
info!(
"Consuming routing key {} as queue {}, tag {}",
queue_name,
queue_name,
consumer.tag()
);
let delegate = Delegate::new(
F::create(
db.clone(),
connection.clone(),
channel.clone(),
client.clone(),
)
.await,
);
consumer.set_delegate(delegate);
channel
}

View File

@@ -1,4 +1,4 @@
use revolt_database::{util::reference::Reference, Database};
use revolt_database::{util::reference::Reference, Database, AMQP};
use revolt_result::{create_error, Result};
use rocket::State;
use rocket_empty::EmptyResponse;
@@ -10,6 +10,7 @@ use rocket_empty::EmptyResponse;
#[delete("/<webhook_id>/<token>/<message_id>")]
pub async fn webhook_delete_message(
db: &State<Database>,
amqp: &State<AMQP>,
webhook_id: Reference<'_>,
token: String,
message_id: Reference<'_>,
@@ -23,5 +24,5 @@ pub async fn webhook_delete_message(
return Err(create_error!(CannotDeleteMessage));
}
message.delete(db).await.map(|_| EmptyResponse)
message.delete(db, Some(amqp)).await.map(|_| EmptyResponse)
}

View File

@@ -2,6 +2,7 @@ use iso8601_timestamp::Timestamp;
use revolt_config::config;
use revolt_database::{
tasks::process_embeds::queue, util::reference::Reference, Database, Message, PartialMessage,
AMQP,
};
use revolt_models::v0::{self, DataEditMessage, Embed};
use revolt_models::validator::Validate;
@@ -15,6 +16,7 @@ use rocket::{serde::json::Json, State};
#[patch("/<webhook_id>/<token>/<message_id>", data = "<data>")]
pub async fn webhook_edit_message(
db: &State<Database>,
amqp: &State<AMQP>,
webhook_id: Reference<'_>,
token: String,
message_id: Reference<'_>,
@@ -73,7 +75,7 @@ pub async fn webhook_edit_message(
partial.embeds = Some(new_embeds);
message.update(db, partial, vec![]).await?;
message.update(db, Some(amqp), partial, vec![]).await?;
// Queue up a task for processing embeds
if let Some(content) = edit.content {