Remove vector index flusher in favor of factory methods for SearchFlusher (#26359)

GitOrigin-RevId: eb3f6568ea5280ad6579bf8e61235d4eb22c75be

Author: sjudd

crates/database/src/index_workers/search_flusher.rs

@@ -135,7 +135,7 @@ pub struct SearchIndexLimits {
 }
 
 impl<RT: Runtime, T: SearchIndexConfigParser + 'static> SearchFlusher<RT, T> {
-    pub fn new(
+    pub(crate) fn new(
         runtime: RT,
         database: Database<RT>,
         reader: Arc<dyn PersistenceReader>,
@@ -170,6 +170,10 @@ impl<RT: Runtime, T: SearchIndexConfigParser + 'static> SearchFlusher<RT, T> {
         }
     }
 
+    /// Run one step of the flusher's main loop.
+    ///
+    /// Returns a map of IndexName to number of documents indexed for each
+    /// index that was built.
     pub async fn step(&mut self) -> anyhow::Result<(BTreeMap<TabletIndexName, u64>, Token)> {
         let mut metrics = BTreeMap::new();
 
@@ -239,7 +243,7 @@ impl<RT: Runtime, T: SearchIndexConfigParser + 'static> SearchFlusher<RT, T> {
     }
 
     /// Compute the set of indexes that need to be backfilled.
-    pub async fn needs_backfill(&self) -> anyhow::Result<(Vec<IndexBuild<T::IndexType>>, Token)> {
+    async fn needs_backfill(&self) -> anyhow::Result<(Vec<IndexBuild<T::IndexType>>, Token)> {
         let mut to_build = vec![];
 
         let mut tx = self.database.begin(Identity::system()).await?;
@@ -313,7 +317,7 @@ impl<RT: Runtime, T: SearchIndexConfigParser + 'static> SearchFlusher<RT, T> {
         Ok((to_build, tx.into_token()?))
     }
 
-    pub async fn build_multipart_segment(
+    async fn build_multipart_segment(
         &self,
         job: &IndexBuild<T::IndexType>,
         build_index_args: <T::IndexType as SearchIndex>::BuildIndexArgs,
@@ -590,13 +594,13 @@ impl<RT: Runtime, T: SearchIndexConfigParser + 'static> SearchFlusher<RT, T> {
     }
 }
 
-pub struct IndexBuild<T: SearchIndex> {
-    pub index_name: TabletIndexName,
-    pub index_id: IndexId,
-    pub by_id: IndexId,
-    pub metadata_id: ResolvedDocumentId,
-    pub index_config: SearchIndexConfig<T>,
-    pub build_reason: BuildReason,
+pub(crate) struct IndexBuild<T: SearchIndex> {
+    pub(crate) index_name: TabletIndexName,
+    pub(crate) index_id: IndexId,
+    pub(crate) by_id: IndexId,
+    pub(crate) metadata_id: ResolvedDocumentId,
+    pub(crate) index_config: SearchIndexConfig<T>,
+    pub(crate) build_reason: BuildReason,
 }
 
 #[derive(Debug)]

crates/database/src/index_workers/search_worker.rs

@@ -36,8 +36,14 @@ use crate::{
         writer::SearchIndexMetadataWriter,
     },
     metrics::log_worker_starting,
-    text_index_worker::flusher2::TextIndexFlusher2,
-    vector_index_worker::compactor::CompactionConfig,
+    text_index_worker::flusher2::{
+        new_text_flusher,
+        TextIndexFlusher2,
+    },
+    vector_index_worker::{
+        compactor::CompactionConfig,
+        flusher::new_vector_flusher,
+    },
     Database,
     TextIndexFlusher,
     VectorIndexCompactor,
@@ -86,7 +92,7 @@ impl<RT: Runtime> SearchIndexWorker<RT> {
             database.clone(),
             // Wait a bit since vector needs time to bootstrap. Makes startup logs a bit cleaner.
             Duration::from_secs(5),
-            SearchIndexWorker::VectorFlusher(VectorIndexFlusher::new(
+            SearchIndexWorker::VectorFlusher(new_vector_flusher(
                 runtime.clone(),
                 database.clone(),
                 reader.clone(),
@@ -108,7 +114,7 @@ impl<RT: Runtime> SearchIndexWorker<RT> {
             )),
         );
         let text_flusher = if *BUILD_MULTI_SEGMENT_TEXT_INDEXES {
-            SearchIndexWorker::TextFlusher2(TextIndexFlusher2::new(
+            SearchIndexWorker::TextFlusher2(new_text_flusher(
                 runtime.clone(),
                 database.clone(),
                 reader,

crates/database/src/search_and_vector_bootstrap.rs

@@ -608,7 +608,7 @@ mod tests {
             DbFixtures,
             DbFixturesArgs,
         },
-        vector_index_worker::flusher::VectorIndexFlusher,
+        vector_index_worker::flusher::backfill_vector_indexes,
         Database,
         IndexModel,
         SystemMetadataModel,
@@ -632,7 +632,13 @@ mod tests {
 
         let db = reopen_db(&rt, &fixtures).await?;
         add_vector(&db, &index_metadata, [1f32, 2f32]).await?;
-        backfill_vector_indexes(&rt, &db, fixtures.tp.reader(), 0, fixtures.search_storage).await?;
+        backfill_vector_indexes(
+            rt.clone(),
+            db.clone(),
+            fixtures.tp.reader(),
+            fixtures.search_storage,
+        )
+        .await?;
 
         // This is a bit of a hack, backfilling with zero size forces all indexes to be
         // written to disk, which causes our boostrapping process to skip our
@@ -832,7 +838,7 @@ mod tests {
             .add_application_index(index_metadata.clone())
             .await?;
         db.commit(tx).await?;
-        backfill_vector_indexes(rt, db, reader, 1000, storage.clone()).await?;
+        backfill_vector_indexes(rt.clone(), db.clone(), reader, storage.clone()).await?;
         let mut tx = db.begin_system().await?;
         let resolved_index = IndexModel::new(&mut tx)
             .pending_index_metadata(TableNamespace::Global, &index_metadata.name)?
@@ -945,24 +951,6 @@ mod tests {
         Ok(metadata)
     }
 
-    async fn backfill_vector_indexes(
-        rt: &TestRuntime,
-        db: &Database<TestRuntime>,
-        reader: Arc<dyn PersistenceReader>,
-        index_size_soft_limit: usize,
-        storage: Arc<dyn Storage>,
-    ) -> anyhow::Result<()> {
-        VectorIndexFlusher::backfill_all_in_test(
-            rt.clone(),
-            db.clone(),
-            reader,
-            storage,
-            index_size_soft_limit,
-        )
-        .await?;
-        Ok(())
-    }
-
     #[convex_macro::test_runtime]
     async fn test_load_snapshot_without_fast_forward(rt: TestRuntime) -> anyhow::Result<()> {
         let db_fixtures = DbFixtures::new(&rt).await?;

crates/database/src/tests/vector_test_utils.rs

@@ -91,7 +91,11 @@ use crate::{
     test_helpers::DbFixturesArgs,
     vector_index_worker::{
         compactor::CompactionConfig,
-        flusher::VectorIndexFlusher,
+        flusher::{
+            backfill_vector_indexes,
+            new_vector_flusher_for_tests,
+            VectorIndexFlusher,
+        },
     },
     Database,
     IndexModel,
@@ -151,7 +155,7 @@ impl VectorFixtures {
     }
 
     pub async fn backfill(&self) -> anyhow::Result<()> {
-        backfill(
+        backfill_vector_indexes(
             self.rt.clone(),
             self.db.clone(),
             self.reader.clone(),
@@ -256,7 +260,7 @@ impl VectorFixtures {
 
     pub async fn run_compaction_during_flush(&self) -> anyhow::Result<()> {
         let (mut pause, pause_client) = PauseController::new([FLUSH_RUNNING_LABEL]);
-        let mut flusher = VectorIndexFlusher::new_for_tests(
+        let mut flusher = new_vector_flusher_for_tests(
             self.rt.clone(),
             self.db.clone(),
             self.reader.clone(),
@@ -284,7 +288,7 @@ impl VectorFixtures {
         &self,
         full_scan_threshold_kb: usize,
     ) -> anyhow::Result<VectorIndexFlusher<TestRuntime>> {
-        Ok(VectorIndexFlusher::new_for_tests(
+        Ok(new_vector_flusher_for_tests(
             self.rt.clone(),
             self.db.clone(),
             self.reader.clone(),
@@ -301,7 +305,7 @@ impl VectorFixtures {
         &self,
         incremental_part_threshold: usize,
     ) -> anyhow::Result<VectorIndexFlusher<TestRuntime>> {
-        Ok(VectorIndexFlusher::new_for_tests(
+        Ok(new_vector_flusher_for_tests(
             self.rt.clone(),
             self.db.clone(),
             self.reader.clone(),
@@ -449,7 +453,7 @@ pub async fn backfilled_vector_index(
     storage: Arc<dyn Storage>,
 ) -> anyhow::Result<IndexData> {
     let index_data = backfilling_vector_index(&db).await?;
-    backfill(rt, db.clone(), reader, storage).await?;
+    backfill_vector_indexes(rt, db.clone(), reader, storage).await?;
 
     Ok(index_data)
 }
@@ -502,18 +506,6 @@ pub async fn backfilled_vector_index_with_doc(
     Ok(result)
 }
 
-// Avoid using index_test_utils.backfill because vector indexes use storage so
-// we need to pass around a consistent instance of storage but index_test_utils
-// creates a new one each time.
-pub async fn backfill<RT: Runtime>(
-    rt: RT,
-    db: Database<RT>,
-    reader: Arc<dyn PersistenceReader>,
-    storage: Arc<dyn Storage>,
-) -> anyhow::Result<()> {
-    VectorIndexFlusher::backfill_all_in_test(rt.clone(), db.clone(), reader, storage, 10).await
-}
-
 pub(crate) async fn assert_backfilled(
     database: &Database<TestRuntime>,
     namespace: TableNamespace,
@@ -651,7 +643,7 @@ impl<RT: Runtime> VectorSearcher for DeleteOnCompactSearchlight<RT> {
             .delete(self.to_delete.into())
             .await?;
         self.db.commit(tx).await?;
-        backfill(
+        backfill_vector_indexes(
             self.rt.clone(),
             self.db.clone(),
             self.reader.clone(),

crates/database/src/tests/vector_tests.rs

@@ -86,7 +86,10 @@ use crate::{
         IndexData,
         VectorFixtures,
     },
-    vector_index_worker::flusher::VectorIndexFlusher,
+    vector_index_worker::flusher::{
+        backfill_vector_indexes,
+        new_vector_flusher_for_tests,
+    },
     Database,
     IndexModel,
     TableModel,
@@ -213,12 +216,11 @@ impl<RT: Runtime> Scenario<RT> {
     }
 
     async fn backfill(&self) -> anyhow::Result<()> {
-        VectorIndexFlusher::backfill_all_in_test(
+        backfill_vector_indexes(
             self.rt.clone(),
             self.database.clone(),
             self.reader.clone(),
             self.search_storage.clone(),
-            0,
         )
         .await?;
         Ok(())
@@ -759,7 +761,7 @@ async fn test_index_backfill_is_incremental(rt: TestRuntime) -> anyhow::Result<(
     scenario.add_vector_index(false).await?;
 
     // Create flusher
-    let mut flusher = VectorIndexFlusher::new_for_tests(
+    let mut flusher = new_vector_flusher_for_tests(
         rt.clone(),
         scenario.database.clone(),
         scenario.reader.clone(),
@@ -844,7 +846,7 @@ async fn test_incremental_backfill_with_compaction(rt: TestRuntime) -> anyhow::R
     scenario.add_vector_index(false).await?;
 
     // Create flusher
-    let mut flusher = VectorIndexFlusher::new_for_tests(
+    let mut flusher = new_vector_flusher_for_tests(
         rt.clone(),
         scenario.database.clone(),
         scenario.reader.clone(),

crates/database/src/text_index_worker/flusher2.rs

@@ -1,15 +1,11 @@
-use std::{
-    collections::BTreeMap,
-    sync::Arc,
-};
+use std::sync::Arc;
 
 #[cfg(any(test, feature = "testing"))]
 use common::pause::PauseClient;
 use common::{
     knobs::SEARCH_INDEX_SIZE_SOFT_LIMIT,
     persistence::PersistenceReader,
     runtime::Runtime,
-    types::TabletIndexName,
 };
 use search::{
     metrics::SearchType,
@@ -31,13 +27,8 @@ use crate::{
         TextSearchIndex,
     },
     Database,
-    Token,
 };
 
-pub struct TextIndexFlusher2<RT: Runtime> {
-    flusher: SearchFlusher<RT, TextIndexConfigParser>,
-}
-
 pub(crate) struct FlusherBuilder<RT: Runtime> {
     runtime: RT,
     database: Database<RT>,
@@ -101,7 +92,7 @@ impl<RT: Runtime> FlusherBuilder<RT> {
             self.storage.clone(),
             SearchType::Text,
         );
-        let flusher = SearchFlusher::new(
+        SearchFlusher::new(
             self.runtime,
             self.database,
             self.reader,
@@ -115,30 +106,20 @@ impl<RT: Runtime> FlusherBuilder<RT> {
             },
             #[cfg(any(test, feature = "testing"))]
             self.pause_client,
-        );
-
-        TextIndexFlusher2 { flusher }
+        )
     }
 }
 
-impl<RT: Runtime> TextIndexFlusher2<RT> {
-    pub(crate) fn new(
-        runtime: RT,
-        database: Database<RT>,
-        reader: Arc<dyn PersistenceReader>,
-        storage: Arc<dyn Storage>,
-        segment_metadata_fetcher: Arc<dyn SegmentTermMetadataFetcher>,
-    ) -> Self {
-        FlusherBuilder::new(runtime, database, reader, storage, segment_metadata_fetcher).build()
-    }
+pub type TextIndexFlusher2<RT> = SearchFlusher<RT, TextIndexConfigParser>;
 
-    /// Run one step of the IndexFlusher's main loop.
-    ///
-    /// Returns a map of IndexName to number of documents indexed for each
-    /// index that was built.
-    pub(crate) async fn step(&mut self) -> anyhow::Result<(BTreeMap<TabletIndexName, u64>, Token)> {
-        self.flusher.step().await
-    }
+pub(crate) fn new_text_flusher<RT: Runtime>(
+    runtime: RT,
+    database: Database<RT>,
+    reader: Arc<dyn PersistenceReader>,
+    storage: Arc<dyn Storage>,
+    segment_metadata_fetcher: Arc<dyn SegmentTermMetadataFetcher>,
+) -> TextIndexFlusher2<RT> {
+    FlusherBuilder::new(runtime, database, reader, storage, segment_metadata_fetcher).build()
 }
 
 #[cfg(test)]