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

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