diff --git a/concurrency/src/lib.rs b/concurrency/src/lib.rs index 322d2d67e..9d5993e4e 100644 --- a/concurrency/src/lib.rs +++ b/concurrency/src/lib.rs @@ -17,7 +17,9 @@ pub use notification_list::NotificationList; pub use parallel_writer::ParallelVecWriter; pub use resettable_oncelock::ResettableOnceLock; pub use shared_arena::{Handle, SharedArena, SharedRef}; -pub use threadpool::{Scope, ThreadPool, current_num_threads, scope, without_current_pool}; +pub use threadpool::{ + SchedulerMetrics, Scope, ThreadPool, current_num_threads, scope, without_current_pool, +}; #[cfg(test)] mod tests; diff --git a/concurrency/src/threadpool/mod.rs b/concurrency/src/threadpool/mod.rs index 427b7a559..1d7b3ae9b 100644 --- a/concurrency/src/threadpool/mod.rs +++ b/concurrency/src/threadpool/mod.rs @@ -1,12 +1,15 @@ -//! A small scoped thread pool backed by a shared crossbeam channel. +//! A small scoped thread pool with global and worker-local work queues. //! //! [`ThreadPool`] owns a fixed set of primary worker threads. Each worker -//! receives boxed `'static` jobs from a shared channel and exits when the -//! channel is closed. A primary worker that blocks waiting for nested scoped -//! work helps drain queued jobs until that nested scope completes, so the pool -//! does not lose a worker while work it needs may still be queued. [`Scope`] -//! provides the safe scoped API on top of those `'static` jobs by erasing task -//! lifetimes when work is sent to the channel, then waiting for all work in the +//! owns a private deque in addition to receiving boxed `'static` jobs from a +//! shared channel. Local work is pushed and popped at the back, giving the +//! owner depth-first execution; when another worker is stalled, half of the +//! oldest local work is donated from the front to the shared queue. A primary +//! worker that blocks waiting for nested scoped work helps drain its local +//! deque first and then the shared queue until that nested scope completes, so +//! the pool does not lose a worker while work it needs may still be queued. +//! [`Scope`] provides the safe scoped API on top of those `'static` jobs by +//! erasing task lifetimes when work is queued, then waiting for all work in the //! scope before returning to the caller. //! //! The root callback runs on the caller thread, and spawned callbacks run on @@ -95,22 +98,22 @@ use std::{ any::Any, - cell::Cell, + cell::{Cell, UnsafeCell}, + collections::VecDeque, marker::PhantomData, mem, panic::{self, AssertUnwindSafe}, ptr::{self, NonNull}, sync::{ Mutex, - atomic::{AtomicU64, Ordering}, + atomic::{AtomicU64, AtomicUsize, Ordering}, }, thread::{self, JoinHandle}, + time::{Duration, Instant}, }; -#[cfg(test)] -use std::sync::atomic::AtomicUsize; - use crossbeam::channel::{Receiver, Sender, bounded, select_biased, unbounded}; +use crossbeam::utils::CachePadded; const EXPECTED_SHIFT: u32 = 32; const COMPLETED_MASK: u64 = u32::MAX as u64; @@ -124,10 +127,69 @@ type PanicPayload = Box; thread_local! { static CURRENT_POOL: Cell<*const ThreadPoolState> = const { Cell::new(ptr::null()) }; + static CURRENT_WORKER: Cell<*const WorkerContext> = const { Cell::new(ptr::null()) }; static IS_BACKGROUND_WORKER: Cell = const { Cell::new(false) }; static INLINE_SCOPE_HELP_DEPTH: Cell = const { Cell::new(0) }; } +/// A cumulative snapshot of a thread pool's scheduler instrumentation. +/// +/// Counter fields and `stalled_time` are cumulative since the pool was created +/// or [`ThreadPool::reset_scheduler_metrics`] was called. Queue-depth fields +/// are high-water marks over that same period. `stalled_workers` is a gauge at +/// the instant of the snapshot rather than a cumulative value. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub struct SchedulerMetrics { + /// Jobs pushed onto the shared global queue, including donations. + pub global_pushes: u64, + /// Jobs popped from the shared global queue. + pub global_pops: u64, + /// Jobs successfully pushed onto worker-private queues. + pub local_pushes: u64, + /// Jobs popped by their owning worker from a private queue. + pub local_pops: u64, + /// Jobs moved from private queues to the global queue. + pub donated_jobs: u64, + /// Aggregate wall-clock time for which primary workers were stalled. + /// + /// This includes the elapsed portion of stalls that are still in progress + /// when the snapshot is taken, which makes before/after snapshots useful + /// even when the pool is idle at both endpoints. + pub stalled_time: Duration, + /// Number of primary workers currently waiting for global work. + pub stalled_workers: usize, + /// Largest number of jobs observed in any one worker-private queue. + pub max_local_queue_depth: usize, + /// Largest number of jobs observed in the global queue. + pub max_global_queue_depth: usize, +} + +impl SchedulerMetrics { + /// Return the cumulative counter and stalled-time difference from `earlier`. + /// + /// Queue high-water marks cannot in general be subtracted. The returned + /// value therefore retains the later snapshot's high-water marks and its + /// current `stalled_workers` gauge. Call + /// [`ThreadPool::reset_scheduler_metrics`] immediately before an isolated + /// measurement when interval-specific queue high-water marks are needed. + pub fn delta_since(self, earlier: Self) -> Self { + Self { + global_pushes: self.global_pushes.saturating_sub(earlier.global_pushes), + global_pops: self.global_pops.saturating_sub(earlier.global_pops), + local_pushes: self.local_pushes.saturating_sub(earlier.local_pushes), + local_pops: self.local_pops.saturating_sub(earlier.local_pops), + donated_jobs: self.donated_jobs.saturating_sub(earlier.donated_jobs), + stalled_time: self + .stalled_time + .checked_sub(earlier.stalled_time) + .unwrap_or_default(), + stalled_workers: self.stalled_workers, + max_local_queue_depth: self.max_local_queue_depth, + max_global_queue_depth: self.max_global_queue_depth, + } + } +} + /// Return the number of threads in the currently installed thread pool. /// /// Returns `1` when called outside [`ThreadPool::install`], [`ThreadPool::scope`], @@ -222,11 +284,11 @@ where }) } -/// A thread pool using a single shared work queue. +/// A thread pool with a shared queue and a private deque per primary worker. /// /// The pool owns a fixed number of primary workers. A primary worker that -/// blocks waiting for nested scoped work helps drain the shared queue until the -/// nested scope completes. +/// blocks waiting for nested scoped work helps drain its private deque before +/// the shared queue until the nested scope completes. /// /// # Examples /// @@ -268,7 +330,7 @@ impl ThreadPool { let state = Box::new(ThreadPoolState::new(sender, receiver.clone(), thread_count)); let state_ptr = ThreadPoolStatePtr::new(&state); let workers = (0..thread_count) - .map(|_| spawn_worker(receiver.clone(), state_ptr)) + .map(|worker_id| spawn_worker(receiver.clone(), state_ptr, worker_id)) .collect(); Self { state, workers } @@ -288,6 +350,26 @@ impl ThreadPool { self.state.thread_count() } + /// Return a cumulative snapshot of scheduler instrumentation. + /// + /// A snapshot includes the elapsed portion of stalls that are still active + /// at the instant it is taken. Consequently, subtracting a before snapshot + /// from an after snapshot with [`SchedulerMetrics::delta_since`] does not + /// attribute idle time before the first snapshot to the measured interval. + pub fn scheduler_metrics(&self) -> SchedulerMetrics { + self.state.instrumentation.snapshot() + } + + /// Reset cumulative scheduler instrumentation and queue high-water marks. + /// + /// Active worker stalls are split at the reset instant, so idle time before + /// the reset is not charged to later snapshots. For meaningful queue + /// high-water marks, call this while no callbacks are running or being + /// scheduled; this method does not pause the pool. + pub fn reset_scheduler_metrics(&self) { + self.state.instrumentation.reset() + } + /// Run `f` with this thread pool installed as the current pool. /// /// The callback runs on the caller thread. While it is running, @@ -396,10 +478,228 @@ impl Drop for ThreadPool { } } +#[derive(Default)] +struct StallTimerState { + completed: Duration, + started: Option, +} + +struct WorkerInstrumentation { + local_pushes: AtomicU64, + local_pops: AtomicU64, + donated_jobs: AtomicU64, + max_local_queue_depth: AtomicUsize, + stall_timer: Mutex, +} + +impl WorkerInstrumentation { + fn new() -> Self { + Self { + local_pushes: AtomicU64::new(0), + local_pops: AtomicU64::new(0), + donated_jobs: AtomicU64::new(0), + max_local_queue_depth: AtomicUsize::new(0), + stall_timer: Mutex::new(StallTimerState::default()), + } + } +} + +struct SchedulerInstrumentation { + global_pushes: AtomicU64, + global_pops: AtomicU64, + stalled_workers: AtomicUsize, + workers: Vec>, + global_queue_depth: AtomicUsize, + max_global_queue_depth: AtomicUsize, +} + +impl SchedulerInstrumentation { + fn new(thread_count: usize) -> Self { + Self { + global_pushes: AtomicU64::new(0), + global_pops: AtomicU64::new(0), + stalled_workers: AtomicUsize::new(0), + workers: (0..thread_count) + .map(|_| CachePadded::new(WorkerInstrumentation::new())) + .collect(), + global_queue_depth: AtomicUsize::new(0), + max_global_queue_depth: AtomicUsize::new(0), + } + } + + fn snapshot(&self) -> SchedulerMetrics { + let now = Instant::now(); + let stalled_time = self + .workers + .iter() + .map(|worker| { + let timer = worker + .stall_timer + .lock() + .unwrap_or_else(|error| error.into_inner()); + timer.completed.saturating_add( + timer + .started + .map_or(Duration::ZERO, |started| now.duration_since(started)), + ) + }) + .fold(Duration::ZERO, Duration::saturating_add); + let local_pushes = self + .workers + .iter() + .map(|worker| worker.local_pushes.load(Ordering::Relaxed)) + .sum(); + let local_pops = self + .workers + .iter() + .map(|worker| worker.local_pops.load(Ordering::Relaxed)) + .sum(); + let donated_jobs = self + .workers + .iter() + .map(|worker| worker.donated_jobs.load(Ordering::Relaxed)) + .sum(); + let max_local_queue_depth = self + .workers + .iter() + .map(|worker| worker.max_local_queue_depth.load(Ordering::Relaxed)) + .max() + .unwrap_or(0); + + SchedulerMetrics { + global_pushes: self.global_pushes.load(Ordering::Relaxed), + global_pops: self.global_pops.load(Ordering::Relaxed), + local_pushes, + local_pops, + donated_jobs, + stalled_time, + stalled_workers: self.stalled_workers.load(Ordering::Acquire), + max_local_queue_depth, + max_global_queue_depth: self.max_global_queue_depth.load(Ordering::Relaxed), + } + } + + fn reset(&self) { + let now = Instant::now(); + for worker in &self.workers { + let mut timer = worker + .stall_timer + .lock() + .unwrap_or_else(|error| error.into_inner()); + timer.completed = Duration::ZERO; + if timer.started.is_some() { + timer.started = Some(now); + } + drop(timer); + + worker.local_pushes.store(0, Ordering::Relaxed); + worker.local_pops.store(0, Ordering::Relaxed); + worker.donated_jobs.store(0, Ordering::Relaxed); + worker.max_local_queue_depth.store(0, Ordering::Relaxed); + } + + self.global_pushes.store(0, Ordering::Relaxed); + self.global_pops.store(0, Ordering::Relaxed); + self.max_global_queue_depth.store( + self.global_queue_depth.load(Ordering::Acquire), + Ordering::Relaxed, + ); + } + + fn record_global_push(&self) { + self.global_pushes.fetch_add(1, Ordering::Relaxed); + let depth = self.global_queue_depth.fetch_add(1, Ordering::AcqRel) + 1; + self.max_global_queue_depth + .fetch_max(depth, Ordering::Relaxed); + } + + fn cancel_global_push(&self) { + self.global_pushes.fetch_sub(1, Ordering::Relaxed); + self.global_queue_depth.fetch_sub(1, Ordering::AcqRel); + } + + fn record_global_pop(&self) { + self.global_pops.fetch_add(1, Ordering::Relaxed); + self.global_queue_depth.fetch_sub(1, Ordering::AcqRel); + } + + fn record_local_push(&self, worker_id: usize, depth: usize) { + let worker = &self.workers[worker_id]; + worker.local_pushes.fetch_add(1, Ordering::Relaxed); + worker + .max_local_queue_depth + .fetch_max(depth, Ordering::Relaxed); + } + + fn record_local_pop(&self, worker_id: usize) { + self.workers[worker_id] + .local_pops + .fetch_add(1, Ordering::Relaxed); + } + + fn record_donation(&self, worker_id: usize, jobs: usize) { + self.workers[worker_id] + .donated_jobs + .fetch_add(jobs as u64, Ordering::Relaxed); + } + + fn has_stalled_workers(&self) -> bool { + self.stalled_workers.load(Ordering::Acquire) != 0 + } + + fn begin_stall(&self, worker_id: usize) { + let mut timer = self.workers[worker_id] + .stall_timer + .lock() + .unwrap_or_else(|error| error.into_inner()); + debug_assert!(timer.started.is_none()); + timer.started = Some(Instant::now()); + drop(timer); + self.stalled_workers.fetch_add(1, Ordering::AcqRel); + } + + fn end_stall(&self, worker_id: usize) { + let now = Instant::now(); + let mut timer = self.workers[worker_id] + .stall_timer + .lock() + .unwrap_or_else(|error| error.into_inner()); + let started = timer + .started + .take() + .expect("primary worker ended a stall it did not begin"); + timer.completed += now.duration_since(started); + drop(timer); + self.stalled_workers.fetch_sub(1, Ordering::AcqRel); + } +} + +struct StalledWorkerGuard<'pool> { + instrumentation: &'pool SchedulerInstrumentation, + worker_id: usize, +} + +impl<'pool> StalledWorkerGuard<'pool> { + fn new(instrumentation: &'pool SchedulerInstrumentation, worker_id: usize) -> Self { + instrumentation.begin_stall(worker_id); + Self { + instrumentation, + worker_id, + } + } +} + +impl Drop for StalledWorkerGuard<'_> { + fn drop(&mut self) { + self.instrumentation.end_stall(self.worker_id); + } +} + struct ThreadPoolState { sender: Option>, receiver: Receiver, thread_count: usize, + instrumentation: SchedulerInstrumentation, #[cfg(test)] backup_workers_spawned: AtomicUsize, #[cfg(test)] @@ -412,6 +712,7 @@ impl ThreadPoolState { sender: Some(sender), receiver, thread_count, + instrumentation: SchedulerInstrumentation::new(thread_count), #[cfg(test)] backup_workers_spawned: AtomicUsize::new(0), #[cfg(test)] @@ -423,6 +724,38 @@ impl ThreadPoolState { self.thread_count } + fn enqueue_global(&self, job: Job) { + let Some(sender) = self.sender.as_ref() else { + job(); + return; + }; + + self.instrumentation.record_global_push(); + if let Err(error) = sender.send(job) { + self.instrumentation.cancel_global_push(); + error.0(); + } + } + + fn try_recv_global(&self) -> Result { + let result = self.receiver.try_recv(); + if result.is_ok() { + self.instrumentation.record_global_pop(); + } + result + } + + fn recv_global(&self, worker_id: usize) -> Result { + let result = { + let _stalled = StalledWorkerGuard::new(&self.instrumentation, worker_id); + self.receiver.recv() + }; + if result.is_ok() { + self.instrumentation.record_global_pop(); + } + result + } + fn scope<'scope, F, R>(&self, f: F) -> R where F: FnOnce(&Scope<'scope>) -> R, @@ -450,12 +783,15 @@ impl ThreadPoolState { } let Some(_guard) = InlineScopeHelpGuard::try_enter() else { + let worker_id = current_worker_id(self); + donate_all_from_current_worker(self); let _backup = BackupWorker::spawn(self); + let _stalled = worker_id + .map(|worker_id| StalledWorkerGuard::new(&self.instrumentation, worker_id)); receive_scope_completion(done); return; }; - let receiver = self.receiver.clone(); loop { match done.try_recv() { Ok(()) => break, @@ -465,7 +801,12 @@ impl ThreadPoolState { } } - match receiver.try_recv() { + if let Some(job) = pop_current_worker_local(self) { + job(); + continue; + } + + match self.try_recv_global() { Ok(job) => { job(); continue; @@ -477,13 +818,31 @@ impl ThreadPoolState { } } - select_biased! { - recv(done) -> message => { + enum WaitEvent { + Scope(Result<(), crossbeam::channel::RecvError>), + Global(Result), + } + + let worker_id = current_worker_id(self); + let event = { + let _stalled = worker_id + .map(|worker_id| StalledWorkerGuard::new(&self.instrumentation, worker_id)); + select_biased! { + recv(done) -> message => WaitEvent::Scope(message), + recv(self.receiver) -> message => WaitEvent::Global(message), + } + }; + + match event { + WaitEvent::Scope(message) => { expect_scope_completion(message); break; } - recv(receiver) -> message => match message { - Ok(job) => job(), + WaitEvent::Global(message) => match message { + Ok(job) => { + self.instrumentation.record_global_pop(); + job(); + } Err(_) => { receive_scope_completion(done); break; @@ -530,6 +889,161 @@ unsafe impl Send for ThreadPoolStatePtr {} // boxed state remains alive until all worker threads have been joined. unsafe impl Sync for ThreadPoolStatePtr {} +/// State owned and accessed exclusively by one primary worker thread. +/// +/// The deque uses interior mutability so no Rust borrow of the queue is held +/// while a job is running. Jobs may recursively call `spawn_local`, which then +/// reaches the same deque through `CURRENT_WORKER` on that thread. +struct WorkerContext { + pool: ThreadPoolStatePtr, + worker_id: usize, + local: UnsafeCell>, +} + +impl WorkerContext { + fn new(pool: ThreadPoolStatePtr, worker_id: usize) -> Self { + Self { + pool, + worker_id, + local: UnsafeCell::new(VecDeque::new()), + } + } + + fn belongs_to(&self, pool: &ThreadPoolState) -> bool { + ptr::eq(self.pool.0, pool) + } + + fn push_local(&self, pool: &ThreadPoolState, job: Job) { + // SAFETY: a `WorkerContext` is installed only on its owning worker and + // `CURRENT_WORKER` is never made available to another thread. + let depth = unsafe { + let queue = &mut *self.local.get(); + queue.push_back(job); + queue.len() + }; + pool.instrumentation + .record_local_push(self.worker_id, depth); + self.donate_half_if_stalled(pool); + } + + fn pop_local(&self, pool: &ThreadPoolState) -> Option { + self.donate_half_if_stalled(pool); + // The owner takes the newest work, retaining depth-first locality. + // SAFETY: see `push_local`. + let job = unsafe { (&mut *self.local.get()).pop_back() }; + if job.is_some() { + pool.instrumentation.record_local_pop(self.worker_id); + } + job + } + + fn donate_half_if_stalled(&self, pool: &ThreadPoolState) { + if !pool.instrumentation.has_stalled_workers() { + return; + } + + // Donation takes the opposite end from the owner: the oldest siblings + // have the weakest connection to the owner's current depth-first path. + // SAFETY: see `push_local`. + let donated = unsafe { + let queue = &mut *self.local.get(); + let count = queue.len() / 2; + queue.drain(..count).collect::>() + }; + self.publish_donation(pool, donated); + } + + fn donate_all(&self, pool: &ThreadPoolState) { + // The inline-help depth escape hatch blocks this worker and starts a + // backup consumer. Publish everything so private work remains live. + // SAFETY: see `push_local`. + let donated = unsafe { (&mut *self.local.get()).drain(..).collect::>() }; + self.publish_donation(pool, donated); + } + + fn publish_donation(&self, pool: &ThreadPoolState, donated: Vec) { + if donated.is_empty() { + return; + } + + pool.instrumentation + .record_donation(self.worker_id, donated.len()); + for job in donated { + pool.enqueue_global(job); + } + } +} + +fn install_worker_context(worker: &WorkerContext, f: F) -> R +where + F: FnOnce() -> R, +{ + CURRENT_WORKER.with(|current| { + let previous = current.replace(worker); + let _guard = CurrentWorkerGuard { current, previous }; + f() + }) +} + +struct CurrentWorkerGuard<'a> { + current: &'a Cell<*const WorkerContext>, + previous: *const WorkerContext, +} + +impl Drop for CurrentWorkerGuard<'_> { + fn drop(&mut self) { + self.current.set(self.previous); + } +} + +fn with_current_worker(f: F) -> R +where + F: FnOnce(Option<&WorkerContext>) -> R, +{ + CURRENT_WORKER.with(|current| { + let worker = current.get(); + // SAFETY: the pointer is installed from a live `WorkerContext` on this + // same thread and is restored before that context is dropped. + f(unsafe { worker.as_ref() }) + }) +} + +fn current_worker_id(pool: &ThreadPoolState) -> Option { + with_current_worker(|worker| { + worker + .filter(|worker| worker.belongs_to(pool)) + .map(|worker| worker.worker_id) + }) +} + +fn push_current_worker_local(pool: &ThreadPoolState, job: Job) -> Result<(), Job> { + with_current_worker( + |worker| match worker.filter(|worker| worker.belongs_to(pool)) { + Some(worker) => { + worker.push_local(pool, job); + Ok(()) + } + None => Err(job), + }, + ) +} + +fn pop_current_worker_local(pool: &ThreadPoolState) -> Option { + with_current_worker(|worker| { + worker + .filter(|worker| worker.belongs_to(pool)) + .and_then(|worker| worker.pop_local(pool)) + }) +} + +fn donate_all_from_current_worker(pool: &ThreadPoolState) { + with_current_worker(|worker| { + if let Some(worker) = worker.filter(|worker| worker.belongs_to(pool)) { + worker.donate_all(pool); + } + }); +} + fn install_pool(pool: &ThreadPoolState, f: F) -> R where F: FnOnce() -> R, @@ -648,7 +1162,6 @@ where /// assert_eq!(counter.load(Ordering::Relaxed), 1); /// ``` pub struct Scope<'scope> { - sender: Sender, pool: ThreadPoolStatePtr, state: ScopeState, _scope: PhantomData &'scope ()>, @@ -657,18 +1170,18 @@ pub struct Scope<'scope> { impl<'scope> Scope<'scope> { fn new(pool: &ThreadPoolState) -> Self { Self { - sender: pool.sender.as_ref().unwrap().clone(), pool: ThreadPoolStatePtr::new(pool), state: ScopeState::new(), _scope: PhantomData, } } - /// Spawn a callback onto the pool. + /// Spawn a callback onto the pool's global queue. /// /// The callback receives the current scope and may add nested work. It may /// borrow values with the scope lifetime. The scope waits for the callback - /// before returning, so those borrows cannot outlive their owners. + /// before returning, so those borrows cannot outlive their owners. This is + /// retained as a compatibility alias for [`Scope::spawn_global`]. /// /// # Examples /// @@ -691,6 +1204,46 @@ impl<'scope> Scope<'scope> { /// assert_eq!(counter.load(Ordering::Relaxed), 2); /// ``` pub fn spawn(&self, f: F) + where + F: FnOnce(&Scope<'scope>) + Send + 'scope, + { + self.spawn_global(f); + } + + /// Spawn a callback onto the pool's global queue. + /// + /// Global work is visible to every worker immediately and is appropriate + /// for coarse top-level partitioning. + pub fn spawn_global(&self, f: F) + where + F: FnOnce(&Scope<'scope>) + Send + 'scope, + { + let job = self.prepare_job(f); + // SAFETY: this scope was created from a live pool, and the pool cannot + // be dropped while its borrowed `scope` call is active. + unsafe { self.pool.as_ref() }.enqueue_global(job); + } + + /// Spawn a callback onto the current worker's private deque. + /// + /// The owning worker executes local jobs depth-first before consulting the + /// global queue. When another worker is stalled, half of the oldest local + /// jobs are donated to the global queue. Calls made outside a primary + /// worker of this same pool—including calls from the root callback and + /// from a worker of a nested, different pool—fall back to the global queue. + pub fn spawn_local(&self, f: F) + where + F: FnOnce(&Scope<'scope>) + Send + 'scope, + { + let job = self.prepare_job(f); + // SAFETY: see `spawn_global`. + let pool = unsafe { self.pool.as_ref() }; + if let Err(job) = push_current_worker_local(pool, job) { + pool.enqueue_global(job); + } + } + + fn prepare_job(&self, f: F) -> Job where F: FnOnce(&Scope<'scope>) + Send + 'scope, { @@ -717,7 +1270,7 @@ impl<'scope> Scope<'scope> { // SAFETY: every erased job records completion in the scope state, and // `Scope::complete_root_and_wait` waits for all expected completions // before `ThreadPool::scope` returns. - enqueue(&self.sender, unsafe { erase_job_lifetime(job) }); + unsafe { erase_job_lifetime(job) } } fn complete_root_and_wait(&self) { @@ -898,7 +1451,10 @@ impl BackupWorker { select_biased! { recv(shutdown_receiver) -> _ => break, recv(receiver) -> message => match message { - Ok(job) => job(), + Ok(job) => { + pool.instrumentation.record_global_pop(); + job(); + } Err(_) => break, }, } @@ -947,30 +1503,45 @@ impl Drop for BackupWorkerLiveGuard<'_> { } } -fn spawn_worker(receiver: Receiver, pool: ThreadPoolStatePtr) -> JoinHandle<()> { +fn spawn_worker( + receiver: Receiver, + pool: ThreadPoolStatePtr, + worker_id: usize, +) -> JoinHandle<()> { thread::spawn(move || { // SAFETY: each worker is joined before the boxed pool state is dropped. let pool = unsafe { pool.as_ref() }; + let worker = WorkerContext::new(ThreadPoolStatePtr::new(pool), worker_id); install_pool(pool, || { - install_background_worker(|| { - for job in receiver { - job(); - } + install_worker_context(&worker, || { + install_background_worker(|| { + loop { + if let Some(job) = worker.pop_local(pool) { + job(); + continue; + } + + match receiver.try_recv() { + Ok(job) => { + pool.instrumentation.record_global_pop(); + job(); + continue; + } + Err(crossbeam::channel::TryRecvError::Empty) => {} + Err(crossbeam::channel::TryRecvError::Disconnected) => break, + } + + let Ok(job) = pool.recv_global(worker_id) else { + break; + }; + job(); + } + }) }); }); }) } -fn enqueue(sender: &Sender, job: Job) { - match sender.send(job) { - Ok(()) => {} - Err(error) => { - let job = error.0; - job(); - } - } -} - fn completed(value: u64) -> u32 { (value & COMPLETED_MASK) as u32 } diff --git a/concurrency/src/threadpool/tests.rs b/concurrency/src/threadpool/tests.rs index 4b43edf8e..36186df5e 100644 --- a/concurrency/src/threadpool/tests.rs +++ b/concurrency/src/threadpool/tests.rs @@ -5,7 +5,7 @@ use std::{ atomic::{AtomicBool, AtomicUsize, Ordering}, }, thread, - time::Duration, + time::{Duration, Instant}, }; use crate::{ @@ -13,7 +13,10 @@ use crate::{ without_current_pool, }; -use super::{MAX_INLINE_SCOPE_HELP_DEPTH, is_background_worker_thread}; +use super::{ + MAX_INLINE_SCOPE_HELP_DEPTH, ThreadPoolState, ThreadPoolStatePtr, WorkerContext, + is_background_worker_thread, +}; #[test] fn scope_runs_spawned_work_and_waits() { @@ -31,6 +34,213 @@ fn scope_runs_spawned_work_and_waits() { assert_eq!(counter.load(Ordering::Relaxed), 128); } +#[test] +fn spawn_remains_a_global_queue_compatibility_alias() { + let pool = ThreadPool::new(1); + pool.reset_scheduler_metrics(); + + pool.scope(|scope| scope.spawn(|_| {})); + + let metrics = pool.scheduler_metrics(); + assert_eq!(metrics.global_pushes, 1); + assert_eq!(metrics.global_pops, 1); + assert_eq!(metrics.local_pushes, 0); + assert_eq!(metrics.local_pops, 0); +} + +#[test] +fn spawn_local_from_root_falls_back_to_global_queue() { + let pool = ThreadPool::new(1); + pool.reset_scheduler_metrics(); + + pool.scope(|scope| scope.spawn_local(|_| {})); + + let metrics = pool.scheduler_metrics(); + assert_eq!(metrics.global_pushes, 1); + assert_eq!(metrics.global_pops, 1); + assert_eq!(metrics.local_pushes, 0); +} + +#[test] +fn owner_runs_local_work_depth_first_before_global_work() { + let pool = ThreadPool::new(1); + let first_started = Notification::new(); + let release_first = Notification::new(); + let order = Mutex::new(Vec::new()); + + pool.scope(|scope| { + scope.spawn_global(|scope| { + order.lock().unwrap().push("parent"); + first_started.notify(); + release_first.wait(); + scope.spawn_local(|_| order.lock().unwrap().push("local-oldest")); + scope.spawn_local(|_| order.lock().unwrap().push("local-newest")); + }); + + first_started.wait(); + scope.spawn_global(|_| order.lock().unwrap().push("global")); + release_first.notify(); + }); + + assert_eq!( + order.into_inner().unwrap(), + ["parent", "local-newest", "local-oldest", "global"] + ); + let metrics = pool.scheduler_metrics(); + assert_eq!(metrics.local_pushes, 2); + assert_eq!(metrics.local_pops, 2); + assert_eq!(metrics.donated_jobs, 0); +} + +#[test] +fn nested_scope_help_takes_local_work_before_global_work() { + let pool = ThreadPool::new(1); + let parent_started = Notification::new(); + let release_parent = Notification::new(); + let order = Mutex::new(Vec::new()); + + pool.scope(|scope| { + scope.spawn_global(|_| { + order.lock().unwrap().push("parent-start"); + parent_started.notify(); + release_parent.wait(); + + threadpool_scope(|nested| { + nested.spawn_local(|_| order.lock().unwrap().push("nested-local")); + }); + order.lock().unwrap().push("parent-after-nested"); + }); + + parent_started.wait(); + scope.spawn_global(|_| order.lock().unwrap().push("global")); + release_parent.notify(); + }); + + assert_eq!( + order.into_inner().unwrap(), + [ + "parent-start", + "nested-local", + "parent-after-nested", + "global" + ] + ); +} + +#[test] +fn donation_moves_half_of_the_oldest_local_jobs() { + let (sender, receiver) = super::unbounded(); + let state = ThreadPoolState::new(sender, receiver, 1); + let worker = WorkerContext::new(ThreadPoolStatePtr::new(&state), 0); + let order = Arc::new(Mutex::new(Vec::new())); + + // Build the queue directly so one donation decision observes all six jobs. + // This isolates deque-end and half-transfer policy from worker timing. + unsafe { + let queue = &mut *worker.local.get(); + for value in 0..6 { + let order = order.clone(); + queue.push_back(Box::new(move || order.lock().unwrap().push(value))); + } + } + state + .instrumentation + .stalled_workers + .store(1, Ordering::Release); + + worker.donate_half_if_stalled(&state); + state + .instrumentation + .stalled_workers + .store(0, Ordering::Release); + + for _ in 0..3 { + state.try_recv_global().unwrap()(); + } + while let Some(job) = worker.pop_local(&state) { + job(); + } + + assert_eq!(*order.lock().unwrap(), [0, 1, 2, 5, 4, 3]); + let metrics = state.instrumentation.snapshot(); + assert_eq!(metrics.donated_jobs, 3); + assert_eq!(metrics.global_pushes, 3); + assert_eq!(metrics.global_pops, 3); + assert_eq!(metrics.local_pops, 3); +} + +#[test] +fn stalled_worker_signal_causes_public_local_work_to_be_donated() { + let pool = ThreadPool::new(2); + let deadline = Instant::now() + Duration::from_secs(1); + while pool.scheduler_metrics().stalled_workers != 2 { + assert!(Instant::now() < deadline, "workers did not become stalled"); + thread::yield_now(); + } + pool.reset_scheduler_metrics(); + + pool.scope(|scope| { + scope.spawn_global(|scope| { + for _ in 0..16 { + scope.spawn_local(|_| {}); + } + }); + }); + + let metrics = pool.scheduler_metrics(); + assert_eq!(metrics.local_pushes, 16); + assert!(metrics.donated_jobs > 0); + assert_eq!( + metrics.local_pushes, + metrics.local_pops + metrics.donated_jobs + ); + assert_eq!(metrics.global_pushes, 1 + metrics.donated_jobs); + assert_eq!(metrics.global_pushes, metrics.global_pops); + assert!(metrics.max_local_queue_depth >= 2); +} + +#[test] +fn stalled_time_snapshots_include_stalls_in_progress() { + let pool = ThreadPool::new(2); + let deadline = Instant::now() + Duration::from_secs(1); + while pool.scheduler_metrics().stalled_workers != 2 { + assert!(Instant::now() < deadline, "workers did not become stalled"); + thread::yield_now(); + } + + let before = pool.scheduler_metrics(); + thread::sleep(Duration::from_millis(2)); + let delta = pool.scheduler_metrics().delta_since(before); + assert!(delta.stalled_time >= Duration::from_millis(2)); + assert_eq!(delta.stalled_workers, 2); + + pool.reset_scheduler_metrics(); + let after_reset = pool.scheduler_metrics(); + thread::sleep(Duration::from_millis(2)); + let reset_delta = pool.scheduler_metrics().delta_since(after_reset); + assert!(reset_delta.stalled_time >= Duration::from_millis(2)); +} + +#[test] +fn panic_from_local_work_is_propagated_after_the_scope_drains() { + let pool = ThreadPool::new(1); + let sibling_completed = AtomicBool::new(false); + + let result = panic::catch_unwind(panic::AssertUnwindSafe(|| { + pool.scope(|scope| { + scope.spawn_global(|scope| { + scope.spawn_local(|_| sibling_completed.store(true, Ordering::Release)); + // Local jobs are LIFO, so the panic runs first. Its job wrapper + // must record the panic and let the worker drain the sibling. + scope.spawn_local(|_| panic!("local worker panic")); + }); + }); + })); + + assert!(result.is_err()); + assert!(sibling_completed.load(Ordering::Acquire)); +} + #[test] fn scope_allows_borrowing_stack_data() { let pool = ThreadPool::new(2); @@ -377,6 +587,28 @@ fn install_sets_and_restores_current_pool() { assert_eq!(current_num_threads(), 1); } +#[test] +fn spawn_local_does_not_use_a_different_pools_worker_queue() { + let outer = ThreadPool::new(1); + let inner = ThreadPool::new(1); + let completed = AtomicBool::new(false); + inner.reset_scheduler_metrics(); + + outer.scope(|scope| { + scope.spawn_global(|_| { + inner.scope(|scope| { + scope.spawn_local(|_| completed.store(true, Ordering::Release)); + }); + }); + }); + + assert!(completed.load(Ordering::Acquire)); + let metrics = inner.scheduler_metrics(); + assert_eq!(metrics.local_pushes, 0); + assert_eq!(metrics.global_pushes, 1); + assert_eq!(metrics.global_pops, 1); +} + #[test] fn without_current_pool_clears_and_restores_current_pool() { let pool = ThreadPool::new(2); @@ -601,3 +833,31 @@ fn deeply_nested_background_waits_fall_back_to_backup_worker() { assert!(pool.state.backup_workers_spawned() >= 1); assert_eq!(pool.state.backup_workers_live(), 0); } + +#[test] +fn deeply_nested_local_work_is_donated_to_backup_worker() { + fn recurse(depth: usize, completed: &AtomicUsize) { + if depth == 0 { + completed.fetch_add(1, Ordering::Release); + return; + } + + threadpool_scope(|scope| { + scope.spawn_local(move |_| recurse(depth - 1, completed)); + }); + } + + let pool = ThreadPool::new(1); + let completed = AtomicUsize::new(0); + + pool.scope(|scope| { + scope.spawn_global(|_| { + recurse(MAX_INLINE_SCOPE_HELP_DEPTH + 8, &completed); + }); + }); + + assert_eq!(completed.load(Ordering::Acquire), 1); + assert!(pool.state.backup_workers_spawned() >= 1); + assert_eq!(pool.state.backup_workers_live(), 0); + assert!(pool.scheduler_metrics().donated_jobs >= 1); +} diff --git a/core-relations/src/free_join/execute.rs b/core-relations/src/free_join/execute.rs index 39d4b0f2b..35b38873a 100644 --- a/core-relations/src/free_join/execute.rs +++ b/core-relations/src/free_join/execute.rs @@ -32,9 +32,12 @@ use crate::{ frame_update::{FrameUpdates, UpdateInstr}, get_index_from_tableinfo, }, - hash_index::{IndexBase, TupleIndex}, + hash_index::{ColumnIndex, Index, IndexBase, TupleIndex}, offsets::{Offsets, RowId, SortedOffsetSlice, SortedOffsetVector, Subset}, - parallel_heuristics::{action_batch_size, free_join_fork_depth, parallelize_db_level_op}, + parallel_heuristics::{ + MIN_TOP_INDEX_KEYS_PER_WORKER, action_batch_size, free_join_fork_depth, + parallelize_db_level_op, + }, pool::Pooled, query::RuleSet, row_buffer::TaggedRowBuffer, @@ -50,6 +53,17 @@ use super::{ const SMALL_RESIDUAL: usize = 8; +fn top_index_shape_is_eligible( + workers: usize, + leader_keys: usize, + nonempty_shards: usize, + min_keys_per_worker: usize, +) -> bool { + workers > 1 + && leader_keys >= min_keys_per_worker.saturating_mul(workers) + && nonempty_shards >= workers +} + struct SparseColumnIndex { n_keys: usize, n_subsets: usize, @@ -244,18 +258,123 @@ impl SortedColumnIndex { } } -enum DynamicIndex { +/// A table-index slot retained for one `run_rule_set` call. +/// +/// The slot lazily acquires its Arc through the existing fully-refreshing +/// catalog helper on first cached use. Keeping that Arc in an execution-scoped +/// sidecar removes catalog lookups and refcount traffic from recursive join +/// execution without constructing indexes for plan accesses that choose a +/// residual-local strategy at runtime. Initialized slots are dropped before +/// the database resets its indexes during `merge_all`. +enum PreparedIndexSlot { + Tuple(OnceLock), + Column(OnceLock), + /// The table specification forbids a global cache for at least one key + /// column, so execution must use its existing dynamic-index path. + Uncacheable, +} + +fn columns_are_cacheable(info: &TableInfo, cols: &[ColumnId]) -> bool { + cols.iter().all(|col| { + !info + .spec + .uncacheable_columns + .get(*col) + .copied() + .unwrap_or(false) + }) +} + +/// Index handles for one immutable [`JoinStages`] value, positionally aligned +/// with `JoinStages::instrs` and with each stage's scans. +struct PreparedJoinIndexes { + stages: Box<[SmallVec<[PreparedIndexSlot; 4]>]>, +} + +impl PreparedJoinIndexes { + fn new(db: &Database, atoms: &Arc>, stages: &JoinStages) -> Self { + let stages = stages + .instrs + .iter() + .map(|stage| { + let mut handles = SmallVec::new(); + match stage { + JoinStage::Intersect { scans, .. } => { + handles.extend(scans.iter().map(|scan| { + let info = &db.tables[atoms[scan.atom].table]; + if !columns_are_cacheable(info, &[scan.column]) { + PreparedIndexSlot::Uncacheable + } else { + PreparedIndexSlot::Column(OnceLock::new()) + } + })); + } + JoinStage::FusedIntersect { to_intersect, .. } + | JoinStage::FusedIntersectMat { to_intersect, .. } => { + handles.extend(to_intersect.iter().map(|(scan, _)| { + let cols = scan.to_index.vars.as_slice(); + let info = &db.tables[atoms[scan.to_index.atom].table]; + if !columns_are_cacheable(info, cols) { + PreparedIndexSlot::Uncacheable + } else if cols.len() == 1 { + PreparedIndexSlot::Column(OnceLock::new()) + } else { + PreparedIndexSlot::Tuple(OnceLock::new()) + } + })); + } + } + handles + }) + .collect(); + Self { stages } + } + + fn stage(&self, index: usize) -> &[PreparedIndexSlot] { + &self.stages[index] + } +} + +/// Execution-scoped index sidecar mirroring the shape of a logical [`Plan`]. +enum PreparedPlanIndexes { + Single(PreparedJoinIndexes), + Decomposed { + blocks: Vec, + result: PreparedJoinIndexes, + }, +} + +impl PreparedPlanIndexes { + fn new(db: &Database, plan: &Plan) -> Self { + match plan { + Plan::SinglePlan(plan) => { + Self::Single(PreparedJoinIndexes::new(db, &plan.atoms, &plan.stages)) + } + Plan::DecomposedPlan(plan) => Self::Decomposed { + blocks: plan + .stages + .blocks + .iter() + .map(|(stages, _)| PreparedJoinIndexes::new(db, &plan.atoms, stages)) + .collect(), + result: PreparedJoinIndexes::new(db, &plan.atoms, &plan.result_block), + }, + } + } +} + +enum DynamicIndex<'plan> { Cached { /// When Some(range), intersect each subset from the index with this dense range. /// The range is the Dense outer subset known at Prober construction time. intersect_outer: Option, - table: HashIndex, + table: &'plan Index, }, CachedColumn { /// When Some(range), intersect each subset from the index with this dense range. /// The range is the Dense outer subset known at Prober construction time. intersect_outer: Option, - table: HashColumnIndex, + table: &'plan Index, }, Dynamic(TupleIndex), DynamicColumn(Arc), @@ -323,19 +442,19 @@ fn intersect_with_dense_ref<'a>(v: SubsetRef<'a>, range: OffsetRange) -> Option< } } -struct Prober { +struct Prober<'plan> { node: Arc, - ix: DynamicIndex, + ix: DynamicIndex<'plan>, } -impl Prober { +impl Prober<'_> { fn get_subset<'a>(&'a self, key: &'a [Value]) -> Option>> { match &self.ix { DynamicIndex::Cached { intersect_outer, table, } => { - let subset_ref = table.get().unwrap().get_subset(key)?; + let subset_ref = table.get_subset(key)?; let subset = if let Some(range) = intersect_outer { intersect_with_dense_ref(subset_ref, *range)? } else { @@ -348,7 +467,7 @@ impl Prober { table, } => { debug_assert_eq!(key.len(), 1); - let subset_ref = table.get().unwrap().get_subset(&key[0])?; + let subset_ref = table.get_subset(&key[0])?; let subset = if let Some(range) = intersect_outer { intersect_with_dense_ref(subset_ref, *range)? } else { @@ -373,7 +492,7 @@ impl Prober { table, } => { let range = *range; - table.get().unwrap().for_each(|k, v| { + table.for_each(|k, v| { if let Some(res) = intersect_with_dense_ref(v, range) { f(k, PotentiallyStale::maybe_stale(res)) } @@ -382,16 +501,13 @@ impl Prober { DynamicIndex::Cached { intersect_outer: None, table, - } => table - .get() - .unwrap() - .for_each(|k, v| f(k, PotentiallyStale::maybe_stale(v))), + } => table.for_each(|k, v| f(k, PotentiallyStale::maybe_stale(v))), DynamicIndex::CachedColumn { intersect_outer: Some(range), table, } => { let range = *range; - table.get().unwrap().for_each(|k, v| { + table.for_each(|k, v| { if let Some(res) = intersect_with_dense_ref(v, range) { f(&[*k], PotentiallyStale::maybe_stale(res)) } @@ -401,10 +517,7 @@ impl Prober { intersect_outer: None, table, } => { - table - .get() - .unwrap() - .for_each(|k, v| f(&[*k], PotentiallyStale::maybe_stale(v))); + table.for_each(|k, v| f(&[*k], PotentiallyStale::maybe_stale(v))); } DynamicIndex::Dynamic(tab) => { tab.for_each(|k, v| f(k, PotentiallyStale::not_stale(v))); @@ -418,10 +531,98 @@ impl Prober { } } + /// Enumerate one physical shard of a cached table index. + /// + /// Dynamic indexes are deliberately excluded: their storage is private to + /// this prober and has no coarse, pre-existing partition worth scheduling. + fn for_each_shard( + &self, + shard: usize, + mut f: impl FnMut(&[Value], PotentiallyStale), + ) { + match &self.ix { + DynamicIndex::Cached { + intersect_outer, + table, + } => { + let intersect_outer = *intersect_outer; + table.for_each_shard(shard, |key, subset| { + let subset = if let Some(range) = intersect_outer { + let Some(subset) = intersect_with_dense_ref(subset, range) else { + return; + }; + subset + } else { + subset + }; + f(key, PotentiallyStale::maybe_stale(subset)); + }); + } + DynamicIndex::CachedColumn { + intersect_outer, + table, + } => { + let intersect_outer = *intersect_outer; + table.for_each_shard(shard, |key, subset| { + let subset = if let Some(range) = intersect_outer { + let Some(subset) = intersect_with_dense_ref(subset, range) else { + return; + }; + subset + } else { + subset + }; + f(&[*key], PotentiallyStale::maybe_stale(subset)); + }); + } + DynamicIndex::Dynamic(_) + | DynamicIndex::DynamicColumn(_) + | DynamicIndex::SparseColumn(_) => { + unreachable!("only cached indexes expose physical shards") + } + } + } + + /// Return the number of coarse physical partitions when this prober is + /// backed by a cached hash index. + fn shard_count(&self) -> Option { + match &self.ix { + DynamicIndex::Cached { table, .. } => Some(table.shard_count()), + DynamicIndex::CachedColumn { table, .. } => Some(table.shard_count()), + DynamicIndex::Dynamic(_) + | DynamicIndex::DynamicColumn(_) + | DynamicIndex::SparseColumn(_) => None, + } + } + + fn shard_len(&self, shard: usize) -> Option { + match &self.ix { + DynamicIndex::Cached { + intersect_outer: None, + table, + } => Some(table.shard_len(shard)), + DynamicIndex::CachedColumn { + intersect_outer: None, + table, + } => Some(table.shard_len(shard)), + DynamicIndex::Cached { + intersect_outer: Some(_), + .. + } + | DynamicIndex::CachedColumn { + intersect_outer: Some(_), + .. + } + | DynamicIndex::Dynamic(_) + | DynamicIndex::DynamicColumn(_) + | DynamicIndex::SparseColumn(_) => None, + } + } + fn len(&self) -> usize { match &self.ix { - DynamicIndex::Cached { table, .. } => table.get().unwrap().len(), - DynamicIndex::CachedColumn { table, .. } => table.get().unwrap().len(), + DynamicIndex::Cached { table, .. } => table.len(), + DynamicIndex::CachedColumn { table, .. } => table.len(), DynamicIndex::Dynamic(tab) => tab.len(), DynamicIndex::DynamicColumn(tab) => tab.len(), DynamicIndex::SparseColumn(tab) => tab.len(), @@ -459,6 +660,15 @@ impl Database { }; let search_and_apply_timer = Instant::now(); + // Prepare one lazy slot for every table-index access named by the + // immutable plans before any worker starts. A slot acquires and refreshes + // its Arc through the regular catalog helper on first cached use. The + // sidecars are dropped before `merge_all` resets catalog entries. + let prepared_indexes = rule_set + .plans + .values() + .map(|(plan, _, _)| PreparedPlanIndexes::new(self, plan)) + .collect::>(); // let mut rule_reports: HashMap>; let mut rule_reports: HashMap, Vec>; let exec_state = ExecutionState::new(self.read_only_view(), Default::default()); @@ -467,7 +677,9 @@ impl Database { Arc::new(DashMap::default()); let db: &Database = self; egglog_concurrency::scope(|scope| { - for (plan, desc, symbol_map) in rule_set.plans.values() { + for ((plan, desc, symbol_map), prepared_index) in + rule_set.plans.values().zip(&prepared_indexes) + { // TODO: add stats let report_plan = match report_level { ReportLevel::TimeOnly => None, @@ -498,17 +710,24 @@ impl Database { } } - match plan { - Plan::SinglePlan(plan) => { + match (plan, prepared_index) { + (Plan::SinglePlan(plan), PreparedPlanIndexes::Single(prepared)) => { join_state.run_join_stages( &plan.stages, + prepared, &plan.atoms, plan.actions, &mut binding_info, &mut action_buf, ); } - Plan::DecomposedPlan(plan) => { + ( + Plan::DecomposedPlan(plan), + PreparedPlanIndexes::Decomposed { + blocks: prepared_blocks, + result: prepared_result, + }, + ) => { let mut materializations: DenseIdMap< MatId, Arc, RowBuffer>>, @@ -531,8 +750,8 @@ impl Database { ); let mut materializations = Arc::new(materializations); - for (mat_id, stage_block) in - plan.stages.blocks.iter().enumerate() + for (mat_id, (stage_block, prepared_block)) in + plan.stages.blocks.iter().zip(prepared_blocks).enumerate() { let mat_id = MatId::from_usize(mat_id); egglog_concurrency::scope(|stage_scope| { @@ -545,6 +764,7 @@ impl Database { }; join_state.run_join_stages( &stage_block.0, + prepared_block, &plan.atoms, mat_id, &mut binding_info, @@ -570,12 +790,14 @@ impl Database { } join_state.run_join_stages( &plan.result_block, + prepared_result, &plan.atoms, plan.actions, &mut binding_info, &mut action_buf, ); } + _ => unreachable!("prepared plan shape must match logical plan"), } } let search_and_apply_time = search_and_apply_timer.elapsed(); @@ -606,7 +828,9 @@ impl Database { match_counter: match_counter.as_ref(), batches: Default::default(), }; - for (plan, desc, symbol_map) in rule_set.plans.values() { + for ((plan, desc, symbol_map), prepared_index) in + rule_set.plans.values().zip(&prepared_indexes) + { let report_plan = match report_level { ReportLevel::TimeOnly => None, ReportLevel::WithPlan | ReportLevel::StageInfo => { @@ -625,17 +849,24 @@ impl Database { None => break 'eval, } } - match plan { - Plan::SinglePlan(plan) => { + match (plan, prepared_index) { + (Plan::SinglePlan(plan), PreparedPlanIndexes::Single(prepared)) => { join_state.run_join_stages( &plan.stages, + prepared, &plan.atoms, plan.actions, &mut binding_info, &mut action_buf, ); } - Plan::DecomposedPlan(plan) => { + ( + Plan::DecomposedPlan(plan), + PreparedPlanIndexes::Decomposed { + blocks: prepared_blocks, + result: prepared_result, + }, + ) => { let mut materializations = DenseIdMap::with_capacity(plan.stages.blocks.len()); for i in 0..plan.stages.blocks.len() { @@ -654,10 +885,13 @@ impl Database { scratch_val: Default::default(), }; - for (mat_id, stage_block) in plan.stages.blocks.iter().enumerate() { + for (mat_id, (stage_block, prepared_block)) in + plan.stages.blocks.iter().zip(prepared_blocks).enumerate() + { let mat_id = MatId::from_usize(mat_id); join_state.run_join_stages( &stage_block.0, + prepared_block, &plan.atoms, mat_id, &mut binding_info, @@ -673,12 +907,14 @@ impl Database { } join_state.run_join_stages( &plan.result_block, + prepared_result, &plan.atoms, plan.actions, &mut binding_info, &mut action_buf, ); } + _ => unreachable!("prepared plan shape must match logical plan"), } } let search_and_apply_time = search_and_apply_timer.elapsed(); @@ -708,6 +944,10 @@ impl Database { // caused by individual queries. reports[i].num_matches = match_counter.read_matches(plan.actions()); } + // `merge_all` requires each catalog Arc to be uniquely owned so it can + // reset the corresponding ResettableOnceLock. No scoped worker can + // survive to this point. + drop(prepared_indexes); let search_and_apply_time = search_and_apply_timer.elapsed(); let merge_timer = Instant::now(); @@ -752,15 +992,42 @@ struct JoinState<'a> { /// Per-column indexes on a trie node's subset, lazily initialized on first access per column. type ColumnIndexes = IdVec>>; -// Each TrieNode is probed with exactly one column in practice, so we store a single -// (ColumnId, map) pair instead of a per-column IdVec of Mutexes. -// // The child cache (see [`TrieNode::get_cached_trie_node`]): keyed by the bound // value, storing the child node and the edge constraints used to build it. The // stored constraints guard against distinct scans reaching the same // (node, col, value) with different slow constraints; they are almost always // empty, in which case the guard is a cheap length check. -type ChildrenMaps = IdVec, Box<[Constraint]>)>>>; +type ChildEntry = (Arc, Box<[Constraint]>); +pub(crate) type ChildLock = RwLock>; +type ChildrenMaps = IdVec>; + +fn get_or_insert_child( + map: &ChildLock, + value: Value, + edge_cs: &[Constraint], + sub: impl FnOnce() -> Subset, +) -> Arc { + // Optimistic read path: most calls are cache hits, so try a shared lock + // first. A hit is only valid when the edge constraints match. + { + let guard = map.read().unwrap(); + if let Some((node, stored_cs)) = guard.get(&value) + && &**stored_cs == edge_cs + { + return node.clone(); + } + } + // Cache miss (or constraint mismatch): acquire the write lock and insert. + let mut guard = map.write().unwrap(); + if let Some((node, stored_cs)) = guard.get(&value) + && &**stored_cs == edge_cs + { + return node.clone(); + } + let new_node = Arc::new(TrieNode::new(sub())); + guard.insert(value, (new_node.clone(), Box::from(edge_cs))); + new_node +} /// Canonical signature of a trie root: the table plus its sorted header (fast) /// constraints. Distinct signatures get distinct base ids from [`TrieCache`]. @@ -880,11 +1147,10 @@ pub(crate) struct TrieNode { subset: Subset, /// Any cached indexes on this subset. cached_subsets: OnceLock>, - /// Cached child trie nodes, keyed by value. In practice each TrieNode is - /// only ever probed with a single column, so we store one (col, map) pair - /// instead of an IdVec across all columns. When this node is a shared root - /// (or reachable from one), this cache is shared across plans too, so - /// children are shared without any global lookup. + /// Cached child trie nodes, keyed first by column and then by value. Each + /// column's lock is allocated lazily. When this node is a shared root (or + /// reachable from one), this cache is shared across plans too, so children + /// are shared without any global lookup. cached_children: OnceLock>, } @@ -925,6 +1191,15 @@ impl TrieNode { .clone() } + fn child_map(&self, col: ColumnId, arity: usize) -> &ChildLock { + self.cached_children.get_or_init(|| { + let mut vec: Pooled = with_pool_set(|ps| ps.get()); + vec.resize_with(arity, OnceLock::new); + vec + })[col] + .get_or_init(|| RwLock::new(HashMap::default())) + } + /// Return the child node reached by additionally constraining `col = value` /// (and applying `edge_cs`). `sub` computes the child subset and is only /// called on a cache miss. @@ -943,31 +1218,7 @@ impl TrieNode { info: &TableInfo, sub: impl FnOnce() -> Subset, ) -> Arc { - let map = &self.cached_children.get_or_init(|| { - let mut vec: Pooled = with_pool_set(|ps| ps.get()); - vec.resize_with(info.spec.arity(), || RwLock::new(HashMap::default())); - vec - })[col]; - // Optimistic read path: most calls are cache hits, so try a shared lock - // first. A hit is only valid when the edge constraints match. - { - let guard = map.read().unwrap(); - if let Some((node, stored_cs)) = guard.get(&value) - && &**stored_cs == edge_cs - { - return node.clone(); - } - } - // Cache miss (or constraint mismatch): acquire the write lock and insert. - let mut guard = map.write().unwrap(); - if let Some((node, stored_cs)) = guard.get(&value) - && &**stored_cs == edge_cs - { - return node.clone(); - } - let new_node = Arc::new(TrieNode::new(sub())); - guard.insert(value, (new_node.clone(), Box::from(edge_cs))); - new_node + get_or_insert_child(self.child_map(col, info.spec.arity()), value, edge_cs, sub) } } @@ -1012,7 +1263,7 @@ impl BindingInfo { /// Probers returned from [`JoinState::get_index`] will move atom-related state out of the /// [`BindingInfo`]. Once the caller is done using a prober, this method moves it back. - fn move_back(&mut self, atom: AtomId, prober: Prober) { + fn move_back(&mut self, atom: AtomId, prober: Prober<'_>) { self.subsets.insert(atom, prober.node); } @@ -1112,13 +1363,14 @@ impl<'a> JoinState<'a> { Some(subset) } - fn get_index( + fn get_index<'plan>( &self, atoms: &Arc>, atom: AtomId, binding_info: &mut BindingInfo, cols: impl Iterator, - ) -> Prober { + prepared: &'plan PreparedIndexSlot, + ) -> Prober<'plan> { let cols = SmallVec::<[ColumnId; 4]>::from_iter(cols); let trie_node = binding_info.subsets.unwrap_val(atom); let subset = &trie_node.subset; @@ -1132,14 +1384,7 @@ impl<'a> JoinState<'a> { cols[0], )) } else { - let all_cacheable = cols.iter().all(|col| { - !info - .spec - .uncacheable_columns - .get(*col) - .copied() - .unwrap_or(false) - }); + let all_cacheable = columns_are_cacheable(info, &cols); let whole_table = info.table.all(); if let Subset::Dense(range) = subset && all_cacheable @@ -1156,14 +1401,28 @@ impl<'a> JoinState<'a> { // large _or_ it is most of the table, or we already have a cached // index for it, then return it. if cols.len() != 1 { + let PreparedIndexSlot::Tuple(index) = prepared else { + unreachable!("multi-column scan must have a prepared tuple index") + }; + let index = + index.get_or_init(|| get_index_from_tableinfo(info, cols.as_slice())); DynamicIndex::Cached { intersect_outer, - table: get_index_from_tableinfo(info, &cols), + table: index + .get() + .expect("prepared tuple index must already be refreshed"), } } else { + let PreparedIndexSlot::Column(index) = prepared else { + unreachable!("single-column scan must have a prepared column index") + }; + let index = + index.get_or_init(|| get_column_index_from_tableinfo(info, cols[0])); DynamicIndex::CachedColumn { intersect_outer, - table: get_column_index_from_tableinfo(info, cols[0]).clone(), + table: index + .get() + .expect("prepared column index must already be refreshed"), } } } else if cols.len() != 1 { @@ -1178,14 +1437,117 @@ impl<'a> JoinState<'a> { ix: dyn_index, } } - fn get_column_index( + fn get_column_index<'plan>( &self, atoms: &Arc>, binding_info: &mut BindingInfo, atom: AtomId, col: ColumnId, - ) -> Prober { - self.get_index(atoms, atom, binding_info, iter::once(col)) + prepared: &'plan PreparedIndexSlot, + ) -> Prober<'plan> { + self.get_index(atoms, atom, binding_info, iter::once(col), prepared) + } + + /// Describe an eligible cached index that could drive a top-level + /// generic-join intersection. + /// + /// This mirrors the leader selection in [`Self::run_plan`]. Probers move + /// trie nodes out of `binding_info`, so restore every node before returning. + /// Every scan must cover its whole table through an unfiltered cached + /// hash/column index. Requiring this before constructing any prober avoids + /// building a dynamic nonleader index during discovery and then rebuilding + /// it in every shard job. Tiny or badly skewed leaders stay on the existing + /// path as well. + fn top_index_shards( + &self, + stage: &JoinStage, + prepared: &[PreparedIndexSlot], + atoms: &Arc>, + binding_info: &mut BindingInfo, + workers: usize, + ) -> Option> { + let JoinStage::Intersect { scans, .. } = stage else { + return None; + }; + if scans.is_empty() || workers <= 1 { + return None; + } + debug_assert_eq!(scans.len(), prepared.len()); + + // Mirror the cached/unfiltered branch in `get_index` without actually + // constructing any dynamic fallback indexes. + for scan in scans { + let node = &binding_info.subsets[scan.atom]; + let info = &self.db.tables[atoms[scan.atom].table]; + if node.subset.size() <= SMALL_RESIDUAL || !columns_are_cacheable(info, &[scan.column]) + { + return None; + } + let Subset::Dense(subset) = &node.subset else { + return None; + }; + let Subset::Dense(whole_table) = info.table.all() else { + return None; + }; + if subset != &whole_table { + return None; + } + } + + let mut leader = 0; + let mut leader_size = usize::MAX; + let mut probers = Vec::with_capacity(scans.len()); + for (i, (scan, prepared)) in scans.iter().zip(prepared).enumerate() { + let prober = + self.get_column_index(atoms, binding_info, scan.atom, scan.column, prepared); + let size = prober.len(); + if size < leader_size { + leader = i; + leader_size = size; + } + probers.push(prober); + } + let shards = probers[leader].shard_count().and_then(|shard_count| { + let shards = (0..shard_count) + .filter(|shard| probers[leader].shard_len(*shard).unwrap_or(0) != 0) + .collect::>(); + top_index_shape_is_eligible( + workers, + leader_size, + shards.len(), + MIN_TOP_INDEX_KEYS_PER_WORKER, + ) + .then_some(shards) + }); + for (scan, prober) in scans.iter().zip(probers) { + binding_info.move_back(scan.atom, prober); + } + shards + } + + /// Return a coarse partition for the variable already sorted to the top of + /// the join. Later variables are deliberately not promoted: benchmarking + /// found that probing and reordering them was not broadly beneficial. + fn select_top_index_shards( + &self, + stages: &JoinStages, + prepared: &PreparedJoinIndexes, + atoms: &Arc>, + order: &InstrOrder, + binding_info: &mut BindingInfo, + ) -> Option> { + if order.len() == 0 { + return None; + } + + let stage_index = order.get(0); + self.top_index_shards( + &stages.instrs[stage_index], + prepared.stage(stage_index), + atoms, + binding_info, + crate::parallel::current_num_threads(), + ) } /// Runs the free join plan, starting with the header. @@ -1200,6 +1562,7 @@ impl<'a> JoinState<'a> { fn run_join_stages<'buf, A: NumericId + 'buf, BUF: ActionBuffer<'buf, A>>( &self, stages: &'buf JoinStages, + prepared: &'buf PreparedJoinIndexes, atoms: &'buf Arc>, action: A, binding_info: &mut BindingInfo, @@ -1218,13 +1581,73 @@ impl<'a> JoinState<'a> { let mut order = InstrOrder::from_iter(0..stages.instrs.len()); let mut leaf_scans: LeafScans = smallvec::smallvec![false; stages.instrs.len()]; sort_plan_by_size(&mut order, &mut leaf_scans, 0, &stages.instrs, binding_info); + + // A cached top index is already partitioned by hash. Schedule one + // coarse global job per nonempty shard and run each shard's subtree + // serially. Besides avoiding an intermediate key copy, this keeps + // related nested probes on one worker. Selecting this path commits the + // whole subtree to coarse-only parallelism: `recur_global_serial` + // supplies a serial buffer whose recursive work stays inline. Buffers + // that cannot construct an independent partition (the in-place + // executors) decline this path. + if !stages.instrs.is_empty() + && action_buf.supports_global_partition() + && let Some(shards) = + self.select_top_index_shards(stages, prepared, atoms, &order, binding_info) + { + let mut updates = FrameUpdates::with_capacity(0); + for shard in shards { + let db = self.db; + let exec_state_for_factory = self.exec_state.clone(); + let exec_state_for_work = self.exec_state.clone(); + let trie_cache = self.trie_cache.clone(); + action_buf.recur_global_serial( + BorrowedLocalState { + binding_info, + instr_order: &mut order, + leaf_scans: &mut leaf_scans, + updates: &mut updates, + }, + move || exec_state_for_factory.clone(), + move |BorrowedLocalState { + binding_info, + instr_order, + leaf_scans, + .. + }, + buf| { + JoinState { + db, + exec_state: exec_state_for_work, + pool: with_pool_set(|ps| ps.get_pool()), + trie_cache, + } + .run_plan( + stages, + prepared, + atoms, + action, + instr_order, + leaf_scans, + 0, + Some(shard), + binding_info, + buf, + ); + }, + ); + } + return; + } self.run_plan( stages, + prepared, atoms, action, &mut order, &mut leaf_scans, 0, + None, binding_info, action_buf, ); @@ -1232,21 +1655,24 @@ impl<'a> JoinState<'a> { /// The core method for executing a free join plan. /// - /// This method takes the plan, mutable data-structures for variable binding and staging - /// actions, and two indexes: `cur` which is the current stage of the plan to run, and `level` - /// which is the current "fan-out" node we are in. The latter parameter is an experimental - /// index used to detect if we are at the "top" of a plan rather than the "bottom", and is - /// currently used as a heuristic to determine if we should increase parallelism more than the - /// default. + /// This method takes the plan, mutable data structures for variable binding + /// and staging actions, and `cur`, the current stage of the plan. A top-level + /// coarse partition also passes `index_shard`; `Some` is only supplied with + /// a serial action buffer, so no nested parallelism is available in that + /// subtree. Recursive calls clear the value so only the first intersection + /// is restricted to that physical shard, while the serial buffer continues + /// to keep later work inline. #[allow(clippy::too_many_arguments)] fn run_plan<'buf, A: NumericId + 'buf, BUF: ActionBuffer<'buf, A>>( &self, stages: &'buf JoinStages, + prepared: &'buf PreparedJoinIndexes, atoms: &'buf Arc>, action: A, instr_order: &mut InstrOrder, leaf_scans: &mut LeafScans, cur: usize, + index_shard: Option, binding_info: &mut BindingInfo, action_buf: &mut BUF, ) where @@ -1282,7 +1708,10 @@ impl<'a> JoinState<'a> { } // TODO: `supports_parallel_drain`` is a hack because currently // `drain_updates_parallel!`` is a bit slower because of the additional ExecutionState clone. - if cur < free_join_fork_depth() && action_buf.supports_parallel_drain() { + if index_shard.is_none() + && cur < free_join_fork_depth() + && action_buf.supports_parallel_drain() + { drain_updates_parallel!($updates) } else { $updates.drain(|update| match update { @@ -1310,11 +1739,13 @@ impl<'a> JoinState<'a> { } else { self.run_plan( stages, + prepared, atoms, action, instr_order, leaf_scans, cur + 1, + None, binding_info, action_buf, ); @@ -1370,11 +1801,13 @@ impl<'a> JoinState<'a> { } .run_plan( stages, + prepared, atoms, action, instr_order, leaf_scans, cur + 1, + None, binding_info, buf, ); @@ -1403,20 +1836,50 @@ impl<'a> JoinState<'a> { } } + // A sharded top-level job enumerates only its assigned physical index + // shard. Every recursive call clears `index_shard`, so this macro is + // used only by the leading prober of the initial `Intersect` stage. + macro_rules! for_each_leader { + ($prober:expr, $callback:expr) => { + if let Some(shard) = index_shard { + $prober.for_each_shard(shard, $callback) + } else { + $prober.for_each($callback) + } + }; + } + let pool = &self.pool; - match &stages.instrs[instr_order.get(cur)] { + let stage_index = instr_order.get(cur); + let stage = &stages.instrs[stage_index]; + let prepared_indexes = prepared.stage(stage_index); + debug_assert_eq!( + prepared_indexes.len(), + match stage { + JoinStage::Intersect { scans, .. } => scans.len(), + JoinStage::FusedIntersect { to_intersect, .. } + | JoinStage::FusedIntersectMat { to_intersect, .. } => to_intersect.len(), + } + ); + match stage { JoinStage::Intersect { var, scans } => match scans.as_slice() { [] => {} [a] => { if binding_info.has_empty_subset(a.atom) { return; } - let prober = self.get_column_index(atoms, binding_info, a.atom, a.column); + let prober = self.get_column_index( + atoms, + binding_info, + a.atom, + a.column, + &prepared_indexes[0], + ); let info = &self.db.tables[atoms[a.atom].table]; let table = info.table.as_ref(); let has_stale = table.has_stale_rows(); let mut updates = FrameUpdates::with_capacity(cmp::min(chunk_size, cur_size)); - prober.for_each(|val, x| { + for_each_leader!(prober, |val, x| { updates.push_binding(*var, val[0]); if x.size() <= 16 { let sub = refine_subset(x, &a.cs, &table, has_stale, pool); @@ -1448,11 +1911,23 @@ impl<'a> JoinState<'a> { binding_info.move_back(a.atom, prober); } [a, b] => { - let a_prober = self.get_column_index(atoms, binding_info, a.atom, a.column); - let b_prober = self.get_column_index(atoms, binding_info, b.atom, b.column); + let a_prober = self.get_column_index( + atoms, + binding_info, + a.atom, + a.column, + &prepared_indexes[0], + ); + let b_prober = self.get_column_index( + atoms, + binding_info, + b.atom, + b.column, + &prepared_indexes[1], + ); let ((smaller, smaller_scan), (larger, larger_scan)) = - if a_prober.len() < b_prober.len() { + if a_prober.len() <= b_prober.len() { ((&a_prober, a), (&b_prober, b)) } else { ((&b_prober, b), (&a_prober, a)) @@ -1467,7 +1942,7 @@ impl<'a> JoinState<'a> { let small_table = small_info.table.as_ref(); let small_has_stale = small_table.has_stale_rows(); let mut updates = FrameUpdates::with_capacity(cmp::min(chunk_size, cur_size)); - smaller.for_each(|val, small_sub| { + for_each_leader!(smaller, |val, small_sub| { if let Some(large_sub) = larger.get_subset(val) { updates.push_binding(*var, val[0]); if small_sub.size() <= 16 { @@ -1555,9 +2030,14 @@ impl<'a> JoinState<'a> { let mut smallest = 0; let mut smallest_size = usize::MAX; let mut probers = Vec::with_capacity(rest.len()); - for (i, scan) in rest.iter().enumerate() { - let prober = - self.get_column_index(atoms, binding_info, scan.atom, scan.column); + for (i, (scan, prepared)) in rest.iter().zip(prepared_indexes).enumerate() { + let prober = self.get_column_index( + atoms, + binding_info, + scan.atom, + scan.column, + prepared, + ); let size = prober.len(); if size < smallest_size { smallest = i; @@ -1585,7 +2065,7 @@ impl<'a> JoinState<'a> { // Smallest leads the scan let mut updates = FrameUpdates::with_capacity(cmp::min(chunk_size, cur_size)); - probers[smallest].for_each(|key, sub| { + for_each_leader!(probers[smallest], |key, sub| { updates.push_binding(*var, key[0]); for (i, scan) in rest.iter().enumerate() { if i == smallest { @@ -1783,10 +2263,11 @@ impl<'a> JoinState<'a> { spec.to_index.atom, binding_info, spec.to_index.vars.iter().copied(), + &prepared_indexes[i], ), ) }) - .collect::>(); + .collect::); 4]>>(); // Pre-compute has_stale per prober to avoid vtable calls in the hot loop. let index_has_stale: SmallVec<[bool; 4]> = index_probers .iter() @@ -1962,15 +2443,17 @@ impl<'a> JoinState<'a> { let mut updates = FrameUpdates::with_capacity(cmp::min(chunk_size, cur_size)); let probers = to_intersect .iter() - .map(|(spec, _)| { + .zip(prepared_indexes) + .map(|((spec, _), prepared)| { self.get_index( atoms, spec.to_index.atom, binding_info, spec.to_index.vars.iter().copied(), + prepared, ) }) - .collect::>(); + .collect::; 4]>>(); // Pre-compute has_stale per prober to avoid vtable calls in the hot loop. let probers_has_stale: SmallVec<[bool; 4]> = to_intersect .iter() @@ -2125,6 +2608,9 @@ const LOCAL_ACTION_BATCH_SIZE: usize = 128; /// for serial and parallel modes. trait ActionBuffer<'state, A: NumericId>: Send { type AsLocal<'a>: ActionBuffer<'state, A> + where + 'state: 'a; + type AsGlobalSerial<'a>: ActionBuffer<'state, A> where 'state: 'a; @@ -2172,6 +2658,17 @@ trait ActionBuffer<'state, A: NumericId>: Send { work: impl for<'a> FnOnce(BorrowedLocalState<'a>, &mut Self::AsLocal<'a>) + Send + 'state, ); + /// Run one coarse partition as a global pool job, using a buffer whose + /// recursive work and action execution are both serial. Implementations + /// that return `false` from [`Self::supports_global_partition`] execute the + /// callback inline; the scoped implementations enqueue exactly one job. + fn recur_global_serial<'local>( + &mut self, + local: BorrowedLocalState<'local>, + to_exec_state: impl FnMut() -> ExecutionState<'state> + Send + 'state, + work: impl for<'a> FnOnce(BorrowedLocalState<'a>, &mut Self::AsGlobalSerial<'a>) + Send + 'state, + ); + /// The unit at which you should batch updates passed to calls to `recur`, /// potentially depending on the current level of recursion. /// @@ -2188,6 +2685,11 @@ trait ActionBuffer<'state, A: NumericId>: Send { fn supports_parallel_drain(&self) -> bool { true } + + /// Whether this buffer can enqueue independent top-level partitions. + fn supports_global_partition(&self) -> bool { + false + } } /// The action buffer we use if we are executing in a single-threaded @@ -2203,6 +2705,10 @@ impl<'a, 'outer: 'a> ActionBuffer<'a, ActionId> for InPlaceActionBuffer<'outer> = Self where 'a: 'b; + type AsGlobalSerial<'b> + = Self + where + 'a: 'b; fn push_bindings( &mut self, @@ -2248,6 +2754,105 @@ impl<'a, 'outer: 'a> ActionBuffer<'a, ActionId> for InPlaceActionBuffer<'outer> work(local, self) } + fn recur_global_serial<'local>( + &mut self, + local: BorrowedLocalState<'local>, + _to_exec_state: impl FnMut() -> ExecutionState<'a> + Send + 'a, + work: impl for<'b> FnOnce(BorrowedLocalState<'b>, &mut Self) + Send + 'a, + ) { + work(local, self) + } + + fn supports_parallel_drain(&self) -> bool { + false + } +} + +/// Strictly serial action buffer used inside one globally scheduled top-index +/// shard. It shares the rule-set match counter with sibling shards, but executes +/// all recursive join and action work inline. It deliberately has no +/// `needs_flush` flag: its owner unconditionally flushes it once after the shard +/// callback returns, and recursive calls reuse the same buffer synchronously. +struct SerialScopedActionBuffer<'scope> { + rule_set: &'scope RuleSet, + match_counter: Arc, + batches: DenseIdMap, +} + +impl<'scope> SerialScopedActionBuffer<'scope> { + fn new(rule_set: &'scope RuleSet, match_counter: Arc) -> Self { + Self { + rule_set, + match_counter, + batches: Default::default(), + } + } +} + +impl<'scope> ActionBuffer<'scope, ActionId> for SerialScopedActionBuffer<'scope> { + type AsLocal<'a> + = Self + where + 'scope: 'a; + type AsGlobalSerial<'a> + = Self + where + 'scope: 'a; + + fn push_bindings( + &mut self, + action: ActionId, + bindings: &DenseIdMap, + mut to_exec_state: impl FnMut() -> ExecutionState<'scope>, + ) { + let batch_size = action_batch_size(); + let action_state = self + .batches + .get_or_insert(action, || ActionState::new(batch_size)); + action_state.n_runs += 1; + action_state.len += 1; + let action_info = &self.rule_set.actions[action]; + // SAFETY: `used_vars` is constant for the rule and the bindings come + // from this rule's join plan. + unsafe { + action_state.bindings.push(bindings, &action_info.used_vars); + } + if action_state.len >= batch_size { + let succeeded = + to_exec_state().run_instrs(&action_info.instrs, &mut action_state.bindings); + action_state.bindings.clear(); + self.match_counter.inc_matches(action, succeeded); + action_state.len = 0; + } + } + + fn flush(&mut self, exec_state: &mut ExecutionState) { + flush_action_states( + exec_state, + &mut self.batches, + self.rule_set, + self.match_counter.as_ref(), + ); + } + + fn recur<'local>( + &mut self, + local: BorrowedLocalState<'local>, + _to_exec_state: impl FnMut() -> ExecutionState<'scope> + Send + 'scope, + work: impl for<'a> FnOnce(BorrowedLocalState<'a>, &mut Self) + Send + 'scope, + ) { + work(local, self); + } + + fn recur_global_serial<'local>( + &mut self, + local: BorrowedLocalState<'local>, + _to_exec_state: impl FnMut() -> ExecutionState<'scope> + Send + 'scope, + work: impl for<'a> FnOnce(BorrowedLocalState<'a>, &mut Self) + Send + 'scope, + ) { + work(local, self); + } + fn supports_parallel_drain(&self) -> bool { false } @@ -2283,6 +2888,10 @@ impl<'scope> ActionBuffer<'scope, ActionId> for ScopedActionBuffer<'_, 'scope> { = ScopedActionBuffer<'a, 'scope> where 'scope: 'a; + type AsGlobalSerial<'a> + = SerialScopedActionBuffer<'scope> + where + 'scope: 'a; fn push_bindings( &mut self, action: ActionId, @@ -2334,7 +2943,10 @@ impl<'scope> ActionBuffer<'scope, ActionId> for ScopedActionBuffer<'_, 'scope> { let rule_set = self.rule_set; let match_counter = self.match_counter.clone(); let mut inner = local.clone_state(); - self.scope.spawn(move |scope| { + // Keep recursive join work on the current worker when possible. If + // coarse top-index partitioning was unavailable, stalled workers cause + // the private deque to donate older siblings to the global queue. + self.scope.spawn_local(move |scope| { let mut buf: ScopedActionBuffer<'_, 'scope> = ScopedActionBuffer { scope, rule_set, @@ -2354,6 +2966,26 @@ impl<'scope> ActionBuffer<'scope, ActionId> for ScopedActionBuffer<'_, 'scope> { }); } + fn recur_global_serial<'local>( + &mut self, + mut local: BorrowedLocalState<'local>, + mut to_exec_state: impl FnMut() -> ExecutionState<'scope> + Send + 'scope, + work: impl for<'a> FnOnce(BorrowedLocalState<'a>, &mut SerialScopedActionBuffer<'scope>) + + Send + + 'scope, + ) { + let rule_set = self.rule_set; + let match_counter = self.match_counter.clone(); + let mut inner = local.clone_state(); + self.scope.spawn_global(move |_| { + let mut buf = SerialScopedActionBuffer::new(rule_set, match_counter); + work(inner.borrow_mut(), &mut buf); + // Unlike a nested `ScopedActionBuffer`, this buffer is owned by the + // one shard job and is always flushed exactly once before it exits. + buf.flush(&mut to_exec_state()); + }); + } + fn morsel_size(&mut self, _level: usize, _total: usize) -> usize { // Lower morsel size to increase parallelism. match _level { @@ -2361,6 +2993,10 @@ impl<'scope> ActionBuffer<'scope, ActionId> for ScopedActionBuffer<'_, 'scope> { _ => 256, } } + + fn supports_global_partition(&self) -> bool { + true + } } fn expand_binding_sets<'state, A: NumericId, BUF: ActionBuffer<'state, A> + ?Sized>( @@ -2435,6 +3071,10 @@ impl<'a> ActionBuffer<'a, MatId> for InPlaceMaterializer<'a> { = Self where 'a: 'b; + type AsGlobalSerial<'b> + = Self + where + 'a: 'b; fn push_bindings( &mut self, @@ -2480,6 +3120,103 @@ impl<'a> ActionBuffer<'a, MatId> for InPlaceMaterializer<'a> { work(local, self) } + fn recur_global_serial<'local>( + &mut self, + local: BorrowedLocalState<'local>, + _to_exec_state: impl FnMut() -> ExecutionState<'a> + Send + 'a, + work: impl for<'b> FnOnce(BorrowedLocalState<'b>, &mut Self) + Send + 'a, + ) { + work(local, self) + } + + fn supports_parallel_drain(&self) -> bool { + false + } +} + +/// Serial recursive view of a scoped materializer. Sibling top-level shards +/// still share the concurrent output maps, but this buffer never schedules +/// deeper work itself. +struct SerialScopedMaterializer { + specs: Arc>, + materializations: Arc, RowBuffer>>>>, + scratch_key: Vec, + scratch_val: Vec, +} + +impl SerialScopedMaterializer { + fn new( + specs: Arc>, + materializations: Arc, RowBuffer>>>>, + ) -> Self { + Self { + specs, + materializations, + scratch_key: Vec::new(), + scratch_val: Vec::new(), + } + } +} + +impl<'scope> ActionBuffer<'scope, MatId> for SerialScopedMaterializer { + type AsLocal<'a> + = Self + where + 'scope: 'a; + type AsGlobalSerial<'a> + = Self + where + 'scope: 'a; + + fn push_bindings( + &mut self, + mat_id: MatId, + bindings: &DenseIdMap, + _to_exec_state: impl FnMut() -> ExecutionState<'scope>, + ) { + let mat = self.materializations.get(mat_id).expect("invalid mat id"); + let spec = self.specs.get(mat_id).expect("invalid mat id"); + self.scratch_key.clear(); + self.scratch_key + .extend(spec.msg_vars.iter().map(|var| bindings[*var])); + self.scratch_val.clear(); + self.scratch_val + .extend(spec.val_vars.iter().map(|var| bindings[*var])); + if self.scratch_val.is_empty() { + self.scratch_val.push(Value::stale()); + } + match mat.entry(self.scratch_key.clone()) { + Entry::Occupied(mut occupied) => { + occupied.get_mut().add_row(&self.scratch_val); + } + Entry::Vacant(vacant) => { + let mut buffer = RowBuffer::new(usize::max(spec.val_vars.len(), 1)); + buffer.add_row(&self.scratch_val); + vacant.insert(buffer); + } + } + } + + fn flush(&mut self, _exec_state: &mut ExecutionState) {} + + fn recur<'local>( + &mut self, + local: BorrowedLocalState<'local>, + _to_exec_state: impl FnMut() -> ExecutionState<'scope> + Send + 'scope, + work: impl for<'a> FnOnce(BorrowedLocalState<'a>, &mut Self) + Send + 'scope, + ) { + work(local, self); + } + + fn recur_global_serial<'local>( + &mut self, + local: BorrowedLocalState<'local>, + _to_exec_state: impl FnMut() -> ExecutionState<'scope> + Send + 'scope, + work: impl for<'a> FnOnce(BorrowedLocalState<'a>, &mut Self) + Send + 'scope, + ) { + work(local, self); + } + fn supports_parallel_drain(&self) -> bool { false } @@ -2497,6 +3234,10 @@ impl<'scope> ActionBuffer<'scope, MatId> for ScopedMaterializer<'_, 'scope> { = ScopedMaterializer<'a, 'scope> where 'scope: 'a; + type AsGlobalSerial<'a> + = SerialScopedMaterializer + where + 'scope: 'a; fn push_bindings( &mut self, @@ -2546,7 +3287,9 @@ impl<'scope> ActionBuffer<'scope, MatId> for ScopedMaterializer<'_, 'scope> { let specs = self.specs.clone(); let materializations = self.materializations.clone(); let mut inner = local.clone_state(); - scope.spawn(move |scope| { + // Match the action path: preserve locality by default and rely on + // private-deque donation when other workers need fallback work. + scope.spawn_local(move |scope| { let mut buf: ScopedMaterializer<'_, 'scope> = ScopedMaterializer { scope, specs, @@ -2557,6 +3300,25 @@ impl<'scope> ActionBuffer<'scope, MatId> for ScopedMaterializer<'_, 'scope> { work(inner.borrow_mut(), &mut buf); }); } + + fn recur_global_serial<'local>( + &mut self, + mut local: BorrowedLocalState<'local>, + _to_exec_state: impl FnMut() -> ExecutionState<'scope> + Send + 'scope, + work: impl for<'a> FnOnce(BorrowedLocalState<'a>, &mut SerialScopedMaterializer) + Send + 'scope, + ) { + let specs = self.specs.clone(); + let materializations = self.materializations.clone(); + let mut inner = local.clone_state(); + self.scope.spawn_global(move |_| { + let mut buf = SerialScopedMaterializer::new(specs, materializations); + work(inner.borrow_mut(), &mut buf); + }); + } + + fn supports_global_partition(&self) -> bool { + true + } } struct MatchCounter { @@ -2725,7 +3487,7 @@ fn sort_plan_by_size_inner( let mut times_refined = with_pool_set(|ps| ps.get::>()); // Count how many times each atom has been refined so far. - for ins in instrs[..range.start].iter() { + for ins in &instrs[..range.start] { match ins { JoinStage::Intersect { scans, .. } => scans.iter().for_each(|scan| { *times_refined.get_or_default(scan.atom) += 1; @@ -2886,3 +3648,58 @@ impl LocalState { } } } + +#[cfg(test)] +mod top_index_tests { + use std::{ + sync::{ + Arc, Barrier, + atomic::{AtomicUsize, Ordering}, + }, + thread, + }; + + use crate::{common::Value, numeric_id::NumericId, offsets::Subset}; + + use super::{ChildLock, get_or_insert_child, top_index_shape_is_eligible}; + + #[test] + fn top_index_partitioning_rejects_serial_tiny_and_skewed_shapes() { + assert!(!top_index_shape_is_eligible(1, 10_000, 8, 64)); + assert!(!top_index_shape_is_eligible(4, 255, 8, 64)); + assert!(!top_index_shape_is_eligible(4, 10_000, 3, 64)); + assert!(top_index_shape_is_eligible(4, 256, 4, 64)); + assert!(top_index_shape_is_eligible(4, 40, 4, 10)); + } + + #[test] + fn child_cache_returns_one_arc_for_racing_same_key() { + const THREADS: usize = 16; + let map = Arc::new(ChildLock::default()); + let barrier = Arc::new(Barrier::new(THREADS)); + let constructed = Arc::new(AtomicUsize::new(0)); + let value = Value::from_usize(17); + + let handles = (0..THREADS) + .map(|_| { + let map = map.clone(); + let barrier = barrier.clone(); + let constructed = constructed.clone(); + thread::spawn(move || { + barrier.wait(); + get_or_insert_child(&map, value, &[], || { + constructed.fetch_add(1, Ordering::Relaxed); + Subset::empty() + }) + }) + }) + .collect::>(); + let nodes = handles + .into_iter() + .map(|handle| handle.join().unwrap()) + .collect::>(); + + assert_eq!(constructed.load(Ordering::Relaxed), 1); + assert!(nodes[1..].iter().all(|node| Arc::ptr_eq(&nodes[0], node))); + } +} diff --git a/core-relations/src/hash_index/mod.rs b/core-relations/src/hash_index/mod.rs index 839667fe5..0cd2f0cb3 100644 --- a/core-relations/src/hash_index/mod.rs +++ b/core-relations/src/hash_index/mod.rs @@ -1,4 +1,6 @@ //! Hash-based secondary indexes. +#[cfg(test)] +use std::sync::atomic::{AtomicUsize, Ordering}; use std::{ cmp, hash::{Hash, Hasher}, @@ -132,6 +134,25 @@ impl Index { self.table.for_each(f); } + /// Call `f` over the elements stored in one physical index shard. + /// + /// The shards form a disjoint partition of the index keys. Exposing that + /// partition lets query execution create a small number of coarse tasks + /// without first copying keys into a separate work list. + pub(crate) fn for_each_shard(&self, shard: usize, f: impl FnMut(&TI::Key, SubsetRef)) { + self.table.for_each_shard(shard, f); + } + + /// Return the number of physical shards in this index. + pub(crate) fn shard_count(&self) -> usize { + self.table.shard_count() + } + + /// Return the number of distinct keys stored in one physical shard. + pub(crate) fn shard_len(&self, shard: usize) -> usize { + self.table.shard_len(shard) + } + pub(crate) fn len(&self) -> usize { self.table.len() } @@ -181,6 +202,12 @@ pub(crate) trait IndexBase { fn merge_rows(&mut self, buf: &TaggedRowBuffer); /// Call `f` over the elements of the index. fn for_each(&self, f: impl FnMut(&Self::Key, SubsetRef)); + /// Call `f` over the elements of one physical index shard. + fn for_each_shard(&self, shard: usize, f: impl FnMut(&Self::Key, SubsetRef)); + /// The number of physical shards in the index. + fn shard_count(&self) -> usize; + /// The number of distinct keys in one physical shard. + fn shard_len(&self, shard: usize) -> usize; /// The number of keys in the index. fn len(&self) -> usize; @@ -281,6 +308,21 @@ impl IndexBase for ColumnIndex { } } + fn for_each_shard(&self, shard: usize, mut f: impl FnMut(&Self::Key, SubsetRef)) { + let shard = &self.shards[ShardId::from_usize(shard)]; + for (key, subset) in shard.table.iter() { + f(key, subset.as_ref(&shard.subsets)); + } + } + + fn shard_count(&self) -> usize { + self.shards.len() + } + + fn shard_len(&self, shard: usize) -> usize { + self.shards[ShardId::from_usize(shard)].table.len() + } + fn len(&self) -> usize { self.shards.iter().map(|(_, shard)| shard.table.len()).sum() } @@ -720,6 +762,23 @@ impl IndexBase for TupleIndex { } } + fn for_each_shard(&self, shard: usize, mut f: impl FnMut(&Self::Key, SubsetRef)) { + let shard = &self.shards[ShardId::from_usize(shard)]; + for entry in shard.table.hash.iter() { + // SAFETY: entry.key was stored by add_row, so it is always in-bounds. + let key = unsafe { shard.table.keys.get_row_unchecked(entry.key) }; + f(key, entry.vals.as_ref(&shard.subsets)); + } + } + + fn shard_count(&self) -> usize { + self.shards.len() + } + + fn shard_len(&self, shard: usize) -> usize { + self.shards[ShardId::from_usize(shard)].table.hash.len() + } + fn len(&self) -> usize { self.shards .iter() @@ -830,6 +889,8 @@ fn hash_key(key: &[Value]) -> u64 { #[derive(Default)] pub struct IndexCatalog { data: ReadOptimizedLock>, + #[cfg(test)] + get_or_insert_calls: AtomicUsize, } impl IndexCatalog @@ -839,6 +900,8 @@ where pub fn new() -> Self { IndexCatalog { data: ReadOptimizedLock::new(Vec::new()), + #[cfg(test)] + get_or_insert_calls: AtomicUsize::new(0), } } @@ -846,6 +909,8 @@ where let vec = self.data.read().iter().map(f).collect(); IndexCatalog { data: ReadOptimizedLock::new(vec), + #[cfg(test)] + get_or_insert_calls: AtomicUsize::new(0), } } @@ -856,6 +921,8 @@ where } pub fn get_or_insert(&self, k: K, init: impl FnOnce() -> I) -> I { + #[cfg(test)] + self.get_or_insert_calls.fetch_add(1, Ordering::Relaxed); let data = self.data.read(); let entry = data.iter().find(|(k1, _)| k1 == &k); if let Some(entry) = entry { @@ -872,6 +939,11 @@ where } } } + + #[cfg(test)] + pub(crate) fn get_or_insert_calls(&self) -> usize { + self.get_or_insert_calls.load(Ordering::Relaxed) + } } define_id!(BufferIndex, u32, "an index into a subset buffer"); diff --git a/core-relations/src/hash_index/tests.rs b/core-relations/src/hash_index/tests.rs index 35afc1ab8..d80914367 100644 --- a/core-relations/src/hash_index/tests.rs +++ b/core-relations/src/hash_index/tests.rs @@ -226,3 +226,63 @@ fn column_index_rebuild_matches_oracle() { } } } + +#[test] +fn physical_shards_partition_index_iteration() { + ThreadPool::new(4).install(|| { + let rows = (0..257) + .map(|i| vec![v(i), v(i % 37), v(10_000 + i)]) + .collect::>(); + let table = WrappedTable::new(fill_table(rows, 1, None, |old, new| { + assert_eq!(old, new, "unique keys, so no conflicts"); + None + })); + + let mut column = Index::new(vec![ColumnId::new(1)], ColumnIndex::new()); + column.refresh(table.as_ref()); + assert_eq!(column.shard_count(), 8); + let mut whole_column = BTreeMap::new(); + column.for_each(|key, subset| { + whole_column.insert(key.rep(), subset.size()); + }); + let mut by_column_shard = BTreeMap::new(); + for shard in 0..column.shard_count() { + let mut shard_keys = 0; + column.for_each_shard(shard, |key, subset| { + shard_keys += 1; + assert!( + by_column_shard.insert(key.rep(), subset.size()).is_none(), + "a column-index key appeared in more than one shard" + ); + }); + assert_eq!(column.shard_len(shard), shard_keys); + } + assert_eq!(by_column_shard, whole_column); + + let mut tuple = Index::new( + vec![ColumnId::new(1), ColumnId::new(2)], + crate::TupleIndex::new(2), + ); + tuple.refresh(table.as_ref()); + assert_eq!(tuple.shard_count(), 8); + let mut whole_tuple = BTreeMap::new(); + tuple.for_each(|key, subset| { + whole_tuple.insert((key[0].rep(), key[1].rep()), subset.size()); + }); + let mut by_tuple_shard = BTreeMap::new(); + for shard in 0..tuple.shard_count() { + let mut shard_keys = 0; + tuple.for_each_shard(shard, |key, subset| { + shard_keys += 1; + assert!( + by_tuple_shard + .insert((key[0].rep(), key[1].rep()), subset.size()) + .is_none(), + "a tuple-index key appeared in more than one shard" + ); + }); + assert_eq!(tuple.shard_len(shard), shard_keys); + } + assert_eq!(by_tuple_shard, whole_tuple); + }); +} diff --git a/core-relations/src/parallel_heuristics.rs b/core-relations/src/parallel_heuristics.rs index d6861470a..e52bf3cf9 100644 --- a/core-relations/src/parallel_heuristics.rs +++ b/core-relations/src/parallel_heuristics.rs @@ -14,6 +14,10 @@ const DEFAULT_TABLE_OP_CUTOFF: usize = 400_000; const DEFAULT_FREE_JOIN_FORK_DEPTH: usize = 2; const DEFAULT_ACTION_BATCH_SIZE: usize = 8 * 1024; +/// Minimum number of top-index leader keys required per worker before coarse +/// generic-join partitioning is worthwhile. +pub(crate) const MIN_TOP_INDEX_KEYS_PER_WORKER: usize = 16; + static CUTOFFS: OnceLock = OnceLock::new(); struct Cutoffs { diff --git a/core-relations/src/pool/mod.rs b/core-relations/src/pool/mod.rs index d62574490..0695cdddf 100644 --- a/core-relations/src/pool/mod.rs +++ b/core-relations/src/pool/mod.rs @@ -12,7 +12,7 @@ use std::{ use crate::{ AtomId, - free_join::execute::TrieNode, + free_join::execute::ChildLock, numeric_id::{DenseIdMap, IdVec}, }; use fixedbitset::FixedBitSet; @@ -442,7 +442,7 @@ pool_set! { cached_subsets: IdVec>> [ 4 << 20 ], intersected_on: DenseIdMap [ 1 << 20 ], - cached_child: IdVec, Box<[Constraint]>)>>> [ 1 << 20 ], + cached_child: IdVec> [ 1 << 20 ], } } diff --git a/core-relations/src/table/mod.rs b/core-relations/src/table/mod.rs index 8a520f4ef..dbfaca9e0 100644 --- a/core-relations/src/table/mod.rs +++ b/core-relations/src/table/mod.rs @@ -627,7 +627,7 @@ impl SortedWritesTable { let (actual_shard, hc) = hash_code(shard_data, to_remove, n_keys); assert_eq!(actual_shard, shard_id); if let Ok(entry) = shard.find_entry(hc, |entry| { - entry.hashcode == (hc as _) + entry.hashcode == (hc as HashCode) && &data.get_row(entry.row)[0..n_keys] == to_remove }) { let (ent, _) = entry.remove(); @@ -665,7 +665,7 @@ impl SortedWritesTable { let (actual_shard, hc) = hash_code(shard_data, to_remove, self.n_keys); assert_eq!(actual_shard, shard_id); if let Ok(entry) = shard.find_entry(hc, |entry| { - entry.hashcode == (hc as _) + entry.hashcode == (hc as HashCode) && &self.data.get_row(entry.row).unwrap()[0..self.n_keys] == to_remove }) { diff --git a/core-relations/src/tests.rs b/core-relations/src/tests.rs index ec6a14742..9cb9bc6d7 100644 --- a/core-relations/src/tests.rs +++ b/core-relations/src/tests.rs @@ -12,7 +12,10 @@ use crate::{ PlanStrategy, action::WriteVal, common::Value, - free_join::{CounterId, Database, TableId}, + free_join::{ + CounterId, Database, TableId, + plan::{JoinStage, Plan}, + }, make_external_func, query::RuleSetBuilder, table::SortedWritesTable, @@ -210,6 +213,575 @@ fn line_graph_1_test(strat: PlanStrategy) { assert_eq!(expected, got); } +#[test] +fn prepared_plan_indexes_refresh_across_runs_and_clear() { + let mut db = Database::default(); + let make_table = |arity| { + SortedWritesTable::new( + arity, + arity, + None, + vec![], + Box::new(|_, old, new, _| { + assert_eq!(old, new, "test tables have unique keys"); + false + }), + ) + }; + let driver = db.add_table(make_table(2), iter::empty(), iter::empty()); + let target = db.add_table(make_table(3), iter::empty(), iter::empty()); + let allowed = db.add_table(make_table(1), iter::empty(), iter::empty()); + let output = db.add_table(make_table(1), iter::empty(), iter::empty()); + + const KEYS: usize = 4; + const VALUES_PER_KEY: usize = 32; + let populate_inputs = |db: &Database, value_base: usize| { + { + let mut buf = db.new_buffer(target); + for x in 0..KEYS { + for z in 0..VALUES_PER_KEY { + buf.stage_insert(&[v(x), v(10_000 + x), v(value_base + x * 100 + z)]); + } + } + } + { + let mut buf = db.new_buffer(allowed); + for x in 0..KEYS { + for z in 0..VALUES_PER_KEY { + buf.stage_insert(&[v(value_base + x * 100 + z)]); + } + } + } + }; + + { + let mut buf = db.new_buffer(driver); + for x in 0..KEYS { + buf.stage_insert(&[v(x), v(10_000 + x)]); + } + } + populate_inputs(&db, 20_000); + db.merge_all(); + + let mut rsb = db.new_rule_set(); + let mut query = rsb.new_rule(); + query.set_plan_strategy(PlanStrategy::PureSize); + query.set_no_decomp(true); + let x = query.new_var_named("x"); + let y = query.new_var_named("y"); + let z = query.new_var_named("z"); + query.add_atom(driver, &[x.into(), y.into()], &[]).unwrap(); + query + .add_atom(target, &[x.into(), y.into(), z.into()], &[]) + .unwrap(); + query.add_atom(allowed, &[z.into()], &[]).unwrap(); + let mut rule = query.build(); + rule.insert(output, &[z.into()]).unwrap(); + rule.build_with_description("prepared-index-lifecycle"); + let rules = rsb.build(); + + let (plan, _, _) = rules.plans.values().next().unwrap(); + let Plan::SinglePlan(plan) = plan else { + panic!("set_no_decomp must produce a single plan") + }; + assert!( + plan.stages.instrs.iter().any(|stage| matches!( + stage, + JoinStage::FusedIntersect { to_intersect, .. } + if to_intersect + .iter() + .any(|(scan, _)| scan.to_index.vars.len() > 1) + )), + "test must exercise a prepared tuple index" + ); + assert!( + plan.stages + .instrs + .iter() + .any(|stage| matches!(stage, JoinStage::Intersect { scans, .. } if !scans.is_empty())), + "test must exercise a prepared column index" + ); + + let planned_tuple_lookups = plan + .stages + .instrs + .iter() + .map(|stage| match stage { + JoinStage::Intersect { .. } => 0, + JoinStage::FusedIntersect { to_intersect, .. } + | JoinStage::FusedIntersectMat { to_intersect, .. } => to_intersect + .iter() + .filter(|(scan, _)| { + plan.atoms[scan.to_index.atom].table == target && scan.to_index.vars.len() != 1 + }) + .count(), + }) + .sum::(); + let planned_allowed_column_lookups = plan + .stages + .instrs + .iter() + .map(|stage| match stage { + JoinStage::Intersect { scans, .. } => scans + .iter() + .filter(|scan| plan.atoms[scan.atom].table == allowed) + .count(), + JoinStage::FusedIntersect { to_intersect, .. } + | JoinStage::FusedIntersectMat { to_intersect, .. } => to_intersect + .iter() + .filter(|(scan, _)| { + plan.atoms[scan.to_index.atom].table == allowed && scan.to_index.vars.len() == 1 + }) + .count(), + }) + .sum::(); + assert!(planned_tuple_lookups > 0); + assert!(planned_allowed_column_lookups > 0); + + let expected = KEYS * VALUES_PER_KEY; + let tuple_lookups_before = db.get_table_info(target).indexes.get_or_insert_calls(); + let allowed_column_lookups_before = db + .get_table_info(allowed) + .column_indexes + .get_or_insert_calls(); + let first = db.run_rule_set(&rules, ReportLevel::TimeOnly); + assert_eq!(first.num_matches("prepared-index-lifecycle"), expected); + assert_eq!(db.get_table(output).len(), expected); + assert_eq!( + db.get_table_info(target).indexes.get_or_insert_calls() - tuple_lookups_before, + planned_tuple_lookups, + "tuple-index catalog access must be plan-static, not recursion-static" + ); + assert_eq!( + db.get_table_info(allowed) + .column_indexes + .get_or_insert_calls() + - allowed_column_lookups_before, + planned_allowed_column_lookups, + "column-index catalog access must be plan-static, not recursion-static" + ); + + // Reusing the same logical plan must not retain catalog Arcs across the + // merge at the end of the previous run. + let second = db.run_rule_set(&rules, ReportLevel::TimeOnly); + assert_eq!(second.num_matches("prepared-index-lifecycle"), expected); + assert_eq!(db.get_table(output).len(), expected); + + db.clear_table(target); + db.clear_table(allowed); + db.clear_table(output); + populate_inputs(&db, 40_000); + db.merge_all(); + + // Clearing bumps the table generation, so every prepared global index must + // be refreshed before its borrowed reference is exposed to execution. + let after_clear = db.run_rule_set(&rules, ReportLevel::TimeOnly); + assert_eq!( + after_clear.num_matches("prepared-index-lifecycle"), + expected + ); + let output_table = db.get_table(output); + let output_rows = output_table.scan(output_table.all().as_ref()); + assert_eq!(output_rows.len(), expected); + assert!(output_rows.iter().all(|(_, row)| row[0].index() >= 40_000)); +} + +#[test] +fn prepared_plan_indexes_preserve_uncacheable_columns() { + let mut db = Database::default(); + let displaced = db.add_table(DisplacedTable::default(), iter::empty(), iter::empty()); + let make_table = || { + SortedWritesTable::new( + 1, + 1, + None, + vec![], + Box::new(|_, old, new, _| { + assert_eq!(old, new); + false + }), + ) + }; + let representatives = db.add_table(make_table(), iter::empty(), iter::empty()); + let output = db.add_table(make_table(), iter::empty(), iter::empty()); + { + let mut buf = db.new_buffer(displaced); + buf.stage_insert(&[v(0), v(10), v(0)]); + } + db.merge_all(); + let displaced_table = db.get_table(displaced); + let displaced_rows = displaced_table.scan(displaced_table.all().as_ref()); + let displaced_row = displaced_rows.iter().next().unwrap().1; + let displaced_key = displaced_row[0]; + let canonical = displaced_row[1]; + let displaced_timestamp = displaced_row[2]; + { + let mut buf = db.new_buffer(representatives); + buf.stage_insert(&[canonical]); + } + db.merge_all(); + + let mut rsb = db.new_rule_set(); + let mut query = rsb.new_rule(); + query.set_plan_strategy(PlanStrategy::Gj); + query.set_no_decomp(true); + let representative = query.new_var_named("representative"); + query + .add_atom( + displaced, + &[ + displaced_key.into(), + representative.into(), + displaced_timestamp.into(), + ], + &[], + ) + .unwrap(); + query + .add_atom(representatives, &[representative.into()], &[]) + .unwrap(); + let mut rule = query.build(); + rule.insert(output, &[representative.into()]).unwrap(); + rule.build_with_description("uncacheable-prepared-index"); + let rules = rsb.build(); + + let (plan, _, _) = rules.plans.values().next().unwrap(); + let Plan::SinglePlan(plan) = plan else { + panic!("set_no_decomp must produce a single plan") + }; + assert!(plan.stages.instrs.iter().any(|stage| matches!( + stage, + JoinStage::Intersect { scans, .. } + if scans.iter().any(|scan| { + plan.atoms[scan.atom].table == displaced && scan.column == ColumnId::new(1) + }) + ))); + + let catalog_calls_before = db + .get_table_info(displaced) + .column_indexes + .get_or_insert_calls(); + let report = db.run_rule_set(&rules, ReportLevel::TimeOnly); + assert_eq!(report.num_matches("uncacheable-prepared-index"), 1); + assert_eq!( + db.get_table_info(displaced) + .column_indexes + .get_or_insert_calls(), + catalog_calls_before, + "preparation must not create a global index for a dynamic column" + ); + assert_eq!(db.get_table(output).len(), 1); +} + +#[test] +fn gj_top_index_shards_preserve_count_and_materialization() { + gj_top_index_shards_preserve_count_and_materialization_inner(); +} + +fn gj_top_index_shards_preserve_count_and_materialization_inner() { + let pool = egglog_concurrency::ThreadPool::new(4); + pool.install(|| { + let mut db = Database::default(); + let make_table = || { + SortedWritesTable::new( + 2, + 2, + None, + vec![], + Box::new(|_, old, new, _| { + assert_eq!(old, new, "test tables have unique keys"); + false + }), + ) + }; + let left = db.add_table(make_table(), iter::empty(), iter::empty()); + let right = db.add_table(make_table(), iter::empty(), iter::empty()); + let output = db.add_table( + SortedWritesTable::new( + 2, + 2, + None, + vec![], + Box::new(|_, old, new, _| { + assert_eq!(old, new, "materialized pairs are unique"); + false + }), + ), + iter::empty(), + iter::empty(), + ); + + const KEYS: usize = 513; + const DEGREE: usize = 17; + { + let mut buf = db.new_buffer(left); + for x in 0..KEYS { + for y in 0..DEGREE { + buf.stage_insert(&[v(x), v(100_000 + x * DEGREE + y)]); + } + } + } + { + let mut buf = db.new_buffer(right); + for x in 0..KEYS { + for z in 0..DEGREE { + buf.stage_insert(&[v(x), v(200_000 + x * DEGREE + z)]); + } + } + } + db.merge_all(); + + let mut rsb = RuleSetBuilder::new(&mut db); + let mut query = rsb.new_rule(); + query.set_plan_strategy(PlanStrategy::Gj); + let x = query.new_var_named("x"); + let y = query.new_var_named("y"); + let z = query.new_var_named("z"); + query.add_atom(left, &[x.into(), y.into()], &[]).unwrap(); + query.add_atom(right, &[x.into(), z.into()], &[]).unwrap(); + let mut rule = query.build(); + rule.insert(output, &[y.into(), z.into()]).unwrap(); + rule.build_with_description("sharded-gj"); + let rules = rsb.build(); + + pool.reset_scheduler_metrics(); + let report = db.run_rule_set(&rules, ReportLevel::TimeOnly); + let scheduler = pool.scheduler_metrics(); + // The root rule job plus one job for each of the cached index's + // 2 * worker-count shards. Merge work may add more global jobs. + assert!( + scheduler.global_pushes >= 9, + "top index sharding did not activate: {scheduler:?}" + ); + assert_eq!( + scheduler.local_pushes, 0, + "the fixed scheduler policy should keep each shard subtree serial" + ); + let expected = KEYS * DEGREE * DEGREE; + assert_eq!(report.num_matches("sharded-gj"), expected); + + let table = db.get_table(output); + let materialized = table.scan(table.all().as_ref()); + assert_eq!(materialized.len(), expected); + for (_, row) in materialized.iter() { + let y = row[0].index() - 100_000; + let z = row[1].index() - 200_000; + assert_eq!(y / DEGREE, z / DEGREE); + } + }); +} + +#[test] +fn gj_small_top_fallback_uses_local_queue() { + let pool = egglog_concurrency::ThreadPool::new(4); + pool.install(|| { + let mut db = Database::default(); + let make_table = || { + SortedWritesTable::new( + 3, + 3, + None, + vec![], + Box::new(|_, old, new, _| { + assert_eq!(old, new, "test tables have unique keys"); + false + }), + ) + }; + let left = db.add_table(make_table(), iter::empty(), iter::empty()); + let right = db.add_table(make_table(), iter::empty(), iter::empty()); + let output = db.add_table(make_table(), iter::empty(), iter::empty()); + + // The sorted top variable has only two keys and is therefore too small + // for coarse index-shard partitioning, while the database is large + // enough to enable parallel rule execution. + const ROWS: usize = 8192; + for table in [left, right] { + let mut buf = db.new_buffer(table); + for i in 0..ROWS { + buf.stage_insert(&[v(i % 2), v(10_000 + i), v(20_000 + i)]); + } + } + db.merge_all(); + + let mut rsb = RuleSetBuilder::new(&mut db); + let mut query = rsb.new_rule(); + query.set_plan_strategy(PlanStrategy::Gj); + query.set_no_decomp(true); + let x = query.new_var_named("x"); + let y = query.new_var_named("y"); + let z = query.new_var_named("z"); + query + .add_atom(left, &[x.into(), y.into(), z.into()], &[]) + .unwrap(); + query + .add_atom(right, &[x.into(), y.into(), z.into()], &[]) + .unwrap(); + let mut rule = query.build(); + rule.insert(output, &[x.into(), y.into(), z.into()]) + .unwrap(); + rule.build_with_description("local-fallback-gj"); + let rules = rsb.build(); + + pool.reset_scheduler_metrics(); + let report = db.run_rule_set(&rules, ReportLevel::TimeOnly); + let scheduler = pool.scheduler_metrics(); + assert!( + scheduler.local_pushes > 0, + "small-top fallback did not use worker-local tasks: {scheduler:?}" + ); + assert_eq!( + scheduler.local_pushes, + scheduler.local_pops + scheduler.donated_jobs, + "every local task must run on its owner or be donated" + ); + assert_eq!(report.num_matches("local-fallback-gj"), ROWS); + + let table = db.get_table(output); + let materialized = table.scan(table.all().as_ref()); + assert_eq!(materialized.len(), ROWS); + for (_, row) in materialized.iter() { + let i = row[2].index() - 20_000; + assert_eq!(row[0].index(), i % 2); + assert_eq!(row[1].index() - 10_000, i); + } + }); +} + +#[test] +fn gj_decomposed_small_top_materialization_uses_local_queue() { + let pool = egglog_concurrency::ThreadPool::new(4); + pool.install(|| { + let mut db = Database::default(); + let make_table = |arity| { + SortedWritesTable::new( + arity, + arity, + None, + vec![], + Box::new(|_, old, new, _| { + assert_eq!(old, new, "test tables have unique keys"); + false + }), + ) + }; + let left_xa = db.add_table(make_table(2), iter::empty(), iter::empty()); + let left_ab = db.add_table(make_table(2), iter::empty(), iter::empty()); + let left_bx = db.add_table(make_table(2), iter::empty(), iter::empty()); + let right_xc = db.add_table(make_table(2), iter::empty(), iter::empty()); + let right_cd = db.add_table(make_table(2), iter::empty(), iter::empty()); + let right_dx = db.add_table(make_table(2), iter::empty(), iter::empty()); + let output = db.add_table(make_table(5), iter::empty(), iter::empty()); + + // Two cyclic bags joined through x force a decomposed plan. The left + // bag is large, but x has only two keys, so its materialization block + // must use recursive fallback work instead of top-index sharding. + const X_KEYS: usize = 2; + const LEFT_ROWS_PER_KEY: usize = 4096; + for x in 0..X_KEYS { + let mut xa = db.new_buffer(left_xa); + let mut ab = db.new_buffer(left_ab); + let mut bx = db.new_buffer(left_bx); + for i in 0..LEFT_ROWS_PER_KEY { + let a = 10_000 + x * LEFT_ROWS_PER_KEY + i; + let b = 20_000 + x * LEFT_ROWS_PER_KEY + i; + xa.stage_insert(&[v(x), v(a)]); + ab.stage_insert(&[v(a), v(b)]); + bx.stage_insert(&[v(b), v(x)]); + } + } + { + let mut xc = db.new_buffer(right_xc); + let mut cd = db.new_buffer(right_cd); + let mut dx = db.new_buffer(right_dx); + for x in 0..X_KEYS { + let c = 100_000 + x; + let d = 200_000 + x; + xc.stage_insert(&[v(x), v(c)]); + cd.stage_insert(&[v(c), v(d)]); + dx.stage_insert(&[v(d), v(x)]); + } + } + db.merge_all(); + + let mut rsb = RuleSetBuilder::new(&mut db); + let mut query = rsb.new_rule(); + query.set_plan_strategy(PlanStrategy::Gj); + let x = query.new_var_named("x"); + let a = query.new_var_named("a"); + let b = query.new_var_named("b"); + let c = query.new_var_named("c"); + let d = query.new_var_named("d"); + query.add_atom(left_xa, &[x.into(), a.into()], &[]).unwrap(); + query.add_atom(left_ab, &[a.into(), b.into()], &[]).unwrap(); + query.add_atom(left_bx, &[b.into(), x.into()], &[]).unwrap(); + query + .add_atom(right_xc, &[x.into(), c.into()], &[]) + .unwrap(); + query + .add_atom(right_cd, &[c.into(), d.into()], &[]) + .unwrap(); + query + .add_atom(right_dx, &[d.into(), x.into()], &[]) + .unwrap(); + let mut rule = query.build(); + rule.insert(output, &[x.into(), a.into(), b.into(), c.into(), d.into()]) + .unwrap(); + rule.build_with_description("local-materialization-fallback-gj"); + let rules = rsb.build(); + + let (plan, _, _) = rules.plans.values().next().unwrap(); + let Plan::DecomposedPlan(plan) = plan else { + panic!("the two cyclic bags must produce a decomposed plan") + }; + assert!(plan.stages.blocks.len() >= 2); + let first_materialization = &plan.stages.blocks[0].0; + assert!(first_materialization.instrs.len() >= 3); + let Some(JoinStage::Intersect { var, scans }) = first_materialization.instrs.first() else { + panic!("the first materialization block must start by intersecting x") + }; + assert_eq!(*var, x); + assert!( + scans.iter().any(|scan| { + let table = plan.atoms[scan.atom].table; + table == right_xc || table == right_dx + }), + "the top x intersection must include a two-row index, making coarse sharding ineligible" + ); + + pool.reset_scheduler_metrics(); + let report = db.run_rule_set(&rules, ReportLevel::TimeOnly); + let scheduler = pool.scheduler_metrics(); + assert!( + scheduler.local_pushes > 0, + "small-top materialization fallback did not use worker-local tasks: {scheduler:?}" + ); + assert_eq!( + scheduler.local_pushes, + scheduler.local_pops + scheduler.donated_jobs, + "every local task must run on its owner or be donated" + ); + + let expected = X_KEYS * LEFT_ROWS_PER_KEY; + assert_eq!( + report.num_matches("local-materialization-fallback-gj"), + expected + ); + let table = db.get_table(output); + let materialized = table.scan(table.all().as_ref()); + assert_eq!(materialized.len(), expected); + for (_, row) in materialized.iter() { + let x = row[0].index(); + let i = row[1].index() - 10_000 - x * LEFT_ROWS_PER_KEY; + assert!(x < X_KEYS); + assert!(i < LEFT_ROWS_PER_KEY); + assert_eq!(row[2].index(), 20_000 + x * LEFT_ROWS_PER_KEY + i); + assert_eq!(row[3].index(), 100_000 + x); + assert_eq!(row[4].index(), 200_000 + x); + } + }); +} + #[test] fn line_graph_2_fj_puresize() { run_serial_and_parallel(|| line_graph_2_test(PlanStrategy::PureSize));