diff --git a/quickwit/quickwit-common/src/thread_pool.rs b/quickwit/quickwit-common/src/thread_pool.rs index 18201196cf9..8cee9b26ecf 100644 --- a/quickwit/quickwit-common/src/thread_pool.rs +++ b/quickwit/quickwit-common/src/thread_pool.rs @@ -17,7 +17,6 @@ use std::sync::Arc; use futures::{Future, TryFutureExt}; use once_cell::sync::Lazy; -use prometheus::IntGauge; use tokio::sync::oneshot; use tracing::error; @@ -30,8 +29,7 @@ use crate::metrics::{GaugeGuard, IntGaugeVec, OwnedGaugeGuard, new_gauge_vec}; #[derive(Clone)] pub struct ThreadPool { thread_pool: Arc, - ongoing_tasks: IntGauge, - pending_tasks: IntGauge, + name: &'static str, } impl ThreadPool { @@ -47,46 +45,43 @@ impl ThreadPool { let thread_pool = rayon_pool_builder .build() .expect("failed to spawn thread pool"); - let ongoing_tasks = THREAD_POOL_METRICS.ongoing_tasks.with_label_values([name]); - let pending_tasks = THREAD_POOL_METRICS.pending_tasks.with_label_values([name]); ThreadPool { thread_pool: Arc::new(thread_pool), - ongoing_tasks, - pending_tasks, + name, } } + /// Returns the underlying rayon thread pool. + /// + /// Using the underlying Rayon thread pool directly is discouraged, as it + /// will not update the metrics and create an observability blind spot. + /// Unfortunately it is sometimes necessary to provide the raw pool + /// directly to Tantivy. pub fn get_underlying_rayon_thread_pool(&self) -> Arc { self.thread_pool.clone() } - /// Function similar to `tokio::spawn_blocking`. - /// - /// Here are two important differences however: - /// - /// 1) The task runs on a rayon thread pool managed by Quickwit. This pool is specifically used - /// only to run CPU-intensive work and is configured to contain `num_cpus` cores. - /// - /// 2) Before the task is effectively scheduled, we check that the spawner is still interested - /// in its result. - /// - /// It is therefore required to `await` the result of this - /// function to get any work done. - /// - /// This is nice because it makes work that has been scheduled - /// but is not running yet "cancellable". - pub fn run_cpu_intensive( + /// Same as `run_cpu_intensive` but with a caller identifier recorded in the + /// metrics. + pub fn run_cpu_intensive_with_identified_caller( &self, cpu_intensive_fn: F, + caller: &'static str, ) -> impl Future> where F: FnOnce() -> R + Send + 'static, R: Send + 'static, { let span = tracing::Span::current(); - let ongoing_tasks = self.ongoing_tasks.clone(); + let ongoing_tasks = THREAD_POOL_METRICS + .ongoing_tasks + .with_label_values([self.name, caller]); + let pending_tasks = THREAD_POOL_METRICS + .pending_tasks + .with_label_values([self.name, caller]); + let ongoing_tasks = ongoing_tasks.clone(); let mut pending_tasks_guard: OwnedGaugeGuard = - OwnedGaugeGuard::from_gauge(self.pending_tasks.clone()); + OwnedGaugeGuard::from_gauge(pending_tasks.clone()); pending_tasks_guard.add(1i64); let (tx, rx) = oneshot::channel(); self.thread_pool.spawn(move || { @@ -102,6 +97,32 @@ impl ThreadPool { }); rx.map_err(|_| Panicked) } + + /// Function similar to `tokio::spawn_blocking`. + /// + /// Here are two important differences however: + /// + /// 1) The task runs on a rayon thread pool managed by Quickwit. This pool is specifically used + /// only to run CPU-intensive work and is configured to contain `num_cpus` cores. + /// + /// 2) Before the task is effectively scheduled, we check that the spawner is still interested + /// in its result. + /// + /// It is therefore required to `await` the result of this + /// function to get any work done. + /// + /// This is nice because it makes work that has been scheduled + /// but is not running yet "cancellable". + pub fn run_cpu_intensive( + &self, + cpu_intensive_fn: F, + ) -> impl Future> + where + F: FnOnce() -> R + Send + 'static, + R: Send + 'static, + { + self.run_cpu_intensive_with_identified_caller(cpu_intensive_fn, "unknown") + } } /// Run a small (<200ms) CPU-intensive task on a dedicated thread pool with a few threads. @@ -138,8 +159,8 @@ impl fmt::Display for Panicked { impl std::error::Error for Panicked {} struct ThreadPoolMetrics { - ongoing_tasks: IntGaugeVec<1>, - pending_tasks: IntGaugeVec<1>, + ongoing_tasks: IntGaugeVec<2>, + pending_tasks: IntGaugeVec<2>, } impl Default for ThreadPoolMetrics { @@ -150,14 +171,14 @@ impl Default for ThreadPoolMetrics { "number of tasks being currently processed by threads in the thread pool", "thread_pool", &[], - ["pool"], + ["pool", "caller"], ), pending_tasks: new_gauge_vec( "pending_tasks", "number of tasks waiting in the queue before being processed by the thread pool", "thread_pool", &[], - ["pool"], + ["pool", "caller"], ), } } diff --git a/quickwit/quickwit-search/src/leaf.rs b/quickwit/quickwit-search/src/leaf.rs index 0053c33fe99..2cbb330dbdc 100644 --- a/quickwit/quickwit-search/src/leaf.rs +++ b/quickwit/quickwit-search/src/leaf.rs @@ -327,7 +327,7 @@ async fn warm_up_automatons( let mut warm_up_futures = Vec::new(); let cpu_intensive_executor = |task| async { crate::search_thread_pool() - .run_cpu_intensive(task) + .run_cpu_intensive_with_identified_caller(task, "automaton_warmup") .await .map_err(|_| std::io::Error::other("task panicked"))? }; @@ -556,52 +556,50 @@ async fn leaf_search_single_split( let split_clone = split.clone(); let ctx_clone = ctx.clone(); + leaf_search_state_guard.set_state(SplitSearchState::CpuQueue); + let cpu_task = move || { + leaf_search_state_guard.set_state(SplitSearchState::Cpu); + let cpu_start = Instant::now(); + let cpu_thread_pool_wait_microsecs = cpu_start.duration_since(warmup_end); + let _span_guard = span.enter(); + // Our search execution has been scheduled, let's check if we can improve the + // request based on the results of the preceding searches + let Some(simplified_search_request) = + simplify_search_request(search_request, &split_clone, &ctx_clone.split_filter) + else { + leaf_search_state_guard.set_state(SplitSearchState::PrunedAfterWarmup); + return Ok(None); + }; + collector.update_search_param(&simplified_search_request); + let mut leaf_search_response: LeafSearchResponse = + if is_metadata_count_request_with_ast(&query_ast, &simplified_search_request) { + let num_docs = searcher + .num_docs() + .saturating_sub(split_clone.soft_deleted_doc_ids.len() as u64); + get_leaf_resp_from_count(num_docs) + } else if collector.is_count_only() { + let count = query.count(&searcher)? as u64; + get_leaf_resp_from_count(count) + } else { + searcher.search(&query, &collector)? + }; + leaf_search_response.resource_stats = Some(ResourceStats { + cpu_microsecs: cpu_start.elapsed().as_micros() as u64, + short_lived_cache_num_bytes: warmup_size.as_u64(), + split_num_docs, + warmup_microsecs: warmup_duration.as_micros() as u64, + cpu_thread_pool_wait_microsecs: cpu_thread_pool_wait_microsecs.as_micros() as u64, + }); + // splits by outcome are estimated at the (doc mapping) leaf + // response level to account for all early returns, so it is + // left None here + leaf_search_state_guard.set_state(SplitSearchState::Processed); + Result::<_, TantivyError>::Ok(Some((simplified_search_request, leaf_search_response))) + }; let search_request_and_result: Option<(SearchRequest, LeafSearchResponse)> = crate::search_thread_pool() - .run_cpu_intensive(move || { - leaf_search_state_guard.set_state(SplitSearchState::Cpu); - let cpu_start = Instant::now(); - let cpu_thread_pool_wait_microsecs = cpu_start.duration_since(warmup_end); - let _span_guard = span.enter(); - // Our search execution has been scheduled, let's check if we can improve the - // request based on the results of the preceding searches - let Some(simplified_search_request) = - simplify_search_request(search_request, &split_clone, &ctx_clone.split_filter) - else { - leaf_search_state_guard.set_state(SplitSearchState::PrunedAfterWarmup); - return Ok(None); - }; - collector.update_search_param(&simplified_search_request); - let mut leaf_search_response: LeafSearchResponse = - if is_metadata_count_request_with_ast(&query_ast, &simplified_search_request) { - let num_docs = searcher - .num_docs() - .saturating_sub(split_clone.soft_deleted_doc_ids.len() as u64); - get_leaf_resp_from_count(num_docs) - } else if collector.is_count_only() { - let count = query.count(&searcher)? as u64; - get_leaf_resp_from_count(count) - } else { - searcher.search(&query, &collector)? - }; - leaf_search_response.resource_stats = Some(ResourceStats { - cpu_microsecs: cpu_start.elapsed().as_micros() as u64, - short_lived_cache_num_bytes: warmup_size.as_u64(), - split_num_docs, - warmup_microsecs: warmup_duration.as_micros() as u64, - cpu_thread_pool_wait_microsecs: cpu_thread_pool_wait_microsecs.as_micros() - as u64, - }); - // splits by outcome are estimated at the (doc mapping) leaf - // response level to account for all early returns, so it is - // left None here - leaf_search_state_guard.set_state(SplitSearchState::Processed); - Result::<_, TantivyError>::Ok(Some(( - simplified_search_request, - leaf_search_response, - ))) - }) + .run_cpu_intensive_with_identified_caller(cpu_task, "split_search") .await .map_err(|_| { crate::SearchError::Internal(format!("leaf search panicked. split={split_id}")) @@ -1295,7 +1293,10 @@ pub async fn multi_index_leaf_search( } crate::search_thread_pool() - .run_cpu_intensive(|| incremental_merge_collector.finalize().map_err(Into::into)) + .run_cpu_intensive_with_identified_caller( + || incremental_merge_collector.finalize().map_err(Into::into), + "finalize", + ) .instrument(info_span!("incremental_merge_finalize")) .await .context("failed to merge split search responses")? @@ -1477,7 +1478,10 @@ pub async fn single_doc_mapping_leaf_search( let leaf_search_response_reresult: Result, _> = crate::search_thread_pool() - .run_cpu_intensive(|| incremental_merge_collector.finalize()) + .run_cpu_intensive_with_identified_caller( + || incremental_merge_collector.finalize(), + "finalize", + ) .instrument(info_span!("incremental_merge_intermediate")) .await .context("failed to merge split search responses"); diff --git a/quickwit/quickwit-search/src/lib.rs b/quickwit/quickwit-search/src/lib.rs index 8476326bd53..73d9cf05ffe 100644 --- a/quickwit/quickwit-search/src/lib.rs +++ b/quickwit/quickwit-search/src/lib.rs @@ -55,6 +55,7 @@ use quickwit_proto::metastore::{ ListIndexesMetadataRequest, ListSplitsRequest, MetastoreService, MetastoreServiceClient, }; use tantivy::schema::NamedFieldDocument; +use tracing::info; /// Refer to this as `crate::Result`. pub type Result = std::result::Result; @@ -99,7 +100,20 @@ pub type SearcherPool = Pool; fn search_thread_pool() -> &'static ThreadPool { static SEARCH_THREAD_POOL: OnceLock = OnceLock::new(); - SEARCH_THREAD_POOL.get_or_init(|| ThreadPool::new("search", None)) + + SEARCH_THREAD_POOL.get_or_init(|| { + let use_all_cpus = + quickwit_common::get_bool_from_env("QW_SEARCH_THREAD_POOL_USE_ALL_CPUS", false); + let num_threads = if use_all_cpus { + info!("search thread pool configured with default Rayon thread pool size"); + None + } else { + let threads = usize::max(quickwit_common::num_cpus().saturating_sub(1), 1); + info!(threads, "search thread pool configured with one free CPU"); + Some(threads) + }; + ThreadPool::new("search", num_threads) + }) } /// GlobalDocAddress serves as a hit address. diff --git a/quickwit/quickwit-search/src/list_fields.rs b/quickwit/quickwit-search/src/list_fields.rs index b9eec1696b5..44866ad3b71 100644 --- a/quickwit/quickwit-search/src/list_fields.rs +++ b/quickwit/quickwit-search/src/list_fields.rs @@ -341,40 +341,40 @@ pub async fn leaf_list_fields( let mut single_split_list_fields_vec: Vec> = future::try_join_all(single_split_list_fields_futures).await?; - - let fields = search_thread_pool() - .run_cpu_intensive(move || { - for single_split_list_fields in &mut single_split_list_fields_vec { - // This contract is enforced on a different node, etc. so we defensively check that - // the fields are sorted and deduplicated. - if !single_split_list_fields.is_sorted_by(|left, right| { - // Checking on less ensure that this is both sorted AND that there are no - // duplicates - field_order(left, right) == std::cmp::Ordering::Less - }) { - rate_limited_warn!( - limit_per_min = 1, - "contract breach: fields returned by a leaf are not strictly sorted! \ - please report" - ); - make_sorted_and_dedup(single_split_list_fields); - } + let cpu_task = move || { + for single_split_list_fields in &mut single_split_list_fields_vec { + // This contract is enforced on a different node, etc. so we defensively check + // that the fields are sorted and deduplicated. + if !single_split_list_fields.is_sorted_by(|left, right| { + // Checking on less ensure that this is both sorted AND that there are no + // duplicates + field_order(left, right) == std::cmp::Ordering::Less + }) { + rate_limited_warn!( + limit_per_min = 1, + "contract breach: fields returned by a leaf are not strictly sorted! please \ + report" + ); + make_sorted_and_dedup(single_split_list_fields); } + } - let filtered_list_fields_sorted_iters: Vec<_> = single_split_list_fields_vec - .into_iter() - .map(|list_fields_sorted| { - list_fields_sorted.into_iter().filter(|field| { - if field_patterns.is_empty() { - true - } else { - matches_any_pattern(&field.field_name, &field_patterns) - } - }) + let filtered_list_fields_sorted_iters: Vec<_> = single_split_list_fields_vec + .into_iter() + .map(|list_fields_sorted| { + list_fields_sorted.into_iter().filter(|field| { + if field_patterns.is_empty() { + true + } else { + matches_any_pattern(&field.field_name, &field_patterns) + } }) - .collect(); - merge_leaf_list_fields(filtered_list_fields_sorted_iters) - }) + }) + .collect(); + merge_leaf_list_fields(filtered_list_fields_sorted_iters) + }; + let fields = search_thread_pool() + .run_cpu_intensive_with_identified_caller(cpu_task, "leaf_list_fields") .await .context("failed to merge single split list fields")??; Ok(ListFieldsResponse { fields }) @@ -448,13 +448,16 @@ pub async fn root_list_fields( } let leaf_list_fields_protos: Vec = try_join_all(leaf_request_tasks).await?; let fields = search_thread_pool() - .run_cpu_intensive(move || { - let leaf_list_fields = leaf_list_fields_protos - .into_iter() - .map(|leaf_list_fields_proto| leaf_list_fields_proto.fields.into_iter()) - .collect(); - merge_leaf_list_fields(leaf_list_fields) - }) + .run_cpu_intensive_with_identified_caller( + move || { + let leaf_list_fields = leaf_list_fields_protos + .into_iter() + .map(|leaf_list_fields_proto| leaf_list_fields_proto.fields.into_iter()) + .collect(); + merge_leaf_list_fields(leaf_list_fields) + }, + "root_list_fields", + ) .await .context("failed to merge leaf list fields responses")??; diff --git a/quickwit/quickwit-search/src/root.rs b/quickwit/quickwit-search/src/root.rs index a2473c6fcbf..99010f2ea37 100644 --- a/quickwit/quickwit-search/src/root.rs +++ b/quickwit/quickwit-search/src/root.rs @@ -799,10 +799,13 @@ pub(crate) async fn search_partial_hits_phase( leaf_search_responses.into_iter().map(Ok).collect_vec(); let span = info_span!("merge_fruits"); let mut leaf_search_response = crate::search_thread_pool() - .run_cpu_intensive(move || { - let _span_guard = span.enter(); - merge_collector.merge_fruits(leaf_search_results) - }) + .run_cpu_intensive_with_identified_caller( + move || { + let _span_guard = span.enter(); + merge_collector.merge_fruits(leaf_search_results) + }, + "root_merge", + ) .await .context("failed to merge leaf search responses")? .map_err(|error: TantivyError| crate::SearchError::Internal(error.to_string()))?;