Use JoinSet instead of its ordered variation

Author: rdettai-sk

quickwit/quickwit-search/src/join_vec.rs

@@ -41,11 +41,10 @@ use tantivy::directory::FileSlice;
 use tantivy::fastfield::FastFieldReaders;
 use tantivy::schema::Field;
 use tantivy::{DateTime, Index, ReloadPolicy, Searcher, TantivyError, Term};
-use tokio::task::JoinError;
+use tokio::task::{JoinError, JoinSet};
 use tracing::*;
 
 use crate::collector::{IncrementalCollector, make_collector_for_split, make_merge_collector};
-use crate::join_vec::{JoinVec, TryJoinVec};
 use crate::root::is_metadata_count_request_with_ast;
 use crate::search_permit_provider::{SearchPermit, compute_initial_memory_allocation};
 use crate::service::{SearcherContext, deserialize_doc_mapper};
@@ -1180,7 +1179,7 @@ pub async fn multi_index_leaf_search(
     //
     // It is a little bit tricky how to handle which is now the incremental_merge_collector, one
     // per index, e.g. when to merge results and how to avoid lock contention.
-    let mut try_join_vec = TryJoinVec::new();
+    let mut join_set = JoinSet::new();
     for leaf_search_request_ref in leaf_search_request.leaf_requests.into_iter() {
         let index_uri = quickwit_common::uri::Uri::from_str(
             leaf_search_request
@@ -1203,7 +1202,7 @@ pub async fn multi_index_leaf_search(
             })?
             .clone();
 
-        try_join_vec.spawn({
+        join_set.spawn({
             let storage_resolver = storage_resolver.clone();
             let searcher_context = searcher_context.clone();
             let search_request = search_request.clone();
@@ -1224,11 +1223,10 @@ pub async fn multi_index_leaf_search(
         });
     }
 
-    let leaf_responses: Vec<LeafSearchResponse> = try_join_vec.try_join_all().await?;
     let merge_collector = make_merge_collector(&search_request, aggregation_limits)?;
     let mut incremental_merge_collector = IncrementalCollector::new(merge_collector);
-    for result in leaf_responses {
-        incremental_merge_collector.add_result(result)?;
+    while let Some(result) = join_set.join_next().await {
+        incremental_merge_collector.add_result(result??)?;
     }
 
     crate::search_thread_pool()
@@ -1313,8 +1311,8 @@ pub async fn single_doc_mapping_leaf_search(
         .search_permit_provider
         .get_permits(permit_sizes)
         .await;
-
-    let mut join_vec = JoinVec::new_with_capacity(split_with_req.len());
+    let mut join_set = JoinSet::new();
+    let mut split_with_task_id = Vec::with_capacity(split_with_req.len());
     for ((split, mut request), permit_fut) in
         split_with_req.into_iter().zip(permit_futures.into_iter())
     {
@@ -1326,8 +1324,8 @@ pub async fn single_doc_mapping_leaf_search(
         if !can_be_better && !run_all_splits {
             continue;
         }
-        join_vec.spawn(
-            split.split_id.clone(),
+        let split_id = split.split_id.clone();
+        let handle = join_set.spawn(
             leaf_search_single_split_wrapper(
                 request,
                 searcher_context.clone(),
@@ -1341,20 +1339,25 @@ pub async fn single_doc_mapping_leaf_search(
             )
             .in_current_span(),
         );
+        split_with_task_id.push((split_id, handle.id()));
     }
 
     // TODO we could cancel running splits when !run_all_splits and the running split can no
     // longer give better results after some other split answered.
     let mut split_search_join_errors: Vec<(String, JoinError)> = Vec::new();
 
-    let leaf_search_single_split_join_handles = join_vec.join_all().await;
-    for (split, leaf_search_join_result) in leaf_search_single_split_join_handles {
+    while let Some(leaf_search_join_result) = join_set.join_next().await {
         // splits that did not panic were already added to the collector
         if let Err(join_error) = leaf_search_join_result {
             if join_error.is_cancelled() {
                 // An explicit task cancellation is not an error.
                 continue;
             }
+            let position = split_with_task_id
+                .iter()
+                .position(|(_, task_id)| *task_id == join_error.id())
+                .unwrap();
+            let (split, _) = split_with_task_id.remove(position);
             if join_error.is_panic() {
                 error!(split=%split, "leaf search task panicked");
             } else {

quickwit/quickwit-search/src/leaf.rs

@@ -23,7 +23,6 @@ mod collector;
 mod error;
 mod fetch_docs;
 mod find_trace_ids_collector;
-mod join_vec;
 mod leaf;
 mod leaf_cache;
 mod list_fields;