Add consume batching delay

apps/labrinth/src/search/incremental.rs

@@ -1,6 +1,6 @@
 pub mod consume;
 
-use std::{mem, sync::Arc};
+use std::{collections::HashSet, mem, sync::Arc, time::Duration};
 
 use rdkafka::{producer::FutureRecord, util::Timeout};
 use serde::Serialize;
@@ -13,31 +13,29 @@ use crate::{
 
 pub const SEARCH_PROJECT_INDEX_QUEUE_TOPIC: &str =
     "public.labrinth.search-project-index-queue.v1";
+const QUEUE_FLUSH_INTERVAL: Duration = Duration::from_secs(10);
 
 #[derive(Clone)]
 pub struct IncrementalSearchQueue {
-    operations: Arc<Mutex<Vec<SearchIndexOperation>>>,
+    project_ids: Arc<Mutex<HashSet<ProjectId>>>,
     kafka_client: actix_web::web::Data<KafkaClientState>,
 }
 
 impl IncrementalSearchQueue {
     pub fn new(kafka_client: actix_web::web::Data<KafkaClientState>) -> Self {
         Self {
-            operations: Arc::new(Mutex::new(Vec::new())),
+            project_ids: Arc::new(Mutex::new(HashSet::new())),
             kafka_client,
         }
     }
 
     pub async fn push(&self, project_id: ProjectId) {
-        self.operations
-            .lock()
-            .await
-            .push(SearchIndexOperation { project_id });
+        self.project_ids.lock().await.insert(project_id);
     }
 
     pub async fn run(self) {
         loop {
-            tokio::time::sleep(KAFKA_OPERATION_INTERVAL).await;
+            tokio::time::sleep(QUEUE_FLUSH_INTERVAL).await;
 
             if let Err(err) = self.drain().await {
                 tracing::error!(
@@ -48,22 +46,20 @@ impl IncrementalSearchQueue {
     }
 
     pub async fn drain(&self) -> eyre::Result<()> {
-        let operations = {
-            let mut operations = self.operations.lock().await;
-            mem::take(&mut *operations)
+        let project_ids = {
+            let mut project_ids = self.project_ids.lock().await;
+            mem::take(&mut *project_ids)
         };
 
-        if operations.is_empty() {
+        if project_ids.is_empty() {
             return Ok(());
         }
 
-        let mut operations = operations.into_iter();
-        while let Some(operation) = operations.next() {
+        let mut project_ids = project_ids.into_iter();
+        while let Some(project_id) = project_ids.next() {
             let event = KafkaEvent::new(
                 SEARCH_PROJECT_INDEX_QUEUE_TOPIC,
-                SearchProjectIndexQueueEventData {
-                    project_id: operation.project_id,
-                },
+                SearchProjectIndexQueueEventData { project_id },
             );
             let event_id = event.event_metadata.event_id;
             let key = event_id.to_string();
@@ -78,9 +74,9 @@ impl IncrementalSearchQueue {
                 .send(record, Timeout::After(KAFKA_OPERATION_INTERVAL))
                 .await
             {
-                let mut queued_operations = self.operations.lock().await;
-                queued_operations.push(operation);
-                queued_operations.extend(operations);
+                let mut queued_project_ids = self.project_ids.lock().await;
+                queued_project_ids.insert(project_id);
+                queued_project_ids.extend(project_ids);
 
                 return Err(err.into());
             }
@@ -90,11 +86,6 @@ impl IncrementalSearchQueue {
     }
 }
 
-#[derive(Debug, Clone)]
-pub struct SearchIndexOperation {
-    pub project_id: ProjectId,
-}
-
 #[derive(Debug, Serialize)]
 pub struct SearchProjectIndexQueueEventData {
     pub project_id: ProjectId,

apps/labrinth/src/search/incremental/consume.rs

@@ -1,16 +1,21 @@
 use actix_web::web;
 use eyre::WrapErr;
-use futures::FutureExt;
+use futures::never::Never;
 use rdkafka::{
     Message,
     consumer::{CommitMode, Consumer, StreamConsumer},
     message::BorrowedMessage,
 };
 use serde::Deserialize;
-use std::collections::HashSet;
+use std::{
+    collections::HashSet,
+    time::{Duration, Instant},
+};
+use tracing::info;
 
 use crate::{
     database::{PgPool, redis::RedisPool},
+    env::ENV,
     models::ids::ProjectId,
     search::{
         SearchBackend, incremental::SEARCH_PROJECT_INDEX_QUEUE_TOPIC,
@@ -22,8 +27,6 @@ use crate::{
     },
 };
 
-const BATCH_SIZE: usize = 100;
-
 pub async fn run(
     ro_pool: PgPool,
     redis_pool: RedisPool,
@@ -60,25 +63,57 @@ async fn consume(
     search_backend: &dyn SearchBackend,
     consumer: &StreamConsumer,
 ) -> eyre::Result<()> {
+    // keep buffer capacity (pre-)allocated
+    let mut messages = Vec::with_capacity(1024);
     loop {
-        let mut messages = Vec::with_capacity(BATCH_SIZE);
-        messages.push(
-            consumer
-                .recv()
-                .await
-                .wrap_err("failed to receive Kafka message")?,
+        messages.clear();
+
+        // wait for a first message to come in...
+        let first_message = consumer
+            .recv()
+            .await
+            .wrap_err("failed to receive Kafka message")?;
+        messages.push(first_message);
+
+        let delay = Duration::from_secs(
+            ENV.SEARCH_INCREMENTAL_INDEX_BATCH_DELAY_SECONDS,
+        );
+        info!(
+            "Received initial Kafka message; waiting {delay:.2?} for more to batch",
         );
 
-        while messages.len() < BATCH_SIZE {
-            let Some(message) = consumer.recv().now_or_never() else {
-                break;
-            };
-
-            messages.push(message.wrap_err("failed to receive Kafka message")?);
+        // ..then wait a while for more messages to batch up
+        // so that we can process a big batch to reindex
+        //
+        // do a little trick with an `AsyncFnMut` closure
+        // so that we can explicitly specify the return type
+        let mut collect_more_messages = async || -> eyre::Result<Never> {
+            loop {
+                let message = consumer
+                    .recv()
+                    .await
+                    .wrap_err("failed to receive Kafka message")?;
+                messages.push(message);
+            }
+        };
+        match tokio::time::timeout(delay, collect_more_messages()).await {
+            Err(_elapsed) => {}
+            Ok(Err(err)) => {
+                return Err(
+                    err.wrap_err("failed to receive more Kafka messages")
+                );
+            }
         }
 
-        consume_batch(ro_pool, redis_pool, search_backend, consumer, messages)
-            .await?;
+        info!("Consuming batch of {} messages", messages.len());
+        consume_batch(
+            ro_pool,
+            redis_pool,
+            search_backend,
+            consumer,
+            messages.drain(..),
+        )
+        .await?;
     }
 }
 
@@ -87,8 +122,10 @@ async fn consume_batch(
     redis_pool: &RedisPool,
     search_backend: &dyn SearchBackend,
     consumer: &StreamConsumer,
-    messages: Vec<BorrowedMessage<'_>>,
+    messages: impl IntoIterator<Item = BorrowedMessage<'_>>,
 ) -> eyre::Result<()> {
+    let start = Instant::now();
+
     let mut project_ids = Vec::new();
     let mut seen_project_ids = HashSet::new();
     let mut messages_to_commit = Vec::new();
@@ -131,19 +168,19 @@ async fn consume_batch(
         messages_to_commit.push(message);
     }
 
-    if project_ids.is_empty() {
-        return Ok(());
-    }
-
-    tracing::info!(
+    info!(
         kafka.message_count = messages_to_commit.len(),
-        project_count = project_ids.len(),
-        "Consumed incremental search index event batch"
+        "Read all Kafka messages in {:.2?}, found {} projects to reindex",
+        start.elapsed(),
+        project_ids.len(),
     );
+    let start = Instant::now();
 
-    reindex_projects(ro_pool, redis_pool, search_backend, &project_ids)
-        .await
-        .wrap_err("failed to reindex project batch")?;
+    if !project_ids.is_empty() {
+        reindex_projects(ro_pool, redis_pool, search_backend, &project_ids)
+            .await
+            .wrap_err("failed to reindex project batch")?;
+    }
 
     for message in messages_to_commit {
         consumer
@@ -151,6 +188,12 @@ async fn consume_batch(
             .wrap_err("failed to commit Kafka message")?;
     }
 
+    info!(
+        "Reindexed {} projects in {:.2?}",
+        project_ids.len(),
+        start.elapsed()
+    );
+
     Ok(())
 }