diff --git a/crates/hotblocks/Cargo.toml b/crates/hotblocks/Cargo.toml index ca23e6a..c3e89dd 100644 --- a/crates/hotblocks/Cargo.toml +++ b/crates/hotblocks/Cargo.toml @@ -50,5 +50,9 @@ sqd-hotblocks-harness = { path = "../hotblocks-harness" } tempfile = { workspace = true } tokio = { workspace = true, features = ["full", "test-util"] } +[[bench]] +name = "query_pool" +harness = false + [lints] workspace = true diff --git a/crates/hotblocks/benches/query_pool.rs b/crates/hotblocks/benches/query_pool.rs new file mode 100644 index 0000000..2c0937b --- /dev/null +++ b/crates/hotblocks/benches/query_pool.rs @@ -0,0 +1,392 @@ +//! Open-loop model of the query pool: how long light query tasks wait while a +//! rare heavy one runs, for each way to schedule them. +//! +//! Tasks arrive from a thread outside the pool, as `POOL.spawn` gets them from +//! tokio. Each splits its work like the legacy engine: `par_iter` over scans, +//! `par_iter` over columns and a `rayon::join` tree over pages, with a +//! sequential part around it. Work is a calibrated spin, so CPU contention +//! slows it as it would slow a real scan. +//! +//! Policies: +//! shared one pool of N threads (the default) +//! groups pools of N/G threads; a task goes to the group with the fewest tasks +//! thread-per-task N one-thread pools that take tasks from one FIFO queue +//! (`--query-thread-per-task`) +//! +//! A polars share sends that part of each task's parallel work to a separate +//! N-thread pool with `install`, as polars does with its global pool. With +//! `shared` that pool is the query pool itself. +//! +//! cargo bench -p sqd-hotblocks --bench query_pool # full sweep, ~25 min +//! SQD_BENCH_SECS=20 cargo bench -p sqd-hotblocks --bench query_pool +//! +//! Needs N free cores: on a busy machine the spin calibration and every +//! result drift. + +use std::{ + cell::Cell, + hint::black_box, + sync::{ + Arc, Mutex, OnceLock, + atomic::{AtomicUsize, Ordering}, + mpsc + }, + time::{Duration, Instant} +}; + +use rayon::{ThreadPool, ThreadPoolBuilder, prelude::*}; + +const THREADS: usize = 8; +const LIGHT_RATE: f64 = 280.0; +const HEAVY_RATE: f64 = 0.3; +/// CPU milliseconds of a light task by quantile, log-linear in between. +const LIGHT_CPU_MS: &[(f64, f64)] = &[ + (0.0, 0.3), + (0.05, 0.6), + (0.5, 1.3), + (0.75, 2.8), + (0.9, 11.0), + (0.95, 25.0), + (0.99, 150.0), + (0.999, 550.0), + (1.0, 1500.0) +]; +/// A heavy task's CPU milliseconds, log-uniform in this range. +const HEAVY_CPU_MS: (f64, f64) = (1000.0, 4000.0); + +static ITERS_PER_US: OnceLock = OnceLock::new(); + +fn spin_iters(n: u64) -> u64 { + let mut x = 1u64; + for _ in 0..n { + x = black_box(x.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407)); + } + x +} + +fn calibrate() -> f64 { + let n = 200_000_000; + spin_iters(n / 10); + let started = Instant::now(); + black_box(spin_iters(n)); + n as f64 / (started.elapsed().as_secs_f64() * 1e6) +} + +fn spin(us: f64) { + if us > 0.0 { + black_box(spin_iters((us * ITERS_PER_US.get().unwrap()) as u64)); + } +} + +#[derive(Clone, Copy)] +struct Shape { + seq_us: f64, + scans: usize, + columns: usize, + pages: usize, + page_us: f64, + polars_scans: usize +} + +impl Shape { + fn new(cpu_us: f64, seq_share: f64, scans: usize, columns: usize, page_us: f64, polars_share: f64) -> Self { + let par_us = cpu_us * (1.0 - seq_share); + let pages = ((par_us / (scans * columns) as f64 / page_us).round() as usize).max(1); + Shape { + seq_us: cpu_us * seq_share, + scans, + columns, + pages, + page_us: par_us / (scans * columns * pages) as f64, + polars_scans: (scans as f64 * polars_share).round() as usize + } + } + + fn execute(&self, polars: Option<&ThreadPool>) { + spin(self.seq_us / 2.0); + + let scan = |_| { + (0..self.columns) + .into_par_iter() + .for_each(|_| pages(self.pages, self.page_us)) + }; + (0..self.scans - self.polars_scans).into_par_iter().for_each(scan); + match polars { + Some(pool) => pool.install(|| (0..self.polars_scans).into_par_iter().for_each(scan)), + None => (0..self.polars_scans).into_par_iter().for_each(scan) + } + + spin(self.seq_us / 2.0); + } +} + +fn pages(n: usize, page_us: f64) { + if n <= 1 { + spin(page_us) + } else { + let mid = n / 2 + n % 2; + rayon::join(|| pages(mid, page_us), || pages(n - mid, page_us)); + } +} + +#[derive(Clone, Copy, Debug)] +enum Policy { + Shared, + Groups(usize), + ThreadPerTask +} + +type Job = Box; + +struct Group { + pool: ThreadPool, + in_flight: AtomicUsize +} + +struct Scheduler { + groups: Vec, + queue: Option>, + polars: Option, + next: AtomicUsize +} + +fn pool(threads: usize) -> ThreadPool { + ThreadPoolBuilder::new().num_threads(threads).build().unwrap() +} + +impl Scheduler { + fn new(policy: Policy) -> Self { + match policy { + Policy::Shared => Self::groups(1, THREADS, None), + Policy::Groups(g) => Self::groups(g, THREADS / g, Some(pool(THREADS))), + Policy::ThreadPerTask => { + let (sender, receiver) = mpsc::channel::(); + let receiver = Arc::new(Mutex::new(receiver)); + let mut scheduler = Self::groups(THREADS, 1, Some(pool(THREADS))); + for group in &scheduler.groups { + let receiver = receiver.clone(); + group.pool.spawn(move || { + loop { + let job = receiver.lock().unwrap().recv(); + let Ok(job) = job else { return }; + job(); + } + }); + } + scheduler.queue = Some(sender); + scheduler + } + } + } + + fn groups(n: usize, threads: usize, polars: Option) -> Self { + Self { + groups: (0..n) + .map(|_| Group { + pool: pool(threads), + in_flight: AtomicUsize::new(0) + }) + .collect(), + queue: None, + polars, + next: AtomicUsize::new(0) + } + } + + fn least_loaded(&self) -> usize { + let n = self.groups.len(); + let start = self.next.fetch_add(1, Ordering::Relaxed) % n; + (0..n) + .map(|i| (start + i) % n) + .min_by_key(|&i| self.groups[i].in_flight.load(Ordering::Relaxed)) + .unwrap() + } +} + +thread_local! { + static NESTED: Cell> = const { Cell::new(None) }; +} + +#[derive(Clone, Copy)] +struct Record { + heavy: bool, + spawned: Duration, + queue: Duration, + latency: Duration, + /// Run time without the tasks nested in it. + own: Duration +} + +struct Rng(u64); + +impl Rng { + fn next_f64(&mut self) -> f64 { + self.0 ^= self.0 << 13; + self.0 ^= self.0 >> 7; + self.0 ^= self.0 << 17; + (self.0 >> 11) as f64 / (1u64 << 53) as f64 + } + + fn exp(&mut self, rate: f64) -> f64 { + -(1.0 - self.next_f64()).ln() / rate + } +} + +fn light_cpu_ms(u: f64) -> f64 { + for w in LIGHT_CPU_MS.windows(2) { + let ((q0, v0), (q1, v1)) = (w[0], w[1]); + if u <= q1 { + let f = (u - q0) / (q1 - q0); + return (v0.ln() + f * (v1.ln() - v0.ln())).exp(); + } + } + LIGHT_CPU_MS.last().unwrap().1 +} + +struct Run { + policy: Policy, + heavy_rate: f64, + polars_share: f64 +} + +fn run(case: &Run, secs: f64, warmup: f64) -> Vec { + let scheduler = Arc::new(Scheduler::new(case.policy)); + let records = Arc::new(Mutex::new(Vec::new())); + let in_flight = Arc::new(AtomicUsize::new(0)); + let mut rng = Rng(0x9E3779B97F4A7C15); + let rate = LIGHT_RATE + case.heavy_rate; + let begin = Instant::now(); + let end = Duration::from_secs_f64(secs + warmup); + let mut next_at = Duration::ZERO; + + while next_at < end { + let now = begin.elapsed(); + if next_at > now { + let ahead = next_at - now; + if ahead > Duration::from_micros(1500) { + std::thread::sleep(ahead - Duration::from_micros(1000)); + } else { + std::thread::yield_now(); + } + continue; + } + + let heavy = rng.next_f64() < case.heavy_rate / rate; + let shape = if heavy { + let (lo, hi) = HEAVY_CPU_MS; + let ms = (lo.ln() + rng.next_f64() * (hi.ln() - lo.ln())).exp(); + Shape::new(ms * 1000.0, 0.1, 3, 10, 200.0, case.polars_share) + } else { + let ms = light_cpu_ms(rng.next_f64()); + Shape::new(ms * 1000.0, 0.5, 2, 4, 100.0, case.polars_share.min(0.5)) + }; + + let group = scheduler.least_loaded(); + scheduler.groups[group].in_flight.fetch_add(1, Ordering::Relaxed); + in_flight.fetch_add(1, Ordering::Relaxed); + let (scheduler2, records2, in_flight2) = (scheduler.clone(), records.clone(), in_flight.clone()); + let spawned = begin.elapsed(); + let spawned_at = Instant::now(); + let job: Job = Box::new(move || { + let started = Instant::now(); + let outer = NESTED.replace(Some(Duration::ZERO)); + shape.execute(scheduler2.polars.as_ref()); + let total = started.elapsed(); + let nested = NESTED.replace(outer.map(|n| n + total)).unwrap_or_default(); + + scheduler2.groups[group].in_flight.fetch_sub(1, Ordering::Relaxed); + records2.lock().unwrap().push(Record { + heavy, + spawned, + queue: started - spawned_at, + latency: started + total - spawned_at, + own: total.saturating_sub(nested) + }); + in_flight2.fetch_sub(1, Ordering::Relaxed); + }); + match &scheduler.queue { + Some(queue) => queue.send(job).unwrap(), + None => scheduler.groups[group].pool.spawn(job) + } + + next_at += Duration::from_secs_f64(rng.exp(rate)); + } + + while in_flight.load(Ordering::Relaxed) > 0 { + std::thread::sleep(Duration::from_millis(5)); + } + let warm = Duration::from_secs_f64(warmup); + let records = records.lock().unwrap(); + records.iter().copied().filter(|r| r.spawned >= warm).collect() +} + +struct Stats { + mean: f64, + p99: f64, + p999: f64, + max: f64 +} + +fn stats(records: &[Record], get: impl Fn(&Record) -> Duration) -> Stats { + let mut ms: Vec = records.iter().map(|r| get(r).as_secs_f64() * 1000.0).collect(); + ms.sort_by(|a, b| a.partial_cmp(b).unwrap()); + let at = |p: f64| ms[((ms.len() as f64 * p) as usize).min(ms.len() - 1)]; + Stats { + mean: ms.iter().sum::() / ms.len() as f64, + p99: at(0.99), + p999: at(0.999), + max: *ms.last().unwrap() + } +} + +fn main() { + let secs: f64 = std::env::var("SQD_BENCH_SECS").map_or(120.0, |s| s.parse().unwrap()); + let warmup = (secs / 24.0).max(1.0); + ITERS_PER_US.set(calibrate()).unwrap(); + + let cases = [ + (Policy::Shared, 0.0, 0.0), + (Policy::ThreadPerTask, 0.0, 0.0), + (Policy::Shared, HEAVY_RATE, 0.0), + (Policy::Groups(2), HEAVY_RATE, 0.0), + (Policy::Groups(8), HEAVY_RATE, 0.0), + (Policy::ThreadPerTask, HEAVY_RATE, 0.0), + (Policy::ThreadPerTask, HEAVY_RATE, 1.0 / 3.0), + (Policy::ThreadPerTask, HEAVY_RATE, 2.0 / 3.0) + ]; + + println!( + "query_pool: {THREADS} threads, {LIGHT_RATE} light tasks/s, {secs} s per run, {:.0} spin iterations/us", + ITERS_PER_US.get().unwrap() + ); + println!( + "| heavy/s | policy | polars share | light queue mean / p99 / p99.9 ms | light latency mean / p99 / p99.9 ms | heavy own mean ms | heavy latency mean / max ms |" + ); + println!("|---|---|---|---|---|---|---|"); + for (policy, heavy_rate, polars_share) in cases { + let case = Run { + policy, + heavy_rate, + polars_share + }; + let records = run(&case, secs, warmup); + let (heavy, light): (Vec, Vec) = records.into_iter().partition(|r| r.heavy); + + let queue = stats(&light, |r| r.queue); + let latency = stats(&light, |r| r.latency); + let (heavy_own, heavy_latency) = if heavy.is_empty() { + ("-".to_string(), "-".to_string()) + } else { + let own = stats(&heavy, |r| r.own); + let latency = stats(&heavy, |r| r.latency); + ( + format!("{:.0}", own.mean), + format!("{:.0} / {:.0}", latency.mean, latency.max) + ) + }; + println!( + "| {heavy_rate} | {policy:?} | {polars_share:.2} | {:.1} / {:.0} / {:.0} | {:.1} / {:.0} / {:.0} | {heavy_own} | {heavy_latency} |", + queue.mean, queue.p99, queue.p999, latency.mean, latency.p99, latency.p999 + ); + } +} diff --git a/crates/hotblocks/spec/13-interface-binding.md b/crates/hotblocks/spec/13-interface-binding.md index 240a153..05b7935 100644 --- a/crates/hotblocks/spec/13-interface-binding.md +++ b/crates/hotblocks/spec/13-interface-binding.md @@ -33,7 +33,7 @@ annotated with the relevant GAP. | METADATA | `GET /datasets/{id}/metadata` | start block, real-time flag, aliases | | GET-RETENTION | `GET /datasets/{id}/retention` | current policy JSON | | SET-RETENTION | `POST /datasets/{id}/retention` | policy JSON; only for `External` datasets, else `FORBIDDEN` (403) | -| GET-SETTINGS | `GET /settings` | runtime settings JSON (`{"shadow":{…},"slowTasks":{…}}`) | +| GET-SETTINGS | `GET /settings` | runtime settings JSON (`{"shadow":{…},"slowTasks":{…},"query":{…}}`) | | UPDATE-SETTINGS | `PATCH /settings` | JSON Merge Patch (RFC 7386); answers the settings after it. An unknown key, a removed key or an invalid value is a 400 and changes nothing. In memory and per process: a restart returns to the flags | | observability | `GET /metrics` (+ engine-diagnostic routes) | OB surface, text formats | | readiness | `GET /ready` | rotation gate (OB-8), distinct from the `/` liveness signal: 503 for the whole pre-drain grace window so the orchestrator withdraws the endpoint before anything closes (LIV-12). Process-level only — per-dataset readability (LIV-5c) is still absent (GAP-7) | diff --git a/crates/hotblocks/src/cli.rs b/crates/hotblocks/src/cli.rs index 5ff7034..8bde07a 100644 --- a/crates/hotblocks/src/cli.rs +++ b/crates/hotblocks/src/cli.rs @@ -16,7 +16,7 @@ use crate::{ metrics::{DatasetMetricsCollector, RocksDbCollector}, query::{QueryService, QueryServiceRef}, query_engine::MetadataCatalogs, - settings::{RuntimeSettings, Settings, SettingsRef, ShadowSettings, SlowTaskSettings}, + settings::{QuerySettings, RuntimeSettings, Settings, SettingsRef, ShadowSettings, SlowTaskSettings}, shadow::{Shadow, ShadowRef}, types::DBRef }; @@ -200,7 +200,17 @@ pub struct CLI { /// Slow-task log lines per 5 minutes, at most one per query shape; 0 turns /// them off. #[arg(long, value_name = "N", default_value = "20")] - pub slow_task_log_limit: usize + pub slow_task_log_limit: usize, + + /// Run each query task on one of `--query-threads` single-thread workers + /// that take tasks from one queue, instead of the shared query pool. The + /// parallel work inside a task then stays on its thread, so a heavy query + /// holds one worker and not the whole pool. Polars keeps its own pool of the + /// same size. + /// + /// A starting value: `PATCH /settings` changes it until a restart. + #[arg(long)] + pub query_thread_per_task: bool } pub struct App { @@ -228,6 +238,9 @@ impl CLI { slow_tasks: SlowTaskSettings { threshold_ms: self.slow_task_threshold_ms, log_limit: self.slow_task_log_limit + }, + query: QuerySettings { + thread_per_task: self.query_thread_per_task } }) .map(Arc::new)?; @@ -287,7 +300,7 @@ impl CLI { if let Some(ms) = self.query_urgency { builder.set_urgency(ms); } - let service = builder.build(); + let service = builder.build()?; metrics_registry.register_collector(Box::new(service.metrics_collector())); Arc::new(service) diff --git a/crates/hotblocks/src/metrics.rs b/crates/hotblocks/src/metrics.rs index 9ee7358..d6cd071 100644 --- a/crates/hotblocks/src/metrics.rs +++ b/crates/hotblocks/src/metrics.rs @@ -327,8 +327,13 @@ static QUERY_TASK_QUEUE_DURATION: LazyLock> = static QUERY_TASK_RUN_DURATION: LazyLock> = LazyLock::new(|| Family::new_with_constructor(|| Histogram::new(exponential_buckets(0.0001, 2.0, 18)))); static QUERY_TASKS_QUEUED: LazyLock = LazyLock::new(Default::default); +static QUERY_THREAD_PER_TASK: LazyLock = LazyLock::new(Default::default); static QUERY_SLOW_TASKS: LazyLock> = LazyLock::new(Default::default); +pub(crate) fn set_query_thread_per_task(on: bool) { + QUERY_THREAD_PER_TASK.set(on as i64); +} + pub(crate) fn report_query_task_spawned() { QUERY_TASKS_QUEUED.inc(); } @@ -760,6 +765,11 @@ pub fn build_metrics_registry() -> Registry { of running tasks", QUERY_TASKS_QUEUED.clone() ); + registry.register( + "query_thread_per_task", + "1 when new query tasks run on a single thread each, 0 when they share the query pool", + QUERY_THREAD_PER_TASK.clone() + ); registry.register( "query_slow_tasks", "Query tasks that held a query pool worker longer than the slow-task threshold, not \ diff --git a/crates/hotblocks/src/query/executor.rs b/crates/hotblocks/src/query/executor.rs index 61e72de..902befb 100644 --- a/crates/hotblocks/src/query/executor.rs +++ b/crates/hotblocks/src/query/executor.rs @@ -3,17 +3,21 @@ use std::{ panic::{AssertUnwindSafe, catch_unwind}, sync::{ Arc, - atomic::{AtomicUsize, Ordering} + atomic::{AtomicUsize, Ordering}, + mpsc }, time::{Duration, Instant} }; +use parking_lot::Mutex; + use crate::{ errors::QueryTaskPanicked, metrics::{ COMPLETED_QUERIES, QueryStep, report_query_task_finished, report_query_task_spawned, report_query_task_started, report_query_too_many_tasks_error - } + }, + settings::SettingsRef }; thread_local! { @@ -36,30 +40,117 @@ impl TaskTime { } } +type Job = Box; + +/// Single-thread workers that take query tasks from one FIFO queue. Each is the +/// only thread of its own rayon pool, so the `par_iter` and `join` inside a task +/// run on the task's thread. In the shared pool they spread over every worker, +/// and a new task waits in the injector until a heavy query has no parallel work +/// left: rayon takes injected jobs last. +struct TaskQueue { + sender: mpsc::Sender, + _workers: Vec +} + +impl TaskQueue { + fn new(threads: usize) -> anyhow::Result { + let (sender, receiver) = mpsc::channel::(); + let receiver = Arc::new(Mutex::new(receiver)); + + let workers = (0..threads) + .map(|i| { + let worker = rayon::ThreadPoolBuilder::new() + .num_threads(1) + .thread_name(move |_| format!("query-task-{i}")) + .build()?; + let receiver = receiver.clone(); + worker.spawn(move || { + loop { + // The guard must drop before the job runs, or the workers run one at a time. + let job = receiver.lock().recv(); + let Ok(job) = job else { return }; + job(); + } + }); + Ok(worker) + }) + .collect::>()?; + + Ok(Self { + sender, + _workers: workers + }) + } +} + +/// Where a query task runs, chosen when its slot is taken. +#[derive(Clone)] +enum TaskPool { + Shared, + ThreadPerTask(Arc) +} + +impl TaskPool { + fn spawn(&self, job: impl FnOnce() + Send + 'static) { + match self { + Self::Shared => sqd_polars::POOL.spawn(job), + Self::ThreadPerTask(queue) => queue + .sender + .send(Box::new(job)) + .expect("the workers outlive the queue's sender") + } + } +} + #[derive(Clone)] pub struct QueryExecutor { // number of concurrent queries in_flight: Arc, // limit for concurrent queries max_pending_tasks: usize, - urgency: usize + urgency: usize, + settings: SettingsRef, + task_queue: Arc } impl QueryExecutor { - pub fn new(max_pending_tasks: usize, urgency: usize) -> Self { - Self { + pub fn new(max_pending_tasks: usize, urgency: usize, settings: SettingsRef) -> anyhow::Result { + Self::with_task_workers( + max_pending_tasks, + urgency, + settings, + sqd_polars::POOL.current_num_threads() + ) + } + + fn with_task_workers( + max_pending_tasks: usize, + urgency: usize, + settings: SettingsRef, + task_workers: usize + ) -> anyhow::Result { + let task_queue = TaskQueue::new(task_workers)?; + Ok(Self { in_flight: Arc::new(AtomicUsize::new(0)), max_pending_tasks, - urgency - } + urgency, + settings, + task_queue: Arc::new(task_queue) + }) } pub fn get_slot(&self) -> Option { let active_queries = self.in_flight.fetch_add(1, Ordering::SeqCst); if active_queries < self.max_pending_tasks { + let pool = if self.settings.query().thread_per_task { + TaskPool::ThreadPerTask(self.task_queue.clone()) + } else { + TaskPool::Shared + }; Some(QuerySlot { in_flight: self.in_flight.clone(), - urgency: self.urgency + urgency: self.urgency, + pool }) } else { self.in_flight.fetch_sub(1, Ordering::SeqCst); @@ -75,7 +166,8 @@ impl QueryExecutor { pub struct QuerySlot { in_flight: Arc, - urgency: usize + urgency: usize, + pool: TaskPool } impl Drop for QuerySlot { @@ -95,8 +187,8 @@ impl QuerySlot { time.min(100) } - /// Run `task` on the query pool, then `report` its result and time on the - /// same worker. A task that panics is not reported. + /// Run `task` on the query pool or a single-thread worker, then `report` its + /// result and time on the same thread. A task that panics is not reported. pub async fn run(self, step: QueryStep, task: F, report: P) -> Result where F: FnOnce(&Self) -> R + Send + 'static, @@ -107,7 +199,8 @@ impl QuerySlot { let spawned = Instant::now(); report_query_task_spawned(); - sqd_polars::POOL.spawn(move || { + let pool = self.pool.clone(); + pool.spawn(move || { let started = Instant::now(); report_query_task_started(step, started.duration_since(spawned)); let outer = NESTED.replace(Some(Duration::ZERO)); @@ -150,24 +243,107 @@ impl QueryExecutorCollector { #[cfg(test)] mod tests { + use std::{collections::HashSet, thread}; + + use rayon::prelude::*; + use super::*; + use crate::settings::{QuerySettings, RuntimeSettings, Settings, ShadowSettings, SlowTaskSettings}; + + fn executor(max_pending_tasks: usize, thread_per_task: bool) -> QueryExecutor { + let settings = RuntimeSettings::new(Settings { + shadow: ShadowSettings { + percent: 0.0, + slow_ratio: 2.0, + slow_margin_ms: 50, + log_limit: 0 + }, + slow_tasks: SlowTaskSettings { + threshold_ms: 500, + log_limit: 0 + }, + query: QuerySettings { thread_per_task } + }) + .unwrap(); + // A worker for every task a test admits, whatever the size of the polars pool. + QueryExecutor::with_task_workers(max_pending_tasks, 1, Arc::new(settings), max_pending_tasks).unwrap() + } #[tokio::test] async fn worker_panic_returns_an_error_and_releases_the_slot() { - let executor = QueryExecutor::new(1, 1); - let result = executor + for thread_per_task in [false, true] { + let executor = executor(1, thread_per_task); + let result = executor + .get_slot() + .expect("the first task must get a slot") + .run(QueryStep::First, |_| panic!("injected query panic"), |_: &(), _| {}) + .await; + + assert!(result.is_err()); + assert!(executor.get_slot().is_some(), "the panicked task must release its slot"); + } + } + + #[tokio::test] + async fn parallel_work_inside_a_task_stays_on_its_thread() { + let executor = executor(1, true); + let threads = executor .get_slot() .expect("the first task must get a slot") - .run(QueryStep::First, |_| panic!("injected query panic"), |_: &(), _| {}) - .await; + .run( + QueryStep::First, + |_| { + // Work long enough for idle workers to steal it. + (0..64) + .into_par_iter() + .map(|_| { + thread::sleep(Duration::from_millis(1)); + thread::current().id() + }) + .collect::>() + }, + |_, _| {} + ) + .await + .unwrap(); + + assert_eq!(threads.len(), 1); + } + + #[tokio::test] + async fn a_task_runs_while_another_holds_its_worker() { + let executor = executor(2, true); + let (start, started) = tokio::sync::oneshot::channel(); + let (release, held) = std::sync::mpsc::channel::<()>(); + + let holder = executor.get_slot().expect("the first task must get a slot").run( + QueryStep::First, + move |_| { + start.send(()).unwrap(); + held.recv().unwrap() + }, + |_, _| {} + ); + let holder = tokio::spawn(holder); + started.await.unwrap(); + + let other = + executor + .get_slot() + .expect("the second task must get a slot") + .run(QueryStep::First, |_| (), |_, _| {}); + tokio::time::timeout(Duration::from_secs(10), other) + .await + .expect("the task waited for the worker of another") + .unwrap(); - assert!(result.is_err()); - assert!(executor.get_slot().is_some(), "the panicked task must release its slot"); + release.send(()).unwrap(); + holder.await.unwrap().unwrap(); } #[tokio::test] async fn a_task_is_reported_after_its_caller_is_gone() { - let executor = QueryExecutor::new(1, 1); + let executor = executor(1, false); let (start, started) = tokio::sync::oneshot::channel(); let (finish, finished) = std::sync::mpsc::channel::<()>(); let (report, reported) = std::sync::mpsc::channel(); diff --git a/crates/hotblocks/src/query/service.rs b/crates/hotblocks/src/query/service.rs index 2907bfd..e13d17c 100644 --- a/crates/hotblocks/src/query/service.rs +++ b/crates/hotblocks/src/query/service.rs @@ -67,16 +67,16 @@ impl QueryServiceBuilder { self } - pub fn build(&self) -> QueryService { - QueryService { + pub fn build(&self) -> anyhow::Result { + Ok(QueryService { db: self.db.clone(), - executor: QueryExecutor::new(self.max_pending_tasks, self.urgency), + executor: QueryExecutor::new(self.max_pending_tasks, self.urgency, self.settings.clone())?, slow_tasks: Arc::new(SlowTasks::new(self.settings.clone())), wait_slots: WaitSlots { waiters: AtomicUsize::new(0), limit: self.max_data_waiters } - } + }) } } diff --git a/crates/hotblocks/src/query/slow_tasks.rs b/crates/hotblocks/src/query/slow_tasks.rs index 9b161d0..0bada3a 100644 --- a/crates/hotblocks/src/query/slow_tasks.rs +++ b/crates/hotblocks/src/query/slow_tasks.rs @@ -132,7 +132,7 @@ mod tests { use super::*; use crate::{ metrics::build_metrics_registry, - settings::{RuntimeSettings, Settings, ShadowSettings, SlowTaskSettings} + settings::{QuerySettings, RuntimeSettings, Settings, ShadowSettings, SlowTaskSettings} }; #[test] @@ -147,7 +147,8 @@ mod tests { slow_tasks: SlowTaskSettings { threshold_ms: 500, log_limit: 0 - } + }, + query: QuerySettings { thread_per_task: false } }) .unwrap(); let slow_tasks = SlowTasks::new(Arc::new(settings)); diff --git a/crates/hotblocks/src/settings.rs b/crates/hotblocks/src/settings.rs index 57dc10b..7d24a36 100644 --- a/crates/hotblocks/src/settings.rs +++ b/crates/hotblocks/src/settings.rs @@ -10,7 +10,7 @@ use serde::{Deserialize, Serialize}; use serde_json::Value; use tracing::info; -use crate::metrics::set_shadow_percent; +use crate::metrics::{set_query_thread_per_task, set_shadow_percent}; pub type SettingsRef = Arc; @@ -18,7 +18,8 @@ pub type SettingsRef = Arc; #[serde(rename_all = "camelCase", deny_unknown_fields)] pub struct Settings { pub shadow: ShadowSettings, - pub slow_tasks: SlowTaskSettings + pub slow_tasks: SlowTaskSettings, + pub query: QuerySettings } #[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)] @@ -56,6 +57,14 @@ impl SlowTaskSettings { } } +#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct QuerySettings { + /// Run each new query task on a single thread instead of the shared query + /// pool. A task already queued keeps the pool it was sent to. + pub thread_per_task: bool +} + impl Settings { fn validate(&self) -> anyhow::Result<()> { let percent = self.shadow.percent; @@ -92,6 +101,10 @@ impl RuntimeSettings { self.current.read().slow_tasks } + pub fn query(&self) -> QuerySettings { + self.current.read().query + } + /// Apply a JSON Merge Patch (RFC 7386). A patch that leaves the settings /// invalid, names an unknown key or removes a key changes nothing. pub fn update(&self, patch: Value) -> anyhow::Result { @@ -113,6 +126,7 @@ impl RuntimeSettings { /// Settings that are also exported as metrics. fn publish(settings: &Settings) { set_shadow_percent(settings.shadow.percent); + set_query_thread_per_task(settings.query.thread_per_task); } /// The merge drops a `null` key before deserialization can reject it as @@ -167,7 +181,8 @@ mod tests { slow_tasks: SlowTaskSettings { threshold_ms: 500, log_limit: 20 - } + }, + query: QuerySettings { thread_per_task: false } }) .unwrap() } @@ -190,6 +205,7 @@ mod tests { assert_eq!(updated.shadow, expected); assert_eq!(settings.shadow(), expected); assert_eq!(updated.slow_tasks, before.slow_tasks); + assert_eq!(updated.query, before.query); } #[test] @@ -203,6 +219,7 @@ mod tests { json!({"shadow": {"percent": 5, "percnt": null}}), json!({"shadow": {"percent": null}}), json!({"shadow": {"logLimit": -1}}), + json!({"query": {"threadPerTask": 1}}), json!({"other": {}}), json!(5) ] { diff --git a/crates/hotblocks/tests/query_task_metrics.rs b/crates/hotblocks/tests/query_task_metrics.rs index 2ab434e..db81b27 100644 --- a/crates/hotblocks/tests/query_task_metrics.rs +++ b/crates/hotblocks/tests/query_task_metrics.rs @@ -3,8 +3,8 @@ use std::sync::Arc; -use anyhow::{Result, ensure}; -use serde_json::json; +use anyhow::{Result, bail, ensure}; +use serde_json::{Value, json}; use sqd_hotblocks_harness::{ Harness, HarnessConfig, Outcome, chain::{Chain, Evm}, @@ -15,6 +15,7 @@ const START: u64 = 1_000; const QUEUE: &str = "hotblocks_query_task_queue_seconds_count"; const RUN: &str = "hotblocks_query_task_run_seconds_count"; const SLOW: &str = "hotblocks_query_slow_tasks_total"; +const THREAD_PER_TASK: &str = "hotblocks_query_thread_per_task"; #[tokio::test(flavor = "multi_thread")] async fn a_query_task_is_timed_in_the_queue_and_on_a_worker() -> Result<()> { @@ -102,3 +103,36 @@ async fn a_first_task_that_fails_is_still_counted() -> Result<()> { assert_eq!(metrics.get(SLOW, Some(("step", "first"))), Some(1.0)); Ok(()) } + +#[tokio::test(flavor = "multi_thread")] +async fn a_query_answers_the_same_with_a_thread_per_task_and_the_mode_changes_at_runtime() -> Result<()> { + let mut cfg = HarnessConfig::from_block(env!("CARGO_BIN_EXE_sqd-hotblocks"), Arc::new(Evm), START); + cfg.sut_args.push("--query-thread-per-task".to_string()); + let mut h = Harness::start(cfg).await?; + h.produce(10)?; + h.settle().await?; + let query = Evm.scan_query(START, None, None); + + let blocks = async |h: &Harness| -> Result> { + match h.client.query(&query).await? { + Outcome::Ok { blocks, .. } => Ok(blocks.into_iter().map(|block| block.raw).collect()), + outcome => bail!("{query} did not return data: {outcome:?}") + } + }; + + assert_eq!(h.client.metrics().await?.get(THREAD_PER_TASK, None), Some(1.0)); + let single = blocks(&h).await?; + + assert_eq!( + h.client + .update_settings(&json!({"query": {"threadPerTask": false}})) + .await?, + 200 + ); + assert_eq!(h.client.metrics().await?.get(THREAD_PER_TASK, None), Some(0.0)); + let shared = blocks(&h).await?; + + assert_eq!(shared.len(), 10); + assert_eq!(single, shared); + Ok(()) +}