delete/upsert project/version messages

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 =
+                    let queue_search = self
-                        self.incremental_search_queue.push(*project_id);
+                        .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,6 +75,7 @@ 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,
@@ -88,7 +89,21 @@ pub async fn clear_project_cache_and_queue_search(
         redis,
     )
     .await?;
-    search_state.queue.push(project_id.into()).await;
+    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<_>>();
+
+    search_state
+        .queue
+        .push(project_id.into(), version_ids)
+        .await;
 
     Ok(())
 }
@@ -1135,6 +1150,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 +1165,9 @@ pub async fn project_edit_internal(
         new_project.status.map(|status| status.is_searchable()),
     ) {
         search_state
-            .backend
+            .queue
-            .remove_documents(
+            .push_project_removal(project_item.inner.id.into())
-                &project_item
+            .await;
-                    .versions
-                    .into_iter()
-                    .map(|x| x.into())
-                    .collect::<Vec<_>>(),
-            )
-            .await
-            .wrap_internal_err("failed to remove documents")?;
     }
 
     Ok(HttpResponse::NoContent().body(""))
@@ -1640,6 +1649,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 +1872,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 +1988,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 +2176,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 +2383,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 +2524,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 +2667,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(
+        db_models::DBProject::clear_cache(
-            &redis,
-            &search_state,
             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)

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

@@ -148,6 +148,7 @@ pub async fn version_create(
     } else if let Ok((_, project_id)) = &result {
         transaction.commit().await?;
         super::projects::clear_project_cache_and_queue_search(
+            &client,
             &redis,
             &search_state,
             *project_id,
@@ -565,7 +566,7 @@ pub async fn upload_file_to_version(
     let result = upload_file_to_version_inner(
         req,
         &mut payload,
-        client,
+        client.clone(),
         &mut transaction,
         redis.clone(),
         &**file_host,
@@ -591,6 +592,7 @@ pub async fn upload_file_to_version(
     } else if let Ok((_, project_id)) = &result {
         transaction.commit().await?;
         super::projects::clear_project_cache_and_queue_search(
+            &client,
             &redis,
             &search_state,
             *project_id,

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

@@ -762,6 +762,7 @@ pub async fn version_edit_helper(
             database::models::DBVersion::clear_cache(&version_item, &redis)
                 .await?;
             super::projects::clear_project_cache_and_queue_search(
+                &pool,
                 &redis,
                 &search_state,
                 version_item.inner.project_id,
@@ -1098,6 +1099,7 @@ pub async fn version_delete(
     transaction.commit().await?;
 
     super::projects::clear_project_cache_and_queue_search(
+        &pool,
         &redis,
         &search_state,
         version.inner.project_id,

apps/labrinth/src/search/incremental.rs

@@ -1,13 +1,18 @@
 pub mod consume;
 
-use std::{collections::HashSet, mem, sync::Arc, time::Duration};
+use std::{
+    collections::{HashMap, HashSet},
+    mem,
+    sync::Arc,
+    time::Duration,
+};
 
 use rdkafka::{producer::FutureRecord, util::Timeout};
 use serde::Serialize;
 use tokio::sync::Mutex;
 
 use crate::{
-    models::ids::ProjectId,
+    models::ids::{ProjectId, VersionId},
     util::kafka::{KAFKA_OPERATION_INTERVAL, KafkaClientState, KafkaEvent},
 };
 
@@ -17,20 +22,36 @@ const QUEUE_FLUSH_INTERVAL: Duration = Duration::from_secs(10);
 
 #[derive(Clone)]
 pub struct IncrementalSearchQueue {
-    project_ids: Arc<Mutex<HashSet<ProjectId>>>,
+    operations: Arc<Mutex<PendingSearchIndexOperations>>,
     kafka_client: actix_web::web::Data<KafkaClientState>,
 }
 
 impl IncrementalSearchQueue {
     pub fn new(kafka_client: actix_web::web::Data<KafkaClientState>) -> Self {
         Self {
-            project_ids: Arc::new(Mutex::new(HashSet::new())),
+            operations: Arc::new(Mutex::new(
+                PendingSearchIndexOperations::default(),
+            )),
             kafka_client,
         }
     }
 
-    pub async fn push(&self, project_id: ProjectId) {
+    pub async fn push(
-        self.project_ids.lock().await.insert(project_id);
+        &self,
+        project_id: ProjectId,
+        version_ids: impl IntoIterator<Item = VersionId>,
+    ) {
+        self.operations
+            .lock()
+            .await
+            .push_project_change(project_id, version_ids);
+    }
+
+    pub async fn push_project_removal(&self, project_id: ProjectId) {
+        self.operations
+            .lock()
+            .await
+            .push_project_removal(project_id);
     }
 
     pub async fn run(self) {
@@ -46,20 +67,20 @@ impl IncrementalSearchQueue {
     }
 
     pub async fn drain(&self) -> eyre::Result<()> {
-        let project_ids = {
+        let operations = {
-            let mut project_ids = self.project_ids.lock().await;
+            let mut operations = self.operations.lock().await;
-            mem::take(&mut *project_ids)
+            mem::take(&mut *operations)
         };
 
-        if project_ids.is_empty() {
+        if operations.is_empty() {
             return Ok(());
         }
 
-        let mut project_ids = project_ids.into_iter();
+        let mut operations = operations.into_events().into_iter();
-        while let Some(project_id) = project_ids.next() {
+        while let Some(operation) = operations.next() {
             let event = KafkaEvent::new(
                 SEARCH_PROJECT_INDEX_QUEUE_TOPIC,
-                SearchProjectIndexQueueEventData { project_id },
+                operation.clone(),
             );
             let event_id = event.event_metadata.event_id;
             let key = event_id.to_string();
@@ -74,9 +95,11 @@ impl IncrementalSearchQueue {
                 .send(record, Timeout::After(KAFKA_OPERATION_INTERVAL))
                 .await
             {
-                let mut queued_project_ids = self.project_ids.lock().await;
+                let mut queued_operations = self.operations.lock().await;
-                queued_project_ids.insert(project_id);
+                queued_operations.push_event(operation);
-                queued_project_ids.extend(project_ids);
+                for operation in operations {
+                    queued_operations.push_event(operation);
+                }
 
                 return Err(err.into());
             }
@@ -86,7 +109,76 @@ impl IncrementalSearchQueue {
     }
 }
 
-#[derive(Debug, Serialize)]
+#[derive(Default)]
-pub struct SearchProjectIndexQueueEventData {
+struct PendingSearchIndexOperations {
-    pub project_id: ProjectId,
+    changed_projects: HashMap<ProjectId, HashSet<VersionId>>,
+    removed_project_ids: HashSet<ProjectId>,
+}
+
+impl PendingSearchIndexOperations {
+    fn is_empty(&self) -> bool {
+        self.changed_projects.is_empty() && self.removed_project_ids.is_empty()
+    }
+
+    fn push_project_change(
+        &mut self,
+        project_id: ProjectId,
+        version_ids: impl IntoIterator<Item = VersionId>,
+    ) {
+        if !self.removed_project_ids.contains(&project_id) {
+            self.changed_projects
+                .entry(project_id)
+                .or_default()
+                .extend(version_ids);
+        }
+    }
+
+    fn push_project_removal(&mut self, project_id: ProjectId) {
+        self.changed_projects.remove(&project_id);
+        self.removed_project_ids.insert(project_id);
+    }
+
+    fn push_event(&mut self, event: SearchProjectIndexQueueEventData) {
+        match event {
+            SearchProjectIndexQueueEventData::ProjectChange {
+                project_id,
+                version_ids,
+            } => self.push_project_change(project_id, version_ids),
+            SearchProjectIndexQueueEventData::ProjectRemoval { project_id } => {
+                self.push_project_removal(project_id)
+            }
+        }
+    }
+
+    fn into_events(self) -> Vec<SearchProjectIndexQueueEventData> {
+        let mut events = Vec::with_capacity(
+            self.changed_projects.len() + self.removed_project_ids.len(),
+        );
+
+        events.extend(self.removed_project_ids.into_iter().map(|project_id| {
+            SearchProjectIndexQueueEventData::ProjectRemoval { project_id }
+        }));
+        events.extend(self.changed_projects.into_iter().map(
+            |(project_id, version_ids)| {
+                SearchProjectIndexQueueEventData::ProjectChange {
+                    project_id,
+                    version_ids: version_ids.into_iter().collect(),
+                }
+            },
+        ));
+
+        events
+    }
+}
+
+#[derive(Debug, Clone, Serialize)]
+#[serde(tag = "type", rename_all = "snake_case")]
+pub enum SearchProjectIndexQueueEventData {
+    ProjectChange {
+        project_id: ProjectId,
+        version_ids: Vec<VersionId>,
+    },
+    ProjectRemoval {
+        project_id: ProjectId,
+    },
 }

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

@@ -15,7 +15,7 @@ use tracing::{Instrument, info, info_span};
 use crate::{
     database::{PgPool, redis::RedisPool},
     env::ENV,
-    models::ids::ProjectId,
+    models::ids::{ProjectId, VersionId},
     search::{
         SearchBackend, incremental::SEARCH_PROJECT_INDEX_QUEUE_TOPIC,
         indexing::index_project_documents,
@@ -128,8 +128,9 @@ async fn consume_batch(
 ) -> eyre::Result<()> {
     let start = Instant::now();
 
-    let mut project_ids = Vec::new();
+    let mut project_ids_to_change = HashSet::new();
-    let mut seen_project_ids = HashSet::new();
+    let mut project_ids_to_remove = HashSet::new();
+    let mut version_ids_to_change = HashSet::new();
     let mut messages_to_commit = Vec::new();
 
     for message in messages {
@@ -164,24 +165,94 @@ async fn consume_batch(
             }
         };
 
-        if seen_project_ids.insert(event.project_id) {
+        match event.into_data() {
-            project_ids.push(event.project_id);
+            SearchProjectIndexQueueEventData::ProjectChange {
+                project_id,
+                version_ids,
+            } => {
+                project_ids_to_change.insert(project_id);
+                version_ids_to_change.extend(version_ids);
+            }
+            SearchProjectIndexQueueEventData::ProjectRemoval { project_id } => {
+                project_ids_to_remove.insert(project_id);
+            }
         }
         messages_to_commit.push(message);
     }
 
+    project_ids_to_change
+        .retain(|project_id| !project_ids_to_remove.contains(project_id));
+
+    let project_ids_to_change =
+        project_ids_to_change.into_iter().collect::<Vec<_>>();
+    let project_ids_to_remove =
+        project_ids_to_remove.into_iter().collect::<Vec<_>>();
+    let version_ids_to_change =
+        version_ids_to_change.into_iter().collect::<Vec<_>>();
+
     info!(
         kafka.message_count = messages_to_commit.len(),
-        "Read all Kafka messages in {:.2?}, found {} projects to reindex",
+        "Read all Kafka messages in {:.2?}, found {} projects to change, {} versions to change, and {} projects to remove",
         start.elapsed(),
-        project_ids.len(),
+        project_ids_to_change.len(),
+        version_ids_to_change.len(),
+        project_ids_to_remove.len(),
     );
     let start = Instant::now();
 
-    if !project_ids.is_empty() {
+    if !project_ids_to_remove.is_empty() {
-        reindex_projects(ro_pool, redis_pool, search_backend, &project_ids)
+        let operation_start = Instant::now();
+        info!(
+            project_count = project_ids_to_remove.len(),
+            "Removing project documents"
+        );
+        search_backend
+            .remove_project_documents(&project_ids_to_remove)
             .await
-            .wrap_err("failed to reindex project batch")?;
+            .wrap_err("failed to remove project documents")?;
+        info!(
+            project_count = project_ids_to_remove.len(),
+            "Removed project documents in {:.2?}",
+            operation_start.elapsed()
+        );
+    }
+
+    if !version_ids_to_change.is_empty() {
+        let operation_start = Instant::now();
+        info!(
+            version_count = version_ids_to_change.len(),
+            "Removing changed version documents"
+        );
+        search_backend
+            .remove_documents(&version_ids_to_change)
+            .await
+            .wrap_err("failed to remove changed version documents")?;
+        info!(
+            version_count = version_ids_to_change.len(),
+            "Removed changed version documents in {:.2?}",
+            operation_start.elapsed()
+        );
+    }
+
+    if !project_ids_to_change.is_empty() {
+        let operation_start = Instant::now();
+        info!(
+            project_count = project_ids_to_change.len(),
+            "Indexing changed projects"
+        );
+        index_changed_projects(
+            ro_pool,
+            redis_pool,
+            search_backend,
+            &project_ids_to_change,
+        )
+        .await
+        .wrap_err("failed to index changed project batch")?;
+        info!(
+            project_count = project_ids_to_change.len(),
+            "Indexed changed projects in {:.2?}",
+            operation_start.elapsed()
+        );
     }
 
     for message in messages_to_commit {
@@ -191,8 +262,9 @@ async fn consume_batch(
     }
 
     info!(
-        "Reindexed {} projects in {:.2?}",
+        "Changed {} projects and removed {} projects in {:.2?}",
-        project_ids.len(),
+        project_ids_to_change.len(),
+        project_ids_to_remove.len(),
         start.elapsed()
     );
 
@@ -218,6 +290,18 @@ pub async fn reindex_projects(
     search_backend.remove_project_documents(project_ids).await?;
 
     info!("Creating project documents");
+    index_changed_projects(ro_pool, redis_pool, search_backend, project_ids)
+        .await?;
+
+    Ok(())
+}
+
+async fn index_changed_projects(
+    ro_pool: &PgPool,
+    redis_pool: &RedisPool,
+    search_backend: &dyn SearchBackend,
+    project_ids: &[ProjectId],
+) -> eyre::Result<()> {
     let documents = index_project_documents(ro_pool, redis_pool, project_ids)
         .instrument(info_span!("index", batch_size = project_ids.len()))
         .await
@@ -236,6 +320,34 @@ pub async fn reindex_projects(
 }
 
 #[derive(Debug, Deserialize)]
-struct SearchProjectIndexQueueEvent {
+#[serde(untagged)]
-    project_id: ProjectId,
+enum SearchProjectIndexQueueEvent {
+    Current(SearchProjectIndexQueueEventData),
+    Legacy { project_id: ProjectId },
+}
+
+impl SearchProjectIndexQueueEvent {
+    fn into_data(self) -> SearchProjectIndexQueueEventData {
+        match self {
+            Self::Current(data) => data,
+            Self::Legacy { project_id } => {
+                SearchProjectIndexQueueEventData::ProjectChange {
+                    project_id,
+                    version_ids: Vec::new(),
+                }
+            }
+        }
+    }
+}
+
+#[derive(Debug, Deserialize)]
+#[serde(tag = "type", rename_all = "snake_case")]
+enum SearchProjectIndexQueueEventData {
+    ProjectChange {
+        project_id: ProjectId,
+        version_ids: Vec<VersionId>,
+    },
+    ProjectRemoval {
+        project_id: ProjectId,
+    },
 }

apps/labrinth/src/search/indexing.rs

@@ -7,7 +7,7 @@ use itertools::Itertools;
 use regex::Regex;
 use std::collections::HashMap;
 use std::sync::LazyLock;
-use tracing::{debug, info, warn};
+use tracing::{info, warn};
 
 use crate::database::PgPool;
 use crate::database::models::loader_fields::{