[ENH] Add a safety cutoff to the rust log service.

If a collection on the rust-based log has more than a
configuration-configurable number of records, shut it down hard no
matter what.  This happens when compaction gets backlogged.  By default,
no more than 1_000_000 can be on the log uncompacted.  This is just a
placeholder that we can change.

Author: rescrv

rust/log-service/src/lib.rs

@@ -2,8 +2,9 @@
 
 use std::cmp::Ordering;
 use std::collections::HashMap;
+use std::collections::HashSet;
 use std::future::Future;
-use std::sync::Arc;
+use std::sync::{Arc, Mutex};
 use std::time::{Duration, Instant, SystemTime};
 
 use bytes::Bytes;
@@ -264,6 +265,7 @@ pub struct Rollup {
     pub advance_to: LogPosition,
     pub reinsert: Vec<DirtyMarker>,
     pub compactable: Vec<CollectionInfo>,
+    pub backpressure: Vec<CollectionUuid>,
 }
 
 //////////////////////////////////////// RollupPerCollection ///////////////////////////////////////
@@ -361,6 +363,7 @@ impl DirtyMarker {
         retrieve_cursor: impl Fn(Arc<Storage>, CollectionUuid) -> F2,
         markers: &[(LogPosition, DirtyMarker)],
         record_count_threshold: u64,
+        record_count_backpressure: u64,
         reinsert_threshold: u64,
         timeout_us: u64,
         metrics: &Metrics,
@@ -424,7 +427,6 @@ impl DirtyMarker {
             })
             .map(
                 |(collection_id, storage, retrieve_manifest, retrieve_cursor)| async move {
-                    // We play a funny game of Ok(Ok(_)) to force try_join_all to not short-circuit.
                     let cursor = (*retrieve_cursor)(Arc::clone(&storage), collection_id);
                     let manifest = (*retrieve_manifest)(Arc::clone(&storage), collection_id);
                     let (cursor, manifest) = futures::future::join(cursor, manifest).await;
@@ -454,6 +456,7 @@ impl DirtyMarker {
             .flat_map(Result::ok)
             .collect::<HashMap<_, _>>();
         let mut uncompacted = 0u64;
+        let mut backpressure = vec![];
         let compactable = compactable
             .into_iter()
             .filter_map(|collection_id| {
@@ -494,6 +497,9 @@ impl DirtyMarker {
                             manifest.maximum_log_position(),
                         );
                         uncompacted += maximum_log_position - cursor.position;
+                        if maximum_log_position - cursor.position >= record_count_backpressure {
+                            backpressure.push(*collection_id);
+                        }
                         if maximum_log_position - cursor.position >= record_count_threshold {
                             Some(CollectionInfo {
                                 collection_id: collection_id.to_string(),
@@ -542,6 +548,7 @@ impl DirtyMarker {
             advance_to,
             reinsert,
             compactable,
+            backpressure,
         }))
     }
 
@@ -631,10 +638,32 @@ pub struct LogServer {
     open_logs: Arc<StateHashTable<LogKey, LogStub>>,
     dirty_log: Arc<LogWriter>,
     compacting: tokio::sync::Mutex<()>,
+    backpressure: Mutex<Arc<HashSet<CollectionUuid>>>,
     cache: Option<Box<dyn chroma_cache::PersistentCache<String, CachedParquetFragment>>>,
     metrics: Metrics,
 }
 
+impl LogServer {
+    fn set_backpressure(&self, to_pressure: &[CollectionUuid]) {
+        let mut new_backpressure = Arc::new(HashSet::from_iter(to_pressure.iter().cloned()));
+        // SAFETY(rescrv): Mutex poisoning.
+        let mut backpressure = self.backpressure.lock().unwrap();
+        std::mem::swap(&mut *backpressure, &mut new_backpressure);
+    }
+
+    fn check_for_backpressure(&self, collection_id: CollectionUuid) -> Result<(), Status> {
+        let backpressure = {
+            // SAFETY(rescrv): Mutex poisoning.
+            let backpressure = self.backpressure.lock().unwrap();
+            Arc::clone(&backpressure)
+        };
+        if backpressure.contains(&collection_id) {
+            return Err(Status::resource_exhausted("log needs compaction; too full"));
+        }
+        Ok(())
+    }
+}
+
 #[async_trait::async_trait]
 impl LogService for LogServer {
     async fn push_logs(
@@ -652,6 +681,7 @@ impl LogService for LogServer {
         if push_logs.records.is_empty() {
             return Err(Status::invalid_argument("Too few records"));
         }
+        self.check_for_backpressure(collection_id)?;
         let span = tracing::info_span!("push_logs");
 
         async move {
@@ -832,6 +862,7 @@ impl LogService for LogServer {
         let source_collection_id = Uuid::parse_str(&request.source_collection_id)
             .map(CollectionUuid)
             .map_err(|_| Status::invalid_argument("Failed to parse collection id"))?;
+        self.check_for_backpressure(source_collection_id)?;
         let target_collection_id = Uuid::parse_str(&request.target_collection_id)
             .map(CollectionUuid)
             .map_err(|_| Status::invalid_argument("Failed to parse collection id"))?;
@@ -1000,6 +1031,7 @@ impl LogService for LogServer {
                 self.config.record_count_threshold,
                 request.min_compaction_size,
             ),
+            self.config.num_records_before_backpressure,
             self.config.reinsert_threshold,
             self.config.timeout_us,
             &self.metrics,
@@ -1020,6 +1052,7 @@ impl LogService for LogServer {
                     .map_err(|err| Status::unavailable(err.to_string()))
             })
             .collect::<Result<Vec<_>, _>>()?;
+        self.set_backpressure(&rollup.backpressure);
         if rollup.advance_to < cursor.position {
             tracing::error!(
                 "advance_to went back in time: {:?} -> {:?}",
@@ -1409,6 +1442,8 @@ pub struct LogServerConfig {
     pub cache: Option<CacheConfig>,
     #[serde(default = "LogServerConfig::default_record_count_threshold")]
     pub record_count_threshold: u64,
+    #[serde(default = "LogServerConfig::default_num_records_before_backpressure")]
+    pub num_records_before_backpressure: u64,
     #[serde(default = "LogServerConfig::default_reinsert_threshold")]
     pub reinsert_threshold: u64,
     #[serde(default = "LogServerConfig::default_timeout_us")]
@@ -1421,6 +1456,11 @@ impl LogServerConfig {
         100
     }
 
+    /// one million records on the log.
+    fn default_num_records_before_backpressure() -> u64 {
+        1_000_000
+    }
+
     /// force compaction if a candidate comes up ten times.
     fn default_reinsert_threshold() -> u64 {
         10
@@ -1442,6 +1482,7 @@ impl Default for LogServerConfig {
             reader: LogReaderOptions::default(),
             cache: None,
             record_count_threshold: Self::default_record_count_threshold(),
+            num_records_before_backpressure: Self::default_num_records_before_backpressure(),
             reinsert_threshold: Self::default_reinsert_threshold(),
             timeout_us: Self::default_timeout_us(),
         }
@@ -1483,12 +1524,14 @@ impl Configurable<LogServerConfig> for LogServer {
         let dirty_log = Arc::new(dirty_log);
         let compacting = tokio::sync::Mutex::new(());
         let metrics = Metrics::new(opentelemetry::global::meter("chroma"));
+        let backpressure = Mutex::new(Arc::new(HashSet::default()));
         Ok(Self {
             config: config.clone(),
             open_logs: Arc::new(StateHashTable::default()),
             storage,
             dirty_log,
             compacting,
+            backpressure,
             cache,
             metrics,
         })
@@ -1592,6 +1635,7 @@ mod tests {
             &markers,
             1,
             1,
+            1,
             86_400_000_000,
             &Metrics::new(opentelemetry::global::meter("chroma")),
         )
@@ -1654,6 +1698,7 @@ mod tests {
             &markers,
             3,
             1,
+            1,
             86_400_000_000,
             &Metrics::new(opentelemetry::global::meter("chroma")),
         )
@@ -1738,6 +1783,7 @@ mod tests {
             &markers,
             3,
             1,
+            1,
             86_400_000_000,
             &Metrics::new(opentelemetry::global::meter("chroma")),
         )
@@ -1852,6 +1898,7 @@ mod tests {
             &markers,
             3,
             1,
+            1,
             86_400_000_000,
             &Metrics::new(opentelemetry::global::meter("chroma")),
         )
@@ -1864,6 +1911,76 @@ mod tests {
         assert!(rollup.reinsert[0].collection_id() == collection_id_blocking);
     }
 
+    #[tokio::test]
+    async fn dirty_marker_backpressure() {
+        // Test that the dirty marker gives proper backpressure.
+        let storage = chroma_storage::test_storage();
+        let now = SystemTime::now()
+            .duration_since(SystemTime::UNIX_EPOCH)
+            .map_err(|_| wal3::Error::Internal)
+            .unwrap()
+            .as_micros() as u64;
+        let collection_id = CollectionUuid::new();
+        let markers = vec![(
+            LogPosition::from_offset(1),
+            DirtyMarker::MarkDirty {
+                collection_id,
+                log_position: LogPosition::from_offset(1),
+                num_records: 1_000_000,
+                reinsert_count: 0,
+                initial_insertion_epoch_us: now,
+            },
+        )];
+        let rollup = DirtyMarker::rollup(
+            Arc::new(storage),
+            |_, collection_id| async move {
+                if collection_id == collection_id {
+                    Ok(Some(Manifest {
+                        writer: "TODO".to_string(),
+                        acc_bytes: 0,
+                        setsum: Setsum::default(),
+                        snapshots: vec![],
+                        fragments: vec![Fragment {
+                            seq_no: FragmentSeqNo(1),
+                            num_bytes: 0,
+                            path: "TODO".to_string(),
+                            setsum: Setsum::default(),
+                            start: LogPosition::from_offset(1),
+                            limit: LogPosition::from_offset(1_000_001),
+                        }],
+                    }))
+                } else {
+                    unreachable!("we aren't testing this case");
+                }
+            },
+            |_, collection_id| async move {
+                if collection_id == collection_id {
+                    Ok(Some(Witness::default_etag_with_cursor(Cursor {
+                        position: LogPosition::from_offset(1),
+                        epoch_us: 0,
+                        writer: "TODO".to_string(),
+                    })))
+                } else {
+                    unreachable!("we aren't testing this case");
+                }
+            },
+            &markers,
+            1,
+            1,
+            1,
+            86_400_000_000,
+            &Metrics::new(opentelemetry::global::meter("chroma")),
+        )
+        .await
+        .unwrap()
+        .unwrap();
+        assert_eq!(LogPosition::from_offset(2), rollup.advance_to);
+        assert_eq!(1, rollup.compactable.len());
+        assert_eq!(1, rollup.reinsert.len());
+        assert_eq!(1, rollup.backpressure.len());
+        assert_eq!(collection_id, rollup.backpressure[0]);
+    }
+
     #[test]
     fn unsafe_constants() {
         assert!(STABLE_PREFIX.is_valid());