fix: search indexing performance and batching

apps/labrinth/.env.docker-compose

@@ -25,6 +25,8 @@ ELASTICSEARCH_INDEX_PREFIX=labrinth
 ELASTICSEARCH_USERNAME=elastic
 ELASTICSEARCH_PASSWORD=elastic
 SEARCH_INDEX_CHUNK_SIZE=5000
+SEARCH_INCREMENTAL_INDEX_BATCH_DELAY_SECONDS=5
+SEARCH_INCREMENTAL_INDEX_BATCH_MAX_SIZE=1000
 TYPESENSE_URL=http://localhost:8108
 TYPESENSE_API_KEY=modrinth
 TYPESENSE_INDEX_PREFIX=labrinth

apps/labrinth/.env.local

@@ -43,6 +43,8 @@ ELASTICSEARCH_USERNAME=
 ELASTICSEARCH_PASSWORD=
 
 SEARCH_INDEX_CHUNK_SIZE=5000
+SEARCH_INCREMENTAL_INDEX_BATCH_DELAY_SECONDS=5
+SEARCH_INCREMENTAL_INDEX_BATCH_MAX_SIZE=1000
 TYPESENSE_URL=http://localhost:8108
 TYPESENSE_API_KEY=modrinth
 TYPESENSE_INDEX_PREFIX=labrinth

apps/labrinth/src/background_task.rs

@@ -19,7 +19,7 @@ use crate::util::anrok;
 use actix_web::web;
 use clap::ValueEnum;
 use eyre::WrapErr;
-use tracing::info;
+use tracing::{info, instrument};
 
 #[derive(ValueEnum, Debug, Copy, Clone, PartialEq, Eq)]
 #[clap(rename_all = "kebab_case")]
