feat: use bulk indexing for existing messages

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

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<Message>) -> Result<(), Error> {
+        let mut ops = Vec::<BulkOperation<Value>>::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);

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);
         }
     }
-}
+
+    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.")
+}

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");