From d568b4369ef2fe3d773429a0efe54d52b03a2b3d Mon Sep 17 00:00:00 2001 From: Zomatree Date: Sat, 21 Mar 2026 06:59:56 +0000 Subject: [PATCH] feat: use bulk indexing for existing messages Signed-off-by: Zomatree --- crates/core/search/src/client.rs | 33 +++++++++++++++++++++++----- crates/daemons/searchd/src/index.rs | 34 ++++++++++++++++++++++++----- crates/daemons/searchd/src/main.rs | 5 ++--- 3 files changed, 58 insertions(+), 14 deletions(-) diff --git a/crates/core/search/src/client.rs b/crates/core/search/src/client.rs index b2d69d29..9b128a90 100644 --- a/crates/core/search/src/client.rs +++ b/crates/core/search/src/client.rs @@ -1,13 +1,10 @@ use std::fmt::Display; use elasticsearch::{ - CreateParts, DeleteByQueryParts, DeleteParts, Elasticsearch, IndexParts, SearchParts, - auth::Credentials, - http::{ + BulkOperation, BulkParts, CreateParts, DeleteByQueryParts, DeleteParts, Elasticsearch, IndexParts, SearchParts, auth::Credentials, http::{ response::Exception, transport::{SingleNodeConnectionPool, TransportBuilder}, - }, - indices::{IndicesCreateParts, IndicesDeleteParts}, + }, indices::{IndicesCreateParts, IndicesDeleteParts} }; use elasticsearch_dsl::{FieldSort, Query, Search, SearchResponse, Sort}; use revolt_database::{Database, Message}; @@ -262,6 +259,32 @@ impl ElasticsearchClient { Value::Object(map) } + pub async fn bulk_index_messages(&self, db: &Database, messages: Vec) -> Result<(), Error> { + let mut ops = Vec::>::new(); + + for message in messages { + let id = message.id.clone(); + let source = self.create_message_source(db, message); + + ops.push(BulkOperation::create(source).id(id).into()); + }; + + let exception = self + .inner + .bulk(BulkParts::Index("messages")) + .body(ops) + .send() + .await? + .exception() + .await?; + + if let Some(exception) = exception { + Err(exception.into()) + } else { + Ok(()) + } + } + pub async fn index_message(&self, db: &Database, message: Message) -> Result<(), Error> { let id = message.id.clone(); let source = self.create_message_source(db, message); diff --git a/crates/daemons/searchd/src/index.rs b/crates/daemons/searchd/src/index.rs index ab942371..fe06e0c6 100644 --- a/crates/daemons/searchd/src/index.rs +++ b/crates/daemons/searchd/src/index.rs @@ -1,14 +1,36 @@ use revolt_config::capture_error; -use revolt_database::{AMQP, Database}; +use revolt_database::Database; +use revolt_search::ElasticsearchClient; +pub async fn index_existing_messages(db: Database, client: ElasticsearchClient) { + log::info!("Starting bulk indexing."); -pub async fn index_existing_messages(db: Database, amqp: AMQP) { - let mut generator = db.fetch_all_messages().await.expect("Database query failed"); + let mut generator = db + .fetch_all_messages() + .await + .expect("Database query failed"); + + let mut chunk = Vec::new(); while let Some(message) = generator.next().await { - if let Err(e) = amqp.new_message_search(message).await { - log::error!("Error pushing message to RabbitMQ: {e}"); + chunk.push(message); + + if chunk.len() >= 1000 + && let Err(e) = client + .bulk_index_messages(&db, std::mem::take(&mut chunk)) + .await + { + log::error!("Error bulk indexing messages: {e}"); capture_error(&e); } } -} \ No newline at end of file + + if !chunk.is_empty() + && let Err(e) = client.bulk_index_messages(&db, chunk).await + { + log::error!("Error bulk indexing messages: {e}"); + capture_error(&e); + } + + log::info!("Finished bulk indexing.") +} diff --git a/crates/daemons/searchd/src/main.rs b/crates/daemons/searchd/src/main.rs index 2097f389..b087e0c2 100644 --- a/crates/daemons/searchd/src/main.rs +++ b/crates/daemons/searchd/src/main.rs @@ -7,7 +7,7 @@ use amqprs::{ consumer::AsyncConsumer, }; use revolt_config::{Settings, config, configure}; -use revolt_database::{AMQP, DatabaseInfo}; +use revolt_database::DatabaseInfo; use revolt_search::ElasticsearchClient; use tokio::{signal::ctrl_c, spawn}; @@ -88,8 +88,7 @@ async fn _main() { let mut task = None; if std::env::var("INDEX_ALL_MESSAGES").as_deref().is_ok_and(|v| v == "1") { - let amqp = AMQP::new_auto().await; - task = Some(spawn(index::index_existing_messages(db, amqp))); + task = Some(spawn(index::index_existing_messages(db, client.clone()))); } ctrl_c().await.expect("Failed to wait for ctrl-c");