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!(