Use try_join helpers in incremental_index (#26425)

GitOrigin-RevId: 336d5b1d63ec2a3c5a05165b2e09d22809d33cf3

Author: sjudd

crates/common/src/runtime/mod.rs

@@ -112,10 +112,21 @@ pub async fn try_join_buffered<
         + Send
         + 'static,
 ) -> anyhow::Result<C> {
-    stream::iter(tasks.map(|task| try_join(rt.clone(), name, task)))
-        .buffered(JOIN_BUFFER_SIZE)
-        .try_collect()
-        .await
+    assert_send(
+        stream::iter(tasks.map(|task| assert_send(try_join(rt.clone(), name, assert_send(task)))))
+            .buffered(JOIN_BUFFER_SIZE)
+            .try_collect(),
+    )
+    .await
+}
+
+// Work around "higher-ranked lifetime errors" due to the borrow checker's
+// inability (bug) to determine that some futures are in fact send.  See
+// https://github.com/rust-lang/rust/issues/102211#issuecomment-1367900125
+fn assert_send<'a, T>(
+    fut: impl 'a + Send + Future<Output = T>,
+) -> impl 'a + Send + Future<Output = T> {
+    fut
 }
 
 pub async fn try_join_buffer_unordered<
@@ -129,10 +140,12 @@ pub async fn try_join_buffer_unordered<
         + Send
         + 'static,
 ) -> anyhow::Result<C> {
-    stream::iter(tasks.map(|task| try_join(rt.clone(), name, task)))
-        .buffer_unordered(JOIN_BUFFER_SIZE)
-        .try_collect()
-        .await
+    assert_send(
+        stream::iter(tasks.map(|task| try_join(rt.clone(), name, task)))
+            .buffer_unordered(JOIN_BUFFER_SIZE)
+            .try_collect(),
+    )
+    .await
 }
 
 pub async fn try_join<RT: Runtime, T: Send + 'static>(

crates/database/src/index_workers/index_meta.rs

