[ENH] Purge dirty log in background at the end of scheduled compaction (#4915)

## Description of changes

_Summarize the changes made by this PR._

- Improvements & Bug fixes
  - Updates the `verify_and_enrich_collections` method in compaction scheduler, so that it does not directly emit purge dirty log requests. Instead, it is stored in a temporary field in the scheduler, and will be periodically drained by the compaction manager
  - Updates the compaction manager, so that it will drain the deleted collections field in the compaction scheduler at the end of the scheduled compaction. The purge dirty log will run as a background task
  - Updates the purge dirty log impl so that it can take in multiple collection ids at a time, instead of only one at a time
- New functionality
  - Add a `PurgeDirtyLog` operator that invokes `<log_client>.purge_dirty_logs(...)` with a timeout.

## Test plan

_How are these changes tested?_

- [ ] Tests pass locally with `pytest` for python, `yarn test` for js, `cargo test` for rust

## Documentation Changes

_Are all docstrings for user-facing APIs updated if required? Do we need to make documentation changes in the [docs section](https://github.com/chroma-core/chroma/tree/main/docs/docs.trychroma.com)?_

Author: Sicheng-Pan

idl/chromadb/proto/logservice.proto

@@ -90,7 +90,7 @@ message UpdateCollectionLogOffsetResponse {
 }
 
 message PurgeDirtyForCollectionRequest {
-  string collection_id = 1;
+  repeated string collection_ids = 1;
 }
 
 message PurgeDirtyForCollectionResponse {

rust/log-service/src/bin/chroma-purge-dirty-log.rs

@@ -23,7 +23,7 @@ async fn main() {
     let mut client = LogServiceClient::new(logservice);
     let _resp = client
         .purge_dirty_for_collection(PurgeDirtyForCollectionRequest {
-            collection_id: collection_id.to_string(),
+            collection_ids: vec![collection_id.to_string()],
         })
         .await
         .expect("purge-dirty-log request should succeed");

rust/log-service/src/lib.rs

@@ -3,6 +3,7 @@
 use std::collections::hash_map::Entry;
 use std::collections::HashMap;
 use std::collections::HashSet;
+use std::str::FromStr;
 use std::sync::Arc;
 use std::time::{Duration, Instant, SystemTime};
 
@@ -1524,19 +1525,27 @@ impl LogServer {
         );
         async move {
             let request = request.into_inner();
-            let collection_id = Uuid::parse_str(&request.collection_id)
-                .map(CollectionUuid)
-                .map_err(|_| Status::invalid_argument("Failed to parse collection id"))?;
-            tracing::info!("purge_dirty_for_collection {collection_id}");
-            let dirty_marker = DirtyMarker::Purge { collection_id };
-            let dirty_marker_json = serde_json::to_string(&dirty_marker)
+            let collection_ids = request
+                .collection_ids
+                .iter()
+                .map(|id| CollectionUuid::from_str(id))
+                .collect::<Result<Vec<_>, _>>()
                 .map_err(|err| {
-                    tracing::error!("Failed to serialize dirty marker: {}", err);
-                    wal3::Error::Internal
+                    Status::invalid_argument(format!("Failed to parse collection id: {err}"))
+                })?;
+            tracing::info!("Purging collections in dirty log: [{collection_ids:?}]");
+            let dirty_marker_json_blobs = collection_ids
+                .into_iter()
+                .map(|collection_id| {
+                    serde_json::to_string(&DirtyMarker::Purge { collection_id })
+                        .map(String::into_bytes)
                 })
-                .map_err(|err| Status::new(err.code().into(), err.to_string()))?;
+                .collect::<Result<_, _>>()
+                .map_err(|err| {
+                    Status::internal(format!("Failed to serialize dirty marker: {err}"))
+                })?;
             self.dirty_log
-                .append(Vec::from(dirty_marker_json))
+                .append_many(dirty_marker_json_blobs)
                 .await
                 .map_err(|err| Status::new(err.code().into(), err.to_string()))?;
             Ok(Response::new(PurgeDirtyForCollectionResponse {}))

rust/log/src/grpc_log.rs

@@ -671,7 +671,7 @@ impl GrpcLog {
 
     pub(super) async fn purge_dirty_for_collection(
         &mut self,
-        collection_id: CollectionUuid,
+        collection_ids: Vec<CollectionUuid>,
     ) -> Result<(), GrpcPurgeDirtyForCollectionError> {
         let Some(assigner) = self.alt_client_assigner.as_mut() else {
             return Ok(());
@@ -681,14 +681,18 @@ impl GrpcLog {
         for client in assigner.all().into_iter() {
             let mut client = client.clone();
             let limiter = Arc::clone(&limiter);
+            let collection_ids_clone = collection_ids.clone();
             let request = async move {
                 // NOTE(rescrv): This can never fail and the result is to fail open.  Don't
                 // error-check.
                 let _permit = limiter.acquire().await;
                 client
                     .purge_dirty_for_collection(chroma_proto::PurgeDirtyForCollectionRequest {
                         // NOTE(rescrv):  Use the untyped string representation of the collection ID.
-                        collection_id: collection_id.0.to_string(),
+                        collection_ids: collection_ids_clone
+                            .iter()
+                            .map(ToString::to_string)
+                            .collect(),
                     })
                     .await
                     .map_err(GrpcPurgeDirtyForCollectionError::FailedToPurgeDirty)

rust/log/src/log.rs

@@ -161,12 +161,12 @@ impl Log {
     #[tracing::instrument(skip(self), err(Display))]
     pub async fn purge_dirty_for_collection(
         &mut self,
-        collection_id: CollectionUuid,
+        collection_ids: Vec<CollectionUuid>,
     ) -> Result<(), Box<dyn ChromaError>> {
         match self {
             Log::Sqlite(_) => unimplemented!("not implemented for sqlite"),
             Log::Grpc(log) => Ok(log
-                .purge_dirty_for_collection(collection_id)
+                .purge_dirty_for_collection(collection_ids)
                 .await
                 .map_err(|err| Box::new(err) as Box<dyn ChromaError>)?),
             Log::InMemory(_) => unimplemented!("not implemented for in memory"),

rust/sysdb/src/sysdb.rs

@@ -476,12 +476,12 @@ impl SysDb {
 
     pub async fn get_last_compaction_time(
         &mut self,
-        tanant_ids: Vec<String>,
+        tenant_ids: Vec<String>,
     ) -> Result<Vec<Tenant>, GetLastCompactionTimeError> {
         match self {
-            SysDb::Grpc(grpc) => grpc.get_last_compaction_time(tanant_ids).await,
+            SysDb::Grpc(grpc) => grpc.get_last_compaction_time(tenant_ids).await,
             SysDb::Sqlite(_) => todo!(),
-            SysDb::Test(test) => test.get_last_compaction_time(tanant_ids).await,
+            SysDb::Test(test) => test.get_last_compaction_time(tenant_ids).await,
         }
     }
 

rust/sysdb/src/test_sysdb.rs

@@ -239,11 +239,10 @@ impl TestSysDb {
         let inner = self.inner.lock();
         let mut tenants = Vec::new();
         for tenant_id in tenant_ids {
-            let last_compaction_time = match inner.tenant_last_compaction_time.get(&tenant_id) {
-                Some(last_compaction_time) => *last_compaction_time,
-                None => {
-                    return Err(GetLastCompactionTimeError::TenantNotFound);
-                }
+            let Some(last_compaction_time) =
+                inner.tenant_last_compaction_time.get(&tenant_id).cloned()
+            else {
+                continue;
             };
             tenants.push(Tenant {
                 id: tenant_id,

rust/worker/src/compactor/compaction_manager.rs

@@ -4,6 +4,10 @@ use super::OneOffCompactMessage;
 use super::RebuildMessage;
 use crate::compactor::types::ScheduledCompactMessage;
 use crate::config::CompactionServiceConfig;
+use crate::execution::operators::purge_dirty_log::PurgeDirtyLog;
+use crate::execution::operators::purge_dirty_log::PurgeDirtyLogError;
+use crate::execution::operators::purge_dirty_log::PurgeDirtyLogInput;
+use crate::execution::operators::purge_dirty_log::PurgeDirtyLogOutput;
 use crate::execution::orchestration::CompactOrchestrator;
 use crate::execution::orchestration::CompactionResponse;
 use async_trait::async_trait;
@@ -18,8 +22,10 @@ use chroma_memberlist::memberlist_provider::Memberlist;
 use chroma_segment::spann_provider::SpannProvider;
 use chroma_storage::Storage;
 use chroma_sysdb::SysDb;
+use chroma_system::wrap;
 use chroma_system::Dispatcher;
 use chroma_system::Orchestrator;
+use chroma_system::TaskResult;
 use chroma_system::{Component, ComponentContext, ComponentHandle, Handler, System};
 use chroma_types::CollectionUuid;
 use futures::stream::FuturesUnordered;
@@ -58,6 +64,7 @@ pub(crate) struct CompactionManager {
     max_compaction_size: usize,
     max_partition_size: usize,
     fetch_log_batch_size: u32,
+    purge_dirty_log_timeout_seconds: u64,
     on_next_memberlist_signal: Option<oneshot::Sender<()>>,
 }
 
@@ -92,6 +99,7 @@ impl CompactionManager {
         max_compaction_size: usize,
         max_partition_size: usize,
         fetch_log_batch_size: u32,
+        purge_dirty_log_timeout_seconds: u64,
     ) -> Self {
         CompactionManager {
             system,
@@ -110,6 +118,7 @@ impl CompactionManager {
             max_partition_size,
             on_next_memberlist_signal: None,
             fetch_log_batch_size,
+            purge_dirty_log_timeout_seconds,
         }
     }
 
@@ -187,7 +196,7 @@ impl CompactionManager {
     }
 
     #[instrument(name = "CompactionManager::rebuild_batch")]
-    pub(crate) async fn rebuild_batch(&mut self, collection_ids: Vec<CollectionUuid>) {
+    pub(crate) async fn rebuild_batch(&mut self, collection_ids: &[CollectionUuid]) {
         let _ = collection_ids
             .iter()
             .map(|id| self.compact(*id, true))
@@ -196,6 +205,41 @@ impl CompactionManager {
             .await;
     }
 
+    #[instrument(name = "CompactionManager::purge_dirty_log", skip(ctx))]
+    pub(crate) async fn purge_dirty_log(&mut self, ctx: &ComponentContext<Self>) {
+        let deleted_collection_uuids = self.scheduler.drain_deleted_collections();
+        if deleted_collection_uuids.is_empty() {
+            tracing::info!("Skipping purge dirty log because there is no deleted collections");
+            return;
+        }
+        let purge_dirty_log = PurgeDirtyLog {
+            log_client: self.log.clone(),
+            timeout: Duration::from_secs(self.purge_dirty_log_timeout_seconds),
+        };
+        let purge_dirty_log_input = PurgeDirtyLogInput {
+            collection_uuids: deleted_collection_uuids.clone(),
+        };
+        let purge_dirty_log_task = wrap(
+            Box::new(purge_dirty_log),
+            purge_dirty_log_input,
+            ctx.receiver(),
+        );
+        let Some(mut dispatcher) = self.dispatcher.clone() else {
+            tracing::error!("Unable to create background task to purge dirty log: Dispatcher is not set for compaction manager");
+            return;
+        };
+        if let Err(err) = dispatcher
+            .send(purge_dirty_log_task, Some(Span::current()))
+            .await
+        {
+            tracing::error!("Unable to create background task to purge dirty log: {err}");
+            return;
+        };
+        tracing::info!(
+            "Purging dirty logs for deleted collections: [{deleted_collection_uuids:?}]",
+        );
+    }
+
     pub(crate) fn set_dispatcher(&mut self, dispatcher: ComponentHandle<Dispatcher>) {
         self.dispatcher = Some(dispatcher);
     }
@@ -240,6 +284,7 @@ impl Configurable<(CompactionServiceConfig, System)> for CompactionManager {
         let max_compaction_size = config.compactor.max_compaction_size;
         let max_partition_size = config.compactor.max_partition_size;
         let fetch_log_batch_size = config.compactor.fetch_log_batch_size;
+        let purge_dirty_log_timeout_seconds = config.compactor.purge_dirty_log_timeout_seconds;
         let mut disabled_collections =
             HashSet::with_capacity(config.compactor.disabled_collections.len());
         for collection_id_str in &config.compactor.disabled_collections {
@@ -298,6 +343,7 @@ impl Configurable<(CompactionServiceConfig, System)> for CompactionManager {
             max_compaction_size,
             max_partition_size,
             fetch_log_batch_size,
+            purge_dirty_log_timeout_seconds,
         ))
     }
 }
@@ -342,6 +388,7 @@ impl Handler<ScheduledCompactMessage> for CompactionManager {
     ) {
         tracing::info!("CompactionManager: Performing scheduled compaction");
         let _ = self.compact_batch().await;
+        self.purge_dirty_log(ctx).await;
 
         // Compaction is done, schedule the next compaction
         ctx.scheduler.schedule(
@@ -382,7 +429,11 @@ impl Handler<RebuildMessage> for CompactionManager {
             "Rebuild started for collections: {:?}",
             message.collection_ids
         );
-        self.rebuild_batch(message.collection_ids).await;
+        self.rebuild_batch(&message.collection_ids).await;
+        tracing::info!(
+            "Rebuild completed for collections: {:?}",
+            message.collection_ids
+        );
     }
 }
 
@@ -400,6 +451,21 @@ impl Handler<Memberlist> for CompactionManager {
     }
 }
 
+#[async_trait]
+impl Handler<TaskResult<PurgeDirtyLogOutput, PurgeDirtyLogError>> for CompactionManager {
+    type Result = ();
+
+    async fn handle(
+        &mut self,
+        message: TaskResult<PurgeDirtyLogOutput, PurgeDirtyLogError>,
+        _ctx: &ComponentContext<CompactionManager>,
+    ) {
+        if let Err(err) = message.into_inner() {
+            tracing::error!("Error when purging dirty log: {err}");
+        }
+    }
+}
+
 pub struct RegisterOnReadySignal {
     pub on_ready_tx: oneshot::Sender<()>,
 }
@@ -618,6 +684,7 @@ mod tests {
         let max_compaction_size = 1000;
         let max_partition_size = 1000;
         let fetch_log_batch_size = 100;
+        let purge_dirty_log_timeout_seconds = 60;
 
         // Set assignment policy
         let mut assignment_policy = Box::new(RendezvousHashingAssignmentPolicy::default());
@@ -682,6 +749,7 @@ mod tests {
             max_compaction_size,
             max_partition_size,
             fetch_log_batch_size,
+            purge_dirty_log_timeout_seconds,
         );
 
         let dispatcher = Dispatcher::new(DispatcherConfig {

rust/worker/src/compactor/config.rs

@@ -18,6 +18,8 @@ pub struct CompactorConfig {
     pub disabled_collections: Vec<String>,
     #[serde(default = "CompactorConfig::default_fetch_log_batch_size")]
     pub fetch_log_batch_size: u32,
+    #[serde(default = "CompactorConfig::default_purge_dirty_log_timeout_seconds")]
+    pub purge_dirty_log_timeout_seconds: u64,
 }
 
 impl CompactorConfig {
@@ -52,6 +54,10 @@ impl CompactorConfig {
     fn default_fetch_log_batch_size() -> u32 {
         100
     }
+
+    fn default_purge_dirty_log_timeout_seconds() -> u64 {
+        60
+    }
 }
 
 impl Default for CompactorConfig {
@@ -65,6 +71,8 @@ impl Default for CompactorConfig {
             max_partition_size: CompactorConfig::default_max_partition_size(),
             disabled_collections: CompactorConfig::default_disabled_collections(),
             fetch_log_batch_size: CompactorConfig::default_fetch_log_batch_size(),
+            purge_dirty_log_timeout_seconds:
+                CompactorConfig::default_purge_dirty_log_timeout_seconds(),
         }
     }
 }