@@ -50,6 +50,7 @@ pub enum BackgroundTask {
 
 impl BackgroundTask {
     #[allow(clippy::too_many_arguments)]
+    #[instrument(skip_all, fields(background_task = ?self))]
     pub async fn run(
         self,
         pool: PgPool,

apps/labrinth/src/env.rs

@@ -38,7 +38,7 @@ macro_rules! vars {
                     )]
                     let $field: Option<$ty> = {
                         let mut default = None::<$ty>;
-                        $( default = Some({ $default }.into()); )?
+                        $( default = Some(<$ty>::from({ $default })); )?
 
                         match parse_value::<$ty>(stringify!($field), default) {
                             Ok(value) => Some(value),
@@ -160,9 +160,12 @@ 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;
+    SEARCH_INCREMENTAL_INDEX_BATCH_MAX_SIZE: usize = 1000usize;
     TYPESENSE_URL: String = "http://localhost:8108";
     TYPESENSE_API_KEY: String = "modrinth";
     TYPESENSE_INDEX_PREFIX: String = "labrinth";
+    TYPESENSE_DELETE_BATCH_SIZE: usize = 10_000usize;
 
     // storage
     STORAGE_BACKEND: crate::file_hosting::FileHostKind = crate::file_hosting::FileHostKind::Local;

apps/labrinth/src/lib.rs

@@ -145,30 +145,6 @@ pub fn app_setup(
     ));
 
     if enable_background_tasks {
-        // The interval in seconds at which the local database is indexed
-        // for searching.  Defaults to 1 hour if unset.
-        let local_index_interval =
-            Duration::from_secs(ENV.LOCAL_INDEX_INTERVAL);
-        let pool_ref = pool.clone();
-        let redis_pool_ref = redis_pool.clone();
-        let search_backend_ref = search_backend.clone();
-        scheduler.run(local_index_interval, move || {
-            let pool_ref = pool_ref.clone();
-            let redis_pool_ref = redis_pool_ref.clone();
-            let search_backend = search_backend_ref.clone();
-            async move {
-                if let Err(err) = background_task::index_search(
-                    pool_ref,
-                    redis_pool_ref,
-                    search_backend,
-                )
-                .await
-                {
-                    warn!("Failed to index search: {err:?}");
-                }
-            }
-        });
-
         // Changes statuses of scheduled projects/versions
         let pool_ref = pool.clone();
         // TODO: Clear cache when these are run

apps/labrinth/src/queue/server_ping.rs

@@ -1,9 +1,9 @@
 use crate::database::DBProject;
-use crate::database::models::DBProjectId;
+use crate::database::models::{DBProjectId, DBVersionId};
 use crate::database::redis::RedisPool;
 use crate::env::ENV;
 use crate::models::exp;
-use crate::models::ids::ProjectId;
+use crate::models::ids::{ProjectId, VersionId};
 use crate::models::projects::ProjectStatus;
 use crate::search::incremental::IncrementalSearchQueue;
 use crate::{database::PgPool, util::error::Context};
@@ -175,14 +175,26 @@ impl ServerPingQueue {
                 }
 
                 if updated_project {
+                    let version_ids = sqlx::query_scalar!(
+                        "SELECT id FROM versions WHERE mod_id = $1",
+                        DBProjectId::from(*project_id) as DBProjectId,
+                    )
+                    .fetch_all(&self.db)
+                    .await
+                    .wrap_err("failed to fetch project version IDs")?
+                    .into_iter()
+                    .map(|version_id| VersionId::from(DBVersionId(version_id)))
+                    .collect::<Vec<_>>();
+
                     let clear_cache = DBProject::clear_cache(
                         (*project_id).into(),
                         None,
                         None,
                         &self.redis,
                     );
-                    let queue_search =
-                        self.incremental_search_queue.push(*project_id);
+                    let queue_search = self
+                        .incremental_search_queue
+                        .push(*project_id, version_ids);
 
                     let (clear_cache_result, _) =
                         join(clear_cache, queue_search).await;

apps/labrinth/src/routes/internal/moderation/tech_review.rs

@@ -1142,6 +1142,7 @@ async fn submit_report(
 
     if verdict == DelphiVerdict::Unsafe {
         crate::routes::v3::projects::clear_project_cache_and_queue_search(
+            &pool,
             &redis,
             &search_state,
             project_id,

apps/labrinth/src/routes/v3/organizations.rs

@@ -802,6 +802,7 @@ pub async fn organization_delete(
 
     for project_id in organization_project_ids {
         super::projects::clear_project_cache_and_queue_search(
+            &pool,
             &redis,
             &search_state,
             project_id,
@@ -969,6 +970,7 @@ pub async fn organization_projects_add(
         )
         .await?;
         super::projects::clear_project_cache_and_queue_search(
+            &pool,
             &redis,
             &search_state,
             project_item.inner.id,
@@ -1159,6 +1161,7 @@ pub async fn organization_projects_remove(
         )
         .await?;
         super::projects::clear_project_cache_and_queue_search(
+            &pool,
             &redis,
             &search_state,
             project_item.inner.id,

apps/labrinth/src/routes/v3/project_creation.rs

@@ -349,6 +349,7 @@ pub async fn project_create_internal(
     } else {
         transaction.commit().await?;
         super::projects::clear_project_cache_and_queue_search(
+            &client,
             &redis,
             &search_state,
             project_id.into(),
@@ -407,6 +408,7 @@ pub async fn project_create_with_id(
     } else {
         transaction.commit().await?;
         super::projects::clear_project_cache_and_queue_search(
+            &client,
             &redis,
             &search_state,
             project_id.into(),

apps/labrinth/src/routes/v3/project_creation/new.rs

@@ -342,6 +342,7 @@ pub async fn create(
         .wrap_internal_err("failed to commit transaction")?;
 
     super::super::projects::clear_project_cache_and_queue_search(
+        &db,
         &redis,
         &search_state,
         project_id.into(),

apps/labrinth/src/routes/v3/projects.rs

@@ -75,11 +75,42 @@ pub fn utoipa_config(
 }
 
 pub async fn clear_project_cache_and_queue_search(
+    pool: &PgPool,
     redis: &RedisPool,
     search_state: &SearchState,
     project_id: db_ids::DBProjectId,
     slug: Option<String>,
     clear_dependencies: Option<bool>,
+) -> Result<(), ApiError> {
+    let version_ids = sqlx::query_scalar!(
+        "SELECT id FROM versions WHERE mod_id = $1",
+        project_id as db_ids::DBProjectId,
+    )
+    .fetch_all(pool)
+    .await
+    .wrap_internal_err("failed to fetch project version IDs")?
+    .into_iter()
+    .map(|version_id| VersionId::from(db_ids::DBVersionId(version_id)))
+    .collect::<Vec<_>>();
+
+    clear_project_cache_and_queue_search_versions(
+        redis,
+        search_state,
+        project_id,
+        slug,
+        clear_dependencies,
+        version_ids,
+    )
+    .await
+}
+
+pub async fn clear_project_cache_and_queue_search_versions(
+    redis: &RedisPool,
+    search_state: &SearchState,
+    project_id: db_ids::DBProjectId,
+    slug: Option<String>,
+    clear_dependencies: Option<bool>,
+    version_ids: impl IntoIterator<Item = VersionId>,
 ) -> Result<(), ApiError> {
     db_models::DBProject::clear_cache(
         project_id,
@@ -88,7 +119,11 @@ pub async fn clear_project_cache_and_queue_search(
         redis,
     )
     .await?;
-    search_state.queue.push(project_id.into()).await;
+
+    search_state
+        .queue
+        .push(project_id.into(), version_ids)
+        .await;
 
     Ok(())
 }
@@ -1135,6 +1170,7 @@ pub async fn project_edit_internal(
     transaction.commit().await?;
 
     clear_project_cache_and_queue_search(
+        &pool,
         &redis,
         &search_state,
         project_item.inner.id,
@@ -1149,16 +1185,9 @@ pub async fn project_edit_internal(
         new_project.status.map(|status| status.is_searchable()),
     ) {
         search_state
-            .backend
-            .remove_documents(
-                &project_item
-                    .versions
-                    .into_iter()
-                    .map(|x| x.into())
-                    .collect::<Vec<_>>(),
-            )
-            .await
-            .wrap_internal_err("failed to remove documents")?;
+            .queue
+            .push_project_removal(project_item.inner.id.into())
+            .await;
     }
 
     Ok(HttpResponse::NoContent().body(""))
@@ -1640,6 +1669,7 @@ pub async fn projects_edit(
 
     for (project_id, slug) in changed_projects {
         clear_project_cache_and_queue_search(
+            &pool,
             &redis,
             &search_state,
             project_id,
@@ -1862,6 +1892,7 @@ pub async fn project_icon_edit_internal(
 
     transaction.commit().await?;
     clear_project_cache_and_queue_search(
+        &pool,
         &redis,
         &search_state,
         project_item.inner.id,
@@ -1977,6 +2008,7 @@ pub async fn delete_project_icon_internal(
 
     transaction.commit().await?;
     clear_project_cache_and_queue_search(
+        &pool,
         &redis,
         &search_state,
         project_item.inner.id,
@@ -2164,6 +2196,7 @@ pub async fn add_gallery_item_internal(
 
     transaction.commit().await?;
     clear_project_cache_and_queue_search(
+        &pool,
         &redis,
         &search_state,
         project_item.inner.id,
@@ -2370,6 +2403,7 @@ pub async fn edit_gallery_item_internal(
     transaction.commit().await?;
 
     clear_project_cache_and_queue_search(
+        &pool,
         &redis,
         &search_state,
         project_item.inner.id,
@@ -2510,6 +2544,7 @@ pub async fn delete_gallery_item_internal(
     transaction.commit().await?;
 
     clear_project_cache_and_queue_search(
+        &pool,
         &redis,
         &search_state,
         project_item.inner.id,
@@ -2652,27 +2687,18 @@ pub async fn project_delete_internal(
         .await
         .wrap_internal_err("failed to commit transaction")?;
 
-    search_state
-        .backend
-        .remove_documents(
-            &project
-                .versions
-                .into_iter()
-                .map(|x| x.into())
-                .collect::<Vec<_>>(),
-        )
-        .await
-        .wrap_internal_err("failed to remove project version documents")?;
-
     if result.is_some() {
-        clear_project_cache_and_queue_search(
-            &redis,
-            &search_state,
+        db_models::DBProject::clear_cache(
             project.inner.id,
             project.inner.slug,
             None,
+            &redis,
         )
         .await?;
+        search_state
+            .queue
+            .push_project_removal(project.inner.id.into())
+            .await;
         Ok(())
     } else {
         Err(ApiError::NotFound)