Skip to content

Shrink searcher thread pool #74

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
rdettai-sk merged 1 commit into from
Jul 1, 2026
Merged

Shrink searcher thread pool #74

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
79 changes: 50 additions & 29 deletions quickwit/quickwit-common/src/thread_pool.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -30,8 +29,7 @@ use crate::metrics::{GaugeGuard, IntGaugeVec, OwnedGaugeGuard, new_gauge_vec};
#[derive(Clone)]
pub struct ThreadPool {
thread_pool: Arc<rayon::ThreadPool>,
ongoing_tasks: IntGauge,
pending_tasks: IntGauge,
name: &'static str,
}

impl ThreadPool {
Expand All @@ -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<rayon::ThreadPool> {
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<F, R>(
/// Same as `run_cpu_intensive` but with a caller identifier recorded in the
/// metrics.
pub fn run_cpu_intensive_with_identified_caller<F, R>(
Comment thread
rdettai-sk marked this conversation as resolved.
&self,
cpu_intensive_fn: F,
caller: &'static str,
) -> impl Future<Output = Result<R, Panicked>>
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 || {
Expand All @@ -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<F, R>(
&self,
cpu_intensive_fn: F,
) -> impl Future<Output = Result<R, Panicked>>
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.
Expand Down Expand Up @@ -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 {
Expand All @@ -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"],
),
}
}
Expand Down
96 changes: 50 additions & 46 deletions quickwit/quickwit-search/src/leaf.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"))?
};
Expand Down Expand Up @@ -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}"))
Expand Down Expand Up @@ -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")?
Expand Down Expand Up @@ -1477,7 +1478,10 @@ pub async fn single_doc_mapping_leaf_search(

let leaf_search_response_reresult: Result<Result<LeafSearchResponse, _>, _> =
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");
Expand Down
16 changes: 15 additions & 1 deletion quickwit/quickwit-search/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<T>`.
pub type Result<T> = std::result::Result<T, SearchError>;
Expand Down Expand Up @@ -99,7 +100,20 @@ pub type SearcherPool = Pool<SocketAddr, SearchServiceClient>;

fn search_thread_pool() -> &'static ThreadPool {
static SEARCH_THREAD_POOL: OnceLock<ThreadPool> = 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.
Expand Down
79 changes: 41 additions & 38 deletions quickwit/quickwit-search/src/list_fields.rs
View file
Open in desktop

@rdettai-sk rdettai-sk Jun 30, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lot of diff for nothing (just changed run_cpu_intensive to run_cpu_intensive_with_identified_caller)

Original file line number Diff line number Diff line change
Expand Up @@ -341,40 +341,40 @@ pub async fn leaf_list_fields(

let mut single_split_list_fields_vec: Vec<Vec<ListFieldsEntryResponse>> =
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 })
Expand Down Expand Up @@ -448,13 +448,16 @@ pub async fn root_list_fields(
}
let leaf_list_fields_protos: Vec<ListFieldsResponse> = 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")??;

Expand Down
Loading
Loading