From 7b45a6465300854d4c246564467883aa42f8a8a2 Mon Sep 17 00:00:00 2001 From: LinZhihao-723 Date: Tue, 1 Sep 2026 15:17:04 -0400 Subject: [PATCH 1/5] feat(spider-scheduler): Add the resource-group-aware round-robin scheduler core. Ports the prototype's tick loop: collect the inbound poll's formatted results, fold them into the job registry and the per-resource-group scheduling states, refill the dispatch queues under the dynamic admission threshold, and retire the jobs that ran dry. The core is not wired to `SchedulerCore` or the scheduler config yet, so nothing selects it and no behaviour changes. --- .../implementation.rs | 630 ++++++++++++++ .../inbound_queue_reader.rs | 4 +- .../resource_group_round_robin/mod.rs | 41 +- .../resource_group_round_robin/tests.rs | 779 ++++++++++++++++++ 4 files changed, 1427 insertions(+), 27 deletions(-) create mode 100644 components/spider-scheduler/src/core_impl/resource_group_round_robin/implementation.rs create mode 100644 components/spider-scheduler/src/core_impl/resource_group_round_robin/tests.rs diff --git a/components/spider-scheduler/src/core_impl/resource_group_round_robin/implementation.rs b/components/spider-scheduler/src/core_impl/resource_group_round_robin/implementation.rs new file mode 100644 index 00000000..25867ea5 --- /dev/null +++ b/components/spider-scheduler/src/core_impl/resource_group_round_robin/implementation.rs @@ -0,0 +1,630 @@ +//! The implementation of the resource-group-aware round-robin scheduler core. See the parent +//! module's documentation for the scheduling policy and configuration. + +use std::collections::HashMap; +use std::collections::HashSet; +use std::num::NonZeroU64; +use std::num::NonZeroUsize; +use std::time::Duration; + +use serde::Deserialize; +use spider_core::session::SessionTracker; +use spider_core::types::id::JobId; +use spider_core::types::id::ResourceGroupId; +use spider_core::types::id::SessionId; +use spider_core::types::id::TaskId; +use tokio::select; +use tokio_util::sync::CancellationToken; + +use super::dispatch_queue::DispatchQueueRegistry; +use super::inbound_queue_reader::FinalizedJob; +use super::inbound_queue_reader::ReadyBatch; +use super::inbound_queue_reader::RgInboundPollState; +use super::inbound_queue_reader::RgInboundQueueReader; +use super::inbound_queue_reader::format_finalized_jobs; +use super::inbound_queue_reader::format_ready_job_batches; +use super::job_registry::JobKey; +use super::job_registry::JobRegistry; +use super::job_registry::UpsertOutcome; +use super::scheduling_state::FinalizeKind; +use super::scheduling_state::MakeAssignmentError; +use super::scheduling_state::RgSchedulingState; +use crate::core::TaskAssignmentIdIssuer; +use crate::error::SchedulerError; +use crate::storage_client::SchedulerStorageClient; +use crate::types::InboundEntry; +use crate::types::TaskAssignment; + +/// The configuration of the resource-group-aware round-robin scheduler core. +#[derive(Clone, Debug, Deserialize)] +pub(super) struct RgRoundRobinConfig { + /// The total dispatch buffer size shared by all resource groups. + pub(super) dispatch_queue_capacity: NonZeroUsize, + + /// The number of active jobs each resource group may hold, applied per group rather than as a + /// global budget. + pub(super) active_job_list_capacity: NonZeroUsize, + + /// The capacity of the total pending ready tasks buffered in the scheduler. + pub(super) ready_task_capacity: NonZeroUsize, + + /// The capacity of the total pending commit-ready tasks buffered in the scheduler. + pub(super) commit_ready_task_capacity: NonZeroUsize, + + /// The capacity of the total pending cleanup-ready tasks buffered in the scheduler. + pub(super) cleanup_ready_task_capacity: NonZeroUsize, + + /// The maximum time (in milliseconds) that the scheduler will wait for the storage server to + /// fill the inbound-queue reading request. + pub(super) storage_poll_timeout_ms: u64, + + /// The time (in milliseconds) that the scheduler will spend on each tick. + pub(super) tick_interval_ms: NonZeroU64, +} + +/// The resource-group-aware round-robin scheduler core created from a [`RgRoundRobinConfig`]. +/// +/// The core owns all of the state it decides with -- job entries in a generational arena, per +/// resource group scheduling states in an append-only vector, and the dispatch queue registry the +/// execution-manager-facing service reads from -- so nothing it holds across an await point is +/// thread-bound. +/// +/// # Type Parameters +/// +/// * `SchedulerStorageClientType` - The storage client used to poll the inbound queue. +/// +/// # Note +/// +/// All member variables are marked `pub(super)` to allow the test module to inspect the internal +/// states. +pub(super) struct RgRoundRobin { + pub(super) global_task_set: HashSet<(JobId, TaskId)>, + pub(super) finalized_jobs: HashSet, + pub(super) job_registry: JobRegistry, + + /// The scheduling states of every resource group the core has seen this session. + /// + /// Append-only within a session: a group is never removed individually, so a position in this + /// vector is stable until [`Self::apply_session_bump`] flushes the whole of it. + pub(super) rg_states: Vec, + + pub(super) rg_id_to_index_map: HashMap, + pub(super) active_rg_list: Vec, + pub(super) last_served_rg: Option, + + pub(super) config: RgRoundRobinConfig, + pub(super) dispatch_queue_registry: DispatchQueueRegistry, + pub(super) session_tracker: SessionTracker, + pub(super) id_issuer: TaskAssignmentIdIssuer, + pub(super) inbound_queue_reader: RgInboundQueueReader, + pub(super) reschedule_queue_reader: tokio::sync::mpsc::UnboundedReceiver, + pub(super) cancellation_token: CancellationToken, +} + +impl + RgRoundRobin +{ + /// Factory function. + /// + /// Creates a core owning a freshly created dispatch queue registry and the session tracker that + /// stamps every group the registry creates. + /// + /// # Returns + /// + /// A newly created core with no buffered task and no in-flight inbound poll. + pub(super) fn new( + storage_client: SchedulerStorageClientType, + reschedule_queue_reader: tokio::sync::mpsc::UnboundedReceiver, + id_issuer: TaskAssignmentIdIssuer, + cancellation_token: CancellationToken, + config: RgRoundRobinConfig, + ) -> Self { + let session_tracker = SessionTracker::new(SessionId::default()); + let dispatch_queue_registry = DispatchQueueRegistry::new(session_tracker.clone()); + Self { + global_task_set: HashSet::new(), + finalized_jobs: HashSet::new(), + job_registry: JobRegistry::new(), + rg_states: Vec::new(), + rg_id_to_index_map: HashMap::new(), + active_rg_list: Vec::new(), + last_served_rg: None, + config, + dispatch_queue_registry, + session_tracker, + id_issuer, + inbound_queue_reader: RgInboundQueueReader::new(storage_client), + reschedule_queue_reader, + cancellation_token, + } + } + + /// # Returns + /// + /// A handle over the core's dispatch queue registry, from which the execution-manager-facing + /// service reads. + pub(super) fn dispatch_queue_registry(&self) -> DispatchQueueRegistry { + self.dispatch_queue_registry.clone() + } + + /// Runs the scheduling loop until the cancellation token is triggered. + /// + /// Each iteration executes one [`Self::tick`] and then sleeps for the remainder of the + /// configured tick interval. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * Forwards [`Self::tick`]'s return values on failure. + pub(super) async fn run(mut self) -> Result<(), SchedulerError> { + tracing::info!( + config = ? self.config, + init_session_id = self.session_tracker.current(), + "Resource-group-aware round-robin scheduler started." + ); + let tick_interval = Duration::from_millis(self.config.tick_interval_ms.get()); + loop { + let now = tokio::time::Instant::now(); + let cancellation_token = self.cancellation_token.clone(); + select! { + () = cancellation_token.cancelled() => { + tracing::info!( + "Resource-group-aware round-robin scheduler cancelled. Shutting down." + ); + return Ok(()); + } + result = self.tick() => { + result.inspect_err(|err| tracing::error!( + err = % err, + "Resource-group-aware round-robin scheduler exits on error." + ))?; + } + } + let sleep_time = tick_interval.saturating_sub(now.elapsed()); + if sleep_time.is_zero() { + tokio::task::yield_now().await; + } else { + tokio::time::sleep(sleep_time).await; + } + } + } + + /// Executes one tick of the scheduling loop. + /// + /// Processing the polling results is skipped while a storage poll is still in flight, but the + /// dispatch queues are refilled from already-buffered tasks on every tick regardless. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * Forwards [`RgInboundQueueReader::try_collect_result`]'s return values on failure. + /// * Forwards [`Self::start_inbound_poll`]'s return values on failure. + /// * Forwards [`Self::fill_dispatch_queues`]'s return values on failure. + pub(super) async fn tick(&mut self) -> Result<(), SchedulerError> { + let poll_state = self + .inbound_queue_reader + .try_collect_result(self.session_tracker.current()) + .await?; + match poll_state { + RgInboundPollState::Ready { + session_id, + ready_result, + commit_ready_result, + cleanup_ready_result, + } => { + if session_id != self.session_tracker.current() { + self.apply_session_bump(session_id); + } + let rescheduled_entries = self.drain_reschedule_queue(session_id); + + let rg_updates = self.process_polling_results( + commit_ready_result, + cleanup_ready_result, + ready_result, + rescheduled_entries, + ); + + self.apply_rg_updates(rg_updates); + self.start_inbound_poll()?; + } + RgInboundPollState::NotStarted => self.start_inbound_poll()?, + RgInboundPollState::Pending => (), + } + + let jobs_to_retire = self.fill_dispatch_queues()?; + self.retire_jobs(jobs_to_retire); + Ok(()) + } + + /// Discards every piece of state published in a session older than `new_session_id`. + /// + /// Clearing the dedup set is load-bearing rather than tidy: storage replays its ready tasks + /// after a bump, and a stale dedup entry would drop a replayed task while the registry no + /// longer holds anything to schedule it from. + /// + /// [`Self::rg_states`], [`Self::rg_index`], and [`Self::active_rg_list`] must be cleared as one + /// operation, and this is the only place any of them is cleared. Positions in `rg_states` carry + /// no generation, so an index that outlives the flush does not fail: it resolves against the + /// new session's states, either out of bounds or -- once the new session has re-created a few + /// groups -- silently against the wrong group. Nothing in the type system checks this. + /// + /// [`Self::dispatch_queue_registry`] must be cleared together with them, and that too is a + /// correctness requirement rather than tidiness. A group's queue closes only once every sender + /// has been dropped, and a scheduling state's write side is one of them; a state that survived + /// the registry's flush would therefore hold a write side onto a queue whose readers are gone, + /// and publishing into a closed queue is fatal to the core. Clearing the registry is also what + /// discards the hints published in the session being left behind. + fn apply_session_bump(&mut self, new_session_id: SessionId) { + let previous_session_id = self.session_tracker.current(); + if self.session_tracker.try_advance(new_session_id) { + tracing::info!( + from = previous_session_id, + to = new_session_id, + num_resource_groups = self.dispatch_queue_registry.len(), + num_jobs = self.job_registry.len(), + "Storage session bumped. Flushing the core." + ); + } else { + tracing::error!( + from = previous_session_id, + to = new_session_id, + "Storage reported a session no newer than the tracked one. Flushing the core \ + anyway, but it keeps serving the tracked session." + ); + } + + self.dispatch_queue_registry.clear(); + self.rg_states.clear(); + self.rg_id_to_index_map.clear(); + self.active_rg_list.clear(); + self.last_served_rg = None; + self.job_registry.clear(); + self.global_task_set.clear(); + self.finalized_jobs.clear(); + } + + /// Drains the reschedule queue, dropping assignments published in a session other than + /// `session_id`. + /// + /// # Returns + /// + /// The rescheduled assignments, in the same form as the entries drained from the inbound queue. + fn drain_reschedule_queue(&mut self, session_id: SessionId) -> Vec { + let mut entries = Vec::new(); + while let Ok(assignment) = self.reschedule_queue_reader.try_recv() { + if assignment.session_id != session_id { + continue; + } + entries.push(InboundEntry { + resource_group_id: assignment.resource_group_id, + job_id: assignment.job_id, + task_id: assignment.task_id, + }); + } + entries + } + + /// Folds the tick's ready tasks into the finalized job table, the global task set, and the job + /// registry. + /// + /// Finalizations are processed before regular tasks, so a regular task arriving in the same + /// batch as its job's finalization is discarded rather than scheduled. + /// + /// # Returns + /// + /// The per-resource-group updates the tick produced. + fn process_polling_results( + &mut self, + commit_ready_jobs: Vec, + cleanup_ready_jobs: Vec, + ready_batches: Vec, + rescheduled_entries: Vec, + ) -> HashMap { + let mut rescheduled_commit_entries = Vec::new(); + let mut rescheduled_cleanup_entries = Vec::new(); + let mut rescheduled_regular_entries = Vec::new(); + for entry in rescheduled_entries { + match entry.task_id { + TaskId::Commit => rescheduled_commit_entries.push(entry), + TaskId::Cleanup => rescheduled_cleanup_entries.push(entry), + TaskId::Index(_) => rescheduled_regular_entries.push(entry), + } + } + + let mut commit_ready = commit_ready_jobs; + commit_ready.extend(format_finalized_jobs(rescheduled_commit_entries)); + let mut cleanup_ready = cleanup_ready_jobs; + cleanup_ready.extend(format_finalized_jobs(rescheduled_cleanup_entries)); + let mut batches = ready_batches; + batches.extend(format_ready_job_batches(rescheduled_regular_entries)); + + let mut rg_updates: HashMap = HashMap::new(); + for (finalized_jobs, kind) in [ + (commit_ready, FinalizeKind::Commit), + (cleanup_ready, FinalizeKind::Cleanup), + ] { + for finalized_job in finalized_jobs { + if !self.finalized_jobs.insert(finalized_job.job_id) { + continue; + } + // The job's still-buffered regular tasks will never be published, so they must + // leave the dedup set with it or nothing would ever remove them. + if let Some(mut job_entry) = + self.job_registry.remove_by_job_id(finalized_job.job_id) + { + for task_index in job_entry.take_ready_tasks() { + self.global_task_set + .remove(&(finalized_job.job_id, TaskId::Index(task_index))); + } + } + rg_updates + .entry(finalized_job.resource_group_id) + .or_default() + .finalized + .push((finalized_job.job_id, kind)); + } + } + + for batch in batches { + let ReadyBatch { + resource_group_id, + job_id, + mut task_indices, + } = batch; + if self.finalized_jobs.contains(&job_id) { + continue; + } + let global_task_set = &mut self.global_task_set; + task_indices + .retain(|task_index| global_task_set.insert((job_id, TaskId::Index(*task_index)))); + if task_indices.is_empty() { + continue; + } + if let UpsertOutcome::New(job_key) = self.job_registry.upsert(job_id, task_indices) { + rg_updates + .entry(resource_group_id) + .or_default() + .new_jobs + .push(job_key); + } + } + + rg_updates + } + + /// Applies the tick's per-resource-group updates to the scheduling states, activating every + /// group the updates touch. + fn apply_rg_updates(&mut self, rg_updates: HashMap) { + for (rg_id, update) in rg_updates { + let state_index = self.get_or_create_state(rg_id); + let rg_state = &mut self.rg_states[state_index]; + for (job_id, kind) in update.finalized { + rg_state.push_finalization(job_id, kind); + } + for job_key in update.new_jobs { + rg_state.place_new_job(job_key); + } + let activated = !rg_state.is_active; + rg_state.is_active = true; + if activated { + self.active_rg_list.push(state_index); + } + } + } + + /// # Returns + /// + /// The position of `rg_id`'s scheduling state in [`Self::rg_states`], appending a state built + /// against the write side of the group's dispatch queue if the core has none. + fn get_or_create_state(&mut self, rg_id: ResourceGroupId) -> usize { + if let Some(state_index) = self.rg_id_to_index_map.get(&rg_id) { + return *state_index; + } + + let writer = self + .dispatch_queue_registry + .get_dispatch_queue_writer(rg_id); + let state_index = self.rg_states.len(); + self.rg_states.push(RgSchedulingState::new( + rg_id, + writer, + self.config.active_job_list_capacity.get(), + )); + self.rg_id_to_index_map.insert(rg_id, state_index); + state_index + } + + /// Publishes assignments into the per-resource-group dispatch queues under the admission + /// policy, then deactivates the groups that ran out of work. + /// + /// # Returns + /// + /// The jobs that exhausted their downgrade budget and must be retired, on success. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * [`SchedulerError::DispatchQueueClosed`] if a group's dispatch queue or the broadcast queue + /// is closed, in which case the assignments the core makes can no longer reach an execution + /// manager. + fn fill_dispatch_queues(&mut self) -> Result, SchedulerError> { + let mut jobs_to_retire = Vec::new(); + if self.active_rg_list.is_empty() { + return Ok(jobs_to_retire); + } + + // The decision loop needs the scheduling states and the job arena borrowed mutably at the + // same time, which the borrow checker accepts only for bindings that name distinct fields. + let Self { + global_task_set, + job_registry, + rg_states, + active_rg_list, + last_served_rg, + config, + session_tracker, + id_issuer, + .. + } = self; + + let mut rg_rr_list = Vec::with_capacity(active_rg_list.len()); + let mut occupancy = 0; + let mut last_served_index = None; + for (index, state_index) in active_rg_list.iter().enumerate() { + let rg_state = &rg_states[*state_index]; + occupancy += rg_state.dispatch_queue_size(); + if Some(rg_state.rg_id) == *last_served_rg { + last_served_index = Some(index); + } + rg_rr_list.push(*state_index); + } + // Bounding by the free space measured here is what makes the loop terminate: the queues + // drain concurrently, so the true free space only ever grows. + let mut free = config + .dispatch_queue_capacity + .get() + .saturating_sub(occupancy); + // Rotating the arm rather than the list keeps the same group from always being visited + // first, which matters because `free` shrinks as the tick proceeds. + let mut arm = last_served_index.map_or(0, |index| (index + 1) % rg_rr_list.len()); + + for state_index in &rg_rr_list { + rg_states[*state_index].promote_pending_jobs(job_registry, &mut jobs_to_retire); + } + + let session_id = session_tracker.current(); + let mut exhausted_states = Vec::new(); + while 0 != free && !rg_rr_list.is_empty() { + let state_index = rg_rr_list[arm]; + let rg_state = &mut rg_states[state_index]; + let result = rg_state.try_make_assignment( + free, + session_id, + id_issuer, + job_registry, + &mut jobs_to_retire, + ); + match result { + Ok((job_id, task_id)) => { + global_task_set.remove(&(job_id, task_id)); + free -= 1; + *last_served_rg = Some(rg_state.rg_id); + arm = (arm + 1) % rg_rr_list.len(); + } + Err(err) => { + match err { + MakeAssignmentError::NoTask => exhausted_states.push(state_index), + MakeAssignmentError::DispatchQueueFull => (), + MakeAssignmentError::DispatchQueueClosed => { + return Err(SchedulerError::DispatchQueueClosed); + } + } + rg_rr_list.swap_remove(arm); + // `swap_remove` moved the tail element into this slot, so advancing the arm + // here would skip it. + if arm == rg_rr_list.len() { + arm = 0; + } + } + } + } + + for state_index in &*active_rg_list { + rg_states[*state_index].apply_downgrades(job_registry); + } + self.deactivate_exhausted_states(exhausted_states); + + Ok(jobs_to_retire) + } + + /// Takes every exhausted group that also holds no assignment off the active resource group + /// list. + /// + /// The empty-queue condition is required for correctness: free space is summed over the active + /// list alone, so a deactivated group still holding assignments would hide its occupancy and + /// let the core over-admit. + fn deactivate_exhausted_states(&mut self, exhausted_states: Vec) { + for state_index in exhausted_states { + let rg_state = &mut self.rg_states[state_index]; + if rg_state.has_schedulable_task() || 0 != rg_state.dispatch_queue_size() { + continue; + } + rg_state.is_active = false; + if let Some(position) = self + .active_rg_list + .iter() + .position(|active_index| *active_index == state_index) + { + self.active_rg_list.swap_remove(position); + } + } + } + + /// Drops the registry's entry for every job that ran out of downgrade lives. + fn retire_jobs(&mut self, jobs_to_retire: Vec) { + for job_key in jobs_to_retire { + self.job_registry.remove(job_key); + } + } + + /// Starts the next inbound poll, sizing each lane's fetch count by the buffer capacity that + /// lane still has left. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * Forwards [`RgInboundQueueReader::start`]'s return values on failure. + fn start_inbound_poll(&mut self) -> Result<(), SchedulerError> { + let (num_commit_ready, num_cleanup_ready) = self.count_buffered_finalizations(); + let max_ready_entries = self + .config + .ready_task_capacity + .get() + .saturating_sub(self.global_task_set.len()); + let max_commit_ready_entries = self + .config + .commit_ready_task_capacity + .get() + .saturating_sub(num_commit_ready); + let max_cleanup_ready_entries = self + .config + .cleanup_ready_task_capacity + .get() + .saturating_sub(num_cleanup_ready); + + self.inbound_queue_reader.start( + Duration::from_millis(self.config.storage_poll_timeout_ms), + max_ready_entries, + max_commit_ready_entries, + max_cleanup_ready_entries, + ) + } + + /// # Returns + /// + /// A tuple containing: + /// + /// * The number of buffered commit tasks. + /// * The number of buffered cleanup tasks. + fn count_buffered_finalizations(&self) -> (usize, usize) { + let mut num_commit_ready = 0; + let mut num_cleanup_ready = 0; + for rg_state in &self.rg_states { + let (num_commits, num_cleanups) = rg_state.num_buffered_finalize_tasks(); + num_commit_ready += num_commits; + num_cleanup_ready += num_cleanups; + } + (num_commit_ready, num_cleanup_ready) + } +} + +/// The updates one tick produced for a single resource group. +#[derive(Default)] +struct RgUpdate { + finalized: Vec<(JobId, FinalizeKind)>, + new_jobs: Vec, +} diff --git a/components/spider-scheduler/src/core_impl/resource_group_round_robin/inbound_queue_reader.rs b/components/spider-scheduler/src/core_impl/resource_group_round_robin/inbound_queue_reader.rs index 8d9e9eb9..2fb5292c 100644 --- a/components/spider-scheduler/src/core_impl/resource_group_round_robin/inbound_queue_reader.rs +++ b/components/spider-scheduler/src/core_impl/resource_group_round_robin/inbound_queue_reader.rs @@ -70,7 +70,7 @@ pub(super) type RgInboundPollState = InboundPollState) -> Vec { +pub(super) fn format_ready_job_batches(entries: Vec) -> Vec { let mut batches: HashMap)> = HashMap::new(); for entry in entries { let TaskId::Index(task_index) = entry.task_id else { @@ -100,7 +100,7 @@ fn format_ready_job_batches(entries: Vec) -> Vec { /// # Returns /// /// One finalized job per entry drained from a finalization lane, in the order they were drained. -fn format_finalized_jobs(entries: Vec) -> Vec { +pub(super) fn format_finalized_jobs(entries: Vec) -> Vec { entries .into_iter() .map(|entry| FinalizedJob { diff --git a/components/spider-scheduler/src/core_impl/resource_group_round_robin/mod.rs b/components/spider-scheduler/src/core_impl/resource_group_round_robin/mod.rs index c295f61b..80be0e36 100644 --- a/components/spider-scheduler/src/core_impl/resource_group_round_robin/mod.rs +++ b/components/spider-scheduler/src/core_impl/resource_group_round_robin/mod.rs @@ -4,44 +4,35 @@ //! of round-robin: the outer level interleaves resource groups, while the inner level interleaves //! active jobs within each resource group. -// The dispatch queues have no consumer outside their own tests until the rest of the core lands, so -// every item they expose reads as dead. `expect` rather than `allow`: once `implementation.rs` uses -// the queues, this attribute becomes unfulfilled and the compiler flags it for removal. +// Only the write side of the dispatch queues has a consumer: the read side and the hints steering +// general execution managers are drained by the dispatch service, which has not landed yet. +// `expect` rather than `allow`: once that service reads through the queues, this attribute becomes +// unfulfilled and the compiler flags it for removal. #[cfg_attr( not(test), expect( dead_code, - reason = "the core and the dispatch service that consume the queues have not landed yet" + reason = "the dispatch service that reads from the queues has not landed yet" ) )] mod dispatch_queue; -// The formatter has no consumer until the rest of the core lands, so every item it exposes reads as -// dead. `expect` rather than `allow`: once `implementation.rs` polls with the formatter, this -// attribute becomes unfulfilled and the compiler flags it for removal. -#[expect( - dead_code, - reason = "the core that polls with the inbound-poll result formatter has not landed yet" -)] -mod inbound_queue_reader; - -// The registry has no consumer outside tests until the rest of the core lands, so every item it -// exposes reads as dead. `expect` rather than `allow`: once `implementation.rs` uses the registry, -// this attribute becomes unfulfilled and the compiler flags it for removal. +// The core has no consumer until the seam that implements `SchedulerCore` over it lands, so every +// item it and the modules it decides with expose reads as dead. `expect` rather than `allow`: once +// the seam constructs the core, this attribute becomes unfulfilled and the compiler flags it for +// removal. #[cfg_attr( not(test), expect( dead_code, - reason = "the core that consumes the registry has not landed yet" + reason = "the `SchedulerCore` implementation that runs the core has not landed yet" ) )] -mod job_registry; +mod implementation; -// The scheduling state has no consumer until the rest of the core lands, so every item it exposes -// reads as dead. `expect` rather than `allow`: once `implementation.rs` uses the state, this -// attribute becomes unfulfilled and the compiler flags it for removal. -#[expect( - dead_code, - reason = "the core that consumes the scheduling state has not landed yet" -)] +mod inbound_queue_reader; +mod job_registry; mod scheduling_state; + +#[cfg(test)] +mod tests; diff --git a/components/spider-scheduler/src/core_impl/resource_group_round_robin/tests.rs b/components/spider-scheduler/src/core_impl/resource_group_round_robin/tests.rs new file mode 100644 index 00000000..db9b9b9e --- /dev/null +++ b/components/spider-scheduler/src/core_impl/resource_group_round_robin/tests.rs @@ -0,0 +1,779 @@ +//! Unit tests for the resource-group-aware round-robin scheduler core. +//! +//! The tests drive [`RgRoundRobin::tick`] and the structures it decides with directly. Whatever an +//! execution manager would do -- draining a group's queue -- a test body does by hand through the +//! same entry points the dispatch service will use. + +/// Drives ticks on `core` until `predicate` holds, failing the calling test if it does not hold +/// within [`TICK_DEADLINE`]. +/// +/// A macro rather than an async function so that the predicate is an expression rather than a +/// closure, and may therefore borrow whatever the tick mutates. +macro_rules! tick_until { + ($core:expr, $predicate:expr) => {{ + let deadline = tokio::time::Instant::now() + TICK_DEADLINE; + loop { + $core.tick().await?; + if $predicate { + break; + } + if deadline < tokio::time::Instant::now() { + ::anyhow::bail!("the core did not reach the expected state in time"); + } + tokio::time::sleep(TICK_RETRY_INTERVAL).await; + } + }}; +} + +use std::collections::HashSet; +use std::num::NonZeroU64; +use std::num::NonZeroUsize; +use std::time::Duration; + +use anyhow::bail; +use spider_core::session::SessionTracker; +use spider_core::task::TaskIndex; +use spider_core::types::id::JobId; +use spider_core::types::id::ResourceGroupId; +use spider_core::types::id::SessionId; +use spider_core::types::id::TaskAssignmentId; +use spider_core::types::id::TaskId; +use tokio_util::sync::CancellationToken; + +use super::dispatch_queue::DispatchQueueRegistry; +use super::dispatch_queue::RgDispatchQueueReader; +use super::implementation::RgRoundRobin; +use super::implementation::RgRoundRobinConfig; +use super::job_registry::JobKey; +use super::job_registry::JobRegistry; +use super::job_registry::UpsertOutcome; +use super::scheduling_state::RgSchedulingState; +use crate::SchedulerError; +use crate::TaskAssignment; +use crate::core::TaskAssignmentIdIssuer; +use crate::core_impl::inbound_queue_reader::test_harness::DEFAULT_SESSION_ID; +use crate::core_impl::inbound_queue_reader::test_harness::MockStorageClient; +use crate::core_impl::inbound_queue_reader::test_harness::make_entry; + +/// The number of active jobs a resource group may hold. +const ACTIVE_JOB_LIST_CAPACITY: usize = 4; + +/// The first resource group a test seeds. +const RG_A: ResourceGroupId = ResourceGroupId::from(0); + +/// The second resource group a test seeds. +const RG_B: ResourceGroupId = ResourceGroupId::from(1); + +/// The third resource group a test seeds. +const RG_C: ResourceGroupId = ResourceGroupId::from(2); + +/// The session a test bumps the mock storage to. +const NEXT_SESSION_ID: SessionId = DEFAULT_SESSION_ID + 1; + +/// The storage poll timeout every test runs with. The mock storage never blocks on it. +const STORAGE_POLL_TIMEOUT_MS: u64 = 10; + +/// The longest a test waits for the ticks it drives to reach the state it expects. +const TICK_DEADLINE: Duration = Duration::from_secs(10); + +/// The interval between two ticks driven by [`tick_until`]. +const TICK_RETRY_INTERVAL: Duration = Duration::from_millis(2); + +/// A core wired to a mock storage and to the dispatch structures a test inspects, driven by manual +/// [`RgRoundRobin::tick`] calls. +struct CoreFixture { + core: RgRoundRobin, + storage: MockStorageClient, + dispatch_queue_registry: DispatchQueueRegistry, + session_tracker: SessionTracker, + reschedule_queue_writer: tokio::sync::mpsc::UnboundedSender, + active_job_list_capacity: usize, +} + +impl CoreFixture { + /// Factory function. + /// + /// # Returns + /// + /// A newly created fixture whose core holds no buffered task and no active resource group. + fn new(config: RgRoundRobinConfig, storage: MockStorageClient) -> Self { + let active_job_list_capacity = config.active_job_list_capacity.get(); + let (reschedule_queue_writer, reschedule_queue_reader) = + tokio::sync::mpsc::unbounded_channel(); + let core = RgRoundRobin::new( + storage.clone(), + reschedule_queue_reader, + TaskAssignmentIdIssuer::new(), + CancellationToken::new(), + config, + ); + let dispatch_queue_registry = core.dispatch_queue_registry(); + let session_tracker = core.session_tracker.clone(); + Self { + core, + storage, + dispatch_queue_registry, + session_tracker, + reschedule_queue_writer, + active_job_list_capacity, + } + } + + /// Puts `rg_id` on the core's active resource group list, appending its scheduling state built + /// against the group's dispatch queue endpoints if the core has none. + /// + /// # Returns + /// + /// The position of the group's scheduling state in the core's state vector. + fn activate_group(&mut self, rg_id: ResourceGroupId) -> usize { + if let Some(state_index) = self.core.rg_index.get(&rg_id) { + return *state_index; + } + + let mut rg_state = RgSchedulingState::new( + rg_id, + self.dispatch_queue_registry + .get_dispatch_queue_writer(rg_id), + self.active_job_list_capacity, + ); + rg_state.is_active = true; + let state_index = self.core.rg_states.len(); + self.core.rg_states.push(rg_state); + self.core.rg_index.insert(rg_id, state_index); + self.core.active_rg_list.push(state_index); + state_index + } + + /// Registers a job of `num_tasks` buffered ready tasks against `rg_id`, exactly as a completed + /// inbound poll carrying that job would, and activates the group. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * Forwards [`make_job_entry`]'s return values on failure. + fn seed_job( + &mut self, + rg_id: ResourceGroupId, + job_id: JobId, + num_tasks: usize, + ) -> anyhow::Result<()> { + for task_index in 0..num_tasks { + self.core + .global_task_set + .insert((job_id, TaskId::Index(task_index))); + } + let job_key = make_job_entry(&mut self.core.job_registry, job_id, num_tasks)?; + let state_index = self.activate_group(rg_id); + self.core.rg_states[state_index].place_new_job(job_key); + Ok(()) + } + + /// Puts `num_assignments` assignments straight into `rg_id`'s dispatch queue, standing in for + /// the occupancy a group carries over from earlier ticks. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * Forwards [`super::dispatch_queue::RgDispatchQueueWriter::try_send`]'s return values on + /// failure. + fn preload_queue(&self, rg_id: ResourceGroupId, num_assignments: usize) -> anyhow::Result<()> { + let writer = self + .dispatch_queue_registry + .get_dispatch_queue_writer(rg_id); + for index in 0..num_assignments { + writer.try_send(make_assignment( + rg_id, + JobId::from(u64::MAX), + TaskId::Index(index), + self.session_tracker.current(), + ))?; + } + Ok(()) + } + + /// # Returns + /// + /// The number of assignments currently queued for `rg_id`. + fn queue_len(&self, rg_id: ResourceGroupId) -> usize { + self.dispatch_queue_registry + .get_dispatch_queue_writer(rg_id) + .queue_len() + } + + /// # Returns + /// + /// The read side of `rg_id`'s dispatch queue, which a test hands to [`drain_reader`] to play a + /// pinned execution manager. + fn reader(&self, rg_id: ResourceGroupId) -> RgDispatchQueueReader { + self.dispatch_queue_registry + .get_dispatch_queue_reader(rg_id) + } +} + +#[tokio::test] +async fn the_rotation_arm_persists_across_ticks() -> anyhow::Result<()> { + const DISPATCH_QUEUE_CAPACITY: usize = 2; + + let mut fixture = CoreFixture::new( + make_config(DISPATCH_QUEUE_CAPACITY, ACTIVE_JOB_LIST_CAPACITY), + MockStorageClient::new(), + ); + for (index, rg_id) in [RG_A, RG_B, RG_C].into_iter().enumerate() { + fixture.seed_job(rg_id, JobId::from(u64::try_from(index)?), 8)?; + } + + fixture.core.tick().await?; + assert_eq!(occupancies_of(&fixture, 3), vec![1, 1, 0]); + assert_eq!(fixture.core.last_served_rg, Some(RG_B)); + + for rg_id in [RG_A, RG_B, RG_C] { + drain_reader(&fixture.reader(rg_id)).await; + } + + fixture.core.tick().await?; + assert_eq!(occupancies_of(&fixture, 3), vec![1, 0, 1]); + assert_eq!(fixture.core.last_served_rg, Some(RG_A)); + Ok(()) +} + +#[tokio::test] +async fn dropping_an_exhausted_group_does_not_skip_the_group_moved_into_its_slot() +-> anyhow::Result<()> { + const DISPATCH_QUEUE_CAPACITY: usize = 1; + + let mut fixture = CoreFixture::new( + make_config(DISPATCH_QUEUE_CAPACITY, ACTIVE_JOB_LIST_CAPACITY), + MockStorageClient::new(), + ); + fixture.seed_job(RG_A, JobId::from(0), 8)?; + fixture.activate_group(RG_B); + fixture.seed_job(RG_C, JobId::from(2), 8)?; + + // The arm starts on the group that has nothing to schedule, so the tick's single assignment is + // won by whichever group the removal moves into the vacated slot. + fixture.core.last_served_rg = Some(RG_A); + fixture.core.tick().await?; + + assert_eq!(occupancies_of(&fixture, 3), vec![0, 0, 1]); + assert_eq!(fixture.core.last_served_rg, Some(RG_C)); + assert_eq!(fixture.core.active_rg_list.len(), 2); + Ok(()) +} + +#[tokio::test] +async fn an_exhausted_group_stays_active_until_its_dispatch_queue_drains() -> anyhow::Result<()> { + const DISPATCH_QUEUE_CAPACITY: usize = 4; + + let mut fixture = CoreFixture::new( + make_config(DISPATCH_QUEUE_CAPACITY, ACTIVE_JOB_LIST_CAPACITY), + MockStorageClient::new(), + ); + fixture.activate_group(RG_A); + fixture.preload_queue(RG_A, 1)?; + + fixture.core.tick().await?; + assert_eq!(fixture.core.active_rg_list.len(), 1); + assert!(is_active(&fixture, RG_A)); + + assert_eq!(drain_reader(&fixture.reader(RG_A)).await.len(), 1); + fixture.core.tick().await?; + assert_eq!(fixture.core.active_rg_list, Vec::::new()); + assert!(!is_active(&fixture, RG_A)); + Ok(()) +} + +#[tokio::test] +async fn dispatching_and_retirement_run_while_a_storage_poll_is_in_flight() -> anyhow::Result<()> { + const DISPATCH_QUEUE_CAPACITY: usize = 4; + // The lone group's admission threshold caps it at half the dispatch buffer per tick. + const NUM_TASKS_PER_TICK: usize = DISPATCH_QUEUE_CAPACITY / 2; + const NUM_TASKS: usize = 2 * NUM_TASKS_PER_TICK; + + let storage = MockStorageClient::new(); + storage.gate_ready_lane(); + let mut fixture = CoreFixture::new( + make_config(DISPATCH_QUEUE_CAPACITY, ACTIVE_JOB_LIST_CAPACITY), + storage, + ); + fixture.seed_job(RG_A, JobId::from(0), NUM_TASKS)?; + fixture.seed_job(RG_B, JobId::from(1), 0)?; + + // The first tick starts the poll the gate holds, so every later tick finds it still in flight. + fixture.core.tick().await?; + assert_eq!(fixture.queue_len(RG_A), NUM_TASKS_PER_TICK); + assert_eq!(fixture.core.job_registry.len(), 2); + + assert_eq!( + drain_reader(&fixture.reader(RG_A)).await.len(), + NUM_TASKS_PER_TICK + ); + fixture.core.tick().await?; + assert_eq!(fixture.queue_len(RG_A), NUM_TASKS_PER_TICK); + assert_eq!(fixture.core.global_task_set, HashSet::new()); + + fixture.core.tick().await?; + assert_eq!(fixture.core.job_registry.len(), 1); + assert_eq!(fixture.core.active_rg_list.len(), 1); + assert_eq!(fixture.storage.num_polls().0, 0); + Ok(()) +} + +#[tokio::test] +async fn a_session_bump_clears_the_dedup_set_and_the_finalized_job_table() -> anyhow::Result<()> { + const DISPATCH_QUEUE_CAPACITY: usize = 8; + const SENTINEL_JOB_ID: JobId = JobId::from(4096); + + let storage = MockStorageClient::new(); + storage.push_ready_batch( + DEFAULT_SESSION_ID, + vec![ + make_entry(RG_A, JobId::from(0), TaskId::Index(0)), + make_entry(RG_A, JobId::from(0), TaskId::Index(1)), + ], + ); + let mut fixture = CoreFixture::new( + make_config(DISPATCH_QUEUE_CAPACITY, ACTIVE_JOB_LIST_CAPACITY), + storage, + ); + tick_until!(fixture.core, 2 == fixture.queue_len(RG_A)); + + fixture + .core + .global_task_set + .insert((SENTINEL_JOB_ID, TaskId::Index(0))); + fixture.core.finalized_jobs.insert(SENTINEL_JOB_ID); + + fixture + .storage + .push_ready_batch(NEXT_SESSION_ID, Vec::new()); + tick_until!( + fixture.core, + NEXT_SESSION_ID == fixture.session_tracker.current() + ); + + assert!( + !fixture + .core + .global_task_set + .contains(&(SENTINEL_JOB_ID, TaskId::Index(0))) + ); + assert!(!fixture.core.finalized_jobs.contains(&SENTINEL_JOB_ID)); + Ok(()) +} + +#[tokio::test] +async fn a_session_bump_readmits_the_tasks_storage_replays() -> anyhow::Result<()> { + const DISPATCH_QUEUE_CAPACITY: usize = 1; + + let replayed_entries = vec![ + make_entry(RG_A, JobId::from(0), TaskId::Index(0)), + make_entry(RG_A, JobId::from(0), TaskId::Index(1)), + ]; + let storage = MockStorageClient::new(); + storage.push_ready_batch(DEFAULT_SESSION_ID, replayed_entries.clone()); + let mut fixture = CoreFixture::new( + make_config(DISPATCH_QUEUE_CAPACITY, ACTIVE_JOB_LIST_CAPACITY), + storage, + ); + tick_until!(fixture.core, 1 == fixture.queue_len(RG_A)); + + // The buffer holds one assignment, so the job's second task is still buffered in the core and + // still in the dedup set when the bump lands. + let published = drain_reader(&fixture.reader(RG_A)).await; + assert_eq!(published.len(), 1); + assert_eq!(published[0].session_id, DEFAULT_SESSION_ID); + assert_eq!(fixture.core.global_task_set.len(), 1); + + fixture + .storage + .push_ready_batch(NEXT_SESSION_ID, replayed_entries); + let mut replayed = Vec::new(); + let deadline = tokio::time::Instant::now() + TICK_DEADLINE; + while replayed.len() < 2 { + fixture.core.tick().await?; + // A poll issued before the bump still lands under the old session, so the assignments it + // produces are stale by construction and are not part of the replay. + replayed.extend( + drain_reader(&fixture.reader(RG_A)) + .await + .into_iter() + .filter(|assignment| assignment.session_id == NEXT_SESSION_ID), + ); + if deadline < tokio::time::Instant::now() { + bail!("storage's replayed tasks were not re-admitted: {replayed:?}"); + } + tokio::time::sleep(TICK_RETRY_INTERVAL).await; + } + + let replayed_task_ids: HashSet = replayed + .iter() + .map(|assignment| assignment.task_id) + .collect(); + assert_eq!( + replayed_task_ids, + HashSet::from([TaskId::Index(0), TaskId::Index(1)]) + ); + Ok(()) +} + +#[tokio::test] +async fn a_rescheduled_assignment_is_readmitted() -> anyhow::Result<()> { + const DISPATCH_QUEUE_CAPACITY: usize = 8; + const LOST_JOB_ID: JobId = JobId::from(0); + const LOST_TASK_ID: TaskId = TaskId::Index(9); + + let mut fixture = CoreFixture::new( + make_config(DISPATCH_QUEUE_CAPACITY, ACTIVE_JOB_LIST_CAPACITY), + MockStorageClient::new(), + ); + let lost = make_assignment( + RG_A, + LOST_JOB_ID, + LOST_TASK_ID, + fixture.session_tracker.current(), + ); + fixture.reschedule_queue_writer.send(lost)?; + + tick_until!(fixture.core, 1 == fixture.queue_len(RG_A)); + + let redispatched = drain_reader(&fixture.reader(RG_A)).await; + assert_eq!(redispatched.len(), 1); + assert_eq!(redispatched[0].task_id, LOST_TASK_ID); + assert_eq!(redispatched[0].job_id, LOST_JOB_ID); + Ok(()) +} + +#[tokio::test] +async fn a_closed_dispatch_queue_fails_the_tick() -> anyhow::Result<()> { + const DISPATCH_QUEUE_CAPACITY: usize = 4; + + let mut fixture = CoreFixture::new( + make_config(DISPATCH_QUEUE_CAPACITY, ACTIVE_JOB_LIST_CAPACITY), + MockStorageClient::new(), + ); + fixture.seed_job(RG_A, JobId::from(0), 4)?; + fixture.dispatch_queue_registry.close_dispatch_queue(RG_A); + + let err = fixture + .core + .tick() + .await + .expect_err("a closed dispatch queue fails the tick"); + let SchedulerError::DispatchQueueClosed = err else { + bail!("the tick failed with something other than a closed dispatch queue: {err:?}"); + }; + Ok(()) +} + +#[tokio::test] +async fn a_closed_broadcast_queue_fails_the_tick() -> anyhow::Result<()> { + const DISPATCH_QUEUE_CAPACITY: usize = 4; + + let mut fixture = CoreFixture::new( + make_config(DISPATCH_QUEUE_CAPACITY, ACTIVE_JOB_LIST_CAPACITY), + MockStorageClient::new(), + ); + fixture.seed_job(RG_A, JobId::from(0), 4)?; + fixture.dispatch_queue_registry.close_broadcast_queue(); + + let err = fixture + .core + .tick() + .await + .expect_err("a closed broadcast queue fails the tick"); + let SchedulerError::DispatchQueueClosed = err else { + bail!("the tick failed with something other than a closed dispatch queue: {err:?}"); + }; + + // Both closures fail the tick with the same error, so the queue is what tells them apart: this + // assignment reached the group's queue first and lost only the hint covering it. + assert_eq!(fixture.queue_len(RG_A), 1); + Ok(()) +} + +#[tokio::test] +async fn the_scheduling_loop_stops_when_it_is_cancelled() -> anyhow::Result<()> { + const DISPATCH_QUEUE_CAPACITY: usize = 4; + + let (_reschedule_queue_writer, reschedule_queue_reader) = + tokio::sync::mpsc::unbounded_channel(); + let cancellation_token = CancellationToken::new(); + let core = RgRoundRobin::new( + MockStorageClient::new(), + reschedule_queue_reader, + TaskAssignmentIdIssuer::new(), + cancellation_token.clone(), + make_config(DISPATCH_QUEUE_CAPACITY, ACTIVE_JOB_LIST_CAPACITY), + ); + let scheduler_handle = tokio::task::spawn(core.run()); + + cancellation_token.cancel(); + tokio::time::timeout(TICK_DEADLINE, scheduler_handle).await???; + Ok(()) +} + +#[tokio::test] +async fn one_tick_leaves_every_backlogged_group_at_the_dynamic_threshold() -> anyhow::Result<()> { + let expected_share = ADMISSION_DISPATCH_QUEUE_CAPACITY / (NUM_BACKLOGGED_GROUPS + 1); + let mut fixture = new_admission_fixture(); + seed_backlogged_groups(&mut fixture, NUM_BACKLOGGED_GROUPS)?; + + fixture.core.tick().await?; + + let occupancies = occupancies_of(&fixture, NUM_BACKLOGGED_GROUPS); + let occupancy: usize = occupancies.iter().sum(); + let free = ADMISSION_DISPATCH_QUEUE_CAPACITY - occupancy; + for (index, group_occupancy) in occupancies.iter().enumerate() { + assert!( + group_occupancy.abs_diff(expected_share) <= SHARE_TOLERANCE, + "group {index} holds {group_occupancy} assignments, expected about {expected_share}: \ + {occupancies:?}" + ); + } + assert!( + free.abs_diff(expected_share) <= SHARE_TOLERANCE, + "the tick left {free} free, expected about {expected_share}: {occupancies:?}" + ); + + let published: usize = + NUM_BACKLOGGED_GROUPS * NUM_TASKS_PER_JOB - fixture.core.global_task_set.len(); + assert_eq!(published, occupancy); + Ok(()) +} + +#[tokio::test] +async fn no_group_is_batch_filled_while_another_waits() -> anyhow::Result<()> { + let mut fixture = new_admission_fixture(); + seed_backlogged_groups(&mut fixture, NUM_BACKLOGGED_GROUPS)?; + + fixture.core.tick().await?; + + let occupancies = occupancies_of(&fixture, NUM_BACKLOGGED_GROUPS); + let most = *occupancies + .iter() + .max() + .expect("the tick served at least one group"); + let least = *occupancies + .iter() + .min() + .expect("the tick served at least one group"); + assert!( + most - least <= 2, + "the rotation is not interleaved at quantum 1: {occupancies:?}" + ); + assert!( + most < ADMISSION_DISPATCH_QUEUE_CAPACITY / 2, + "a single group took half the dispatch buffer: {occupancies:?}" + ); + assert!( + least > 0, + "a backlogged group was left at zero: {occupancies:?}" + ); + Ok(()) +} + +#[tokio::test] +async fn a_newly_active_group_is_admitted_against_a_backlogged_incumbent() -> anyhow::Result<()> { + const INCUMBENT_OCCUPANCY: usize = 100; + + let mut fixture = new_admission_fixture(); + seed_backlogged_groups(&mut fixture, 2)?; + fixture.preload_queue(RG_A, INCUMBENT_OCCUPANCY)?; + + fixture.core.tick().await?; + + let occupancies = occupancies_of(&fixture, 2); + let free = ADMISSION_DISPATCH_QUEUE_CAPACITY - occupancies.iter().sum::(); + assert!( + occupancies[1] >= ADMISSION_DISPATCH_QUEUE_CAPACITY / 8, + "the newly active group was starved by the incumbent: {occupancies:?}" + ); + assert!( + occupancies[0] - occupancies[1] < INCUMBENT_OCCUPANCY, + "the incumbent's head start grew instead of shrinking: {occupancies:?}" + ); + assert!( + free >= ADMISSION_DISPATCH_QUEUE_CAPACITY / 8, + "the tick left only {free} free for the next group to arrive: {occupancies:?}" + ); + Ok(()) +} + +#[tokio::test] +async fn a_lone_group_takes_no_more_than_half_the_dispatch_buffer() -> anyhow::Result<()> { + let mut fixture = new_admission_fixture(); + seed_backlogged_groups(&mut fixture, 1)?; + + fixture.core.tick().await?; + + let occupancy = fixture.queue_len(RG_A); + assert!( + occupancy <= ADMISSION_DISPATCH_QUEUE_CAPACITY / 2, + "the only active group holds {occupancy} of {ADMISSION_DISPATCH_QUEUE_CAPACITY} \ + assignments" + ); + assert!( + occupancy >= ADMISSION_DISPATCH_QUEUE_CAPACITY / 4, + "the only active group holds only {occupancy} of {ADMISSION_DISPATCH_QUEUE_CAPACITY} \ + assignments" + ); + Ok(()) +} + +/// The dispatch buffer capacity `B` of the admission tests, matching the design document's worked +/// example. +const ADMISSION_DISPATCH_QUEUE_CAPACITY: usize = 256; + +/// The number of backlogged resource groups `N` in the equilibrium tests. +const NUM_BACKLOGGED_GROUPS: usize = 5; + +/// The number of ready tasks each backlogged group is seeded with, more than any single tick may +/// publish. +const NUM_TASKS_PER_JOB: usize = 512; + +/// How far a group's occupancy may sit from the equilibrium share before the rotation is considered +/// broken. The staircase these tests exist to catch misses by an order of magnitude: batch filling +/// five groups against `B = 256` yields `64, 64, 64, 64, 0` with no free space at all. +const SHARE_TOLERANCE: usize = 6; + +/// # Returns +/// +/// A core config with the given capacities, whose remaining tunables never throttle a test. +/// +/// # Panics +/// +/// Panics if either capacity is zero. +fn make_config( + dispatch_queue_capacity: usize, + active_job_list_capacity: usize, +) -> RgRoundRobinConfig { + RgRoundRobinConfig { + dispatch_queue_capacity: NonZeroUsize::new(dispatch_queue_capacity) + .expect("the dispatch queue capacity is non-zero"), + active_job_list_capacity: NonZeroUsize::new(active_job_list_capacity) + .expect("the active job list capacity is non-zero"), + ready_task_capacity: NonZeroUsize::new(16_384).expect("16384 is non-zero"), + commit_ready_task_capacity: NonZeroUsize::new(64).expect("64 is non-zero"), + cleanup_ready_task_capacity: NonZeroUsize::new(64).expect("64 is non-zero"), + storage_poll_timeout_ms: STORAGE_POLL_TIMEOUT_MS, + tick_interval_ms: NonZeroU64::new(1).expect("1 is non-zero"), + } +} + +/// # Returns +/// +/// A newly created fixture over a mock storage with no scripted batch, so a tick's decisions come +/// from the tasks the test seeded and from nothing else. +fn new_admission_fixture() -> CoreFixture { + CoreFixture::new( + make_config(ADMISSION_DISPATCH_QUEUE_CAPACITY, ACTIVE_JOB_LIST_CAPACITY), + MockStorageClient::new(), + ) +} + +/// Seeds the first `num_groups` resource groups with one job each, backlogged with more ready tasks +/// than a single tick may publish. +/// +/// # Errors +/// +/// Returns an error if: +/// +/// * Forwards [`CoreFixture::seed_job`]'s return values on failure. +/// * Forwards [`u64::try_from`]'s return values on failure. +fn seed_backlogged_groups(fixture: &mut CoreFixture, num_groups: usize) -> anyhow::Result<()> { + for index in 0..num_groups { + let raw_id = u64::try_from(index)?; + fixture.seed_job( + ResourceGroupId::from(raw_id), + JobId::from(raw_id), + NUM_TASKS_PER_JOB, + )?; + } + Ok(()) +} + +/// Registers a job of `num_tasks` buffered ready tasks, whose task indices are `0..num_tasks`. +/// +/// # Returns +/// +/// The key of the registered job, which still needs a scheduling position. +/// +/// # Errors +/// +/// Returns an error if: +/// +/// * [`anyhow::Error`] if the job is already registered. +fn make_job_entry( + registry: &mut JobRegistry, + job_id: JobId, + num_tasks: usize, +) -> anyhow::Result { + let task_indices: Vec = (0..num_tasks).collect(); + let UpsertOutcome::New(job_key) = registry.upsert(job_id, task_indices) else { + bail!("job {job_id} is already registered"); + }; + Ok(job_key) +} + +/// # Returns +/// +/// An assignment of `task_id` carrying an ID no assignment the core publishes can collide with. +fn make_assignment( + rg_id: ResourceGroupId, + job_id: JobId, + task_id: TaskId, + session_id: SessionId, +) -> TaskAssignment { + TaskAssignment { + id: TaskAssignmentId::from(u64::MAX), + resource_group_id: rg_id, + job_id, + task_id, + session_id, + } +} + +/// # Returns +/// +/// The number of assignments queued for each of the first `num_groups` resource groups, indexed by +/// resource group ID. +fn occupancies_of(fixture: &CoreFixture, num_groups: usize) -> Vec { + (0..num_groups) + .map(|index| { + let raw_id = u64::try_from(index).expect("a group index fits a raw ID"); + fixture.queue_len(ResourceGroupId::from(raw_id)) + }) + .collect() +} + +/// # Returns +/// +/// Whether the core still holds `rg_id` on its active resource group list. +/// +/// # Panics +/// +/// Panics if the core has no scheduling state for `rg_id`. +fn is_active(fixture: &CoreFixture, rg_id: ResourceGroupId) -> bool { + let state_index = *fixture + .core + .rg_index + .get(&rg_id) + .expect("the core holds a scheduling state for the group"); + fixture.core.rg_states[state_index].is_active +} + +/// Takes every assignment currently queued behind `reader`, playing a pinned execution manager, +/// which leaves the group's hint counter untouched. +/// +/// # Returns +/// +/// The assignments taken, in dispatch order. +async fn drain_reader(reader: &RgDispatchQueueReader) -> Vec { + let mut assignments = Vec::new(); + while let Some(assignment) = reader.recv_pinned(Duration::ZERO).await { + assignments.push(assignment); + } + assignments +} From a98e98f00254e7e9d6bcc6e0fef58d92299ffe95 Mon Sep 17 00:00:00 2001 From: LinZhihao-723 Date: Tue, 1 Sep 2026 18:51:07 -0400 Subject: [PATCH 2/5] Reviewed the core implementation. need some more time to review and think about. Otherwise has been polished. Haven't looked into test cases yet. --- .../src/core_impl/inbound_queue_reader.rs | 32 ++ .../implementation.rs | 304 ++++++++---- .../scheduling_state.rs | 31 +- .../resource_group_round_robin/tests.rs | 438 +++++++++++++++++- 4 files changed, 683 insertions(+), 122 deletions(-) diff --git a/components/spider-scheduler/src/core_impl/inbound_queue_reader.rs b/components/spider-scheduler/src/core_impl/inbound_queue_reader.rs index fcde55ab..d4972e26 100644 --- a/components/spider-scheduler/src/core_impl/inbound_queue_reader.rs +++ b/components/spider-scheduler/src/core_impl/inbound_queue_reader.rs @@ -336,6 +336,7 @@ pub(super) mod test_harness { use std::sync::Mutex; use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicU64; + use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::time::Duration; @@ -415,6 +416,9 @@ pub(super) mod test_harness { num_cleanup_ready_polls: AtomicU64::new(0), is_ready_lane_gated: AtomicBool::new(false), ready_lane_gate: Semaphore::new(0), + last_ready_limit: AtomicUsize::new(0), + last_commit_ready_limit: AtomicUsize::new(0), + last_cleanup_ready_limit: AtomicUsize::new(0), }), } } @@ -464,6 +468,22 @@ pub(super) mod test_harness { ) } + /// # Returns + /// + /// A tuple containing the entry limit each lane was given by the poll it served most + /// recently, or 0 for a lane that has served none: + /// + /// * The regular-task lane's entry limit. + /// * The commit-task lane's entry limit. + /// * The cleanup-task lane's entry limit. + pub fn last_poll_limits(&self) -> (usize, usize, usize) { + ( + self.inner.last_ready_limit.load(Ordering::Relaxed), + self.inner.last_commit_ready_limit.load(Ordering::Relaxed), + self.inner.last_cleanup_ready_limit.load(Ordering::Relaxed), + ) + } + /// Holds every subsequent regular-task poll until [`Self::admit_ready_poll`] releases it. pub fn gate_ready_lane(&self) { self.inner @@ -528,6 +548,9 @@ pub(super) mod test_harness { .forget(); } self.inner.num_ready_polls.fetch_add(1, Ordering::Relaxed); + self.inner + .last_ready_limit + .store(max_items, Ordering::Relaxed); Ok(self.serve_batch(&self.inner.ready_batches, max_items)) } @@ -539,6 +562,9 @@ pub(super) mod test_harness { self.inner .num_commit_ready_polls .fetch_add(1, Ordering::Relaxed); + self.inner + .last_commit_ready_limit + .store(max_items, Ordering::Relaxed); Ok(self.serve_batch(&self.inner.commit_ready_batches, max_items)) } @@ -550,6 +576,9 @@ pub(super) mod test_harness { self.inner .num_cleanup_ready_polls .fetch_add(1, Ordering::Relaxed); + self.inner + .last_cleanup_ready_limit + .store(max_items, Ordering::Relaxed); Ok(self.serve_batch(&self.inner.cleanup_ready_batches, max_items)) } @@ -643,6 +672,9 @@ pub(super) mod test_harness { num_cleanup_ready_polls: AtomicU64, is_ready_lane_gated: AtomicBool, ready_lane_gate: Semaphore, + last_ready_limit: AtomicUsize, + last_commit_ready_limit: AtomicUsize, + last_cleanup_ready_limit: AtomicUsize, } } diff --git a/components/spider-scheduler/src/core_impl/resource_group_round_robin/implementation.rs b/components/spider-scheduler/src/core_impl/resource_group_round_robin/implementation.rs index 25867ea5..495bf398 100644 --- a/components/spider-scheduler/src/core_impl/resource_group_round_robin/implementation.rs +++ b/components/spider-scheduler/src/core_impl/resource_group_round_robin/implementation.rs @@ -3,9 +3,11 @@ use std::collections::HashMap; use std::collections::HashSet; +use std::collections::VecDeque; use std::num::NonZeroU64; use std::num::NonZeroUsize; use std::time::Duration; +use std::time::Instant; use serde::Deserialize; use spider_core::session::SessionTracker; @@ -58,17 +60,17 @@ pub(super) struct RgRoundRobinConfig { /// fill the inbound-queue reading request. pub(super) storage_poll_timeout_ms: u64, - /// The time (in milliseconds) that the scheduler will spend on each tick. + /// The time (in milliseconds) that the scheduler will spend on each tick. If the tick spends + /// less than the configured interval, the core will sleep for the remainder. pub(super) tick_interval_ms: NonZeroU64, + + /// The time (in seconds) that a job may remain in the finalized job table before the scheduler + /// drops it from the table. + pub(super) finalized_job_expiration_timeout_sec: u64, } /// The resource-group-aware round-robin scheduler core created from a [`RgRoundRobinConfig`]. /// -/// The core owns all of the state it decides with -- job entries in a generational arena, per -/// resource group scheduling states in an append-only vector, and the dispatch queue registry the -/// execution-manager-facing service reads from -- so nothing it holds across an await point is -/// thread-bound. -/// /// # Type Parameters /// /// * `SchedulerStorageClientType` - The storage client used to poll the inbound queue. @@ -78,8 +80,12 @@ pub(super) struct RgRoundRobinConfig { /// All member variables are marked `pub(super)` to allow the test module to inspect the internal /// states. pub(super) struct RgRoundRobin { - pub(super) global_task_set: HashSet<(JobId, TaskId)>, + pub(super) global_task_set: GlobalTaskSet, pub(super) finalized_jobs: HashSet, + + /// The insertion time of every job in [`Self::finalized_jobs`], in insertion order. + pub(super) finalized_job_queue: VecDeque<(JobId, Instant)>, + pub(super) job_registry: JobRegistry, /// The scheduling states of every resource group the core has seen this session. @@ -122,8 +128,9 @@ impl let session_tracker = SessionTracker::new(SessionId::default()); let dispatch_queue_registry = DispatchQueueRegistry::new(session_tracker.clone()); Self { - global_task_set: HashSet::new(), + global_task_set: GlobalTaskSet::new(), finalized_jobs: HashSet::new(), + finalized_job_queue: VecDeque::new(), job_registry: JobRegistry::new(), rg_states: Vec::new(), rg_id_to_index_map: HashMap::new(), @@ -200,14 +207,17 @@ impl /// Returns an error if: /// /// * Forwards [`RgInboundQueueReader::try_collect_result`]'s return values on failure. + /// * Forwards [`Self::apply_session_bump`]'s return values on failure. /// * Forwards [`Self::start_inbound_poll`]'s return values on failure. - /// * Forwards [`Self::fill_dispatch_queues`]'s return values on failure. + /// * Forwards [`Self::publish_task_assignments_into_dispatch_queues`]'s return values on + /// failure. + /// * Forwards [`Self::retire_jobs`]'s return values on failure. pub(super) async fn tick(&mut self) -> Result<(), SchedulerError> { - let poll_state = self + match self .inbound_queue_reader .try_collect_result(self.session_tracker.current()) - .await?; - match poll_state { + .await? + { RgInboundPollState::Ready { session_id, ready_result, @@ -215,7 +225,7 @@ impl cleanup_ready_result, } => { if session_id != self.session_tracker.current() { - self.apply_session_bump(session_id); + self.apply_session_bump(session_id)?; } let rescheduled_entries = self.drain_reschedule_queue(session_id); @@ -233,49 +243,39 @@ impl RgInboundPollState::Pending => (), } - let jobs_to_retire = self.fill_dispatch_queues()?; - self.retire_jobs(jobs_to_retire); + let jobs_to_retire = self.publish_task_assignments_into_dispatch_queues()?; + self.retire_jobs(jobs_to_retire)?; + self.retire_expired_finalized_jobs(); Ok(()) } /// Discards every piece of state published in a session older than `new_session_id`. /// - /// Clearing the dedup set is load-bearing rather than tidy: storage replays its ready tasks - /// after a bump, and a stale dedup entry would drop a replayed task while the registry no - /// longer holds anything to schedule it from. - /// - /// [`Self::rg_states`], [`Self::rg_index`], and [`Self::active_rg_list`] must be cleared as one - /// operation, and this is the only place any of them is cleared. Positions in `rg_states` carry - /// no generation, so an index that outlives the flush does not fail: it resolves against the - /// new session's states, either out of bounds or -- once the new session has re-created a few - /// groups -- silently against the wrong group. Nothing in the type system checks this. - /// - /// [`Self::dispatch_queue_registry`] must be cleared together with them, and that too is a - /// correctness requirement rather than tidiness. A group's queue closes only once every sender - /// has been dropped, and a scheduling state's write side is one of them; a state that survived - /// the registry's flush would therefore hold a write side onto a queue whose readers are gone, - /// and publishing into a closed queue is fatal to the core. Clearing the registry is also what - /// discards the hints published in the session being left behind. - fn apply_session_bump(&mut self, new_session_id: SessionId) { + /// # Errors + /// + /// Returns an error if: + /// + /// * [`SchedulerError::InvalidSessionId`] if `new_session_id` is no newer than the tracked + /// session. Storage only ever moves its session forward, so a session that does not advance + /// the tracker is unreachable in a healthy deployment. + fn apply_session_bump(&mut self, new_session_id: SessionId) -> Result<(), SchedulerError> { let previous_session_id = self.session_tracker.current(); - if self.session_tracker.try_advance(new_session_id) { - tracing::info!( - from = previous_session_id, - to = new_session_id, - num_resource_groups = self.dispatch_queue_registry.len(), - num_jobs = self.job_registry.len(), - "Storage session bumped. Flushing the core." - ); - } else { + if !self.session_tracker.try_advance(new_session_id) { tracing::error!( from = previous_session_id, to = new_session_id, - "Storage reported a session no newer than the tracked one. Flushing the core \ - anyway, but it keeps serving the tracked session." + "Storage reported a session no newer than the tracked one." ); + return Err(SchedulerError::InvalidSessionId(new_session_id)); } + tracing::info!( + from = previous_session_id, + to = new_session_id, + num_resource_groups = self.dispatch_queue_registry.len(), + num_jobs = self.job_registry.len(), + "Storage session bumped. Flushing the core." + ); - self.dispatch_queue_registry.clear(); self.rg_states.clear(); self.rg_id_to_index_map.clear(); self.active_rg_list.clear(); @@ -283,6 +283,10 @@ impl self.job_registry.clear(); self.global_task_set.clear(); self.finalized_jobs.clear(); + self.finalized_job_queue.clear(); + self.dispatch_queue_registry.clear(); + + Ok(()) } /// Drains the reschedule queue, dropping assignments published in a session other than @@ -306,11 +310,7 @@ impl entries } - /// Folds the tick's ready tasks into the finalized job table, the global task set, and the job - /// registry. - /// - /// Finalizations are processed before regular tasks, so a regular task arriving in the same - /// batch as its job's finalization is discarded rather than scheduled. + /// Processes the inbound polling results along with the assignments to reschedule. /// /// # Returns /// @@ -346,17 +346,25 @@ impl (cleanup_ready, FinalizeKind::Cleanup), ] { for finalized_job in finalized_jobs { - if !self.finalized_jobs.insert(finalized_job.job_id) { + // A job reaches at most one of each finalization, so the dedup key is the + // finalization rather than the job: a cleanup that follows a commit is a distinct + // task and must still be scheduled. + if !self + .global_task_set + .insert(finalized_job.job_id, TaskId::from(kind)) + { continue; } - // The job's still-buffered regular tasks will never be published, so they must - // leave the dedup set with it or nothing would ever remove them. - if let Some(mut job_entry) = - self.job_registry.remove_by_job_id(finalized_job.job_id) + // Only the first finalization has a registry entry to drop: the job's + // still-buffered regular tasks will never be published, so they must leave the + // dedup set with it or nothing would ever remove them. + if self.mark_job_finalized(finalized_job.job_id) + && let Some(mut job_entry) = + self.job_registry.remove_by_job_id(finalized_job.job_id) { for task_index in job_entry.take_ready_tasks() { self.global_task_set - .remove(&(finalized_job.job_id, TaskId::Index(task_index))); + .remove(finalized_job.job_id, TaskId::Index(task_index)); } } rg_updates @@ -378,7 +386,7 @@ impl } let global_task_set = &mut self.global_task_set; task_indices - .retain(|task_index| global_task_set.insert((job_id, TaskId::Index(*task_index)))); + .retain(|task_index| global_task_set.insert(job_id, TaskId::Index(*task_index))); if task_indices.is_empty() { continue; } @@ -406,11 +414,11 @@ impl for job_key in update.new_jobs { rg_state.place_new_job(job_key); } - let activated = !rg_state.is_active; - rg_state.is_active = true; - if activated { - self.active_rg_list.push(state_index); + if rg_state.is_active { + continue; } + rg_state.is_active = true; + self.active_rg_list.push(state_index); } } @@ -450,7 +458,9 @@ impl /// * [`SchedulerError::DispatchQueueClosed`] if a group's dispatch queue or the broadcast queue /// is closed, in which case the assignments the core makes can no longer reach an execution /// manager. - fn fill_dispatch_queues(&mut self) -> Result, SchedulerError> { + fn publish_task_assignments_into_dispatch_queues( + &mut self, + ) -> Result, SchedulerError> { let mut jobs_to_retire = Vec::new(); if self.active_rg_list.is_empty() { return Ok(jobs_to_retire); @@ -509,7 +519,7 @@ impl ); match result { Ok((job_id, task_id)) => { - global_task_set.remove(&(job_id, task_id)); + global_task_set.remove(job_id, task_id); free -= 1; *last_served_rg = Some(rg_state.rg_id); arm = (arm + 1) % rg_rr_list.len(); @@ -563,10 +573,59 @@ impl } } - /// Drops the registry's entry for every job that ran out of downgrade lives. - fn retire_jobs(&mut self, jobs_to_retire: Vec) { + /// Records that `job_id` has reached a finalizing state, so that its later regular tasks are + /// discarded rather than scheduled. + /// + /// A job that reaches both of its finalizations is recorded once: the table gates the job's + /// regular tasks, which the first finalization already settles. + /// + /// # Returns + /// + /// Whether this is the job's first finalization. + fn mark_job_finalized(&mut self, job_id: JobId) -> bool { + if !self.finalized_jobs.insert(job_id) { + return false; + } + self.finalized_job_queue.push_back((job_id, Instant::now())); + true + } + + /// Drops the job registry's entry for every job that ran out of downgrade lives. + /// + /// A key that no longer resolves is skipped rather than reported: the job it referred to was + /// removed when it finalized, and the key was buffered before that. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * [`SchedulerError::Internal`] if a retired job still buffers a ready task. + fn retire_jobs(&mut self, jobs_to_retire: Vec) -> Result<(), SchedulerError> { for job_key in jobs_to_retire { - self.job_registry.remove(job_key); + let Some(job_entry) = self.job_registry.remove(job_key) else { + continue; + }; + if job_entry.has_ready_task() { + return Err(SchedulerError::Internal(format!( + "retired job {:?} still buffers ready tasks", + job_entry.job_id() + ))); + } + } + + Ok(()) + } + + /// Drops every expired entry from the finalized job table. + fn retire_expired_finalized_jobs(&mut self) { + let expiration_time = Duration::from_secs(self.config.finalized_job_expiration_timeout_sec); + while let Some((job_id, insertion_time)) = self.finalized_job_queue.front() { + if insertion_time.elapsed() <= expiration_time { + break; + } + tracing::info!(job_id = ? job_id, "Finalized job table entry expired."); + self.finalized_jobs.remove(job_id); + self.finalized_job_queue.pop_front(); } } @@ -579,22 +638,21 @@ impl /// /// * Forwards [`RgInboundQueueReader::start`]'s return values on failure. fn start_inbound_poll(&mut self) -> Result<(), SchedulerError> { - let (num_commit_ready, num_cleanup_ready) = self.count_buffered_finalizations(); let max_ready_entries = self .config .ready_task_capacity .get() - .saturating_sub(self.global_task_set.len()); + .saturating_sub(self.global_task_set.num_ready()); let max_commit_ready_entries = self .config .commit_ready_task_capacity .get() - .saturating_sub(num_commit_ready); + .saturating_sub(self.global_task_set.num_commit_ready()); let max_cleanup_ready_entries = self .config .cleanup_ready_task_capacity .get() - .saturating_sub(num_cleanup_ready); + .saturating_sub(self.global_task_set.num_cleanup_ready()); self.inbound_queue_reader.start( Duration::from_millis(self.config.storage_poll_timeout_ms), @@ -603,22 +661,106 @@ impl max_cleanup_ready_entries, ) } +} + +/// The core's set of buffered tasks, carrying a running count of the tasks each inbound-queue lane +/// contributed. +/// +/// # Note +/// +/// [`Self::tasks`] is marked `pub(super)` to allow the test module to inspect the internal state. +pub(super) struct GlobalTaskSet { + pub(super) tasks: HashSet<(JobId, TaskId)>, + num_ready: usize, + num_commit_ready: usize, + num_cleanup_ready: usize, +} + +impl GlobalTaskSet { + /// Factory function. + /// + /// # Returns + /// + /// A newly created set. + pub(super) fn new() -> Self { + Self { + tasks: HashSet::new(), + num_ready: 0, + num_commit_ready: 0, + num_cleanup_ready: 0, + } + } + + /// Buffers `job_id`'s `task_id`, counting it against the lane it arrived on. + /// + /// # Returns + /// + /// Whether this call buffered the task. A task already buffered is not counted twice. + pub(super) fn insert(&mut self, job_id: JobId, task_id: TaskId) -> bool { + if !self.tasks.insert((job_id, task_id)) { + return false; + } + *self.lane_count_mut(task_id) += 1; + true + } + + /// Takes `job_id`'s `task_id` out of the buffer, discounting it from the lane it arrived on. + /// + /// A task that is not buffered leaves every count untouched. + pub(super) fn remove(&mut self, job_id: JobId, task_id: TaskId) { + if !self.tasks.remove(&(job_id, task_id)) { + return; + } + // The count and the set membership only ever move together, so the guard above is what + // makes this decrement unable to underflow. + *self.lane_count_mut(task_id) -= 1; + } + + /// Empties the buffer, zeroing every lane count with it. + pub(super) fn clear(&mut self) { + self.tasks.clear(); + self.num_ready = 0; + self.num_commit_ready = 0; + self.num_cleanup_ready = 0; + } + + /// # Returns + /// + /// The number of buffered tasks, across every lane. + pub(super) fn len(&self) -> usize { + self.tasks.len() + } + + /// # Returns + /// + /// The number of buffered regular tasks. + pub(super) const fn num_ready(&self) -> usize { + self.num_ready + } + + /// # Returns + /// + /// The number of buffered commit tasks. + pub(super) const fn num_commit_ready(&self) -> usize { + self.num_commit_ready + } + + /// # Returns + /// + /// The number of buffered cleanup tasks. + pub(super) const fn num_cleanup_ready(&self) -> usize { + self.num_cleanup_ready + } /// # Returns /// - /// A tuple containing: - /// - /// * The number of buffered commit tasks. - /// * The number of buffered cleanup tasks. - fn count_buffered_finalizations(&self) -> (usize, usize) { - let mut num_commit_ready = 0; - let mut num_cleanup_ready = 0; - for rg_state in &self.rg_states { - let (num_commits, num_cleanups) = rg_state.num_buffered_finalize_tasks(); - num_commit_ready += num_commits; - num_cleanup_ready += num_cleanups; + /// The count of the lane `task_id` arrives on. + const fn lane_count_mut(&mut self, task_id: TaskId) -> &mut usize { + match task_id { + TaskId::Index(_) => &mut self.num_ready, + TaskId::Commit => &mut self.num_commit_ready, + TaskId::Cleanup => &mut self.num_cleanup_ready, } - (num_commit_ready, num_cleanup_ready) } } diff --git a/components/spider-scheduler/src/core_impl/resource_group_round_robin/scheduling_state.rs b/components/spider-scheduler/src/core_impl/resource_group_round_robin/scheduling_state.rs index 3059e097..680ac520 100644 --- a/components/spider-scheduler/src/core_impl/resource_group_round_robin/scheduling_state.rs +++ b/components/spider-scheduler/src/core_impl/resource_group_round_robin/scheduling_state.rs @@ -73,8 +73,6 @@ pub(super) struct RgSchedulingState { pub(super) is_active: bool, finalize_queue: VecDeque<(JobId, FinalizeKind)>, - num_buffered_commits: usize, - num_buffered_cleanups: usize, writer: RgDispatchQueueWriter, active_to_pending_downgrade_buffer: Vec, pending_downgrade_buffer: Vec, @@ -99,8 +97,6 @@ impl RgSchedulingState { rr_arm: 0, is_active: false, finalize_queue: VecDeque::new(), - num_buffered_commits: 0, - num_buffered_cleanups: 0, writer, active_to_pending_downgrade_buffer: Vec::new(), pending_downgrade_buffer: Vec::new(), @@ -124,23 +120,9 @@ impl RgSchedulingState { || !self.pending_jobs.is_empty() } - /// # Returns - /// - /// A tuple containing: - /// - /// * The number of commit tasks the group has buffered. - /// * The number of cleanup tasks the group has buffered. - pub(super) const fn num_buffered_finalize_tasks(&self) -> (usize, usize) { - (self.num_buffered_commits, self.num_buffered_cleanups) - } - /// Records that `job_id` has reached the finalization named by `kind`. pub(super) fn push_finalization(&mut self, job_id: JobId, kind: FinalizeKind) { self.finalize_queue.push_back((job_id, kind)); - match kind { - FinalizeKind::Commit => self.num_buffered_commits += 1, - FinalizeKind::Cleanup => self.num_buffered_cleanups += 1, - } } /// Gives a newly registered job its scheduling position in this group. @@ -243,16 +225,7 @@ impl RgSchedulingState { /// /// The finalization to schedule, or [`None`] if the group owes no finalization. fn pop_finalization(&mut self) -> Option<(JobId, FinalizeKind)> { - let (job_id, kind) = self.finalize_queue.pop_front()?; - match kind { - FinalizeKind::Commit => { - self.num_buffered_commits = self.num_buffered_commits.saturating_sub(1); - } - FinalizeKind::Cleanup => { - self.num_buffered_cleanups = self.num_buffered_cleanups.saturating_sub(1); - } - } - Some((job_id, kind)) + self.finalize_queue.pop_front() } /// Takes the next regular task to dispatch out of the job that buffers it, rotating the arm @@ -836,7 +809,7 @@ mod tests { fixture.try_make(FREE_SPACE), Err(MakeAssignmentError::DispatchQueueClosed) ); - assert_eq!(fixture.state.num_buffered_finalize_tasks(), (0, 0)); + assert!(!fixture.state.has_schedulable_task()); } #[test] diff --git a/components/spider-scheduler/src/core_impl/resource_group_round_robin/tests.rs b/components/spider-scheduler/src/core_impl/resource_group_round_robin/tests.rs index db9b9b9e..87e36d49 100644 --- a/components/spider-scheduler/src/core_impl/resource_group_round_robin/tests.rs +++ b/components/spider-scheduler/src/core_impl/resource_group_round_robin/tests.rs @@ -1,8 +1,8 @@ //! Unit tests for the resource-group-aware round-robin scheduler core. //! -//! The tests drive [`RgRoundRobin::tick`] and the structures it decides with directly. Whatever an -//! execution manager would do -- draining a group's queue -- a test body does by hand through the -//! same entry points the dispatch service will use. +//! The tests drive [`RgRoundRobin::tick`] and the structures it decides with directly. +//! Whatever an execution manager would do -- draining a group's queue -- a test body does by hand +//! through the same entry points the dispatch service will use. /// Drives ticks on `core` until `predicate` holds, failing the calling test if it does not hold /// within [`TICK_DEADLINE`]. @@ -73,6 +73,26 @@ const NEXT_SESSION_ID: SessionId = DEFAULT_SESSION_ID + 1; /// The storage poll timeout every test runs with. The mock storage never blocks on it. const STORAGE_POLL_TIMEOUT_MS: u64 = 10; +/// The regular-task buffer capacity every test runs with. +const READY_TASK_CAPACITY: usize = 16_384; + +/// The commit-task buffer capacity every test runs with. +const COMMIT_READY_TASK_CAPACITY: usize = 64; + +/// The cleanup-task buffer capacity every test runs with. +const CLEANUP_READY_TASK_CAPACITY: usize = 64; + +/// A finalized job table expiry long enough that no test sweeps an entry unless it asks to. +const NEVER_EXPIRING_TIMEOUT_SEC: u64 = 6 * 60 * 60; + +/// The finalized job table expiry an expiry test runs with. +const SHORT_EXPIRATION_TIMEOUT_SEC: u64 = 1; + +/// How long an expiry test waits before the tick that must sweep an entry stamped +/// [`SHORT_EXPIRATION_TIMEOUT_SEC`] seconds ago. A sleep is a lower bound on the time that passes, +/// so the margin only has to cover the sweep reading a monotonic clock, never scheduling delay. +const EXPIRATION_WAIT: Duration = Duration::from_millis(1_200); + /// The longest a test waits for the ticks it drives to reach the state it expects. const TICK_DEADLINE: Duration = Duration::from_secs(10); @@ -126,7 +146,7 @@ impl CoreFixture { /// /// The position of the group's scheduling state in the core's state vector. fn activate_group(&mut self, rg_id: ResourceGroupId) -> usize { - if let Some(state_index) = self.core.rg_index.get(&rg_id) { + if let Some(state_index) = self.core.rg_id_to_index_map.get(&rg_id) { return *state_index; } @@ -139,7 +159,7 @@ impl CoreFixture { rg_state.is_active = true; let state_index = self.core.rg_states.len(); self.core.rg_states.push(rg_state); - self.core.rg_index.insert(rg_id, state_index); + self.core.rg_id_to_index_map.insert(rg_id, state_index); self.core.active_rg_list.push(state_index); state_index } @@ -161,7 +181,7 @@ impl CoreFixture { for task_index in 0..num_tasks { self.core .global_task_set - .insert((job_id, TaskId::Index(task_index))); + .insert(job_id, TaskId::Index(task_index)); } let job_key = make_job_entry(&mut self.core.job_registry, job_id, num_tasks)?; let state_index = self.activate_group(rg_id); @@ -311,7 +331,7 @@ async fn dispatching_and_retirement_run_while_a_storage_poll_is_in_flight() -> a ); fixture.core.tick().await?; assert_eq!(fixture.queue_len(RG_A), NUM_TASKS_PER_TICK); - assert_eq!(fixture.core.global_task_set, HashSet::new()); + assert_eq!(fixture.core.global_task_set.tasks, HashSet::new()); fixture.core.tick().await?; assert_eq!(fixture.core.job_registry.len(), 1); @@ -342,7 +362,7 @@ async fn a_session_bump_clears_the_dedup_set_and_the_finalized_job_table() -> an fixture .core .global_task_set - .insert((SENTINEL_JOB_ID, TaskId::Index(0))); + .insert(SENTINEL_JOB_ID, TaskId::Index(0)); fixture.core.finalized_jobs.insert(SENTINEL_JOB_ID); fixture @@ -357,6 +377,7 @@ async fn a_session_bump_clears_the_dedup_set_and_the_finalized_job_table() -> an !fixture .core .global_task_set + .tasks .contains(&(SENTINEL_JOB_ID, TaskId::Index(0))) ); assert!(!fixture.core.finalized_jobs.contains(&SENTINEL_JOB_ID)); @@ -493,6 +514,347 @@ async fn a_closed_broadcast_queue_fails_the_tick() -> anyhow::Result<()> { Ok(()) } +#[tokio::test] +async fn an_expired_finalized_job_leaves_the_table_while_a_fresh_one_stays() -> anyhow::Result<()> { + const DISPATCH_QUEUE_CAPACITY: usize = 8; + const EXPIRING_JOB_ID: JobId = JobId::from(0); + const FRESH_JOB_ID: JobId = JobId::from(1); + + let storage = MockStorageClient::new(); + storage.push_commit_ready_batch( + DEFAULT_SESSION_ID, + vec![make_entry(RG_A, EXPIRING_JOB_ID, TaskId::Commit)], + ); + let mut fixture = CoreFixture::new( + make_expiring_config( + DISPATCH_QUEUE_CAPACITY, + ACTIVE_JOB_LIST_CAPACITY, + SHORT_EXPIRATION_TIMEOUT_SEC, + ), + storage, + ); + tick_until!( + fixture.core, + fixture.core.finalized_jobs.contains(&EXPIRING_JOB_ID) + ); + assert_eq!(finalized_job_ids(&fixture), vec![EXPIRING_JOB_ID]); + + tokio::time::sleep(EXPIRATION_WAIT).await; + fixture.storage.push_commit_ready_batch( + DEFAULT_SESSION_ID, + vec![make_entry(RG_A, FRESH_JOB_ID, TaskId::Commit)], + ); + tick_until!( + fixture.core, + fixture.core.finalized_jobs.contains(&FRESH_JOB_ID) + ); + + assert_eq!(fixture.core.finalized_jobs, HashSet::from([FRESH_JOB_ID])); + assert_eq!(finalized_job_ids(&fixture), vec![FRESH_JOB_ID]); + Ok(()) +} + +#[tokio::test] +async fn a_cleanup_is_scheduled_after_the_same_job_committed() -> anyhow::Result<()> { + const DISPATCH_QUEUE_CAPACITY: usize = 8; + const JOB_ID: JobId = JobId::from(0); + + let storage = MockStorageClient::new(); + storage.push_commit_ready_batch( + DEFAULT_SESSION_ID, + vec![make_entry(RG_A, JOB_ID, TaskId::Commit)], + ); + let mut fixture = CoreFixture::new( + make_config(DISPATCH_QUEUE_CAPACITY, ACTIVE_JOB_LIST_CAPACITY), + storage, + ); + tick_until!(fixture.core, 1 == fixture.queue_len(RG_A)); + + fixture.storage.push_cleanup_ready_batch( + DEFAULT_SESSION_ID, + vec![make_entry(RG_A, JOB_ID, TaskId::Cleanup)], + ); + tick_until!(fixture.core, 2 == fixture.queue_len(RG_A)); + + let task_ids: Vec = drain_reader(&fixture.reader(RG_A)) + .await + .into_iter() + .map(|assignment| assignment.task_id) + .collect(); + assert_eq!(task_ids, vec![TaskId::Commit, TaskId::Cleanup]); + Ok(()) +} + +#[tokio::test] +async fn a_repeated_finalization_is_scheduled_once() -> anyhow::Result<()> { + const DISPATCH_QUEUE_CAPACITY: usize = 8; + const JOB_ID: JobId = JobId::from(0); + + let storage = MockStorageClient::new(); + storage.push_commit_ready_batch( + DEFAULT_SESSION_ID, + vec![make_entry(RG_A, JOB_ID, TaskId::Commit)], + ); + let mut fixture = CoreFixture::new( + make_config(DISPATCH_QUEUE_CAPACITY, ACTIVE_JOB_LIST_CAPACITY), + storage, + ); + tick_until!(fixture.core, 1 == fixture.queue_len(RG_A)); + + fixture.storage.push_commit_ready_batch( + DEFAULT_SESSION_ID, + vec![make_entry(RG_A, JOB_ID, TaskId::Commit)], + ); + fixture.core.tick().await?; + fixture.core.tick().await?; + + assert_eq!(fixture.queue_len(RG_A), 1); + Ok(()) +} + +#[tokio::test] +async fn an_expired_finalization_readmits_the_jobs_later_tasks() -> anyhow::Result<()> { + const DISPATCH_QUEUE_CAPACITY: usize = 8; + const JOB_ID: JobId = JobId::from(0); + const LATE_TASK_ID: TaskId = TaskId::Index(0); + + let storage = MockStorageClient::new(); + storage.push_commit_ready_batch( + DEFAULT_SESSION_ID, + vec![make_entry(RG_A, JOB_ID, TaskId::Commit)], + ); + let mut fixture = CoreFixture::new( + make_expiring_config( + DISPATCH_QUEUE_CAPACITY, + ACTIVE_JOB_LIST_CAPACITY, + SHORT_EXPIRATION_TIMEOUT_SEC, + ), + storage, + ); + tick_until!(fixture.core, 1 == fixture.queue_len(RG_A)); + let finalization = drain_reader(&fixture.reader(RG_A)).await; + assert_eq!(finalization.len(), 1); + assert_eq!(finalization[0].task_id, TaskId::Commit); + + let num_polls_before = fixture.storage.num_polls().0; + fixture.storage.push_ready_batch( + DEFAULT_SESSION_ID, + vec![make_entry(RG_A, JOB_ID, LATE_TASK_ID)], + ); + tick_until!( + fixture.core, + num_polls_before + 2 <= fixture.storage.num_polls().0 + ); + assert_eq!(fixture.queue_len(RG_A), 0); + assert_eq!(fixture.core.global_task_set.tasks, HashSet::new()); + assert_eq!(fixture.core.job_registry.len(), 0); + + tokio::time::sleep(EXPIRATION_WAIT).await; + tick_until!(fixture.core, fixture.core.finalized_jobs.is_empty()); + + fixture.storage.push_ready_batch( + DEFAULT_SESSION_ID, + vec![make_entry(RG_A, JOB_ID, LATE_TASK_ID)], + ); + tick_until!(fixture.core, 1 == fixture.queue_len(RG_A)); + let readmitted = drain_reader(&fixture.reader(RG_A)).await; + assert_eq!(readmitted.len(), 1); + assert_eq!(readmitted[0].task_id, LATE_TASK_ID); + Ok(()) +} + +#[tokio::test] +async fn a_session_bump_empties_the_finalized_job_table_and_its_queue() -> anyhow::Result<()> { + const DISPATCH_QUEUE_CAPACITY: usize = 8; + const JOB_ID: JobId = JobId::from(0); + + let storage = MockStorageClient::new(); + storage.push_commit_ready_batch( + DEFAULT_SESSION_ID, + vec![make_entry(RG_A, JOB_ID, TaskId::Commit)], + ); + let mut fixture = CoreFixture::new( + make_config(DISPATCH_QUEUE_CAPACITY, ACTIVE_JOB_LIST_CAPACITY), + storage, + ); + tick_until!(fixture.core, fixture.core.finalized_jobs.contains(&JOB_ID)); + assert_eq!(finalized_job_ids(&fixture), vec![JOB_ID]); + + fixture + .storage + .push_ready_batch(NEXT_SESSION_ID, Vec::new()); + tick_until!( + fixture.core, + NEXT_SESSION_ID == fixture.session_tracker.current() + ); + + assert_eq!(fixture.core.finalized_jobs, HashSet::new()); + assert_eq!(finalized_job_ids(&fixture), Vec::::new()); + Ok(()) +} + +#[tokio::test] +async fn publishing_an_assignment_discounts_the_lane_that_buffered_it() -> anyhow::Result<()> { + // One assignment preloaded into the group's queue leaves the tick no free space, so every task + // the batches deliver stays buffered until the test drains the queue. + const DISPATCH_QUEUE_CAPACITY: usize = 1; + const REGULAR_JOB_ID: JobId = JobId::from(0); + const COMMIT_JOB_ID: JobId = JobId::from(1); + const NUM_REGULAR_TASKS: usize = 2; + + let storage = MockStorageClient::new(); + storage.push_ready_batch( + DEFAULT_SESSION_ID, + vec![ + make_entry(RG_A, REGULAR_JOB_ID, TaskId::Index(0)), + make_entry(RG_A, REGULAR_JOB_ID, TaskId::Index(1)), + ], + ); + storage.push_commit_ready_batch( + DEFAULT_SESSION_ID, + vec![make_entry(RG_A, COMMIT_JOB_ID, TaskId::Commit)], + ); + let mut fixture = CoreFixture::new( + make_config(DISPATCH_QUEUE_CAPACITY, ACTIVE_JOB_LIST_CAPACITY), + storage, + ); + fixture.preload_queue(RG_A, DISPATCH_QUEUE_CAPACITY)?; + + tick_until!( + fixture.core, + fixture.core.finalized_jobs.contains(&COMMIT_JOB_ID) + ); + assert_eq!(lane_counts(&fixture), (NUM_REGULAR_TASKS, 1, 0)); + assert_eq!( + fixture.core.global_task_set.len(), + NUM_REGULAR_TASKS + 1, + "every buffered task is counted exactly once" + ); + + // Draining the preloaded assignment frees exactly one slot, so the next tick publishes exactly + // one assignment: the finalization, which outranks the regular tasks. + assert_eq!(drain_reader(&fixture.reader(RG_A)).await.len(), 1); + fixture.core.tick().await?; + let published = drain_reader(&fixture.reader(RG_A)).await; + assert_eq!(published.len(), 1); + assert_eq!(published[0].task_id, TaskId::Commit); + assert_eq!(lane_counts(&fixture), (NUM_REGULAR_TASKS, 0, 0)); + + fixture.core.tick().await?; + let published = drain_reader(&fixture.reader(RG_A)).await; + assert_eq!(published.len(), 1); + assert_eq!(published[0].job_id, REGULAR_JOB_ID); + assert_eq!(lane_counts(&fixture), (NUM_REGULAR_TASKS - 1, 0, 0)); + Ok(()) +} + +#[tokio::test] +async fn the_inbound_poll_is_sized_from_the_lane_counters() -> anyhow::Result<()> { + const DISPATCH_QUEUE_CAPACITY: usize = 1; + const NUM_REGULAR_TASKS: usize = 3; + const NUM_COMMIT_READY_JOBS: usize = 2; + const NUM_CLEANUP_READY_JOBS: usize = 2; + + let storage = MockStorageClient::new(); + storage.push_ready_batch( + DEFAULT_SESSION_ID, + vec![ + make_entry(RG_A, JobId::from(0), TaskId::Index(0)), + make_entry(RG_A, JobId::from(0), TaskId::Index(1)), + make_entry(RG_A, JobId::from(0), TaskId::Index(2)), + ], + ); + storage.push_commit_ready_batch( + DEFAULT_SESSION_ID, + vec![ + make_entry(RG_A, JobId::from(1), TaskId::Commit), + make_entry(RG_A, JobId::from(2), TaskId::Commit), + ], + ); + storage.push_cleanup_ready_batch( + DEFAULT_SESSION_ID, + vec![ + make_entry(RG_A, JobId::from(3), TaskId::Cleanup), + make_entry(RG_A, JobId::from(4), TaskId::Cleanup), + ], + ); + let mut fixture = CoreFixture::new( + make_config(DISPATCH_QUEUE_CAPACITY, ACTIVE_JOB_LIST_CAPACITY), + storage, + ); + fixture.preload_queue(RG_A, DISPATCH_QUEUE_CAPACITY)?; + + let num_finalized_jobs = NUM_COMMIT_READY_JOBS + NUM_CLEANUP_READY_JOBS; + tick_until!( + fixture.core, + num_finalized_jobs == fixture.core.finalized_jobs.len() + ); + assert_eq!( + lane_counts(&fixture), + ( + NUM_REGULAR_TASKS, + NUM_COMMIT_READY_JOBS, + NUM_CLEANUP_READY_JOBS + ) + ); + + // Every later poll is sized from the same counts, because nothing is published while the + // preloaded assignment holds the buffer full. + let (_, num_commit_ready_polls_before, _) = fixture.storage.num_polls(); + tick_until!( + fixture.core, + num_commit_ready_polls_before < fixture.storage.num_polls().1 + ); + assert_eq!( + fixture.storage.last_poll_limits(), + ( + READY_TASK_CAPACITY - NUM_REGULAR_TASKS, + COMMIT_READY_TASK_CAPACITY - NUM_COMMIT_READY_JOBS, + CLEANUP_READY_TASK_CAPACITY - NUM_CLEANUP_READY_JOBS + ) + ); + Ok(()) +} + +#[tokio::test] +async fn a_session_bump_zeroes_every_lane_counter() -> anyhow::Result<()> { + const DISPATCH_QUEUE_CAPACITY: usize = 1; + + let storage = MockStorageClient::new(); + storage.push_ready_batch( + DEFAULT_SESSION_ID, + vec![make_entry(RG_A, JobId::from(0), TaskId::Index(0))], + ); + storage.push_commit_ready_batch( + DEFAULT_SESSION_ID, + vec![make_entry(RG_A, JobId::from(1), TaskId::Commit)], + ); + storage.push_cleanup_ready_batch( + DEFAULT_SESSION_ID, + vec![make_entry(RG_A, JobId::from(2), TaskId::Cleanup)], + ); + let mut fixture = CoreFixture::new( + make_config(DISPATCH_QUEUE_CAPACITY, ACTIVE_JOB_LIST_CAPACITY), + storage, + ); + fixture.preload_queue(RG_A, DISPATCH_QUEUE_CAPACITY)?; + + tick_until!(fixture.core, 2 == fixture.core.finalized_jobs.len()); + assert_eq!(lane_counts(&fixture), (1, 1, 1)); + + fixture + .storage + .push_ready_batch(NEXT_SESSION_ID, Vec::new()); + tick_until!( + fixture.core, + NEXT_SESSION_ID == fixture.session_tracker.current() + ); + + assert_eq!(lane_counts(&fixture), (0, 0, 0)); + assert_eq!(fixture.core.global_task_set.tasks, HashSet::new()); + Ok(()) +} + #[tokio::test] async fn the_scheduling_loop_stops_when_it_is_cancelled() -> anyhow::Result<()> { const DISPATCH_QUEUE_CAPACITY: usize = 4; @@ -648,17 +1010,41 @@ const SHARE_TOLERANCE: usize = 6; fn make_config( dispatch_queue_capacity: usize, active_job_list_capacity: usize, +) -> RgRoundRobinConfig { + make_expiring_config( + dispatch_queue_capacity, + active_job_list_capacity, + NEVER_EXPIRING_TIMEOUT_SEC, + ) +} + +/// # Returns +/// +/// A core config as [`make_config`] builds it, sweeping its finalized job table after +/// `finalized_job_expiration_timeout_sec`. +/// +/// # Panics +/// +/// Panics if either capacity is zero. +fn make_expiring_config( + dispatch_queue_capacity: usize, + active_job_list_capacity: usize, + finalized_job_expiration_timeout_sec: u64, ) -> RgRoundRobinConfig { RgRoundRobinConfig { dispatch_queue_capacity: NonZeroUsize::new(dispatch_queue_capacity) .expect("the dispatch queue capacity is non-zero"), active_job_list_capacity: NonZeroUsize::new(active_job_list_capacity) .expect("the active job list capacity is non-zero"), - ready_task_capacity: NonZeroUsize::new(16_384).expect("16384 is non-zero"), - commit_ready_task_capacity: NonZeroUsize::new(64).expect("64 is non-zero"), - cleanup_ready_task_capacity: NonZeroUsize::new(64).expect("64 is non-zero"), + ready_task_capacity: NonZeroUsize::new(READY_TASK_CAPACITY) + .expect("the regular-task buffer capacity is non-zero"), + commit_ready_task_capacity: NonZeroUsize::new(COMMIT_READY_TASK_CAPACITY) + .expect("the commit-task buffer capacity is non-zero"), + cleanup_ready_task_capacity: NonZeroUsize::new(CLEANUP_READY_TASK_CAPACITY) + .expect("the cleanup-task buffer capacity is non-zero"), storage_poll_timeout_ms: STORAGE_POLL_TIMEOUT_MS, tick_interval_ms: NonZeroU64::new(1).expect("1 is non-zero"), + finalized_job_expiration_timeout_sec, } } @@ -748,6 +1134,34 @@ fn occupancies_of(fixture: &CoreFixture, num_groups: usize) -> Vec { .collect() } +/// # Returns +/// +/// A tuple containing the core's per-lane counts of buffered tasks: +/// +/// * The number of buffered regular tasks. +/// * The number of buffered commit tasks. +/// * The number of buffered cleanup tasks. +fn lane_counts(fixture: &CoreFixture) -> (usize, usize, usize) { + let global_task_set = &fixture.core.global_task_set; + ( + global_task_set.num_ready(), + global_task_set.num_commit_ready(), + global_task_set.num_cleanup_ready(), + ) +} + +/// # Returns +/// +/// The jobs the core's finalized job table holds, in the order they were finalized. +fn finalized_job_ids(fixture: &CoreFixture) -> Vec { + fixture + .core + .finalized_job_queue + .iter() + .map(|(job_id, _)| *job_id) + .collect() +} + /// # Returns /// /// Whether the core still holds `rg_id` on its active resource group list. @@ -758,7 +1172,7 @@ fn occupancies_of(fixture: &CoreFixture, num_groups: usize) -> Vec { fn is_active(fixture: &CoreFixture, rg_id: ResourceGroupId) -> bool { let state_index = *fixture .core - .rg_index + .rg_id_to_index_map .get(&rg_id) .expect("the core holds a scheduling state for the group"); fixture.core.rg_states[state_index].is_active From ec70421d947dff25e638630b3f647e62783e223e Mon Sep 17 00:00:00 2001 From: LinZhihao-723 Date: Tue, 1 Sep 2026 20:02:44 -0400 Subject: [PATCH 3/5] Reviewed the assginment making. --- .../implementation.rs | 78 +++++++++---------- .../resource_group_round_robin/tests.rs | 42 +++++----- 2 files changed, 58 insertions(+), 62 deletions(-) diff --git a/components/spider-scheduler/src/core_impl/resource_group_round_robin/implementation.rs b/components/spider-scheduler/src/core_impl/resource_group_round_robin/implementation.rs index 495bf398..165b2ff5 100644 --- a/components/spider-scheduler/src/core_impl/resource_group_round_robin/implementation.rs +++ b/components/spider-scheduler/src/core_impl/resource_group_round_robin/implementation.rs @@ -94,7 +94,7 @@ pub(super) struct RgRoundRobin, - pub(super) rg_id_to_index_map: HashMap, + pub(super) rg_id_to_idx_map: HashMap, pub(super) active_rg_list: Vec, pub(super) last_served_rg: Option, @@ -133,7 +133,7 @@ impl finalized_job_queue: VecDeque::new(), job_registry: JobRegistry::new(), rg_states: Vec::new(), - rg_id_to_index_map: HashMap::new(), + rg_id_to_idx_map: HashMap::new(), active_rg_list: Vec::new(), last_served_rg: None, config, @@ -277,7 +277,7 @@ impl ); self.rg_states.clear(); - self.rg_id_to_index_map.clear(); + self.rg_id_to_idx_map.clear(); self.active_rg_list.clear(); self.last_served_rg = None; self.job_registry.clear(); @@ -406,8 +406,8 @@ impl /// group the updates touch. fn apply_rg_updates(&mut self, rg_updates: HashMap) { for (rg_id, update) in rg_updates { - let state_index = self.get_or_create_state(rg_id); - let rg_state = &mut self.rg_states[state_index]; + let state_idx = self.get_or_create_state(rg_id); + let rg_state = &mut self.rg_states[state_idx]; for (job_id, kind) in update.finalized { rg_state.push_finalization(job_id, kind); } @@ -418,7 +418,7 @@ impl continue; } rg_state.is_active = true; - self.active_rg_list.push(state_index); + self.active_rg_list.push(state_idx); } } @@ -427,21 +427,21 @@ impl /// The position of `rg_id`'s scheduling state in [`Self::rg_states`], appending a state built /// against the write side of the group's dispatch queue if the core has none. fn get_or_create_state(&mut self, rg_id: ResourceGroupId) -> usize { - if let Some(state_index) = self.rg_id_to_index_map.get(&rg_id) { - return *state_index; + if let Some(state_idx) = self.rg_id_to_idx_map.get(&rg_id) { + return *state_idx; } let writer = self .dispatch_queue_registry .get_dispatch_queue_writer(rg_id); - let state_index = self.rg_states.len(); + let state_idx = self.rg_states.len(); self.rg_states.push(RgSchedulingState::new( rg_id, writer, self.config.active_job_list_capacity.get(), )); - self.rg_id_to_index_map.insert(rg_id, state_index); - state_index + self.rg_id_to_idx_map.insert(rg_id, state_idx); + state_idx } /// Publishes assignments into the per-resource-group dispatch queues under the admission @@ -480,70 +480,66 @@ impl .. } = self; - let mut rg_rr_list = Vec::with_capacity(active_rg_list.len()); + let mut rr_candidates = Vec::with_capacity(active_rg_list.len()); let mut occupancy = 0; - let mut last_served_index = None; - for (index, state_index) in active_rg_list.iter().enumerate() { - let rg_state = &rg_states[*state_index]; + let mut last_served_idx = None; + for (rr_idx, state_idx) in active_rg_list.iter().enumerate() { + let rg_state = &mut rg_states[*state_idx]; occupancy += rg_state.dispatch_queue_size(); if Some(rg_state.rg_id) == *last_served_rg { - last_served_index = Some(index); + last_served_idx = Some(rr_idx); } - rg_rr_list.push(*state_index); + rg_state.promote_pending_jobs(job_registry, &mut jobs_to_retire); + rr_candidates.push(*state_idx); } + // Bounding by the free space measured here is what makes the loop terminate: the queues // drain concurrently, so the true free space only ever grows. let mut free = config .dispatch_queue_capacity .get() .saturating_sub(occupancy); + // Rotating the arm rather than the list keeps the same group from always being visited // first, which matters because `free` shrinks as the tick proceeds. - let mut arm = last_served_index.map_or(0, |index| (index + 1) % rg_rr_list.len()); - - for state_index in &rg_rr_list { - rg_states[*state_index].promote_pending_jobs(job_registry, &mut jobs_to_retire); - } + let mut arm = last_served_idx.map_or(0, |idx| (idx + 1) % rr_candidates.len()); let session_id = session_tracker.current(); let mut exhausted_states = Vec::new(); - while 0 != free && !rg_rr_list.is_empty() { - let state_index = rg_rr_list[arm]; - let rg_state = &mut rg_states[state_index]; - let result = rg_state.try_make_assignment( + while 0 != free && !rr_candidates.is_empty() { + if arm == rr_candidates.len() { + arm = 0; + } + let state_idx = rr_candidates[arm]; + let rg_state = &mut rg_states[state_idx]; + match rg_state.try_make_assignment( free, session_id, id_issuer, job_registry, &mut jobs_to_retire, - ); - match result { + ) { Ok((job_id, task_id)) => { global_task_set.remove(job_id, task_id); free -= 1; *last_served_rg = Some(rg_state.rg_id); - arm = (arm + 1) % rg_rr_list.len(); + arm += 1; } Err(err) => { match err { - MakeAssignmentError::NoTask => exhausted_states.push(state_index), + MakeAssignmentError::NoTask => exhausted_states.push(state_idx), MakeAssignmentError::DispatchQueueFull => (), MakeAssignmentError::DispatchQueueClosed => { return Err(SchedulerError::DispatchQueueClosed); } } - rg_rr_list.swap_remove(arm); - // `swap_remove` moved the tail element into this slot, so advancing the arm - // here would skip it. - if arm == rg_rr_list.len() { - arm = 0; - } + rr_candidates.swap_remove(arm); } } } - for state_index in &*active_rg_list { - rg_states[*state_index].apply_downgrades(job_registry); + for state_idx in &*active_rg_list { + rg_states[*state_idx].apply_downgrades(job_registry); } self.deactivate_exhausted_states(exhausted_states); @@ -557,8 +553,8 @@ impl /// list alone, so a deactivated group still holding assignments would hide its occupancy and /// let the core over-admit. fn deactivate_exhausted_states(&mut self, exhausted_states: Vec) { - for state_index in exhausted_states { - let rg_state = &mut self.rg_states[state_index]; + for state_idx in exhausted_states { + let rg_state = &mut self.rg_states[state_idx]; if rg_state.has_schedulable_task() || 0 != rg_state.dispatch_queue_size() { continue; } @@ -566,7 +562,7 @@ impl if let Some(position) = self .active_rg_list .iter() - .position(|active_index| *active_index == state_index) + .position(|active_idx| *active_idx == state_idx) { self.active_rg_list.swap_remove(position); } diff --git a/components/spider-scheduler/src/core_impl/resource_group_round_robin/tests.rs b/components/spider-scheduler/src/core_impl/resource_group_round_robin/tests.rs index 87e36d49..a2042a11 100644 --- a/components/spider-scheduler/src/core_impl/resource_group_round_robin/tests.rs +++ b/components/spider-scheduler/src/core_impl/resource_group_round_robin/tests.rs @@ -146,8 +146,8 @@ impl CoreFixture { /// /// The position of the group's scheduling state in the core's state vector. fn activate_group(&mut self, rg_id: ResourceGroupId) -> usize { - if let Some(state_index) = self.core.rg_id_to_index_map.get(&rg_id) { - return *state_index; + if let Some(state_idx) = self.core.rg_id_to_idx_map.get(&rg_id) { + return *state_idx; } let mut rg_state = RgSchedulingState::new( @@ -157,11 +157,11 @@ impl CoreFixture { self.active_job_list_capacity, ); rg_state.is_active = true; - let state_index = self.core.rg_states.len(); + let state_idx = self.core.rg_states.len(); self.core.rg_states.push(rg_state); - self.core.rg_id_to_index_map.insert(rg_id, state_index); - self.core.active_rg_list.push(state_index); - state_index + self.core.rg_id_to_idx_map.insert(rg_id, state_idx); + self.core.active_rg_list.push(state_idx); + state_idx } /// Registers a job of `num_tasks` buffered ready tasks against `rg_id`, exactly as a completed @@ -184,8 +184,8 @@ impl CoreFixture { .insert(job_id, TaskId::Index(task_index)); } let job_key = make_job_entry(&mut self.core.job_registry, job_id, num_tasks)?; - let state_index = self.activate_group(rg_id); - self.core.rg_states[state_index].place_new_job(job_key); + let state_idx = self.activate_group(rg_id); + self.core.rg_states[state_idx].place_new_job(job_key); Ok(()) } @@ -202,11 +202,11 @@ impl CoreFixture { let writer = self .dispatch_queue_registry .get_dispatch_queue_writer(rg_id); - for index in 0..num_assignments { + for idx in 0..num_assignments { writer.try_send(make_assignment( rg_id, JobId::from(u64::MAX), - TaskId::Index(index), + TaskId::Index(idx), self.session_tracker.current(), ))?; } @@ -240,8 +240,8 @@ async fn the_rotation_arm_persists_across_ticks() -> anyhow::Result<()> { make_config(DISPATCH_QUEUE_CAPACITY, ACTIVE_JOB_LIST_CAPACITY), MockStorageClient::new(), ); - for (index, rg_id) in [RG_A, RG_B, RG_C].into_iter().enumerate() { - fixture.seed_job(rg_id, JobId::from(u64::try_from(index)?), 8)?; + for (idx, rg_id) in [RG_A, RG_B, RG_C].into_iter().enumerate() { + fixture.seed_job(rg_id, JobId::from(u64::try_from(idx)?), 8)?; } fixture.core.tick().await?; @@ -887,10 +887,10 @@ async fn one_tick_leaves_every_backlogged_group_at_the_dynamic_threshold() -> an let occupancies = occupancies_of(&fixture, NUM_BACKLOGGED_GROUPS); let occupancy: usize = occupancies.iter().sum(); let free = ADMISSION_DISPATCH_QUEUE_CAPACITY - occupancy; - for (index, group_occupancy) in occupancies.iter().enumerate() { + for (idx, group_occupancy) in occupancies.iter().enumerate() { assert!( group_occupancy.abs_diff(expected_share) <= SHARE_TOLERANCE, - "group {index} holds {group_occupancy} assignments, expected about {expected_share}: \ + "group {idx} holds {group_occupancy} assignments, expected about {expected_share}: \ {occupancies:?}" ); } @@ -1069,8 +1069,8 @@ fn new_admission_fixture() -> CoreFixture { /// * Forwards [`CoreFixture::seed_job`]'s return values on failure. /// * Forwards [`u64::try_from`]'s return values on failure. fn seed_backlogged_groups(fixture: &mut CoreFixture, num_groups: usize) -> anyhow::Result<()> { - for index in 0..num_groups { - let raw_id = u64::try_from(index)?; + for idx in 0..num_groups { + let raw_id = u64::try_from(idx)?; fixture.seed_job( ResourceGroupId::from(raw_id), JobId::from(raw_id), @@ -1127,8 +1127,8 @@ fn make_assignment( /// resource group ID. fn occupancies_of(fixture: &CoreFixture, num_groups: usize) -> Vec { (0..num_groups) - .map(|index| { - let raw_id = u64::try_from(index).expect("a group index fits a raw ID"); + .map(|idx| { + let raw_id = u64::try_from(idx).expect("a group index fits a raw ID"); fixture.queue_len(ResourceGroupId::from(raw_id)) }) .collect() @@ -1170,12 +1170,12 @@ fn finalized_job_ids(fixture: &CoreFixture) -> Vec { /// /// Panics if the core has no scheduling state for `rg_id`. fn is_active(fixture: &CoreFixture, rg_id: ResourceGroupId) -> bool { - let state_index = *fixture + let state_idx = *fixture .core - .rg_id_to_index_map + .rg_id_to_idx_map .get(&rg_id) .expect("the core holds a scheduling state for the group"); - fixture.core.rg_states[state_index].is_active + fixture.core.rg_states[state_idx].is_active } /// Takes every assignment currently queued behind `reader`, playing a pinned execution manager, From ffe73bf52cab7e9f6c5caa80618b16785287fb55 Mon Sep 17 00:00:00 2001 From: LinZhihao-723 Date: Thu, 3 Sep 2026 22:53:44 -0400 Subject: [PATCH 4/5] Cleanup the harness fixture. --- .../resource_group_round_robin/tests.rs | 743 +++++++++--------- 1 file changed, 375 insertions(+), 368 deletions(-) diff --git a/components/spider-scheduler/src/core_impl/resource_group_round_robin/tests.rs b/components/spider-scheduler/src/core_impl/resource_group_round_robin/tests.rs index a2042a11..cbb2455c 100644 --- a/components/spider-scheduler/src/core_impl/resource_group_round_robin/tests.rs +++ b/components/spider-scheduler/src/core_impl/resource_group_round_robin/tests.rs @@ -1,30 +1,10 @@ //! Unit tests for the resource-group-aware round-robin scheduler core. //! -//! The tests drive [`RgRoundRobin::tick`] and the structures it decides with directly. -//! Whatever an execution manager would do -- draining a group's queue -- a test body does by hand -//! through the same entry points the dispatch service will use. - -/// Drives ticks on `core` until `predicate` holds, failing the calling test if it does not hold -/// within [`TICK_DEADLINE`]. -/// -/// A macro rather than an async function so that the predicate is an expression rather than a -/// closure, and may therefore borrow whatever the tick mutates. -macro_rules! tick_until { - ($core:expr, $predicate:expr) => {{ - let deadline = tokio::time::Instant::now() + TICK_DEADLINE; - loop { - $core.tick().await?; - if $predicate { - break; - } - if deadline < tokio::time::Instant::now() { - ::anyhow::bail!("the core did not reach the expected state in time"); - } - tokio::time::sleep(TICK_RETRY_INTERVAL).await; - } - }}; -} +//! The tests drive [`RgRoundRobin::tick`] and the structures it decides with directly. Whatever an +//! execution manager would do -- draining a group's queue -- a test body does by hand through the +//! same entry points the dispatch service will use. +use std::collections::HashMap; use std::collections::HashSet; use std::num::NonZeroU64; use std::num::NonZeroUsize; @@ -41,11 +21,8 @@ use spider_core::types::id::TaskId; use tokio_util::sync::CancellationToken; use super::dispatch_queue::DispatchQueueRegistry; -use super::dispatch_queue::RgDispatchQueueReader; use super::implementation::RgRoundRobin; use super::implementation::RgRoundRobinConfig; -use super::job_registry::JobKey; -use super::job_registry::JobRegistry; use super::job_registry::UpsertOutcome; use super::scheduling_state::RgSchedulingState; use crate::SchedulerError; @@ -55,8 +32,26 @@ use crate::core_impl::inbound_queue_reader::test_harness::DEFAULT_SESSION_ID; use crate::core_impl::inbound_queue_reader::test_harness::MockStorageClient; use crate::core_impl::inbound_queue_reader::test_harness::make_entry; -/// The number of active jobs a resource group may hold. -const ACTIVE_JOB_LIST_CAPACITY: usize = 4; +/// Drives ticks on `core` until `predicate` holds, failing the calling test if it does not hold +/// within [`TICK_DEADLINE`]. +/// +/// A macro rather than an async function so that the predicate is an expression rather than a +/// closure, and may therefore borrow whatever the tick mutates. +macro_rules! tick_until { + ($core:expr, $predicate:expr) => {{ + let deadline = tokio::time::Instant::now() + TICK_DEADLINE; + loop { + $core.tick().await?; + if $predicate { + break; + } + if deadline < tokio::time::Instant::now() { + ::anyhow::bail!("the core did not reach the expected state in time"); + } + tokio::time::sleep(TICK_RETRY_INTERVAL).await; + } + }}; +} /// The first resource group a test seeds. const RG_A: ResourceGroupId = ResourceGroupId::from(0); @@ -70,37 +65,56 @@ const RG_C: ResourceGroupId = ResourceGroupId::from(2); /// The session a test bumps the mock storage to. const NEXT_SESSION_ID: SessionId = DEFAULT_SESSION_ID + 1; -/// The storage poll timeout every test runs with. The mock storage never blocks on it. -const STORAGE_POLL_TIMEOUT_MS: u64 = 10; - -/// The regular-task buffer capacity every test runs with. -const READY_TASK_CAPACITY: usize = 16_384; - -/// The commit-task buffer capacity every test runs with. -const COMMIT_READY_TASK_CAPACITY: usize = 64; +/// The config every test starts from, naming in its own literal only the fields it varies, with the +/// following properties: +/// +/// * Its finalized job table expiry outlasts any test run, so no test sweeps the table unless it +/// names a shorter `finalized_job_expiration_timeout_sec`. +/// * Its dispatch queue capacity is a placeholder every test overrides. +/// * Its storage poll timeout is arbitrary because the mock storage never blocks on one. +const BASE_CONFIG: RgRoundRobinConfig = RgRoundRobinConfig { + dispatch_queue_capacity: nonzero_usize(4), + active_job_list_capacity: nonzero_usize(4), + ready_task_capacity: nonzero_usize(16_384), + commit_ready_task_capacity: nonzero_usize(64), + cleanup_ready_task_capacity: nonzero_usize(64), + storage_poll_timeout_ms: 10, + tick_interval_ms: nonzero_u64(1), + finalized_job_expiration_timeout_sec: 6 * 60 * 60, +}; -/// The cleanup-task buffer capacity every test runs with. -const CLEANUP_READY_TASK_CAPACITY: usize = 64; +/// The longest a test waits for the ticks it drives to reach the state it expects. +const TICK_DEADLINE: Duration = Duration::from_secs(10); -/// A finalized job table expiry long enough that no test sweeps an entry unless it asks to. -const NEVER_EXPIRING_TIMEOUT_SEC: u64 = 6 * 60 * 60; +/// The interval between two ticks driven by [`tick_until`]. +const TICK_RETRY_INTERVAL: Duration = Duration::from_millis(2); /// The finalized job table expiry an expiry test runs with. const SHORT_EXPIRATION_TIMEOUT_SEC: u64 = 1; /// How long an expiry test waits before the tick that must sweep an entry stamped -/// [`SHORT_EXPIRATION_TIMEOUT_SEC`] seconds ago. A sleep is a lower bound on the time that passes, -/// so the margin only has to cover the sweep reading a monotonic clock, never scheduling delay. +/// [`SHORT_EXPIRATION_TIMEOUT_SEC`] seconds ago. const EXPIRATION_WAIT: Duration = Duration::from_millis(1_200); -/// The longest a test waits for the ticks it drives to reach the state it expects. -const TICK_DEADLINE: Duration = Duration::from_secs(10); +/// The dispatch buffer capacity of the admission tests. +const ADMISSION_DISPATCH_QUEUE_CAPACITY: usize = 256; -/// The interval between two ticks driven by [`tick_until`]. -const TICK_RETRY_INTERVAL: Duration = Duration::from_millis(2); +/// The number of backlogged resource groups in the equilibrium tests. +const NUM_BACKLOGGED_GROUPS: usize = 5; + +/// The number of ready tasks each backlogged group is seeded with, more than any single tick may +/// publish. +const NUM_TASKS_PER_JOB: usize = 512; /// A core wired to a mock storage and to the dispatch structures a test inspects, driven by manual /// [`RgRoundRobin::tick`] calls. +/// +/// Its methods serve three purposes: +/// +/// * Establishing the state a test drives its ticks from. +/// * Reporting the core's internal state for a test to assert on. +/// * Handing out the read side of a group's queue, so a test can drain it as an execution manager +/// would. struct CoreFixture { core: RgRoundRobin, storage: MockStorageClient, @@ -139,16 +153,35 @@ impl CoreFixture { } } - /// Puts `rg_id` on the core's active resource group list, appending its scheduling state built - /// against the group's dispatch queue endpoints if the core has none. + /// # Returns + /// + /// A newly created fixture over a mock storage with no scripted batch, so a tick's decisions + /// come from the tasks the test seeded and from nothing else. + fn new_admission() -> Self { + Self::new( + RgRoundRobinConfig { + dispatch_queue_capacity: nonzero_usize(ADMISSION_DISPATCH_QUEUE_CAPACITY), + ..BASE_CONFIG + }, + MockStorageClient::new(), + ) + } + + /// Appends `rg_id`'s scheduling state, built against the group's dispatch queue endpoints, and + /// puts it on the core's active resource group list. /// /// # Returns /// /// The position of the group's scheduling state in the core's state vector. - fn activate_group(&mut self, rg_id: ResourceGroupId) -> usize { - if let Some(state_idx) = self.core.rg_id_to_idx_map.get(&rg_id) { - return *state_idx; - } + /// + /// # Panics + /// + /// Panics if the core already holds a scheduling state for `rg_id`. + fn create_and_activate_group(&mut self, rg_id: ResourceGroupId) -> usize { + assert!( + !self.core.rg_id_to_idx_map.contains_key(&rg_id), + "the resource group must not already have a scheduling state" + ); let mut rg_state = RgSchedulingState::new( rg_id, @@ -167,25 +200,42 @@ impl CoreFixture { /// Registers a job of `num_tasks` buffered ready tasks against `rg_id`, exactly as a completed /// inbound poll carrying that job would, and activates the group. /// - /// # Errors - /// - /// Returns an error if: + /// # Panics /// - /// * Forwards [`make_job_entry`]'s return values on failure. - fn seed_job( - &mut self, - rg_id: ResourceGroupId, - job_id: JobId, - num_tasks: usize, - ) -> anyhow::Result<()> { + /// Panics if the job is already registered, or if the group already has a scheduling state. + fn seed_job(&mut self, rg_id: ResourceGroupId, job_id: JobId, num_tasks: usize) { for task_index in 0..num_tasks { self.core .global_task_set .insert(job_id, TaskId::Index(task_index)); } - let job_key = make_job_entry(&mut self.core.job_registry, job_id, num_tasks)?; - let state_idx = self.activate_group(rg_id); + let task_indices: Vec = (0..num_tasks).collect(); + let UpsertOutcome::New(job_key) = self.core.job_registry.upsert(job_id, task_indices) + else { + panic!("job {job_id} is already registered"); + }; + let state_idx = self.create_and_activate_group(rg_id); self.core.rg_states[state_idx].place_new_job(job_key); + } + + /// Seeds the first `num_groups` resource groups with one job each, backlogged with more ready + /// tasks than a single tick may publish. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * Forwards [`Self::seed_job`]'s return values on failure. + /// * Forwards [`u64::try_from`]'s return values on failure. + fn seed_backlogged_groups(&mut self, num_groups: usize) -> anyhow::Result<()> { + for idx in 0..num_groups { + let raw_id = u64::try_from(idx)?; + self.seed_job( + ResourceGroupId::from(raw_id), + JobId::from(raw_id), + NUM_TASKS_PER_JOB, + ); + } Ok(()) } @@ -203,7 +253,7 @@ impl CoreFixture { .dispatch_queue_registry .get_dispatch_queue_writer(rg_id); for idx in 0..num_assignments { - writer.try_send(make_assignment( + writer.try_send(make_unused_assignment( rg_id, JobId::from(u64::MAX), TaskId::Index(idx), @@ -222,13 +272,116 @@ impl CoreFixture { .queue_len() } + /// Takes every assignment currently queued for `rg_id`, playing a pinned execution manager, + /// which leaves the group's hint counter untouched. + /// /// # Returns /// - /// The read side of `rg_id`'s dispatch queue, which a test hands to [`drain_reader`] to play a - /// pinned execution manager. - fn reader(&self, rg_id: ResourceGroupId) -> RgDispatchQueueReader { - self.dispatch_queue_registry - .get_dispatch_queue_reader(rg_id) + /// The assignments taken, in dispatch order. + async fn drain_reader(&self, rg_id: ResourceGroupId) -> Vec { + let reader = self + .dispatch_queue_registry + .get_dispatch_queue_reader(rg_id); + let mut assignments = Vec::new(); + while let Some(assignment) = reader.recv_pinned(Duration::ZERO).await { + assignments.push(assignment); + } + assignments + } + + /// # Returns + /// + /// The number of assignments queued for every resource group the core holds a scheduling state + /// for, keyed by resource group ID. + fn occupancies(&self) -> HashMap { + self.core + .rg_states + .iter() + .map(|rg_state| (rg_state.rg_id, rg_state.dispatch_queue_size())) + .collect() + } + + /// # Returns + /// + /// A tuple containing the core's per-lane counts of buffered tasks: + /// + /// * The number of buffered regular tasks. + /// * The number of buffered commit tasks. + /// * The number of buffered cleanup tasks. + fn lane_counts(&self) -> (usize, usize, usize) { + let global_task_set = &self.core.global_task_set; + ( + global_task_set.num_ready(), + global_task_set.num_commit_ready(), + global_task_set.num_cleanup_ready(), + ) + } + + /// # Returns + /// + /// The jobs the core's finalized job table holds, in the order they were finalized. + fn finalized_job_ids(&self) -> Vec { + self.core + .finalized_job_queue + .iter() + .map(|(job_id, _)| *job_id) + .collect() + } + + /// # Returns + /// + /// Whether the core still holds `rg_id` on its active resource group list. + /// + /// # Panics + /// + /// Panics if the core has no scheduling state for `rg_id`. + fn is_active(&self, rg_id: ResourceGroupId) -> bool { + let state_idx = *self + .core + .rg_id_to_idx_map + .get(&rg_id) + .expect("the core holds a scheduling state for the group"); + self.core.rg_states[state_idx].is_active + } +} + +/// # Returns +/// +/// `capacity` as a [`NonZeroUsize`]. +/// +/// # Panics +/// +/// Panics if `capacity` is zero. +const fn nonzero_usize(capacity: usize) -> NonZeroUsize { + NonZeroUsize::new(capacity).expect("a test config capacity must be non-zero") +} + +/// # Returns +/// +/// `interval` as a [`NonZeroU64`]. +/// +/// # Panics +/// +/// Panics if `interval` is zero. +const fn nonzero_u64(interval: u64) -> NonZeroU64 { + NonZeroU64::new(interval).expect("a test config interval must be non-zero") +} + +/// # Returns +/// +/// An assignment of `task_id` carrying an ID no assignment the core publishes can collide with. +const fn make_unused_assignment( + rg_id: ResourceGroupId, + job_id: JobId, + task_id: TaskId, + session_id: SessionId, +) -> TaskAssignment { + TaskAssignment { + id: TaskAssignmentId::from(u64::MAX), + resource_group_id: rg_id, + job_id, + task_id, + session_id, } } @@ -237,23 +390,32 @@ async fn the_rotation_arm_persists_across_ticks() -> anyhow::Result<()> { const DISPATCH_QUEUE_CAPACITY: usize = 2; let mut fixture = CoreFixture::new( - make_config(DISPATCH_QUEUE_CAPACITY, ACTIVE_JOB_LIST_CAPACITY), + RgRoundRobinConfig { + dispatch_queue_capacity: nonzero_usize(DISPATCH_QUEUE_CAPACITY), + ..BASE_CONFIG + }, MockStorageClient::new(), ); for (idx, rg_id) in [RG_A, RG_B, RG_C].into_iter().enumerate() { - fixture.seed_job(rg_id, JobId::from(u64::try_from(idx)?), 8)?; + fixture.seed_job(rg_id, JobId::from(u64::try_from(idx)?), 8); } fixture.core.tick().await?; - assert_eq!(occupancies_of(&fixture, 3), vec![1, 1, 0]); + assert_eq!( + fixture.occupancies(), + HashMap::from([(RG_A, 1), (RG_B, 1), (RG_C, 0)]) + ); assert_eq!(fixture.core.last_served_rg, Some(RG_B)); for rg_id in [RG_A, RG_B, RG_C] { - drain_reader(&fixture.reader(rg_id)).await; + fixture.drain_reader(rg_id).await; } fixture.core.tick().await?; - assert_eq!(occupancies_of(&fixture, 3), vec![1, 0, 1]); + assert_eq!( + fixture.occupancies(), + HashMap::from([(RG_A, 1), (RG_B, 0), (RG_C, 1)]) + ); assert_eq!(fixture.core.last_served_rg, Some(RG_A)); Ok(()) } @@ -264,19 +426,25 @@ async fn dropping_an_exhausted_group_does_not_skip_the_group_moved_into_its_slot const DISPATCH_QUEUE_CAPACITY: usize = 1; let mut fixture = CoreFixture::new( - make_config(DISPATCH_QUEUE_CAPACITY, ACTIVE_JOB_LIST_CAPACITY), + RgRoundRobinConfig { + dispatch_queue_capacity: nonzero_usize(DISPATCH_QUEUE_CAPACITY), + ..BASE_CONFIG + }, MockStorageClient::new(), ); - fixture.seed_job(RG_A, JobId::from(0), 8)?; - fixture.activate_group(RG_B); - fixture.seed_job(RG_C, JobId::from(2), 8)?; + fixture.seed_job(RG_A, JobId::from(0), 8); + fixture.create_and_activate_group(RG_B); + fixture.seed_job(RG_C, JobId::from(2), 8); // The arm starts on the group that has nothing to schedule, so the tick's single assignment is // won by whichever group the removal moves into the vacated slot. fixture.core.last_served_rg = Some(RG_A); fixture.core.tick().await?; - assert_eq!(occupancies_of(&fixture, 3), vec![0, 0, 1]); + assert_eq!( + fixture.occupancies(), + HashMap::from([(RG_A, 0), (RG_B, 0), (RG_C, 1)]) + ); assert_eq!(fixture.core.last_served_rg, Some(RG_C)); assert_eq!(fixture.core.active_rg_list.len(), 2); Ok(()) @@ -287,20 +455,23 @@ async fn an_exhausted_group_stays_active_until_its_dispatch_queue_drains() -> an const DISPATCH_QUEUE_CAPACITY: usize = 4; let mut fixture = CoreFixture::new( - make_config(DISPATCH_QUEUE_CAPACITY, ACTIVE_JOB_LIST_CAPACITY), + RgRoundRobinConfig { + dispatch_queue_capacity: nonzero_usize(DISPATCH_QUEUE_CAPACITY), + ..BASE_CONFIG + }, MockStorageClient::new(), ); - fixture.activate_group(RG_A); + fixture.create_and_activate_group(RG_A); fixture.preload_queue(RG_A, 1)?; fixture.core.tick().await?; assert_eq!(fixture.core.active_rg_list.len(), 1); - assert!(is_active(&fixture, RG_A)); + assert!(fixture.is_active(RG_A)); - assert_eq!(drain_reader(&fixture.reader(RG_A)).await.len(), 1); + assert_eq!(fixture.drain_reader(RG_A).await.len(), 1); fixture.core.tick().await?; assert_eq!(fixture.core.active_rg_list, Vec::::new()); - assert!(!is_active(&fixture, RG_A)); + assert!(!fixture.is_active(RG_A)); Ok(()) } @@ -314,21 +485,21 @@ async fn dispatching_and_retirement_run_while_a_storage_poll_is_in_flight() -> a let storage = MockStorageClient::new(); storage.gate_ready_lane(); let mut fixture = CoreFixture::new( - make_config(DISPATCH_QUEUE_CAPACITY, ACTIVE_JOB_LIST_CAPACITY), + RgRoundRobinConfig { + dispatch_queue_capacity: nonzero_usize(DISPATCH_QUEUE_CAPACITY), + ..BASE_CONFIG + }, storage, ); - fixture.seed_job(RG_A, JobId::from(0), NUM_TASKS)?; - fixture.seed_job(RG_B, JobId::from(1), 0)?; + fixture.seed_job(RG_A, JobId::from(0), NUM_TASKS); + fixture.seed_job(RG_B, JobId::from(1), 0); // The first tick starts the poll the gate holds, so every later tick finds it still in flight. fixture.core.tick().await?; assert_eq!(fixture.queue_len(RG_A), NUM_TASKS_PER_TICK); assert_eq!(fixture.core.job_registry.len(), 2); - assert_eq!( - drain_reader(&fixture.reader(RG_A)).await.len(), - NUM_TASKS_PER_TICK - ); + assert_eq!(fixture.drain_reader(RG_A).await.len(), NUM_TASKS_PER_TICK); fixture.core.tick().await?; assert_eq!(fixture.queue_len(RG_A), NUM_TASKS_PER_TICK); assert_eq!(fixture.core.global_task_set.tasks, HashSet::new()); @@ -354,7 +525,10 @@ async fn a_session_bump_clears_the_dedup_set_and_the_finalized_job_table() -> an ], ); let mut fixture = CoreFixture::new( - make_config(DISPATCH_QUEUE_CAPACITY, ACTIVE_JOB_LIST_CAPACITY), + RgRoundRobinConfig { + dispatch_queue_capacity: nonzero_usize(DISPATCH_QUEUE_CAPACITY), + ..BASE_CONFIG + }, storage, ); tick_until!(fixture.core, 2 == fixture.queue_len(RG_A)); @@ -395,14 +569,17 @@ async fn a_session_bump_readmits_the_tasks_storage_replays() -> anyhow::Result<( let storage = MockStorageClient::new(); storage.push_ready_batch(DEFAULT_SESSION_ID, replayed_entries.clone()); let mut fixture = CoreFixture::new( - make_config(DISPATCH_QUEUE_CAPACITY, ACTIVE_JOB_LIST_CAPACITY), + RgRoundRobinConfig { + dispatch_queue_capacity: nonzero_usize(DISPATCH_QUEUE_CAPACITY), + ..BASE_CONFIG + }, storage, ); tick_until!(fixture.core, 1 == fixture.queue_len(RG_A)); // The buffer holds one assignment, so the job's second task is still buffered in the core and // still in the dedup set when the bump lands. - let published = drain_reader(&fixture.reader(RG_A)).await; + let published = fixture.drain_reader(RG_A).await; assert_eq!(published.len(), 1); assert_eq!(published[0].session_id, DEFAULT_SESSION_ID); assert_eq!(fixture.core.global_task_set.len(), 1); @@ -417,7 +594,8 @@ async fn a_session_bump_readmits_the_tasks_storage_replays() -> anyhow::Result<( // A poll issued before the bump still lands under the old session, so the assignments it // produces are stale by construction and are not part of the replay. replayed.extend( - drain_reader(&fixture.reader(RG_A)) + fixture + .drain_reader(RG_A) .await .into_iter() .filter(|assignment| assignment.session_id == NEXT_SESSION_ID), @@ -446,10 +624,13 @@ async fn a_rescheduled_assignment_is_readmitted() -> anyhow::Result<()> { const LOST_TASK_ID: TaskId = TaskId::Index(9); let mut fixture = CoreFixture::new( - make_config(DISPATCH_QUEUE_CAPACITY, ACTIVE_JOB_LIST_CAPACITY), + RgRoundRobinConfig { + dispatch_queue_capacity: nonzero_usize(DISPATCH_QUEUE_CAPACITY), + ..BASE_CONFIG + }, MockStorageClient::new(), ); - let lost = make_assignment( + let lost = make_unused_assignment( RG_A, LOST_JOB_ID, LOST_TASK_ID, @@ -459,7 +640,7 @@ async fn a_rescheduled_assignment_is_readmitted() -> anyhow::Result<()> { tick_until!(fixture.core, 1 == fixture.queue_len(RG_A)); - let redispatched = drain_reader(&fixture.reader(RG_A)).await; + let redispatched = fixture.drain_reader(RG_A).await; assert_eq!(redispatched.len(), 1); assert_eq!(redispatched[0].task_id, LOST_TASK_ID); assert_eq!(redispatched[0].job_id, LOST_JOB_ID); @@ -471,10 +652,13 @@ async fn a_closed_dispatch_queue_fails_the_tick() -> anyhow::Result<()> { const DISPATCH_QUEUE_CAPACITY: usize = 4; let mut fixture = CoreFixture::new( - make_config(DISPATCH_QUEUE_CAPACITY, ACTIVE_JOB_LIST_CAPACITY), + RgRoundRobinConfig { + dispatch_queue_capacity: nonzero_usize(DISPATCH_QUEUE_CAPACITY), + ..BASE_CONFIG + }, MockStorageClient::new(), ); - fixture.seed_job(RG_A, JobId::from(0), 4)?; + fixture.seed_job(RG_A, JobId::from(0), 4); fixture.dispatch_queue_registry.close_dispatch_queue(RG_A); let err = fixture @@ -493,10 +677,13 @@ async fn a_closed_broadcast_queue_fails_the_tick() -> anyhow::Result<()> { const DISPATCH_QUEUE_CAPACITY: usize = 4; let mut fixture = CoreFixture::new( - make_config(DISPATCH_QUEUE_CAPACITY, ACTIVE_JOB_LIST_CAPACITY), + RgRoundRobinConfig { + dispatch_queue_capacity: nonzero_usize(DISPATCH_QUEUE_CAPACITY), + ..BASE_CONFIG + }, MockStorageClient::new(), ); - fixture.seed_job(RG_A, JobId::from(0), 4)?; + fixture.seed_job(RG_A, JobId::from(0), 4); fixture.dispatch_queue_registry.close_broadcast_queue(); let err = fixture @@ -526,18 +713,18 @@ async fn an_expired_finalized_job_leaves_the_table_while_a_fresh_one_stays() -> vec![make_entry(RG_A, EXPIRING_JOB_ID, TaskId::Commit)], ); let mut fixture = CoreFixture::new( - make_expiring_config( - DISPATCH_QUEUE_CAPACITY, - ACTIVE_JOB_LIST_CAPACITY, - SHORT_EXPIRATION_TIMEOUT_SEC, - ), + RgRoundRobinConfig { + dispatch_queue_capacity: nonzero_usize(DISPATCH_QUEUE_CAPACITY), + finalized_job_expiration_timeout_sec: SHORT_EXPIRATION_TIMEOUT_SEC, + ..BASE_CONFIG + }, storage, ); tick_until!( fixture.core, fixture.core.finalized_jobs.contains(&EXPIRING_JOB_ID) ); - assert_eq!(finalized_job_ids(&fixture), vec![EXPIRING_JOB_ID]); + assert_eq!(fixture.finalized_job_ids(), vec![EXPIRING_JOB_ID]); tokio::time::sleep(EXPIRATION_WAIT).await; fixture.storage.push_commit_ready_batch( @@ -550,7 +737,7 @@ async fn an_expired_finalized_job_leaves_the_table_while_a_fresh_one_stays() -> ); assert_eq!(fixture.core.finalized_jobs, HashSet::from([FRESH_JOB_ID])); - assert_eq!(finalized_job_ids(&fixture), vec![FRESH_JOB_ID]); + assert_eq!(fixture.finalized_job_ids(), vec![FRESH_JOB_ID]); Ok(()) } @@ -565,7 +752,10 @@ async fn a_cleanup_is_scheduled_after_the_same_job_committed() -> anyhow::Result vec![make_entry(RG_A, JOB_ID, TaskId::Commit)], ); let mut fixture = CoreFixture::new( - make_config(DISPATCH_QUEUE_CAPACITY, ACTIVE_JOB_LIST_CAPACITY), + RgRoundRobinConfig { + dispatch_queue_capacity: nonzero_usize(DISPATCH_QUEUE_CAPACITY), + ..BASE_CONFIG + }, storage, ); tick_until!(fixture.core, 1 == fixture.queue_len(RG_A)); @@ -576,7 +766,8 @@ async fn a_cleanup_is_scheduled_after_the_same_job_committed() -> anyhow::Result ); tick_until!(fixture.core, 2 == fixture.queue_len(RG_A)); - let task_ids: Vec = drain_reader(&fixture.reader(RG_A)) + let task_ids: Vec = fixture + .drain_reader(RG_A) .await .into_iter() .map(|assignment| assignment.task_id) @@ -596,7 +787,10 @@ async fn a_repeated_finalization_is_scheduled_once() -> anyhow::Result<()> { vec![make_entry(RG_A, JOB_ID, TaskId::Commit)], ); let mut fixture = CoreFixture::new( - make_config(DISPATCH_QUEUE_CAPACITY, ACTIVE_JOB_LIST_CAPACITY), + RgRoundRobinConfig { + dispatch_queue_capacity: nonzero_usize(DISPATCH_QUEUE_CAPACITY), + ..BASE_CONFIG + }, storage, ); tick_until!(fixture.core, 1 == fixture.queue_len(RG_A)); @@ -624,15 +818,15 @@ async fn an_expired_finalization_readmits_the_jobs_later_tasks() -> anyhow::Resu vec![make_entry(RG_A, JOB_ID, TaskId::Commit)], ); let mut fixture = CoreFixture::new( - make_expiring_config( - DISPATCH_QUEUE_CAPACITY, - ACTIVE_JOB_LIST_CAPACITY, - SHORT_EXPIRATION_TIMEOUT_SEC, - ), + RgRoundRobinConfig { + dispatch_queue_capacity: nonzero_usize(DISPATCH_QUEUE_CAPACITY), + finalized_job_expiration_timeout_sec: SHORT_EXPIRATION_TIMEOUT_SEC, + ..BASE_CONFIG + }, storage, ); tick_until!(fixture.core, 1 == fixture.queue_len(RG_A)); - let finalization = drain_reader(&fixture.reader(RG_A)).await; + let finalization = fixture.drain_reader(RG_A).await; assert_eq!(finalization.len(), 1); assert_eq!(finalization[0].task_id, TaskId::Commit); @@ -657,7 +851,7 @@ async fn an_expired_finalization_readmits_the_jobs_later_tasks() -> anyhow::Resu vec![make_entry(RG_A, JOB_ID, LATE_TASK_ID)], ); tick_until!(fixture.core, 1 == fixture.queue_len(RG_A)); - let readmitted = drain_reader(&fixture.reader(RG_A)).await; + let readmitted = fixture.drain_reader(RG_A).await; assert_eq!(readmitted.len(), 1); assert_eq!(readmitted[0].task_id, LATE_TASK_ID); Ok(()) @@ -674,11 +868,14 @@ async fn a_session_bump_empties_the_finalized_job_table_and_its_queue() -> anyho vec![make_entry(RG_A, JOB_ID, TaskId::Commit)], ); let mut fixture = CoreFixture::new( - make_config(DISPATCH_QUEUE_CAPACITY, ACTIVE_JOB_LIST_CAPACITY), + RgRoundRobinConfig { + dispatch_queue_capacity: nonzero_usize(DISPATCH_QUEUE_CAPACITY), + ..BASE_CONFIG + }, storage, ); tick_until!(fixture.core, fixture.core.finalized_jobs.contains(&JOB_ID)); - assert_eq!(finalized_job_ids(&fixture), vec![JOB_ID]); + assert_eq!(fixture.finalized_job_ids(), vec![JOB_ID]); fixture .storage @@ -689,7 +886,7 @@ async fn a_session_bump_empties_the_finalized_job_table_and_its_queue() -> anyho ); assert_eq!(fixture.core.finalized_jobs, HashSet::new()); - assert_eq!(finalized_job_ids(&fixture), Vec::::new()); + assert_eq!(fixture.finalized_job_ids(), Vec::::new()); Ok(()) } @@ -715,7 +912,10 @@ async fn publishing_an_assignment_discounts_the_lane_that_buffered_it() -> anyho vec![make_entry(RG_A, COMMIT_JOB_ID, TaskId::Commit)], ); let mut fixture = CoreFixture::new( - make_config(DISPATCH_QUEUE_CAPACITY, ACTIVE_JOB_LIST_CAPACITY), + RgRoundRobinConfig { + dispatch_queue_capacity: nonzero_usize(DISPATCH_QUEUE_CAPACITY), + ..BASE_CONFIG + }, storage, ); fixture.preload_queue(RG_A, DISPATCH_QUEUE_CAPACITY)?; @@ -724,7 +924,7 @@ async fn publishing_an_assignment_discounts_the_lane_that_buffered_it() -> anyho fixture.core, fixture.core.finalized_jobs.contains(&COMMIT_JOB_ID) ); - assert_eq!(lane_counts(&fixture), (NUM_REGULAR_TASKS, 1, 0)); + assert_eq!(fixture.lane_counts(), (NUM_REGULAR_TASKS, 1, 0)); assert_eq!( fixture.core.global_task_set.len(), NUM_REGULAR_TASKS + 1, @@ -733,18 +933,18 @@ async fn publishing_an_assignment_discounts_the_lane_that_buffered_it() -> anyho // Draining the preloaded assignment frees exactly one slot, so the next tick publishes exactly // one assignment: the finalization, which outranks the regular tasks. - assert_eq!(drain_reader(&fixture.reader(RG_A)).await.len(), 1); + assert_eq!(fixture.drain_reader(RG_A).await.len(), 1); fixture.core.tick().await?; - let published = drain_reader(&fixture.reader(RG_A)).await; + let published = fixture.drain_reader(RG_A).await; assert_eq!(published.len(), 1); assert_eq!(published[0].task_id, TaskId::Commit); - assert_eq!(lane_counts(&fixture), (NUM_REGULAR_TASKS, 0, 0)); + assert_eq!(fixture.lane_counts(), (NUM_REGULAR_TASKS, 0, 0)); fixture.core.tick().await?; - let published = drain_reader(&fixture.reader(RG_A)).await; + let published = fixture.drain_reader(RG_A).await; assert_eq!(published.len(), 1); assert_eq!(published[0].job_id, REGULAR_JOB_ID); - assert_eq!(lane_counts(&fixture), (NUM_REGULAR_TASKS - 1, 0, 0)); + assert_eq!(fixture.lane_counts(), (NUM_REGULAR_TASKS - 1, 0, 0)); Ok(()) } @@ -779,7 +979,10 @@ async fn the_inbound_poll_is_sized_from_the_lane_counters() -> anyhow::Result<() ], ); let mut fixture = CoreFixture::new( - make_config(DISPATCH_QUEUE_CAPACITY, ACTIVE_JOB_LIST_CAPACITY), + RgRoundRobinConfig { + dispatch_queue_capacity: nonzero_usize(DISPATCH_QUEUE_CAPACITY), + ..BASE_CONFIG + }, storage, ); fixture.preload_queue(RG_A, DISPATCH_QUEUE_CAPACITY)?; @@ -790,7 +993,7 @@ async fn the_inbound_poll_is_sized_from_the_lane_counters() -> anyhow::Result<() num_finalized_jobs == fixture.core.finalized_jobs.len() ); assert_eq!( - lane_counts(&fixture), + fixture.lane_counts(), ( NUM_REGULAR_TASKS, NUM_COMMIT_READY_JOBS, @@ -808,9 +1011,9 @@ async fn the_inbound_poll_is_sized_from_the_lane_counters() -> anyhow::Result<() assert_eq!( fixture.storage.last_poll_limits(), ( - READY_TASK_CAPACITY - NUM_REGULAR_TASKS, - COMMIT_READY_TASK_CAPACITY - NUM_COMMIT_READY_JOBS, - CLEANUP_READY_TASK_CAPACITY - NUM_CLEANUP_READY_JOBS + BASE_CONFIG.ready_task_capacity.get() - NUM_REGULAR_TASKS, + BASE_CONFIG.commit_ready_task_capacity.get() - NUM_COMMIT_READY_JOBS, + BASE_CONFIG.cleanup_ready_task_capacity.get() - NUM_CLEANUP_READY_JOBS ) ); Ok(()) @@ -834,13 +1037,16 @@ async fn a_session_bump_zeroes_every_lane_counter() -> anyhow::Result<()> { vec![make_entry(RG_A, JobId::from(2), TaskId::Cleanup)], ); let mut fixture = CoreFixture::new( - make_config(DISPATCH_QUEUE_CAPACITY, ACTIVE_JOB_LIST_CAPACITY), + RgRoundRobinConfig { + dispatch_queue_capacity: nonzero_usize(DISPATCH_QUEUE_CAPACITY), + ..BASE_CONFIG + }, storage, ); fixture.preload_queue(RG_A, DISPATCH_QUEUE_CAPACITY)?; tick_until!(fixture.core, 2 == fixture.core.finalized_jobs.len()); - assert_eq!(lane_counts(&fixture), (1, 1, 1)); + assert_eq!(fixture.lane_counts(), (1, 1, 1)); fixture .storage @@ -850,7 +1056,7 @@ async fn a_session_bump_zeroes_every_lane_counter() -> anyhow::Result<()> { NEXT_SESSION_ID == fixture.session_tracker.current() ); - assert_eq!(lane_counts(&fixture), (0, 0, 0)); + assert_eq!(fixture.lane_counts(), (0, 0, 0)); assert_eq!(fixture.core.global_task_set.tasks, HashSet::new()); Ok(()) } @@ -867,7 +1073,10 @@ async fn the_scheduling_loop_stops_when_it_is_cancelled() -> anyhow::Result<()> reschedule_queue_reader, TaskAssignmentIdIssuer::new(), cancellation_token.clone(), - make_config(DISPATCH_QUEUE_CAPACITY, ACTIVE_JOB_LIST_CAPACITY), + RgRoundRobinConfig { + dispatch_queue_capacity: nonzero_usize(DISPATCH_QUEUE_CAPACITY), + ..BASE_CONFIG + }, ); let scheduler_handle = tokio::task::spawn(core.run()); @@ -878,20 +1087,26 @@ async fn the_scheduling_loop_stops_when_it_is_cancelled() -> anyhow::Result<()> #[tokio::test] async fn one_tick_leaves_every_backlogged_group_at_the_dynamic_threshold() -> anyhow::Result<()> { + /// How far a group's occupancy may sit from the equilibrium share before the rotation is + /// considered broken. The staircase these tests exist to catch misses by an order of + /// magnitude: batch filling five groups against `B = 256` yields `64, 64, 64, 64, 0` with + /// no free space at all. + const SHARE_TOLERANCE: usize = 6; + let expected_share = ADMISSION_DISPATCH_QUEUE_CAPACITY / (NUM_BACKLOGGED_GROUPS + 1); - let mut fixture = new_admission_fixture(); - seed_backlogged_groups(&mut fixture, NUM_BACKLOGGED_GROUPS)?; + let mut fixture = CoreFixture::new_admission(); + fixture.seed_backlogged_groups(NUM_BACKLOGGED_GROUPS)?; fixture.core.tick().await?; - let occupancies = occupancies_of(&fixture, NUM_BACKLOGGED_GROUPS); - let occupancy: usize = occupancies.iter().sum(); + let occupancies = fixture.occupancies(); + let occupancy: usize = occupancies.values().sum(); let free = ADMISSION_DISPATCH_QUEUE_CAPACITY - occupancy; - for (idx, group_occupancy) in occupancies.iter().enumerate() { + for (rg_id, group_occupancy) in &occupancies { assert!( group_occupancy.abs_diff(expected_share) <= SHARE_TOLERANCE, - "group {idx} holds {group_occupancy} assignments, expected about {expected_share}: \ - {occupancies:?}" + "group {rg_id:?} holds {group_occupancy} assignments, expected about \ + {expected_share}: {occupancies:?}" ); } assert!( @@ -907,18 +1122,18 @@ async fn one_tick_leaves_every_backlogged_group_at_the_dynamic_threshold() -> an #[tokio::test] async fn no_group_is_batch_filled_while_another_waits() -> anyhow::Result<()> { - let mut fixture = new_admission_fixture(); - seed_backlogged_groups(&mut fixture, NUM_BACKLOGGED_GROUPS)?; + let mut fixture = CoreFixture::new_admission(); + fixture.seed_backlogged_groups(NUM_BACKLOGGED_GROUPS)?; fixture.core.tick().await?; - let occupancies = occupancies_of(&fixture, NUM_BACKLOGGED_GROUPS); + let occupancies = fixture.occupancies(); let most = *occupancies - .iter() + .values() .max() .expect("the tick served at least one group"); let least = *occupancies - .iter() + .values() .min() .expect("the tick served at least one group"); assert!( @@ -940,20 +1155,20 @@ async fn no_group_is_batch_filled_while_another_waits() -> anyhow::Result<()> { async fn a_newly_active_group_is_admitted_against_a_backlogged_incumbent() -> anyhow::Result<()> { const INCUMBENT_OCCUPANCY: usize = 100; - let mut fixture = new_admission_fixture(); - seed_backlogged_groups(&mut fixture, 2)?; + let mut fixture = CoreFixture::new_admission(); + fixture.seed_backlogged_groups(2)?; fixture.preload_queue(RG_A, INCUMBENT_OCCUPANCY)?; fixture.core.tick().await?; - let occupancies = occupancies_of(&fixture, 2); - let free = ADMISSION_DISPATCH_QUEUE_CAPACITY - occupancies.iter().sum::(); + let occupancies = fixture.occupancies(); + let free = ADMISSION_DISPATCH_QUEUE_CAPACITY - occupancies.values().sum::(); assert!( - occupancies[1] >= ADMISSION_DISPATCH_QUEUE_CAPACITY / 8, + occupancies[&RG_B] >= ADMISSION_DISPATCH_QUEUE_CAPACITY / 8, "the newly active group was starved by the incumbent: {occupancies:?}" ); assert!( - occupancies[0] - occupancies[1] < INCUMBENT_OCCUPANCY, + occupancies[&RG_A] - occupancies[&RG_B] < INCUMBENT_OCCUPANCY, "the incumbent's head start grew instead of shrinking: {occupancies:?}" ); assert!( @@ -965,8 +1180,8 @@ async fn a_newly_active_group_is_admitted_against_a_backlogged_incumbent() -> an #[tokio::test] async fn a_lone_group_takes_no_more_than_half_the_dispatch_buffer() -> anyhow::Result<()> { - let mut fixture = new_admission_fixture(); - seed_backlogged_groups(&mut fixture, 1)?; + let mut fixture = CoreFixture::new_admission(); + fixture.seed_backlogged_groups(1)?; fixture.core.tick().await?; @@ -983,211 +1198,3 @@ async fn a_lone_group_takes_no_more_than_half_the_dispatch_buffer() -> anyhow::R ); Ok(()) } - -/// The dispatch buffer capacity `B` of the admission tests, matching the design document's worked -/// example. -const ADMISSION_DISPATCH_QUEUE_CAPACITY: usize = 256; - -/// The number of backlogged resource groups `N` in the equilibrium tests. -const NUM_BACKLOGGED_GROUPS: usize = 5; - -/// The number of ready tasks each backlogged group is seeded with, more than any single tick may -/// publish. -const NUM_TASKS_PER_JOB: usize = 512; - -/// How far a group's occupancy may sit from the equilibrium share before the rotation is considered -/// broken. The staircase these tests exist to catch misses by an order of magnitude: batch filling -/// five groups against `B = 256` yields `64, 64, 64, 64, 0` with no free space at all. -const SHARE_TOLERANCE: usize = 6; - -/// # Returns -/// -/// A core config with the given capacities, whose remaining tunables never throttle a test. -/// -/// # Panics -/// -/// Panics if either capacity is zero. -fn make_config( - dispatch_queue_capacity: usize, - active_job_list_capacity: usize, -) -> RgRoundRobinConfig { - make_expiring_config( - dispatch_queue_capacity, - active_job_list_capacity, - NEVER_EXPIRING_TIMEOUT_SEC, - ) -} - -/// # Returns -/// -/// A core config as [`make_config`] builds it, sweeping its finalized job table after -/// `finalized_job_expiration_timeout_sec`. -/// -/// # Panics -/// -/// Panics if either capacity is zero. -fn make_expiring_config( - dispatch_queue_capacity: usize, - active_job_list_capacity: usize, - finalized_job_expiration_timeout_sec: u64, -) -> RgRoundRobinConfig { - RgRoundRobinConfig { - dispatch_queue_capacity: NonZeroUsize::new(dispatch_queue_capacity) - .expect("the dispatch queue capacity is non-zero"), - active_job_list_capacity: NonZeroUsize::new(active_job_list_capacity) - .expect("the active job list capacity is non-zero"), - ready_task_capacity: NonZeroUsize::new(READY_TASK_CAPACITY) - .expect("the regular-task buffer capacity is non-zero"), - commit_ready_task_capacity: NonZeroUsize::new(COMMIT_READY_TASK_CAPACITY) - .expect("the commit-task buffer capacity is non-zero"), - cleanup_ready_task_capacity: NonZeroUsize::new(CLEANUP_READY_TASK_CAPACITY) - .expect("the cleanup-task buffer capacity is non-zero"), - storage_poll_timeout_ms: STORAGE_POLL_TIMEOUT_MS, - tick_interval_ms: NonZeroU64::new(1).expect("1 is non-zero"), - finalized_job_expiration_timeout_sec, - } -} - -/// # Returns -/// -/// A newly created fixture over a mock storage with no scripted batch, so a tick's decisions come -/// from the tasks the test seeded and from nothing else. -fn new_admission_fixture() -> CoreFixture { - CoreFixture::new( - make_config(ADMISSION_DISPATCH_QUEUE_CAPACITY, ACTIVE_JOB_LIST_CAPACITY), - MockStorageClient::new(), - ) -} - -/// Seeds the first `num_groups` resource groups with one job each, backlogged with more ready tasks -/// than a single tick may publish. -/// -/// # Errors -/// -/// Returns an error if: -/// -/// * Forwards [`CoreFixture::seed_job`]'s return values on failure. -/// * Forwards [`u64::try_from`]'s return values on failure. -fn seed_backlogged_groups(fixture: &mut CoreFixture, num_groups: usize) -> anyhow::Result<()> { - for idx in 0..num_groups { - let raw_id = u64::try_from(idx)?; - fixture.seed_job( - ResourceGroupId::from(raw_id), - JobId::from(raw_id), - NUM_TASKS_PER_JOB, - )?; - } - Ok(()) -} - -/// Registers a job of `num_tasks` buffered ready tasks, whose task indices are `0..num_tasks`. -/// -/// # Returns -/// -/// The key of the registered job, which still needs a scheduling position. -/// -/// # Errors -/// -/// Returns an error if: -/// -/// * [`anyhow::Error`] if the job is already registered. -fn make_job_entry( - registry: &mut JobRegistry, - job_id: JobId, - num_tasks: usize, -) -> anyhow::Result { - let task_indices: Vec = (0..num_tasks).collect(); - let UpsertOutcome::New(job_key) = registry.upsert(job_id, task_indices) else { - bail!("job {job_id} is already registered"); - }; - Ok(job_key) -} - -/// # Returns -/// -/// An assignment of `task_id` carrying an ID no assignment the core publishes can collide with. -fn make_assignment( - rg_id: ResourceGroupId, - job_id: JobId, - task_id: TaskId, - session_id: SessionId, -) -> TaskAssignment { - TaskAssignment { - id: TaskAssignmentId::from(u64::MAX), - resource_group_id: rg_id, - job_id, - task_id, - session_id, - } -} - -/// # Returns -/// -/// The number of assignments queued for each of the first `num_groups` resource groups, indexed by -/// resource group ID. -fn occupancies_of(fixture: &CoreFixture, num_groups: usize) -> Vec { - (0..num_groups) - .map(|idx| { - let raw_id = u64::try_from(idx).expect("a group index fits a raw ID"); - fixture.queue_len(ResourceGroupId::from(raw_id)) - }) - .collect() -} - -/// # Returns -/// -/// A tuple containing the core's per-lane counts of buffered tasks: -/// -/// * The number of buffered regular tasks. -/// * The number of buffered commit tasks. -/// * The number of buffered cleanup tasks. -fn lane_counts(fixture: &CoreFixture) -> (usize, usize, usize) { - let global_task_set = &fixture.core.global_task_set; - ( - global_task_set.num_ready(), - global_task_set.num_commit_ready(), - global_task_set.num_cleanup_ready(), - ) -} - -/// # Returns -/// -/// The jobs the core's finalized job table holds, in the order they were finalized. -fn finalized_job_ids(fixture: &CoreFixture) -> Vec { - fixture - .core - .finalized_job_queue - .iter() - .map(|(job_id, _)| *job_id) - .collect() -} - -/// # Returns -/// -/// Whether the core still holds `rg_id` on its active resource group list. -/// -/// # Panics -/// -/// Panics if the core has no scheduling state for `rg_id`. -fn is_active(fixture: &CoreFixture, rg_id: ResourceGroupId) -> bool { - let state_idx = *fixture - .core - .rg_id_to_idx_map - .get(&rg_id) - .expect("the core holds a scheduling state for the group"); - fixture.core.rg_states[state_idx].is_active -} - -/// Takes every assignment currently queued behind `reader`, playing a pinned execution manager, -/// which leaves the group's hint counter untouched. -/// -/// # Returns -/// -/// The assignments taken, in dispatch order. -async fn drain_reader(reader: &RgDispatchQueueReader) -> Vec { - let mut assignments = Vec::new(); - while let Some(assignment) = reader.recv_pinned(Duration::ZERO).await { - assignments.push(assignment); - } - assignments -} From df284501c966414dc6d6bb16982880e5363f5205 Mon Sep 17 00:00:00 2001 From: LinZhihao-723 Date: Tue, 8 Sep 2026 17:32:01 -0400 Subject: [PATCH 5/5] Done with test case polishing. --- .../resource_group_round_robin/tests.rs | 221 ++++++++++++++++++ 1 file changed, 221 insertions(+) diff --git a/components/spider-scheduler/src/core_impl/resource_group_round_robin/tests.rs b/components/spider-scheduler/src/core_impl/resource_group_round_robin/tests.rs index cbb2455c..d23f5084 100644 --- a/components/spider-scheduler/src/core_impl/resource_group_round_robin/tests.rs +++ b/components/spider-scheduler/src/core_impl/resource_group_round_robin/tests.rs @@ -1198,3 +1198,224 @@ async fn a_lone_group_takes_no_more_than_half_the_dispatch_buffer() -> anyhow::R ); Ok(()) } + +#[tokio::test] +async fn a_finalization_drains_the_jobs_buffered_tasks() -> anyhow::Result<()> { + const DISPATCH_QUEUE_CAPACITY: usize = 1; + const JOB_ID: JobId = JobId::from(0); + const NUM_TASKS: usize = 4; + + let storage = MockStorageClient::new(); + storage.push_ready_batch( + DEFAULT_SESSION_ID, + vec![ + make_entry(RG_A, JOB_ID, TaskId::Index(0)), + make_entry(RG_A, JOB_ID, TaskId::Index(1)), + make_entry(RG_A, JOB_ID, TaskId::Index(2)), + make_entry(RG_A, JOB_ID, TaskId::Index(3)), + ], + ); + let mut fixture = CoreFixture::new( + RgRoundRobinConfig { + dispatch_queue_capacity: nonzero_usize(DISPATCH_QUEUE_CAPACITY), + ..BASE_CONFIG + }, + storage, + ); + + tick_until!(fixture.core, 1 == fixture.queue_len(RG_A)); + assert_eq!(fixture.lane_counts(), (NUM_TASKS - 1, 0, 0)); + assert_eq!(fixture.core.job_registry.len(), 1); + + fixture.storage.push_commit_ready_batch( + DEFAULT_SESSION_ID, + vec![make_entry(RG_A, JOB_ID, TaskId::Commit)], + ); + tick_until!(fixture.core, fixture.core.finalized_jobs.contains(&JOB_ID)); + + // The commit is buffered like any other task until its assignment publishes, which the full + // dispatch queue holds off, so it is what the dedup set is left holding. + assert_eq!( + fixture.core.global_task_set.tasks, + HashSet::from([(JOB_ID, TaskId::Commit)]) + ); + assert_eq!(fixture.lane_counts(), (0, 1, 0)); + assert_eq!(fixture.core.job_registry.len(), 0); + Ok(()) +} + +#[tokio::test] +async fn a_task_offered_twice_across_polls_is_admitted_once() -> anyhow::Result<()> { + const DISPATCH_QUEUE_CAPACITY: usize = 1; + const JOB_ID: JobId = JobId::from(0); + const NUM_TASKS_PER_BATCH: usize = 2; + /// Twice the assignments one batch is worth, each tick publishing at most one, so a task + /// admitted a second time would still have a round to surface in. + const NUM_PUBLISHING_TICKS: usize = 2 * NUM_TASKS_PER_BATCH; + + let batch = vec![ + make_entry(RG_A, JOB_ID, TaskId::Index(0)), + make_entry(RG_A, JOB_ID, TaskId::Index(1)), + ]; + let storage = MockStorageClient::new(); + storage.push_ready_batch(DEFAULT_SESSION_ID, batch.clone()); + let mut fixture = CoreFixture::new( + RgRoundRobinConfig { + dispatch_queue_capacity: nonzero_usize(DISPATCH_QUEUE_CAPACITY), + ..BASE_CONFIG + }, + storage, + ); + // A published task leaves the dedup set by design, so only a dispatch queue that stays full for + // the whole test makes the second offer's rejection the reason nothing is admitted twice. + fixture.preload_queue(RG_A, DISPATCH_QUEUE_CAPACITY)?; + + tick_until!(fixture.core, NUM_TASKS_PER_BATCH == fixture.lane_counts().0); + + let num_polls_before = fixture.storage.num_polls().0; + fixture.storage.push_ready_batch(DEFAULT_SESSION_ID, batch); + tick_until!( + fixture.core, + num_polls_before + 2 <= fixture.storage.num_polls().0 + ); + + assert_eq!(fixture.lane_counts(), (NUM_TASKS_PER_BATCH, 0, 0)); + assert_eq!(fixture.queue_len(RG_A), DISPATCH_QUEUE_CAPACITY); + + // The dedup set counts a task once however often it was admitted, so what the core buffered for + // the job is only visible in what it goes on to publish once the buffer frees up. + assert_eq!( + fixture.drain_reader(RG_A).await.len(), + DISPATCH_QUEUE_CAPACITY + ); + let mut published = Vec::new(); + for _ in 0..NUM_PUBLISHING_TICKS { + fixture.core.tick().await?; + published.extend( + fixture + .drain_reader(RG_A) + .await + .into_iter() + .map(|assignment| assignment.task_id), + ); + } + assert_eq!(published, vec![TaskId::Index(0), TaskId::Index(1)]); + Ok(()) +} + +#[tokio::test] +async fn a_second_batch_for_a_registered_job_appends_to_it() -> anyhow::Result<()> { + const DISPATCH_QUEUE_CAPACITY: usize = 1; + const JOB_ID: JobId = JobId::from(0); + const NUM_TASKS_PER_BATCH: usize = 2; + + let storage = MockStorageClient::new(); + storage.push_ready_batch( + DEFAULT_SESSION_ID, + vec![ + make_entry(RG_A, JOB_ID, TaskId::Index(0)), + make_entry(RG_A, JOB_ID, TaskId::Index(1)), + ], + ); + let mut fixture = CoreFixture::new( + RgRoundRobinConfig { + dispatch_queue_capacity: nonzero_usize(DISPATCH_QUEUE_CAPACITY), + ..BASE_CONFIG + }, + storage, + ); + // Nothing may publish, or the job would run dry and leave the registry before the second batch + // arrives. + fixture.preload_queue(RG_A, DISPATCH_QUEUE_CAPACITY)?; + + tick_until!(fixture.core, NUM_TASKS_PER_BATCH == fixture.lane_counts().0); + + fixture.storage.push_ready_batch( + DEFAULT_SESSION_ID, + vec![ + make_entry(RG_A, JOB_ID, TaskId::Index(2)), + make_entry(RG_A, JOB_ID, TaskId::Index(3)), + ], + ); + tick_until!( + fixture.core, + 2 * NUM_TASKS_PER_BATCH == fixture.lane_counts().0 + ); + + assert_eq!(fixture.core.job_registry.len(), 1); + assert_eq!(fixture.lane_counts(), (2 * NUM_TASKS_PER_BATCH, 0, 0)); + Ok(()) +} + +#[tokio::test] +async fn a_rescheduled_assignment_of_a_stale_session_is_dropped() -> anyhow::Result<()> { + const DISPATCH_QUEUE_CAPACITY: usize = 8; + const LOST_JOB_ID: JobId = JobId::from(0); + const LOST_TASK_ID: TaskId = TaskId::Index(9); + + let storage = MockStorageClient::new(); + storage.push_ready_batch(NEXT_SESSION_ID, Vec::new()); + let mut fixture = CoreFixture::new( + RgRoundRobinConfig { + dispatch_queue_capacity: nonzero_usize(DISPATCH_QUEUE_CAPACITY), + ..BASE_CONFIG + }, + storage, + ); + tick_until!( + fixture.core, + NEXT_SESSION_ID == fixture.session_tracker.current() + ); + + let num_polls_before = fixture.storage.num_polls().0; + fixture + .reschedule_queue_writer + .send(make_unused_assignment( + RG_A, + LOST_JOB_ID, + LOST_TASK_ID, + DEFAULT_SESSION_ID, + ))?; + tick_until!( + fixture.core, + num_polls_before + 2 <= fixture.storage.num_polls().0 + ); + + assert_eq!(fixture.queue_len(RG_A), 0); + assert_eq!(fixture.core.job_registry.len(), 0); + assert_eq!(fixture.core.global_task_set.tasks, HashSet::new()); + Ok(()) +} + +#[tokio::test] +async fn a_drained_group_is_reactivated_by_a_later_poll() -> anyhow::Result<()> { + const DISPATCH_QUEUE_CAPACITY: usize = 4; + const JOB_ID: JobId = JobId::from(0); + + let mut fixture = CoreFixture::new( + RgRoundRobinConfig { + dispatch_queue_capacity: nonzero_usize(DISPATCH_QUEUE_CAPACITY), + ..BASE_CONFIG + }, + MockStorageClient::new(), + ); + fixture.create_and_activate_group(RG_A); + fixture.preload_queue(RG_A, 1)?; + + fixture.core.tick().await?; + assert_eq!(fixture.drain_reader(RG_A).await.len(), 1); + fixture.core.tick().await?; + assert_eq!(fixture.core.active_rg_list, Vec::::new()); + assert!(!fixture.is_active(RG_A)); + let num_rg_states = fixture.core.rg_states.len(); + + fixture.storage.push_ready_batch( + DEFAULT_SESSION_ID, + vec![make_entry(RG_A, JOB_ID, TaskId::Index(0))], + ); + tick_until!(fixture.core, fixture.is_active(RG_A)); + + assert_eq!(fixture.core.active_rg_list.len(), 1); + assert_eq!(fixture.core.rg_states.len(), num_rg_states); + Ok(()) +}