diff --git a/quickwit/quickwit-search/src/leaf.rs b/quickwit/quickwit-search/src/leaf.rs index f8aa60cf98f..8dafb55db14 100644 --- a/quickwit/quickwit-search/src/leaf.rs +++ b/quickwit/quickwit-search/src/leaf.rs @@ -41,7 +41,7 @@ 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}; @@ -1202,8 +1202,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 leaf_request_tasks = Vec::new(); - + let mut leaf_request_futures = 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 @@ -1226,7 +1225,7 @@ pub async fn multi_index_leaf_search( })? .clone(); - let leaf_request_future = tokio::spawn({ + leaf_request_futures.spawn({ let storage_resolver = storage_resolver.clone(); let searcher_context = searcher_context.clone(); let search_request = search_request.clone(); @@ -1241,24 +1240,20 @@ pub async fn multi_index_leaf_search( doc_mapper, aggregation_limits, ) + .in_current_span() .await } - .in_current_span() }); - leaf_request_tasks.push(leaf_request_future); } - let leaf_responses: Vec> = tokio::time::timeout( - searcher_context.searcher_config.request_timeout(), - try_join_all(leaf_request_tasks), - ) - .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 { - match result { - Ok(result) => { - incremental_merge_collector.add_result(result)?; + while let Some(leaf_response_join_result) = leaf_request_futures.join_next().await { + // abort the search on join errors + let leaf_response_result = leaf_response_join_result?; + match leaf_response_result { + Ok(leaf_response) => { + incremental_merge_collector.add_result(leaf_response)?; } Err(err) => { incremental_merge_collector.add_failed_split(SplitSearchError { @@ -1349,9 +1344,6 @@ pub async fn single_doc_mapping_leaf_search( let split_filter = Arc::new(RwLock::new(split_filter)); - let mut leaf_search_single_split_join_handles: Vec<(String, tokio::task::JoinHandle<()>)> = - Vec::with_capacity(split_with_req.len()); - let merge_collector = make_merge_collector(&request, aggregations_limits.clone())?; let incremental_merge_collector = IncrementalCollector::new(merge_collector); let incremental_merge_collector = Arc::new(Mutex::new(incremental_merge_collector)); @@ -1379,6 +1371,8 @@ pub async fn single_doc_mapping_leaf_search( split_filter: split_filter.clone(), }); + let mut split_search_futures = JoinSet::new(); + let mut task_id_to_split_id_map = HashMap::with_capacity(split_with_req.len()); for ((split, search_request), permit_fut) in split_with_req.into_iter().zip(permit_futures.into_iter()) { @@ -1394,41 +1388,39 @@ pub async fn single_doc_mapping_leaf_search( leaf_search_state_guard.set_state(SplitSearchState::PrunedBeforeWarmup); continue; }; - - leaf_search_single_split_join_handles.push(( - split.split_id.clone(), - tokio::spawn( - leaf_search_single_split_wrapper( - simplified_search_request, - leaf_search_context.clone(), - index_storage.clone(), - split, - leaf_split_search_permit, - aggregations_limits.clone(), - ) - .in_current_span(), - ), - )); + let split_id = split.split_id.clone(); + let handle = split_search_futures.spawn( + leaf_search_single_split_wrapper( + simplified_search_request, + leaf_search_context.clone(), + index_storage.clone(), + split, + leaf_split_search_permit, + aggregations_limits.clone(), + ) + .in_current_span(), + ); + task_id_to_split_id_map.insert(handle.id(), split_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(); - // There is no need to use `join_all`, as these are spawned tasks. - for (split, leaf_search_join_handle) in leaf_search_single_split_join_handles { + while let Some(leaf_search_join_result) = split_search_futures.join_next().await { // splits that did not panic were already added to the collector - if let Err(join_error) = leaf_search_join_handle.await { + if let Err(join_error) = leaf_search_join_result { if join_error.is_cancelled() { // An explicit task cancellation is not an error. continue; } + let split_id = task_id_to_split_id_map.get(&join_error.id()).unwrap(); if join_error.is_panic() { - error!(split=%split, "leaf search task panicked"); + error!(split=%split_id, "leaf search task panicked"); } else { - error!(split=%split, "please report: leaf search was not cancelled, and could not extract panic. this should never happen"); + error!(split=%split_id, "please report: leaf search was not cancelled, and could not extract panic. this should never happen"); } - split_search_join_errors.push((split, join_error)); + split_search_join_errors.push((split_id.clone(), join_error)); } } diff --git a/quickwit/quickwit-search/src/service.rs b/quickwit/quickwit-search/src/service.rs index 88dba389c83..3a8088cd61a 100644 --- a/quickwit/quickwit-search/src/service.rs +++ b/quickwit/quickwit-search/src/service.rs @@ -186,7 +186,7 @@ impl SearchService for SearchServiceImpl { .map(|req| req.split_offsets.len()) .sum::(); - LeafSearchMetricsFuture { + let tracked_future = LeafSearchMetricsFuture { tracked: multi_index_leaf_search( self.searcher_context.clone(), leaf_search_request, @@ -195,8 +195,9 @@ impl SearchService for SearchServiceImpl { start: Instant::now(), targeted_splits: num_splits, status: None, - } - .await + }; + let timeout = self.searcher_context.searcher_config.request_timeout(); + tokio::time::timeout(timeout, tracked_future).await? } async fn fetch_docs(