[ENH]: Refactor compactor

Author: tanujnay112

Cargo.lock

@@ -1,11 +1,24 @@
 use std::collections::HashMap;
 use std::fmt::Debug;
 
-use chroma_error::ChromaError;
+use chroma_error::{ChromaError, ErrorCodes};
 use chroma_types::{CollectionUuid, LogRecord};
+use thiserror::Error;
 
 use crate::types::CollectionInfo;
 
+#[derive(Error, Debug)]
+pub enum InMemoryLogError {
+    #[error("Failed to update log offset (simulated failure)")]
+    UpdateOffsetFailed,
+}
+
+impl ChromaError for InMemoryLogError {
+    fn code(&self) -> ErrorCodes {
+        ErrorCodes::Internal
+    }
+}
+
 // This is used for testing only, it represents a log record that is stored in memory
 // internal to a mock log implementation
 #[derive(Clone)]
@@ -32,16 +45,22 @@ impl Debug for InternalLogRecord {
 pub struct InMemoryLog {
     collection_to_log: HashMap<CollectionUuid, Vec<InternalLogRecord>>,
     offsets: HashMap<CollectionUuid, i64>,
+    fail_update_offset: bool,
 }
 
 impl InMemoryLog {
     pub fn new() -> InMemoryLog {
         InMemoryLog {
             collection_to_log: HashMap::new(),
             offsets: HashMap::new(),
+            fail_update_offset: false,
         }
     }
 
+    pub fn set_fail_update_offset(&mut self, fail: bool) {
+        self.fail_update_offset = fail;
+    }
+
     pub fn add_log(&mut self, collection_id: CollectionUuid, log: InternalLogRecord) {
         let logs = self.collection_to_log.entry(collection_id).or_default();
         // Ensure that the log offset is correct. Since we only use the InMemoryLog for testing,
@@ -123,8 +142,12 @@ impl InMemoryLog {
         &mut self,
         collection_id: CollectionUuid,
         new_offset: i64,
-    ) {
+    ) -> Result<(), Box<dyn ChromaError>> {
+        if self.fail_update_offset {
+            return Err(Box::new(InMemoryLogError::UpdateOffsetFailed));
+        }
         self.offsets.insert(collection_id, new_offset);
+        Ok(())
     }
 
     pub(super) async fn scout_logs(

rust/log/src/in_memory_log.rs

@@ -166,8 +166,7 @@ impl Log {
                 .map_err(|e| Box::new(e) as Box<dyn ChromaError>),
             Log::InMemory(log) => {
                 log.update_collection_log_offset(collection_id, new_offset)
-                    .await;
-                Ok(())
+                    .await
             }
         }
     }

rust/log/src/log.rs

@@ -115,6 +115,33 @@ pub fn add_delete_generator(offset: usize) -> OperationRecord {
     }
 }
 
+/// Adds new record and deletes from the start every 6 records`
+///
+/// # Illustration for head of log
+/// [Add 1], [Del 1], [Add 2], [Del 2], [Add 3], [Del 3], [Add 4] ...
+pub fn add_delete_net_zero_generator(offset: usize) -> OperationRecord {
+    if offset % 2 == 1 {
+        OperationRecord {
+            id: int_as_id(offset / 2),
+            embedding: None,
+            encoding: None,
+            metadata: None,
+            document: None,
+            operation: Operation::Delete,
+        }
+    } else {
+        let int_id = offset / 2;
+        OperationRecord {
+            id: int_as_id(int_id),
+            embedding: Some(random_embedding(TEST_EMBEDDING_DIMENSION)),
+            encoding: None,
+            metadata: Some(modulo_metadata(int_id)),
+            document: Some(modulo_document(int_id)),
+            operation: Operation::Add,
+        }
+    }
+}
+
 #[async_trait]
 pub trait LoadFromGenerator<L: LogGenerator> {
     async fn populate_with_generator(&mut self, log_count: usize, generator: L);

rust/log/src/test.rs

@@ -475,6 +475,17 @@ impl TestSysDb {
                 return Err(FlushCompactionError::CollectionNotFound);
             }
             let collection = collection.unwrap();
+
+            // Check for stale version (optimistic concurrency control)
+            if collection.version > collection_version {
+                return Err(FlushCompactionError::FailedToFlushCompaction(
+                    tonic::Status::failed_precondition(format!(
+                        "Collection version is stale: expected {}, but collection is at version {}",
+                        collection_version, collection.version
+                    )),
+                ));
+            }
+
             let mut collection = collection.clone();
             collection.log_position = log_position;
             new_collection_version = collection_version + 1;