max batch size

This commit is contained in:
aecsocket
2026-06-26 18:24:10 +01:00
parent 7821b6c1cd
commit 003ceb9918
5 changed files with 25 additions and 15 deletions
+1
View File
@@ -26,6 +26,7 @@ ELASTICSEARCH_USERNAME=elastic
ELASTICSEARCH_PASSWORD=elastic ELASTICSEARCH_PASSWORD=elastic
SEARCH_INDEX_CHUNK_SIZE=5000 SEARCH_INDEX_CHUNK_SIZE=5000
SEARCH_INCREMENTAL_INDEX_BATCH_DELAY_SECONDS=5 SEARCH_INCREMENTAL_INDEX_BATCH_DELAY_SECONDS=5
SEARCH_INCREMENTAL_INDEX_BATCH_MAX_SIZE=1000
TYPESENSE_URL=http://localhost:8108 TYPESENSE_URL=http://localhost:8108
TYPESENSE_API_KEY=modrinth TYPESENSE_API_KEY=modrinth
TYPESENSE_INDEX_PREFIX=labrinth TYPESENSE_INDEX_PREFIX=labrinth
+1
View File
@@ -44,6 +44,7 @@ ELASTICSEARCH_PASSWORD=
SEARCH_INDEX_CHUNK_SIZE=5000 SEARCH_INDEX_CHUNK_SIZE=5000
SEARCH_INCREMENTAL_INDEX_BATCH_DELAY_SECONDS=5 SEARCH_INCREMENTAL_INDEX_BATCH_DELAY_SECONDS=5
SEARCH_INCREMENTAL_INDEX_BATCH_MAX_SIZE=1000
TYPESENSE_URL=http://localhost:8108 TYPESENSE_URL=http://localhost:8108
TYPESENSE_API_KEY=modrinth TYPESENSE_API_KEY=modrinth
TYPESENSE_INDEX_PREFIX=labrinth TYPESENSE_INDEX_PREFIX=labrinth
+2 -1
View File
@@ -38,7 +38,7 @@ macro_rules! vars {
)] )]
let $field: Option<$ty> = { let $field: Option<$ty> = {
let mut default = None::<$ty>; let mut default = None::<$ty>;
$( default = Some({ $default }.into()); )? $( default = Some(<$ty>::from({ $default })); )?
match parse_value::<$ty>(stringify!($field), default) { match parse_value::<$ty>(stringify!($field), default) {
Ok(value) => Some(value), Ok(value) => Some(value),
@@ -161,6 +161,7 @@ vars! {
SEARCH_BACKEND: crate::search::SearchBackendKind = crate::search::SearchBackendKind::Typesense; SEARCH_BACKEND: crate::search::SearchBackendKind = crate::search::SearchBackendKind::Typesense;
SEARCH_INDEX_CHUNK_SIZE: i64 = 5000i64; SEARCH_INDEX_CHUNK_SIZE: i64 = 5000i64;
SEARCH_INCREMENTAL_INDEX_BATCH_DELAY_SECONDS: u64 = 5u64; SEARCH_INCREMENTAL_INDEX_BATCH_DELAY_SECONDS: u64 = 5u64;
SEARCH_INCREMENTAL_INDEX_BATCH_MAX_SIZE: usize = 1000usize;
TYPESENSE_URL: String = "http://localhost:8108"; TYPESENSE_URL: String = "http://localhost:8108";
TYPESENSE_API_KEY: String = "modrinth"; TYPESENSE_API_KEY: String = "modrinth";
TYPESENSE_INDEX_PREFIX: String = "labrinth"; TYPESENSE_INDEX_PREFIX: String = "labrinth";
@@ -83,21 +83,24 @@ async fn consume(
); );
// ..then wait a while for more messages to batch up // ..then wait a while for more messages to batch up
// so that we can process a big batch to reindex // so that we can process a big batch to reindex.
// we stop until either we've reached the max batch size,
// or we've waited enough time - whichever is first.
// //
// do a little trick with an `AsyncFnMut` closure // do a little trick with an `AsyncFnMut` closure
// so that we can explicitly specify the return type // so that we can explicitly specify the return type
let mut collect_more_messages = async || -> eyre::Result<Never> { let mut collect_more_messages = async || -> eyre::Result<()> {
loop { while messages.len() < ENV.SEARCH_INCREMENTAL_INDEX_BATCH_MAX_SIZE {
let message = consumer let message = consumer
.recv() .recv()
.await .await
.wrap_err("failed to receive Kafka message")?; .wrap_err("failed to receive Kafka message")?;
messages.push(message); messages.push(message);
} }
eyre::Ok(())
}; };
match tokio::time::timeout(delay, collect_more_messages()).await { match tokio::time::timeout(delay, collect_more_messages()).await {
Err(_elapsed) => {} Ok(Ok(())) | Err(_) => {}
Ok(Err(err)) => { Ok(Err(err)) => {
return Err( return Err(
err.wrap_err("failed to receive more Kafka messages") err.wrap_err("failed to receive more Kafka messages")
@@ -224,6 +227,8 @@ pub async fn reindex_projects(
) )
})?; })?;
info!("Fetched all project documents, indexing into backend");
search_backend.index_documents(&documents).await?; search_backend.index_documents(&documents).await?;
Ok(()) Ok(())
+12 -10
View File
@@ -161,6 +161,8 @@ pub async fn index_project_documents(
.await .await
.wrap_err("failed to fetch project")?; .wrap_err("failed to fetch project")?;
info!("Fetched partial projects");
build_search_documents(pool, redis, db_projects).await build_search_documents(pool, redis, db_projects).await
} }
@@ -180,7 +182,7 @@ async fn build_search_documents(
.await .await
.wrap_err("failed to fetch query context")?; .wrap_err("failed to fetch query context")?;
debug!("Indexing local dependencies!"); info!("Indexing local dependencies!");
let dependencies: DashMap<DBProjectId, Vec<SearchProjectDependency>> = let dependencies: DashMap<DBProjectId, Vec<SearchProjectDependency>> =
sqlx::query!( sqlx::query!(
@@ -234,7 +236,7 @@ async fn build_search_documents(
ordering: i64, ordering: i64,
} }
debug!("Indexing local gallery!"); info!("Indexing local gallery!");
let mods_gallery: DashMap<DBProjectId, Vec<PartialGallery>> = sqlx::query!( let mods_gallery: DashMap<DBProjectId, Vec<PartialGallery>> = sqlx::query!(
" "
@@ -260,7 +262,7 @@ async fn build_search_documents(
) )
.await?; .await?;
debug!("Indexing local categories!"); info!("Indexing local categories!");
let categories: DashMap<DBProjectId, Vec<(String, bool)>> = sqlx::query!( let categories: DashMap<DBProjectId, Vec<(String, bool)>> = sqlx::query!(
" "
@@ -283,10 +285,10 @@ async fn build_search_documents(
) )
.await?; .await?;
debug!("Indexing local versions!"); info!("Indexing local versions!");
let mut versions = index_versions(pool, project_ids.clone()).await?; let mut versions = index_versions(pool, project_ids.clone()).await?;
debug!("Indexing local org owners!"); info!("Indexing local org owners!");
let mods_org_owners: DashMap<DBProjectId, ProjectOwner> = sqlx::query!( let mods_org_owners: DashMap<DBProjectId, ProjectOwner> = sqlx::query!(
" "
@@ -311,7 +313,7 @@ async fn build_search_documents(
}) })
.await?; .await?;
debug!("Indexing local team owners!"); info!("Indexing local team owners!");
let mods_team_owners: DashMap<DBProjectId, ProjectOwner> = sqlx::query!( let mods_team_owners: DashMap<DBProjectId, ProjectOwner> = sqlx::query!(
" "
@@ -335,7 +337,7 @@ async fn build_search_documents(
}) })
.await?; .await?;
debug!("Getting all loader fields!"); info!("Getting all loader fields!");
let loader_fields: Vec<QueryLoaderField> = sqlx::query!( let loader_fields: Vec<QueryLoaderField> = sqlx::query!(
" "
SELECT DISTINCT id, field, field_type, enum_type, min_val, max_val, optional SELECT DISTINCT id, field, field_type, enum_type, min_val, max_val, optional
@@ -356,7 +358,7 @@ async fn build_search_documents(
.await?; .await?;
let loader_fields: Vec<&QueryLoaderField> = loader_fields.iter().collect(); let loader_fields: Vec<&QueryLoaderField> = loader_fields.iter().collect();
debug!("Getting all loader field enum values!"); info!("Getting all loader field enum values!");
let loader_field_enum_values: Vec<QueryLoaderFieldEnumValue> = let loader_field_enum_values: Vec<QueryLoaderFieldEnumValue> =
sqlx::query!( sqlx::query!(
@@ -378,7 +380,7 @@ async fn build_search_documents(
.try_collect() .try_collect()
.await?; .await?;
debug!("Indexing loaders, project types!"); info!("Indexing loaders, project types!");
let mut uploads = Vec::new(); let mut uploads = Vec::new();
let total_len = db_projects.len(); let total_len = db_projects.len();
@@ -387,7 +389,7 @@ async fn build_search_documents(
count += 1; count += 1;
if count % 1000 == 0 { if count % 1000 == 0 {
debug!("projects index prog: {count}/{total_len}"); info!("projects index prog: {count}/{total_len}");
} }
let Some(( let Some((
_, _,