@@ -103,7 +103,8 @@ pub trait SearchIndex: Clone + Debug {
 
     fn estimate_document_size(schema: &Self::Schema, doc: &ResolvedDocument) -> u64;
 
-    async fn build_disk_index(
+    async fn build_disk_index<RT: Runtime>(
+        rt: &RT,
         schema: &Self::Schema,
         index_path: &PathBuf,
         documents: DocumentStream<'_>,

crates/database/src/index_workers/search_flusher.rs

@@ -564,6 +564,7 @@ impl<RT: Runtime, T: SearchIndexConfigParser + 'static> SearchFlusher<RT, T> {
             params.database.retention_validator(),
         );
         let new_segment = T::IndexType::build_disk_index(
+            &params.runtime,
             &qdrant_schema,
             &index_path,
             documents,

crates/database/src/text_index_worker/text_meta.rs

@@ -208,7 +208,8 @@ impl SearchIndex for TextSearchIndex {
         schema.estimate_size(doc)
     }
 
-    async fn build_disk_index(
+    async fn build_disk_index<RT: Runtime>(
+        rt: &RT,
         schema: &Self::Schema,
         index_path: &PathBuf,
         documents: DocumentStream<'_>,
@@ -222,6 +223,7 @@ impl SearchIndex for TextSearchIndex {
         let revision_stream = Box::pin(stream_revision_pairs(documents, &reader));
 
         build_new_segment(
+            rt,
             revision_stream,
             schema.clone(),
             index_path,

crates/database/src/vector_index_worker/vector_meta.rs

@@ -216,7 +216,8 @@ impl SearchIndex for VectorSearchIndex {
         schema.estimate_vector_size() as u64
     }
 
-    async fn build_disk_index(
+    async fn build_disk_index<RT: Runtime>(
+        _rt: &RT,
         schema: &Self::Schema,
         index_path: &PathBuf,
         documents: DocumentStream<'_>,

crates/search/src/incremental_index.rs

@@ -18,15 +18,13 @@ use common::{
     bounded_thread_pool::BoundedThreadPool,
     id_tracker::StaticIdTracker,
     persistence::DocumentRevisionStream,
-    runtime::Runtime,
+    runtime::{
+        try_join_buffer_unordered,
+        Runtime,
+    },
     types::ObjectKey,
 };
-use futures::{
-    stream,
-    FutureExt,
-    StreamExt,
-    TryStreamExt,
-};
+use futures::TryStreamExt;
 use storage::{
     Storage,
     StorageExt,
@@ -303,7 +301,8 @@ impl PreviousTextSegments {
 /// Note the descending order requirement can be relaxed if the caller can
 /// guarantee that no deletes will be present in the stream. A caller can do so
 /// when providing this function with a stream from table iterator for example.
-pub async fn build_new_segment(
+pub async fn build_new_segment<RT: Runtime>(
+    rt: &RT,
     revision_stream: DocumentRevisionStream<'_>,
     tantivy_schema: TantivySearchIndexSchema,
     dir: &Path,
@@ -405,6 +404,7 @@ pub async fn build_new_segment(
     );
 
     let segments_term_metadata = get_all_segment_term_metadata(
+        rt,
         search_storage,
         segment_term_metadata_fetcher,
         term_deletes_by_segment,
@@ -447,7 +447,8 @@ fn get_size(path: &PathBuf) -> anyhow::Result<u64> {
     std::fs::read_dir(path)?.try_fold(0, |acc, curr| Ok(acc + get_size(&curr?.path())?))
 }
 
-async fn get_all_segment_term_metadata(
+async fn get_all_segment_term_metadata<RT: Runtime>(
+    rt: &RT,
     storage: Arc<dyn Storage>,
     segment_term_metadata_fetcher: Arc<dyn SegmentTermMetadataFetcher>,
     term_deletes_by_segment: BTreeMap<ObjectKey, TermDeletionsByField>,
@@ -470,12 +471,9 @@ async fn get_all_segment_term_metadata(
             }
         },
     );
-    // TODO: Use `try_join_buffer_unordered` helper here. We couldn't figure out the
-    // higher-ranked lifetime error.
-    let segments_term_metadata: Vec<_> = stream::iter(segment_term_metadata_futs)
-        .buffer_unordered(10)
-        .try_collect::<Vec<_>>()
-        .await?;
+    let segments_term_metadata: Vec<_> =
+        try_join_buffer_unordered(rt.clone(), "text_term_metadata", segment_term_metadata_futs)
+            .await?;
     Ok(segments_term_metadata)
 }
 
@@ -545,15 +543,21 @@ pub async fn fetch_compact_and_upload_text_segment<RT: Runtime>(
     segments: Vec<FragmentedTextStorageKeys>,
 ) -> anyhow::Result<FragmentedTextSegment> {
     let _storage = storage.clone();
-    let opened_segments: Vec<_> = stream::iter(segments)
-        .map(move |segment| fetch_text_segment(cache.clone(), storage.clone(), segment).boxed())
-        .buffer_unordered(20)
-        .and_then(|paths: TextSegmentPaths| {
+    let opened_segments = try_join_buffer_unordered(
+        rt.clone(),
+        "text_segment_merge",
+        segments.into_iter().map(move |segment| {
             let pool = blocking_thread_pool.clone();
-            async move { pool.execute(|| open_text_segment_for_merge(paths)).await? }
-        })
-        .try_collect()
-        .await?;
+            let storage = storage.clone();
+            let cache = cache.clone();
+            let segment = segment.clone();
+            async move {
+                let paths = fetch_text_segment(cache, storage, segment).await?;
+                pool.execute(|| open_text_segment_for_merge(paths)).await?
+            }
+        }),
+    )
+    .await?;
 
     let dir = TempDir::new()?;
     let new_segment = merge_segments(opened_segments, dir.path()).await?;

crates/search/src/searcher/searcher.rs

@@ -1734,9 +1734,10 @@ mod tests {
         let revision_stream = futures::stream::iter(revisions).boxed();
         let mut previous_segments = PreviousTextSegments::default();
         let storage: Arc<dyn Storage> = Arc::new(LocalDirStorage::new(rt.clone())?);
-        let segment_term_metadata_fetcher = Arc::new(InProcessSearcher::new(rt).await?);
+        let segment_term_metadata_fetcher = Arc::new(InProcessSearcher::new(rt.clone()).await?);
 
         let new_segment = build_new_segment(
+            &rt,
             revision_stream,
             schema.clone(),
             test_dir.path(),
@@ -1929,9 +1930,10 @@ mod tests {
         // tests because they don't use the same directory that the indexes are stored
         // in, which means we must use empty PreviousTextSegments in these tests.
         let storage: Arc<dyn Storage> = Arc::new(LocalDirStorage::new(rt.clone())?);
-        let segment_term_metadata_fetcher = Arc::new(InProcessSearcher::new(rt).await?);
+        let segment_term_metadata_fetcher = Arc::new(InProcessSearcher::new(rt.clone()).await?);
         let mut previous_segments = PreviousTextSegments::default();
         let new_segment = build_new_segment(
+            &rt,
             revision_stream,
             schema.clone(),
             index_dir,