[ENH] Batch get info from sysdb for compaction scheduler

Author: sicheng

rust/worker/src/compactor/scheduler.rs

@@ -1,4 +1,4 @@
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
 use std::str::FromStr;
 
 use chroma_config::assignment::assignment_policy::AssignmentPolicy;
@@ -96,94 +96,94 @@ impl Scheduler {
         &mut self,
         collections: Vec<CollectionInfo>,
     ) -> Vec<CollectionRecord> {
-        let mut collection_records = Vec::new();
-        for collection_info in collections {
-            if self
-                .disabled_collections
-                .contains(&collection_info.collection_id)
-            {
-                tracing::info!(
-                    "Ignoring collection: {:?} because it disabled for compaction",
-                    collection_info.collection_id
-                );
-                continue;
-            }
-            // TODO: add a cache to avoid fetching the same collection multiple times
-            let result = self
-                .sysdb
-                .get_collections(GetCollectionsOptions {
-                    collection_id: Some(collection_info.collection_id),
-                    ..Default::default()
-                })
-                .await;
+        let mut disabled_collections = Vec::new();
+        let collection_id_to_log_info = collections
+            .into_iter()
+            .filter_map(|collection| {
+                if self
+                    .disabled_collections
+                    .contains(&collection.collection_id)
+                {
+                    disabled_collections.push(collection);
+                    None
+                } else {
+                    Some((collection.collection_id, collection))
+                }
+            })
+            .collect::<HashMap<_, _>>();
 
-            match result {
-                Ok(collection) => {
-                    if collection.is_empty() {
-                        self.deleted_collections
-                            .insert(collection_info.collection_id);
-                        continue;
-                    }
+        if !disabled_collections.is_empty() {
+            tracing::info!(
+                "Disabled collections are excluded from compaction: {disabled_collections:?}"
+            );
+        }
 
-                    // TODO: make querying the last compaction time in batch
-                    let log_position_in_collection = collection[0].log_position;
-                    let tenant_ids = vec![collection[0].tenant.clone()];
-                    let tenant = self.sysdb.get_last_compaction_time(tenant_ids).await;
-
-                    let last_compaction_time = match tenant {
-                        Ok(tenant) => {
-                            if tenant.is_empty() {
-                                tracing::info!(
-                                    "Ignoring collection: {:?}",
-                                    collection_info.collection_id
-                                );
-                                continue;
-                            }
-                            tenant[0].last_compaction_time
-                        }
-                        Err(e) => {
-                            tracing::error!("Error: {:?}", e);
-                            // Ignore this collection id for this compaction iteration
-                            tracing::info!(
-                                "Ignoring collection: {:?}",
-                                collection_info.collection_id
-                            );
-                            continue;
-                        }
-                    };
-
-                    let mut offset = collection_info.first_log_offset;
-                    // offset in log is the first offset in the log that has not been compacted. Note that
-                    // since the offset is the first offset of log we get from the log service, we should
-                    // use this offset to pull data from the log service.
-                    if log_position_in_collection + 1 < offset {
-                        panic!(
-                            "offset in sysdb ({}) is less than offset in log ({}) for {}",
-                            log_position_in_collection + 1,
-                            offset,
-                            collection[0].collection_id,
-                        )
-                    } else {
-                        // The offset in sysdb is the last offset that has been compacted.
-                        // We need to start from the next offset.
-                        offset = log_position_in_collection + 1;
-                    }
+        let collection_id_to_sysdb_info = match self
+            .sysdb
+            .get_collections(GetCollectionsOptions {
+                collection_ids: Some(collection_id_to_log_info.keys().cloned().collect()),
+                ..Default::default()
+            })
+            .await
+        {
+            Ok(collections) => collections
+                .into_iter()
+                .map(|collection| (collection.collection_id, collection))
+                .collect::<HashMap<_, _>>(),
+            Err(err) => {
+                tracing::error!("Unable to fetch collection information from sysdb: {err}");
+                return Vec::new();
+            }
+        };
 
-                    collection_records.push(CollectionRecord {
-                        collection_id: collection[0].collection_id,
-                        tenant_id: collection[0].tenant.clone(),
-                        last_compaction_time,
-                        first_record_time: collection_info.first_log_ts,
-                        offset,
-                        collection_version: collection[0].version,
-                        collection_logical_size_bytes: collection[0].size_bytes_post_compaction,
-                    });
-                }
-                Err(e) => {
-                    tracing::error!("Error: {:?}", e);
+        let sysdb_tenants = collection_id_to_sysdb_info
+            .values()
+            .map(|collection| collection.tenant.clone())
+            .collect();
+
+        let tenant_to_last_compaction_time =
+            match self.sysdb.get_last_compaction_time(sysdb_tenants).await {
+                Ok(tenants) => tenants
+                    .into_iter()
+                    .map(|tenant| (tenant.id, tenant.last_compaction_time))
+                    .collect::<HashMap<_, _>>(),
+                Err(err) => {
+                    tracing::error!(
+                        "Unable to fetch tenant last compaction time from sysdb: {err}"
+                    );
+                    return Vec::new();
                 }
+            };
+
+        let mut collection_records = Vec::new();
+        for (collection_id, log_info) in collection_id_to_log_info {
+            let Some(sysdb_info) = collection_id_to_sysdb_info.get(&collection_id) else {
+                self.deleted_collections.insert(collection_id);
+                continue;
+            };
+            let Some(last_compaction_time) = tenant_to_last_compaction_time
+                .get(&sysdb_info.tenant)
+                .cloned()
+            else {
+                tracing::error!("Unable to find last compaction info for tenant [{}]. Skipping compaction for collection [{collection_id}]", &sysdb_info.tenant);
+                continue;
+            };
+            let next_log_to_compact = sysdb_info.log_position.saturating_add(1);
+            if next_log_to_compact < log_info.first_log_offset {
+                tracing::error!("Next log to compact for collection [{collection_id}] at [{next_log_to_compact}] is already purged in log service: the first available log is at [{}]", log_info.first_log_offset);
+                continue;
             }
+            collection_records.push(CollectionRecord {
+                collection_id,
+                tenant_id: sysdb_info.tenant.clone(),
+                last_compaction_time,
+                first_record_time: log_info.first_log_ts,
+                offset: next_log_to_compact,
+                collection_version: sysdb_info.version,
+                collection_logical_size_bytes: sysdb_info.size_bytes_post_compaction,
+            });
         }
+
         collection_records
     }
 
@@ -505,8 +505,7 @@ mod tests {
     }
 
     #[tokio::test]
-    #[should_panic(expected = "is less than offset")]
-    async fn test_scheduler_panic() {
+    async fn test_scheduler_invalid_offsets() {
         let mut log = Log::InMemory(InMemoryLog::new());
         let in_memory_log = match log {
             Log::InMemory(ref mut in_memory_log) => in_memory_log,
@@ -600,17 +599,13 @@ mod tests {
                 },
             },
         );
-        let _ = log
-            .update_collection_log_offset(&tenant_1, collection_uuid_1, 2)
-            .await;
-
         let mut sysdb = SysDb::Test(TestSysDb::new());
 
         match sysdb {
             SysDb::Test(ref mut sysdb) => {
                 sysdb.add_collection(collection_1);
                 let last_compaction_time_1 = 2;
-                sysdb.add_tenant_last_compaction_time(tenant_1, last_compaction_time_1);
+                sysdb.add_tenant_last_compaction_time(tenant_1.clone(), last_compaction_time_1);
             }
             _ => panic!("Invalid sysdb type"),
         }
@@ -626,6 +621,27 @@ mod tests {
         let mut assignment_policy = Box::new(RendezvousHashingAssignmentPolicy::default());
         assignment_policy.set_members(vec![my_member.member_id.clone()]);
 
+        let mut scheduler = Scheduler::new(
+            my_member.member_id.clone(),
+            log.clone(),
+            sysdb.clone(),
+            scheduler_policy.clone(),
+            max_concurrent_jobs,
+            1,
+            assignment_policy.clone(),
+            HashSet::new(),
+        );
+
+        scheduler.set_memberlist(vec![my_member.clone()]);
+        scheduler.schedule().await;
+        let jobs = scheduler.get_jobs();
+        assert_eq!(jobs.count(), 1);
+
+        // Update compaction offset in log so that it's ahead of the value in sysdb
+        let _ = log
+            .update_collection_log_offset(&tenant_1, collection_uuid_1, 2)
+            .await;
+
         let mut scheduler = Scheduler::new(
             my_member.member_id.clone(),
             log,
@@ -639,5 +655,7 @@ mod tests {
 
         scheduler.set_memberlist(vec![my_member.clone()]);
         scheduler.schedule().await;
+        let jobs = scheduler.get_jobs();
+        assert_eq!(jobs.count(), 0);
     }
 }