Use TextFlusher2 in tests (#27337)

This PR converts tests to use `TextFlusher2` instead of `TextFlusher`, which enables us to remove the single-segment build path.

GitOrigin-RevId: be6a9788d62a3f289affa904bed5c3c1d1fcb8e2

Author: emmaling27

Cargo.lock

@@ -76,7 +76,7 @@ use crate::{
 /// conflicting writes that may have happened due to concurrent modifications in
 /// the flusher and compactor.
 #[derive(Clone)]
-pub(crate) struct SearchIndexMetadataWriter<RT: Runtime, T: SearchIndex> {
+pub struct SearchIndexMetadataWriter<RT: Runtime, T: SearchIndex> {
     inner: Arc<Mutex<Inner<RT, T>>>,
 }
 

crates/database/src/index_workers/writer.rs

@@ -51,7 +51,7 @@ mod table_iteration;
 pub mod test_helpers;
 #[cfg(test)]
 pub mod tests;
-mod text_index_worker;
+pub mod text_index_worker;
 
 pub use execution_size::FunctionExecutionSize;
 pub use index_worker::IndexWorker;

crates/database/src/lib.rs

@@ -580,7 +580,6 @@ mod tests {
         types::{
             IndexId,
             IndexName,
-            TabletIndexName,
             WriteTimestamp,
         },
     };
@@ -614,13 +613,13 @@ mod tests {
             DbFixtures,
             DbFixturesArgs,
         },
+        text_index_worker::flusher2::new_text_flusher_for_tests,
         vector_index_worker::flusher::backfill_vector_indexes,
         Database,
         IndexModel,
         SystemMetadataModel,
         TableModel,
         TestFacingModel,
-        TextIndexFlusher,
         Transaction,
         UserFacingModel,
     };
@@ -961,8 +960,7 @@ mod tests {
     async fn test_load_snapshot_without_fast_forward(rt: TestRuntime) -> anyhow::Result<()> {
         let db_fixtures = DbFixtures::new(&rt).await?;
         let db = &db_fixtures.db;
-        let (index_id, _) =
-            create_new_search_index(&rt, db, db_fixtures.search_storage.clone()).await?;
+        let (index_id, _) = create_new_search_index(&rt, &db_fixtures).await?;
 
         let mut tx = db.begin_system().await.unwrap();
         add_document(
@@ -991,8 +989,7 @@ mod tests {
     async fn test_load_snapshot_with_fast_forward(rt: TestRuntime) -> anyhow::Result<()> {
         let db_fixtures = DbFixtures::new(&rt).await?;
         let db = &db_fixtures.db;
-        let (index_id, _) =
-            create_new_search_index(&rt, db, db_fixtures.search_storage.clone()).await?;
+        let (index_id, _) = create_new_search_index(&rt, &db_fixtures).await?;
 
         rt.advance_time(Duration::from_secs(10)).await;
 
@@ -1041,8 +1038,7 @@ mod tests {
     ) -> anyhow::Result<()> {
         let db_fixtures = DbFixtures::new(&rt).await?;
         let db = &db_fixtures.db;
-        let (index_id, index_doc) =
-            create_new_search_index(&rt, db, db_fixtures.search_storage.clone()).await?;
+        let (index_id, index_doc) = create_new_search_index(&rt, &db_fixtures).await?;
 
         // We shouldn't ever fast forward across an update in real life, but doing so
         // and verifying we don't read the document is a simple way to verify we
@@ -1083,14 +1079,10 @@ mod tests {
 
     #[convex_macro::test_runtime]
     async fn test_load_fast_forward_ts(rt: TestRuntime) -> anyhow::Result<()> {
-        let DbFixtures {
-            tp,
-            db,
-            search_storage,
-            ..
-        } = DbFixtures::new(&rt).await?;
-        let (index_id, index_doc) =
-            create_new_search_index(&rt, &db, search_storage.clone()).await?;
+        let db_fixtures = DbFixtures::new(&rt).await?;
+        let (index_id, index_doc) = create_new_search_index(&rt, &db_fixtures).await?;
+        let db = db_fixtures.db;
+        let tp = db_fixtures.tp;
         let mut tx = db.begin_system().await?;
         let mut model = IndexWorkerMetadataModel::new(&mut tx);
         let (metadata_id, mut metadata) = model
@@ -1128,7 +1120,7 @@ mod tests {
         let db = &db_fixtures.db;
         let search_storage = db_fixtures.search_storage.clone();
         // Add a search index at t0 to make bootstrapping start at t0
-        create_new_search_index(&rt, db, search_storage.clone()).await?;
+        create_new_search_index(&rt, &db_fixtures).await?;
         // Add a vector index to a table with a vector already in it
         add_vector_by_table(db, table(), [1f32, 2f32]).await?;
         add_and_enable_vector_index(&rt, db, db_fixtures.tp.reader(), search_storage).await?;
@@ -1154,7 +1146,7 @@ mod tests {
         )
         .await?;
         db.commit(tx).await?;
-        create_new_search_index(&rt, db, search_storage).await?;
+        create_new_search_index(&rt, &db_fixtures).await?;
         // Bootstrap
         reopen_db(&rt, &db_fixtures).await?;
         Ok(())
@@ -1173,9 +1165,15 @@ mod tests {
 
     async fn create_new_search_index<RT: Runtime>(
         rt: &RT,
-        db: &Database<RT>,
-        search_storage: Arc<dyn Storage>,
+        db_fixtures: &DbFixtures<RT>,
     ) -> anyhow::Result<(IndexId, ParsedDocument<TabletIndexMetadata>)> {
+        let DbFixtures {
+            tp,
+            db,
+            search_storage,
+            build_index_args,
+            ..
+        } = db_fixtures;
         let table_name: TableName = "test".parse()?;
         let mut tx = db.begin_system().await?;
         TableModel::new(&mut tx)
@@ -1191,21 +1189,14 @@ mod tests {
             .await?;
         db.commit(tx).await?;
 
-        let snapshot = db.latest_snapshot()?;
-        let table_id = snapshot
-            .table_mapping()
-            .namespace(TableNamespace::test_user())
-            .id(&"test".parse()?)?
-            .tablet_id;
-        let index_name = TabletIndexName::new(table_id, "by_text".parse()?)?;
-        TextIndexFlusher::build_index_in_test(
-            index_name.clone(),
-            "test".parse()?,
+        let mut flusher = new_text_flusher_for_tests(
             rt.clone(),
             db.clone(),
+            tp.reader(),
             search_storage.clone(),
-        )
-        .await?;
+            build_index_args.segment_term_metadata_fetcher.clone(),
+        );
+        flusher.step().await?;
 
         let index_name = IndexName::new(table_name, "by_text".parse()?)?;
         let mut tx = db.begin_system().await?;

crates/database/src/search_and_vector_bootstrap.rs

@@ -11,7 +11,10 @@ use common::{
 };
 use events::usage::NoOpUsageEventLogger;
 use search::{
-    searcher::SearcherStub,
+    searcher::{
+        InProcessSearcher,
+        SearcherStub,
+    },
     Searcher,
 };
 use storage::{
@@ -20,6 +23,7 @@ use storage::{
 };
 
 use crate::{
+    text_index_worker::BuildTextIndexArgs,
     Database,
     ShutdownSignal,
     Transaction,
@@ -31,6 +35,7 @@ pub struct DbFixtures<RT: Runtime> {
     pub db: Database<RT>,
     pub searcher: Arc<dyn Searcher>,
     pub search_storage: Arc<dyn Storage>,
+    pub build_index_args: BuildTextIndexArgs,
 }
 
 pub struct DbFixturesArgs {
@@ -89,11 +94,16 @@ impl<RT: Runtime> DbFixtures<RT> {
                 .into_join_future()
                 .await?;
         }
+        let build_index_args = BuildTextIndexArgs {
+            search_storage: search_storage.clone(),
+            segment_term_metadata_fetcher: Arc::new(InProcessSearcher::new(rt.clone()).await?),
+        };
         Ok(Self {
             tp,
             db,
             searcher,
             search_storage,
+            build_index_args,
         })
     }
 }

crates/database/src/test_helpers/db_fixtures.rs

@@ -157,6 +157,7 @@ impl Scenario {
             search_storage,
             searcher,
             tp,
+            build_index_args,
         } = DbFixtures::new_with_args(
             &rt,
             DbFixturesArgs {
@@ -182,11 +183,6 @@ impl Scenario {
             .await?;
         database.commit(tx).await?;
 
-        let build_index_args = BuildTextIndexArgs {
-            search_storage: search_storage.clone(),
-            segment_term_metadata_fetcher: Arc::new(InProcessSearcher::new(rt.clone()).await?),
-        };
-
         let mut self_ = Self {
             rt,
             database,
@@ -210,6 +206,7 @@ impl Scenario {
             searcher,
             search_storage,
             tp,
+            build_index_args,
         } = DbFixtures::new_with_args(
             &self.rt,
             DbFixturesArgs {
@@ -226,6 +223,7 @@ impl Scenario {
         self.searcher = searcher;
         self.search_storage = search_storage;
         self.tp = tp;
+        self.build_index_args = build_index_args;
         Ok(())
     }
 

crates/database/src/tests/randomized_search_tests.rs

@@ -354,18 +354,6 @@ impl<RT: Runtime> TextIndexFlusher<RT> {
         worker.build_one(job).await?;
         Ok(())
     }
-
-    /// Backfills all search indexes that are in a "backfilling" state.
-    #[cfg(any(test, feature = "testing"))]
-    pub async fn backfill_all_in_test(
-        runtime: RT,
-        database: Database<RT>,
-        storage: Arc<dyn Storage>,
-    ) -> anyhow::Result<()> {
-        let mut worker = TextIndexFlusher::new(runtime, database, storage);
-        worker.step().await?;
-        Ok(())
-    }
 }
 
 #[derive(Clone)]

crates/database/src/text_index_worker/flusher.rs

@@ -137,6 +137,36 @@ impl<RT: Runtime> FlusherBuilder<RT> {
 
 pub type TextIndexFlusher2<RT> = SearchFlusher<RT, TextSearchIndex>;
 
+#[allow(unused)]
+#[cfg(any(test, feature = "testing"))]
+pub fn new_text_flusher_for_tests<RT: Runtime>(
+    runtime: RT,
+    database: Database<RT>,
+    reader: Arc<dyn PersistenceReader>,
+    storage: Arc<dyn Storage>,
+    segment_metadata_fetcher: Arc<dyn SegmentTermMetadataFetcher>,
+) -> TextIndexFlusher2<RT> {
+    let writer = SearchIndexMetadataWriter::new(
+        runtime.clone(),
+        database.clone(),
+        reader.clone(),
+        storage.clone(),
+        BuildTextIndexArgs {
+            search_storage: storage.clone(),
+            segment_term_metadata_fetcher: segment_metadata_fetcher.clone(),
+        },
+    );
+    FlusherBuilder::new(
+        runtime,
+        database,
+        reader,
+        storage,
+        segment_metadata_fetcher,
+        writer,
+    )
+    .build()
+}
+
 pub(crate) fn new_text_flusher<RT: Runtime>(
     runtime: RT,
     database: Database<RT>,

crates/database/src/text_index_worker/flusher2.rs

@@ -56,13 +56,13 @@ use database::{
         DbFixtures,
         DbFixturesArgs,
     },
+    text_index_worker::flusher2::backfill_text_indexes,
     vector_index_worker::flusher::backfill_vector_indexes,
     BootstrapComponentsModel,
     Database,
     FollowerRetentionManager,
     IndexModel,
     IndexWorker,
-    TextIndexFlusher,
     Transaction,
 };
 use file_storage::TransactionalFileStorage;
@@ -808,11 +808,15 @@ impl<RT: Runtime, P: Persistence + Clone> UdfTest<RT, P> {
         Ok(())
     }
 
-    pub async fn backfill_search_indexes(&self) -> anyhow::Result<()> {
-        TextIndexFlusher::backfill_all_in_test(
+    pub async fn backfill_text_indexes(&self) -> anyhow::Result<()> {
+        let segment_term_metadata_fetcher =
+            Arc::new(InProcessSearcher::new(self.rt.clone()).await?);
+        backfill_text_indexes(
             self.rt.clone(),
             self.database.clone(),
+            self.persistence.reader(),
             self.search_storage.clone(),
+            segment_term_metadata_fetcher,
         )
         .await?;
         self.enable_backfilled_indexes().await

crates/isolate/src/test_helpers.rs

@@ -46,7 +46,7 @@ async fn add_and_backfill_search_index(
     t: &UdfTest<TestRuntime, TestPersistence>,
 ) -> anyhow::Result<()> {
     add_search_index(t).await?;
-    t.backfill_search_indexes().await?;
+    t.backfill_text_indexes().await?;
 
     Ok(())
 }