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 new file mode 100644 index 00000000..165b2ff5 --- /dev/null +++ b/components/spider-scheduler/src/core_impl/resource_group_round_robin/implementation.rs @@ -0,0 +1,768 @@ +//! 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::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; +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. 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`]. +/// +/// # 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: 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. + /// + /// 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_idx_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: GlobalTaskSet::new(), + finalized_jobs: HashSet::new(), + finalized_job_queue: VecDeque::new(), + job_registry: JobRegistry::new(), + rg_states: Vec::new(), + rg_id_to_idx_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::apply_session_bump`]'s return values on failure. + /// * Forwards [`Self::start_inbound_poll`]'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> { + match self + .inbound_queue_reader + .try_collect_result(self.session_tracker.current()) + .await? + { + 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.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`. + /// + /// # 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::error!( + from = previous_session_id, + to = new_session_id, + "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.rg_states.clear(); + self.rg_id_to_idx_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(); + self.finalized_job_queue.clear(); + self.dispatch_queue_registry.clear(); + + Ok(()) + } + + /// 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 + } + + /// Processes the inbound polling results along with the assignments to reschedule. + /// + /// # 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 { + // 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; + } + // 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)); + } + } + 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_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); + } + for job_key in update.new_jobs { + rg_state.place_new_job(job_key); + } + if rg_state.is_active { + continue; + } + rg_state.is_active = true; + self.active_rg_list.push(state_idx); + } + } + + /// # 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_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_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_idx_map.insert(rg_id, state_idx); + state_idx + } + + /// 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 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); + } + + // 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 rr_candidates = Vec::with_capacity(active_rg_list.len()); + let mut occupancy = 0; + 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_idx = Some(rr_idx); + } + 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_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 && !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, + ) { + Ok((job_id, task_id)) => { + global_task_set.remove(job_id, task_id); + free -= 1; + *last_served_rg = Some(rg_state.rg_id); + arm += 1; + } + Err(err) => { + match err { + MakeAssignmentError::NoTask => exhausted_states.push(state_idx), + MakeAssignmentError::DispatchQueueFull => (), + MakeAssignmentError::DispatchQueueClosed => { + return Err(SchedulerError::DispatchQueueClosed); + } + } + rr_candidates.swap_remove(arm); + } + } + } + + for state_idx in &*active_rg_list { + rg_states[*state_idx].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_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; + } + rg_state.is_active = false; + if let Some(position) = self + .active_rg_list + .iter() + .position(|active_idx| *active_idx == state_idx) + { + self.active_rg_list.swap_remove(position); + } + } + } + + /// 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 { + 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(); + } + } + + /// 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 max_ready_entries = self + .config + .ready_task_capacity + .get() + .saturating_sub(self.global_task_set.num_ready()); + let max_commit_ready_entries = self + .config + .commit_ready_task_capacity + .get() + .saturating_sub(self.global_task_set.num_commit_ready()); + let max_cleanup_ready_entries = self + .config + .cleanup_ready_task_capacity + .get() + .saturating_sub(self.global_task_set.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, + ) + } +} + +/// 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 + /// + /// 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, + } + } +} + +/// 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/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 new file mode 100644 index 00000000..d23f5084 --- /dev/null +++ b/components/spider-scheduler/src/core_impl/resource_group_round_robin/tests.rs @@ -0,0 +1,1421 @@ +//! 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. + +use std::collections::HashMap; +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::implementation::RgRoundRobin; +use super::implementation::RgRoundRobinConfig; +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; + +/// 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); + +/// 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 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 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); + +/// 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. +const EXPIRATION_WAIT: Duration = Duration::from_millis(1_200); + +/// The dispatch buffer capacity of the admission tests. +const ADMISSION_DISPATCH_QUEUE_CAPACITY: usize = 256; + +/// 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, + 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, + } + } + + /// # 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. + /// + /// # 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, + self.dispatch_queue_registry + .get_dispatch_queue_writer(rg_id), + self.active_job_list_capacity, + ); + rg_state.is_active = true; + let state_idx = self.core.rg_states.len(); + self.core.rg_states.push(rg_state); + 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 + /// inbound poll carrying that job would, and activates the group. + /// + /// # Panics + /// + /// 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 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(()) + } + + /// 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 idx in 0..num_assignments { + writer.try_send(make_unused_assignment( + rg_id, + JobId::from(u64::MAX), + TaskId::Index(idx), + 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() + } + + /// Takes every assignment currently queued for `rg_id`, playing a pinned execution manager, + /// which leaves the group's hint counter untouched. + /// + /// # Returns + /// + /// 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, + } +} + +#[tokio::test] +async fn the_rotation_arm_persists_across_ticks() -> anyhow::Result<()> { + const DISPATCH_QUEUE_CAPACITY: usize = 2; + + let mut fixture = CoreFixture::new( + 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.core.tick().await?; + 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] { + fixture.drain_reader(rg_id).await; + } + + fixture.core.tick().await?; + 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(()) +} + +#[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( + RgRoundRobinConfig { + dispatch_queue_capacity: nonzero_usize(DISPATCH_QUEUE_CAPACITY), + ..BASE_CONFIG + }, + MockStorageClient::new(), + ); + 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!( + 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(()) +} + +#[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( + 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.core.active_rg_list.len(), 1); + assert!(fixture.is_active(RG_A)); + + 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)); + 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( + 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); + + // 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!(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()); + + 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( + RgRoundRobinConfig { + dispatch_queue_capacity: nonzero_usize(DISPATCH_QUEUE_CAPACITY), + ..BASE_CONFIG + }, + 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 + .tasks + .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( + 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 = 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); + + 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( + fixture + .drain_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( + RgRoundRobinConfig { + dispatch_queue_capacity: nonzero_usize(DISPATCH_QUEUE_CAPACITY), + ..BASE_CONFIG + }, + MockStorageClient::new(), + ); + let lost = make_unused_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 = 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); + 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( + RgRoundRobinConfig { + dispatch_queue_capacity: nonzero_usize(DISPATCH_QUEUE_CAPACITY), + ..BASE_CONFIG + }, + 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( + RgRoundRobinConfig { + dispatch_queue_capacity: nonzero_usize(DISPATCH_QUEUE_CAPACITY), + ..BASE_CONFIG + }, + 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 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( + 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!(fixture.finalized_job_ids(), 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!(fixture.finalized_job_ids(), 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( + RgRoundRobinConfig { + dispatch_queue_capacity: nonzero_usize(DISPATCH_QUEUE_CAPACITY), + ..BASE_CONFIG + }, + 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 = fixture + .drain_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( + RgRoundRobinConfig { + dispatch_queue_capacity: nonzero_usize(DISPATCH_QUEUE_CAPACITY), + ..BASE_CONFIG + }, + 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( + 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 = fixture.drain_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 = fixture.drain_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( + 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!(fixture.finalized_job_ids(), 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!(fixture.finalized_job_ids(), 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( + RgRoundRobinConfig { + dispatch_queue_capacity: nonzero_usize(DISPATCH_QUEUE_CAPACITY), + ..BASE_CONFIG + }, + storage, + ); + fixture.preload_queue(RG_A, DISPATCH_QUEUE_CAPACITY)?; + + tick_until!( + fixture.core, + fixture.core.finalized_jobs.contains(&COMMIT_JOB_ID) + ); + assert_eq!(fixture.lane_counts(), (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!(fixture.drain_reader(RG_A).await.len(), 1); + fixture.core.tick().await?; + let published = fixture.drain_reader(RG_A).await; + assert_eq!(published.len(), 1); + assert_eq!(published[0].task_id, TaskId::Commit); + assert_eq!(fixture.lane_counts(), (NUM_REGULAR_TASKS, 0, 0)); + + fixture.core.tick().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!(fixture.lane_counts(), (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( + RgRoundRobinConfig { + dispatch_queue_capacity: nonzero_usize(DISPATCH_QUEUE_CAPACITY), + ..BASE_CONFIG + }, + 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!( + fixture.lane_counts(), + ( + 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(), + ( + 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(()) +} + +#[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( + 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!(fixture.lane_counts(), (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!(fixture.lane_counts(), (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; + + 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(), + RgRoundRobinConfig { + dispatch_queue_capacity: nonzero_usize(DISPATCH_QUEUE_CAPACITY), + ..BASE_CONFIG + }, + ); + 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<()> { + /// 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 = CoreFixture::new_admission(); + fixture.seed_backlogged_groups(NUM_BACKLOGGED_GROUPS)?; + + fixture.core.tick().await?; + + let occupancies = fixture.occupancies(); + let occupancy: usize = occupancies.values().sum(); + let free = ADMISSION_DISPATCH_QUEUE_CAPACITY - occupancy; + for (rg_id, group_occupancy) in &occupancies { + assert!( + group_occupancy.abs_diff(expected_share) <= SHARE_TOLERANCE, + "group {rg_id:?} 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 = CoreFixture::new_admission(); + fixture.seed_backlogged_groups(NUM_BACKLOGGED_GROUPS)?; + + fixture.core.tick().await?; + + let occupancies = fixture.occupancies(); + let most = *occupancies + .values() + .max() + .expect("the tick served at least one group"); + let least = *occupancies + .values() + .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 = CoreFixture::new_admission(); + fixture.seed_backlogged_groups(2)?; + fixture.preload_queue(RG_A, INCUMBENT_OCCUPANCY)?; + + fixture.core.tick().await?; + + let occupancies = fixture.occupancies(); + let free = ADMISSION_DISPATCH_QUEUE_CAPACITY - occupancies.values().sum::(); + assert!( + occupancies[&RG_B] >= ADMISSION_DISPATCH_QUEUE_CAPACITY / 8, + "the newly active group was starved by the incumbent: {occupancies:?}" + ); + assert!( + occupancies[&RG_A] - occupancies[&RG_B] < 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 = CoreFixture::new_admission(); + fixture.seed_backlogged_groups(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(()) +} + +#[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(()) +}