From afb9f52cfe1c4854391eccb263781e6a1ed4569a Mon Sep 17 00:00:00 2001 From: aecsocket <43144841+aecsocket@users.noreply.github.com> Date: Fri, 26 Jun 2026 16:50:00 +0100 Subject: [PATCH] Add consume batching delay --- apps/labrinth/.env.docker-compose | 1 + apps/labrinth/.env.local | 1 + apps/labrinth/src/env.rs | 1 + apps/labrinth/src/search/incremental.rs | 41 +++---- .../src/search/incremental/consume.rs | 101 +++++++++++++----- 5 files changed, 91 insertions(+), 54 deletions(-) diff --git a/apps/labrinth/.env.docker-compose b/apps/labrinth/.env.docker-compose index 097a47c640..288e1aa0fc 100644 --- a/apps/labrinth/.env.docker-compose +++ b/apps/labrinth/.env.docker-compose @@ -25,6 +25,7 @@ ELASTICSEARCH_INDEX_PREFIX=labrinth ELASTICSEARCH_USERNAME=elastic ELASTICSEARCH_PASSWORD=elastic SEARCH_INDEX_CHUNK_SIZE=5000 +SEARCH_INCREMENTAL_INDEX_BATCH_DELAY_SECONDS=5 TYPESENSE_URL=http://localhost:8108 TYPESENSE_API_KEY=modrinth TYPESENSE_INDEX_PREFIX=labrinth diff --git a/apps/labrinth/.env.local b/apps/labrinth/.env.local index 7925e52a1e..7843c066f8 100644 --- a/apps/labrinth/.env.local +++ b/apps/labrinth/.env.local @@ -43,6 +43,7 @@ ELASTICSEARCH_USERNAME= ELASTICSEARCH_PASSWORD= SEARCH_INDEX_CHUNK_SIZE=5000 +SEARCH_INCREMENTAL_INDEX_BATCH_DELAY_SECONDS=5 TYPESENSE_URL=http://localhost:8108 TYPESENSE_API_KEY=modrinth TYPESENSE_INDEX_PREFIX=labrinth diff --git a/apps/labrinth/src/env.rs b/apps/labrinth/src/env.rs index bf99822156..3001a40d0c 100644 --- a/apps/labrinth/src/env.rs +++ b/apps/labrinth/src/env.rs @@ -160,6 +160,7 @@ vars! { // search SEARCH_BACKEND: crate::search::SearchBackendKind = crate::search::SearchBackendKind::Typesense; SEARCH_INDEX_CHUNK_SIZE: i64 = 5000i64; + SEARCH_INCREMENTAL_INDEX_BATCH_DELAY_SECONDS: u64 = 5u64; TYPESENSE_URL: String = "http://localhost:8108"; TYPESENSE_API_KEY: String = "modrinth"; TYPESENSE_INDEX_PREFIX: String = "labrinth"; diff --git a/apps/labrinth/src/search/incremental.rs b/apps/labrinth/src/search/incremental.rs index 07c5a05b55..88a9f1d6a2 100644 --- a/apps/labrinth/src/search/incremental.rs +++ b/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>>, + project_ids: Arc>>, kafka_client: actix_web::web::Data, } impl IncrementalSearchQueue { pub fn new(kafka_client: actix_web::web::Data) -> 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, diff --git a/apps/labrinth/src/search/incremental/consume.rs b/apps/labrinth/src/search/incremental/consume.rs index 69e3b30f54..040285a895 100644 --- a/apps/labrinth/src/search/incremental/consume.rs +++ b/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 { + 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>, + messages: impl IntoIterator>, ) -> 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(()) }