delete/upsert project/version messages

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) {
-        self.project_ids.lock().await.insert(project_id);
+    pub async fn push(
+        &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 mut project_ids = self.project_ids.lock().await;
-            mem::take(&mut *project_ids)
+        let operations = {
+            let mut operations = self.operations.lock().await;
+            mem::take(&mut *operations)
         };
 
-        if project_ids.is_empty() {
+        if operations.is_empty() {
             return Ok(());
         }
 
-        let mut project_ids = project_ids.into_iter();
-        while let Some(project_id) = project_ids.next() {
+        let mut operations = operations.into_events().into_iter();
+        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;
-                queued_project_ids.insert(project_id);
-                queued_project_ids.extend(project_ids);
+                let mut queued_operations = self.operations.lock().await;
+                queued_operations.push_event(operation);
+                for operation in operations {
+                    queued_operations.push_event(operation);
+                }
 
                 return Err(err.into());
             }
@@ -86,7 +109,76 @@ impl IncrementalSearchQueue {
     }
 }
 
-#[derive(Debug, Serialize)]
-pub struct SearchProjectIndexQueueEventData {
-    pub project_id: ProjectId,
+#[derive(Default)]
+struct PendingSearchIndexOperations {
+    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 seen_project_ids = HashSet::new();
+    let mut project_ids_to_change = 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) {
-            project_ids.push(event.project_id);
+        match event.into_data() {
+            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() {
-        reindex_projects(ro_pool, redis_pool, search_backend, &project_ids)
+    if !project_ids_to_remove.is_empty() {
+        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?}",
-        project_ids.len(),
+        "Changed {} projects and removed {} projects in {:.2?}",
+        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 {
-    project_id: ProjectId,
+#[serde(untagged)]
+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::{