From ee3d7bd69cfb80952552b5b80e089f8e7dd02c40 Mon Sep 17 00:00:00 2001 From: Remi Dettai Date: Thu, 3 Sep 2026 22:06:43 +0200 Subject: [PATCH 1/7] Add CPU scheduler --- quickwit/quickwit-common/src/thread_pool.rs | 333 ---------- .../quickwit-common/src/thread_pool/mod.rs | 189 ++++++ .../src/thread_pool/regular_pool.rs | 178 ++++++ .../src/thread_pool/scheduler.rs | 605 ++++++++++++++++++ .../src/thread_pool/search_pool.rs | 119 ++++ quickwit/quickwit-search/src/leaf.rs | 82 ++- quickwit/quickwit-search/src/lib.rs | 8 +- quickwit/quickwit-search/src/list_fields.rs | 4 +- quickwit/quickwit-search/src/root.rs | 2 +- 9 files changed, 1157 insertions(+), 363 deletions(-) delete mode 100644 quickwit/quickwit-common/src/thread_pool.rs create mode 100644 quickwit/quickwit-common/src/thread_pool/mod.rs create mode 100644 quickwit/quickwit-common/src/thread_pool/regular_pool.rs create mode 100644 quickwit/quickwit-common/src/thread_pool/scheduler.rs create mode 100644 quickwit/quickwit-common/src/thread_pool/search_pool.rs diff --git a/quickwit/quickwit-common/src/thread_pool.rs b/quickwit/quickwit-common/src/thread_pool.rs deleted file mode 100644 index f1ed58309ad..00000000000 --- a/quickwit/quickwit-common/src/thread_pool.rs +++ /dev/null @@ -1,333 +0,0 @@ -// Copyright 2021-Present Datadog, Inc. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use std::fmt; -use std::sync::Arc; - -use futures::{Future, TryFutureExt}; -use once_cell::sync::Lazy; -use tokio::sync::oneshot; -use tracing::error; - -use crate::metrics::{ - Histogram, HistogramTimer, HistogramVec, IntGauge, IntGaugeVec, OwnedGaugeGuard, - exponential_buckets, new_gauge_vec, new_histogram_vec, -}; - -/// An executor backed by a thread pool to run CPU-intensive tasks. -/// -/// tokio::spawn_blocking should only used for IO-bound tasks, as it has not limit on its -/// thread count. -#[derive(Clone)] -pub struct ThreadPool { - thread_pool: Arc, - name: &'static str, -} - -impl ThreadPool { - pub fn new(name: &'static str, num_threads_opt: Option) -> ThreadPool { - let mut rayon_pool_builder = rayon::ThreadPoolBuilder::new() - .thread_name(move |thread_id| format!("quickwit-{name}-{thread_id}")) - .panic_handler(move |_my_panic| { - error!("task running in the quickwit {name} thread pool panicked"); - }); - if let Some(num_threads) = num_threads_opt { - rayon_pool_builder = rayon_pool_builder.num_threads(num_threads); - } - let thread_pool = rayon_pool_builder - .build() - .expect("failed to spawn thread pool"); - ThreadPool { - thread_pool: Arc::new(thread_pool), - name, - } - } - - /// Returns a Tantivy [`tantivy::Executor`] backed by this thread pool. - /// - /// Tasks that Tantivy schedules through it are tracked by metrics. - pub fn get_executor( - &self, - caller: &'static str, - cost_class: &'static str, - ) -> tantivy::Executor { - tantivy::Executor::InstrumentedThreadPool( - self.thread_pool.clone(), - Arc::new(ThreadPoolTaskInstrumentation { - pool_name: self.name, - caller, - cost_class, - }), - ) - } - - /// Same as `run_cpu_intensive` but with a caller identifier recorded in the - /// metrics. - pub fn run_cpu_intensive_with_extra_tags( - &self, - cpu_intensive_fn: F, - caller: &'static str, - cost_class: &'static str, - ) -> impl Future> - where - F: FnOnce() -> R + Send + 'static, - R: Send + 'static, - { - let span = tracing::Span::current(); - let queued_task = QueuedTask::new(self.name, caller, cost_class); - let (tx, rx) = oneshot::channel(); - self.thread_pool.spawn(move || { - if tx.is_closed() { - // dropping `queued_task` still records the time it spent queued - return; - } - let _guard = span.enter(); - let running_task = queued_task.start(); - let result = cpu_intensive_fn(); - drop(running_task); - let _ = tx.send(result); - }); - 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_extra_tags(cpu_intensive_fn, "unknown", "NA") - } -} - -/// Tracks a task submitted to a [`ThreadPool`] while it waits in the queue. -/// -/// Dropping it without calling [`Self::start`] records the time spent queued but -/// no run time, which is what happens to tasks cancelled before they run. -struct QueuedTask { - ongoing_tasks: IntGauge, - run_time: Histogram, - pending_tasks_guard: OwnedGaugeGuard, - queue_wait_timer: HistogramTimer, -} - -impl QueuedTask { - /// Must be called when submitting the task, not once a worker picks it up, - /// for the queue wait time to be measured correctly. - fn new(pool_name: &'static str, caller: &'static str, cost_class: &'static str) -> QueuedTask { - let labels = [pool_name, caller, cost_class]; - let mut pending_tasks_guard = OwnedGaugeGuard::from_gauge( - THREAD_POOL_METRICS.pending_tasks.with_label_values(labels), - ); - pending_tasks_guard.add(1i64); - QueuedTask { - ongoing_tasks: THREAD_POOL_METRICS.ongoing_tasks.with_label_values(labels), - run_time: THREAD_POOL_METRICS.run_time_secs.with_label_values(labels), - pending_tasks_guard, - queue_wait_timer: THREAD_POOL_METRICS - .queue_wait_time_secs - .with_label_values(labels) - .start_timer(), - } - } - - fn start(self) -> RunningTaskGuard { - drop(self.pending_tasks_guard); - self.queue_wait_timer.observe_duration(); - let mut ongoing_tasks_guard = OwnedGaugeGuard::from_gauge(self.ongoing_tasks); - ongoing_tasks_guard.add(1i64); - RunningTaskGuard { - _ongoing_tasks_guard: ongoing_tasks_guard, - _run_timer: self.run_time.start_timer(), - } - } -} - -/// Tracks a task that is running. Dropping it records the run time. -struct RunningTaskGuard { - _ongoing_tasks_guard: OwnedGaugeGuard, - _run_timer: HistogramTimer, -} - -/// Records tasks that Tantivy schedules on a [`ThreadPool`] into the same -/// metrics as the ones scheduled by Quickwit itself. -struct ThreadPoolTaskInstrumentation { - pool_name: &'static str, - caller: &'static str, - cost_class: &'static str, -} - -impl tantivy::TaskInstrumentation for ThreadPoolTaskInstrumentation { - fn enqueue(&self) -> Box { - Box::new(QueuedTask::new( - self.pool_name, - self.caller, - self.cost_class, - )) - } -} - -impl tantivy::EnqueuedTask for QueuedTask { - fn run(self: Box) -> Box { - Box::new(self.start()) - } -} - -impl tantivy::RunningTask for RunningTaskGuard {} - -/// Run a small (<200ms) CPU-intensive task on a dedicated thread pool with a few threads. -/// -/// When running blocking io (or side-effects in general), prefer using `tokio::spawn_blocking` -/// instead. When running long tasks or a set of tasks that you expect to take more than 33% of -/// your vCPUs, use a dedicated thread/runtime or executor instead. -/// -/// Disclaimer: The function will no be executed if the Future is dropped. -#[must_use = "run_cpu_intensive will not run if the future it returns is dropped"] -pub fn run_cpu_intensive(cpu_intensive_fn: F) -> impl Future> -where - F: FnOnce() -> R + Send + 'static, - R: Send + 'static, -{ - static SMALL_TASK_EXECUTOR: std::sync::OnceLock = std::sync::OnceLock::new(); - SMALL_TASK_EXECUTOR - .get_or_init(|| { - let num_threads: usize = (crate::num_cpus() / 3).max(2); - ThreadPool::new("small_tasks", Some(num_threads)) - }) - .run_cpu_intensive(cpu_intensive_fn) -} - -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -pub struct Panicked; - -impl fmt::Display for Panicked { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - write!(f, "scheduled task panicked") - } -} - -impl std::error::Error for Panicked {} - -struct ThreadPoolMetrics { - ongoing_tasks: IntGaugeVec<3>, - pending_tasks: IntGaugeVec<3>, - queue_wait_time_secs: HistogramVec<3>, - run_time_secs: HistogramVec<3>, -} - -/// From 1ms to ~32.768s -fn wait_and_run_time_buckets() -> Vec { - exponential_buckets(0.001, 2.0, 16).unwrap() -} - -impl Default for ThreadPoolMetrics { - fn default() -> Self { - ThreadPoolMetrics { - ongoing_tasks: new_gauge_vec( - "ongoing_tasks", - "number of tasks being currently processed by threads in the thread pool", - "thread_pool", - &[], - ["pool", "caller", "cost_class"], - ), - pending_tasks: new_gauge_vec( - "pending_tasks", - "number of tasks waiting in the queue before being processed by the thread pool", - "thread_pool", - &[], - ["pool", "caller", "cost_class"], - ), - queue_wait_time_secs: new_histogram_vec( - "queue_wait_time_secs", - "amount of time a task waited in the queue before being picked up by a thread in \ - the thread pool", - "thread_pool", - &[], - ["pool", "caller", "cost_class"], - wait_and_run_time_buckets(), - ), - run_time_secs: new_histogram_vec( - "run_time_secs", - "amount of time spent actually running a task on a thread pool worker, once it \ - has been picked up from the queue", - "thread_pool", - &[], - ["pool", "caller", "cost_class"], - wait_and_run_time_buckets(), - ), - } - } -} - -static THREAD_POOL_METRICS: Lazy = Lazy::new(ThreadPoolMetrics::default); - -#[cfg(test)] -mod tests { - use std::sync::Arc; - use std::sync::atomic::{AtomicU64, Ordering}; - use std::time::Duration; - - use super::*; - - #[tokio::test] - async fn test_run_cpu_intensive() { - assert_eq!(run_cpu_intensive(|| 1).await, Ok(1)); - } - - #[tokio::test] - async fn test_run_cpu_intensive_panicks() { - assert!(run_cpu_intensive(|| panic!("")).await.is_err()); - } - - #[tokio::test] - async fn test_run_cpu_intensive_panicks_do_not_shrink_thread_pool() { - for _ in 0..100 { - assert!(run_cpu_intensive(|| panic!("")).await.is_err()); - } - } - - #[tokio::test] - async fn test_run_cpu_intensive_abort() { - let counter: Arc = Default::default(); - let mut futures = Vec::new(); - for _ in 0..1_000 { - let counter_clone = counter.clone(); - let fut = run_cpu_intensive(move || { - std::thread::sleep(Duration::from_millis(5)); - counter_clone.fetch_add(1, Ordering::SeqCst) - }); - // The first few num_cores tasks should run, but the other should get cancelled. - futures.push(tokio::time::timeout(Duration::from_millis(1), fut)); - } - futures::future::join_all(futures).await; - assert!(counter.load(Ordering::SeqCst) < 100); - } -} diff --git a/quickwit/quickwit-common/src/thread_pool/mod.rs b/quickwit/quickwit-common/src/thread_pool/mod.rs new file mode 100644 index 00000000000..8d9e86e269e --- /dev/null +++ b/quickwit/quickwit-common/src/thread_pool/mod.rs @@ -0,0 +1,189 @@ +// Copyright 2021-Present Datadog, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use futures::{Future, TryFutureExt}; +use once_cell::sync::Lazy; +use tokio::sync::oneshot; + +use crate::metrics::{ + Histogram, HistogramTimer, HistogramVec, IntGauge, IntGaugeVec, OwnedGaugeGuard, + exponential_buckets, new_gauge_vec, new_histogram_vec, +}; + +pub mod scheduler; + +mod regular_pool; +mod search_pool; + +pub use regular_pool::{Panicked, ThreadPool, run_cpu_intensive}; +pub use search_pool::SearchThreadPool; + +/// Wraps `cpu_intensive_fn` with cancellation-awareness and queue/run-time +/// metrics tracking, then hands the resulting job to `dispatch` for actual +/// scheduling (a raw rayon spawn, or one of the scheduler's queues). +fn spawn_traced( + pool_name: &'static str, + caller: &'static str, + cost_class: &'static str, + cpu_intensive_fn: F, + dispatch: impl FnOnce(Box), +) -> impl Future> +where + F: FnOnce() -> R + Send + 'static, + R: Send + 'static, +{ + let span = tracing::Span::current(); + let queued_task = QueuedTask::new(pool_name, caller, cost_class); + let (tx, rx) = oneshot::channel(); + dispatch(Box::new(move || { + if tx.is_closed() { + // dropping `queued_task` still records the time it spent queued + return; + } + let _guard = span.enter(); + let running_task = queued_task.start(); + let result = cpu_intensive_fn(); + drop(running_task); + let _ = tx.send(result); + })); + rx.map_err(|_| Panicked) +} + +/// Tracks a task submitted to a [`ThreadPool`] while it waits in the queue. +/// +/// Dropping it without calling [`Self::start`] records the time spent queued but +/// no run time, which is what happens to tasks cancelled before they run. +struct QueuedTask { + ongoing_tasks: IntGauge, + run_time: Histogram, + pending_tasks_guard: OwnedGaugeGuard, + queue_wait_timer: HistogramTimer, +} + +impl QueuedTask { + /// Must be called when submitting the task, not once a worker picks it up, + /// for the queue wait time to be measured correctly. + fn new(pool_name: &'static str, caller: &'static str, cost_class: &'static str) -> QueuedTask { + let labels = [pool_name, caller, cost_class]; + let mut pending_tasks_guard = OwnedGaugeGuard::from_gauge( + THREAD_POOL_METRICS.pending_tasks.with_label_values(labels), + ); + pending_tasks_guard.add(1i64); + QueuedTask { + ongoing_tasks: THREAD_POOL_METRICS.ongoing_tasks.with_label_values(labels), + run_time: THREAD_POOL_METRICS.run_time_secs.with_label_values(labels), + pending_tasks_guard, + queue_wait_timer: THREAD_POOL_METRICS + .queue_wait_time_secs + .with_label_values(labels) + .start_timer(), + } + } + + fn start(self) -> RunningTaskGuard { + drop(self.pending_tasks_guard); + self.queue_wait_timer.observe_duration(); + let mut ongoing_tasks_guard = OwnedGaugeGuard::from_gauge(self.ongoing_tasks); + ongoing_tasks_guard.add(1i64); + RunningTaskGuard { + _ongoing_tasks_guard: ongoing_tasks_guard, + _run_timer: self.run_time.start_timer(), + } + } +} + +/// Tracks a task that is running. Dropping it records the run time. +struct RunningTaskGuard { + _ongoing_tasks_guard: OwnedGaugeGuard, + _run_timer: HistogramTimer, +} + +/// Records tasks that Tantivy schedules on a [`ThreadPool`] into the same +/// metrics as the ones scheduled by Quickwit itself. +struct ThreadPoolTaskInstrumentation { + pool_name: &'static str, + caller: &'static str, + cost_class: &'static str, +} + +impl tantivy::TaskInstrumentation for ThreadPoolTaskInstrumentation { + fn enqueue(&self) -> Box { + Box::new(QueuedTask::new( + self.pool_name, + self.caller, + self.cost_class, + )) + } +} + +impl tantivy::EnqueuedTask for QueuedTask { + fn run(self: Box) -> Box { + Box::new(self.start()) + } +} + +impl tantivy::RunningTask for RunningTaskGuard {} + +struct ThreadPoolMetrics { + ongoing_tasks: IntGaugeVec<3>, + pending_tasks: IntGaugeVec<3>, + queue_wait_time_secs: HistogramVec<3>, + run_time_secs: HistogramVec<3>, +} + +/// From 1ms to ~32.768s +fn wait_and_run_time_buckets() -> Vec { + exponential_buckets(0.001, 2.0, 16).unwrap() +} + +impl Default for ThreadPoolMetrics { + fn default() -> Self { + ThreadPoolMetrics { + ongoing_tasks: new_gauge_vec( + "ongoing_tasks", + "number of tasks being currently processed by threads in the thread pool", + "thread_pool", + &[], + ["pool", "caller", "cost_class"], + ), + pending_tasks: new_gauge_vec( + "pending_tasks", + "number of tasks waiting in the queue before being processed by the thread pool", + "thread_pool", + &[], + ["pool", "caller", "cost_class"], + ), + queue_wait_time_secs: new_histogram_vec( + "queue_wait_time_secs", + "amount of time a task waited in the queue before being picked up by a thread in \ + the thread pool", + "thread_pool", + &[], + ["pool", "caller", "cost_class"], + wait_and_run_time_buckets(), + ), + run_time_secs: new_histogram_vec( + "run_time_secs", + "amount of time spent actually running a task on a thread pool worker, once it \ + has been picked up from the queue", + "thread_pool", + &[], + ["pool", "caller", "cost_class"], + wait_and_run_time_buckets(), + ), + } + } +} + +static THREAD_POOL_METRICS: Lazy = Lazy::new(ThreadPoolMetrics::default); diff --git a/quickwit/quickwit-common/src/thread_pool/regular_pool.rs b/quickwit/quickwit-common/src/thread_pool/regular_pool.rs new file mode 100644 index 00000000000..35981845cbd --- /dev/null +++ b/quickwit/quickwit-common/src/thread_pool/regular_pool.rs @@ -0,0 +1,178 @@ +// Copyright 2021-Present Datadog, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::fmt; +use std::sync::Arc; + +use futures::Future; +use tracing::error; + +use super::ThreadPoolTaskInstrumentation; + +/// An executor backed by a thread pool to run CPU-intensive tasks. +/// +/// tokio::spawn_blocking should only used for IO-bound tasks, as it has not limit on its +/// thread count. +/// +/// Unlike [`super::SearchThreadPool`], dispatches FIFO straight onto the raw +/// rayon pool, with no per-query priority/fairness layer. +#[derive(Clone)] +pub struct ThreadPool { + pub(super) rayon_pool: Arc, + pub(super) name: &'static str, +} + +impl ThreadPool { + pub fn new(name: &'static str, num_threads_opt: Option) -> ThreadPool { + let mut rayon_pool_builder = rayon::ThreadPoolBuilder::new() + .thread_name(move |thread_id| format!("quickwit-{name}-{thread_id}")) + .panic_handler(move |_my_panic| { + error!("task running in the quickwit {name} thread pool panicked"); + }); + if let Some(num_threads) = num_threads_opt { + rayon_pool_builder = rayon_pool_builder.num_threads(num_threads); + } + let rayon_pool = rayon_pool_builder + .build() + .expect("failed to spawn thread pool"); + ThreadPool { + rayon_pool: Arc::new(rayon_pool), + name, + } + } + + /// Returns a Tantivy [`tantivy::Executor`] backed by this thread pool. + /// + /// Tasks that Tantivy schedules through it are tracked by metrics. + pub fn get_executor( + &self, + caller: &'static str, + cost_class: &'static str, + ) -> tantivy::Executor { + tantivy::Executor::InstrumentedThreadPool( + self.rayon_pool.clone(), + Arc::new(ThreadPoolTaskInstrumentation { + pool_name: self.name, + caller, + cost_class, + }), + ) + } + + /// 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, + caller: &'static str, + cost_class: &'static str, + ) -> impl Future> + where + F: FnOnce() -> R + Send + 'static, + R: Send + 'static, + { + super::spawn_traced(self.name, caller, cost_class, cpu_intensive_fn, |job| { + self.rayon_pool.spawn(job) + }) + } +} + +/// Run a small (<200ms) CPU-intensive task on a dedicated thread pool with a few threads. +/// +/// When running blocking io (or side-effects in general), prefer using `tokio::spawn_blocking` +/// instead. When running long tasks or a set of tasks that you expect to take more than 33% of +/// your vCPUs, use a dedicated thread/runtime or executor instead. +/// +/// Disclaimer: The function will no be executed if the Future is dropped. +#[must_use = "run_cpu_intensive will not run if the future it returns is dropped"] +pub fn run_cpu_intensive(cpu_intensive_fn: F) -> impl Future> +where + F: FnOnce() -> R + Send + 'static, + R: Send + 'static, +{ + static SMALL_TASK_EXECUTOR: std::sync::OnceLock = std::sync::OnceLock::new(); + SMALL_TASK_EXECUTOR + .get_or_init(|| { + let num_threads: usize = (crate::num_cpus() / 3).max(2); + ThreadPool::new("small_tasks", Some(num_threads)) + }) + .run_cpu_intensive(cpu_intensive_fn, "unknown", "NA") +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct Panicked; + +impl fmt::Display for Panicked { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "scheduled task panicked") + } +} + +impl std::error::Error for Panicked {} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::sync::atomic::{AtomicU64, Ordering}; + use std::time::Duration; + + use super::*; + + #[tokio::test] + async fn test_run_cpu_intensive() { + assert_eq!(run_cpu_intensive(|| 1).await, Ok(1)); + } + + #[tokio::test] + async fn test_run_cpu_intensive_panicks() { + assert!(run_cpu_intensive(|| panic!("")).await.is_err()); + } + + #[tokio::test] + async fn test_run_cpu_intensive_panicks_do_not_shrink_thread_pool() { + for _ in 0..100 { + assert!(run_cpu_intensive(|| panic!("")).await.is_err()); + } + } + + #[tokio::test] + async fn test_run_cpu_intensive_abort() { + let counter: Arc = Default::default(); + let mut futures = Vec::new(); + for _ in 0..1_000 { + let counter_clone = counter.clone(); + let fut = run_cpu_intensive(move || { + std::thread::sleep(Duration::from_millis(5)); + counter_clone.fetch_add(1, Ordering::SeqCst) + }); + // The first few num_cores tasks should run, but the other should get cancelled. + futures.push(tokio::time::timeout(Duration::from_millis(1), fut)); + } + futures::future::join_all(futures).await; + assert!(counter.load(Ordering::SeqCst) < 100); + } +} diff --git a/quickwit/quickwit-common/src/thread_pool/scheduler.rs b/quickwit/quickwit-common/src/thread_pool/scheduler.rs new file mode 100644 index 00000000000..179d844abf0 --- /dev/null +++ b/quickwit/quickwit-common/src/thread_pool/scheduler.rs @@ -0,0 +1,605 @@ +// Copyright 2021-Present Datadog, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Priority scheduler sitting in front of a [`rayon::ThreadPool`]. +//! +//! Rayon has no notion of task priority: it dispatches in whatever order tasks +//! land in its own local deques / injector queue. This module keeps its own +//! queue and uses rayon only to own OS threads that continuously drain it (the +//! "pump loop" pattern). +//! +//! Two tiers of priority exist: +//! - High priority: always runs first and processed in strict FIFO order. Meant for short and rarer +//! tasks such as merging/finalizing a query's results. +//! - Per-query tasks: tries to be fair among queries, with a bias towards queries that are closer +//! to completion. + +use std::collections::{HashMap, VecDeque}; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex, MutexGuard}; +use std::time::{Duration, Instant}; + +use once_cell::sync::Lazy; +use tracing::error; + +use crate::metrics::{IntCounter, IntGauge, new_counter, new_gauge}; + +/// Identifies a query (leaf search) whose split-processing tasks should be +/// scheduled and fair-shared together. Must be unique among currently active +/// queries: reusing an id while a previous query with that id is still being +/// cleaned up would corrupt its accounting. +#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)] +pub struct QueryId(u64); + +impl QueryId { + /// Allocates a fresh, never-reused `QueryId`. + pub fn next() -> QueryId { + static NEXT_QUERY_ID: AtomicU64 = AtomicU64::new(1); + QueryId(NEXT_QUERY_ID.fetch_add(1, Ordering::Relaxed)) + } +} + +/// One per split registered via [`Scheduler::register_query`]. Dropping it +/// resolves that split -- via normal completion or via being dropped along +/// with a cancelled future -- so the query's entry can never leak regardless +/// of how its caller ends. +#[must_use = "dropping this immediately resolves the split"] +pub struct SchedulerSplitGuard { + scheduler: Arc, + query_id: QueryId, +} + +impl SchedulerSplitGuard { + pub fn query_id(&self) -> QueryId { + self.query_id + } +} + +impl Drop for SchedulerSplitGuard { + fn drop(&mut self) { + self.scheduler.split_resolved(self.query_id); + } +} + +type Job = Box; + +/// How long a pump loop keeps grabbing tasks before handing its worker back +/// to rayon's own scheduler (see [`pump_loop`]). +const PUMP_LOOP_YIELD_INTERVAL: Duration = Duration::from_millis(200); + +/// Per-query scheduling state. +struct QueryState { + /// Tasks submitted for this query that have not yet been dispatched. + ready: VecDeque, + /// Tasks currently executing on a worker. + running_count: usize, + /// Number of this query's splits still waiting on a `SearchPermit`, not + /// yet admitted into warmup/CPU processing. Kept up to date by the + /// caller via [`Scheduler::set_waiting_for_permit`]. + waiting_for_permit: usize, + /// Number of this query's splits not yet resolved through *any* terminal + /// path (processed, pruned, cache hit, ...). Decremented by + /// [`Scheduler::split_resolved`], which also cleans up the query once it + /// reaches zero. + remaining: usize, + /// Used as the final tie-break: older queries win, for fairness/liveness + /// among otherwise-indistinguishable queries. + created_at: Instant, +} + +impl QueryState { + fn priority_key(&self) -> (usize, usize, Instant) { + (self.waiting_for_permit, self.remaining, self.created_at) + } +} + +struct SchedulerState { + /// Always-first, uncapped tasks (e.g. finalize/root_merge). + high_priority_queue: VecDeque, + queries: HashMap, + /// Number of pump loops currently alive, bounded by `num_threads`. + active_pump_workers: usize, +} + +impl SchedulerState { + /// The maximum number of concurrently running tasks any single query may + /// have right now, given how many queries are currently competing for the + /// pool. + fn current_cap(&self, num_threads: usize) -> usize { + let competing_queries = self + .queries + .values() + .filter(|query| !query.ready.is_empty()) + .count(); + num_threads.div_ceil(competing_queries.max(1)) + } + + /// Pops the single highest-priority ready and eligible task, if any, + /// updating `running_count` for its owning query. Called with the lock + /// already held, by a pump loop looking for its next unit of work. + /// + /// Deliberately implemented here rather than on [`Scheduler`]: taking + /// only `&mut self` (no access to `Scheduler::state`, the `Mutex` this is + /// always called with already locked) makes it structurally impossible + /// for this to ever try to re-lock it and deadlock. + fn pick_next(&mut self, num_threads: usize) -> Option { + if let Some(job) = self.high_priority_queue.pop_front() { + return Some(job); + } + let cap = self.current_cap(num_threads); + let best_query_id = self + .queries + .iter() + .filter(|(_, query)| query.running_count < cap && !query.ready.is_empty()) + .min_by_key(|(_, query)| query.priority_key()) + .map(|(query_id, _)| *query_id)?; + let query = self + .queries + .get_mut(&best_query_id) + .expect("query looked up right above must still be present"); + let job = query + .ready + .pop_front() + .expect("query was only selected because its ready queue is non-empty"); + query.running_count += 1; + Some(job) + } +} + +/// A priority scheduler backed by a [`rayon::ThreadPool`]. See the module +/// documentation for the overall design. +pub struct Scheduler { + rayon_pool: Arc, + num_threads: usize, + state: Mutex, +} + +impl Scheduler { + pub fn new(rayon_pool: Arc) -> Arc { + let num_threads = rayon_pool.current_num_threads(); + Arc::new(Scheduler { + rayon_pool, + num_threads, + state: Mutex::new(SchedulerState { + high_priority_queue: VecDeque::new(), + queries: HashMap::new(), + active_pump_workers: 0, + }), + }) + } + + /// Locks `state`, adding how long the calling thread had to wait to + /// acquire it to [`SchedulerMetrics::lock_wait_time_nanos_total`]. The + /// lock only ever guards brief in-memory bookkeeping, so a fast-growing + /// total directly indicates contention. + fn lock_state(&self) -> MutexGuard<'_, SchedulerState> { + let wait_start = Instant::now(); + let guard = self.state.lock().unwrap(); + SCHEDULER_METRICS + .lock_wait_time_nanos_total + .inc_by(wait_start.elapsed().as_nanos() as u64); + guard + } + + /// Registers a new query with its total split count. Must be called + /// exactly once per query, before any [`Self::enqueue_fifo`], + /// [`Self::enqueue_fair`] or [`Self::set_waiting_for_permit`] call for + /// that `query_id`. + /// + /// Returns one guard per split, meant to be wrapped into that split's own + /// cleanup guard on the caller side, so the query's entry can never leak + /// regardless of how each split's processing ends. + pub fn register_query( + self: &Arc, + query_id: QueryId, + total_splits: usize, + ) -> Vec { + let mut state = self.lock_state(); + state.queries.insert( + query_id, + QueryState { + ready: VecDeque::new(), + running_count: 0, + waiting_for_permit: 0, + remaining: total_splits, + created_at: Instant::now(), + }, + ); + SCHEDULER_METRICS.queries.set(state.queries.len() as i64); + drop(state); + (0..total_splits) + .map(|_| SchedulerSplitGuard { + scheduler: self.clone(), + query_id, + }) + .collect() + } + + /// Updates how many of `query_id`'s splits are still waiting on a + /// `SearchPermit`. A query with none left (the common case once + /// admission is done) is preferred over one still mostly permit-gated. + pub fn set_waiting_for_permit(&self, query_id: QueryId, waiting_for_permit: usize) { + let mut state = self.lock_state(); + if let Some(query) = state.queries.get_mut(&query_id) { + query.waiting_for_permit = waiting_for_permit; + } + } + + /// Called by [`SplitGuard::drop`] exactly once per split of `query_id`, + /// decrementing `remaining` and removing the query's entry once it hits + /// zero. + fn split_resolved(&self, query_id: QueryId) { + let mut state = self.lock_state(); + let Some(query) = state.queries.get_mut(&query_id) else { + return; + }; + query.remaining = query.remaining.saturating_sub(1); + if query.remaining == 0 { + state.queries.remove(&query_id); + SCHEDULER_METRICS.queries.set(state.queries.len() as i64); + } + } + + /// Schedules a high priority task: always dispatched before any per-query + /// task, and processed FIFO. Long tasks (>100ms) are not recommended. + pub fn enqueue_fifo(self: &Arc, job: F) + where F: FnOnce() + Send + 'static { + let mut state = self.lock_state(); + state.high_priority_queue.push_back(Box::new(job)); + if state.active_pump_workers < self.num_threads { + state.active_pump_workers += 1; + let scheduler = self.clone(); + self.rayon_pool.spawn(move || pump_loop(&scheduler)); + } + } + + /// Schedules a task belonging to `query_id`. The query must already have + /// been [`Self::register_query`]-ed. + pub fn enqueue_fair(self: &Arc, query_id: QueryId, job: F) + where F: FnOnce() + Send + 'static { + let scheduler = self.clone(); + let wrapped: Job = Box::new(move || { + job(); + scheduler.on_task_complete(query_id); + }); + let mut state = self.lock_state(); + let query = state + .queries + .get_mut(&query_id) + .expect("query must be registered before tasks are enqueued for it"); + query.ready.push_back(wrapped); + if state.active_pump_workers < self.num_threads { + state.active_pump_workers += 1; + let scheduler = self.clone(); + self.rayon_pool.spawn(move || pump_loop(&scheduler)); + } + } + + fn on_task_complete(&self, query_id: QueryId) { + let mut state = self.lock_state(); + if let Some(query) = state.queries.get_mut(&query_id) { + query.running_count = query.running_count.saturating_sub(1); + } + } +} + +/// Runs on a rayon worker thread: repeatedly picks and runs the current best +/// task until none is eligible, then gives up its slot. +/// +/// The exit check and the `active_pump_workers` decrement happen in the same +/// critical section as `pick_next`'s "nothing to do" verdict, so a concurrent +/// `enqueue_fifo`/`enqueue_fair` call always sees an up-to-date count and +/// spawns a replacement if this one is exiting right as new work arrives. +/// +/// Unfortunately, the rayon pool is also used without the scheduler (as a +/// Tantivy Executor), so the pump loop needs to periodically yield to let rayon +/// schedule those tasks. +fn pump_loop(scheduler: &Arc) { + let yield_deadline = Instant::now() + PUMP_LOOP_YIELD_INTERVAL; + loop { + let job = { + let mut state = scheduler.lock_state(); + match state.pick_next(scheduler.num_threads) { + Some(job) => job, + None => { + state.active_pump_workers -= 1; + return; + } + } + }; + if std::panic::catch_unwind(std::panic::AssertUnwindSafe(job)).is_err() { + error!("task running in the thread pool scheduler panicked"); + } + if Instant::now() >= yield_deadline { + let replacement = scheduler.clone(); + scheduler.rayon_pool.spawn(move || pump_loop(&replacement)); + return; + } + } +} + +struct SchedulerMetrics { + /// Number of queries currently registered with the scheduler. + queries: IntGauge, + /// Cumulative time (in nanoseconds) callers have spent waiting to + /// acquire `Scheduler::state`. The lock only ever guards brief in-memory + /// bookkeeping, so a fast-growing total directly indicates contention. + lock_wait_time_nanos_total: IntCounter, +} + +impl Default for SchedulerMetrics { + fn default() -> Self { + SchedulerMetrics { + queries: new_gauge( + "scheduler_queries", + "number of queries currently registered with the CPU scheduler", + "thread_pool", + &[], + ), + lock_wait_time_nanos_total: new_counter( + "scheduler_lock_wait_time_nanos_total", + "cumulative time, in nanoseconds, spent waiting to acquire the CPU scheduler's \ + internal lock", + "thread_pool", + &[], + ), + } + } +} + +static SCHEDULER_METRICS: Lazy = Lazy::new(SchedulerMetrics::default); + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::{Arc, Mutex as StdMutex}; + use std::time::Duration; + + use super::*; + + fn test_scheduler(num_threads: usize) -> Arc { + let rayon_pool = Arc::new( + rayon::ThreadPoolBuilder::new() + .num_threads(num_threads) + .build() + .unwrap(), + ); + Scheduler::new(rayon_pool) + } + + // Polls until `condition` is true or the timeout elapses, to avoid flaky + // sleeps while still bounding worst-case test time. + fn wait_until(mut condition: impl FnMut() -> bool) { + let deadline = Instant::now() + Duration::from_secs(5); + while !condition() { + assert!(Instant::now() < deadline, "condition never became true"); + std::thread::sleep(Duration::from_millis(5)); + } + } + + #[test] + fn test_level0_runs_before_query_tasks() { + let scheduler = test_scheduler(1); + let order: Arc>> = Arc::new(StdMutex::new(Vec::new())); + + let _guards = scheduler.register_query(QueryId(1), 1); + // Block the single worker so both tasks below are enqueued before either runs. + let (release_tx, release_rx) = std::sync::mpsc::channel::<()>(); + scheduler.enqueue_fair(QueryId(1), move || { + release_rx.recv().unwrap(); + }); + wait_until(|| scheduler.state.lock().unwrap().active_pump_workers == 1); + + let order_clone = order.clone(); + scheduler.enqueue_fair(QueryId(1), move || { + order_clone.lock().unwrap().push("query") + }); + let order_clone = order.clone(); + scheduler.enqueue_fifo(move || order_clone.lock().unwrap().push("level0")); + + release_tx.send(()).unwrap(); + wait_until(|| order.lock().unwrap().len() == 2); + assert_eq!(*order.lock().unwrap(), vec!["level0", "query"]); + } + + #[test] + fn test_smaller_remaining_runs_first() { + let scheduler = test_scheduler(1); + let order: Arc>> = Arc::new(StdMutex::new(Vec::new())); + + let _guards1 = scheduler.register_query(QueryId(1), 100); + let _guards2 = scheduler.register_query(QueryId(2), 2); + + let (release_tx, release_rx) = std::sync::mpsc::channel::<()>(); + scheduler.enqueue_fair(QueryId(1), move || { + release_rx.recv().unwrap(); + }); + wait_until(|| scheduler.state.lock().unwrap().active_pump_workers == 1); + + let order_clone = order.clone(); + scheduler.enqueue_fair(QueryId(1), move || { + order_clone.lock().unwrap().push(QueryId(1)) + }); + let order_clone = order.clone(); + scheduler.enqueue_fair(QueryId(2), move || { + order_clone.lock().unwrap().push(QueryId(2)) + }); + + release_tx.send(()).unwrap(); + wait_until(|| order.lock().unwrap().len() == 2); + // query 2 has far fewer remaining splits, so it should be picked first. + assert_eq!(*order.lock().unwrap(), vec![QueryId(2), QueryId(1)]); + } + + #[test] + fn test_cap_ignores_queries_with_no_ready_or_running_work() { + let scheduler = test_scheduler(3); + // Queries 1 and 2 simulate splits still waiting on a `SearchPermit`: + // registered, but with nothing enqueued on the CPU scheduler yet. + let _guards1 = scheduler.register_query(QueryId(1), 100); + let _guards2 = scheduler.register_query(QueryId(2), 100); + let _guards3 = scheduler.register_query(QueryId(3), 100); + + let concurrent_3 = Arc::new(AtomicUsize::new(0)); + let max_concurrent_3 = Arc::new(AtomicUsize::new(0)); + let (release_tx, release_rx) = std::sync::mpsc::channel::<()>(); + let release_rx = Arc::new(StdMutex::new(release_rx)); + for _ in 0..10 { + let concurrent_3 = concurrent_3.clone(); + let max_concurrent_3 = max_concurrent_3.clone(); + let release_rx = release_rx.clone(); + scheduler.enqueue_fair(QueryId(3), move || { + let current = concurrent_3.fetch_add(1, Ordering::SeqCst) + 1; + max_concurrent_3.fetch_max(current, Ordering::SeqCst); + release_rx.lock().unwrap().recv().unwrap(); + concurrent_3.fetch_sub(1, Ordering::SeqCst); + }); + } + + wait_until(|| concurrent_3.load(Ordering::SeqCst) == 3); + // 3 queries are registered, but only query 3 has any ready/running + // work, so it should get the whole pool instead of being capped at + // usable/3 == 1 while the other two threads sit idle. + assert_eq!(max_concurrent_3.load(Ordering::SeqCst), 3); + + for _ in 0..10 { + release_tx.send(()).unwrap(); + } + } + + #[test] + fn test_per_query_cap_shares_the_pool() { + let scheduler = test_scheduler(4); + let _guards1 = scheduler.register_query(QueryId(1), 100); + let _guards2 = scheduler.register_query(QueryId(2), 100); + + let concurrent_1 = Arc::new(AtomicUsize::new(0)); + let max_concurrent_1 = Arc::new(AtomicUsize::new(0)); + let (release_tx_1, release_rx_1) = std::sync::mpsc::channel::<()>(); + let release_rx_1 = Arc::new(StdMutex::new(release_rx_1)); + let (release_tx_2, release_rx_2) = std::sync::mpsc::channel::<()>(); + let release_rx_2 = Arc::new(StdMutex::new(release_rx_2)); + + // Interleave both queries' backlogs (each with more tasks than the + // pool could ever run at once for it alone) so both count as + // competing from the start, splitting cap = 4/2 = 2 between them. + // Giving one query its whole backlog first would let it alone grab + // cap = 4/1 = 4 -- the entire pool -- before the other ever gets a + // chance to compete. + for _ in 0..10 { + let concurrent_1 = concurrent_1.clone(); + let max_concurrent_1 = max_concurrent_1.clone(); + let release_rx_1 = release_rx_1.clone(); + scheduler.enqueue_fair(QueryId(1), move || { + let current = concurrent_1.fetch_add(1, Ordering::SeqCst) + 1; + max_concurrent_1.fetch_max(current, Ordering::SeqCst); + release_rx_1.lock().unwrap().recv().unwrap(); + concurrent_1.fetch_sub(1, Ordering::SeqCst); + }); + let release_rx_2 = release_rx_2.clone(); + scheduler.enqueue_fair(QueryId(2), move || { + release_rx_2.lock().unwrap().recv().unwrap(); + }); + } + + wait_until(|| concurrent_1.load(Ordering::SeqCst) >= 2); + // both queries have a backlog, so cap = 4/2 = 2. + assert!(max_concurrent_1.load(Ordering::SeqCst) <= 2); + + for _ in 0..10 { + release_tx_1.send(()).unwrap(); + } + for _ in 0..10 { + release_tx_2.send(()).unwrap(); + } + } + + #[test] + fn test_pump_loop_yields_periodically_for_directly_injected_rayon_work() { + // A continuous flow of fair-share work keeps pump loops finding more + // ready tasks (there's always another one queued), so without a + // periodic yield they would never return control to rayon's own + // scheduler -- starving anything submitted straight to the same + // rayon pool outside the scheduler (e.g. Tantivy's own internal + // parallelism via `ThreadPool::get_executor`), which only gets + // picked up by a worker that actually returns to rayon's scheduling + // loop. No core is ever permanently sacrificed for this: pump loops + // just periodically hand their worker back and re-queue themselves. + let scheduler = test_scheduler(2); + let _guards = scheduler.register_query(QueryId(1), 100_000); + + // Keep both workers continuously busy with a long stream of short + // tasks, well over one yield interval in total. + for _ in 0..2_000 { + scheduler.enqueue_fair(QueryId(1), || { + std::thread::sleep(Duration::from_millis(1)); + }); + } + wait_until(|| scheduler.state.lock().unwrap().active_pump_workers >= 1); + + let (tx, rx) = std::sync::mpsc::channel(); + scheduler.rayon_pool.spawn(move || tx.send(()).unwrap()); + rx.recv_timeout(Duration::from_secs(2)) + .expect("directly-injected rayon work starved: pump loops never yielded"); + } + + #[test] + fn test_query_state_cleaned_up_once_remaining_reaches_zero() { + let scheduler = test_scheduler(2); + let mut guards = scheduler.register_query(QueryId(1), 2); + assert!( + scheduler + .state + .lock() + .unwrap() + .queries + .contains_key(&QueryId(1)) + ); + + drop(guards.pop().unwrap()); + assert!( + scheduler + .state + .lock() + .unwrap() + .queries + .contains_key(&QueryId(1)) + ); + + drop(guards.pop().unwrap()); + assert!( + !scheduler + .state + .lock() + .unwrap() + .queries + .contains_key(&QueryId(1)) + ); + } + + #[test] + fn test_pump_loops_drain_and_exit() { + let scheduler = test_scheduler(2); + let _guards = scheduler.register_query(QueryId(1), 3); + let ran = Arc::new(AtomicUsize::new(0)); + for _ in 0..3 { + let ran = ran.clone(); + scheduler.enqueue_fair(QueryId(1), move || { + ran.fetch_add(1, Ordering::SeqCst); + }); + } + wait_until(|| ran.load(Ordering::SeqCst) == 3); + wait_until(|| scheduler.state.lock().unwrap().active_pump_workers == 0); + } +} diff --git a/quickwit/quickwit-common/src/thread_pool/search_pool.rs b/quickwit/quickwit-common/src/thread_pool/search_pool.rs new file mode 100644 index 00000000000..fad860c6fe0 --- /dev/null +++ b/quickwit/quickwit-common/src/thread_pool/search_pool.rs @@ -0,0 +1,119 @@ +// Copyright 2021-Present Datadog, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use futures::Future; + +use super::scheduler::{QueryId, Scheduler, SchedulerSplitGuard}; +use super::{Panicked, ThreadPool}; + +/// A [`ThreadPool`] with a per-query fair-share priority scheduler (see +/// [`super::scheduler`]) sitting in front of it, for CPU-intensive tasks that +/// belong to a specific query and must be prioritized/fair-shared against +/// each other. The plain [`ThreadPool`] has no notion of query and dispatches +/// FIFO, which is all one-off or query-less CPU work needs. +#[derive(Clone)] +pub struct SearchThreadPool { + thread_pool: ThreadPool, + scheduler: Arc, +} + +impl SearchThreadPool { + pub fn new(name: &'static str, num_threads_opt: Option) -> SearchThreadPool { + let thread_pool = ThreadPool::new(name, num_threads_opt); + let scheduler = Scheduler::new(thread_pool.rayon_pool.clone()); + SearchThreadPool { + thread_pool, + scheduler, + } + } + + /// Registers a new query for per-query fair-share scheduling. See + /// [`Scheduler::register_query`]. + pub fn register_query( + &self, + query_id: QueryId, + total_splits: usize, + ) -> Vec { + self.scheduler.register_query(query_id, total_splits) + } + + /// See [`Scheduler::set_waiting_for_permit`]. + pub fn set_waiting_for_permit(&self, query_id: QueryId, waiting_for_permit: usize) { + self.scheduler + .set_waiting_for_permit(query_id, waiting_for_permit); + } + + /// Returns a Tantivy [`tantivy::Executor`] backed by this thread pool. + /// + /// Tasks that Tantivy schedules through it are tracked by metrics, but -- + /// unlike [`Self::run_cpu_intensive_fair`] -- bypass the per-query + /// priority scheduler entirely: Tantivy dispatches directly onto the raw + /// rayon pool. + pub fn get_executor( + &self, + caller: &'static str, + cost_class: &'static str, + ) -> tantivy::Executor { + self.thread_pool.get_executor(caller, cost_class) + } + + /// Runs a CPU-intensive task belonging to `query_id`, subject to the + /// per-query fair-share scheduling and priority ordering described in + /// [`super::scheduler`]. `query_id` must already have been registered via + /// [`Self::register_query`]. + pub fn run_cpu_intensive_fair( + &self, + cpu_intensive_fn: F, + query_id: QueryId, + caller: &'static str, + cost_class: &'static str, + ) -> impl Future> + where + F: FnOnce() -> R + Send + 'static, + R: Send + 'static, + { + super::spawn_traced( + self.thread_pool.name, + caller, + cost_class, + cpu_intensive_fn, + move |job| self.scheduler.enqueue_fair(query_id, job), + ) + } + + /// Runs a CPU-intensive task ahead of any per-query fair-share task (see + /// [`super::scheduler`]'s high-priority queue). Meant for short, rare, + /// one-shot work such as finalizing or merging a query's results, and for + /// callers with no query to fair-share against at all. + pub fn run_cpu_intensive( + &self, + cpu_intensive_fn: F, + caller: &'static str, + cost_class: &'static str, + ) -> impl Future> + where + F: FnOnce() -> R + Send + 'static, + R: Send + 'static, + { + super::spawn_traced( + self.thread_pool.name, + caller, + cost_class, + cpu_intensive_fn, + |job| self.scheduler.enqueue_fifo(job), + ) + } +} diff --git a/quickwit/quickwit-search/src/leaf.rs b/quickwit/quickwit-search/src/leaf.rs index 75e665c1746..040259aea43 100644 --- a/quickwit/quickwit-search/src/leaf.rs +++ b/quickwit/quickwit-search/src/leaf.rs @@ -23,6 +23,7 @@ use anyhow::Context; use bytesize::ByteSize; use futures::future::try_join_all; use quickwit_common::pretty::PrettySample; +use quickwit_common::thread_pool::scheduler::{QueryId, SchedulerSplitGuard}; use quickwit_directories::{CachingDirectory, HotDirectory, StorageDirectory}; use quickwit_doc_mapper::{Automaton, DocMapper, FastFieldWarmupInfo, TermRange, WarmupInfo}; use quickwit_proto::search::{ @@ -224,6 +225,7 @@ pub(crate) async fn warmup( searcher: &Searcher, warmup_info: &WarmupInfo, cost_class: QueryCostClass, + query_id: QueryId, ) -> anyhow::Result<()> { debug!(warmup_info=?warmup_info); let warm_up_terms_future = warm_up_terms(searcher, &warmup_info.terms_grouped_by_field) @@ -239,6 +241,7 @@ pub(crate) async fn warmup( searcher, &warmup_info.automatons_grouped_by_field, cost_class, + query_id, ) .instrument(debug_span!("warm_up_automatons")); @@ -339,11 +342,12 @@ async fn warm_up_automatons( searcher: &Searcher, terms_grouped_by_field: &HashMap>, cost_class: QueryCostClass, + query_id: QueryId, ) -> anyhow::Result<()> { let mut warm_up_futures = Vec::new(); let cpu_intensive_executor = |task| async { crate::search_thread_pool() - .run_cpu_intensive_with_extra_tags(task, "automaton_warmup", cost_class.as_label()) + .run_cpu_intensive_fair(task, query_id, "automaton_warmup", cost_class.as_label()) .await .map_err(|_| std::io::Error::other("task panicked"))? }; @@ -462,9 +466,8 @@ async fn leaf_search_single_split( split: SplitIdAndFooterOffsets, aggregations_limits: AggregationLimitsGuard, search_permit: &mut SearchPermit, + mut leaf_search_state_guard: SplitSearchStateGuard, ) -> crate::Result> { - let mut leaf_search_state_guard = - SplitSearchStateGuard::new(ctx.split_outcome_counters.clone()); rewrite_request( &mut search_request, &split, @@ -557,7 +560,13 @@ async fn leaf_search_single_split( let warmup_start = Instant::now(); leaf_search_state_guard.set_state(SplitSearchState::WarmUp); - warmup(&searcher, &warmup_info, ctx.cost_class).await?; + warmup( + &searcher, + &warmup_info, + ctx.cost_class, + leaf_search_state_guard.scheduler_guard.query_id(), + ) + .await?; let warmup_end = Instant::now(); let warmup_duration: Duration = warmup_end.duration_since(warmup_start); let warmup_size = ByteSize(byte_range_cache.get_num_bytes()); @@ -581,6 +590,8 @@ async fn leaf_search_single_split( let ctx_clone = ctx.clone(); + let query_id = leaf_search_state_guard.scheduler_guard.query_id(); + leaf_search_state_guard.set_state(SplitSearchState::CpuQueue); let cpu_task = move || { leaf_search_state_guard.set_state(SplitSearchState::Cpu); @@ -623,7 +634,12 @@ async fn leaf_search_single_split( }; let search_request_and_result: Option<(SearchRequest, LeafSearchResponse)> = crate::search_thread_pool() - .run_cpu_intensive_with_extra_tags(cpu_task, "split_search", ctx.cost_class.as_label()) + .run_cpu_intensive_fair( + cpu_task, + query_id, + "split_search", + ctx.cost_class.as_label(), + ) .await .map_err(|_| { crate::SearchError::Internal(format!("leaf search panicked. split={split_id}")) @@ -1317,7 +1333,7 @@ pub async fn multi_index_leaf_search( } crate::search_thread_pool() - .run_cpu_intensive_with_extra_tags( + .run_cpu_intensive( || incremental_merge_collector.finalize().map_err(Into::into), "finalize", cost_class.as_label(), @@ -1395,6 +1411,14 @@ pub async fn single_doc_mapping_leaf_search( let num_splits = splits.len(); info!(num_docs, num_splits, split_offsets = ?PrettySample::new(&splits, 5)); + let query_id = QueryId::next(); + let scheduler_split_guards = crate::search_thread_pool().register_query(query_id, num_splits); + let split_outcome_counters = Arc::new(SplitSearchOutcomeCounters::new_unregistered()); + let leaf_search_state_guards: Vec = scheduler_split_guards + .into_iter() + .map(|split_guard| SplitSearchStateGuard::new(split_outcome_counters.clone(), split_guard)) + .collect(); + let split_filter = CanSplitDoBetter::from_request(&request, doc_mapper.timestamp_field_name()); let split_with_req = split_filter.optimize(request.clone(), splits)?; @@ -1421,25 +1445,30 @@ pub async fn single_doc_mapping_leaf_search( let leaf_search_context = Arc::new(LeafSearchContext { searcher_context: searcher_context.clone(), - split_outcome_counters: Arc::new(SplitSearchOutcomeCounters::new_unregistered()), incremental_merge_collector: incremental_merge_collector.clone(), doc_mapper: doc_mapper.clone(), split_filter: split_filter.clone(), cost_class, }); + let total_permits = permit_futures.len(); let mut join_set = JoinSet::new(); let mut split_with_task_id = Vec::with_capacity(split_with_req.len()); - for ((split, search_request), permit_fut) in split_with_req.into_iter().zip(permit_futures) { + for (index, (((split, search_request), permit_fut), mut leaf_search_state_guard)) in + split_with_req + .into_iter() + .zip(permit_futures) + .zip(leaf_search_state_guards) + .enumerate() + { let leaf_split_search_permit = permit_fut .instrument(info_span!("waiting_for_leaf_search_split_semaphore")) .await; + crate::search_thread_pool().set_waiting_for_permit(query_id, total_permits - (index + 1)); let Some(simplified_search_request) = simplify_search_request(search_request, &split, &split_filter) else { - let mut leaf_search_state_guard = - SplitSearchStateGuard::new(leaf_search_context.split_outcome_counters.clone()); leaf_search_state_guard.set_state(SplitSearchState::PrunedBeforeWarmup); continue; }; @@ -1452,6 +1481,7 @@ pub async fn single_doc_mapping_leaf_search( split, leaf_split_search_permit, aggregations_limits.clone(), + leaf_search_state_guard, ) .in_current_span(), ); @@ -1499,7 +1529,7 @@ pub async fn single_doc_mapping_leaf_search( let leaf_search_response_reresult: Result, _> = crate::search_thread_pool() - .run_cpu_intensive_with_extra_tags( + .run_cpu_intensive( || incremental_merge_collector.finalize(), "finalize", cost_class.as_label(), @@ -1509,11 +1539,7 @@ pub async fn single_doc_mapping_leaf_search( .context("failed to merge split search responses"); let mut leaf_response = leaf_search_response_reresult??; - leaf_response.splits_by_outcome = Some( - leaf_search_context - .split_outcome_counters - .split_by_outcome(), - ); + leaf_response.splits_by_outcome = Some(split_outcome_counters.split_by_outcome()); Ok(leaf_response) } @@ -1551,19 +1577,25 @@ impl Drop for SplitSearchStateGuard { self.state .inc(&crate::metrics::SEARCH_METRICS.split_search_outcome_total); self.state.inc(&self.local_split_search_outcome_counters); + // Resolving the split with the scheduler happens as `_split_guard` drops. } } struct SplitSearchStateGuard { state: SplitSearchState, local_split_search_outcome_counters: Arc, + scheduler_guard: SchedulerSplitGuard, } impl SplitSearchStateGuard { - pub fn new(local_split_search_outcome_counters: Arc) -> Self { + pub fn new( + local_split_search_outcome_counters: Arc, + scheduler_guard: SchedulerSplitGuard, + ) -> Self { SplitSearchStateGuard { state: SplitSearchState::Start, - local_split_search_outcome_counters: local_split_search_outcome_counters.clone(), + local_split_search_outcome_counters, + scheduler_guard, } } @@ -1574,7 +1606,6 @@ impl SplitSearchStateGuard { struct LeafSearchContext { searcher_context: Arc, - split_outcome_counters: Arc, incremental_merge_collector: Arc>, doc_mapper: Arc, split_filter: Arc>, @@ -1590,6 +1621,7 @@ async fn leaf_search_single_split_wrapper( split: SplitIdAndFooterOffsets, mut search_permit: SearchPermit, aggregations_limits: AggregationLimitsGuard, + leaf_search_state_guard: SplitSearchStateGuard, ) { let timer = crate::SEARCH_METRICS .leaf_search_split_duration_secs @@ -1602,6 +1634,7 @@ async fn leaf_search_single_split_wrapper( split.clone(), aggregations_limits, &mut search_permit, + leaf_search_state_guard, ) .await; @@ -2234,6 +2267,9 @@ mod tests { index_writer.commit().unwrap(); let searcher = index.reader().unwrap().searcher(); + let query_id = QueryId::next(); + let _split_guards = crate::search_thread_pool().register_query(query_id, 1); + // Several valid regexes targeting the same field combine into a single // automaton and warm up successfully. let valid: HashMap> = std::iter::once(( @@ -2245,7 +2281,7 @@ mod tests { )) .collect(); assert!( - warm_up_automatons(&searcher, &valid, QueryCostClass::Regular) + warm_up_automatons(&searcher, &valid, QueryCostClass::Regular, query_id) .await .is_ok() ); @@ -2262,7 +2298,7 @@ mod tests { )) .collect(); assert!( - warm_up_automatons(&searcher, &valid_json, QueryCostClass::Regular) + warm_up_automatons(&searcher, &valid_json, QueryCostClass::Regular, query_id) .await .is_ok() ); @@ -2274,7 +2310,7 @@ mod tests { HashSet::from([Automaton::Regex(None, vec!["(".to_string()])]), )) .collect(); - let error = warm_up_automatons(&searcher, &invalid, QueryCostClass::Regular) + let error = warm_up_automatons(&searcher, &invalid, QueryCostClass::Regular, query_id) .await .unwrap_err() .to_string(); @@ -2289,7 +2325,7 @@ mod tests { HashSet::from([Automaton::Regex(Some(json_path), vec!["(".to_string()])]), )) .collect(); - let error = warm_up_automatons(&searcher, &invalid_json, QueryCostClass::Regular) + let error = warm_up_automatons(&searcher, &invalid_json, QueryCostClass::Regular, query_id) .await .unwrap_err() .to_string(); diff --git a/quickwit/quickwit-search/src/lib.rs b/quickwit/quickwit-search/src/lib.rs index bc35a6fe8a1..a741ecb14f9 100644 --- a/quickwit/quickwit-search/src/lib.rs +++ b/quickwit/quickwit-search/src/lib.rs @@ -49,7 +49,7 @@ mod tests; pub use collector::QuickwitAggregations; use metrics::SEARCH_METRICS; -use quickwit_common::thread_pool::ThreadPool; +use quickwit_common::thread_pool::SearchThreadPool; use quickwit_common::tower::Pool; use quickwit_doc_mapper::DocMapper; use quickwit_proto::metastore::{ @@ -132,11 +132,11 @@ fn compute_search_thread_pool_num_threads() -> Option { Some(threads) } -fn search_thread_pool() -> &'static ThreadPool { - static SEARCH_THREAD_POOL: OnceLock = OnceLock::new(); +fn search_thread_pool() -> &'static SearchThreadPool { + static SEARCH_THREAD_POOL: OnceLock = OnceLock::new(); SEARCH_THREAD_POOL - .get_or_init(|| ThreadPool::new("search", compute_search_thread_pool_num_threads())) + .get_or_init(|| SearchThreadPool::new("search", compute_search_thread_pool_num_threads())) } #[cfg(test)] diff --git a/quickwit/quickwit-search/src/list_fields.rs b/quickwit/quickwit-search/src/list_fields.rs index 59b0b66a8d5..2bf5e98e06d 100644 --- a/quickwit/quickwit-search/src/list_fields.rs +++ b/quickwit/quickwit-search/src/list_fields.rs @@ -370,7 +370,7 @@ pub async fn leaf_list_fields( merge_leaf_list_fields(filtered_list_fields_sorted_iters) }; let fields = search_thread_pool() - .run_cpu_intensive_with_extra_tags( + .run_cpu_intensive( cpu_task, "leaf_list_fields", QueryCostClass::Regular.as_label(), @@ -448,7 +448,7 @@ 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_with_extra_tags( + .run_cpu_intensive( move || { let leaf_list_fields = leaf_list_fields_protos .into_iter() diff --git a/quickwit/quickwit-search/src/root.rs b/quickwit/quickwit-search/src/root.rs index e35bf87f26a..fc4a580e721 100644 --- a/quickwit/quickwit-search/src/root.rs +++ b/quickwit/quickwit-search/src/root.rs @@ -801,7 +801,7 @@ pub(crate) async fn search_partial_hits_phase( let cost_class = query_cost_classifier::classify_serialized(&search_request.query_ast); let span = info_span!("merge_fruits"); let mut leaf_search_response = crate::search_thread_pool() - .run_cpu_intensive_with_extra_tags( + .run_cpu_intensive( move || { let _span_guard = span.enter(); merge_collector.merge_fruits(leaf_search_results) From 5428c388d58adc4d8ee2c800e9ea3bb97a75b14d Mon Sep 17 00:00:00 2001 From: Remi Dettai Date: Fri, 4 Sep 2026 09:09:14 +0200 Subject: [PATCH 2/7] Add warmup queue state to split --- quickwit/quickwit-proto/protos/quickwit/search.proto | 2 ++ .../quickwit-proto/src/codegen/quickwit/quickwit.search.rs | 3 +++ quickwit/quickwit-search/src/leaf.rs | 5 ++++- quickwit/quickwit-search/src/metrics.rs | 6 ++++++ quickwit/quickwit-search/src/metrics_trackers.rs | 3 +++ .../quickwit-serve/src/elasticsearch_api/rest_handler.rs | 1 + 6 files changed, 19 insertions(+), 1 deletion(-) diff --git a/quickwit/quickwit-proto/protos/quickwit/search.proto b/quickwit/quickwit-proto/protos/quickwit/search.proto index 3244f227c6f..60d56479b3a 100644 --- a/quickwit/quickwit-proto/protos/quickwit/search.proto +++ b/quickwit/quickwit-proto/protos/quickwit/search.proto @@ -380,6 +380,8 @@ message SplitsByOutcome { uint64 cancel_cpu_queue = 8; // Cancelled during CPU processing (error or timeout) uint64 cancel_cpu = 9; + // Cancelled while waiting for a search permit + uint64 cancel_warmup_queue = 10; } message ResourceStats { diff --git a/quickwit/quickwit-proto/src/codegen/quickwit/quickwit.search.rs b/quickwit/quickwit-proto/src/codegen/quickwit/quickwit.search.rs index c69a5b597a5..0d75065472e 100644 --- a/quickwit/quickwit-proto/src/codegen/quickwit/quickwit.search.rs +++ b/quickwit/quickwit-proto/src/codegen/quickwit/quickwit.search.rs @@ -309,6 +309,9 @@ pub struct SplitsByOutcome { /// Cancelled during CPU processing (error or timeout) #[prost(uint64, tag = "9")] pub cancel_cpu: u64, + /// Cancelled while waiting for a search permit + #[prost(uint64, tag = "10")] + pub cancel_warmup_queue: u64, } #[derive(serde::Serialize, serde::Deserialize, utoipa::ToSchema)] #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] diff --git a/quickwit/quickwit-search/src/leaf.rs b/quickwit/quickwit-search/src/leaf.rs index 040259aea43..7bd35215e05 100644 --- a/quickwit/quickwit-search/src/leaf.rs +++ b/quickwit/quickwit-search/src/leaf.rs @@ -1464,6 +1464,7 @@ pub async fn single_doc_mapping_leaf_search( let leaf_split_search_permit = permit_fut .instrument(info_span!("waiting_for_leaf_search_split_semaphore")) .await; + leaf_search_state_guard.set_state(SplitSearchState::Start); crate::search_thread_pool().set_waiting_for_permit(query_id, total_permits - (index + 1)); let Some(simplified_search_request) = @@ -1545,6 +1546,7 @@ pub async fn single_doc_mapping_leaf_search( #[derive(Copy, Clone)] enum SplitSearchState { + WarmupQueue, Start, CacheHit, ProcessedFromMetadata, @@ -1559,6 +1561,7 @@ enum SplitSearchState { impl SplitSearchState { pub fn inc(self, counters: &SplitSearchOutcomeCounters) { match self { + SplitSearchState::WarmupQueue => counters.cancel_warmup_queue.inc(), SplitSearchState::Start => counters.cancel_before_warmup.inc(), SplitSearchState::CacheHit => counters.cache_hit.inc(), SplitSearchState::ProcessedFromMetadata => counters.processed_from_metadata.inc(), @@ -1593,7 +1596,7 @@ impl SplitSearchStateGuard { scheduler_guard: SchedulerSplitGuard, ) -> Self { SplitSearchStateGuard { - state: SplitSearchState::Start, + state: SplitSearchState::WarmupQueue, local_split_search_outcome_counters, scheduler_guard, } diff --git a/quickwit/quickwit-search/src/metrics.rs b/quickwit/quickwit-search/src/metrics.rs index c545f60f299..d85af3f14dc 100644 --- a/quickwit/quickwit-search/src/metrics.rs +++ b/quickwit/quickwit-search/src/metrics.rs @@ -43,6 +43,7 @@ fn print_if_not_null( /// /// Cancellation counters cover two scenarios: errors in splits and timeouts. pub struct SplitSearchOutcomeCounters { + pub cancel_warmup_queue: IntCounter, pub cancel_before_warmup: IntCounter, pub cache_hit: IntCounter, pub processed_from_metadata: IntCounter, @@ -56,6 +57,7 @@ pub struct SplitSearchOutcomeCounters { impl fmt::Display for SplitSearchOutcomeCounters { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + print_if_not_null("cancel_warmup_queue", &self.cancel_warmup_queue, f)?; print_if_not_null("cancel_before_warmup", &self.cancel_before_warmup, f)?; print_if_not_null("cache_hit", &self.cache_hit, f)?; print_if_not_null("processed_from_metadata", &self.processed_from_metadata, f)?; @@ -97,6 +99,8 @@ impl SplitSearchOutcomeCounters { pub fn new_from_counter_vec(search_split_outcome_vec: IntCounterVec<1>) -> Self { SplitSearchOutcomeCounters { + cancel_warmup_queue: search_split_outcome_vec + .with_label_values(["cancel_warmup_queue"]), cancel_before_warmup: search_split_outcome_vec .with_label_values(["cancel_before_warmup"]), cache_hit: search_split_outcome_vec.with_label_values(["cache_hit"]), @@ -119,6 +123,7 @@ impl SplitSearchOutcomeCounters { let Self { pruned_before_warmup, pruned_after_warmup, + cancel_warmup_queue, cancel_before_warmup, cancel_warmup, cancel_cpu_queue, @@ -130,6 +135,7 @@ impl SplitSearchOutcomeCounters { SplitsByOutcome { pruned_before_warmup: pruned_before_warmup.get(), pruned_after_warmup: pruned_after_warmup.get(), + cancel_warmup_queue: cancel_warmup_queue.get(), cancel_before_warmup: cancel_before_warmup.get(), cancel_warmup: cancel_warmup.get(), cancel_cpu_queue: cancel_cpu_queue.get(), diff --git a/quickwit/quickwit-search/src/metrics_trackers.rs b/quickwit/quickwit-search/src/metrics_trackers.rs index d04b0b70113..1655aff20ae 100644 --- a/quickwit/quickwit-search/src/metrics_trackers.rs +++ b/quickwit/quickwit-search/src/metrics_trackers.rs @@ -134,6 +134,7 @@ impl std::fmt::Display for SplitsByOutcomeDisp { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { // Destructure to make sure we update this if a state is added let SplitsByOutcome { + cancel_warmup_queue, pruned_before_warmup, pruned_after_warmup, cancel_before_warmup, @@ -146,6 +147,7 @@ impl std::fmt::Display for SplitsByOutcomeDisp { } = self.0; let mut sep = "{"; for (name, val) in [ + ("cancel_warmup_queue", cancel_warmup_queue), ("pruned_before_warmup", pruned_before_warmup), ("pruned_after_warmup", pruned_after_warmup), ("cancel_before_warmup", cancel_before_warmup), @@ -332,6 +334,7 @@ mod tests { fn test_splits_by_outcome_disp_all_fields() { assert_eq!( disp(SplitsByOutcome { + cancel_warmup_queue: 0, pruned_before_warmup: 1, pruned_after_warmup: 2, cancel_before_warmup: 3, diff --git a/quickwit/quickwit-serve/src/elasticsearch_api/rest_handler.rs b/quickwit/quickwit-serve/src/elasticsearch_api/rest_handler.rs index adf8f1df259..af8809cabac 100644 --- a/quickwit/quickwit-serve/src/elasticsearch_api/rest_handler.rs +++ b/quickwit/quickwit-serve/src/elasticsearch_api/rest_handler.rs @@ -1021,6 +1021,7 @@ fn get_relation_from_split_outcome( }; // Destructure to make sure we update this if a state is added. let SplitsByOutcome { + cancel_warmup_queue: _, cancel_before_warmup: _, cancel_warmup: _, cancel_cpu_queue: _, From 313fc601923a8ff9dd54bfc238e35481ea7c92c5 Mon Sep 17 00:00:00 2001 From: Remi Dettai Date: Fri, 4 Sep 2026 09:32:20 +0200 Subject: [PATCH 3/7] Avoid panic on uninitialized query --- .../src/thread_pool/scheduler.rs | 40 ++++++++++++------- 1 file changed, 26 insertions(+), 14 deletions(-) diff --git a/quickwit/quickwit-common/src/thread_pool/scheduler.rs b/quickwit/quickwit-common/src/thread_pool/scheduler.rs index 179d844abf0..4a952d12332 100644 --- a/quickwit/quickwit-common/src/thread_pool/scheduler.rs +++ b/quickwit/quickwit-common/src/thread_pool/scheduler.rs @@ -34,6 +34,7 @@ use once_cell::sync::Lazy; use tracing::error; use crate::metrics::{IntCounter, IntGauge, new_counter, new_gauge}; +use crate::rate_limited_error; /// Identifies a query (leaf search) whose split-processing tasks should be /// scheduled and fair-shared together. Must be unique among currently active @@ -192,19 +193,20 @@ impl Scheduler { guard } - /// Registers a new query with its total split count. Must be called - /// exactly once per query, before any [`Self::enqueue_fifo`], - /// [`Self::enqueue_fair`] or [`Self::set_waiting_for_permit`] call for - /// that `query_id`. + /// Registers a new query with its total split count. Must be called exactly + /// once per query, before any [`Self::enqueue_fair`] or + /// [`Self::set_waiting_for_permit`] call for that `query_id`. /// - /// Returns one guard per split, meant to be wrapped into that split's own - /// cleanup guard on the caller side, so the query's entry can never leak - /// regardless of how each split's processing ends. + /// Returns one guard per split to track the number of remaining splits for + /// the query. pub fn register_query( self: &Arc, query_id: QueryId, total_splits: usize, ) -> Vec { + if total_splits == 0 { + return Vec::new(); + } let mut state = self.lock_state(); state.queries.insert( query_id, @@ -264,8 +266,8 @@ impl Scheduler { } } - /// Schedules a task belonging to `query_id`. The query must already have - /// been [`Self::register_query`]-ed. + /// Schedules a task belonging to `query_id`. The query is expected to + /// already have been [`Self::register_query`]-ed. pub fn enqueue_fair(self: &Arc, query_id: QueryId, job: F) where F: FnOnce() + Send + 'static { let scheduler = self.clone(); @@ -274,11 +276,21 @@ impl Scheduler { scheduler.on_task_complete(query_id); }); let mut state = self.lock_state(); - let query = state - .queries - .get_mut(&query_id) - .expect("query must be registered before tasks are enqueued for it"); - query.ready.push_back(wrapped); + match state.queries.get_mut(&query_id) { + Some(query) => query.ready.push_back(wrapped), + None => { + debug_assert!( + false, + "query must be registered before tasks are enqueued for it" + ); + rate_limited_error!( + limit_per_min = 1, + ?query_id, + "query not registered on the scheduler, fall back to FIFO" + ); + state.high_priority_queue.push_back(wrapped); + } + } if state.active_pump_workers < self.num_threads { state.active_pump_workers += 1; let scheduler = self.clone(); From 46b784f0c6711b88afeceb1da8028f3514bba31a Mon Sep 17 00:00:00 2001 From: Remi Dettai Date: Fri, 4 Sep 2026 10:11:00 +0200 Subject: [PATCH 4/7] Update running count even when panicking --- .../src/thread_pool/regular_pool.rs | 4 +- .../src/thread_pool/scheduler.rs | 48 ++++++++++++++----- 2 files changed, 39 insertions(+), 13 deletions(-) diff --git a/quickwit/quickwit-common/src/thread_pool/regular_pool.rs b/quickwit/quickwit-common/src/thread_pool/regular_pool.rs index 35981845cbd..0c4dd77795e 100644 --- a/quickwit/quickwit-common/src/thread_pool/regular_pool.rs +++ b/quickwit/quickwit-common/src/thread_pool/regular_pool.rs @@ -22,8 +22,8 @@ use super::ThreadPoolTaskInstrumentation; /// An executor backed by a thread pool to run CPU-intensive tasks. /// -/// tokio::spawn_blocking should only used for IO-bound tasks, as it has not limit on its -/// thread count. +/// tokio::spawn_blocking should only be used for IO-bound tasks, as it has not +/// limit on its thread count. /// /// Unlike [`super::SearchThreadPool`], dispatches FIFO straight onto the raw /// rayon pool, with no per-query priority/fairness layer. diff --git a/quickwit/quickwit-common/src/thread_pool/scheduler.rs b/quickwit/quickwit-common/src/thread_pool/scheduler.rs index 4a952d12332..5fc6e2487ab 100644 --- a/quickwit/quickwit-common/src/thread_pool/scheduler.rs +++ b/quickwit/quickwit-common/src/thread_pool/scheduler.rs @@ -73,6 +73,22 @@ impl Drop for SchedulerSplitGuard { } } +/// Ensures the running count of the query is decremented after the job +/// completes even if the job itself panics. +struct RunningCountGuard { + scheduler: Arc, + query_id: QueryId, +} + +impl Drop for RunningCountGuard { + fn drop(&mut self) { + let mut state = self.scheduler.lock_state(); + if let Some(query) = state.queries.get_mut(&self.query_id) { + query.running_count = query.running_count.saturating_sub(1); + } + } +} + type Job = Box; /// How long a pump loop keeps grabbing tasks before handing its worker back @@ -238,9 +254,8 @@ impl Scheduler { } } - /// Called by [`SplitGuard::drop`] exactly once per split of `query_id`, - /// decrementing `remaining` and removing the query's entry once it hits - /// zero. + /// Should be called exactly once per split of `query_id` when we know for + /// sure that the split won't be submitted again to the fair scheduler. fn split_resolved(&self, query_id: QueryId) { let mut state = self.lock_state(); let Some(query) = state.queries.get_mut(&query_id) else { @@ -272,8 +287,11 @@ impl Scheduler { where F: FnOnce() + Send + 'static { let scheduler = self.clone(); let wrapped: Job = Box::new(move || { + let _running_guard = RunningCountGuard { + scheduler, + query_id, + }; job(); - scheduler.on_task_complete(query_id); }); let mut state = self.lock_state(); match state.queries.get_mut(&query_id) { @@ -297,13 +315,6 @@ impl Scheduler { self.rayon_pool.spawn(move || pump_loop(&scheduler)); } } - - fn on_task_complete(&self, query_id: QueryId) { - let mut state = self.lock_state(); - if let Some(query) = state.queries.get_mut(&query_id) { - query.running_count = query.running_count.saturating_sub(1); - } - } } /// Runs on a rayon worker thread: repeatedly picks and runs the current best @@ -425,6 +436,21 @@ mod tests { assert_eq!(*order.lock().unwrap(), vec!["level0", "query"]); } + #[test] + fn test_on_task_complete_runs_even_if_job_panics() { + let scheduler = test_scheduler(1); + let _guards = scheduler.register_query(QueryId(1), 1); + + scheduler.enqueue_fair(QueryId(1), || panic!("boom")); + wait_until(|| { + let state = scheduler.state.lock().unwrap(); + match state.queries.get(&QueryId(1)) { + Some(query) => query.running_count == 0, + None => false, + } + }); + } + #[test] fn test_smaller_remaining_runs_first() { let scheduler = test_scheduler(1); From ae87675d4a6e53dc8992ca181427d33d0bdc8721 Mon Sep 17 00:00:00 2001 From: Remi Dettai Date: Fri, 4 Sep 2026 10:23:46 +0200 Subject: [PATCH 5/7] Fix split outcome accumulation --- .../protos/quickwit/search.proto | 2 +- .../src/codegen/quickwit/quickwit.search.rs | 2 +- quickwit/quickwit-search/src/lib.rs | 32 +++++++++++++------ 3 files changed, 25 insertions(+), 11 deletions(-) diff --git a/quickwit/quickwit-proto/protos/quickwit/search.proto b/quickwit/quickwit-proto/protos/quickwit/search.proto index 60d56479b3a..2ce92172555 100644 --- a/quickwit/quickwit-proto/protos/quickwit/search.proto +++ b/quickwit/quickwit-proto/protos/quickwit/search.proto @@ -368,7 +368,7 @@ message LeafSearchRequest { message SplitsByOutcome { uint64 pruned_before_warmup = 1; uint64 pruned_after_warmup = 2; - // Cancelled before warmup started (error or timeout) + // Cancelled while setting up the warmup (footer error or timeout) uint64 cancel_before_warmup = 3; uint64 processed = 4; uint64 processed_from_metadata = 5; diff --git a/quickwit/quickwit-proto/src/codegen/quickwit/quickwit.search.rs b/quickwit/quickwit-proto/src/codegen/quickwit/quickwit.search.rs index 0d75065472e..3b7e8a9dc23 100644 --- a/quickwit/quickwit-proto/src/codegen/quickwit/quickwit.search.rs +++ b/quickwit/quickwit-proto/src/codegen/quickwit/quickwit.search.rs @@ -290,7 +290,7 @@ pub struct SplitsByOutcome { pub pruned_before_warmup: u64, #[prost(uint64, tag = "2")] pub pruned_after_warmup: u64, - /// Cancelled before warmup started (error or timeout) + /// Cancelled while setting up the warmup (footer error or timeout) #[prost(uint64, tag = "3")] pub cancel_before_warmup: u64, #[prost(uint64, tag = "4")] diff --git a/quickwit/quickwit-search/src/lib.rs b/quickwit/quickwit-search/src/lib.rs index a741ecb14f9..82ae3e0cbf7 100644 --- a/quickwit/quickwit-search/src/lib.rs +++ b/quickwit/quickwit-search/src/lib.rs @@ -511,16 +511,30 @@ pub(crate) fn merge_splits_by_outcome( acc_opt: &mut Option, ) { if let Some(new) = new_opt { + // Destructure to ensure all fields are accounted for. + let SplitsByOutcome { + cancel_warmup_queue, + pruned_before_warmup, + pruned_after_warmup, + cancel_before_warmup, + cancel_warmup, + cancel_cpu_queue, + cancel_cpu, + processed, + processed_from_metadata, + cache_hit, + } = new; if let Some(acc) = acc_opt { - acc.pruned_before_warmup += new.pruned_before_warmup; - acc.pruned_after_warmup += new.pruned_after_warmup; - acc.cancel_before_warmup += new.cancel_before_warmup; - acc.cancel_warmup += new.cancel_warmup; - acc.cancel_cpu_queue += new.cancel_cpu_queue; - acc.cancel_cpu += new.cancel_cpu; - acc.processed += new.processed; - acc.processed_from_metadata += new.processed_from_metadata; - acc.cache_hit += new.cache_hit; + acc.cancel_warmup_queue += cancel_warmup_queue; + acc.pruned_before_warmup += pruned_before_warmup; + acc.pruned_after_warmup += pruned_after_warmup; + acc.cancel_before_warmup += cancel_before_warmup; + acc.cancel_warmup += cancel_warmup; + acc.cancel_cpu_queue += cancel_cpu_queue; + acc.cancel_cpu += cancel_cpu; + acc.processed += processed; + acc.processed_from_metadata += processed_from_metadata; + acc.cache_hit += cache_hit; } else { *acc_opt = Some(new); } From 4d6fb265c84c1265ddda842f4b828277aa5c1362 Mon Sep 17 00:00:00 2001 From: Remi Dettai Date: Fri, 4 Sep 2026 13:06:38 +0200 Subject: [PATCH 6/7] Fix pump yielding --- .../src/thread_pool/scheduler.rs | 42 ++++++++++++------- 1 file changed, 28 insertions(+), 14 deletions(-) diff --git a/quickwit/quickwit-common/src/thread_pool/scheduler.rs b/quickwit/quickwit-common/src/thread_pool/scheduler.rs index 5fc6e2487ab..f552eecc6d5 100644 --- a/quickwit/quickwit-common/src/thread_pool/scheduler.rs +++ b/quickwit/quickwit-common/src/thread_pool/scheduler.rs @@ -19,11 +19,14 @@ //! queue and uses rayon only to own OS threads that continuously drain it (the //! "pump loop" pattern). //! -//! Two tiers of priority exist: +//! Three tiers of priority exist: //! - High priority: always runs first and processed in strict FIFO order. Meant for short and rarer //! tasks such as merging/finalizing a query's results. -//! - Per-query tasks: tries to be fair among queries, with a bias towards queries that are closer -//! to completion. +//! - Per-query: tries to be fair among queries, with a bias towards queries that are closer to +//! completion. +//! - External: tasks submitted to the rayon threadpool without going through the scheduler are +//! executed before all other tasks, but with some latency because they need the pump loops to +//! yield to be picked up. use std::collections::{HashMap, VecDeque}; use std::sync::atomic::{AtomicU64, Ordering}; @@ -91,9 +94,12 @@ impl Drop for RunningCountGuard { type Job = Box; -/// How long a pump loop keeps grabbing tasks before handing its worker back -/// to rayon's own scheduler (see [`pump_loop`]). -const PUMP_LOOP_YIELD_INTERVAL: Duration = Duration::from_millis(200); +/// How long a pump loop keeps grabbing tasks before handing its worker back to +/// rayon's own scheduler (see [`pump_loop`]). +/// +/// Setting this too low adds un-necessary context work, setting this too high +/// adds latency to tasks submitted directly to the rayon threadpool. +const PUMP_LOOP_YIELD_INTERVAL: Duration = Duration::from_millis(100); /// Per-query scheduling state. struct QueryState { @@ -329,7 +335,7 @@ impl Scheduler { /// Tantivy Executor), so the pump loop needs to periodically yield to let rayon /// schedule those tasks. fn pump_loop(scheduler: &Arc) { - let yield_deadline = Instant::now() + PUMP_LOOP_YIELD_INTERVAL; + let mut yield_deadline = Instant::now() + PUMP_LOOP_YIELD_INTERVAL; loop { let job = { let mut state = scheduler.lock_state(); @@ -345,9 +351,18 @@ fn pump_loop(scheduler: &Arc) { error!("task running in the thread pool scheduler panicked"); } if Instant::now() >= yield_deadline { - let replacement = scheduler.clone(); - scheduler.rayon_pool.spawn(move || pump_loop(&replacement)); - return; + // Drain all currently pending externally-submitted work. + loop { + match std::panic::catch_unwind(rayon::yield_now) { + Ok(Some(rayon::Yield::Executed)) => continue, + Ok(_) => break, + Err(_) => { + error!("externally-submitted rayon task panicked while yielding"); + break; + } + } + } + yield_deadline = Instant::now() + PUMP_LOOP_YIELD_INTERVAL; } } } @@ -572,8 +587,7 @@ mod tests { // rayon pool outside the scheduler (e.g. Tantivy's own internal // parallelism via `ThreadPool::get_executor`), which only gets // picked up by a worker that actually returns to rayon's scheduling - // loop. No core is ever permanently sacrificed for this: pump loops - // just periodically hand their worker back and re-queue themselves. + // loop. let scheduler = test_scheduler(2); let _guards = scheduler.register_query(QueryId(1), 100_000); @@ -581,14 +595,14 @@ mod tests { // tasks, well over one yield interval in total. for _ in 0..2_000 { scheduler.enqueue_fair(QueryId(1), || { - std::thread::sleep(Duration::from_millis(1)); + std::thread::sleep(Duration::from_millis(10)); }); } wait_until(|| scheduler.state.lock().unwrap().active_pump_workers >= 1); let (tx, rx) = std::sync::mpsc::channel(); scheduler.rayon_pool.spawn(move || tx.send(()).unwrap()); - rx.recv_timeout(Duration::from_secs(2)) + rx.recv_timeout(Duration::from_secs(1)) .expect("directly-injected rayon work starved: pump loops never yielded"); } From ea5066d9600d0e9d9d1e3310cbabd6a5ae1c824a Mon Sep 17 00:00:00 2001 From: Remi Dettai Date: Fri, 4 Sep 2026 15:04:40 +0200 Subject: [PATCH 7/7] Add missing tests --- .../src/thread_pool/regular_pool.rs | 2 +- .../src/thread_pool/scheduler.rs | 109 ++++++++++-------- .../src/thread_pool/search_pool.rs | 41 +++++++ 3 files changed, 104 insertions(+), 48 deletions(-) diff --git a/quickwit/quickwit-common/src/thread_pool/regular_pool.rs b/quickwit/quickwit-common/src/thread_pool/regular_pool.rs index 0c4dd77795e..c187ff71d8d 100644 --- a/quickwit/quickwit-common/src/thread_pool/regular_pool.rs +++ b/quickwit/quickwit-common/src/thread_pool/regular_pool.rs @@ -38,7 +38,7 @@ impl ThreadPool { let mut rayon_pool_builder = rayon::ThreadPoolBuilder::new() .thread_name(move |thread_id| format!("quickwit-{name}-{thread_id}")) .panic_handler(move |_my_panic| { - error!("task running in the quickwit {name} thread pool panicked"); + error!(pool_name = name, "task running in the thread pool panicked"); }); if let Some(num_threads) = num_threads_opt { rayon_pool_builder = rayon_pool_builder.num_threads(num_threads); diff --git a/quickwit/quickwit-common/src/thread_pool/scheduler.rs b/quickwit/quickwit-common/src/thread_pool/scheduler.rs index f552eecc6d5..6e8162464ce 100644 --- a/quickwit/quickwit-common/src/thread_pool/scheduler.rs +++ b/quickwit/quickwit-common/src/thread_pool/scheduler.rs @@ -351,17 +351,8 @@ fn pump_loop(scheduler: &Arc) { error!("task running in the thread pool scheduler panicked"); } if Instant::now() >= yield_deadline { - // Drain all currently pending externally-submitted work. - loop { - match std::panic::catch_unwind(rayon::yield_now) { - Ok(Some(rayon::Yield::Executed)) => continue, - Ok(_) => break, - Err(_) => { - error!("externally-submitted rayon task panicked while yielding"); - break; - } - } - } + // Drain all currently pending externally-submitted work + while rayon::yield_now() == Some(rayon::Yield::Executed) {} yield_deadline = Instant::now() + PUMP_LOOP_YIELD_INTERVAL; } } @@ -432,13 +423,15 @@ mod tests { let order: Arc>> = Arc::new(StdMutex::new(Vec::new())); let _guards = scheduler.register_query(QueryId(1), 1); - // Block the single worker so both tasks below are enqueued before either runs. + + // Step 1: Block the single worker so both tasks from step 2 stay in the + // queue. let (release_tx, release_rx) = std::sync::mpsc::channel::<()>(); scheduler.enqueue_fair(QueryId(1), move || { release_rx.recv().unwrap(); }); - wait_until(|| scheduler.state.lock().unwrap().active_pump_workers == 1); + // Step 2: Add two tasks that remain queued by the scheduler let order_clone = order.clone(); scheduler.enqueue_fair(QueryId(1), move || { order_clone.lock().unwrap().push("query") @@ -446,6 +439,8 @@ mod tests { let order_clone = order.clone(); scheduler.enqueue_fifo(move || order_clone.lock().unwrap().push("level0")); + // Step 3: Release the blocked worker to validate that the priority + // queue is picked up first release_tx.send(()).unwrap(); wait_until(|| order.lock().unwrap().len() == 2); assert_eq!(*order.lock().unwrap(), vec!["level0", "query"]); @@ -457,12 +452,9 @@ mod tests { let _guards = scheduler.register_query(QueryId(1), 1); scheduler.enqueue_fair(QueryId(1), || panic!("boom")); - wait_until(|| { - let state = scheduler.state.lock().unwrap(); - match state.queries.get(&QueryId(1)) { - Some(query) => query.running_count == 0, - None => false, - } + wait_until(|| match scheduler.lock_state().queries.get(&QueryId(1)) { + Some(query) => query.running_count == 0, + None => false, }); } @@ -474,12 +466,14 @@ mod tests { let _guards1 = scheduler.register_query(QueryId(1), 100); let _guards2 = scheduler.register_query(QueryId(2), 2); + // Step 1: Block the single worker so both tasks from step 2 stay in the + // queue. let (release_tx, release_rx) = std::sync::mpsc::channel::<()>(); scheduler.enqueue_fair(QueryId(1), move || { release_rx.recv().unwrap(); }); - wait_until(|| scheduler.state.lock().unwrap().active_pump_workers == 1); + // Step 2: Add two tasks that remain queued by the scheduler let order_clone = order.clone(); scheduler.enqueue_fair(QueryId(1), move || { order_clone.lock().unwrap().push(QueryId(1)) @@ -489,9 +483,10 @@ mod tests { order_clone.lock().unwrap().push(QueryId(2)) }); + // Step 3: Release the blocked worker to validate that query 2 with + // fewer remaining splits is picked up first. release_tx.send(()).unwrap(); wait_until(|| order.lock().unwrap().len() == 2); - // query 2 has far fewer remaining splits, so it should be picked first. assert_eq!(*order.lock().unwrap(), vec![QueryId(2), QueryId(1)]); } @@ -598,7 +593,7 @@ mod tests { std::thread::sleep(Duration::from_millis(10)); }); } - wait_until(|| scheduler.state.lock().unwrap().active_pump_workers >= 1); + wait_until(|| scheduler.lock_state().active_pump_workers >= 1); let (tx, rx) = std::sync::mpsc::channel(); scheduler.rayon_pool.spawn(move || tx.send(()).unwrap()); @@ -610,34 +605,13 @@ mod tests { fn test_query_state_cleaned_up_once_remaining_reaches_zero() { let scheduler = test_scheduler(2); let mut guards = scheduler.register_query(QueryId(1), 2); - assert!( - scheduler - .state - .lock() - .unwrap() - .queries - .contains_key(&QueryId(1)) - ); + assert!(scheduler.lock_state().queries.contains_key(&QueryId(1))); drop(guards.pop().unwrap()); - assert!( - scheduler - .state - .lock() - .unwrap() - .queries - .contains_key(&QueryId(1)) - ); + assert!(scheduler.lock_state().queries.contains_key(&QueryId(1))); drop(guards.pop().unwrap()); - assert!( - !scheduler - .state - .lock() - .unwrap() - .queries - .contains_key(&QueryId(1)) - ); + assert!(!scheduler.lock_state().queries.contains_key(&QueryId(1))); } #[test] @@ -652,6 +626,47 @@ mod tests { }); } wait_until(|| ran.load(Ordering::SeqCst) == 3); - wait_until(|| scheduler.state.lock().unwrap().active_pump_workers == 0); + wait_until(|| scheduler.lock_state().active_pump_workers == 0); + } + + #[test] + fn test_more_waiting_for_permit_runs_last() { + let scheduler = test_scheduler(1); + let order: Arc>> = Arc::new(StdMutex::new(Vec::new())); + + let _guards1 = scheduler.register_query(QueryId(1), 20); + let _guards2 = scheduler.register_query(QueryId(2), 10); + scheduler.set_waiting_for_permit(QueryId(1), 5); + + // Step 1: Block the single worker so both tasks from step 2 stay in + // the queue. + let (release_tx, release_rx) = std::sync::mpsc::channel::<()>(); + scheduler.enqueue_fair(QueryId(1), move || { + release_rx.recv().unwrap(); + }); + + // Step 2: Add one task per query that remain in the queue. + let order_clone = order.clone(); + scheduler.enqueue_fair(QueryId(1), move || { + order_clone.lock().unwrap().push(QueryId(1)) + }); + let order_clone = order.clone(); + scheduler.enqueue_fair(QueryId(2), move || { + order_clone.lock().unwrap().push(QueryId(2)) + }); + + // Step 3: Release the blocked worker to validate that query 2, which + // has no splits waiting on a permit, is picked up before query 1. + release_tx.send(()).unwrap(); + wait_until(|| order.lock().unwrap().len() == 2); + assert_eq!(*order.lock().unwrap(), vec![QueryId(2), QueryId(1)]); + } + + #[test] + fn test_register_query_with_zero_splits_returns_no_guards() { + let scheduler = test_scheduler(1); + let guards = scheduler.register_query(QueryId(1), 0); + assert!(guards.is_empty()); + assert!(!scheduler.lock_state().queries.contains_key(&QueryId(1))); } } diff --git a/quickwit/quickwit-common/src/thread_pool/search_pool.rs b/quickwit/quickwit-common/src/thread_pool/search_pool.rs index fad860c6fe0..409917f2d53 100644 --- a/quickwit/quickwit-common/src/thread_pool/search_pool.rs +++ b/quickwit/quickwit-common/src/thread_pool/search_pool.rs @@ -117,3 +117,44 @@ impl SearchThreadPool { ) } } + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use super::*; + use crate::thread_pool::scheduler::QueryId; + + #[test] + fn test_externally_submitted_panic_during_yield_does_not_stop_pump_loop() { + crate::setup_logging_for_tests(); + let search_pool = SearchThreadPool::new("test", Some(1)); + let query_id = QueryId::next(); + let _guards = search_pool.register_query(query_id, 100_000); + + // Keep the worker continuously busy so externally-submitted work + // can only ever be serviced through a pump loop's periodic yield. + let mut futures = Vec::with_capacity(2_000); + for _ in 0..2_000 { + futures.push(search_pool.run_cpu_intensive_fair( + || std::thread::sleep(Duration::from_millis(10)), + query_id, + "test", + "test", + )); + } + + // Submitted directly to the raw rayon pool, bypassing the scheduler + // entirely. + search_pool.thread_pool.rayon_pool.spawn(|| panic!("boom")); + + // Confirm the unique pump loop kept yielding to external work afterward. + let (tx, rx) = std::sync::mpsc::channel(); + search_pool + .thread_pool + .rayon_pool + .spawn(move || tx.send(()).unwrap()); + rx.recv_timeout(Duration::from_secs(2)) + .expect("pump loop stopped yielding to external work after the panic"); + } +}