From 40bec9948e6c6450ace78e192e38eb5e3c0ed62e Mon Sep 17 00:00:00 2001 From: Eli Rosenthal Date: Thu, 23 Jul 2026 23:30:49 +0200 Subject: [PATCH] Add demand-driven coarse join scheduling Introduce explicit global and worker-local spawning, local-first nested helping, private queues, and timed half-queue donation when workers stall. Cover the scheduler with deterministic tests and expose queue, donation, and stalled-time instrumentation. For eligible cached top-index scans, preserve the physical shard partitioning and respond to idle workers by publishing bounded coarse batches to the global queue. Each batch executes its complete lower subtree serially for locality, and fixes its grain after the first demand signal to avoid geometric task proliferation. In five-run paired graph-motif tests, this improves Q1 by 0.8%, 6.8%, and 8.7% at 8, 12, and 16 threads and Q2 by 2.4%, 6.6%, and 3.7%. CPU cost is effectively flat except for a 2.6% increase on Q2 at 12 threads, while scheduler idle time drops substantially. --- concurrency/src/threadpool/mod.rs | 156 +++++-- concurrency/src/threadpool/tests.rs | 190 ++++++++- core-relations/src/free_join/execute.rs | 523 ++++++++++++++++++------ 3 files changed, 717 insertions(+), 152 deletions(-) diff --git a/concurrency/src/threadpool/mod.rs b/concurrency/src/threadpool/mod.rs index 1d7b3ae9b..8151d9683 100644 --- a/concurrency/src/threadpool/mod.rs +++ b/concurrency/src/threadpool/mod.rs @@ -3,11 +3,12 @@ //! [`ThreadPool`] owns a fixed set of primary worker threads. Each worker //! 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. +//! owner depth-first execution. At a bounded cadence, a worker with stalled +//! peers packages the oldest half of its private work as one global batch. 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. @@ -120,6 +121,7 @@ const COMPLETED_MASK: u64 = u32::MAX as u64; // Keep inline helping bounded so pathological nested scopes move onto a fresh // stack before exhausting the current worker stack. const MAX_INLINE_SCOPE_HELP_DEPTH: usize = 64; +const LOCAL_DONATION_INTERVAL: Duration = Duration::from_millis(10); type Job = Box; type ScopedJob<'scope> = Box; @@ -150,6 +152,8 @@ pub struct SchedulerMetrics { pub local_pops: u64, /// Jobs moved from private queues to the global queue. pub donated_jobs: u64, + /// Global queue batches created from donated private jobs. + pub donation_batches: u64, /// Aggregate wall-clock time for which primary workers were stalled. /// /// This includes the elapsed portion of stalls that are still in progress @@ -179,6 +183,9 @@ impl SchedulerMetrics { 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), + donation_batches: self + .donation_batches + .saturating_sub(earlier.donation_batches), stalled_time: self .stalled_time .checked_sub(earlier.stalled_time) @@ -488,6 +495,7 @@ struct WorkerInstrumentation { local_pushes: AtomicU64, local_pops: AtomicU64, donated_jobs: AtomicU64, + donation_batches: AtomicU64, max_local_queue_depth: AtomicUsize, stall_timer: Mutex, } @@ -498,6 +506,7 @@ impl WorkerInstrumentation { local_pushes: AtomicU64::new(0), local_pops: AtomicU64::new(0), donated_jobs: AtomicU64::new(0), + donation_batches: AtomicU64::new(0), max_local_queue_depth: AtomicUsize::new(0), stall_timer: Mutex::new(StallTimerState::default()), } @@ -559,6 +568,11 @@ impl SchedulerInstrumentation { .iter() .map(|worker| worker.donated_jobs.load(Ordering::Relaxed)) .sum(); + let donation_batches = self + .workers + .iter() + .map(|worker| worker.donation_batches.load(Ordering::Relaxed)) + .sum(); let max_local_queue_depth = self .workers .iter() @@ -572,6 +586,7 @@ impl SchedulerInstrumentation { local_pushes, local_pops, donated_jobs, + donation_batches, stalled_time, stalled_workers: self.stalled_workers.load(Ordering::Acquire), max_local_queue_depth, @@ -595,6 +610,7 @@ impl SchedulerInstrumentation { worker.local_pushes.store(0, Ordering::Relaxed); worker.local_pops.store(0, Ordering::Relaxed); worker.donated_jobs.store(0, Ordering::Relaxed); + worker.donation_batches.store(0, Ordering::Relaxed); worker.max_local_queue_depth.store(0, Ordering::Relaxed); } @@ -637,14 +653,24 @@ impl SchedulerInstrumentation { .fetch_add(1, Ordering::Relaxed); } - fn record_donation(&self, worker_id: usize, jobs: usize) { - self.workers[worker_id] + fn record_donation(&self, worker_id: usize, jobs: usize, batches: usize) { + let worker = &self.workers[worker_id]; + worker .donated_jobs .fetch_add(jobs as u64, Ordering::Relaxed); + worker + .donation_batches + .fetch_add(batches as u64, Ordering::Relaxed); + } + + fn stalled_workers(&self) -> usize { + // This counter is only an advisory scheduling hint. Queue operations + // provide the synchronization needed to publish and consume jobs. + self.stalled_workers.load(Ordering::Relaxed) } fn has_stalled_workers(&self) -> bool { - self.stalled_workers.load(Ordering::Acquire) != 0 + self.stalled_workers() != 0 } fn begin_stall(&self, worker_id: usize) { @@ -898,6 +924,7 @@ struct WorkerContext { pool: ThreadPoolStatePtr, worker_id: usize, local: UnsafeCell>, + last_donation: Cell, } impl WorkerContext { @@ -906,6 +933,7 @@ impl WorkerContext { pool, worker_id, local: UnsafeCell::new(VecDeque::new()), + last_donation: Cell::new(Instant::now()), } } @@ -913,21 +941,34 @@ impl WorkerContext { ptr::eq(self.pool.0, pool) } + fn local_queue_depth(&self) -> usize { + // SAFETY: a `WorkerContext` is installed only on its owning worker and + // `CURRENT_WORKER` is never made available to another thread. + unsafe { (&*self.local.get()).len() } + } + 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 (depth, was_empty) = unsafe { let queue = &mut *self.local.get(); + let was_empty = queue.is_empty(); queue.push_back(job); - queue.len() + (queue.len(), was_empty) }; + // A worker that has been idle for longer than the interval should + // still get one full collection window. Otherwise its first two local + // jobs can be donated immediately using a stale timestamp. + if was_empty { + self.last_donation.set(Instant::now()); + } pool.instrumentation .record_local_push(self.worker_id, depth); - self.donate_half_if_stalled(pool); + self.donate_half_if_due(pool); } fn pop_local(&self, pool: &ThreadPoolState) -> Option { - self.donate_half_if_stalled(pool); + self.donate_half_if_due(pool); // The owner takes the newest work, retaining depth-first locality. // SAFETY: see `push_local`. let job = unsafe { (&mut *self.local.get()).pop_back() }; @@ -937,20 +978,31 @@ impl WorkerContext { job } - fn donate_half_if_stalled(&self, pool: &ThreadPoolState) { + fn donate_half_if_due(&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. + let now = Instant::now(); + if now.saturating_duration_since(self.last_donation.get()) < LOCAL_DONATION_INTERVAL { + return; + } + + // Donation takes the opposite end from the owner. Package the oldest + // half as one global job so its recipient can process related siblings + // without repeated channel traffic or cache migration. // 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); + if donated.is_empty() { + return; + } + + self.last_donation.set(now); + self.publish_donation_batch(pool, donated); } fn donate_all(&self, pool: &ThreadPoolState) { @@ -958,16 +1010,30 @@ impl WorkerContext { // 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); + self.publish_donation_individually(pool, donated); + } + + fn publish_donation_batch(&self, pool: &ThreadPoolState, donated: Vec) { + if donated.is_empty() { + return; + } + + pool.instrumentation + .record_donation(self.worker_id, donated.len(), 1); + pool.enqueue_global(Box::new(move || { + for job in donated { + job(); + } + })); } - fn publish_donation(&self, pool: &ThreadPoolState, donated: Vec) { + fn publish_donation_individually(&self, pool: &ThreadPoolState, donated: Vec) { if donated.is_empty() { return; } pool.instrumentation - .record_donation(self.worker_id, donated.len()); + .record_donation(self.worker_id, donated.len(), donated.len()); for job in donated { pool.enqueue_global(job); } @@ -1036,6 +1102,14 @@ fn pop_current_worker_local(pool: &ThreadPoolState) -> Option { }) } +fn current_worker_local_queue_depth(pool: &ThreadPoolState) -> usize { + with_current_worker(|worker| { + worker + .filter(|worker| worker.belongs_to(pool)) + .map_or(0, WorkerContext::local_queue_depth) + }) +} + fn donate_all_from_current_worker(pool: &ThreadPoolState) { with_current_worker(|worker| { if let Some(worker) = worker.filter(|worker| worker.belongs_to(pool)) { @@ -1227,10 +1301,11 @@ impl<'scope> Scope<'scope> { /// 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. + /// global queue. At a bounded cadence, a worker with stalled peers packages + /// the oldest half of its private queue as one global batch. 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, @@ -1243,6 +1318,41 @@ impl<'scope> Scope<'scope> { } } + /// Return the number of primary workers currently waiting for global work. + /// + /// This is a cheap, relaxed, advisory snapshot of the scheduler's atomic + /// stalled worker gauge. The value may become stale immediately and must + /// not be used for correctness decisions or as a reservation of idle + /// capacity. Unlike [`ThreadPool::scheduler_metrics`], this method does not + /// collect counters, lock stall timers, or scan the pool's workers. + pub fn stalled_workers(&self) -> usize { + // 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() } + .instrumentation + .stalled_workers() + } + + /// Return whether any primary worker currently appears to need global work. + /// + /// This is the boolean convenience form of [`Scope::stalled_workers`] and + /// has the same advisory, inherently racy semantics. + pub fn has_stalled_workers(&self) -> bool { + // SAFETY: see `Scope::stalled_workers`. + unsafe { self.pool.as_ref() } + .instrumentation + .has_stalled_workers() + } + + /// Return the current worker's private queue depth. + /// + /// The result is exact for the current worker because its deque is private. + /// Calls made outside a primary worker of this pool return zero. + pub fn local_queue_depth(&self) -> usize { + // SAFETY: see `Scope::stalled_workers`. + current_worker_local_queue_depth(unsafe { self.pool.as_ref() }) + } + fn prepare_job(&self, f: F) -> Job where F: FnOnce(&Scope<'scope>) + Send + 'scope, diff --git a/concurrency/src/threadpool/tests.rs b/concurrency/src/threadpool/tests.rs index 36186df5e..b9733d3d5 100644 --- a/concurrency/src/threadpool/tests.rs +++ b/concurrency/src/threadpool/tests.rs @@ -14,8 +14,8 @@ use crate::{ }; use super::{ - MAX_INLINE_SCOPE_HELP_DEPTH, ThreadPoolState, ThreadPoolStatePtr, WorkerContext, - is_background_worker_thread, + LOCAL_DONATION_INTERVAL, MAX_INLINE_SCOPE_HELP_DEPTH, ThreadPoolState, ThreadPoolStatePtr, + WorkerContext, is_background_worker_thread, with_current_worker, }; #[test] @@ -128,7 +128,7 @@ fn nested_scope_help_takes_local_work_before_global_work() { } #[test] -fn donation_moves_half_of_the_oldest_local_jobs() { +fn timed_donation_batches_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); @@ -146,17 +146,18 @@ fn donation_moves_half_of_the_oldest_local_jobs() { state .instrumentation .stalled_workers - .store(1, Ordering::Release); + .store(3, Ordering::Release); + worker + .last_donation + .set(Instant::now().checked_sub(LOCAL_DONATION_INTERVAL).unwrap()); - worker.donate_half_if_stalled(&state); + worker.donate_half_if_due(&state); state .instrumentation .stalled_workers .store(0, Ordering::Release); - for _ in 0..3 { - state.try_recv_global().unwrap()(); - } + state.try_recv_global().unwrap()(); while let Some(job) = worker.pop_local(&state) { job(); } @@ -164,11 +165,167 @@ fn donation_moves_half_of_the_oldest_local_jobs() { 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.donation_batches, 1); + assert_eq!(metrics.global_pushes, 1); + assert_eq!(metrics.global_pops, 1); assert_eq!(metrics.local_pops, 3); } +#[test] +fn timed_donation_keeps_a_single_local_job_for_the_owner() { + let (sender, receiver) = super::unbounded(); + let state = ThreadPoolState::new(sender, receiver, 1); + let worker = WorkerContext::new(ThreadPoolStatePtr::new(&state), 0); + let completed = Arc::new(AtomicBool::new(false)); + + unsafe { + let completed = completed.clone(); + (&mut *worker.local.get()).push_back(Box::new(move || { + completed.store(true, Ordering::Release); + })); + } + state + .instrumentation + .stalled_workers + .store(1, Ordering::Release); + worker + .last_donation + .set(Instant::now().checked_sub(LOCAL_DONATION_INTERVAL).unwrap()); + + worker.donate_half_if_due(&state); + state + .instrumentation + .stalled_workers + .store(0, Ordering::Release); + + worker.pop_local(&state).unwrap()(); + assert!(completed.load(Ordering::Acquire)); + + let metrics = state.instrumentation.snapshot(); + assert_eq!(metrics.donated_jobs, 0); + assert_eq!(metrics.donation_batches, 0); + assert_eq!(metrics.global_pushes, 0); + assert_eq!(metrics.global_pops, 0); + assert_eq!(metrics.local_pops, 1); +} + +#[test] +fn first_local_push_after_idle_starts_a_fresh_donation_window() { + let (sender, receiver) = super::unbounded(); + let state = ThreadPoolState::new(sender, receiver, 1); + let worker = WorkerContext::new(ThreadPoolStatePtr::new(&state), 0); + let stale = Instant::now() + .checked_sub(LOCAL_DONATION_INTERVAL + LOCAL_DONATION_INTERVAL) + .unwrap(); + worker.last_donation.set(stale); + state + .instrumentation + .stalled_workers + .store(1, Ordering::Relaxed); + + worker.push_local(&state, Box::new(|| {})); + + assert!(worker.last_donation.get() > stale); + assert_eq!(worker.local_queue_depth(), 1); + assert_eq!(state.instrumentation.snapshot().donated_jobs, 0); + + state + .instrumentation + .stalled_workers + .store(0, Ordering::Relaxed); + worker.pop_local(&state).unwrap()(); +} + +#[test] +fn donation_interval_suppresses_repeated_batches() { + let (sender, receiver) = super::unbounded(); + let state = ThreadPoolState::new(sender, receiver, 3); + let worker = WorkerContext::new(ThreadPoolStatePtr::new(&state), 0); + worker.push_local(&state, Box::new(|| {})); + worker + .last_donation + .set(Instant::now() + Duration::from_secs(1)); + + state + .instrumentation + .stalled_workers + .store(1, Ordering::Relaxed); + for _ in 1..6 { + worker.push_local(&state, Box::new(|| {})); + } + assert_eq!(worker.local_queue_depth(), 6); + assert_eq!(state.instrumentation.snapshot().donated_jobs, 0); + + worker + .last_donation + .set(Instant::now().checked_sub(LOCAL_DONATION_INTERVAL).unwrap()); + worker.donate_half_if_due(&state); + assert_eq!(worker.local_queue_depth(), 3); + let metrics = state.instrumentation.snapshot(); + assert_eq!(metrics.donated_jobs, 3); + assert_eq!(metrics.donation_batches, 1); + assert_eq!(metrics.global_pushes, 1); + assert_eq!(metrics.max_local_queue_depth, 6); + + worker.donate_half_if_due(&state); + assert_eq!(worker.local_queue_depth(), 3); + assert_eq!(state.instrumentation.snapshot().global_pushes, 1); + + state + .instrumentation + .stalled_workers + .store(0, Ordering::Relaxed); + while let Ok(job) = state.try_recv_global() { + job(); + } + while let Some(job) = worker.pop_local(&state) { + job(); + } +} + +#[test] +fn scope_stalled_worker_accessors_read_the_atomic_gauge() { + let (sender, receiver) = super::unbounded(); + let state = ThreadPoolState::new(sender, receiver, 4); + let scope = Scope::new(&state); + + assert_eq!(scope.stalled_workers(), 0); + assert!(!scope.has_stalled_workers()); + assert_eq!(scope.local_queue_depth(), 0); + + state + .instrumentation + .stalled_workers + .store(3, Ordering::Release); + assert_eq!(scope.stalled_workers(), 3); + assert!(scope.has_stalled_workers()); + + state + .instrumentation + .stalled_workers + .store(0, Ordering::Release); + assert_eq!(scope.stalled_workers(), 0); + assert!(!scope.has_stalled_workers()); +} + +#[test] +fn scope_reports_current_workers_local_queue_depth() { + let pool = ThreadPool::new(1); + let observed_depth = AtomicUsize::new(0); + + pool.scope(|scope| { + assert_eq!(scope.local_queue_depth(), 0); + scope.spawn_global(|scope| { + assert_eq!(scope.local_queue_depth(), 0); + scope.spawn_local(|_| {}); + scope.spawn_local(|_| {}); + observed_depth.store(scope.local_queue_depth(), Ordering::Relaxed); + }); + }); + + assert_eq!(observed_depth.load(Ordering::Relaxed), 2); +} + #[test] fn stalled_worker_signal_causes_public_local_work_to_be_donated() { let pool = ThreadPool::new(2); @@ -181,7 +338,14 @@ fn stalled_worker_signal_causes_public_local_work_to_be_donated() { pool.scope(|scope| { scope.spawn_global(|scope| { - for _ in 0..16 { + scope.spawn_local(|_| {}); + with_current_worker(|worker| { + worker + .unwrap() + .last_donation + .set(Instant::now().checked_sub(LOCAL_DONATION_INTERVAL).unwrap()); + }); + for _ in 1..16 { scope.spawn_local(|_| {}); } }); @@ -194,9 +358,9 @@ fn stalled_worker_signal_causes_public_local_work_to_be_donated() { metrics.local_pushes, metrics.local_pops + metrics.donated_jobs ); - assert_eq!(metrics.global_pushes, 1 + metrics.donated_jobs); + assert_eq!(metrics.global_pushes, 1 + metrics.donation_batches); assert_eq!(metrics.global_pushes, metrics.global_pops); - assert!(metrics.max_local_queue_depth >= 2); + assert!(metrics.max_local_queue_depth >= 1); } #[test] diff --git a/core-relations/src/free_join/execute.rs b/core-relations/src/free_join/execute.rs index af3041a6c..92c4fe99e 100644 --- a/core-relations/src/free_join/execute.rs +++ b/core-relations/src/free_join/execute.rs @@ -359,6 +359,17 @@ fn top_index_shape_is_eligible( && nonempty_shards >= workers } +/// One physical cached-index shard used to drive the top join level. +/// +/// `scan_size` is the exact number of leader keys in the shard. Besides +/// avoiding another lookup during execution, it lets a demand-driven executor +/// coalesce many top bindings into roughly one recursive job per worker. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +struct TopIndexPartition { + shard: usize, + scan_size: usize, +} + /// A key ordinal in a root continuation grid. /// /// Persistent catalog indexes produce an [`IndexPosition`], while round-local @@ -3057,7 +3068,7 @@ impl<'a, 'state, 'exec> JoinState<'a, 'state, 'exec> { atoms: &Arc>, binding_info: &mut BindingInfo<'rows, 'exec>, workers: usize, - ) -> Option> + ) -> Option> where 'exec: 'rows, { @@ -3115,7 +3126,10 @@ impl<'a, 'state, 'exec> JoinState<'a, 'state, 'exec> { } 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) + .filter_map(|shard| { + let scan_size = probers[leader].shard_len(shard).unwrap_or(0); + (scan_size != 0).then_some(TopIndexPartition { shard, scan_size }) + }) .collect::>(); top_index_shape_is_eligible( workers, @@ -3141,7 +3155,7 @@ impl<'a, 'state, 'exec> JoinState<'a, 'state, 'exec> { atoms: &Arc>, order: &InstrOrder, binding_info: &mut BindingInfo<'rows, 'exec>, - ) -> Option> + ) -> Option> where 'exec: 'rows, { @@ -3218,7 +3232,7 @@ impl<'a, 'state, 'exec> JoinState<'a, 'state, 'exec> { self.select_top_index_shards(stages, prepared, atoms, &order, binding_info) { let mut updates = FrameUpdates::with_capacity(0); - for shard in shards { + for partition in shards { let db = self.db; let exec_state_for_factory = self.exec_state; let exec_state_for_work = self.exec_state; @@ -3253,7 +3267,7 @@ impl<'a, 'state, 'exec> JoinState<'a, 'state, 'exec> { instr_order, leaf_scans, 0, - Some(shard), + Some(partition), all_stages, binding_info, buf, @@ -3282,11 +3296,10 @@ impl<'a, 'state, 'exec> JoinState<'a, 'state, 'exec> { /// /// 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. + /// coarse partition also passes `index_partition`; recursive calls clear it + /// so only the first intersection is restricted to that physical shard. + /// The associated top-shard buffer policy, rather than the presence of the + /// partition itself, determines whether later recursive work may fan out. #[allow(clippy::too_many_arguments)] fn run_plan<'plan, 'rows, 'scope, A: NumericId + 'scope, BUF: ActionBuffer<'scope, 'exec, A>>( &self, @@ -3297,7 +3310,7 @@ impl<'a, 'state, 'exec> JoinState<'a, 'state, 'exec> { instr_order: &mut InstrOrder, leaf_scans: &mut LeafScans, cur: usize, - index_shard: Option, + index_partition: Option, remaining_stages: Option, binding_info: &mut BindingInfo<'rows, 'exec>, action_buf: &mut BUF, @@ -3332,8 +3345,10 @@ impl<'a, 'state, 'exec> JoinState<'a, 'state, 'exec> { ); return; } - let chunk_size = action_buf.morsel_size(cur, instr_order.len()); - let mut cur_size = estimate_size(&stages.instrs[instr_order.get(cur)], binding_info); + let mut cur_size = index_partition.map_or_else( + || estimate_size(&stages.instrs[instr_order.get(cur)], binding_info), + |partition| partition.scan_size, + ); if cur_size > 32 && cur % 3 == 1 && cur < instr_order.len() - 1 { // Re-evaluate the remaining suffix after observing the residuals // produced by earlier stages. Packed child families make the @@ -3341,6 +3356,13 @@ impl<'a, 'state, 'exec> JoinState<'a, 'state, 'exec> { sort_plan_by_size(instr_order, leaf_scans, cur, &stages.instrs, binding_info); cur_size = estimate_size(&stages.instrs[instr_order.get(cur)], binding_info); } + // A rejected intersection candidate consumes a leader key without + // producing a frame. Subtracting completed frames therefore maintains + // a conservative upper bound, which is sufficient for choosing a + // locality-oriented morsel size. + let mut remaining_scan_upper_bound = cur_size; + let mut chunk_size = + action_buf.morsel_size(cur, instr_order.len(), remaining_scan_upper_bound); let stage_index = instr_order.get(cur); let stage = &stages.instrs[stage_index]; @@ -3355,62 +3377,6 @@ impl<'a, 'state, 'exec> JoinState<'a, 'state, 'exec> { remaining & !stage_bit }); - // Helper macro (not its own method to appease the borrow checker). - macro_rules! drain_updates { - ($updates:expr) => { - if self.exec_state.should_stop() { - return; - } - // TODO: `supports_parallel_drain`` is a hack because currently - // `drain_updates_parallel!`` is a bit slower because of the additional ExecutionState clone. - 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 { - UpdateInstr::PushBinding(var, val) => { - binding_info.bindings.insert(var, val); - } - UpdateInstr::RefineAtom(atom, subset) => { - binding_info.insert_node(atom, subset); - } - UpdateInstr::RefineAtomDense(atom, range) => { - binding_info.insert_subset(atom, Subset::Dense(range)); - } - UpdateInstr::EndFrame => { - // Inline leaf-level: if cur+1 is the leaf (no more - // join stages), call push_bindings directly without - // a recursive run_plan call, avoiding function call - // overhead + an extra should_stop() check. - if cur + 1 >= instr_order.len() { - action_buf.push_bindings_factorized( - action, - &mut binding_info.bindings, - &binding_info.binding_sets, - self.exec_state, - ); - } else { - self.run_plan( - stages, - prepared, - atoms, - action, - instr_order, - leaf_scans, - cur + 1, - None, - remaining_after_current, - binding_info, - action_buf, - ); - } - } - }) - } - }; - } macro_rules! drain_updates_parallel { ($updates:expr) => {{ if self.exec_state.should_stop() { @@ -3473,14 +3439,83 @@ impl<'a, 'state, 'exec> JoinState<'a, 'state, 'exec> { $updates.clear(); }}; } + // Helper macros (not their own methods to appease the borrow checker). + macro_rules! drain_updates_body { + ($updates:expr) => {{ + if self.exec_state.should_stop() { + return; + } + // 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(cur, index_partition.is_some()) + { + drain_updates_parallel!($updates) + } else { + $updates.drain(|update| match update { + UpdateInstr::PushBinding(var, val) => { + binding_info.bindings.insert(var, val); + } + UpdateInstr::RefineAtom(atom, subset) => { + binding_info.insert_node(atom, subset); + } + UpdateInstr::RefineAtomDense(atom, range) => { + binding_info.insert_subset(atom, Subset::Dense(range)); + } + UpdateInstr::EndFrame => { + // Inline leaf-level: if cur+1 is the leaf (no more + // join stages), call push_bindings directly without + // a recursive run_plan call, avoiding function call + // overhead + an extra should_stop() check. + if cur + 1 >= instr_order.len() { + action_buf.push_bindings_factorized( + action, + &mut binding_info.bindings, + &binding_info.binding_sets, + self.exec_state, + ); + } else { + self.run_plan( + stages, + prepared, + atoms, + action, + instr_order, + leaf_scans, + cur + 1, + None, + remaining_after_current, + binding_info, + action_buf, + ); + } + } + }) + } + }}; + } + macro_rules! drain_updates { + ($updates:expr) => {{ + let completed_frames = $updates.frames(); + drain_updates_body!($updates); + remaining_scan_upper_bound = + remaining_scan_upper_bound.saturating_sub(completed_frames); + chunk_size = + action_buf.morsel_size(cur, instr_order.len(), remaining_scan_upper_bound); + }}; + ($updates:expr, false) => {{ + drain_updates_body!($updates); + }}; + } // A sharded top-level job enumerates only its assigned physical index - // shard. Every recursive call clears `index_shard`, so this macro is + // shard. Every recursive call clears `index_partition`, 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) + if let Some(partition) = index_partition { + $prober.for_each_shard(partition.shard, $callback) } else { $prober.for_each($callback) } @@ -3529,7 +3564,7 @@ impl<'a, 'state, 'exec> JoinState<'a, 'state, 'exec> { drain_updates!(updates); } }); - drain_updates!(updates); + drain_updates!(updates, false); binding_info.move_back(a.atom, prober); } [a, b] => { @@ -3589,7 +3624,7 @@ impl<'a, 'state, 'exec> JoinState<'a, 'state, 'exec> { } } }); - drain_updates!(updates); + drain_updates!(updates, false); binding_info.move_back(a.atom, a_prober); binding_info.move_back(b.atom, b_prober); @@ -3652,7 +3687,7 @@ impl<'a, 'state, 'exec> JoinState<'a, 'state, 'exec> { drain_updates!(updates); } }); - drain_updates!(updates); + drain_updates!(updates, false); } for (spec, prober) in rest.iter().zip(probers.into_iter()) { binding_info.move_back(spec.atom, prober); @@ -3695,7 +3730,7 @@ impl<'a, 'state, 'exec> JoinState<'a, 'state, 'exec> { binding_info.binding_sets.push((vars, Arc::new(buf))); let mut updates = FrameUpdates::with_capacity(1); updates.finish_frame(); - drain_updates!(updates); + drain_updates!(updates, false); binding_info.binding_sets.pop(); binding_info.move_back_node(cover_atom, cover_node); } else { @@ -3748,7 +3783,7 @@ impl<'a, 'state, 'exec> JoinState<'a, 'state, 'exec> { } break; } - drain_updates!(updates); + drain_updates!(updates, false); // Restore the subsets we swapped out. binding_info.move_back_node(cover_atom, cover_node); } @@ -3859,7 +3894,7 @@ impl<'a, 'state, 'exec> JoinState<'a, 'state, 'exec> { // deduping (and hence we can do a straight scan: e.g. when the // cover is binding a superset of the primary key for the // table). - drain_updates!(updates); + drain_updates!(updates, false); // Restore the subsets we swapped out. binding_info.move_back_node(cover_atom, cover_node); for (_, atom, prober) in index_probers { @@ -3946,7 +3981,7 @@ impl<'a, 'state, 'exec> JoinState<'a, 'state, 'exec> { binding_info.binding_sets.push((vars, Arc::new(buf))); let mut updates = FrameUpdates::with_capacity(1); updates.finish_frame(); - drain_updates!(updates); + drain_updates!(updates, false); binding_info.binding_sets.pop(); })(); if restore_materialization { @@ -4113,7 +4148,7 @@ impl<'a, 'state, 'exec> JoinState<'a, 'state, 'exec> { } } - drain_updates!(updates); + drain_updates!(updates, false); for (spec, prober) in to_intersect.iter().zip(probers) { binding_info.move_back(spec.0.to_index.atom, prober); } @@ -4128,6 +4163,74 @@ impl<'a, 'state, 'exec> JoinState<'a, 'state, 'exec> { } const LOCAL_ACTION_BATCH_SIZE: usize = 128; +const EXPERIMENTAL_TOP_SHARD_MORSEL_SIZE: usize = 4 * 1024; +const EXPERIMENTAL_COARSE_MIN_MORSEL_SIZE: usize = 256; +const EXPERIMENTAL_TOP_SHARD_POLICY_ENV: &str = "EGGLOG_EXPERIMENTAL_TOP_SHARD_POLICY"; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum ExperimentalTopShardPolicy { + Serial, + DemandGlobal, + DemandLocal, + CoarseGlobal, +} + +impl ExperimentalTopShardPolicy { + fn parse(value: Option<&str>) -> Result { + match value { + None | Some("serial") => Ok(Self::Serial), + Some("demand-global") => Ok(Self::DemandGlobal), + Some("demand-local") => Ok(Self::DemandLocal), + Some("coarse-global") => Ok(Self::CoarseGlobal), + Some(_) => Err(()), + } + } + + fn split_demand(self, stalled_workers: usize, _local_queue_depth: usize) -> usize { + match self { + Self::Serial => 0, + Self::DemandGlobal | Self::DemandLocal | Self::CoarseGlobal => stalled_workers, + } + } + + fn permits_parallel_drain(self, level: usize, partitioned_top_scan: bool) -> bool { + match self { + Self::Serial => false, + Self::DemandGlobal | Self::DemandLocal => !partitioned_top_scan, + Self::CoarseGlobal => partitioned_top_scan && level == 0, + } + } +} + +fn coarse_scan_morsel_size(scan_size: usize, workers: usize) -> usize { + scan_size + .div_ceil(workers.max(1)) + .max(EXPERIMENTAL_COARSE_MIN_MORSEL_SIZE) +} + +fn cached_coarse_scan_morsel_size( + cached: &mut Option, + scan_size: usize, + workers: usize, +) -> usize { + *cached.get_or_insert_with(|| coarse_scan_morsel_size(scan_size, workers)) +} + +fn experimental_top_shard_policy() -> ExperimentalTopShardPolicy { + static POLICY: OnceLock = OnceLock::new(); + *POLICY.get_or_init(|| match std::env::var(EXPERIMENTAL_TOP_SHARD_POLICY_ENV) { + Ok(value) => ExperimentalTopShardPolicy::parse(Some(&value)).unwrap_or_else(|()| { + panic!( + "{EXPERIMENTAL_TOP_SHARD_POLICY_ENV} must be one of \ + serial, demand-global, demand-local, or coarse-global; got {value:?}" + ) + }), + Err(std::env::VarError::NotPresent) => ExperimentalTopShardPolicy::Serial, + Err(std::env::VarError::NotUnicode(value)) => { + panic!("{EXPERIMENTAL_TOP_SHARD_POLICY_ENV} is not valid UTF-8: {value:?}") + } + }) +} /// A trait used to abstract over different ways of buffering actions together /// before running them. @@ -4214,7 +4317,7 @@ where /// /// As of right now this is just a hard-coded value. We may change it in the /// future to fan out more at higher levels though. - fn morsel_size(&mut self, _level: usize, _total: usize) -> usize { + fn morsel_size(&mut self, _level: usize, _total: usize, _estimated_scan_size: usize) -> usize { 256 } @@ -4222,7 +4325,7 @@ where /// /// When `false`, `drain_updates` will use the serial path even at `cur <= 1`, /// avoiding the per-frame `ExecutionState::clone()` overhead. - fn supports_parallel_drain(&self) -> bool { + fn supports_parallel_drain(&self, _level: usize, _partitioned_top_scan: bool) -> bool { true } @@ -4315,42 +4418,64 @@ where work(inner.borrow_mut(), self) } - fn supports_parallel_drain(&self) -> bool { + fn supports_parallel_drain(&self, _level: usize, _partitioned_top_scan: bool) -> 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> { +/// Action buffer used inside one globally scheduled top-index shard. +/// +/// `Serial` preserves the production behavior exactly. The fine-grained demand +/// policies may hand recursive drains to the global or worker-local queue. +/// `CoarseGlobal` instead groups many top bindings from one cached-index shard +/// into a bounded global job and runs that job's complete subtree serially. +/// Action batches remain local to each spawned child and are flushed before +/// that child exits. The buffer deliberately has no `needs_flush` flag: every +/// owner unconditionally flushes its local batches before returning. +struct TopShardActionBuffer<'inner, 'scope> { + scope: &'inner Scope<'scope>, rule_set: &'scope RuleSet, match_counter: Arc, batches: DenseIdMap, + policy: ExperimentalTopShardPolicy, + // Fixed after the first demand signal so one shard produces a bounded + // number of similarly sized jobs instead of a geometric series. + coarse_morsel_size: Option, } -impl<'scope> SerialScopedActionBuffer<'scope> { - fn new(rule_set: &'scope RuleSet, match_counter: Arc) -> Self { +impl<'inner, 'scope> TopShardActionBuffer<'inner, 'scope> { + fn new( + scope: &'inner Scope<'scope>, + rule_set: &'scope RuleSet, + match_counter: Arc, + policy: ExperimentalTopShardPolicy, + ) -> Self { Self { + scope, rule_set, match_counter, batches: Default::default(), + policy, + coarse_morsel_size: None, } } + + fn split_demand(&self) -> usize { + self.policy + .split_demand(self.scope.stalled_workers(), self.scope.local_queue_depth()) + } } -impl<'scope, 'exec> ActionBuffer<'scope, 'exec, ActionId> for SerialScopedActionBuffer<'scope> +impl<'scope, 'exec> ActionBuffer<'scope, 'exec, ActionId> for TopShardActionBuffer<'_, 'scope> where 'exec: 'scope, { type AsLocal<'a> - = Self + = TopShardActionBuffer<'a, 'scope> where 'scope: 'a; type AsGlobalSerial<'a> - = Self + = TopShardActionBuffer<'a, 'scope> where 'scope: 'a; @@ -4394,13 +4519,56 @@ where &mut self, mut local: BorrowedLocalState<'local, 'rows, 'exec>, subset_clone_plan: SubsetClonePlan<'_>, - _to_exec_state: impl FnMut() -> ExecutionState<'scope> + Send + 'scope, - work: impl for<'a> FnOnce(BorrowedLocalState<'a, 'scope, 'exec>, &mut Self) + Send + 'scope, + mut to_exec_state: impl FnMut() -> ExecutionState<'scope> + Send + 'scope, + work: impl for<'a> FnOnce( + BorrowedLocalState<'a, 'scope, 'exec>, + &mut TopShardActionBuffer<'a, 'scope>, + ) + Send + + 'scope, ) where 'rows: 'scope, { + // `supports_parallel_drain` made the first advisory stalled-state + // observation. Recheck immediately before committing to a queue: a + // sibling shard may have filled the pool in the meantime. + let should_split = self.split_demand() != 0; let mut inner: LocalState<'scope, 'exec> = local.clone_state(subset_clone_plan); - work(inner.borrow_mut(), self); + if !should_split { + let mut inline = TopShardActionBuffer { + scope: self.scope, + rule_set: self.rule_set, + match_counter: self.match_counter.clone(), + batches: mem::take(&mut self.batches), + policy: self.policy, + coarse_morsel_size: self.coarse_morsel_size, + }; + work(inner.borrow_mut(), &mut inline); + self.batches = inline.batches; + self.coarse_morsel_size = inline.coarse_morsel_size; + return; + } + + let rule_set = self.rule_set; + let match_counter = self.match_counter.clone(); + let policy = self.policy; + let child_policy = match policy { + ExperimentalTopShardPolicy::CoarseGlobal => ExperimentalTopShardPolicy::Serial, + policy => policy, + }; + let task = move |scope: &Scope<'scope>| { + let mut buf = TopShardActionBuffer::new(scope, rule_set, match_counter, child_policy); + work(inner.borrow_mut(), &mut buf); + buf.flush(&mut to_exec_state()); + }; + match policy { + ExperimentalTopShardPolicy::Serial => { + unreachable!("serial top-shard policy cannot request a demand split") + } + ExperimentalTopShardPolicy::DemandGlobal | ExperimentalTopShardPolicy::CoarseGlobal => { + self.scope.spawn_global(task) + } + ExperimentalTopShardPolicy::DemandLocal => self.scope.spawn_local(task), + } } fn recur_global_serial<'local, 'rows>( @@ -4408,16 +4576,54 @@ where mut local: BorrowedLocalState<'local, 'rows, 'exec>, subset_clone_plan: SubsetClonePlan<'_>, _to_exec_state: impl FnMut() -> ExecutionState<'scope> + Send + 'scope, - work: impl for<'a> FnOnce(BorrowedLocalState<'a, 'scope, 'exec>, &mut Self) + Send + 'scope, + work: impl for<'a> FnOnce( + BorrowedLocalState<'a, 'scope, 'exec>, + &mut TopShardActionBuffer<'a, 'scope>, + ) + Send + + 'scope, ) where 'rows: 'scope, { let mut inner: LocalState<'scope, 'exec> = local.clone_state(subset_clone_plan); - work(inner.borrow_mut(), self); + let mut inline = TopShardActionBuffer { + scope: self.scope, + rule_set: self.rule_set, + match_counter: self.match_counter.clone(), + batches: mem::take(&mut self.batches), + policy: self.policy, + coarse_morsel_size: self.coarse_morsel_size, + }; + work(inner.borrow_mut(), &mut inline); + self.batches = inline.batches; + self.coarse_morsel_size = inline.coarse_morsel_size; } - fn supports_parallel_drain(&self) -> bool { - false + fn supports_parallel_drain(&self, level: usize, partitioned_top_scan: bool) -> bool { + // This first check avoids state cloning entirely while all workers are + // occupied. `recur` performs a second advisory check before spawning. + let eligible = self + .policy + .permits_parallel_drain(level, partitioned_top_scan); + eligible && self.split_demand() != 0 + } + + fn morsel_size(&mut self, level: usize, _total: usize, estimated_scan_size: usize) -> usize { + match self.policy { + ExperimentalTopShardPolicy::DemandGlobal | ExperimentalTopShardPolicy::DemandLocal + if level == 1 => + { + EXPERIMENTAL_TOP_SHARD_MORSEL_SIZE + } + ExperimentalTopShardPolicy::CoarseGlobal if level == 0 && self.split_demand() != 0 => { + let workers = crate::parallel::current_num_threads(); + cached_coarse_scan_morsel_size( + &mut self.coarse_morsel_size, + estimated_scan_size, + workers, + ) + } + _ => 256, + } } } @@ -4455,7 +4661,7 @@ where where 'scope: 'a; type AsGlobalSerial<'a> - = SerialScopedActionBuffer<'scope> + = TopShardActionBuffer<'a, 'scope> where 'scope: 'a; fn push_bindings( @@ -4544,7 +4750,7 @@ where mut to_exec_state: impl FnMut() -> ExecutionState<'scope> + Send + 'scope, work: impl for<'a> FnOnce( BorrowedLocalState<'a, 'scope, 'exec>, - &mut SerialScopedActionBuffer<'scope>, + &mut TopShardActionBuffer<'a, 'scope>, ) + Send + 'scope, ) where @@ -4553,8 +4759,9 @@ where let rule_set = self.rule_set; let match_counter = self.match_counter.clone(); let mut inner = local.clone_state(subset_clone_plan); - self.scope.spawn_global(move |_| { - let mut buf = SerialScopedActionBuffer::new(rule_set, match_counter); + let policy = experimental_top_shard_policy(); + self.scope.spawn_global(move |scope| { + let mut buf = TopShardActionBuffer::new(scope, rule_set, match_counter, policy); 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. @@ -4562,7 +4769,7 @@ where }); } - fn morsel_size(&mut self, _level: usize, _total: usize) -> usize { + fn morsel_size(&mut self, _level: usize, _total: usize, _estimated_scan_size: usize) -> usize { // Lower morsel size to increase parallelism. match _level { 0 if _total > 2 => 32, @@ -4719,14 +4926,16 @@ where work(inner.borrow_mut(), self) } - fn supports_parallel_drain(&self) -> bool { + fn supports_parallel_drain(&self, _level: usize, _partitioned_top_scan: bool) -> bool { false } } -/// Serial recursive view of a scoped materializer. Sibling top-level shards +/// 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. +/// deeper work itself. The experimental demand policy is intentionally limited +/// to direct action execution; materialization remains an unchanged +/// serial-per-shard control. struct SerialScopedMaterializer { specs: Arc>, materializations: Arc, RowBuffer>>>>, @@ -4818,7 +5027,7 @@ where work(inner.borrow_mut(), self); } - fn supports_parallel_drain(&self) -> bool { + fn supports_parallel_drain(&self, _level: usize, _partitioned_top_scan: bool) -> bool { false } } @@ -5314,12 +5523,13 @@ mod top_index_tests { }; use super::{ - AccessId, BindingInfo, CatalogContinuation, ContinuationPosition, Database, InstrOrder, - LazyArenaHandle, PreparedIndexKind, PreparedIndexSlot, PreparedIndexState, - PreparedIndexStateId, PreparedJoinIndexes, PreparedTailMasks, RootContinuationCache, - RootProjection, SmallColumnIndex, SmallColumnSink, TrieNode, for_each_stage_atom, - materialization_is_live_in_tail, packed_child_shape_in_tail, scan_atom_tail_use, - sort_plan_by_size_inner, top_index_shape_is_eligible, + AccessId, BindingInfo, CatalogContinuation, ContinuationPosition, Database, + ExperimentalTopShardPolicy, InstrOrder, LazyArenaHandle, PreparedIndexKind, + PreparedIndexSlot, PreparedIndexState, PreparedIndexStateId, PreparedJoinIndexes, + PreparedTailMasks, RootContinuationCache, RootProjection, SmallColumnIndex, + SmallColumnSink, TrieNode, cached_coarse_scan_morsel_size, coarse_scan_morsel_size, + for_each_stage_atom, materialization_is_live_in_tail, packed_child_shape_in_tail, + scan_atom_tail_use, sort_plan_by_size_inner, top_index_shape_is_eligible, }; use crate::free_join::packed_trie::ChildShape; @@ -5347,6 +5557,87 @@ mod top_index_tests { assert!(handle.handle.get().is_some()); } + #[test] + fn parses_experimental_top_shard_policy_without_mutating_process_env() { + assert_eq!( + ExperimentalTopShardPolicy::parse(None), + Ok(ExperimentalTopShardPolicy::Serial) + ); + assert_eq!( + ExperimentalTopShardPolicy::parse(Some("serial")), + Ok(ExperimentalTopShardPolicy::Serial) + ); + assert_eq!( + ExperimentalTopShardPolicy::parse(Some("demand-global")), + Ok(ExperimentalTopShardPolicy::DemandGlobal) + ); + assert_eq!( + ExperimentalTopShardPolicy::parse(Some("demand-local")), + Ok(ExperimentalTopShardPolicy::DemandLocal) + ); + assert_eq!( + ExperimentalTopShardPolicy::parse(Some("coarse-global")), + Ok(ExperimentalTopShardPolicy::CoarseGlobal) + ); + assert_eq!(ExperimentalTopShardPolicy::parse(Some("")), Err(())); + assert_eq!(ExperimentalTopShardPolicy::parse(Some("demand")), Err(())); + } + + #[test] + fn experimental_top_shard_policy_splits_only_on_demand() { + assert_eq!(ExperimentalTopShardPolicy::Serial.split_demand(3, 0), 0); + assert_eq!( + ExperimentalTopShardPolicy::DemandGlobal.split_demand(0, 100), + 0 + ); + assert_eq!( + ExperimentalTopShardPolicy::DemandGlobal.split_demand(3, 100), + 3 + ); + assert_eq!( + ExperimentalTopShardPolicy::DemandLocal.split_demand(0, 100), + 0 + ); + assert_eq!( + ExperimentalTopShardPolicy::DemandLocal.split_demand(3, 0), + 3 + ); + assert_eq!( + ExperimentalTopShardPolicy::DemandLocal.split_demand(3, 100), + 3 + ); + assert_eq!( + ExperimentalTopShardPolicy::CoarseGlobal.split_demand(0, 100), + 0 + ); + assert_eq!( + ExperimentalTopShardPolicy::CoarseGlobal.split_demand(3, 100), + 3 + ); + } + + #[test] + fn coarse_scan_morsels_are_bounded_and_cover_the_scan() { + assert_eq!(coarse_scan_morsel_size(0, 0), 256); + assert_eq!(coarse_scan_morsel_size(100, 8), 256); + assert_eq!(coarse_scan_morsel_size(2048, 8), 256); + assert_eq!(coarse_scan_morsel_size(2049, 8), 257); + assert_eq!(coarse_scan_morsel_size(4096, 3), 1366); + + let mut cached = None; + assert_eq!(cached_coarse_scan_morsel_size(&mut cached, 4096, 3), 1366); + assert_eq!(cached_coarse_scan_morsel_size(&mut cached, 512, 3), 1366); + } + + #[test] + fn coarse_policy_only_partitions_the_cached_top_scan() { + assert!(ExperimentalTopShardPolicy::DemandGlobal.permits_parallel_drain(1, false)); + assert!(!ExperimentalTopShardPolicy::DemandGlobal.permits_parallel_drain(0, true)); + assert!(ExperimentalTopShardPolicy::CoarseGlobal.permits_parallel_drain(0, true)); + assert!(!ExperimentalTopShardPolicy::CoarseGlobal.permits_parallel_drain(1, false)); + assert!(!ExperimentalTopShardPolicy::CoarseGlobal.permits_parallel_drain(1, true)); + } + #[test] fn continuation_positions_remain_compact() { assert_eq!(